2 files changed, 149 insertions, 19 deletions

diff --git a/src/database/database.cpp b/src/database/database.cpp
index 76d3a2ab..03451c05 100644
--- a/src/database/database.cpp
+++ b/src/database/database.cpp
@@ -756,6 +756,18 @@ auto Database::dbGetPage(const Continuation& c) -> Result<T>* {
   return new Result<T>(std::move(records), next_page, prev_page);
 }
 
+auto Database::dbCount(const Continuation& c) -> size_t {
+  std::unique_ptr<leveldb::Iterator> it{
+      db_->NewIterator(leveldb::ReadOptions{})};
+  size_t count = 0;
+  for (it->Seek({c.start_key.data(), c.start_key.size()});
+       it->Valid() && it->key().starts_with({c.prefix.data(), c.prefix.size()});
+       it->Next()) {
+    count++;
+  }
+  return count;
+}
+
 template auto Database::dbGetPage<Track>(const Continuation& c)
     -> Result<Track>*;
 template auto Database::dbGetPage<std::pmr::string>(const Continuation& c)
@@ -880,6 +892,12 @@ Iterator::Iterator(const Iterator& other)
       current_pos_(other.current_pos_),
       prev_pos_(other.prev_pos_) {}
 
+Iterator& Iterator::operator=(const Iterator& other) {
+  current_pos_ = other.current_pos_;
+  prev_pos_ = other.prev_pos_;
+  return *this;
+}
+
 auto Iterator::Next(Callback cb) -> void {
   auto db = db_.lock();
   if (!db) {
@@ -910,24 +928,35 @@ auto Iterator::NextSync() -> std::optional<IndexRecord> {
   if (!db) {
     return {};
   }
-  return db->worker_task_
-      ->Dispatch<std::optional<IndexRecord>>(
-          [=]() -> std::optional<IndexRecord> {
-            std::lock_guard lock{pos_mutex_};
-            if (!current_pos_) {
-              return {};
-            }
-            std::unique_ptr<Result<IndexRecord>> res{
-                db->dbGetPage<IndexRecord>(*current_pos_)};
-            prev_pos_ = current_pos_;
-            current_pos_ = res->next_page();
-            if (!res || res->values().empty() || !res->values()[0]) {
-              ESP_LOGI(kTag, "dropping empty result");
-              return {};
-            }
-            return *res->values()[0];
-          })
-      .get();
+  std::lock_guard lock{pos_mutex_};
+  if (!current_pos_) {
+    return {};
+  }
+  std::unique_ptr<Result<IndexRecord>> res{
+      db->dbGetPage<IndexRecord>(*current_pos_)};
+  prev_pos_ = current_pos_;
+  current_pos_ = res->next_page();
+  if (!res || res->values().empty() || !res->values()[0]) {
+    ESP_LOGI(kTag, "dropping empty result");
+    return {};
+  }
+  return *res->values()[0];
+}
+
+auto Iterator::PeekSync() -> std::optional<IndexRecord> {
+  auto db = db_.lock();
+  if (!db) {
+    return {};
+  }
+  auto pos = current_pos_;
+  if (!pos) {
+    return {};
+  }
+  std::unique_ptr<Result<IndexRecord>> res{db->dbGetPage<IndexRecord>(*pos)};
+  if (!res || res->values().empty() || !res->values()[0]) {
+    return {};
+  }
+  return *res->values()[0];
 }
 
 auto Iterator::Prev(Callback cb) -> void {
@@ -950,8 +979,82 @@ auto Iterator::Prev(Callback cb) -> void {
   });
 }
 
+auto Iterator::Size() const -> size_t {
+  auto db = db_.lock();
+  if (!db) {
+    return {};
+  }
+  std::optional<Continuation> pos = current_pos_;
+  if (!pos) {
+    return 0;
+  }
+  return db->dbCount(*pos);
+}
+
 auto Iterator::InvokeNull(Callback cb) -> void {
   std::invoke(cb, std::optional<IndexRecord>{});
 }
 
+TrackIterator::TrackIterator(const Iterator& it) : db_(it.db_), levels_() {
+  if (it.current_pos_) {
+    levels_.push_back(it);
+  }
+  NextLeaf();
+}
+
+TrackIterator::TrackIterator(const TrackIterator& other)
+    : db_(other.db_), levels_(other.levels_) {}
+
+TrackIterator& TrackIterator::operator=(TrackIterator&& other) {
+  levels_ = std::move(other.levels_);
+  return *this;
+}
+
+auto TrackIterator::Next() -> std::optional<TrackId> {
+  std::optional<TrackId> next{};
+  while (!next && !levels_.empty()) {
+    auto next_record = levels_.back().NextSync();
+    if (!next_record) {
+      levels_.pop_back();
+      NextLeaf();
+      continue;
+    }
+    // May still be nullopt_t; hence the loop.
+    next = next_record->track();
+  }
+  return next;
+}
+
+auto TrackIterator::Size() const -> size_t {
+  size_t size = 0;
+  TrackIterator copy{*this};
+  while (!copy.levels_.empty()) {
+    size += copy.levels_.back().Size();
+    copy.levels_.pop_back();
+    copy.NextLeaf();
+  }
+  return size;
+}
+
+auto TrackIterator::NextLeaf() -> void {
+  while (!levels_.empty()) {
+    ESP_LOGI(kTag, "check next candidate");
+    Iterator& candidate = levels_.back();
+    auto next = candidate.PeekSync();
+    if (!next) {
+      ESP_LOGI(kTag, "candidate is empty");
+      levels_.pop_back();
+      continue;
+    }
+    if (!next->track()) {
+      ESP_LOGI(kTag, "candidate is a branch");
+      candidate.NextSync();
+      levels_.push_back(Iterator{db_, next->Expand(1).value()});
+      continue;
+    }
+    ESP_LOGI(kTag, "candidate is a leaf");
+    break;
+  }
+}
+
 }  // namespace database
