Skip to content

Commit

Permalink
Use map refs instead of mutable map fields
Browse files Browse the repository at this point in the history
  • Loading branch information
@nholland94
nholland94 committed Nov 22, 2023
1 parent 97e6b95 commit efdf699
Showing 1 changed file with 39 additions and 37 deletions.
76 changes: 39 additions & 37 deletions src/lib/merkle_mask/masking_merkle_tree.ml
View file
Original file line number Diff line number Diff line change
Expand Up @@ -40,12 +40,12 @@ module Make (Inputs : Inputs_intf.S) = struct

type t =
{ uuid : Uuid.Stable.V1.t
; mutable accounts : Account.t Location_binable.Map.t
; mutable token_owners : Account_id.t Token_id.Map.t
; accounts : Account.t Location_binable.Map.t ref
; token_owners : Account_id.t Token_id.Map.t ref
; mutable parent : Parent.t
; detached_parent_signal : Detached_parent_signal.t
; mutable hashes : Hash.t Addr.Map.t
; mutable locations : Location.t Account_id.Map.t
; hashes : Hash.t Addr.Map.t ref
; locations : Location.t Account_id.Map.t ref
; mutable current_location : Location.t option
; depth : int
}
Expand All @@ -57,10 +57,10 @@ module Make (Inputs : Inputs_intf.S) = struct
{ uuid = Uuid_unix.create ()
; parent = Error __LOC__
; detached_parent_signal = Async.Ivar.create ()
; accounts = Location_binable.Map.empty
; token_owners = Token_id.Map.empty
; hashes = Addr.Map.empty
; locations = Account_id.Map.empty
; accounts = ref Location_binable.Map.empty
; token_owners = ref Token_id.Map.empty
; hashes = ref Addr.Map.empty
; locations = ref Account_id.Map.empty
; current_location = None
; depth
}
Expand Down Expand Up @@ -133,11 +133,12 @@ module Make (Inputs : Inputs_intf.S) = struct

