Use sparse ledger for staged ledger diff application

src/lib/mina_base/pending_coinbase.ml

@@ -766,17 +766,16 @@ module Make_str (A : Wire_types.Concrete) = struct
             Type_equal.t ) =
         Type_equal.T
 
-      module M = Sparse_ledger_lib.Sparse_ledger.Make (Hash) (Stack_id) (Stack)
+      module M =
+        Sparse_ledger_lib.Sparse_ledger.Make (Hash) (Stack_id)
+          (struct
+            include Stack
+
+            let empty = lazy empty
+          end)
 
       [%%define_locally
-      M.
-        ( of_hash
-        , get_exn
-        , path_exn
-        , set_exn
-        , find_index_exn
-        , add_path
-        , merkle_root )]
+      M.(of_hash, get_exn, path_exn, set_exn, find_index, add_path, merkle_root)]
     end
 
     module Checked = struct
@@ -1055,7 +1054,10 @@ module Make_str (A : Wire_types.Concrete) = struct
       in
       let root_hash = hash_at_level depth in
       { Poly.tree =
-          make_tree (Merkle_tree.of_hash ~depth root_hash) Stack_id.zero
+          make_tree
+            (Merkle_tree.of_hash ~depth root_hash
+               ~current_location:(* Hack: unused*) None )
+            Stack_id.zero
       ; pos_list = []
       ; new_pos = Stack_id.zero
       }
@@ -1073,7 +1075,7 @@ module Make_str (A : Wire_types.Concrete) = struct
       try_with (fun () -> Merkle_tree.path_exn t.tree index)
 
     let find_index (t : t) key =
-      try_with (fun () -> Merkle_tree.find_index_exn t.tree key)
+      try_with (fun () -> Option.value_exn @@ Merkle_tree.find_index t.tree key)
 
     let next_index ~depth (t : t) =
       if
@@ -1110,15 +1112,17 @@ module Make_str (A : Wire_types.Concrete) = struct
         Option.value ~default:Stack_id.zero (curr_stack_id t)
       in
       Or_error.try_with (fun () ->
-          let index = Merkle_tree.find_index_exn t.tree prev_stack_id in
+          let index =
+            Option.value_exn @@ Merkle_tree.find_index t.tree prev_stack_id
+          in
           Merkle_tree.get_exn t.tree index )
 
     let latest_stack (t : t) ~is_new_stack =
       let open Or_error.Let_syntax in
       let key = latest_stack_id t ~is_new_stack in
       let%bind res =
         Or_error.try_with (fun () ->
-            let index = Merkle_tree.find_index_exn t.tree key in
+            let index = Option.value_exn @@ Merkle_tree.find_index t.tree key in
             Merkle_tree.get_exn t.tree index )
       in
       if is_new_stack then

src/lib/mina_base/sparse_ledger_base.ml

@@ -28,6 +28,8 @@ module Account = struct
   include Account
 
   let data_hash = Fn.compose Ledger_hash.of_digest Account.digest
+
+  let empty = lazy empty
 end
 
 module Global_state = struct
@@ -50,7 +52,8 @@ type account_state = [ `Added | `Existed ] [@@deriving equal]
     This ledger has an invalid root hash, and cannot be used except as a
     placeholder.
 *)
-let empty ~depth () = M.of_hash ~depth Outside_hash_image.t
+let empty ~depth () =
+  M.of_hash ~depth ~current_location:None Outside_hash_image.t
 
 module L = struct
   type t = M.t ref
@@ -67,22 +70,21 @@ module L = struct
 
   let location_of_account : t -> Account_id.t -> location option =
    fun t id ->
-    try
-      let loc = M.find_index_exn !t id in
-      let account = M.get_exn !t loc in
-      if Public_key.Compressed.(equal empty account.public_key) then None
-      else Some loc
-    with _ -> None
+    match M.find_index !t id with
+    | None ->
+        None
+    | Some loc ->
+        let account = M.get_exn !t loc in
+        if Public_key.Compressed.(equal empty account.public_key) then None
+        else Some loc
 
   let set : t -> location -> Account.t -> unit =
    fun t loc a -> t := M.set_exn !t loc a
 
   let get_or_create_exn :
       t -> Account_id.t -> account_state * Account.t * location =
    fun t id ->
