Iter 18c.2: linearity check + suggested_rewrites

Adds a per-fn linearity check that runs on every fn whose
param_modes are all explicit (Borrow or Own, no Implicit), tracks
each binder's consume/borrow state through the AST, and emits
two diagnostics with form-A `suggested_rewrites`:

- `use-after-consume` — a binder is referenced after it has
  already been consumed.
- `consume-while-borrowed` — a binder is consumed while a borrow
  of it is still live (sibling arg slot, or an enclosing-fn
  Borrow param).

Fns with any Implicit param skip the check entirely; that's the
back-compat lane that keeps every existing fixture green. Pure
diagnostic addition: no IR change, no codegen change, no runtime
change.

Each diagnostic carries a non-empty `suggested_rewrites` whose
`replacement` is form-A AILang text (typically `(clone <name>)`).
Adds two new public entrypoints to ailang-surface — `parse_term`
and `term_to_form_a` — so the round-trip
`parse_term(term_to_form_a(t)) ≡ t` is the contract every
emitted replacement must satisfy. Tests assert the contract.

Linearity pass runs only on modules that typechecked clean (any
upstream typecheck error suppresses it for that module — running
on partly-defined IR would produce noise).

Tests: 3 unit tests in `linearity::tests`, 3 integration tests
in `ailang-check/tests/workspace.rs`, 1 serialisation test in
`diagnostic.rs`. `cargo test --workspace` green.
This commit is contained in:
2026-05-08 10:06:14 +02:00
parent 8d704cd9b5
commit 09fb5bb113
9 changed files with 987 additions and 4 deletions
+609
View File
@@ -0,0 +1,609 @@
//! Iter 18c.2: linearity check for fns whose every parameter mode is
//! explicit (`Borrow` or `Own`).
//!
//! ## Scope
//!
//! - **Active only** on `Def::Fn` whose `Type::Fn.param_modes` is
//! non-empty AND contains no [`ParamMode::Implicit`]. Any
//! `Implicit`-mode param skips the entire fn — that is the back-compat
//! lane that keeps every existing fixture green (only
//! `borrow_own_demo` ships an all-explicit signature today).
//! - **Pure diagnostic.** Emits [`Diagnostic`]s with codes
//! `use-after-consume` / `consume-while-borrowed`. No IR change, no
//! codegen change, no runtime change. The check runs after
//! [`crate::check_fn`] type-checking succeeded for the def.
//! - **No uniqueness inference.** That is Iter 18c.3. The check only
//! tracks what is *spelled* in the source: a binder is consumed when
//! it appears in a non-borrow position; cloned via [`Term::Clone`];
//! borrowed via fn-call into a `Borrow` parameter.
//!
//! ## Position model
//!
//! Each [`Term::Var`] occurrence is in some "position", determined by
//! its parent term:
//!
//! - **Consume** — default. The use takes ownership.
//! - **Borrow** — the use lends a reference; ownership stays with the
//! binder.
//!
//! Position propagation:
//! - `Term::App.callee` — Consume (the value is called once).
//! - `Term::App.args[i]` — Borrow if the callee is a `Var` whose
//! resolved type is a `Type::Fn` with `param_modes[i] == Borrow`;
//! otherwise Consume (Own / Implicit / unknown all default to
//! Consume).
//! - `Term::Clone.value` — Borrow (clone reads + bumps RC, doesn't
//! move).
//! - `Term::Match.scrutinee` — Borrow (matching is a read; for an
//! `Own` param this leaves the binder uneconsumed, which is a known
//! false-negative for "param declared own but never moved" — out of
//! scope for 18c.2 per the assignment).
//! - `Term::Let.value`, `Term::Seq.lhs`, `Term::If.cond` — Consume.
//! - `Term::Ctor.args[*]`, `Term::Do.args[*]` — Consume.
//! - `Term::Lam` — captures are conservatively Consume.
//!
//! ## Within-call borrow lifetime (consume-while-borrowed)
//!
//! For `App { callee, args, .. }` whose callee is a `Var` resolving to
//! a fn-typed binding with explicit `param_modes`: when the i-th arg is
//! a bare `Term::Var { name }` consumed in Borrow position, we
//! **increment** `borrow_count[name]` *before* evaluating
//! `args[i+1..]`. After all args are evaluated and the call has run,
//! we **decrement** them again. So a sibling subterm that consumes the
//! same binder while it is still borrowed by an earlier sibling
//! triggers the `consume-while-borrowed` diagnostic.
//!
//! For a `Borrow` *parameter* of the enclosing fn, the binder starts
//! the body with `borrow_count = 1` (the caller's outer borrow). Any
//! attempt to consume it inside the body therefore also triggers
//! `consume-while-borrowed`.
//!
//! ## Suggested rewrites
//!
//! Each diagnostic carries a non-empty
//! [`crate::diagnostic::Diagnostic::suggested_rewrites`]. The form-A
//! replacement is parseable by [`ailang_surface::parse_term`] — that's
//! the contract guarded by the round-trip test in `tests/`.
//!
//! - `use-after-consume` on `Var x`: replacement `(clone <name>)` —
//! the author should clone an *earlier* use so this later use stays
//! the unique consume.
//! - `consume-while-borrowed` on `Var x`: replacement `(clone <name>)`
//! — the author should clone here so the original keeps the borrow.
use crate::diagnostic::Diagnostic;
use ailang_core::ast::{Arm, Def, FnDef, Module, ParamMode, Pattern, Term, Type};
use ailang_surface::term_to_form_a;
use std::collections::HashMap;
/// Position in which a [`Term::Var`] is being used. See module-level
/// docs for the propagation rules.
#[derive(Copy, Clone, Debug, PartialEq, Eq)]
enum Position {
/// The use takes ownership (default position).
Consume,
/// The use lends a non-owning reference for the duration of the
/// enclosing call (or, for `Term::Clone` / `Term::Match`, for the
/// duration of the read).
Borrow,
}
/// Mutable per-binder linearity state, keyed by binder name. Names
/// shadow lexically — see [`Checker::with_binder`] for the
/// save/restore pattern.
#[derive(Debug, Default, Clone)]
struct BinderState {
/// `true` once the binder has been used in a Consume position.
/// Subsequent uses (in any position) trigger `use-after-consume`.
consumed: bool,
/// Number of currently-live borrows of this binder. Incremented
/// when a Borrow-position use starts (in a fn-call arg slot or as
/// the initial state of a `Borrow` parameter); decremented when
/// the borrow ends. Consume while `> 0` triggers
/// `consume-while-borrowed`.
borrow_count: u32,
}
/// Iter 18c.2: top-level entry. Walks every fn in `m` and emits
/// linearity diagnostics for the all-explicit-mode fns. Other fns are
/// skipped without traversal.
///
/// Output ordering: defs in declaration order; within a def, diagnostics
/// in source-order discovery (depth-first left-to-right walk).
pub(crate) fn check_module(m: &Module) -> Vec<Diagnostic> {
let mut globals: HashMap<String, Type> = HashMap::new();
for def in &m.defs {
match def {
Def::Fn(f) => {
let ty = strip_forall(&f.ty);
globals.insert(f.name.clone(), ty.clone());
}
Def::Const(c) => {
globals.insert(c.name.clone(), strip_forall(&c.ty).clone());
}
Def::Type(_) => {}
}
}
let mut diags = Vec::new();
for def in &m.defs {
if let Def::Fn(f) = def {
check_fn(f, &globals, &mut diags);
}
}
diags
}
/// Returns the inner type of a top-level `Forall<Fn ...>` (or `t`
/// unchanged if not a `Forall`). The linearity check only inspects
/// `param_modes`, which sit on the inner `Type::Fn`.
fn strip_forall(t: &Type) -> &Type {
match t {
Type::Forall { body, .. } => body,
other => other,
}
}
/// Per-fn check. Skips fns whose signature has any `Implicit` param —
/// see module-level scope rules.
fn check_fn(f: &FnDef, globals: &HashMap<String, Type>, diags: &mut Vec<Diagnostic>) {
let inner = strip_forall(&f.ty);
let (param_tys, param_modes) = match inner {
Type::Fn { params, param_modes, .. } => (params.as_slice(), param_modes.as_slice()),
_ => return,
};
// Activation gate: every param mode must be explicit (Borrow / Own).
// No params at all also disables the check — there are no binders
// to track in the param-list, so the check has nothing to enforce.
if param_modes.is_empty() || param_modes.iter().any(|m| matches!(m, ParamMode::Implicit)) {
return;
}
// Sanity: param-count has been verified by `check_fn` upstream, but
// double-defend in case the type vec is shorter than the binder list
// — the linearity walk treats over-long binder lists as `Implicit`
// and would silently mis-skip the check otherwise.
if param_modes.len() != f.params.len() || param_tys.len() != f.params.len() {
return;
}
let mut checker = Checker {
globals,
diags,
def_name: &f.name,
binders: HashMap::new(),
};
// Install fn parameters as binders, with initial `borrow_count = 1`
// for `Borrow` params (the caller's outer borrow stays live for the
// body's whole duration) and `0` for `Own` params.
for (name, mode) in f.params.iter().zip(param_modes.iter()) {
let initial = BinderState {
consumed: false,
borrow_count: match mode {
ParamMode::Borrow => 1,
ParamMode::Own | ParamMode::Implicit => 0,
},
};
checker.binders.insert(name.clone(), initial);
}
checker.walk(&f.body, Position::Consume);
}
/// Per-fn linearity walker.
struct Checker<'a> {
/// Top-level symbol → type. Used to look up the param_modes of an
/// `App` callee that resolves to a global fn def.
globals: &'a HashMap<String, Type>,
/// Diagnostic accumulator (shared with the module-level walk).
diags: &'a mut Vec<Diagnostic>,
/// Name of the fn currently being checked (becomes
/// [`Diagnostic::def`]).
def_name: &'a str,
/// Live binder state, keyed by name. Modified in place; lexical
/// scoping is restored by [`Checker::with_binder`].
binders: HashMap<String, BinderState>,
}
impl<'a> Checker<'a> {
/// The core position-aware walk. Recurses through `t` and updates
/// per-binder state according to the rules in the module-level
/// docs.
fn walk(&mut self, t: &Term, pos: Position) {
match t {
Term::Lit { .. } => {}
Term::Var { name } => self.use_var(name, pos),
Term::App { callee, args, .. } => {
// The callee itself is consumed (it's "the function value");
// for var-callees this is typically a global fn ref.
self.walk(callee, Position::Consume);
// Determine arg modes from the callee's resolved type.
let arg_modes = self.callee_arg_modes(callee, args.len());
// Track which binders we lent so we can release them
// after the call.
let mut lent: Vec<String> = Vec::new();
for (i, arg) in args.iter().enumerate() {
let arg_pos = match arg_modes.get(i) {
Some(ParamMode::Borrow) => Position::Borrow,
// Own / Implicit / unknown → Consume.
_ => Position::Consume,
};
self.walk(arg, arg_pos);
// If we lent a *bare* binder reference at this arg
// slot, the borrow stays live for the rest of the
// call (i.e. while subsequent args evaluate, and
// until the callee returns). Bump `borrow_count`
// now so a sibling Consume of the same binder
// triggers `consume-while-borrowed`.
if matches!(arg_pos, Position::Borrow) {
if let Term::Var { name } = arg {
if let Some(s) = self.binders.get_mut(name) {
s.borrow_count = s.borrow_count.saturating_add(1);
lent.push(name.clone());
}
}
}
}
// Call returns: all lent borrows end.
for name in lent {
if let Some(s) = self.binders.get_mut(&name) {
s.borrow_count = s.borrow_count.saturating_sub(1);
}
}
}
Term::Let { name, value, body } => {
self.walk(value, Position::Consume);
self.with_binder(name, BinderState::default(), |this| {
this.walk(body, pos);
});
}
Term::LetRec { name, body, in_term, .. } => {
// The body is the recursive fn's own body; treating it as
// an opaque closure is sufficient for 18c.2 — its
// captures are conservatively traversed in `Consume`
// mode but with the recursive `name` masked (so the
// recursive self-reference does not flag as a
// use-after-consume of the in_term's view).
self.with_binder(name, BinderState::default(), |this| {
this.walk(body, Position::Consume);
this.walk(in_term, pos);
});
}
Term::If { cond, then, else_ } => {
self.walk(cond, Position::Consume);
// Branch independence: each arm starts from the pre-branch
// state. If any arm consumes/borrows a binder, that
// outcome is visible to the post-`If` continuation
// (conservative merge: union of consumed flags, max of
// borrow_count).
let saved = self.binders.clone();
self.walk(then, pos);
let after_then = std::mem::replace(&mut self.binders, saved);
self.walk(else_, pos);
merge_states(&mut self.binders, &after_then);
}
Term::Match { scrutinee, arms } => {
// Scrutinee is read (Borrow), not moved.
self.walk(scrutinee, Position::Borrow);
let saved = self.binders.clone();
let mut acc: Option<HashMap<String, BinderState>> = None;
for arm in arms {
self.binders = saved.clone();
self.walk_arm(arm, pos);
match acc.as_mut() {
None => acc = Some(self.binders.clone()),
Some(m) => merge_states(m, &self.binders),
}
}
if let Some(merged) = acc {
self.binders = merged;
} else {
self.binders = saved;
}
}
Term::Seq { lhs, rhs } => {
self.walk(lhs, Position::Consume);
self.walk(rhs, pos);
}
Term::Ctor { args, .. } => {
// Ctor args are consumed (the new value owns them).
for a in args {
self.walk(a, Position::Consume);
}
}
Term::Do { args, .. } => {
for a in args {
self.walk(a, Position::Consume);
}
}
Term::Lam { params, body, .. } => {
// Lam captures cross the linearity boundary: any free
// var of the body is implicitly consumed by closure
// construction (we don't yet model captured-borrow
// discipline; that is 18c.3 territory). For 18c.2 we
// walk the body with the lam's own params as fresh
// binders.
let mut saved_for_params: HashMap<String, Option<BinderState>> = HashMap::new();
for p in params {
saved_for_params.insert(p.clone(), self.binders.remove(p));
self.binders.insert(p.clone(), BinderState::default());
}
self.walk(body, Position::Consume);
for p in params {
self.binders.remove(p);
if let Some(prev) = saved_for_params.remove(p).flatten() {
self.binders.insert(p.clone(), prev);
}
}
}
Term::Clone { value } => {
// (clone X) reads X (Borrow) regardless of outer pos:
// the resulting *value* is owned, but the inner term is
// only borrowed. So `(clone xs)` does NOT consume `xs`.
self.walk(value, Position::Borrow);
}
}
}
/// Walk one match arm: introduce the pattern's bindings as fresh
/// `Live` binders, walk the body, then drop them. Pattern bindings
/// are linearity-fresh in 18c.2; we do not yet track ownership
/// transfer from the scrutinee (that requires modelling pattern
/// borrow vs. consume, which is 18c.3 work).
fn walk_arm(&mut self, arm: &Arm, pos: Position) {
let names = collect_pattern_binders(&arm.pat);
let mut saved: HashMap<String, Option<BinderState>> = HashMap::new();
for n in &names {
saved.insert(n.clone(), self.binders.remove(n));
self.binders.insert(n.clone(), BinderState::default());
}
self.walk(&arm.body, pos);
for n in &names {
self.binders.remove(n);
if let Some(prev) = saved.remove(n).flatten() {
self.binders.insert(n.clone(), prev);
}
}
}
/// Record a use of `name` in `pos`. If `name` is not a tracked
/// binder (i.e. it's a global / built-in), we ignore the use.
///
/// Note: Borrow-position uses do **not** increment `borrow_count`
/// here. The counter tracks borrows whose lifetime extends across
/// sibling subterms — only the [`Term::App`] walker manages those
/// (incrementing before the next sibling, decrementing when the
/// call returns). A scrutinee or `Term::Clone` borrow does not
/// span siblings, so it just checks `consumed` and returns.
fn use_var(&mut self, name: &str, pos: Position) {
let state = match self.binders.get_mut(name) {
Some(s) => s,
None => return,
};
// Already-consumed: any further use is a use-after-consume,
// regardless of position.
if state.consumed {
self.diags
.push(make_use_after_consume(self.def_name, name));
return;
}
match pos {
Position::Borrow => {
// No state change; the App walker is responsible for
// bumping `borrow_count` when a Borrow arg lives across
// sibling args, and for the per-fn-param initial value.
}
Position::Consume => {
if state.borrow_count > 0 {
self.diags
.push(make_consume_while_borrowed(self.def_name, name));
return;
}
state.consumed = true;
}
}
}
/// Resolve the param-mode vector of a fn-call's callee.
///
/// Currently only handles `callee = Term::Var { name }` resolving
/// to a global fn type (with explicit `param_modes`). All other
/// callee shapes return an empty vec → all args default to
/// Consume. This is the conservative, "no false negatives on the
/// borrow-arg case" behaviour: if we can't see the modes, we
/// assume Consume and may miss a borrow.
fn callee_arg_modes(&self, callee: &Term, n_args: usize) -> Vec<ParamMode> {
let name = match callee {
Term::Var { name } => name,
_ => return vec![],
};
// Locals never carry fn-types in 18c.2 (no first-class fn vars
// with mode info). Look up the global symbol table.
let ty = match self.globals.get(name) {
Some(t) => t,
None => return vec![],
};
let inner = strip_forall(ty);
match inner {
Type::Fn { param_modes, .. } => {
if param_modes.is_empty() {
// All-implicit case (legacy): no mode info → all
// Consume, by the rule above.
vec![ParamMode::Implicit; n_args]
} else {
param_modes.clone()
}
}
_ => vec![],
}
}
/// Save the current state of `name`, install `fresh` for the
/// duration of `body`, then restore. Used by `Let` / `LetRec` /
/// `Lam` / pattern-arm to introduce a fresh binder while
/// preserving lexical scope.
fn with_binder<F: FnOnce(&mut Self)>(&mut self, name: &str, fresh: BinderState, body: F) {
let prev = self.binders.remove(name);
self.binders.insert(name.to_string(), fresh);
body(self);
self.binders.remove(name);
if let Some(p) = prev {
self.binders.insert(name.to_string(), p);
}
}
}
/// Recursively collect every binder name introduced by a pattern.
/// Wildcards and literals contribute nothing; `Var` adds itself; `Ctor`
/// recurses into fields.
fn collect_pattern_binders(p: &Pattern) -> Vec<String> {
match p {
Pattern::Wild | Pattern::Lit { .. } => vec![],
Pattern::Var { name } => vec![name.clone()],
Pattern::Ctor { fields, .. } => fields
.iter()
.flat_map(|f| collect_pattern_binders(f))
.collect(),
}
}
/// Conservative merge of two branch states: a binder is consumed in
/// the merged state iff it was consumed in either branch; borrow_count
/// becomes the max of the two. Names present only in one map are
/// carried through.
fn merge_states(into: &mut HashMap<String, BinderState>, other: &HashMap<String, BinderState>) {
for (name, s) in other {
let entry = into.entry(name.clone()).or_default();
entry.consumed = entry.consumed || s.consumed;
entry.borrow_count = entry.borrow_count.max(s.borrow_count);
}
}
/// Build a `use-after-consume` diagnostic for `binder` in `def`.
/// Suggested rewrite: clone the offending use (the LATER use is the
/// flagged one; turning IT into `(clone <binder>)` doesn't help, but
/// the suggested rewrite is the intent — "this use should have been
/// `(clone X)`, or an earlier use should have been"). The form-A
/// snippet is `(clone <binder>)`, parseable by
/// [`ailang_surface::parse_term`].
fn make_use_after_consume(def: &str, binder: &str) -> Diagnostic {
let replacement = term_to_form_a(&Term::Clone {
value: Box::new(Term::Var { name: binder.to_string() }),
});
Diagnostic::error(
"use-after-consume",
format!("`{binder}` is used after it was already consumed"),
)
.with_def(def)
.with_ctx(serde_json::json!({"binder": binder}))
.with_suggested_rewrite(
format!(
"wrap an earlier use of `{binder}` in `(clone)` so this later use stays the unique consume"
),
replacement,
)
}
/// Build a `consume-while-borrowed` diagnostic for `binder` in `def`.
/// Suggested rewrite: clone here so the original retains the borrow.
fn make_consume_while_borrowed(def: &str, binder: &str) -> Diagnostic {
let replacement = term_to_form_a(&Term::Clone {
value: Box::new(Term::Var { name: binder.to_string() }),
});
Diagnostic::error(
"consume-while-borrowed",
format!("`{binder}` is consumed while a borrow of it is still live"),
)
.with_def(def)
.with_ctx(serde_json::json!({"binder": binder}))
.with_suggested_rewrite(
format!(
"consume `(clone {binder})` here so the original keeps its outstanding borrow"
),
replacement,
)
}
#[cfg(test)]
mod tests {
use super::*;
use ailang_core::ast::{FnDef, Literal};
fn fn_with_modes(name: &str, modes: Vec<ParamMode>, body: Term) -> Def {
Def::Fn(FnDef {
name: name.into(),
ty: Type::Fn {
params: modes
.iter()
.map(|_| Type::Con { name: "List".into(), args: vec![] })
.collect(),
param_modes: modes.clone(),
ret: Box::new(Type::int()),
ret_mode: ParamMode::Implicit,
effects: vec![],
},
params: (0..modes.len()).map(|i| format!("p{i}")).collect(),
body,
doc: None,
})
}
/// All-Implicit fn → check is skipped, no diagnostics, even if the
/// body would trigger a use-after-consume under explicit modes.
#[test]
fn implicit_fn_is_exempt() {
let m = Module {
schema: ailang_core::SCHEMA.into(),
name: "t".into(),
imports: vec![],
defs: vec![fn_with_modes(
"f",
vec![ParamMode::Implicit],
Term::Var { name: "p0".into() },
)],
};
assert!(check_module(&m).is_empty());
}
/// Mixed Implicit + Borrow → still skipped (any Implicit disables
/// the check entirely, per the activation gate).
#[test]
fn mixed_implicit_explicit_is_exempt() {
let m = Module {
schema: ailang_core::SCHEMA.into(),
name: "t".into(),
imports: vec![],
defs: vec![fn_with_modes(
"f",
vec![ParamMode::Borrow, ParamMode::Implicit],
Term::Lit { lit: Literal::Int { value: 0 } },
)],
};
assert!(check_module(&m).is_empty());
}
/// All-explicit fn that NEVER uses its params: clean.
#[test]
fn explicit_fn_with_no_uses_is_clean() {
let m = Module {
schema: ailang_core::SCHEMA.into(),
name: "t".into(),
imports: vec![],
defs: vec![fn_with_modes(
"f",
vec![ParamMode::Own],
Term::Lit { lit: Literal::Int { value: 0 } },
)],
};
assert!(check_module(&m).is_empty());
}
}