zkDSL compiler: make it posible to use placeholders in files

Author: TomWambsgans

crates/lean_compiler/src/lib.rs

@@ -1,3 +1,5 @@
+use std::collections::BTreeMap;
+
 use lean_vm::*;
 
 use crate::{
@@ -12,8 +14,41 @@ pub mod ir;
 mod lang;
 mod parser;
 
-pub fn compile_program(filepath: &str, program: String) -> Bytecode {
-    let parsed_program = parse_program(filepath, &program).unwrap();
+#[derive(Debug, Clone)]
+pub enum ProgramSource {
+    Raw(String),
+    Filepath(String),
+}
+
+impl ProgramSource {
+    pub fn get_content(&self, flags: &CompilationFlags) -> Result<String, std::io::Error> {
+        match self {
+            ProgramSource::Raw(src) => {
+                let mut result = src.clone();
+                for (key, value) in flags.replacements.iter() {
+                    result = result.replace(key, value);
+                }
+                Ok(result)
+            }
+            ProgramSource::Filepath(fp) => {
+                let mut result = std::fs::read_to_string(fp)?;
+                for (key, value) in flags.replacements.iter() {
+                    result = result.replace(key, value);
+                }
+                Ok(result)
+            }
+        }
+    }
+}
+
+#[derive(Debug, Clone, Default)]
+pub struct CompilationFlags {
+    /// useful for placeholder replacements in source code
+    pub replacements: BTreeMap<String, String>,
+}
+
+pub fn compile_program_with_flags(input: &ProgramSource, flags: CompilationFlags) -> Bytecode {
+    let parsed_program = parse_program(input, flags).unwrap();
     // println!("Parsed program: {}", parsed_program.to_string());
     let function_locations = parsed_program.function_locations.clone();
     let source_code = parsed_program.source_code.clone();
@@ -33,14 +68,17 @@ pub fn compile_program(filepath: &str, program: String) -> Bytecode {
     // compiled
 }
 
+pub fn compile_program(input: &ProgramSource) -> Bytecode {
+    compile_program_with_flags(input, CompilationFlags::default())
+}
+
 pub fn compile_and_run(
-    filepath: &str,
-    program: String,
+    source: &ProgramSource,
     (public_input, private_input): (&[F], &[F]),
     no_vec_runtime_memory: usize, // size of the "non-vectorized" runtime memory
     profiler: bool,
 ) {
-    let bytecode = compile_program(filepath, program);
+    let bytecode = compile_program(source);
     let summary = execute_bytecode(
         &bytecode,
         (public_input, private_input),

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

@@ -3,6 +3,7 @@ use crate::parser::{
     error::{ParseResult, SemanticError},
     grammar::{ParsePair, Rule},
 };
+use crate::{CompilationFlags, ProgramSource};
 use std::collections::{BTreeMap, BTreeSet};
 
 pub mod expression;
@@ -37,20 +38,28 @@ pub struct ParseContext {
     pub imported_filepaths: BTreeSet<String>,
     /// Next unused file ID
     pub next_file_id: usize,
+    /// Compilation flags
+    pub flags: CompilationFlags,
 }
 
 impl ParseContext {
-    pub fn new(current_filepath: &str, current_source_code: &str) -> Self {
-        Self {
+    pub fn new(input: &ProgramSource, flags: CompilationFlags) -> Result<Self, SemanticError> {
+        let current_source_code = input.get_content(&flags).unwrap();
+        let (current_filepath, imported_filepaths) = match input {
+            ProgramSource::Raw(_) => ("<raw_input>".to_string(), BTreeSet::new()),
+            ProgramSource::Filepath(fp) => (fp.clone(), [fp.clone()].into_iter().collect()),
+        };
+        Ok(Self {
             constants: BTreeMap::new(),
             const_arrays: BTreeMap::new(),
             trash_var_count: 0,
-            current_filepath: current_filepath.to_string(),
+            current_filepath,
             current_file_id: 0,
-            imported_filepaths: BTreeSet::new(),
-            current_source_code: current_source_code.to_string(),
+            imported_filepaths,
+            current_source_code,
             next_file_id: 1,
-        }
+            flags,
+        })
     }
 
     /// Adds a scalar constant to the context.
@@ -103,12 +112,6 @@ impl ParseContext {
     }
 }
 
-impl Default for ParseContext {
-    fn default() -> Self {
-        Self::new("<string>", "")
-    }
-}
-
 /// Core trait for all parsers in the system.
 pub trait Parse<T>: Sized {
     /// Parses the given input into the target type.

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

@@ -1,6 +1,7 @@
 use super::function::FunctionParser;
 use super::literal::ConstantDeclarationParser;
 use crate::{
+    CompilationFlags, ProgramSource,
     lang::{Program, SourceLocation},
     parser::{
         error::{ParseError, ParseResult, SemanticError},
@@ -55,14 +56,9 @@ impl Parse<Program> for ProgramParser {
                         let file_id = ctx.get_next_file_id();
                         ctx.current_file_id = file_id;
                         filepaths.insert(file_id, filepath.clone());
-                        let input = std::fs::read_to_string(filepath.clone()).map_err(|_| {
-                            SemanticError::with_context(
-                                format!("Imported file not found: {filepath}"),
-                                "import declaration",
-                            )
-                        })?;
-                        source_code.insert(file_id, input.clone());
-                        let subprogram = parse_program_helper(filepath.as_str(), input.as_str(), ctx)?;
+                        ctx.current_source_code = ProgramSource::Filepath(filepath).get_content(&ctx.flags).unwrap();
+                        source_code.insert(file_id, ctx.current_source_code.clone());
+                        let subprogram = parse_program_helper(ctx)?;
                         functions.extend(subprogram.functions);
                         function_locations.extend(subprogram.function_locations);
                         source_code.extend(subprogram.source_code);
@@ -138,22 +134,17 @@ impl Parse<String> for ImportStatementParser {
     }
 }
 
-fn parse_program_helper(filepath: &str, input: &str, ctx: &mut ParseContext) -> Result<Program, ParseError> {
+fn parse_program_helper(ctx: &mut ParseContext) -> Result<Program, ParseError> {
     // Preprocess source to remove comments
-    let processed_input = lexer::preprocess_source(input);
+    let processed_input = lexer::preprocess_source(&ctx.current_source_code);
 
     // Parse grammar into AST nodes
     let program_pair = parse_source(&processed_input)?;
 
     // Parse into semantic structures
-    ctx.current_filepath = filepath.to_string();
-    ctx.current_source_code = input.to_string();
-    ctx.imported_filepaths.insert(filepath.to_string());
     ProgramParser.parse(program_pair, ctx)
 }
 
-pub fn parse_program(filepath: &str, input: &str) -> Result<Program, ParseError> {
-    let mut ctx = ParseContext::new(filepath, input);
-    ctx.imported_filepaths.insert(filepath.to_string());
-    parse_program_helper(filepath, input, &mut ctx)
+pub fn parse_program(input: &ProgramSource, flags: CompilationFlags) -> Result<Program, ParseError> {
+    parse_program_helper(&mut ParseContext::new(input, flags)?)
 }

crates/lean_compiler/tests/program_0.snark

@@ -0,0 +1,5 @@
+import "asdfasdfadsfasdf.snark";
+
+fn main() {
+    return;
+}

crates/lean_compiler/tests/program_1.snark

@@ -0,0 +1,10 @@
+ import "bar.snark";
+import "foo.snark";
+
+fn bar() {
+    return;
+}
+
+fn main() {
+    return;
+}

crates/lean_compiler/tests/program_2.snark

@@ -0,0 +1,8 @@
+ import "bar.snark";
+import "foo.snark";
+
+const FOO = 5;
+
+fn main() {
+    return;
+}

crates/lean_compiler/tests/program_3.snark

@@ -0,0 +1,6 @@
+import "foo.snark";
+import "foo.snark";
+
+fn main() {
+    return;
+}

crates/lean_compiler/tests/program_4.snark

@@ -0,0 +1,5 @@
+import "circular_import.snark";
+
+fn main() {
+    return;
+}

crates/lean_compiler/tests/program_5.snark

@@ -0,0 +1,8 @@
+import "bar.snark";
+import "foo.snark";
+
+fn main() {
+    x = bar(FOO);
+    assert x == 6;
+    return;
+}