Fix up Display implementations for AST nodes.

Get details right like parens depending on precedence,
semicolons after statements only when warranted, etc.
I haven't extensively tested this yet, so there are likely
places where it does the wrong thing, but this will be
a slightly more principled implementation we can fix up
as needed.

compiler/ast/src/access/array_access.rs

@@ -35,7 +35,7 @@ pub struct ArrayAccess {
 
 impl fmt::Display for ArrayAccess {
     fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
-        write!(f, "{}.{}", self.array, self.index)
+        write!(f, "{}[{}]", self.array, self.index)
     }
 }
 

compiler/ast/src/access/associated_function_access.rs

@@ -17,6 +17,7 @@
 use crate::{Expression, Identifier, Node, NodeID};
 use leo_span::Span;
 
+use itertools::Itertools as _;
 use serde::{Deserialize, Serialize};
 use std::fmt;
 
@@ -37,7 +38,7 @@ pub struct AssociatedFunction {
 
 impl fmt::Display for AssociatedFunction {
     fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
-        write!(f, "{}::{}", self.variant, self.name)
+        write!(f, "{}::{}({})", self.variant, self.name, self.arguments.iter().format(", "))
     }
 }
 

compiler/ast/src/expressions/array.rs

@@ -16,6 +16,8 @@
 
 use super::*;
 
+use itertools::Itertools as _;
+
 /// An array expression, e.g., `[true, false, true, false]`.
 #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
 pub struct ArrayExpression {
@@ -29,7 +31,7 @@ pub struct ArrayExpression {
 
 impl fmt::Display for ArrayExpression {
     fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
-        write!(f, "[{}]", self.elements.iter().map(|x| x.to_string()).collect::<Vec<_>>().join(","))
+        write!(f, "[{}]", self.elements.iter().format(", "))
     }
 }
 

compiler/ast/src/expressions/binary.rs

@@ -17,6 +17,8 @@
 use super::*;
 use leo_span::{Symbol, sym};
 
+use std::cmp::Ordering;
+
 /// A binary operator.
 ///
 /// Precedence is defined in the parser.
@@ -177,7 +179,78 @@ pub struct BinaryExpression {
 
 impl fmt::Display for BinaryExpression {
     fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
-        write!(f, "{} {} {}", self.left, self.op, self.right)
+        use Associativity::*;
+        use BinaryOperation::*;
+
+        if matches!(
+            self.op,
+            AddWrapped
+                | DivWrapped
+                | Mod
+                | MulWrapped
+                | Nand
+                | Nor
+                | PowWrapped
+                | RemWrapped
+                | ShlWrapped
+                | ShrWrapped
+                | SubWrapped
+        ) {
+            if self.left.precedence() < 20 {
+                write!(f, "({})", self.left)?;
+            } else {
+                write!(f, "{}", self.left)?;
+            }
+            write!(f, ".{}({})", self.op, self.right)
+        } else {
+            let my_precedence = self.precedence();
+            let my_associativity = self.associativity();
+            match (self.left.precedence().cmp(&my_precedence), my_associativity, self.left.associativity()) {
+                (Ordering::Greater, _, _) | (Ordering::Equal, Left, Left) => write!(f, "{}", self.left)?,
+                _ => write!(f, "({})", self.left)?,
+            }
+            write!(f, " {} ", self.op)?;
+            match (self.right.precedence().cmp(&my_precedence), my_associativity, self.right.associativity()) {
+                (Ordering::Greater, _, _) | (Ordering::Equal, Right, Right) => write!(f, "{}", self.right)?,
+                _ => write!(f, "({})", self.right)?,
+            }
+            Ok(())
+        }
+    }
+}
+
+impl BinaryExpression {
+    pub(crate) fn precedence(&self) -> u32 {
+        use BinaryOperation::*;
+
+        match self.op {
+            BitwiseOr => 1,
+            BitwiseAnd => 2,
+            Eq | Neq | Lt | Gt | Lte | Gte => 3,
+            Or => 4,
+            Xor => 5,
+            And => 6,
+            Shl => 7,
+            Shr => 8,
+            Add | Sub => 9,
+            Mul | Div | Rem => 10,
+            Pow => 11,
+            AddWrapped | DivWrapped | Mod | MulWrapped | Nand | Nor | PowWrapped | RemWrapped | ShlWrapped
+            | ShrWrapped | SubWrapped => 20,
+        }
+    }
+
+    pub(crate) fn associativity(&self) -> Associativity {
+        use Associativity::*;
+        use BinaryOperation::*;
+
+        match self.op {
+            Pow => Right,
+            BitwiseOr | BitwiseAnd | Eq | Neq | Lt | Gt | Lte | Gte | Or | Xor | And | Shl | Shr | Add | Sub | Mul
+            | Div | Rem => Left,
+            AddWrapped | DivWrapped | Mod | MulWrapped | Nand | Nor | PowWrapped | RemWrapped | ShlWrapped
+            | ShrWrapped | SubWrapped => None,
+        }
     }
 }
 

compiler/ast/src/expressions/cast.rs

@@ -33,7 +33,12 @@ pub struct CastExpression {
 
 impl fmt::Display for CastExpression {
     fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
-        write!(f, "({} as {})", self.expression, self.type_)
+        if self.expression.precedence() < 12 {
+            write!(f, "({})", self.expression)?;
+        } else {
+            write!(f, "{}", self.expression)?;
+        }
+        write!(f, " as {}", self.type_)
     }
 }
 

compiler/ast/src/expressions/mod.rs

@@ -202,3 +202,27 @@ impl fmt::Display for Expression {
         }
     }
 }