-    let loc = M.find_index_exn !t id in
-    let account = M.get_exn !t loc in
-    if Public_key.Compressed.(equal empty account.public_key) then (
+    let create_account loc account =
       let public_key = Account_id.public_key id in
       let account' : Account.t =
         { account with
@@ -91,22 +93,39 @@ module L = struct
         ; token_id = Account_id.token_id id
         }
       in
-      set t loc account' ;
-      (`Added, account', loc) )
-    else (`Existed, account, loc)
+      set t loc account' ; account'
+    in
+    match M.find_index !t id with
+    | None ->
+        let loc, new_t = M.allocate_index !t id in
+        t := new_t ;
+        let account' = create_account loc (Lazy.force Account.empty) in
+        (`Added, account', loc)
+    | Some loc ->
+        let account = M.get_exn !t loc in
+        if Public_key.Compressed.(equal empty account.public_key) then
+          let account' = create_account loc account in
+          (`Added, account', loc)
+        else (`Existed, account, loc)
 
   let get_or_create t id = Or_error.try_with (fun () -> get_or_create_exn t id)
 
   let get_or_create_account :
       t -> Account_id.t -> Account.t -> (account_state * location) Or_error.t =
    fun t id to_set ->
     Or_error.try_with (fun () ->
-        let loc = M.find_index_exn !t id in
-        let a = M.get_exn !t loc in
-        if Public_key.Compressed.(equal empty a.public_key) then (
-          set t loc to_set ;
-          (`Added, loc) )
-        else (`Existed, loc) )
+        match M.find_index !t id with
+        | None ->
+            let loc, new_ledger = M.allocate_index !t id in
+            t := new_ledger ;
+            set t loc to_set ;
+            (`Added, loc)
+        | Some loc ->
+            let a = M.get_exn !t loc in
+            if Public_key.Compressed.(equal empty a.public_key) then (
+              set t loc to_set ;
+              (`Added, loc) )
+            else (`Existed, loc) )
 
   let create_new_account t id to_set =
     get_or_create_account t id to_set |> Or_error.map ~f:ignore
@@ -151,13 +170,15 @@ M.
   , get_exn
   , path_exn
   , set_exn
-  , find_index_exn
+  , allocate_index
+  , find_index
   , add_path
   , merkle_root
   , iteri )]
 
-let of_root ~depth (h : Ledger_hash.t) =
-  of_hash ~depth (Ledger_hash.of_digest (h :> Random_oracle.Digest.t))
+let of_root ~depth ~current_location (h : Ledger_hash.t) =
+  of_hash ~depth ~current_location
+    (Ledger_hash.of_digest (h :> Random_oracle.Digest.t))
 
 let get_or_initialize_exn account_id t idx =
   let account = get_exn t idx in
