e9976f3 | plum | pyrossh

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

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

e9976f3 — Peter John 2026-11-01T20:34:14+05:30

T5: checker auto-dispatch

Files changed (3) hide show

plum-checker/src/lib.rs +20 -5
plum-checker/src/unions.rs +105 -1
plum-checker/tests/unions_dispatch_test.rs +108 -0

plum-checker/src/lib.rs CHANGED Viewed

@@ -750,7 +750,7 @@ fn unifyErr(t1: &PlumType, t2: &PlumType, ctx: &CheckCtx) -> Result<(), String>
 pub fn unify(t1: &PlumType, t2: &PlumType, ctx: &CheckCtx) -> Result<(), String> {
     match (t1, t2) {
         (PlumType::TVar(_), _) | (_, PlumType::TVar(_)) => Ok(()),
-        (PlumType::TErr(..), _) | (_, PlumType::TErr(..)) => unifyErr(t1, t2, ctx),
+        (PlumType::TErr(..), _) | (_, PlumType::TErr(..)) => unions::nested(|| unifyErr(t1, t2, ctx)),
         (PlumType::TInt, PlumType::TInt) => Ok(()),
         (PlumType::TFloat, PlumType::TFloat) => Ok(()),
         (PlumType::TBool, PlumType::TBool) => Ok(()),
@@ -774,7 +774,7 @@ pub fn unify(t1: &PlumType, t2: &PlumType, ctx: &CheckCtx) -> Result<(), String>
         // Structural generics: same base, arguments unify pairwise.
         (PlumType::TApp(a, xs), PlumType::TApp(b, ys)) if a == b && xs.len() == ys.len() => {
             for (x, y) in xs.iter().zip(ys.iter()) {
-                unify(x, y, ctx).map_err(|e| {
+                unions::nested(|| unify(x, y, ctx)).map_err(|e| {
                     format!("type mismatch: expected {}, found {} ({})", t1, t2, e)
                 })?;
             }
@@ -827,9 +827,9 @@ pub fn unify(t1: &PlumType, t2: &PlumType, ctx: &CheckCtx) -> Result<(), String>
         }
         (PlumType::TFun(ps1, r1), PlumType::TFun(ps2, r2)) if ps1.len() == ps2.len() => {
             for (p1, p2) in ps1.iter().zip(ps2.iter()) {
-                unify(p1, p2, ctx)?;
+                unions::nested(|| unify(p1, p2, ctx))?;
             }
-            unify(r1, r2, ctx)
+            unions::nested(|| 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(()),
@@ -3245,6 +3245,7 @@ fn inferMethodCall(
     ctx: &CheckCtx,
     expected: Option<&PlumType>,
 ) -> Result<PlumType, String> {
+    let call_addr = CUR_EXPR.with(|c| c.get());
     // `e.name()` on an error value (`catch e`) is its text.
     if matches!(obj_ty, PlumType::TErrName(_)) && name == "name" && args.is_empty() {
         return Ok(PlumType::TStr);
@@ -3310,6 +3311,16 @@ fn inferMethodCall(
         }
         return Ok(PlumType::TVar("_".to_string()));
     };
+    if let Some(r) = unions::dispatchSig(obj_ty, name, call_addr, ctx) {
+        let (params, ret) = r?;
+        let mut binds = Binds::new();
+        if let Some(exp) = expected {
+            bindMatch(&ret, exp, &mut binds);
+        }
+        let label = format!("method '{}.{}'", class_name, name);
+        checkCallArgs(&label, &params, args, params.len(), env, ctx, &mut binds)?;
+        return Ok(substVars(&ret, &binds));
+    }
     let Some(sig) = ctx.methods.get(&(class_name.to_string(), name.to_string())) else {
         match plum_core::forms::seqMethodKind(name) {
             Some(false) => {
@@ -3679,7 +3690,11 @@ fn inferExprInner(
             // were never part of this check.
             for part in &s.parts {
                 if let ast::StringPart::Interp(e) = part {
+                    let ity = inferExprWith(e, env, ctx, None);
+                    if let Ok(t) = &ity {
+                        unions::recordInterpolation(e, t, ctx);
+                    }
-                    if let Ok(PlumType::TApp(n, targs)) = inferExprWith(e, env, ctx, None) {
+                    if let Ok(PlumType::TApp(n, targs)) = ity {
                         let gens = enumGenericsOf(&n, ctx);
                         if gens.len() == targs.len()
                             && ctx.methods.contains_key(&(n.clone(), "toStr".to_string()))

plum-checker/src/unions.rs CHANGED Viewed

@@ -10,6 +10,16 @@ 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) };
+    /// Nesting depth inside type arguments / fn types: conversion applies only at depth 0.
+    static DEPTH: Cell<usize> = const { Cell::new(0) };
+}
+/// Runs `f` (a unify of a nested type position), where no member conversion applies.
+pub fn nested<R>(f: impl FnOnce() -> R) -> R {
+    DEPTH.with(|d| d.set(d.get() + 1));
+    let r = f();
+    DEPTH.with(|d| d.set(d.get() - 1));
+    r
 }
 fn enumOf(ty: &PlumType) -> Option<(&str, &[PlumType])> {
@@ -53,7 +63,7 @@ pub fn isUnion(enumName: &str, ctx: &CheckCtx) -> bool {
 /// 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()) {
+    if IN_WRAP.with(|c| c.get()) || DEPTH.with(|d| d.get()) > 0 {
         return None;
     }
     if matches!(actual, PlumType::TVar(_) | PlumType::TParam(_) | PlumType::TErr(..)) {
@@ -278,3 +288,97 @@ pub fn copyWriteError(object: &ast::Expr, env: &TypeEnv) -> Option<String> {
         "{name} is a copy of the {member} inside {en}; assign it to a local to change it"
     ))
 }
+// ---- auto-dispatch -----------------------------------------------------------
+/// Every variant of `enumName` in tag order: `(variant name, member type or None for a tag)`.
+fn allVariants(enumName: &str, args: &[PlumType], ctx: &CheckCtx) -> Vec<(String, Option<PlumType>)> {
+    let prefix = format!("{enumName}__");
+    let mut vs: Vec<(&String, &EnumVariantInfo)> =
+        ctx.enum_variants.iter().filter(|(_, i)| i.enum_name == enumName).collect();
+    vs.sort_by_key(|(_, i)| i.tag);
+    let ms = members(enumName, args, ctx);
+    vs.into_iter()
+        .map(|(n, _)| {
+            let m = if n.starts_with(&prefix) {
+                ms.iter().find(|(v, _)| v == n).map(|(_, t)| t.clone())
+            } else {
+                None
+            };
+            (n.clone(), m)
+        })
+        .collect()
+}
+/// One member's instantiated method signature, if its type defines the method.
+fn memberSig(mt: &PlumType, name: &str, ctx: &CheckCtx) -> Option<PlumType> {
+    let recv = crate::methodReceiverName(mt)?;
+    let sig = ctx.methods.get(&(recv.clone(), name.to_string()))?;
+    let owner = match mt {
+        PlumType::TApp(n, args) if *n == recv => crate::appSubst(crate::enumGenericsOf(n, ctx), args),
+        _ => BTreeMap::new(),
+    };
+    Some(crate::instSigMap(sig, &owner).0)
+}
+fn memberLabel(variant: &str, enumName: &str) -> String {
+    variant.strip_prefix(&format!("{enumName}__")).unwrap_or(variant).to_string()
+}
+/// For `obj.name(..)` on a union enum with no method of its own: the common
+/// signature of the members' methods, and the dispatch is recorded. `None`
+/// when the receiver is not a union or defines `name` itself.
+pub fn dispatchSig(
+    obj: &PlumType,
+    name: &str,
+    callAddr: usize,
+    ctx: &CheckCtx,
+) -> Option<Result<(Vec<PlumType>, PlumType), String>> {
+    let (en, args) = enumOf(obj)?;
+    if ctx.methods.contains_key(&(en.to_string(), name.to_string())) || !isUnion(en, ctx) {
+        return None;
+    }
+    let mut first: Option<(String, PlumType)> = None;
+    let mut variants = Vec::new();
+    for (v, m) in allVariants(en, args, ctx) {
+        let label = memberLabel(&v, en);
+        let Some(mt) = m else {
+            return Some(Err(format!("{en} has no method '{name}': member {label} is a tag")));
+        };
+        let Some(sig) = memberSig(&mt, name, ctx) else {
+            return Some(Err(format!("{en} has no method '{name}': member {label} doesn't define it")));
+        };
+        match &first {
+            None => first = Some((label, sig)),
+            Some((l0, s0)) => {
+                if unify(s0, &sig, ctx).is_err() {
+                    return Some(Err(format!(
+                        "{en}.{name}: members disagree: {l0}.{name} is {s0}, {label}.{name} is {sig}"
+                    )));
+                }
+            }
+        }
+        variants.push(v);
+    }
+    let (_, sig) = first?;
+    plan::recordDispatch(callAddr, variants);
+    let PlumType::TFun(params, ret) = sig else { return None };
+    Some(Ok((params, *ret)))
+}
+/// `"{e}"` on a union enum: records a `toStr` dispatch when every member has `toStr`.
+pub fn recordInterpolation(e: &ast::Expr, ty: &PlumType, ctx: &CheckCtx) {
+    let Some((en, args)) = enumOf(ty) else { return };
+    if ctx.methods.contains_key(&(en.to_string(), "toStr".to_string())) || !isUnion(en, ctx) {
+        return;
+    }
+    let vs = allVariants(en, args, ctx);
+    let mut out = Vec::new();
+    for (v, m) in vs {
+        match m {
+            Some(mt) if memberSig(&mt, "toStr", ctx).is_some() => out.push(v),
+            _ => return,
+        }
+    }
+    plan::recordDispatch(e as *const ast::Expr as usize, out);
+}

plum-checker/tests/unions_dispatch_test.rs ADDED Viewed

@@ -0,0 +1,108 @@
+#![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) =
+  fun area(self) -> Float = self.r
+  fun toStr(self) -> Str = \"c\"
+record Rectangle(w: Float) =
+  fun area(self) -> Float = self.w
+  fun toStr(self) -> Str = \"r\"
+record Odd(w: Float) =
+  fun area(self) -> Int = 1
+record NoArea(w: Float)
+enum Shape = Circle | Rectangle
+enum Own =
+  | Circle
+  | Rectangle
+  fun area(self) -> Float = 0.0
+enum WithTag = Circle | Plain
+enum Missing = Circle | NoArea
+enum Clash = Circle | Odd
+record Wrap[T](v: T) =
+  fun toStr(self) -> Str = \"w\"
+enum Js = Circle | Wrap[Js]
+";
+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 errs(r: Result<(), Vec<String>>) -> String {
+    r.unwrap_err().join(" | ")
+}
+#[test]
+fn dispatch_ok() {
+    let (r, p) = run("fun f(s: Shape) -> Float =\n  s.area()");
+    r.unwrap();
+    assert_eq!(p.dispatches.len(), 1);
+    let v = p.dispatches.values().next().unwrap();
+    assert_eq!(v, &vec!["Shape__Circle".to_string(), "Shape__Rectangle".to_string()]);
+    let (r, _) = run("fun f(s: Shape) -> Int =\n  x := s.area()\n  take(x)\n\nfun take(i: Int) -> Int =\n  i");
+    assert!(r.is_err(), "area() is Float, not Int");
+}
+#[test]
+fn own_method_wins() {
+    let (r, p) = run("fun f(s: Own) -> Float =\n  s.area()");
+    r.unwrap();
+    assert!(p.dispatches.is_empty());
+}
+#[test]
+fn tag_member_error() {
+    let (r, _) = run("fun f(s: WithTag) -> Float =\n  s.area()");
+    assert!(errs(r).contains("WithTag has no method 'area': member Plain is a tag"));
+}
+#[test]
+fn missing_method_error() {
+    let (r, _) = run("fun f(s: Missing) -> Float =\n  s.area()");
+    assert!(errs(r).contains("Missing has no method 'area': member NoArea doesn't define it"));
+}
+#[test]
+fn signature_mismatch_error() {
+    let (r, _) = run("fun f(s: Clash) -> Float =\n  s.area()");
+    let e = errs(r);
+    assert!(e.contains("Clash.area: members disagree: Circle.area is "), "{e}");
+    assert!(e.contains(", Odd.area is "), "{e}");
+}
+#[test]
+fn interpolation_toStr() {
+    let (r, p) = run("fun f(s: Shape) -> Str =\n  \"{s}\"");
+    r.unwrap();
+    assert_eq!(p.dispatches.len(), 1);
+    let (r, p) = run("fun f(s: Missing) -> Str =\n  \"{s}\"");
+    let _ = r;
+    assert!(p.dispatches.is_empty());
+}
+#[test]
+fn generic_member_dispatch() {
+    let (r, p) = run("fun f(j: Js) -> Str =\n  j.toStr()");
+    r.unwrap();
+    assert_eq!(p.dispatches.values().next().unwrap().len(), 2);
+}