Skip to content

Commit

Permalink
feat: add Expr methods: as_tuple, as_parenthesized, as_index,…
Browse files Browse the repository at this point in the history
… `as_if` (#5726)

# Description

## Problem

Part of #5668

## Summary

Just four, to keep PRs small.

## Additional Context


## Documentation

Check one:
- [ ] No documentation needed.
- [ ] Documentation included in this PR.
- [x] **[For Experimental Features]** Documentation to be submitted in a
separate PR.

# PR Checklist

- [x] I have tested the changes locally.
- [x] I have formatted the changes with [Prettier](https://prettier.io/)
and/or `cargo fmt` on default settings.
  • Loading branch information
asterite authored Aug 15, 2024
1 parent 5a3d241 commit f57a7b2
Show file tree
Hide file tree
Showing 3 changed files with 106 additions and 2 deletions.
76 changes: 74 additions & 2 deletions compiler/noirc_frontend/src/hir/comptime/interpreter/builtin.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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),
Expand Down Expand Up @@ -770,6 +773,71 @@ fn expr_as_function_call(
})
}

// fn as_if(self) -> Option<(Expr, Expr, Option<Expr>)>
fn expr_as_if(
arguments: Vec<(Value, Location)>,
return_type: Type,
location: Location,
) -> IResult<Value> {
expr_as(arguments, return_type.clone(), location, |expr| {
if let ExpressionKind::If(if_expr) = expr {
// Get the type of `Option<Expr>`
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<Expr>
fn expr_as_index(
arguments: Vec<(Value, Location)>,
return_type: Type,
location: Location,
) -> IResult<Value> {
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<Value> {
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<F>(
arguments: Vec<(Value, Location)>,
Expand All @@ -781,8 +849,12 @@ where
F: FnOnce(ExpressionKind) -> Option<Value>,
{
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)
}

Expand Down
9 changes: 9 additions & 0 deletions noir_stdlib/src/meta/expr.nr
Original file line number Diff line number Diff line change
Expand Up @@ -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<Expr>)> {}

#[builtin(expr_as_index)]
fn as_index(self) -> Option<(Expr, Expr)> {}

#[builtin(expr_as_tuple)]
fn as_tuple(self) -> Option<[Expr]> {}
}
23 changes: 23 additions & 0 deletions test_programs/compile_success_empty/comptime_exp/src/main.nr
Original file line number Diff line number Diff line change
@@ -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());
}
}

0 comments on commit f57a7b2

Please sign in to comment.