Merge pull request #1003 from 0xPolygonMiden/jjcnn-ast-pretty-printer

AST formatter/pretty-printer

assembly/src/ast/format.rs

@@ -0,0 +1,134 @@
+use super::{
+    CodeBody, FormattableNode, InvokedProcsMap, LibraryPath, ProcedureAst, ProcedureId,
+    ProcedureName, Vec,
+};
+use core::fmt;
+
+const INDENT_STRING: &str = "    ";
+
+/// Context for the Ast formatter
+///
+/// The context keeps track of the current indentation level, as well as the declared and imported
+/// procedures in the program/module being formatted.
+pub struct AstFormatterContext<'a> {
+    indent_level: usize,
+    local_procs: &'a Vec<ProcedureAst>,
+    imported_procs: &'a InvokedProcsMap,
+}
+
+impl<'a> AstFormatterContext<'a> {
+    pub fn new(
+        local_procs: &'a Vec<ProcedureAst>,
+        imported_procs: &'a InvokedProcsMap,
+    ) -> AstFormatterContext<'a> {
+        Self {
+            indent_level: 0,
+            local_procs,
+            imported_procs,
+        }
+    }
+
+    /// Build a context for the inner scope, e.g., the body of a while loop
+    pub fn inner_scope_context(&self) -> Self {
+        Self {
+            indent_level: self.indent_level + 1,
+            local_procs: self.local_procs,
+            imported_procs: self.imported_procs,
+        }
+    }
+
+    /// Add indentation to the current line in the formatter
+    pub fn indent(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        for _ in 0..self.indent_level {
+            write!(f, "{INDENT_STRING}")?;
+        }
+        Ok(())
+    }
+
+    /// Get the name of the local procedure with the given index.
+    ///
+    /// # Panics
+    /// Panics if the index is not associated with a procedure name
+    pub fn local_proc(&self, index: usize) -> &ProcedureName {
+        assert!(index < self.local_procs.len(), "Local procedure with index {index} not found");
+        &self.local_procs[index].name
+    }
+
+    /// Get the name of the imported procedure with the given id/hash.
+    ///
+    /// # Panics
+    /// Panics if the id/hash is not associated with an imported procedure
+    pub fn imported_proc(&self, id: &ProcedureId) -> &(ProcedureName, LibraryPath) {
+        self.imported_procs
+            .get(id)
+            .expect("Imported procedure with id/hash {id} not found")
+    }
+}
+
+// FORMATTING OF PROCEDURES
+// ================================================================================================
+pub struct FormattableProcedureAst<'a> {
+    proc: &'a ProcedureAst,
+    context: &'a AstFormatterContext<'a>,
+}
+
+impl<'a> FormattableProcedureAst<'a> {
+    pub fn new(
+        proc: &'a ProcedureAst,
+        context: &'a AstFormatterContext<'a>,
+    ) -> FormattableProcedureAst<'a> {
+        Self { proc, context }
+    }
+}
+
+impl fmt::Display for FormattableProcedureAst<'_> {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        // Docs
+        self.context.indent(f)?;
+        if let Some(ref doc) = self.proc.docs {
+            writeln!(f, "#! {doc}")?;
+        }
+        // Procedure header
+        self.context.indent(f)?;
+        if self.proc.is_export {
+            write!(f, "export.")?;
+        } else {
+            write!(f, "proc.")?;
+        }
+        writeln!(f, "{}.{}", self.proc.name, self.proc.num_locals)?;
+        // Body
+        write!(
+            f,
+            "{}",
+            FormattableCodeBody::new(&self.proc.body, &self.context.inner_scope_context())
+        )?;
+        // Procedure footer
+        self.context.indent(f)?;
+        writeln!(f, "end")
+    }
+}
+
+// FORMATTING OF CODE BODIES
+// ================================================================================================
+pub struct FormattableCodeBody<'a> {
+    body: &'a CodeBody,
+    context: &'a AstFormatterContext<'a>,
+}
+
+impl<'a> FormattableCodeBody<'a> {
+    pub fn new(
+        body: &'a CodeBody,
+        context: &'a AstFormatterContext<'a>,
+    ) -> FormattableCodeBody<'a> {
+        Self { body, context }
+    }
+}
+
+impl fmt::Display for FormattableCodeBody<'_> {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        for node in self.body.nodes() {
+            write!(f, "{}", FormattableNode::new(node, self.context))?;
+        }
+        Ok(())
+    }
+}

assembly/src/ast/imports.rs

@@ -1,14 +1,13 @@
 use super::{
-    BTreeMap, ByteReader, ByteWriter, Deserializable, DeserializationError, LibraryPath,
-    ParsingError, ProcedureId, ProcedureName, Serializable, String, ToString, Token, TokenStream,
-    Vec, MAX_IMPORTS, MAX_INVOKED_IMPORTED_PROCS,
+    BTreeMap, ByteReader, ByteWriter, Deserializable, DeserializationError, InvokedProcsMap,
+    LibraryPath, ParsingError, ProcedureId, ProcedureName, Serializable, String, ToString, Token,
+    TokenStream, Vec, MAX_IMPORTS, MAX_INVOKED_IMPORTED_PROCS,
 };
 
 // TYPE ALIASES
 // ================================================================================================
 
 type ImportedModulesMap = BTreeMap<String, LibraryPath>;
