feat: added a draft timelocked staking api implementation

crates/sui-json-rpc-api/src/transaction_builder.rs

@@ -249,4 +249,36 @@ pub trait TransactionBuilder {
         /// the gas budget, the transaction will fail if the gas cost exceed the budget
         gas_budget: BigInt<u64>,
     ) -> RpcResult<TransactionBlockBytes>;
+
+    /// Add timelocked stake to a validator's staking pool using multiple balances and amount.
+    #[method(name = "requestAddTimelockedStake")]
+    async fn request_add_timelocked_stake(
+        &self,
+        /// the transaction signer's Sui address
+        signer: SuiAddress,
+        /// TimeLock<Balance<SUI>> object to stake
+        locked_balances: Vec<ObjectID>,
+        /// stake amount
+        amount: Option<BigInt<u64>>,
+        /// the validator's Sui address
+        validator: SuiAddress,
+        /// gas object to be used in this transaction
+        gas: ObjectID,
+        /// the gas budget, the transaction will fail if the gas cost exceed the budget
+        gas_budget: BigInt<u64>,
+    ) -> RpcResult<TransactionBlockBytes>;
+
+    /// Withdraw timelocked stake from a validator's staking pool.
+    #[method(name = "requestWithdrawTimelockedStake")]
+    async fn request_withdraw_timelocked_stake(
+        &self,
+        /// the transaction signer's Sui address
+        signer: SuiAddress,
+        /// TimelockedStakedSui object ID
+        timelocked_staked_sui: ObjectID,
+        /// gas object to be used in this transaction
+        gas: ObjectID,
+        /// the gas budget, the transaction will fail if the gas cost exceed the budget
+        gas_budget: BigInt<u64>,
+    ) -> RpcResult<TransactionBlockBytes>;
 }

crates/sui-json-rpc/src/transaction_builder_api.rs

@@ -314,6 +314,44 @@ impl TransactionBuilderServer for TransactionBuilderApi {
                 .await?,
         )?)
     }
