Merge pull request #4179 from roc-lang/i4150

Support inference of non-unary tag union function

crates/compiler/can/src/copy.rs

@@ -60,8 +60,6 @@ trait CopyEnv {
 
     fn clone_name(&mut self, name: SubsIndex<Lowercase>) -> SubsIndex<Lowercase>;
 
-    fn clone_tag_name(&mut self, tag_name: SubsIndex<TagName>) -> SubsIndex<TagName>;
-
     fn clone_field_names(&mut self, field_names: SubsSlice<Lowercase>) -> SubsSlice<Lowercase>;
 
     fn clone_tag_names(&mut self, tag_names: SubsSlice<TagName>) -> SubsSlice<TagName>;
@@ -95,11 +93,6 @@ impl CopyEnv for Subs {
         name
     }
 
-    #[inline(always)]
-    fn clone_tag_name(&mut self, tag_name: SubsIndex<TagName>) -> SubsIndex<TagName> {
-        tag_name
-    }
-
     #[inline(always)]
     fn clone_field_names(&mut self, field_names: SubsSlice<Lowercase>) -> SubsSlice<Lowercase> {
         field_names
@@ -150,11 +143,6 @@ impl<'a> CopyEnv for AcrossSubs<'a> {
         SubsIndex::push_new(&mut self.target.field_names, self.source[name].clone())
     }
 
-    #[inline(always)]
-    fn clone_tag_name(&mut self, tag_name: SubsIndex<TagName>) -> SubsIndex<TagName> {
-        SubsIndex::push_new(&mut self.target.tag_names, self.source[tag_name].clone())
-    }
-
     #[inline(always)]
     fn clone_field_names(&mut self, field_names: SubsSlice<Lowercase>) -> SubsSlice<Lowercase> {
         SubsSlice::extend_new(
@@ -935,12 +923,13 @@ fn deep_copy_type_vars<C: CopyEnv>(
                         Structure(RecursiveTagUnion(new_rec_var, new_union_tags, new_ext_var))
                     })
                 }
-                FunctionOrTagUnion(tag_name, symbol, ext_var) => {
+                FunctionOrTagUnion(tag_names, symbols, ext_var) => {
                     let new_ext_var = descend_var!(ext_var);
-                    let new_tag_name = env.clone_tag_name(tag_name);
+                    let new_tag_names = env.clone_tag_names(tag_names);
+                    let new_symbols = env.clone_lambda_names(symbols);
                     perform_clone!(Structure(FunctionOrTagUnion(
-                        new_tag_name,
-                        symbol,
+                        new_tag_names,
+                        new_symbols,
                         new_ext_var
                     )))
                 }

crates/compiler/derive_key/src/encoding.rs

@@ -2,7 +2,7 @@ use roc_module::{
     ident::{Lowercase, TagName},
     symbol::Symbol,
 };
-use roc_types::subs::{Content, FlatType, Subs, Variable};
+use roc_types::subs::{Content, FlatType, GetSubsSlice, Subs, Variable};
 
 use crate::{
     util::{check_derivable_ext_var, debug_name_record},
@@ -107,9 +107,14 @@ impl FlatEncodable {
 
                     Ok(Key(FlatEncodableKey::TagUnion(tag_names_and_payload_sizes)))
                 }
-                FlatType::FunctionOrTagUnion(name_index, _, _) => Ok(Key(
-                    FlatEncodableKey::TagUnion(vec![(subs[name_index].clone(), 0)]),
-                )),
+                FlatType::FunctionOrTagUnion(names_index, _, _) => {
+                    Ok(Key(FlatEncodableKey::TagUnion(
+                        subs.get_subs_slice(names_index)
+                            .iter()
+                            .map(|t| (t.clone(), 0))
+                            .collect(),
+                    )))
+                }
                 FlatType::EmptyRecord => Ok(Key(FlatEncodableKey::Record(vec![]))),
                 FlatType::EmptyTagUnion => Ok(Key(FlatEncodableKey::TagUnion(vec![]))),
                 //

crates/compiler/mono/src/layout.rs

@@ -12,7 +12,7 @@ use roc_problem::can::RuntimeError;
 use roc_target::{PtrWidth, TargetInfo};
 use roc_types::num::NumericRange;
 use roc_types::subs::{
-    self, Content, FlatType, GetSubsSlice, Label, OptVariable, RecordFields, Subs, UnionTags,
+    self, Content, FlatType, GetSubsSlice, Label, OptVariable, RecordFields, Subs,
     UnsortedUnionLabels, Variable,
 };
 use roc_types::types::{gather_fields_unsorted_iter, RecordField, RecordFieldsError};
@@ -3152,16 +3152,18 @@ fn layout_from_flat_type<'a>(
 
             layout_from_non_recursive_union(env, &tags).map(Ok)
         }
-        FunctionOrTagUnion(tag_name, _, ext_var) => {
+        FunctionOrTagUnion(tag_names, _, ext_var) => {
             debug_assert!(
                 ext_var_is_empty_tag_union(subs, ext_var),
                 "If ext_var wasn't empty, this wouldn't be a FunctionOrTagUnion!"
             );
 
-            let union_tags = UnionTags::from_tag_name_index(tag_name);
-            let (tags, _) = union_tags.unsorted_tags_and_ext(subs, ext_var);
+            let tag_names = subs.get_subs_slice(tag_names);
+            let unsorted_tags = UnsortedUnionLabels {
+                tags: tag_names.iter().map(|t| (t, &[] as &[Variable])).collect(),
+            };
 
-            layout_from_non_recursive_union(env, &tags).map(Ok)
+            layout_from_non_recursive_union(env, &unsorted_tags).map(Ok)
         }
         RecursiveTagUnion(rec_var, tags, ext_var) => {
             let (tags, ext_var) = tags.unsorted_tags_and_ext(subs, ext_var);

crates/compiler/solve/src/solve.rs

@@ -2299,10 +2299,11 @@ fn type_to_variable<'a>(
                     unreachable!("we assert that the ext var is empty; otherwise we'd already know it was a tag union!");
                 }
 
-                let slice = SubsIndex::new(subs.tag_names.len() as u32);
-                subs.tag_names.push(tag_name.clone());
+                let tag_names = SubsSlice::extend_new(&mut subs.tag_names, [tag_name.clone()]);
+                let symbols = SubsSlice::extend_new(&mut subs.closure_names, [*symbol]);
 
-                let content = Content::Structure(FlatType::FunctionOrTagUnion(slice, *symbol, ext));
+                let content =
+                    Content::Structure(FlatType::FunctionOrTagUnion(tag_names, symbols, ext));
 
                 register_with_known_var(subs, destination, rank, pools, content)
             }

crates/compiler/solve/tests/solve_expr.rs

@@ -7842,6 +7842,20 @@ mod solve_expr {
         );
     }
 
+    #[test]
+    fn dispatch_tag_union_function_inferred() {
+        infer_eq_without_problem(
+            indoc!(
+                r#"
+                g = if Bool.true then A else B
+
+                g ""
+                "#
+            ),
+            "[A Str, B Str]*",
+        );
+    }
+
     #[test]
     fn check_char_as_u8() {
         infer_eq_without_problem(

crates/compiler/test_gen/src/gen_tags.rs

@@ -2004,3 +2004,21 @@ fn match_on_result_with_uninhabited_error_branch() {
         RocStr
     );
 }
+
+#[test]
+#[cfg(any(feature = "gen-llvm", feature = "gen-wasm"))]
+fn dispatch_tag_union_function_inferred() {
+    assert_evals_to!(
+        indoc!(
+            r#"
+            g = \b -> if b then H else J
+
+            when P ((g Bool.true) "") ((g Bool.false) "") is
+                P (H _) (J _) -> "okay"
+                _ -> "FAIL"
+            "#
+        ),
+        RocStr::from("okay"),
+        RocStr
+    );
+}

crates/compiler/types/src/pretty_print.rs

@@ -1156,11 +1156,15 @@ fn write_flat_type<'a>(
             )
         }
 
-        FunctionOrTagUnion(tag_name, _, ext_var) => {
+        FunctionOrTagUnion(tag_names, _, ext_var) => {
             buf.push('[');
 
             let mut tags: MutMap<TagName, _> = MutMap::default();
-            tags.insert(subs[*tag_name].clone(), vec![]);
+            tags.extend(
+                subs.get_subs_slice(*tag_names)
+                    .iter()
+                    .map(|t| (t.clone(), vec![])),
+            );
             let ext_content = write_sorted_tags(env, ctx, subs, buf, &tags, *ext_var);
 
             buf.push(']');
@@ -1241,8 +1245,12 @@ pub fn chase_ext_tag_union(
             push_union(subs, tags, fields);
             chase_ext_tag_union(subs, *ext_var, fields)
         }
-        Content::Structure(FunctionOrTagUnion(tag_name, _, ext_var)) => {
-            fields.push((subs[*tag_name].clone(), vec![]));
+        Content::Structure(FunctionOrTagUnion(tag_names, _, ext_var)) => {
+            fields.extend(
+                subs.get_subs_slice(*tag_names)
+                    .iter()
+                    .map(|t| (t.clone(), vec![])),
+            );
 
             chase_ext_tag_union(subs, *ext_var, fields)
         }

crates/compiler/types/src/subs.rs

@@ -962,13 +962,13 @@ fn subs_fmt_flat_type(this: &FlatType, subs: &Subs, f: &mut fmt::Formatter) -> f
 
             write!(f, "]<{:?}>", new_ext)
         }
-        FlatType::FunctionOrTagUnion(tagname_index, symbol, ext) => {
-            let tagname: &TagName = &subs[*tagname_index];
+        FlatType::FunctionOrTagUnion(tagnames, symbol, ext) => {
+            let tagnames: &[TagName] = subs.get_subs_slice(*tagnames);
 
             write!(
                 f,
                 "FunctionOrTagUnion({:?}, {:?}, {:?})",
-                tagname, symbol, ext
+                tagnames, symbol, ext
             )
         }
         FlatType::RecursiveTagUnion(rec, tags, ext) => {
@@ -2424,7 +2424,12 @@ pub enum FlatType {
     Func(VariableSubsSlice, Variable, Variable),
     Record(RecordFields, Variable),
     TagUnion(UnionTags, Variable),
-    FunctionOrTagUnion(SubsIndex<TagName>, Symbol, Variable),
+
+    /// `A` might either be a function
+    ///   x -> A x : a -> [A a, B a, C a]
+    /// or a tag `[A, B, C]`
+    FunctionOrTagUnion(SubsSlice<TagName>, SubsSlice<Symbol>, Variable),
+
     RecursiveTagUnion(Variable, UnionTags, Variable),
     Erroneous(SubsIndex<Problem>),
     EmptyRecord,
@@ -3881,12 +3886,12 @@ fn flat_type_to_err_type(
             }
         }
 
-        FunctionOrTagUnion(tag_name, _, ext_var) => {
-            let tag_name = subs[tag_name].clone();
+        FunctionOrTagUnion(tag_names, _, ext_var) => {
+            let tag_names = subs.get_subs_slice(tag_names);
 
-            let mut err_tags = SendMap::default();
+            let mut err_tags: SendMap<TagName, Vec<_>> = SendMap::default();
 
-            err_tags.insert(tag_name, vec![]);
+            err_tags.extend(tag_names.iter().map(|t| (t.clone(), vec![])));
 
             match var_to_err_type(subs, state, ext_var).unwrap_structural_alias() {
                 ErrorType::TagUnion(sub_tags, sub_ext) => {
@@ -4202,8 +4207,8 @@ impl StorageSubs {
                 Self::offset_tag_union(offsets, *union_tags),
                 Self::offset_variable(offsets, *ext),
             ),
-            FlatType::FunctionOrTagUnion(tag_name, symbol, ext) => FlatType::FunctionOrTagUnion(
-                Self::offset_tag_name_index(offsets, *tag_name),
+            FlatType::FunctionOrTagUnion(tag_names, symbol, ext) => FlatType::FunctionOrTagUnion(
+                Self::offset_tag_name_slice(offsets, *tag_names),
                 *symbol,
                 Self::offset_variable(offsets, *ext),
             ),
@@ -4295,13 +4300,13 @@ impl StorageSubs {
         record_fields
     }
 
-    fn offset_tag_name_index(
+    fn offset_tag_name_slice(
         offsets: &StorageSubsOffsets,
-        mut tag_name: SubsIndex<TagName>,
-    ) -> SubsIndex<TagName> {
-        tag_name.index += offsets.tag_names;
+        mut tag_names: SubsSlice<TagName>,
+    ) -> SubsSlice<TagName> {
+        tag_names.start += offsets.tag_names;
 
-        tag_name
+        tag_names
     }
 
     fn offset_variable(offsets: &StorageSubsOffsets, variable: Variable) -> Variable {
@@ -4542,12 +4547,22 @@ fn storage_copy_var_to_help(env: &mut StorageCopyVarToEnv<'_>, var: Variable) ->
                     TagUnion(union_tags, new_ext)
                 }
 
-                FunctionOrTagUnion(tag_name, symbol, ext_var) => {
-                    let new_tag_name = SubsIndex::new(env.target.tag_names.len() as u32);
+                FunctionOrTagUnion(tag_names, symbols, ext_var) => {
+                    let new_tag_names = SubsSlice::extend_new(
+                        &mut env.target.tag_names,
+                        env.source.get_subs_slice(tag_names).iter().cloned(),
+                    );
 
-                    env.target.tag_names.push(env.source[tag_name].clone());
+                    let new_symbols = SubsSlice::extend_new(
+                        &mut env.target.closure_names,
+                        env.source.get_subs_slice(symbols).iter().cloned(),
+                    );
 
-                    FunctionOrTagUnion(new_tag_name, symbol, storage_copy_var_to_help(env, ext_var))
+                    FunctionOrTagUnion(
+                        new_tag_names,
+                        new_symbols,
+                        storage_copy_var_to_help(env, ext_var),
+                    )
                 }
 
                 RecursiveTagUnion(rec_var, tags, ext_var) => {
@@ -4981,14 +4996,20 @@ fn copy_import_to_help(env: &mut CopyImportEnv<'_>, max_rank: Rank, var: Variabl
                     TagUnion(union_tags, new_ext)
                 }
 
-                FunctionOrTagUnion(tag_name, symbol, ext_var) => {
-                    let new_tag_name = SubsIndex::new(env.target.tag_names.len() as u32);
+                FunctionOrTagUnion(tag_names, symbols, ext_var) => {
+                    let new_tag_names = SubsSlice::extend_new(
+                        &mut env.target.tag_names,
+                        env.source.get_subs_slice(tag_names).iter().cloned(),
+                    );
 
-                    env.target.tag_names.push(env.source[tag_name].clone());
+                    let new_symbols = SubsSlice::extend_new(
+                        &mut env.target.closure_names,
+                        env.source.get_subs_slice(symbols).iter().cloned(),
+                    );
 
                     FunctionOrTagUnion(
-                        new_tag_name,
-                        symbol,
+                        new_tag_names,
+                        new_symbols,
                         copy_import_to_help(env, max_rank, ext_var),
                     )
                 }

crates/compiler/unify/src/unify.rs

@@ -2657,13 +2657,13 @@ fn unify_flat_type<M: MetaCollector>(
 
             outcome
         }
-        (FunctionOrTagUnion(tag_name, tag_symbol, ext), Func(args, closure, ret)) => {
+        (FunctionOrTagUnion(tag_names, tag_symbols, ext), Func(args, closure, ret)) => {
             unify_function_or_tag_union_and_func(
                 env,
                 pool,
                 ctx,
-                tag_name,
-                *tag_symbol,
+                *tag_names,
+                *tag_symbols,
                 *ext,
                 *args,
                 *ret,
@@ -2671,13 +2671,13 @@ fn unify_flat_type<M: MetaCollector>(
                 true,
             )
         }
-        (Func(args, closure, ret), FunctionOrTagUnion(tag_name, tag_symbol, ext)) => {
+        (Func(args, closure, ret), FunctionOrTagUnion(tag_names, tag_symbols, ext)) => {
             unify_function_or_tag_union_and_func(
                 env,
                 pool,
                 ctx,
-                tag_name,
-                *tag_symbol,
+                *tag_names,
+                *tag_symbols,
                 *ext,
                 *args,
                 *ret,
@@ -2685,50 +2685,61 @@ fn unify_flat_type<M: MetaCollector>(
                 false,
             )
         }
-        (FunctionOrTagUnion(tag_name_1, _, ext1), FunctionOrTagUnion(tag_name_2, _, ext2)) => {
-            let tag_name_1_ref = &env.subs[*tag_name_1];
-            let tag_name_2_ref = &env.subs[*tag_name_2];
-
-            if tag_name_1_ref == tag_name_2_ref {
-                let outcome = unify_pool(env, pool, *ext1, *ext2, ctx.mode);
-                if outcome.mismatches.is_empty() {
-                    let content = *env.subs.get_content_without_compacting(ctx.second);
-                    merge(env, ctx, content)
-                } else {
-                    outcome
-                }
-            } else {
-                let tags1 = UnionTags::from_tag_name_index(*tag_name_1);
-                let tags2 = UnionTags::from_tag_name_index(*tag_name_2);
-
-                unify_tag_unions(env, pool, ctx, tags1, *ext1, tags2, *ext2, Rec::None)
-            }
-        }
-        (TagUnion(tags1, ext1), FunctionOrTagUnion(tag_name, _, ext2)) => {
-            let tags2 = UnionTags::from_tag_name_index(*tag_name);
+        (
+            FunctionOrTagUnion(tag_names_1, tag_symbols_1, ext1),
+            FunctionOrTagUnion(tag_names_2, tag_symbols_2, ext2),
+        ) => unify_two_function_or_tag_unions(
+            env,
+            pool,
+            ctx,
+            *tag_names_1,
+            *tag_symbols_1,
+            *ext1,
+            *tag_names_2,
+            *tag_symbols_2,
+            *ext2,
+        ),
+        (TagUnion(tags1, ext1), FunctionOrTagUnion(tag_names, _, ext2)) => {
+            let empty_tag_var_slices = SubsSlice::extend_new(
+                &mut env.subs.variable_slices,
+                std::iter::repeat(Default::default()).take(tag_names.len()),
+            );
+            let tags2 = UnionTags::from_slices(*tag_names, empty_tag_var_slices);
 
             unify_tag_unions(env, pool, ctx, *tags1, *ext1, tags2, *ext2, Rec::None)
         }
-        (FunctionOrTagUnion(tag_name, _, ext1), TagUnion(tags2, ext2)) => {
-            let tags1 = UnionTags::from_tag_name_index(*tag_name);
+        (FunctionOrTagUnion(tag_names, _, ext1), TagUnion(tags2, ext2)) => {
+            let empty_tag_var_slices = SubsSlice::extend_new(
+                &mut env.subs.variable_slices,
+                std::iter::repeat(Default::default()).take(tag_names.len()),
+            );
+            let tags1 = UnionTags::from_slices(*tag_names, empty_tag_var_slices);
 
             unify_tag_unions(env, pool, ctx, tags1, *ext1, *tags2, *ext2, Rec::None)
         }
 
-        (RecursiveTagUnion(recursion_var, tags1, ext1), FunctionOrTagUnion(tag_name, _, ext2)) => {
+        (RecursiveTagUnion(recursion_var, tags1, ext1), FunctionOrTagUnion(tag_names, _, ext2)) => {
             // this never happens in type-correct programs, but may happen if there is a type error
             debug_assert!(is_recursion_var(env.subs, *recursion_var));
 
-            let tags2 = UnionTags::from_tag_name_index(*tag_name);
+            let empty_tag_var_slices = SubsSlice::extend_new(
+                &mut env.subs.variable_slices,
+                std::iter::repeat(Default::default()).take(tag_names.len()),
+            );
+            let tags2 = UnionTags::from_slices(*tag_names, empty_tag_var_slices);
             let rec = Rec::Left(*recursion_var);
 
             unify_tag_unions(env, pool, ctx, *tags1, *ext1, tags2, *ext2, rec)
         }
 
-        (FunctionOrTagUnion(tag_name, _, ext1), RecursiveTagUnion(recursion_var, tags2, ext2)) => {
+        (FunctionOrTagUnion(tag_names, _, ext1), RecursiveTagUnion(recursion_var, tags2, ext2)) => {
             debug_assert!(is_recursion_var(env.subs, *recursion_var));
 
-            let tags1 = UnionTags::from_tag_name_index(*tag_name);
+            let empty_tag_var_slices = SubsSlice::extend_new(
+                &mut env.subs.variable_slices,
+                std::iter::repeat(Default::default()).take(tag_names.len()),
+            );
+            let tags1 = UnionTags::from_slices(*tag_names, empty_tag_var_slices);
             let rec = Rec::Right(*recursion_var);
 
             unify_tag_unions(env, pool, ctx, tags1, *ext1, *tags2, *ext2, rec)
@@ -3133,17 +3144,20 @@ fn unify_function_or_tag_union_and_func<M: MetaCollector>(
     env: &mut Env,
     pool: &mut Pool,
     ctx: &Context,
-    tag_name_index: &SubsIndex<TagName>,
-    tag_symbol: Symbol,
+    tag_names_slice: SubsSlice<TagName>,
+    tag_fn_lambdas: SubsSlice<Symbol>,
     tag_ext: Variable,
     function_arguments: VariableSubsSlice,
     function_return: Variable,
     function_lambda_set: Variable,
     left: bool,
 ) -> Outcome<M> {
-    let tag_name = env.subs[*tag_name_index].clone();
+    let tag_names = env.subs.get_subs_slice(tag_names_slice).to_vec();
 
-    let union_tags = UnionTags::insert_slices_into_subs(env.subs, [(tag_name, function_arguments)]);
+    let union_tags = UnionTags::insert_slices_into_subs(
+        env.subs,
+        tag_names.into_iter().map(|tag| (tag, function_arguments)),
+    );
     let content = Content::Structure(FlatType::TagUnion(union_tags, tag_ext));
 
     let new_tag_union_var = fresh(env, pool, ctx, content);
@@ -3155,7 +3169,14 @@ fn unify_function_or_tag_union_and_func<M: MetaCollector>(
     };
 
     {
-        let union_tags = UnionLambdas::tag_without_arguments(env.subs, tag_symbol);
+        let lambda_names = env.subs.get_subs_slice(tag_fn_lambdas).to_vec();
+        let new_lambda_names = SubsSlice::extend_new(&mut env.subs.closure_names, lambda_names);
+        let empty_captures_slices = SubsSlice::extend_new(
+            &mut env.subs.variable_slices,
+            std::iter::repeat(Default::default()).take(new_lambda_names.len()),
+        );
+        let union_tags = UnionLambdas::from_slices(new_lambda_names, empty_captures_slices);
+
         let ambient_function_var = if left { ctx.first } else { ctx.second };
         let lambda_set_content = LambdaSet(self::LambdaSet {
             solved: union_tags,
@@ -3196,3 +3217,53 @@ fn unify_function_or_tag_union_and_func<M: MetaCollector>(
 
     outcome
 }
+
+#[allow(clippy::too_many_arguments)]
+fn unify_two_function_or_tag_unions<M: MetaCollector>(
+    env: &mut Env,
+    pool: &mut Pool,
+    ctx: &Context,
+    tag_names_1: SubsSlice<TagName>,
+    tag_symbols_1: SubsSlice<Symbol>,
+    ext1: Variable,
+    tag_names_2: SubsSlice<TagName>,
+    tag_symbols_2: SubsSlice<Symbol>,
+    ext2: Variable,
+) -> Outcome<M> {
+    let merged_tags = {
+        let mut all_tags: Vec<_> = (env.subs.get_subs_slice(tag_names_1).iter())
+            .chain(env.subs.get_subs_slice(tag_names_2))
+            .cloned()
+            .collect();
+        all_tags.sort();
+        all_tags.dedup();
+        SubsSlice::extend_new(&mut env.subs.tag_names, all_tags)
+    };
+    let merged_lambdas = {
+        let mut all_lambdas: Vec<_> = (env.subs.get_subs_slice(tag_symbols_1).iter())
+            .chain(env.subs.get_subs_slice(tag_symbols_2))
+            .cloned()
+            .collect();
+        all_lambdas.sort();
+        all_lambdas.dedup();
+        SubsSlice::extend_new(&mut env.subs.closure_names, all_lambdas)
+    };
+
+    let mut outcome = unify_pool(env, pool, ext1, ext2, ctx.mode);
+    if !outcome.mismatches.is_empty() {
+        return outcome;
+    }
+
+    let merge_outcome = merge(
+        env,
+        ctx,
+        Content::Structure(FlatType::FunctionOrTagUnion(
+            merged_tags,
+            merged_lambdas,
+            ext1,
+        )),
+    );
+
+    outcome.union(merge_outcome);
+    outcome
+}

crates/repl_eval/src/eval.rs

@@ -450,8 +450,8 @@ fn jit_to_ast_help<'a, A: ReplApp<'a>>(
                         payload_vars,
                     ))
                 }
-                Content::Structure(FlatType::FunctionOrTagUnion(tag_name, _, _)) => {
-                    let tag_name = &env.subs[*tag_name];
+                Content::Structure(FlatType::FunctionOrTagUnion(tag_names, _, _)) => {
+                    let tag_name = &env.subs.get_subs_slice(*tag_names)[0];
 
                     Ok(single_tag_union_to_ast(
                         env,
@@ -631,8 +631,8 @@ fn addr_to_ast<'a, M: ReplAppMemory>(
                 let (tag_name, payload_vars) = unpack_single_element_tag_union(env.subs, *tags);
                 single_tag_union_to_ast(env, mem, addr, field_layouts, tag_name, payload_vars)
             }
-            Content::Structure(FlatType::FunctionOrTagUnion(tag_name, _, _)) => {
-                let tag_name = &env.subs[*tag_name];
+            Content::Structure(FlatType::FunctionOrTagUnion(tag_names, _, _)) => {
+                let tag_name = &env.subs.get_subs_slice(*tag_names)[0];
                 single_tag_union_to_ast(env, mem, addr, field_layouts, tag_name, &[])
             }
             Content::Structure(FlatType::EmptyRecord) => {
@@ -1200,6 +1200,19 @@ fn bool_to_ast<'a, M: ReplAppMemory>(
 
                     tag_name_to_expr(env, tag_name)
                 }
+                FlatType::FunctionOrTagUnion(tags, _, _) if tags.len() == 2 => {
+                    let tags = env.subs.get_subs_slice(*tags);
+                    let tag_name_1 = &tags[0];
+                    let tag_name_2 = &tags[1];
+
+                    let tag_name = if value {
+                        max_by_key(tag_name_1, tag_name_2, |n| n.as_ident_str())
+                    } else {
+                        min_by_key(tag_name_1, tag_name_2, |n| n.as_ident_str())
+                    };
+
+                    tag_name_to_expr(env, tag_name)
+                }
                 other => {
                     unreachable!("Unexpected FlatType {:?} in bool_to_ast", other);
                 }
@@ -1293,6 +1306,37 @@ fn byte_to_ast<'a, M: ReplAppMemory>(
                         _ => unreachable!("invalid union variant for a Byte!"),
                     }
                 }
+                FlatType::FunctionOrTagUnion(tags, _, _) => {
+                    // anything with fewer tags is not a byte
+                    debug_assert!(tags.len() > 2);
+
+                    let tags_vec: std::vec::Vec<(TagName, std::vec::Vec<Variable>)> = env
+                        .subs
+                        .get_subs_slice(*tags)
+                        .iter()
+                        .map(|t| (t.clone(), vec![]))
+                        .collect();
+
+                    let union_variant = {
+                        let mut layout_env = layout::Env::from_components(
+                            &mut env.layout_cache,
+                            env.subs,
+                            env.arena,
+                            env.target_info,
+                        );
+                        union_sorted_tags_pub(&mut layout_env, tags_vec, None)
+                    };
+
+                    match union_variant {
+                        UnionVariant::ByteUnion(tagnames) => {
+                            let tag_name = &tagnames[value as usize].expect_tag_ref();
+                            let tag_expr = tag_name_to_expr(env, tag_name);
+                            let loc_tag_expr = Loc::at_zero(tag_expr);
+                            Expr::Apply(env.arena.alloc(loc_tag_expr), &[], CalledVia::Space)
+                        }
+                        _ => unreachable!("invalid union variant for a Byte!"),
+                    }
+                }
                 other => {
                     unreachable!("Unexpected FlatType {:?} in byte_to_ast", other);
                 }