Add tx_input_owner() function to std. (FuelLabs#2215)

* feat: add tx_input_message_owner() function

* feat: also get owner from InputMessage

* fix: typo

* feat: allow getting owner of coin & message inputs

* refactor: use new tx function for input owner

* docs: improve documentation

* test: update tests with todos

* refactor: change tx_input_owner to accept an index

* refactor: use new tx_input_owner and handle Option::None case

* fix: switch offsets  to u32 type

* fix: remove redundent Option wrapper

* fix: use index rather than ptr to get type

* test: update types in test contract

* test: fix test

* fix: remove unneeded check and error

* fix: change integers in tx_input_owner to u64

Co-authored-by: Mohammad Fawaz <[email protected]>

sway-lib-std/src/chain/auth.sw

@@ -49,10 +49,10 @@ pub fn msg_sender() -> Result<Identity, AuthError> {
 fn coins_owner() -> Result<Identity, AuthError> {
     let target_input_type = 0u8;
     let inputs_count = tx_inputs_count();
-
     let mut candidate = Option::None::<Address>();
     let mut i = 0;
 
+    // Note: `inputs_count` is guaranteed to be at least 1 for any valid tx.
     while i < inputs_count {
         let input_type = tx_input_type(i);
         if input_type != target_input_type {
@@ -61,28 +61,28 @@ fn coins_owner() -> Result<Identity, AuthError> {
             i += 1;
         } else {
             // type == InputCoin
-            let input_owner = Option::Some(tx_input_coin_owner(i));
+            let input_owner = tx_input_owner(i);
             if candidate.is_none() {
                 // This is the first input seen of the correct type.
                 candidate = input_owner;
                 i += 1;
             } else {
                 // Compare current coin owner to candidate.
-                // `candidate` and `input_owner` must be `Option::Some` at this point,
-                // so can unwrap safely.
+                // `candidate` and `input_owner` must be `Option::Some`
+                // at this point, so we can unwrap safely.
                 if input_owner.unwrap() == candidate.unwrap() {
                     // Owners are a match, continue looping.
                     i += 1;
                 } else {
                     // Owners don't match. Return Err.
+                    // TODO: Use break keyword when possible
                     i = inputs_count;
                     return Result::Err(AuthError::InputsNotAllOwnedBySameAddress);
                 };
             };
-        };
+        }
     }
 
     // `candidate` must be `Option::Some` at this point, so can unwrap safely.
-    // Note: `inputs_count` is guaranteed to be at least 1 for any valid tx.
     Result::Ok(Identity::Address(candidate.unwrap()))
 }

sway-lib-std/src/tx.sw

@@ -7,6 +7,7 @@ use ::context::registers::instrs_start;
 use ::contract_id::ContractId;
 use ::intrinsics::is_reference_type;
 use ::mem::read;
+use ::option::Option;
 
 ////////////////////////////////////////
 // Transaction fields
@@ -195,6 +196,34 @@ pub fn tx_input_type_from_pointer(ptr: u64) -> u8 {
     }
 }
 
+/// If the input's type is `InputCoin` or `InputMessage`,
+/// return the owner as an Option::Some(owner).
+/// Otherwise, returns Option::None.
+pub fn tx_input_owner(index: u64) -> Option<Address> {
+    let type = tx_input_type(index);
+    let owner_offset = match type {
+        // 0 is the `Coin` Input type
+        0u8 => {
+            // Need to skip over six words, so add 8*6=48
+            48
+        },
+        // 2 is the `Message` Input type
+        2u8 => {
+            // Need to skip over eighteen words, so add 8*18=144
+            144
+        },
+        _ => {
+            return Option::None;
+        },
+    };
+
+    let ptr = tx_input_pointer(index);
+    Option::Some(~Address::from(b256_from_pointer_offset(
+        ptr,
+        owner_offset
+    )))
+}
+
 /// Get the type of an input at a given index
 pub fn tx_input_type(index: u64) -> u8 {
     let ptr = tx_input_pointer(index);
@@ -216,13 +245,6 @@ pub fn b256_from_pointer_offset(pointer: u64, offset: u64) -> b256 {
         buffer: b256
     }
 }
-/// If the input's type is `InputCoin`, return the owner.
-/// Otherwise, undefined behavior.
-pub fn tx_input_coin_owner(index: u64) -> Address {
-    let input_ptr = tx_input_pointer(index);
-    // Need to skip over six words, so offset is 8*6=48
-    ~Address::from(b256_from_pointer_offset(input_ptr, 48))
-}
 
 ////////////////////////////////////////
 // Inputs > Predicate

test/src/sdk-harness/test_artifacts/tx_contract/src/main.sw

@@ -1,6 +1,7 @@
 contract;
 
 use std::address::Address;
+use std::option::Option;
 use std::tx::*;
 
 abi TxContractTest {
@@ -17,6 +18,7 @@ abi TxContractTest {
     fn get_tx_receipts_root() -> b256;
     fn get_tx_script_start_pointer() -> u64;
 
+    fn get_tx_input_type_from_ptr(ptr: u64) -> u8;
     fn get_tx_input_pointer(index: u64) -> u64;
     fn get_tx_input_type(ptr: u64) -> u8;
     fn get_tx_input_coin_owner(index: u64) -> Address;
@@ -66,11 +68,18 @@ impl TxContractTest for Contract {
     fn get_tx_input_pointer(index: u64) -> u64 {
         tx_input_pointer(index)
     }
-    fn get_tx_input_type(ptr: u64) -> u8 {
+    fn get_tx_input_type_from_ptr(ptr: u64) -> u8 {
         tx_input_type_from_pointer(ptr)
     }
+    fn get_tx_input_type(index: u64) -> u8 {
+        tx_input_type(index)
+    }
+    // TODO: Add test for getting InputMessage owner when we have InputMessages
+    // fn get_tx_input_message_owner(index: u64) -> Address {
+    //     tx_input_owner(index)
+    // }
     fn get_tx_input_coin_owner(index: u64) -> Address {
-        tx_input_coin_owner(index)
+        tx_input_owner(index).unwrap()
     }
     fn get_tx_output_pointer(index: u64) -> u64 {
         tx_output_pointer(index)

test/src/sdk-harness/test_projects/tx_fields/mod.rs

@@ -207,8 +207,9 @@ async fn can_get_tx_input_type() {
         .call()
         .await
         .unwrap();
+
     let result = contract_instance
-        .get_tx_input_type(result_ptr.value)
+        .get_tx_input_type_from_ptr(result_ptr.value)
         .call()
         .await
         .unwrap();
@@ -222,23 +223,25 @@ async fn can_get_tx_input_type() {
         .await
         .unwrap();
     let result = contract_instance
-        .get_tx_input_type(result_ptr.value)
+        .get_tx_input_type_from_ptr(result_ptr.value)
         .call()
         .await
         .unwrap();
     assert_eq!(result.value, input_type);
 }
 
+// TODO: Add tests for getting InputMessage owner, type when InputMessages land.
 #[tokio::test]
 async fn can_get_tx_input_coin_owner() {
     let (contract_instance, _, wallet) = get_contracts().await;
 
-    let result = contract_instance
+    let owner_result = contract_instance
         .get_tx_input_coin_owner(1)
         .call()
         .await
         .unwrap();
-    assert_eq!(result.value, wallet.address());
+
+    assert_eq!(owner_result.value, wallet.address());
 }
 
 #[tokio::test]