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