diff --git a/src/lib/merkle_mask/masking_merkle_tree.ml b/src/lib/merkle_mask/masking_merkle_tree.ml index 62d93113394e..db07d05499a7 100644 --- a/src/lib/merkle_mask/masking_merkle_tree.ml +++ b/src/lib/merkle_mask/masking_merkle_tree.ml @@ -72,6 +72,7 @@ module Make (Inputs : Inputs_intf.S) = struct and for a few ancestors. This is used as a lookup cache. *) ; mutable accumulated : (accumulated_t[@sexp.opaque]) option + ; mutable unhashed_accounts : (Hash.t * Location.t) list } [@@deriving sexp] @@ -93,6 +94,7 @@ module Make (Inputs : Inputs_intf.S) = struct ; depth ; accumulated = None ; maps = empty_maps () + ; unhashed_accounts = [] } let get_uuid { uuid; _ } = uuid @@ -201,15 +203,186 @@ module Make (Inputs : Inputs_intf.S) = struct Option.iter t.accumulated ~f:(fun { current; next; _ } -> f current ; f next ) - let self_set_hash t address hash = + let self_set_hash_impl t address hash = update_maps t ~f:(fun maps -> maps.hashes <- Map.set maps.hashes ~key:address ~data:hash ) + let path_batch_impl ~fixup_path ~self_lookup ~base_lookup locations = + let self_paths = + List.map locations ~f:(fun location -> + let address = Location.to_path_exn location in + self_lookup address + |> Option.value_map + ~default:(Either.Second (location, address)) + ~f:Either.first ) + in + let all_parent_paths = + let locs = + List.filter_map self_paths ~f:(function + | Either.First _ -> + None + | Either.Second (location, _) -> + Some location ) + in + if List.is_empty locs then [] else base_lookup locs + in + let f parent_paths = function + | Either.First path -> + (parent_paths, path) + | Either.Second (_, address) -> + let path = fixup_path ~address (List.hd_exn parent_paths) in + (List.tl_exn parent_paths, path) + in + snd @@ List.fold_map ~init:all_parent_paths ~f self_paths + + let rec self_path_impl ~element ~depth address = + let height = Addr.height ~ledger_depth:depth address in + if height >= depth then Some [] + else + let%bind.Option el = element height address in + let%bind.Option parent_address = Addr.parent address |> Or_error.ok in + let%map.Option rest = self_path_impl ~element ~depth parent_address in + el :: rest + + let empty_hash = + Empty_hashes.extensible_cache (module Hash) ~init_hash:Hash.empty_account + + let self_path_get_hash ~hashes ~current_location height address = + match Map.find hashes address with + | Some hash -> + Some hash + | None -> + let is_empty = + match current_location with + | None -> + true + | Some current_location -> + let current_address = Location.to_path_exn current_location in + Addr.is_further_right ~than:current_address address + in + if is_empty then Some (empty_hash height) else None + + let self_merkle_path ~hashes ~current_location = + let element height address = + let sibling = Addr.sibling address in + let dir = Location.last_direction address in + let%map.Option sibling_hash = + self_path_get_hash ~hashes ~current_location height sibling + in + Direction.map dir ~left:(`Left sibling_hash) ~right:(`Right sibling_hash) + in + self_path_impl ~element + + (* fixup_merkle_path patches a Merkle path reported by the parent, + overriding with hashes which are stored in the mask *) + let fixup_merkle_path ~hashes ~address:init = + let f address = + (* first element in the path contains hash at sibling of address *) + let sibling_mask_hash = Map.find hashes (Addr.sibling address) in + let parent_addr = Addr.parent_exn address in + let open Option in + function + | `Left h -> + (parent_addr, `Left (value sibling_mask_hash ~default:h)) + | `Right h -> + (parent_addr, `Right (value sibling_mask_hash ~default:h)) + in + Fn.compose snd @@ List.fold_map ~init ~f + + let compute_merge_hashes : + (Hash.t * Addr.t * [ `Left of Hash.t | `Right of Hash.t ] list) list + -> (Addr.t * Hash.t) list = + let process_pair height = function + | (lh, laddr, `Left _ :: lpath), (rh, _, `Right _ :: _rpath) -> + (* Assertion: lpath == _rpath *) + let parent = Addr.parent_exn laddr in + let h = Hash.merge ~height lh rh in + (h, parent, lpath) + | _ -> + failwith "compute_merge_hashes: unexpected match of nodes" + in + let process_single height (self_hash, addr, path) = + let parent = Addr.parent_exn addr in + match path with + | `Left sibling_hash :: rest -> + let new_hash = Hash.merge ~height self_hash sibling_hash in + (new_hash, parent, rest) + | `Right sibling_hash :: rest -> + let new_hash = Hash.merge ~height sibling_hash self_hash in + (new_hash, parent, rest) + | _ -> + failwith "compute_merge_hashes: path is empty" + in + let converge height task = + let reversed, mlast = + List.fold task ~init:([], None) + ~f:(fun (processed, mprev) ((_, addr, _) as el) -> + match mprev with + | None -> + (processed, Some el) + | Some ((_, prev_addr, _) as prev) + when Addr.(equal @@ sibling prev_addr) addr -> + (process_pair height (prev, el) :: processed, None) + | Some prev -> + (process_single height prev :: processed, Some el) ) + in + List.rev_append reversed + @@ Option.(map ~f:(process_single height) mlast |> to_list) + in + let rec impl acc height task = + let acc' = + List.unordered_append (List.map ~f:(fun (a, b, _) -> (b, a)) task) acc + in + match task with + | [] | [ (_, _, []) ] -> + acc' + | _ -> + impl acc' (height + 1) (converge height task) + in + impl [] 0 + + let finalize_hashes_do t unhashed_accounts = + let merkle_path_batch = + let { hashes; _ }, ancestor = maps_and_ancestor t in + path_batch_impl + ~base_lookup:(Base.merkle_path_batch ancestor) + ~self_lookup: + (self_merkle_path ~current_location:t.current_location + ~depth:t.depth ~hashes ) + ~fixup_path:(fixup_merkle_path ~hashes) + (List.map ~f:snd unhashed_accounts) + in + (* let _task = *) + List.map2_exn + ~f:(fun (h, loc) p -> (h, Location.to_path_exn loc, p)) + unhashed_accounts merkle_path_batch + |> List.stable_sort ~compare:(fun (_, a, _) (_, b, _) -> + Addr.compare a b ) + |> List.remove_consecutive_duplicates ~which_to_keep:`Last + ~equal:(fun (_, a, _) (_, b, _) -> Addr.equal a b) + |> compute_merge_hashes + |> List.iter ~f:(Tuple2.uncurry @@ self_set_hash_impl t) + + let finalize_hashes t = + let unhashed_accounts = t.unhashed_accounts in + if not @@ List.is_empty unhashed_accounts then ( + t.unhashed_accounts <- [] ; + finalize_hashes_do t unhashed_accounts ) + + let self_set_hash t address hash = + finalize_hashes t ; + self_set_hash_impl t address hash + let set_inner_hash_at_addr_exn t address hash = assert_is_attached t ; assert (Addr.depth address <= t.depth) ; self_set_hash t address hash + let hashes_and_ancestor t = + finalize_hashes t ; + let { hashes; _ }, ancestor = maps_and_ancestor t in + (hashes, ancestor) + let self_set_location t account_id location = update_maps t ~f:(fun maps -> maps.locations <- @@ -303,44 +476,6 @@ module Make (Inputs : Inputs_intf.S) = struct in self_find_or_batch_lookup self_find Base.get_batch t - let empty_hash = - Empty_hashes.extensible_cache (module Hash) ~init_hash:Hash.empty_account - - let self_path_get_hash ~hashes ~current_location height address = - match Map.find hashes address with - | Some hash -> - Some hash - | None -> - let is_empty = - match current_location with - | None -> - true - | Some current_location -> - let current_address = Location.to_path_exn current_location in - Addr.is_further_right ~than:current_address address - in - if is_empty then Some (empty_hash height) else None - - let rec self_path_impl ~element ~depth address = - let height = Addr.height ~ledger_depth:depth address in - if height >= depth then Some [] - else - let%bind.Option el = element height address in - let%bind.Option parent_address = Addr.parent address |> Or_error.ok in - let%map.Option rest = self_path_impl ~element ~depth parent_address in - el :: rest - - let self_merkle_path ~hashes ~current_location = - let element height address = - let sibling = Addr.sibling address in - let dir = Location.last_direction address in - let%map.Option sibling_hash = - self_path_get_hash ~hashes ~current_location height sibling - in - Direction.map dir ~left:(`Left sibling_hash) ~right:(`Right sibling_hash) - in - self_path_impl ~element - let self_wide_merkle_path ~hashes ~current_location = let element height address = let sibling = Addr.sibling address in @@ -357,22 +492,6 @@ module Make (Inputs : Inputs_intf.S) = struct in self_path_impl ~element - (* fixup_merkle_path patches a Merkle path reported by the parent, - overriding with hashes which are stored in the mask *) - let fixup_merkle_path ~hashes ~address:init = - let f address = - (* first element in the path contains hash at sibling of address *) - let sibling_mask_hash = Map.find hashes (Addr.sibling address) in - let parent_addr = Addr.parent_exn address in - let open Option in - function - | `Left h -> - (parent_addr, `Left (value sibling_mask_hash ~default:h)) - | `Right h -> - (parent_addr, `Right (value sibling_mask_hash ~default:h)) - in - Fn.compose snd @@ List.fold_map ~init ~f - (* fixup_merkle_path patches a Merkle path reported by the parent, overriding with hashes which are stored in the mask *) let fixup_wide_merkle_path ~hashes ~address:init = @@ -401,9 +520,9 @@ module Make (Inputs : Inputs_intf.S) = struct let merkle_path_at_addr_exn t address = assert_is_attached t ; - let maps, ancestor = maps_and_ancestor t in + let hashes, ancestor = hashes_and_ancestor t in match - self_merkle_path ~depth:t.depth ~hashes:maps.hashes + self_merkle_path ~depth:t.depth ~hashes ~current_location:t.current_location address with | Some path -> @@ -412,7 +531,7 @@ module Make (Inputs : Inputs_intf.S) = struct let parent_merkle_path = Base.merkle_path_at_addr_exn ancestor address in - fixup_merkle_path ~hashes:maps.hashes parent_merkle_path ~address + fixup_merkle_path ~hashes parent_merkle_path ~address let merkle_path_at_index_exn t index = merkle_path_at_addr_exn t (Addr.of_int_exn ~ledger_depth:t.depth index) @@ -420,73 +539,31 @@ module Make (Inputs : Inputs_intf.S) = struct let merkle_path t location = merkle_path_at_addr_exn t (Location.to_path_exn location) - let path_batch_impl ~fixup_path ~self_lookup ~base_lookup t locations = + let merkle_path_batch t = assert_is_attached t ; - let maps, ancestor = maps_and_ancestor t in - let self_paths = - List.map locations ~f:(fun location -> - let address = Location.to_path_exn location in - self_lookup ~hashes:maps.hashes ~current_location:t.current_location - ~depth:t.depth address - |> Option.value_map - ~default:(Either.Second (location, address)) - ~f:Either.first ) - in - let all_parent_paths = - let locs = - List.filter_map self_paths ~f:(function - | Either.First _ -> - None - | Either.Second (location, _) -> - Some location ) - in - if List.is_empty locs then [] else base_lookup ancestor locs - in - let f parent_paths = function - | Either.First path -> - (parent_paths, path) - | Either.Second (_, address) -> - let path = - fixup_path ~hashes:maps.hashes ~address (List.hd_exn parent_paths) - in - (List.tl_exn parent_paths, path) - in - snd @@ List.fold_map ~init:all_parent_paths ~f self_paths - - let merkle_path_batch = - path_batch_impl ~base_lookup:Base.merkle_path_batch - ~self_lookup:self_merkle_path ~fixup_path:fixup_merkle_path - - let wide_merkle_path_batch = - path_batch_impl ~base_lookup:Base.wide_merkle_path_batch - ~self_lookup:self_wide_merkle_path ~fixup_path:fixup_wide_merkle_path - - (* 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 *) - let addresses_and_hashes_from_merkle_path_exn merkle_path starting_address - starting_hash : (Addr.t * Hash.t) list = - let get_addresses_hashes height accum node = - let last_address, last_hash = List.hd_exn accum in - let next_address = Addr.parent_exn last_address in - let next_hash = - match node with - | `Left sibling_hash -> - Hash.merge ~height last_hash sibling_hash - | `Right sibling_hash -> - Hash.merge ~height sibling_hash last_hash - in - (next_address, next_hash) :: accum - in - List.foldi merkle_path - ~init:[ (starting_address, starting_hash) ] - ~f:get_addresses_hashes + let hashes, ancestor = hashes_and_ancestor t in + path_batch_impl + ~base_lookup:(Base.merkle_path_batch ancestor) + ~self_lookup: + (self_merkle_path ~current_location:t.current_location ~depth:t.depth + ~hashes ) + ~fixup_path:(fixup_merkle_path ~hashes) + + let wide_merkle_path_batch t = + assert_is_attached t ; + let hashes, ancestor = hashes_and_ancestor t in + path_batch_impl + ~base_lookup:(Base.wide_merkle_path_batch ancestor) + ~self_lookup: + (self_wide_merkle_path ~current_location:t.current_location + ~depth:t.depth ~hashes ) + ~fixup_path:(fixup_wide_merkle_path ~hashes) (* use mask Merkle root, if it exists, else get from parent *) let merkle_root t = assert_is_attached t ; - let maps, ancestor = maps_and_ancestor t in - match Map.find maps.hashes (Addr.root ()) with + let hashes, ancestor = hashes_and_ancestor t in + match Map.find hashes (Addr.root ()) with | Some hash -> hash | None -> @@ -514,15 +591,8 @@ module Make (Inputs : Inputs_intf.S) = struct | None -> t.current_location <- None ) ; (* update hashes *) - let account_address = Location.to_path_exn location in let account_hash = Hash.empty_account in - let merkle_path = merkle_path t location in - let addresses_and_hashes = - addresses_and_hashes_from_merkle_path_exn merkle_path account_address - account_hash - in - List.iter addresses_and_hashes ~f:(fun (addr, hash) -> - self_set_hash t addr hash ) + t.unhashed_accounts <- (account_hash, location) :: t.unhashed_accounts let set_account_unsafe t location account = assert_is_attached t ; @@ -539,15 +609,8 @@ module Make (Inputs : Inputs_intf.S) = struct assert_is_attached t ; set_account_unsafe t location account ; (* Update merkle path. *) - let account_address = Location.to_path_exn location in let account_hash = Hash.hash_account account in - let merkle_path = merkle_path t location in - let addresses_and_hashes = - addresses_and_hashes_from_merkle_path_exn merkle_path account_address - account_hash - in - List.iter addresses_and_hashes ~f:(fun (addr, hash) -> - self_set_hash t addr hash ) + t.unhashed_accounts <- (account_hash, location) :: t.unhashed_accounts (* if the mask's parent sets an account, we can prune an entry in the mask if the account in the parent is the same in the mask *) @@ -568,8 +631,8 @@ module Make (Inputs : Inputs_intf.S) = struct parent *) let get_hash t addr = assert_is_attached t ; - let maps, ancestor = maps_and_ancestor t in - match Map.find maps.hashes addr with + let hashes, ancestor = hashes_and_ancestor t in + match Map.find hashes addr with | Some hash -> Some hash | None -> ( @@ -587,10 +650,10 @@ module Make (Inputs : Inputs_intf.S) = struct let get_hash_batch_exn t locations = assert_is_attached t ; - let maps, ancestor = maps_and_ancestor t in + let hashes, ancestor = hashes_and_ancestor t in let self_hashes_rev = List.rev_map locations ~f:(fun location -> - (location, Map.find maps.hashes (Location.to_path_exn location)) ) + (location, Map.find hashes (Location.to_path_exn location)) ) in let parent_locations_rev = List.filter_map self_hashes_rev ~f:(fun (location, hash) -> @@ -618,6 +681,7 @@ module Make (Inputs : Inputs_intf.S) = struct let parent = get_parent t in let old_root_hash = merkle_root t in let account_data = Map.to_alist t.maps.accounts in + finalize_hashes t ; Base.set_batch ~hash_cache:t.maps.hashes parent account_data ; t.maps.accounts <- Location_binable.Map.empty ; t.maps.hashes <- Addr.Map.empty ; @@ -648,6 +712,7 @@ module Make (Inputs : Inputs_intf.S) = struct ; next = maps_copy acc.next ; current = maps_copy acc.current } ) + ; unhashed_accounts = t.unhashed_accounts } let last_filled t =