Change input predicate offset to message & coin (#174)

Resolves #166
FuelLabs · Aug 2, 2022 · 9b5bde8 · 9b5bde8
1 parent 9347abc
commit 9b5bde8
Show file tree

Hide file tree

Showing 8 changed files with 107 additions and 59 deletions.
diff --git a/src/builder.rs b/src/builder.rs
@@ -12,12 +12,28 @@ use alloc::vec::Vec;
 pub struct TransactionBuilder {
     tx: Transaction,
 
+    should_prepare_script: bool,
+    should_prepare_predicate: bool,
+
     // We take the key by reference so this lib won't have the responsibility to properly zeroize
     // the keys
     sign_keys: Vec<SecretKey>,
 }
 
 impl TransactionBuilder {
+    fn with_tx(tx: Transaction) -> Self {
+        let should_prepare_script = false;
+        let should_prepare_predicate = false;
+        let sign_keys = Vec::new();
+
+        Self {
+            tx,
+            should_prepare_script,
+            should_prepare_predicate,
+            sign_keys,
+        }
+    }
+
     pub fn create(bytecode: Witness, salt: Salt, storage_slots: Vec<StorageSlot>) -> Self {
         let mut tx = Transaction::create(
             Default::default(),
@@ -33,9 +49,7 @@ impl TransactionBuilder {
 
         tx._set_bytecode(bytecode);
 
-        let sign_keys = Vec::new();
-
-        Self { tx, sign_keys }
+        Self::with_tx(tx)
     }
 
     pub fn script(script: Vec<u8>, script_data: Vec<u8>) -> Self {
@@ -49,9 +63,22 @@ impl TransactionBuilder {
             Default::default(),
             Default::default(),
         );
-        let sign_keys = Vec::new();
 
-        Self { tx, sign_keys }
+        let mut slf = Self::with_tx(tx);
+
+        slf.prepare_script(true);
+
+        slf
+    }
+
+    pub fn prepare_script(&mut self, should_prepare_script: bool) -> &mut Self {
+        self.should_prepare_script = should_prepare_script;
+        self
+    }
+
+    pub fn prepare_predicate(&mut self, should_prepare_predicate: bool) -> &mut Self {
+        self.should_prepare_predicate = should_prepare_predicate;
+        self
     }
 
     pub fn sign_keys(&self) -> &[SecretKey] {
@@ -144,8 +171,22 @@ impl TransactionBuilder {
         self
     }
 
+    fn prepare_finalize(&mut self) {
+        if self.should_prepare_predicate {
+            self.tx.prepare_init_predicate();
+        }
+
+        if self.should_prepare_script {
+            self.tx
+                .prepare_init_script()
+                .expect("failed to prepare script");
+        }
+    }
+
     #[cfg(feature = "std")]
     pub fn finalize(&mut self) -> Transaction {
+        self.prepare_finalize();
+
         let mut tx = core::mem::take(&mut self.tx);
 
         self.sign_keys.iter().for_each(|k| tx.sign_inputs(k));
@@ -157,6 +198,8 @@ impl TransactionBuilder {
 
     #[cfg(feature = "std")]
     pub fn finalize_without_signature(&mut self) -> Transaction {
+        self.prepare_finalize();
+
         let mut tx = core::mem::take(&mut self.tx);
 
         tx.precompute_metadata();

diff --git a/src/checked_transaction.rs b/src/checked_transaction.rs
@@ -250,7 +250,7 @@ impl CheckedTransaction {
 
     /// Prepare the transaction for VM initialization for script execution
     #[cfg(feature = "std")]
-    pub fn prepare_init_script<F>(&mut self) -> io::Result<&mut Self> {
+    pub fn prepare_init_script(&mut self) -> io::Result<&mut Self> {
         self.transaction.prepare_init_script()?;
         Ok(self)
     }

diff --git a/src/transaction.rs b/src/transaction.rs
@@ -475,11 +475,7 @@ impl Transaction {
     /// should already set by the client beforehand.
     #[cfg(feature = "std")]
     pub fn prepare_init_script(&mut self) -> io::Result<&mut Self> {
-        let outputs = match self {
-            Transaction::Script { outputs, .. } | Transaction::Create { outputs, .. } => outputs,
-        };
-
-        outputs
+        self._outputs_mut()
             .iter_mut()
             .try_for_each(|o| o.prepare_init_script())?;
 

diff --git a/src/transaction/metadata.rs b/src/transaction/metadata.rs
@@ -13,7 +13,7 @@ use alloc::vec::Vec;
 pub struct Metadata {
     id: Bytes32,
     script_data_offset: Option<usize>,
-    input_coin_predicate_offset: Vec<Option<(usize, usize)>>,
+    input_predicate_offset: Vec<Option<(usize, usize)>>,
     inputs_offset: Vec<usize>,
     outputs_offset: Vec<usize>,
     witnesses_offset: Vec<usize>,
@@ -23,15 +23,15 @@ impl Metadata {
     pub const fn new(
         id: Bytes32,
         script_data_offset: Option<usize>,
-        input_coin_predicate_offset: Vec<Option<(usize, usize)>>,
+        input_predicate_offset: Vec<Option<(usize, usize)>>,
         inputs_offset: Vec<usize>,
         outputs_offset: Vec<usize>,
         witnesses_offset: Vec<usize>,
     ) -> Self {
         Self {
             id,
             script_data_offset,
-            input_coin_predicate_offset,
+            input_predicate_offset,
             inputs_offset,
             outputs_offset,
             witnesses_offset,
@@ -46,11 +46,8 @@ impl Metadata {
         self.script_data_offset
     }
 
-    pub fn input_coin_predicate_offset(&self, index: usize) -> Option<(usize, usize)> {
-        self.input_coin_predicate_offset
-            .get(index)
-            .copied()
-            .flatten()
+    pub fn input_predicate_offset(&self, index: usize) -> Option<(usize, usize)> {
+        self.input_predicate_offset.get(index).copied().flatten()
     }
 
     pub fn inputs_offset(&self, index: usize) -> Option<usize> {
@@ -79,11 +76,11 @@ impl Transaction {
         let id = self._id();
 
         let script_data_offset = self._script_data_offset();
-        let input_coin_predicate_offset = self
+        let input_predicate_offset = self
             .inputs()
             .iter()
             .enumerate()
-            .map(|(i, _)| self._input_coin_predicate_offset(i))
+            .map(|(i, _)| self._input_predicate_offset(i))
             .collect();
 
         let offset = self.inputs_offset();
@@ -125,7 +122,7 @@ impl Transaction {
         let metadata = Metadata::new(
             id,
             script_data_offset,
-            input_coin_predicate_offset,
+            input_predicate_offset,
             inputs_offset,
             outputs_offset,
             witnesses_offset,

diff --git a/src/transaction/offset.rs b/src/transaction/offset.rs
@@ -1,5 +1,5 @@
 use super::{TRANSACTION_CREATE_FIXED_SIZE, TRANSACTION_SCRIPT_FIXED_SIZE};
-use crate::{Input, Metadata, StorageSlot, Transaction};
+use crate::{Metadata, StorageSlot, Transaction};
 
 use fuel_types::bytes::{self, SizedBytes};
 use fuel_types::Bytes32;
@@ -31,21 +31,18 @@ impl Transaction {
     /// Return the offset/length tuple for an input predicate indexed by `index`.
     ///
     /// Return `None` either if `index` is invalid, or if its not a [`Input::CoinPredicate`]
-    pub fn input_coin_predicate_offset(&self, index: usize) -> Option<(usize, usize)> {
+    pub fn input_predicate_offset(&self, index: usize) -> Option<(usize, usize)> {
         self.metadata()
-            .map(|m| m.input_coin_predicate_offset(index))
-            .unwrap_or_else(|| self._input_coin_predicate_offset(index))
+            .map(|m| m.input_predicate_offset(index))
+            .unwrap_or_else(|| self._input_predicate_offset(index))
     }
 
-    pub(crate) fn _input_coin_predicate_offset(&self, index: usize) -> Option<(usize, usize)> {
-        self.input_offset(index)
-            .map(|ofs| ofs + Input::coin_predicate_offset())
-            .and_then(|ofs| {
-                self.inputs()
-                    .get(index)
-                    .and_then(Input::coin_predicate_len)
-                    .map(|l| (ofs, l))
-            })
+    pub(crate) fn _input_predicate_offset(&self, index: usize) -> Option<(usize, usize)> {
+        self.inputs().get(index).and_then(|i| {
+            i.predicate_offset()
+                .and_then(|p| self.input_offset(index).map(|o| p + o))
+                .zip(i.predicate_len().map(bytes::padded_len_usize))
+        })
     }
 
     /// Return the serialized bytes offset of the input with the provided index

diff --git a/src/transaction/types/input.rs b/src/transaction/types/input.rs
@@ -442,24 +442,6 @@ impl Input {
         INPUT_MESSAGE_FIXED_SIZE
     }
 
-    pub fn coin_predicate_len(&self) -> Option<usize> {
-        match self {
-            Input::CoinPredicate { predicate, .. } => Some(bytes::padded_len(predicate.as_slice())),
-
-            _ => None,
-        }
-    }
-
-    pub fn coin_predicate_data_offset(&self) -> Option<usize> {
-        match self {
-            Input::CoinPredicate { predicate, .. } => {
-                Some(Self::coin_predicate_offset() + bytes::padded_len(predicate.as_slice()))
-            }
-
-            _ => None,
-        }
-    }
-
     pub const fn balance_root(&self) -> Option<&Bytes32> {
         match self {
             Input::Contract { balance_root, .. } => Some(balance_root),

diff --git a/tests/bytes.rs b/tests/bytes.rs
@@ -642,7 +642,7 @@ fn transaction() {
 }
 
 #[test]
-fn create_input_coin_data_offset() {
+fn create_input_data_offset() {
     let rng = &mut StdRng::seed_from_u64(8586);
 
     let gas_price = 100;
@@ -691,6 +691,19 @@ fn create_input_coin_data_offset() {
         rng.gen(),
         rng.gen(),
         predicate.clone(),
+        predicate_data.clone(),
+    );
+
+    let data = generate_bytes(rng);
+    let input_message = Input::message_predicate(
+        rng.gen(),
+        rng.gen(),
+        rng.gen(),
+        rng.gen(),
+        rng.gen(),
+        owner,
+        data,
+        predicate.clone(),
         predicate_data,
     );
 
@@ -700,9 +713,13 @@ fn create_input_coin_data_offset() {
             for outputs in outputs.iter() {
                 for witnesses in witnesses.iter() {
                     let mut inputs = inputs.clone();
-                    let last_input = inputs.len();
+
+                    let input_coin_idx = inputs.len();
                     inputs.push(input_coin.clone());
 
+                    let input_message_idx = inputs.len();
+                    inputs.push(input_message.clone());
+
                     dbg!(&inputs);
 
                     let mut tx = Transaction::create(
@@ -728,11 +745,25 @@ fn create_input_coin_data_offset() {
                         .expect("Failed to serialize input");
 
                     let (offset, len) = tx
-                        .input_coin_predicate_offset(last_input)
+                        .input_predicate_offset(input_coin_idx)
+                        .expect("Failed to fetch offset");
+
+                    let (offset_p, _) = tx_p
+                        .input_predicate_offset(input_coin_idx)
+                        .expect("Failed to fetch offset from tx with precomputed metadata!");
+
+                    assert_eq!(offset, offset_p);
+                    assert_eq!(
+                        predicate.as_slice(),
+                        &buffer[offset..offset + len][..predicate.len()]
+                    );
+
+                    let (offset, len) = tx
+                        .input_predicate_offset(input_message_idx)
                         .expect("Failed to fetch offset");
 
                     let (offset_p, _) = tx_p
-                        .input_coin_predicate_offset(last_input)
+                        .input_predicate_offset(input_message_idx)
                         .expect("Failed to fetch offset from tx with precomputed metadata!");
 
                     assert_eq!(offset, offset_p);
@@ -858,10 +889,10 @@ fn script_input_coin_data_offset() {
                         );
 
                         let (offset, len) = tx
-                            .input_coin_predicate_offset(offset)
+                            .input_predicate_offset(offset)
                             .expect("Failed to fetch offset");
 
-                        assert_ne!(predicate.len(), len);
+                        assert_ne!(bytes::padded_len(&predicate), predicate.len());
                         assert_eq!(bytes::padded_len(&predicate), len);
 
                         assert_eq!(

diff --git a/tests/prepared_init.rs b/tests/prepared_init.rs
@@ -10,6 +10,7 @@ fn output_message_prepare_init_zeroes_recipient_and_amount() {
     let zeroed = Output::message(Address::zeroed(), 0);
 
     let tx = TransactionBuilder::script(vec![], vec![])
+        .prepare_script(false)
         .add_output(message)
         .finalize();
 
@@ -40,6 +41,7 @@ fn output_variable_prepare_init_zeroes_recipient_and_amount() {
     let zeroed = Output::variable(Address::zeroed(), 0, AssetId::zeroed());
 
     let tx = TransactionBuilder::script(vec![], vec![])
+        .prepare_script(false)
         .add_output(variable)
         .finalize();