-type InvokedProcsMap = BTreeMap<ProcedureId, (ProcedureName, LibraryPath)>;
 
 // MODULE IMPORTS
 // ================================================================================================
@@ -89,6 +88,11 @@ impl ModuleImports {
         self.imports.values().collect()
     }
 
+    /// Returns a reference to the invoked procedure map which maps procedure IDs to their names.
+    pub fn invoked_procs(&self) -> &InvokedProcsMap {
+        &self.invoked_procs
+    }
+
     // STATE MUTATORS
     // --------------------------------------------------------------------------------------------
 

assembly/src/ast/mod.rs

@@ -9,17 +9,21 @@ use super::{
     Serializable, SliceReader, StarkField, String, ToString, Token, TokenStream, Vec,
     MAX_LABEL_LEN,
 };
-use core::{iter, str::from_utf8};
+use core::{fmt, iter, str::from_utf8};
 use vm_core::utils::bound_into_included_u64;
 
 pub use super::tokens::SourceLocation;
 
 mod nodes;
+use nodes::FormattableNode;
 pub use nodes::{AdviceInjectorNode, Instruction, Node};
 
 mod code_body;
 pub use code_body::CodeBody;
 
+mod format;
+use format::*;
+
 mod imports;
 pub use imports::ModuleImports;
 
@@ -66,6 +70,7 @@ const MAX_STACK_WORD_OFFSET: u8 = 12;
 type LocalProcMap = BTreeMap<ProcedureName, (u16, ProcedureAst)>;
 type LocalConstMap = BTreeMap<String, u64>;
 type ReExportedProcMap = BTreeMap<ProcedureName, ProcReExport>;
+type InvokedProcsMap = BTreeMap<ProcedureId, (ProcedureName, LibraryPath)>;
 
 // EXECUTABLE PROGRAM AST
 // ================================================================================================
@@ -211,7 +216,6 @@ impl ProgramAst {
 
         let local_procs = sort_procs_into_vec(context.local_procs);
         let (nodes, locations) = body.into_parts();
-
         Ok(Self::new(nodes, local_procs)?
             .with_source_locations(locations, start)
             .with_import_info(import_info))
@@ -325,6 +329,46 @@ impl ProgramAst {
     }
 }
 
+impl fmt::Display for ProgramAst {
+    /// Writes this [ProgramAst] as formatted MASM code into the formatter.
+    ///
+    /// The formatted code puts each instruction on a separate line and preserves correct indentation
+    /// for instruction blocks.
+    ///
+    /// # Panics
+    /// Panics if import info is not associated with this program.
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        assert!(self.import_info.is_some(), "Program imports not instantiated");
+
+        // Imports
+        if let Some(ref info) = self.import_info {
+            let paths = info.import_paths();
+            for path in paths.iter() {
+                writeln!(f, "use.{path}")?;
+            }
+            if !paths.is_empty() {
+                writeln!(f)?;
+            }
+        }
+
+        let tmp_procs = InvokedProcsMap::new();
+        let invoked_procs =
+            self.import_info.as_ref().map(|info| info.invoked_procs()).unwrap_or(&tmp_procs);
+
+        let context = AstFormatterContext::new(&self.local_procs, invoked_procs);
+
+        // Local procedures
+        for proc in self.local_procs.iter() {
+            writeln!(f, "{}", FormattableProcedureAst::new(proc, &context))?;
+        }
+
+        // Main progrma
+        writeln!(f, "begin")?;
+        write!(f, "{}", FormattableCodeBody::new(&self.body, &context.inner_scope_context()))?;
+        writeln!(f, "end")
+    }
+}
+
 // MODULE AST
 // ================================================================================================
 
@@ -594,6 +638,54 @@ impl ModuleAst {
     }
 }
 
+impl fmt::Display for ModuleAst {
+    /// Writes this [ModuleAst] as formatted MASM code into the formatter.
+    ///
+    /// The formatted code puts each instruction on a separate line and preserves correct indentation
+    /// for instruction blocks.
+    ///
+    /// # Panics
+    /// Panics if import info is not associated with this module.
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        assert!(self.import_info.is_some(), "Program imports not instantiated");
+
+        // Docs
+        if let Some(ref doc) = self.docs {
+            writeln!(f, "#! {doc}")?;
+            writeln!(f)?;
+        }
+
+        // Imports
+        if let Some(ref info) = self.import_info {
+            let paths = info.import_paths();
+            for path in paths.iter() {
+                writeln!(f, "use.{path}")?;
+            }
+            if !paths.is_empty() {
+                writeln!(f)?;
+            }
+        }
+
+        // Re-exports
+        for proc in self.reexported_procs.iter() {
+            writeln!(f, "export.{}", proc.name)?;
+            writeln!(f)?;
+        }
+
+        // Local procedures
+        let tmp_procs = InvokedProcsMap::new();
+        let invoked_procs =
+            self.import_info.as_ref().map(|info| info.invoked_procs()).unwrap_or(&tmp_procs);
+
+        let context = AstFormatterContext::new(&self.local_procs, invoked_procs);
+
+        for proc in self.local_procs.iter() {
+            writeln!(f, "{}", FormattableProcedureAst::new(proc, &context))?;
+        }
+        Ok(())
+    }
+}
+
 // PROCEDURE AST
 // ================================================================================================