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