convert enum tag unions to mono

compiler/mono/src/expr.rs

@@ -353,6 +353,25 @@ fn from_can<'a>(
                 elems: elems.into_bump_slice(),
             }
         }
+
+        Tag {
+            variant_var,
+            ext_var,
+            name,
+            arguments,
+        } => {
+            let subs = &env.subs;
+            let arena = env.arena;
+
+            match Layout::from_var(arena, variant_var, subs) {
+                Ok(Layout::Builtin(Builtin::Bool(_smaller, larger))) => Expr::Bool(name == larger),
+                Ok(Layout::Builtin(Builtin::Byte(tags))) => match tags.get(&name) {
+                    Some(v) => Expr::Byte(*v),
+                    None => panic!("Tag name is not part of the type"),
+                },
+                _ => panic!(),
+            }
+        }
         other => panic!("TODO convert canonicalized {:?} to ll::Expr", other),
     }
 }

compiler/mono/src/layout.rs

@@ -19,6 +19,8 @@ pub enum Layout<'a> {
 pub enum Builtin<'a> {
     Int64,
     Float64,
+    Bool(TagName, TagName),
+    Byte(MutMap<TagName, u8>),
     Str,
     Map(&'a Layout<'a>, &'a Layout<'a>),
     Set(&'a Layout<'a>),
@@ -81,6 +83,8 @@ impl<'a> Layout<'a> {
 impl<'a> Builtin<'a> {
     const I64_SIZE: u32 = std::mem::size_of::<i64>() as u32;
     const F64_SIZE: u32 = std::mem::size_of::<f64>() as u32;
+    const BOOL_SIZE: u32 = std::mem::size_of::<bool>() as u32;
+    const BYTE_SIZE: u32 = std::mem::size_of::<u8>() as u32;
 
     /// Number of machine words in an empty one of these
     const STR_WORDS: u32 = 3;
@@ -94,6 +98,8 @@ impl<'a> Builtin<'a> {
         match self {
             Int64 => Builtin::I64_SIZE,
             Float64 => Builtin::F64_SIZE,
+            Bool(_, _) => Builtin::BOOL_SIZE,
+            Byte(_) => Builtin::BYTE_SIZE,
             Str => Builtin::STR_WORDS * pointer_size,
             Map(_, _) => Builtin::MAP_WORDS * pointer_size,
             Set(_) => Builtin::SET_WORDS * pointer_size,
@@ -243,10 +249,37 @@ fn layout_from_flat_type<'a>(
                     }
                 }
                 _ => {
+                    // Check if we can turn this tag union into an enum
+                    // TODO rather than the arguments being empty, check whether their layout has size 0.
+                    if tags.len() <= 256 && tags.iter().all(|(_, args)| args.is_empty()) {
+                        if tags.len() <= 2 {
+                            // Up to 2 enum tags can be stored (in theory) in one bit
+                            let mut it = tags.keys();
+                            let a: TagName = it.next().unwrap().clone();
+                            let b: TagName = it.next().unwrap().clone();
+
+                            if a < b {
+                                Ok(Layout::Builtin(Builtin::Bool(a, b)))
+                            } else {
+                                Ok(Layout::Builtin(Builtin::Bool(b, a)))
+                            }
+                        } else {
+                            // up to 256 enum tags can be stored in a byte
+                            let mut counter = 0u8;
+                            let mut tag_to_u8 = MutMap::default();
+
+                            for (name, _) in tags {
+                                tag_to_u8.insert(name, counter);
+                                counter += 1;
+                            }
+                            Ok(Layout::Builtin(Builtin::Byte(tag_to_u8)))
+                        }
+                    } else {
                         panic!("TODO handle a tag union with mutliple tags: {:?}", tags);
                     }
                 }
             }
+        }
         RecursiveTagUnion(_, _, _) => {
             panic!("TODO make Layout for non-empty Tag Union");
         }
