ed7f16c | plum | pyrossh

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

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

ed7f16c — Peter John 2026-11-02T19:46:13+05:30

T5: field access through a trait

Files changed (4) hide show

plum-checker/src/unions.rs +76 -3
plum-checker/tests/trait_fields_tests.rs +223 -0
plum-core/src/union_lower.rs +52 -0
plum-core/src/unions.rs +25 -4

plum-checker/src/unions.rs CHANGED Viewed

@@ -554,9 +554,38 @@ pub fn restorePlan(p: &plan::UnionPlan) {
     for (a, vs) in &p.dispatches {
         plan::recordDispatch(*a, vs.clone());
     }
+    for (a, vs) in &p.fieldReads {
+        plan::recordFieldRead(*a, vs.clone());
+    }
+    for (a, vs) in &p.fieldWrites {
+        plan::recordFieldWrite(*a, vs.clone());
+    }
+}
+/// Names of the fields every member of the union has an accessor for.
+fn commonFields(ms: &[(String, PlumType)], ctx: &CheckCtx) -> Option<Vec<String>> {
+    let mut common: Option<Vec<String>> = None;
+    for (_, mt) in ms {
+        let recv = crate::methodReceiverName(mt)?;
+        let mut mine: Vec<String> = ctx
+            .methods
+            .keys()
+            .filter(|(r, n)| *r == recv && n.starts_with("__get_"))
+            .map(|(_, n)| n["__get_".len()..].to_string())
+            .collect();
+        mine.sort();
+        common = Some(match common {
+            None => mine,
+            Some(c) => c.into_iter().filter(|f| mine.contains(f)).collect(),
+        });
+    }
+    common.filter(|c| !c.is_empty())
 }
