MassaLabs: Implement a Intermediate Representation to improve the compilation process

.gitignore

@@ -1,6 +1,6 @@
 # Generated by Cargo
 # will have compiled files and executables
-/target/
+**/target/
 
 # Remove system config files generated by mac os
 **/.DS_Store
@@ -17,4 +17,4 @@ Cargo.lock
 
 # Ignore examples codegen generated files
 **/examples/*.rs
-**/examples/*.masm
+**/examples/*.masm

Cargo.toml

@@ -3,7 +3,8 @@ members = [
     "air-script",
     "parser",
     "pass",
-    "ir",
+    "mir",
+    "air",
     "codegen/masm",
     "codegen/winterfell",
 ]

air-script/Cargo.toml

@@ -17,7 +17,8 @@ name = "airc"
 path = "src/main.rs"
 
 [dependencies]
-air-ir = { package = "air-ir", path = "../ir", version = "0.4" }
+mir = { package = "mir", path = "../mir", version = "0.4" }
+air-ir = { package = "air-ir", path = "../air", version = "0.4" }
 air-parser = { package = "air-parser", path = "../parser", version = "0.4" }
 air-pass = { package = "air-pass", path = "../pass", version = "0.1" }
 air-codegen-masm = { package = "air-codegen-masm", path = "../codegen/masm", version = "0.4" }

air-script/src/cli/transpile.rs

@@ -21,6 +21,11 @@ impl Target {
         }
     }
 }
+#[derive(Copy, Clone, PartialEq, Eq, PartialOrd, Ord, ValueEnum)]
+pub enum Pipeline {
+    WithMIR,
+    WithoutMIR,
+}
 
 #[derive(Args)]
 pub struct Transpile {
@@ -40,28 +45,55 @@ pub struct Transpile {
         help = "Defines the target language, defaults to Winterfell"
     )]
     target: Option<Target>,
+
+    #[arg(
+        short,
+        long,
+        help = "Defines the compilation pipeline (WithMIR or WithoutMIR), defaults to WithMIR"
+    )]
+    pipeline: Option<Pipeline>,
 }
 
 impl Transpile {
     pub fn execute(&self) -> Result<(), String> {
         println!("============================================================");
-        println!("Transpiling...");
 
         let input_path = &self.input;
 
         let codemap = Arc::new(CodeMap::new());
         let emitter = Arc::new(DefaultEmitter::new(ColorChoice::Auto));
         let diagnostics = DiagnosticsHandler::new(Default::default(), codemap.clone(), emitter);
 
+        let pipeline = self.pipeline.unwrap_or(Pipeline::WithMIR);
         // Parse from file to internal representation
-        let air = air_parser::parse_file(&diagnostics, codemap, input_path)
-            .map_err(CompileError::Parse)
-            .and_then(|ast| {
-                let mut pipeline = air_parser::transforms::ConstantPropagation::new(&diagnostics)
-                    .chain(air_parser::transforms::Inlining::new(&diagnostics))
-                    .chain(air_ir::passes::AstToAir::new(&diagnostics));
-                pipeline.run(ast)
-            });
+        let air = match pipeline {
+            Pipeline::WithMIR => {
+                println!("Transpiling with Mir pipeline...");
+                air_parser::parse_file(&diagnostics, codemap, input_path)
+                    .map_err(CompileError::Parse)
+                    .and_then(|ast| {
+                        let mut pipeline =
+                            air_parser::transforms::ConstantPropagation::new(&diagnostics)
+                                .chain(mir::passes::AstToMir::new(&diagnostics))
+                                .chain(mir::passes::Inlining::new(&diagnostics))
+                                .chain(mir::passes::Unrolling::new(&diagnostics))
+                                .chain(air_ir::passes::MirToAir::new(&diagnostics));
+                        pipeline.run(ast)
+                    })
+            }
+            Pipeline::WithoutMIR => {
+                println!("Transpiling without Mir pipeline...");
+                air_parser::parse_file(&diagnostics, codemap, input_path)
+                    .map_err(CompileError::Parse)
+                    .and_then(|ast| {
+                        let mut pipeline =
+                            air_parser::transforms::ConstantPropagation::new(&diagnostics)
+                                .chain(air_parser::transforms::Inlining::new(&diagnostics))
+                                .chain(air_ir::passes::AstToAir::new(&diagnostics));
+                        pipeline.run(ast)
+                    })
+            }
+        };
 
         match air {
             Ok(air) => {

air-script/tests/codegen/helpers.rs

@@ -10,6 +10,10 @@ pub enum Target {
     Winterfell,
     Masm,
 }
+pub enum Pipeline {
+    WithMIR,
+    WithoutMIR,
+}
 
 pub struct Test {
     input_path: String,
@@ -19,20 +23,34 @@ impl Test {
         Test { input_path }
     }
 
-    pub fn transpile(&self, target: Target) -> Result<String, CompileError> {
+    pub fn transpile(&self, target: Target, pipeline: Pipeline) -> Result<String, CompileError> {
         let codemap = Arc::new(CodeMap::new());
         let emitter = Arc::new(DefaultEmitter::new(ColorChoice::Auto));
         let diagnostics = DiagnosticsHandler::new(Default::default(), codemap.clone(), emitter);
 
         // Parse from file to internal representation
-        let air = air_parser::parse_file(&diagnostics, codemap, &self.input_path)
-            .map_err(CompileError::Parse)
-            .and_then(|ast| {
-                let mut pipeline = air_parser::transforms::ConstantPropagation::new(&diagnostics)
-                    .chain(air_parser::transforms::Inlining::new(&diagnostics))
-                    .chain(air_ir::passes::AstToAir::new(&diagnostics));
-                pipeline.run(ast)
-            })?;
+        let air = match pipeline {
+            Pipeline::WithMIR => air_parser::parse_file(&diagnostics, codemap, &self.input_path)
+                .map_err(CompileError::Parse)
+                .and_then(|ast| {
+                    let mut pipeline =
+                        air_parser::transforms::ConstantPropagation::new(&diagnostics)
+                            .chain(mir::passes::AstToMir::new(&diagnostics))
+                            .chain(mir::passes::Inlining::new(&diagnostics))
+                            .chain(mir::passes::Unrolling::new(&diagnostics))
+                            .chain(air_ir::passes::MirToAir::new(&diagnostics));
+                    pipeline.run(ast)
+                })?,
+            Pipeline::WithoutMIR => air_parser::parse_file(&diagnostics, codemap, &self.input_path)
+                .map_err(CompileError::Parse)
+                .and_then(|ast| {
+                    let mut pipeline =
+                        air_parser::transforms::ConstantPropagation::new(&diagnostics)
+                            .chain(air_parser::transforms::Inlining::new(&diagnostics))
+                            .chain(air_ir::passes::AstToAir::new(&diagnostics));
+                    pipeline.run(ast)
+                })?,
+        };
 
         let backend: Box<dyn CodeGenerator<Output = String>> = match target {
             Target::Winterfell => Box::new(air_codegen_winter::CodeGenerator),

air-script/tests/codegen/masm_with_mir.rs

@@ -0,0 +1,228 @@
+use super::helpers::{Pipeline, Target, Test};
+use expect_test::expect_file;
+
+// tests_wo_mir
+// ================================================================================================
+
+#[test]
+fn aux_trace() {
+    let generated_masm = Test::new("tests/aux_trace/aux_trace.air".to_string())
+        .transpile(Target::Masm, Pipeline::WithMIR)
+        .unwrap();
+
+    let expected = expect_file!["../aux_trace/aux_trace.masm"];
+    expected.assert_eq(&generated_masm);
+}
+
+#[test]
+fn binary() {
+    let generated_masm = Test::new("tests/binary/binary.air".to_string())
+        .transpile(Target::Masm, Pipeline::WithMIR)
+        .unwrap();
+
+    let expected = expect_file!["../binary/binary.masm"];
+    expected.assert_eq(&generated_masm);
+}
+
+#[test]
+fn periodic_columns() {
+    let generated_masm = Test::new("tests/periodic_columns/periodic_columns.air".to_string())
+        .transpile(Target::Masm, Pipeline::WithMIR)
+        .unwrap();
+
+    let expected = expect_file!["../periodic_columns/periodic_columns.masm"];
+    expected.assert_eq(&generated_masm);
+}
+
+#[test]
+fn pub_inputs() {
+    let generated_masm = Test::new("tests/pub_inputs/pub_inputs.air".to_string())
+        .transpile(Target::Masm, Pipeline::WithMIR)
+        .unwrap();
+
+    let expected = expect_file!["../pub_inputs/pub_inputs.masm"];
+    expected.assert_eq(&generated_masm);
+}
+
+#[test]
+fn system() {
+    let generated_masm = Test::new("tests/system/system.air".to_string())
+        .transpile(Target::Masm, Pipeline::WithMIR)
+        .unwrap();
+
+    let expected = expect_file!["../system/system.masm"];
+    expected.assert_eq(&generated_masm);
+}
+
+#[test]
+fn bitwise() {
+    let generated_masm = Test::new("tests/bitwise/bitwise.air".to_string())
+        .transpile(Target::Masm, Pipeline::WithMIR)
+        .unwrap();
+
+    let expected = expect_file!["../bitwise/bitwise.masm"];
+    expected.assert_eq(&generated_masm);
+}
+
+#[test]
+fn constants() {
+    let generated_masm = Test::new("tests/constants/constants.air".to_string())
+        .transpile(Target::Masm, Pipeline::WithMIR)
+        .unwrap();
+
+    let expected = expect_file!["../constants/constants.masm"];
+    expected.assert_eq(&generated_masm);
+}
+
+#[test]
+fn constant_in_range() {
+    let generated_masm = Test::new("tests/constant_in_range/constant_in_range.air".to_string())
+        .transpile(Target::Masm, Pipeline::WithMIR)
+        .unwrap();
+
+    let expected = expect_file!["../constant_in_range/constant_in_range.masm"];
+    expected.assert_eq(&generated_masm);
+}
+
+#[test]
+fn evaluators() {
+    let generated_masm = Test::new("tests/evaluators/evaluators.air".to_string())
+        .transpile(Target::Masm, Pipeline::WithMIR)
+        .unwrap();
+
+    let expected = expect_file!["../evaluators/evaluators.masm"];
+    expected.assert_eq(&generated_masm);
+}
+
+#[test]
+fn functions_simple() {
+    let generated_masm = Test::new("tests/functions/functions_simple.air".to_string())
+        .transpile(Target::Masm, Pipeline::WithMIR)
+        .unwrap();
+
+    let expected = expect_file!["../functions/functions_simple.masm"];
+    expected.assert_eq(&generated_masm);
+}
+
+#[test]
+fn functions_simple_inlined() {
+    // make sure that the constraints generated using inlined functions are the same as the ones
+    // generated using regular functions
+    let generated_masm = Test::new("tests/functions/inlined_functions_simple.air".to_string())
+        .transpile(Target::Masm, Pipeline::WithMIR)
+        .unwrap();
+    let expected = expect_file!["../functions/functions_simple.masm"];
+    expected.assert_eq(&generated_masm);
+}
+
+#[test]
+fn functions_complex() {
+    let generated_masm = Test::new("tests/functions/functions_complex.air".to_string())
+        .transpile(Target::Masm, Pipeline::WithMIR)
+        .unwrap();
+
+    let expected = expect_file!["../functions/functions_complex_with_mir.masm"];
+    expected.assert_eq(&generated_masm);
+}
+
+#[test]
+fn variables() {
+    let generated_masm = Test::new("tests/variables/variables.air".to_string())
+        .transpile(Target::Masm, Pipeline::WithMIR)
+        .unwrap();
+
+    let expected = expect_file!["../variables/variables.masm"];
+    expected.assert_eq(&generated_masm);
+}
+
+#[test]
+fn trace_col_groups() {
+    let generated_masm = Test::new("tests/trace_col_groups/trace_col_groups.air".to_string())
+        .transpile(Target::Masm, Pipeline::WithMIR)
+        .unwrap();
+
+    let expected = expect_file!["../trace_col_groups/trace_col_groups.masm"];
+    expected.assert_eq(&generated_masm);
+}
+
+#[test]
+fn indexed_trace_access() {
+    let generated_masm =
+        Test::new("tests/indexed_trace_access/indexed_trace_access.air".to_string())
+            .transpile(Target::Masm, Pipeline::WithMIR)
+            .unwrap();
+
+    let expected = expect_file!["../indexed_trace_access/indexed_trace_access.masm"];
+    expected.assert_eq(&generated_masm);
+}
+
+#[test]
+#[ignore] // TODO: There is some non-determinism in the IR creation, unskip this test once it is fixed
+fn random_values() {
+    let generated_masm = Test::new("tests/random_values/random_values_simple.air".to_string())
+        .transpile(Target::Masm, Pipeline::WithMIR)
+        .unwrap();
+    let expected = expect_file!["../random_values/random_values.masm"];
+    expected.assert_eq(&generated_masm);
+
+    let generated_masm = Test::new("tests/random_values/random_values_bindings.air".to_string())
+        .transpile(Target::Masm, Pipeline::WithMIR)
+        .unwrap();
+    let expected = expect_file!["../random_values/random_values.masm"];
+    expected.assert_eq(&generated_masm);
+}
+
+#[test]
+fn list_comprehension() {
+    let generated_masm = Test::new("tests/list_comprehension/list_comprehension.air".to_string())
+        .transpile(Target::Masm, Pipeline::WithMIR)
+        .unwrap();
+
+    let expected = expect_file!["../list_comprehension/list_comprehension_with_mir.masm"];
+    expected.assert_eq(&generated_masm);
+}
+
+#[test]
+fn list_folding() {
+    let generated_masm = Test::new("tests/list_folding/list_folding.air".to_string())
+        .transpile(Target::Masm, Pipeline::WithMIR)
+        .unwrap();
+
+    let expected = expect_file!["../list_folding/list_folding_with_mir.masm"];
+    expected.assert_eq(&generated_masm);
+}
+
+#[test]
+#[ignore] // TODO: There is some non-determinism in the IR creation, unskip this test once it is fixed
+fn selectors() {
+    let generated_masm = Test::new("tests/selectors/selectors.air".to_string())
+        .transpile(Target::Masm, Pipeline::WithMIR)
+        .unwrap();
+    let expected = expect_file!["../selectors/selectors.masm"];
+    expected.assert_eq(&generated_masm);
+
+    let generated_masm = Test::new("tests/selectors/selectors_with_evaluators.air".to_string())
+        .transpile(Target::Masm, Pipeline::WithMIR)
+        .unwrap();
+    let expected = expect_file!["../selectors/selectors.masm"];
+    expected.assert_eq(&generated_masm);
+}
+
+#[test]
+fn constraint_comprehension() {
+    let generated_masm =
+        Test::new("tests/constraint_comprehension/constraint_comprehension.air".to_string())
+            .transpile(Target::Masm, Pipeline::WithMIR)
+            .unwrap();
+
+    let expected = expect_file!["../constraint_comprehension/constraint_comprehension.masm"];
+    expected.assert_eq(&generated_masm);
+
+    let generated_masm =
+        Test::new("tests/constraint_comprehension/cc_with_evaluators.air".to_string())
+            .transpile(Target::Masm, Pipeline::WithMIR)
+            .unwrap();
+
+    let expected = expect_file!["../constraint_comprehension/constraint_comprehension.masm"];
+    expected.assert_eq(&generated_masm);
+}