@@ -301,13 +334,15 @@ fn flatten_union(
 
     match subs.get_without_compacting(ext_var).content {
         Structure(EmptyTagUnion) => (),
-        Structure(TagUnion(new_tags, new_ext_var)) => {
+        Structure(TagUnion(new_tags, new_ext_var))
+        | Structure(RecursiveTagUnion(_, new_tags, new_ext_var)) => {
             for (tag_name, vars) in new_tags {
                 tags.insert(tag_name, vars);
             }
 
             flatten_union(tags, new_ext_var, subs)
         }
+        Alias(_, _, actual) => flatten_union(tags, actual, subs),
         invalid => {
             panic!("Compiler error: flatten_union got an ext_var in a tag union that wasn't itself a tag union; instead, it was: {:?}", invalid);
         }
@@ -329,6 +364,7 @@ fn flatten_record(fields: &mut MutMap<Lowercase, Variable>, ext_var: Variable, s
 
             flatten_record(fields, new_ext_var, subs)
         }
+        Alias(_, _, actual) => flatten_record(fields, actual, subs),
         invalid => {
             panic!("Compiler error: flatten_record encountered an ext_var in a record that wasn't itself a record; instead, it was: {:?}", invalid);
         }

compiler/mono/tests/test_mono.rs

@@ -11,10 +11,11 @@ mod helpers;
 // Test monomorphization
 #[cfg(test)]
 mod test_mono {
-    use crate::helpers::{can_expr, infer_expr, CanExprOut};
+    use crate::helpers::{can_expr, infer_expr, test_home, CanExprOut};
     use bumpalo::Bump;
     use roc_collections::all::MutMap;
-    use roc_module::symbol::Symbol;
+    use roc_module::ident::TagName::*;
+    use roc_module::symbol::{Interns, Symbol};
     use roc_mono::expr::Expr::{self, *};
     use roc_mono::layout::{Builtin, Layout};
     use roc_types::subs::Subs;
@@ -22,6 +23,13 @@ mod test_mono {
     // HELPERS
 
     fn compiles_to(src: &str, expected: Expr<'_>) {
+        compiles_to_with_interns(src, |_| expected)
+    }
+
+    fn compiles_to_with_interns<'a, F>(src: &str, get_expected: F)
+    where
+        F: FnOnce(Interns) -> Expr<'a>,
+    {
         let arena = Bump::new();
         let CanExprOut {
             loc_expr,
@@ -50,7 +58,7 @@ mod test_mono {
             &mut ident_ids,
         );
 
-        assert_eq!(mono_expr, expected);
+        assert_eq!(mono_expr, get_expected(interns));
     }
 
     #[test]
@@ -63,10 +71,62 @@ mod test_mono {
         compiles_to("0.5", Float(0.5));
     }
 
+    //    #[test]
+    //    fn bool_literal() {
+    //        compiles_to_with_interns(
+    //            r#"
+    //            x : Bool
+    //            x = True
+    //
+    //            x
+    //            "#,
+    //            |interns| {
+    //                let home = test_home();
+    //                let var_x = interns.symbol(home, "x".into());
+    //
+    //                let stores = [(
+    //                    var_x,
+    //                    Layout::Builtin(Builtin::Bool(Global("False".into()), Global("True".into()))),
+    //                    Bool(true),
+    //                )];
+    //
+    //                let load = Load(var_x);
+    //
+    //                Store(&stores, &load)
+    //            },
+    //        );
+    //    }
+
+    //    #[test]
+    //    fn two_element_enum() {
+    //        compiles_to(
+    //            r#"
+    //            x : [ Yes, No ]
+    //            x = No
+    //
+    //            x
+    //            "#,
+    //            Int(32),
+    //        );
+    //    }
+    //
+    //    #[test]
+    //    fn three_element_enum() {
+    //        compiles_to(
+    //            r#"
+    //            # this test is brought to you by fruits.com!
+    //            x : [ Apple, Orange, Banana ]
+    //            x = Orange
+    //
+    //            x
+    //            "#,
+    //            Int(32),
+    //        );
+    //    }
+
     #[test]
     fn set_unique_int_list() {
-        compiles_to(
+        compiles_to("List.getUnsafe (List.set [ 12, 9, 7, 3 ] 1 42) 1", {
-            "List.getUnsafe (List.set [ 12, 9, 7, 3 ] 1 42) 1",
             CallByName(
                 Symbol::LIST_GET_UNSAFE,
                 &vec![
@@ -91,7 +151,7 @@ mod test_mono {
                     ),
                     (Int(1), Layout::Builtin(Builtin::Int64)),
                 ],
-            ),
+            )
-        );
+        });
     }
 }