Merge pull request #135 from Lay3rLabs/release_stake/splitter

Implement release_stake for splitter

Author: ismellike

contracts/splitter/src/contract.rs

@@ -250,7 +250,8 @@ mod execute {
         nodes::events::{
             common::{GenericDistributeRewardsEvent, GenericWithdrawRewardsEvent, RestakingEvent},
             splitter::{
-                DistributeRewardsEvent, ForceUnstakeEvent, StakeEvent, WithdrawRewardsEvent,
+                DistributeRewardsEvent, ForceUnstakeEvent, ReleaseStakeEvent, StakeEvent,
+                WithdrawRewardsEvent,
             },
         },
     };
@@ -399,7 +400,11 @@ mod execute {
         for (consumer, _security_level) in consumer_security_map {
             let mut output = OUTPUTS.load(deps.storage, (consumer, &denom))?;
             output.amount = output.amount.checked_sub(amount)?;
-            OUTPUTS.save(deps.storage, (consumer, &denom), &output)?;
+            if output.amount.is_zero() {
+                OUTPUTS.remove(deps.storage, (consumer, &denom));
+            } else {
+                OUTPUTS.save(deps.storage, (consumer, &denom), &output)?;
+            }
 
             messages.push(consumer::force_unstake(consumer, &staker, &denom, amount)?)
         }
@@ -472,17 +477,61 @@ mod execute {
     }
 
     pub fn release_stake(
-        _deps: DepsMut,
+        deps: DepsMut,
         _env: Env,
         info: MessageInfo,
-        _staker: String,
-        _amount: Uint128,
-        _denom: String,
+        staker: String,
+        amount: Uint128,
+        denom: String,
     ) -> Result<Response, ContractError> {
         nonpayable(&info)?;
+        let staker = deps.api.addr_validate(&staker)?;
+        let releasing_consumer = info.sender;
+
+        // Update input tally
+        let mut input = INPUTS.load(deps.storage, &denom)?;
+        input.amount = input.amount.checked_sub(amount)?;
+        // Do not remove, because we hold other commitment info here
+        INPUTS.save(deps.storage, &denom, &input)?;
+
+        // Force unstake for all other consumers
+        let config = CONFIG.load(deps.storage)?;
+        let mut msgs = vec![];
+
+        let map = create_consumer_security_map(&config.flows, &denom);
+
+        for (consumer, _) in map {
+            if consumer != releasing_consumer {
+                msgs.push(consumer::force_unstake(consumer, &staker, &denom, amount)?);
+            }
+
+            // Update their output tallies
+            let mut output = OUTPUTS.load(deps.storage, (consumer, &denom))?;
+            output.amount = output.amount.checked_sub(amount)?;
+            if output.amount.is_zero() {
+                OUTPUTS.remove(deps.storage, (consumer, &denom));
+            } else {
+                OUTPUTS.save(deps.storage, (consumer, &denom), &output)?;
+            }
+        }
+
+        // Add release message to provider
+        msgs.push(provider::release_stake(
+            &config.provider,
+            &staker,
+            &denom,
+            amount,
+        )?);
+
+        let event = ReleaseStakeEvent {
+            common: RestakingEvent {
+                staker_addr: staker,
+                amount,
+                denom,
+            },
+        };
 
-        // TODO: add event for it
-        todo!("From Meeting: Logic here is tricky, so let's save this for later")
+        Ok(Response::new().add_messages(msgs).add_event(event.into()))
     }
 
     pub fn distribute_rewards(

contracts/splitter/src/tests/common.rs

@@ -1399,3 +1399,167 @@ where
     let total_power_response = weighting.total_power_at_height(None).unwrap();
     assert_eq!(total_power_response.power, power_response_2.power);
 }
+
+pub fn test_release_stake<C>(chain: C, denom: String)
+where
+    C: CwEnv + AltSigner,
+    C::Sender: Addressable,
+{
+    let commitment = Commitment {
+        denom: denom.clone(),
+        security: SecurityLevel {
+            max_penalty_percent: 25,
+            unbonding_period: 1000,
+        },
+    };
+
+    // Create 4 flows that all use the same reference token
+    let flows = vec![
+        SplitFlow {
+            consumer: nodes::Counterparty::Caller {},
+            shares: vec![commitment.clone()],
+        },
+        SplitFlow {
+            consumer: nodes::Counterparty::Caller {},
+            shares: vec![commitment.clone()],
+        },
+        SplitFlow {
+            consumer: nodes::Counterparty::Caller {},
+            shares: vec![commitment.clone()],
+        },
+        SplitFlow {
+            consumer: nodes::Counterparty::Caller {},
+            shares: vec![commitment.clone()],
+        },
+    ];
+
+    // Set up contracts
+    let weighting = TokenWeightingContract::new(TOKEN_WEIGHTING_CONTRACT_ID, chain.clone());
+    weighting.upload().unwrap();
+
+    let staking = TokenStakingContract::new("staking", chain.clone());
+    staking.upload().unwrap();
+
+    let mut fan_in = FanInContract::new(FAN_IN_CONTRACT_ID, chain.clone());
+    fan_in.upload().unwrap();
+
+    let splitter = SplitterContract::new(SPLITTER_CONTRACT_ID, chain.clone());
+    splitter.upload().unwrap();
+
+    // Configure flows with weighting contracts
+    let mut configured_flows = flows;
+    for flow in configured_flows.iter_mut() {
+        let terms = RelativeWeights {
+            reference: commitment.clone(),
+            secondary: vec![], // No secondary tokens in this test
+        };
+
+        flow.consumer = nodes::Counterparty::Create {
+            code_id: weighting.code_id().unwrap(),
+            msg: serde_json::to_value(token_weighting::InstantiateMsg {
+                terms,
+                reward_denoms: None,
+                provider: nodes::Counterparty::Caller {},
+                slasher: None,
+                admin: None,
+            })
+            .unwrap(),
+            label: "Token Weighting".to_string(),
+            admin: nodes::ContractAdmin::None(),
+        }
+    }
+
+    // Instantiate the splitter with fan-in as provider
+    let msg = InstantiateMsg {
+        provider: nodes::Counterparty::Create {
+            code_id: fan_in.code_id().unwrap(),
+            msg: serde_json::to_value(fan_in::InstantiateMsg {
+                manager: chain.sender().to_string(),
+                consumer: nodes::Counterparty::Caller {},
+            })
+            .unwrap(),
+            label: "Fan In".to_string(),
+            admin: nodes::ContractAdmin::None(),
+        },
+        flows: configured_flows,
+    };
+    splitter.instantiate(&msg, None, None).unwrap();
+
+    // Set up fan-in
+    let config = splitter.config().unwrap();
+    let provider = Addr::unchecked(config.provider);
+    fan_in.set_default_address(&provider);
+
+    // Set up staking contract
+    staking
+        .instantiate(
+            &rstk_token_staking::msg::InstantiateMsg {
+                commitment: commitment.clone(),
+                consumer: nodes::Counterparty::Existing {
+                    contract: fan_in.addr_str().unwrap(),
+                },
+            },
+            None,
+            None,
+        )
+        .unwrap();
+
+    fan_in
+        .add_input(staking.addr_str().unwrap(), commitment.clone(), None)
+        .unwrap();
+
+    // Get consumers for testing
+    let consumers_response = splitter.consumers().unwrap();
+    assert_eq!(consumers_response.consumers.len(), 4);
+
+    // Initial stake
+    staking.stake(&coins(1000, &denom)).unwrap();
+    chain.next_block().unwrap();
+
+    // Verify initial power for all consumers
+    for consumer_info in consumers_response.consumers.iter() {
+        weighting.set_address(&Addr::unchecked(&consumer_info.consumer));
+
+        let power = weighting
+            .voting_power_at_height(chain.sender(), None)
+            .unwrap();
+        assert_eq!(power.power, Uint128::new(1000));
+
+        let total_power = weighting.total_power_at_height(None).unwrap();
+        assert_eq!(total_power.power, Uint128::new(1000));
+    }
+
+    // Unstake and verify power update
+    staking.unstake(coin(400, denom.clone())).unwrap();
+    chain.next_block().unwrap();
+
+    // Verify updated power for all consumers
+    for consumer_info in consumers_response.consumers.iter() {
+        weighting.set_address(&Addr::unchecked(&consumer_info.consumer));
+
+        let power = weighting
+            .voting_power_at_height(chain.sender(), None)
+            .unwrap();
+        assert_eq!(power.power, Uint128::new(600));
+
+        let total_power = weighting.total_power_at_height(None).unwrap();
+        assert_eq!(total_power.power, Uint128::new(600));
+    }
+
+    // Complete unstaking process
+    chain.wait_seconds(1001).unwrap();
+    staking.withdraw_unbonded(None).unwrap();
+
+    // Final verification that power remains correct
+    for consumer_info in consumers_response.consumers.iter() {
+        weighting.set_address(&Addr::unchecked(&consumer_info.consumer));
+
+        let power = weighting
+            .voting_power_at_height(chain.sender(), None)
+            .unwrap();
+        assert_eq!(power.power, Uint128::new(600));
+
+        let total_power = weighting.total_power_at_height(None).unwrap();
+        assert_eq!(total_power.power, Uint128::new(600));
+    }
+}

contracts/splitter/src/tests/multi.rs

@@ -213,3 +213,14 @@ fn test_rewards_flow() {
         assert_eq!(final_balance, Uint128::new(3000));
     }
 }
