Skip to content

Commit

Permalink
Let statements
Browse files Browse the repository at this point in the history
  • Loading branch information
chriseth committed Sep 30, 2024
1 parent a8b1406 commit b24ee27
Show file tree
Hide file tree
Showing 2 changed files with 57 additions and 8 deletions.
34 changes: 26 additions & 8 deletions jit-compiler/src/codegen.rs
Original file line number Diff line number Diff line change
Expand Up @@ -5,10 +5,10 @@ use powdr_ast::{
analyzed::{Analyzed, Expression, FunctionValueDefinition, PolynomialReference, Reference},
parsed::{
display::quote,
types::{ArrayType, FunctionType, Type, TypeScheme},
types::{ArrayType, FunctionType, TupleType, Type, TypeScheme},
ArrayLiteral, BinaryOperation, BinaryOperator, BlockExpression, FunctionCall, IfExpression,
IndexAccess, LambdaExpression, MatchArm, MatchExpression, Number, Pattern,
StatementInsideBlock, UnaryOperation,
IndexAccess, LambdaExpression, LetStatementInsideBlock, MatchArm, MatchExpression, Number,
Pattern, StatementInsideBlock, UnaryOperation,
},
};
use powdr_number::{BigInt, BigUint, FieldElement, LargeInt};
Expand Down Expand Up @@ -261,7 +261,7 @@ impl<'a, T: FieldElement> CodeGenerator<'a, T> {
.join(", ")
)
}
Expression::String(_, s) => quote(s),
Expression::String(_, s) => format!("{}.to_string()", quote(s)),
Expression::Tuple(_, items) => format!(
"({})",
items
Expand Down Expand Up @@ -309,9 +309,27 @@ impl<'a, T: FieldElement> CodeGenerator<'a, T> {
}

fn format_statement(&mut self, s: &StatementInsideBlock<Expression>) -> Result<String, String> {
Err(format!(
"Compiling statements inside blocks is not yet implemented: {s}"
))
Ok(match s {
StatementInsideBlock::LetStatement(LetStatementInsideBlock { pattern, ty, value }) => {
let Some(value) = value else {
return Err(format!(
"Column creating 'let'-statements not yet supported: {s}"
));
};
let value = self.format_expr(value)?;
let var_name = "scrutinee__";
let ty = ty
.as_ref()
.map(|ty| format!(": {}", map_type(ty)))
.unwrap_or_default();

let (vars, code) = check_pattern(var_name, pattern)?;
// TODO if we want to explicitly specify the type, we need to exchange the non-captured
// parts by `()`.
format!("let {vars} = (|{var_name}{ty}| {code})({value}).unwrap();\n",)
}
StatementInsideBlock::Expression(e) => format!("{};\n", self.format_expr(e)?),
})
}

/// Returns a string expression evaluating to the value of the symbol.
Expand Down Expand Up @@ -448,7 +466,7 @@ fn map_type(ty: &Type) -> String {
Type::String => "String".to_string(),
Type::Expr => "Expr".to_string(),
Type::Array(ArrayType { base, length: _ }) => format!("Vec<{}>", map_type(base)),
Type::Tuple(_) => todo!(),
Type::Tuple(TupleType { items }) => format!("({})", items.iter().map(map_type).join(", ")),
Type::Function(ft) => format!(
"fn({}) -> {}",
ft.params.iter().map(map_type).join(", "),
Expand Down
31 changes: 31 additions & 0 deletions jit-compiler/tests/execution.rs
Original file line number Diff line number Diff line change
Expand Up @@ -158,3 +158,34 @@ fn match_array() {
assert_eq!(f.call(5), 6);
assert_eq!(f.call(6), 7);
}

#[test]
fn let_simple() {
let f = compile(
r#"let f: int -> int = |x| {
let a = 1;
let b = a + 9;
b - 9 + x
};"#,
"f",
);

assert_eq!(f.call(0), 1);
assert_eq!(f.call(1), 2);
assert_eq!(f.call(2), 3);
assert_eq!(f.call(3), 4);
}

#[test]
fn let_complex() {
let f = compile(
r#"let f: int -> int = |x| {
let (a, b, (_, d)) = (1, 2, ("abc", [x, 5]));
a + b + d[0] + d[1]
};"#,
"f",
);

assert_eq!(f.call(0), 8);
assert_eq!(f.call(1), 9);
}

0 comments on commit b24ee27

Please sign in to comment.