From 034bd8aa8b3e627e8f6b371e126feb130d34ad62 Mon Sep 17 00:00:00 2001
From: Lucas Randazzo <lucas@nomadic-labs.com>
Date: Thu, 29 Feb 2024 16:25:38 +0100
Subject: [PATCH 01/10] Proto/tests/scenarios: make Scenarios_bake for bake

---
 .../lib_protocol/test/helpers/scenario.ml     |   1 +
 .../test/helpers/scenario_bake.ml             | 335 ++++++++++++++++++
 .../test/helpers/scenario_base.ml             | 325 -----------------
 .../test/helpers/scenario_begin.ml            |   1 +
 .../lib_protocol/test/helpers/scenario_op.ml  |   1 +
 5 files changed, 338 insertions(+), 325 deletions(-)
 create mode 100644 src/proto_alpha/lib_protocol/test/helpers/scenario_bake.ml

diff --git a/src/proto_alpha/lib_protocol/test/helpers/scenario.ml b/src/proto_alpha/lib_protocol/test/helpers/scenario.ml
index 6d493167b1ea..38a2e0be669a 100644
--- a/src/proto_alpha/lib_protocol/test/helpers/scenario.ml
+++ b/src/proto_alpha/lib_protocol/test/helpers/scenario.ml
@@ -14,3 +14,4 @@ include Scenario_op
 include Scenario_dsl
 include Scenario_begin
 include Scenario_constants
