merkle range with test

1 files changed, 38 insertions, 12 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