From 565dbe06dfcbdf4a2b0679bc08e34c1f1b3fdaa1 Mon Sep 17 00:00:00 2001
From: Lucas Randazzo <lucas@nomadic-labs.com>
Date: Mon, 3 Mar 2025 17:59:09 +0100
Subject: [PATCH 1/7] Proto: add storage for a second consensus key

---
 src/proto_alpha/lib_protocol/storage.ml  | 18 ++++++++++++++++++
 src/proto_alpha/lib_protocol/storage.mli | 16 +++++++++++++++-
 2 files changed, 33 insertions(+), 1 deletion(-)

diff --git a/src/proto_alpha/lib_protocol/storage.ml b/src/proto_alpha/lib_protocol/storage.ml
index c88f8dd04805..d58e13733864 100644
--- a/src/proto_alpha/lib_protocol/storage.ml
+++ b/src/proto_alpha/lib_protocol/storage.ml
@@ -335,6 +335,14 @@ module Contract = struct
       end)
       (Signature.Public_key)
 
+  module Companion_key =
+    Indexed_context.Make_map
+      (Registered)
+      (struct
+        let name = ["companion_key"; "active"]
+      end)
+      (Bls.Public_key)
+
   module Staking_parameters =
     Indexed_context.Make_map
       (Registered)
@@ -1350,6 +1358,15 @@ module Cycle = struct
          (Make_index (Contract_repr.Index))
       (Signature.Public_key)
 