+include Scenario_bake
diff --git a/src/proto_alpha/lib_protocol/test/helpers/scenario_bake.ml b/src/proto_alpha/lib_protocol/test/helpers/scenario_bake.ml
new file mode 100644
index 000000000000..af1b7913a48c
--- /dev/null
+++ b/src/proto_alpha/lib_protocol/test/helpers/scenario_bake.ml
@@ -0,0 +1,335 @@
+(*****************************************************************************)
+(*                                                                           *)
+(* SPDX-License-Identifier: MIT                                              *)
+(* Copyright (c) 2024 Nomadic Labs, <contact@nomadic-labs.com>               *)
+(*                                                                           *)
+(*****************************************************************************)
+
+open State_account
+open State
+open Scenario_dsl
+open Log_helpers
+open Scenario_base
+open Adaptive_issuance_helpers
+
+(** After baking and applying rewards in state *)
+let check_all_balances block state : unit tzresult Lwt.t =
+  let open Lwt_result_syntax in
+  let State.{account_map; total_supply; _} = state in
+  let* actual_total_supply = Context.get_total_supply (B block) in
+  let*! r1 =
+    String.Map.fold_s
+      (fun name account acc ->
+        log_debug_balance name account_map ;
+        let* () = log_debug_rpc_balance name (Implicit account.pkh) block in
+        let*! r =
+          assert_balance_check ~loc:__LOC__ (B block) name account_map
+        in
+        Assert.join_errors r acc)
+      account_map
+      Result.return_unit
+  in
+  let*! r2 =
+    Assert.equal
+      ~loc:__LOC__
+      Tez.equal
+      "Total supplies do not match"
+      Tez.pp
+      actual_total_supply
+      total_supply
+  in
+  Assert.join_errors r1 r2
+
+let check_issuance_rpc block : unit tzresult Lwt.t =
+  let open Lwt_result_syntax in
+  (* We assume one block per minute *)
+  let* rewards_per_block = Context.get_issuance_per_minute (B block) in
+  let* total_supply = Context.get_total_supply (B block) in
+  let* expected_issuance = Context.get_ai_expected_issuance (B block) in
+  let* () =
+    match expected_issuance with
+    | ei :: _ ->
+        (* We assume only the fixed portion is issued *)
+        Assert.equal_tez
+          ~loc:__LOC__
+          rewards_per_block
+          ei.baking_reward_fixed_portion
+    | _ -> failwith "expected_issuance rpc: unexpected value"
+  in
+  let* yearly_rate = Context.get_ai_current_yearly_rate (B block) in
+  let* yearly_rate_exact = Context.get_ai_current_yearly_rate_exact (B block) in
+  let yr = float_of_string yearly_rate in
+  let yre = Q.to_float yearly_rate_exact in
+  (* Precision for yearly rate is 0.001 *)
+  let* () =
+    Assert.equal
+      ~loc:__LOC__
+      (fun x y -> Float.(abs (x -. y) <= 0.001))
+      "Yearly rate (float)"
+      Format.pp_print_float
+      yr
+      yre
+  in
+  (* Divided by 525_600 minutes per year, x100 because rpc returns a pct *)
+  let issuance_from_rate =
+    Tez.(
+      mul_q total_supply Q.(div yearly_rate_exact ~$525_600_00)
+      |> of_q ~round:`Down)
+  in
+  let* () =
+    Assert.equal
+      ~loc:__LOC__
+      Tez.equal
+      "Issuance"
+      Tez.pp
+      rewards_per_block
+      issuance_from_rate
+  in
+  return_unit
+
+(** Bake a block, with the given baker and the given operations. *)
+let bake ?baker : t -> t tzresult Lwt.t =
+ fun (block, state) ->
+  let open Lwt_result_wrap_syntax in
+  let policy =
+    match baker with
+    | None -> state.baking_policy
+    | Some baker ->
+        let {pkh; _} =
+          try State.find_account baker state
+          with Not_found ->
+            Log.info
+              ~color:warning_color
+              "Invalid baker: %s not found. Aborting"
+              baker ;
+            assert false
+        in
+        Some (Block.By_account pkh)
+  in
+  let* baker, _, _, _ = Block.get_next_baker ?policy block in
+  let baker_name, {contract = baker_contract; _} =
+    State.find_account_from_pkh baker state
+  in
+  Log.info
+    ~color:time_color
+    "Baking level %d with %s"
+    (Int32.to_int (Int32.succ Block.(block.header.shell.level)))
+    baker_name ;
+  let current_cycle = Block.current_cycle block in
+  let adaptive_issuance_vote =
+    if state.activate_ai then
+      Protocol.Alpha_context.Per_block_votes.Per_block_vote_on
+    else Per_block_vote_pass
+  in
+  let* () = check_issuance_rpc block in
+  let state, operations = State.pop_pending_operations state in
+  let* block, state =
+    let* block', _metadata =
+      Block.bake_with_metadata ?policy ~adaptive_issuance_vote ~operations block
+    in
+    if state.burn_rewards then
+      (* Incremental mode *)
+      let* i =
+        Incremental.begin_construction ?policy ~adaptive_issuance_vote block
+      in
+      let* block_rewards = Context.get_issuance_per_minute (B block') in
+      let ctxt = Incremental.alpha_ctxt i in
+      let*@ context, _ =
+        Protocol.Alpha_context.Token.transfer
+          ctxt
+          (`Contract baker_contract)
+          `Burned
+          block_rewards
+      in
+      let i = Incremental.set_alpha_ctxt i context in
+      let* i = List.fold_left_es Incremental.add_operation i operations in
+      let* block = Incremental.finalize_block i in
+      let state = State.apply_burn block_rewards baker_name state in
+      return (block, state)
+    else return (block', state)
+  in
+  let* state = State.apply_rewards ~baker:baker_name block state in
+  (* First block of a new cycle *)
+  let new_current_cycle = Block.current_cycle block in
+  let* state =
+    if Protocol.Alpha_context.Cycle.(current_cycle = new_current_cycle) then
+      return state
+    else (
+      Log.info
+        ~color:time_color
+        "Cycle %d"
+        (Protocol.Alpha_context.Cycle.to_int32 new_current_cycle |> Int32.to_int) ;
+      return @@ State.apply_new_cycle new_current_cycle state)
+  in
+  (* Dawn of a new cycle *)
+  let* state =
+    if not (Block.last_block_of_cycle block) then return state
+    else State.apply_end_cycle current_cycle block state
+  in
+  let* () = check_all_balances block state in
+  return (block, state)
+
+(** Bake until a cycle is reached, using [bake] instead of [Block.bake]
+    Should be slower because checks balances at the end of every block (avoidable in some cases) *)
+let bake_until_next_cycle : t -> t tzresult Lwt.t =
+ fun (init_block, init_state) ->
+  let open Lwt_result_syntax in
+  let current_cycle = Block.current_cycle init_block in
+  let rec step (old_block, old_state) =
+    let step_cycle = Block.current_cycle old_block in
+    if Protocol.Alpha_context.Cycle.(step_cycle > current_cycle) then
+      return (old_block, old_state)
+    else
+      let* new_block, new_state = bake (old_block, old_state) in
+      step (new_block, new_state)
+  in
+  step (init_block, init_state)
+
+(* ======== Operations ======== *)
+
+(** Bake a single block *)
+let next_block =
+  exec (fun input ->
+      Log.info ~color:action_color "[Next block]" ;
+      bake input)
+
+(** Bake a single block with a specific baker *)
+let next_block_with_baker baker =
+  exec (fun input ->
+      Log.info ~color:action_color "[Next block (baker %s)]" baker ;
+      bake ~baker input)
+
+(** Bake until the end of a cycle *)
+let next_cycle_ input =
+  Log.info ~color:action_color "[Next cycle]" ;
+  bake_until_next_cycle input
+
+(** Bake until the end of a cycle *)
+let next_cycle = exec next_cycle_
+
+(** Executes an operation: f should return a new state and a list of operations, which are then applied *)
+let exec_op f =
+  let open Lwt_result_syntax in
+  Action
+    (fun ((block, _state) as input) ->
+      let* state, ops = f input in
+      let state = State.add_pending_operations ops state in
+      return (block, state))
+  --> next_block
+
+(** Waiting functions *)
+let rec wait_n_cycles n =
+  if n <= 0 then noop
+  else if n = 1 then next_cycle
+  else wait_n_cycles (n - 1) --> next_cycle
+
+let rec wait_n_blocks n =
+  if n <= 0 then noop
+  else if n = 1 then next_block
+  else wait_n_blocks (n - 1) --> next_block
+
+let wait_cycle_f_es (condition : t -> t -> bool tzresult Lwt.t) :
+    (t, t) scenarios =
+  let open Lwt_result_syntax in
+  exec (fun init_t ->
+      let rec bake_while t =
+        let* b = condition init_t t in
+        if b then return t
+        else
+          let* t = next_cycle_ t in
+          bake_while t
+      in
+      bake_while init_t)
+
+(** Waits until [condition init_t current_t] is fulfilled.
+    It is checked on the first block of every cycle. If it returns false,
+    another cycle is baked, until it succeeds.
+*)
+let wait_cycle_f (condition : t -> t -> bool) : (t, t) scenarios =
+  let open Lwt_result_syntax in
+  exec (fun init_t ->
+      let rec bake_while t =
+        if condition init_t t then return t
+        else
+          let* t = next_cycle_ t in
+          bake_while t
+      in
+      bake_while init_t)
+
+(** Wait until we are in a cycle satisfying the given condition.
+    Fails if AI_activation is requested and AI is not set to be activated in the future. *)
+let wait_cycle_until condition =
+  let to_, done_ =
+    let rec get_names condition =
+      match condition with
+      | `AI_activation -> ("AI activation", "AI activated")
+      | `delegate_parameters_activation ->
+          ("delegate parameters activation", "delegate parameters activated")
+      | `And (cond1, cond2) ->
+          let to1, done1 = get_names cond1 in
+          let to2, done2 = get_names cond2 in
+          (to1 ^ " and " ^ to2, done1 ^ " and " ^ done2)
+    in
+    get_names condition
+  in
+  let condition (init_block, init_state) =
+    let open Lwt_result_syntax in
+    let rec stopper condition =
+      match condition with
+      | `AI_activation ->
+          fun (block, _state) ->
+            if init_state.State.activate_ai then
+              let* launch_cycle = get_launch_cycle ~loc:__LOC__ init_block in
+              let current_cycle = Block.current_cycle block in
+              return Cycle.(current_cycle >= launch_cycle)
+            else assert false
+      | `delegate_parameters_activation ->
+          fun (block, _state) ->
+            let init_cycle = Block.current_cycle init_block in
+            let cycles_to_wait =
+              init_state.constants.delegate_parameters_activation_delay
+            in
+            return
+              Cycle.(Block.current_cycle block >= add init_cycle cycles_to_wait)
+      | `And (cond1, cond2) ->
+          let stop1 = stopper cond1 in
+          let stop2 = stopper cond2 in
+          fun (block, state) ->
+            let* b1 = stop1 (block, state) in
+            let* b2 = stop2 (block, state) in
+            return (b1 && b2)
+    in
+    stopper condition
+  in
+  log ~color:time_color "Fast forward to %s" to_
+  --> wait_cycle_f_es condition
+  --> log ~color:event_color "%s" done_
+
+(** Wait until AI activates.
+    Fails if AI is not set to be activated in the future. *)
+let wait_ai_activation = wait_cycle_until `AI_activation
+
+(** wait delegate_parameters_activation_delay cycles  *)
+let wait_delegate_parameters_activation =
+  wait_cycle_until `delegate_parameters_activation
+
+let wait_n_cycles_f (n_cycles : t -> int) =
+  let condition ((init_block, _init_state) as t_init)
+      ((current_block, _current_state) as _t_current) =
+    let n = n_cycles t_init in
+    let init_cycle = Block.current_cycle init_block in
+    let current_cycle = Block.current_cycle current_block in
+    Cycle.(current_cycle >= add init_cycle n)
+  in
+  wait_cycle_f condition
+
+let wait_n_cycles_f_es (n_cycles : t -> int tzresult Lwt.t) =
+  let open Lwt_result_syntax in
+  let condition ((init_block, _init_state) as t_init)
+      ((current_block, _current_state) as _t_current) =
+    let* n = n_cycles t_init in
+    let init_cycle = Block.current_cycle init_block in
+    let current_cycle = Block.current_cycle current_block in
+    return Cycle.(current_cycle >= add init_cycle n)
+  in
+  wait_cycle_f_es condition
diff --git a/src/proto_alpha/lib_protocol/test/helpers/scenario_base.ml b/src/proto_alpha/lib_protocol/test/helpers/scenario_base.ml
index cc461201c2a3..d29abf488145 100644
--- a/src/proto_alpha/lib_protocol/test/helpers/scenario_base.ml
+++ b/src/proto_alpha/lib_protocol/test/helpers/scenario_base.ml
@@ -13,7 +13,6 @@ open State_account
 open State
 open Scenario_dsl
 open Log_helpers
-open Adaptive_issuance_helpers
 
 (** For [assert_failure], when expected error does not match the actual error. *)
 type error += Unexpected_error
@@ -30,181 +29,6 @@ let log ?(level = Cli.Logs.Info) ?color format =
           return_unit))
     format
 
-(* ======== Baking ======== *)
-
-(** After baking and applying rewards in state *)
-let check_all_balances block state : unit tzresult Lwt.t =
-  let open Lwt_result_syntax in
-  let State.{account_map; total_supply; _} = state in
-  let* actual_total_supply = Context.get_total_supply (B block) in
-  let*! r1 =
-    String.Map.fold_s
-      (fun name account acc ->
-        log_debug_balance name account_map ;
-        let* () = log_debug_rpc_balance name (Implicit account.pkh) block in
-        let*! r =
-          assert_balance_check ~loc:__LOC__ (B block) name account_map
-        in
-        Assert.join_errors r acc)
-      account_map
-      Result.return_unit
-  in
-  let*! r2 =
-    Assert.equal
-      ~loc:__LOC__
-      Tez.equal
-      "Total supplies do not match"
-      Tez.pp
-      actual_total_supply
-      total_supply
-  in
-  Assert.join_errors r1 r2
-
-let check_issuance_rpc block : unit tzresult Lwt.t =
-  let open Lwt_result_syntax in
-  (* We assume one block per minute *)
-  let* rewards_per_block = Context.get_issuance_per_minute (B block) in
-  let* total_supply = Context.get_total_supply (B block) in
-  let* expected_issuance = Context.get_ai_expected_issuance (B block) in
-  let* () =
-    match expected_issuance with
-    | ei :: _ ->
-        (* We assume only the fixed portion is issued *)
-        Assert.equal_tez
-          ~loc:__LOC__
-          rewards_per_block
-          ei.baking_reward_fixed_portion
-    | _ -> failwith "expected_issuance rpc: unexpected value"
-  in
-  let* yearly_rate = Context.get_ai_current_yearly_rate (B block) in
-  let* yearly_rate_exact = Context.get_ai_current_yearly_rate_exact (B block) in
-  let yr = float_of_string yearly_rate in
-  let yre = Q.to_float yearly_rate_exact in
-  (* Precision for yearly rate is 0.001 *)
-  let* () =
-    Assert.equal
-      ~loc:__LOC__
-      (fun x y -> Float.(abs (x -. y) <= 0.001))
-      "Yearly rate (float)"
-      Format.pp_print_float
-      yr
-      yre
-  in
-  (* Divided by 525_600 minutes per year, x100 because rpc returns a pct *)
-  let issuance_from_rate =
-    Tez.(
-      mul_q total_supply Q.(div yearly_rate_exact ~$525_600_00)
-      |> of_q ~round:`Down)
-  in
-  let* () =
-    Assert.equal
-      ~loc:__LOC__
-      Tez.equal
-      "Issuance"
-      Tez.pp
-      rewards_per_block
-      issuance_from_rate
-  in
-  return_unit
-
-(** Bake a block, with the given baker and the given operations. *)
-let bake ?baker : t -> t tzresult Lwt.t =
- fun (block, state) ->
-  let open Lwt_result_wrap_syntax in
-  let policy =
-    match baker with
-    | None -> state.baking_policy
-    | Some baker ->
-        let {pkh; _} =
-          try State.find_account baker state
-          with Not_found ->
-            Log.info
-              ~color:warning_color
-              "Invalid baker: %s not found. Aborting"
-              baker ;
-            assert false
-        in
-        Some (Block.By_account pkh)
-  in
-  let* baker, _, _, _ = Block.get_next_baker ?policy block in
-  let baker_name, {contract = baker_contract; _} =
-    State.find_account_from_pkh baker state
-  in
-  Log.info
-    ~color:time_color
-    "Baking level %d with %s"
-    (Int32.to_int (Int32.succ Block.(block.header.shell.level)))
-    baker_name ;
-  let current_cycle = Block.current_cycle block in
-  let adaptive_issuance_vote =
-    if state.activate_ai then
-      Protocol.Alpha_context.Per_block_votes.Per_block_vote_on
-    else Per_block_vote_pass
-  in
-  let* () = check_issuance_rpc block in
-  let state, operations = State.pop_pending_operations state in
-  let* block, state =
-    let* block', _metadata =
-      Block.bake_with_metadata ?policy ~adaptive_issuance_vote ~operations block
-    in
-    if state.burn_rewards then
-      (* Incremental mode *)
-      let* i =
-        Incremental.begin_construction ?policy ~adaptive_issuance_vote block
-      in
-      let* block_rewards = Context.get_issuance_per_minute (B block') in
-      let ctxt = Incremental.alpha_ctxt i in
-      let*@ context, _ =
-        Protocol.Alpha_context.Token.transfer
-          ctxt
-          (`Contract baker_contract)
-          `Burned
-          block_rewards
-      in
-      let i = Incremental.set_alpha_ctxt i context in
-      let* i = List.fold_left_es Incremental.add_operation i operations in
-      let* block = Incremental.finalize_block i in
-      let state = State.apply_burn block_rewards baker_name state in
-      return (block, state)
-    else return (block', state)
-  in
-  let* state = State.apply_rewards ~baker:baker_name block state in
-  (* First block of a new cycle *)
-  let new_current_cycle = Block.current_cycle block in
-  let* state =
-    if Protocol.Alpha_context.Cycle.(current_cycle = new_current_cycle) then
-      return state
-    else (
-      Log.info
-        ~color:time_color
-        "Cycle %d"
-        (Protocol.Alpha_context.Cycle.to_int32 new_current_cycle |> Int32.to_int) ;
-      return @@ State.apply_new_cycle new_current_cycle state)
-  in
-  (* Dawn of a new cycle *)
-  let* state =
-    if not (Block.last_block_of_cycle block) then return state
-    else State.apply_end_cycle current_cycle block state
-  in
-  let* () = check_all_balances block state in
-  return (block, state)
-
-(** Bake until a cycle is reached, using [bake] instead of [Block.bake]
-    Should be slower because checks balances at the end of every block (avoidable in some cases) *)
-let bake_until_next_cycle : t -> t tzresult Lwt.t =
- fun (init_block, init_state) ->
-  let open Lwt_result_syntax in
-  let current_cycle = Block.current_cycle init_block in
-  let rec step (old_block, old_state) =
-    let step_cycle = Block.current_cycle old_block in
-    if Protocol.Alpha_context.Cycle.(step_cycle > current_cycle) then
-      return (old_block, old_state)
-    else
-      let* new_block, new_state = bake (old_block, old_state) in
-      step (new_block, new_state)
-  in
-  step (init_block, init_state)
-
 (* ======== State updates ======== *)
 
 (** Sets the de facto baker for all future blocks *)
@@ -311,38 +135,6 @@ let check_rate_evolution (f : Q.t -> Q.t -> bool) : (t, t) scenarios =
               Q.pp_print
               new_rate)
 
-(* ======== Operations ======== *)
-
-(** Bake a single block *)
-let next_block =
-  exec (fun input ->
-      Log.info ~color:action_color "[Next block]" ;
-      bake input)
-
-(** Bake a single block with a specific baker *)
-let next_block_with_baker baker =
-  exec (fun input ->
-      Log.info ~color:action_color "[Next block (baker %s)]" baker ;
-      bake ~baker input)
-
-(** Bake until the end of a cycle *)
-let next_cycle_ input =
-  Log.info ~color:action_color "[Next cycle]" ;
-  bake_until_next_cycle input
-
-(** Bake until the end of a cycle *)
-let next_cycle = exec next_cycle_
-
-(** Executes an operation: f should return a new state and a list of operations, which are then applied *)
-let exec_op f =
-  let open Lwt_result_syntax in
-  Action
-    (fun ((block, _state) as input) ->
-      let* state, ops = f input in
-      let state = State.add_pending_operations ops state in
-      return (block, state))
-  --> next_block
-
 (* ======== Misc functions ========*)
 
 let check_failure_aux ?expected_error :
@@ -437,120 +229,3 @@ let check_balance_field src_name field amount : (t, t) scenarios =
             check src_total
       in
       return state)
-
-(** Waiting functions *)
-let rec wait_n_cycles n =
-  if n <= 0 then noop
-  else if n = 1 then next_cycle
-  else wait_n_cycles (n - 1) --> next_cycle
-
-let rec wait_n_blocks n =
-  if n <= 0 then noop
-  else if n = 1 then next_block
-  else wait_n_blocks (n - 1) --> next_block
-
-let wait_cycle_f_es (condition : t -> t -> bool tzresult Lwt.t) :
-    (t, t) scenarios =
-  let open Lwt_result_syntax in
-  exec (fun init_t ->
-      let rec bake_while t =
-        let* b = condition init_t t in
-        if b then return t
-        else
-          let* t = next_cycle_ t in
-          bake_while t
-      in
-      bake_while init_t)
-
-(** Waits until [condition init_t current_t] is fulfilled.
-    It is checked on the first block of every cycle. If it returns false,
-    another cycle is baked, until it succeeds.
-*)
-let wait_cycle_f (condition : t -> t -> bool) : (t, t) scenarios =
-  let open Lwt_result_syntax in
-  exec (fun init_t ->
-      let rec bake_while t =
-        if condition init_t t then return t
-        else
-          let* t = next_cycle_ t in
-          bake_while t
-      in
-      bake_while init_t)
-
-(** Wait until we are in a cycle satisfying the given condition.
-    Fails if AI_activation is requested and AI is not set to be activated in the future. *)
-let wait_cycle_until condition =
-  let to_, done_ =
-    let rec get_names condition =
-      match condition with
-      | `AI_activation -> ("AI activation", "AI activated")
-      | `delegate_parameters_activation ->
-          ("delegate parameters activation", "delegate parameters activated")
-      | `And (cond1, cond2) ->
-          let to1, done1 = get_names cond1 in
-          let to2, done2 = get_names cond2 in
-          (to1 ^ " and " ^ to2, done1 ^ " and " ^ done2)
-    in
-    get_names condition
-  in
-  let condition (init_block, init_state) =
-    let open Lwt_result_syntax in
-    let rec stopper condition =
-      match condition with
-      | `AI_activation ->
-          fun (block, _state) ->
-            if init_state.State.activate_ai then
-              let* launch_cycle = get_launch_cycle ~loc:__LOC__ init_block in
-              let current_cycle = Block.current_cycle block in
-              return Cycle.(current_cycle >= launch_cycle)
-            else assert false
-      | `delegate_parameters_activation ->
-          fun (block, _state) ->
-            let init_cycle = Block.current_cycle init_block in
-            let cycles_to_wait =
-              init_state.constants.delegate_parameters_activation_delay
-            in
-            return
-              Cycle.(Block.current_cycle block >= add init_cycle cycles_to_wait)
-      | `And (cond1, cond2) ->
-          let stop1 = stopper cond1 in
-          let stop2 = stopper cond2 in
-          fun (block, state) ->
-            let* b1 = stop1 (block, state) in
-            let* b2 = stop2 (block, state) in
-            return (b1 && b2)
-    in
-    stopper condition
-  in
-  log ~color:time_color "Fast forward to %s" to_
-  --> wait_cycle_f_es condition
-  --> log ~color:event_color "%s" done_
-
-(** Wait until AI activates.
-    Fails if AI is not set to be activated in the future. *)
-let wait_ai_activation = wait_cycle_until `AI_activation
-
-(** wait delegate_parameters_activation_delay cycles  *)
-let wait_delegate_parameters_activation =
-  wait_cycle_until `delegate_parameters_activation
-
-let wait_n_cycles_f (n_cycles : t -> int) =
-  let condition ((init_block, _init_state) as t_init)
-      ((current_block, _current_state) as _t_current) =
-    let n = n_cycles t_init in
-    let init_cycle = Block.current_cycle init_block in
-    let current_cycle = Block.current_cycle current_block in
-    Cycle.(current_cycle >= add init_cycle n)
-  in
-  wait_cycle_f condition
-
-let wait_n_cycles_f_es (n_cycles : t -> int tzresult Lwt.t) =
-  let open Lwt_result_syntax in
-  let condition ((init_block, _init_state) as t_init)
-      ((current_block, _current_state) as _t_current) =
-    let* n = n_cycles t_init in
-    let init_cycle = Block.current_cycle init_block in
-    let current_cycle = Block.current_cycle current_block in
-    return Cycle.(current_cycle >= add init_cycle n)
-  in
-  wait_cycle_f_es condition
diff --git a/src/proto_alpha/lib_protocol/test/helpers/scenario_begin.ml b/src/proto_alpha/lib_protocol/test/helpers/scenario_begin.ml
index 5c21b92705f0..8fedb8d9901a 100644
--- a/src/proto_alpha/lib_protocol/test/helpers/scenario_begin.ml
+++ b/src/proto_alpha/lib_protocol/test/helpers/scenario_begin.ml
@@ -7,6 +7,7 @@
 
 open State_account
 open Scenario_dsl
+open Scenario_bake
 open Scenario_base
 open Log_helpers
 open Adaptive_issuance_helpers
diff --git a/src/proto_alpha/lib_protocol/test/helpers/scenario_op.ml b/src/proto_alpha/lib_protocol/test/helpers/scenario_op.ml
index 8c684e882f45..bedba04f7bba 100644
--- a/src/proto_alpha/lib_protocol/test/helpers/scenario_op.ml
+++ b/src/proto_alpha/lib_protocol/test/helpers/scenario_op.ml
@@ -14,6 +14,7 @@ open State_account
 open Adaptive_issuance_helpers
 open Scenario_dsl
 open Scenario_base
+open Scenario_bake
 open Tez_helpers.Ez_tez
 
 (** Set delegate parameters for the given delegate *)
-- 
GitLab


From fddb6919c1494030b25eb938a0ee31a6783fb896 Mon Sep 17 00:00:00 2001
From: Lucas Randazzo <lucas@nomadic-labs.com>
Date: Tue, 5 Mar 2024 11:50:24 +0100
Subject: [PATCH 02/10] Proto/tests/scenarios: new files for flags logic

---
 src/proto_alpha/lib_protocol/test/helpers/state_ai_flags.ml | 6 ++++++
 .../lib_protocol/test/helpers/state_ai_flags.mli            | 6 ++++++
 2 files changed, 12 insertions(+)
 create mode 100644 src/proto_alpha/lib_protocol/test/helpers/state_ai_flags.ml
 create mode 100644 src/proto_alpha/lib_protocol/test/helpers/state_ai_flags.mli

diff --git a/src/proto_alpha/lib_protocol/test/helpers/state_ai_flags.ml b/src/proto_alpha/lib_protocol/test/helpers/state_ai_flags.ml
new file mode 100644
index 000000000000..0a966492a536
--- /dev/null
+++ b/src/proto_alpha/lib_protocol/test/helpers/state_ai_flags.ml
@@ -0,0 +1,6 @@
+(*****************************************************************************)
+(*                                                                           *)
+(* SPDX-License-Identifier: MIT                                              *)
+(* Copyright (c) 2024 Nomadic Labs, <contact@nomadic-labs.com>               *)
+(*                                                                           *)
+(*****************************************************************************)
diff --git a/src/proto_alpha/lib_protocol/test/helpers/state_ai_flags.mli b/src/proto_alpha/lib_protocol/test/helpers/state_ai_flags.mli
new file mode 100644
index 000000000000..0a966492a536
--- /dev/null
+++ b/src/proto_alpha/lib_protocol/test/helpers/state_ai_flags.mli
@@ -0,0 +1,6 @@
+(*****************************************************************************)
+(*                                                                           *)
+(* SPDX-License-Identifier: MIT                                              *)
+(* Copyright (c) 2024 Nomadic Labs, <contact@nomadic-labs.com>               *)
+(*                                                                           *)
+(*****************************************************************************)
-- 
GitLab


From b0961c467be70e5519bff40af7f7af6dbff11ebc Mon Sep 17 00:00:00 2001
From: Lucas Randazzo <lucas@nomadic-labs.com>
Date: Tue, 5 Mar 2024 13:52:18 +0100
Subject: [PATCH 03/10] Proto/tests: AI_Activation module

---
 .../test/helpers/scenario_bake.ml             |  3 ++
 .../test/helpers/scenario_begin.ml            |  4 ++
 .../lib_protocol/test/helpers/state.ml        |  1 +
 .../test/helpers/state_ai_flags.ml            | 51 +++++++++++++++++++
 .../test/helpers/state_ai_flags.mli           | 19 +++++++
 5 files changed, 78 insertions(+)

diff --git a/src/proto_alpha/lib_protocol/test/helpers/scenario_bake.ml b/src/proto_alpha/lib_protocol/test/helpers/scenario_bake.ml
index af1b7913a48c..fc8ce9340cbc 100644
--- a/src/proto_alpha/lib_protocol/test/helpers/scenario_bake.ml
+++ b/src/proto_alpha/lib_protocol/test/helpers/scenario_bake.ml
@@ -148,6 +148,9 @@ let bake ?baker : t -> t tzresult Lwt.t =
       return (block, state)
     else return (block', state)
   in
+  let* state =
+    State_ai_flags.AI_Activation.check_activation_cycle block state
+  in
   let* state = State.apply_rewards ~baker:baker_name block state in
   (* First block of a new cycle *)
   let new_current_cycle = Block.current_cycle block in
diff --git a/src/proto_alpha/lib_protocol/test/helpers/scenario_begin.ml b/src/proto_alpha/lib_protocol/test/helpers/scenario_begin.ml
index 8fedb8d9901a..2a20a38e6985 100644
--- a/src/proto_alpha/lib_protocol/test/helpers/scenario_begin.ml
+++ b/src/proto_alpha/lib_protocol/test/helpers/scenario_begin.ml
@@ -152,7 +152,11 @@ let begin_test ?(burn_rewards = false) delegates_name_list :
             pending_operations = [];
             pending_slashes = [];
             double_signings = [];
+            ai_activation_cycle = None;
           }
       in
       let* () = check_all_balances block state in
+      let* state =
+        State_ai_flags.AI_Activation.check_activation_cycle block state
+      in
       return (block, state))
