zkDSL new feature: constant arrays (compile time)

Author: TomWambsgans

crates/lean_compiler/src/a_simplify_lang.rs

@@ -4,7 +4,8 @@ WHITESPACE = _{ " " | "\t" | "\n" | "\r" }
 program = { SOI ~ constant_declaration* ~ function+ ~ EOI }
 
 // Constants
-constant_declaration = { "const" ~ identifier ~ "=" ~ expression ~ ";" }
+constant_declaration = { "const" ~ identifier ~ "=" ~ (array_literal | expression) ~ ";" }
+array_literal = { "[" ~ (expression ~ ("," ~ expression)*)? ~ "]" }
 
 // Functions
 function = { pragma? ~ "fn" ~ identifier ~ "(" ~ parameter_list? ~ ")" ~ inlined_statement? ~ return_count? ~ "{" ~ statement* ~ "}" }
@@ -89,17 +90,19 @@ primary = {
     "(" ~ expression ~ ")" | 
     log2_ceil_expr |
     next_multiple_of_expr |
+    len_expr |
     array_access_expr |
     var_or_constant
 }
 log2_ceil_expr = { "log2_ceil" ~ "(" ~ expression ~ ")" }
 next_multiple_of_expr = { "next_multiple_of" ~ "(" ~ expression ~ "," ~ expression ~ ")" }
+len_expr = { "len" ~ "(" ~ identifier ~ ")" }
 array_access_expr = { identifier ~ "[" ~ expression ~ "]" }
 
 // Basic elements
 var_or_constant = { constant_value | identifier }
-constant_value = { number | "public_input_start" }
+constant_value = { number | "public_input_start" | "pointer_to_zero_vector" | "pointer_to_one_vector" }
 
 // Lexical elements
 identifier = @{ (ASCII_ALPHA | "_") ~ (ASCII_ALPHANUMERIC | "_")* }
-number = @{ ASCII_DIGIT+ }
+number = @{ ASCII_DIGIT+ }

crates/lean_compiler/src/grammar.pest

@@ -10,6 +10,7 @@ use crate::{F, ir::HighLevelOperation};
 #[derive(Debug, Clone)]
 pub struct Program {
     pub functions: BTreeMap<String, Function>,
+    pub const_arrays: BTreeMap<String, Vec<usize>>,
 }
 
 #[derive(Debug, Clone)]
