Compression everywhere (Poseidon2)

crates/lean_compiler/snark_lib.py

@@ -62,8 +62,9 @@ def pop(self):
 
 # Built-in constants
 ZERO_VEC_PTR = 0
-ONE_VEC_PTR = 16
-NONRESERVED_PROGRAM_INPUT_START = 58
+SAMPLING_DOMAIN_SEPARATOR_PTR = 16
+ONE_VEC_PTR = 24
+NONRESERVED_PROGRAM_INPUT_START = 66
 
 
 def poseidon16(left, right, output, mode):
@@ -83,6 +84,10 @@ def log2_ceil(x: int) -> int:
     return math.ceil(math.log2(x))
 
 
+def div_ceil(a: int, b: int) -> int:
+    return (a + b - 1) // b
+
+
 def next_multiple_of(x: int, n: int) -> int:
     return x + (n - x % n) % n
 

crates/lean_compiler/src/b_compile_intermediate.rs

@@ -470,7 +470,7 @@ fn compile_lines(
             SimpleLine::Precompile { table, args, .. } => {
                 match table {
                     Table::DotProduct(_) => assert_eq!(args.len(), 5),
-                    Table::Poseidon16(_) => assert_eq!(args.len(), 4),
+                    Table::Poseidon16(_) => assert_eq!(args.len(), 3),
                     Table::Execution(_) => unreachable!(),
                 }
                 // if arg_c is constant, create a variable (in memory) to hold it

crates/lean_compiler/src/grammar.pest

@@ -118,6 +118,7 @@ primary = {
     lambda_expr |
     log2_ceil_expr |
     next_multiple_of_expr |
+    div_ceil_expr |
     saturating_sub_expr |
     len_expr |
     array_access_expr |
@@ -134,6 +135,7 @@ vec_element = { vec_literal | expression }
 function_call_expr = { identifier ~ "(" ~ tuple_expression? ~ ")" }
 log2_ceil_expr = { "log2_ceil" ~ "(" ~ expression ~ ")" }
 next_multiple_of_expr = { "next_multiple_of" ~ "(" ~ expression ~ "," ~ expression ~ ")" }
+div_ceil_expr = { "div_ceil" ~ "(" ~ expression ~ "," ~ expression ~ ")" }
 saturating_sub_expr = { "saturating_sub" ~ "(" ~ expression ~ "," ~ expression ~ ")" }
 len_expr = { "len" ~ "(" ~ len_argument ~ ")" }
 len_argument = { identifier ~ ("[" ~ expression ~ "]")* }

crates/lean_compiler/src/ir/instruction.rs

@@ -99,11 +99,7 @@ impl IntermediateInstruction {
                 arg_b,
                 res: arg_a,
             },
-            MathOperation::Exp
-            | MathOperation::Mod
-            | MathOperation::NextMultipleOf
-            | MathOperation::SaturatingSub
-            | MathOperation::Log2Ceil => {
+            _ => {
                 unreachable!()
             }
         }

crates/lean_compiler/src/lang.rs

@@ -331,6 +331,8 @@ pub enum MathOperation {
     NextMultipleOf,
     /// saturating subtraction
     SaturatingSub,
+    /// Integer division with ceiling
+    DivCeil,
 }
 
 impl TryFrom<MathOperation> for Operation {
@@ -357,6 +359,7 @@ impl Display for MathOperation {
             Self::Log2Ceil => write!(f, "log2_ceil"),
             Self::NextMultipleOf => write!(f, "next_multiple_of"),
             Self::SaturatingSub => write!(f, "saturating_sub"),
+            Self::DivCeil => write!(f, "div_ceil"),
         }
     }
 }
@@ -375,7 +378,8 @@ impl MathOperation {
             | Self::Exp
             | Self::Mod
             | Self::NextMultipleOf
-            | Self::SaturatingSub => 2,
+            | Self::SaturatingSub
+            | Self::DivCeil => 2,
         }
     }
     pub fn eval(&self, args: &[F]) -> F {
@@ -397,6 +401,7 @@ impl MathOperation {
                 F::from_usize(res)
             }
             Self::SaturatingSub => F::from_usize(args[0].to_usize().saturating_sub(args[1].to_usize())),
+            Self::DivCeil => F::from_usize(args[0].to_usize().div_ceil(args[1].to_usize())),
         }
     }
 }

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

@@ -29,6 +29,7 @@ impl Parse<Expression> for ExpressionParser {
             Rule::exp_expr => MathOperation::Exp.parse(pair, ctx),
             Rule::log2_ceil_expr => MathOperation::Log2Ceil.parse(pair, ctx),
             Rule::next_multiple_of_expr => MathOperation::NextMultipleOf.parse(pair, ctx),
+            Rule::div_ceil_expr => MathOperation::DivCeil.parse(pair, ctx),
             Rule::saturating_sub_expr => MathOperation::SaturatingSub.parse(pair, ctx),
             Rule::var_or_constant => Ok(Expression::Value(VarOrConstantParser.parse(pair, ctx)?)),
             Rule::array_access_expr => ArrayAccessParser.parse(pair, ctx),

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

@@ -1,4 +1,6 @@
-use lean_vm::{NONRESERVED_PROGRAM_INPUT_START, ONE_VEC_PTR, PRIVATE_INPUT_START_PTR, ZERO_VEC_PTR};
+use lean_vm::{
+    NONRESERVED_PROGRAM_INPUT_START, ONE_VEC_PTR, PRIVATE_INPUT_START_PTR, SAMPLING_DOMAIN_SEPARATOR_PTR, ZERO_VEC_PTR,
+};
 use multilinear_toolkit::prelude::*;
 
 use super::expression::ExpressionParser;
@@ -134,6 +136,9 @@ impl VarOrConstantParser {
             "PRIVATE_INPUT_START_PTR" => Ok(SimpleExpr::Constant(ConstExpression::from(PRIVATE_INPUT_START_PTR))),
             "ZERO_VEC_PTR" => Ok(SimpleExpr::Constant(ConstExpression::from(ZERO_VEC_PTR))),
             "ONE_VEC_PTR" => Ok(SimpleExpr::Constant(ConstExpression::from(ONE_VEC_PTR))),
+            "SAMPLING_DOMAIN_SEPARATOR_PTR" => Ok(SimpleExpr::Constant(ConstExpression::from(
+                SAMPLING_DOMAIN_SEPARATOR_PTR,
+            ))),
             _ => {
                 // Check if it's a const array (error case - can't use array as value)
                 if ctx.get_const_array(text).is_some() {