From f8c3af9e09d27d3af807eb83b866cf29325b09b4 Mon Sep 17 00:00:00 2001
From: Mehdi Bouaziz <mehdi.bouaziz@nomadic-labs.com>
Date: Wed, 7 Jun 2023 15:42:44 +0200
Subject: [PATCH 1/7] Proto/Apply: convert apply_stake to Lwt_result_syntax

---
 src/proto_alpha/lib_protocol/apply.ml | 32 +++++++++++++++------------
 1 file changed, 18 insertions(+), 14 deletions(-)

diff --git a/src/proto_alpha/lib_protocol/apply.ml b/src/proto_alpha/lib_protocol/apply.ml
index 6fb2d793b294..6b8421d66f77 100644
--- a/src/proto_alpha/lib_protocol/apply.ml
+++ b/src/proto_alpha/lib_protocol/apply.ml
@@ -321,23 +321,27 @@ let apply_transaction_to_implicit ~ctxt ~sender ~amount ~pkh ~before_operation =
   return (ctxt, result, [])
 
 let apply_stake ~ctxt ~sender ~amount ~destination ~before_operation =
+  let open Lwt_result_syntax in
   let contract = Contract.Implicit destination in
   (* Staking of zero is forbidden. *)
-  error_when Tez.(amount = zero) (Empty_transaction contract) >>?= fun () ->
-  error_unless
-    Signature.Public_key_hash.(sender = destination)
-    Invalid_self_transaction_destination
-  >>?= fun () ->
-  Contract.is_delegate ctxt sender >>=? fun is_delegate ->
-  error_unless is_delegate Staking_for_nondelegate_while_costaking_disabled
-  >>?= fun () ->
+  let*? () = error_when Tez.(amount = zero) (Empty_transaction contract) in
+  let*? () =
+    error_unless
+      Signature.Public_key_hash.(sender = destination)
+      Invalid_self_transaction_destination
+  in
+  let* is_delegate = Contract.is_delegate ctxt sender in
+  let*? () =
+    error_unless is_delegate Staking_for_nondelegate_while_costaking_disabled
+  in
   let delegate = sender in
-  Token.transfer
-    ctxt
-    (`Contract (Contract.Implicit sender))
-    (`Frozen_deposits delegate)
-    amount
-  >>=? fun (ctxt, balance_updates) ->
+  let* ctxt, balance_updates =
+    Token.transfer
+      ctxt
+      (`Contract (Contract.Implicit sender))
+      (`Frozen_deposits delegate)
+      amount
+  in
   (* Since [delegate] is an already existing delegate, it is already allocated. *)
   let allocated_destination_contract = false in
   let result =
-- 
GitLab


From ea6e22b249e030a111d5da9cb7e51d5b9650bdf5 Mon Sep 17 00:00:00 2001
From: Mehdi Bouaziz <mehdi.bouaziz@nomadic-labs.com>
Date: Wed, 7 Jun 2023 15:45:36 +0200
Subject: [PATCH 2/7] Proto/Apply/apply_stake: expend Contract.is_delegate

---
 src/proto_alpha/lib_protocol/apply.ml | 59 ++++++++++++++-------------
 1 file changed, 30 insertions(+), 29 deletions(-)

diff --git a/src/proto_alpha/lib_protocol/apply.ml b/src/proto_alpha/lib_protocol/apply.ml
index 6b8421d66f77..be68faf9e35b 100644
--- a/src/proto_alpha/lib_protocol/apply.ml
+++ b/src/proto_alpha/lib_protocol/apply.ml
@@ -330,35 +330,36 @@ let apply_stake ~ctxt ~sender ~amount ~destination ~before_operation =
       Signature.Public_key_hash.(sender = destination)
       Invalid_self_transaction_destination
   in
