iota_transaction_builder/
stake.rs

1// Copyright (c) Mysten Labs, Inc.
2// Modifications Copyright (c) 2024 IOTA Stiftung
3// SPDX-License-Identifier: Apache-2.0
4
5use anyhow::{Ok, anyhow, ensure};
6use iota_types::{
7    IOTA_SYSTEM_PACKAGE_ID,
8    base_types::{IotaAddress, ObjectID, ObjectType},
9    governance::{ADD_STAKE_MUL_COIN_FUN_NAME, WITHDRAW_STAKE_FUN_NAME},
10    iota_system_state::IOTA_SYSTEM_MODULE_NAME,
11    programmable_transaction_builder::ProgrammableTransactionBuilder,
12    timelock::timelocked_staking::{
13        ADD_TIMELOCKED_STAKE_FUN_NAME, TIMELOCKED_STAKING_MODULE_NAME,
14        WITHDRAW_TIMELOCKED_STAKE_FUN_NAME,
15    },
16    transaction::{CallArg, Command, ObjectArg, TransactionData},
17};
18
19use crate::TransactionBuilder;
20
21impl TransactionBuilder {
22    /// Add stake to a validator's staking pool using multiple IOTA coins.
23    pub async fn request_add_stake(
24        &self,
25        signer: IotaAddress,
26        mut coins: Vec<ObjectID>,
27        amount: impl Into<Option<u64>>,
28        validator: IotaAddress,
29        gas: impl Into<Option<ObjectID>>,
30        gas_budget: u64,
31    ) -> anyhow::Result<TransactionData> {
32        let gas_price = self.0.get_reference_gas_price().await?;
33        let gas = self
34            .select_gas(signer, gas, gas_budget, coins.clone(), gas_price)
35            .await?;
36
37        let mut obj_vec = vec![];
38        let coin = coins
39            .pop()
40            .ok_or_else(|| anyhow!("Coins input should contain at lease one coin object."))?;
41        let (oref, coin_type) = self.get_object_ref_and_type(coin).await?;
42
43        let ObjectType::Struct(type_) = &coin_type else {
44            return Err(anyhow!("Provided object [{coin}] is not a move object."));
45        };
46        ensure!(
47            type_.is_coin(),
48            "Expecting either Coin<T> input coin objects. Received [{type_}]"
49        );
50
51        for coin in coins {
52            let (oref, type_) = self.get_object_ref_and_type(coin).await?;
53            ensure!(
54                type_ == coin_type,
55                "All coins should be the same type, expecting {coin_type}, got {type_}."
56            );
57            obj_vec.push(ObjectArg::ImmOrOwnedObject(oref))
58        }
59        obj_vec.push(ObjectArg::ImmOrOwnedObject(oref));
60
61        let pt = {
62            let mut builder = ProgrammableTransactionBuilder::new();
63            let arguments = vec![
64                builder.input(CallArg::IOTA_SYSTEM_MUT).unwrap(),
65                builder.make_obj_vec(obj_vec)?,
66                builder
67                    .input(CallArg::Pure(bcs::to_bytes(&amount.into())?))
68                    .unwrap(),
69                builder
70                    .input(CallArg::Pure(bcs::to_bytes(&validator)?))
71                    .unwrap(),
72            ];
73            builder.command(Command::move_call(
74                IOTA_SYSTEM_PACKAGE_ID,
75                IOTA_SYSTEM_MODULE_NAME.to_owned(),
76                ADD_STAKE_MUL_COIN_FUN_NAME.to_owned(),
77                vec![],
78                arguments,
79            ));
80            builder.finish()
81        };
82        Ok(TransactionData::new_programmable(
83            signer,
84            vec![gas],
85            pt,
86            gas_budget,
87            gas_price,
88        ))
89    }
90
91    /// Withdraw stake from a validator's staking pool.
92    pub async fn request_withdraw_stake(
93        &self,
94        signer: IotaAddress,
95        staked_iota: ObjectID,
96        gas: impl Into<Option<ObjectID>>,
97        gas_budget: u64,
98    ) -> anyhow::Result<TransactionData> {
99        let staked_iota = self.get_object_ref(staked_iota).await?;
100        let gas_price = self.0.get_reference_gas_price().await?;
101        let gas = self
102            .select_gas(signer, gas, gas_budget, vec![], gas_price)
103            .await?;
104        TransactionData::new_move_call(
105            signer,
106            IOTA_SYSTEM_PACKAGE_ID,
107            IOTA_SYSTEM_MODULE_NAME.to_owned(),
108            WITHDRAW_STAKE_FUN_NAME.to_owned(),
109            vec![],
110            gas,
111            vec![
112                CallArg::IOTA_SYSTEM_MUT,
113                CallArg::Object(ObjectArg::ImmOrOwnedObject(staked_iota)),
114            ],
115            gas_budget,
116            gas_price,
117        )
118    }
119
120    /// Add stake to a validator's staking pool using a timelocked IOTA coin.
121    pub async fn request_add_timelocked_stake(
122        &self,
123        signer: IotaAddress,
124        locked_balance: ObjectID,
125        validator: IotaAddress,
126        gas: ObjectID,
127        gas_budget: u64,
128    ) -> anyhow::Result<TransactionData> {
129        let gas_price = self.0.get_reference_gas_price().await?;
130        let gas = self
131            .select_gas(signer, Some(gas), gas_budget, vec![], gas_price)
132            .await?;
133
134        let (oref, locked_balance_type) = self.get_object_ref_and_type(locked_balance).await?;
135
136        let ObjectType::Struct(type_) = &locked_balance_type else {
137            anyhow::bail!("Provided object [{locked_balance}] is not a move object.");
138        };
139        ensure!(
140            type_.is_timelocked_balance(),
141            "Expecting either TimeLock<Balance<T>> input objects. Received [{type_}]"
142        );
143
144        let pt = {
145            let mut builder = ProgrammableTransactionBuilder::new();
146            let arguments = vec![
147                builder.input(CallArg::IOTA_SYSTEM_MUT)?,
148                builder.input(CallArg::Object(ObjectArg::ImmOrOwnedObject(oref)))?,
149                builder.input(CallArg::Pure(bcs::to_bytes(&validator)?))?,
150            ];
151            builder.command(Command::move_call(
152                IOTA_SYSTEM_PACKAGE_ID,
153                TIMELOCKED_STAKING_MODULE_NAME.to_owned(),
154                ADD_TIMELOCKED_STAKE_FUN_NAME.to_owned(),
155                vec![],
156                arguments,
157            ));
158            builder.finish()
159        };
160        Ok(TransactionData::new_programmable(
161            signer,
162            vec![gas],
163            pt,
164            gas_budget,
165            gas_price,
166        ))
167    }
168
169    /// Withdraw timelocked stake from a validator's staking pool.
170    pub async fn request_withdraw_timelocked_stake(
171        &self,
172        signer: IotaAddress,
173        timelocked_staked_iota: ObjectID,
174        gas: ObjectID,
175        gas_budget: u64,
176    ) -> anyhow::Result<TransactionData> {
177        let timelocked_staked_iota = self.get_object_ref(timelocked_staked_iota).await?;
178        let gas_price = self.0.get_reference_gas_price().await?;
179        let gas = self
180            .select_gas(signer, Some(gas), gas_budget, vec![], gas_price)
181            .await?;
182        TransactionData::new_move_call(
183            signer,
184            IOTA_SYSTEM_PACKAGE_ID,
185            TIMELOCKED_STAKING_MODULE_NAME.to_owned(),
186            WITHDRAW_TIMELOCKED_STAKE_FUN_NAME.to_owned(),
187            vec![],
188            gas,
189            vec![
190                CallArg::IOTA_SYSTEM_MUT,
191                CallArg::Object(ObjectArg::ImmOrOwnedObject(timelocked_staked_iota)),
192            ],
193            gas_budget,
194            gas_price,
195        )
196    }
197}