+  module Pending_companion_keys =
+    Make_indexed_data_storage
+      (Make_subcontext (Registered) (Indexed_context.Raw_context)
+         (struct
+           let name = ["pending_companion_keys"]
+         end))
+         (Make_index (Contract_repr.Index))
+      (Bls.Public_key)
+
   module Pending_staking_parameters =
     Make_indexed_data_storage
       (Make_subcontext (Registered) (Indexed_context.Raw_context)
@@ -1363,6 +1380,7 @@ end
 module Already_denounced = Cycle.Already_denounced
 module Dal_already_denounced = Cycle.Dal_already_denounced
 module Pending_consensus_keys = Cycle.Pending_consensus_keys
+module Pending_companion_keys = Cycle.Pending_companion_keys
 module Pending_staking_parameters = Cycle.Pending_staking_parameters
 
 module Stake = struct
diff --git a/src/proto_alpha/lib_protocol/storage.mli b/src/proto_alpha/lib_protocol/storage.mli
index aa8226cd3c67..bf6fa9d67747 100644
--- a/src/proto_alpha/lib_protocol/storage.mli
+++ b/src/proto_alpha/lib_protocol/storage.mli
@@ -183,6 +183,13 @@ module Contract : sig
        and type value = Signature.Public_key.t
        and type t := Raw_context.t
 
+  (** The active companion key of a delegate *)
+  module Companion_key :
+    Indexed_data_storage
+      with type key = Contract_repr.t
+       and type value = Bls.Public_key.t
+       and type t := Raw_context.t
+
   (** The delegate of a contract, if any. *)
   module Delegate :
     Indexed_data_storage
@@ -468,7 +475,7 @@ module Delegates :
     with type t := Raw_context.t
      and type elt = Signature.Public_key_hash.t
 
-(** Set of all active consensus keys in cycle `current + consensus_rights_delay + 1` *)
+(** Set of all active consensus and companion keys in cycle `current + consensus_rights_delay + 1` *)
 module Consensus_keys :
   Data_set_storage
     with type t := Raw_context.t
@@ -481,6 +488,13 @@ module Pending_consensus_keys :
      and type key = Contract_repr.t
      and type value = Signature.public_key
 
+(** The pending companion key of a delegate at the given cycle *)
+module Pending_companion_keys :
+  Indexed_data_storage
+    with type t := Raw_context.t * Cycle_repr.t
+     and type key = Contract_repr.t
+     and type value = Bls.Public_key.t
+
 (** All denunciations of the current and previous cycles that will have an effect
     (slashing, reward), i.e. all below 100%, deferred to the end of their
     slashing period. *)
-- 
GitLab


From 736648950483c4fb0aa4440ac44fb3f8418c7369 Mon Sep 17 00:00:00 2001
From: Lucas Randazzo <lucas@nomadic-labs.com>
Date: Tue, 4 Mar 2025 11:56:30 +0100
Subject: [PATCH 2/7] Proto: pending companion updates

---
 .../lib_protocol/delegate_consensus_key.ml    | 34 +++++++++++++++----
 .../lib_protocol/delegate_consensus_key.mli   |  6 ++++
 2 files changed, 34 insertions(+), 6 deletions(-)

diff --git a/src/proto_alpha/lib_protocol/delegate_consensus_key.ml b/src/proto_alpha/lib_protocol/delegate_consensus_key.ml
index 004dd0d6816c..4aec61d4536b 100644
--- a/src/proto_alpha/lib_protocol/delegate_consensus_key.ml
+++ b/src/proto_alpha/lib_protocol/delegate_consensus_key.ml
@@ -71,6 +71,10 @@ let () =
     (function Invalid_consensus_key_update_tz4 pk -> Some pk | _ -> None)
     (fun pk -> Invalid_consensus_key_update_tz4 pk)
 
+type 'a kind =
+  | Consensus : Signature.public_key kind
+  | Companion : Bls.Public_key.t option kind
+
 type pk = Raw_context.consensus_pk = {
   delegate : Signature.Public_key_hash.t;
   consensus_pk : Signature.Public_key.t;
@@ -151,7 +155,8 @@ let active_key ctxt delegate =
   let* pk = active_pubkey ctxt delegate in
   return (pkh pk)
 
-let raw_pending_updates ctxt ?up_to_cycle delegate =
+let raw_pending_updates (type a) ctxt ?up_to_cycle ~(kind : a kind) delegate :
+    (Cycle_repr.t * a) list tzresult Lwt.t =
   let open Lwt_result_syntax in
   let relevant_cycles =
     let level = Raw_context.current_level ctxt in
@@ -168,23 +173,40 @@ let raw_pending_updates ctxt ?up_to_cycle delegate =
     Cycle_repr.(first_cycle ---> last_cycle)
   in
   let delegate = Contract_repr.Implicit delegate in
+  let find (type a) (ctxt, cycle) delegate (kind : a kind) :
+      a option tzresult Lwt.t =
+    match kind with
+    | Consensus -> Storage.Pending_consensus_keys.find (ctxt, cycle) delegate
+    | Companion ->
+        let* pk = Storage.Pending_companion_keys.find (ctxt, cycle) delegate in
+        return (Option.map Option.some pk)
+  in
   List.filter_map_es
     (fun cycle ->
-      let* pending_for_cycle =
-        Storage.Pending_consensus_keys.find (ctxt, cycle) delegate
-      in
+      let* pending_for_cycle = find (ctxt, cycle) delegate kind in
       pending_for_cycle |> Option.map (fun pk -> (cycle, pk)) |> return)
     relevant_cycles
 
 let pending_updates ctxt delegate =
   let open Lwt_result_syntax in
-  let* updates = raw_pending_updates ctxt delegate in
+  let* updates = raw_pending_updates ctxt ~kind:Consensus delegate in
   return
     (List.map (fun (c, pk) -> (c, Signature.Public_key.hash pk, pk)) updates)
 
+let pending_companion_updates ctxt delegate =
+  let open Lwt_result_syntax in
+  let* updates = raw_pending_updates ctxt ~kind:Companion delegate in
+  return
+    (List.filter_map
+       (fun (c, pk) ->
+         Option.map (fun pk -> (c, Bls.Public_key.hash pk, pk)) pk)
+       updates)
+
 let raw_active_pubkey_for_cycle ctxt delegate cycle =
   let open Lwt_result_syntax in
-  let* pendings = raw_pending_updates ctxt ~up_to_cycle:cycle delegate in
+  let* pendings =
+    raw_pending_updates ctxt ~up_to_cycle:cycle ~kind:Consensus delegate
+  in
   let* active = active_pubkey ctxt delegate in
   let current_cycle = (Raw_context.current_level ctxt).cycle in
   match List.hd (List.rev pendings) with
diff --git a/src/proto_alpha/lib_protocol/delegate_consensus_key.mli b/src/proto_alpha/lib_protocol/delegate_consensus_key.mli
index 920ec51e028c..60b1ca2a21c8 100644
--- a/src/proto_alpha/lib_protocol/delegate_consensus_key.mli
+++ b/src/proto_alpha/lib_protocol/delegate_consensus_key.mli
@@ -88,6 +88,12 @@ val pending_updates :
   tzresult
   Lwt.t
 
+(** Returns the list of pending companion key updates in upcoming cycles. *)
+val pending_companion_updates :
+  Raw_context.t ->
+  Signature.Public_key_hash.t ->
+  (Cycle_repr.t * Bls.Public_key_hash.t * Bls.Public_key.t) list tzresult Lwt.t
+
 (** Register a consensus-key update. *)
 val register_update :
   Raw_context.t ->
-- 
GitLab


From f3d190b99684d6d19f64744a6f4e39eecb8bfce7 Mon Sep 17 00:00:00 2001
From: Lucas Randazzo <lucas@nomadic-labs.com>
Date: Tue, 11 Mar 2025 17:18:52 +0100
Subject: [PATCH 3/7] Proto: add companion fields in consensus pk type

Note: the type Delegate_consensus_key.t does not need to be updated,
but it could easily be done if necessary
---
 src/proto_alpha/lib_plugin/RPC.ml             | 10 ++-
 src/proto_alpha/lib_plugin/mempool.ml         | 20 +++++-
 .../lib_protocol/alpha_context.mli            |  2 +
 .../lib_protocol/delegate_consensus_key.ml    | 62 +++++++++++++------
 .../lib_protocol/delegate_consensus_key.mli   |  2 +
 src/proto_alpha/lib_protocol/raw_context.ml   | 28 ++++++---
 src/proto_alpha/lib_protocol/raw_context.mli  |  2 +
 7 files changed, 96 insertions(+), 30 deletions(-)

diff --git a/src/proto_alpha/lib_plugin/RPC.ml b/src/proto_alpha/lib_plugin/RPC.ml
index 6449f4f7da32..2494ee70f45b 100644
--- a/src/proto_alpha/lib_plugin/RPC.ml
+++ b/src/proto_alpha/lib_plugin/RPC.ml
@@ -3654,7 +3654,13 @@ module Baking_rights = struct
       else
         let* ( ctxt,
                _slot,
-               {Consensus_key.consensus_pkh; delegate; consensus_pk = _} ) =
+               {
+                 Consensus_key.consensus_pkh;
+                 delegate;
+                 consensus_pk = _;
+                 companion_pk = _;
+                 companion_pkh = _;
+               } ) =
           Stake_distribution.baking_rights_owner ctxt level ~round
         in
         let*? timestamp =
@@ -3866,6 +3872,8 @@ module Attestation_rights = struct
                  Consensus_key.delegate;
                  consensus_pk = _;
                  consensus_pkh = consensus_key;
+                 companion_pk = _;
+                 companion_pkh = _;
                },
                attestation_power,
                _dal_power )
