Merge branch 'feature/batch-account-lookups' into feature/batch-merkl…

…e-path-lookups

src/lib/merkle_ledger/database.ml

@@ -693,32 +693,12 @@ module Make (Inputs : Inputs_intf) :
         Location.Hash (Location.to_path_exn location)
       else location
     in
-    assert (Location.is_hash location) ;
-    let rev_locations, rev_directions =
-      let rec loop k loc_acc dir_acc =
-        if Location.height ~ledger_depth:mdb.depth k >= mdb.depth then
-          (loc_acc, dir_acc)
-        else
-          let sibling = Location.sibling k in
-          let sibling_dir = Location.last_direction (Location.to_path_exn k) in
-          loop (Location.parent k) (sibling :: loc_acc) (sibling_dir :: dir_acc)
-      in
-      loop location [] []
-    in
-    let rev_hashes = get_hash_batch mdb rev_locations in
-    let rec loop directions hashes acc =
-      match (directions, hashes) with
-      | [], [] ->
-          acc
-      | direction :: directions, hash :: hashes ->
-          let dir =
-            Direction.map direction ~left:(`Left hash) ~right:(`Right hash)
-          in
-          loop directions hashes (dir :: acc)
-      | _ ->
-          failwith "Mismatched lengths"
+    let dependency_locs, dependency_dirs =
+      List.unzip (Location.merkle_path_dependencies_exn location)
     in
-    loop rev_directions rev_hashes []
+    let dependency_hashes = get_hash_batch mdb dependency_locs in
+    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 =

src/lib/merkle_ledger/location.ml

@@ -147,6 +147,20 @@ module T = struct
     | Right ->
         (sibling, base)
 
+  let merkle_path_dependencies_exn (location : t) : (t * Direction.t) list =
+    let rec loop k acc =
+      if Addr.depth k = 0 then acc
+      else
+        let sibling = Hash (Addr.sibling k) in
+        let sibling_dir = last_direction k in
+        loop (Addr.parent_exn k) ((sibling, sibling_dir) :: acc)
+    in
+    match location with
+    | Hash addr ->
+        List.rev (loop addr [])
+    | _ ->
+        failwith "can only get merkle path dependencies of a hash location"
+
   type location = t [@@deriving sexp, compare]
 
   include Comparable.Make (struct

src/lib/merkle_ledger/location_intf.ml

@@ -48,5 +48,7 @@ module type S = sig
 
   val order_siblings : t -> 'a -> 'a -> 'a * 'a
 
+  val merkle_path_dependencies_exn : t -> (t * Direction.t) list
+
   include Comparable.S with type t := t
 end

src/lib/merkle_mask/masking_merkle_tree.ml

@@ -183,46 +183,33 @@ module Make (Inputs : Inputs_intf.S) = struct
       | None ->
           Base.get (get_parent t) location
 
-    let get_batch t locations =
-      assert_is_attached t ;
-      let locations_with_locations_rev, leftover_locations_rev =
-        let rec go locations locations_with_locations_rev leftover_locations_rev
-            =
-          match locations with
-          | [] ->
-              (locations_with_locations_rev, leftover_locations_rev)
-          | location :: locations -> (
-              match self_find_account t location with
-              | None ->
-                  go locations
-                    ((location, None) :: locations_with_locations_rev)
-                    (location :: leftover_locations_rev)
-              | Some account ->
-                  go locations
-                    ((location, Some account) :: locations_with_locations_rev)
-                    leftover_locations_rev )
-        in
-        go locations [] []
+    let self_find_or_batch_lookup self_find lookup_parent t ids =
+      assert_is_attached t ;
+      let self_found_or_none =
+        List.map ids ~f:(fun id -> (id, self_find t id))
       in
-      let leftover_location_accounts_rev =
-        Base.get_batch (get_parent t) leftover_locations_rev
+      let not_found =
+        List.filter_map self_found_or_none ~f:(function
+          | id, None ->
+              Some id
+          | _ ->
+              None )
       in
-      let rec go locations_with_locations_rev leftover_locations_rev accounts =
-        match (locations_with_locations_rev, leftover_locations_rev) with
-        | [], _ ->
-            accounts
-        | ( (location, None) :: locations_with_locations_rev
-          , (_location, account) :: leftover_locations_rev ) ->
-            go locations_with_locations_rev leftover_locations_rev
-              ((location, account) :: accounts)
-        | ( (location, Some account) :: locations_with_locations_rev
-          , leftover_locations_rev ) ->
-            go locations_with_locations_rev leftover_locations_rev
-              ((location, Some account) :: accounts)
-        | _ :: _, [] ->
-            assert false
+      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
-      go locations_with_locations_rev leftover_location_accounts_rev []
+      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 *)
@@ -572,46 +559,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 account_ids_with_locations_rev, leftover_account_ids_rev =
-        let rec go account_ids account_ids_with_locations_rev
-            leftover_account_ids_rev =
-          match account_ids with
-          | [] ->
-              (account_ids_with_locations_rev, leftover_account_ids_rev)
-          | account_id :: account_ids -> (
-              match self_find_location t account_id with
-              | None ->
-                  go account_ids
-                    ((account_id, None) :: account_ids_with_locations_rev)
-                    (account_id :: leftover_account_ids_rev)
-              | Some loc ->
-                  go account_ids
-                    ((account_id, Some loc) :: account_ids_with_locations_rev)
-                    leftover_account_ids_rev )
-        in
-        go account_ids [] []
-      in
-      let leftover_account_id_locs_rev =
-        Base.location_of_account_batch (get_parent t) leftover_account_ids_rev
-      in
-      let rec go account_ids_with_locations_rev leftover_account_ids_rev locs =
-        match (account_ids_with_locations_rev, leftover_account_ids_rev) with
-        | [], _ ->
-            locs
-        | ( (account_id, None) :: account_ids_with_locations_rev
-          , (_account_id, loc) :: leftover_account_ids_rev ) ->
-            go account_ids_with_locations_rev leftover_account_ids_rev
-              ((account_id, loc) :: locs)
-        | ( (account_id, Some loc) :: account_ids_with_locations_rev
-          , leftover_account_ids_rev ) ->
-            go account_ids_with_locations_rev leftover_account_ids_rev
-              ((account_id, Some loc) :: locs)
-        | _ :: _, [] ->
-            assert false
-      in
-      go account_ids_with_locations_rev leftover_account_id_locs_rev []
+    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

