improved debug printing of layouts

compiler/mono/src/ir.rs

@@ -136,8 +136,15 @@ pub enum SelfRecursive {
     SelfRecursive(JoinPointId),
 }
 
+#[derive(Clone, Copy, Debug, PartialEq)]
+pub enum Parens {
+    NotNeeded,
+    InTypeParam,
+    InFunction,
+}
+
 impl<'a> Proc<'a> {
-    pub fn to_doc<'b, D, A>(&'b self, alloc: &'b D, _parens: bool) -> DocBuilder<'b, D, A>
+    pub fn to_doc<'b, D, A>(&'b self, alloc: &'b D, _parens: Parens) -> DocBuilder<'b, D, A>
     where
         D: DocAllocator<'b, A>,
         D::Doc: Clone,
@@ -148,6 +155,21 @@ impl<'a> Proc<'a> {
             .iter()
             .map(|(_, symbol)| symbol_to_doc(alloc, *symbol));
 
+        if PRETTY_PRINT_IR_SYMBOLS {
+            alloc
+                .text("procedure : ")
+                .append(symbol_to_doc(alloc, self.name))
+                .append(" ")
+                .append(self.ret_layout.to_doc(alloc, Parens::NotNeeded))
+                .append(alloc.hardline())
+                .append(alloc.text("procedure = "))
+                .append(symbol_to_doc(alloc, self.name))
+                .append(" (")
+                .append(alloc.intersperse(args_doc, ", "))
+                .append("):")
+                .append(alloc.hardline())
+                .append(self.body.to_doc(alloc).indent(4))
+        } else {
             alloc
                 .text("procedure ")
                 .append(symbol_to_doc(alloc, self.name))
@@ -157,11 +179,12 @@ impl<'a> Proc<'a> {
                 .append(alloc.hardline())
                 .append(self.body.to_doc(alloc).indent(4))
         }
+    }
 
     pub fn to_pretty(&self, width: usize) -> String {
         let allocator = BoxAllocator;
         let mut w = std::vec::Vec::new();
-        self.to_doc::<_, ()>(&allocator, false)
+        self.to_doc::<_, ()>(&allocator, Parens::NotNeeded)
             .1
             .render(width, &mut w)
             .unwrap();
@@ -2460,8 +2483,7 @@ fn specialize_naked_symbol<'a>(
                 match hole {
                     Stmt::Jump(_, _) => todo!("not sure what to do in this case yet"),
                     _ => {
-                        let expr =
-                            call_by_pointer(env, procs, layout_cache, symbol, layout.clone());
+                        let expr = call_by_pointer(env, procs, symbol, layout.clone());
                         let new_symbol = env.unique_symbol();
                         return Stmt::Let(
                             new_symbol,
@@ -3549,7 +3571,7 @@ pub fn with_hole<'a>(
                     // TODO should the let have layout Pointer?
                     Stmt::Let(
                         assigned,
-                        call_by_pointer(env, procs, layout_cache, name, layout.clone()),
+                        call_by_pointer(env, procs, name, layout.clone()),
                         layout,
                         hole,
                     )
@@ -3746,13 +3768,7 @@ pub fn with_hole<'a>(
                         }
                     }
 
-                    let expr = call_by_pointer(
-                        env,
-                        procs,
-                        layout_cache,
-                        name,
-                        function_ptr_layout.clone(),
-                    );
+                    let expr = call_by_pointer(env, procs, name, function_ptr_layout.clone());
 
                     stmt = Stmt::Let(
                         function_pointer,
@@ -3778,7 +3794,7 @@ pub fn with_hole<'a>(
                             // TODO should the let have layout Pointer?
                             Stmt::Let(
                                 assigned,
-                                call_by_pointer(env, procs, layout_cache, name, layout.clone()),
+                                call_by_pointer(env, procs, name, layout.clone()),
                                 layout,
                                 hole,
                             )
@@ -5359,7 +5375,7 @@ fn handle_variable_aliasing<'a>(
             .from_var(env.arena, variable, env.subs)
             .unwrap();
 
-        let expr = call_by_pointer(env, procs, layout_cache, right, layout.clone());
+        let expr = call_by_pointer(env, procs, right, layout.clone());
         Stmt::Let(left, expr, layout, env.arena.alloc(result))
     } else {
         substitute_in_exprs(env.arena, &mut result, left, right);
@@ -5407,7 +5423,7 @@ fn reuse_function_symbol<'a>(
 
                     // an imported symbol is always a function pointer:
                     // either it's a function, or a top-level 0-argument thunk
-                    let expr = call_by_pointer(env, procs, layout_cache, original, layout.clone());
+                    let expr = call_by_pointer(env, procs, original, layout.clone());
                     return Stmt::Let(symbol, expr, layout, env.arena.alloc(result));
                 }
                 _ => {
@@ -5505,13 +5521,7 @@ fn reuse_function_symbol<'a>(
                         }
                     }
 
-                    let expr = call_by_pointer(
-                        env,
-                        procs,
-                        layout_cache,
-                        original,
-                        function_ptr_layout.clone(),
-                    );
+                    let expr = call_by_pointer(env, procs, original, function_ptr_layout.clone());
 
                     stmt = Stmt::Let(
                         function_pointer,
@@ -5533,7 +5543,7 @@ fn reuse_function_symbol<'a>(
 
                     Stmt::Let(
                         symbol,
-                        call_by_pointer(env, procs, layout_cache, original, layout.clone()),
+                        call_by_pointer(env, procs, original, layout.clone()),
                         layout,
                         env.arena.alloc(result),
                     )
@@ -5613,7 +5623,6 @@ where
 fn call_by_pointer<'a>(
     env: &mut Env<'a, '_>,
     procs: &mut Procs<'a>,
-    layout_cache: &mut LayoutCache<'a>,
     symbol: Symbol,
     layout: Layout<'a>,
 ) -> Expr<'a> {

compiler/mono/src/layout.rs

@@ -1,3 +1,4 @@
+use crate::ir::Parens;
 use bumpalo::collections::Vec;
 use bumpalo::Bump;
 use roc_collections::all::{default_hasher, MutMap, MutSet};
@@ -6,6 +7,7 @@ use roc_module::symbol::{Interns, Symbol};
 use roc_types::subs::{Content, FlatType, Subs, Variable};
 use roc_types::types::RecordField;
 use std::collections::HashMap;
+use ven_pretty::{DocAllocator, DocBuilder};
 
 pub const MAX_ENUM_SIZE: usize = (std::mem::size_of::<u8>() * 8) as usize;
 const GENERATE_NULLABLE: bool = true;
@@ -63,6 +65,34 @@ pub enum UnionLayout<'a> {
     },
 }
 
+impl<'a> UnionLayout<'a> {
+    pub fn to_doc<'b, D, A>(&'b self, alloc: &'b D, _parens: Parens) -> DocBuilder<'b, D, A>
+    where
+        D: DocAllocator<'b, A>,
+        D::Doc: Clone,
+        A: Clone,
+    {
+        use UnionLayout::*;
+
+        match self {
+            NonRecursive(tags) => {
+                let tags_doc = tags.iter().map(|fields| {
+                    alloc.text("C ").append(alloc.intersperse(
+                        fields.iter().map(|x| x.to_doc(alloc, Parens::InTypeParam)),
+                        ", ",
+                    ))
+                });
+
+                alloc
+                    .text("[")
+                    .append(alloc.intersperse(tags_doc, ", "))
+                    .append(alloc.text("]"))
+            }
+            _ => alloc.text("TODO"),
+        }
+    }
+}
+
 #[derive(Clone, Debug, PartialEq, Eq, Hash)]
 pub struct ClosureLayout<'a> {
     /// the layout that this specific closure captures
@@ -654,6 +684,51 @@ impl<'a> Layout<'a> {
             FunctionPointer(_, _) | Pointer(_) => false,
         }
     }
