Merge pull request #14522 from MinaProtocol/feature/batch-account-loc…

…ation-lookups

Batch account location lookups for sparse ledger

src/lib/merkle_ledger/location.ml

@@ -147,17 +147,23 @@ module T = struct
     | Right ->
         (sibling, base)
 
+  (* Returns a reverse of traversal path from top of the tree to the location
+     (direction to take and sibling's hash).contents
+
+     By reverse it means that head of returned list contains direction from
+     location's parent to the location along with the location's sibling.
+  *)
   let merkle_path_dependencies_exn (location : t) : (t * Direction.t) list =
-    let rec loop k acc =
-      if Addr.depth k = 0 then acc
+    let rec loop k =
+      if Addr.depth k = 0 then []
       else
         let sibling = Hash (Addr.sibling k) in
-        let sibling_dir = last_direction k in
-        loop (Addr.parent_exn k) ((sibling, sibling_dir) :: acc)
+        let dir = last_direction k in
+        (sibling, dir) :: loop (Addr.parent_exn k)
     in
     match location with
     | Hash addr ->
-        List.rev (loop addr [])
+        loop addr
     | _ ->
         failwith "can only get merkle path dependencies of a hash location"
 

src/lib/merkle_mask/masking_merkle_tree.ml

@@ -183,17 +183,33 @@ module Make (Inputs : Inputs_intf.S) = struct
       | None ->
           Base.get (get_parent t) location
 
-    let get_batch t locations =
+    let self_find_or_batch_lookup self_find lookup_parent t ids =
       assert_is_attached t ;
-      let found_accounts, leftover_locations =
-        List.partition_map locations ~f:(fun location ->
-            match self_find_account t location with
-            | Some account ->
-                Either.first (location, Some account)
-            | None ->
-                Either.second location )
+      let self_found_or_none =
+        List.map ids ~f:(fun id -> (id, self_find t id))
+      in
+      let not_found =
+        List.filter_map self_found_or_none ~f:(function
+          | id, None ->
+              Some id
+          | _ ->
+              None )
+      in
+      let from_parent = lookup_parent (get_parent t) not_found in
+      let _, res =
+        List.fold_map self_found_or_none ~init:from_parent
+          ~f:(fun from_parent (id, self_found) ->
+            match (self_found, from_parent) with
+            | None, r :: rest ->
+                (rest, r)
+            | Some _, _ ->
+                (from_parent, (id, self_found))
+            | _ ->
+                failwith "unexpected number of results from DB" )
       in
-      found_accounts @ Base.get_batch (get_parent t) leftover_locations
+      res
+
+    let get_batch = self_find_or_batch_lookup self_find_account Base.get_batch
 
     (* fixup_merkle_path patches a Merkle path reported by the parent,
        overriding with hashes which are stored in the mask *)
@@ -533,18 +549,9 @@ module Make (Inputs : Inputs_intf.S) = struct
       | None ->
           Base.location_of_account (get_parent t) account_id
 
-    let location_of_account_batch t account_ids =
-      assert_is_attached t ;
-      let found_locations, leftover_account_ids =
-        List.partition_map account_ids ~f:(fun account_id ->
-            match self_find_location t account_id with
-            | Some location ->
-                Either.first (account_id, Some location)
-            | None ->
-                Either.second account_id )
-      in
-      found_locations
-      @ Base.location_of_account_batch (get_parent t) leftover_account_ids
+    let location_of_account_batch =
+      self_find_or_batch_lookup self_find_location
+        Base.location_of_account_batch
 
     (* not needed for in-memory mask; in the database, it's currently a NOP *)
     let make_space_for t =

src/lib/mina_ledger/sparse_ledger.ml

@@ -8,10 +8,11 @@ let of_ledger_root ledger =
 
 let of_ledger_subset_exn (oledger : Ledger.t) keys =
   let ledger = Ledger.copy oledger in
+  let locations = Ledger.location_of_account_batch ledger keys in
   let _, sparse =
-    List.fold keys
-      ~f:(fun (new_keys, sl) key ->
-        match Ledger.location_of_account ledger key with
+    List.fold locations
+      ~f:(fun (new_keys, sl) (key, loc) ->
+        match loc with
         | Some loc ->
             ( new_keys
             , add_path sl