refactor: make calculate_num_of_elements a ParamType method

FuelLabs · Mar 6, 2023 · 23cf3ca · 23cf3ca
1 parent 4d431bf
commit 23cf3ca
Show file tree

Hide file tree

Showing 2 changed files with 26 additions and 23 deletions.
diff --git a/packages/fuels-core/src/abi_decoder.rs b/packages/fuels-core/src/abi_decoder.rs
@@ -50,7 +50,7 @@ impl ABIDecoder {
         if !param_type.contains_no_nested_heap_types() {
             return Err(error!(
                 InvalidData,
-                "Type {param_type:?} contains nested vectors, this is not supported."
+                "Type {param_type:?} contains nested heap types (`Vector` or `Bytes`), this is not supported."
             ));
         }
         Ok(Self::decode_param(param_type, bytes)?.token)
@@ -86,7 +86,7 @@ impl ABIDecoder {
     }
 
     fn decode_bytes(bytes: &[u8]) -> Result<DecodeResult> {
-        let num_of_elements = calculate_num_of_elements(&ParamType::Byte, bytes, true)?;
+        let num_of_elements = ParamType::Bytes.calculate_num_of_elements(bytes)?;
         let (tokens, bytes_read) =
             Self::decode_multiple(&vec![ParamType::Byte; num_of_elements], bytes)?;
 
@@ -105,7 +105,8 @@ impl ABIDecoder {
         })
     }
     fn decode_vector(param_type: &ParamType, bytes: &[u8]) -> Result<DecodeResult> {
-        let num_of_elements = calculate_num_of_elements(param_type, bytes, false)?;
+        let num_of_elements =
+            ParamType::Vector(Box::from(param_type.clone())).calculate_num_of_elements(bytes)?;
         let (tokens, bytes_read) = Self::decode_multiple(vec![param_type; num_of_elements], bytes)?;
 
         Ok(DecodeResult {
@@ -160,7 +161,7 @@ impl ABIDecoder {
     }
 
     fn decode_raw_slice(bytes: &[u8]) -> Result<DecodeResult> {
-        let num_of_elements = calculate_num_of_elements(&ParamType::U64, bytes, false)?;
+        let num_of_elements = ParamType::RawSlice.calculate_num_of_elements(bytes)?;
         let (tokens, bytes_read) =
             Self::decode_multiple(&vec![ParamType::U64; num_of_elements], bytes)?;
         let elements = tokens
@@ -361,25 +362,6 @@ fn skip(slice: &[u8], num_bytes: usize) -> Result<&[u8]> {
     }
 }
 
-fn calculate_num_of_elements(
-    param_type: &ParamType,
-    bytes: &[u8],
-    is_byte_packed: bool,
-) -> Result<usize> {
-    let memory_size = if is_byte_packed {
-        param_type.compute_encoding_width()
-    } else {
-        param_type.compute_encoding_width() * WORD_SIZE
-    };
-    if bytes.len() % memory_size != 0 {
-        return Err(error!(
-            InvalidData,
-            "The bytes provided do not correspond to a Vec<{:?}> got: {:?}", param_type, bytes
-        ));
-    }
-    Ok(bytes.len() / memory_size)
-}
-
 #[cfg(test)]
 mod tests {
     use std::vec;

diff --git a/packages/fuels-types/src/param_types.rs b/packages/fuels-types/src/param_types.rs
@@ -70,6 +70,27 @@ impl ParamType {
         }
     }
 
+    pub fn calculate_num_of_elements(&self, bytes: &[u8]) -> Result<usize> {
+        let memory_size = match self {
+            // The `Bytes` type in the VM uses byte-packing, don't multiply by WORD_SIZE
+            ParamType::Bytes => Ok(ParamType::U8.compute_encoding_width()),
+            ParamType::Vector(param_type) => Ok(param_type.compute_encoding_width() * WORD_SIZE),
+            ParamType::RawSlice => Ok(ParamType::U64.compute_encoding_width() * WORD_SIZE),
+            _ => Err(error!(
+                InvalidData,
+                "Method `calculate_num_of_elements` called on an incompatible type: {self:?}"
+            )),
+        }
+        .unwrap();
+        if bytes.len() % memory_size != 0 {
+            return Err(error!(
+                InvalidData,
+                "The bytes provided do not correspond to a {:?} got: {:?}", self, bytes
+            ));
+        }
+        Ok(bytes.len() / memory_size)
+    }
+
     pub fn contains_no_nested_heap_types(&self) -> bool {
         match &self {
             ParamType::Vector(param_type) | ParamType::Array(param_type, ..) => {