Feature/fast iteration (#5)

whilo · web-flow · commit 1013771e21be · 2026-03-14T12:02:51.000-07:00
* Add lookupGE, ForwardCursor, and Seq leaf-safe optimization

lookupGE: O(log n) ceiling lookup with zero allocations. Returns the
first stored element &gt;= key, or nil.

ForwardCursor: mutable forward-only cursor for efficient sequential
ceiling lookups. Amortized O(1) per seekGE when keys arrive in
ascending order. Supports next() for O(1) sequential advance and
current() for non-advancing read. 

Seq leafSafe: skips per-element comparator bounds check when the
entire current leaf is within the iteration range.

Also adds searchFirstFrom to ANode (binary search with lower bound
hint) and CLJS lookup-ge support.

* Fix CLJS type inference warnings in branch and btset

Add ^Branch type hints to merge/merge-split params, ^BTSet to helper
functions, and ^js/^BTSet to seek's seq field access to eliminate
"Cannot infer target type" warnings during CLJS compilation.

* Remove lookupGE/ForwardCursor from public Clojure API

Keep Java methods for internal use by datahike but do not expose
them in the Clojure namespace. Mark lookupGE, ForwardCursor, and
forwardCursor() as internal API subject to change in Javadoc.

Tests updated to use interop directly via a local helper.
diff --git a/src-clojure/org/replikativ/persistent_sorted_set/branch.cljs b/src-clojure/org/replikativ/persistent_sorted_set/branch.cljs
@@ -467,7 +467,7 @@
                                     acc))]
                      (set! _measure result)
                      result)))))
-  (merge [this next]
+  (merge [this ^Branch next]
     (let [sc1 subtree-count
           sc2 (.-subtree-count next)
           new-sc (if (and (>= sc1 0) (>= sc2 0)) (+ sc1 sc2) -1)
@@ -489,7 +489,7 @@
                new-sc
                new-measure
                settings)))
-  (merge-split [this next]
+  (merge-split [this ^Branch next]
     (let [;; Ensure children arrays exist
           c1 (ensure-children this)
           c2 (ensure-children next)
diff --git a/src-clojure/org/replikativ/persistent_sorted_set/btset.cljs b/src-clojure/org/replikativ/persistent_sorted_set/btset.cljs
@@ -216,6 +216,29 @@
                      cmp  (or cmp (.-comparator set))]
                  (await (node/lookup root (.-storage set) key cmp opts))))))
 