diff --git a/src/proto_alpha/lib_plugin/mempool.ml b/src/proto_alpha/lib_plugin/mempool.ml
index 402b0b6c57e2..02dd68d80ac3 100644
--- a/src/proto_alpha/lib_plugin/mempool.ml
+++ b/src/proto_alpha/lib_plugin/mempool.ml
@@ -814,7 +814,15 @@ let sources_from_operation ctxt
         Stake_distribution.slot_owner ctxt level consensus_content.slot
       in
       match slot_owner with
-      | Ok (_ctxt, {delegate; consensus_pkh; consensus_pk = _}) ->
+      | Ok
+          ( _ctxt,
+            {
+              delegate;
+              consensus_pkh;
+              consensus_pk = _;
+              companion_pkh = _;
+              companion_pk = _;
+            } ) ->
           return @@ map_pkh_env [delegate; consensus_pkh]
       | Error _ -> return_nil)
   | Single (Attestations_aggregate {consensus_content; committee}) ->
@@ -824,7 +832,15 @@ let sources_from_operation ctxt
           (fun acc slot ->
             let* slot_owner = Stake_distribution.slot_owner ctxt level slot in
             match slot_owner with
-            | Ok (_ctxt, {delegate; consensus_pkh; consensus_pk = _}) ->
+            | Ok
+                ( _ctxt,
+                  {
+                    delegate;
+                    consensus_pkh;
+                    consensus_pk = _;
+                    companion_pkh = _;
+                    companion_pk = _;
+                  } ) ->
                 return (delegate :: consensus_pkh :: acc)
             | Error _ -> return acc)
           []
diff --git a/src/proto_alpha/lib_protocol/alpha_context.mli b/src/proto_alpha/lib_protocol/alpha_context.mli
index 7891ddf1cf0d..97bc7c722db7 100644
--- a/src/proto_alpha/lib_protocol/alpha_context.mli
+++ b/src/proto_alpha/lib_protocol/alpha_context.mli
@@ -2239,6 +2239,8 @@ module Consensus_key : sig
     delegate : Signature.Public_key_hash.t;
     consensus_pk : Signature.Public_key.t;
     consensus_pkh : Signature.Public_key_hash.t;