-  let* is_delegate = Contract.is_delegate ctxt sender in
-  let*? () =
-    error_unless is_delegate Staking_for_nondelegate_while_costaking_disabled
-  in
-  let delegate = sender in
-  let* ctxt, balance_updates =
-    Token.transfer
-      ctxt
-      (`Contract (Contract.Implicit sender))
-      (`Frozen_deposits delegate)
-      amount
-  in
-  (* Since [delegate] is an already existing delegate, it is already allocated. *)
-  let allocated_destination_contract = false in
-  let result =
-    Transaction_to_contract_result
-      {
-        storage = None;
-        lazy_storage_diff = None;
-        balance_updates;
-        ticket_receipt = [];
-        originated_contracts = [];
-        consumed_gas = Gas.consumed ~since:before_operation ~until:ctxt;
-        storage_size = Z.zero;
-        paid_storage_size_diff = Z.zero;
-        allocated_destination_contract;
-      }
-  in
-  return (ctxt, result, [])
+  let* delegate_opt = Contract.Delegate.find ctxt contract in
+  match delegate_opt with
+  | None -> tzfail Staking_for_nondelegate_while_costaking_disabled
+  | Some delegate when Signature.Public_key_hash.(delegate <> sender) ->
+      tzfail Staking_for_nondelegate_while_costaking_disabled
+  | Some delegate ->
+      let* ctxt, balance_updates =
+        Token.transfer
+          ctxt
+          (`Contract (Contract.Implicit sender))
+          (`Frozen_deposits delegate)
+          amount
+      in
+      (* Since [delegate] is an already existing delegate, it is already allocated. *)
+      let allocated_destination_contract = false in
+      let result =
+        Transaction_to_contract_result
+          {
+            storage = None;
+            lazy_storage_diff = None;
+            balance_updates;
+            ticket_receipt = [];
+            originated_contracts = [];
+            consumed_gas = Gas.consumed ~since:before_operation ~until:ctxt;
+            storage_size = Z.zero;
+            paid_storage_size_diff = Z.zero;
+            allocated_destination_contract;
+          }
+      in
+      return (ctxt, result, [])
 
 let apply_finalize_unstake ~ctxt ~sender ~amount ~destination ~before_operation
     =
-- 
GitLab


From a9facd59e24c8f07a87727ead581f0f46b483bfe Mon Sep 17 00:00:00 2001
From: Mehdi Bouaziz <mehdi.bouaziz@nomadic-labs.com>
Date: Wed, 7 Jun 2023 15:49:40 +0200
Subject: [PATCH 3/7] Proto/Apply/apply_stake: move to Staking

---
 src/proto_alpha/lib_protocol/apply.ml    | 6 +-----
 src/proto_alpha/lib_protocol/staking.ml  | 8 ++++++++
 src/proto_alpha/lib_protocol/staking.mli | 9 +++++++++
 3 files changed, 18 insertions(+), 5 deletions(-)

diff --git a/src/proto_alpha/lib_protocol/apply.ml b/src/proto_alpha/lib_protocol/apply.ml
index be68faf9e35b..a623e7956d67 100644
--- a/src/proto_alpha/lib_protocol/apply.ml
+++ b/src/proto_alpha/lib_protocol/apply.ml
@@ -337,11 +337,7 @@ let apply_stake ~ctxt ~sender ~amount ~destination ~before_operation =
       tzfail Staking_for_nondelegate_while_costaking_disabled
   | Some delegate ->
       let* ctxt, balance_updates =
-        Token.transfer
-          ctxt
-          (`Contract (Contract.Implicit sender))
-          (`Frozen_deposits delegate)
-          amount
+        Staking.stake ctxt ~sender ~delegate amount
       in
       (* Since [delegate] is an already existing delegate, it is already allocated. *)
       let allocated_destination_contract = false in
diff --git a/src/proto_alpha/lib_protocol/staking.ml b/src/proto_alpha/lib_protocol/staking.ml
index 9ed8e30864e9..dd87fc46faca 100644
--- a/src/proto_alpha/lib_protocol/staking.ml
+++ b/src/proto_alpha/lib_protocol/staking.ml
@@ -22,8 +22,16 @@
 (* DEALINGS IN THE SOFTWARE.                                                 *)
 (*                                                                           *)
 (*****************************************************************************)
