Sparse ledger changes to support dynamic account creation

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,7 +198,7 @@ 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 idx = Option.value_exn (find_index t account_id) in
   let _, account = get_or_initialize_exn account_id t idx in
   Account.has_locked_tokens ~global_slot account
 
@@ -203,7 +224,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/sparse_ledger.ml

@@ -4,7 +4,12 @@ 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
@@ -35,6 +40,10 @@ let of_ledger_index_subset_exn (ledger : Ledger.Any_ledger.witness) 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))