Integrate SSA pass with compiler

compiler/ast/src/statements/assign/mod.rs

@@ -47,12 +47,12 @@ pub enum AssignOperation {
     BitXor,
     /// Shift right assignment.
     Shr,
-    /// Signed shift right assignment.
-    ShrSigned,
+    // /// Signed shift right assignment.
+    // ShrSigned,
     /// Shift left assignment.
     Shl,
-    /// Modulus / remainder assignment.
-    Mod,
+    // /// Modulus / remainder assignment.
+    // Mod,
 }
 
 impl AsRef<str> for AssignOperation {
@@ -70,9 +70,9 @@ impl AsRef<str> for AssignOperation {
             AssignOperation::BitAnd => "&=",
             AssignOperation::BitXor => "^=",
             AssignOperation::Shr => ">>=",
-            AssignOperation::ShrSigned => ">>>=",
+            // AssignOperation::ShrSigned => ">>>=",
             AssignOperation::Shl => "<<=",
-            AssignOperation::Mod => "%=",
+            // AssignOperation::Mod => "%=",
         }
     }
 }

compiler/compiler/src/compiler.rs

@@ -98,18 +98,12 @@ impl<'a> Compiler<'a> {
         let mut ast: leo_ast::Ast = leo_parser::parse_ast(self.handler, &prg_sf.src, prg_sf.start_pos)?;
         ast = ast.set_program_name(self.program_name.clone());
         ast = ast.set_network(self.network.clone());
+        self.ast = ast;
 
         if self.output_options.initial_ast {
-            // Write the AST snapshot post parsing.
-            if self.output_options.spans_enabled {
-                ast.to_json_file(self.output_directory.clone(), "initial_ast.json")?;
-            } else {
-                ast.to_json_file_without_keys(self.output_directory.clone(), "initial_ast.json", &["span"])?;
-            }
+            self.write_ast_to_json("initial_ast.json")?;
         }
 
-        self.ast = ast;
-
         Ok(())
     }
 
@@ -171,6 +165,17 @@ impl<'a> Compiler<'a> {
         Ok(symbol_table)
     }
 
+    /// Runs the static single assignment pass.
+    pub fn static_single_assignment_pass(&mut self) -> Result<()> {
+        self.ast = StaticSingleAssigner::do_pass((std::mem::take(&mut self.ast), self.handler))?;
+
+        if self.output_options.ssa_ast {
+            self.write_ast_to_json("ssa_ast.json")?;
+        }
+
+        Ok(())
+    }
+
     /// Runs the compiler stages.
     pub fn compiler_stages(&mut self) -> Result<SymbolTable> {
         let st = self.symbol_table_pass()?;
@@ -178,6 +183,9 @@ impl<'a> Compiler<'a> {
 
         // TODO: Make this pass optional.
         let st = self.loop_unrolling_pass(st)?;
+
+        self.static_single_assignment_pass()?;
+
         Ok(st)
     }
 

compiler/compiler/src/options.rs

@@ -24,4 +24,6 @@ pub struct OutputOptions {
     pub initial_input_ast: bool,
     /// If enabled writes the AST after loop unrolling.
     pub unrolled_ast: bool,
+    /// If enabled write the AST after static single assignment.
+    pub ssa_ast: bool,
 }

compiler/compiler/src/test.rs

@@ -51,6 +51,7 @@ fn new_compiler(handler: &Handler, main_file_path: PathBuf) -> Compiler<'_> {
             initial_input_ast: true,
             initial_ast: true,
             unrolled_ast: true,
+            ssa_ast: true,
         }),
     )
 }
@@ -106,6 +107,7 @@ struct CompileOutput {
     pub output: Vec<OutputItem>,
     pub initial_ast: String,
     pub unrolled_ast: String,
+    pub ssa_ast: String,
 }
 
 /// Get the path of the `input_file` given in `input` into `list`.
@@ -139,6 +141,7 @@ fn compile_and_process<'a>(parsed: &'a mut Compiler<'a>) -> Result<SymbolTable,
     let st = parsed.symbol_table_pass()?;
     let st = parsed.type_checker_pass(st)?;
     let st = parsed.loop_unrolling_pass(st)?;
+    parsed.static_single_assignment_pass()?;
     Ok(st)
 }
 
