merkle range with test

2 files changed, 81 insertions, 13 deletions

diff --git a/lib/kernel/Merkle.ml b/lib/kernel/Merkle.ml
index 714f5ba..4882ec4 100644
--- a/lib/kernel/Merkle.ml
+++ b/lib/kernel/Merkle.ml
@@ -35,6 +35,8 @@ module type TREE = functor (S : Abstract.Storage.STORAGE) (K : KEY) -> sig
   val lookup : S.transaction -> K.t -> node -> (address option, Concepts.Condition.condition) result
 
   val fold_left :
+    ?from:K.t ->
+    ?until:(K.t -> bool) ->
     S.transaction ->
     ('c -> K.t -> address -> 'c) ->
     'c ->
@@ -341,19 +343,43 @@ functor
             (fun node (key, value) -> insert tx key value node)
             empty bindings )
 
-    let rec fold_left tx f acc =
+    let drop i ft =
+      if i < BatFingerTree.size ft then snd (BatFingerTree.split_at ft i) else BatFingerTree.empty
+
+    let fold_left ?from ?(until = fun _ -> false) tx f acc node =
       let open Utilities.Result in
-      function
-      | Leaf {keys; values} ->
-          Utilities.FingerTree.zip keys values
-          |> BatFingerTree.fold_left (fun acc (k, v) -> f acc k v) acc
-          |> Result.ok
-      | Trunk {children; _} ->
-          Utilities.FingerTree.foldm_left
-            (fun acc addr ->
-              let* node = find' tx addr in
-              fold_left tx f acc node )
-            acc children
+      let rec walk from acc = function
+        | Leaf {keys; values} ->
+            let i = Option.fold ~none:0 ~some:(fun key -> fst (lookup1 keys key)) from in
+            Utilities.FingerTree.zip (drop i keys) (drop i values)
+            |> BatFingerTree.fold_left
+                 (fun step (k, v) ->
+                   match step with
+                   | `Continue acc -> if until k then `Stop acc else `Continue (f acc k v)
+                   | `Stop _ -> step )
+                 (`Continue acc)
+            |> Result.ok
+        | Trunk {keys; children} ->
+            let i =
+              Option.fold ~none:0
+                ~some:(fun key ->
+                  let i, found = lookup1 keys key in
+                  if found then i + 1 else i )
+                from
+            in
+            Utilities.FingerTree.foldm_left
+              (fun (from, step) addr ->
+                match step with
+                | `Stop _ -> Ok (from, step)
+                | `Continue acc ->
+                    let* node = find' tx addr in
+                    let* step = walk from acc node in
+                    Ok (None, step) )
+              (from, `Continue acc)
+              (drop i children)
+            |> Result.map snd
+      in
+      walk from acc node |> Result.map (function `Continue acc | `Stop acc -> acc)
 
     let keys tx node =
       fold_left tx (fun acc k _ -> BatFingerTree.snoc acc k) BatFingerTree.empty node
diff --git a/test/unit/kernel/Merkle.ml b/test/unit/kernel/Merkle.ml
index 6ac0f63..b54f8d3 100644
--- a/test/unit/kernel/Merkle.ml
+++ b/test/unit/kernel/Merkle.ml
@@ -7,6 +7,8 @@ module Make (S : Abstract.Storage.STORAGE) (C : Helpers.Storage.CONFIGURATOR) =
   module T =
     Rnt.Kernel.Merkle.Interface (S) (Rnt.Kernel.Merkle.StringKey) (Rnt.Kernel.Merkle.StringKey)
 
+  module M = Rnt.Kernel.Merkle.Make (S) (Rnt.Kernel.Merkle.StringKey)
+
   (* TODO: be a bit more comprehensive (ideally, we want to test splits as well) *)
   let insert_and_lookup conn =
     let v1, v2, v3, v3', bogus =
@@ -146,13 +148,53 @@ module Make (S : Abstract.Storage.STORAGE) (C : Helpers.Storage.CONFIGURATOR) =
     check bool "removing every key leaves the empty tree" true
       (Concepts.Hash.hash_equals emptied (T.hash_of T.empty))
 
+  let range conn =
+    let key i = Printf.sprintf "k%04d" i in
+    let all = List.init 500 key in
+    let from_on k = List.filter (fun k' -> k' >= k) all in
+    let whole, present, absent, before, after, bounded, stopped =
+      begin
+        let open Utilities.Result in
+        let* tx = S.start conn in
+        let scrambled = List.init 500 (fun i -> key (i * 7919 mod 500)) in
+        let* root =
+          List.fold_left
+            (fun node k -> node |>> M.insert tx k (Concepts.Hash.hash_of_bytes (Bytes.of_string k)))
+            (Ok M.empty) scrambled
+        in
+        let walk ?from ?until () =
+          M.fold_left ?from ?until tx (fun acc k _ -> k :: acc) [] root |> Result.map List.rev
+        in
+        let* whole = walk () in
+        let* present = walk ~from:"k0200" () in
+        let* absent = walk ~from:"k0200a" () in
+        let* before = walk ~from:"a" () in
+        let* after = walk ~from:"z" () in
+        let* bounded = walk ~from:"k0200" ~until:(fun k -> k >= "k0300") () in
+        let* stopped = walk ~until:(fun _ -> true) () in
+        let* () = S.abort tx in
+        Ok (whole, present, absent, before, after, bounded, stopped)
+      end
+      |> Helpers.condition_as_failure
+    in
+    check (list string) "a walk with no range visits every key in order" all whole;
+    check (list string) "a walk from a stored key starts at it" (from_on "k0200") present;
+    check (list string) "a walk from an absent key starts at the next one" (from_on "k0201") absent;
+    check (list string) "a walk from before the first key visits every key" all before;
+    check (list string) "a walk from past the last key visits nothing" [] after;
+    check (list string) "a walk stops before the first key until accepts"
+      (List.init 100 (fun i -> key (200 + i)))
+      bounded;
+    check (list string) "a walk that stops at once visits nothing" [] stopped
+
   let suite =
     ( "kernel/merkle",
       [ test_case "insert-and-lookup" `Quick (H.with_connection insert_and_lookup "merkle-test");
         test_case "batching" `Quick (H.with_connection batching "merkle-test");
         test_case "persistence" `Quick (H.with_connection persistence "merkle-test");
         test_case "iteration" `Quick (H.with_connection iteration "merkle-test");
-        test_case "removal" `Quick (H.with_connection removal "merkle-test") ] )
+        test_case "removal" `Quick (H.with_connection removal "merkle-test");
+        test_case "range" `Quick (H.with_connection range "merkle-test") ] )
 end
 
 module LMDB = Make (Rnt.Backend.Storage.LMDB) (Helpers.Storage.LMDB_Configurator)