diff --git a/src/database/include/database.hpp b/src/database/include/database.hpp
index 36f734b8..263153fb 100644
--- a/src/database/include/database.hpp
+++ b/src/database/include/database.hpp
@@ -12,6 +12,7 @@
 #include <future>
 #include <memory>
 #include <optional>
+#include <stack>
 #include <string>
 #include <utility>
 #include <vector>
@@ -168,6 +169,8 @@ class Database {
   template <typename T>
   auto dbGetPage(const Continuation& c) -> Result<T>*;
 
+  auto dbCount(const Continuation& c) -> size_t;
+
   template <typename T>
   auto ParseRecord(const leveldb::Slice& key, const leveldb::Slice& val)
       -> std::shared_ptr<T>;
@@ -193,7 +196,9 @@ class Iterator {
  public:
   Iterator(std::weak_ptr<Database>, const IndexInfo&);
   Iterator(std::weak_ptr<Database>, const Continuation&);
-  Iterator(const Iterator &);
+  Iterator(const Iterator&);
+
+  Iterator& operator=(const Iterator& other);
 
   auto database() const { return db_; }
 
@@ -204,8 +209,13 @@ class Iterator {
 
   auto Prev(Callback) -> void;
 
+  auto PeekSync() -> std::optional<IndexRecord>;
+
+  auto Size() const -> size_t;
 
  private:
+  friend class TrackIterator;
+
   auto InvokeNull(Callback) -> void;
 
   std::weak_ptr<Database> db_;
@@ -215,4 +225,21 @@ class Iterator {
   std::optional<Continuation> prev_pos_;
 };
 
+class TrackIterator {
+ public:
+  TrackIterator(const Iterator&);
+  TrackIterator(const TrackIterator&);
+
+  TrackIterator& operator=(TrackIterator&& other);
+
+  auto Next() -> std::optional<TrackId>;
+  auto Size() const -> size_t;
+
+ private:
+  auto NextLeaf() -> void;
+
+  std::weak_ptr<Database> db_;
+  std::vector<Iterator> levels_;
+};
+
 }  // namespace database