Alpha: maintain the set of delegates with frozen balance

src/proto_alpha/lib_protocol/src/delegate_storage.ml

@@ -345,11 +345,13 @@ let get_frozen_deposit ctxt contract cycle =
   | None -> return Tez_repr.zero
   | Some frozen -> return frozen
 
-let credit_frozen_deposit ctxt contract cycle amount =
+let credit_frozen_deposit ctxt delegate cycle amount =
+  let contract = Contract_repr.implicit_contract delegate in
   get_frozen_deposit ctxt contract cycle >>=? fun old_amount ->
   Lwt.return Tez_repr.(old_amount +? amount) >>=? fun new_amount ->
   Storage.Contract.Frozen_deposits.init_set
     (ctxt, contract) cycle new_amount >>= fun ctxt ->
+  Storage.Delegates_with_frozen_balance.add (ctxt, cycle) delegate >>= fun ctxt ->
   return ctxt
 
 let freeze_deposit ctxt delegate amount =
@@ -361,25 +363,26 @@ let freeze_deposit ctxt delegate amount =
     (record_trace (Balance_too_low_for_deposit { delegate; deposit = amount; balance })
        Tez_repr.(balance -? amount)) >>=? fun new_balance ->
   Storage.Contract.Balance.set ctxt contract new_balance >>=? fun ctxt ->
-  credit_frozen_deposit ctxt contract cycle amount
+  credit_frozen_deposit ctxt delegate cycle amount
 
 let get_frozen_fees ctxt contract cycle =
   Storage.Contract.Frozen_fees.get_option (ctxt, contract) cycle >>=? function
   | None -> return Tez_repr.zero
   | Some frozen -> return frozen
 
-let credit_frozen_fees ctxt contract cycle amount =
+let credit_frozen_fees ctxt delegate cycle amount =
+  let contract = Contract_repr.implicit_contract delegate in
   get_frozen_fees ctxt contract cycle >>=? fun old_amount ->
   Lwt.return Tez_repr.(old_amount +? amount) >>=? fun new_amount ->
   Storage.Contract.Frozen_fees.init_set
     (ctxt, contract) cycle new_amount >>= fun ctxt ->
+  Storage.Delegates_with_frozen_balance.add (ctxt, cycle) delegate >>= fun ctxt ->
   return ctxt
 
 let freeze_fees ctxt delegate amount =
   let { Level_repr.cycle ; _ } = Level_storage.current ctxt in
-  let contract = Contract_repr.implicit_contract delegate in
   Roll_storage.Delegate.add_amount ctxt delegate amount >>=? fun ctxt ->
-  credit_frozen_fees ctxt contract cycle amount
+  credit_frozen_fees ctxt delegate cycle amount
 
 let burn_fees ctxt delegate cycle amount =
   let contract = Contract_repr.implicit_contract delegate in
@@ -403,17 +406,18 @@ let get_frozen_rewards ctxt contract cycle =
   | None -> return Tez_repr.zero
   | Some frozen -> return frozen
 
-let credit_frozen_rewards ctxt contract cycle amount =
+let credit_frozen_rewards ctxt delegate cycle amount =
+  let contract = Contract_repr.implicit_contract delegate in
   get_frozen_rewards ctxt contract cycle >>=? fun old_amount ->
   Lwt.return Tez_repr.(old_amount +? amount) >>=? fun new_amount ->
   Storage.Contract.Frozen_rewards.init_set
     (ctxt, contract) cycle new_amount >>= fun ctxt ->
+  Storage.Delegates_with_frozen_balance.add (ctxt, cycle) delegate >>= fun ctxt ->
   return ctxt
 
 let freeze_rewards ctxt delegate amount =
   let { Level_repr.cycle ; _ } = Level_storage.current ctxt in
-  let contract = Contract_repr.implicit_contract delegate in
-  credit_frozen_rewards ctxt contract cycle amount
+  credit_frozen_rewards ctxt delegate cycle amount
 
 let burn_rewards ctxt delegate cycle amount =
   let contract = Contract_repr.implicit_contract delegate in
@@ -467,13 +471,14 @@ let cycle_end ctxt last_cycle unrevealed =
   match Cycle_repr.sub last_cycle preserved with
   | None -> return (ctxt, balance_updates, [])
   | Some unfrozen_cycle ->
-      Storage.Delegates.fold ctxt
+      Storage.Delegates_with_frozen_balance.fold (ctxt, unfrozen_cycle)
         ~init:(Ok (ctxt, balance_updates))
         ~f:(fun delegate acc ->
             Lwt.return acc >>=? fun (ctxt, bus) ->
             unfreeze ctxt
               delegate unfrozen_cycle >>=? fun (ctxt, balance_updates) ->
             return (ctxt, balance_updates @ bus)) >>=? fun (ctxt, balance_updates) ->
+      Storage.Delegates_with_frozen_balance.clear (ctxt, unfrozen_cycle) >>= fun ctxt ->
       Storage.Active_delegates_with_rolls.fold ctxt
         ~init:(Ok (ctxt, []))
         ~f:(fun delegate acc ->

src/proto_alpha/lib_protocol/src/storage.ml

@@ -311,6 +311,17 @@ module Active_delegates_with_rolls =
     (Make_subcontext(Raw_context)(struct let name = ["active_delegates_with_rolls"] end))
     (Make_index(Signature.Public_key_hash))
 
+module Delegates_with_frozen_balance_index =
+  Make_indexed_subcontext
+    (Make_subcontext(Raw_context)
+       (struct let name = ["delegates_with_frozen_balance"] end))
+    (Make_index(Cycle_repr.Index))
+
+module Delegates_with_frozen_balance =
+  Make_data_set_storage
+    (Delegates_with_frozen_balance_index.Raw_context)
+    (Make_index(Signature.Public_key_hash))
+
 (** Rolls *)
 
 module Cycle = struct

src/proto_alpha/lib_protocol/src/storage.mli

@@ -123,7 +123,8 @@ module Contract : sig
      and type value = Tez_repr.t
      and type t := Raw_context.t
 
-  (** Frozen balance, see 'delegate_storage.mli' for more explanation *)
+  (** Frozen balance, see 'delegate_storage.mli' for more explanation.
+      Always update `Delegates_with_frozen_balance` accordingly. *)
   module Frozen_deposits : Indexed_data_storage
     with type key = Cycle_repr.t
      and type value = Tez_repr.t
@@ -225,6 +226,11 @@ module Active_delegates_with_rolls : Data_set_storage
   with type t := Raw_context.t
    and type elt = Signature.Public_key_hash.t
 
+(** Set of all the delegates with frozen rewards/bonds/fees for a given cycle. *)
+module Delegates_with_frozen_balance : Data_set_storage
+  with type t = Raw_context.t * Cycle_repr.t
+   and type elt = Signature.Public_key_hash.t
+
 (** Votes *)
 
 module Vote : sig