Add more builtin functions

danhper · Jul 12, 2024 · 0a8d566 · 0a8d566
1 parent 62e15e2
commit 0a8d566
Show file tree

Hide file tree

Showing 3 changed files with 44 additions and 25 deletions.
diff --git a/src/interpreter/builtin_functions.rs b/src/interpreter/builtin_functions.rs
@@ -90,10 +90,21 @@ fn concat_arrays(arr: Vec<Value>, args: &[Value]) -> Result<Vec<Value>> {
     }
 }
 
+fn concat_bytes(bytes: Vec<u8>, args: &[Value]) -> Result<Vec<u8>> {
+    if let Some(Value::Bytes(other)) = args.first() {
+        let mut new_bytes = bytes.clone();
+        new_bytes.extend(other.clone());
+        Ok(new_bytes)
+    } else {
+        bail!("cannot concat {:?} with {:?}", bytes, args)
+    }
+}
+
 fn concat(value: &Value, args: &[Value]) -> Result<Value> {
     match value {
         Value::Str(s) => concat_strings(s.clone(), args).map(Value::Str),
         Value::Array(arr) => concat_arrays(arr.clone(), args).map(Value::Array),
+        Value::Bytes(b) => concat_bytes(b.clone(), args).map(Value::Bytes),
         _ => bail!("cannot concat {}", value),
     }
 }
@@ -257,6 +268,7 @@ pub enum BuiltinFunction {
     Min(Type),
     Max(Type),
     Concat(Box<Value>),
+    Length(Box<Value>),
     Decode(String, JsonAbi),
     Map(Vec<Value>, Type),
     Keys(HashableIndexMap<Value, Value>),
@@ -278,6 +290,7 @@ impl fmt::Display for BuiltinFunction {
             Self::Balance(addr) => write!(f, "{}.balance", addr),
             Self::Format(v) => write!(f, "{}.format", v),
             Self::Concat(s) => write!(f, "{}.concat", s),
+            Self::Length(s) => write!(f, "{}.length", s),
             Self::Mul(v) => write!(f, "{}.mul", v),
             Self::Div(v) => write!(f, "{}.div", v),
             Self::Min(t) => write!(f, "{}.min", t),
@@ -327,7 +340,9 @@ impl BuiltinFunction {
 
             (v, "format") => Self::Format(Box::new(v.clone())),
 
-            (v @ (Value::Str(_) | Value::Array(_)), "concat") => Self::Concat(Box::new(v.clone())),
+            (v @ (Value::Str(_) | Value::Array(_) | Value::Bytes(_)), "concat") => {
+                Self::Concat(Box::new(v.clone()))
+            }
 
             (v @ Value::Uint(..) | v @ Value::Int(..), "mul") => Self::Mul(Box::new(v.clone())),
             (v @ Value::Uint(..) | v @ Value::Int(..), "div") => Self::Div(Box::new(v.clone())),
@@ -336,10 +351,15 @@ impl BuiltinFunction {
                 values.clone(),
                 Type::Tuple(values.iter().map(Value::get_type).collect()),
             ),
+
             (Value::Array(values), "map") => {
                 let arr_type = values.first().map_or(Type::Uint(256), Value::get_type);
                 Self::Map(values.clone(), Type::Array(Box::new(arr_type)))
             }
+            (
+                v @ (Value::Array(_) | Value::Bytes(_) | Value::Str(_) | Value::Tuple(_)),
+                "length",
+            ) => Self::Length(Box::new(v.clone())),
 
             (Value::TypeObject(Type::Type(t)), "max") if t.is_int() => Self::Max(*t.clone()),
             (Value::TypeObject(Type::Type(t)), "min") if t.is_int() => Self::Min(*t.clone()),
@@ -376,6 +396,7 @@ impl BuiltinFunction {
             | Self::Block(_)
             | Self::Min(_)
             | Self::Max(_)
+            | Self::Length(_)
             | Self::ReadReceipt(_, _)
             | Self::Keys(_) => true,
             Self::Directive(d) => d.is_property(),
@@ -405,6 +426,8 @@ impl BuiltinFunction {
             Self::Max(t) => t.max(),
             Self::Min(t) => t.min(),
 
+            Self::Length(v) => v.len().map(|v| Value::Uint(U256::from(v), 256)),
+
             Self::Keys(values) => Ok(Value::Array(values.0.keys().cloned().collect_vec())),
 
             Self::Keccak256 => keccak256(args),

diff --git a/src/interpreter/types.rs b/src/interpreter/types.rs
@@ -420,8 +420,14 @@ impl Type {
             }
             Type::Transaction => vec!["getReceipt".to_string()],
             Type::TransactionReceipt => Receipt::keys(),
-            Type::Array(_) => vec!["concat".to_string()],
-            Type::String => vec!["concat".to_string()],
+            Type::Array(_) => vec![
+                "concat".to_string(),
+                "length".to_string(),
+                "map".to_string(),
+            ],
+            Type::String => vec!["concat".to_string(), "length".to_string()],
+
+            Type::Bytes => vec!["concat".to_string(), "length".to_string()],
 
             Type::Mapping(_, _) => vec!["keys".to_string()],
 

diff --git a/src/interpreter/value.rs b/src/interpreter/value.rs
@@ -16,7 +16,7 @@ use super::{
     types::{ContractInfo, HashableIndexMap, Receipt, Type},
 };
 
-#[derive(Debug, Clone, Hash, Eq)]
+#[derive(Debug, Clone, Hash, PartialEq, Eq)]
 pub enum Value {
     Null,
     Bool(bool),
@@ -217,27 +217,6 @@ where
     }
 }
 
-impl PartialEq for Value {
-    fn eq(&self, other: &Self) -> bool {
-        match (self, other) {
-            (Value::Bool(a), Value::Bool(b)) => a == b,
-            (Value::Int(a, _), Value::Int(b, _)) => a == b,
-            (Value::Uint(a, _), Value::Uint(b, _)) => a == b,
-            (Value::Int(a, _), Value::Uint(b, _)) => *a == I256::from_raw(*b),
-            (Value::Uint(a, _), Value::Int(b, _)) => I256::from_raw(*a) == *b,
-            (Value::Str(a), Value::Str(b)) => a == b,
-            (Value::Addr(a), Value::Addr(b)) => a == b,
-            (Value::FixBytes(a, _), Value::FixBytes(b, _)) => a == b,
-            (Value::Bytes(a), Value::Bytes(b)) => a == b,
-            (Value::Tuple(a), Value::Tuple(b)) => a == b,
-            (Value::Array(a), Value::Array(b)) => a == b,
-            (Value::TypeObject(a), Value::TypeObject(b)) => a == b,
-            (Value::Contract(_, a), Value::Contract(_, b)) => a == b,
-            _ => false,
-        }
-    }
-}
-
 impl PartialOrd for Value {
     fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
         match (self, other) {
@@ -289,6 +268,17 @@ impl Value {
         }
     }
 
+    pub fn len(&self) -> Result<usize> {
+        let len = match self {
+            Value::Array(items) => items.len(),
+            Value::Tuple(items) => items.len(),
+            Value::Bytes(b) => b.len(),
+            Value::Str(s) => s.len(),
+            v => bail!("{} is not iterable", v.get_type()),
+        };
+        Ok(len)
+    }
+
     pub fn set_index(&mut self, index: &Value, value: Value) -> Result<()> {
         match self {
             Value::Array(items) => {