e02a80a | plum | pyrossh

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

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

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

T6: applyPlan rewrite

Files changed (2) hide show

plum-core/src/union_lower.rs +593 -0
plum-core/tests/union_lower_test.rs +329 -0

plum-core/src/union_lower.rs ADDED Viewed

@@ -0,0 +1,593 @@
+//! Rewrites a program in place from a `UnionPlan`: wraps member values in their
+//! hidden variants, rewrites patterns, narrows subjects inside match arms and
+//! synthesizes dispatch methods on enums. Post-order (children first), so every
+//! recorded address is still valid when its node is visited.
+use crate::ast::*;
+use crate::unions::UnionPlan;
+use std::collections::{HashMap, HashSet};
+/// Applies `plan` to `source`.
+pub fn applyPlan(source: &mut Source, plan: &UnionPlan) {
+    let mut lw = Lowerer { plan, pending: Vec::new(), needs: Vec::new() };
+    for item in &mut source.items {
+        match item {
+            Item::Fn(f) => match &mut f.body {
+                FnBody::Expr(e) => lw.expr(e),
+                FnBody::Block(b) => lw.block(b),
+                FnBody::Extern => {}
+            },
+            Item::Const(c) => lw.expr(&mut c.value),
+            Item::Test(t) => lw.block(&mut t.body),
+            _ => {}
+        }
+    }
+    // Items are only appended after the traversal: pushing moves addresses.
+    let mut counter = 0usize;
+    let mut done: HashSet<(String, String)> = HashSet::new();
+    for (variants, method) in std::mem::take(&mut lw.needs) {
+        let Some(enumName) = variants.first().and_then(|v| v.split("__").next()) else { continue };
+        let enumName = enumName.to_string();
+        if !done.insert((enumName.clone(), method.clone())) {
+            continue;
+        }
+        if let Some(f) = synthesize(source, &enumName, &variants, &method, &mut counter) {
+            source.items.push(Item::Fn(f));
+        }
+    }
+}
+struct Lowerer<'a> {
+    plan: &'a UnionPlan,
+    /// Narrowings (from, to) found in the condition being visited.
+    pending: Vec<(String, String)>,
+    /// (variants, method name) of every dispatch seen.
+    needs: Vec<(Vec<String>, String)>,
+}
+fn addr<T>(t: &T) -> usize {
+    t as *const T as usize
+}
+impl<'a> Lowerer<'a> {
+    fn block(&mut self, b: &mut Block) {
+        for s in &mut b.stmts {
+            self.stmt(s);
+        }
+    }
+    fn stmt(&mut self, s: &mut Stmt) {
+        let mark = self.pending.len();
+        match s {
+            Stmt::Assign(a) => {
+                for t in &mut a.targets {
+                    if let AssignTarget::Field(o, _) = t {
+                        self.expr(o);
+                    }
+                }
+                for v in &mut a.values {
+                    self.expr(v);
+                }
+            }
+            Stmt::Assert(c) => self.expr(&mut c.cond),
+            Stmt::ForLoop(f) => {
+                self.expr(&mut f.source);
+                for c in &mut f.clauses {
+                    match c {
+                        ForClause::If(e) | ForClause::Bind(_, e) => self.expr(e),
+                    }
+                }
+                self.block(&mut f.body);
+            }
+            Stmt::While(w) => {
+                self.expr(&mut w.condition);
+                let pairs = self.pending.split_off(mark);
+                self.block(&mut w.body);
+                renameBlock(&mut w.body, &pairs);
+            }
+            Stmt::If(i) => {
+                self.expr(&mut i.condition);
+                let pairs = self.pending.split_off(mark);
+                self.block(&mut i.body);
+                renameBlock(&mut i.body, &pairs);
+                for ei in &mut i.else_ifs {
+                    self.expr(&mut ei.condition);
+                    let pairs = self.pending.split_off(mark);
+                    self.block(&mut ei.body);
+                    renameBlock(&mut ei.body, &pairs);
+                }
+                if let Some(b) = &mut i.else_ {
+                    self.block(b);
+                }
+            }
+            Stmt::Match(m) => self.matchStmt(m),
+            Stmt::Return(Some(e)) | Stmt::Crash(e) | Stmt::Expr(e, ..) => self.expr(e),
+            Stmt::Return(None) | Stmt::Break | Stmt::Continue | Stmt::Todo => {}
+        }
+        self.pending.truncate(mark);
+    }
+    fn matchStmt(&mut self, m: &mut Match) {
+        for s in &mut m.subjects {
+            self.expr(s);
+        }
+        let subjects = m.subjects.clone();
+        for c in &mut m.cases {
+            self.caseArm(c, &subjects);
+        }
+    }
+    /// Visits one case (children first), applies its pattern rewrite and its
+    /// narrowing. Returns the (from, to) rename of a narrowing.
+    fn caseArm(&mut self, c: &mut Case, subjects: &[Expr]) -> Option<(String, String)> {
+        let caseAddr = addr(c);
+        for p in &mut c.patterns {
+            self.pattern(p);
+        }
+        if let Some(g) = &mut c.guard {
+            self.expr(g);
+        }
+        self.block(&mut c.body);
+        let (subject, variant, fresh) = self.plan.narrowings.get(&caseAddr)?.clone();
+        let idx = subjects.iter().position(|s| match s {
+            Expr::Var(n) => *n == subject,
+            Expr::Self_ => subject == "self",
+            _ => false,
+        })?;
+        let p = c.patterns.get_mut(idx)?;
+        match p {
+            CasePattern::Name(_) => {
+                *p = CasePattern::Class { name: variant.clone(), fields: vec![CasePattern::Name(fresh.clone())] };
+            }
+            CasePattern::Class { name, fields }
+                if *name == variant && matches!(fields.as_slice(), [CasePattern::Wildcard]) =>
+            {
+                fields[0] = CasePattern::Name(fresh.clone());
+            }
+            _ => return None,
+        }
+        let pair = vec![(subject.clone(), fresh.clone())];
+        if let Some(g) = &mut c.guard {
+            renameExpr(g, &pair);
+        }
+        renameBlock(&mut c.body, &pair);
+        Some((subject, fresh))
+    }
+    fn pattern(&mut self, p: &mut CasePattern) {
+        let a = addr(p);
+        if let CasePattern::Class { fields, .. } = p {
+            for f in fields {
+                self.pattern(f);
+            }
+        }
+        if let Some(variant) = self.plan.patterns.get(&a) {
+            let inner = match p {
+                CasePattern::Name(_) => CasePattern::Wildcard,
+                _ => std::mem::replace(p, CasePattern::Wildcard),
+            };
+            *p = CasePattern::Class { name: variant.clone(), fields: vec![inner] };
+        }
+    }
+    fn args(&mut self, args: &mut [Arg]) {
+        for a in args {
+            match a {
+                Arg::Positional(e) | Arg::Keyword { value: e, .. } => self.expr(e),
+            }
+        }
+    }
+    fn expr(&mut self, e: &mut Expr) {
+        let a = addr(e);
+        match e {
+            Expr::Binary(b) => {
+                self.expr(&mut b.left);
+                self.expr(&mut b.right);
+            }
+            Expr::Unary(u) => self.expr(&mut u.operand),
+            Expr::Bool(b) => {
+                self.expr(&mut b.left);
+                let mark = self.pending.len();
+                let _ = mark;
+                // Narrowings found on the left are visible on the right of `&&`.
+                let leftPairs: Vec<(String, String)> = self.pending.clone();
+                self.expr(&mut b.right);
+                if b.op == BoolOp::And && !leftPairs.is_empty() {
+                    renameExpr(&mut b.right, &leftPairs);
+                }
+            }
+            Expr::Not(x) | Expr::Try(x) | Expr::Paren(x) => self.expr(x),
+            Expr::Compare(c) => {
+                self.expr(&mut c.left);
+                self.expr(&mut c.right);
+            }
+            Expr::InlineIf(i) => {
+                let mark = self.pending.len();
+                self.expr(&mut i.condition);
+                let pairs = self.pending.split_off(mark);
+                self.expr(&mut i.then);
+                renameExpr(&mut i.then, &pairs);
+                self.expr(&mut i.else_);
+            }
+            Expr::FnCall(c) => self.args(&mut c.args),
+            Expr::Yield(es) => es.iter_mut().for_each(|x| self.expr(x)),
+            Expr::ClassCall(c) => {
+                for f in &mut c.fields {
+                    self.expr(&mut f.value);
+                }
+                if let Some(s) = &mut c.spread {
+                    self.expr(s);
+                }
+            }
+            Expr::Attribute(at) => {
+                self.expr(&mut at.object);
+                if let AttrKind::Method(c) = &mut at.attr {
+                    self.args(&mut c.args);
+                }
+            }
+            Expr::Bind(b) => {
+                self.expr(&mut b.value);
+                let subjects = [b.value.clone()];
+                if let Some(pair) = self.caseArm(&mut b.case, &subjects) {
+                    self.pending.push(pair);
+                }
+            }
+            Expr::Bang(b) => self.expr(&mut b.value),
+            Expr::Catch(c) => {
+                self.expr(&mut c.value);
+                match &mut c.handler {
+                    CatchHandler::Default(d) => self.expr(d),
+                    CatchHandler::Block { body, .. } => self.block(body),
+                }
+            }
+            Expr::String(s) => {
+                for p in &mut s.parts {
+                    if let StringPart::Interp(x) = p {
+                        self.expr(x);
+                    }
+                }
+            }
+            Expr::Closure(c) => self.block(&mut c.body),
+            Expr::ErrorValue(_)
+            | Expr::Int(_)
+            | Expr::Float(_)
+            | Expr::BoolLit(_)
+            | Expr::Self_
+            | Expr::Var(_)
+            | Expr::TypeName(_) => {}
+        }
+        if let Some(variants) = self.plan.dispatches.get(&a) {
+            let method = match e {
+                Expr::Attribute(at) => match &at.attr {
+                    AttrKind::Method(c) => c.name.clone(),
+                    AttrKind::Field(_) => "toStr".to_string(),
+                },
+                _ => "toStr".to_string(),
+            };
+            self.needs.push((variants.clone(), method));
+        }
+        if let Some(variant) = self.plan.wraps.get(&a) {
+            let inner = std::mem::replace(e, Expr::Int(0));
+            *e = Expr::FnCall(FnCall {
+                name: variant.clone(),
+                args: vec![Arg::Positional(inner)],
+                line: 0,
+                column: 0,
+            });
+        }
+    }
+}
+// ---- scope-aware renaming of a narrowed subject ---------------------------
+type Pairs = [(String, String)];
+fn renameBlock(b: &mut Block, pairs: &Pairs) {
+    for (from, to) in pairs {
+        renameStmts(&mut b.stmts, from, to);
+    }
+}
+fn renameExpr(e: &mut Expr, pairs: &Pairs) {
+    for (from, to) in pairs {
+        rnExpr(e, from, to);
+    }
+}
+fn patBinds(p: &CasePattern, name: &str) -> bool {
+    match p {
+        CasePattern::Name(n) => n == name,
+        CasePattern::Class { fields, .. } => fields.iter().any(|f| patBinds(f, name)),
+        _ => false,
+    }
+}
+fn condBinds(e: &Expr, name: &str) -> bool {
+    match e {
+        Expr::Bind(b) => patBinds(b.pattern(), name),
+        Expr::Bool(b) => condBinds(&b.left, name) || condBinds(&b.right, name),
+        Expr::Paren(x) => condBinds(x, name),
+        _ => false,
+    }
+}
+/// Renames in a statement list; a later `name := ...` declaration shadows the rest.
+fn renameStmts(stmts: &mut [Stmt], from: &str, to: &str) {
+    for s in stmts {
+        rnStmt(s, from, to);
+        if let Stmt::Assign(a) = s {
+            if a.declare && a.targets.iter().any(|t| matches!(t, AssignTarget::Var(n) if n == from)) {
+                return;
+            }
+        }
+    }
+}
+fn rnBlock(b: &mut Block, from: &str, to: &str) {
+    renameStmts(&mut b.stmts, from, to);
+}
+fn rnStmt(s: &mut Stmt, from: &str, to: &str) {
+    match s {
+        Stmt::Assign(a) => {
+            for v in &mut a.values {
+                rnExpr(v, from, to);
+            }
+            for t in &mut a.targets {
+                match t {
+                    AssignTarget::Field(o, _) => rnExpr(o, from, to),
+                    AssignTarget::Var(n) if !a.declare && n == from => *n = to.to_string(),
+                    _ => {}
+                }
+            }
+        }
+        Stmt::Assert(c) => rnExpr(&mut c.cond, from, to),
+        Stmt::ForLoop(f) => {
+            rnExpr(&mut f.source, from, to);
+            for c in &mut f.clauses {
+                match c {
+                    ForClause::If(e) | ForClause::Bind(_, e) => rnExpr(e, from, to),
+                }
+            }
+            if !f.binders.iter().any(|b| b == from) {
+                rnBlock(&mut f.body, from, to);
+            }
+        }
+        Stmt::While(w) => {
+            rnExpr(&mut w.condition, from, to);
+            if !condBinds(&w.condition, from) {
+                rnBlock(&mut w.body, from, to);
+            }
+        }
+        Stmt::If(i) => {
+            rnExpr(&mut i.condition, from, to);
+            if !condBinds(&i.condition, from) {
+                rnBlock(&mut i.body, from, to);
+            }
+            for ei in &mut i.else_ifs {
+                rnExpr(&mut ei.condition, from, to);
+                if !condBinds(&ei.condition, from) {
+                    rnBlock(&mut ei.body, from, to);
+                }
+            }
+            if let Some(b) = &mut i.else_ {
+                rnBlock(b, from, to);
+            }
+        }
+        Stmt::Match(m) => {
+            for x in &mut m.subjects {
+                rnExpr(x, from, to);
+            }
+            for c in &mut m.cases {
+                rnCase(c, from, to);
+            }
+        }
+        Stmt::Return(Some(e)) | Stmt::Crash(e) | Stmt::Expr(e, ..) => rnExpr(e, from, to),
+        Stmt::Return(None) | Stmt::Break | Stmt::Continue | Stmt::Todo => {}
+    }
+}
+fn rnCase(c: &mut Case, from: &str, to: &str) {
+    if c.patterns.iter().any(|p| patBinds(p, from)) {
+        return;
+    }
+    if let Some(g) = &mut c.guard {
+        rnExpr(g, from, to);
+    }
+    rnBlock(&mut c.body, from, to);
+}
+fn rnArgs(args: &mut [Arg], from: &str, to: &str) {
+    for a in args {
+        match a {
+            Arg::Positional(e) | Arg::Keyword { value: e, .. } => rnExpr(e, from, to),
+        }
+    }
+}
+fn rnExpr(e: &mut Expr, from: &str, to: &str) {
+    match e {
+        Expr::Var(n) if n == from => *n = to.to_string(),
+        Expr::Self_ if from == "self" => *e = Expr::Var(to.to_string()),
+        Expr::Binary(b) => {
+            rnExpr(&mut b.left, from, to);
+            rnExpr(&mut b.right, from, to);
+        }
+        Expr::Unary(u) => rnExpr(&mut u.operand, from, to),
+        Expr::Bool(b) => {
+            rnExpr(&mut b.left, from, to);
+            if !(b.op == BoolOp::And && condBinds(&b.left, from)) {
+                rnExpr(&mut b.right, from, to);
+            }
+        }
+        Expr::Not(x) | Expr::Try(x) | Expr::Paren(x) => rnExpr(x, from, to),
+        Expr::Compare(c) => {
+            rnExpr(&mut c.left, from, to);
+            rnExpr(&mut c.right, from, to);
+        }
+        Expr::InlineIf(i) => {
+            rnExpr(&mut i.condition, from, to);
+            if !condBinds(&i.condition, from) {
+                rnExpr(&mut i.then, from, to);
+            }
+            rnExpr(&mut i.else_, from, to);
+        }
+        Expr::FnCall(c) => rnArgs(&mut c.args, from, to),
+        Expr::Yield(es) => es.iter_mut().for_each(|x| rnExpr(x, from, to)),
+        Expr::ClassCall(c) => {
+            for f in &mut c.fields {
+                rnExpr(&mut f.value, from, to);
+            }
+            if let Some(s) = &mut c.spread {
+                rnExpr(s, from, to);
+            }
+        }
+        Expr::Attribute(at) => {
+            rnExpr(&mut at.object, from, to);
+            if let AttrKind::Method(c) = &mut at.attr {
+                rnArgs(&mut c.args, from, to);
+            }
+        }
+        Expr::Bind(b) => {
+            rnExpr(&mut b.value, from, to);
+            if !patBinds(b.pattern(), from) {
+                if let Some(g) = &mut b.case.guard {
+                    rnExpr(g, from, to);
+                }
+            }
+        }
+        Expr::Bang(b) => rnExpr(&mut b.value, from, to),
+        Expr::Catch(c) => {
+            rnExpr(&mut c.value, from, to);
+            match &mut c.handler {
+                CatchHandler::Default(d) => rnExpr(d, from, to),
+                CatchHandler::Block { binder, body } => {
+                    if binder != from {
+                        rnBlock(body, from, to);
+                    }
+                }
+            }
+        }
+        Expr::String(s) => {
+            for p in &mut s.parts {
+                if let StringPart::Interp(x) = p {
+                    rnExpr(x, from, to);
+                }
+            }
+        }
+        Expr::Closure(c) => {
+            if !c.params.iter().any(|p| p == from) {
+                rnBlock(&mut c.body, from, to);
+            }
+        }
+        _ => {}
+    }
+}
+// ---- dispatch method synthesis ---------------------------------------------
+fn substType(t: &Type, map: &HashMap<String, Type>) -> Type {
+    if t.generics.is_empty() {
+        if let Some(r) = map.get(&t.name) {
+            return r.clone();
+        }
+    }
+    Type { name: t.name.clone(), generics: t.generics.iter().map(|g| substType(g, map)).collect() }
+}
+fn substParam(p: &ParamType, map: &HashMap<String, Type>) -> ParamType {
+    match p {
+        ParamType::Type(t) => ParamType::Type(substType(t, map)),
+        ParamType::Variadic(t) => ParamType::Variadic(substType(t, map)),
+        ParamType::Seq(t) => ParamType::Seq(substType(t, map)),
+        ParamType::Fn(ps, r) => ParamType::Fn(
+            ps.iter().map(|t| substType(t, map)).collect(),
+            r.as_ref().map(|t| Box::new(substType(t, map))),
+        ),
+    }
+}
+/// Builds `fun <method>(self, ...) = match self` over `variants`, with the
+/// signature of the first member type that declares `method`.
+fn synthesize(
+    source: &Source,
+    enumName: &str,
+    variants: &[String],
+    method: &str,
+    counter: &mut usize,
+) -> Option<Fn> {
+    let exists = source.items.iter().any(|i| {
+        matches!(i, Item::Fn(f) if f.type_param.as_deref() == Some(enumName) && f.name == method)
+    });
+    if exists {
+        return None;
+    }
+    let enumItem = source.items.iter().find_map(|i| match i {
+        Item::Enum(e) if e.name == enumName => Some(e),
+        _ => None,
+    })?;
+    // Signature: first member that declares the method.
+    let mut sig: Option<(&Fn, HashMap<String, Type>)> = None;
+    for vname in variants {
+        let Some(member) = enumItem.variants.iter().find(|v| &v.name == vname).and_then(|v| v.fields.first()) else {
+            continue;
+        };
+        let decl = source.items.iter().find_map(|i| match i {
+            Item::Fn(f) if f.type_param.as_deref() == Some(member.name.as_str()) && f.name == method => Some(f),
+            _ => None,
+        });
+        if let Some(f) = decl {
+            let mut map = HashMap::new();
+            let generics = source.items.iter().find_map(|i| match i {
+                Item::Enum(me) if me.name == member.name => Some(&me.generics),
+                _ => None,
+            });
+            if let Some(gs) = generics {
+                for (g, arg) in gs.iter().zip(member.generics.iter()) {
+                    map.insert(g.name.clone(), arg.clone());
+                }
+            }
+            sig = Some((f, map));
+            break;
+        }
+    }
+    let (decl, map) = sig?;
+    let params: Vec<Param> = decl
+        .params
+        .iter()
+        .map(|p| Param { name: p.name.clone(), ty: substParam(&p.ty, &map), default: p.default.clone() })
+        .collect();
+    let cases: Vec<Case> = variants
+        .iter()
+        .map(|v| {
+            *counter += 1;
+            let d = format!("d__{counter}");
+            let args = params.iter().map(|p| Arg::Positional(Expr::Var(p.name.clone()))).collect();
+            let call = Expr::Attribute(Box::new(AttributeExpr {
+                object: Expr::Var(d.clone()),
+                attr: AttrKind::Method(FnCall { name: method.to_string(), args, line: 0, column: 0 }),
+            }));
+            Case {
+                patterns: vec![CasePattern::Class { name: v.clone(), fields: vec![CasePattern::Name(d)] }],
+                guard: None,
+                body: Block { stmts: vec![Stmt::Expr(call, 0, 0)] },
+            }
+        })
+        .collect();
+    Some(Fn {
+        name: method.to_string(),
+        file: decl.file.clone(),
+        line: 0,
+        column: 0,
+        type_param: Some(enumName.to_string()),
+        is_extern: false,
+        is_star: false,
+        has_self: true,
+        params,
+        yields: None,
+        returns: decl.returns.as_ref().map(|t| substType(t, &map)),
+        body: FnBody::Block(Block {
+            stmts: vec![Stmt::Match(Match { subjects: vec![Expr::Self_], cases })],
+        }),
+    })
+}