+
+#[test]
+fn release_stake() {
+    let chain = MockBech32::new(BECH_PREFIX);
+
+    chain
+        .add_balance(&chain.sender, coins(10000, DENOM))
+        .unwrap();
+
+    super::common::test_release_stake(chain, DENOM.to_string());
+}

packages/apis/src/nodes/events/splitter.rs

@@ -189,6 +189,39 @@ impl From<DistributeRewardsEvent> for Event {
     }
 }
 
+#[derive(Debug, Clone, PartialEq)]
+pub struct ReleaseStakeEvent {
+    pub common: RestakingEvent,
+}
+
+impl TypedEvent for ReleaseStakeEvent {
+    const NAME: &'static str = "splitter_release_stake_event";
+}
+
+impl TryFrom<&Event> for ReleaseStakeEvent {
+    type Error = EventError;
+
+    fn try_from(event: &Event) -> Result<Self, EventError> {
+        if !Self::is_type(&event.ty) {
+            return Err(EventError::IncorrectEventType {
+                expected: Self::NAME.to_string(),
+                actual: event.clone().ty,
+            });
+        }
+
+        let common = RestakingEvent::try_from(event)?;
+
+        Ok(Self { common })
+    }
+}
+
+impl From<ReleaseStakeEvent> for Event {
+    fn from(value: ReleaseStakeEvent) -> Self {
+        let event = Event::new(ReleaseStakeEvent::NAME);
+        value.common.add_to_event(event)
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;