29ad747 | plum | pyrossh

git clone https://git.pyrossh.dev/plum

A statically typed, imperative programming language inspired by rust, python

29ad747 — Peter John 2026-10-25T16:06:45+05:30

feat(codegen): $Tag runtime representation, tag patterns, ? rewrap, crash import

Files changed (2) hide show

plum-wasm-codegen/src/lib.rs +554 -8
plum-wasm-codegen/tests/codegen_tests.rs +205 -0

plum-wasm-codegen/src/lib.rs CHANGED Viewed

@@ -454,6 +454,12 @@ struct LocalCtx<'a> {
     /// through to the next statement instead of skipping/exiting — see the
     /// `plum_compiler_bugs_found_2026_08` memory.
     loop_extra_depth: Cell<u32>,
+    /// The declared return type of the function being compiled (`None` inside a
+    /// closure body) — `?` builds the OUTER function's `Err` variant from it.
+    ret_type: Option<PlumType>,
+    /// Match-subject scratch local -> its tag-union type, for subjects whose
+    /// static type is a `TTags` (their `Name`/`Class` arms test `$Tag.id`).
+    match_subject_types: RefCell<HashMap<u32, PlumType>>,
 }
 fn fnKey(f: &ast::Fn) -> String {
@@ -517,6 +523,138 @@ thread_local! {
     static CURRENT_CONSTS: RefCell<HashMap<String, ast::Expr>> = RefCell::new(HashMap::new());
 }
+thread_local! {
+    /// Tag name -> runtime `$Tag.id`. Seeded (sorted) from every tag appearing in
+    /// the program's signatures/variants at the top of `compileSourceInner`;
+    /// `tagId` appends any stray tag literal never mentioned in a type.
+    static CURRENT_TAG_IDS: RefCell<std::collections::BTreeMap<String, i32>> =
+        const { RefCell::new(std::collections::BTreeMap::new()) };
+}
+/// Import key of the `plum::crash(Str)` host function in `func_ids`.
+const CRASH_FN_KEY: &str = "__crash";
+fn tagId(name: &str) -> i32 {
+    CURRENT_TAG_IDS.with(|c| {
+        let mut m = c.borrow_mut();
+        if let Some(id) = m.get(name) {
+            return *id;
+        }
+        let id = m.len() as i32;
+        m.insert(name.to_string(), id);
+        id
+    })
+}
+/// Collects every tag name inside every `[...]` tag-union segment of `s`
+/// (a canonical type or mangled name), recursing into payload types.
+fn collectTagNames(s: &str, out: &mut std::collections::BTreeSet<String>) {
+    let chars: Vec<char> = s.chars().collect();
+    let mut i = 0;
+    while i < chars.len() {
+        if chars[i] == '[' && chars.get(i + 1) != Some(&']') {
+            let start = i;
+            let mut depth = 0;
+            while i < chars.len() {
+                match chars[i] {
+                    '[' => depth += 1,
+                    ']' => depth -= 1,
+                    _ => {}
+                }
+                i += 1;
+                if depth == 0 {
+                    break;
+                }
+            }
+            let inner: String = chars[start + 1..i.saturating_sub(1)].iter().collect();
+            let mut item = String::new();
+            let mut d = 0;
+            let mut items = Vec::new();
+            for ch in inner.chars() {
+                match ch {
+                    '(' | '[' => d += 1,
+                    ')' | ']' => d -= 1,
+                    _ => {}
+                }
+                if ch == ',' && d == 0 {
+                    items.push(std::mem::take(&mut item));
+                } else {
+                    item.push(ch);
+                }
+            }
+            items.push(item);
+            for it in items {
+                let name = it.split('(').next().unwrap_or("").trim().to_string();
+                if !name.is_empty() {
+                    out.insert(name);
+                }
+                if let Some(pos) = it.find('(') {
+                    collectTagNames(&it[pos + 1..], out);
+                }
+            }
+        } else {
+            i += 1;
+        }
+    }
+}
+/// Pushes the boxed form of the payload value on top of the stack (typed `ty`)
+/// as an `anyref`: `Int`/`Float`/`Byte` get a `$BoxI64`/`$BoxF64`/`$BoxI32`,
+/// reference values are stored as-is.
+fn boxTagPayload(ty: &PlumType, body: &mut Vec<u8>, ctx: &LocalCtx) {
+    match plumTypeToValtype(ty) {
+        ValType::I64 => Instruction::StructNew(ctx.gc_types.box_i64_idx).encode(body),
+        ValType::F64 => Instruction::StructNew(ctx.gc_types.box_f64_idx).encode(body),
+        ValType::I32 => Instruction::StructNew(ctx.gc_types.box_i32_idx).encode(body),
+        _ => {}
+    }
+}
+/// Inverse of `boxTagPayload`: an `anyref` payload on the stack becomes a `ty` value.
+fn unboxTagPayload(ty: &PlumType, body: &mut Vec<u8>, ctx: &LocalCtx) {
+    let unbox = |idx: u32, body: &mut Vec<u8>| {
+        Instruction::RefCastNonNull(HeapType::Concrete(idx)).encode(body);
+        Instruction::StructGet {
+            struct_type_index: idx,
+            field_index: 0,
+        }
+        .encode(body);
+    };
+    match plumTypeToValtype(ty) {
+        ValType::I64 => unbox(ctx.gc_types.box_i64_idx, body),
+        ValType::F64 => unbox(ctx.gc_types.box_f64_idx, body),
+        ValType::I32 => unbox(ctx.gc_types.box_i32_idx, body),
+        ValType::Ref(rt) => {
+            if let HeapType::Concrete(_) = rt.heap_type {
+                Instruction::RefCastNullable(rt.heap_type).encode(body);
+            }
+        }
+        _ => {}
+    }
+}
+/// Builds a tag value (`NotFound`, `Bad(x)`): `id` plus the optional payload,
+/// boxed by `boxTagPayload`, as a `$Tag` struct.
+fn compileTagLiteral(
+    name: &str,
+    payload: Option<&ast::Expr>,
+    body: &mut Vec<u8>,
+    ctx: &LocalCtx,
+    state: &mut ModuleState,
+) -> Result<(), String> {
+    Instruction::I32Const(tagId(name)).encode(body);
+    match payload {
+        Some(e) => {
+            let ty = inferLocalType(e, ctx);
+            compileExpr(e, body, ctx, state)?;
+            boxTagPayload(&ty, body, ctx);
+        }
+        None => Instruction::RefNull(RefType::ANYREF.heap_type).encode(body),
+    }
+    Instruction::StructNew(ctx.gc_types.tag_type_idx).encode(body);
+    Ok(())
+}
 fn withGcTypes<R>(f: impl FnOnce(&GcTypeRegistry) -> R) -> R {
     CURRENT_GC_TYPES.with(|c| {
         let borrow = c.borrow();
@@ -570,6 +708,10 @@ fn astTypeToWasm(name: &str) -> Option<ValType> {
         // this name, orphaned/unused) class struct type instead of the shared
         // `array<i8>` ref every `[]Byte` value actually is.
         "ByteSlice" => Some(plumTypeToValtype(&PlumType::TByteSlice)),
+        // A structural tag union `[A,B(Str)]` (but not the `[]Byte` slice spelling).
+        other if other.starts_with('[') && !other.starts_with("[]") => {
+            Some(withGcTypes(|r| gcRef(r.tag_type_idx)))
+        }
         other => Some(plumTypeToValtype(&PlumType::TNamed(other.to_string()))),
     }
 }
@@ -649,7 +791,8 @@ fn plumTypeToValtype(t: &PlumType) -> ValType {
         // where a closure's exact struct type isn't being constructed/destructured
         // directly (e.g. deciding a local's storage class), matching `TVariadic`'s
         // existing "not modeled as a concrete shape here" treatment.
+        PlumType::TTags(_) => withGcTypes(|r| gcRef(r.tag_type_idx)),
-        PlumType::TFun(_, _) | PlumType::TVariadic(_) | PlumType::TTags(_) => ValType::Ref(RefType::ANYREF),
+        PlumType::TFun(_, _) | PlumType::TVariadic(_) => ValType::Ref(RefType::ANYREF),
         PlumType::TVar(_) | PlumType::TUnit => ValType::I64,
     }
 }
@@ -689,6 +832,13 @@ pub struct GcTypeRegistry {
     /// harmlessly never read as such) until/unless the program actually
     /// instantiates `Array` at least once.
     pub array_type_idx: u32,
+    /// The single `$Tag {id: i32, payload: anyref}` struct every structural
+    /// tag-union value (`NotFound`, `Bad(x)`) is represented as.
+    pub tag_type_idx: u32,
+    /// Boxes for non-reference tag payloads (`Int`/`Float`/`Byte`).
+    pub box_i64_idx: u32,
+    pub box_f64_idx: u32,
+    pub box_i32_idx: u32,
 }
 /// Whether `name` is `Array` itself or one of its monomorphized specializations
@@ -755,7 +905,8 @@ fn plumTypeToGcValtype(t: &PlumType, registry: &GcTypeRegistry) -> ValType {
         // TFun (closures) and TVariadic get their own concrete representation once
         // Task 2 (closures/variadic calls) lands — `anyref` is a safe, valid-but-not-
         // yet-meaningful placeholder in the meantime, since nothing consumes it yet.
+        PlumType::TTags(_) => gcRef(registry.tag_type_idx),
-        PlumType::TFun(_, _) | PlumType::TVariadic(_) | PlumType::TTags(_) => ValType::Ref(RefType::ANYREF),
+        PlumType::TFun(_, _) | PlumType::TVariadic(_) => ValType::Ref(RefType::ANYREF),
         PlumType::TVar(_) | PlumType::TUnit => ValType::I64,
     }
 }
@@ -781,6 +932,10 @@ fn buildGcTypeRegistry(
         /// The single shared `array<anyref>` type every `Array[T]`
         /// specialization aliases — see `isArraySpecialization`.
         ArrayRef,
+        Tag,
+        BoxI64,
+        BoxF64,
+        BoxI32,
     }
     // Pass 1: assign every entry a slot (and therefore a type index) up front, before
@@ -790,6 +945,8 @@ fn buildGcTypeRegistry(
     let mut enum_super_type_idx: HashMap<String, u32> = HashMap::new();
     let mut variant_type_idx: HashMap<String, u32> = HashMap::new();
     let mut array_ref_slot: Option<u32> = None;
+    let mut erased_enum_super: HashMap<String, u32> = HashMap::new();
+    let mut erased_variant_idx: HashMap<String, u32> = HashMap::new();
     // `Bool` is an ordinary enum now (`type Bool = | True | False` in
     // `libs/std/Bool.plum`) — registered exactly like any other enum
@@ -824,10 +981,31 @@ fn buildGcTypeRegistry(
             }
             continue;
         }
+        // Specializations that differ only in their tag-union segments
+        // (`Try$Int$[A]` / `Try$Int$[A,B]`) share ONE set of wasm types — every
+        // tag value is the same `$Tag` struct — so a later one is aliased onto
+        // the first one's slots instead of getting its own.
+        let erased_enum = plum_checker::erasedEnumName(enum_name);
+        if let Some(existing_super) = erased_enum_super.get(&erased_enum) {
+            enum_super_type_idx.insert(enum_name.clone(), *existing_super);
+            for vname in variant_names {
+                let ev = plum_checker::erasedEnumName(vname);
+                if let Some(idx) = erased_variant_idx.get(&ev) {
+                    variant_type_idx.insert(vname.clone(), *idx);
+                    continue;
+                }
+                variant_type_idx.insert(vname.clone(), slots.len() as u32);
+                erased_variant_idx.insert(ev, slots.len() as u32);
+                slots.push(Slot::Variant(vname.clone()));
+            }
+            continue;
+        }
         enum_super_type_idx.insert(enum_name.clone(), slots.len() as u32);
+        erased_enum_super.insert(erased_enum, slots.len() as u32);
         slots.push(Slot::EnumSuper(enum_name.clone()));
         for vname in variant_names {
             variant_type_idx.insert(vname.clone(), slots.len() as u32);
+            erased_variant_idx.insert(plum_checker::erasedEnumName(vname), slots.len() as u32);
             slots.push(Slot::Variant(vname.clone()));
         }
         // A RECORD-shaped enum (the `type X = ...` replacement, e.g. `enum Cat =
@@ -842,6 +1020,15 @@ fn buildGcTypeRegistry(
         }
     }
