Implement tag union discriminant extraction for byte- and newtype-variants

crates/compiler/mono/src/ir.rs

@@ -5446,38 +5446,10 @@ pub fn with_hole<'a>(
                 }
                 TagDiscriminant => {
                     debug_assert_eq!(arg_symbols.len(), 1);
-                    let x = arg_symbols[0];
-                    let x_layout = layout_cache.from_var(arena, args[0].0, env.subs).expect(
-                        "TagDiscriminant only built in derived impls, which must type-check",
-                    );
+                    let tag = arg_symbols[0];
+                    let tag_var = args[0].0;
 
-                    let x_union_layout = match x_layout {
-                        Layout::Union(un) => un,
-                        _ => internal_error!("TagDiscriminant can only apply to tags"),
-                    };
-                    let tag_id_layout = x_union_layout.tag_id_layout();
-                    debug_assert!(tag_id_layout == Layout::u8() || tag_id_layout == Layout::u16());
-
-                    let tag_id_sym = env.unique_symbol();
-
-                    let cast_to_u16 = self::Call {
-                        call_type: CallType::LowLevel {
-                            op: LowLevel::NumIntCast,
-                            update_mode: env.next_update_mode_id(),
-                        },
-                        arguments: env.arena.alloc([tag_id_sym]),
-                    };
-
-                    let hole = Stmt::Let(assigned, Expr::Call(cast_to_u16), Layout::u16(), hole);
-                    Stmt::Let(
-                        tag_id_sym,
-                        Expr::GetTagId {
-                            structure: x,
-                            union_layout: x_union_layout,
-                        },
-                        tag_id_layout,
-                        env.arena.alloc(hole),
-                    )
+                    build_tag_discriminant_expr(env, layout_cache, assigned, tag, tag_var, hole)
                 }
                 _ => {
                     let call = self::Call {
@@ -5504,6 +5476,76 @@ pub fn with_hole<'a>(
     }
 }
 
+fn build_tag_discriminant_expr<'a>(
+    env: &mut Env<'a, '_>,
+    layout_cache: &mut LayoutCache<'a>,
+    assigned: Symbol,
+    tag_sym: Symbol,
+    tag_var: Variable,
+    hole: &'a Stmt<'a>,
+) -> Stmt<'a> {
+    use crate::layout::UnionVariant::*;
+
+    let union_variant = {
+        let mut layout_env =
+            layout::Env::from_components(layout_cache, env.subs, env.arena, env.target_info);
+        crate::layout::union_sorted_tags(&mut layout_env, tag_var)
+            .expect("TagDiscriminant only built in derived impls, which must type-check")
+    };
+
+    let build_cast_to_u16 = |env: &mut Env<'a, '_>, tag_id_sym| self::Call {
+        call_type: CallType::LowLevel {
+            op: LowLevel::NumIntCast,
+            update_mode: env.next_update_mode_id(),
+        },
+        arguments: env.arena.alloc([tag_id_sym]),
+    };
+
+    match union_variant {
+        Never | Unit | Newtype { .. } | NewtypeByVoid { .. } => {
+            // All trivial unions decay into `0` for their discriminant.
+            Stmt::Let(
+                assigned,
+                Expr::Literal(Literal::Int(0u128.to_ne_bytes())),
+                Layout::u16(),
+                hole,
+            )
+        }
+        BoolUnion { .. } | ByteUnion(..) => {
+            // These are represented as integers at runtime, so we just need to cast them.
+            let cast_to_u16 = build_cast_to_u16(env, tag_sym);
+
+            Stmt::Let(assigned, Expr::Call(cast_to_u16), Layout::u16(), hole)
+        }
+        Wrapped(..) => {
+            let tag_layout = layout_cache.from_var(env.arena, tag_var, env.subs).unwrap();
+            let tag_union_layout = match tag_layout {
+                Layout::Union(un) => un,
+                _ => internal_error!(
+                    "Somehow this is a `Wrapped` variant, but its layout is not a union"
+                ),
+            };
+
+            let tag_id_layout = tag_union_layout.tag_id_layout();
+            debug_assert!(tag_id_layout == Layout::u8() || tag_id_layout == Layout::u16());
+
+            let tag_id_sym = env.unique_symbol();
+            let cast_to_u16 = build_cast_to_u16(env, tag_id_sym);
+
+            let hole = Stmt::Let(assigned, Expr::Call(cast_to_u16), Layout::u16(), hole);
+            Stmt::Let(
+                tag_id_sym,
+                Expr::GetTagId {
+                    structure: tag_sym,
+                    union_layout: tag_union_layout,
+                },
+                tag_id_layout,
+                env.arena.alloc(hole),
+            )
+        }
+    }
+}
+
 #[inline(always)]
 fn late_resolve_ability_specialization<'a>(
     env: &mut Env<'a, '_>,

crates/compiler/mono/src/layout.rs

@@ -879,15 +879,8 @@ impl<'a> UnionLayout<'a> {
         }
     }
 
-    pub fn tag_id_layout(&self) -> Layout<'a> {
-        // TODO is it beneficial to return a more specific layout?
-        // e.g. Layout::bool() and Layout::VOID
-        match self.discriminant() {
-            Discriminant::U0 => Layout::u8(),
-            Discriminant::U1 => Layout::u8(),
-            Discriminant::U8 => Layout::u8(),
-            Discriminant::U16 => Layout::u16(),
-        }
+    pub fn tag_id_layout(&self) -> Layout<'static> {
+        self.discriminant().layout()
     }
 
     fn stores_tag_id_in_pointer_bits(tags: &[&[Layout<'a>]], target_info: TargetInfo) -> bool {
@@ -1138,6 +1131,17 @@ impl Discriminant {
     pub const fn alignment_bytes(&self) -> u32 {
         self.stack_size()
     }
+
+    pub const fn layout(&self) -> Layout<'static> {
+        // TODO is it beneficial to return a more specific layout?
+        // e.g. Layout::bool() and Layout::VOID
+        match self {
+            Discriminant::U0 => Layout::u8(),
+            Discriminant::U1 => Layout::u8(),
+            Discriminant::U8 => Layout::u8(),
+            Discriminant::U16 => Layout::u16(),
+        }
+    }
 }
 
 /// Custom type so we can get the numeric representation of a symbol in tests (so `#UserApp.3`
@@ -2713,7 +2717,7 @@ impl<'a> Layout<'a> {
         Layout::Builtin(Builtin::Int(IntWidth::U8))
     }
 
-    pub fn u16() -> Layout<'a> {
+    pub const fn u16() -> Layout<'a> {
         Layout::Builtin(Builtin::Int(IntWidth::U16))
     }