@@ -12,36 +12,27 @@ let of_ledger_subset_exn (oledger : Ledger.t) keys =
   let non_empty_locations = List.filter_map ~f:snd locations in
   let accounts = Ledger.get_batch ledger non_empty_locations in
   let merkle_paths = Ledger.merkle_path_batch ledger non_empty_locations in
-  let _, sparse =
-    let rec go (new_keys, sl) locations accounts merkle_paths =
-      match locations with
-      | [] ->
-          (new_keys, sl)
-      | (key, Some _loc) :: locations -> (
+  let sl, _, _ =
+    List.fold locations
+      ~init:(of_ledger_root ledger, accounts, merkle_paths)
+      ~f:(fun (sl, accounts, merkle_paths) (key, location) ->
+        if Option.is_some location then
           match (accounts, merkle_paths) with
           | (_, account) :: accounts, merkle_path :: merkle_paths ->
-              go
-                ( new_keys
-                , add_path sl merkle_path key
-                    ( account
-                    |> Option.value_exn ?here:None ?error:None ?message:None )
-                )
-                locations accounts merkle_paths
+              let sl = add_path sl merkle_path key (Option.value_exn account) in
+              (sl, accounts, merkle_paths)
           | _ ->
-              assert false )
-      | (key, None) :: locations ->
-          let path, acct = Ledger.create_empty_exn ledger key in
-          go
-            (key :: new_keys, add_path sl path key acct)
-            locations accounts merkle_paths
-    in
-    go ([], of_ledger_root ledger) locations accounts merkle_paths
+              failwith "unexpected number of non empty accounts or merkle paths"
+        else
+          let path, account = Ledger.create_empty_exn ledger key in
+          let sl = add_path sl path key account in
+          (sl, accounts, merkle_paths) )
   in
   Debug_assert.debug_assert (fun () ->
       [%test_eq: Ledger_hash.t]
         (Ledger.merkle_root ledger)
-        ((merkle_root sparse :> Random_oracle.Digest.t) |> Ledger_hash.of_hash) ) ;
-  sparse
+        ((merkle_root sl :> Random_oracle.Digest.t) |> Ledger_hash.of_hash) ) ;
+  sl
 
 let of_ledger_index_subset_exn (ledger : Ledger.Any_ledger.witness) indexes =
   List.fold indexes