diff --git a/src/proto_alpha/lib_protocol/alpha_context.ml b/src/proto_alpha/lib_protocol/alpha_context.ml index 2896f64cd2bc8a5685212599b634ccdc2d54d8c4..df9752cc6165e0f20bda25c41c4b7a84fce4c83f 100644 --- a/src/proto_alpha/lib_protocol/alpha_context.ml +++ b/src/proto_alpha/lib_protocol/alpha_context.ml @@ -629,8 +629,12 @@ module Stake_distribution = struct let slot_owner = Delegate_sampler.slot_owner + let stake_info = Delegate_sampler.stake_info + let load_sampler_for_cycle = Delegate_sampler.load_sampler_for_cycle + let load_stake_info_for_cycle = Delegate_sampler.load_stake_info_for_cycle + let get_total_frozen_stake ctxt cycle = let open Lwt_result_syntax in let* total_stake = Stake_storage.get_total_active_stake ctxt cycle in diff --git a/src/proto_alpha/lib_protocol/alpha_context.mli b/src/proto_alpha/lib_protocol/alpha_context.mli index 4956a80d3f1e79316a235ec7de604dd8548798d0..585fcce1d9f3a13b3708e5d4f0d922d0cb75f8eb 100644 --- a/src/proto_alpha/lib_protocol/alpha_context.mli +++ b/src/proto_alpha/lib_protocol/alpha_context.mli @@ -5196,9 +5196,17 @@ module Stake_distribution : sig val slot_owner : context -> Level.t -> Slot.t -> (context * Consensus_key.pk) tzresult Lwt.t + val stake_info : + context -> + Level.t -> + (context * Int64.t * (Consensus_key.pk * Int64.t) list) tzresult Lwt.t + (** See {!Delegate_sampler.load_sampler_for_cycle}. *) val load_sampler_for_cycle : context -> Cycle.t -> context tzresult Lwt.t + (** See {!Delegate_sampler.load_stake_info_for_cycle}. *) + val load_stake_info_for_cycle : context -> Cycle.t -> context tzresult Lwt.t + val get_total_frozen_stake : context -> Cycle.t -> Tez.t tzresult Lwt.t module For_RPC : sig diff --git a/src/proto_alpha/lib_protocol/delegate_sampler.ml b/src/proto_alpha/lib_protocol/delegate_sampler.ml index b73a69e6768bc8c8e71fe359a37da5bfb4b2741e..b271dec004aaac09e077a6dc648ab9e95db96e15 100644 --- a/src/proto_alpha/lib_protocol/delegate_sampler.ml +++ b/src/proto_alpha/lib_protocol/delegate_sampler.ml @@ -164,6 +164,40 @@ let load_sampler_for_cycle ctxt cycle = in return ctxt +(** [stake_info_for_cycle ctxt cycle] reads the stake info for [cycle] from + [ctxt] if it has been previously initialized. Otherwise it initializes + the sampler and caches it in [ctxt] with + [Raw_context.set_stake_info_for_cycle]. *) +let stake_info_for_cycle ctxt cycle = + let open Lwt_result_syntax in + let read ctxt = + let* total_stake = Stake_storage.get_total_active_stake ctxt cycle in + let total_stake = Stake_repr.staking_weight total_stake in + let* stakes_pkh = Stake_storage.get_selected_distribution ctxt cycle in + let* stakes_pk = + List.rev_map_es + (fun (pkh, stake) -> + let+ pk = + Delegate_consensus_key.active_pubkey_for_cycle ctxt pkh cycle + in + (pk, Stake_repr.staking_weight stake)) + stakes_pkh + in + return (total_stake, stakes_pk) + in + Raw_context.stake_info_for_cycle ~read ctxt cycle + +let stake_info ctxt level = + let cycle = level.Level_repr.cycle in + stake_info_for_cycle ctxt cycle + +let load_stake_info_for_cycle ctxt cycle = + let open Lwt_result_syntax in + let* ctxt, (_ : int64), (_ : (Raw_context.consensus_pk * int64) list) = + stake_info_for_cycle ctxt cycle + in + return ctxt + let get_delegate_stake_from_staking_balance ctxt delegate staking_balance = let open Lwt_result_syntax in let* staking_parameters = @@ -230,7 +264,12 @@ let select_distribution_for_cycle ctxt cycle = let state = Sampler.create stakes_pk in let* ctxt = Delegate_sampler_state.init ctxt cycle state in (* pre-allocate the sampler *) - Lwt.return (Raw_context.init_sampler_for_cycle ctxt cycle seed state) + let*? ctxt = Raw_context.init_sampler_for_cycle ctxt cycle seed state in + (* pre-allocate the raw stake distribution info *) + let*? ctxt = + Raw_context.init_stake_info_for_cycle ctxt cycle total_stake stakes_pk + in + return ctxt let select_new_distribution_at_cycle_end ctxt ~new_cycle = let consensus_rights_delay = Constants_storage.consensus_rights_delay ctxt in diff --git a/src/proto_alpha/lib_protocol/delegate_sampler.mli b/src/proto_alpha/lib_protocol/delegate_sampler.mli index 1afecee5c0eb8b81dda5a8818fbb47462985c062..6dd5bd5f9adc68c3801cf9390a31182e7f811638 100644 --- a/src/proto_alpha/lib_protocol/delegate_sampler.mli +++ b/src/proto_alpha/lib_protocol/delegate_sampler.mli @@ -61,6 +61,21 @@ val baking_rights_owner : val load_sampler_for_cycle : Raw_context.t -> Cycle_repr.t -> Raw_context.t tzresult Lwt.t +val stake_info : + Raw_context.t -> + Level_repr.t -> + (Raw_context.t * Int64.t * (Delegate_consensus_key.pk * Int64.t) list) + tzresult + Lwt.t + +(** [load_stake_info_for_cycle ctxt cycle] caches the stake info + for [cycle] in [ctxt]. If the stake info was already cached, + then [ctxt] is returned unchanged. + + This function has the same effect on [ctxt] as {!stake_info} *) +val load_stake_info_for_cycle : + Raw_context.t -> Cycle_repr.t -> Raw_context.t tzresult Lwt.t + val select_new_distribution_at_cycle_end : Raw_context.t -> new_cycle:Cycle_repr.t -> Raw_context.t tzresult Lwt.t diff --git a/src/proto_alpha/lib_protocol/delegate_storage.mli b/src/proto_alpha/lib_protocol/delegate_storage.mli index a5ef9fd9f9e26fcfa5526be27e9de777bf7be092..697b8736d96f3fcdc7b4c48fefaeb1a2bca3eaf1 100644 --- a/src/proto_alpha/lib_protocol/delegate_storage.mli +++ b/src/proto_alpha/lib_protocol/delegate_storage.mli @@ -101,7 +101,7 @@ val fold : (** List all registered delegates. *) val list : Raw_context.t -> Signature.Public_key_hash.t list Lwt.t -(** Returns a delegate's initial frozen deposits at the beginning of cycle. *) +(** Returns a delegate's initial frozen deposits at the beginning of the current cycle. *) val initial_frozen_deposits : Raw_context.t -> Signature.public_key_hash -> Tez_repr.t tzresult Lwt.t diff --git a/src/proto_alpha/lib_protocol/main.ml b/src/proto_alpha/lib_protocol/main.ml index d93e978a5f95eb19d6a466173b911226fd2f5338..064aa4e283793995963eaeedff1210052f6a182e 100644 --- a/src/proto_alpha/lib_protocol/main.ml +++ b/src/proto_alpha/lib_protocol/main.ml @@ -206,12 +206,19 @@ let init_consensus_rights_for_mempool ctxt ~predecessor_level = [slot_owner] itself). *) let cycle = (Level.current ctxt).cycle in let* ctxt = Stake_distribution.load_sampler_for_cycle ctxt cycle in + let* ctxt = Stake_distribution.load_stake_info_for_cycle ctxt cycle in (* If the cycle has changed between the grandparent level and the current level, we also initialize the sampler for that cycle. That way, all three allowed levels are covered. *) match Level.pred ctxt predecessor_level with | Some gp_level when Cycle.(gp_level.cycle <> cycle) -> - Stake_distribution.load_sampler_for_cycle ctxt gp_level.cycle + let* ctxt = + Stake_distribution.load_sampler_for_cycle ctxt gp_level.cycle + in + let* ctxt = + Stake_distribution.load_stake_info_for_cycle ctxt gp_level.cycle + in + return ctxt | Some _ | None -> return ctxt let prepare_ctxt ctxt mode ~(predecessor : Block_header.shell_header) = diff --git a/src/proto_alpha/lib_protocol/raw_context.ml b/src/proto_alpha/lib_protocol/raw_context.ml index fb81643cf2d064e754ca8805789c497732ad24d2..5f01a7c46e16884d946b515a6111247091bb38d7 100644 --- a/src/proto_alpha/lib_protocol/raw_context.ml +++ b/src/proto_alpha/lib_protocol/raw_context.ml @@ -283,6 +283,11 @@ type back = { non_consensus_operations_rev : Operation_hash.t list; dictator_proposal_seen : bool; sampler_state : (Seed_repr.seed * consensus_pk Sampler.t) Cycle_repr.Map.t; + (* [stake_info] maps cycles to a pair [(total_weight, distribution)], where + [total_weight] is the total active staking weight for that cycle, and [distribution] + is a list associating consensus keys with their respective staking weight, ordered + lexicographically by their delegate public key hash. *) + stake_info : (int64 * (consensus_pk * int64) list) Cycle_repr.Map.t; stake_distribution_for_current_cycle : Stake_repr.t Signature.Public_key_hash.Map.t option; reward_coeff_for_current_cycle : Q.t; @@ -355,6 +360,8 @@ let[@inline] dictator_proposal_seen ctxt = ctxt.back.dictator_proposal_seen let[@inline] sampler_state ctxt = ctxt.back.sampler_state +let[@inline] stake_info ctxt = ctxt.back.stake_info + let[@inline] reward_coeff_for_current_cycle ctxt = ctxt.back.reward_coeff_for_current_cycle @@ -403,6 +410,9 @@ let[@inline] update_dictator_proposal_seen ctxt dictator_proposal_seen = let[@inline] update_sampler_state ctxt sampler_state = update_back ctxt {ctxt.back with sampler_state} +let[@inline] update_stake_info ctxt stake_info = + update_back ctxt {ctxt.back with stake_info} + let[@inline] update_reward_coeff_for_current_cycle ctxt reward_coeff_for_current_cycle = update_back ctxt {ctxt.back with reward_coeff_for_current_cycle} @@ -422,6 +432,8 @@ type error += Stake_distribution_not_set (* `Branch *) type error += Sampler_already_set of Cycle_repr.t (* `Permanent *) +type error += Stake_info_already_set of Cycle_repr.t (* `Permanent *) + let () = let open Data_encoding in register_error_kind @@ -480,7 +492,23 @@ let () = c) (obj1 (req "cycle" Cycle_repr.encoding)) (function Sampler_already_set c -> Some c | _ -> None) - (fun c -> Sampler_already_set c) + (fun c -> Sampler_already_set c) ; + register_error_kind + `Permanent + ~id:"stake_info_already_set" + ~title:"Stake already set" + ~description: + "Internal error: Raw_context.set_stake_info_for_cycle was called twice \ + for a given cycle" + ~pp:(fun ppf c -> + Format.fprintf + ppf + "Internal error: stake info already set for cycle %a." + Cycle_repr.pp + c) + (obj1 (req "cycle" Cycle_repr.encoding)) + (function Stake_info_already_set c -> Some c | _ -> None) + (fun c -> Stake_info_already_set c) let fresh_internal_nonce ctxt = let open Result_syntax in @@ -896,6 +924,7 @@ let prepare ~level ~predecessor_timestamp ~timestamp ~adaptive_issuance_enable non_consensus_operations_rev = []; dictator_proposal_seen = false; sampler_state = Cycle_repr.Map.empty; + stake_info = Cycle_repr.Map.empty; stake_distribution_for_current_cycle = None; reward_coeff_for_current_cycle = Q.one; sc_rollup_current_messages; @@ -1906,6 +1935,41 @@ let sampler_for_cycle ~read ctxt cycle = let ctxt = update_sampler_state ctxt map in return (ctxt, seed, state) +let sort_stakes_pk_for_stake_info stakes_pk = + (* The stakes_pk is supposedly already sorted by decreasing stake, from + the call to get_selected_distribution when it was initialized. + We sort them here by lexicographical order on the pkh of the delegate instead. + *) + List.sort + (fun ((consensus_pk1 : consensus_pk), _) (consensus_pk2, _) -> + Signature.Public_key_hash.compare + consensus_pk1.delegate + consensus_pk2.delegate) + stakes_pk + +let init_stake_info_for_cycle ctxt cycle total_stake stakes_pk = + let open Result_syntax in + let map = stake_info ctxt in + if Cycle_repr.Map.mem cycle map then tzfail (Stake_info_already_set cycle) + else + let stakes_pk = sort_stakes_pk_for_stake_info stakes_pk in + let total_stake = Stake_repr.staking_weight total_stake in + let map = Cycle_repr.Map.add cycle (total_stake, stakes_pk) map in + let ctxt = update_stake_info ctxt map in + return ctxt + +let stake_info_for_cycle ~read ctxt cycle = + let open Lwt_result_syntax in + let map = stake_info ctxt in + match Cycle_repr.Map.find cycle map with + | Some (total_stake, stakes_pk) -> return (ctxt, total_stake, stakes_pk) + | None -> + let* total_stake, stakes_pk = read ctxt in + let stakes_pk = sort_stakes_pk_for_stake_info stakes_pk in + let map = Cycle_repr.Map.add cycle (total_stake, stakes_pk) map in + let ctxt = update_stake_info ctxt map in + return (ctxt, total_stake, stakes_pk) + let find_stake_distribution_for_current_cycle ctxt = ctxt.back.stake_distribution_for_current_cycle diff --git a/src/proto_alpha/lib_protocol/raw_context.mli b/src/proto_alpha/lib_protocol/raw_context.mli index b597e4dcdbdc3c1a7d0def4e6ad602d4aa2cc55d..69c4986b20459f258d387c0c33d8e2ba52abf850 100644 --- a/src/proto_alpha/lib_protocol/raw_context.mli +++ b/src/proto_alpha/lib_protocol/raw_context.mli @@ -296,6 +296,30 @@ val sampler_for_cycle : Cycle_repr.t -> (t * Seed_repr.seed * consensus_pk Sampler.t) tzresult Lwt.t +(** [init_stake_info_for_cycle ctxt cycle total_stake stakes_pk] caches the stakes + of the active delegates for [cycle] in memory for quick access. + + @return [Error Stake_info_already_set] if the info was already + cached. *) +val init_stake_info_for_cycle : + t -> + Cycle_repr.t -> + Stake_repr.t -> + (consensus_pk * Int64.t) list -> + t tzresult + +(** [stake_info_for_cycle ~read ctxt cycle] returns the stakes + for [cycle]. The stake info is read in memory if + [init_stake_info_for_cycle] or [stake_info_for_cycle] was previously + called for the same [cycle]. Otherwise, it is read "on-disk" with + the [read] function and then cached in [ctxt] like + [init_stake_info_for_cycle]. *) +val stake_info_for_cycle : + read:(t -> (Int64.t * (consensus_pk * int64) list) tzresult Lwt.t) -> + t -> + Cycle_repr.t -> + (t * Int64.t * (consensus_pk * int64) list) tzresult Lwt.t + (* The stake distribution is stored both in [t] and in the cache. It may be sufficient to only store it in the cache. *) val stake_distribution_for_current_cycle :