8a21edd | plum | pyrossh

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

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

8a21edd — Peter John 2026-11-01T20:31:34+05:30

T4: checker patterns + narrowing

Files changed (3) hide show

plum-checker/src/lib.rs +43 -9
plum-checker/src/unions.rs +148 -0
plum-checker/tests/unions_pattern_test.rs +109 -0

plum-checker/src/lib.rs CHANGED Viewed

@@ -629,7 +629,7 @@ const RETURN_KEY: &str = "$return";
 const STAR_KEY: &str = "$star";
 /// `TypeEnv` key prefix marking a name as a top-level function (a local of the
 /// same name shadows it instead of being a redeclaration).
-const GLOBAL_KEY: &str = "$global:";
+pub(crate) const GLOBAL_KEY: &str = "$global:";
 pub fn plumTypeFromAst(ty: &ast::Type) -> PlumType {
     if ty.name == ast::ERROR_UNION_TYPE_NAME && !ty.generics.is_empty() {
         return errUnionFromAst(ty, plumTypeFromAst);
@@ -910,6 +910,8 @@ pub struct EnumVariantInfo {
     /// The owning enum's generic parameter names (`[K, V]` for `Map`) —
     /// empty for a non-generic enum and for non-generic uses.
     pub generics: Vec<String>,
+    /// The owning declaration is a `struct` (copy semantics).
+    pub is_value: bool,
 }
 /// Enum variant name -> its info, e.g. `"Red" -> { enum_name: "Color", tag: 1, field_types: [] }`.
 pub type EnumVariants = BTreeMap<String, EnumVariantInfo>;
@@ -1002,6 +1004,7 @@ pub fn buildGlobalTables(source: &ast::Source) -> (TypeEnv, MethodEnv, EnumVaria
                             field_names,
                             values: v.values.clone(),
                             generics: enumGenericParams(e),
+                            is_value: e.isValue,
                         },
                     );
                 }
@@ -2243,6 +2246,10 @@ fn checkStmtInner(
                     }
                     ast::AssignTarget::Field(object, field_name) => {
                         let label = format!("{}.{}", describeTargetObject(object), field_name);
+                        if let Some(msg) = unions::copyWriteError(object, env) {
+                            errors.push(CheckError::at(fn_pos.0, fn_pos.1, format!("fn '{}': assign '{}': {}", fn_name, label, msg)));
+                            continue;
+                        }
                         let field_expected = inferExpr(object, env, ctx)
                             .ok()
                             .and_then(|o| fieldTypeOf(&o, field_name, ctx).ok());
@@ -2451,6 +2458,8 @@ enum CoverFail {
 struct CoverCtor {
     name: String,
     fields: Vec<PlumType>,
+    /// A union member: covered by `Name` or `Name(..)` with irrefutable sub-patterns.
+    member: bool,
 }
 static COVER_WILDCARD: ast::CasePattern = ast::CasePattern::Wildcard;
@@ -2486,6 +2495,7 @@ fn coverCtors(ty: &PlumType, ctx: &CheckCtx) -> Option<(String, Vec<CoverCtor>)>
                 .map(|n| CoverCtor {
                     name: n.to_string(),
                     fields: Vec::new(),
+                    member: false,
                 })
                 .collect(),
         )),
@@ -2496,6 +2506,7 @@ fn coverCtors(ty: &PlumType, ctx: &CheckCtx) -> Option<(String, Vec<CoverCtor>)>
                 .map(|n| CoverCtor {
                     name: n.clone(),
                     fields: Vec::new(),
+                    member: false,
                 })
                 .collect(),
         )),