@@ -177,9 +198,12 @@ let get_or_initialize_exn account_id t idx =
   else (`Existed, account)
 
 let has_locked_tokens_exn ~global_slot ~account_id t =
-  let idx = find_index_exn t account_id in
-  let _, account = get_or_initialize_exn account_id t idx in
-  Account.has_locked_tokens ~global_slot account
+  match find_index t account_id with
+  | Some idx ->
+      let _, account = get_or_initialize_exn account_id t idx in
+      Account.has_locked_tokens ~global_slot account
+  | None ->
+      false
 
 let merkle_root t = Ledger_hash.of_hash (merkle_root t :> Random_oracle.Digest.t)
 
@@ -203,7 +227,18 @@ let handler t =
           ledger := set_exn !ledger idx account ;
           respond (Provide ())
       | Ledger_hash.Find_index pk ->
-          let index = find_index_exn !ledger pk in
+          let index =
+            match find_index !ledger pk with
+            | Some index ->
+                index
+            | None ->
+                let index, new_ledger = allocate_index !ledger pk in
+                let new_ledger =
+                  set_exn new_ledger index (Lazy.force Account.empty)
+                in
+                ledger := new_ledger ;
+                index
+          in
           respond (Provide index)
       | _ ->
           unhandled )

src/lib/mina_base/sparse_ledger_base.mli

@@ -42,9 +42,11 @@ val set_exn : t -> int -> Account.t -> t
 val path_exn :
   t -> int -> [ `Left of Ledger_hash.t | `Right of Ledger_hash.t ] list
 
-val find_index_exn : t -> Account_id.t -> int
+val allocate_index : t -> Account_id.t -> int * t
 
-val of_root : depth:int -> Ledger_hash.t -> t
+val find_index : t -> Account_id.t -> int option
+
+val of_root : depth:int -> current_location:int option -> Ledger_hash.t -> t
 
 (** Create a new 'empty' ledger.
     This ledger has an invalid root hash, and cannot be used except as a

src/lib/mina_ledger/dune

@@ -53,4 +53,5 @@
    unsigned_extended
    with_hash
    ppx_version.runtime
+   sparse_ledger_lib
 ))

src/lib/mina_ledger/sparse_ledger.ml

@@ -4,50 +4,105 @@ include Sparse_ledger_base
 module GS = Global_state
 
 let of_ledger_root ledger =
-  of_root ~depth:(Ledger.depth ledger) (Ledger.merkle_root ledger)
+  of_root ~depth:(Ledger.depth ledger)
+    ~current_location:
+      ( Option.map ~f:(fun x ->
+            Ledger.Location.Addr.to_int @@ Ledger.Location.to_path_exn x )
+      @@ Ledger.last_filled ledger )
+    (Ledger.merkle_root ledger)
 
 let of_ledger_subset_exn (oledger : Ledger.t) keys =
   let ledger = Ledger.copy oledger in
   let locations = Ledger.location_of_account_batch ledger keys in
   let non_empty_locations = List.filter_map ~f:snd locations in
   let accounts = Ledger.get_batch ledger non_empty_locations in
   let merkle_paths = Ledger.merkle_path_batch ledger non_empty_locations in
-  let _, sparse =
-    let rec go (new_keys, sl) locations accounts merkle_paths =
+  let uses_last, sparse =
+    let rec go (uses_last, sl) locations accounts merkle_paths =
       match locations with
       | [] ->
-          (new_keys, sl)
+          (uses_last, sl)
       | (key, Some _loc) :: locations -> (
           match (accounts, merkle_paths) with
           | (_, account) :: accounts, merkle_path :: merkle_paths ->
               go
-                ( new_keys
+                ( uses_last
                 , add_path sl merkle_path key
                     ( account
                     |> Option.value_exn ?here:None ?error:None ?message:None )
                 )
                 locations accounts merkle_paths
           | _ ->
               assert false )
-      | (key, None) :: locations ->
-          let path, acct = Ledger.create_empty_exn ledger key in
-          go
-            (key :: new_keys, add_path sl path key acct)
-            locations accounts merkle_paths
+      | (_, None) :: locations ->
+          go (true, sl) locations accounts merkle_paths
     in
-    go ([], of_ledger_root ledger) locations accounts merkle_paths
+    go (false, of_ledger_root ledger) locations accounts merkle_paths
+  in
+  (* TODO: With some care over the external contract this satisfies, this could
+     be batched with the account and path reads above.
+  *)
+  let sparse =
+    if uses_last then
+      match Ledger.last_filled ledger with
+      | Some loc ->
+          let account =
+            Ledger.get ledger loc
+            |> Option.value_exn ?here:None ?error:None ?message:None
+          in
+          add_path sparse
+            (Ledger.merkle_path ledger loc)
+            (Account.identifier account)
+            account
+      | None ->
+          sparse
+    else sparse
   in
   Debug_assert.debug_assert (fun () ->
       [%test_eq: Ledger_hash.t]
         (Ledger.merkle_root ledger)
         ((merkle_root sparse :> Random_oracle.Digest.t) |> Ledger_hash.of_hash) ) ;
   sparse
 
+let sparse_ledger_subset_exn (oledger : t) keys =
+  let ledger =
+    of_root ~depth:(depth oledger) ~current_location:oledger.current_location
+      (merkle_root oledger)
+  in
+  let uses_last, ledger =
+    List.fold ~init:(false, ledger) keys ~f:(fun (uses_last, ledger) key ->
+        match find_index oledger key with
+        | Some index ->
+            let path = path_exn oledger index in
+            let account = get_exn oledger index in
+            (uses_last, add_path ledger path key account)
+        | None ->
+            (true, ledger) )
+  in
+  let ledger =
+    if uses_last then
+      match ledger.current_location with
+      | Some loc ->
+          let account = Mina_base.Sparse_ledger_base.get_exn oledger loc in
+          add_path ledger
+            (Mina_base.Sparse_ledger_base.path_exn oledger loc)
+            (Account.identifier account)
+            account
+      | None ->
+          ledger
+    else ledger
+  in
+  ledger
+
 let of_ledger_index_subset_exn (ledger : Ledger.Any_ledger.witness) indexes =
   List.fold indexes
     ~init:
       (of_root
          ~depth:(Ledger.Any_ledger.M.depth ledger)
+         ~current_location:
+           ( Option.map ~f:(fun x ->
+                 Ledger.Location.Addr.to_int @@ Ledger.Location.to_path_exn x )
+           @@ Ledger.Any_ledger.M.last_filled ledger )
          (Ledger.Any_ledger.M.merkle_root ledger) )
     ~f:(fun acc i ->
       let account = Ledger.Any_ledger.M.get_at_index_exn ledger i in

src/lib/sparse_ledger_lib/dune

@@ -13,6 +13,8 @@
    bin_prot.shape
    result
    ppx_version.runtime
+   ;; local libraries
+   empty_hashes
  )
  (preprocess
   (pps ppx_jane ppx_compare ppx_deriving_yojson ppx_version))