Integrate mask accumulation into ledger lifecycle

src/lib/merkle_mask/maskable_merkle_tree.ml

@@ -157,67 +157,68 @@ module Make (Inputs : Inputs_intf) = struct
     Uuid.Table.add_multi registered_masks ~key:(get_uuid t) ~data:attached_mask ;
     attached_mask
 
-  let rec unregister_mask_exn ?(grandchildren = `Check) ~loc
-      (mask : Mask.Attached.t) : Mask.unattached =
+  let rec iter_descendants ~f uuid =
+    List.iter
+      (Hashtbl.find registered_masks uuid |> Option.value ~default:[])
+      ~f:(fun child_mask ->
+        if f child_mask then
+          iter_descendants ~f (Mask.Attached.get_uuid child_mask) )
+
+  let unregister_mask_error_msg ~uuid ~parent_uuid suffix =
+    sprintf "Couldn't unregister mask with UUID %s from parent %s, %s"
+      (Uuid.to_string_hum uuid)
+      (Uuid.to_string_hum parent_uuid)
+      suffix
+
+  let unregister_mask_exn_do ?trigger_signal mask =
+    let uuid = Mask.Attached.get_uuid mask in
     let parent_uuid = Mask.Attached.get_parent mask |> get_uuid in
-    let error_msg suffix =
-      sprintf "Couldn't unregister mask with UUID %s from parent %s, %s"
-        (Mask.Attached.get_uuid mask |> Uuid.to_string_hum)
-        (Uuid.to_string_hum parent_uuid)
-        suffix
-    in
-    let trigger_detach_signal =
-      match grandchildren with
-      | `Check | `Recursive ->
-          true
-      | `I_promise_I_am_reparenting_this_mask ->
-          false
-    in
-    ( match grandchildren with
-    | `Check -> (
-        match Hashtbl.find registered_masks (Mask.Attached.get_uuid mask) with
-        | Some children ->
-            failwith @@ error_msg
-            @@ sprintf
-                 !"mask has children that must be unregistered first: %{sexp: \
-                   Uuid.t list}"
-                 (List.map ~f:Mask.Attached.get_uuid children)
-        | None ->
-            () )
-    | `I_promise_I_am_reparenting_this_mask ->
-        ()
-    | `Recursive ->
-        (* You must not retain any references to children of the mask we're
-           unregistering if you pass `Recursive, so this is only used in
-           with_ephemeral_ledger. *)
-        List.iter
-          ( Hashtbl.find registered_masks (Mask.Attached.get_uuid mask)
-          |> Option.value ~default:[] )
-          ~f:(fun child_mask ->
-            ignore
-            @@ unregister_mask_exn ~loc ~grandchildren:`Recursive child_mask )
-    ) ;
+    let error_msg = unregister_mask_error_msg ~uuid ~parent_uuid in
     match Uuid.Table.find registered_masks parent_uuid with
     | None ->
         failwith @@ error_msg "parent not in registered_masks"
     | Some masks ->
-        ( match List.find masks ~f:(fun m -> phys_equal m mask) with
-        | None ->
-            failwith @@ error_msg "mask not registered with that parent"
-        | Some _ -> (
-            let bad, good =
-              List.partition_tf masks ~f:(fun m -> phys_equal m mask)
-            in
-            assert (List.length bad = 1) ;
-            match good with
-            | [] ->
-                (* no other masks for this maskable *)
-                Uuid.Table.remove registered_masks parent_uuid
-            | other_masks ->
-                Uuid.Table.set registered_masks ~key:parent_uuid
-                  ~data:other_masks ) ) ;
-        Mask.Attached.unset_parent ~trigger_signal:trigger_detach_signal ~loc
-          mask
+        let bad, good =
+          List.partition_tf masks ~f:(fun m -> phys_equal m mask)
+        in
+        if List.length bad <> 1 then
+          failwith @@ error_msg "mask not registered with that parent" ;
+        if List.is_empty good then
+          (* no other masks for this maskable *)
+          Uuid.Table.remove registered_masks parent_uuid
+        else Uuid.Table.set registered_masks ~key:parent_uuid ~data:good ;
+        Mask.Attached.unset_parent ?trigger_signal mask
+
+  let unregister_mask_exn ?(grandchildren = `Check) ~loc (mask : Mask.Attached.t)
+      : Mask.unattached =
+    let uuid = Mask.Attached.get_uuid mask in
+    let parent_uuid = Mask.Attached.get_parent mask |> get_uuid in
+    let error_msg = unregister_mask_error_msg ~uuid ~parent_uuid in
+    let trigger_signal =
+      match grandchildren with
+      | `Check -> (
+          match Hashtbl.find registered_masks (Mask.Attached.get_uuid mask) with
+          | Some children ->
+              failwith @@ error_msg
+              @@ sprintf
+                   !"mask has children that must be unregistered first: \
+                     %{sexp: Uuid.t list}"
+                   (List.map ~f:Mask.Attached.get_uuid children)
+          | None ->
+              true )
+      | `I_promise_I_am_reparenting_this_mask ->
+          false
+      | `Recursive ->
+          (* You must not retain any references to children of the mask we're
+             unregistering if you pass `Recursive, so this is only used in
+             with_ephemeral_ledger. *)
+          iter_descendants uuid
+            ~f:
+              ( Fn.compose (Fn.const true)
+              @@ unregister_mask_exn_do ~trigger_signal:true ~loc ) ;
+          true
+    in
+    unregister_mask_exn_do ~trigger_signal ~loc mask
 
   (** a set calls the Base implementation set, notifies registered mask childen *)
   let set t location account =

src/lib/merkle_mask/masking_merkle_tree.ml

