Perform account hashing in finalize_hashes call

MinaProtocol · Aug 25, 2024 · 433f6b0 · 433f6b0
1 parent 6a3a91f
commit 433f6b0
Showing 1 changed file with 12 additions and 10 deletions.
diff --git a/src/lib/merkle_mask/masking_merkle_tree.ml b/src/lib/merkle_mask/masking_merkle_tree.ml
@@ -93,7 +93,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 : (Hash.t * Location.t) list
+    ; mutable unhashed_accounts : (Account.t option * Location.t) list
     }
   [@@deriving sexp]
 
@@ -284,7 +284,10 @@ module Make (Inputs : Inputs_intf.S) = struct
       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
+           ( 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) ->
@@ -333,7 +336,10 @@ module Make (Inputs : Inputs_intf.S) = struct
         | _ ->
             impl acc' (height + 1) (converge height task)
       in
-      impl [] 0
+      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 merkle_path_batch =
@@ -348,7 +354,7 @@ module Make (Inputs : Inputs_intf.S) = struct
       in
       (* let _task = *)
       List.map2_exn
-        ~f:(fun (h, loc) p -> (h, Location.to_path_exn loc, p))
+        ~f:(fun (a, loc) p -> (a, Location.to_path_exn loc, p))
         unhashed_accounts merkle_path_batch
       |> List.stable_sort ~compare:(fun (_, a, _) (_, b, _) ->
              Addr.compare a b )
@@ -608,9 +614,7 @@ module Make (Inputs : Inputs_intf.S) = struct
                 t.current_location <- Some prev_loc
             | None ->
                 t.current_location <- None ) ;
-      (* update hashes *)
-      let account_hash = Hash.empty_account in
-      t.unhashed_accounts <- (account_hash, location) :: t.unhashed_accounts
+      t.unhashed_accounts <- (None, location) :: t.unhashed_accounts
 
     let set_account_unsafe t location account =
       assert_is_attached t ;
@@ -626,9 +630,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_hash = Hash.hash_account account in
-      t.unhashed_accounts <- (account_hash, location) :: t.unhashed_accounts
+      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