7df4164 | plum | pyrossh

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

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

7df4164 — Peter John 2026-11-01T18:28:00+05:30

T3: checker coercions

Files changed (3) hide show

plum-checker/src/lib.rs +12 -49
plum-checker/src/unions.rs +132 -0
plum-checker/tests/unions_coerce_test.rs +97 -0

plum-checker/src/lib.rs CHANGED Viewed

@@ -3,6 +3,7 @@
 #![allow(non_snake_case)]
 pub mod types;
+pub mod unions;
 use plum_core::ast;
 use std::cell::{Cell, RefCell};
@@ -830,6 +831,8 @@ pub fn unify(t1: &PlumType, t2: &PlumType, ctx: &CheckCtx) -> Result<(), String>
             }
             unify(r1, r2, ctx)
         }
+        // A member value converts to its union enum (`Circle` where `Shape` is expected).
+        _ if unions::wrapVariantFor(t1, t2, ctx).is_some() => Ok(()),
         _ => Err(format!("type mismatch: expected {}, found {}", t1, t2)),
     }
 }
@@ -965,21 +968,6 @@ pub fn buildGlobalTables(source: &ast::Source) -> (TypeEnv, MethodEnv, EnumVaria
     let mut enum_variants: EnumVariants = BTreeMap::new();
     let mut enum_params: EnumParams = BTreeMap::new();
-    // First pass: find every "record-shaped" enum — exactly one variant, whose
-    // name equals the enum's own name (the `type X = ...` replacement shape,
-    // e.g. `enum Cat = | Cat(name: Str)`) — so a bare variant elsewhere naming
-    // one can wrap it as sugar (see the bare-wrap arm below), regardless of
-    // which declaration appears first in the file.
-    let record_shaped: std::collections::BTreeSet<&str> = source
-        .items
-        .iter()
-        .filter_map(|item| match item {
-            ast::Item::Enum(e) if e.variants.len() == 1 && e.variants[0].name == e.name => {
-                Some(e.name.as_str())
-            }
-            _ => None,
-        })
-        .collect();
     for item in &source.items {
         match item {
             ast::Item::Enum(e) => {
@@ -999,19 +987,6 @@ pub fn buildGlobalTables(source: &ast::Source) -> (TypeEnv, MethodEnv, EnumVaria
                             shared_field_types.clone(),
                             e.params.iter().map(|p| p.name.clone()).collect(),
                         )
-                    } else if v.fields.is_empty()
-                        && e.name != v.name
-                        && record_shaped.contains(v.name.as_str())
-                    {
-                        // Bare-type variant sugar: `| Cat` (no `[...]`) naming a
-                        // DIFFERENT, already-declared record-shaped enum means "this
-                        // variant wraps one payload field of that whole record" —
-                        // shorthand for the equivalent `| Cat[Cat]`, so `enum Animal
-                        // = | Cat | Dog` needs no separate wrapper tag distinct from
-                        // the record it holds. Excludes `e.name == v.name` (a plain
-                        // payload-free singleton like `enum Foo = | Foo`, which is
-                        // already exactly what it says — nothing else to wrap).
-                        (vec![PlumType::TNamed(v.name.clone())], Vec::new())
                     } else {
                         (
                             v.fields.iter().map(conv).collect(),
@@ -2738,8 +2713,8 @@ fn checkMatch(
 }
 /// Type-checks a `ClassCall`'s fields against its class's declared field types and
-/// returns the BARE class type (`TNamed(call.type_name)`, never the enum a
+/// returns the BARE class type (`TNamed(call.type_name)`; it converts to a
-/// bare-type variant sugar wraps it in — see `bareVariantWrapType`). Split out from
+/// union enum on demand, see `unions`). Split out from
 /// `inferExpr`'s `ClassCall` arm so the `FnCall` variant-arg check below can type
 /// the class's OWN inner construction (`FantasyBook(...)` nested inside the
 /// desugared `FantasyBook(FantasyBook(...))`) without re-triggering the wrap.
@@ -2900,20 +2875,6 @@ fn inferClassCallRaw(
     }
 }
-/// A class construction call (`FantasyBook(title: ..., ...)`) whose name is ALSO a
-/// bare-type enum variant (`enum Book = | FantasyBook`, see `buildGlobalTables`)
-/// infers as the enum it wraps, not the bare class — mirroring the AST-level
-/// desugaring codegen does (`FantasyBook(...)`
-/// there literally becomes `Book`'s `FantasyBook(FantasyBook(...))` constructor).
-fn bareVariantWrapType(type_name: &str, ctx: &CheckCtx) -> PlumType {
-    match ctx.enum_variants.get(type_name) {
-        Some(info) if info.field_types == [PlumType::TNamed(type_name.to_string())] => {
-            PlumType::TNamed(info.enum_name.clone())
-        }
-        _ => PlumType::TNamed(type_name.to_string()),
-    }
-}
 /// Checks a single case pattern against the type of the subject it matches, binding any
 /// new names it introduces into `env`. Constructor-payload sub-patterns (`Next(x)`) bind
 /// against that variant's declared field types (see `EnumVariantInfo::field_types`).
@@ -3502,6 +3463,7 @@ pub fn inferExprWith(
     if let Ok(t) = &r {
         recordType(expr, t);
         recordErrCoercion(expr, t, expected, ctx);
+        unions::recordMemberWrap(expr, t, expected, ctx);
         // A call returning another fn's inferred `!T` is an edge of the call graph.
         if let PlumType::TErr(_, ErrSet::Inferred(k)) = t {
             if !k.is_empty() {
@@ -3969,6 +3931,11 @@ fn inferExprInner(
                     if matches!(tt, PlumType::TErr(..)) { &et } else { &tt }
                 ));
             }
+            if tt != et {
+                if let Some(joined) = unions::joinBranches(&t.then, &tt, &t.else_, &et, ctx) {
+                    return joined.map_err(|e| format!("inline if branches: {}", e));
+                }
+            }
             unify(&tt, &et, ctx).map_err(|e| format!("inline if branches: {}", e))?;
             Ok(tt)
         }
@@ -4075,11 +4042,7 @@ fn inferExprInner(
             }
         }
         ast::Expr::ClassCall(call) => {
-            let bare = inferClassCallRaw(call, env, ctx, expected)?;
+            inferClassCallRaw(call, env, ctx, expected)
-            Ok(match bare {
-                PlumType::TNamed(n) => bareVariantWrapType(&n, ctx),
-                other => other,
-            })
         }
         ast::Expr::Attribute(attr) => {
             let obj_ty = resolveRecorded(&inferExpr(&attr.object, env, ctx)?);

plum-checker/src/unions.rs ADDED Viewed

@@ -0,0 +1,132 @@
+//! Union enums in the checker: a member value converts to its enum, and
+//! branches of different members join to their common enum.
+use crate::types::{self, PlumType};
+use crate::{unify, CheckCtx};
+use plum_core::ast;
+use plum_core::unions as plan;
+use std::cell::Cell;
+use std::collections::{BTreeMap, BTreeSet};
+thread_local! {
+    /// Set while a member field is being unified, so conversions do not nest.
+    static IN_WRAP: Cell<bool> = const { Cell::new(false) };
+}
+fn enumOf(ty: &PlumType) -> Option<(&str, &[PlumType])> {
+    match ty {
+        PlumType::TNamed(n) => Some((n, &[])),
+        PlumType::TApp(n, args) => Some((n, args)),
+        _ => None,
+    }
+}
+/// `(hidden variant name, member type)` of every member of `enumName`, with the
+/// enum's generics substituted by `args` when they match in number.
+fn members(enumName: &str, args: &[PlumType], ctx: &CheckCtx) -> Vec<(String, PlumType)> {
+    let prefix = format!("{enumName}__");
+    ctx.enum_variants
+        .iter()
+        .filter(|(n, i)| {
+            i.enum_name == enumName
+                && n.starts_with(&prefix)
+                && i.field_types.len() == 1
+                && i.field_names.is_empty()
+        })
+        .map(|(n, i)| {
+            let f = &i.field_types[0];
+            let f = if !args.is_empty() && args.len() == i.generics.len() {
+                let s: BTreeMap<String, PlumType> =
+                    i.generics.iter().cloned().zip(args.iter().cloned()).collect();
+                types::substitute(f, &s)
+            } else {
+                f.clone()
+            };
+            (n.clone(), f)
+        })
+        .collect()
+}
+/// The hidden variant that converts `actual` into the union enum `expected`.
+pub fn wrapVariantFor(expected: &PlumType, actual: &PlumType, ctx: &CheckCtx) -> Option<String> {
+    if IN_WRAP.with(|c| c.get()) {
+        return None;
+    }
+    if matches!(actual, PlumType::TVar(_) | PlumType::TParam(_) | PlumType::TErr(..)) {
+        return None;
+    }
+    let (name, args) = enumOf(expected)?;
+    if enumOf(actual).is_some_and(|(a, _)| a == name) {
+        return None;
+    }
+    let ms = members(name, args, ctx);
+    if ms.is_empty() {
+        return None;
+    }
+    IN_WRAP.with(|c| c.set(true));
+    let found = ms.into_iter().find(|(_, t)| unify(t, actual, ctx).is_ok());
+    IN_WRAP.with(|c| c.set(false));
+    found.map(|(v, _)| v)
+}
+/// Records a wrap when `expr` (of type `actual`) lands where `expected` is a union enum.
+pub fn recordMemberWrap(expr: &ast::Expr, actual: &PlumType, expected: Option<&PlumType>, ctx: &CheckCtx) {
+    let Some(expected) = expected else { return };
+    if let Some(v) = wrapVariantFor(expected, actual, ctx) {
+        plan::recordWrap(expr as *const ast::Expr as usize, v);
+    }
+}
+/// Enums that `ty` is a member of, plus `ty`'s own enum.
+fn candidates(ty: &PlumType, ctx: &CheckCtx) -> BTreeSet<String> {
+    let mut out = BTreeSet::new();
+    for (vn, info) in ctx.enum_variants.iter() {
+        if info.field_types.len() == 1
+            && info.field_names.is_empty()
+            && vn.starts_with(&format!("{}__", info.enum_name))
+        {
+            let ok = match (&info.field_types[0], ty) {
+                (PlumType::TApp(a, _), PlumType::TApp(b, _)) | (PlumType::TNamed(a), PlumType::TNamed(b)) => a == b,
+                (a, b) => a == b,
+            };
+            if ok {
+                out.insert(info.enum_name.clone());
+            }
+        }
+    }
+    if let Some((n, _)) = enumOf(ty) {
+        if !members(n, &[], ctx).is_empty() {
+            out.insert(n.to_string());
+        }
+    }
+    out
+}
+/// Joins two branch types through their one common union enum. `None` when
+/// neither branch involves unions (the ordinary rules apply).
+pub fn joinBranches(
+    a: &ast::Expr,
+    ta: &PlumType,
+    b: &ast::Expr,
+    tb: &PlumType,
+    ctx: &CheckCtx,
+) -> Option<Result<PlumType, String>> {
+    let ca = candidates(ta, ctx);
+    let cb = candidates(tb, ctx);
+    if ca.is_empty() || cb.is_empty() {
+        return None;
+    }
+    let common: Vec<&String> = ca.intersection(&cb).collect();
+    if common.len() != 1 {
+        return Some(Err(format!(
+            "cannot join {} and {}: no single enum has both members",
+            ta, tb
+        )));
+    }
+    let e = PlumType::TNamed(common[0].clone());
+    for (x, tx) in [(a, ta), (b, tb)] {
+        if let Some(v) = wrapVariantFor(&e, tx, ctx) {
+            plan::recordWrap(x as *const ast::Expr as usize, v);
+        }
+    }
+    Some(Ok(e))
+}

plum-checker/tests/unions_coerce_test.rs ADDED Viewed

@@ -0,0 +1,97 @@
+#![allow(non_snake_case)]
+use plum_checker::checkSource;
+use plum_core::unions::{beginPlan, resolveMembers, takePlan, UnionPlan};
+use plum_core::AstParser;
+const DEFS: &str = "\
+record Circle(r: Float)
+record Square(s: Float)
+record Oval(a: Float)
+enum Shape = Circle | Square
+enum Round = Circle | Oval
+enum Json = Float | Str
+enum Tree[T] = Node[T] | Leaf
+record Node[T](v: T)
+";
+fn run(body: &str) -> (Result<(), Vec<String>>, UnionPlan) {
+    let src = format!("{DEFS}\n{body}\n");
+    let mut parser = tree_sitter::Parser::new();
+    parser.set_language(&tree_sitter_plum::LANGUAGE.into()).unwrap();
+    let tree = parser.parse(&src, None).unwrap();
+    assert!(!tree.root_node().has_error(), "parse errors:\n{src}");
+    let mut source = AstParser::new(&src).parseSource(tree.root_node());
+    let errs = resolveMembers(&mut source);
+    assert!(errs.is_empty(), "{errs:?}");
+    beginPlan();
+    let r = checkSource(&source).map_err(|e| e.into_iter().map(|x| x.message).collect());
+    (r, takePlan())
+}
+fn wraps(p: &UnionPlan) -> Vec<String> {
+    let mut v: Vec<String> = p.wraps.values().cloned().collect();
+    v.sort();
+    v
+}
+#[test]
+fn literal_into_json() {
+    let (r, p) = run("fun f() -> Json =\n  3.0");
+    r.unwrap();
+    assert_eq!(wraps(&p), vec!["Json__Float"]);
+}
+#[test]
+fn wrap_at_arg_return_push() {
+    let (r, p) = run("fun take(s: Shape) -> Float =\n  1.0\n\nfun go() -> Float =\n  take(Circle(r: 1.0))");
+    r.unwrap();
+    assert_eq!(wraps(&p), vec!["Shape__Circle"]);
+    let (r, p) = run("fun mk() -> Shape =\n  Square(s: 2.0)");
+    r.unwrap();
+    assert_eq!(wraps(&p), vec!["Shape__Square"]);
+    let (r, p) = run("fun many(xs: ...Shape) -> Int =\n  1\n\nfun go() -> Int =\n  many(Circle(r: 1.0), Square(s: 1.0))");
+    r.unwrap();
+    assert_eq!(wraps(&p), vec!["Shape__Circle", "Shape__Square"]);
+}
+#[test]
+fn wrap_reassign_and_generic() {
+    let (r, p) = run("fun pick() -> Shape =\n  Circle(r: 1.0)\n\nfun go() -> Int =\n  s := pick()\n  s = Square(s: 1.0)\n  1");
+    r.unwrap();
+    assert_eq!(wraps(&p), vec!["Shape__Circle", "Shape__Square"]);
+    let (r, p) = run("fun mk() -> Tree[Int] =\n  Node(v: 3)");
+    r.unwrap();
+    assert_eq!(wraps(&p), vec!["Tree__Node"]);
+}
+#[test]
+fn join_if_branches() {
+    let (r, p) = run("fun go(c: Bool) -> Int =\n  x := if c then Circle(r: 1.0) else Square(s: 1.0)\n  1");
+    r.unwrap();
+    assert_eq!(wraps(&p), vec!["Shape__Circle", "Shape__Square"]);
+}
+#[test]
+fn join_ambiguous() {
+    let (r, _) = run("fun go(c: Bool) -> Int =\n  x := if c then Square(s: 1.0) else Oval(a: 1.0)\n  1");
+    assert!(r.unwrap_err().join("|").contains("cannot join Square and Oval: no single enum has both members"));
+}
+#[test]
+fn no_wrap_same_type() {
+    let (r, p) = run("fun take(c: Circle) -> Int =\n  1\n\nfun go() -> Int =\n  take(Circle(r: 1.0))");
+    r.unwrap();
+    assert!(p.wraps.is_empty());
+    let (r, p) = run("fun mk() -> Circle =\n  Circle(r: 1.0)");
+    r.unwrap();
+    assert!(p.wraps.is_empty());
+}
+#[test]
+fn member_alone_is_its_own_type() {
+    let (r, p) = run("fun mk() -> Circle =\n  Circle(r: 1.0)");
+    r.unwrap();
+    assert!(p.wraps.is_empty());
+    let (r, _) = run("fun mk() -> Shape =\n  Oval(a: 1.0)");
+    assert!(r.is_err());
+}