-/// Field access through a trait value (T5 implements). `None`: not a trait field access.
+/// Field access through a trait value: the field's type, with the read (or
+/// write, when `write`) recorded for `union_lower` to turn into an accessor
+/// dispatch. `None`: `obj` is not a trait. `exprAddr` is the address of the
+/// `Attribute` expression for a read, of the assign statement for a write.
 pub fn traitField(
     obj: &PlumType,
     field: &str,
@@ -564,6 +593,50 @@ pub fn traitField(
     write: bool,
     ctx: &CheckCtx,
 ) -> Option<Result<PlumType, String>> {
+    let (en, args) = enumOf(obj)?;
+    let ms = members(en, args, ctx);
+    if ms.is_empty() {
+        return None;
+    }
+    let fields = commonFields(&ms, ctx)?;
+    if !fields.iter().any(|f| f == field) {
+        return Some(Err(format!("trait {en} has no field {field}")));
+    }
+    let acc = plum_core::traits::getter(field);
+    let ty = match memberSig(&ms[0].1, &acc, ctx) {
+        Some(PlumType::TFun(_, ret)) => *ret,
+        _ => return Some(Err(format!("trait {en} has no field {field}"))),
+    };
+    let variants: Vec<String> = ms.into_iter().map(|(v, _)| v).collect();
+    if write {
+        plan::recordFieldWrite(exprAddr, variants);
+    } else {
-    let _ = (obj, field, exprAddr, write, ctx);
+        plan::recordFieldRead(exprAddr, variants);
+    }
+    Some(Ok(ty))
+}
+/// For a field-assignment target `object.x = ..` whose `object` (already
+/// inferred, type `objTy`) is itself a trait field read of a struct type:
+/// the error to report, since writing into a struct copy would be lost.
+pub fn nestedTraitWrite(
+    object: &ast::Expr,
+    objTy: &PlumType,
+    ctx: &CheckCtx,
+) -> Option<String> {
+    if !plan::hasFieldRead(object as *const ast::Expr as usize) {
-    None
+        return None;
+    }
+    let PlumType::TNamed(n) = objTy else { return None };
+    if !ctx.enum_variants.get(n).is_some_and(|i| i.is_value) {
+        return None;
+    }
+    let ast::Expr::Attribute(at) = object else { return None };
+    let ast::AttrKind::Field(f) = &at.attr else { return None };
+    let base = match &at.object {
+        ast::Expr::Var(v) => v.clone(),
+        ast::Expr::Self_ => "self".to_string(),
+        _ => "e".to_string(),
+    };
+    Some(format!("assign the whole field: {base}.{f} = ..."))
 }

plum-checker/tests/trait_fields_tests.rs ADDED Viewed

@@ -0,0 +1,223 @@
+#![allow(non_snake_case)]
+//! Trait field reads/writes: the checker's `traitField` plus `union_lower`'s
+//! rewrite of the recorded reads and writes into `__get_f` / `__set_f` dispatch.
+use plum_checker::types::PlumType;
+use plum_checker::unions::{nestedTraitWrite, traitField};
+use plum_checker::{
+    buildGlobalTables, buildMinRequiredArgs, buildPrimitiveGenerics, buildStaticMethods,
+    checkSource, CheckCtx,
+};
+use plum_core::ast::*;
+use plum_core::traits::lowerTraits;
+use plum_core::union_lower::applyPlan;
+use plum_core::unions::{beginPlan, resolveMembers, takePlan};
+use plum_core::AstParser;
+const DEFS: &str = "\
+struct Vec2(x: Float, y: Float)
+trait Ent(pos: Vec2, z: Float) =
+  fun draw(self) -> Float
+record Hut(Ent, k: Float) =
+  fun draw(self) -> Float = self.z
+record Orc(Ent) =
+  fun draw(self) -> Float = self.z + 1.0
+record Circle(r: Float)
+record Square(s: Float)
+enum Shape = Circle | Square
+";
+fn lowered(body: &str) -> Source {
+    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(), "{}", tree.root_node().to_sexp());
+    let mut s = AstParser::new(&src).parseSourceRaw(tree.root_node());
+    assert_eq!(lowerTraits(&mut s), Vec::<String>::new());
+    let errs = resolveMembers(&mut s);
+    assert!(errs.is_empty(), "{errs:?}");
+    s
+}
+fn withCtx<R>(s: &Source, f: impl FnOnce(&CheckCtx) -> R) -> R {
+    let (_, methods, enum_variants, enum_params) = buildGlobalTables(s);
+    let min_required = buildMinRequiredArgs(s);
+    let primitive_generics = buildPrimitiveGenerics(s);
+    let statics = buildStaticMethods(s);
+    let ctx = CheckCtx {
+        methods: &methods,
+        enum_variants: &enum_variants,
+        enum_params: &enum_params,
+        min_required: &min_required,
+        primitive_generics: &primitive_generics,
+        statics: Some(&statics),
+    };
+    f(&ctx)
+}
+fn named(n: &str) -> PlumType {
+    PlumType::TNamed(n.to_string())
+}
+fn fnBody<'a>(s: &'a Source, name: &str) -> &'a FnBody {
+    s.items
+        .iter()
+        .find_map(|i| match i {
+            Item::Fn(f) if f.name == name && f.type_param.is_none() => Some(&f.body),
+            _ => None,
+        })
+        .unwrap()
+}
+fn firstStmt<'a>(s: &'a Source, name: &str) -> &'a Stmt {
+    let FnBody::Block(b) = fnBody(s, name) else { panic!("block") };
+    &b.stmts[0]
+}
+fn hasMethod(s: &Source, recv: &str, name: &str) -> bool {
+    s.items.iter().any(
+        |i| matches!(i, Item::Fn(f) if f.type_param.as_deref() == Some(recv) && f.name == name),
+    )
+}
+#[test]
+fn read_type_and_plan() {
+    let s = lowered("");
+    withCtx(&s, |ctx| {
+        beginPlan();
+        let r = traitField(&named("Ent"), "pos", 7, false, ctx);
+        assert_eq!(r, Some(Ok(named("Vec2"))));
+        let p = takePlan();
+        assert_eq!(p.fieldReads[&7], vec!["Ent__Hut", "Ent__Orc"]);
+        assert!(p.fieldWrites.is_empty());
+        beginPlan();
+        assert_eq!(traitField(&named("Ent"), "z", 9, true, ctx), Some(Ok(PlumType::TFloat)));
+        let p = takePlan();
+        assert_eq!(p.fieldWrites[&9], vec!["Ent__Hut", "Ent__Orc"]);
+    });
+}
+#[test]
+fn unknown_field_and_non_trait() {
+    let s = lowered("");
+    withCtx(&s, |ctx| {
+        beginPlan();
+        assert_eq!(
+            traitField(&named("Ent"), "nope", 1, false, ctx),
+            Some(Err("trait Ent has no field nope".to_string()))
+        );
+        assert_eq!(traitField(&named("Shape"), "r", 1, false, ctx), None);
+        assert_eq!(traitField(&named("Hut"), "z", 1, false, ctx), None);
+        assert_eq!(traitField(&PlumType::TInt, "z", 1, false, ctx), None);
+        assert!(takePlan().fieldReads.is_empty());
+    });
+}
+#[test]
+fn nested_write_rejected() {
+    let s = lowered("fun f(e: Ent) =\n  e.pos.x = 1.0\n");
+    // The target is Field(Attribute(e, pos), x).
+    let Stmt::Assign(a) = firstStmt(&s, "f") else { panic!() };
+    let AssignTarget::Field(obj, _) = &a.targets[0] else { panic!() };
+    withCtx(&s, |ctx| {
+        beginPlan();
+        let objAddr = &**obj as *const Expr as usize;
+        assert_eq!(nestedTraitWrite(obj, &named("Vec2"), ctx), None); // nothing read yet
+        traitField(&named("Ent"), "pos", objAddr, false, ctx).unwrap().unwrap();
+        assert_eq!(
+            nestedTraitWrite(obj, &named("Vec2"), ctx),
+            Some("assign the whole field: e.pos = ...".to_string())
+        );
+        // A non-struct field type (a record) is shared, so writing into it is fine.
+        assert_eq!(nestedTraitWrite(obj, &named("Hut"), ctx), None);
+        takePlan();
+    });
+}
+/// `get` / `set` / `bump` lowered by hand-recorded plans (the checker call sites
+/// live in lib.rs); the result still type-checks.
+#[test]
+fn lowering_rewrites_reads_and_writes() {
+    let mut s = lowered(
+        "fun get(e: Ent) -> Float =\n  e.z\n\nfun set(e: Ent) =\n  e.z = 5.0\n\nfun bump(e: Ent) =\n  e.z = e.z + 1.0\n",
+    );
+    withCtx(&s, |ctx| {
+        beginPlan();
+        let getAddr = match fnBody(&s, "get") {
+            FnBody::Expr(e) => e as *const Expr as usize,
+            FnBody::Block(b) => match &b.stmts[0] {
+                Stmt::Return(Some(e)) | Stmt::Expr(e, ..) => e as *const Expr as usize,
+                _ => panic!(),
+            },
+            _ => panic!(),
+        };
+        traitField(&named("Ent"), "z", getAddr, false, ctx).unwrap().unwrap();
+        let setAddr = firstStmt(&s, "set") as *const Stmt as usize;
+        traitField(&named("Ent"), "z", setAddr, true, ctx).unwrap().unwrap();
+        // bump: read inside the value, write on the statement.
+        let bumpStmt = firstStmt(&s, "bump");
+        let Stmt::Assign(a) = bumpStmt else { panic!() };
+        let Expr::Binary(bin) = &a.values[0] else { panic!() };
+        traitField(&named("Ent"), "z", &bin.left as *const Expr as usize, false, ctx)
+            .unwrap()
+            .unwrap();
+        traitField(&named("Ent"), "z", bumpStmt as *const Stmt as usize, true, ctx)
+            .unwrap()
+            .unwrap();
+    });
+    let plan = takePlan();
+    applyPlan(&mut s, &plan);
+    // Accessor dispatch exists on the trait union.
+    assert!(hasMethod(&s, "Ent", "__get_z"));
+    assert!(hasMethod(&s, "Ent", "__set_z"));
+    // `set` is now a call `e.__set_z(5.0)`.
+    let Stmt::Expr(Expr::Attribute(at), ..) = firstStmt(&s, "set") else {
+        panic!("{:?}", firstStmt(&s, "set"))
+    };
+    let AttrKind::Method(c) = &at.attr else { panic!() };
+    assert_eq!(c.name, "__set_z");
+    // `bump` is `e.__set_z(e.__get_z() + 1.0)`: get then set.
+    let Stmt::Expr(Expr::Attribute(at), ..) = firstStmt(&s, "bump") else { panic!() };
+    let AttrKind::Method(c) = &at.attr else { panic!() };
+    assert_eq!(c.name, "__set_z");
+    let Arg::Positional(Expr::Binary(b)) = &c.args[0] else { panic!() };
+    let Expr::Attribute(g) = &b.left else { panic!() };
+    assert!(matches!(&g.attr, AttrKind::Method(m) if m.name == "__get_z"));
+    // The rewritten program is well typed.
+    checkSource(&s).unwrap_or_else(|e| panic!("{:?}", e.iter().map(|x| &x.message).collect::<Vec<_>>()));
+}
+#[test]
+fn plain_union_unaffected() {
+    let mut s = lowered("fun f(x: Shape) -> Shape =\n  x\n");
+    let before = s.clone();
+    applyPlan(&mut s, &takePlan());
+    assert_eq!(s, before);
+    withCtx(&s, |ctx| {
+        beginPlan();
+        assert_eq!(traitField(&named("Shape"), "r", 3, false, ctx), None);
+        assert!(takePlan().isEmpty());
+    });
+}
+/// Through the checker proper (needs the `traitField` call sites in lib.rs and
+/// `UnionPlan::isEmpty` in `checkUnions`).
+#[test]
+#[ignore = "needs T4 call sites"]
+fn typed_source_lowers_field_access() {
+    let s = lowered(
+        "fun bump(e: Ent) -> Float =\n  e.z = e.z + 1.0\n  e.z\n\nfun nest(e: Ent) =\n  e.pos.x = 1.0\n",
+    );
+    let errs = plum_checker::typedSource(&s).err().expect("nested write must be rejected");
+    assert!(errs.iter().any(|e| e.message.contains("assign the whole field: e.pos = ...")), "{:?}", errs.iter().map(|e| &e.message).collect::<Vec<_>>());
+    let mut ok = s.clone();
+    ok.items.retain(|i| !matches!(i, Item::Fn(f) if f.name == "nest"));
+    let info = plum_checker::typedSource(&ok).unwrap_or_else(|e| panic!("{:?}", e.iter().map(|x| &x.message).collect::<Vec<_>>()));
+    let low = info.lowered.expect("lowered copy");
+    assert!(hasMethod(&low, "Ent", "__get_z") && hasMethod(&low, "Ent", "__set_z"));
+}