@@ -219,6 +222,7 @@ fn run_test(test: Test, handler: &Handler, err_buf: &BufferEmitter) -> Result<Va
 
     let initial_ast = hash_file("/tmp/output/initial_ast.json");
     let unrolled_ast = hash_file("/tmp/output/unrolled_ast.json");
+    let ssa_ast = hash_file("/tmp/output/ssa_ast.json");
 
     if fs::read_dir("/tmp/output").is_ok() {
         fs::remove_dir_all(Path::new("/tmp/output")).expect("Error failed to clean up output dir.");
@@ -228,6 +232,7 @@ fn run_test(test: Test, handler: &Handler, err_buf: &BufferEmitter) -> Result<Va
         output: output_items,
         initial_ast,
         unrolled_ast,
+        ssa_ast,
     };
     Ok(serde_yaml::to_value(&final_output).expect("serialization failed"))
 }

compiler/passes/src/lib.rs

@@ -25,6 +25,9 @@ pub use self::pass::*;
 pub mod loop_unrolling;
 pub use self::loop_unrolling::*;
 
+pub mod static_single_assignment;
+pub use static_single_assignment::*;
+
 pub mod symbol_table;
 pub use symbol_table::*;
 

compiler/passes/src/static_single_assignment/rename_statement.rs

@@ -29,9 +29,13 @@ impl<'a> StatementReconstructor for StaticSingleAssigner<'a> {
     /// Reconstructs the `DefinitionStatement` into an `AssignStatement`, renaming the left-hand-side as appropriate.
     fn reconstruct_definition(&mut self, definition: DefinitionStatement) -> Statement {
         self.is_lhs = true;
-        let identifier = match self.reconstruct_identifier(definition.variable_name.identifier).0 {
-            Expression::Identifier(identifier) => identifier,
-            _ => unreachable!("`reconstruct_identifier` will always return an `Identifier`."),
+        // TODO: Change to support a vector of identifiers.
+        let identifier = match definition.variable_names.len() == 1 {
+            true => match self.reconstruct_identifier(definition.variable_names[0].identifier).0 {
+                Expression::Identifier(identifier) => identifier,
+                _ => unreachable!("`reconstruct_identifier` will always return an `Identifier`."),
+            },
+            false => unreachable!("DefinitionStatement should have only one variable name."),
         };
         self.is_lhs = false;
 
@@ -100,17 +104,19 @@ impl<'a> StatementReconstructor for StaticSingleAssigner<'a> {
 
         // Instantiate a `RenameTable` for the else-block.
         self.push();
-        let next = conditional
-            .next
-            .map(|statement| Box::new(match *statement {
+        let next = conditional.next.map(|statement| {
+            Box::new(match *statement {
                 // The `ConditionalStatement` must be reconstructed as a `Block` statement to ensure that appropriate statements are produced.
-                Statement::Conditional(stmt) => { self.reconstruct_statement(Statement::Block(Block {
+                Statement::Conditional(stmt) => self.reconstruct_statement(Statement::Block(Block {
                     statements: vec![Statement::Conditional(stmt)],
-                    span: Default::default()
-                }))}
-                Statement::Block(stmt) => { self.reconstruct_statement(Statement::Block(stmt)) }
-                _ => unreachable!("`ConditionalStatement`s next statement must be a `ConditionalStatement` or a `Block`."),
-            }));
+                    span: Default::default(),
+                })),
+                Statement::Block(stmt) => self.reconstruct_statement(Statement::Block(stmt)),
+                _ => unreachable!(
+                    "`ConditionalStatement`s next statement must be a `ConditionalStatement` or a `Block`."
+                ),
+            })
+        });
 
         let else_table = self.pop();
 
@@ -195,7 +201,12 @@ impl<'a> StatementReconstructor for StaticSingleAssigner<'a> {
                         Expression::Identifier(..) | Expression::Literal(..) => {
                             self.reconstruct_conditional(conditional_statement)
                         }
-                        Expression::Binary(..) | Expression::Unary(..) | Expression::Ternary(..) => {
+                        Expression::Access(..)
+                        | Expression::Circuit(..)
+                        | Expression::Tuple(..)
+                        | Expression::Binary(..)
+                        | Expression::Unary(..)
+                        | Expression::Ternary(..) => {
                             // Create a fresh variable name for the condition.
                             let symbol = Symbol::intern(&format!("cond${}", self.get_unique_id()));
                             self.rename_table.update(symbol, symbol);

leo/commands/build.rs

@@ -52,6 +52,8 @@ pub struct BuildOptions {
     pub enable_initial_ast_snapshot: bool,
     #[structopt(long, help = "Writes AST snapshot of the unrolled AST.")]
     pub enable_unrolled_ast_snapshot: bool,
+    #[structopt(long, help = "Writes AST snapshot of the SSA AST.")]
+    pub enable_ssa_ast_snapshot: bool,
     // Note: This is currently made optional since code generation is just a prototype.
     #[structopt(
         long,
@@ -67,11 +69,13 @@ impl From<BuildOptions> for OutputOptions {
             initial_input_ast: options.enable_initial_input_ast_snapshot,
             initial_ast: options.enable_initial_ast_snapshot,
             unrolled_ast: options.enable_unrolled_ast_snapshot,
+            ssa_ast: options.enable_ssa_ast_snapshot,
         };
         if options.enable_all_ast_snapshots {
             out_options.initial_input_ast = true;
             out_options.initial_ast = true;
             out_options.unrolled_ast = true;
+            out_options.ssa_ast = true;
         }
 
         out_options

tests/test-framework/benches/leo_compiler.rs

@@ -40,6 +40,8 @@ enum BenchMode {
     Type,
     /// Benchmarks loop unrolling.
     Unroll,
+    /// Benchmarks static single assignment.
+    Ssa,
     /// Benchmarks all the above stages.
     Full,
 }
@@ -103,6 +105,7 @@ impl Sample {
             BenchMode::Symbol => self.bench_symbol_table(c),
             BenchMode::Type => self.bench_type_checker(c),
             BenchMode::Unroll => self.bench_loop_unroller(c),
+            BenchMode::Ssa => self.bench_ssa(c),
             BenchMode::Full => self.bench_full(c),
         }
     }
@@ -209,6 +212,38 @@ impl Sample {
         });
     }
 
+    fn bench_ssa(&self, c: &mut Criterion) {
+        c.bench_function(&format!("full {}", self.name), |b| {
+            // Iter custom is used so we can use custom timings around the compiler stages.
+            // This way we can only time the necessary stages.
+            b.iter_custom(|iters| {
+                let mut time = Duration::default();
+                for _ in 0..iters {
+                    SESSION_GLOBALS.set(&SessionGlobals::default(), || {
+                        let handler = BufEmitter::new_handler();
+                        let mut compiler = new_compiler(&handler);
+                        let (input, name) = self.data();
+                        compiler
+                            .parse_program_from_string(input, name)
+                            .expect("Failed to parse program");
+                        let symbol_table = compiler.symbol_table_pass().expect("failed to generate symbol table");
+                        let symbol_table = compiler
+                            .type_checker_pass(symbol_table)
+                            .expect("failed to run type check pass");
+                        compiler
+                            .loop_unrolling_pass(symbol_table)
+                            .expect("failed to run loop unrolling pass");
+                        let start = Instant::now();
+                        let out = compiler.static_single_assignment_pass();
+                        time += start.elapsed();
+                        out.expect("failed to run ssa pass")
+                    });
+                }
+                time
+            })
+        });
+    }
+
     fn bench_full(&self, c: &mut Criterion) {
         c.bench_function(&format!("full {}", self.name), |b| {
             // Iter custom is used so we can use custom timings around the compiler stages.
@@ -231,6 +266,9 @@ impl Sample {
                         compiler
                             .loop_unrolling_pass(symbol_table)
                             .expect("failed to run loop unrolling pass");
+                        compiler
+                            .static_single_assignment_pass()
+                            .expect("failed to run ssa pass");
                         time += start.elapsed();
                     });
                 }
@@ -252,6 +290,7 @@ bench!(bench_parse, BenchMode::Parse);
 bench!(bench_symbol, BenchMode::Symbol);
 bench!(bench_type, BenchMode::Type);
 bench!(bench_unroll, BenchMode::Unroll);
+bench!(bench_ssa, BenchMode::Ssa);
 bench!(bench_full, BenchMode::Full);
 
 criterion_group!(
@@ -262,6 +301,7 @@ criterion_group!(
         bench_symbol,
         bench_type,
         bench_unroll,
+        bench_ssa,
         bench_full
 );
 criterion_main!(benches);