1use anyhow::{Context, bail};
8use indexmap::IndexMap;
9use move_core_types::{ident_str, identifier::Identifier, language_storage::TypeTag};
10use mysten_common::ZipDebugEqIteratorExt;
11use serde::Serialize;
12
13use crate::{
14 SUI_FRAMEWORK_PACKAGE_ID,
15 base_types::{FullObjectID, FullObjectRef, ObjectID, ObjectRef, SuiAddress},
16 move_package::PACKAGE_MODULE_NAME,
17 transaction::{
18 Argument, CallArg, Command, FundsWithdrawalArg, ObjectArg, ProgrammableTransaction,
19 SharedObjectMutability,
20 },
21};
22
23#[cfg(test)]
24#[path = "unit_tests/programmable_transaction_builder_tests.rs"]
25mod programmable_transaction_builder_tests;
26
27#[derive(PartialEq, Eq, Hash)]
28enum BuilderArg {
29 Object(ObjectID),
30 Pure(Vec<u8>),
31 ForcedNonUniquePure(usize),
32 FundsWithdraw(usize),
33}
34
35#[derive(Default)]
36pub struct ProgrammableTransactionBuilder {
37 inputs: IndexMap<BuilderArg, CallArg>,
38 commands: Vec<Command>,
39}
40
41impl ProgrammableTransactionBuilder {
42 pub fn new() -> Self {
43 Self::default()
44 }
45
46 pub fn finish(self) -> ProgrammableTransaction {
47 let Self { inputs, commands } = self;
48 let inputs = inputs.into_values().collect();
49 ProgrammableTransaction { inputs, commands }
50 }
51
52 pub fn pure_bytes(&mut self, bytes: Vec<u8>, force_separate: bool) -> Argument {
53 let arg = if force_separate {
54 BuilderArg::ForcedNonUniquePure(self.inputs.len())
55 } else {
56 BuilderArg::Pure(bytes.clone())
57 };
58 let (i, _) = self.inputs.insert_full(arg, CallArg::Pure(bytes));
59 Argument::Input(i as u16)
60 }
61
62 pub fn pure<T: Serialize>(&mut self, value: T) -> anyhow::Result<Argument> {
63 Ok(self.pure_bytes(
64 bcs::to_bytes(&value).context("Serializing pure argument.")?,
65 false,
66 ))
67 }
68
69 pub fn force_separate_pure<T: Serialize>(&mut self, value: T) -> anyhow::Result<Argument> {
71 Ok(self.pure_bytes(
72 bcs::to_bytes(&value).context("Serializing pure argument.")?,
73 true,
74 ))
75 }
76
77 pub fn obj(&mut self, obj_arg: ObjectArg) -> anyhow::Result<Argument> {
78 let id = obj_arg.id();
79 let obj_arg = if let Some(old_value) = self.inputs.get(&BuilderArg::Object(id)) {
80 let old_obj_arg = match old_value {
81 CallArg::Pure(_) => anyhow::bail!("invariant violation! object has pure argument"),
82 CallArg::Object(arg) => arg,
83 CallArg::FundsWithdrawal(_) => {
84 anyhow::bail!("invariant violation! object has balance withdraw argument")
85 }
86 };
87 match (old_obj_arg, obj_arg) {
88 (
89 ObjectArg::SharedObject {
90 id: id1,
91 initial_shared_version: v1,
92 mutability: mut1,
93 },
94 ObjectArg::SharedObject {
95 id: id2,
96 initial_shared_version: v2,
97 mutability: mut2,
98 },
99 ) if v1 == &v2 => {
100 anyhow::ensure!(
101 id1 == &id2 && id == id2,
102 "invariant violation! object has id does not match call arg"
103 );
104 ObjectArg::SharedObject {
105 id,
106 initial_shared_version: v2,
107 mutability: if mut1 == &SharedObjectMutability::Mutable
108 || mut2 == SharedObjectMutability::Mutable
109 {
110 SharedObjectMutability::Mutable
111 } else {
112 mut2
113 },
114 }
115 }
116 (old_obj_arg, obj_arg) => {
117 anyhow::ensure!(
118 old_obj_arg == &obj_arg,
119 "Mismatched Object argument kind for object {id}. \
120 {old_value:?} is not compatible with {obj_arg:?}"
121 );
122 obj_arg
123 }
124 }
125 } else {
126 obj_arg
127 };
128 let (i, _) = self
129 .inputs
130 .insert_full(BuilderArg::Object(id), CallArg::Object(obj_arg));
131 Ok(Argument::Input(i as u16))
132 }
133
134 pub fn funds_withdrawal(&mut self, arg: FundsWithdrawalArg) -> anyhow::Result<Argument> {
135 let (i, _) = self.inputs.insert_full(
136 BuilderArg::FundsWithdraw(self.inputs.len()),
137 CallArg::FundsWithdrawal(arg),
138 );
139 Ok(Argument::Input(i as u16))
140 }
141
142 pub fn input(&mut self, call_arg: CallArg) -> anyhow::Result<Argument> {
143 match call_arg {
144 CallArg::Pure(bytes) => Ok(self.pure_bytes(bytes, false)),
145 CallArg::Object(obj) => self.obj(obj),
146 CallArg::FundsWithdrawal(arg) => self.funds_withdrawal(arg),
147 }
148 }
149
150 pub fn make_obj_vec(
151 &mut self,
152 objs: impl IntoIterator<Item = ObjectArg>,
153 ) -> anyhow::Result<Argument> {
154 let make_vec_args = objs
155 .into_iter()
156 .map(|obj| self.obj(obj))
157 .collect::<Result<_, _>>()?;
158 Ok(self.command(Command::MakeMoveVec(None, make_vec_args)))
159 }
160
161 pub fn command(&mut self, command: Command) -> Argument {
162 let i = self.commands.len();
163 self.commands.push(command);
164 Argument::Result(i as u16)
165 }
166
167 pub fn move_call(
169 &mut self,
170 package: ObjectID,
171 module: Identifier,
172 function: Identifier,
173 type_arguments: Vec<TypeTag>,
174 call_args: Vec<CallArg>,
175 ) -> anyhow::Result<()> {
176 let arguments = call_args
177 .into_iter()
178 .map(|a| self.input(a))
179 .collect::<Result<_, _>>()?;
180 self.command(Command::move_call(
181 package,
182 module,
183 function,
184 type_arguments,
185 arguments,
186 ));
187 Ok(())
188 }
189
190 pub fn programmable_move_call(
191 &mut self,
192 package: ObjectID,
193 module: Identifier,
194 function: Identifier,
195 type_arguments: Vec<TypeTag>,
196 arguments: Vec<Argument>,
197 ) -> Argument {
198 self.command(Command::move_call(
199 package,
200 module,
201 function,
202 type_arguments,
203 arguments,
204 ))
205 }
206
207 pub fn publish_upgradeable(
208 &mut self,
209 modules: Vec<Vec<u8>>,
210 dep_ids: Vec<ObjectID>,
211 ) -> Argument {
212 self.command(Command::Publish(modules, dep_ids))
213 }
214
215 pub fn publish_immutable(&mut self, modules: Vec<Vec<u8>>, dep_ids: Vec<ObjectID>) {
216 let cap = self.publish_upgradeable(modules, dep_ids);
217 self.commands.push(Command::move_call(
218 SUI_FRAMEWORK_PACKAGE_ID,
219 PACKAGE_MODULE_NAME.to_owned(),
220 ident_str!("make_immutable").to_owned(),
221 vec![],
222 vec![cap],
223 ));
224 }
225
226 pub fn upgrade(
227 &mut self,
228 current_package_object_id: ObjectID,
229 upgrade_ticket: Argument,
230 transitive_deps: Vec<ObjectID>,
231 modules: Vec<Vec<u8>>,
232 ) -> Argument {
233 self.command(Command::Upgrade(
234 modules,
235 transitive_deps,
236 current_package_object_id,
237 upgrade_ticket,
238 ))
239 }
240
241 pub fn transfer_arg(&mut self, recipient: SuiAddress, arg: Argument) {
242 self.transfer_args(recipient, vec![arg])
243 }
244
245 pub fn transfer_args(&mut self, recipient: SuiAddress, args: Vec<Argument>) {
246 let rec_arg = self.pure(recipient).unwrap();
247 self.commands.push(Command::TransferObjects(args, rec_arg));
248 }
249
250 pub fn transfer_object(
251 &mut self,
252 recipient: SuiAddress,
253 full_object_ref: FullObjectRef,
254 ) -> anyhow::Result<()> {
255 let rec_arg = self.pure(recipient).unwrap();
256 let obj_arg = self.obj(match full_object_ref.0 {
257 FullObjectID::Fastpath(_) => {
258 ObjectArg::ImmOrOwnedObject(full_object_ref.as_object_ref())
259 }
260 FullObjectID::Consensus((id, initial_shared_version)) => ObjectArg::SharedObject {
261 id,
262 initial_shared_version,
263 mutability: SharedObjectMutability::Mutable,
264 },
265 });
266 self.commands
267 .push(Command::TransferObjects(vec![obj_arg?], rec_arg));
268 Ok(())
269 }
270
271 pub fn transfer_sui(&mut self, recipient: SuiAddress, amount: Option<u64>) {
272 let rec_arg = self.pure(recipient).unwrap();
273 let coin_arg = if let Some(amount) = amount {
274 let amt_arg = self.pure(amount).unwrap();
275 self.command(Command::SplitCoins(Argument::GasCoin, vec![amt_arg]))
276 } else {
277 Argument::GasCoin
278 };
279 self.command(Command::TransferObjects(vec![coin_arg], rec_arg));
280 }
281
282 pub fn pay_all_sui(&mut self, recipient: SuiAddress) {
283 let rec_arg = self.pure(recipient).unwrap();
284 self.command(Command::TransferObjects(vec![Argument::GasCoin], rec_arg));
285 }
286
287 pub fn pay_sui(
289 &mut self,
290 recipients: Vec<SuiAddress>,
291 amounts: Vec<u64>,
292 ) -> anyhow::Result<()> {
293 self.pay_impl(recipients, amounts, Argument::GasCoin)
294 }
295
296 pub fn split_coin(&mut self, recipient: SuiAddress, coin: ObjectRef, amounts: Vec<u64>) {
297 let coin_arg = self.obj(ObjectArg::ImmOrOwnedObject(coin)).unwrap();
298 let amounts_len = amounts.len();
299 let amt_args = amounts.into_iter().map(|a| self.pure(a).unwrap()).collect();
300 let result = self.command(Command::SplitCoins(coin_arg, amt_args));
301 let Argument::Result(result) = result else {
302 panic!("self.command should always give a Argument::Result");
303 };
304
305 let recipient = self.pure(recipient).unwrap();
306 self.command(Command::TransferObjects(
307 (0..amounts_len)
308 .map(|i| Argument::NestedResult(result, i as u16))
309 .collect(),
310 recipient,
311 ));
312 }
313
314 pub fn merge_coins(&mut self, target: ObjectRef, coins: Vec<ObjectRef>) -> anyhow::Result<()> {
316 let target_arg = self.obj(ObjectArg::ImmOrOwnedObject(target))?;
317 let coin_args = coins
318 .into_iter()
319 .map(|coin| self.obj(ObjectArg::ImmOrOwnedObject(coin)).unwrap())
320 .collect::<Vec<_>>();
321 self.command(Command::MergeCoins(target_arg, coin_args));
322 Ok(())
323 }
324
325 pub fn smash_coins(&mut self, coins: Vec<ObjectRef>) -> anyhow::Result<Argument> {
328 let mut coins = coins.into_iter();
329 let Some(target) = coins.next() else {
330 bail!("coins vector is empty");
331 };
332 let target_arg = self.obj(ObjectArg::ImmOrOwnedObject(target))?;
333 let coin_args = coins
334 .map(|coin| self.obj(ObjectArg::ImmOrOwnedObject(coin)).unwrap())
335 .collect::<Vec<_>>();
336 self.command(Command::MergeCoins(target_arg, coin_args));
337 Ok(target_arg)
338 }
339
340 pub fn pay(
343 &mut self,
344 coins: Vec<ObjectRef>,
345 recipients: Vec<SuiAddress>,
346 amounts: Vec<u64>,
347 ) -> anyhow::Result<()> {
348 let mut coins = coins.into_iter();
349 let Some(coin) = coins.next() else {
350 anyhow::bail!("coins vector is empty");
351 };
352 let coin_arg = self.obj(ObjectArg::ImmOrOwnedObject(coin))?;
353 let merge_args: Vec<_> = coins
354 .map(|c| self.obj(ObjectArg::ImmOrOwnedObject(c)))
355 .collect::<Result<_, _>>()?;
356 if !merge_args.is_empty() {
357 self.command(Command::MergeCoins(coin_arg, merge_args));
358 }
359 self.pay_impl(recipients, amounts, coin_arg)
360 }
361
362 fn pay_impl(
363 &mut self,
364 recipients: Vec<SuiAddress>,
365 amounts: Vec<u64>,
366 coin: Argument,
367 ) -> anyhow::Result<()> {
368 if recipients.len() != amounts.len() {
369 anyhow::bail!(
370 "Recipients and amounts mismatch. Got {} recipients but {} amounts",
371 recipients.len(),
372 amounts.len()
373 )
374 }
375 if amounts.is_empty() {
376 return Ok(());
377 }
378
379 let mut recipient_map: IndexMap<SuiAddress, Vec<usize>> = IndexMap::new();
382 let mut amt_args = Vec::with_capacity(recipients.len());
383 for (i, (recipient, amount)) in recipients.into_iter().zip_debug_eq(amounts).enumerate() {
384 recipient_map.entry(recipient).or_default().push(i);
385 amt_args.push(self.pure(amount)?);
386 }
387 let Argument::Result(split_primary) = self.command(Command::SplitCoins(coin, amt_args))
388 else {
389 panic!("self.command should always give a Argument::Result")
390 };
391 for (recipient, split_secondaries) in recipient_map {
392 let rec_arg = self.pure(recipient).unwrap();
393 let coins = split_secondaries
394 .into_iter()
395 .map(|j| Argument::NestedResult(split_primary, j as u16))
396 .collect();
397 self.command(Command::TransferObjects(coins, rec_arg));
398 }
399 Ok(())
400 }
401}