diff --git a/src/proto_alpha/lib_protocol/test/helpers/state.ml b/src/proto_alpha/lib_protocol/test/helpers/state.ml
index 75671b168efa..29ddc78d480d 100644
--- a/src/proto_alpha/lib_protocol/test/helpers/state.ml
+++ b/src/proto_alpha/lib_protocol/test/helpers/state.ml
@@ -32,6 +32,7 @@ type t = {
   pending_slashes :
     (Signature.Public_key_hash.t * Protocol.Denunciations_repr.item) list;
   double_signings : double_signing_state list;
+  ai_activation_cycle : Protocol.Alpha_context.Cycle.t option;
 }
 
 (** Expected number of cycles before staking parameters get applied *)
diff --git a/src/proto_alpha/lib_protocol/test/helpers/state_ai_flags.ml b/src/proto_alpha/lib_protocol/test/helpers/state_ai_flags.ml
index 0a966492a536..fd99e2f95d17 100644
--- a/src/proto_alpha/lib_protocol/test/helpers/state_ai_flags.ml
+++ b/src/proto_alpha/lib_protocol/test/helpers/state_ai_flags.ml
@@ -4,3 +4,54 @@
 (* Copyright (c) 2024 Nomadic Labs, <contact@nomadic-labs.com>               *)
 (*                                                                           *)
 (*****************************************************************************)
+
+(** This module handles the logic of flags for AI/AS in the [State] *)
+
+open State
+
+module AI_Activation = struct
+  (** This module is responsible for the field [state.ai_activation_cycle],
+      which depends on three protocol parameters:
+      [adaptive_issuance.force_activation], [adaptive_issuance.activation_vote_enable], and
+      [adaptive_issuance.launch_ema_threshold]. *)
+
+  (** AI can be activated with both flags set to false if the threshold is set to 0 *)
+  let enabled state =
+    state.constants.adaptive_issuance.force_activation
+    || state.constants.adaptive_issuance.activation_vote_enable
+    || Compare.Int32.(
+         state.constants.adaptive_issuance.launch_ema_threshold = 0l)
+
+  let set_activation_cycle block state block_launch_cycle =
+    let current_cycle = Block.current_cycle block in
+    let offset =
+      if state.constants.adaptive_issuance.force_activation then 0
+      else
+        1 + state.constants.consensus_rights_delay
+        + Protocol.Constants_repr.max_slashing_period
+    in
+    assert (
+      Protocol.Alpha_context.Cycle.(
+        add current_cycle offset = block_launch_cycle)) ;
+    {state with ai_activation_cycle = Some block_launch_cycle}
+
+  (** Check the activation_cycle is only ever set once.
+      Run every block *)
+  let check_activation_cycle block state =
+    let open Lwt_result_syntax in
+    let open Protocol.Alpha_context in
+    let* block_launch_cycle =
+      Context.get_adaptive_issuance_launch_cycle (B block)
+    in
+    match (enabled state, state.ai_activation_cycle, block_launch_cycle) with
+    | _, None, None -> return state
+    | true, Some x, Some y ->
+        (* Activation cycle cannot be changed *)
+        if Cycle.(x = y) then return state else assert false
+    | _, Some _, None -> (* Activation cycle cannot be unset *) assert false
+    | false, _, Some _ ->
+        (* AI cannot be activated if [enabled] is false *)
+        assert false
+    | true, None, Some block_launch_cycle ->
+        return @@ set_activation_cycle block state block_launch_cycle
+end
diff --git a/src/proto_alpha/lib_protocol/test/helpers/state_ai_flags.mli b/src/proto_alpha/lib_protocol/test/helpers/state_ai_flags.mli
index 0a966492a536..26361dbb865a 100644
--- a/src/proto_alpha/lib_protocol/test/helpers/state_ai_flags.mli
+++ b/src/proto_alpha/lib_protocol/test/helpers/state_ai_flags.mli
@@ -4,3 +4,22 @@
 (* Copyright (c) 2024 Nomadic Labs, <contact@nomadic-labs.com>               *)
 (*                                                                           *)
 (*****************************************************************************)
+
+(** This module handles the logic of flags for AI/AS in the [State] *)
+
+module AI_Activation : sig
+  (** This module takes care of the flags [force_activation] and
+      [activation_vote_enable], and updates and check the record field
+      [ai_activation_cycle] in the state. *)
+
+  (** [AI_Activation] is enabled iff either of the above flags are set to true,
+      or the vote threshold is set to 0 (regardless of votes).
+      "Enabled" here means that AI *can* be activated (either by vote
+      or by force), but does not mean that AI is activated *)
+  val enabled : State.t -> bool
+
+  (** Checks the [ai_activation_cycle] is set as expected.
+      Run at the beginning, and after every block to check the activation
+      cycle is set only once ever. *)
+  val check_activation_cycle : Block.t -> State.t -> State.t tzresult Lwt.t
+end
-- 
GitLab


From 2cfbddee5ef211035026c27b026d86674caec34c Mon Sep 17 00:00:00 2001
From: Lucas Randazzo <lucas@nomadic-labs.com>
Date: Tue, 5 Mar 2024 14:31:03 +0100
Subject: [PATCH 04/10] Proto/tests: AI check module

