Review comments addressed

0xPolygonMiden · Jul 18, 2023 · 9735e5f · 9735e5f
1 parent 93cc387
commit 9735e5f
Show file tree

Hide file tree

Showing 3 changed files with 43 additions and 35 deletions.
diff --git a/assembly/src/ast/imports.rs b/assembly/src/ast/imports.rs
@@ -88,6 +88,7 @@ impl ModuleImports {
         self.imports.values().collect()
     }
 
+    /// Return the invoked procedures map from procedure ids to procedure names
     pub fn invoked_procs(&self) -> &InvokedProcsMap {
         &self.invoked_procs
     }

diff --git a/assembly/src/ast/mod.rs b/assembly/src/ast/mod.rs
@@ -15,13 +15,14 @@ use vm_core::utils::bound_into_included_u64;
 pub use super::tokens::SourceLocation;
 
 mod nodes;
-pub use nodes::{AdviceInjectorNode, FormattableNode, Instruction, Node};
+use nodes::FormattableNode;
+pub use nodes::{AdviceInjectorNode, Instruction, Node};
 
 mod code_body;
 pub use code_body::CodeBody;
 
 mod format;
-pub use format::*;
+use format::*;
 
 mod imports;
 pub use imports::ModuleImports;
@@ -329,6 +330,8 @@ impl ProgramAst {
 }
 
 impl fmt::Display for ProgramAst {
+    /// Displays the ProgramAst as formatted MASM code with indentation etc.
+    ///
     /// # Panics
     ///
     /// Import info must be added before the program can be displayed.
@@ -347,11 +350,8 @@ impl fmt::Display for ProgramAst {
         }
 
         let tmp_procs = InvokedProcsMap::new();
-        let invoked_procs = if let Some(ref info) = self.import_info {
-            info.invoked_procs()
-        } else {
-            &tmp_procs
-        };
+        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);
 
@@ -637,6 +637,8 @@ impl ModuleAst {
 }
 
 impl fmt::Display for ModuleAst {
+    /// Displays the ModuleAst as formatted MASM code with indentation etc.
+    ///
     /// # Panics
     ///
     /// Import info must be added before the program can be displayed.
@@ -668,21 +670,13 @@ impl fmt::Display for ModuleAst {
 
         // Local procedures
         let tmp_procs = InvokedProcsMap::new();
-        let invoked_procs = if let Some(ref info) = self.import_info {
-            info.invoked_procs()
-        } else {
-            &tmp_procs
-        };
+        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);
-        let mut first = true;
 
         for proc in self.local_procs.iter() {
-            if !first {
-                writeln!(f)?;
-            }
-            write!(f, "{}", FormattableProcedureAst::new(proc, &context))?;
-            first = false;
+            writeln!(f, "{}", FormattableProcedureAst::new(proc, &context))?;
         }
         Ok(())
     }

diff --git a/assembly/src/ast/nodes/format.rs b/assembly/src/ast/nodes/format.rs
@@ -25,23 +25,36 @@ impl fmt::Display for FormattableNode<'_> {
                 false_case,
             } => {
                 self.context.indent(f)?;
-                writeln!(f, "if.true")?;
-
-                write!(
-                    f,
-                    "{}",
-                    FormattableCodeBody::new(true_case, &self.context.inner_scope_context())
-                )?;
-
-                self.context.indent(f)?;
-                writeln!(f, "else")?;
-
-                write!(
-                    f,
-                    "{}",
-                    FormattableCodeBody::new(false_case, &self.context.inner_scope_context())
-                )?;
-
+                if true_case.nodes().is_empty() {
+                    // No true branch - only output false branch
+                    writeln!(f, "if.false")?;
+                    write!(
+                        f,
+                        "{}",
+                        FormattableCodeBody::new(false_case, &self.context.inner_scope_context())
+                    )?;
+                } else {
+                    writeln!(f, "if.true")?;
+                    write!(
+                        f,
+                        "{}",
+                        FormattableCodeBody::new(true_case, &self.context.inner_scope_context())
+                    )?;
+                    if false_case.nodes().is_empty() {
+                        // No false branch - don't output else branch
+                        self.context.indent(f)?;
+                        writeln!(f, "else")?;
+
+                        write!(
+                            f,
+                            "{}",
+                            FormattableCodeBody::new(
+                                false_case,
+                                &self.context.inner_scope_context()
+                            )
+                        )?;
+                    }
+                }
                 self.context.indent(f)?;
                 writeln!(f, "end")
             }