diff --git a/src/lib/merkle_mask/masking_merkle_tree.ml b/src/lib/merkle_mask/masking_merkle_tree.ml index 220de8b0dac8..c6a8889624b8 100644 --- a/src/lib/merkle_mask/masking_merkle_tree.ml +++ b/src/lib/merkle_mask/masking_merkle_tree.ml @@ -185,15 +185,44 @@ module Make (Inputs : Inputs_intf.S) = struct let get_batch t locations = 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 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 [] [] + in + let leftover_location_accounts_rev = + Base.get_batch (get_parent t) leftover_locations_rev + 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 in - found_accounts @ Base.get_batch (get_parent t) leftover_locations + go locations_with_locations_rev leftover_location_accounts_rev [] (* fixup_merkle_path patches a Merkle path reported by the parent, overriding with hashes which are stored in the mask *)