+
+    pub fn to_doc<'b, D, A>(&'b self, alloc: &'b D, parens: Parens) -> DocBuilder<'b, D, A>
+    where
+        D: DocAllocator<'b, A>,
+        D::Doc: Clone,
+        A: Clone,
+    {
+        use Layout::*;
+
+        match self {
+            Builtin(builtin) => builtin.to_doc(alloc, parens),
+            PhantomEmptyStruct => alloc.text("{}"),
+            Struct(fields) => {
+                let fields_doc = fields.iter().map(|x| x.to_doc(alloc, parens));
+
+                alloc
+                    .text("{")
+                    .append(alloc.intersperse(fields_doc, ", "))
+                    .append(alloc.text("}"))
+            }
+            Union(union_layout) => union_layout.to_doc(alloc, parens),
+            RecursivePointer => alloc.text("*self"),
+            FunctionPointer(args, result) => {
+                let args_doc = args.iter().map(|x| x.to_doc(alloc, Parens::InFunction));
+
+                alloc
+                    .intersperse(args_doc, ", ")
+                    .append(alloc.text(" -> "))
+                    .append(result.to_doc(alloc, Parens::InFunction))
+            }
+            Closure(args, closure_layout, result) => {
+                let args_doc = args.iter().map(|x| x.to_doc(alloc, Parens::InFunction));
+
+                let bom = closure_layout.layout.to_doc(alloc, Parens::NotNeeded);
+
+                alloc
+                    .intersperse(args_doc, ", ")
+                    .append(alloc.text(" {| "))
+                    .append(bom)
+                    .append(" |} -> ")
+                    .append(result.to_doc(alloc, Parens::InFunction))
+            }
+            Pointer(_) => todo!(),
+        }
+    }
 }
 
 /// Avoid recomputing Layout from Variable multiple times.
