From 11d2a2daabb9f02ef4f00615e10baa2f83fce1cc Mon Sep 17 00:00:00 2001
From: "Dr. Patrick Urbanke" <patrick@getml.com>
Date: Wed, 9 Apr 2025 16:16:51 +0200
Subject: [PATCH] Implemented the sqlite Iterator

---
 include/sqlgen/sqlite/Iterator.hpp | 20 ++++++----
 src/sqlgen/sqlite/Iterator.cpp     | 59 ++++++++++++++++++++++++++++++
 src/sqlgen_sqlite.cpp              |  1 +
 3 files changed, 73 insertions(+), 7 deletions(-)
 create mode 100644 src/sqlgen/sqlite/Iterator.cpp

diff --git a/include/sqlgen/sqlite/Iterator.hpp b/include/sqlgen/sqlite/Iterator.hpp
index dd7f037..985d63e 100644
--- a/include/sqlgen/sqlite/Iterator.hpp
+++ b/include/sqlgen/sqlite/Iterator.hpp
@@ -19,26 +19,32 @@ class Iterator : public sqlgen::IteratorBase {
   using StmtPtr = Ref<sqlite3_stmt>;
 
  public:
-  Iterator(const StmtPtr& _stmt, const ConnPtr& _conn)
-      : rownum_(0), stmt_(_stmt), conn_(_conn) {}
+  Iterator(const StmtPtr& _stmt, const ConnPtr& _conn);
 
-  ~Iterator() = default;
+  ~Iterator();
 
   /// Whether the end of the available data has been reached.
-  bool end() const final { return false; }
+  bool end() const final;
 
   /// Returns the next batch of rows.
   /// If _batch_size is greater than the number of rows left, returns all
   /// of the rows left.
   Result<std::vector<std::vector<std::optional<std::string>>>> next(
-      const size_t _batch_size) final {
-    return error("TODO");
-  }
+      const size_t _batch_size) final;
 
  private:
+  void step() { end_ = (sqlite3_step(stmt_.get()) != SQLITE_ROW); }
+
+ private:
+  /// Whether the end is reached.
+  bool end_;
+
   /// The current rownumber.
   size_t rownum_;
 
+  /// The number of columns.
+  int num_cols_;
+
   /// The prepared statement. Note that we have
   /// declared it before conn_, meaning it will be destroyed first.
   StmtPtr stmt_;
diff --git a/src/sqlgen/sqlite/Iterator.cpp b/src/sqlgen/sqlite/Iterator.cpp
new file mode 100644
index 0000000..c0c5475
--- /dev/null
+++ b/src/sqlgen/sqlite/Iterator.cpp
@@ -0,0 +1,59 @@
+#include "sqlgen/sqlite/Iterator.hpp"
+
+#include <ranges>
+#include <rfl.hpp>
+#include <sstream>
+
+#include "sqlgen/internal/collect/vector.hpp"
+#include "sqlgen/internal/strings/strings.hpp"
+#include "sqlgen/sqlite/Iterator.hpp"
+
+namespace sqlgen::sqlite {
+
+Iterator::Iterator(const StmtPtr& _stmt, const ConnPtr& _conn)
+    : end_(false),
+      rownum_(0),
+      num_cols_(sqlite3_column_count(_stmt.get())),
+      stmt_(_stmt),
+      conn_(_conn) {
+  step();
+}
+
+Iterator::~Iterator() = default;
+
+bool Iterator::end() const { return end_; }
+
+Result<std::vector<std::vector<std::optional<std::string>>>> Iterator::next(
+    const size_t _batch_size) {
+  if (end()) {
+    return error("End is reached.");
+  }
+
+  std::vector<std::vector<std::optional<std::string>>> batch;
+
+  for (size_t i = 0; i < _batch_size; ++i) {
+    std::vector<std::optional<std::string>> new_row;
+
+    for (int j = 0; j < num_cols_; ++j) {
+      auto ptr = sqlite3_column_text(stmt_.get(), j);
+      if (ptr) {
+        new_row.emplace_back(
+            std::string(std::launder(reinterpret_cast<const char*>(ptr))));
+      } else {
+        new_row.emplace_back(std::nullopt);
+      }
+    }
+
+    batch.emplace_back(std::move(new_row));
+
+    step();
+
+    if (end()) {
+      return batch;
+    }
+  }
+
+  return batch;
+}
+
+}  // namespace sqlgen::sqlite
diff --git a/src/sqlgen_sqlite.cpp b/src/sqlgen_sqlite.cpp
index ec44eec..3008ade 100644
--- a/src/sqlgen_sqlite.cpp
+++ b/src/sqlgen_sqlite.cpp
@@ -1 +1,2 @@
 #include "sqlgen/sqlite/Connection.cpp"
+#include "sqlgen/sqlite/Iterator.cpp"