plum-core/src/union_lower.rs CHANGED Viewed

@@ -64,6 +64,7 @@ impl<'a> Lowerer<'a> {
     fn stmt(&mut self, s: &mut Stmt) {
         let mark = self.pending.len();
+        let stmtAddr = addr(s);
         match s {
             Stmt::Assign(a) => {
                 for t in &mut a.targets {
@@ -110,9 +111,46 @@ impl<'a> Lowerer<'a> {
             Stmt::Return(Some(e)) | Stmt::Crash(e) | Stmt::Expr(e, ..) => self.expr(e),
             Stmt::Return(None) | Stmt::Break | Stmt::Continue | Stmt::Todo => {}
         }
+        self.lowerFieldWrite(s, stmtAddr);
         self.pending.truncate(mark);
     }
+    /// `e.f = v` on a trait value becomes the dispatch call `e.__set_f(v)`.
+    fn lowerFieldWrite(&mut self, s: &mut Stmt, stmtAddr: usize) {
+        let Stmt::Assign(a) = s else { return };
+        let inner = addr(a);
+        let Some(variants) = self
+            .plan
+            .fieldWrites
+            .get(&stmtAddr)
+            .or_else(|| self.plan.fieldWrites.get(&inner))
+            .cloned()
+        else {
+            return;
+        };
+        if a.targets.len() != 1 || a.values.len() != 1 || a.declare {
+            return;
+        }
+        let AssignTarget::Field(_, f) = &a.targets[0] else { return };
+        let method = crate::traits::setter(f);
+        let AssignTarget::Field(obj, _) = a.targets.remove(0) else { return };
+        let value = a.values.remove(0);
+        self.needs.push((variants, method.clone()));
+        *s = Stmt::Expr(
+            Expr::Attribute(Box::new(AttributeExpr {
+                object: *obj,
+                attr: AttrKind::Method(FnCall {
+                    name: method,
+                    args: vec![Arg::Positional(value)],
+                    line: 0,
+                    column: 0,
+                }),
+            })),
+            0,
+            0,
+        );
+    }
     fn matchStmt(&mut self, m: &mut Match) {
         for s in &mut m.subjects {
             self.expr(s);
@@ -268,6 +306,20 @@ impl<'a> Lowerer<'a> {
             | Expr::Var(_)
             | Expr::TypeName(_) => {}
         }
+        if let Some(variants) = self.plan.fieldReads.get(&a) {
+            if let Expr::Attribute(at) = e {
+                if let AttrKind::Field(f) = &at.attr {
+                    let method = crate::traits::getter(f);
+                    self.needs.push((variants.clone(), method.clone()));
+                    at.attr = AttrKind::Method(FnCall {
+                        name: method,
+                        args: vec![],
+                        line: 0,
+                        column: 0,
+                    });
+                }
+            }
+        }
         if let Some(variants) = self.plan.dispatches.get(&a) {
             let method = match e {
                 Expr::Attribute(at) => match &at.attr {

plum-core/src/unions.rs CHANGED Viewed

@@ -153,6 +153,22 @@ pub struct UnionPlan {
     pub narrowings: HashMap<usize, (String, String, String)>,
     /// call address -> hidden variants to dispatch over
     pub dispatches: HashMap<usize, Vec<String>>,
+    /// trait field read (`Attribute(e, f)`) address -> hidden variants to dispatch `__get_f` over
+    pub fieldReads: HashMap<usize, Vec<String>>,
+    /// trait field write (assign statement) address -> hidden variants to dispatch `__set_f` over
+    pub fieldWrites: HashMap<usize, Vec<String>>,
+}
+impl UnionPlan {
+    /// Nothing was recorded.
+    pub fn isEmpty(&self) -> bool {
+        self.wraps.is_empty()
+            && self.patterns.is_empty()
+            && self.narrowings.is_empty()
+            && self.dispatches.is_empty()
+            && self.fieldReads.is_empty()
+            && self.fieldWrites.is_empty()
+    }
 }
 thread_local! {
@@ -197,11 +213,16 @@ pub fn takePlan() -> UnionPlan {
     PLAN.with(|p| std::mem::take(&mut *p.borrow_mut()))
 }
-/// Records a trait field read at `exprAddr` (T5 implements).
+/// Records a trait field read at `exprAddr`.
 pub fn recordFieldRead(exprAddr: usize, variants: Vec<String>) {
-    let _ = (exprAddr, variants);
+    PLAN.with(|p| p.borrow_mut().fieldReads.insert(exprAddr, variants));
 }
-/// Records a trait field write at `stmtAddr` (T5 implements).
+/// Records a trait field write at `stmtAddr` (the address of the assign statement
+/// or of its `Assign`).
 pub fn recordFieldWrite(stmtAddr: usize, variants: Vec<String>) {
-    let _ = (stmtAddr, variants);
+    PLAN.with(|p| p.borrow_mut().fieldWrites.insert(stmtAddr, variants));
+}
+/// Whether a trait field read is recorded at `exprAddr`.
+pub fn hasFieldRead(exprAddr: usize) -> bool {
+    PLAN.with(|p| p.borrow().fieldReads.contains_key(&exprAddr))
 }