diff --git a/compiler/noirc_frontend/src/hir/comptime/interpreter/builtin.rs b/compiler/noirc_frontend/src/hir/comptime/interpreter/builtin.rs index 5698cc8d3bd..c3528434a72 100644 --- a/compiler/noirc_frontend/src/hir/comptime/interpreter/builtin.rs +++ b/compiler/noirc_frontend/src/hir/comptime/interpreter/builtin.rs @@ -48,6 +48,9 @@ impl<'local, 'context> Interpreter<'local, 'context> { "array_len" => array_len(interner, arguments, location), "as_slice" => as_slice(interner, arguments, location), "expr_as_function_call" => expr_as_function_call(arguments, return_type, location), + "expr_as_if" => expr_as_if(arguments, return_type, location), + "expr_as_index" => expr_as_index(arguments, return_type, location), + "expr_as_tuple" => expr_as_tuple(arguments, return_type, location), "is_unconstrained" => Ok(Value::Bool(true)), "function_def_name" => function_def_name(interner, arguments, location), "function_def_parameters" => function_def_parameters(interner, arguments, location), @@ -770,6 +773,71 @@ fn expr_as_function_call( }) } +// fn as_if(self) -> Option<(Expr, Expr, Option)> +fn expr_as_if( + arguments: Vec<(Value, Location)>, + return_type: Type, + location: Location, +) -> IResult { + expr_as(arguments, return_type.clone(), location, |expr| { + if let ExpressionKind::If(if_expr) = expr { + // Get the type of `Option` + let option_type = extract_option_generic_type(return_type.clone()); + let Type::Tuple(option_types) = option_type else { + panic!("Expected the return type option generic arg to be a tuple"); + }; + assert_eq!(option_types.len(), 3); + let alternative_option_type = option_types[2].clone(); + + let alternative = + option(alternative_option_type, if_expr.alternative.map(|e| Value::Expr(e.kind))); + + Some(Value::Tuple(vec![ + Value::Expr(if_expr.condition.kind), + Value::Expr(if_expr.consequence.kind), + alternative.ok()?, + ])) + } else { + None + } + }) +} + +// fn as_index(self) -> Option +fn expr_as_index( + arguments: Vec<(Value, Location)>, + return_type: Type, + location: Location, +) -> IResult { + expr_as(arguments, return_type, location, |expr| { + if let ExpressionKind::Index(index_expr) = expr { + Some(Value::Tuple(vec![ + Value::Expr(index_expr.collection.kind), + Value::Expr(index_expr.index.kind), + ])) + } else { + None + } + }) +} + +// fn as_tuple(self) -> Option<[Expr]> +fn expr_as_tuple( + arguments: Vec<(Value, Location)>, + return_type: Type, + location: Location, +) -> IResult { + expr_as(arguments, return_type, location, |expr| { + if let ExpressionKind::Tuple(expressions) = expr { + let expressions = expressions.into_iter().map(|expr| Value::Expr(expr.kind)).collect(); + let typ = Type::Slice(Box::new(Type::Quoted(QuotedType::Expr))); + Some(Value::Slice(expressions, typ)) + } else { + None + } + }) +} + // Helper function for implementing the `expr_as_...` functions. fn expr_as( arguments: Vec<(Value, Location)>, @@ -781,8 +849,12 @@ where F: FnOnce(ExpressionKind) -> Option, { let self_argument = check_one_argument(arguments, location)?; - let expr = get_expr(self_argument)?; - let option_value = f(expr); + let mut expression_kind = get_expr(self_argument)?; + while let ExpressionKind::Parenthesized(expression) = expression_kind { + expression_kind = expression.kind; + } + + let option_value = f(expression_kind); option(return_type, option_value) } diff --git a/noir_stdlib/src/meta/expr.nr b/noir_stdlib/src/meta/expr.nr index 54681632543..2a463ef4d86 100644 --- a/noir_stdlib/src/meta/expr.nr +++ b/noir_stdlib/src/meta/expr.nr @@ -3,4 +3,13 @@ use crate::option::Option; impl Expr { #[builtin(expr_as_function_call)] fn as_function_call(self) -> Option<(Expr, [Expr])> {} + + #[builtin(expr_as_if)] + fn as_if(self) -> Option<(Expr, Expr, Option)> {} + + #[builtin(expr_as_index)] + fn as_index(self) -> Option<(Expr, Expr)> {} + + #[builtin(expr_as_tuple)] + fn as_tuple(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 index 8b6f7b480c7..81531310ff5 100644 --- a/test_programs/compile_success_empty/comptime_exp/src/main.nr +++ b/test_programs/compile_success_empty/comptime_exp/src/main.nr @@ -1,8 +1,31 @@ fn main() { comptime { + // Check Expr::as_function_call let expr = quote { foo(bar) }.as_expr().unwrap(); let (_function, args) = expr.as_function_call().unwrap(); assert_eq(args.len(), 1); + + // Check Expr::as_index + let expr = quote { foo[bar] }.as_expr().unwrap(); + let _ = expr.as_index().unwrap(); + + // Check Expr::as_tuple + let expr = quote { (1, 2) }.as_expr().unwrap(); + let tuple_exprs = expr.as_tuple().unwrap(); + assert_eq(tuple_exprs.len(), 2); + + // Check Expr::as_if + let expr = quote { if 1 { 2 } }.as_expr().unwrap(); + let (_condition, _consequence, alternative) = expr.as_if().unwrap(); + assert(alternative.is_none()); + + let expr = quote { if 1 { 2 } else { 3 } }.as_expr().unwrap(); + let (_condition, _consequence, alternative) = expr.as_if().unwrap(); + assert(alternative.is_some()); + + // Check parenthesized expression is automatically unwrapped + let expr = quote { ((if 1 { 2 })) }.as_expr().unwrap(); + assert(expr.as_if().is_some()); } }