1use 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 false,
59 ))
60 }
61
62 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 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, 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 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 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 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 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}