From 5ea21d46fb7241ed60312105dfa672456e31d6b0 Mon Sep 17 00:00:00 2001 From: Lucas Randazzo Date: Tue, 26 Nov 2024 14:17:09 +0100 Subject: [PATCH 1/2] Proto: add stake_info in context back This follows the same pattern and execution paths as sampler_state. Note that the Stake_distribution storage was already cached. --- src/proto_alpha/lib_protocol/alpha_context.ml | 2 + .../lib_protocol/alpha_context.mli | 3 + .../lib_protocol/delegate_sampler.ml | 37 ++++++++++- .../lib_protocol/delegate_sampler.mli | 6 ++ src/proto_alpha/lib_protocol/main.ml | 9 ++- src/proto_alpha/lib_protocol/raw_context.ml | 66 ++++++++++++++++++- src/proto_alpha/lib_protocol/raw_context.mli | 24 +++++++ 7 files changed, 144 insertions(+), 3 deletions(-) diff --git a/src/proto_alpha/lib_protocol/alpha_context.ml b/src/proto_alpha/lib_protocol/alpha_context.ml index 2896f64cd2bc..22b085ab1e78 100644 --- a/src/proto_alpha/lib_protocol/alpha_context.ml +++ b/src/proto_alpha/lib_protocol/alpha_context.ml @@ -631,6 +631,8 @@ module Stake_distribution = struct 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 4956a80d3f1e..9a8476635fc5 100644 --- a/src/proto_alpha/lib_protocol/alpha_context.mli +++ b/src/proto_alpha/lib_protocol/alpha_context.mli @@ -5199,6 +5199,9 @@ module Stake_distribution : sig (** 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 b73a69e6768b..f7325bad5a5c 100644 --- a/src/proto_alpha/lib_protocol/delegate_sampler.ml +++ b/src/proto_alpha/lib_protocol/delegate_sampler.ml @@ -164,6 +164,36 @@ 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 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 +260,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 1afecee5c0eb..71693673eaec 100644 --- a/src/proto_alpha/lib_protocol/delegate_sampler.mli +++ b/src/proto_alpha/lib_protocol/delegate_sampler.mli @@ -61,6 +61,12 @@ val baking_rights_owner : val load_sampler_for_cycle : Raw_context.t -> Cycle_repr.t -> Raw_context.t 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. *) +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/main.ml b/src/proto_alpha/lib_protocol/main.ml index d93e978a5f95..064aa4e28379 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 fb81643cf2d0..5f01a7c46e16 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 b597e4dcdbdc..69c4986b2045 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 : -- GitLab From 49457769988a857ab6643d1dd1a6a366ec600ec2 Mon Sep 17 00:00:00 2001 From: Lucas Randazzo Date: Tue, 26 Nov 2024 16:17:34 +0100 Subject: [PATCH 2/2] Proto: Add stake_info in alpha_context A way to get it from the back. I suppose it is also an alternative way to get the total staking weight, which is used in some places. I cannot ensure that this value is available for other things than attesting, though, so let's keep it like this for now. --- src/proto_alpha/lib_protocol/alpha_context.ml | 2 ++ src/proto_alpha/lib_protocol/alpha_context.mli | 5 +++++ src/proto_alpha/lib_protocol/delegate_sampler.ml | 4 ++++ src/proto_alpha/lib_protocol/delegate_sampler.mli | 11 ++++++++++- src/proto_alpha/lib_protocol/delegate_storage.mli | 2 +- 5 files changed, 22 insertions(+), 2 deletions(-) diff --git a/src/proto_alpha/lib_protocol/alpha_context.ml b/src/proto_alpha/lib_protocol/alpha_context.ml index 22b085ab1e78..df9752cc6165 100644 --- a/src/proto_alpha/lib_protocol/alpha_context.ml +++ b/src/proto_alpha/lib_protocol/alpha_context.ml @@ -629,6 +629,8 @@ 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 diff --git a/src/proto_alpha/lib_protocol/alpha_context.mli b/src/proto_alpha/lib_protocol/alpha_context.mli index 9a8476635fc5..585fcce1d9f3 100644 --- a/src/proto_alpha/lib_protocol/alpha_context.mli +++ b/src/proto_alpha/lib_protocol/alpha_context.mli @@ -5196,6 +5196,11 @@ 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 diff --git a/src/proto_alpha/lib_protocol/delegate_sampler.ml b/src/proto_alpha/lib_protocol/delegate_sampler.ml index f7325bad5a5c..b271dec004aa 100644 --- a/src/proto_alpha/lib_protocol/delegate_sampler.ml +++ b/src/proto_alpha/lib_protocol/delegate_sampler.ml @@ -187,6 +187,10 @@ let stake_info_for_cycle ctxt cycle = 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) = diff --git a/src/proto_alpha/lib_protocol/delegate_sampler.mli b/src/proto_alpha/lib_protocol/delegate_sampler.mli index 71693673eaec..6dd5bd5f9adc 100644 --- a/src/proto_alpha/lib_protocol/delegate_sampler.mli +++ b/src/proto_alpha/lib_protocol/delegate_sampler.mli @@ -61,9 +61,18 @@ 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. *) + 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 diff --git a/src/proto_alpha/lib_protocol/delegate_storage.mli b/src/proto_alpha/lib_protocol/delegate_storage.mli index a5ef9fd9f9e2..697b8736d96f 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 -- GitLab