stake-contract: allow stake topup and partial unstake

consensus/src/user/provisioners.rs

@@ -112,7 +112,7 @@ impl Provisioners {
 
     /// Adds a provisioner with stake.
     ///
-    /// It appends the stake if the given provisioner already exists.
+    /// If the provisioner already exists, no action is performed.
     pub fn add_member_with_stake(
         &mut self,
         pubkey_bls: PublicKey,
@@ -136,6 +136,28 @@ impl Provisioners {
         self.members.insert(pubkey_bls, stake)
     }
 
+    /// Subtract `amount` from a staker, returning the stake left
+    ///
+    /// Return None if the entry was not found or `amount` is higher than
+    /// current stake
+    pub fn sub_stake(
+        &mut self,
+        pubkey_bls: &PublicKey,
+        amount: u64,
+    ) -> Option<u64> {
+        let stake = self.members.get_mut(pubkey_bls)?;
+        if stake.value() < amount {
+            None
+        } else {
+            stake.subtract(amount);
+            let left = stake.value();
+            if left == 0 {
+                self.members.remove(pubkey_bls);
+            }
+            Some(left)
+        }
+    }
+
     pub fn remove_stake(&mut self, pubkey_bls: &PublicKey) -> Option<Stake> {
         self.members.remove(pubkey_bls)
     }

consensus/src/user/stake.rs

@@ -52,4 +52,12 @@ impl Stake {
     pub fn change_eligibility(&mut self, new_value: u64) {
         self.eligible_since = new_value;
     }
+
+    /// Add an amount to the stake
+    pub fn add(&mut self, add: u64) {
+        // The value is the LUX representation of the stake,
+        // someone should own more than 18B DUSK in order to overflow
+        // (way more than our max supply)
+        self.value += add
+    }
 }

contracts/stake/src/state.rs

@@ -56,28 +56,17 @@ impl StakeState {
 
     pub fn stake(&mut self, stake: Stake) {
         let value = stake.value();
-        let keys = *stake.keys();
-        let account = &keys.account;
         let signature = *stake.signature();
 
         if stake.chain_id() != self.chain_id() {
             panic!("The stake must target the correct chain");
         }
 
-        let prev_stake = self.get_stake(account).copied();
-        let (loaded_stake, loaded_keys) = self.load_or_create_stake_mut(&keys);
+        let account = stake.keys().account;
+        let prev_stake = self.get_stake(&stake.keys().account).copied();
+        let (loaded_stake, keys) = self.load_or_create_stake_mut(stake.keys());
 
-        if loaded_stake.amount.is_some() {
-            panic!("Can't stake twice for the same key");
-        }
-
-        // Update the funds key with the newly provided one
-        // This operation will rollback if the signature is invalid
-        *loaded_keys = keys;
-
-        // ensure the stake is at least the minimum and that there isn't an
-        // amount staked already
-        if value < MINIMUM_STAKE {
+        if loaded_stake.amount.is_none() && value < MINIMUM_STAKE {
             panic!("The staked value is lower than the minimum amount!");
         }
 
@@ -95,16 +84,32 @@ impl StakeState {
             rusk_abi::call::<_, ()>(TRANSFER_CONTRACT, "deposit", &value)
                 .expect("Depositing funds into contract should succeed");
 
+        let block_height = rusk_abi::block_height();
         // update the state accordingly
-        loaded_stake.amount =
-            Some(StakeAmount::new(value, rusk_abi::block_height()));
-
-        rusk_abi::emit("stake", StakeEvent { keys, value });
+        let stake_event = match &mut loaded_stake.amount {
+            Some(amount) => {
+                let locked = if block_height >= amount.eligibility {
+                    value / 10
+                } else {
+                    // No penalties applied if the stake is not eligible yet
+                    0
+                };
+                let value = value - locked;
+                amount.locked += locked;
+                amount.value += value;
+                StakeEvent::new(*keys, value).locked(locked)
+            }
+            amount => {
+                let _ = amount.insert(StakeAmount::new(value, block_height));
+                StakeEvent::new(*keys, value)
+            }
+        };
+        rusk_abi::emit("stake", stake_event);
 
-        let key = account.to_bytes();
+        let key = keys.account.to_bytes();
         self.previous_block_state
             .entry(key)
-            .or_insert_with(|| (prev_stake, *account));
+            .or_insert((prev_stake, account));
     }
 
     pub fn unstake(&mut self, unstake: Withdraw) {
@@ -118,15 +123,15 @@ impl StakeState {
             .expect("A stake should exist in the map to be unstaked!");
         let prev_stake = Some(*loaded_stake);
 
-        // ensure there is a value staked, and that the withdrawal is exactly
-        // the same amount
+        // ensure there is a value staked, and that the withdrawal is not
+        // greater than the available funds
         let stake = loaded_stake
             .amount
-            .as_ref()
+            .as_mut()
             .expect("There must be an amount to unstake");
 
-        if value != stake.total_funds() {
-            panic!("Value withdrawn different from staked amount");
+        if value > stake.total_funds() {
+            panic!("Value to unstake higher than the staked amount");
         }
 
         // check signature is correct
@@ -144,13 +149,26 @@ impl StakeState {
             rusk_abi::call(TRANSFER_CONTRACT, "withdraw", transfer_withdraw)
                 .expect("Withdrawing stake should succeed");
 
-        // update the state accordingly
-        loaded_stake.amount = None;
-
-        rusk_abi::emit("unstake", StakeEvent { keys: *keys, value });
-
-        if loaded_stake.reward == 0 {
-            self.stakes.remove(&unstake.account().to_bytes());
+        let stake_event = if value > stake.value {
+            let from_locked = value - stake.value;
+            let from_stake = stake.value;
+            stake.value = 0;
+            stake.locked -= from_locked;
+            StakeEvent::new(*keys, from_stake).locked(from_locked)
+        } else {
+            stake.value -= value;
+            StakeEvent::new(*keys, value)
+        };
+
+        rusk_abi::emit("unstake", stake_event);
+        if stake.total_funds() == 0 {
+            // update the state accordingly
+            loaded_stake.amount = None;
+            if loaded_stake.reward == 0 {
+                self.stakes.remove(&unstake.account().to_bytes());
+            }
+        } else if stake.total_funds() < MINIMUM_STAKE {
+            panic!("Stake left is lower than minimum stake");
         }
 
         let key = account.to_bytes();
@@ -197,7 +215,7 @@ impl StakeState {
 
         // update the state accordingly
         loaded_stake.reward -= value;
-        rusk_abi::emit("withdraw", StakeEvent { keys: *keys, value });
+        rusk_abi::emit("withdraw", StakeEvent::new(*keys, value));
 
         if loaded_stake.reward == 0 && loaded_stake.amount.is_none() {
             self.stakes.remove(&account.to_bytes());
@@ -227,23 +245,25 @@ impl StakeState {
         self.stakes.insert(keys.account.to_bytes(), (stake, keys));
     }
 
-    /// Gets a mutable reference to the stake of a given `keys`. If said stake
-    /// doesn't exist, a default one is inserted and a mutable reference
-    /// returned.
+    /// Gets a mutable reference to the stake of a given `keys`.
+    ///
+    /// If said stake doesn't exist, a default one is inserted and a mutable
+    /// reference returned.
+    ///
+    /// # Panics
+    /// Panics if the provided keys doesn't match the existing (if any)
     pub(crate) fn load_or_create_stake_mut(
         &mut self,
         keys: &StakeKeys,
     ) -> &mut (StakeData, StakeKeys) {
         let key = keys.account.to_bytes();
-        let is_missing = self.stakes.get(&key).is_none();
-
-        if is_missing {
-            let stake = StakeData::EMPTY;
-            self.stakes.insert(key, (stake, *keys));
-        }
 
-        // SAFETY: unwrap is ok since we're sure we inserted an element
-        self.stakes.get_mut(&key).unwrap()
+        self.stakes
+            .entry(key)
+            .and_modify(|(_, loaded_keys)| {
+                assert_eq!(keys, loaded_keys, "Keys mismatch")
+            })
+            .or_insert_with(|| (StakeData::EMPTY, *keys))
     }
 
     /// Rewards a `account` with the given `value`.

execution-core/src/stake.rs

@@ -9,7 +9,7 @@
 use alloc::vec::Vec;
 
 use bytecheck::CheckBytes;
-use dusk_bytes::{DeserializableSlice, Serializable, Write};
+use dusk_bytes::Serializable;
 use rkyv::{Archive, Deserialize, Serialize};
 
 use crate::signatures::bls::{
@@ -197,8 +197,41 @@ impl Withdraw {
 pub struct StakeEvent {
     /// Keys associated to the event.
     pub keys: StakeKeys,
-    /// Value of the relevant operation, be it `stake`, `unstake`,`withdraw`
+    /// Effective value of the relevant operation, be it `stake`,
+    /// `unstake`,`withdraw`
     pub value: u64,
+    /// The locked amount involved in the operation (e.g., for `stake` or
+    /// `unstake`). Defaults to zero for operations that do not involve
+    /// locking.
+    pub locked: u64,
+}
+
+impl StakeEvent {
+    /// Creates a new `StakeEvent` with the specified keys and value.
+    ///
+    /// ### Parameters
+    /// - `keys`: The keys associated with the stake event.
+    /// - `value`: The effective value of the operation (e.g., `stake`,
+    ///   `unstake`, `withdraw`).
+    ///
+    /// The `locked` amount is initialized to zero by default.
+    #[must_use]
+    pub fn new(keys: StakeKeys, value: u64) -> Self {
+        Self {
+            keys,
+            value,
+            locked: 0,
+        }
+    }
+    /// Sets the locked amount for the `StakeEvent`.
+    ///
+    /// ### Parameters
+    /// - `locked`: The locked amount associated with the operation.
+    #[must_use]
+    pub fn locked(mut self, locked: u64) -> Self {
+        self.locked = locked;
+        self
+    }
 }
 
 /// Event emitted after a slash operation is performed.
@@ -324,67 +357,17 @@ impl StakeData {
     pub const fn eligibility_from_height(block_height: u64) -> u64 {
         StakeAmount::eligibility_from_height(block_height)
     }
-}
-
-const STAKE_DATA_SIZE: usize =
-    u8::SIZE + StakeAmount::SIZE + u64::SIZE + u8::SIZE + u8::SIZE;
-
-impl Serializable<STAKE_DATA_SIZE> for StakeData {
-    type Error = dusk_bytes::Error;
-
-    fn from_bytes(buf: &[u8; Self::SIZE]) -> Result<Self, Self::Error>
-    where
-        Self: Sized,
-    {
-        let mut buf = &buf[..];
-
-        // if the tag is zero we skip the bytes
-        let tag = u8::from_reader(&mut buf)?;
-        let amount = match tag {
-            0 => {
-                buf = &buf[..StakeAmount::SIZE];
-                None
-            }
-            _ => Some(StakeAmount::from_reader(&mut buf)?),
-        };
-
-        let reward = u64::from_reader(&mut buf)?;
 
-        let faults = u8::from_reader(&mut buf)?;
-        let hard_faults = u8::from_reader(&mut buf)?;
-
-        Ok(Self {
-            amount,
-            reward,
-            faults,
-            hard_faults,
-        })
-    }
-
-    #[allow(unused_must_use)]
-    fn to_bytes(&self) -> [u8; Self::SIZE] {
-        const ZERO_AMOUNT: [u8; StakeAmount::SIZE] = [0u8; StakeAmount::SIZE];
-
-        let mut buf = [0u8; Self::SIZE];
-        let mut writer = &mut buf[..];
-
-        match &self.amount {
-            None => {
-                writer.write(&0u8.to_bytes());
-                writer.write(&ZERO_AMOUNT);
-            }
-            Some(amount) => {
-                writer.write(&1u8.to_bytes());
-                writer.write(&amount.to_bytes());
-            }
-        }
-
-        writer.write(&self.reward.to_bytes());
-
-        writer.write(&self.faults.to_bytes());
-        writer.write(&self.hard_faults.to_bytes());
-
-        buf
+    /// Check if there is no amount left to withdraw
+    ///
+    /// Return true if both stake and rewards are 0
+    pub fn is_empty(&self) -> bool {
+        let stake = self
+            .amount
+            .as_ref()
+            .map(StakeAmount::total_funds)
+            .unwrap_or_default();
+        self.reward + stake == 0
     }
 }
 
@@ -441,41 +424,6 @@ impl StakeAmount {
     }
 }
 
-const STAKE_AMOUNT_SIZE: usize = u64::SIZE + u64::SIZE + u64::SIZE;
-
-impl Serializable<STAKE_AMOUNT_SIZE> for StakeAmount {
-    type Error = dusk_bytes::Error;
-
-    fn from_bytes(buf: &[u8; Self::SIZE]) -> Result<Self, Self::Error>
-    where
-        Self: Sized,
-    {
-        let mut buf = &buf[..];
-
-        let value = u64::from_reader(&mut buf)?;
-        let locked = u64::from_reader(&mut buf)?;
-        let eligibility = u64::from_reader(&mut buf)?;
-
-        Ok(Self {
-            value,
-            locked,
-            eligibility,
-        })
-    }
-
-    #[allow(unused_must_use)]
-    fn to_bytes(&self) -> [u8; Self::SIZE] {
-        let mut buf = [0u8; Self::SIZE];
-        let mut writer = &mut buf[..];
-
-        writer.write(&self.value.to_bytes());
-        writer.write(&self.locked.to_bytes());
-        writer.write(&self.eligibility.to_bytes());
-
-        buf
-    }
-}
-
 /// Used in a `reward` call to reward a given account with an amount of Dusk,
 /// and emitted as an event, once a reward succeeds.
 #[derive(Debug, Clone, PartialEq, Eq, Archive, Serialize, Deserialize)]