---
 .../lib_protocol/test/helpers/scenario_bake.ml           | 6 +++++-
 .../lib_protocol/test/helpers/state_ai_flags.ml          | 9 +++++++++
 .../lib_protocol/test/helpers/state_ai_flags.mli         | 7 +++++++
 3 files changed, 21 insertions(+), 1 deletion(-)

diff --git a/src/proto_alpha/lib_protocol/test/helpers/scenario_bake.ml b/src/proto_alpha/lib_protocol/test/helpers/scenario_bake.ml
index fc8ce9340cbc..2f0aa9b686fc 100644
--- a/src/proto_alpha/lib_protocol/test/helpers/scenario_bake.ml
+++ b/src/proto_alpha/lib_protocol/test/helpers/scenario_bake.ml
@@ -310,7 +310,11 @@ let wait_cycle_until condition =
 
 (** Wait until AI activates.
     Fails if AI is not set to be activated in the future. *)
-let wait_ai_activation = wait_cycle_until `AI_activation
+let wait_ai_activation =
+  wait_cycle_until `AI_activation
+  --> exec_unit (fun (block, state) ->
+          assert (State_ai_flags.AI.enabled block state) ;
+          return_unit)
 
 (** wait delegate_parameters_activation_delay cycles  *)
 let wait_delegate_parameters_activation =
diff --git a/src/proto_alpha/lib_protocol/test/helpers/state_ai_flags.ml b/src/proto_alpha/lib_protocol/test/helpers/state_ai_flags.ml
index fd99e2f95d17..e191a105cb3c 100644
--- a/src/proto_alpha/lib_protocol/test/helpers/state_ai_flags.ml
+++ b/src/proto_alpha/lib_protocol/test/helpers/state_ai_flags.ml
@@ -55,3 +55,12 @@ module AI_Activation = struct
     | true, None, Some block_launch_cycle ->
         return @@ set_activation_cycle block state block_launch_cycle
 end
+
+module AI = struct
+  let enabled (block : Block.t) (state : State.t) =
+    match state.ai_activation_cycle with
+    | None -> false
+    | Some activation_cycle ->
+        let current_cycle = Block.current_cycle block in
+        Protocol.Alpha_context.Cycle.(current_cycle >= activation_cycle)
+end
diff --git a/src/proto_alpha/lib_protocol/test/helpers/state_ai_flags.mli b/src/proto_alpha/lib_protocol/test/helpers/state_ai_flags.mli
index 26361dbb865a..20f3df44813d 100644
--- a/src/proto_alpha/lib_protocol/test/helpers/state_ai_flags.mli
+++ b/src/proto_alpha/lib_protocol/test/helpers/state_ai_flags.mli
@@ -23,3 +23,10 @@ module AI_Activation : sig
       cycle is set only once ever. *)
   val check_activation_cycle : Block.t -> State.t -> State.t tzresult Lwt.t
 end
+
+module AI : sig
+  (** This module only checks (for now) if AI is activated or not *)
+
+  (** AI is enabled iff the activation cycle is set and passed *)
+  val enabled : Block.t -> State.t -> bool
+end
-- 
GitLab


From bb9414e1defcd7026ab1e91f02a49621cc18ce90 Mon Sep 17 00:00:00 2001
From: Lucas Randazzo <lucas@nomadic-labs.com>
Date: Tue, 5 Mar 2024 14:37:36 +0100
Subject: [PATCH 05/10] Proto/tests: move end_cycle and new_cycle in bake

---
 .../test/helpers/scenario_bake.ml             | 58 ++++++++++++++++++-
 .../lib_protocol/test/helpers/state.ml        | 54 -----------------
 2 files changed, 56 insertions(+), 56 deletions(-)

diff --git a/src/proto_alpha/lib_protocol/test/helpers/scenario_bake.ml b/src/proto_alpha/lib_protocol/test/helpers/scenario_bake.ml
index 2f0aa9b686fc..65c3932e3958 100644
--- a/src/proto_alpha/lib_protocol/test/helpers/scenario_bake.ml
+++ b/src/proto_alpha/lib_protocol/test/helpers/scenario_bake.ml
@@ -12,6 +12,60 @@ open Log_helpers
 open Scenario_base
 open Adaptive_issuance_helpers
 
+(** Applies when baking the last block of a cycle *)
+let apply_end_cycle current_cycle block state : State.t tzresult Lwt.t =
+  let open Lwt_result_syntax in
+  Log.debug ~color:time_color "Ending cycle %a" Cycle.pp current_cycle ;
+  let* launch_cycle_opt =
+    Context.get_adaptive_issuance_launch_cycle (B block)
+  in
+  (* Apply all slashes *)
+  let state = apply_all_slashes_at_cycle_end current_cycle state in
+  (* Sets initial frozen for future cycle *)
+  let state =
+    update_map
+      ~f:
+        (update_frozen_rights_cycle
+           (Cycle.add
+              current_cycle
+              (state.constants.consensus_rights_delay + 1)))
+      state
+  in
+  (* Apply autostaking *)
+  let*? state =
+    if not state.constants.adaptive_issuance.autostaking_enable then Ok state
+    else
+      match launch_cycle_opt with
+      | Some launch_cycle when Cycle.(current_cycle >= launch_cycle) -> Ok state
+      | None | Some _ ->
+          String.Map.fold_e
+            (fun name account state ->
+              apply_autostake ~name ~old_cycle:current_cycle account state)
+            state.account_map
+            state
+  in
+  (* Apply parameter changes *)
+  let state, param_requests =
+    List.fold_left
+      (fun (state, remaining_requests) (name, params, wait) ->
+        if wait > 0 then (state, (name, params, wait - 1) :: remaining_requests)
+        else
+          let src = find_account name state in
+          let state =
+            update_account name {src with parameters = params} state
+          in
+          (state, remaining_requests))
+      (state, [])
+      state.param_requests
+  in
+  return {state with param_requests}
+
+(** Applies when baking the first block of a cycle.
+      Technically nothing special happens, but we need to update the unslashable unstakes
+      since it's done lazily *)
+let apply_new_cycle new_cycle state : State.t =
+  apply_unslashable_for_all new_cycle state
+
 (** After baking and applying rewards in state *)
 let check_all_balances block state : unit tzresult Lwt.t =
   let open Lwt_result_syntax in
@@ -162,12 +216,12 @@ let bake ?baker : t -> t tzresult Lwt.t =
         ~color:time_color
         "Cycle %d"
         (Protocol.Alpha_context.Cycle.to_int32 new_current_cycle |> Int32.to_int) ;
-      return @@ State.apply_new_cycle new_current_cycle state)
+      return @@ apply_new_cycle new_current_cycle state)
   in
   (* Dawn of a new cycle *)
   let* state =
     if not (Block.last_block_of_cycle block) then return state
-    else State.apply_end_cycle current_cycle block state
+    else apply_end_cycle current_cycle block state
   in
   let* () = check_all_balances block state in
   return (block, state)
diff --git a/src/proto_alpha/lib_protocol/test/helpers/state.ml b/src/proto_alpha/lib_protocol/test/helpers/state.ml
index 29ddc78d480d..3e8e496a8d98 100644
--- a/src/proto_alpha/lib_protocol/test/helpers/state.ml
+++ b/src/proto_alpha/lib_protocol/test/helpers/state.ml
@@ -336,57 +336,3 @@ let apply_autostake ~name ~old_cycle
         state)
     in
     return new_state
