diff --git a/src/lib/crypto/kimchi_bindings/stubs/Cargo.lock b/src/lib/crypto/kimchi_bindings/stubs/Cargo.lock index 14e0ca17966..5809faf03de 100644 --- a/src/lib/crypto/kimchi_bindings/stubs/Cargo.lock +++ b/src/lib/crypto/kimchi_bindings/stubs/Cargo.lock @@ -668,6 +668,7 @@ dependencies = [ "rand", "rand_core", "rayon", + "rmp-serde 1.1.1", "serde", "serde_with", "sha2", diff --git a/src/lib/crypto/kimchi_bindings/wasm/Cargo.lock b/src/lib/crypto/kimchi_bindings/wasm/Cargo.lock index bfafd70a5a5..e1030235b3d 100644 --- a/src/lib/crypto/kimchi_bindings/wasm/Cargo.lock +++ b/src/lib/crypto/kimchi_bindings/wasm/Cargo.lock @@ -774,6 +774,7 @@ dependencies = [ "rand", "rand_core", "rayon", + "rmp-serde", "serde", "serde_with 1.14.0", "sha2", diff --git a/src/lib/crypto/proof-systems b/src/lib/crypto/proof-systems index 553795286d4..23e0787a701 160000 --- a/src/lib/crypto/proof-systems +++ b/src/lib/crypto/proof-systems @@ -1 +1 @@ -Subproject commit 553795286d4561aa5d7e928ed1e3555e3a4a81be +Subproject commit 23e0787a701e6c6c5e45ff3d567bb2c80b6f5e22 diff --git a/src/lib/merkle_mask/masking_merkle_tree.ml b/src/lib/merkle_mask/masking_merkle_tree.ml index d696870a791..a0b64ad1dcd 100644 --- a/src/lib/merkle_mask/masking_merkle_tree.ml +++ b/src/lib/merkle_mask/masking_merkle_tree.ml @@ -81,6 +81,18 @@ module Make (Inputs : Inputs_intf.S) = struct *) } + (** Type for an account that wasn't yet hashed. + + Manual definitions of sexp are needed to avoid warning 4 in autogenerated code + *) + type unhashed_account_t = Account.t option * Location.t + + let sexp_of_unhashed_account_t = + Tuple2.sexp_of_t (Option.sexp_of_t Account.sexp_of_t) Location.sexp_of_t + + let unhashed_account_t_of_sexp = + Tuple2.t_of_sexp (Option.t_of_sexp Account.t_of_sexp) Location.t_of_sexp + type t = { uuid : Uuid.Stable.V1.t ; mutable parent : Parent.t @@ -93,6 +105,7 @@ module Make (Inputs : Inputs_intf.S) = struct This is used as a lookup cache. *) ; mutable accumulated : (accumulated_t[@sexp.opaque]) option ; mutable is_committing : bool + ; mutable unhashed_accounts : unhashed_account_t list } [@@deriving sexp] @@ -115,6 +128,7 @@ module Make (Inputs : Inputs_intf.S) = struct ; accumulated = None ; maps = empty_maps ; is_committing = false + ; unhashed_accounts = [] } let get_uuid { uuid; _ } = uuid @@ -181,21 +195,6 @@ module Make (Inputs : Inputs_intf.S) = struct | None -> (t.maps, get_parent t) - (** Either copies accumulated or initializes it with the parent being used as the [base]. *) - let to_accumulated t = - actualize_accumulated t ; - match (t.accumulated, t.parent) with - | Some { base; detached_next_signal; next; current }, _ -> - { base; detached_next_signal; next; current } - | None, Ok base -> - { base - ; next = t.maps - ; current = t.maps - ; detached_next_signal = t.detached_parent_signal - } - | None, Error loc -> - raise (Dangling_parent_reference (t.uuid, loc)) - let get_uuid t = assert_is_attached t ; t.uuid let get_directory t = @@ -210,15 +209,206 @@ module Make (Inputs : Inputs_intf.S) = struct acc.current <- f acc.current ; acc.next <- f acc.next ) - let self_set_hash t address hash = + let self_set_hash_impl t address hash = update_maps t ~f:(fun maps -> { maps with 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 : + ( Account.t option + * 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 + let hash_account = + Option.value_map ~default:Hash.empty_account ~f:Hash.hash_account + in + Fn.compose (impl [] 0) (List.map ~f:(Tuple3.map_fst ~f:hash_account)) + + let finalize_hashes_do t unhashed_accounts = + let with_merkle_path_batch accs = + 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 accs) + |> List.map2_exn accs ~f:(fun (a, loc) p -> + (a, Location.to_path_exn loc, p) ) + in + let on_snd f (_, a) (_, b) = f a b in + List.stable_sort ~compare:(on_snd Location.compare) unhashed_accounts + |> List.remove_consecutive_duplicates ~which_to_keep:`First + ~equal:(on_snd Location.equal) + |> with_merkle_path_batch |> 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 ) + + (** Either copies accumulated or initializes it with the parent being used as the [base]. *) + let to_accumulated t = + finalize_hashes t ; + actualize_accumulated t ; + match (t.accumulated, t.parent) with + | Some { base; detached_next_signal; next; current }, _ -> + { base; detached_next_signal; next; current } + | None, Ok base -> + { base + ; next = t.maps + ; current = t.maps + ; detached_next_signal = t.detached_parent_signal + } + | None, Error loc -> + raise (Dangling_parent_reference (t.uuid, loc)) + + 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 with @@ -254,8 +444,8 @@ module Make (Inputs : Inputs_intf.S) = struct parent *) let get t location = assert_is_attached t ; - let maps, ancestor = maps_and_ancestor t in - match Map.find maps.accounts location with + let { accounts; _ }, ancestor = maps_and_ancestor t in + match Map.find accounts location with | Some account -> Some account | None -> @@ -314,44 +504,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 @@ -368,22 +520,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 = @@ -412,9 +548,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 -> @@ -423,7 +559,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) @@ -431,73 +567,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 -> @@ -527,16 +621,7 @@ module Make (Inputs : Inputs_intf.S) = struct t.current_location <- Some prev_loc | 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 <- (None, location) :: t.unhashed_accounts let set_account_unsafe t location account = assert_is_attached t ; @@ -553,16 +638,7 @@ module Make (Inputs : Inputs_intf.S) = struct let set t location account = 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 <- (Some account, 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 @@ -586,8 +662,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 -> ( @@ -598,10 +674,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) -> @@ -631,6 +707,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 ; let hash_cache = t.maps.hashes in t.maps <- empty_maps ; Base.set_batch ~hash_cache parent account_data ; @@ -661,6 +738,7 @@ module Make (Inputs : Inputs_intf.S) = struct ; current = acc.current } ) ; is_committing = false + ; unhashed_accounts = t.unhashed_accounts } let last_filled t = @@ -748,8 +826,8 @@ module Make (Inputs : Inputs_intf.S) = struct let token_owner t tid = assert_is_attached t ; - let maps, ancestor = maps_and_ancestor t in - match Map.find maps.token_owners tid with + let { token_owners; _ }, ancestor = maps_and_ancestor t in + match Map.find token_owners tid with | Some id -> Some id | None -> @@ -757,18 +835,18 @@ module Make (Inputs : Inputs_intf.S) = struct let token_owners (t : t) : Account_id.Set.t = assert_is_attached t ; - let maps, ancestor = maps_and_ancestor t in + let { token_owners; _ }, ancestor = maps_and_ancestor t in let mask_owners = - Map.fold maps.token_owners ~init:Account_id.Set.empty + Map.fold token_owners ~init:Account_id.Set.empty ~f:(fun ~key:_tid ~data:owner acc -> Set.add acc owner) in Set.union mask_owners (Base.token_owners ancestor) let tokens t pk = assert_is_attached t ; - let maps, ancestor = maps_and_ancestor t in + let { locations; _ }, ancestor = maps_and_ancestor t in let mask_tokens = - Map.keys maps.locations + Map.keys locations |> List.filter_map ~f:(fun aid -> if Key.equal pk (Account_id.public_key aid) then Some (Account_id.token_id aid) @@ -926,8 +1004,8 @@ module Make (Inputs : Inputs_intf.S) = struct let foldi_with_ignored_accounts t ignored_accounts ~init ~f = assert_is_attached t ; - let maps, ancestor = maps_and_ancestor t in - let locations_and_accounts = Map.to_alist maps.accounts in + let { accounts; _ }, ancestor = maps_and_ancestor t in + let locations_and_accounts = Map.to_alist accounts in (* parent should ignore accounts in this mask *) let mask_accounts = List.map locations_and_accounts ~f:(fun (_loc, acct) -> @@ -975,6 +1053,7 @@ module Make (Inputs : Inputs_intf.S) = struct let address_in_mask t addr = assert_is_attached t ; + finalize_hashes t ; Option.is_some (Map.find t.maps.hashes addr) let current_location t = t.current_location @@ -1031,6 +1110,7 @@ module Make (Inputs : Inputs_intf.S) = struct assert (Int.equal t.depth (Base.depth parent)) ; t.parent <- Ok parent ; t.current_location <- Attached.last_filled t ; + Attached.finalize_hashes t ; (* If [t.accumulated] isn't empty, then this mask had a parent before and now we just reparent it (which may only happen if both old and new parents have the same merkle root (and some masks in between may have been removed),