Implement post-order codegen for structs and records

compiler/compiler/src/compiler.rs

@@ -203,9 +203,9 @@ impl<'a> Compiler<'a> {
     }
 
     /// Runs the compiler stages.
-    pub fn compiler_stages(&mut self) -> Result<SymbolTable> {
+    pub fn compiler_stages(&mut self) -> Result<(SymbolTable, StructGraph)> {
         let st = self.symbol_table_pass()?;
-        let (st, _struct_graph) = self.type_checker_pass(st)?;
+        let (st, struct_graph) = self.type_checker_pass(st)?;
 
         // TODO: Make this pass optional.
         let st = self.loop_unrolling_pass(st)?;
@@ -215,16 +215,16 @@ impl<'a> Compiler<'a> {
 
         self.flattening_pass(&st, assigner)?;
 
-        Ok(st)
+        Ok((st, struct_graph))
     }
 
     /// Returns a compiled Leo program and prints the resulting bytecode.
     // TODO: Remove when code generation is ready to be integrated into the compiler.
     pub fn compile_and_generate_instructions(&mut self) -> Result<(SymbolTable, String)> {
         self.parse_program()?;
-        let symbol_table = self.compiler_stages()?;
+        let (symbol_table, struct_graph) = self.compiler_stages()?;
 
-        let bytecode = CodeGenerator::do_pass((&self.ast, &symbol_table))?;
+        let bytecode = CodeGenerator::do_pass((&self.ast, &symbol_table, &struct_graph))?;
 
         Ok((symbol_table, bytecode))
     }
@@ -232,7 +232,7 @@ impl<'a> Compiler<'a> {
     /// Returns a compiled Leo program.
     pub fn compile(&mut self) -> Result<SymbolTable> {
         self.parse_program()?;
-        self.compiler_stages()
+        self.compiler_stages().map(|(st, _)| st)
     }
 
     /// Writes the AST to a JSON file.

compiler/compiler/src/test.rs

@@ -194,14 +194,14 @@ fn temp_dir() -> PathBuf {
 
 fn compile_and_process<'a>(parsed: &'a mut Compiler<'a>) -> Result<String, LeoError> {
     let st = parsed.symbol_table_pass()?;
-    let (st, _struct_graph) = parsed.type_checker_pass(st)?;
+    let (st, struct_graph) = parsed.type_checker_pass(st)?;
     let st = parsed.loop_unrolling_pass(st)?;
     let assigner = parsed.static_single_assignment_pass(&st)?;
 
     parsed.flattening_pass(&st, assigner)?;
 
     // Compile Leo program to bytecode.
-    let bytecode = CodeGenerator::do_pass((&parsed.ast, &st))?;
+    let bytecode = CodeGenerator::do_pass((&parsed.ast, &st, &struct_graph))?;
 
     Ok(bytecode)
 }

compiler/passes/src/code_generation/generator.rs

@@ -15,6 +15,7 @@
 // along with the Leo library. If not, see <https://www.gnu.org/licenses/>.
 
 use crate::SymbolTable;
+use crate::StructGraph;
 
 use leo_ast::Function;
 use leo_span::Symbol;
@@ -24,6 +25,8 @@ use indexmap::IndexMap;
 pub struct CodeGenerator<'a> {
     /// The symbol table for the program.
     pub(crate) symbol_table: &'a SymbolTable,
+    /// The struct dependency graph for the program.
+    pub(crate) struct_graph: &'a StructGraph,
     /// A counter to track the next available register.
     pub(crate) next_register: u64,
     /// Reference to the current function.
@@ -42,10 +45,11 @@ pub struct CodeGenerator<'a> {
 
 impl<'a> CodeGenerator<'a> {
     /// Initializes a new `CodeGenerator`.
-    pub fn new(symbol_table: &'a SymbolTable) -> Self {
+    pub fn new(symbol_table: &'a SymbolTable, struct_graph: &'a StructGraph) -> Self {
         // Initialize variable mapping.
         Self {
             symbol_table,
+            struct_graph,
             next_register: 0,
             current_function: None,
             variable_mapping: IndexMap::new(),

compiler/passes/src/code_generation/mod.rs

@@ -26,16 +26,17 @@ mod visit_statements;
 mod visit_type;
 
 use crate::{Pass, SymbolTable};
+use crate::{StructGraph};
 
 use leo_ast::Ast;
 use leo_errors::Result;
 
 impl<'a> Pass for CodeGenerator<'a> {
-    type Input = (&'a Ast, &'a SymbolTable);
+    type Input = (&'a Ast, &'a SymbolTable, &'a StructGraph);
     type Output = Result<String>;
 
-    fn do_pass((ast, symbol_table): Self::Input) -> Self::Output {
+    fn do_pass((ast, symbol_table, struct_graph): Self::Input) -> Self::Output {
-        let mut generator = Self::new(symbol_table);
+        let mut generator = Self::new(symbol_table, struct_graph);
         let bytecode = generator.visit_program(ast.as_repr());
 
         Ok(bytecode)

compiler/passes/src/code_generation/visit_program.rs

@@ -16,11 +16,11 @@
 
 use crate::CodeGenerator;
 
-use leo_ast::{functions, CallType, Function, Identifier, Mapping, Mode, Program, ProgramScope, Struct, Type};
+use leo_ast::{functions, CallType, Function, Mapping, Mode, Program, ProgramScope, Struct, Type};
 
 use indexmap::IndexMap;
 use itertools::Itertools;
-use leo_span::sym;
+use leo_span::{sym, Symbol};
 use std::fmt::Write as _;
 
 impl<'a> CodeGenerator<'a> {
@@ -53,12 +53,19 @@ impl<'a> CodeGenerator<'a> {
         // Newline separator.
         program_string.push('\n');
 
-        // Visit each `Struct` or `Record` in the Leo AST and produce a Aleo struct.
+        // Get the post-order ordering of the composite data types.
+        // Note that the unwrap is safe since type checking guarantees that the struct dependency graph is acyclic.
+        let order = self.struct_graph.post_order().unwrap();
+
+        // Visit each `Struct` or `Record` in the post-ordering and produce an Aleo struct or record.
         program_string.push_str(
-            &program_scope
+            &order
-                .structs
+                .into_iter()
-                .values()
+                .map(|name| {
-                .map(|struct_| self.visit_struct_or_record(struct_))
+                    // Note that this unwrap is safe since type checking guarantees that all structs are declared.
+                    let struct_ = program_scope.structs.get(&name).unwrap();
+                    self.visit_struct_or_record(struct_)
+                })
                 .join("\n"),
         );
 
@@ -104,7 +111,7 @@ impl<'a> CodeGenerator<'a> {
         program_string
     }
 
-    fn visit_import(&mut self, import_name: &'a Identifier, import_program: &'a Program) -> String {
+    fn visit_import(&mut self, import_name: &'a Symbol, import_program: &'a Program) -> String {
         // Load symbols into composite mapping.
         let _import_program_string = self.visit_program(import_program);
         // todo: We do not need the import program string because we generate instructions for imports separately during leo build.

tests/test-framework/benches/leo_compiler.rs

@@ -41,6 +41,7 @@ enum BenchMode {
     Ssa,
     /// Benchmarks flattening.
     Flatten,
+    // TODO: Benchmark code generation
     /// Benchmarks all the above stages.
     Full,
 }