+
+    async fn request_add_timelocked_stake(
+        &self,
+        signer: SuiAddress,
+        locked_balances: Vec<ObjectID>,
+        amount: Option<BigInt<u64>>,
+        validator: SuiAddress,
+        gas: ObjectID,
+        gas_budget: BigInt<u64>,
+    ) -> RpcResult<TransactionBlockBytes> {
+        let amount = amount.map(|a| *a);
+        Ok(TransactionBlockBytes::from_data(
+            self.0
+                .request_add_timelocked_stake(
+                    signer,
+                    locked_balances,
+                    amount,
+                    validator,
+                    gas,
+                    *gas_budget,
+                )
+                .await?,
+        )?)
+    }
+
+    async fn request_withdraw_timelocked_stake(
+        &self,
+        signer: SuiAddress,
+        timelocked_staked_sui: ObjectID,
+        gas: ObjectID,
+        gas_budget: BigInt<u64>,
+    ) -> RpcResult<TransactionBlockBytes> {
+        Ok(TransactionBlockBytes::from_data(
+            self.0
+                .request_withdraw_timelocked_stake(signer, timelocked_staked_sui, gas, *gas_budget)
+                .await?,
+        )?)
+    }
 }
 
 impl SuiRpcModule for TransactionBuilderApi {

crates/sui-source-validation-service/tests/tests.rs

@@ -5,10 +5,10 @@ use expect_test::expect;
 use reqwest::Client;
 use std::fs;
 use std::io::Read;
-#[cfg(target_os = "windows")]
-use std::os::windows::fs::FileExt;
 #[cfg(not(target_os = "windows"))]
 use std::os::unix::fs::FileExt;
+#[cfg(target_os = "windows")]
+use std::os::windows::fs::FileExt;
 use std::path::PathBuf;
 use std::sync::{Arc, RwLock};
 use sui::client_commands::{SuiClientCommandResult, SuiClientCommands};

crates/sui-transaction-builder/src/lib.rs

@@ -29,10 +29,16 @@ use sui_types::move_package::MovePackage;
 use sui_types::object::{Object, Owner};
 use sui_types::programmable_transaction_builder::ProgrammableTransactionBuilder;
 use sui_types::sui_system_state::SUI_SYSTEM_MODULE_NAME;
+use sui_types::timelocked_staking::{
+    ADD_TIMELOCKED_STAKE_MUL_BAL_FUN_NAME, TIMELOCKED_STAKING_MODULE_NAME,
+    WITHDRAW_TIMELOCKED_STAKE_FUN_NAME,
+};
 use sui_types::transaction::{
     Argument, CallArg, Command, InputObjectKind, ObjectArg, TransactionData, TransactionKind,
 };
-use sui_types::{coin, fp_ensure, SUI_FRAMEWORK_PACKAGE_ID, SUI_SYSTEM_PACKAGE_ID};
+use sui_types::{
+    coin, fp_ensure, STARDUST_PACKAGE_ID, SUI_FRAMEWORK_PACKAGE_ID, SUI_SYSTEM_PACKAGE_ID,
+};
 
 #[async_trait]
 pub trait DataReader {
@@ -775,6 +781,104 @@ impl TransactionBuilder {
         )
     }
 
+    pub async fn request_add_timelocked_stake(
+        &self,
+        signer: SuiAddress,
+        mut locked_balances: Vec<ObjectID>,
+        amount: Option<u64>,
+        validator: SuiAddress,
+        gas: ObjectID,
+        gas_budget: u64,
+    ) -> anyhow::Result<TransactionData> {
+        let gas_price = self.0.get_reference_gas_price().await?;
+        let gas = self
+            .select_gas(signer, Some(gas), gas_budget, vec![], gas_price)
+            .await?;
+
+        let mut obj_vec = vec![];
+        let locked_balance = locked_balances
+            .pop()
+            .ok_or_else(|| anyhow!("Locked balances input should contain at lease one object."))?;
+        let (oref, locked_balance_type) = self.get_object_ref_and_type(locked_balance).await?;
+
+        let ObjectType::Struct(type_) = &locked_balance_type else {
+            return Err(anyhow!(
+                "Provided object [{locked_balance}] is not a move object."
+            ));
+        };
+        ensure!(
+            type_.is_timelocked_balance(),
+            "Expecting either TimeLock<Balance<T>> input objects. Received [{type_}]"
+        );
+
+        for locked_balance in locked_balances {
+            let (oref, type_) = self.get_object_ref_and_type(locked_balance).await?;
+            ensure!(
+                type_ == locked_balance_type,
+                "All coins should be the same type, expecting {locked_balance_type}, got {type_}."
+            );
+            obj_vec.push(ObjectArg::ImmOrOwnedObject(oref))
+        }
+        obj_vec.push(ObjectArg::ImmOrOwnedObject(oref));
+
+        let pt = {
+            let mut builder = ProgrammableTransactionBuilder::new();
+            let arguments = vec![
+                builder.input(CallArg::SUI_SYSTEM_MUT).unwrap(),
+                builder.make_obj_vec(obj_vec)?,
+                builder
+                    .input(CallArg::Pure(bcs::to_bytes(&amount)?))
+                    .unwrap(),
+                builder
+                    .input(CallArg::Pure(bcs::to_bytes(&validator)?))
+                    .unwrap(),
+            ];
+            builder.command(Command::move_call(
+                STARDUST_PACKAGE_ID,
+                TIMELOCKED_STAKING_MODULE_NAME.to_owned(),
+                ADD_TIMELOCKED_STAKE_MUL_BAL_FUN_NAME.to_owned(),
+                vec![],
+                arguments,
+            ));
+            builder.finish()
+        };
+        Ok(TransactionData::new_programmable(
+            signer,
+            vec![gas],
+            pt,
+            gas_budget,
+            gas_price,
+        ))
+    }
+
+    pub async fn request_withdraw_timelocked_stake(
+        &self,
+        signer: SuiAddress,
+        timelocked_staked_sui: ObjectID,
+        gas: ObjectID,
+        gas_budget: u64,
+    ) -> anyhow::Result<TransactionData> {
+        let timelocked_staked_sui = self.get_object_ref(timelocked_staked_sui).await?;
+        let gas_price = self.0.get_reference_gas_price().await?;
+        let gas = self
+            .select_gas(signer, Some(gas), gas_budget, vec![], gas_price)
+            .await?;
+        TransactionData::new_move_call(
+            signer,
+            SUI_SYSTEM_PACKAGE_ID,
+            TIMELOCKED_STAKING_MODULE_NAME.to_owned(),
+            WITHDRAW_TIMELOCKED_STAKE_FUN_NAME.to_owned(),
+            vec![],
+            gas,
+            vec![
+                CallArg::SUI_SYSTEM_MUT,
+                CallArg::Object(ObjectArg::ImmOrOwnedObject(timelocked_staked_sui)),
+            ],
+            gas_budget,
+            gas_price,
+        )
+    }
+
     // TODO: we should add retrial to reduce the transaction building error rate
     pub async fn get_object_ref(&self, object_id: ObjectID) -> anyhow::Result<ObjectRef> {
         self.get_object_ref_and_type(object_id)

crates/sui-types/src/base_types.rs

@@ -33,6 +33,7 @@ use crate::parse_sui_struct_tag;
 use crate::signature::GenericSignature;
 use crate::sui_serde::Readable;
 use crate::sui_serde::{to_sui_struct_tag_string, HexAccountAddress};
+use crate::timelock;
 use crate::transaction::Transaction;
 use crate::transaction::VerifiedTransaction;
 use crate::zk_login_authenticator::ZkLoginAuthenticator;
@@ -333,6 +334,15 @@ impl MoveObjectType {
         }
     }
 
+    pub fn is_timelocked_balance(&self) -> bool {
+        match &self.0 {
+            MoveObjectType_::GasCoin | MoveObjectType_::StakedSui | MoveObjectType_::Coin(_) => {
+                false
+            }
+            MoveObjectType_::Other(s) => timelock::is_timelocked_balance(s),
+        }
+    }
+
     pub fn try_extract_field_name(&self, type_: &DynamicFieldType) -> SuiResult<TypeTag> {
         match &self.0 {
             MoveObjectType_::GasCoin | MoveObjectType_::StakedSui | MoveObjectType_::Coin(_) => {

crates/sui-types/src/lib.rs

@@ -74,6 +74,8 @@ pub mod signature;
 pub mod storage;
 pub mod sui_serde;
 pub mod sui_system_state;
+pub mod timelock;
+pub mod timelocked_staking;
 pub mod transaction;
 pub mod transfer;
 pub mod type_resolver;

crates/sui-types/src/timelock.rs

@@ -0,0 +1,117 @@
+// Copyright (c) Mysten Labs, Inc.
+// SPDX-License-Identifier: Apache-2.0
+
+use crate::{balance::Balance, base_types::ObjectID, id::UID, STARDUST_ADDRESS};
+use move_core_types::{
+    //annotated_value::{MoveFieldLayout, MoveStructLayout, MoveTypeLayout},
+    ident_str,
+    identifier::IdentStr,
+    language_storage::{StructTag, TypeTag},
+};
+use schemars::JsonSchema;
+use serde::{Deserialize, Serialize};
+
+pub const TIMELOCK_MODULE_NAME: &IdentStr = ident_str!("timelock");
+pub const TIMELOCK_STRUCT_NAME: &IdentStr = ident_str!("TimeLock");
+
+// Rust version of the Move stardust::timelock::TimeLock type.
+#[derive(Debug, Serialize, Deserialize, Clone, JsonSchema, Eq, PartialEq)]
+pub struct TimeLock<T> {
+    pub id: UID,
+    /// The locked object.
+    pub locked: T,
+    /// This is the epoch time stamp of when the lock expires.
+    pub expire_timestamp_ms: u64,
+}
+
+impl<'de, T> TimeLock<T>
+where
+    T: Serialize + Deserialize<'de>,
+{
+    /// Constructor.
+    pub fn new(id: UID, locked: T, expire_timestamp_ms: u64) -> Self {
+        Self {
+            id,
+            locked,
+            expire_timestamp_ms,
+        }
+    }
+
+    /// The `TimeLock` type accessor.
+    pub fn type_(type_param: TypeTag) -> StructTag {
+        StructTag {
+            address: STARDUST_ADDRESS,
+            name: TIMELOCK_STRUCT_NAME.to_owned(),
+            module: TIMELOCK_MODULE_NAME.to_owned(),
+            type_params: vec![type_param],
+        }
+    }
+
+    pub fn id(&self) -> &ObjectID {
+        self.id.object_id()
+    }
+
+    pub fn locked(&self) -> &T {
+        &self.locked
+    }
+
+    pub fn expire_timestamp_ms(&self) -> u64 {
+        self.expire_timestamp_ms
+    }
+
+    /// Create a `TimeLock` from BCS bytes.
+    pub fn from_bcs_bytes(content: &'de [u8]) -> Result<Self, bcs::Error> {
+        bcs::from_bytes(content)
+    }
+
+    /// Serialize a `TimeLock` as a `Vec<u8>` of BCS.
+    pub fn to_bcs_bytes(&self) -> Vec<u8> {
+        bcs::to_bytes(&self).unwrap()
+    }
+
+    // TODO
+    // pub fn layout(type_param: TypeTag) -> MoveStructLayout {
+    //     MoveStructLayout {
+    //         type_: Self::type_(type_param.clone()),
+    //         fields: vec![
+    //             MoveFieldLayout::new(
+    //                 ident_str!("id").to_owned(),
+    //                 MoveTypeLayout::Struct(UID::layout()),
+    //             ),
+    //             // MoveFieldLayout::new(
+    //             //     ident_str!("locked").to_owned(),
+    //             //     MoveTypeLayout::Struct(locked.),
+    //             // ),
+    //             MoveFieldLayout::new(
+    //                 ident_str!("expire_timestamp_ms").to_owned(),
+    //                 MoveTypeLayout::U64,
+    //             ),
+    //         ],
+    //     }
+    // }
+}
+
+/// Is this other StructTag representing a TimeLock?
+pub fn is_timelock(other: &StructTag) -> bool {
+    other.address == STARDUST_ADDRESS
+        && other.module.as_ident_str() == TIMELOCK_MODULE_NAME
+        && other.name.as_ident_str() == TIMELOCK_STRUCT_NAME
+}
+
+/// Is this other StructTag representing a TimeLock<Balance<T>>?
+pub fn is_timelocked_balance(other: &StructTag) -> bool {
+    if !is_timelock(other) {
+        return false;
+    }
+
+    if other.type_params.len() != 1 {
+        return false;
+    }
+
+    let param = &other.type_params[0];
+
+    match param {
+        TypeTag::Struct(tag) => Balance::is_balance(tag),
+        _ => false,
+    }
+}

crates/sui-types/src/timelocked_staking.rs

@@ -0,0 +1,8 @@
+// Copyright (c) Mysten Labs, Inc.
+// SPDX-License-Identifier: Apache-2.0
+
+use move_core_types::{ident_str, identifier::IdentStr};
+
+pub const TIMELOCKED_STAKING_MODULE_NAME: &IdentStr = ident_str!("timelocked_staking");
+pub const ADD_TIMELOCKED_STAKE_MUL_BAL_FUN_NAME: &IdentStr = ident_str!("request_add_stake");
+pub const WITHDRAW_TIMELOCKED_STAKE_FUN_NAME: &IdentStr = ident_str!("request_withdraw_stake");

crates/sui/tests/cli_tests.rs

@@ -3,12 +3,12 @@
 
 use std::collections::BTreeSet;
 use std::io::Read;
-use std::str::FromStr;
-use std::{fmt::Write, fs::read_dir, path::PathBuf, str, thread, time::Duration};
-#[cfg(target_os = "windows")]
-use std::os::windows::fs::FileExt;
 #[cfg(not(target_os = "windows"))]
 use std::os::unix::fs::FileExt;
+#[cfg(target_os = "windows")]
+use std::os::windows::fs::FileExt;
+use std::str::FromStr;
+use std::{fmt::Write, fs::read_dir, path::PathBuf, str, thread, time::Duration};
 
 use expect_test::expect;
 use move_package::BuildConfig as MoveBuildConfig;