Fix type checker literal union upcasting issues

- Extend infer_arithmetic's base_ty to handle Ty::Union by folding over
  members with a promote helper, fixing InvalidBinaryOp for union-of-literal
  arithmetic (e.g. `1 + if(true){2} else {3}`)
- Make widen_fresh() recursive into Union, List, Map, Optional so container
  literals at unannotated let-bindings widen correctly (e.g. `[1,2,3]` infers
  as int[] instead of (1|2|3)[])
- Fix propagate_copies in MIR cleanup to guard against inlining phi-like
  temps that have multiple definitions
- Add subtype unit tests for union-of-literals scenarios
- Un-ignore 6 tests across soundness, if_else, and match_basics

Author: hellovai

baml_language/crates/baml_compiler2_mir/src/cleanup.rs

@@ -275,6 +275,51 @@ fn collect_place_index_locals(body: &MirFunctionBody) -> HashSet<Local> {
     set
 }
 
+/// Count definition sites (assignments) for each local across all blocks.
+///
+/// A local is "single-definition" if it appears as an assignment destination
+/// in exactly one statement. Locals defined in multiple branches (e.g., a temp
+/// that is assigned in both arms of an if-else) have a count > 1 and must not
+/// be constant-propagated.
+fn count_local_defs(body: &MirFunctionBody) -> Vec<usize> {
+    let mut defs = vec![0usize; body.locals.len()];
+
+    for block in &body.blocks {
+        for stmt in &block.statements {
+            if let crate::StatementKind::Assign { destination, .. } = &stmt.kind {
+                // Walk the place to find the root local being defined.
+                let mut place = destination;
+                loop {
+                    match place {
+                        Place::Local(l) => {
+                            defs[l.0] += 1;
+                            break;
+                        }
+                        Place::Field { base, .. } => place = base,
+                        Place::Index { base, .. } => place = base,
+                    }
+                }
+            }
+        }
+        // Call destinations also count as definitions.
+        if let Some(Terminator::Call { destination, .. }) = &block.terminator {
+            let mut place = destination;
+            loop {
+                match place {
+                    Place::Local(l) => {
+                        defs[l.0] += 1;
+                        break;
+                    }
+                    Place::Field { base, .. } => place = base,
+                    Place::Index { base, .. } => place = base,
+                }
+            }
+        }
+    }
+
+    defs
+}
+
 fn count_local_uses(body: &MirFunctionBody) -> Vec<usize> {
     let mut uses = vec![0usize; body.locals.len()];
 
@@ -436,6 +481,7 @@ fn count_in_terminator(term: &Terminator, uses: &mut [usize]) {
 fn propagate_copies(body: &mut MirFunctionBody, arity: usize) {
     // Build substitution map: Local -> replacement Operand
     let uses = count_local_uses(body);
+    let defs = count_local_defs(body);
     // Locals used as the `index` field of a `Place::Index` cannot be replaced
     // with constants — that field is typed `Local`, not `Operand`. Collect them
     // so we can exclude them from constant inlining below.
@@ -448,8 +494,13 @@ fn propagate_copies(body: &mut MirFunctionBody, arity: usize) {
     // SAFETY: Only propagate unnamed locals (compiler temporaries from
     // lower_to_operand / builder.temp()). Named locals (user variables from
     // AstStmt::Let) can be reassigned via AstStmt::Assign or AstStmt::AssignOp,
-    // making propagation unsound. Unnamed temps are always fresh single-definition
-    // locals, so this is safe.
+    // making propagation unsound.
+    //
+    // We additionally require defs[dest] == 1 to guard against phi-like temps
+    // that are assigned in multiple branches (e.g., the result temp of an
+    // if-else used directly as an arithmetic operand). Such temps have a
+    // single use-site but two definition-sites; propagating the last-seen
+    // constant would silently use the wrong branch value.
     for block in &body.blocks {
         for stmt in &block.statements {
             if let crate::StatementKind::Assign {
@@ -462,6 +513,11 @@ fn propagate_copies(body: &mut MirFunctionBody, arity: usize) {
                     continue;
                 }
 
+                // Skip locals with multiple definition sites (phi-like).
+                if defs[dest.0] != 1 {
+                    continue;
+                }
+
                 match operand {
                     Operand::Copy(Place::Local(src)) if src.0 >= 1 && src.0 <= arity => {
                         // Copy of param — substitute
@@ -470,9 +526,9 @@ fn propagate_copies(body: &mut MirFunctionBody, arity: usize) {
                     Operand::Constant(c)
                         if uses[dest.0] == 1 && !used_as_place_index.contains(dest) =>
                     {
-                        // Single-use constant — inline. Skip locals that appear
-                        // as a Place::Index index, since that position can only
-                        // hold a Local, not a Constant.
+                        // Single-use, single-definition constant — inline. Skip
+                        // locals that appear as a Place::Index index, since that
+                        // position can only hold a Local, not a Constant.
                         subst.insert(*dest, Operand::Constant(c.clone()));
                     }
                     _ => {}

baml_language/crates/baml_compiler2_tir/src/builder.rs

@@ -3644,13 +3644,36 @@ impl<'db> TypeInferenceBuilder<'db> {
     /// String concatenation is only valid for `Add`; other arithmetic ops on
     /// strings are invalid and return `Unknown` (triggering an error upstream).
     fn infer_arithmetic(op: baml_compiler2_ast::BinaryOp, lhs: &Ty, rhs: &Ty) -> Ty {
-        let base_ty = |ty: &Ty| -> Option<PrimitiveType> {
+        fn promote(a: PrimitiveType, b: &PrimitiveType) -> Option<PrimitiveType> {
+            if a == *b {
+                return Some(a);
+            }
+            match (&a, &b) {
+                (PrimitiveType::Int, PrimitiveType::Float)
+                | (PrimitiveType::Float, PrimitiveType::Int) => Some(PrimitiveType::Float),
+                _ => None,
+            }
+        }
+
+        fn base_ty(ty: &Ty) -> Option<PrimitiveType> {
             match ty {
                 Ty::Primitive(p, _) => Some(p.clone()),
                 Ty::Literal(lit, _, _) => Some(PrimitiveType::from_literal(lit)),
+                Ty::Union(members, _) => {
+                    let mut result: Option<PrimitiveType> = None;
+                    for m in members {
+                        let p = base_ty(m)?;
+                        result = Some(match result {
+                            None => p,
+                            Some(existing) => promote(existing, &p)?,
+                        });
+                    }
+                    result
+                }
                 _ => None,
             }
-        };
+        }
+
         match (base_ty(lhs), base_ty(rhs)) {
             (Some(PrimitiveType::Float), _) | (_, Some(PrimitiveType::Float)) => {
                 Ty::Primitive(PrimitiveType::Float, TyAttr::default())

baml_language/crates/baml_compiler2_tir/src/normalize.rs

@@ -1543,4 +1543,59 @@ mod tests {
             "map<string, int> (evolving)"
         );
     }
+
+    // ── Literal-union subtype tests ─────────────────────────────────────────
+
+    #[test]
+    fn test_union_of_int_literals_subtype_of_int() {
+        let aliases = HashMap::new();
+        let sub = Ty::Union(
+            vec![
+                Ty::Literal(LiteralValue::Int(1), Freshness::Fresh, TyAttr::default()),
+                Ty::Literal(LiteralValue::Int(2), Freshness::Fresh, TyAttr::default()),
+                Ty::Literal(LiteralValue::Int(3), Freshness::Fresh, TyAttr::default()),
+            ],
+            TyAttr::default(),
+        );
+        let sup = Ty::Primitive(PrimitiveType::Int, TyAttr::default());
+        assert!(is_subtype_of(&sub, &sup, &aliases));
+    }
+
+    #[test]
+    fn test_list_of_literal_union_subtype_of_list_int() {
+        let aliases = HashMap::new();
+        let sub = Ty::List(
+            Box::new(Ty::Union(
+                vec![
+                    Ty::Literal(LiteralValue::Int(1), Freshness::Fresh, TyAttr::default()),
+                    Ty::Literal(LiteralValue::Int(2), Freshness::Fresh, TyAttr::default()),
+                ],
+                TyAttr::default(),
+            )),
+            TyAttr::default(),
+        );
+        let sup = Ty::List(
+            Box::new(Ty::Primitive(PrimitiveType::Int, TyAttr::default())),
+            TyAttr::default(),
+        );
+        assert!(is_subtype_of(&sub, &sup, &aliases));
+    }
+
+    #[test]
+    fn test_union_of_int_and_float_literals_subtype_of_float() {
+        let aliases = HashMap::new();
+        let sub = Ty::Union(
+            vec![
+                Ty::Literal(LiteralValue::Int(1), Freshness::Fresh, TyAttr::default()),
+                Ty::Literal(
+                    LiteralValue::Float("2.0".to_string()),
+                    Freshness::Fresh,
+                    TyAttr::default(),
+                ),
+            ],
+            TyAttr::default(),
+        );
+        let sup = Ty::Primitive(PrimitiveType::Float, TyAttr::default());
+        assert!(is_subtype_of(&sub, &sup, &aliases));
+    }
 }

baml_language/crates/baml_compiler2_tir/src/ty.rs

@@ -282,6 +282,41 @@ pub enum Freshness {
 /// Re-export `baml_base::Literal` as `LiteralValue` for backward compatibility.
 pub type LiteralValue = baml_base::Literal;
 
+/// Flatten, deduplicate, and collapse a vec of widened types into a single `Ty`.
+///
+/// After `widen_fresh()` has run on each union member, multiple members may
+/// have widened to the same primitive (e.g. `[Literal(1,Fresh), Literal(2,Fresh)]`
+/// both become `Primitive(Int)`). This helper deduplicates and collapses:
+/// - Flattens nested unions one level
+/// - Deduplicates by `PartialEq`
+/// - Unwraps singletons
+fn dedup_and_collapse(types: Vec<Ty>) -> Ty {
+    let mut members: Vec<Ty> = Vec::new();
+    for ty in types {
+        match ty {
+            Ty::Union(inner, _) => {
+                for m in inner {
+                    if !members.contains(&m) {
+                        members.push(m);
+                    }
+                }
+            }
+            _ => {
+                if !members.contains(&ty) {
+                    members.push(ty);
+                }
+            }
+        }
+    }
+    match members.len() {
+        0 => Ty::Never {
+            attr: TyAttr::default(),
+        },
+        1 => members.into_iter().next().unwrap(),
+        _ => Ty::Union(members, TyAttr::default()),
+    }
+}
+
 impl Ty {
     /// Access the `TyAttr` on this type.
     pub fn attr(&self) -> &TyAttr {
@@ -343,12 +378,26 @@ impl Ty {
     ///
     /// Called at mutable binding sites (`let` without annotation).
     /// Regular (non-fresh) literals pass through unchanged.
+    ///
+    /// Recurses into `Union`, `List`, `Map`, and `Optional` so that compound
+    /// types like `(1 | 2 | 3)[]` widen to `int[]` at unannotated bindings.
     #[must_use]
     pub fn widen_fresh(self) -> Ty {
         match self {
             Ty::Literal(lit, Freshness::Fresh, attr) => {
                 Ty::Primitive(PrimitiveType::from_literal(&lit), attr)
             }
+            Ty::Union(members, _attr) => {
+                let widened: Vec<Ty> = members.into_iter().map(Ty::widen_fresh).collect();
+                dedup_and_collapse(widened)
+            }
+            Ty::List(inner, attr) => Ty::List(Box::new((*inner).widen_fresh()), attr),
+            Ty::Map(k, v, attr) => Ty::Map(
+                Box::new((*k).widen_fresh()),
+                Box::new((*v).widen_fresh()),
+                attr,
+            ),
+            Ty::Optional(inner, attr) => Ty::Optional(Box::new((*inner).widen_fresh()), attr),
             other => other,
         }
     }

baml_language/crates/baml_tests/projects/literal_union_arithmetic/main.baml

@@ -0,0 +1,33 @@
+function ifElseAdd() -> int {
+  1 + if (true) { 2 } else { 3 }
+}
+
+function matchAdd(x: int) -> int {
+  match (x) {
+    1 => 10,
+    2 => 20,
+    _ => 0
+  } + 1
+}
+
+function bothSidesUnion() -> int {
+  (if (true) { 1 } else { 2 }) + (if (false) { 3 } else { 4 })
+}
+
+function subtractUnion() -> int {
+  100 - if (true) { 7 } else { 3 }
+}
+
+function loopMatchAccum() -> int {
+  let sum = 0;
+  let i = 0;
+  while (i < 3) {
+    sum = sum + match (i) {
+      0 => 10,
+      1 => 20,
+      _ => 30
+    };
+    i = i + 1;
+  }
+  return sum;
+}

baml_language/crates/baml_tests/projects/literal_union_widening/main.baml

@@ -0,0 +1,19 @@
+function arrayOfInts() -> int[] {
+  let x = [1, 2, 3];
+  return x;
+}
+
+function mapOfInts() -> map<string, int> {
+  let m = {"a": 1, "b": 2};
+  return m;
+}
+
+function nestedArray() -> int[][] {
+  let x = [[1, 2], [3, 4]];
+  return x;
+}
+
+function annotatedLiteralUnion() -> int {
+  let x: 1 | 2 = 1;
+  return x;
+}