Add Expr::as_repeated_array/slice_element

noir-lang · Aug 18, 2024 · f4ba855 · f4ba855
1 parent da2f3b6
commit f4ba855
Show file tree

Hide file tree

Showing 3 changed files with 70 additions and 2 deletions.
diff --git a/compiler/noirc_frontend/src/hir/comptime/interpreter/builtin.rs b/compiler/noirc_frontend/src/hir/comptime/interpreter/builtin.rs
@@ -55,6 +55,12 @@ impl<'local, 'context> Interpreter<'local, 'context> {
             "expr_as_index" => expr_as_index(arguments, return_type, location),
             "expr_as_integer" => expr_as_integer(arguments, return_type, location),
             "expr_as_member_access" => expr_as_member_access(arguments, return_type, location),
+            "expr_as_repeated_element_array" => {
+                expr_as_repeated_element_array(arguments, return_type, location)
+            }
+            "expr_as_repeated_element_slice" => {
+                expr_as_repeated_element_slice(arguments, return_type, location)
+            }
             "expr_as_slice" => expr_as_slice(arguments, return_type, location),
             "expr_as_tuple" => expr_as_tuple(arguments, return_type, location),
             "expr_as_unary_op" => expr_as_unary_op(arguments, return_type, location),
@@ -924,6 +930,44 @@ fn expr_as_member_access(
     })
 }
 
+// fn as_repeated_element_array(self) -> Option<(Expr, Expr)>
+fn expr_as_repeated_element_array(
+    arguments: Vec<(Value, Location)>,
+    return_type: Type,
+    location: Location,
+) -> IResult<Value> {
+    expr_as(arguments, return_type, location, |expr| {
+        if let ExpressionKind::Literal(Literal::Array(ArrayLiteral::Repeated {
+            repeated_element,
+            length,
+        })) = expr
+        {
+            Some(Value::Tuple(vec![Value::Expr(repeated_element.kind), Value::Expr(length.kind)]))
+        } else {
+            None
+        }
+    })
+}
+
+// fn as_repeated_element_slice(self) -> Option<(Expr, Expr)>
+fn expr_as_repeated_element_slice(
+    arguments: Vec<(Value, Location)>,
+    return_type: Type,
+    location: Location,
+) -> IResult<Value> {
+    expr_as(arguments, return_type, location, |expr| {
+        if let ExpressionKind::Literal(Literal::Slice(ArrayLiteral::Repeated {
+            repeated_element,
+            length,
+        })) = expr
+        {
+            Some(Value::Tuple(vec![Value::Expr(repeated_element.kind), Value::Expr(length.kind)]))
+        } else {
+            None
+        }
+    })
+}
+
 // fn as_slice(self) -> Option<[Expr]>
 fn expr_as_slice(
     arguments: Vec<(Value, Location)>,

diff --git a/noir_stdlib/src/meta/expr.nr b/noir_stdlib/src/meta/expr.nr
@@ -27,6 +27,12 @@ impl Expr {
     #[builtin(expr_as_member_access)]
     fn as_member_access(self) -> Option<(Expr, Quoted)> {}
 
+    #[builtin(expr_as_repeated_element_array)]
+    fn as_repeated_element_array(self) -> Option<(Expr, Expr)> {}
+
+    #[builtin(expr_as_repeated_element_slice)]
+    fn as_repeated_element_slice(self) -> Option<(Expr, Expr)> {}
+
     #[builtin(expr_as_slice)]
     fn as_slice(self) -> Option<[Expr]> {}
 

diff --git a/test_programs/compile_success_empty/comptime_exp/src/main.nr b/test_programs/compile_success_empty/comptime_exp/src/main.nr
@@ -73,14 +73,32 @@ fn main() {
         assert_eq(name, quote { bar });
 
         // Check Expr::as_array
-        let expr = quote { [1, 2, 3] }.as_expr().unwrap();
+        let expr = quote { [1, 2, 4] }.as_expr().unwrap();
         let elems = expr.as_array().unwrap();
         assert_eq(elems.len(), 3);
+        assert_eq(elems[0].as_integer().unwrap(), (1, false));
+        assert_eq(elems[1].as_integer().unwrap(), (2, false));
+        assert_eq(elems[2].as_integer().unwrap(), (4, false));
 
         // Check Expr::as_slice
-        let expr = quote { &[1, 2, 3] }.as_expr().unwrap();
+        let expr = quote { &[1, 3, 5] }.as_expr().unwrap();
         let elems = expr.as_slice().unwrap();
         assert_eq(elems.len(), 3);
+        assert_eq(elems[0].as_integer().unwrap(), (1, false));
+        assert_eq(elems[1].as_integer().unwrap(), (3, false));
+        assert_eq(elems[2].as_integer().unwrap(), (5, false));
+
+        // Check Expr::as_repeated_element_array
+        let expr = quote { [1; 3] }.as_expr().unwrap();
+        let (expr, length) = expr.as_repeated_element_array().unwrap();
+        assert_eq(expr.as_integer().unwrap(), (1, false));
+        assert_eq(length.as_integer().unwrap(), (3, false));
+
+        // Check Expr::as_repeated_element_slice
+        let expr = quote { &[1; 3] }.as_expr().unwrap();
+        let (expr, length) = expr.as_repeated_element_slice().unwrap();
+        assert_eq(expr.as_integer().unwrap(), (1, false));
+        assert_eq(length.as_integer().unwrap(), (3, false));
     }
 }