+
+#[derive(Clone, Copy, Eq, PartialEq)]
+pub(crate) enum Associativity {
+    Left,
+    Right,
+    None,
+}
+
+impl Expression {
+    pub(crate) fn precedence(&self) -> u32 {
+        use Expression::*;
+        match self {
+            Binary(e) => e.precedence(),
+            Cast(_) => 12,
+            Ternary(_) => 14,
+            Access(_) | Array(_) | Call(_) | Err(_) | Identifier(_) | Literal(_) | Locator(_) | Struct(_)
+            | Tuple(_) | Unary(_) | Unit(_) => 20,
+        }
+    }
+
+    pub(crate) fn associativity(&self) -> Associativity {
+        if let Expression::Binary(bin) = self { bin.associativity() } else { Associativity::None }
+    }
+}

compiler/ast/src/expressions/struct_init.rs

@@ -17,6 +17,8 @@
 use super::*;
 use leo_span::sym;
 
+use itertools::Itertools as _;
+
 /// An initializer for a single field / variable of a struct initializer expression.
 /// That is, in `Foo { bar: 42, baz }`, this is either `bar: 42`, or `baz`.
 #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
@@ -90,7 +92,15 @@ impl StructExpression {
 
 impl fmt::Display for StructExpression {
     fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
-        write!(f, "{{{}}}", self.members.iter().map(|x| x.to_string()).collect::<Vec<_>>().join(", "))
+        write!(f, "{} {{", self.name)?;
+        if !self.members.is_empty() {
+            write!(f, " ")?;
+        }
+        write!(f, "{}", self.members.iter().format(", "))?;
+        if !self.members.is_empty() {
+            write!(f, " ")?;
+        }
+        write!(f, "}}")
     }
 }
 

compiler/ast/src/expressions/ternary.rs

@@ -33,7 +33,21 @@ pub struct TernaryExpression {
 
 impl fmt::Display for TernaryExpression {
     fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
-        write!(f, "({} ? {} : {})", self.condition, self.if_true, self.if_false)
+        if self.condition.precedence() > 14 {
+            write!(f, "{}", self.condition)?;
+        } else {
+            write!(f, "({})", self.condition)?;
+        }
+
+        write!(f, " ? {} : ", self.if_true)?;
+
+        if self.if_false.precedence() > 14 {
+            write!(f, "{}", self.if_false)?;
+        } else {
+            write!(f, "({})", self.if_false)?;
+        }
+
+        Ok(())
     }
 }
 