+    companion_pk : Bls.Public_key.t option;
+    companion_pkh : Bls.Public_key_hash.t option;
   }
 
   type t = {
diff --git a/src/proto_alpha/lib_protocol/delegate_consensus_key.ml b/src/proto_alpha/lib_protocol/delegate_consensus_key.ml
index 4aec61d4536b..adab017eb5e6 100644
--- a/src/proto_alpha/lib_protocol/delegate_consensus_key.ml
+++ b/src/proto_alpha/lib_protocol/delegate_consensus_key.ml
@@ -79,6 +79,8 @@ type pk = Raw_context.consensus_pk = {
   delegate : Signature.Public_key_hash.t;
   consensus_pk : Signature.Public_key.t;
   consensus_pkh : Signature.Public_key_hash.t;
+  companion_pk : Bls.Public_key.t option;
+  companion_pkh : Bls.Public_key_hash.t option;
 }
 
 type t = {
@@ -95,7 +97,15 @@ let encoding =
        (req "delegate" Signature.Public_key_hash.encoding)
        (req "consensus_pkh" Signature.Public_key_hash.encoding)
 
-let pkh {delegate; consensus_pkh; consensus_pk = _} = {delegate; consensus_pkh}
+let pkh
+    {
+      delegate;
+      consensus_pkh;
+      companion_pkh = _;
+      consensus_pk = _;
+      companion_pk = _;
+    } =
+  {delegate; consensus_pkh}
 
 let zero =
   {
@@ -142,13 +152,21 @@ let init ctxt delegate pk =
   let*! ctxt = set_used ctxt pkh in
   Storage.Contract.Consensus_key.init ctxt (Contract_repr.Implicit delegate) pk
 
+let active_pubkey_kind (type a) ctxt ~(kind : a kind) delegate :
+    a tzresult Lwt.t =
+  match kind with
+  | Consensus ->
+      Storage.Contract.Consensus_key.get ctxt (Contract_repr.Implicit delegate)
+  | Companion ->
+      Storage.Contract.Companion_key.find ctxt (Contract_repr.Implicit delegate)
+
 let active_pubkey ctxt delegate =
   let open Lwt_result_syntax in
-  let* pk =
-    Storage.Contract.Consensus_key.get ctxt (Contract_repr.Implicit delegate)
-  in
-  let pkh = Signature.Public_key.hash pk in
-  return {consensus_pk = pk; consensus_pkh = pkh; delegate}
+  let* consensus_pk = active_pubkey_kind ctxt ~kind:Consensus delegate in
+  let consensus_pkh = Signature.Public_key.hash consensus_pk in
+  let* companion_pk = active_pubkey_kind ctxt ~kind:Companion delegate in
+  let companion_pkh = Option.map Bls.Public_key.hash companion_pk in
+  return {consensus_pk; consensus_pkh; companion_pk; companion_pkh; delegate}
 
 let active_key ctxt delegate =
   let open Lwt_result_syntax in
@@ -202,25 +220,31 @@ let pending_companion_updates ctxt delegate =
          Option.map (fun pk -> (c, Bls.Public_key.hash pk, pk)) pk)
        updates)
 
-let raw_active_pubkey_for_cycle ctxt delegate cycle =
+let raw_active_pubkey_for_cycle ctxt ~kind delegate cycle =
   let open Lwt_result_syntax in
-  let* pendings =
-    raw_pending_updates ctxt ~up_to_cycle:cycle ~kind:Consensus delegate
-  in
-  let* active = active_pubkey ctxt delegate in
+  let* pendings = raw_pending_updates ctxt ~up_to_cycle:cycle ~kind delegate in
+  let* active_pk = active_pubkey_kind ctxt ~kind delegate in
   let current_cycle = (Raw_context.current_level ctxt).cycle in
   match List.hd (List.rev pendings) with
-  | None -> return (current_cycle, active.consensus_pk)
+  | None -> return (current_cycle, active_pk)
   | Some (cycle, pk) -> return (cycle, pk)
 
 let active_pubkey_for_cycle ctxt delegate cycle =
   let open Lwt_result_syntax in
-  let+ _, consensus_pk = raw_active_pubkey_for_cycle ctxt delegate cycle in
-  {
-    consensus_pk;
-    consensus_pkh = Signature.Public_key.hash consensus_pk;
-    delegate;
-  }
+  let* _, consensus_pk =
+    raw_active_pubkey_for_cycle ctxt ~kind:Consensus delegate cycle
+  in
+  let* _, companion_pk =
+    raw_active_pubkey_for_cycle ctxt ~kind:Companion delegate cycle
+  in
+  return
+    {
+      consensus_pk;
+      consensus_pkh = Signature.Public_key.hash consensus_pk;
+      companion_pk;
+      companion_pkh = Option.map Bls.Public_key.hash companion_pk;
+      delegate;
+    }
 
 let register_update ctxt delegate pk =
   let open Lwt_result_syntax in
@@ -231,7 +255,7 @@ let register_update ctxt delegate pk =
   in
   let* () =
     let* first_active_cycle, active_pubkey =
-      raw_active_pubkey_for_cycle ctxt delegate update_cycle
+      raw_active_pubkey_for_cycle ctxt ~kind:Consensus delegate update_cycle
     in
     fail_when
       Signature.Public_key.(pk = active_pubkey)
diff --git a/src/proto_alpha/lib_protocol/delegate_consensus_key.mli b/src/proto_alpha/lib_protocol/delegate_consensus_key.mli
index 60b1ca2a21c8..f4ba0fde8fa0 100644
--- a/src/proto_alpha/lib_protocol/delegate_consensus_key.mli
+++ b/src/proto_alpha/lib_protocol/delegate_consensus_key.mli
@@ -39,6 +39,8 @@ type pk = Raw_context.consensus_pk = {
   delegate : Signature.Public_key_hash.t;
   consensus_pk : Signature.Public_key.t;
   consensus_pkh : Signature.Public_key_hash.t;
+  companion_pk : Bls.Public_key.t option;
+  companion_pkh : Bls.Public_key_hash.t option;
 }
 
 (** The public key hash of a consensus key and the associated delegate. *)
diff --git a/src/proto_alpha/lib_protocol/raw_context.ml b/src/proto_alpha/lib_protocol/raw_context.ml
index e8ebe72b7a3b..fb81643cf2d0 100644
--- a/src/proto_alpha/lib_protocol/raw_context.ml
+++ b/src/proto_alpha/lib_protocol/raw_context.ml
@@ -60,24 +60,36 @@ type consensus_pk = {
   delegate : Signature.Public_key_hash.t;
   consensus_pk : Signature.Public_key.t;
   consensus_pkh : Signature.Public_key_hash.t;
+  companion_pk : Bls.Public_key.t option;
+  companion_pkh : Bls.Public_key_hash.t option;
 }
 
 let consensus_pk_encoding =
   let open Data_encoding in
   conv
-    (fun {delegate; consensus_pk; consensus_pkh} ->
-      if Signature.Public_key_hash.equal consensus_pkh delegate then
-        (consensus_pk, None)
-      else (consensus_pk, Some delegate))
-    (fun (consensus_pk, delegate) ->
+    (fun {
+           delegate;
+           consensus_pk;
+           consensus_pkh;
+           companion_pk;
+           companion_pkh = _;
+         } ->
+      let delegate =
+        if Signature.Public_key_hash.equal consensus_pkh delegate then None
+        else Some delegate
+      in
+      (consensus_pk, delegate, companion_pk))
+    (fun (consensus_pk, delegate, companion_pk) ->
       let consensus_pkh = Signature.Public_key.hash consensus_pk in
+      let companion_pkh = Option.map Bls.Public_key.hash companion_pk in
       let delegate =
         match delegate with None -> consensus_pkh | Some del -> del
       in
-      {delegate; consensus_pk; consensus_pkh})
-    (obj2
+      {delegate; consensus_pk; consensus_pkh; companion_pk; companion_pkh})
+    (obj3
        (req "consensus_pk" Signature.Public_key.encoding)
-       (opt "delegate" Signature.Public_key_hash.encoding))
+       (opt "delegate" Signature.Public_key_hash.encoding)
+       (opt "companion_pk" Bls.Public_key.encoding))
 
 module Raw_consensus = struct
   (** Consensus operations are indexed by their [initial slots]. Given
diff --git a/src/proto_alpha/lib_protocol/raw_context.mli b/src/proto_alpha/lib_protocol/raw_context.mli
index c4b0a70e30c2..b597e4dcdbdc 100644
--- a/src/proto_alpha/lib_protocol/raw_context.mli
+++ b/src/proto_alpha/lib_protocol/raw_context.mli
@@ -264,6 +264,8 @@ type consensus_pk = {
   delegate : Signature.Public_key_hash.t;
   consensus_pk : Signature.Public_key.t;
   consensus_pkh : Signature.Public_key_hash.t;
+  companion_pk : Bls.Public_key.t option;
+  companion_pkh : Bls.Public_key_hash.t option;
 }
 
 val consensus_pk_encoding : consensus_pk Data_encoding.t
-- 
GitLab


From 55d6a6b75ea5fe383d22a71163162452c5a6382a Mon Sep 17 00:00:00 2001
From: Lucas Randazzo <lucas@nomadic-labs.com>
Date: Tue, 11 Mar 2025 17:46:02 +0100
Subject: [PATCH 4/7] Proto: sampler stitching

---
 src/proto_alpha/lib_protocol/init_storage.ml | 23 +++++++++++
 src/proto_alpha/lib_protocol/sampler.ml      |  5 +++
 src/proto_alpha/lib_protocol/sampler.mli     |  3 ++
 src/proto_alpha/lib_protocol/storage.ml      | 42 ++++++++++++++++++++
 src/proto_alpha/lib_protocol/storage.mli     | 14 +++++++
 5 files changed, 87 insertions(+)

diff --git a/src/proto_alpha/lib_protocol/init_storage.ml b/src/proto_alpha/lib_protocol/init_storage.ml
index 45dddf66cdcc..180b22224b07 100644
--- a/src/proto_alpha/lib_protocol/init_storage.ml
+++ b/src/proto_alpha/lib_protocol/init_storage.ml
@@ -125,6 +125,28 @@ let prepare ctxt ~level ~predecessor_timestamp ~timestamp =
 (* Start of code to remove at r022 automatic protocol snapshot *)
 
 (* Please add here any code that should be removed at the r022 automatic protocol snapshot *)
+let update_delegate_sampler_state_value_type_from_R_to_S ctxt =
+  let open Lwt_result_syntax in
+  Storage.Delegate_sampler_state_up_to_R.fold
+    ctxt
+    ~order:`Undefined
+    ~init:(ok ctxt)
+    ~f:(fun cycle sampler ctxt ->
+      let*? ctxt in
+      let sampler =
+        Sampler.map
+          (fun {Storage.delegate; consensus_pk; consensus_pkh} ->
+            {
+              Delegate_consensus_key.delegate;
+              consensus_pk;
+              consensus_pkh;
+              companion_pk = None;
+              companion_pkh = None;
+            })
+          sampler
+      in
+      let* ctxt = Storage.Delegate_sampler_state.update ctxt cycle sampler in
+      return ctxt)
 
 (* End of code to remove at r022 automatic protocol snapshot *)
 
@@ -235,6 +257,7 @@ let prepare_first_block chain_id ctxt ~typecheck_smart_contract
         let* ctxt =
           Sc_rollup_refutation_storage.migrate_clean_refutation_games ctxt
         in
+        let* ctxt = update_delegate_sampler_state_value_type_from_R_to_S ctxt in
         return (ctxt, [])
     (* End of alpha predecessor stitching. Comment used for automatic snapshot *)
   in
diff --git a/src/proto_alpha/lib_protocol/sampler.ml b/src/proto_alpha/lib_protocol/sampler.ml
index 34eb5adda73b..207797902888 100644
--- a/src/proto_alpha/lib_protocol/sampler.ml
+++ b/src/proto_alpha/lib_protocol/sampler.ml
@@ -62,6 +62,8 @@ module type S = sig
   val sample : 'a t -> (int_bound:int -> mass_bound:mass -> int * mass) -> 'a
 
   val encoding : 'a Data_encoding.t -> 'a t Data_encoding.t
+
+  val map : ('a -> 'b) -> 'a t -> 'b t
 end
 
 module Make (Mass : SMass) : S with type mass = Mass.t = struct
@@ -174,6 +176,9 @@ module Make (Mass : SMass) : S with type mass = Mass.t = struct
          (req "support" (array_encoding enc))
          (req "p" mass_array_encoding)
          (req "alias" int_array_encoding))
+
+  let map f ({support; _} as sampler) =
+    {sampler with support = FallbackArray.map f support}
 end
 
 module Internal_for_tests = struct
diff --git a/src/proto_alpha/lib_protocol/sampler.mli b/src/proto_alpha/lib_protocol/sampler.mli
index 00cf20a7768c..4fc30c0cdf35 100644
--- a/src/proto_alpha/lib_protocol/sampler.mli
+++ b/src/proto_alpha/lib_protocol/sampler.mli
@@ -61,6 +61,9 @@ module type S = sig
 
   (** [encoding e] constructs an encoding for ['a t] given an encoding for ['a]. *)
   val encoding : 'a Data_encoding.t -> 'a t Data_encoding.t
+
+  (** [map f t] changes the return values of the sampler without changing the distribution. *)
+  val map : ('a -> 'b) -> 'a t -> 'b t
 end
 
 (**/**)
diff --git a/src/proto_alpha/lib_protocol/storage.ml b/src/proto_alpha/lib_protocol/storage.ml
index d58e13733864..647a0824dc24 100644
--- a/src/proto_alpha/lib_protocol/storage.ml
+++ b/src/proto_alpha/lib_protocol/storage.ml
@@ -1253,6 +1253,41 @@ module Cycle = struct
       end)
       (Stake_repr)
 
+  type consensus_pk_in_R = {
+    delegate : Signature.Public_key_hash.t;
+    consensus_pk : Signature.Public_key.t;
+    consensus_pkh : Signature.Public_key_hash.t;
+  }
+
+  let consensus_pk_encoding_in_R =
+    let open Data_encoding in
+    conv
+      (fun {delegate; consensus_pk; consensus_pkh} ->
+        if Signature.Public_key_hash.equal consensus_pkh delegate then
+          (consensus_pk, None)
+        else (consensus_pk, Some delegate))
+      (fun (consensus_pk, delegate) ->
+        let consensus_pkh = Signature.Public_key.hash consensus_pk in
+        let delegate =
+          match delegate with None -> consensus_pkh | Some del -> del
+        in
+        {delegate; consensus_pk; consensus_pkh})
+      (obj2
+         (req "consensus_pk" Signature.Public_key.encoding)
+         (opt "delegate" Signature.Public_key_hash.encoding))
+
+  module Delegate_sampler_state_up_to_R =
+    Indexed_context.Make_map
+      (Ghost)
+      (struct
+        let name = ["delegate_sampler_state"]
+      end)
+      (struct
+        type t = consensus_pk_in_R Sampler.t
+
+        let encoding = Sampler.encoding consensus_pk_encoding_in_R
+      end)
+
   module Delegate_sampler_state =
     Indexed_context.Make_map
       (Registered)
@@ -1405,6 +1440,13 @@ module Stake = struct
   module Total_active_stake = Cycle.Total_active_stake
 end
 
+type consensus_pk_in_R = Cycle.consensus_pk_in_R = {
+  delegate : Signature.Public_key_hash.t;
+  consensus_pk : Signature.Public_key.t;
+  consensus_pkh : Signature.Public_key_hash.t;
+}
+
+module Delegate_sampler_state_up_to_R = Cycle.Delegate_sampler_state_up_to_R
 module Delegate_sampler_state = Cycle.Delegate_sampler_state
 module Issuance_bonus = Cycle.Issuance_bonus
 module Issuance_coeff = Cycle.Issuance_coeff
diff --git a/src/proto_alpha/lib_protocol/storage.mli b/src/proto_alpha/lib_protocol/storage.mli
index bf6fa9d67747..ba64b60b6237 100644
--- a/src/proto_alpha/lib_protocol/storage.mli
+++ b/src/proto_alpha/lib_protocol/storage.mli
@@ -591,6 +591,20 @@ module Stake : sig
        and type t := Raw_context.t
 end
 
+type consensus_pk_in_R = {
+  delegate : Signature.Public_key_hash.t;
+  consensus_pk : Signature.Public_key.t;
+  consensus_pkh : Signature.Public_key_hash.t;
+}
+
+(** State of the sampler used to select delegates up to R.
+    Remove in S. *)
+module Delegate_sampler_state_up_to_R :
+  Indexed_data_storage
+    with type key = Cycle_repr.t
+     and type value = consensus_pk_in_R Sampler.t
+     and type t := Raw_context.t
+
 (** State of the sampler used to select delegates. Managed synchronously
     with [Stake.Selected_distribution_for_cycle]. *)
 module Delegate_sampler_state :
-- 
GitLab


From 1fa1ac47cb31c44c6909af1119198e3495c74b9d Mon Sep 17 00:00:00 2001
From: Lucas Randazzo <lucas@nomadic-labs.com>
Date: Wed, 12 Mar 2025 17:40:12 +0100
Subject: [PATCH 5/7] Proto/consensus: add register_companion_key

---
 .../lib_protocol/delegate_consensus_key.ml    | 46 ++++++++++++++++++-
 .../lib_protocol/delegate_consensus_key.mli   |  7 +++
 2 files changed, 51 insertions(+), 2 deletions(-)

diff --git a/src/proto_alpha/lib_protocol/delegate_consensus_key.ml b/src/proto_alpha/lib_protocol/delegate_consensus_key.ml
index adab017eb5e6..d80a4b650e33 100644
--- a/src/proto_alpha/lib_protocol/delegate_consensus_key.ml
+++ b/src/proto_alpha/lib_protocol/delegate_consensus_key.ml
@@ -125,9 +125,12 @@ let pp ppf {delegate; consensus_pkh} =
 
 (* Invariant:
       No two delegates use the same active consensus key at a given time.
+      No two delegates use the same active companion key at a given time.
+      An active key is either a consensus key, or a companion key: it
+   cannot be both at once.
 
-   To ensure that, {!Storage.Consensus_keys} contains keys that will be active
-   at cycle `current + consensus_rights_delay + 1`.
+   To ensure that, {!Storage.Consensus_keys} contains keys (both consensus and
+   companion) that will be active at cycle `current + consensus_rights_delay + 1`.
 *)
 
 let check_unused ctxt pkh =
@@ -276,6 +279,45 @@ let register_update ctxt delegate pk =
   in
   return ctxt
 
+let register_update_companion ctxt delegate pk =
+  let open Lwt_result_syntax in
+  let update_cycle =
+    let current_level = Raw_context.current_level ctxt in
+    let cycles_delay = Constants_storage.consensus_key_activation_delay ctxt in
+    Cycle_repr.add current_level.cycle (cycles_delay + 1)
+  in
+  let* () =
+    let* first_active_cycle, active_companion_pubkey =
+      raw_active_pubkey_for_cycle ctxt ~kind:Companion delegate update_cycle
+    in
+    match active_companion_pubkey with
+    | None -> return_unit
+    | Some active_pubkey ->
+        fail_when
+          Bls.Public_key.(pk = active_pubkey)
+          (Invalid_consensus_key_update_noop first_active_cycle)
+  in
+  let pkh : Signature.public_key_hash =
+    Signature.Bls (Bls.Public_key.hash pk)
+  in
+  let* () = check_unused ctxt pkh in
+  let*! ctxt = set_used ctxt pkh in
+  let* {companion_pkh = old_pkh; _} =
+    active_pubkey_for_cycle ctxt delegate update_cycle
+  in
+  let*! ctxt =
+    match old_pkh with
+    | None -> Lwt.return ctxt
+    | Some old_pkh -> set_unused ctxt (Signature.Bls old_pkh)
+  in
+  let*! ctxt =
+    Storage.Pending_companion_keys.add
+      (ctxt, update_cycle)
+      (Contract_repr.Implicit delegate)
+      pk
+  in
+  return ctxt
+
 let activate ctxt ~new_cycle =
   let open Lwt_syntax in
   let* ctxt =
diff --git a/src/proto_alpha/lib_protocol/delegate_consensus_key.mli b/src/proto_alpha/lib_protocol/delegate_consensus_key.mli
index f4ba0fde8fa0..55355f6b211f 100644
--- a/src/proto_alpha/lib_protocol/delegate_consensus_key.mli
+++ b/src/proto_alpha/lib_protocol/delegate_consensus_key.mli
@@ -103,5 +103,12 @@ val register_update :
   Signature.Public_key.t ->
   Raw_context.t tzresult Lwt.t
 
+(** Register a companion-key update. *)
+val register_update_companion :
+  Raw_context.t ->
+  Signature.Public_key_hash.t ->
+  Bls.Public_key.t ->
+  Raw_context.t tzresult Lwt.t
+
 (** Activate consensus keys at the beginning of cycle [new_cycle]. *)
 val activate : Raw_context.t -> new_cycle:Cycle_repr.t -> Raw_context.t Lwt.t
-- 
GitLab


From ac9bb26ecf9bd52d28fe7c2814b709c128d5f9a9 Mon Sep 17 00:00:00 2001
From: Lucas Randazzo <lucas@nomadic-labs.com>
Date: Wed, 12 Mar 2025 17:41:29 +0100
Subject: [PATCH 6/7] Proto/consensus: activate pending companion keys

---
 .../lib_protocol/delegate_consensus_key.ml            | 11 ++++++++++-
 .../lib_protocol/delegate_consensus_key.mli           |  2 +-
 2 files changed, 11 insertions(+), 2 deletions(-)

diff --git a/src/proto_alpha/lib_protocol/delegate_consensus_key.ml b/src/proto_alpha/lib_protocol/delegate_consensus_key.ml
index d80a4b650e33..5d9bd8fb6dcc 100644
--- a/src/proto_alpha/lib_protocol/delegate_consensus_key.ml
+++ b/src/proto_alpha/lib_protocol/delegate_consensus_key.ml
@@ -328,4 +328,13 @@ let activate ctxt ~new_cycle =
       ~f:(fun delegate pk ctxt ->
         Storage.Contract.Consensus_key.add ctxt delegate pk)
   in
-  Storage.Pending_consensus_keys.clear (ctxt, new_cycle)
+  let* ctxt = Storage.Pending_consensus_keys.clear (ctxt, new_cycle) in
+  let* ctxt =
+    Storage.Pending_companion_keys.fold
+      (ctxt, new_cycle)
+      ~order:`Undefined
+      ~init:ctxt
+      ~f:(fun delegate pk ctxt ->
+        Storage.Contract.Companion_key.add ctxt delegate pk)
+  in
+  Storage.Pending_companion_keys.clear (ctxt, new_cycle)
diff --git a/src/proto_alpha/lib_protocol/delegate_consensus_key.mli b/src/proto_alpha/lib_protocol/delegate_consensus_key.mli
index 55355f6b211f..41348c23c75a 100644
--- a/src/proto_alpha/lib_protocol/delegate_consensus_key.mli
+++ b/src/proto_alpha/lib_protocol/delegate_consensus_key.mli
@@ -110,5 +110,5 @@ val register_update_companion :
   Bls.Public_key.t ->
   Raw_context.t tzresult Lwt.t
 
-(** Activate consensus keys at the beginning of cycle [new_cycle]. *)
+(** Activate consensus and companion keys at the beginning of cycle [new_cycle]. *)
 val activate : Raw_context.t -> new_cycle:Cycle_repr.t -> Raw_context.t Lwt.t
-- 
GitLab


From f51983553579267227a702b658bef2a20ab7c960 Mon Sep 17 00:00:00 2001
From: Lucas Randazzo <lucas@nomadic-labs.com>
Date: Fri, 14 Mar 2025 14:31:09 +0100
Subject: [PATCH 7/7] Proto: update changelog

---
 docs/protocols/alpha.rst | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/docs/protocols/alpha.rst b/docs/protocols/alpha.rst
index 4a4795cf36e3..d221f48d6fab 100644
--- a/docs/protocols/alpha.rst
+++ b/docs/protocols/alpha.rst
@@ -64,3 +64,8 @@ Internal
 --------
 
 - Added ``octez_riscv_pvm`` as a dependency for the protocol environment (:gl:`!15918`)
+
+- Added ``companion_key`` for delegates. These keys are used to build attestations for tz4
+  addresses, and are registered the same way as consensus keys. They are used at validation
+  time for the consensus, so the type of the storage ``delegate_sampler_state`` had to be
+  updated. (MR :gl:`!17213`)
-- 
GitLab