Allow users to defer persistence when constructing tree nodes

3 files changed, 87 insertions, 16 deletions

diff --git a/lib/kernel/Merkle.ml b/lib/kernel/Merkle.ml
index b2e341a..2566941 100644
--- a/lib/kernel/Merkle.ml
+++ b/lib/kernel/Merkle.ml
@@ -16,15 +16,16 @@ module type TREE = functor (S : Abstract.Storage.STORAGE) (K : KEY) -> sig
 
   type node
 
+  val with_batch: S.transaction -> (unit -> (node, Concepts.Condition.condition) result) -> (node, Concepts.Condition.condition) result
+
   val find : S.transaction -> address -> (node option, Concepts.Condition.condition) result
+  val persist : S.transaction -> node -> (address, Concepts.Condition.condition) result
 
   val empty : node
   val empty_under : S.transaction -> (address, Concepts.Condition.condition) result
 
   val hash_of : node -> address
 
-  val persist : S.transaction -> node -> (address, Concepts.Condition.condition) result
-
   val insert : S.transaction -> K.t -> address -> node -> (node, Concepts.Condition.condition) result
   val remove : S.transaction -> K.t -> node -> (node, Concepts.Condition.condition) result
   val lookup : S.transaction -> K.t -> node -> (address option, Concepts.Condition.condition) result
@@ -121,14 +122,57 @@ module Make : TREE = functor (S : Abstract.Storage.STORAGE) (K : KEY) -> struct
      be used for intermediates and the like. *)
   let hash_of node = to_blob node |> Concepts.Hash.hash_of_blob
 
-  let find tx node =
+  type _ Effect.t += Intern: (address * node * Concepts.Blob.t) -> unit Effect.t | Retrieve: address -> (node * Concepts.Blob.t) option Effect.t
+
+  let rec flush tx buffer addr =
     let open Utilities.Result in
-    let* data = S.get tx (S.Hash node) in
-    match data with
-    | None -> Ok None
-    | Some data ->
-       let* node = from_blob data in
-       Ok (Some node)
+    match BatMap.find_opt addr buffer with
+    | None -> Ok ()
+    | Some (node, data) ->
+       let* _ = S.put tx (S.Hash addr) data in
+       match node with
+       | Leaf _ -> Ok ()
+       | Trunk { children; _ } ->
+          BatFingerTree.map (fun addr -> flush tx buffer addr) children
+          |> Utilities.FingerTree.sequence
+          |> Result.map ignore
+
+  let with_batch tx f =
+    let open Utilities.Result in
+    let buffer = ref BatMap.empty in
+    try
+      let* node = f () in
+      let* _ = flush tx !buffer (hash_of node) in
+      Ok node
+    with
+    | effect (Intern (addr, node, data)), k ->
+       buffer := BatMap.add addr (node, data) !buffer;
+       Effect.Deep.continue k ()
+    | effect (Retrieve addr), k ->
+       Effect.Deep.continue k (BatMap.find_opt addr !buffer)
+
+  let find tx addr =
+    let open Utilities.Result in
+    try Effect.perform (Retrieve addr) |> Option.map fst |> Result.ok
+    with
+    | Effect.Unhandled _ ->
+       let* data = S.get tx (S.Hash addr) in
+       match data with
+       | None -> Ok None
+       | Some data ->
+          let* node = from_blob data in
+          Ok (Some node)
+
+  let persist tx node =
+    let data = to_blob node in
+    let addr = Concepts.Hash.hash_of_blob data in
+    begin
+      try Ok (Effect.perform (Intern (addr, node, data)))
+      with
+      | Effect.Unhandled _ ->
+         S.put tx (S.Hash addr) data
+    end
+    |> Result.map (fun () -> addr)
 
   let find' tx addr =
     let open Utilities.Result in
@@ -147,13 +191,6 @@ module Make : TREE = functor (S : Abstract.Storage.STORAGE) (K : KEY) -> struct
        let* child = BatFingerTree.get children (if found then i+1 else i) |> find' tx in
        lookup tx key child
 
-  let persist tx node =
-    let open Utilities.Result in
-    let data = to_blob node in
-    let addr = Concepts.Hash.hash_of_blob data in
-    let* () = S.put tx (S.Hash addr) data in
-    Ok addr
-
   let empty_under tx = persist tx empty
 
   type op = Update of address * node | Split of address * K.t * address
@@ -311,6 +348,8 @@ module type INTERFACE = functor (S : Abstract.Storage.STORAGE) (K : KEY) (V : VA
   type address = Concepts.Hash.hash
   type node
 
+  val with_batch: S.transaction -> (unit -> (node, Concepts.Condition.condition) result) -> (node, Concepts.Condition.condition) result
+
   val find : S.transaction -> address -> (node option, Concepts.Condition.condition) result
 
   val empty : node
@@ -359,6 +398,7 @@ module Interface : INTERFACE = functor (S : Abstract.Storage.STORAGE) (K : KEY)
   let empty = T.empty
   let empty_under = T.empty_under
   let hash_of = T.hash_of
+  let with_batch = T.with_batch
 
   let insert tx k v node =
     let* addr = intern tx v in
diff --git a/test/test_rnt.ml b/test/test_rnt.ml
index 7bad887..b80de90 100644
--- a/test/test_rnt.ml
+++ b/test/test_rnt.ml
@@ -1,5 +1,6 @@
 let all_suites =
   Unit.Backend.Storage.suites ()
+  @ Unit.Kernel.Merkle.suites ()
   @ Unit.Kernel.SubstantialRelation.suites ()
   @ Integration.Registration.suites ()
 
diff --git a/test/unit/kernel/Merkle.ml b/test/unit/kernel/Merkle.ml
index d5a075e..85c756c 100644
--- a/test/unit/kernel/Merkle.ml
+++ b/test/unit/kernel/Merkle.ml
@@ -34,6 +34,35 @@ module Make (S : Abstract.Storage.STORAGE) (C : Helpers.Storage.CONFIGURATOR) =
     check (option string) "that in-memory reads from a value replacement work properly" (Some "toodles") v3';
     check (option string) "that in-memory reads from a non-existent key returns nothing" None bogus
 
+  let batching conn =
+    let intermediate, final =
+      begin
+        let open Utilities.Result in
+        let* tx = S.start conn in
+        let intermediate = ref None in
+        let* final = T.with_batch tx (fun () ->
+                         let* i = T.empty |> T.insert tx "tree" "maple" in
+                         intermediate := Some (T.hash_of i);
+                         T.insert tx "tree" "spruce" i) in
+        let* () = S.commit tx in
+        Ok (Option.get !intermediate, T.hash_of final)
+      end
+      |> Helpers.condition_as_failure
+    in
+    let i_node, f_node =
+      begin
+        let open Utilities.Result in
+        let* tx = S.start conn in
+        let* intermediate = T.find tx intermediate in
+        let* final = T.find tx final in
+        Ok (intermediate, final)
+      end
+      |> Helpers.condition_as_failure
+    in
+    check bool "that the intermediate node does not get persisted" true (Option.is_none i_node);
+    check bool "that the final node does get persisted" true (Option.is_some f_node)
+
+
   let persistence conn =
     let addr =
       begin
@@ -92,6 +121,7 @@ module Make (S : Abstract.Storage.STORAGE) (C : Helpers.Storage.CONFIGURATOR) =
   let suite prefix =
     ( "merkle/" ^ prefix,
       [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")])
 end