Repository: huff-language/huff-rs
Pull request #119 (2022-06-23 21:47:27)

feat+fix(huff_utils+huff_parser): abi.rs support for arrays
## Overview

`abi::FunctionParamType` was missing support for arrays.

Note that if the array has a size of 0, it means it's unbounded.

When writing a test for the `str` to `FunctionParamType` conversions, I encountered a bug in `huff_parser` which used a `TokenKind::Comma` instead of the proper `TokenKind::ArrayType` if we defined a n-dimensional array that is preceding another arg.



Patch:
diff --git a/huff_lexer/src/lib.rs b/huff_lexer/src/lib.rs
index e474da4..cea3965 100644
--- a/huff_lexer/src/lib.rs
+++ b/huff_lexer/src/lib.rs
@@ -477,7 +477,6 @@ impl<'a> Iterator for Lexer<'a> {
                                     .split(&raw_type)
                                     .map(|x| x.replace(']', ""))
                                     .collect();
-
                                 let mut size_vec: Vec<usize> = Vec::new();
                                 // go over all array sizes
                                 let sizes = words.get(1..words.len()).unwrap();
diff --git a/huff_parser/src/lib.rs b/huff_parser/src/lib.rs
index df831aa..174ee3f 100644
--- a/huff_parser/src/lib.rs
+++ b/huff_parser/src/lib.rs
@@ -870,8 +870,12 @@ impl Parser {
         match self.current_token.kind.clone() {
             TokenKind::PrimitiveType(prim) => Ok(self.parse_primitive_type(prim)?),
             TokenKind::ArrayType(prim, _) => {
+                // The trick is that when we parse the primitive type
+                // of the array, it will consume the current token which is the ArrayType.
+                // So we have to preserve the token before parsing (and matching thus consuming).
+                let token = self.current_token.kind.clone();
                 let _ = self.parse_primitive_type(prim);
-                Ok(self.match_kind(self.current_token.kind.clone())?)
+                Ok(token)
             }
             kind => Err(ParserError {
                 kind: ParserErrorKind::InvalidArgs(kind),
diff --git a/huff_utils/src/abi.rs b/huff_utils/src/abi.rs
index 08ff814..a8caf45 100644
--- a/huff_utils/src/abi.rs
+++ b/huff_utils/src/abi.rs
@@ -245,20 +245,29 @@ pub enum FunctionParamType {
     Bool,
     /// A String
     String,
-    /// An array of parameters
-    Array(Box<FunctionParamType>),
+    /// Array ; uint256[2][] => Array(Uint(256), [2, 0])
+    Array(Box<FunctionParamType>, Vec<usize>),
     /// Fixed number of bytes
     FixedBytes(usize),
-    /// Fixed size array of parameters
-    FixedArray(Box<FunctionParamType>, usize),
     /// A tuple of parameters
     Tuple(Vec<FunctionParamType>),
 }
 
 impl FunctionParamType {
     /// Convert string to type
-    pub fn convert_string_to_type(string: &str) -> Self {
+    pub fn convert_string_to_type(string: &str) -> Result<Self, Box<dyn std::error::Error>> {
         let input = string.to_string().to_lowercase();
+        let split_input: Vec<&str> = input.split('[').collect();
+        if split_input.len() > 1 {
+            let mut cleaned: Vec<String> = split_input
+                .iter()
+                .map(|x| x.replace(']', ""))
+                .map(|x| if x.is_empty() { "0".to_owned() } else { x })
+                .collect();
+            let func_type = FunctionParamType::convert_string_to_type(&cleaned.remove(0))?;
+            let sizes: Vec<usize> = cleaned.iter().map(|x| x.parse::<usize>().unwrap()).collect();
+            return Ok(Self::Array(Box::new(func_type), sizes))
+        }
         if input.starts_with("uint") {
             // Default to 256 if no size
             let size = match input.get(4..input.len()) {
@@ -268,7 +277,7 @@ impl FunctionParamType {
                 },
                 None => 256,
             };
-            return Self::Uint(size)
+            return Ok(Self::Uint(size))
         }
         if input.starts_with("int") {
             // Default to 256 if no size
@@ -279,38 +288,38 @@ impl FunctionParamType {
                 },
                 None => 256,
             };
-            return Self::Int(size)
+            return Ok(Self::Int(size))
         }
         if input.starts_with("bytes") && input.len() != 5 {
             let size = input.get(5..input.len()).unwrap().parse::<usize>().unwrap();
-            return Self::FixedBytes(size)
+            return Ok(Self::FixedBytes(size))
         }
         if input.starts_with("bool") {
-            return Self::Bool
+            return Ok(Self::Bool)
         }
         if input.starts_with("address") {
-            return Self::Address
+            return Ok(Self::Address)
         }
         if input.starts_with("string") {
-            return Self::String
+            return Ok(Self::String)
         }
         if input == "bytes" {
-            Self::Bytes
+            Ok(Self::Bytes)
         } else {
             tracing::error!("Failed to create FunctionParamType from string: {}", string);
-            panic!("{}", format!("Failed to create FunctionParamType from string: {}", string))
+            Err(format!("Failed to create FunctionParamType from string: {}", string))?
         }
     }
 }
 
 impl From<&str> for FunctionParamType {
     fn from(string: &str) -> Self {
-        FunctionParamType::convert_string_to_type(string)
+        FunctionParamType::convert_string_to_type(string).unwrap()
     }
 }
 
 impl From<String> for FunctionParamType {
     fn from(string: String) -> Self {
-        FunctionParamType::convert_string_to_type(&string)
+        FunctionParamType::convert_string_to_type(&string).unwrap()
     }
 }


Test patch:
diff --git a/huff_parser/tests/abi.rs b/huff_parser/tests/abi.rs
new file mode 100644
index 0000000..7197dd2
--- /dev/null
+++ b/huff_parser/tests/abi.rs
@@ -0,0 +1,25 @@
+use huff_lexer::*;
+use huff_parser::*;
+use huff_utils::{abi::*, prelude::*};
+
+#[test]
+fn build_abi_from_ast() {
+    let source = "#define function test(uint256[2][],string) view returns(uint256)";
+
+    let flattened_source = FullFileSource { source, file: None, spans: vec![] };
+    let lexer = Lexer::new(flattened_source);
+    let tokens = lexer
+        .into_iter()
+        .map(|x| x.unwrap())
+        .filter(|x| !matches!(x.kind, TokenKind::Whitespace))
+        .collect::<Vec<Token>>();
+    let mut parser = Parser::new(tokens, None);
+    let contract = parser.parse().unwrap();
+    let abi = Abi::from(contract);
+
+    assert_eq!(
+        abi.functions.get("test").unwrap().inputs[0].kind,
+        FunctionParamType::Array(Box::new(FunctionParamType::Uint(256)), vec![2, 0])
+    );
+    assert_eq!(abi.functions.get("test").unwrap().inputs[1].kind, FunctionParamType::String);
+}
diff --git a/huff_utils/tests/abi.rs b/huff_utils/tests/abi.rs
new file mode 100644
index 0000000..ba6c03d
--- /dev/null
+++ b/huff_utils/tests/abi.rs
@@ -0,0 +1,34 @@
+use huff_utils::abi::*;
+use std::collections::HashMap;
+
+#[test]
+fn convert_function_param_type() {
+    let func_types = [
+        "bool",
+        "string",
+        "address",
+        "uint198",
+        "bytes32",
+        "bytes",
+        "uint[][]",
+        "address[5][]",
+        "bool[][1][][2]",
+    ];
+
+    let expected_fn_types = HashMap::from([
+        (0, FunctionParamType::Bool),
+        (1, FunctionParamType::String),
+        (2, FunctionParamType::Address),
+        (3, FunctionParamType::Uint(198)),
+        (4, FunctionParamType::FixedBytes(32)),
+        (5, FunctionParamType::Bytes),
+        (6, FunctionParamType::Array(Box::new(FunctionParamType::Uint(256)), vec![0, 0])),
+        (7, FunctionParamType::Array(Box::new(FunctionParamType::Address), vec![5, 0])),
+        (8, FunctionParamType::Array(Box::new(FunctionParamType::Bool), vec![0, 1, 0, 2])),
+    ]);
+
+    for (index, func_type) in func_types.into_iter().enumerate() {
+        let ac_func_type = FunctionParamType::convert_string_to_type(func_type).unwrap();
+        assert_eq!(ac_func_type, *expected_fn_types.get(&index).unwrap());
+    }
+}