Merge pull request #14528 from MinaProtocol/feature/batch-merkle-path…

…-lookups

Batch merkle_path lookups in Sparse_ledger

src/lib/merkle_ledger/any_ledger.ml

@@ -120,6 +120,8 @@ module Make_base (Inputs : Inputs_intf) :
 
     let merkle_path (T ((module Base), t)) = Base.merkle_path t
 
+    let merkle_path_batch (T ((module Base), t)) = Base.merkle_path_batch t
+
     let merkle_root (T ((module Base), t)) = Base.merkle_root t
 
     let index_of_account_exn (T ((module Base), t)) =

src/lib/merkle_ledger/base_ledger_intf.ml

@@ -134,6 +134,8 @@ module type S = sig
 
   val merkle_path_at_index_exn : t -> int -> Path.t
 
+  val merkle_path_batch : t -> Location.t list -> Path.t list
+
   val remove_accounts_exn : t -> account_id list -> unit
 
   (** Triggers when the ledger has been detached and should no longer be

src/lib/merkle_ledger/database.ml

@@ -700,6 +700,58 @@ module Make (Inputs : Inputs_intf) :
     List.map2_exn dependency_dirs dependency_hashes ~f:(fun dir hash ->
         Direction.map dir ~left:(`Left hash) ~right:(`Right hash) )
 
+  let merkle_path_batch mdb locations =
+    let locations =
+      List.map locations ~f:(fun location ->
+          if Location.is_account location then
+            Location.Hash (Location.to_path_exn location)
+          else (
+            assert (Location.is_hash location) ;
+            location ) )
+    in
+    let rev_locations, rev_directions, rev_lengths =
+      let rec loop locations loc_acc dir_acc length_acc =
+        match (locations, length_acc) with
+        | [], _ :: length_acc ->
+            (loc_acc, dir_acc, length_acc)
+        | k :: locations, length :: length_acc ->
+            if Location.height ~ledger_depth:mdb.depth k >= mdb.depth then
+              loop locations loc_acc dir_acc (0 :: length :: length_acc)
+            else
+              let sibling = Location.sibling k in
+              let sibling_dir =
+                Location.last_direction (Location.to_path_exn k)
+              in
+              loop
+                (Location.parent k :: locations)
+                (sibling :: loc_acc) (sibling_dir :: dir_acc)
+                ((length + 1) :: length_acc)
+        | _ ->
+            assert false
+      in
+      loop locations [] [] [ 0 ]
+    in
+    let rev_hashes = get_hash_batch mdb rev_locations in
+    let rec loop directions hashes lengths acc =
+      match (directions, hashes, lengths, acc) with
+      | [], [], [], _ (* actually [] *) :: acc_tl ->
+          acc_tl
+      | _, _, 0 :: lengths, _ ->
+          loop directions hashes lengths ([] :: acc)
+      | ( direction :: directions
+        , hash :: hashes
+        , length :: lengths
+        , acc_hd :: acc_tl ) ->
+          let dir =
+            Direction.map direction ~left:(`Left hash) ~right:(`Right hash)
+          in
+          loop directions hashes ((length - 1) :: lengths)
+            ((dir :: acc_hd) :: acc_tl)
+      | _ ->
+          failwith "Mismatched lengths"
+    in
+    loop rev_directions rev_hashes rev_lengths [ [] ]
+
   let merkle_path_at_addr_exn t addr = merkle_path t (Location.Hash addr)
 
   let merkle_path_at_index_exn t index =

src/lib/merkle_ledger/null_ledger.ml

@@ -64,6 +64,8 @@ end = struct
     in
     loop location
 
+  let merkle_path_batch t locations = List.map ~f:(merkle_path t) locations
+
   let merkle_root t = empty_hash_at_height t.depth
 
   let merkle_path_at_addr_exn t addr = merkle_path t (Location.Hash addr)

src/lib/merkle_mask/masking_merkle_tree.ml

@@ -260,6 +260,16 @@ module Make (Inputs : Inputs_intf.S) = struct
       let parent_merkle_path = Base.merkle_path (get_parent t) location in
       fixup_merkle_path t parent_merkle_path address
 
+    let merkle_path_batch t locations =
+      assert_is_attached t ;
+      let addresses = List.map ~f:Location.to_path_exn locations in
+      let parent_merkle_paths =
+        Base.merkle_path_batch (get_parent t) locations
+      in
+      List.map2_exn
+        ~f:(fun path address -> fixup_merkle_path t path address)
+        parent_merkle_paths addresses
+
     (* given a Merkle path corresponding to a starting address, calculate
        addresses and hashes for each node affected by the starting hash; that is,
        along the path from the account address to root *)

src/lib/mina_ledger/sparse_ledger.ml

@@ -11,26 +11,25 @@ let of_ledger_subset_exn (oledger : Ledger.t) keys =
   let locations = Ledger.location_of_account_batch ledger keys in
   let non_empty_locations = List.filter_map ~f:snd locations in
   let accounts = Ledger.get_batch ledger non_empty_locations in
-  let sl, _ =
+  let merkle_paths = Ledger.merkle_path_batch ledger non_empty_locations in
+  let sl, _, _ =
     List.fold locations
-      ~init:(of_ledger_root ledger, accounts)
-      ~f:(fun (sl, accounts) (key, location) ->
+      ~init:(of_ledger_root ledger, accounts, merkle_paths)
+      ~f:(fun (sl, accounts, merkle_paths) (key, location) ->
         match location with
-        | Some loc -> (
-            match accounts with
-            | (_, account) :: rest ->
+        | Some _loc -> (
+            match (accounts, merkle_paths) with
+            | (_, account) :: rest, merkle_path :: rest_merkle_paths ->
                 let sl =
-                  add_path sl
-                    (Ledger.merkle_path ledger loc)
-                    key (Option.value_exn account)
+                  add_path sl merkle_path key (Option.value_exn account)
                 in
-                (sl, rest)
-            | [] ->
+                (sl, rest, rest_merkle_paths)
+            | _ ->
                 failwith "unexpected number of non empty accounts" )
         | None ->
             let path, account = Ledger.create_empty_exn ledger key in
             let sl = add_path sl path key account in
-            (sl, accounts) )
+            (sl, accounts, merkle_paths) )
   in
   Debug_assert.debug_assert (fun () ->
       [%test_eq: Ledger_hash.t]