plum-core/tests/union_lower_test.rs ADDED Viewed

@@ -0,0 +1,329 @@
+#![allow(non_snake_case)]
+use plum_core::ast::*;
+use plum_core::union_lower::applyPlan;
+use plum_core::unions::UnionPlan;
+use plum_core::AstParser;
+fn parse(src: &str) -> Source {
+    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());
+    AstParser::new(src).parseSource(tree.root_node())
+}
+fn fnNamed<'a>(s: &'a Source, name: &str) -> &'a Fn {
+    s.items
+        .iter()
+        .find_map(|i| match i {
+            Item::Fn(f) if f.name == name => Some(f),
+            _ => None,
+        })
+        .unwrap()
+}
+fn stmts<'a>(s: &'a Source, name: &str) -> &'a Vec<Stmt> {
+    match &fnNamed(s, name).body {
+        FnBody::Block(b) => &b.stmts,
+        _ => panic!("block body expected"),
+    }
+}
+fn exprOf(s: &Stmt) -> &Expr {
+    match s {
+        Stmt::Expr(e, ..) => e,
+        _ => panic!("expr stmt expected"),
+    }
+}
+fn matchOf(s: &Stmt) -> &Match {
+    match s {
+        Stmt::Match(m) => m,
+        _ => panic!("match expected"),
+    }
+}
+fn addr<T>(t: &T) -> usize {
+    t as *const T as usize
+}
+fn varName(e: &Expr) -> &str {
+    match e {
+        Expr::Var(n) => n,
+        _ => panic!("var expected, got {e:?}"),
+    }
+}
+/// `obj.field` -> the object's variable name.
+fn fieldObj(e: &Expr) -> &str {
+    match e {
+        Expr::Attribute(a) => varName(&a.object),
+        _ => panic!("attribute expected, got {e:?}"),
+    }
+}
+const SHAPES: &str = "struct Circle(r: Float) =\n  fun area(self) -> Float = 1.0\n  fun scale(self, k: Float, j: Float) -> Float = 2.0\n  fun toStr(self) -> Str = \"c\"\nstruct Square(w: Float) =\n  fun area(self) -> Float = 2.0\n  fun scale(self, k: Float, j: Float) -> Float = 3.0\n  fun toStr(self) -> Str = \"s\"\nenum Shape = Circle | Square\n";
+fn withShapes(rest: &str) -> Source {
+    parse(&format!("{SHAPES}{rest}"))
+}
+/// The parser's own positional variant call, with positions zeroed.
+fn parsedCall(src: &str) -> Expr {
+    let s = parse(&format!("fun g() =\n  {src}\n"));
+    let mut e = exprOf(&stmts(&s, "g")[0]).clone();
+    if let Expr::FnCall(c) = &mut e {
+        c.line = 0;
+        c.column = 0;
+    }
+    e
+}
+#[test]
+fn wrap_builds_the_parsers_variant_call() {
+    let mut s = parse("fun f() =\n  Circle(r: 2.0)\n");
+    let a = addr(exprOf(&stmts(&s, "f")[0]));
+    let mut plan = UnionPlan::default();
+    plan.wraps.insert(a, "ShapeCircle".to_string());
+    applyPlan(&mut s, &plan);
+    assert_eq!(*exprOf(&stmts(&s, "f")[0]), parsedCall("ShapeCircle(Circle(r: 2.0))"));
+}
+#[test]
+fn pattern_class_and_bare() {
+    let mut s = parse("fun f(s: Shape) =\n  match s\n    Circle(r) =>\n      r\n    Square =>\n      0.0\n");
+    let m = matchOf(&stmts(&s, "f")[0]);
+    let a0 = addr(&m.cases[0].patterns[0]);
+    let a1 = addr(&m.cases[1].patterns[0]);
+    let mut plan = UnionPlan::default();
+    plan.patterns.insert(a0, "Shape__Circle".into());
+    plan.patterns.insert(a1, "Shape__Square".into());
+    applyPlan(&mut s, &plan);
+    let m = matchOf(&stmts(&s, "f")[0]);
+    assert_eq!(
+        m.cases[0].patterns[0],
+        CasePattern::Class {
+            name: "Shape__Circle".into(),
+            fields: vec![CasePattern::Class { name: "Circle".into(), fields: vec![CasePattern::Name("r".into())] }]
+        }
+    );
+    assert_eq!(
+        m.cases[1].patterns[0],
+        CasePattern::Class { name: "Shape__Square".into(), fields: vec![CasePattern::Wildcard] }
+    );
+}
+fn narrowPlan(s: &Source, fnName: &str, subject: &str, variant: &str, fresh: &str) -> UnionPlan {
+    let m = matchOf(&stmts(s, fnName)[0]);
+    let mut plan = UnionPlan::default();
+    plan.narrowings.insert(addr(&m.cases[0]), (subject.into(), variant.into(), fresh.into()));
+    plan
+}
+#[test]
+fn narrow_renames_guard_and_body() {
+    let mut s = parse("fun f(s: Shape) =\n  match s\n    Circle =>\n      s.r\n    Square =>\n      s.w\n");
+    let plan = narrowPlan(&s, "f", "s", "Shape__Circle", "s__n1");
+    applyPlan(&mut s, &plan);
+    let m = matchOf(&stmts(&s, "f")[0]);
+    assert_eq!(
+        m.cases[0].patterns[0],
+        CasePattern::Class { name: "Shape__Circle".into(), fields: vec![CasePattern::Name("s__n1".into())] }
+    );
+    assert_eq!(fieldObj(exprOf(&m.cases[0].body.stmts[0])), "s__n1");
+    // the other arm is untouched
+    assert_eq!(fieldObj(exprOf(&m.cases[1].body.stmts[0])), "s");
+}
+#[test]
+fn narrow_closure_capture() {
+    let mut s = parse("fun f(s: Shape) =\n  match s\n    Circle =>\n      g := |x| s.r\n      g(1)\n");
+    let plan = narrowPlan(&s, "f", "s", "Shape__Circle", "s__n1");
+    applyPlan(&mut s, &plan);
+    let m = matchOf(&stmts(&s, "f")[0]);
+    let Stmt::Assign(a) = &m.cases[0].body.stmts[0] else { panic!() };
+    let Expr::Closure(c) = &a.values[0] else { panic!() };
+    assert_eq!(fieldObj(exprOf(&c.body.stmts[0])), "s__n1");
+}
+#[test]
+fn narrow_respects_shadowing() {
+    let mut s = parse(
+        "fun f(s: Shape) =\n  match s\n    Circle =>\n      s.r\n      h := |s| s.w\n      s := 5\n      s.r\n",
+    );
+    let plan = narrowPlan(&s, "f", "s", "Shape__Circle", "s__n1");
+    applyPlan(&mut s, &plan);
+    let m = matchOf(&stmts(&s, "f")[0]);
+    let body = &m.cases[0].body.stmts;
+    assert_eq!(fieldObj(exprOf(&body[0])), "s__n1");
+    let Stmt::Assign(a) = &body[1] else { panic!() };
+    let Expr::Closure(c) = &a.values[0] else { panic!() };
+    assert_eq!(fieldObj(exprOf(&c.body.stmts[0])), "s", "closure param shadows");
+    assert_eq!(fieldObj(exprOf(&body[3])), "s", "after `s := 5`");
+}
+#[test]
+fn narrow_self() {
+    let mut s = parse("fun area(self) =\n  match self\n    Circle =>\n      self.r\n    Square =>\n      0.0\n");
+    let plan = narrowPlan(&s, "area", "self", "Shape__Circle", "self__n1");
+    applyPlan(&mut s, &plan);
+    let m = matchOf(&stmts(&s, "area")[0]);
+    assert_eq!(fieldObj(exprOf(&m.cases[0].body.stmts[0])), "self__n1");
+}
+#[test]
+fn narrow_in_pattern_if() {
+    let mut s = parse("fun f(s: Shape) =\n  if Circle := s\n    s.r\n  else\n    s.w\n");
+    let Stmt::If(i) = &stmts(&s, "f")[0] else { panic!() };
+    let Expr::Bind(b) = &i.condition else { panic!("bind expected: {:?}", i.condition) };
+    let mut plan = UnionPlan::default();
+    plan.narrowings.insert(addr(&b.case), ("s".into(), "Shape__Circle".into(), "s__n1".into()));
+    applyPlan(&mut s, &plan);
+    let Stmt::If(i) = &stmts(&s, "f")[0] else { panic!() };
+    assert_eq!(fieldObj(exprOf(&i.body.stmts[0])), "s__n1");
+    assert_eq!(fieldObj(exprOf(&i.else_.as_ref().unwrap().stmts[0])), "s");
+    let Expr::Bind(b) = &i.condition else { panic!() };
+    assert_eq!(
+        *b.pattern(),
+        CasePattern::Class { name: "Shape__Circle".into(), fields: vec![CasePattern::Name("s__n1".into())] }
+    );
+}
+fn shapeFn<'a>(s: &'a Source, method: &str) -> &'a Fn {
+    s.items
+        .iter()
+        .find_map(|i| match i {
+            Item::Fn(f) if f.name == method && f.type_param.as_deref() == Some("Shape") => Some(f),
+            _ => None,
+        })
+        .unwrap_or_else(|| panic!("no synthesized {method}"))
+}
+fn assertDispatchBody(f: &Fn, method: &str, argNames: &[&str]) {
+    let FnBody::Block(b) = &f.body else { panic!() };
+    let m = matchOf(&b.stmts[0]);
+    assert_eq!(m.subjects, vec![Expr::Self_]);
+    assert_eq!(m.cases.len(), 2);
+    for (k, (case, v)) in m.cases.iter().zip(["Shape__Circle", "Shape__Square"]).enumerate() {
+        let d = format!("d__{}", k + 1);
+        assert_eq!(
+            case.patterns[0],
+            CasePattern::Class { name: v.into(), fields: vec![CasePattern::Name(d.clone())] }
+        );
+        let Expr::Attribute(a) = exprOf(&case.body.stmts[0]) else { panic!() };
+        assert_eq!(varName(&a.object), d);
+        let AttrKind::Method(c) = &a.attr else { panic!() };
+        assert_eq!(c.name, method);
+        let args: Vec<&str> = c
+            .args
+            .iter()
+            .map(|a| match a {
+                Arg::Positional(e) => varName(e),
+                _ => panic!(),
+            })
+            .collect();
+        assert_eq!(args, argNames);
+    }
+}
+#[test]
+fn dispatch_leaves_call_and_adds_method() {
+    let mut s = withShapes("fun f(s: Shape) -> Float =\n  s.area()\n");
+    let before = exprOf(&stmts(&s, "f")[0]).clone();
+    let a = addr(exprOf(&stmts(&s, "f")[0]));
+    let mut plan = UnionPlan::default();
+    plan.dispatches.insert(a, vec!["Shape__Circle".into(), "Shape__Square".into()]);
+    applyPlan(&mut s, &plan);
+    assert_eq!(*exprOf(&stmts(&s, "f")[0]), before);
+    let f = shapeFn(&s, "area");
+    assert!(f.has_self && f.params.is_empty());
+    assert_eq!(f.returns, Some(Type { name: "Float".into(), generics: vec![] }));
+    assertDispatchBody(f, "area", &[]);
+}
+#[test]
+fn dispatch_passes_two_args() {
+    let mut s = withShapes("fun f(s: Shape) -> Float =\n  s.scale(1.0, 2.0)\n");
+    let a = addr(exprOf(&stmts(&s, "f")[0]));
+    let mut plan = UnionPlan::default();
+    plan.dispatches.insert(a, vec!["Shape__Circle".into(), "Shape__Square".into()]);
+    applyPlan(&mut s, &plan);
+    let f = shapeFn(&s, "scale");
+    assert_eq!(f.params.len(), 2);
+    assert_eq!(f.params[0].name, "k");
+    assertDispatchBody(f, "scale", &["k", "j"]);
+}
+#[test]
+fn dispatch_in_interpolation() {
+    let mut s = withShapes("fun f(s: Shape) -> Str =\n  \"{s}\"\n");
+    let Expr::String(se) = exprOf(&stmts(&s, "f")[0]) else { panic!() };
+    let StringPart::Interp(inner) = &se.parts[0] else { panic!() };
+    let a = addr(inner);
+    let mut plan = UnionPlan::default();
+    plan.dispatches.insert(a, vec!["Shape__Circle".into(), "Shape__Square".into()]);
+    applyPlan(&mut s, &plan);
+    let Expr::String(se) = exprOf(&stmts(&s, "f")[0]) else { panic!() };
+    let StringPart::Interp(inner) = &se.parts[0] else { panic!() };
+    assert_eq!(varName(inner), "s");
+    assertDispatchBody(shapeFn(&s, "toStr"), "toStr", &[]);
+}
+#[test]
+fn dispatch_once_per_method() {
+    let mut s = withShapes("fun f(s: Shape) -> Float =\n  s.area() + s.area()\n");
+    let Expr::Binary(b) = exprOf(&stmts(&s, "f")[0]) else { panic!() };
+    let (a0, a1) = (addr(&b.left), addr(&b.right));
+    let mut plan = UnionPlan::default();
+    for a in [a0, a1] {
+        plan.dispatches.insert(a, vec!["Shape__Circle".into(), "Shape__Square".into()]);
+    }
+    applyPlan(&mut s, &plan);
+    let n = s
+        .items
+        .iter()
+        .filter(|i| matches!(i, Item::Fn(f) if f.name == "area" && f.type_param.as_deref() == Some("Shape")))
+        .count();
+    assert_eq!(n, 1);
+}
+#[test]
+fn nested_wrap_inside_dispatch_arg() {
+    let mut s = withShapes("fun f(s: Shape) -> Float =\n  s.scale(1.0, 2.0)\n");
+    // dispatch on the call, wrap on its first argument
+    let e = exprOf(&stmts(&s, "f")[0]);
+    let Expr::Attribute(a) = e else { panic!() };
+    let AttrKind::Method(c) = &a.attr else { panic!() };
+    let Arg::Positional(arg0) = &c.args[0] else { panic!() };
+    let mut plan = UnionPlan::default();
+    plan.dispatches.insert(addr(e), vec!["Shape__Circle".into(), "Shape__Square".into()]);
+    plan.wraps.insert(addr(arg0), "Shape__Circle".into());
+    applyPlan(&mut s, &plan);
+    let Expr::Attribute(a) = exprOf(&stmts(&s, "f")[0]) else { panic!() };
+    let AttrKind::Method(c) = &a.attr else { panic!() };
+    let Arg::Positional(Expr::FnCall(w)) = &c.args[0] else { panic!("wrap expected") };
+    assert_eq!(w.name, "Shape__Circle");
+    shapeFn(&s, "scale");
+}
+#[test]
+fn dispatch_substitutes_member_generics() {
+    let mut s = parse(
+        "record Box[T](v: T) =\n  fun get(self, d: T) -> T = d\nenum Opt = Box[Int] | Nothing\nfun f(o: Opt) -> Int =\n  o.get(1)\n",
+    );
+    let a = addr(exprOf(&stmts(&s, "f")[0]));
+    let mut plan = UnionPlan::default();
+    plan.dispatches.insert(a, vec!["Opt__Box".into()]);
+    applyPlan(&mut s, &plan);
+    let f = s
+        .items
+        .iter()
+        .find_map(|i| match i {
+            Item::Fn(f) if f.type_param.as_deref() == Some("Opt") => Some(f),
+            _ => None,
+        })
+        .unwrap();
+    let int = Type { name: "Int".into(), generics: vec![] };
+    assert_eq!(f.params[0].ty, ParamType::Type(int.clone()));
+    assert_eq!(f.returns, Some(int));
+}