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

iota_types/
programmable_transaction_builder.rs