-
-(** Applies when baking the last block of a cycle *)
-let apply_end_cycle current_cycle block state : t tzresult Lwt.t =
-  let open Lwt_result_syntax in
-  Log.debug ~color:time_color "Ending cycle %a" Cycle.pp current_cycle ;
-  let* launch_cycle_opt =
-    Context.get_adaptive_issuance_launch_cycle (B block)
-  in
-  (* Apply all slashes *)
-  let state = apply_all_slashes_at_cycle_end current_cycle state in
-  (* Sets initial frozen for future cycle *)
-  let state =
-    update_map
-      ~f:
-        (update_frozen_rights_cycle
-           (Cycle.add
-              current_cycle
-              (state.constants.consensus_rights_delay + 1)))
-      state
-  in
-  (* Apply autostaking *)
-  let*? state =
-    if not state.constants.adaptive_issuance.autostaking_enable then Ok state
-    else
-      match launch_cycle_opt with
-      | Some launch_cycle when Cycle.(current_cycle >= launch_cycle) -> Ok state
-      | None | Some _ ->
-          String.Map.fold_e
-            (fun name account state ->
-              apply_autostake ~name ~old_cycle:current_cycle account state)
-            state.account_map
-            state
-  in
-  (* Apply parameter changes *)
-  let state, param_requests =
-    List.fold_left
-      (fun (state, remaining_requests) (name, params, wait) ->
-        if wait > 0 then (state, (name, params, wait - 1) :: remaining_requests)
-        else
-          let src = find_account name state in
-          let state =
-            update_account name {src with parameters = params} state
-          in
-          (state, remaining_requests))
-      (state, [])
-      state.param_requests
-  in
-  return {state with param_requests}
-
-(** Applies when baking the first block of a cycle.
-      Technically nothing special happens, but we need to update the unslashable unstakes
-      since it's done lazily *)
-let apply_new_cycle new_cycle state : t =
-  apply_unslashable_for_all new_cycle state
-- 
GitLab


From 1b4e9e25fc9e7733cf7addc2a697c9d7c2b6fbe0 Mon Sep 17 00:00:00 2001
From: Lucas Randazzo <lucas@nomadic-labs.com>
Date: Tue, 5 Mar 2024 14:53:46 +0100
Subject: [PATCH 06/10] Proto/tests: Autostake module (refactor)

---
 .../test/helpers/scenario_bake.ml             |  16 +--
 .../lib_protocol/test/helpers/state.ml        |  99 ---------------
 .../test/helpers/state_ai_flags.ml            | 119 ++++++++++++++++++
 .../test/helpers/state_ai_flags.mli           |  12 ++
 4 files changed, 132 insertions(+), 114 deletions(-)

diff --git a/src/proto_alpha/lib_protocol/test/helpers/scenario_bake.ml b/src/proto_alpha/lib_protocol/test/helpers/scenario_bake.ml
index 65c3932e3958..3275e403a119 100644
--- a/src/proto_alpha/lib_protocol/test/helpers/scenario_bake.ml
+++ b/src/proto_alpha/lib_protocol/test/helpers/scenario_bake.ml
@@ -16,9 +16,6 @@ open Adaptive_issuance_helpers
 let apply_end_cycle current_cycle block state : State.t tzresult Lwt.t =
   let open Lwt_result_syntax in
   Log.debug ~color:time_color "Ending cycle %a" Cycle.pp current_cycle ;
-  let* launch_cycle_opt =
-    Context.get_adaptive_issuance_launch_cycle (B block)
-  in
   (* Apply all slashes *)
   let state = apply_all_slashes_at_cycle_end current_cycle state in
   (* Sets initial frozen for future cycle *)
@@ -32,18 +29,7 @@ let apply_end_cycle current_cycle block state : State.t tzresult Lwt.t =
       state
   in
   (* Apply autostaking *)
-  let*? state =
-    if not state.constants.adaptive_issuance.autostaking_enable then Ok state
-    else
-      match launch_cycle_opt with
-      | Some launch_cycle when Cycle.(current_cycle >= launch_cycle) -> Ok state
-      | None | Some _ ->
-          String.Map.fold_e
-            (fun name account state ->
-              apply_autostake ~name ~old_cycle:current_cycle account state)
-            state.account_map
-            state
-  in
+  let*? state = State_ai_flags.Autostake.run_at_cycle_end block state in
   (* Apply parameter changes *)
   let state, param_requests =
     List.fold_left
diff --git a/src/proto_alpha/lib_protocol/test/helpers/state.ml b/src/proto_alpha/lib_protocol/test/helpers/state.ml
index 3e8e496a8d98..7ce5d0ab6055 100644
--- a/src/proto_alpha/lib_protocol/test/helpers/state.ml
+++ b/src/proto_alpha/lib_protocol/test/helpers/state.ml
@@ -237,102 +237,3 @@ let add_pending_operations operations state =
 
 let pop_pending_operations state =
   ({state with pending_operations = []}, state.pending_operations)
-
-let log_model_autostake name pkh old_cycle op ~optimal amount =
-  Log.debug
-    "Model Autostaking: at end of cycle %a, %s(%a) to reach optimal stake %a \
-     %s %a"
-    Cycle.pp
-    old_cycle
-    name
-    Signature.Public_key_hash.pp
-    pkh
-    Tez.pp
-    optimal
-    op
-    Tez.pp
-    (Tez.of_mutez amount)
-
-let apply_autostake ~name ~old_cycle
-    ({
-       pkh;
-       contract = _;
-       delegate;
-       parameters = _;
-       liquid;
-       bonds = _;
-       frozen_deposits;
-       unstaked_frozen;
-       unstaked_finalizable;
-       staking_delegator_numerator = _;
-       staking_delegate_denominator = _;
-       frozen_rights = _;
-       slashed_cycles = _;
-     } :
-      account_state) state =
-  let open Result_syntax in
-  if Some name <> delegate then (
-    Log.debug
-      "Model Autostaking: %s <> %s, noop@."
-      name
-      (Option.value ~default:"None" delegate) ;
-    return state)
-  else
-    let* current_liquid_delegated = liquid_delegated ~name state in
-    let current_frozen = Frozen_tez.total_current frozen_deposits in
-    let current_unstaked_frozen_delegated =
-      Unstaked_frozen.sum_current unstaked_frozen
-    in
-    let current_unstaked_final_delegated =
-      Unstaked_finalizable.total unstaked_finalizable
-    in
-    let power =
-      Tez.(
-        current_liquid_delegated +! current_frozen
-        +! current_unstaked_frozen_delegated +! current_unstaked_final_delegated
-        |> to_mutez |> Z.of_int64)
-    in
-    let optimal =
-      Tez.of_z
-        (Z.cdiv
-           power
-           (Z.of_int (state.constants.limit_of_delegation_over_baking + 1)))
-    in
-    let autostaked =
-      Int64.(sub (Tez.to_mutez optimal) (Tez.to_mutez current_frozen))
-    in
-    let state = apply_unslashable (Cycle.succ old_cycle) name state in
-    let state = apply_finalize name state in
-    (* stake or unstake *)
-    let new_state =
-      if autostaked > 0L then (
-        log_model_autostake ~optimal name pkh old_cycle "stake" autostaked ;
-        apply_stake
-          Tez.(min liquid (of_mutez autostaked))
-          (Cycle.succ old_cycle)
-          name
-          state)
-      else if autostaked < 0L then (
-        log_model_autostake
-          ~optimal
-          name
-          pkh
-          old_cycle
-          "unstake"
-          (Int64.neg autostaked) ;
-        apply_unstake
-          (Cycle.succ old_cycle)
-          (Tez_helpers.of_mutez Int64.(neg autostaked))
-          name
-          state)
-      else (
-        log_model_autostake
-          ~optimal
-          name
-          pkh
-          old_cycle
-          "only finalize"
-          autostaked ;
-        state)
-    in
-    return new_state
diff --git a/src/proto_alpha/lib_protocol/test/helpers/state_ai_flags.ml b/src/proto_alpha/lib_protocol/test/helpers/state_ai_flags.ml
index e191a105cb3c..3ac91f59132b 100644
--- a/src/proto_alpha/lib_protocol/test/helpers/state_ai_flags.ml
+++ b/src/proto_alpha/lib_protocol/test/helpers/state_ai_flags.ml
@@ -8,6 +8,7 @@
 (** This module handles the logic of flags for AI/AS in the [State] *)
 
 open State
+open State_account
 
 module AI_Activation = struct
   (** This module is responsible for the field [state.ai_activation_cycle],
@@ -64,3 +65,121 @@ module AI = struct
         let current_cycle = Block.current_cycle block in
         Protocol.Alpha_context.Cycle.(current_cycle >= activation_cycle)
 end
+
+module Autostake = struct
+  let enabled (block : Block.t) (state : State.t) =
+    (not (AI.enabled block state))
+    && state.constants.adaptive_issuance.autostaking_enable
+
+  let log_model_autostake name pkh old_cycle op ~optimal amount =
+    let open Protocol.Alpha_context in
+    Log.debug
+      "Model Autostaking: at end of cycle %a, %s(%a) to reach optimal stake %a \
+       %s %a"
+      Cycle.pp
+      old_cycle
+      name
+      Signature.Public_key_hash.pp
+      pkh
+      Tez.pp
+      optimal
+      op
+      Tez.pp
+      (Tez_helpers.of_mutez amount)
+
+  let apply_autostake ~name ~old_cycle
+      ({
+         pkh;
+         contract = _;
+         delegate;
+         parameters = _;
+         liquid;
+         bonds = _;
+         frozen_deposits;
+         unstaked_frozen;
+         unstaked_finalizable;
+         staking_delegator_numerator = _;
+         staking_delegate_denominator = _;
+         frozen_rights = _;
+         slashed_cycles = _;
+       } :
+        account_state) state =
+    let open Result_syntax in
+    if Some name <> delegate then (
+      Log.debug
+        "Model Autostaking: %s <> %s, noop@."
+        name
+        (Option.value ~default:"None" delegate) ;
+      return state)
+    else
+      let* current_liquid_delegated = liquid_delegated ~name state in
+      let current_frozen = Frozen_tez.total_current frozen_deposits in
+      let current_unstaked_frozen_delegated =
+        Unstaked_frozen.sum_current unstaked_frozen
+      in
+      let current_unstaked_final_delegated =
+        Unstaked_finalizable.total unstaked_finalizable
+      in
+      let power =
+        Tez.(
+          current_liquid_delegated +! current_frozen
+          +! current_unstaked_frozen_delegated
+          +! current_unstaked_final_delegated
+          |> to_mutez |> Z.of_int64)
+      in
+      let optimal =
+        Tez.of_z
+          (Z.cdiv
+             power
+             (Z.of_int (state.constants.limit_of_delegation_over_baking + 1)))
+      in
+      let autostaked =
+        Int64.(sub (Tez.to_mutez optimal) (Tez.to_mutez current_frozen))
+      in
+      let state = State.apply_unslashable (Cycle.succ old_cycle) name state in
+      let state = State.apply_finalize name state in
+      (* stake or unstake *)
+      let new_state =
+        if autostaked > 0L then (
+          log_model_autostake ~optimal name pkh old_cycle "stake" autostaked ;
+          State.apply_stake
+            Tez.(min liquid (of_mutez autostaked))
+            (Cycle.succ old_cycle)
+            name
+            state)
+        else if autostaked < 0L then (
+          log_model_autostake
+            ~optimal
+            name
+            pkh
+            old_cycle
+            "unstake"
+            (Int64.neg autostaked) ;
+          State.apply_unstake
+            (Cycle.succ old_cycle)
+            (Tez_helpers.of_mutez Int64.(neg autostaked))
+            name
+            state)
+        else (
+          log_model_autostake
+            ~optimal
+            name
+            pkh
+            old_cycle
+            "only finalize"
+            autostaked ;
+          state)
+      in
+      return new_state
+
+  let run_at_cycle_end block state =
+    let open Result_syntax in
+    if enabled block state then
+      let current_cycle = Block.current_cycle block in
+      String.Map.fold_e
+        (fun name account state ->
+          apply_autostake ~name ~old_cycle:current_cycle account state)
+        state.account_map
+        state
+    else return state
+end
diff --git a/src/proto_alpha/lib_protocol/test/helpers/state_ai_flags.mli b/src/proto_alpha/lib_protocol/test/helpers/state_ai_flags.mli
index 20f3df44813d..79b7c4b44d41 100644
--- a/src/proto_alpha/lib_protocol/test/helpers/state_ai_flags.mli
+++ b/src/proto_alpha/lib_protocol/test/helpers/state_ai_flags.mli
@@ -30,3 +30,15 @@ module AI : sig
   (** AI is enabled iff the activation cycle is set and passed *)
   val enabled : Block.t -> State.t -> bool
 end
+
+module Autostake : sig
+  (** This module takes care of autostaking when it is enabled *)
+
+  (** Autostaking is enabled iff the flag [autostaking_enable] is true and
+      AI is not activated ([AI.enabled = false]). *)
+  val enabled : Block.t -> State.t -> bool
+
+  (** Runs the autostake operations at cycle end. Does nothing if
+      [enabled = false]. *)
+  val run_at_cycle_end : Block.t -> State.t -> State.t tzresult
+end
-- 
GitLab


From a9f61dcb3ade7381f1b3d3cbeb7d60e6ca7bb136 Mon Sep 17 00:00:00 2001
From: Lucas Randazzo <lucas@nomadic-labs.com>
Date: Tue, 5 Mar 2024 16:09:29 +0100
Subject: [PATCH 07/10] Proto/tests: move slashing functions in
 slashing_helpers

---
 .../test/helpers/scenario_bake.ml             |   4 +-
 .../test/helpers/slashing_helpers.ml          | 170 ++++++++++++++++++
 .../test/helpers/slashing_helpers.mli         |   4 +
 .../lib_protocol/test/helpers/state.ml        |  50 ------
 .../test/helpers/state_account.ml             | 117 ------------
 5 files changed, 177 insertions(+), 168 deletions(-)

diff --git a/src/proto_alpha/lib_protocol/test/helpers/scenario_bake.ml b/src/proto_alpha/lib_protocol/test/helpers/scenario_bake.ml
index 3275e403a119..43d368cb3aeb 100644
--- a/src/proto_alpha/lib_protocol/test/helpers/scenario_bake.ml
+++ b/src/proto_alpha/lib_protocol/test/helpers/scenario_bake.ml
@@ -17,7 +17,9 @@ let apply_end_cycle current_cycle block state : State.t tzresult Lwt.t =
   let open Lwt_result_syntax in
   Log.debug ~color:time_color "Ending cycle %a" Cycle.pp current_cycle ;
   (* Apply all slashes *)
-  let state = apply_all_slashes_at_cycle_end current_cycle state in
+  let state =
+    Slashing_helpers.apply_all_slashes_at_cycle_end current_cycle state
+  in
   (* Sets initial frozen for future cycle *)
   let state =
     update_map
diff --git a/src/proto_alpha/lib_protocol/test/helpers/slashing_helpers.ml b/src/proto_alpha/lib_protocol/test/helpers/slashing_helpers.ml
index 97f7a712dd7f..f8a35271f581 100644
--- a/src/proto_alpha/lib_protocol/test/helpers/slashing_helpers.ml
+++ b/src/proto_alpha/lib_protocol/test/helpers/slashing_helpers.ml
@@ -109,3 +109,173 @@ module Full_denunciation = struct
       list1
       list2
 end
+
+let apply_slashing_account
+    ( culprit,
+      Protocol.Denunciations_repr.{rewarded; misbehaviour; operation_hash = _}
+    ) constants (account_map : State_account.account_map) =
+  let open State_account in
+  let find_account_name_from_pkh_exn pkh account_map =
+    match
+      Option.map
+        fst
+        String.Map.(
+          choose
+          @@ filter
+               (fun _ account ->
+                 Signature.Public_key_hash.equal pkh account.State_account.pkh)
+               account_map)
+    with
+    | None -> assert false
+    | Some x -> x
+  in
+  let slashed_cycle =
+    Block.current_cycle_of_level
+      ~blocks_per_cycle:
+        constants.Protocol.Alpha_context.Constants.Parametric.blocks_per_cycle
+      ~current_level:(Protocol.Raw_level_repr.to_int32 misbehaviour.level)
+  in
+  let culprit_name = find_account_name_from_pkh_exn culprit account_map in
+  let rewarded_name = find_account_name_from_pkh_exn rewarded account_map in
+  let slashed_pct =
+    match misbehaviour.kind with
+    | Double_baking ->
+        constants
+          .Protocol.Alpha_context.Constants.Parametric
+           .percentage_of_frozen_deposits_slashed_per_double_baking
+    | Double_attesting | Double_preattesting ->
+        constants.percentage_of_frozen_deposits_slashed_per_double_attestation
+  in
+  let get_total_supply acc_map =
+    String.Map.fold
+      (fun _name
+           {
+             State_account.pkh = _;
+             contract = _;
+             delegate = _;
+             parameters = _;
+             liquid;
+             bonds;
+             frozen_deposits;
+             unstaked_frozen;
+             unstaked_finalizable;
+             staking_delegator_numerator = _;
+             staking_delegate_denominator = _;
+             frozen_rights = _;
+             slashed_cycles = _;
+           }
+           tot ->
+        Tez.(
+          liquid +! bonds
+          +! Frozen_tez.total_current frozen_deposits
+          +! Unstaked_frozen.sum_current unstaked_frozen
+          +! Unstaked_finalizable.total unstaked_finalizable
+          +! tot))
+      acc_map
+      Tez.zero
+  in
+  let total_before_slash = get_total_supply account_map in
+  let slash_culprit
+      ({frozen_deposits; unstaked_frozen; frozen_rights; _} as acc) =
+    let base_rights =
+      CycleMap.find slashed_cycle frozen_rights
+      |> Option.value ~default:Tez.zero
+    in
+    let frozen_deposits, slashed_frozen =
+      Frozen_tez.slash base_rights slashed_pct frozen_deposits
+    in
+    let slashed_pct_q = Protocol.Percentage.to_q slashed_pct in
+    let slashed_pct = Q.(100 // 1 * slashed_pct_q |> to_int) in
+    let unstaked_frozen, slashed_unstaked =
+      Unstaked_frozen.slash
+        ~slashable_deposits_period:constants.consensus_rights_delay
+        slashed_cycle
+        slashed_pct
+        unstaked_frozen
+    in
+    ( {acc with frozen_deposits; unstaked_frozen},
+      slashed_frozen :: slashed_unstaked )
+  in
+  let culprit_account =
+    String.Map.find culprit_name account_map
+    |> Option.value_f ~default:(fun () ->
+           fail_account_not_found "apply_slashing" culprit_name)
+  in
+  let slashed_culprit_account, total_slashed = slash_culprit culprit_account in
+  let account_map =
+    update_account
+      ~f:(fun _ -> slashed_culprit_account)
+      culprit_name
+      account_map
+  in
+  let total_after_slash = get_total_supply account_map in
+  let portion_reward =
+    constants.adaptive_issuance.global_limit_of_staking_over_baking + 2
+  in
+  (* For each container slashed, the snitch gets a reward transferred. It gets rounded
+     down each time *)
+  let reward_to_snitch =
+    List.map
+      (fun x -> Tez.mul_q x Q.(1 // portion_reward) |> Tez.of_q ~round:`Down)
+      total_slashed
+    |> List.fold_left Tez.( +! ) Tez.zero
+  in
+  let account_map =
+    add_liquid_rewards reward_to_snitch rewarded_name account_map
+  in
+  let actual_total_burnt_amount =
+    Tez.(total_before_slash -! total_after_slash -! reward_to_snitch)
+  in
+  (account_map, actual_total_burnt_amount)
+
+let apply_slashing_state
+    ( culprit,
+      Protocol.Denunciations_repr.{rewarded; misbehaviour; operation_hash} )
+    (state : State.t) : State.t * Tez_helpers.t =
+  let account_map, total_burnt =
+    apply_slashing_account
+      (culprit, {rewarded; misbehaviour; operation_hash})
+      state.constants
+      state.account_map
+  in
+  (* TODO: add culprit's stakers *)
+  let log_updates =
+    List.map
+      (fun x -> fst @@ State.find_account_from_pkh x state)
+      [culprit; rewarded]
+  in
+  let state = State.update_map ~log_updates ~f:(fun _ -> account_map) state in
+  (state, total_burnt)
+
+let apply_all_slashes_at_cycle_end current_cycle (state : State.t) : State.t =
+  let to_slash_later, to_slash_now =
+    if
+      not
+        state.constants
+          .Protocol.Alpha_context.Constants.Parametric.adaptive_issuance
+          .ns_enable
+    then ([], state.pending_slashes)
+    else
+      List.partition
+        (fun (_, Protocol.Denunciations_repr.{misbehaviour; _}) ->
+          let cycle =
+            Block.current_cycle_of_level
+              ~blocks_per_cycle:
+                state.constants
+                  .Protocol.Alpha_context.Constants.Parametric.blocks_per_cycle
+              ~current_level:
+                (Protocol.Raw_level_repr.to_int32 misbehaviour.level)
+          in
+          Protocol.Alpha_context.Cycle.(cycle = current_cycle))
+        state.pending_slashes
+  in
+  let state, total_burnt =
+    List.fold_left
+      (fun (acc_state, acc_total) x ->
+        let state, burnt = apply_slashing_state x acc_state in
+        (state, Tez_helpers.(acc_total +! burnt)))
+      (state, Tez_helpers.zero)
+      to_slash_now
+  in
+  let total_supply = Tez_helpers.(state.total_supply -! total_burnt) in
+  {state with pending_slashes = to_slash_later; total_supply}
diff --git a/src/proto_alpha/lib_protocol/test/helpers/slashing_helpers.mli b/src/proto_alpha/lib_protocol/test/helpers/slashing_helpers.mli
index 98fadb0c8763..42954ca0ed8d 100644
--- a/src/proto_alpha/lib_protocol/test/helpers/slashing_helpers.mli
+++ b/src/proto_alpha/lib_protocol/test/helpers/slashing_helpers.mli
@@ -49,3 +49,7 @@ module Full_denunciation : sig
   val check_same_lists_any_order :
     loc:string -> t list -> t list -> unit tzresult Lwt.t
 end
+
+(** Applies all slashes at cycle end in the state *)
+val apply_all_slashes_at_cycle_end :
+  Protocol.Alpha_context.Cycle.t -> State.t -> State.t
diff --git a/src/proto_alpha/lib_protocol/test/helpers/state.ml b/src/proto_alpha/lib_protocol/test/helpers/state.ml
index 7ce5d0ab6055..a890a5e3d5ce 100644
--- a/src/proto_alpha/lib_protocol/test/helpers/state.ml
+++ b/src/proto_alpha/lib_protocol/test/helpers/state.ml
@@ -171,56 +171,6 @@ let apply_rewards ~(baker : string) block (state : t) : t tzresult Lwt.t =
     let*? total_supply = Tez.(total_supply +? delta_rewards) in
     return {state with last_level_rewards = current_level; total_supply}
 
-let apply_slashing
-    ( culprit,
-      Protocol.Denunciations_repr.{rewarded; misbehaviour; operation_hash} )
-    (state : t) : t * Tez.t =
-  let account_map, total_burnt =
-    apply_slashing
-      (culprit, {rewarded; misbehaviour; operation_hash})
-      state.constants
-      state.account_map
-  in
-  (* TODO: add culprit's stakers *)
-  let log_updates =
-    List.map (fun x -> fst @@ find_account_from_pkh x state) [culprit; rewarded]
-  in
-  let state = update_map ~log_updates ~f:(fun _ -> account_map) state in
-  (state, total_burnt)
-
-let apply_all_slashes_at_cycle_end current_cycle (state : t) : t =
-  let to_slash_later, to_slash_now =
-    if
-      not
-        state.constants
-          .Protocol.Alpha_context.Constants.Parametric.adaptive_issuance
-          .ns_enable
-    then ([], state.pending_slashes)
-    else
-      List.partition
-        (fun (_, Protocol.Denunciations_repr.{misbehaviour; _}) ->
-          let cycle =
-            Block.current_cycle_of_level
-              ~blocks_per_cycle:
-                state.constants
-                  .Protocol.Alpha_context.Constants.Parametric.blocks_per_cycle
-              ~current_level:
-                (Protocol.Raw_level_repr.to_int32 misbehaviour.level)
-          in
-          Cycle.(cycle = current_cycle))
-        state.pending_slashes
-  in
-  let state, total_burnt =
-    List.fold_left
-      (fun (acc_state, acc_total) x ->
-        let state, burnt = apply_slashing x acc_state in
-        (state, Tez.(acc_total +! burnt)))
-      (state, Tez.zero)
-      to_slash_now
-  in
-  let total_supply = Tez.(state.total_supply -! total_burnt) in
-  {state with pending_slashes = to_slash_later; total_supply}
-
 (** Given an account name and new account state, updates [state] accordingly
       Preferably use other specific update functions *)
 let update_account (account_name : string) (value : account_state) (state : t) :
diff --git a/src/proto_alpha/lib_protocol/test/helpers/state_account.ml b/src/proto_alpha/lib_protocol/test/helpers/state_account.ml
index a995df7853ff..e831bfb24350 100644
--- a/src/proto_alpha/lib_protocol/test/helpers/state_account.ml
+++ b/src/proto_alpha/lib_protocol/test/helpers/state_account.ml
@@ -367,123 +367,6 @@ let apply_finalize staker_name account_map =
         account_map
         account_map
 
-let apply_slashing
-    ( culprit,
-      Protocol.Denunciations_repr.{rewarded; misbehaviour; operation_hash = _}
-    ) constants account_map =
-  let find_account_name_from_pkh_exn pkh account_map =
-    match
-      Option.map
-        fst
-        String.Map.(
-          choose
-          @@ filter
-               (fun _ account ->
-                 Signature.Public_key_hash.equal pkh account.pkh)
-               account_map)
-    with
-    | None -> assert false
-    | Some x -> x
-  in
-  let slashed_cycle =
-    Block.current_cycle_of_level
-      ~blocks_per_cycle:
-        constants.Protocol.Alpha_context.Constants.Parametric.blocks_per_cycle
-      ~current_level:(Protocol.Raw_level_repr.to_int32 misbehaviour.level)
-  in
-  let culprit_name = find_account_name_from_pkh_exn culprit account_map in
-  let rewarded_name = find_account_name_from_pkh_exn rewarded account_map in
-  let slashed_pct =
-    match misbehaviour.kind with
-    | Double_baking ->
-        constants
-          .Protocol.Alpha_context.Constants.Parametric
-           .percentage_of_frozen_deposits_slashed_per_double_baking
-    | Double_attesting | Double_preattesting ->
-        constants.percentage_of_frozen_deposits_slashed_per_double_attestation
-  in
-  let get_total_supply acc_map =
-    String.Map.fold
-      (fun _name
-           {
-             pkh = _;
-             contract = _;
-             delegate = _;
-             parameters = _;
-             liquid;
-             bonds;
-             frozen_deposits;
-             unstaked_frozen;
-             unstaked_finalizable;
-             staking_delegator_numerator = _;
-             staking_delegate_denominator = _;
-             frozen_rights = _;
-             slashed_cycles = _;
-           }
-           tot ->
-        Tez.(
-          liquid +! bonds
-          +! Frozen_tez.total_current frozen_deposits
-          +! Unstaked_frozen.sum_current unstaked_frozen
-          +! Unstaked_finalizable.total unstaked_finalizable
-          +! tot))
-      acc_map
-      Tez.zero
-  in
-  let total_before_slash = get_total_supply account_map in
-  let slash_culprit
-      ({frozen_deposits; unstaked_frozen; frozen_rights; _} as acc) =
-    let base_rights =
-      CycleMap.find slashed_cycle frozen_rights
-      |> Option.value ~default:Tez.zero
-    in
-    let frozen_deposits, slashed_frozen =
-      Frozen_tez.slash base_rights slashed_pct frozen_deposits
-    in
-    let slashed_pct_q = Protocol.Percentage.to_q slashed_pct in
-    let slashed_pct = Q.(100 // 1 * slashed_pct_q |> to_int) in
-    let unstaked_frozen, slashed_unstaked =
-      Unstaked_frozen.slash
-        ~slashable_deposits_period:constants.consensus_rights_delay
-        slashed_cycle
-        slashed_pct
-        unstaked_frozen
-    in
-    ( {acc with frozen_deposits; unstaked_frozen},
-      slashed_frozen :: slashed_unstaked )
-  in
-  let culprit_account =
-    String.Map.find culprit_name account_map
-    |> Option.value_f ~default:(fun () ->
-           fail_account_not_found "apply_slashing" culprit_name)
-  in
-  let slashed_culprit_account, total_slashed = slash_culprit culprit_account in
-  let account_map =
-    update_account
-      ~f:(fun _ -> slashed_culprit_account)
-      culprit_name
-      account_map
-  in
-  let total_after_slash = get_total_supply account_map in
-  let portion_reward =
-    constants.adaptive_issuance.global_limit_of_staking_over_baking + 2
-  in
-  (* For each container slashed, the snitch gets a reward transferred. It gets rounded
-     down each time *)
-  let reward_to_snitch =
-    List.map
-      (fun x -> Tez.mul_q x Q.(1 // portion_reward) |> Tez.of_q ~round:`Down)
-      total_slashed
-    |> List.fold_left Tez.( +! ) Tez.zero
-  in
-  let account_map =
-    add_liquid_rewards reward_to_snitch rewarded_name account_map
-  in
-  let actual_total_burnt_amount =
-    Tez.(total_before_slash -! total_after_slash -! reward_to_snitch)
-  in
-  (account_map, actual_total_burnt_amount)
-
 (* Given cycle is the cycle for which the rights are computed, usually current +
    consensus rights delay *)
 let update_frozen_rights_cycle cycle account_map =
-- 
GitLab


From 155506b86438f9426992e1ea137cc0be4a7494c3 Mon Sep 17 00:00:00 2001
From: Lucas Randazzo <lucas@nomadic-labs.com>
Date: Tue, 5 Mar 2024 16:32:05 +0100
Subject: [PATCH 08/10] Proto/tests: introduce the NS module

---
 .../test/helpers/scenario_bake.ml             | 11 +++++++---
 .../test/helpers/slashing_helpers.ml          | 21 ++++++++++++-------
 .../test/helpers/slashing_helpers.mli         |  2 +-
 .../test/helpers/state_ai_flags.ml            | 17 +++++++++++++++
 .../test/helpers/state_ai_flags.mli           | 13 ++++++++++++
 5 files changed, 53 insertions(+), 11 deletions(-)

diff --git a/src/proto_alpha/lib_protocol/test/helpers/scenario_bake.ml b/src/proto_alpha/lib_protocol/test/helpers/scenario_bake.ml
index 43d368cb3aeb..82cf6ff707e3 100644
--- a/src/proto_alpha/lib_protocol/test/helpers/scenario_bake.ml
+++ b/src/proto_alpha/lib_protocol/test/helpers/scenario_bake.ml
@@ -13,12 +13,16 @@ open Scenario_base
 open Adaptive_issuance_helpers
 
 (** Applies when baking the last block of a cycle *)
-let apply_end_cycle current_cycle block state : State.t tzresult Lwt.t =
+let apply_end_cycle current_cycle previous_block block state :
+    State.t tzresult Lwt.t =
   let open Lwt_result_syntax in
   Log.debug ~color:time_color "Ending cycle %a" Cycle.pp current_cycle ;
   (* Apply all slashes *)
   let state =
-    Slashing_helpers.apply_all_slashes_at_cycle_end current_cycle state
+    Slashing_helpers.apply_all_slashes_at_cycle_end
+      current_cycle
+      previous_block
+      state
   in
   (* Sets initial frozen for future cycle *)
   let state =
@@ -133,6 +137,7 @@ let check_issuance_rpc block : unit tzresult Lwt.t =
 let bake ?baker : t -> t tzresult Lwt.t =
  fun (block, state) ->
   let open Lwt_result_wrap_syntax in
+  let previous_block = block in
   let policy =
     match baker with
     | None -> state.baking_policy
@@ -209,7 +214,7 @@ let bake ?baker : t -> t tzresult Lwt.t =
   (* Dawn of a new cycle *)
   let* state =
     if not (Block.last_block_of_cycle block) then return state
-    else apply_end_cycle current_cycle block state
+    else apply_end_cycle current_cycle previous_block block state
   in
   let* () = check_all_balances block state in
   return (block, state)
diff --git a/src/proto_alpha/lib_protocol/test/helpers/slashing_helpers.ml b/src/proto_alpha/lib_protocol/test/helpers/slashing_helpers.ml
index f8a35271f581..fcb0cf10bbba 100644
--- a/src/proto_alpha/lib_protocol/test/helpers/slashing_helpers.ml
+++ b/src/proto_alpha/lib_protocol/test/helpers/slashing_helpers.ml
@@ -113,8 +113,10 @@ end
 let apply_slashing_account
     ( culprit,
       Protocol.Denunciations_repr.{rewarded; misbehaviour; operation_hash = _}
-    ) constants (account_map : State_account.account_map) =
+    ) (block_before_slash : Block.t) (state : State.t) =
   let open State_account in
+  let constants = state.constants in
+  let (account_map : State_account.account_map) = state.account_map in
   let find_account_name_from_pkh_exn pkh account_map =
     match
       Option.map
@@ -144,7 +146,9 @@ let apply_slashing_account
           .Protocol.Alpha_context.Constants.Parametric
            .percentage_of_frozen_deposits_slashed_per_double_baking
     | Double_attesting | Double_preattesting ->
-        constants.percentage_of_frozen_deposits_slashed_per_double_attestation
+        State_ai_flags.NS.get_double_attestation_slashing_percentage
+          block_before_slash
+          state
   in
   let get_total_supply acc_map =
     String.Map.fold
@@ -231,12 +235,12 @@ let apply_slashing_account
 let apply_slashing_state
     ( culprit,
       Protocol.Denunciations_repr.{rewarded; misbehaviour; operation_hash} )
-    (state : State.t) : State.t * Tez_helpers.t =
+    block_before_slash (state : State.t) : State.t * Tez_helpers.t =
   let account_map, total_burnt =
     apply_slashing_account
       (culprit, {rewarded; misbehaviour; operation_hash})
-      state.constants
-      state.account_map
+      block_before_slash
+      state
   in
   (* TODO: add culprit's stakers *)
   let log_updates =
@@ -247,7 +251,8 @@ let apply_slashing_state
   let state = State.update_map ~log_updates ~f:(fun _ -> account_map) state in
   (state, total_burnt)
 
-let apply_all_slashes_at_cycle_end current_cycle (state : State.t) : State.t =
+let apply_all_slashes_at_cycle_end current_cycle (block_before_slash : Block.t)
+    (state : State.t) : State.t =
   let to_slash_later, to_slash_now =
     if
       not
@@ -272,7 +277,9 @@ let apply_all_slashes_at_cycle_end current_cycle (state : State.t) : State.t =
   let state, total_burnt =
     List.fold_left
       (fun (acc_state, acc_total) x ->
-        let state, burnt = apply_slashing_state x acc_state in
+        let state, burnt =
+          apply_slashing_state x block_before_slash acc_state
+        in
         (state, Tez_helpers.(acc_total +! burnt)))
       (state, Tez_helpers.zero)
       to_slash_now
diff --git a/src/proto_alpha/lib_protocol/test/helpers/slashing_helpers.mli b/src/proto_alpha/lib_protocol/test/helpers/slashing_helpers.mli
index 42954ca0ed8d..ed9e7cb6fb0b 100644
--- a/src/proto_alpha/lib_protocol/test/helpers/slashing_helpers.mli
+++ b/src/proto_alpha/lib_protocol/test/helpers/slashing_helpers.mli
@@ -52,4 +52,4 @@ end
 
 (** Applies all slashes at cycle end in the state *)
 val apply_all_slashes_at_cycle_end :
-  Protocol.Alpha_context.Cycle.t -> State.t -> State.t
+  Protocol.Alpha_context.Cycle.t -> Block.t -> State.t -> State.t
diff --git a/src/proto_alpha/lib_protocol/test/helpers/state_ai_flags.ml b/src/proto_alpha/lib_protocol/test/helpers/state_ai_flags.ml
index 3ac91f59132b..56bdaa82bf51 100644
--- a/src/proto_alpha/lib_protocol/test/helpers/state_ai_flags.ml
+++ b/src/proto_alpha/lib_protocol/test/helpers/state_ai_flags.ml
@@ -183,3 +183,20 @@ module Autostake = struct
         state
     else return state
 end
+
+module NS = struct
+  let enabled (block : Block.t) (state : State.t) =
+    AI.enabled block state && state.constants.adaptive_issuance.ns_enable
+
+  let get_double_attestation_slashing_percentage block_before_slash state =
+    (* We need to get the block before the slash, because after the slash,
+       the context gets rid of the required Seed to recompute the rights
+       for the misbehaving delegates. *)
+    if not (enabled block_before_slash state) then
+      state.constants
+        .percentage_of_frozen_deposits_slashed_per_double_attestation
+    else
+      (* TODO *)
+      state.constants
+        .percentage_of_frozen_deposits_slashed_per_double_attestation
+end
diff --git a/src/proto_alpha/lib_protocol/test/helpers/state_ai_flags.mli b/src/proto_alpha/lib_protocol/test/helpers/state_ai_flags.mli
index 79b7c4b44d41..97015571ee00 100644
--- a/src/proto_alpha/lib_protocol/test/helpers/state_ai_flags.mli
+++ b/src/proto_alpha/lib_protocol/test/helpers/state_ai_flags.mli
@@ -42,3 +42,16 @@ module Autostake : sig
       [enabled = false]. *)
   val run_at_cycle_end : Block.t -> State.t -> State.t tzresult
 end
+
+module NS : sig
+  (** This module takes care of the new adaptive slashing mechanism.*)
+
+  (** It is enabled iff the flag [ns_enable] is set to true, and AI is
+      also enabled. *)
+  val enabled : Block.t -> State.t -> bool
+
+  (** Whatever the value of the flag is, this function returns the
+      slashing value for a given double attestation *)
+  val get_double_attestation_slashing_percentage :
+    Block.t -> State.t -> Protocol.Percentage.t
+end
-- 
GitLab


From 04004f4a0144d0e96ed8958fd972d0d89b871b56 Mon Sep 17 00:00:00 2001
From: Lucas Randazzo <lucas@nomadic-labs.com>
Date: Tue, 5 Mar 2024 17:38:57 +0100
Subject: [PATCH 09/10] Proto/tests: implement New slashing

---
 .../lib_protocol/test/helpers/context.mli     |  2 +
 .../test/helpers/scenario_bake.ml             |  2 +-
 .../test/helpers/slashing_helpers.ml          | 42 +++++++++++--------
 .../test/helpers/slashing_helpers.mli         |  2 +-
 .../test/helpers/state_ai_flags.ml            | 42 ++++++++++++++++---
 .../test/helpers/state_ai_flags.mli           |  6 ++-
 6 files changed, 70 insertions(+), 26 deletions(-)

diff --git a/src/proto_alpha/lib_protocol/test/helpers/context.mli b/src/proto_alpha/lib_protocol/test/helpers/context.mli
index 05737069baf7..1bfb3e326666 100644
--- a/src/proto_alpha/lib_protocol/test/helpers/context.mli
+++ b/src/proto_alpha/lib_protocol/test/helpers/context.mli
@@ -30,6 +30,8 @@ open Environment
 
 type t = B of Block.t | I of Incremental.t
 
+val get_alpha_ctxt : ?policy:Block.baker_policy -> t -> context tzresult Lwt.t
+
 val branch : t -> Block_hash.t
 
 val pred_branch : t -> Block_hash.t
diff --git a/src/proto_alpha/lib_protocol/test/helpers/scenario_bake.ml b/src/proto_alpha/lib_protocol/test/helpers/scenario_bake.ml
index 82cf6ff707e3..2f66162f1ea1 100644
--- a/src/proto_alpha/lib_protocol/test/helpers/scenario_bake.ml
+++ b/src/proto_alpha/lib_protocol/test/helpers/scenario_bake.ml
@@ -18,7 +18,7 @@ let apply_end_cycle current_cycle previous_block block state :
   let open Lwt_result_syntax in
   Log.debug ~color:time_color "Ending cycle %a" Cycle.pp current_cycle ;
   (* Apply all slashes *)
-  let state =
+  let* state =
     Slashing_helpers.apply_all_slashes_at_cycle_end
       current_cycle
       previous_block
diff --git a/src/proto_alpha/lib_protocol/test/helpers/slashing_helpers.ml b/src/proto_alpha/lib_protocol/test/helpers/slashing_helpers.ml
index fcb0cf10bbba..77a72558f599 100644
--- a/src/proto_alpha/lib_protocol/test/helpers/slashing_helpers.ml
+++ b/src/proto_alpha/lib_protocol/test/helpers/slashing_helpers.ml
@@ -110,10 +110,11 @@ module Full_denunciation = struct
       list2
 end
 
-let apply_slashing_account
+let apply_slashing_account all_denunciations_to_apply
     ( culprit,
       Protocol.Denunciations_repr.{rewarded; misbehaviour; operation_hash = _}
     ) (block_before_slash : Block.t) (state : State.t) =
+  let open Lwt_result_syntax in
   let open State_account in
   let constants = state.constants in
   let (account_map : State_account.account_map) = state.account_map in
@@ -139,16 +140,19 @@ let apply_slashing_account
   in
   let culprit_name = find_account_name_from_pkh_exn culprit account_map in
   let rewarded_name = find_account_name_from_pkh_exn rewarded account_map in
-  let slashed_pct =
+  let* slashed_pct =
     match misbehaviour.kind with
     | Double_baking ->
-        constants
-          .Protocol.Alpha_context.Constants.Parametric
-           .percentage_of_frozen_deposits_slashed_per_double_baking
+        return
+          constants
+            .Protocol.Alpha_context.Constants.Parametric
+             .percentage_of_frozen_deposits_slashed_per_double_baking
     | Double_attesting | Double_preattesting ->
         State_ai_flags.NS.get_double_attestation_slashing_percentage
+          all_denunciations_to_apply
           block_before_slash
           state
+          misbehaviour
   in
   let get_total_supply acc_map =
     String.Map.fold
@@ -230,14 +234,17 @@ let apply_slashing_account
   let actual_total_burnt_amount =
     Tez.(total_before_slash -! total_after_slash -! reward_to_snitch)
   in
-  (account_map, actual_total_burnt_amount)
+  return (account_map, actual_total_burnt_amount)
 
-let apply_slashing_state
+let apply_slashing_state all_denunciations_to_apply
     ( culprit,
       Protocol.Denunciations_repr.{rewarded; misbehaviour; operation_hash} )
-    block_before_slash (state : State.t) : State.t * Tez_helpers.t =
-  let account_map, total_burnt =
+    block_before_slash (state : State.t) :
+    (State.t * Tez_helpers.t) tzresult Lwt.t =
+  let open Lwt_result_syntax in
+  let* account_map, total_burnt =
     apply_slashing_account
+      all_denunciations_to_apply
       (culprit, {rewarded; misbehaviour; operation_hash})
       block_before_slash
       state
@@ -249,10 +256,11 @@ let apply_slashing_state
       [culprit; rewarded]
   in
   let state = State.update_map ~log_updates ~f:(fun _ -> account_map) state in
-  (state, total_burnt)
+  return (state, total_burnt)
 
 let apply_all_slashes_at_cycle_end current_cycle (block_before_slash : Block.t)
-    (state : State.t) : State.t =
+    (state : State.t) : State.t tzresult Lwt.t =
+  let open Lwt_result_syntax in
   let to_slash_later, to_slash_now =
     if
       not
@@ -274,15 +282,15 @@ let apply_all_slashes_at_cycle_end current_cycle (block_before_slash : Block.t)
           Protocol.Alpha_context.Cycle.(cycle = current_cycle))
         state.pending_slashes
   in
-  let state, total_burnt =
-    List.fold_left
+  let* state, total_burnt =
+    List.fold_left_es
       (fun (acc_state, acc_total) x ->
-        let state, burnt =
-          apply_slashing_state x block_before_slash acc_state
+        let* state, burnt =
+          apply_slashing_state to_slash_now x block_before_slash acc_state
         in
-        (state, Tez_helpers.(acc_total +! burnt)))
+        return (state, Tez_helpers.(acc_total +! burnt)))
       (state, Tez_helpers.zero)
       to_slash_now
   in
   let total_supply = Tez_helpers.(state.total_supply -! total_burnt) in
-  {state with pending_slashes = to_slash_later; total_supply}
+  return {state with pending_slashes = to_slash_later; total_supply}
diff --git a/src/proto_alpha/lib_protocol/test/helpers/slashing_helpers.mli b/src/proto_alpha/lib_protocol/test/helpers/slashing_helpers.mli
index ed9e7cb6fb0b..2a9a2c368ce7 100644
--- a/src/proto_alpha/lib_protocol/test/helpers/slashing_helpers.mli
+++ b/src/proto_alpha/lib_protocol/test/helpers/slashing_helpers.mli
@@ -52,4 +52,4 @@ end
 
 (** Applies all slashes at cycle end in the state *)
 val apply_all_slashes_at_cycle_end :
-  Protocol.Alpha_context.Cycle.t -> Block.t -> State.t -> State.t
+  Protocol.Alpha_context.Cycle.t -> Block.t -> State.t -> State.t tzresult Lwt.t
diff --git a/src/proto_alpha/lib_protocol/test/helpers/state_ai_flags.ml b/src/proto_alpha/lib_protocol/test/helpers/state_ai_flags.ml
index 56bdaa82bf51..84cc0bebc605 100644
--- a/src/proto_alpha/lib_protocol/test/helpers/state_ai_flags.ml
+++ b/src/proto_alpha/lib_protocol/test/helpers/state_ai_flags.ml
@@ -188,15 +188,45 @@ module NS = struct
   let enabled (block : Block.t) (state : State.t) =
     AI.enabled block state && state.constants.adaptive_issuance.ns_enable
 
-  let get_double_attestation_slashing_percentage block_before_slash state =
+  let get_double_attestation_slashing_percentage all_denunciations_to_apply
+      block_before_slash state (misbehaviour : Protocol.Misbehaviour_repr.t) =
+    let open Lwt_result_wrap_syntax in
     (* We need to get the block before the slash, because after the slash,
        the context gets rid of the required Seed to recompute the rights
        for the misbehaving delegates. *)
     if not (enabled block_before_slash state) then
-      state.constants
-        .percentage_of_frozen_deposits_slashed_per_double_attestation
+      return
+        state.constants
+          .percentage_of_frozen_deposits_slashed_per_double_attestation
     else
-      (* TODO *)
-      state.constants
-        .percentage_of_frozen_deposits_slashed_per_double_attestation
+      let* alpha_ctxt =
+        Context.(
+          get_alpha_ctxt ?policy:state.baking_policy (B block_before_slash))
+      in
+      let raw_ctxt =
+        Protocol.Alpha_context.Internal_for_tests.to_raw alpha_ctxt
+      in
+      let level =
+        Protocol.Level_repr.level_from_raw
+          ~cycle_eras:(Protocol.Raw_context.cycle_eras raw_ctxt)
+          misbehaviour.level
+      in
+      let delegates =
+        List.filter
+          (fun (_, (den : Protocol.Denunciations_repr.item)) ->
+            Compare.Int.(
+              Protocol.Misbehaviour_repr.compare misbehaviour den.misbehaviour
+              = 0))
+          all_denunciations_to_apply
+        |> List.map fst
+        |> List.sort_uniq Signature.Public_key_hash.compare
+      in
+      let*@ _, pct =
+        Protocol.Slash_percentage.get
+          raw_ctxt
+          ~kind:misbehaviour.kind
+          ~level
+          delegates
+      in
+      return pct
 end
diff --git a/src/proto_alpha/lib_protocol/test/helpers/state_ai_flags.mli b/src/proto_alpha/lib_protocol/test/helpers/state_ai_flags.mli
index 97015571ee00..4c8d303a46f9 100644
--- a/src/proto_alpha/lib_protocol/test/helpers/state_ai_flags.mli
+++ b/src/proto_alpha/lib_protocol/test/helpers/state_ai_flags.mli
@@ -53,5 +53,9 @@ module NS : sig
   (** Whatever the value of the flag is, this function returns the
       slashing value for a given double attestation *)
   val get_double_attestation_slashing_percentage :
-    Block.t -> State.t -> Protocol.Percentage.t
+    (Signature.public_key_hash * Protocol.Denunciations_repr.item) list ->
+    Block.t ->
+    State.t ->
+    Protocol.Misbehaviour_repr.t ->
+    Protocol.Percentage.t tzresult Lwt.t
 end
-- 
GitLab


From 90aaf20a7d2502a1a6278114ab688896f9b13919 Mon Sep 17 00:00:00 2001
From: Lucas Randazzo <lucas@nomadic-labs.com>
Date: Thu, 7 Mar 2024 10:30:36 +0100
Subject: [PATCH 10/10] Proto/tests: add Delayed_slashing module

---
 .../test/helpers/slashing_helpers.ml          | 20 +-----------------
 .../test/helpers/state_ai_flags.ml            | 21 +++++++++++++++++++
 .../test/helpers/state_ai_flags.mli           | 19 +++++++++++++++++
 .../integration/test_scenario_slashing.ml     |  2 +-
 4 files changed, 42 insertions(+), 20 deletions(-)

diff --git a/src/proto_alpha/lib_protocol/test/helpers/slashing_helpers.ml b/src/proto_alpha/lib_protocol/test/helpers/slashing_helpers.ml
index 77a72558f599..52fbe46f7ef3 100644
--- a/src/proto_alpha/lib_protocol/test/helpers/slashing_helpers.ml
+++ b/src/proto_alpha/lib_protocol/test/helpers/slashing_helpers.ml
@@ -262,25 +262,7 @@ let apply_all_slashes_at_cycle_end current_cycle (block_before_slash : Block.t)
     (state : State.t) : State.t tzresult Lwt.t =
   let open Lwt_result_syntax in
   let to_slash_later, to_slash_now =
-    if
-      not
-        state.constants
-          .Protocol.Alpha_context.Constants.Parametric.adaptive_issuance
-          .ns_enable
-    then ([], state.pending_slashes)
-    else
-      List.partition
-        (fun (_, Protocol.Denunciations_repr.{misbehaviour; _}) ->
-          let cycle =
-            Block.current_cycle_of_level
-              ~blocks_per_cycle:
-                state.constants
-                  .Protocol.Alpha_context.Constants.Parametric.blocks_per_cycle
-              ~current_level:
-                (Protocol.Raw_level_repr.to_int32 misbehaviour.level)
-          in
-          Protocol.Alpha_context.Cycle.(cycle = current_cycle))
-        state.pending_slashes
+    State_ai_flags.Delayed_slashing.partition_slashes state current_cycle
   in
   let* state, total_burnt =
     List.fold_left_es
diff --git a/src/proto_alpha/lib_protocol/test/helpers/state_ai_flags.ml b/src/proto_alpha/lib_protocol/test/helpers/state_ai_flags.ml
index 84cc0bebc605..e3a2a63e3fcb 100644
--- a/src/proto_alpha/lib_protocol/test/helpers/state_ai_flags.ml
+++ b/src/proto_alpha/lib_protocol/test/helpers/state_ai_flags.ml
@@ -184,6 +184,27 @@ module Autostake = struct
     else return state
 end
 
+module Delayed_slashing = struct
+  let enabled (state : State.t) = state.constants.adaptive_issuance.ns_enable
+
+  (* Returns a pair, fst is the delayed slashes, snd is the slashes to apply now *)
+  let partition_slashes state current_cycle =
+    if not (enabled state) then ([], state.pending_slashes)
+    else
+      List.partition
+        (fun (_, Protocol.Denunciations_repr.{misbehaviour; _}) ->
+          let cycle =
+            Block.current_cycle_of_level
+              ~blocks_per_cycle:
+                state.constants
+                  .Protocol.Alpha_context.Constants.Parametric.blocks_per_cycle
+              ~current_level:
+                (Protocol.Raw_level_repr.to_int32 misbehaviour.level)
+          in
+          Protocol.Alpha_context.Cycle.(cycle = current_cycle))
+        state.pending_slashes
+end
+
 module NS = struct
   let enabled (block : Block.t) (state : State.t) =
     AI.enabled block state && state.constants.adaptive_issuance.ns_enable
diff --git a/src/proto_alpha/lib_protocol/test/helpers/state_ai_flags.mli b/src/proto_alpha/lib_protocol/test/helpers/state_ai_flags.mli
index 4c8d303a46f9..04e28933d234 100644
--- a/src/proto_alpha/lib_protocol/test/helpers/state_ai_flags.mli
+++ b/src/proto_alpha/lib_protocol/test/helpers/state_ai_flags.mli
@@ -43,6 +43,25 @@ module Autostake : sig
   val run_at_cycle_end : Block.t -> State.t -> State.t tzresult
 end
 
+module Delayed_slashing : sig
+  (** This module takes care of choosing the denunciations that need to be
+      applied at the end of a cycle. It depends on the flag [ns_enable]. *)
+
+  (** [Delayed_slashing] is enabled iff [ns_enable = true]. *)
+  val enabled : State.t -> bool
+
+  (** [partition_slashes s cycle] returns a pair [(l1,l2)] of lists of slashes,
+      partitioned from the [state.pending_slashes]. [l2] is the list of slashes to
+      apply at the end of the given [cycle], and [l1] is the rest (which should
+      usually replace [state.pending_slashes])
+      *)
+  val partition_slashes :
+    State.t ->
+    Protocol.Alpha_context.Cycle.t ->
+    (Signature.Public_key_hash.t * Protocol.Denunciations_repr.item) list
+    * (Signature.Public_key_hash.t * Protocol.Denunciations_repr.item) list
+end
+
 module NS : sig
   (** This module takes care of the new adaptive slashing mechanism.*)
 
diff --git a/src/proto_alpha/lib_protocol/test/integration/test_scenario_slashing.ml b/src/proto_alpha/lib_protocol/test/integration/test_scenario_slashing.ml
index d85ad2ae0fd5..ab14e041bc7d 100644
--- a/src/proto_alpha/lib_protocol/test/integration/test_scenario_slashing.ml
+++ b/src/proto_alpha/lib_protocol/test/integration/test_scenario_slashing.ml
@@ -73,7 +73,7 @@ let test_simple_slash =
        --> next_cycle
        --> assert_failure
              (exec_unit (fun (_block, state) ->
-                  if state.State.constants.adaptive_issuance.ns_enable then
+                  if State_ai_flags.Delayed_slashing.enabled state then
                     failwith "ns_enable = true: slash not applied yet"
                   else return_unit)
              --> check_snapshot_balances "before slash")
-- 
GitLab