(* don't rely on a particular implementation *)
let self_find_hash t address =
assert_is_attached t ; Map.find t.hashes address
assert_is_attached t ;
Map.find !(t.hashes) address

let self_set_hash t address hash =
assert_is_attached t ;
t.hashes <- Map.set t.hashes ~key:address ~data:hash
t.hashes := Map.set !(t.hashes) ~key:address ~data:hash

let set_inner_hash_at_addr_exn t address hash =
assert_is_attached t ;
Expand All @@ -147,11 +148,11 @@ module Make (Inputs : Inputs_intf.S) = struct
(* don't rely on a particular implementation *)
let self_find_location t account_id =
assert_is_attached t ;
Map.find t.locations account_id
Map.find !(t.locations) account_id

let self_set_location t account_id location =
assert_is_attached t ;
t.locations <- Map.set t.locations ~key:account_id ~data:location ;
t.locations := Map.set !(t.locations) ~key:account_id ~data:location ;
(* if account is at a hitherto-unused location, that
becomes the current location
*)
Expand All @@ -165,11 +166,11 @@ module Make (Inputs : Inputs_intf.S) = struct
(* don't rely on a particular implementation *)
let self_find_account t location =
assert_is_attached t ;
Map.find t.accounts location
Map.find !(t.accounts) location

let self_set_account t location account =
assert_is_attached t ;
t.accounts <- Map.set t.accounts ~key:location ~data:account ;
t.accounts := Map.set !(t.accounts) ~key:location ~data:account ;
self_set_location t (Account.identifier account) location

(* a read does a lookup in the account_tbl; if that fails, delegate to
Expand Down Expand Up @@ -515,15 +516,15 @@ module Make (Inputs : Inputs_intf.S) = struct
assert_is_attached t ;
(* remove account and key from tables *)
let account = Option.value_exn (self_find_account t location) in
t.accounts <- Map.remove t.accounts location ;
t.accounts := Map.remove !(t.accounts) location ;
(* Update token info. *)
let account_id = Account.identifier account in
t.token_owners <-
Token_id.Map.remove t.token_owners
t.token_owners :=
Token_id.Map.remove !(t.token_owners)
(Account_id.derive_token_id ~owner:account_id) ;
(* TODO : use stack database to save unused location, which can be used
when allocating a location *)
t.locations <- Map.remove t.locations account_id ;
t.locations := Map.remove !(t.locations) account_id ;
(* reuse location if possible *)
Option.iter t.current_location ~f:(fun curr_loc ->
if Location.equal location curr_loc then
Expand All @@ -548,8 +549,8 @@ module Make (Inputs : Inputs_intf.S) = struct
self_set_account t location account ;
(* Update token info. *)
let account_id = Account.identifier account in
t.token_owners <-
Map.set t.token_owners
t.token_owners :=
Map.set !(t.token_owners)
~key:(Account_id.derive_token_id ~owner:account_id)
~data:account_id

Expand Down Expand Up @@ -645,10 +646,10 @@ module Make (Inputs : Inputs_intf.S) = struct
let commit t =
assert_is_attached t ;
let old_root_hash = merkle_root t in
let account_data = Map.to_alist t.accounts in
let account_data = Map.to_alist !(t.accounts) in
Base.set_batch (get_parent t) account_data ;
t.accounts <- Location_binable.Map.empty ;
t.hashes <- Addr.Map.empty ;
t.accounts := Location_binable.Map.empty ;
t.hashes := Addr.Map.empty ;
Debug_assert.debug_assert (fun () ->
[%test_result: Hash.t]
~message:
Expand All @@ -666,10 +667,10 @@ module Make (Inputs : Inputs_intf.S) = struct
{ uuid = Uuid_unix.create ()
; parent = Ok (get_parent t)
; detached_parent_signal = Async.Ivar.create ()
; accounts = t.accounts
; token_owners = t.token_owners
; locations = t.locations
; hashes = t.hashes
; accounts = ref !(t.accounts)
; token_owners = ref !(t.token_owners)
; locations = ref !(t.locations)
; hashes = ref !(t.hashes)
; current_location = t.current_location
; depth = t.depth
}
Expand Down Expand Up @@ -731,13 +732,14 @@ module Make (Inputs : Inputs_intf.S) = struct
let set_location_batch ~last_location t account_to_location_list =
t.current_location <- Some last_location ;
Mina_stdlib.Nonempty_list.iter account_to_location_list
~f:(fun (key, data) -> t.locations <- Map.set t.locations ~key ~data)
~f:(fun (key, data) ->
t.locations := Map.set !(t.locations) ~key ~data )

let set_raw_account_batch t locations_and_accounts =
List.iter locations_and_accounts ~f:(fun (location, account) ->
let account_id = Account.identifier account in
t.token_owners <-
Map.set t.token_owners
t.token_owners :=
Map.set !(t.token_owners)
~key:(Account_id.derive_token_id ~owner:account_id)
~data:account_id ;
self_set_account t location account )
Expand All @@ -754,7 +756,7 @@ module Make (Inputs : Inputs_intf.S) = struct

let token_owner t tid =
assert_is_attached t ;
match Map.find t.token_owners tid with
match Map.find !(t.token_owners) tid with
| Some id ->
Some id
| None ->
Expand All @@ -763,15 +765,15 @@ module Make (Inputs : Inputs_intf.S) = struct
let token_owners (t : t) : Account_id.Set.t =
assert_is_attached t ;
let mask_owners =
Map.fold t.token_owners ~init:Account_id.Set.empty
Map.fold !(t.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 (get_parent t))

let tokens t pk =
assert_is_attached t ;
let mask_tokens =
Map.keys t.locations
Map.keys !(t.locations)
|> List.filter_map ~f:(fun aid ->
if Key.equal pk (Account_id.public_key aid) then
Some (Account_id.token_id aid)
Expand Down Expand Up @@ -935,9 +937,9 @@ module Make (Inputs : Inputs_intf.S) = struct
as sometimes this is desired behavior *)
let close t =
assert_is_attached t ;
t.accounts <- Location_binable.Map.empty ;
t.hashes <- Addr.Map.empty ;
t.locations <- Account_id.Map.empty ;
t.accounts := Location_binable.Map.empty ;
t.hashes := Addr.Map.empty ;
t.locations := Account_id.Map.empty ;
Async.Ivar.fill_if_empty t.detached_parent_signal ()

let index_of_account_exn t key =
Expand Down Expand Up @@ -981,7 +983,7 @@ module Make (Inputs : Inputs_intf.S) = struct

let foldi_with_ignored_accounts t ignored_accounts ~init ~f =
assert_is_attached t ;
let locations_and_accounts = Map.to_alist t.accounts in
let locations_and_accounts = Map.to_alist !(t.accounts) in
(* parent should ignore accounts in this mask *)
let mask_accounts =
List.map locations_and_accounts ~f:(fun (_loc, acct) ->
Expand Down

0 comments on commit efdf699

Please sign in to comment.