compiler/ast/src/expressions/unary.rs

@@ -98,7 +98,12 @@ pub struct UnaryExpression {
 
 impl fmt::Display for UnaryExpression {
     fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
-        write!(f, "({}).{}()", self.receiver, self.op)
+        if self.receiver.precedence() < 20 {
+            write!(f, "({})", self.receiver)?;
+        } else {
+            write!(f, "{}", self.receiver)?;
+        }
+        write!(f, ".{}()", self.op)
     }
 }
 

compiler/ast/src/functions/input.rs

@@ -37,7 +37,11 @@ pub struct Input {
 
 impl Input {
     fn format(&self, f: &mut fmt::Formatter) -> fmt::Result {
-        write!(f, "{} {}: {}", self.mode, self.identifier, self.type_)
+        if self.mode == Mode::None {
+            write!(f, "{}: {}", self.identifier, self.type_)
+        } else {
+            write!(f, "{} {}: {}", self.mode, self.identifier, self.type_)
+        }
     }
 
     pub fn identifier(&self) -> &Identifier {

compiler/ast/src/functions/mod.rs

@@ -32,9 +32,10 @@ pub use output::*;
 pub mod mode;
 pub use mode::*;
 
-use crate::{Block, FunctionStub, Identifier, Node, NodeID, TupleType, Type};
+use crate::{Block, FunctionStub, Identifier, Indent, Node, NodeID, TupleType, Type};
 use leo_span::{Span, Symbol};
 
+use itertools::Itertools as _;
 use serde::{Deserialize, Serialize};
 use std::fmt;
 
@@ -95,28 +96,6 @@ impl Function {
     pub fn name(&self) -> Symbol {
         self.identifier.name
     }
-
-    ///
-    /// Private formatting method used for optimizing [fmt::Debug] and [fmt::Display] implementations.
-    ///
-    fn format(&self, f: &mut fmt::Formatter) -> fmt::Result {
-        match self.variant {
-            Variant::Inline => write!(f, "inline ")?,
-            Variant::Function | Variant::AsyncFunction => write!(f, "function ")?,
-            Variant::Transition | Variant::AsyncTransition => write!(f, "transition ")?,
-        }
-        write!(f, "{}", self.identifier)?;
-
-        let parameters = self.input.iter().map(|x| x.to_string()).collect::<Vec<_>>().join(",");
-        let returns = match self.output.len() {
-            0 => "()".to_string(),
-            1 => self.output[0].to_string(),
-            _ => self.output.iter().map(|x| x.to_string()).collect::<Vec<_>>().join(","),
-        };
-        write!(f, "({parameters}) -> {returns} {}", self.block)?;
-
-        Ok(())
-    }
 }
 
 impl From<FunctionStub> for Function {
@@ -137,13 +116,38 @@ impl From<FunctionStub> for Function {
 
 impl fmt::Debug for Function {
     fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
-        self.format(f)
+        write!(f, "{}", self)
     }
 }
 
 impl fmt::Display for Function {
     fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
-        self.format(f)
+        match self.variant {
+            Variant::Inline => write!(f, "inline ")?,
+            Variant::Function => write!(f, "function ")?,
+            Variant::AsyncFunction => write!(f, "async function ")?,
+            Variant::Transition => write!(f, "transition ")?,
+            Variant::AsyncTransition => write!(f, "asyc transition ")?,
+        }
+        write!(f, "{}({})", self.identifier, self.input.iter().format(", "))?;
+
+        match self.output.len() {
+            0 => {}
+            1 => {
+                if !matches!(self.output[0].type_, Type::Unit) {
+                    write!(f, " -> {}", self.output[0])?;
+                }
+            }
+            _ => {
+                write!(f, " -> {}", self.output.iter().format(", "))?;
+            }
+        }
+
+        writeln!(f, " {{")?;
+        for stmt in self.block.statements.iter() {
+            writeln!(f, "{}{}", Indent(stmt), stmt.semicolon())?;
+        }
+        write!(f, "}}")
     }
 }
 

compiler/ast/src/indent_display.rs

@@ -0,0 +1,55 @@
+// Copyright (C) 2019-2024 Aleo Systems Inc.
+// This file is part of the Leo library.
+
+// The Leo library is free software: you can redistribute it and/or modify
+// it under the terms of the GNU General Public License as published by
+// the Free Software Foundation, either version 3 of the License, or
+// (at your option) any later version.
+
+// The Leo library is distributed in the hope that it will be useful,
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+// GNU General Public License for more details.
+
+// You should have received a copy of the GNU General Public License
+// along with the Leo library. If not, see <https://www.gnu.org/licenses/>.
+
+use std::{fmt, fmt::Write};
+
+/// Implements `Display` by putting 4 spaces in front of each line
+/// of `T`'s output.
+pub struct Indent<T>(pub T);
+
+impl<T: fmt::Display> fmt::Display for Indent<T> {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        write!(IndentWriter { f, new_line: true }, "{}", self.0)
+    }
+}
+
+const SPACES: &str = "    ";
+
+struct IndentWriter<'a, 'b> {
+    new_line: bool,
+    f: &'b mut fmt::Formatter<'a>,
+}
+
+impl Write for IndentWriter<'_, '_> {
+    fn write_str(&mut self, s: &str) -> fmt::Result {
+        let mut iter = s.lines().peekable();
+
+        while let Some(line) = iter.next() {
+            if self.new_line {
+                self.f.write_str(SPACES)?;
+            }
+            self.f.write_str(line)?;
+            if iter.peek().is_some() || s.ends_with('\n') {
+                self.f.write_str("\n")?;
+                self.new_line = true;
+            } else {
+                self.new_line = false;
+            }
+        }
+
+        Ok(())
+    }
+}

compiler/ast/src/lib.rs

@@ -40,6 +40,9 @@ pub use self::functions::*;
 pub mod groups;
 pub use self::groups::*;
 
+mod indent_display;
+use indent_display::*;
+
 pub mod mapping;
 pub use self::mapping::*;
 

compiler/ast/src/program/mod.rs

@@ -43,15 +43,13 @@ pub struct Program {
 impl fmt::Display for Program {
     fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
         for (id, _import) in self.imports.iter() {
-            writeln!(f, "import {id}.leo;")?;
+            writeln!(f, "import {id}.aleo;")?;
         }
         for (_, stub) in self.stubs.iter() {
-            stub.fmt(f)?;
-            writeln!(f,)?;
+            writeln!(f, "{}", stub)?;
         }
         for (_, program_scope) in self.program_scopes.iter() {
-            program_scope.fmt(f)?;
-            writeln!(f,)?;
+            writeln!(f, "{}", program_scope)?;
         }
         Ok(())
     }

compiler/ast/src/program/program_scope.rs

@@ -16,7 +16,7 @@
 
 //! A Leo program scope consists of struct, function, and mapping definitions.
 
-use crate::{Composite, ConstDeclaration, Function, Mapping, ProgramId, Stub};
+use crate::{Composite, ConstDeclaration, Function, Indent, Mapping, ProgramId, Stub};
 
 use leo_span::{Span, Symbol};
 use serde::{Deserialize, Serialize};
@@ -60,17 +60,17 @@ impl fmt::Display for ProgramScope {
     fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
         writeln!(f, "program {} {{", self.program_id)?;
         for (_, const_decl) in self.consts.iter() {
-            writeln!(f, "    {const_decl}")?;
+            writeln!(f, "{};", Indent(const_decl))?;
         }
         for (_, struct_) in self.structs.iter() {
-            writeln!(f, "    {struct_}")?;
+            writeln!(f, "{}", Indent(struct_))?;
         }
         for (_, mapping) in self.mappings.iter() {
-            writeln!(f, "    {mapping}")?;
+            writeln!(f, "{};", Indent(mapping))?;
         }
         for (_, function) in self.functions.iter() {
-            writeln!(f, "    {function}")?;
+            writeln!(f, "{}", Indent(function))?;
         }
-        Ok(())
+        write!(f, "}}")
     }
 }

compiler/ast/src/statement/assert.rs

@@ -46,9 +46,9 @@ pub struct AssertStatement {
 impl fmt::Display for AssertStatement {
     fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
         match self.variant {
-            AssertVariant::Assert(ref expr) => write!(f, "assert({expr});"),
-            AssertVariant::AssertEq(ref expr1, ref expr2) => write!(f, "assert_eq({expr1}, {expr2});"),
-            AssertVariant::AssertNeq(ref expr1, ref expr2) => write!(f, "assert_neq({expr1}, {expr2});"),
+            AssertVariant::Assert(ref expr) => write!(f, "assert({expr})"),
+            AssertVariant::AssertEq(ref expr1, ref expr2) => write!(f, "assert_eq({expr1}, {expr2})"),
+            AssertVariant::AssertNeq(ref expr1, ref expr2) => write!(f, "assert_neq({expr1}, {expr2})"),
         }
     }
 }

compiler/ast/src/statement/block.rs

@@ -14,7 +14,7 @@
 // You should have received a copy of the GNU General Public License
 // along with the Leo library. If not, see <https://www.gnu.org/licenses/>.
 
-use crate::{Node, NodeID, Statement};
+use crate::{Indent, Node, NodeID, Statement};
 use leo_span::Span;
 
 use serde::{Deserialize, Serialize};
@@ -34,10 +34,8 @@ pub struct Block {
 impl fmt::Display for Block {
     fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
         writeln!(f, "{{")?;
-        if self.statements.is_empty() {
-            writeln!(f, "\t")?;
-        } else {
-            self.statements.iter().try_for_each(|statement| writeln!(f, "\t{statement};"))?;
+        for stmt in self.statements.iter() {
+            writeln!(f, "{}{}", Indent(stmt), stmt.semicolon())?;
         }
         write!(f, "}}")
     }

compiler/ast/src/statement/conditional.rs

@@ -14,7 +14,7 @@
 // You should have received a copy of the GNU General Public License
 // along with the Leo library. If not, see <https://www.gnu.org/licenses/>.
 
-use crate::{Block, Expression, Node, NodeID, Statement};
+use crate::{Block, Expression, Indent, Node, NodeID, Statement};
 use leo_span::Span;
 
 use serde::{Deserialize, Serialize};
@@ -37,11 +37,25 @@ pub struct ConditionalStatement {
 
 impl fmt::Display for ConditionalStatement {
     fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
-        write!(f, "if ({}) {}", self.condition, self.then)?;
-        match self.otherwise.as_ref() {
-            Some(n_or_e) => write!(f, " else {n_or_e}"),
-            None => write!(f, ""),
+        writeln!(f, "if {} {{", self.condition)?;
+        for stmt in self.then.statements.iter() {
+            writeln!(f, "{}{}", Indent(stmt), stmt.semicolon())?;
         }
+        match self.otherwise.as_deref() {
+            None => write!(f, "}}")?,
+            Some(Statement::Block(block)) => {
+                writeln!(f, "}} else {{")?;
+                for stmt in block.statements.iter() {
+                    writeln!(f, "{}{}", Indent(stmt), stmt.semicolon())?;
+                }
+                write!(f, "}}")?;
+            }
+            Some(Statement::Conditional(cond)) => {
+                write!(f, "}} else {cond}")?;
+            }
+            Some(_) => panic!("`otherwise` of a `ConditionalStatement` must be a block or conditional."),
+        }
+        Ok(())
     }
 }
 

compiler/ast/src/statement/iteration.rs

@@ -14,7 +14,7 @@
 // You should have received a copy of the GNU General Public License
 // along with the Leo library. If not, see <https://www.gnu.org/licenses/>.
 
-use crate::{Block, Expression, Identifier, Node, NodeID, Type, Value};
+use crate::{Block, Expression, Identifier, Indent, Node, NodeID, Type, Value};
 
 use leo_span::Span;
 
@@ -52,7 +52,11 @@ pub struct IterationStatement {
 impl fmt::Display for IterationStatement {
     fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
         let eq = if self.inclusive { "=" } else { "" };
-        write!(f, "for {} in {}..{eq}{} {}", self.variable, self.start, self.stop, self.block)
+        writeln!(f, "for {}: {} in {}..{eq}{} {{", self.variable, self.type_, self.start, self.stop)?;
+        for stmt in self.block.statements.iter() {
+            writeln!(f, "{}{}", Indent(stmt), stmt.semicolon())?;
+        }
+        writeln!(f, "}}")
     }
 }
 

compiler/ast/src/statement/mod.rs

@@ -81,6 +81,12 @@ impl Statement {
     pub fn dummy(span: Span, id: NodeID) -> Self {
         Self::Block(Block { statements: Vec::new(), span, id })
     }
+
+    pub(crate) fn semicolon(&self) -> &'static str {
+        use Statement::*;
+
+        if matches!(self, Block(..) | Conditional(..) | Iteration(..)) { "" } else { ";" }
+    }
 }
 
 impl fmt::Display for Statement {

compiler/ast/src/struct/mod.rs

@@ -17,7 +17,7 @@
 pub mod member;
 pub use member::*;
 
-use crate::{Identifier, Mode, Node, NodeID, Type};
+use crate::{Identifier, Indent, Mode, Node, NodeID, Type};
 use leo_span::{Span, Symbol};
 
 use itertools::Itertools;
@@ -137,9 +137,9 @@ impl fmt::Display for Composite {
         f.write_str(if self.is_record { "record" } else { "struct" })?;
         writeln!(f, " {} {{ ", self.identifier)?;
         for field in self.members.iter() {
-            writeln!(f, "        {field}")?;
+            writeln!(f, "{},", Indent(field))?;
         }
-        write!(f, "    }}")
+        write!(f, "}}")
     }
 }
 

compiler/ast/src/stub/mod.rs

@@ -19,7 +19,7 @@
 pub mod function_stub;
 pub use function_stub::*;
 
-use crate::{Composite, ConstDeclaration, Identifier, Mapping, NodeID, ProgramId};
+use crate::{Composite, ConstDeclaration, Identifier, Indent, Mapping, NodeID, ProgramId};
 use leo_span::{Span, Symbol};
 use serde::{Deserialize, Serialize};
 use std::fmt;
@@ -65,18 +65,17 @@ impl fmt::Display for Stub {
     fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
         writeln!(f, "stub {} {{", self.stub_id)?;
         for import in self.imports.iter() {
-            writeln!(f, "    import {import}")?;
+            writeln!(f, "    import {import};")?;
         }
         for (_, mapping) in self.mappings.iter() {
-            writeln!(f, "    {mapping}")?;
+            writeln!(f, "{};", Indent(mapping))?;
         }
         for (_, struct_) in self.structs.iter() {
-            writeln!(f, "    {struct_}")?;
+            writeln!(f, "{}", Indent(struct_))?;
         }
         for (_, function) in self.functions.iter() {
-            writeln!(f, "    {function}")?;
+            writeln!(f, "{}", Indent(function))?;
         }
-        writeln!(f, "}}")?;
-        Ok(())
+        write!(f, "}}")
     }
 }

tests/expectations/compiler/finalize/unknown_mapping_operation_fail.out

@@ -16,7 +16,7 @@ Error [ETYC0372009]: Mapping::has_key is not a valid core function.
      |
   12 |          let has_key: bool = Mapping::has_key(account, receiver);
      |                              ^^^^^^^
-Error [ETYC0372014]: Mapping::has_key is not a valid core function call.
+Error [ETYC0372014]: Mapping::has_key(account, receiver) is not a valid core function call.
     --> compiler-test:12:30
      |
   12 |          let has_key: bool = Mapping::has_key(account, receiver);