use a custom type to store Union variants

compiler/mono/src/expr.rs

@@ -1,4 +1,4 @@
-use crate::layout::{Builtin, Layout, MAX_ENUM_SIZE};
+use crate::layout::{Builtin, Layout};
 use crate::pattern::{Ctor, Guard};
 use bumpalo::collections::Vec;
 use bumpalo::Bump;
@@ -682,58 +682,57 @@ fn from_can<'a>(
             arguments: args,
             ..
         } => {
+            use crate::layout::UnionVariant::*;
             let arena = env.arena;
 
-            let (is_enum_candidate, sorted_tag_layouts) = crate::layout::union_sorted_tags(
+            let variant = crate::layout::union_sorted_tags(
                 env.arena,
                 variant_var,
                 env.subs,
                 env.pointer_size,
             );
 
-            let union_size = sorted_tag_layouts.len() as u8;
+            match variant {
+                Never => panic!("TODO cannot create a value of a type with no constructors"),
+                Unit => Expr::Struct(&[]),
+                BoolUnion { ttrue, .. } => Expr::Bool(tag_name == ttrue),
+                ByteUnion(tag_names) => {
+                    let tag_id = tag_names
+                        .iter()
+                        .position(|key| key == &tag_name)
+                        .expect("tag must be in its own type");
 
-            let (tag_id, (_, argument_layouts)) = sorted_tag_layouts
+                    Expr::Byte(tag_id as u8)
-                .iter()
+                }
-                .enumerate()
+                Unwrapped(field_layouts) => {
-                .find(|(_, (key, _))| key == &tag_name)
+                    let field_exprs = args
-                .expect("tag must be in its own type");
+                        .into_iter()
+                        .map(|(_, arg)| from_can(env, arg.value, procs, None));
+
+                    let mut field_tuples = Vec::with_capacity_in(field_layouts.len(), arena);
+
+                    for tuple in field_exprs.zip(field_layouts.into_iter()) {
+                        field_tuples.push(tuple)
+                    }
+
+                    Expr::Struct(field_tuples.into_bump_slice())
+                }
+                Wrapped(sorted_tag_layouts) => {
+                    let union_size = sorted_tag_layouts.len() as u8;
+                    let (tag_id, (_, argument_layouts)) = sorted_tag_layouts
+                        .iter()
+                        .enumerate()
+                        .find(|(_, (key, _))| key == &tag_name)
+                        .expect("tag must be in its own type");
 
-            match sorted_tag_layouts.len() {
-                2 if is_enum_candidate => Expr::Bool(tag_id != 0),
-                3..=MAX_ENUM_SIZE if is_enum_candidate => Expr::Byte(tag_id as u8),
-                len => {
                     let mut arguments = Vec::with_capacity_in(args.len(), arena);
 
-                    let is_unwrapped = len == 1;
+                    let it = std::iter::once(Expr::Int(tag_id as i64)).chain(
+                        args.into_iter()
+                            .map(|(_, arg)| from_can(env, arg.value, procs, None)),
+                    );
 
-                    // when the tag is not unwrapped, the layout will contain a slot for the tag
+                    for (arg_layout, arg_expr) in argument_layouts.iter().zip(it) {
-                    // discriminant, but it's not yet part of the argument expressions. So we need
-                    // to conditionally insert it. We don't want to collect() into a data
-                    // structure, and thus must conditionally pick one of two iterators. that's a
-                    // bit tricky, as you can see below. Based on
-                    // https://stackoverflow.com/questions/29760668/conditionally-iterate-over-one-of-several-possible-iterators
-                    let argument_exprs = {
-                        let mut iter_1 = None;
-                        let mut iter_2 = None;
-
-                        let shared = args
-                            .into_iter()
-                            .map(|(_, arg)| from_can(env, arg.value, procs, None));
-
-                        if is_unwrapped {
-                            iter_1 = Some(shared);
-                        } else {
-                            iter_2 = Some(std::iter::once(Expr::Int(tag_id as i64)).chain(shared));
-                        }
-
-                        iter_1
-                            .into_iter()
-                            .flatten()
-                            .chain(iter_2.into_iter().flatten())
-                    };
-
-                    for (arg_layout, arg_expr) in argument_layouts.iter().zip(argument_exprs) {
                         arguments.push((arg_expr, arg_layout.clone()));
                     }
 

compiler/mono/src/layout.rs

@@ -328,32 +328,94 @@ pub fn record_fields_btree<'a>(
     }
 }
 
+pub enum UnionVariant<'a> {
+    Never,
+    Unit,
+    BoolUnion { ttrue: TagName, ffalse: TagName },
+    ByteUnion(Vec<'a, TagName>),
+    Unwrapped(Vec<'a, Layout<'a>>),
+    Wrapped(Vec<'a, (TagName, &'a [Layout<'a>])>),
+}
+
 pub fn union_sorted_tags<'a>(
     arena: &'a Bump,
     var: Variable,
     subs: &Subs,
     pointer_size: u32,
-) -> (bool, Vec<'a, (TagName, &'a [Layout<'a>])>) {
+) -> UnionVariant<'a> {
     let mut tags_vec = std::vec::Vec::new();
     match roc_types::pretty_print::chase_ext_tag_union(subs, var, &mut tags_vec) {
         Ok(()) | Err((_, Content::FlexVar(_))) => {
-            // for this union be be an enum, none of the tags may have any arguments
+            union_sorted_tags_help(arena, tags_vec, subs, pointer_size)
-            let is_enum_candidate =
+        }
-                tags_vec.len() <= MAX_ENUM_SIZE && tags_vec.iter().all(|(_, args)| args.is_empty());
+        Err(other) => panic!("invalid content in record variable: {:?}", other),
+    }
+}
 
-            let is_unwrapped = tags_vec.len() == 1;
+fn union_sorted_tags_help<'a>(
-            // collect into btreemap to sort
+    arena: &'a Bump,
-            tags_vec.sort();
+    mut tags_vec: std::vec::Vec<(TagName, std::vec::Vec<Variable>)>,
+    subs: &Subs,
+    pointer_size: u32,
+) -> UnionVariant<'a> {
+    // for this union be be an enum, none of the tags may have any arguments
+    let has_no_arguments = tags_vec.iter().all(|(_, args)| args.is_empty());
 
+    // sort up-front, make sure the ordering stays intact!
+    tags_vec.sort();
+
+    match tags_vec.len() {
+        0 => {
+            // trying to instantiate a type with no values
+            UnionVariant::Never
+        }
+        1 if has_no_arguments => {
+            // a unit type
+            UnionVariant::Unit
+        }
+        2 if has_no_arguments => {
+            // type can be stored in a boolean
+
+            // tags_vec is sorted,
+            let ttrue = tags_vec.remove(1).0;
+            let ffalse = tags_vec.remove(0).0;
+
+            UnionVariant::BoolUnion { ffalse, ttrue }
+        }
+        3..=MAX_ENUM_SIZE if has_no_arguments => {
+            // type can be stored in a byte
+            // needs the sorted tag names to determine the tag_id
+            let mut tag_names = Vec::with_capacity_in(tags_vec.len(), arena);
+
+            for (label, _) in tags_vec {
+                tag_names.push(label);
+            }
+
+            UnionVariant::ByteUnion(tag_names)
+        }
+        1 => {
+            // just one tag in the union (but with arguments) can be a struct
+            let mut layouts = Vec::with_capacity_in(tags_vec.len(), arena);
+
+            let arguments = tags_vec.remove(0).1;
+
+            for var in arguments.iter() {
+                let layout = Layout::from_var(arena, *var, subs, pointer_size)
+                    .expect("invalid layout from var");
+                layouts.push(layout);
+            }
+            UnionVariant::Unwrapped(layouts)
+        }
+        _ => {
+            // default path
             let mut result = Vec::with_capacity_in(tags_vec.len(), arena);
-            for (tag_name, arguments) in tags_vec {
-                let mut arg_layouts =
-                    Vec::with_capacity_in(arguments.len() + !is_unwrapped as usize, arena);
 
-                // if not unwrapped, add a slot for the tag discriminant
+            for (tag_name, arguments) in tags_vec {
-                if !is_unwrapped && !is_enum_candidate {
+                // resverse space for the tag discriminant
-                    arg_layouts.push(Layout::Builtin(Builtin::Int64))
+                let mut arg_layouts = Vec::with_capacity_in(arguments.len() + 1, arena);
-                }
+
+                // add the tag discriminant
+                arg_layouts.push(Layout::Builtin(Builtin::Int64));
 
                 for var in arguments {
                     let layout = Layout::from_var(arena, var, subs, pointer_size)
@@ -363,10 +425,8 @@ pub fn union_sorted_tags<'a>(
 
                 result.push((tag_name, arg_layouts.into_bump_slice()));
             }
-
+            UnionVariant::Wrapped(result)
-            (is_enum_candidate, result)
         }
-        Err(other) => panic!("invalid content in record variable: {:?}", other),
     }
 }