@@ -49,7 +49,15 @@ module Make (Inputs : Inputs_intf.S) = struct
   let maps_copy { accounts; token_owners; hashes; locations } =
     { accounts; token_owners; hashes; locations }
 
-  type accumulated_t = { accumulated_maps : maps_t; ancestor : Base.t }
+  type accumulated_t =
+    { current : maps_t
+    ; next : maps_t
+    ; base : Base.t
+    ; detached_next_signal : Detached_parent_signal.t
+          (* Ivar for mask from which next was started being built
+             When it's fulfilled, "next" becomes "current".
+          *)
+    }
 
   type t =
     { uuid : Uuid.Stable.V1.t
@@ -67,19 +75,21 @@ module Make (Inputs : Inputs_intf.S) = struct
 
   type unattached = t [@@deriving sexp]
 
+  let empty_maps () =
+    { accounts = Location_binable.Map.empty
+    ; token_owners = Token_id.Map.empty
+    ; hashes = Addr.Map.empty
+    ; locations = Account_id.Map.empty
+    }
+
   let create ~depth () =
     { uuid = Uuid_unix.create ()
     ; parent = Error __LOC__
     ; detached_parent_signal = Async.Ivar.create ()
     ; current_location = None
     ; depth
     ; accumulated = None
-    ; maps =
-        { accounts = Location_binable.Map.empty
-        ; token_owners = Token_id.Map.empty
-        ; hashes = Addr.Map.empty
-        ; locations = Account_id.Map.empty
-        }
+    ; maps = empty_maps ()
     }
 
   let get_uuid { uuid; _ } = uuid
@@ -113,10 +123,18 @@ module Make (Inputs : Inputs_intf.S) = struct
 
     let to_accumulated t =
       match (t.accumulated, t.parent) with
-      | Some { accumulated_maps; ancestor }, _ ->
-          { ancestor; accumulated_maps = maps_copy accumulated_maps }
-      | None, Ok ancestor ->
-          { ancestor; accumulated_maps = maps_copy t.maps }
+      | Some { base; detached_next_signal; next; current }, _ ->
+          { base
+          ; detached_next_signal
+          ; next = maps_copy next
+          ; current = maps_copy current
+          }
+      | None, Ok base ->
+          { base
+          ; next = maps_copy t.maps
+          ; current = maps_copy t.maps
+          ; detached_next_signal = t.detached_parent_signal
+          }
       | None, Error loc ->
           raise (Dangling_parent_reference (t.uuid, loc))
 
@@ -133,9 +151,9 @@ module Make (Inputs : Inputs_intf.S) = struct
     let unset_parent ?(trigger_signal = true) ~loc t =
       assert (Result.is_ok t.parent) ;
       t.parent <- Error loc ;
-      t.accumulated <- None ;
-      if trigger_signal then
-        Async.Ivar.fill_if_empty t.detached_parent_signal () ;
+      if trigger_signal then (
+        t.accumulated <- None ;
+        Async.Ivar.fill_if_empty t.detached_parent_signal () ) ;
       t
 
     let assert_is_attached t =
@@ -151,9 +169,19 @@ module Make (Inputs : Inputs_intf.S) = struct
       assert_is_attached t ; Result.ok_or_failwith opt
 
     let maps_and_ancestor t =
+      Option.iter t.accumulated
+        ~f:(fun { detached_next_signal; next; base; current = _ } ->
+          if Async.Ivar.is_full detached_next_signal then
+            t.accumulated <-
+              Some
+                { next = empty_maps ()
+                ; current = next
+                ; detached_next_signal = t.detached_parent_signal
+                ; base
+                } ) ;
       match t.accumulated with
-      | Some { accumulated_maps; ancestor } ->
-          (accumulated_maps, ancestor)
+      | Some { current; base; _ } ->
+          (current, base)
       | None ->
           (t.maps, get_parent t)
 
@@ -167,7 +195,8 @@ module Make (Inputs : Inputs_intf.S) = struct
 
     let update_maps ~f t =
       f t.maps ;
-      Option.iter t.accumulated ~f:(fun { accumulated_maps = ms; _ } -> f ms)
+      Option.iter t.accumulated ~f:(fun { current; next; _ } ->
+          f current ; f next )
 
     let self_set_hash t address hash =
       update_maps t ~f:(fun maps ->
@@ -614,7 +643,10 @@ module Make (Inputs : Inputs_intf.S) = struct
       ; maps = maps_copy t.maps
       ; accumulated =
           Option.map t.accumulated ~f:(fun acc ->
-              { acc with accumulated_maps = maps_copy acc.accumulated_maps } )
+              { acc with
+                next = maps_copy acc.next
+              ; current = maps_copy acc.current
+              } )
       }
 
     let last_filled t =

src/lib/mina_ledger/ledger.ml

@@ -271,8 +271,9 @@ module Ledger_inner = struct
 
   let packed t = Any_ledger.cast (module Mask.Attached) t
 
-  let register_mask ?accumulated t mask =
-    Maskable.register_mask ?accumulated (packed t) mask
+  let register_mask t =
+    let accumulated = Mask.Attached.to_accumulated t in
+    Maskable.register_mask ~accumulated (packed t)
 
   let unsafe_preload_accounts_from_parent =
     Maskable.unsafe_preload_accounts_from_parent

src/lib/mina_ledger/ledger.mli

@@ -102,8 +102,7 @@ val of_database : Db.t -> t
 (** This is not _really_ copy, merely a stop-gap until we remove usages of copy in our codebase. What this actually does is creates a new empty mask on top of the current ledger *)
 val copy : t -> t
 
-val register_mask :
-  ?accumulated:Mask.accumulated_t -> t -> Mask.t -> Mask.Attached.t
+val register_mask : t -> Mask.t -> Mask.Attached.t
 
 val commit : Mask.Attached.t -> unit