Skip to content

Commit cdc1dff

Browse files
TomWambsgansTomWambsgans
committed
new syntax: next_multiple_of
1 parent 5d58ad1 commit cdc1dff

File tree

5 files changed

+106
-1
lines changed

5 files changed

+106
-1
lines changed

‎crates/lean_compiler/src/a_simplify_lang.rs‎

Lines changed: 33 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -406,6 +406,10 @@ fn check_expr_scoping(expr: &Expression, ctx: &Context) {
406406
Expression::Log2Ceil { value } => {
407407
check_expr_scoping(value, ctx);
408408
}
409+
Expression::NextMultipleOf { value, multiple } => {
410+
check_expr_scoping(value, ctx);
411+
check_expr_scoping(multiple, ctx);
412+
}
409413
}
410414
}
411415

@@ -536,7 +540,7 @@ fn simplify_lines(
536540
arg1: right,
537541
});
538542
}
539-
Expression::Log2Ceil { .. } => unreachable!(),
543+
Expression::Log2Ceil { .. } | Expression::NextMultipleOf { .. } => unreachable!(),
540544
},
541545
Line::ArrayAssign { array, index, value } => {
542546
handle_array_assignment(
@@ -1021,6 +1025,18 @@ fn simplify_expr(
10211025
value: Box::new(const_value),
10221026
})
10231027
}
1028+
Expression::NextMultipleOf { value, multiple } => {
1029+
let const_value = simplify_expr(value, lines, counters, array_manager, const_malloc)
1030+
.as_constant()
1031+
.unwrap();
1032+
let const_multiple = simplify_expr(multiple, lines, counters, array_manager, const_malloc)
1033+
.as_constant()
1034+
.unwrap();
1035+
SimpleExpr::Constant(ConstExpression::NextMultipleOf {
1036+
value: Box::new(const_value),
1037+
multiple: Box::new(const_multiple),
1038+
})
1039+
}
10241040
}
10251041
}
10261042

@@ -1184,6 +1200,10 @@ fn inline_expr(expr: &mut Expression, args: &BTreeMap<Var, SimpleExpr>, inlining
11841200
Expression::Log2Ceil { value } => {
11851201
inline_expr(value, args, inlining_count);
11861202
}
1203+
Expression::NextMultipleOf { value, multiple } => {
1204+
inline_expr(value, args, inlining_count);
1205+
inline_expr(multiple, args, inlining_count);
1206+
}
11871207
}
11881208
}
11891209

@@ -1344,6 +1364,10 @@ fn vars_in_expression(expr: &Expression) -> BTreeSet<Var> {
13441364
Expression::Log2Ceil { value } => {
13451365
vars.extend(vars_in_expression(value));
13461366
}
1367+
Expression::NextMultipleOf { value, multiple } => {
1368+
vars.extend(vars_in_expression(value));
1369+
vars.extend(vars_in_expression(multiple));
1370+
}
13471371
}
13481372
vars
13491373
}
@@ -1504,6 +1528,10 @@ fn replace_vars_for_unroll_in_expr(
15041528
Expression::Log2Ceil { value } => {
15051529
replace_vars_for_unroll_in_expr(value, iterator, unroll_index, iterator_value, internal_vars);
15061530
}
1531+
Expression::NextMultipleOf { value, multiple } => {
1532+
replace_vars_for_unroll_in_expr(value, iterator, unroll_index, iterator_value, internal_vars);
1533+
replace_vars_for_unroll_in_expr(multiple, iterator, unroll_index, iterator_value, internal_vars);
1534+
}
15071535
}
15081536
}
15091537

@@ -1988,6 +2016,10 @@ fn replace_vars_by_const_in_expr(expr: &mut Expression, map: &BTreeMap<Var, F>)
19882016
Expression::Log2Ceil { value } => {
19892017
replace_vars_by_const_in_expr(value, map);
19902018
}
2019+
Expression::NextMultipleOf { value, multiple } => {
2020+
replace_vars_by_const_in_expr(value, map);
2021+
replace_vars_by_const_in_expr(multiple, map);
2022+
}
19912023
}
19922024
}
19932025

‎crates/lean_compiler/src/grammar.pest‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -88,10 +88,12 @@ exp_expr = { primary ~ ("**" ~ primary)* }
8888
primary = {
8989
"(" ~ expression ~ ")" |
9090
log2_ceil_expr |
91+
next_multiple_of_expr |
9192
array_access_expr |
9293
var_or_constant
9394
}
9495
log2_ceil_expr = { "log2_ceil" ~ "(" ~ expression ~ ")" }
96+
next_multiple_of_expr = { "next_multiple_of" ~ "(" ~ expression ~ "," ~ expression ~ ")" }
9597
array_access_expr = { identifier ~ "[" ~ expression ~ "]" }
9698

9799
// Basic elements

‎crates/lean_compiler/src/lang.rs‎

Lines changed: 38 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -117,6 +117,10 @@ pub enum ConstExpression {
117117
Log2Ceil {
118118
value: Box<Self>,
119119
},
120+
NextMultipleOf {
121+
value: Box<Self>,
122+
multiple: Box<Self>,
123+
},
120124
}
121125