@@ -878,6 +953,47 @@ impl<'a> Builtin<'a> {
             Str | Dict(_, _) | Set(_) => true,
         }
     }
+
+    pub fn to_doc<'b, D, A>(&'b self, alloc: &'b D, _parens: Parens) -> DocBuilder<'b, D, A>
+    where
+        D: DocAllocator<'b, A>,
+        D::Doc: Clone,
+        A: Clone,
+    {
+        use Builtin::*;
+
+        match self {
+            Int128 => alloc.text("Int128"),
+            Int64 => alloc.text("Int64"),
+            Int32 => alloc.text("Int32"),
+            Int16 => alloc.text("Int16"),
+            Int8 => alloc.text("Int8"),
+            Int1 => alloc.text("Int1"),
+            Usize => alloc.text("Usize"),
+            Float128 => alloc.text("Float128"),
+            Float64 => alloc.text("Float64"),
+            Float32 => alloc.text("Float32"),
+            Float16 => alloc.text("Float16"),
+
+            EmptyStr => alloc.text("EmptyStr"),
+            EmptyList => alloc.text("EmptyList"),
+            EmptyDict => alloc.text("EmptyDict"),
+            EmptySet => alloc.text("EmptySet"),
+
+            Str => alloc.text("Str"),
+            List(_, layout) => alloc
+                .text("List ")
+                .append(layout.to_doc(alloc, Parens::InTypeParam)),
+            Set(layout) => alloc
+                .text("Set ")
+                .append(layout.to_doc(alloc, Parens::InTypeParam)),
+            Dict(key_layout, value_layout) => alloc
+                .text("Dict ")
+                .append(key_layout.to_doc(alloc, Parens::InTypeParam))
+                .append(" ")
+                .append(value_layout.to_doc(alloc, Parens::InTypeParam)),
+        }
+    }
 }
 
 fn layout_from_flat_type<'a>(