@@ -288,8 +289,18 @@ impl From<Var> for Expression {
 }
 
 impl Expression {
-    pub fn naive_eval(&self) -> Option<F> {
-        self.eval_with(&|value: &SimpleExpr| value.as_constant()?.naive_eval(), &|_, _| None)
+    pub fn naive_eval(&self, const_arrays: &BTreeMap<String, Vec<usize>>) -> Option<F> {
+        self.eval_with(
+            &|value: &SimpleExpr| value.as_constant()?.naive_eval(),
+            &|arr, index| {
+                let SimpleExpr::Var(name) = arr else {
+                    return None;
+                };
+                let index_usize = index.to_usize();
+                let array = const_arrays.get(name)?;
+                Some(F::from_usize(*array.get(index_usize)?))
+            },
+        )
     }
 
     pub fn eval_with<ValueFn, ArrayFn>(&self, value_fn: &ValueFn, array_fn: &ArrayFn) -> Option<F>
@@ -417,10 +428,15 @@ pub enum Line {
 pub struct Context {
     /// A list of lexical scopes, innermost scope last.
     pub scopes: Vec<Scope>,
+    /// A mapping from constant array names to their values.
+    pub const_arrays: BTreeMap<String, Vec<usize>>,
 }
 
 impl Context {
     pub fn defines(&self, var: &Var) -> bool {
+        if self.const_arrays.contains_key(var) {
+            return true;
+        }
         for scope in self.scopes.iter() {
             if scope.vars.contains(var) {
                 return true;
@@ -680,7 +696,19 @@ impl Display for Line {
 
 impl Display for Program {
     fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
-        let mut first = true;
+        // Print const arrays
+        for (name, values) in &self.const_arrays {
+            write!(f, "const {name} = [")?;
+            for (i, v) in values.iter().enumerate() {
+                if i > 0 {
+                    write!(f, ", ")?;
+                }
+                write!(f, "{v}")?;
+            }
+            writeln!(f, "];")?;
+        }
+
+        let mut first = self.const_arrays.is_empty();
         for function in self.functions.values() {
             if !first {
                 writeln!(f)?;

crates/lean_compiler/src/lang.rs

@@ -2,7 +2,7 @@ use super::literal::VarOrConstantParser;
 use super::{Parse, ParseContext, next_inner_pair};
 use crate::{
     ir::HighLevelOperation,
-    lang::Expression,
+    lang::{ConstExpression, ConstantValue, Expression, SimpleExpr},
     parser::{
         error::{ParseError, ParseResult, SemanticError},
         grammar::{ParsePair, Rule},
@@ -76,6 +76,7 @@ impl Parse<Expression> for PrimaryExpressionParser {
             Rule::array_access_expr => ArrayAccessParser::parse(inner, ctx),
             Rule::log2_ceil_expr => Log2CeilParser::parse(inner, ctx),
             Rule::next_multiple_of_expr => NextMultipleOfParser::parse(inner, ctx),
+            Rule::len_expr => LenParser::parse(inner, ctx),
             _ => Err(SemanticError::new("Invalid primary expression").into()),
         }
     }
@@ -124,3 +125,25 @@ impl Parse<Expression> for NextMultipleOfParser {
         })
     }
 }
+
+/// Parser for len() expressions on const arrays.
+pub struct LenParser;
+
+impl Parse<Expression> for LenParser {
+    fn parse(pair: ParsePair<'_>, ctx: &mut ParseContext) -> ParseResult<Expression> {
+        let mut inner = pair.into_inner();
+        let ident = next_inner_pair(&mut inner, "len argument")?.as_str();
+
+        if let Some(arr) = ctx.get_const_array(ident) {
+            Ok(Expression::Value(SimpleExpr::Constant(ConstExpression::Value(
+                ConstantValue::Scalar(arr.len()),
+            ))))
+        } else {
+            Err(SemanticError::with_context(
+                format!("len() argument '{ident}' is not a const array"),
+                "len expression",
+            )
+            .into())
+        }
+    }
+}

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

@@ -1,5 +1,5 @@
 use super::expression::ExpressionParser;
-use super::{Parse, ParseContext, next_inner_pair};
+use super::{Parse, ParseContext, ParsedConstant, next_inner_pair};
 use crate::{
     F,
     lang::{ConstExpression, ConstantValue, SimpleExpr},
@@ -14,30 +14,61 @@ use utils::ToUsize;
 /// Parser for constant declarations.
 pub struct ConstantDeclarationParser;
 
-impl Parse<(String, usize)> for ConstantDeclarationParser {
-    fn parse(pair: ParsePair<'_>, ctx: &mut ParseContext) -> ParseResult<(String, usize)> {
+impl Parse<(String, ParsedConstant)> for ConstantDeclarationParser {
+    fn parse(pair: ParsePair<'_>, ctx: &mut ParseContext) -> ParseResult<(String, ParsedConstant)> {
         let mut inner = pair.into_inner();
         let name = next_inner_pair(&mut inner, "constant name")?.as_str().to_string();
         let value_pair = next_inner_pair(&mut inner, "constant value")?;
 
-        // Parse the expression and evaluate it
-        let expr = ExpressionParser::parse(value_pair, ctx)?;
+        match value_pair.as_rule() {
+            Rule::array_literal => {
+                let values: Vec<usize> = value_pair
+                    .into_inner()
+                    .map(|expr_pair| {
+                        let expr = ExpressionParser::parse(expr_pair, ctx).unwrap();
+                        expr.eval_with(
+                            &|simple_expr| match simple_expr {
+                                SimpleExpr::Constant(cst) => cst.naive_eval(),
+                                SimpleExpr::Var(var) => ctx.get_constant(var).map(F::from_usize),
+                                SimpleExpr::ConstMallocAccess { .. } => None,
+                            },
+                            &|_, _| None,
+                        )
+                        .ok_or_else(|| {
+                            SemanticError::with_context(
+                                format!("Failed to evaluate array element in constant: {name}"),
+                                "constant declaration",
+                            )
+                        })
+                        .map(|f| f.to_usize())
+                    })
+                    .collect::<Result<Vec<_>, _>>()?;
+                Ok((name, ParsedConstant::Array(values)))
+            }
+            _ => {
+                // Parse the expression and evaluate it
+                let expr = ExpressionParser::parse(value_pair, ctx)?;
 
-        let value = expr
-            .eval_with(
-                &|simple_expr| match simple_expr {
-                    SimpleExpr::Constant(cst) => cst.naive_eval(),
-                    SimpleExpr::Var(var) => ctx.get_constant(var).map(F::from_usize),
-                    SimpleExpr::ConstMallocAccess { .. } => None, // Not allowed in constants
-                },
-                &|_, _| None,
-            )
-            .ok_or_else(|| {
-                SemanticError::with_context(format!("Failed to evaluate constant: {name}"), "constant declaration")
-            })?
-            .to_usize();
+                let value = expr
+                    .eval_with(
+                        &|simple_expr| match simple_expr {
+                            SimpleExpr::Constant(cst) => cst.naive_eval(),
+                            SimpleExpr::Var(var) => ctx.get_constant(var).map(F::from_usize),
+                            SimpleExpr::ConstMallocAccess { .. } => None,
+                        },
+                        &|_, _| None,
+                    )
+                    .ok_or_else(|| {
+                        SemanticError::with_context(
+                            format!("Failed to evaluate constant: {name}"),
+                            "constant declaration",
+                        )
+                    })?
+                    .to_usize();
 
-        Ok((name, value))
+                Ok((name, ParsedConstant::Scalar(value)))
+            }
+        }
     }
 }
 
@@ -73,6 +104,15 @@ impl VarOrConstantParser {
                 ConstantValue::PointerToOneVector,
             ))),
             _ => {
+                // Check if it's a const array (error case - can't use array as value)
+                if ctx.get_const_array(text).is_some() {
+                    return Err(SemanticError::with_context(
+                        format!("Cannot use const array '{text}' as a value directly (use indexing or len())"),
+                        "variable reference",
+                    )
+                    .into());
+                }
+
                 // Try to resolve as defined constant
                 if let Some(value) = ctx.get_constant(text) {
                     Ok(SimpleExpr::Constant(ConstExpression::Value(ConstantValue::Scalar(

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

@@ -10,11 +10,20 @@ pub mod literal;
 pub mod program;
 pub mod statement;
 
+/// Represents a parsed constant value (scalar or array).
+#[derive(Debug, Clone)]
+pub enum ParsedConstant {
+    Scalar(usize),
+    Array(Vec<usize>),
+}
+
 /// Core parsing context that all parsers share.
 #[derive(Debug)]
 pub struct ParseContext {
-    /// Compile-time constants defined in the program
+    /// Compile-time scalar constants defined in the program
     pub constants: BTreeMap<String, usize>,
+    /// Compile-time array constants defined in the program
+    pub const_arrays: BTreeMap<String, Vec<usize>>,
     /// Counter for generating unique trash variable names
     pub trash_var_count: usize,
 }
@@ -23,27 +32,47 @@ impl ParseContext {
     pub const fn new() -> Self {
         Self {
             constants: BTreeMap::new(),
+            const_arrays: BTreeMap::new(),
             trash_var_count: 0,
         }
     }
 
-    /// Adds a constant to the context.
+    /// Adds a scalar constant to the context.
     pub fn add_constant(&mut self, name: String, value: usize) -> Result<(), SemanticError> {
-        if self.constants.insert(name.clone(), value).is_some() {
+        if self.constants.contains_key(&name) || self.const_arrays.contains_key(&name) {
+            Err(SemanticError::with_context(
+                format!("Defined multiple times: {name}"),
+                "constant declaration",
+            ))
+        } else {
+            self.constants.insert(name, value);
+            Ok(())
+        }
+    }
+
+    /// Adds an array constant to the context.
+    pub fn add_const_array(&mut self, name: String, values: Vec<usize>) -> Result<(), SemanticError> {
+        if self.constants.contains_key(&name) || self.const_arrays.contains_key(&name) {
             Err(SemanticError::with_context(
-                format!("Multiply defined constant: {name}"),
+                format!("Defined multiple times: {name}"),
                 "constant declaration",
             ))
         } else {
+            self.const_arrays.insert(name, values);
             Ok(())
         }
     }
 
-    /// Looks up a constant value.
+    /// Looks up a scalar constant value.
     pub fn get_constant(&self, name: &str) -> Option<usize> {
         self.constants.get(name).copied()
     }
 
+    /// Looks up an array constant.
+    pub fn get_const_array(&self, name: &str) -> Option<&Vec<usize>> {
+        self.const_arrays.get(name)
+    }
+
     /// Generates a unique trash variable name.
     pub fn next_trash_var(&mut self) -> String {
         self.trash_var_count += 1;

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

@@ -1,6 +1,6 @@
 use super::function::FunctionParser;
 use super::literal::ConstantDeclarationParser;
-use super::{Parse, ParseContext};
+use super::{Parse, ParseContext, ParsedConstant};
 use crate::{
     lang::Program,
     parser::{
@@ -23,7 +23,10 @@ impl Parse<(Program, BTreeMap<usize, String>)> for ProgramParser {
             match item.as_rule() {
                 Rule::constant_declaration => {
                     let (name, value) = ConstantDeclarationParser::parse(item, &mut ctx)?;
-                    ctx.add_constant(name, value)?;
+                    match value {
+                        ParsedConstant::Scalar(v) => ctx.add_constant(name, v)?,
+                        ParsedConstant::Array(arr) => ctx.add_const_array(name, arr)?,
+                    }
                 }
                 Rule::function => {
                     let location = item.line_col().0;
@@ -45,6 +48,12 @@ impl Parse<(Program, BTreeMap<usize, String>)> for ProgramParser {
             }
         }
 
-        Ok((Program { functions }, function_locations))
+        Ok((
+            Program {
+                functions,
+                const_arrays: ctx.const_arrays,
+            },
+            function_locations,
+        ))
     }
 }

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

@@ -729,3 +729,40 @@ fn test_next_multiple_of() {
     "#;
     compile_and_run(program.to_string(), (&[], &[]), DEFAULT_NO_VEC_RUNTIME_MEMORY, false);
 }
+
+#[test]
+fn test_const_array() {
+    let program = r#"
+    const FIVE = 5;
+    const ARR = [4, FIVE, 4 + 2, 3 * 2 + 1];
+    fn main() {
+        for i in 1..len(ARR) unroll {
+            x = i + 4;
+            assert ARR[i] == x;
+        }
+        four = 4;
+        assert len(ARR) == four;
+        res = func(2);
+        six = 6;
+        assert res == six;
+        nothing(ARR[0]);
+        mem_arr = malloc(len(ARR));
+        for i in 0..len(ARR) unroll {
+            mem_arr[i] = ARR[i];
+        }
+        for i in 0..ARR[0] {
+            print(2**ARR[0]);
+        }
+        print(2**ARR[1]);
+        return;
+    }
+
+    fn func(const x) -> 1 {
+        return ARR[x];
+    }
+    fn nothing(x) {
+        return;
+    }
+    "#;
+    compile_and_run(program.to_string(), (&[], &[]), DEFAULT_NO_VEC_RUNTIME_MEMORY, false);
+}