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 iota_sdk_types::{
11    Address, Argument, Command, Identifier, ObjectId, ProgrammableTransaction, TypeTag,
12};
13use serde::Serialize;
14
15use crate::{
16    base_types::ObjectRef,
17    transaction::{CallArg, SharedObjectRef},
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: impl Into<CallArg>) -> anyhow::Result<Argument> {
72        let obj_arg: CallArg = obj_arg.into();
73        let id = *obj_arg
74            .object_id_opt()
75            .ok_or_else(|| anyhow::anyhow!("expected object CallArg, found pure argument"))?;
76        let obj_arg = if let Some(old_value) = self.inputs.get(&BuilderArg::Object(id)) {
77            match (old_value.as_opt_shared(), obj_arg.as_opt_shared()) {
78                (
79                    Some(&SharedObjectRef {
80                        object_id: id1,
81                        initial_shared_version: v1,
82                        mutable: mut1,
83                    }),
84                    Some(&SharedObjectRef {
85                        object_id: id2,
86                        initial_shared_version: v2,
87                        mutable: mut2,
88                    }),
89                ) if v1 == v2 => {
90                    anyhow::ensure!(
91                        id1 == id2 && id == id2,
92                        "invariant violation! object has id does not match call arg"
93                    );
94                    CallArg::Shared(SharedObjectRef::new(id, v2, mut1 || mut2))
95                }
96                _ => {
97                    anyhow::ensure!(
98                        *old_value == obj_arg,
99                        "Mismatched Object argument kind for object {id}. \
100                        {old_value:?} is not compatible with {obj_arg:?}"
101                    );
102                    obj_arg
103                }
104            }
105        } else {
106            obj_arg
107        };
108        let (i, _) = self.inputs.insert_full(BuilderArg::Object(id), obj_arg);
109        Ok(Argument::Input(i as u16))
110    }
111
112    pub fn input(&mut self, call_arg: CallArg) -> anyhow::Result<Argument> {
113        match call_arg {
114            CallArg::Pure(value) => Ok(self.pure_bytes(value, /* force separate */ false)),
115            CallArg::ImmutableOrOwned(_) | CallArg::Shared(_) | CallArg::Receiving(_) => {
116                self.obj(call_arg)
117            }
118            _ => unimplemented!("a new CallArg enum variant was added and needs to be handled"),
119        }
120    }
121
122    pub fn make_obj_vec<T: Into<CallArg>>(
123        &mut self,
124        objs: impl IntoIterator<Item = T>,
125    ) -> anyhow::Result<Argument> {
126        let make_vec_args = objs
127            .into_iter()
128            .map(|obj| self.obj(obj.into()))
129            .collect::<Result<_, _>>()?;
130        Ok(self.command(Command::new_make_move_vector(None, make_vec_args)))
131    }
132
133    pub fn command(&mut self, command: Command) -> Argument {
134        let i = self.commands.len();
135        self.commands.push(command);
136        Argument::Result(i as u16)
137    }
138
139    /// Will fail to generate if given an empty ObjVec
140    pub fn move_call(
141        &mut self,
142        package: ObjectId,
143        module: Identifier,
144        function: Identifier,
145        type_arguments: Vec<TypeTag>,
146        call_args: Vec<CallArg>,
147    ) -> anyhow::Result<()> {
148        let arguments = call_args
149            .into_iter()
150            .map(|a| self.input(a))
151            .collect::<Result<_, _>>()?;
152        self.command(Command::new_move_call(
153            package,
154            module,
155            function,
156            type_arguments,
157            arguments,
158        ));
159        Ok(())
160    }
161
162    pub fn programmable_move_call(
163        &mut self,
164        package: ObjectId,
165        module: Identifier,
166        function: Identifier,
167        type_arguments: Vec<TypeTag>,
168        arguments: Vec<Argument>,
169    ) -> Argument {
170        self.command(Command::new_move_call(
171            package,
172            module,
173            function,
174            type_arguments,
175            arguments,
176        ))
177    }
178
179    pub fn publish_upgradeable(
180        &mut self,
181        modules: Vec<Vec<u8>>,
182        dep_ids: Vec<ObjectId>,
183    ) -> Argument {
184        self.command(Command::new_publish(modules, dep_ids))
185    }
186
187    pub fn publish_immutable(&mut self, modules: Vec<Vec<u8>>, dep_ids: Vec<ObjectId>) {
188        let cap = self.publish_upgradeable(modules, dep_ids);
189        self.commands.push(Command::new_move_call(
190            ObjectId::FRAMEWORK,
191            Identifier::PACKAGE_MODULE,
192            Identifier::from_static("make_immutable"),
193            vec![],
194            vec![cap],
195        ));
196    }
197
198    pub fn upgrade(
199        &mut self,
200        current_package_object_id: ObjectId,
201        upgrade_ticket: Argument,
202        transitive_deps: Vec<ObjectId>,
203        modules: Vec<Vec<u8>>,
204    ) -> Argument {
205        self.command(Command::new_upgrade(
206            modules,
207            transitive_deps,
208            current_package_object_id,
209            upgrade_ticket,
210        ))
211    }
212
213    pub fn transfer_arg(&mut self, recipient: Address, arg: Argument) {
214        self.transfer_args(recipient, vec![arg])
215    }
216
217    pub fn transfer_args(&mut self, recipient: Address, args: Vec<Argument>) {
218        let rec_arg = self.pure(recipient).unwrap();
219        self.commands
220            .push(Command::new_transfer_objects(args, rec_arg));
221    }
222
223    pub fn transfer_object(
224        &mut self,
225        recipient: Address,
226        object_ref: ObjectRef,
227    ) -> anyhow::Result<()> {
228        let rec_arg = self.pure(recipient).unwrap();
229        let obj_arg = self.obj(CallArg::ImmutableOrOwned(object_ref))?;
230        self.commands
231            .push(Command::new_transfer_objects(vec![obj_arg], rec_arg));
232        Ok(())
233    }
234
235    pub fn transfer_iota(&mut self, recipient: Address, amount: Option<u64>) {
236        let rec_arg = self.pure(recipient).unwrap();
237        let coin_arg = if let Some(amount) = amount {
238            let amt_arg = self.pure(amount).unwrap();
239            self.command(Command::new_split_coins(Argument::Gas, vec![amt_arg]))
240        } else {
241            Argument::Gas
242        };
243        self.command(Command::new_transfer_objects(vec![coin_arg], rec_arg));
244    }
245
246    pub fn pay_all_iota(&mut self, recipient: Address) {
247        let rec_arg = self.pure(recipient).unwrap();
248        self.command(Command::new_transfer_objects(vec![Argument::Gas], rec_arg));
249    }
250
251    /// Will fail to generate if recipients and amounts do not have the same
252    /// lengths
253    pub fn pay_iota(&mut self, recipients: Vec<Address>, amounts: Vec<u64>) -> anyhow::Result<()> {
254        self.pay_impl(recipients, amounts, Argument::Gas)
255    }
256
257    pub fn split_coin(&mut self, recipient: Address, coin: ObjectRef, amounts: Vec<u64>) {
258        let coin_arg = self.obj(CallArg::ImmutableOrOwned(coin)).unwrap();
259        let amounts_len = amounts.len();
260        let amt_args = amounts.into_iter().map(|a| self.pure(a).unwrap()).collect();
261        let result = self.command(Command::new_split_coins(coin_arg, amt_args));
262        let Argument::Result(result) = result else {
263            panic!("self.command should always give a Argument::Result");
264        };
265
266        let recipient = self.pure(recipient).unwrap();
267        self.command(Command::new_transfer_objects(
268            (0..amounts_len)
269                .map(|i| Argument::NestedResult(result, i as u16))
270                .collect(),
271            recipient,
272        ));
273    }
274
275    /// Will fail to generate if recipients and amounts do not have the same
276    /// lengths. Or if coins is empty
277    pub fn pay(
278        &mut self,
279        coins: Vec<ObjectRef>,
280        recipients: Vec<Address>,
281        amounts: Vec<u64>,
282    ) -> anyhow::Result<()> {
283        let mut coins = coins.into_iter();
284        let Some(coin) = coins.next() else {
285            anyhow::bail!("coins vector is empty");
286        };
287        let coin_arg = self.obj(CallArg::ImmutableOrOwned(coin))?;
288        let merge_args: Vec<_> = coins
289            .map(|c| self.obj(CallArg::ImmutableOrOwned(c)))
290            .collect::<Result<_, _>>()?;
291        if !merge_args.is_empty() {
292            self.command(Command::new_merge_coins(coin_arg, merge_args));
293        }
294        self.pay_impl(recipients, amounts, coin_arg)
295    }
296
297    fn pay_impl(
298        &mut self,
299        recipients: Vec<Address>,
300        amounts: Vec<u64>,
301        coin: Argument,
302    ) -> anyhow::Result<()> {
303        if recipients.len() != amounts.len() {
304            anyhow::bail!(
305                "Recipients and amounts mismatch. Got {} recipients but {} amounts",
306                recipients.len(),
307                amounts.len()
308            )
309        }
310        if amounts.is_empty() {
311            return Ok(());
312        }
313
314        // collect recipients in the case where they are non-unique in order
315        // to minimize the number of transfers that must be performed
316        let mut recipient_map: IndexMap<Address, Vec<usize>> = IndexMap::new();
317        let mut amt_args = Vec::with_capacity(recipients.len());
318        for (i, (recipient, amount)) in recipients.into_iter().zip(amounts).enumerate() {
319            recipient_map.entry(recipient).or_default().push(i);
320            amt_args.push(self.pure(amount)?);
321        }
322        let Argument::Result(split_primary) =
323            self.command(Command::new_split_coins(coin, amt_args))
324        else {
325            panic!("self.command should always give a Argument::Result")
326        };
327        for (recipient, split_secondaries) in recipient_map {
328            let rec_arg = self.pure(recipient).unwrap();
329            let coins = split_secondaries
330                .into_iter()
331                .map(|j| Argument::NestedResult(split_primary, j as u16))
332                .collect();
333            self.command(Command::new_transfer_objects(coins, rec_arg));
334        }
335        Ok(())
336    }
337}
iota_types/programmable_transaction_builder.rs

iota_types/
programmable_transaction_builder.rs
