add "saturating_sub" in zkDSL (constant expression)

Author: TomWambsgans

crates/lean_compiler/src/grammar.pest

@@ -91,12 +91,14 @@ primary = {
     "(" ~ expression ~ ")" | 
     log2_ceil_expr |
     next_multiple_of_expr |
+    saturating_sub_expr |
     len_expr |
     array_access_expr |
     var_or_constant
 }
 log2_ceil_expr = { "log2_ceil" ~ "(" ~ expression ~ ")" }
 next_multiple_of_expr = { "next_multiple_of" ~ "(" ~ expression ~ "," ~ expression ~ ")" }
+saturating_sub_expr = { "saturating_sub" ~ "(" ~ expression ~ "," ~ expression ~ ")" }
 len_expr = { "len" ~ "(" ~ identifier ~ ")" }
 array_access_expr = { identifier ~ "[" ~ expression ~ "]" }
 

crates/lean_compiler/src/lang.rs

@@ -257,13 +257,15 @@ pub enum Expression {
 pub enum MathExpr {
     Log2Ceil,
     NextMultipleOf,
+    SaturatingSub,
 }
 
 impl Display for MathExpr {
     fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
         match self {
             Self::Log2Ceil => write!(f, "log2_ceil"),
             Self::NextMultipleOf => write!(f, "next_multiple_of"),
+            Self::SaturatingSub => write!(f, "saturating_sub"),
         }
     }
 }
@@ -273,6 +275,7 @@ impl MathExpr {
         match self {
             Self::Log2Ceil => 1,
             Self::NextMultipleOf => 2,
+            Self::SaturatingSub => 2,
         }
     }
     pub fn eval(&self, args: &[F]) -> F {
@@ -291,6 +294,10 @@ impl MathExpr {
                 let res = value_usize.next_multiple_of(multiple_usize);
                 F::from_usize(res)
             }
+            Self::SaturatingSub => {
+                assert_eq!(args.len(), 2);
+                F::from_usize(args[0].to_usize().saturating_sub(args[1].to_usize()))
+            }
         }
     }
 }

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

@@ -77,6 +77,7 @@ impl Parse<Expression> for PrimaryExpressionParser {
             Rule::array_access_expr => ArrayAccessParser.parse(inner, ctx),
             Rule::log2_ceil_expr => MathExpr::Log2Ceil.parse(inner, ctx),
             Rule::next_multiple_of_expr => MathExpr::NextMultipleOf.parse(inner, ctx),
+            Rule::saturating_sub_expr => MathExpr::SaturatingSub.parse(inner, ctx),
             Rule::len_expr => LenParser.parse(inner, ctx),
             _ => Err(SemanticError::new("Invalid primary expression").into()),
         }

crates/lean_compiler/tests/test_compiler.rs

@@ -881,3 +881,26 @@ fn intertwined_unrolled_loops_and_const_function_arguments() {
     "#;
     compile_and_run(program.to_string(), (&[], &[]), DEFAULT_NO_VEC_RUNTIME_MEMORY, false);
 }
+
+#[test]
+fn test_const_fibonacci() {
+    let program = r#"
+    fn main() {
+        res = fib(8);
+        assert res == 21;
+        return;
+    }
+    fn fib(const n) -> 1 {
+        if n == 0 {
+            return 0;
+        }
+        if n == 1 {
+            return 1;
+        }
+        a = fib(saturating_sub(n, 1));
+        b = fib(saturating_sub(n, 2));
+        return a + b;
+    }
+    "#;
+    compile_and_run(program.to_string(), (&[], &[]), DEFAULT_NO_VEC_RUNTIME_MEMORY, false);
+}