+    let tag_type_idx = slots.len() as u32;
+    slots.push(Slot::Tag);
+    let box_i64_idx = slots.len() as u32;
+    slots.push(Slot::BoxI64);
+    let box_f64_idx = slots.len() as u32;
+    slots.push(Slot::BoxF64);
+    let box_i32_idx = slots.len() as u32;
+    slots.push(Slot::BoxI32);
     // The registry is fully index-complete after pass 1 (every name has an assigned
     // slot) even though no field lists exist yet — safe to hand to `plumTypeToGcValtype`
     // for pass 2's field-type resolution.
@@ -853,6 +1040,10 @@ fn buildGcTypeRegistry(
         closure_type_idx: 0,
         variadic_array_type_idx: HashMap::new(),
         array_type_idx: array_ref_slot.unwrap_or(0),
+        tag_type_idx,
+        box_i64_idx,
+        box_f64_idx,
+        box_i32_idx,
     };
     // Pass 2: build the real SubType for every slot, now that every cross-reference
@@ -939,6 +1130,29 @@ fn buildGcTypeRegistry(
                     },
                 }
             }
+            Slot::Tag => SubType {
+                is_final: true,
+                supertype_idx: None,
+                composite_type: CompositeType {
+                    inner: CompositeInnerType::Struct(StructType {
+                        fields: vec![
+                            FieldType {
+                                element_type: StorageType::Val(ValType::I32),
+                                mutable: false,
+                            },
+                            FieldType {
+                                element_type: StorageType::Val(ValType::Ref(RefType::ANYREF)),
+                                mutable: false,
+                            },
+                        ]
+                        .into(),
+                    }),
+                    shared: false,
+                },
+            },
+            Slot::BoxI64 => boxSubType(ValType::I64),
+            Slot::BoxF64 => boxSubType(ValType::F64),
+            Slot::BoxI32 => boxSubType(ValType::I32),
             Slot::ArrayRef => SubType {
                 is_final: true,
                 supertype_idx: None,
@@ -958,6 +1172,23 @@ fn buildGcTypeRegistry(
     registry
 }
+fn boxSubType(vt: ValType) -> SubType {
+    SubType {
+        is_final: true,
+        supertype_idx: None,
+        composite_type: CompositeType {
+            inner: CompositeInnerType::Struct(StructType {
+                fields: vec![FieldType {
+                    element_type: StorageType::Val(vt),
+                    mutable: false,
+                }]
+                .into(),
+            }),
+            shared: false,
+        },
+    }
+}
 /// Maps an `ast::ParamType::Fn(params, ret)` to the wasm signature of the *closure
 /// function* it compiles to: an implicit leading `env: anyref`, then one param per
 /// declared param type, returning `ret`'s val type (or nothing for `Unit`).
@@ -1232,6 +1463,31 @@ fn compileSourceInner(
     let (global_env, methods, enum_variants, enum_params) = plum_checker::buildGlobalTables(source);
     let min_required = plum_checker::buildMinRequiredArgs(source);
+    // Intern every tag name the program's types mention, in sorted order.
+    {
+        let mut names = std::collections::BTreeSet::new();
+        for (k, info) in &enum_variants {
+            collectTagNames(k, &mut names);
+            collectTagNames(&info.enum_name, &mut names);
+            for t in &info.field_types {
+                collectTagNames(&t.to_string(), &mut names);
+            }
+        }
+        for scheme in global_env.values() {
+            collectTagNames(&scheme.body.to_string(), &mut names);
+        }
+        for t in methods.values() {
+            collectTagNames(&t.to_string(), &mut names);
+        }
+        CURRENT_TAG_IDS.with(|c| {
+            *c.borrow_mut() = names
+                .into_iter()
+                .enumerate()
+                .map(|(i, n)| (n, i as i32))
+                .collect();
+        });
+    }
     let mut module = WasmModule::new();
     let mut gc_types = buildGcTypeRegistry(&mut module, source, &enum_variants, &enum_params);
@@ -1317,6 +1573,21 @@ fn compileSourceInner(
         }
     }
+    // `plum::crash(Str)`: same wasm type as `printLn`; imported only when the
+    // program actually contains a `crash` statement.
+    {
+        let dbg = format!("{:?}", source.items);
+        let uses_crash = dbg
+            .match_indices("Crash(")
+            .any(|(i, _)| i == 0 || matches!(dbg.as_bytes()[i - 1], b'[' | b'(' | b' '));
+        if uses_crash {
+            let str_vt = plumTypeToValtype(&PlumType::TStr);
+            let type_idx = module.addType(&[str_vt], &[]);
+            let func_idx = module.addImport("plum", "crash", type_idx);
+            func_ids.insert(CRASH_FN_KEY.to_string(), func_idx);
+        }
+    }
     // Pre-allocate one instance of every payload-free variant (True/False/None/...)
     // as a global, populated once by a `start` function rather than reconstructed on
     // every reference — see this migration plan's Decision 2. `struct.new` isn't
@@ -3094,6 +3365,18 @@ impl<'a> Collector<'a> {
     /// sub-pattern found *nested* inside another constructor pattern's fields (the
     /// outermost, per-subject pattern doesn't need one — see `nested_class_scratch`).
     fn collectPattern(&mut self, pat: &ast::CasePattern, ty: &PlumType) {
+        // Tag-union subject/field: uppercase names are tags (no binding, no
+        // narrowing scratch slot); a tag's payload sub-pattern binds to its type.
+        if let PlumType::TTags(tags) = ty {
+            if isTagPattern(pat) {
+                if let ast::CasePattern::Class { name, fields } = pat {
+                    if let (Some(Some(pty)), Some(f)) = (tags.get(name), fields.first()) {
+                        self.collectPattern(f, pty);
+                    }
+                }
+                return;
+            }
+        }
         match pat {
             ast::CasePattern::Name(n) => {
                 let is_variant = n.chars().next().map(|c| c.is_uppercase()).unwrap_or(false)
@@ -3473,6 +3756,8 @@ fn compileFnBody(
         type_env: RefCell::new(base_env),
         closure_local_sigs: RefCell::new(HashMap::new()),
         loop_extra_depth: Cell::new(0),
+        ret_type: f.returns.as_ref().map(|r| plum_checker::plumTypeFromAst(r)),
+        match_subject_types: RefCell::new(HashMap::new()),
     };
     let result_vt = retTypeToWasm(f.returns.as_ref());
@@ -3561,7 +3846,9 @@ fn compileStmtInValuePosition(
 ) -> Result<(), String> {
     match stmt {
         ast::Stmt::Expr(e, ..) => compileExpr(e, body, ctx, state),
-        ast::Stmt::Return(_) | ast::Stmt::Todo => compileStmt(stmt, body, ctx, state),
+        ast::Stmt::Return(_) | ast::Stmt::Todo | ast::Stmt::Crash(_) => {
+            compileStmt(stmt, body, ctx, state)
+        }
         ast::Stmt::If(if_) => compileIf(if_, Some(result_vt), body, ctx, state),
         ast::Stmt::Match(m) => compileMatch(m, body, ctx, state, Some(result_vt)),
         _ => Err(
@@ -3993,8 +4280,18 @@ fn compileStmt(
             Instruction::Unreachable.encode(body);
             Instruction::End.encode(body);
         }
-        // Placeholder until T3 implements the `plum::crash` import.
+        // `crash msg` — hand the message to the host's `plum::crash(Str)`
+        // import (registered only when the program uses `crash`, see
+        // `CRASH_FN_KEY`), then trap.
+        ast::Stmt::Crash(msg) => {
+            compileExpr(msg, body, ctx, state)?;
+            let crash_idx = *ctx
+                .func_ids
+                .get(CRASH_FN_KEY)
+                .ok_or_else(|| "internal codegen error: crash import missing".to_string())?;
+            Instruction::Call(crash_idx).encode(body);
-        ast::Stmt::Crash(_) => Instruction::Unreachable.encode(body),
+            Instruction::Unreachable.encode(body);
+        }
         ast::Stmt::Todo => {
             // Marks an unimplemented body — trap rather than silently continuing.
             Instruction::Unreachable.encode(body);
@@ -4131,6 +4428,83 @@ fn matchCoversEveryEnumVariant(
     !tuples_seen.is_empty() && tuples_seen.len() == total_combinations
 }
+/// An uppercase `Name`/`Class` pattern — against a tag-union subject, a tag arm.
+fn isTagPattern(pat: &ast::CasePattern) -> bool {
+    let upper = |n: &str| n.chars().next().is_some_and(|c| c.is_uppercase());
+    match pat {
+        ast::CasePattern::Name(n) => upper(n),
+        ast::CasePattern::Class { name, .. } => upper(name),
+        _ => false,
+    }
+}
+/// Compiles one tag arm (`A =>` / `B(s) =>`) against a `$Tag` value that
+/// `load` pushes: tests `id`, binds/unboxes the payload, then runs `on_match`,
+/// else `on_fail`.
+#[allow(clippy::too_many_arguments)]
+fn compileTagPattern(
+    pat: &ast::CasePattern,
+    tags: &std::collections::BTreeMap<String, Option<Box<PlumType>>>,
+    load: &dyn Fn(&mut Vec<u8>),
+    result_vt: Option<ValType>,
+    body: &mut Vec<u8>,
+    ctx: &LocalCtx,
+    state: &mut ModuleState,
+    on_match: &mut dyn FnMut(&mut Vec<u8>, &mut ModuleState) -> Result<(), String>,
+    on_fail: &mut dyn FnMut(&mut Vec<u8>, &mut ModuleState) -> Result<(), String>,
+) -> Result<(), String> {
+    let (name, sub) = match pat {
+        ast::CasePattern::Name(n) => (n.as_str(), None),
+        ast::CasePattern::Class { name, fields } => (name.as_str(), fields.first()),
+        _ => unreachable!("isTagPattern checked"),
+    };
+    let payload_ty = tags
+        .get(name)
+        .ok_or_else(|| format!("codegen: tag '{}' is not part of the matched tag union", name))?
+        .as_ref()
+        .map(|b| (**b).clone());
+    let tag_idx = ctx.gc_types.tag_type_idx;
+    load(body);
+    Instruction::StructGet {
+        struct_type_index: tag_idx,
+        field_index: 0,
+    }
+    .encode(body);
+    Instruction::I32Const(tagId(name)).encode(body);
+    Instruction::I32Eq.encode(body);
+    Instruction::If(blockTypeFor(result_vt)).encode(body);
+    match (sub, &payload_ty) {
+        (None, _) | (Some(ast::CasePattern::Wildcard), _) => {}
+        (Some(ast::CasePattern::Name(n)), Some(pty)) if !isTagPattern(&ast::CasePattern::Name(n.clone())) => {
+            let idx = ctx.locals.get(n).copied().ok_or_else(|| {
+                format!("internal codegen error: missing binding local '{}'", n)
+            })?;
+            load(body);
+            Instruction::StructGet {
+                struct_type_index: tag_idx,
+                field_index: 1,
+            }
+            .encode(body);
+            unboxTagPayload(pty, body, ctx);
+            Instruction::LocalSet(idx).encode(body);
+            ctx.type_env
+                .borrow_mut()
+                .insert(n.clone(), TypeScheme::mono(pty.clone()));
+        }
+        (Some(_), _) => {
+            return Err(format!(
+                "codegen: unsupported pattern inside tag '{}' payload (use a name or '_')",
+                name
+            ))
+        }
+    }
+    on_match(body, state)?;
+    Instruction::Else.encode(body);
+    on_fail(body, state)?;
+    Instruction::End.encode(body);
+    Ok(())
+}
 fn compileMatch(
     m: &ast::Match,
     body: &mut Vec<u8>,
@@ -4151,6 +4525,11 @@ fn compileMatch(
         let subject_ty = inferLocalType(subject, ctx);
         let subject_vt = plumTypeToValtype(&subject_ty);
         let scratch_local = ctx.match_scratch_base + base_slot + i as u32;
+        if matches!(subject_ty, PlumType::TTags(_)) {
+            ctx.match_subject_types
+                .borrow_mut()
+                .insert(scratch_local, subject_ty.clone());
+        }
         compileExpr(subject, body, ctx, state)?;
         Instruction::LocalSet(scratch_local).encode(body);
         all_subjects.push((subject_vt, scratch_local));
@@ -4166,8 +4545,14 @@ fn compileMatch(
     // would reallocate every nested `CasePattern::Class` node at a new address and
     // break `nested_class_scratch`'s pointer-identity-keyed lookup.
     let subject_vts: Vec<ValType> = all_subjects.iter().map(|(vt, _)| *vt).collect();
+    // A match over a tag union with no catch-all was proven exhaustive by the checker.
+    let tag_union_match = m.subjects.len() == 1
+        && ctx
+            .match_subject_types
+            .borrow()
+            .contains_key(&all_subjects[0].1);
-    let exhaustive_fallback =
+    let exhaustive_fallback = result_vt.is_some()
-        result_vt.is_some() && matchCoversEveryEnumVariant(&m.cases, &subject_vts, ctx);
+        && (tag_union_match || matchCoversEveryEnumVariant(&m.cases, &subject_vts, ctx));
     compileMatchArmsMulti(
         &m.cases,
@@ -4277,6 +4662,44 @@ fn compileCasePositions(
     }
     let pat = &case.patterns[pos];
     let (subject_vt, scratch_local) = all_subjects[pos];
+    let tag_subject = ctx.match_subject_types.borrow().get(&scratch_local).cloned();
+    if let Some(PlumType::TTags(tags)) = &tag_subject {
+        if isTagPattern(pat) {
+            return compileTagPattern(
+                pat,
+                tags,
+                &|b: &mut Vec<u8>| Instruction::LocalGet(scratch_local).encode(b),
+                result_vt,
+                body,
+                ctx,
+                state,
+                &mut |body, state| {
+                    compileCasePositions(
+                        case,
+                        pos + 1,
+                        all_subjects,
+                        rest,
+                        result_vt,
+                        exhaustive_fallback,
+                        body,
+                        ctx,
+                        state,
+                    )
+                },
+                &mut |body, state| {
+                    compileMatchArmsMulti(
+                        rest,
+                        all_subjects,
+                        result_vt,
+                        exhaustive_fallback,
+                        body,
+                        ctx,
+                        state,
+                    )
+                },
+            );
+        }
+    }
     match pat {
         ast::CasePattern::Wildcard => {
             // Always matches this position; move on to the next one.
@@ -4318,7 +4741,9 @@ fn compileCasePositions(
                 Instruction::LocalSet(idx).encode(body);
                 ctx.type_env.borrow_mut().insert(
                     n.clone(),
+                    TypeScheme::mono(
-                    TypeScheme::mono(plumTypeFromValtypeHint(subject_vt)),
+                        tag_subject.unwrap_or_else(|| plumTypeFromValtypeHint(subject_vt)),
+                    ),
                 );
                 // A binding always matches this position; move on to the next one.
                 compileCasePositions(
@@ -4588,6 +5013,55 @@ fn compileFieldPatterns(
     let pat = &fields[fpos];
     let field_ty = &field_types[fpos];
+    if let PlumType::TTags(tags) = field_ty {
+        if isTagPattern(pat) {
+            return compileTagPattern(
+                pat,
+                tags,
+                &|b: &mut Vec<u8>| {
+                    Instruction::LocalGet(container_local).encode(b);
+                    Instruction::StructGet {
+                        struct_type_index: container_type_idx,
+                        field_index: fpos as u32,
+                    }
+                    .encode(b);
+                },
+                result_vt,
+                body,
+                ctx,
+                state,
+                &mut |body, state| {
+                    compileFieldPatterns(
+                        fields,
+                        field_types,
+                        fpos + 1,
+                        container_local,
+                        container_type_idx,
+                        rest,
+                        all_subjects,
+                        result_vt,
+                        exhaustive_fallback,
+                        body,
+                        ctx,
+                        state,
+                        on_match,
+                    )
+                },
+                &mut |body, state| {
+                    compileMatchArmsMulti(
+                        rest,
+                        all_subjects,
+                        result_vt,
+                        exhaustive_fallback,
+                        body,
+                        ctx,
+                        state,
+                    )
+                },
+            );
+        }
+    }
     match pat {
         ast::CasePattern::Wildcard => compileFieldPatterns(
             fields,
@@ -4951,6 +5425,35 @@ fn compileExpr(
             Instruction::If(BlockType::Result(success_vt)).encode(body);
             Instruction::LocalGet(scratch_local).encode(body);
             Instruction::RefCastNonNull(HeapType::Concrete(failure_idx)).encode(body);
+            // Failure branch: take the inner `Err` payload and build the OUTER
+            // function's `Err` variant from it (the Ok types may differ).
+            let outer_err = match &ctx.ret_type {
+                Some(PlumType::TNamed(outer_enum)) if *outer_enum != enum_name => {
+                    plum_checker::tryOperatorVariants(outer_enum, &cctx)
+                        .map(|(_, _, fname, _)| fname)
+                }
+                _ => None,
+            };
+            if let Some(outer_err_name) = outer_err {
+                let outer_idx = *ctx
+                    .gc_types
+                    .variant_type_idx
+                    .get(outer_err_name)
+                    .ok_or_else(|| {
+                        format!(
+                            "codegen: variant '{}' missing from the GC type registry",
+                            outer_err_name
+                        )
+                    })?;
+                if outer_idx != failure_idx {
+                    Instruction::StructGet {
+                        struct_type_index: failure_idx,
+                        field_index: 0,
+                    }
+                    .encode(body);
+                    Instruction::StructNew(outer_idx).encode(body);
+                }
+            }
             Instruction::Return.encode(body);
             Instruction::Else.encode(body);
             Instruction::LocalGet(scratch_local).encode(body);
@@ -5116,6 +5619,30 @@ fn compileExpr(
         }
         ast::Expr::Compare(c) => {
             let left_ty = inferLocalType(&c.left, ctx);
+            // Tag unions compare by `$Tag.id` (the checker rejects payload tags).
+            if matches!(left_ty, PlumType::TTags(_)) {
+                let tag_idx = ctx.gc_types.tag_type_idx;
+                for side in [&c.left, &c.right] {
+                    compileExpr(side, body, ctx, state)?;
+                    Instruction::StructGet {
+                        struct_type_index: tag_idx,
+                        field_index: 0,
+                    }
+                    .encode(body);
+                }
+                match &c.op {
+                    ast::CmpOp::Eq => Instruction::I32Eq.encode(body),
+                    ast::CmpOp::Neq | ast::CmpOp::NotEq2 => Instruction::I32Ne.encode(body),
+                    other => {
+                        return Err(format!(
+                            "codegen: '{:?}' is not supported between tag unions",
+                            other
+                        ))
+                    }
+                }
+                pushBoolRefFromI32Flag(body, ctx);
+                return Ok(());
+            }
             compileExpr(&c.left, body, ctx, state)?;
             compileExpr(&c.right, body, ctx, state)?;
             match left_ty {
@@ -5282,6 +5809,18 @@ fn compileExpr(
                 compileExpr(&synthetic, body, ctx, state)?;
             } else if let Some(info) = ctx.enum_variants.get(&call.name) {
                 compileVariantConstruction(info, call, expr, body, ctx, state)?;
+            } else if !ctx.func_ids.contains_key(&call.name)
+                && call.args.len() <= 1
+                && call.name.chars().next().is_some_and(|c| c.is_uppercase())
+                && matches!(inferLocalType(expr, ctx), PlumType::TTags(_))
+            {
+                // `Bad(x)` / `Bad()` — an undeclared uppercase name is a tag literal.
+                let payload = call.args.first().map(|a| match a {
+                    ast::Arg::Positional(e) => e,
+                    ast::Arg::Keyword { value, .. } => value,
+                    ast::Arg::Pair { value, .. } => value,
+                });
+                compileTagLiteral(&call.name, payload, body, ctx, state)?;
             } else {
                 fn argExprOf(arg: &ast::Arg) -> &ast::Expr {
                     match arg {
@@ -5382,6 +5921,11 @@ fn compileExpr(
                     let const_value = CURRENT_CONSTS.with(|c| c.borrow().get(n).cloned());
                     match const_value {
                         Some(value) => compileExpr(&value, body, ctx, state)?,
+                        None if n.chars().next().is_some_and(|c| c.is_uppercase())
+                            && matches!(inferLocalType(expr, ctx), PlumType::TTags(_)) =>
+                        {
+                            compileTagLiteral(n, None, body, ctx, state)?
+                        }
                         None => {
                             return Err(format!(
                                 "codegen: type name '{}' is not yet supported as a value",
@@ -6200,6 +6744,8 @@ fn compileClosureBody(
         type_env: RefCell::new(base_env),
         closure_local_sigs: RefCell::new(HashMap::new()),
         loop_extra_depth: Cell::new(0),
+        ret_type: None,
+        match_subject_types: RefCell::new(HashMap::new()),
     };
     compileBlockAsFnBody(&cl.body, &mut body, &local_ctx, state, info.ret_vt)?;

plum-wasm-codegen/tests/codegen_tests.rs CHANGED Viewed

@@ -1032,3 +1032,208 @@ fn wasmtimeGcConfigSupportsWideningConcreteStructRefToAnyrefAndCastingBack() {
         "concrete struct ref should widen to anyref implicitly and cast back correctly"
     );
 }
+// ---------------------------------------------------------------------------
+// Roc-style structural error tags (`Try[T, E]` with `E = [A, B(Str)]`).
+// ---------------------------------------------------------------------------
+const TRY_PRELUDE: &str = "\
+enum Bool =
+  | True
+  | False
+enum Try[T, E] =
+  | Ok(T)
+  | Err(E)
+";
+/// Compiles `TRY_PRELUDE + body`, validates it, runs `main() -> Int` and
+/// returns its result.
+fn runTryMain(body: &str) -> i64 {
+    let src = format!("{}{}", TRY_PRELUDE, body);
+    let bytes = assertValid(&src);
+    let engine = gcEngine();
+    let module = wasmtime::Module::new(&engine, &bytes).unwrap();
+    let mut store = wasmtime::Store::new(&engine, ());
+    let instance = wasmtime::Instance::new(&mut store, &module, &[]).unwrap();
+    let main = instance
+        .get_typed_func::<(), i64>(&mut store, "main")
+        .unwrap();
+    main.call(&mut store, ()).unwrap()
+}
+#[test]
+fn int_payload_roundtrip() {
+    let got = runTryMain(
+        "\
+fun make() -> Try[Int, [Big(Int), NotFound]] =
+  Err(Big(9007199254740993))
+fun main() -> Int =
+  match make()
+    Ok(v) => v
+    Err(NotFound) => 0
+    Err(Big(n)) => n
+",
+    );
+    assert_eq!(got, 9007199254740993);
+}
+#[test]
+fn float_payload_roundtrip() {
+    let got = runTryMain(
+        "\
+fun make() -> Try[Int, [Frac(Float), NotFound]] =
+  Err(Frac(0.1))
+fun main() -> Int =
+  match make()
+    Ok(v) => v
+    Err(NotFound) => 0
+    Err(Frac(f)) =>
+      if f == 0.1
+        return 1
+      return 2
+",
+    );
+    assert_eq!(got, 1);
+}
+#[test]
+fn nested_err_tag_pattern() {
+    let prog = |ctor: &str| {
+        format!(
+            "\
+fun make() -> Try[Int, [Bad(Int), NotFound]] =
+  {}
+fun main() -> Int =
+  match make()
+    Ok(v) => v
+    Err(NotFound) => 0
+    Err(Bad(n)) => n
+",
+            ctor
+        )
+    };
+    assert_eq!(runTryMain(&prog("Ok(7)")), 7);
+    assert_eq!(runTryMain(&prog("Err(NotFound)")), 0);
+    assert_eq!(runTryMain(&prog("Err(Bad(42))")), 42);
+}
+#[test]
+fn question_mark_different_ok_types() {
+    // Inner Try[Int,[A]] inside fn returning Try[Str,[A,B]]: failure rewraps
+    // into the OUTER Err variant (different Ok type), success continues.
+    let prog = |inner: &str| {
+        format!(
+            "\
+fun inner() -> Try[Int, [A]] =
+  {}
+fun outer() -> Try[Str, [A, B]] =
+  x := inner()?
+  Ok(\"fine\")
+fun main() -> Int =
+  match outer()
+    Ok(_) => 1
+    Err(A) => 2
+    Err(B) => 3
+",
+            inner
+        )
+    };
+    assert_eq!(runTryMain(&prog("Err(A)")), 2);
+    assert_eq!(runTryMain(&prog("Ok(5)")), 1);
+}
+#[test]
+fn tag_eq() {
+    let got = runTryMain(
+        "\
+fun main() -> Int =
+  a := A == A
+  b := A == B
+  if a
+    if b
+      return 0
+    return 1
+  return 2
+",
+    );
+    assert_eq!(got, 1);
+}
+#[test]
+fn match_on_bare_tag_union() {
+    let got = runTryMain(
+        "\
+fun pick(n: Int) -> [Small, Big(Int)] =
+  if n < 10
+    return Small
+  return Big(n)
+fun main() -> Int =
+  match pick(77)
+    Small => 0
+    Big(n) => n
+",
+    );
+    assert_eq!(got, 77);
+}
+#[test]
+fn crash_traps_and_calls_import() {
+    let src = format!(
+        "{}fun main() -> Int =\n  crash \"boom\"\n",
+        TRY_PRELUDE
+    );
+    let bytes = assertValid(&src);
+    let engine = gcEngine();
+    let module = wasmtime::Module::new(&engine, &bytes).unwrap();
+    assert!(
+        module.imports().any(|i| i.module() == "plum" && i.name() == "crash"),
+        "crash import must be emitted when crash is used"
+    );
+    let plain = assertValid("fun main() -> Int =\n  1\n");
+    let plain_module = wasmtime::Module::new(&engine, &plain).unwrap();
+    assert!(
+        plain_module.imports().all(|i| i.name() != "crash"),
+        "crash import must NOT be emitted when unused"
+    );
+}
+fn countTypes(bytes: &[u8]) -> usize {
+    let mut n = 0;
+    for payload in wasmparser::Parser::new(0).parse_all(bytes) {
+        if let wasmparser::Payload::TypeSection(s) = payload.unwrap() {
+            for g in s {
+                n += g.unwrap().types().count();
+            }
+        }
+    }
+    n
+}
+#[test]
+fn specializations_share_gc_type() {
+    let prog = |second: &str| {
+        format!(
+            "{}\
+fun f() -> Try[Int, [A]] =
+  Err(A)
+fun g() -> Try[Int, {}] =
+  Err(A)
+fun main() -> Int = 1
+",
+            TRY_PRELUDE, second
+        )
+    };
+    let same = assertValid(&prog("[A]"));
+    let wider = assertValid(&prog("[A, B]"));
+    assert_eq!(countTypes(&same), countTypes(&wider));
+}