@@ -2512,7 +2523,10 @@ fn coverCtors(ty: &PlumType, ctx: &CheckCtx) -> Option<(String, Vec<CoverCtor>)>
             Some((
                 enum_name.to_string(),
                 vs.into_iter()
+                    .map(|(k, i)| match unions::memberBase(k, i) {
+                        Some(base) => CoverCtor { name: base, fields: Vec::new(), member: true },
-                    .map(|(k, i)| CoverCtor {
+                        None => CoverCtor {
+                        member: false,
                         name: k.to_string(),
                         // A discriminant enum's field types are its shared
                         // params, which patterns never destructure.
@@ -2521,6 +2535,7 @@ fn coverCtors(ty: &PlumType, ctx: &CheckCtx) -> Option<(String, Vec<CoverCtor>)>
                         } else {
                             variantFields(i, ty)
                         },
+                        },
                     })
                     .collect(),
             ))
@@ -2572,7 +2587,11 @@ fn coverRows<'a>(
                 if n != c.name {
                     continue;
                 }
+                if c.member {
+                    if !fs.iter().all(patCatchAll) {
+                        continue;
+                    }
-                if let ast::CasePattern::Name(_) = r[0] {
+                } else if let ast::CasePattern::Name(_) = r[0] {
                     // Bare `Circle` against a payload variant: tag check only.
                     row.extend(std::iter::repeat(&COVER_WILDCARD).take(c.fields.len()));
                 } else if fs.len() == c.fields.len() {
@@ -2661,12 +2680,19 @@ fn checkMatch(
         let mut case_env = env.clone();
         if case.patterns.len() == subject_types.len() {
             for (pat, sty) in case.patterns.iter().zip(subject_types.iter()) {
+                let narrow = if case.patterns.len() == 1 {
+                    unions::narrowTarget(&m.subjects[0], pat, sty, env, ctx)
+                } else {
+                    None
+                };
                 if let Err(msg) = checkPattern(pat, sty, &mut case_env, ctx) {
                     errors.push(CheckError::at(
                         fn_pos.0,
                         fn_pos.1,
                         format!("fn '{}': match case: {}", fn_name, msg),
                     ));
+                } else {
+                    unions::applyNarrow(narrow, case, &mut case_env, ctx);
                 }
             }
         } else {
@@ -2878,7 +2904,7 @@ fn inferClassCallRaw(
 /// 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`).
-fn checkPattern(
+pub(crate) fn checkPattern(
     pat: &ast::CasePattern,
     subject_ty: &PlumType,
     env: &mut TypeEnv,
@@ -2897,6 +2923,9 @@ fn checkPattern(
             }
         }
     }
+    if let Some(r) = unions::checkMemberPattern(pat, subject_ty, env, ctx) {
+        return r;
+    }
     match pat {
         ast::CasePattern::Wildcard => Ok(()),
         ast::CasePattern::Int(_) => unify(subject_ty, &PlumType::TInt, ctx),
@@ -3088,10 +3117,11 @@ pub(crate) fn lookupFieldType(
                 // codegen will catch a genuine mismatch, exactly as it always has).
                 // Distinguish by whether any variant belonging to this enum name
                 // declares named fields at all.
+                [] if unions::isUnion(class_name, ctx)
-                [] if ctx
+                    || ctx
-                    .enum_variants
+                        .enum_variants
-                    .values()
+                        .values()
-                    .any(|info| info.enum_name == class_name && !info.field_names.is_empty()) =>
+                        .any(|info| info.enum_name == class_name && !info.field_names.is_empty()) =>
                 {
                     Err(format!(
                         "no field '{}' on type '{}'",
@@ -3510,7 +3540,11 @@ fn bindCondEnv(cond: &ast::Expr, env: &TypeEnv, ctx: &CheckCtx) -> Result<TypeEn
             };
             benv.insert(n.clone(), TypeScheme::mono(t));
         }
+        pat => {
+            let narrow = unions::narrowTarget(&b.value, pat, &ty, env, ctx);
-        pat => checkPattern(pat, &ty, &mut benv, ctx)?,
+            checkPattern(pat, &ty, &mut benv, ctx)?;
+            unions::applyNarrow(narrow, &b.case, &mut benv, ctx);
+        }
     }
     if let Some(g) = &b.case.guard {
         let gt = inferExpr(g, &benv, ctx)?;

plum-checker/src/unions.rs CHANGED Viewed

@@ -46,6 +46,11 @@ fn members(enumName: &str, args: &[PlumType], ctx: &CheckCtx) -> Vec<(String, Pl
         .collect()
 }
+/// Whether `enumName` has any type members.
+pub fn isUnion(enumName: &str, ctx: &CheckCtx) -> bool {
+    !members(enumName, &[], ctx).is_empty()
+}
 /// 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()) {
@@ -130,3 +135,146 @@ pub fn joinBranches(
     }
     Some(Ok(e))
 }
+// ---- patterns and narrowing ------------------------------------------------
+use crate::types::{TypeEnv, TypeScheme};
+use crate::{checkPattern, EnumVariantInfo};
+thread_local! {
+    /// Address of a top-level pattern that narrows (it records a narrowing, not a pattern).
+    static NARROW_TOP: Cell<usize> = const { Cell::new(0) };
+}
+const COPY_KEY: &str = "$union-copy:";
+/// The member base name of a hidden variant (`Shape__Circle` -> `Circle`).
+pub fn memberBase(variant: &str, info: &EnumVariantInfo) -> Option<String> {
+    if info.field_types.len() != 1 || !info.field_names.is_empty() {
+        return None;
+    }
+    variant
+        .strip_prefix(&format!("{}__", info.enum_name))
+        .map(|b| b.to_string())
+}
+/// `(hidden variant, member type)` of the member of `subject`'s enum named `name`.
+fn memberNamed(subject: &PlumType, name: &str, ctx: &CheckCtx) -> Option<(String, PlumType)> {
+    let (en, args) = enumOf(subject)?;
+    let prefix = format!("{en}__");
+    members(en, args, ctx)
+        .into_iter()
+        .find(|(v, _)| v.strip_prefix(&prefix) == Some(name))
+}
+/// Checks a pattern that names a member of the subject's union enum. `None`
+/// when it does not (the ordinary rules apply).
+pub fn checkMemberPattern(
+    pat: &ast::CasePattern,
+    subject: &PlumType,
+    env: &mut TypeEnv,
+    ctx: &CheckCtx,
+) -> Option<Result<(), String>> {
+    let addr = pat as *const ast::CasePattern as usize;
+    match pat {
+        ast::CasePattern::Name(n) if n.chars().next().is_some_and(|c| c.is_uppercase()) => {
+            let (v, _) = memberNamed(subject, n, ctx)?;
+            if NARROW_TOP.with(|c| c.get()) != addr {
+                plan::recordPattern(addr, v);
+            }
+            Some(Ok(()))
+        }
+        ast::CasePattern::Class { name, fields } => {
+            let (v, mt) = memberNamed(subject, name, ctx)?;
+            plan::recordPattern(addr, v);
+            let tys: Vec<PlumType> = match enumOf(&mt) {
+                Some((m, _)) => match ctx.enum_variants.get(m) {
+                    Some(info) if info.enum_name == m => crate::variantFields(info, &mt),
+                    _ => vec![mt.clone()],
+                },
+                None => vec![mt.clone()],
+            };
+            if fields.len() != tys.len() {
+                return Some(Err(format!(
+                    "constructor pattern '{}' expects {} field(s), got {}",
+                    name,
+                    tys.len(),
+                    fields.len()
+                )));
+            }
+            for (f, t) in fields.iter().zip(tys.iter()) {
+                if let Err(e) = checkPattern(f, t, env, ctx) {
+                    return Some(Err(e));
+                }
+            }
+            Some(Ok(()))
+        }
+        _ => None,
+    }
+}
+pub struct Narrow {
+    subject: String,
+    variant: String,
+    ty: PlumType,
+    enumName: String,
+    pat: usize,
+}
+/// A bare member pattern over a local/param/`self` subject narrows it. Marks
+/// the pattern so `checkMemberPattern` does not also record it.
+pub fn narrowTarget(
+    subject: &ast::Expr,
+    pat: &ast::CasePattern,
+    subject_ty: &PlumType,
+    env: &TypeEnv,
+    ctx: &CheckCtx,
+) -> Option<Narrow> {
+    let ast::CasePattern::Name(n) = pat else { return None };
+    let name = match subject {
+        ast::Expr::Var(v) if env.contains_key(v) && !env.contains_key(&format!("{}{}", crate::GLOBAL_KEY, v)) => v.clone(),
+        ast::Expr::Self_ => "self".to_string(),
+        _ => return None,
+    };
+    let (variant, ty) = memberNamed(subject_ty, n, ctx)?;
+    let addr = pat as *const ast::CasePattern as usize;
+    NARROW_TOP.with(|c| c.set(addr));
+    Some(Narrow { subject: name, variant, ty, enumName: enumOf(subject_ty)?.0.to_string(), pat: addr })
+}
+/// After a pattern check: records the narrowing and retypes the subject.
+pub fn applyNarrow(n: Option<Narrow>, case: &ast::Case, env: &mut TypeEnv, ctx: &CheckCtx) {
+    let Some(n) = n else { return };
+    NARROW_TOP.with(|c| {
+        if c.get() == n.pat {
+            c.set(0)
+        }
+    });
+    plan::recordNarrow(case as *const ast::Case as usize, n.subject.clone(), n.variant.clone());
+    let isStruct = enumOf(&n.ty)
+        .and_then(|(m, _)| ctx.enum_variants.get(m))
+        .is_some_and(|i| i.is_value);
+    env.insert(n.subject.clone(), TypeScheme::mono(n.ty.clone()));
+    let key = format!("{COPY_KEY}{}", n.subject);
+    if isStruct {
+        let member = enumOf(&n.ty).map(|x| x.0).unwrap_or("");
+        env.insert(key, TypeScheme::mono(PlumType::TNamed(format!("{member}|{}", n.enumName))));
+    } else {
+        env.remove(&key);
+    }
+}
+/// The error for `subject.field = ..` where `subject` is a narrowed struct copy.
+pub fn copyWriteError(object: &ast::Expr, env: &TypeEnv) -> Option<String> {
+    let name = match object {
+        ast::Expr::Var(v) => v.as_str(),
+        ast::Expr::Self_ => "self",
+        _ => return None,
+    };
+    let s = env.get(&format!("{COPY_KEY}{name}"))?;
+    let PlumType::TNamed(t) = &*s.body else { return None };
+    let (member, en) = t.split_once('|')?;
+    Some(format!(
+        "{name} is a copy of the {member} inside {en}; assign it to a local to change it"
+    ))
+}

plum-checker/tests/unions_pattern_test.rs ADDED Viewed

@@ -0,0 +1,109 @@
+#![allow(non_snake_case)]
+use plum_checker::checkSource;
+use plum_core::unions::{beginPlan, resolveMembers, takePlan, UnionPlan};
+use plum_core::AstParser;
+const DEFS: &str = "\
+struct Circle(r: Float)
+struct Rectangle(w: Float, h: Float)
+record Pill(len: Float)
+enum Shape = Circle | Rectangle
+enum Cap = Pill | Circle
+enum Jv = JvNull | Float | Int
+";
+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 destructure_member() {
+    let (r, p) = run("fun f(s: Shape) -> Float =\n  match s\n    Circle(r) => r\n    Rectangle(w, h) => w * h");
+    r.unwrap();
+    let mut v: Vec<_> = p.patterns.values().cloned().collect();
+    v.sort();
+    assert_eq!(v, vec!["Shape__Circle", "Shape__Rectangle"]);
+    assert!(p.narrowings.is_empty());
+    let (r, _) = run("fun take(i: Int) -> Int =\n  i\n\nfun f(s: Shape) -> Int =\n  match s\n    Circle(r) => take(r)\n    Rectangle(w, h) => 1");
+    assert!(r.is_err(), "r is a Float, not an Int");
+}
+#[test]
+fn narrow_match() {
+    let (r, p) = run("fun f(s: Shape) -> Float =\n  match s\n    Circle => s.r\n    Rectangle => s.w");
+    r.unwrap();
+    assert_eq!(p.narrowings.len(), 2);
+    let mut v: Vec<_> = p.narrowings.values().map(|x| x.1.clone()).collect();
+    v.sort();
+    assert_eq!(v, vec!["Shape__Circle", "Shape__Rectangle"]);
+    assert!(p.narrowings.values().all(|x| x.0 == "s"));
+    // after the match, s is a Shape again
+    let (r, _) = run("fun f(s: Shape) -> Float =\n  match s\n    Circle => 1.0\n    Rectangle => 2.0\n  s.r");
+    assert!(r.is_err());
+    // outside its arm s has no w
+    let (r, _) = run("fun f(s: Shape) -> Float =\n  match s\n    Circle => s.w\n    Rectangle => 2.0");
+    assert!(r.is_err());
+}
+#[test]
+fn narrow_struct_write_error() {
+    let (r, _) = run("fun f(s: Shape) -> Int =\n  match s\n    Circle =>\n      s.r = 1.0\n    Rectangle => 0\n  0");
+    assert!(
+        errs(r).contains("s is a copy of the Circle inside Shape; assign it to a local to change it")
+    );
+}
+#[test]
+fn narrow_record_write_ok() {
+    let (r, _) = run("fun f(s: Cap) -> Int =\n  match s\n    Pill =>\n      s.len = 1.0\n    Circle => 0\n  0");
+    r.unwrap();
+}
+#[test]
+fn exhaustive_members() {
+    let (r, _) = run("fun f(s: Shape) -> Float =\n  match s\n    Circle => 1.0\n    Rectangle => 2.0");
+    r.unwrap();
+    let (r, _) = run("fun f(s: Shape) -> Float =\n  match s\n    Circle => 1.0");
+    let e = errs(r);
+    assert!(e.contains("non-exhaustive match on Shape: missing Rectangle"), "{e}");
+    let (r, _) = run("fun f(s: Jv) -> Int =\n  match s\n    JvNull => 0\n    Float => 1\n    Int => 2");
+    r.unwrap();
+}
+#[test]
+fn expr_subject_no_narrow() {
+    let (r, p) = run("fun g() -> Shape =\n  Circle(r: 1.0)\n\nfun f() -> Int =\n  match g()\n    Circle => 1\n    Rectangle => 2");
+    r.unwrap();
+    assert!(p.narrowings.is_empty());
+    assert_eq!(p.patterns.len(), 2);
+}
+#[test]
+fn pattern_if() {
+    let (r, p) = run("fun f(s: Shape) -> Float =\n  if Circle := s\n    return s.r\n  0.0");
+    r.unwrap();
+    assert_eq!(p.narrowings.len(), 1);
+    let (r, _) = run("fun f(s: Shape) -> Float =\n  if Circle := s\n    return 1.0\n  s.r");
+    assert!(r.is_err());
+}
+#[test]
+fn primitive_member_narrows() {
+    let (r, p) = run("fun f(j: Jv) -> Float =\n  match j\n    Float => j + 1.0\n    Int => 0.0\n    JvNull => 0.0");
+    r.unwrap();
+    assert_eq!(p.narrowings.len(), 2);
+}