From 2e7d3198a7766ffa6b6ca561d9ec32990a1a14c7 Mon Sep 17 00:00:00 2001
From: Daniel Perez <daniel@perez.sh>
Date: Fri, 2 Aug 2024 12:37:29 +0100
Subject: [PATCH] Add slicing on bytes and strings

---
 CHANGELOG.md                   |  1 +
 src/interpreter/interpreter.rs | 18 +++++--------
 src/interpreter/value.rs       | 48 ++++++++++++++++++++++++++++++++++
 3 files changed, 55 insertions(+), 12 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 9b47eba..eabb878 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -15,6 +15,7 @@
 * Allow to select version of Eclair when installing using install script
 * Add `repl.loadKeystore` to load keystore from a file created by `cast`
 * Allow conversion between different fixed-size bytes types (e.g. bytes32 -> bytes4 or vice-versa)
+* Allow slicing on bytes and strings
 
 ### Bug fixes
 
diff --git a/src/interpreter/interpreter.rs b/src/interpreter/interpreter.rs
index d5d246a..4b75617 100644
--- a/src/interpreter/interpreter.rs
+++ b/src/interpreter/interpreter.rs
@@ -463,22 +463,16 @@ pub fn evaluate_expression(env: &mut Env, expr: Box<Expression>) -> BoxFuture<'_
             }
 
             Expression::ArraySlice(_, arr_expr, start_expr, end_expr) => {
-                let (values, type_) = match evaluate_expression(env, arr_expr).await? {
-                    Value::Array(v, t) => (v, t),
-                    v => bail!("invalid type for slice, expected tuple, got {}", v),
-                };
+                let value = evaluate_expression(env, arr_expr).await?;
                 let start = match start_expr {
-                    Some(expr) => evaluate_expression(env, expr).await?.as_usize()?,
-                    None => 0,
+                    Some(expr) => Some(evaluate_expression(env, expr).await?.as_usize()?),
+                    None => None,
                 };
                 let end = match end_expr {
-                    Some(expr) => evaluate_expression(env, expr).await?.as_usize()?,
-                    None => values.len(),
+                    Some(expr) => Some(evaluate_expression(env, expr).await?.as_usize()?),
+                    None => None,
                 };
-                if end > values.len() {
-                    bail!("end index out of bounds");
-                }
-                Ok(Value::Array(values[start..end].to_vec(), type_.clone()))
+                value.slice(start, end)
             }
 
             Expression::Add(_, lhs, rhs) => _eval_binop(env, lhs, rhs, Value::add).await,
diff --git a/src/interpreter/value.rs b/src/interpreter/value.rs
index 76633d5..f8075b8 100644
--- a/src/interpreter/value.rs
+++ b/src/interpreter/value.rs
@@ -530,6 +530,29 @@ impl Value {
             }
         }
     }
+
+    pub fn slice(&self, start: Option<usize>, end: Option<usize>) -> Result<Value> {
+        let start = start.unwrap_or(0);
+        let end = end.unwrap_or(self.len()?);
+        if end > self.len()? {
+            bail!("index out of bounds")
+        }
+        match self {
+            Value::Array(items, t) => {
+                let items = items[start..end].to_vec();
+                Ok(Value::Array(items, t.clone()))
+            }
+            Value::Bytes(bytes) => {
+                let bytes = bytes[start..end].to_vec();
+                Ok(Value::Bytes(bytes))
+            }
+            Value::Str(s) => {
+                let s = s.chars().skip(start).take(end - start).collect();
+                Ok(Value::Str(s))
+            }
+            _ => bail!("{} is not sliceable", self.get_type()),
+        }
+    }
 }
 
 impl Add for Value {
@@ -649,4 +672,29 @@ mod tests {
         let fix_bytes = B256::from_slice(&bytes);
         assert_eq!(value, Value::FixBytes(fix_bytes, 4));
     }
+
+    #[test]
+    fn test_slice() {
+        let array = Value::Array(
+            vec![Value::from(1u64), Value::from(2u64), Value::from(3u64)],
+            Box::new(Type::Int(256)),
+        );
+        let slice = array.slice(Some(1), Some(2)).unwrap();
+        assert_eq!(
+            slice,
+            Value::Array(vec![Value::from(2u64)], Box::new(Type::Int(256)))
+        );
+
+        let bytes = Value::Bytes(vec![1, 2, 3]);
+        let slice = bytes.slice(Some(1), Some(2)).unwrap();
+        assert_eq!(slice, Value::Bytes(vec![2]));
+
+        let bytes = Value::Bytes(vec![1, 2, 3]);
+        let slice = bytes.slice(Some(1), None).unwrap();
+        assert_eq!(slice, Value::Bytes(vec![2, 3]));
+
+        let str = Value::Str("hello".to_string());
+        let slice = str.slice(Some(1), Some(3)).unwrap();
+        assert_eq!(slice, Value::Str("el".to_string()));
+    }
 }