iota_types/
programmable_transaction_builder.rs

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