diff --git a/compiler/mono/src/expr.rs b/compiler/mono/src/expr.rs index 8b285a324a..f65a290837 100644 --- a/compiler/mono/src/expr.rs +++ b/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 - .iter() - .enumerate() - .find(|(_, (key, _))| key == &tag_name) - .expect("tag must be in its own type"); + Expr::Byte(tag_id as u8) + } + Unwrapped(field_layouts) => { + let field_exprs = args + .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 - // 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) { + for (arg_layout, arg_expr) in argument_layouts.iter().zip(it) { arguments.push((arg_expr, arg_layout.clone())); } diff --git a/compiler/mono/src/layout.rs b/compiler/mono/src/layout.rs index 93b7812032..0184420621 100644 --- a/compiler/mono/src/layout.rs +++ b/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 - let is_enum_candidate = - tags_vec.len() <= MAX_ENUM_SIZE && tags_vec.iter().all(|(_, args)| args.is_empty()); + union_sorted_tags_help(arena, tags_vec, subs, pointer_size) + } + Err(other) => panic!("invalid content in record variable: {:?}", other), + } +} - let is_unwrapped = tags_vec.len() == 1; - // collect into btreemap to sort - tags_vec.sort(); +fn union_sorted_tags_help<'a>( + arena: &'a Bump, + mut tags_vec: std::vec::Vec<(TagName, std::vec::Vec)>, + 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 - if !is_unwrapped && !is_enum_candidate { - arg_layouts.push(Layout::Builtin(Builtin::Int64)) - } + for (tag_name, arguments) in tags_vec { + // resverse space for the tag discriminant + 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())); } - - (is_enum_candidate, result) + UnionVariant::Wrapped(result) } - Err(other) => panic!("invalid content in record variable: {:?}", other), } }