+(defn lookup-ge
+  "Look up the first element >= key (ceiling/GE lookup).
+   Returns nil if no element >= key exists. O(log n)."
+  [^BTSet set key cmp {:keys [sync?] :or {sync? true} :as opts}]
+  (async+sync sync?
+              (async
+               (let [root (await (-root set opts))
+                     cmp  (or cmp (.-comparator set))
+                     storage (.-storage set)]
+                 (loop [node root]
+                   (let [keys (.-keys node)
+                         len  (arrays/alength keys)]
+                     (if (== len 0)
+                       nil
+                       (let [idx (binary-search-l cmp keys (dec len) key)]
+                         (if (>= idx len)
+                           nil
+                           (if (instance? Branch node)
+                             (let [child-node (await (branch/child node storage idx opts))]
+                               (recur child-node))
+                             ;; Leaf — return the key at idx
+                             (arrays/aget keys idx)))))))))))
+
 (defn measure
   "Get the aggregated statistics for the entire set."
   [^BTSet set {:keys [sync?] :or {sync? true} :as opts}]
@@ -247,11 +270,11 @@
 
 (def ^:const EMPTY_PATH (js* "0n"))
 
-(defn- bits-per-level [set]
+(defn- bits-per-level [^BTSet set]
   (let [bf (get (.-settings set) :branching-factor)]
     (Math/ceil (Math/log2 bf))))
 
-(defn- max-len [set]
+(defn- max-len [^BTSet set]
   (get (.-settings set) :branching-factor))
 
 (defn- min-len [set]
@@ -1383,7 +1406,7 @@
    (assert (some? seq))
    (if (fn? arg)
      (seek seq key arg {:sync? true})
-     (seek seq key (.-comparator (.-set seq)) arg)))
+     (seek seq key (.-comparator ^BTSet (.-set ^js seq)) arg)))
   ([seq key cmp {:keys [sync?] :or {sync? true} :as opts}]
    (assert (some? seq))
    (assert (fn? cmp))
diff --git a/src-java/org/replikativ/persistent_sorted_set/ANode.java b/src-java/org/replikativ/persistent_sorted_set/ANode.java
@@ -117,6 +117,34 @@ public int searchFirst(Key key, Comparator<Key> cmp) {
     return low;
   }
 
+  /**
+   * Like searchFirst but starts from a known lower bound.
+   * For forward cursors where we know the target is at or after startIdx.
+   * Uses linear scan for small distances (<=8), binary search otherwise.
+   */
+  public int searchFirstFrom(Key key, Comparator<Key> cmp, int startIdx) {
+    int remaining = _len - startIdx;
+    if (remaining <= 0) return _len;
+    // Linear scan for small distances — avoids binary search overhead
+    if (remaining <= 8) {
+      for (int i = startIdx; i < _len; i++) {
+        if (cmp.compare(_keys[i], key) >= 0) return i;
+      }
+      return _len;
+    }
+    // Binary search from startIdx
+    int low = startIdx, high = _len;
+    while (low < high) {
+      int mid = (high + low) >>> 1;
+      int d = cmp.compare(_keys[mid], key);
+      if (d < 0)
+        low = mid + 1;
+      else
+        high = mid;
+    }
+    return low;
+  }
+
   public int searchLast(Key key, Comparator<Key> cmp) {
     int low = 0, high = _len;
     while (low < high) {
diff --git a/src-java/org/replikativ/persistent_sorted_set/PersistentSortedSet.java b/src-java/org/replikativ/persistent_sorted_set/PersistentSortedSet.java
@@ -644,6 +644,210 @@ public Key lookup(Object key, Comparator<Key> cmp) {
     }
   }
 
+  /**
+   * Look up the first element >= key (ceiling/GE lookup).
+   * O(log n) with zero allocations — no Seq chain created.
+   * Returns null if no element >= key exists.
+   *
+   * <p><b>Internal API</b> — subject to change without notice.
+   * Not exposed in the public Clojure namespace.
+   */
+  public Key lookupGE(Object key) {
+    return lookupGE(key, _cmp);
+  }
+
+  public Key lookupGE(Object key, Comparator<Key> cmp) {
+    ANode<Key, Address> node = root();
+
+    if (node.len() == 0) {
+      return null;
+    }
+
+    while (true) {
+      int idx = node.searchFirst((Key) key, cmp);
+      if (idx >= node._len) {
+        return null;
+      }
+
+      if (node instanceof Branch) {
+        node = ((Branch<Key, Address>) node).child(_storage, idx);
+      } else {
+        return node._keys[idx];
+      }
+    }
+  }
+
+  /**
+   * Mutable forward-only cursor for efficient sequential lookupGE.
+   * For sorted lookup keys, amortized O(1) per lookup instead of O(log n).
+   * Not thread-safe. Must only be used for forward (ascending) seeks.
+   *
+   * <p><b>Internal API</b> — this class is subject to change or removal
+   * without notice. It is not exposed in the public Clojure namespace.
+   * External consumers should use {@link #lookupGE} or {@code slice} instead.
+   *
+   * <p>Usage:
+   * <pre>
+   *   ForwardCursor c = pss.forwardCursor();
+   *   Key result1 = c.seekGE(key1);  // O(log n) — first seek
+   *   Key result2 = c.seekGE(key2);  // O(1) if in same leaf, else O(siblings skipped)
+   * </pre>
+   */
+  public class ForwardCursor {
+    private final Comparator<Key> _cursorCmp;
+    private ANode<Key, Address> _leaf;     // current leaf node
+    private int _leafIdx;                   // current index within leaf
+    // Stack for tree traversal (height levels)
+    // For a tree of height H, we need H-1 branch levels above the leaf.
+    private Branch<Key, Address>[] _branches;
+    private int[] _branchIdxs;
+    private int _depth;                     // number of branch levels (0 for leaf-only)
+
+    @SuppressWarnings("unchecked")
+    ForwardCursor(Comparator<Key> cmp) {
+      _cursorCmp = cmp;
+      ANode<Key, Address> root = root();
+      if (root.len() == 0) {
+        _leaf = null;
+        _depth = 0;
+        return;
+      }
+      // Compute tree height
+      int height = 0;
+      ANode<Key, Address> n = root;
+      while (n instanceof Branch) {
+        height++;
+        n = ((Branch<Key, Address>) n).child(_storage, 0);
+      }
+      _depth = height;
+      _branches = new Branch[height];
+      _branchIdxs = new int[height];
+      // Position at start (leftmost leaf)
+      n = root;
+      for (int level = 0; level < height; level++) {
+        _branches[level] = (Branch<Key, Address>) n;
+        _branchIdxs[level] = 0;
+        n = ((Branch<Key, Address>) n).child(_storage, 0);
+      }
+      _leaf = n;
+      _leafIdx = 0;
+    }
+
+    /**
+     * Seek forward to first element >= key.
+     * Keys MUST be passed in ascending order across calls.
+     * Returns null if no element >= key exists.
+     */
+    public Key seekGE(Key key) {
+      if (_leaf == null) return null;
+
+      // Fast path: key is within current leaf
+      if (_cursorCmp.compare(key, _leaf.maxKey()) <= 0) {
+        // Search from current position — target is always >= _leafIdx since keys are ascending
+        int idx = _leaf.searchFirstFrom(key, _cursorCmp, _leafIdx);
+        if (idx < _leaf._len) {
+          _leafIdx = idx;
+          return _leaf._keys[idx];
+        }
+      }
+
+      // Need to advance to a later leaf.
+      // Walk up the branch stack to find a branch that contains our key,
+      // then walk back down.
+      int level = _depth - 1; // start from immediate parent of leaf
+      while (level >= 0) {
+        Branch<Key, Address> branch = _branches[level];
+        int bi = _branchIdxs[level] + 1; // advance past current child
+        // Linear scan forward through siblings (amortized O(1))
+        while (bi < branch._len) {
+          if (_cursorCmp.compare(key, branch._keys[bi]) <= 0) {
+            // key <= this child's maxKey, so answer is in this subtree
+            _branchIdxs[level] = bi;
+            // Walk down to leaf
+            ANode<Key, Address> node = branch.child(_storage, bi);
+            for (int d = level + 1; d < _depth; d++) {
+              _branches[d] = (Branch<Key, Address>) node;
+              int childIdx = node.searchFirst(key, _cursorCmp);
+              if (childIdx >= node._len) childIdx = node._len - 1;
+              _branchIdxs[d] = childIdx;
+              node = ((Branch<Key, Address>) node).child(_storage, childIdx);
+            }
+            _leaf = node;
+            int idx = _leaf.searchFirst(key, _cursorCmp);
+            if (idx < _leaf._len) {
+              _leafIdx = idx;
+              return _leaf._keys[idx];
+            }
+            // Key exceeds this leaf — continue to next sibling at this level
+            bi++;
+            continue;
+          }
+          bi++;
+        }
+        level--; // go up one level
+      }
+
+      // Exhausted all branches
+      _leaf = null;
+      return null;
+    }
+
+    /**
+     * Advance cursor to the next element and return it.
+     * Returns null if no more elements exist.
+     * O(1) within a leaf, amortized O(1) across leaves.
+     */
+    public Key next() {
+      if (_leaf == null) return null;
+      _leafIdx++;
+      if (_leafIdx < _leaf._len) {
+        return _leaf._keys[_leafIdx];
+      }
+      // Need to advance to next leaf via branch stack
+      for (int level = _depth - 1; level >= 0; level--) {
+        int bi = _branchIdxs[level] + 1;
+        if (bi < _branches[level]._len) {
+          _branchIdxs[level] = bi;
+          ANode<Key, Address> node = _branches[level].child(_storage, bi);
+          for (int d = level + 1; d < _depth; d++) {
+            _branches[d] = (Branch<Key, Address>) node;
+            _branchIdxs[d] = 0;
+            node = ((Branch<Key, Address>) node).child(_storage, 0);
+          }
+          _leaf = node;
+          _leafIdx = 0;
+          return _leaf._keys[0];
+        }
+      }
+      _leaf = null;
+      return null;
+    }
+
+    /**
+     * Return the current element without advancing.
+     * Returns null if the cursor is exhausted or not yet positioned.
+     */
+    public Key current() {
+      if (_leaf == null || _leafIdx < 0 || _leafIdx >= _leaf._len) return null;
+      return _leaf._keys[_leafIdx];
+    }
+
+  }
+
+  /**
+   * Create a forward cursor positioned at the start of the set.
+   * Use for efficient sequential lookupGE with ascending keys.
+   *
+   * <p><b>Internal API</b> — subject to change without notice.
+   */
+  public ForwardCursor forwardCursor() {
+    return new ForwardCursor(_cmp);
+  }
+
+  public ForwardCursor forwardCursor(Comparator<Key> cmp) {
+    return new ForwardCursor(cmp);
+  }
+
   // IEditableCollection
   public PersistentSortedSet asTransient() {
     if (editable()) {
diff --git a/src-java/org/replikativ/persistent_sorted_set/Seq.java b/src-java/org/replikativ/persistent_sorted_set/Seq.java
@@ -13,6 +13,8 @@ public class Seq extends ASeq implements IReduce, Reversible, IChunkedSeq, ISeek
   final Comparator _cmp;
   final boolean _asc;
   final int _version;
+  // True when the entire current leaf is within bounds — skip over() check
+  boolean _leafSafe;
 
   Seq(IPersistentMap meta, PersistentSortedSet set, Seq parent, ANode node, int idx, Object keyTo, Comparator cmp, boolean asc, int version) {
     super(meta);
@@ -24,6 +26,20 @@ public class Seq extends ASeq implements IReduce, Reversible, IChunkedSeq, ISeek
     _cmp    = cmp;
     _asc    = asc;
     _version = version;
+    _leafSafe = computeLeafSafe();
+  }
+
+  // Check if all remaining elements in the current leaf are within bounds.
+  // Avoids per-element comparator calls during within-leaf iteration.
+  boolean computeLeafSafe() {
+    if (_keyTo == null || _cmp == null) return true;
+    if (_asc) {
+      // All elements up to leaf end are valid if maxKey <= _keyTo
+      return _cmp.compare(_node.maxKey(), _keyTo) <= 0;
+    } else {
+      // All elements down to leaf start are valid if minKey >= _keyTo
+      return _cmp.compare(_node.minKey(), _keyTo) >= 0;
+    }
   }
 
   void checkVersion() {
@@ -47,25 +63,27 @@ boolean advance() {
     if (_asc) {
       if (_idx < _node._len - 1) {
         _idx++;
-        return !over();
+        return _leafSafe || !over();
       } else if (_parent != null) {
         _parent = _parent.next();
         if (_parent != null) {
           _node = _parent.child();
           _idx = 0;
-          return !over();
+          _leafSafe = computeLeafSafe();
+          return _leafSafe || !over();
         }
       }
     } else { // !_asc
       if (_idx > 0) {
         _idx--;
-        return !over();
+        return _leafSafe || !over();
       } else if (_parent != null) {
         _parent = _parent.next();
         if (_parent != null) {
           _node = _parent.child();
           _idx = _node._len - 1;
-          return !over();
+          _leafSafe = computeLeafSafe();
+          return _leafSafe || !over();
         }
       }
     }
diff --git a/test-clojure/org/replikativ/persistent_sorted_set/test/replace_and_lookup.cljc b/test-clojure/org/replikativ/persistent_sorted_set/test/replace_and_lookup.cljc
