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