[script-composer] Add infer functionality, fix multiple return values

aptos-move/script-composer/src/builder.rs

@@ -317,14 +317,35 @@ impl TransactionComposer {
                 .collect::<Vec<_>>()
         };
 
-        for return_ty in expected_returns_ty {
+        let mut returned_arguments = vec![];
+        let num_of_calls = self.calls.len() as u16;
+
+        for (idx, return_ty) in expected_returns_ty.into_iter().enumerate() {
+            let ability = BinaryIndexedView::Script(self.builder.as_script())
+                .abilities(&return_ty, &[])
+                .map_err(|_| anyhow!("Failed to calculate ability for type"))?;
+
             let local_idx = self.locals_ty.len() as u16;
             self.locals_ty.push(return_ty);
             self.locals_availability.push(true);
             returns.push(local_idx);
+
+            // For values that has drop and copy ability, use copy by default to avoid calling copy manually
+            // on the client side.
+            returned_arguments.push(CallArgument::PreviousResult(PreviousResult {
+                operation_type: if ability.has_drop() && ability.has_copy() {
+                    ArgumentOperation::Copy
+                } else {
+                    ArgumentOperation::Move
+                },
+                call_idx: num_of_calls,
+                return_idx: idx as u16,
+            }));
         }
 
-        let num_of_calls = self.calls.len() as u16;
+        if self.parameters.len() + self.locals_ty.len() > u8::MAX as usize {
+            bail!("Too many locals being allocated, please truncate the transaction");
+        }
 
         self.calls.push(BuilderCall {
             type_args: type_arguments,
@@ -335,15 +356,7 @@ impl TransactionComposer {
             type_tags: ty_args,
         });
 
-        Ok((0..returns.len())
-            .map(|idx| {
-                CallArgument::PreviousResult(PreviousResult {
-                    operation_type: ArgumentOperation::Move,
-                    call_idx: num_of_calls,
-                    return_idx: idx as u16,
-                })
-            })
-            .collect())
+        Ok(returned_arguments)
     }
 
     fn check_drop_at_end(&self) -> anyhow::Result<()> {
@@ -390,11 +403,11 @@ impl TransactionComposer {
             }
 
             // Storing return values
-            for arg in call.returns {
+            for arg in call.returns.iter().rev() {
                 script
                     .code
                     .code
-                    .push(Bytecode::StLoc((arg + parameters_count) as u8));
+                    .push(Bytecode::StLoc((*arg + parameters_count) as u8));
             }
         }
         script.code.code.push(Bytecode::Ret);

aptos-move/script-composer/src/tests/mod.rs

@@ -301,7 +301,9 @@ fn test_module() {
 
     run_txn(builder, &mut h);
 
-    // Create a copyable value and move it twice
+    // Create a droppable and copyable value and move it twice. This is ok because the builder
+    // will use copy instruction instead of move instruction for values that are both copyable
+    // and droppable.
     let mut builder = TransactionComposer::single_signer();
     load_module(&mut builder, &h, "0x1::batched_execution");
     let returns_1 = builder
@@ -329,7 +331,7 @@ fn test_module() {
         )
         .unwrap();
 
-    assert!(builder
+    builder
         .add_batched_call(
             "0x1::batched_execution".to_string(),
             "consume_copyable_value".to_string(),
@@ -339,6 +341,47 @@ fn test_module() {
                 CallArgument::new_bytes(MoveValue::U8(10).simple_serialize().unwrap()),
             ],
         )
+        .unwrap();
+
+    run_txn(builder, &mut h);
+
+    // Create a droppable value and move it twice
+    let mut builder = TransactionComposer::single_signer();
+    load_module(&mut builder, &h, "0x1::batched_execution");
+    let returns_1 = builder
+        .add_batched_call(
+            "0x1::batched_execution".to_string(),
+            "create_droppable_value".to_string(),
+            vec![],
+            vec![CallArgument::new_bytes(
+                MoveValue::U8(10).simple_serialize().unwrap(),
+            )],
+        )
+        .unwrap()
+        .pop()
+        .unwrap();
+
+    builder
+        .add_batched_call(
+            "0x1::batched_execution".to_string(),
+            "consume_droppable_value".to_string(),
+            vec![],
+            vec![
+                returns_1.clone(),
+                CallArgument::new_bytes(MoveValue::U8(10).simple_serialize().unwrap()),
+            ],
+        )
+        .unwrap();
+    assert!(builder
+        .add_batched_call(
+            "0x1::batched_execution".to_string(),
+            "consume_droppable_value".to_string(),
+            vec![],
+            vec![
+                returns_1,
+                CallArgument::new_bytes(MoveValue::U8(10).simple_serialize().unwrap()),
+            ],
+        )
         .is_err());
 
     // Copying a non-copyable value should return error on call.
@@ -635,4 +678,31 @@ fn test_module() {
             ],
         )
         .is_err());
+
+    // Test functions with multiple return values.
+    // Create a copyable value and copy it twice
+    let mut builder = TransactionComposer::single_signer();
+    load_module(&mut builder, &h, "0x1::batched_execution");
+    let returns = builder
+        .add_batched_call(
+            "0x1::batched_execution".to_string(),
+            "multiple_returns".to_string(),
+            vec![],
+            vec![],
+        )
+        .unwrap();
+
+    builder
+        .add_batched_call(
+            "0x1::batched_execution".to_string(),
+            "consume_non_droppable_value".to_string(),
+            vec![],
+            vec![
+                returns[1].clone(),
+                CallArgument::new_bytes(MoveValue::U8(1).simple_serialize().unwrap()),
+            ],
+        )
+        .unwrap();
+
+    run_txn(builder, &mut h);
 }

aptos-move/script-composer/src/tests/test_modules/sources/test.move

@@ -10,7 +10,7 @@ module 0x1::batched_execution {
         val: u8,
     }
 
-    struct CopyableValue has copy {
+    struct CopyableValue has drop, copy {
         val: u8,
     }
 
@@ -73,4 +73,8 @@ module 0x1::batched_execution {
         let GenericNonDroppableValue { val } = v;
         assert!(val == expected_val, 10);
     }
+
+    public fun multiple_returns(): (DroppableValue, NonDroppableValue) {
+        return (DroppableValue { val: 0 }, NonDroppableValue { val: 1} )
+    }
 }