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

iota_types/
programmable_transaction_builder.rs