122126
impl From<usize> for ConstExpression {
@@ -148,6 +152,14 @@ impl TryFrom<Expression> for ConstExpression {
148152
value: Box::new(value_expr),
149153
})
150154
}
155+
Expression::NextMultipleOf { value, multiple } => {
156+
let value_expr = Self::try_from(*value)?;
157+
let multiple_expr = Self::try_from(*multiple)?;
158+
Ok(Self::NextMultipleOf {
159+
value: Box::new(value_expr),
160+
multiple: Box::new(multiple_expr),
161+
})
162+
}
151163
}
152164
}
153165
}
@@ -185,6 +197,14 @@ impl ConstExpression {
185197
let value = value.eval_with(func)?;
186198
Some(F::from_usize(log2_ceil_usize(value.to_usize())))
187199
}
200+
Self::NextMultipleOf { value, multiple } => {
201+
let value = value.eval_with(func)?;
202+
let multiple = multiple.eval_with(func)?;
203+
let value_usize = value.to_usize();
204+
let multiple_usize = multiple.to_usize();
205+
let res = value_usize.next_multiple_of(multiple_usize);
206+
Some(F::from_usize(res))
207+
}
188208
}
189209
}
190210

@@ -249,6 +269,10 @@ pub enum Expression {
249269
Log2Ceil {
250270
value: Box<Expression>,
251271
}, // only for const expressions
272+
NextMultipleOf {
273+
value: Box<Expression>,
274+
multiple: Box<Expression>,
275+
}, // only for const expressions
252276
}
253277

254278
impl From<SimpleExpr> for Expression {
@@ -284,6 +308,14 @@ impl Expression {
284308
let value = value.eval_with(value_fn, array_fn)?;
285309
Some(F::from_usize(log2_ceil_usize(value.to_usize())))
286310
}
311+
Self::NextMultipleOf { value, multiple } => {
312+
let value = value.eval_with(value_fn, array_fn)?;
313+
let multiple = multiple.eval_with(value_fn, array_fn)?;
314+
let value_usize = value.to_usize();
315+
let multiple_usize = multiple.to_usize();
316+
let res = value_usize.next_multiple_of(multiple_usize);
317+
Some(F::from_usize(res))
318+
}
287319
}
288320
}
289321

@@ -417,6 +449,9 @@ impl Display for Expression {
417449
Self::Log2Ceil { value } => {
418450
write!(f, "log2_ceil({value})")
419451
}
452+
Self::NextMultipleOf { value, multiple } => {
453+
write!(f, "next_multiple_of({value}, {multiple})")
454+
}
420455
}
421456
}
422457
}
@@ -630,6 +665,9 @@ impl Display for ConstExpression {
630665
Self::Log2Ceil { value } => {
631666
write!(f, "log2_ceil({value})")
632667
}
668+
Self::NextMultipleOf { value, multiple } => {
669+
write!(f, "next_multiple_of({value}, {multiple})")
670+
}
633671
}
634672
}
635673
}

‎crates/lean_compiler/src/parser/parsers/expression.rs‎

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -75,6 +75,7 @@ impl Parse<Expression> for PrimaryExpressionParser {
7575
}
7676
Rule::array_access_expr => ArrayAccessParser::parse(inner, ctx),
7777
Rule::log2_ceil_expr => Log2CeilParser::parse(inner, ctx),
78+
Rule::next_multiple_of_expr => NextMultipleOfParser::parse(inner, ctx),
7879
_ => Err(SemanticError::new("Invalid primary expression").into()),
7980
}
8081
}
@@ -107,3 +108,19 @@ impl Parse<Expression> for Log2CeilParser {
107108
Ok(Expression::Log2Ceil { value: Box::new(expr) })
108109
}
109110
}
111+
112+
/// Parser for next_multiple_of function calls.
113+
pub struct NextMultipleOfParser;
114+
115+
impl Parse<Expression> for NextMultipleOfParser {
116+
fn parse(pair: ParsePair<'_>, ctx: &mut ParseContext) -> ParseResult<Expression> {
117+
let mut inner = pair.into_inner();
118+
let value = ExpressionParser::parse(next_inner_pair(&mut inner, "next_multiple_of value")?, ctx)?;
119+
let multiple = ExpressionParser::parse(next_inner_pair(&mut inner, "next_multiple_of multiple")?, ctx)?;
120+
121+
Ok(Expression::NextMultipleOf {
122+
value: Box::new(value),
123+
multiple: Box::new(multiple),
124+
})
125+
}
126+
}

‎crates/lean_compiler/tests/test_compiler.rs‎

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -713,3 +713,19 @@ fn test_debug_assert_lt_fail() {
713713
"#;
714714
compile_and_run(program.to_string(), (&[], &[]), DEFAULT_NO_VEC_RUNTIME_MEMORY, false);
715715
}
716+
717+
#[test]
718+
fn test_next_multiple_of() {
719+
let program = r#"
720+
fn main() {
721+
a = double(next_multiple_of(12, 8));
722+
assert a == 32;
723+
return;
724+
}
725+
726+
fn double(const n) -> 1 {
727+
return next_multiple_of(n, n) * 2;
728+
}
729+
"#;
730+
compile_and_run(program.to_string(), (&[], &[]), DEFAULT_NO_VEC_RUNTIME_MEMORY, false);
731+
}

0 commit comments

Comments
 (0)