sui_types/
programmable_transaction_builder.rs

1// Copyright (c) Mysten Labs, Inc.
2// SPDX-License-Identifier: Apache-2.0
3
4//! Utility for generating programmable transactions, either by specifying a command or for
5//! migrating legacy transactions
6
7use anyhow::{Context, bail};
8use indexmap::IndexMap;
9use move_core_types::{ident_str, identifier::Identifier, language_storage::TypeTag};
10use mysten_common::ZipDebugEqIteratorExt;
11use serde::Serialize;
12
13use crate::{
14    SUI_FRAMEWORK_PACKAGE_ID,
15    base_types::{FullObjectID, FullObjectRef, ObjectID, ObjectRef, SuiAddress},
16    move_package::PACKAGE_MODULE_NAME,
17    transaction::{
18        Argument, CallArg, Command, FundsWithdrawalArg, ObjectArg, ProgrammableTransaction,
19        SharedObjectMutability,
20    },
21};
22
23#[cfg(test)]
24#[path = "unit_tests/programmable_transaction_builder_tests.rs"]
25mod programmable_transaction_builder_tests;
26
27#[derive(PartialEq, Eq, Hash)]
28enum BuilderArg {
29    Object(ObjectID),
30    Pure(Vec<u8>),
31    ForcedNonUniquePure(usize),
32    FundsWithdraw(usize),
33}
34
35#[derive(Default)]
36pub struct ProgrammableTransactionBuilder {
37    inputs: IndexMap<BuilderArg, CallArg>,
38    commands: Vec<Command>,
39}
40
41impl ProgrammableTransactionBuilder {
42    pub fn new() -> Self {
43        Self::default()
44    }
45
46    pub fn finish(self) -> ProgrammableTransaction {
47        let Self { inputs, commands } = self;
48        let inputs = inputs.into_values().collect();
49        ProgrammableTransaction { inputs, commands }
50    }
51
52    pub fn pure_bytes(&mut self, bytes: Vec<u8>, force_separate: bool) -> Argument {
53        let arg = if force_separate {
54            BuilderArg::ForcedNonUniquePure(self.inputs.len())
55        } else {
56            BuilderArg::Pure(bytes.clone())
57        };
58        let (i, _) = self.inputs.insert_full(arg, CallArg::Pure(bytes));
59        Argument::Input(i as u16)
60    }
61
62    pub fn pure<T: Serialize>(&mut self, value: T) -> anyhow::Result<Argument> {
63        Ok(self.pure_bytes(
64            bcs::to_bytes(&value).context("Serializing pure argument.")?,
65            /* force separate */ false,
66        ))
67    }
68
69    /// Like pure but forces a separate input entry
70    pub fn force_separate_pure<T: Serialize>(&mut self, value: T) -> anyhow::Result<Argument> {
71        Ok(self.pure_bytes(
72            bcs::to_bytes(&value).context("Serializing pure argument.")?,
73            /* force separate */ true,
74        ))
75    }
76
77    pub fn obj(&mut self, obj_arg: ObjectArg) -> anyhow::Result<Argument> {
78        let id = obj_arg.id();
79        let obj_arg = if let Some(old_value) = self.inputs.get(&BuilderArg::Object(id)) {
80            let old_obj_arg = match old_value {
81                CallArg::Pure(_) => anyhow::bail!("invariant violation! object has pure argument"),
82                CallArg::Object(arg) => arg,
83                CallArg::FundsWithdrawal(_) => {
84                    anyhow::bail!("invariant violation! object has balance withdraw argument")
85                }
86            };
87            match (old_obj_arg, obj_arg) {
88                (
89                    ObjectArg::SharedObject {
90                        id: id1,
91                        initial_shared_version: v1,
92                        mutability: mut1,
93                    },
94                    ObjectArg::SharedObject {
95                        id: id2,
96                        initial_shared_version: v2,
97                        mutability: mut2,
98                    },
99                ) if v1 == &v2 => {
100                    anyhow::ensure!(
101                        id1 == &id2 && id == id2,
102                        "invariant violation! object has id does not match call arg"
103                    );
104                    ObjectArg::SharedObject {
105                        id,
106                        initial_shared_version: v2,
107                        mutability: if mut1 == &SharedObjectMutability::Mutable
108                            || mut2 == SharedObjectMutability::Mutable
109                        {
110                            SharedObjectMutability::Mutable
111                        } else {
112                            mut2
113                        },
114                    }
115                }
116                (old_obj_arg, obj_arg) => {
117                    anyhow::ensure!(
118                        old_obj_arg == &obj_arg,
119                        "Mismatched Object argument kind for object {id}. \
120                        {old_value:?} is not compatible with {obj_arg:?}"
121                    );
122                    obj_arg
123                }
124            }
125        } else {
126            obj_arg
127        };
128        let (i, _) = self
129            .inputs
130            .insert_full(BuilderArg::Object(id), CallArg::Object(obj_arg));
131        Ok(Argument::Input(i as u16))
132    }
133
134    pub fn funds_withdrawal(&mut self, arg: FundsWithdrawalArg) -> anyhow::Result<Argument> {
135        let (i, _) = self.inputs.insert_full(
136            BuilderArg::FundsWithdraw(self.inputs.len()),
137            CallArg::FundsWithdrawal(arg),
138        );
139        Ok(Argument::Input(i as u16))
140    }
141
142    pub fn input(&mut self, call_arg: CallArg) -> anyhow::Result<Argument> {
143        match call_arg {
144            CallArg::Pure(bytes) => Ok(self.pure_bytes(bytes, /* force separate */ false)),
145            CallArg::Object(obj) => self.obj(obj),
146            CallArg::FundsWithdrawal(arg) => self.funds_withdrawal(arg),
147        }
148    }
149
150    pub fn make_obj_vec(
151        &mut self,
152        objs: impl IntoIterator<Item = ObjectArg>,
153    ) -> anyhow::Result<Argument> {
154        let make_vec_args = objs
155            .into_iter()
156            .map(|obj| self.obj(obj))
157            .collect::<Result<_, _>>()?;
158        Ok(self.command(Command::MakeMoveVec(None, make_vec_args)))
159    }
160
161    pub fn command(&mut self, command: Command) -> Argument {
162        let i = self.commands.len();
163        self.commands.push(command);
164        Argument::Result(i as u16)
165    }
166
167    /// Will fail to generate if given an empty ObjVec
168    pub fn move_call(
169        &mut self,
170        package: ObjectID,
171        module: Identifier,
172        function: Identifier,
173        type_arguments: Vec<TypeTag>,
174        call_args: Vec<CallArg>,
175    ) -> anyhow::Result<()> {
176        let arguments = call_args
177            .into_iter()
178            .map(|a| self.input(a))
179            .collect::<Result<_, _>>()?;
180        self.command(Command::move_call(
181            package,
182            module,
183            function,
184            type_arguments,
185            arguments,
186        ));
187        Ok(())
188    }
189
190    pub fn programmable_move_call(
191        &mut self,
192        package: ObjectID,
193        module: Identifier,
194        function: Identifier,
195        type_arguments: Vec<TypeTag>,
196        arguments: Vec<Argument>,
197    ) -> Argument {
198        self.command(Command::move_call(
199            package,
200            module,
201            function,
202            type_arguments,
203            arguments,
204        ))
205    }
206
207    pub fn publish_upgradeable(
208        &mut self,
209        modules: Vec<Vec<u8>>,
210        dep_ids: Vec<ObjectID>,
211    ) -> Argument {
212        self.command(Command::Publish(modules, dep_ids))
213    }
214
215    pub fn publish_immutable(&mut self, modules: Vec<Vec<u8>>, dep_ids: Vec<ObjectID>) {
216        let cap = self.publish_upgradeable(modules, dep_ids);
217        self.commands.push(Command::move_call(
218            SUI_FRAMEWORK_PACKAGE_ID,
219            PACKAGE_MODULE_NAME.to_owned(),
220            ident_str!("make_immutable").to_owned(),
221            vec![],
222            vec![cap],
223        ));
224    }
225
226    pub fn upgrade(
227        &mut self,
228        current_package_object_id: ObjectID,
229        upgrade_ticket: Argument,
230        transitive_deps: Vec<ObjectID>,
231        modules: Vec<Vec<u8>>,
232    ) -> Argument {
233        self.command(Command::Upgrade(
234            modules,
235            transitive_deps,
236            current_package_object_id,
237            upgrade_ticket,
238        ))
239    }
240
241    pub fn transfer_arg(&mut self, recipient: SuiAddress, arg: Argument) {
242        self.transfer_args(recipient, vec![arg])
243    }
244
245    pub fn transfer_args(&mut self, recipient: SuiAddress, args: Vec<Argument>) {
246        let rec_arg = self.pure(recipient).unwrap();
247        self.commands.push(Command::TransferObjects(args, rec_arg));
248    }
249
250    pub fn transfer_object(
251        &mut self,
252        recipient: SuiAddress,
253        full_object_ref: FullObjectRef,
254    ) -> anyhow::Result<()> {
255        let rec_arg = self.pure(recipient).unwrap();
256        let obj_arg = self.obj(match full_object_ref.0 {
257            FullObjectID::Fastpath(_) => {
258                ObjectArg::ImmOrOwnedObject(full_object_ref.as_object_ref())
259            }
260            FullObjectID::Consensus((id, initial_shared_version)) => ObjectArg::SharedObject {
261                id,
262                initial_shared_version,
263                mutability: SharedObjectMutability::Mutable,
264            },
265        });
266        self.commands
267            .push(Command::TransferObjects(vec![obj_arg?], rec_arg));
268        Ok(())
269    }
270
271    pub fn transfer_sui(&mut self, recipient: SuiAddress, amount: Option<u64>) {
272        let rec_arg = self.pure(recipient).unwrap();
273        let coin_arg = if let Some(amount) = amount {
274            let amt_arg = self.pure(amount).unwrap();
275            self.command(Command::SplitCoins(Argument::GasCoin, vec![amt_arg]))
276        } else {
277            Argument::GasCoin
278        };
279        self.command(Command::TransferObjects(vec![coin_arg], rec_arg));
280    }
281
282    pub fn pay_all_sui(&mut self, recipient: SuiAddress) {
283        let rec_arg = self.pure(recipient).unwrap();
284        self.command(Command::TransferObjects(vec![Argument::GasCoin], rec_arg));
285    }
286
287    /// Will fail to generate if recipients and amounts do not have the same lengths
288    pub fn pay_sui(
289        &mut self,
290        recipients: Vec<SuiAddress>,
291        amounts: Vec<u64>,
292    ) -> anyhow::Result<()> {
293        self.pay_impl(recipients, amounts, Argument::GasCoin)
294    }
295
296    pub fn split_coin(&mut self, recipient: SuiAddress, coin: ObjectRef, amounts: Vec<u64>) {
297        let coin_arg = self.obj(ObjectArg::ImmOrOwnedObject(coin)).unwrap();
298        let amounts_len = amounts.len();
299        let amt_args = amounts.into_iter().map(|a| self.pure(a).unwrap()).collect();
300        let result = self.command(Command::SplitCoins(coin_arg, amt_args));
301        let Argument::Result(result) = result else {
302            panic!("self.command should always give a Argument::Result");
303        };
304
305        let recipient = self.pure(recipient).unwrap();
306        self.command(Command::TransferObjects(
307            (0..amounts_len)
308                .map(|i| Argument::NestedResult(result, i as u16))
309                .collect(),
310            recipient,
311        ));
312    }
313
314    /// Merge `coins` into the `target` coin.
315    pub fn merge_coins(&mut self, target: ObjectRef, coins: Vec<ObjectRef>) -> anyhow::Result<()> {
316        let target_arg = self.obj(ObjectArg::ImmOrOwnedObject(target))?;
317        let coin_args = coins
318            .into_iter()
319            .map(|coin| self.obj(ObjectArg::ImmOrOwnedObject(coin)).unwrap())
320            .collect::<Vec<_>>();
321        self.command(Command::MergeCoins(target_arg, coin_args));
322        Ok(())
323    }
324
325    /// Merge all `coins` into the first coin in the vector.
326    /// Returns an `Argument` for the first coin.
327    pub fn smash_coins(&mut self, coins: Vec<ObjectRef>) -> anyhow::Result<Argument> {
328        let mut coins = coins.into_iter();
329        let Some(target) = coins.next() else {
330            bail!("coins vector is empty");
331        };
332        let target_arg = self.obj(ObjectArg::ImmOrOwnedObject(target))?;
333        let coin_args = coins
334            .map(|coin| self.obj(ObjectArg::ImmOrOwnedObject(coin)).unwrap())
335            .collect::<Vec<_>>();
336        self.command(Command::MergeCoins(target_arg, coin_args));
337        Ok(target_arg)
338    }
339
340    /// Will fail to generate if recipients and amounts do not have the same lengths.
341    /// Or if coins is empty
342    pub fn pay(
343        &mut self,
344        coins: Vec<ObjectRef>,
345        recipients: Vec<SuiAddress>,
346        amounts: Vec<u64>,
347    ) -> anyhow::Result<()> {
348        let mut coins = coins.into_iter();
349        let Some(coin) = coins.next() else {
350            anyhow::bail!("coins vector is empty");
351        };
352        let coin_arg = self.obj(ObjectArg::ImmOrOwnedObject(coin))?;
353        let merge_args: Vec<_> = coins
354            .map(|c| self.obj(ObjectArg::ImmOrOwnedObject(c)))
355            .collect::<Result<_, _>>()?;
356        if !merge_args.is_empty() {
357            self.command(Command::MergeCoins(coin_arg, merge_args));
358        }
359        self.pay_impl(recipients, amounts, coin_arg)
360    }
361
362    fn pay_impl(
363        &mut self,
364        recipients: Vec<SuiAddress>,
365        amounts: Vec<u64>,
366        coin: Argument,
367    ) -> anyhow::Result<()> {
368        if recipients.len() != amounts.len() {
369            anyhow::bail!(
370                "Recipients and amounts mismatch. Got {} recipients but {} amounts",
371                recipients.len(),
372                amounts.len()
373            )
374        }
375        if amounts.is_empty() {
376            return Ok(());
377        }
378
379        // collect recipients in the case where they are non-unique in order
380        // to minimize the number of transfers that must be performed
381        let mut recipient_map: IndexMap<SuiAddress, Vec<usize>> = IndexMap::new();
382        let mut amt_args = Vec::with_capacity(recipients.len());
383        for (i, (recipient, amount)) in recipients.into_iter().zip_debug_eq(amounts).enumerate() {
384            recipient_map.entry(recipient).or_default().push(i);
385            amt_args.push(self.pure(amount)?);
386        }
387        let Argument::Result(split_primary) = self.command(Command::SplitCoins(coin, amt_args))
388        else {
389            panic!("self.command should always give a Argument::Result")
390        };
391        for (recipient, split_secondaries) in recipient_map {
392            let rec_arg = self.pure(recipient).unwrap();
393            let coins = split_secondaries
394                .into_iter()
395                .map(|j| Argument::NestedResult(split_primary, j as u16))
396                .collect();
397            self.command(Command::TransferObjects(coins, rec_arg));
398        }
399        Ok(())
400    }
401}