+
 open Alpha_context
 
+let stake ctxt ~sender ~delegate amount =
+  Token.transfer
+    ctxt
+    (`Contract (Contract.Implicit sender))
+    (`Frozen_deposits delegate)
+    amount
+
 let finalize_unstake ctxt pkh =
   let open Lwt_result_syntax in
   let contract = Contract.Implicit pkh in
diff --git a/src/proto_alpha/lib_protocol/staking.mli b/src/proto_alpha/lib_protocol/staking.mli
index bf9b986201c0..5d7af413b97d 100644
--- a/src/proto_alpha/lib_protocol/staking.mli
+++ b/src/proto_alpha/lib_protocol/staking.mli
@@ -25,6 +25,15 @@
 
 open Alpha_context
 
+(** [stake ctxt ~sender ~delegate amount] add [amount] as [sender]'s stake
+    to [delegate]. *)
+val stake :
+  context ->
+  sender:public_key_hash ->
+  delegate:public_key_hash ->
+  Tez.t ->
+  (context * Receipt.balance_updates) tzresult Lwt.t
+
 (** [finalize_unstake ctxt pkh] performs the finalization of all unstake
     requests from [pkh] that can be finalized.
     An unstake request can be finalized if it is old enough, specifically the
-- 
GitLab


From 60814e05926d1f7808f89dd5626af5285c2bc821 Mon Sep 17 00:00:00 2001
From: Mehdi Bouaziz <mehdi.bouaziz@nomadic-labs.com>
Date: Wed, 7 Jun 2023 15:53:48 +0200
Subject: [PATCH 4/7] Proto/Staking: split finalize_unstake

---
 src/proto_alpha/lib_protocol/staking.ml | 34 ++++++++++++++-----------
 1 file changed, 19 insertions(+), 15 deletions(-)

diff --git a/src/proto_alpha/lib_protocol/staking.ml b/src/proto_alpha/lib_protocol/staking.ml
index dd87fc46faca..8b939b5dc89a 100644
--- a/src/proto_alpha/lib_protocol/staking.ml
+++ b/src/proto_alpha/lib_protocol/staking.ml
@@ -25,22 +25,8 @@
 
 open Alpha_context
 
-let stake ctxt ~sender ~delegate amount =
-  Token.transfer
-    ctxt
-    (`Contract (Contract.Implicit sender))
-    (`Frozen_deposits delegate)
-    amount
-
-let finalize_unstake ctxt pkh =
+let perform_finalizable_unstake_transfers ctxt contract finalizable =
   let open Lwt_result_syntax in
-  let contract = Contract.Implicit pkh in
-  let* ctxt, finalizable =
-    Unstake_requests
-    .prepare_finalize_unstake_and_save_remaining_unfinalizable_requests
-      ctxt
-      contract
-  in
   List.fold_left_es
     (fun (ctxt, balance_updates) (delegate, cycle, amount) ->
       let+ ctxt, new_balance_updates =
@@ -54,6 +40,17 @@ let finalize_unstake ctxt pkh =
     (ctxt, [])
     finalizable
 
+let finalize_unstake ctxt pkh =
+  let open Lwt_result_syntax in
+  let contract = Contract.Implicit pkh in
+  let* ctxt, finalizable =
+    Unstake_requests
+    .prepare_finalize_unstake_and_save_remaining_unfinalizable_requests
+      ctxt
+      contract
+  in
+  perform_finalizable_unstake_transfers ctxt contract finalizable
+
 let punish_delegate ctxt delegate level mistake ~rewarded =
   let open Lwt_result_syntax in
   let punish =
@@ -82,3 +79,10 @@ let punish_delegate ctxt delegate level mistake ~rewarded =
     Token.transfer_n ctxt to_reward (`Contract rewarded)
   in
   (ctxt, reward_balance_updates @ punish_balance_updates)
+
+let stake ctxt ~sender ~delegate amount =
+  Token.transfer
+    ctxt
+    (`Contract (Contract.Implicit sender))
+    (`Frozen_deposits delegate)
+    amount
-- 
GitLab


From a458861adc11266c02be7e0470d7b83e30748b12 Mon Sep 17 00:00:00 2001
From: Mehdi Bouaziz <mehdi.bouaziz@nomadic-labs.com>
Date: Wed, 7 Jun 2023 16:08:00 +0200
Subject: [PATCH 5/7] Proto/Staking: break
 prepare_finalize_unstake_and_save_remaining_unfinalizable_requests

---
 .../lib_protocol/alpha_context.mli            | 17 +++++++++++++++--
 src/proto_alpha/lib_protocol/staking.ml       | 13 ++++++-------
 .../lib_protocol/unstake_requests_storage.ml  | 19 +++++++------------
 .../lib_protocol/unstake_requests_storage.mli | 16 ++++++++++------
 4 files changed, 38 insertions(+), 27 deletions(-)

diff --git a/src/proto_alpha/lib_protocol/alpha_context.mli b/src/proto_alpha/lib_protocol/alpha_context.mli
index 5d03e0a3f2f3..9e7ca8685b32 100644
--- a/src/proto_alpha/lib_protocol/alpha_context.mli
+++ b/src/proto_alpha/lib_protocol/alpha_context.mli
@@ -4839,8 +4839,21 @@ end
 module Unstake_requests : sig
   type finalizable = (public_key_hash * Cycle.t * Tez.t) list
 
-  val prepare_finalize_unstake_and_save_remaining_unfinalizable_requests :
-    context -> Contract.t -> (context * finalizable) tzresult Lwt.t
+  type stored_requests = private {
+    delegate : public_key_hash;
+    requests : (Cycle.t * Tez.t) list;
+  }
+
+  type prepared_finalize_unstake = {
+    finalizable : finalizable;
+    unfinalizable : stored_requests;
+  }
+
+  val prepare_finalize_unstake :
+    context -> Contract.t -> prepared_finalize_unstake option tzresult Lwt.t
+
+  val update :
+    context -> Contract.t -> stored_requests -> context tzresult Lwt.t
 end
 
 (** This module re-exports definitions from {!Fees_storage}. *)
diff --git a/src/proto_alpha/lib_protocol/staking.ml b/src/proto_alpha/lib_protocol/staking.ml
index 8b939b5dc89a..9c189841dfbc 100644
--- a/src/proto_alpha/lib_protocol/staking.ml
+++ b/src/proto_alpha/lib_protocol/staking.ml
@@ -43,13 +43,12 @@ let perform_finalizable_unstake_transfers ctxt contract finalizable =
 let finalize_unstake ctxt pkh =
   let open Lwt_result_syntax in
   let contract = Contract.Implicit pkh in
-  let* ctxt, finalizable =
-    Unstake_requests
-    .prepare_finalize_unstake_and_save_remaining_unfinalizable_requests
-      ctxt
-      contract
-  in
-  perform_finalizable_unstake_transfers ctxt contract finalizable
+  let* prepared_opt = Unstake_requests.prepare_finalize_unstake ctxt contract in
+  match prepared_opt with
+  | None -> return (ctxt, [])
+  | Some {finalizable; unfinalizable} ->
+      let* ctxt = Unstake_requests.update ctxt contract unfinalizable in
+      perform_finalizable_unstake_transfers ctxt contract finalizable
 
 let punish_delegate ctxt delegate level mistake ~rewarded =
   let open Lwt_result_syntax in
diff --git a/src/proto_alpha/lib_protocol/unstake_requests_storage.ml b/src/proto_alpha/lib_protocol/unstake_requests_storage.ml
index 710aba989bb2..012650903e51 100644
--- a/src/proto_alpha/lib_protocol/unstake_requests_storage.ml
+++ b/src/proto_alpha/lib_protocol/unstake_requests_storage.ml
@@ -26,9 +26,14 @@
 type finalizable =
   (Signature.Public_key_hash.t * Cycle_repr.t * Tez_repr.t) list
 
+type stored_requests = Storage.Unstake_request.t = {
+  delegate : Signature.Public_key_hash.t;
+  requests : (Cycle_repr.t * Tez_repr.t) list;
+}
+
 type prepared_finalize_unstake = {
   finalizable : finalizable;
-  unfinalizable : Storage.Unstake_request.t;
+  unfinalizable : stored_requests;
 }
 
 let z100 = Z.of_int 100
@@ -102,14 +107,4 @@ let prepare_finalize_unstake ctxt contract =
           in
           return_some {finalizable; unfinalizable})
 
-let prepare_finalize_unstake_and_save_remaining_unfinalizable_requests ctxt
-    contract =
-  let open Lwt_result_syntax in
-  let* prepared_opt = prepare_finalize_unstake ctxt contract in
-  match prepared_opt with
-  | None -> return (ctxt, [])
-  | Some {finalizable; unfinalizable} ->
-      let+ ctxt =
-        Storage.Contract.Unstake_requests.update ctxt contract unfinalizable
-      in
-      (ctxt, finalizable)
+let update = Storage.Contract.Unstake_requests.update
diff --git a/src/proto_alpha/lib_protocol/unstake_requests_storage.mli b/src/proto_alpha/lib_protocol/unstake_requests_storage.mli
index f32cbd55e020..6c6043aa52e3 100644
--- a/src/proto_alpha/lib_protocol/unstake_requests_storage.mli
+++ b/src/proto_alpha/lib_protocol/unstake_requests_storage.mli
@@ -31,9 +31,14 @@
 type finalizable =
   (Signature.Public_key_hash.t * Cycle_repr.t * Tez_repr.t) list
 
+type stored_requests = Storage.Unstake_request.t = {
+  delegate : Signature.Public_key_hash.t;
+  requests : (Cycle_repr.t * Tez_repr.t) list;
+}
+
 type prepared_finalize_unstake = {
   finalizable : finalizable;
-  unfinalizable : Storage.Unstake_request.t;
+  unfinalizable : stored_requests;
 }
 
 (** [prepare_finalize_unstake ctxt contract] preprocesses a [finalize_unstake]
@@ -49,10 +54,9 @@ val prepare_finalize_unstake :
   Contract_repr.t ->
   prepared_finalize_unstake option tzresult Lwt.t
 
-(** [prepare_finalize_unstake_and_save_remaining_unfinalizable_requests ctxt contract]
-    calls [prepare_finalize_unstake], saves the remaining [unfinalizable]
-    requests and returns the [finalizable] ones. *)
-val prepare_finalize_unstake_and_save_remaining_unfinalizable_requests :
+(** [update ctxt contract requests] updates unstake requests for [contract]. *)
+val update :
   Raw_context.t ->
   Contract_repr.t ->
-  (Raw_context.t * finalizable) tzresult Lwt.t
+  stored_requests ->
+  Raw_context.t tzresult Lwt.t
-- 
GitLab


From 0ab2419a293d3a83bcc8519c78c1b58055b87636 Mon Sep 17 00:00:00 2001
From: Mehdi Bouaziz <mehdi.bouaziz@nomadic-labs.com>
Date: Wed, 7 Jun 2023 16:12:48 +0200
Subject: [PATCH 6/7] Proto/Staking: finalize_unstake_and_check

---
 src/proto_alpha/lib_protocol/staking.ml | 7 ++++++-
 1 file changed, 6 insertions(+), 1 deletion(-)

diff --git a/src/proto_alpha/lib_protocol/staking.ml b/src/proto_alpha/lib_protocol/staking.ml
index 9c189841dfbc..239b1e82748d 100644
--- a/src/proto_alpha/lib_protocol/staking.ml
+++ b/src/proto_alpha/lib_protocol/staking.ml
@@ -40,16 +40,21 @@ let perform_finalizable_unstake_transfers ctxt contract finalizable =
     (ctxt, [])
     finalizable
 
-let finalize_unstake ctxt pkh =
+let finalize_unstake_and_check ~check_unfinalizable ctxt pkh =
   let open Lwt_result_syntax in
   let contract = Contract.Implicit pkh in
   let* prepared_opt = Unstake_requests.prepare_finalize_unstake ctxt contract in
   match prepared_opt with
   | None -> return (ctxt, [])
   | Some {finalizable; unfinalizable} ->
+      let* () = check_unfinalizable unfinalizable in
       let* ctxt = Unstake_requests.update ctxt contract unfinalizable in
       perform_finalizable_unstake_transfers ctxt contract finalizable
 
+let finalize_unstake ctxt pkh =
+  let check_unfinalizable _unfinalizable = Lwt_result_syntax.return_unit in
+  finalize_unstake_and_check ~check_unfinalizable ctxt pkh
+
 let punish_delegate ctxt delegate level mistake ~rewarded =
   let open Lwt_result_syntax in
   let punish =
-- 
GitLab


From f515e1f158223d076d410b781d55143444533df6 Mon Sep 17 00:00:00 2001
From: Mehdi Bouaziz <mehdi.bouaziz@nomadic-labs.com>
Date: Wed, 7 Jun 2023 16:22:00 +0200
Subject: [PATCH 7/7] Proto/Staking: on stake, do a finalize_unstake but fail
 if there are unfinalizable unstake requests to another delegate

---
 src/proto_alpha/lib_protocol/staking.ml | 50 ++++++++++++++++++++++---
 1 file changed, 45 insertions(+), 5 deletions(-)

diff --git a/src/proto_alpha/lib_protocol/staking.ml b/src/proto_alpha/lib_protocol/staking.ml
index 239b1e82748d..f87d5ffc59ca 100644
--- a/src/proto_alpha/lib_protocol/staking.ml
+++ b/src/proto_alpha/lib_protocol/staking.ml
@@ -25,6 +25,30 @@
 
 open Alpha_context
 
+type error +=
+  | Cannot_stake_with_unfinalizable_unstake_requests_to_another_delegate
+
+let () =
+  let description =
+    "A contract tries to stake to its delegate while having unstake requests \
+     to a previous delegate that cannot be finalized yet. Try again in a later \
+     cycle (no more than preserved_cycles + max_slashing_period)."
+  in
+  register_error_kind
+    `Permanent
+    ~id:
+      "operation.cannot_stake_with_unfinalizable_unstake_requests_to_another_delegate"
+    ~title:
+      "Cannot stake with unfinalizable unstake requests to another delegate"
+    ~description
+    Data_encoding.unit
+    (function
+      | Cannot_stake_with_unfinalizable_unstake_requests_to_another_delegate ->
+          Some ()
+      | _ -> None)
+    (fun () ->
+      Cannot_stake_with_unfinalizable_unstake_requests_to_another_delegate)
+
 let perform_finalizable_unstake_transfers ctxt contract finalizable =
   let open Lwt_result_syntax in
   List.fold_left_es
@@ -85,8 +109,24 @@ let punish_delegate ctxt delegate level mistake ~rewarded =
   (ctxt, reward_balance_updates @ punish_balance_updates)
 
 let stake ctxt ~sender ~delegate amount =
-  Token.transfer
-    ctxt
-    (`Contract (Contract.Implicit sender))
-    (`Frozen_deposits delegate)
-    amount
+  let open Lwt_result_syntax in
+  let check_unfinalizable
+      Unstake_requests.{delegate = unstake_delegate; requests} =
+    match requests with
+    | [] -> return_unit
+    | _ :: _ ->
+        fail_when
+          Signature.Public_key_hash.(delegate <> unstake_delegate)
+          Cannot_stake_with_unfinalizable_unstake_requests_to_another_delegate
+  in
+  let* ctxt, finalize_balance_updates =
+    finalize_unstake_and_check ~check_unfinalizable ctxt sender
+  in
+  let* ctxt, stake_balance_updates =
+    Token.transfer
+      ctxt
+      (`Contract (Contract.Implicit sender))
+      (`Frozen_deposits delegate)
+      amount
+  in
+  return (ctxt, stake_balance_updates @ finalize_balance_updates)
-- 
GitLab