plum

# Closures Implementation Plan

> **For agentic workers:** REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task. Steps use checkbox (`- [ ]`) syntax for tracking.

**Goal:** Make `|params| body` closure literals — full, capturing closures with snapshot-by-value semantics, usable as an ordinary call argument — parse, type-check, and compile to working wasm.

**Architecture:** Grammar wires the already-existing (but unreachable) `closure` rule into expression position and adds a new `fn(...)` function-value type annotation. The AST gains `Expr::Closure` and `ParamType::Fn`. The checker's existing `PlumType::TFun` models a closure's type directly, and calling a closure-typed binding by name already works through the unmodified `FnCall` inference path. Codegen is the substantial part: wasm has no native closures, so every closure literal compiles to its own real wasm function (registered in a new function table, with an implicit first parameter — the captured-environment pointer, mirroring how a method already receives `self`), and a closure *value* at runtime is a single `i32` pointer to a heap-allocated `{table_index, env_pointer}` pair. Calling a closure value loads both fields and does `call_indirect`.

**Tech Stack:** Rust (workspace: `plum-core`, `plum-checker`, `plum-wasm-codegen`), tree-sitter grammar, `wasm-encoder` (already a dependency; this plan uses its `TableSection`/`ElementSection`/`Elements`/`Instruction::CallIndirect`, none of which this codebase uses yet).

## Global Constraints

- **Scope**: closures as expressions, passable as an ordinary call argument (`each(|v| ...)`) — NOT the trailing-closure calling convention (`each() |v| ...`) shown in `libs/std/list.plum`'s aspirational draft; NOT a class field storing a closure value. Capture is snapshot-by-value (a captured variable's value at closure-creation time), not live/shared mutation.
- Function-type annotations use **positional types only** — `fn(Int) -> Bool`, `fn(a) -> b` — no param names inside the annotation.
- Task 5 (closure discovery + codegen) is substantial new algorithmic code, not a modification of existing tested logic. Its design is sound but should be validated primarily through that task's own tests (TDD), the same posture the generics-monomorphization plan took for its hardest task — expect to need judgment calls during implementation, and report BLOCKED/NEEDS_CONTEXT rather than papering over a genuine design gap, exactly as that plan's precedent established.
- Follow existing code style: terse one-line "why" comments only where non-obvious; error messages use the existing `"codegen: ..."` / `"monomorphize: ..."` prefixes as appropriate.
- Every task must leave `cargo test --workspace` and `npx --yes tree-sitter-cli test` (from `tooling/tree-sitter-plum/`) green before moving to the next task.

---

### Task 1: Grammar — wire closures into expression position, add function-value type syntax

**Files:**
- Modify: `tooling/tree-sitter-plum/grammar.js`
- Test: `tooling/tree-sitter-plum/test/corpus/` (new cases)

**Interfaces:**
- Consumes: nothing from other tasks.
- Produces: a `closure` node reachable from `expression`, and a new `fn_value_type` node reachable from `param`'s type field. Neither `plum-core`'s parser nor its AST need any changes yet — that's Task 2.

- [ ] **Step 1: Uncomment `$.closure` in `expression`'s choice list**

Find (in `grammar.js`):

```js
    expression: ($) =>
      choice(
        $.comparison_operator,
        $.not_operator,
        $.boolean_operator,
        // $.closure,
        $.primary_expression,
        $.ternary_expression,
      ),
```

Change to:

```js
    expression: ($) =>
      choice(
        $.comparison_operator,
        $.not_operator,
        $.boolean_operator,
        $.closure,
        $.primary_expression,
        $.ternary_expression,
      ),
```

- [ ] **Step 2: Add the `fn_value_type` rule and wire it into `param`**

Find:

```js
    param: ($) =>
      seq(
        field("name", $.var_identifier),
        ":",
        field("type", choice($.type, $.variadic_type)),
        optional(seq("=", field("value", $.expression))),
      ),
```

Change to:

```js
    param: ($) =>
      seq(
        field("name", $.var_identifier),
        ":",
        field("type", choice($.type, $.variadic_type, $.fn_value_type)),
        optional(seq("=", field("value", $.expression))),
      ),

    fn_value_type: ($) =>
      seq(
        "fn",
        "(",
        field("params", optional(commaSep1($.type))),
        ")",
        optional(seq("->", field("returns", $.type))),
      ),
```

(Place `fn_value_type` near `variadic_type`'s definition, a few lines above `param`. The explicit `"params"`/`"returns"` field names — the same idiom `fn`'s own `field("returns", ...)` already uses — let `parser.rs` (Task 2) distinguish the return type from the param types unambiguously by field, rather than by counting/positional-guessing among same-kind `"type"` children.)

- [ ] **Step 3: Regenerate and run the existing corpus suite**

```bash
cd tooling/tree-sitter-plum && npx --yes tree-sitter-cli generate && npx --yes tree-sitter-cli test
```

Expected: generation succeeds with no unresolved-conflict errors, and all pre-existing corpus cases still pass. (The `closure` rule's own body — `"|" params "|" body` — already exists unchanged; only its reachability and the new `fn_value_type` rule are new. If `generate` reports a conflict — e.g. between `|` as the closure delimiter and `|` as the bitwise-or binary operator in expression position — read the reported example carefully; a `conflicts: [[$.closure, ...]]` entry or precedence adjustment may be needed. Don't guess a fix blindly; the existing `conflicts` array in `grammar.js` (already used for the `fn_call`/`class_call` ambiguity from prior work) is the established idiom for this.)

- [ ] **Step 4: Add corpus cases**

Append to `tooling/tree-sitter-plum/test/corpus/function.txt` (input halves — the next step fills in expected trees):

```
================================================================================
function - closure literal in expression position
================================================================================

useClosure() -> Bool =
  cb = |v|
    True
  cb(5)

--------------------------------------------------------------------------------
================================================================================
function - closure literal with no params
================================================================================

useClosure() -> Bool =
  cb = ||
    True
  cb()

--------------------------------------------------------------------------------
================================================================================
function - function-value type param annotation
================================================================================

each(cb: fn(Int)) -> Bool =
  True

--------------------------------------------------------------------------------
================================================================================
function - function-value type param annotation with generic types and return
================================================================================

each(cb: fn(a) -> b) -> Bool =
  True

--------------------------------------------------------------------------------
```

- [ ] **Step 5: Generate expected trees and verify**

```bash
cd tooling/tree-sitter-plum && npx --yes tree-sitter-cli test -u -f "closure literal" && npx --yes tree-sitter-cli test -u -f "function-value type"
```

Open `test/corpus/function.txt` and confirm each new case's generated tree has no `ERROR`/`MISSING` node, and that the closure literal appears as a single `closure` node (not split), and `fn_value_type` appears as a single node containing its param types and optional return type.

- [ ] **Step 6: Run the full corpus suite and the Rust workspace suite**

```bash
cd tooling/tree-sitter-plum && npx --yes tree-sitter-cli test
cargo test --workspace
```

Expected: both green. (The workspace suite should be unaffected — no Rust code changes yet — but confirm nothing regresses via the parser's generic wrapper-node handling, the same way past grammar-only changes this session were verified.)

- [ ] **Step 7: Commit**

```bash
git add tooling/tree-sitter-plum/grammar.js tooling/tree-sitter-plum/test/corpus/function.txt
git add tooling/tree-sitter-plum/src  # generated parser.c etc, only if tracked — check git status first
git commit -m "feat(tree-sitter-plum): wire closure literals into expression position; add fn(...) type syntax"
```

---

### Task 2: AST + Parser — `Expr::Closure`, `ParamType::Fn`

**Files:**
- Modify: `plum-core/src/ast.rs`
- Modify: `plum-core/src/parser.rs`
- Test: `plum-core/tests/` (a new or extended integration test parsing a closure and a `fn(...)` param — check existing test file conventions, e.g. `formatter_test.rs`, for how this crate structures its own tests; there is no dedicated parser test file today, so add one, `plum-core/tests/parser_test.rs`, following the same `AstParser::new(src)` + `tree_sitter::Parser` pattern already used in `plum-checker`/`plum-wasm-codegen`'s own test helpers)

**Interfaces:**
- Consumes: the `closure`/`fn_value_type` grammar nodes from Task 1.
- Produces:
  ```rust
  pub struct Closure { pub params: Vec<String>, pub body: Block }
  // added to Expr:
  Closure(Box<Closure>),
  // added to ParamType:
  Fn(Vec<Type>, Option<Box<Type>>),
  ```
  Task 3 (checker) and Task 5 (codegen) both match on these.

- [ ] **Step 1: Write failing tests**

Create `plum-core/tests/parser_test.rs`:

```rust
use plum_core::ast::*;
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(), "parse error:\n{}", tree.root_node().to_sexp());
    let ap = AstParser::new(src);
    ap.parse_source(tree.root_node())
}

fn only_fn(source: &Source) -> &Fn {
    source.items.iter().find_map(|i| match i { Item::Fn(f) => Some(f), _ => None }).expect("expected a Fn item")
}

#[test]
fn closure_literal_parses_with_params_and_body() {
    let src = "\
useClosure() -> Bool =
  cb = |v|
    True
  cb(5)
";
    let source = parse(src);
    let f = only_fn(&source);
    let FnBody::Block(block) = &f.body else { panic!("expected a block body") };
    let Stmt::Assign(assign) = &block.stmts[0] else { panic!("expected an assign statement") };
    let Expr::Closure(closure) = &assign.values[0] else { panic!("expected a closure expression, got {:?}", assign.values[0]) };
    assert_eq!(closure.params, vec!["v".to_string()]);
    assert_eq!(closure.body.stmts.len(), 1);
}

#[test]
fn closure_literal_parses_with_no_params() {
    let src = "\
useClosure() -> Bool =
  cb = ||
    True
  cb()
";
    let source = parse(src);
    let f = only_fn(&source);
    let FnBody::Block(block) = &f.body else { panic!("expected a block body") };
    let Stmt::Assign(assign) = &block.stmts[0] else { panic!("expected an assign statement") };
    let Expr::Closure(closure) = &assign.values[0] else { panic!("expected a closure expression") };
    assert!(closure.params.is_empty());
}

#[test]
fn fn_value_type_param_parses_with_positional_types_and_return() {
    let src = "each(cb: fn(Int) -> Bool) -> Bool =\n  True\n";
    let source = parse(src);
    let f = only_fn(&source);
    let ParamType::Fn(param_types, ret) = &f.params[0].ty else { panic!("expected ParamType::Fn, got {:?}", f.params[0].ty) };
    assert_eq!(param_types.len(), 1);
    assert_eq!(param_types[0].name, "Int");
    assert_eq!(ret.as_ref().map(|t| t.name.clone()), Some("Bool".to_string()));
}

#[test]
fn fn_value_type_param_parses_with_no_return() {
    let src = "each(cb: fn(Int)) -> Bool =\n  True\n";
    let source = parse(src);
    let f = only_fn(&source);
    let ParamType::Fn(param_types, ret) = &f.params[0].ty else { panic!("expected ParamType::Fn") };
    assert_eq!(param_types.len(), 1);
    assert!(ret.is_none());
}
```

- [ ] **Step 2: Run to see them fail**

Run: `cargo test -p plum-core --test parser_test`
Expected: fails to compile — `Expr::Closure`/`ParamType::Fn` don't exist yet.

- [ ] **Step 3: Add the AST nodes**

In `plum-core/src/ast.rs`, find the `Expr` enum (it currently ends with variants like `Var(String)`, `TypeName(String)`) and add a new variant:

```rust
    /// `|params| body`
    Closure(Box<Closure>),
```

Add the `Closure` struct near `Block`'s definition:

```rust
#[derive(Debug, Clone, PartialEq)]
pub struct Closure {
    pub params: Vec<String>,
    pub body: Block,
}
```

Find the `ParamType` enum (`Type(Type)`, `Variadic(Type)`) and add:

```rust
    /// `fn(Int, Str) -> Bool` — a function-value type annotation. Positional types
    /// only, no param names (types don't need names).
    Fn(Vec<Type>, Option<Box<Type>>),
```

- [ ] **Step 4: Parse `closure` and `fn_value_type` nodes**

In `plum-core/src/parser.rs`, find `parse_primary_expression`'s match (it currently has arms like `"binary_operator" => ...`, `"fn_call" => ...`) — but `closure` is NOT under `primary_expression` in the grammar, it's a direct alternative of `expression` itself, so it needs handling in `parse_expression` instead. Find:

```rust
    pub fn parse_expression(&self, node: Node) -> Expr {
        let node = self.unwrap_expr_node(node);
        match node.kind() {
            "comparison_operator" => self.parse_compare(node),
            "not_operator" => {
                let arg = node.named_child(0)
                    .map(|n| { let u = self.unwrap_expr_node(n); self.parse_expression(u) })
                    .unwrap_or(Expr::Int(0));
                Expr::Not(Box::new(arg))
            }
            "boolean_operator" => self.parse_bool_op(node),
            "ternary_expression" => self.parse_ternary(node),
            _ => self.parse_primary_expression(node),
        }
    }
```

Change to:

```rust
    pub fn parse_expression(&self, node: Node) -> Expr {
        let node = self.unwrap_expr_node(node);
        match node.kind() {
            "comparison_operator" => self.parse_compare(node),
            "not_operator" => {
                let arg = node.named_child(0)
                    .map(|n| { let u = self.unwrap_expr_node(n); self.parse_expression(u) })
                    .unwrap_or(Expr::Int(0));
                Expr::Not(Box::new(arg))
            }
            "boolean_operator" => self.parse_bool_op(node),
            "ternary_expression" => self.parse_ternary(node),
            "closure" => Expr::Closure(Box::new(self.parse_closure(node))),
            _ => self.parse_primary_expression(node),
        }
    }

    fn parse_closure(&self, node: Node) -> Closure {
        // closure: "|" var_identifier,* "|" body
        let params: Vec<String> = self.children_of_kind(node, "var_identifier")
            .into_iter()
            .map(|n| self.text(n))
            .collect();
        let body = self.children_of_kind(node, "body")
            .into_iter()
            .next()
            .map(|n| self.parse_block(n))
            .unwrap_or(Block { stmts: vec![] });
        Closure { params, body }
    }
```

Find `parse_param` (it currently matches `n.kind() == "variadic_type"` vs. the `else` branch treating everything else as a plain `type`):

```rust
    fn parse_param(&self, node: Node) -> Param {
        // param: var_identifier ":" (type | variadic_type) ("=" expression)?
        let name = node.named_child(0).map(|n| self.text(n)).unwrap_or_default();
        let ty = node.named_child(1).map(|n| {
            if n.kind() == "variadic_type" {
                let inner = n.named_child(0)
                    .map(|t| self.parse_type(t))
                    .unwrap_or(Type { name: String::new(), generics: vec![] });
                ParamType::Variadic(inner)
            } else {
                ParamType::Type(self.parse_type(n))
            }
        }).unwrap_or(ParamType::Type(Type { name: String::new(), generics: vec![] }));
        let default = node.named_child(2).map(|n| {
            let unwrapped = self.unwrap_expr_node(n);
            self.parse_expression(unwrapped)
        });
        Param { name, ty, default }
    }
```

Change the type-dispatch to also handle `fn_value_type`:

```rust
    fn parse_param(&self, node: Node) -> Param {
        // param: var_identifier ":" (type | variadic_type | fn_value_type) ("=" expression)?
        let name = node.named_child(0).map(|n| self.text(n)).unwrap_or_default();
        let ty = node.named_child(1).map(|n| match n.kind() {
            "variadic_type" => {
                let inner = n.named_child(0)
                    .map(|t| self.parse_type(t))
                    .unwrap_or(Type { name: String::new(), generics: vec![] });
                ParamType::Variadic(inner)
            }
            "fn_value_type" => self.parse_fn_value_type(n),
            _ => ParamType::Type(self.parse_type(n)),
        }).unwrap_or(ParamType::Type(Type { name: String::new(), generics: vec![] }));
        let default = node.named_child(2).map(|n| {
            let unwrapped = self.unwrap_expr_node(n);
            self.parse_expression(unwrapped)
        });
        Param { name, ty, default }
    }

    fn parse_fn_value_type(&self, node: Node) -> ParamType {
        // fn_value_type: "fn" "(" field("params", type,*) ")" ("->" field("returns", type))?
        // The "returns" field (if present) is a distinct field from "params", so the
        // two are disambiguated unambiguously by field name, not by counting/position
        // among same-kind "type" children — the same idiom `fn`'s own `returns` field
        // already uses.
        let returns_node = node.child_by_field_name("returns");
        let param_types: Vec<Type> = self.children_of_kind(node, "type")
            .into_iter()
            .filter(|n| Some(*n) != returns_node)
            .map(|n| self.parse_type(n))
            .collect();
        let ret = returns_node.map(|n| Box::new(self.parse_type(n)));
        ParamType::Fn(param_types, ret)
    }
```

- [ ] **Step 5: Run parser tests**

Run: `cargo test -p plum-core --test parser_test`
Expected: all 4 tests pass. If `parse_fn_value_type`'s field extraction needed a grammar tweak (per the note above), go back and apply it, re-running Task 1's corpus tests too.

- [ ] **Step 6: Run the full workspace suite**

Run: `cargo test --workspace`
Expected: green (new AST variants are additive; nothing existing matches on `Expr`/`ParamType` exhaustively in a way that would fail to compile — verify this by checking for compile errors from non-exhaustive `match` arms in `plum-checker`/`plum-wasm-codegen`, and add a minimal `_ => ...` fallback arm or explicit handling wherever the compiler flags one).

- [ ] **Step 7: Commit**

```bash
git add plum-core/src/ast.rs plum-core/src/parser.rs plum-core/tests/parser_test.rs
git commit -m "feat(plum-core): parse closure literals and fn(...) type annotations"
```

---

### Task 3: Checker — infer a closure's type; confirm closure calls type-check

**Files:**
- Modify: `plum-checker/src/lib.rs`
- Test: `plum-checker/tests/checker_tests.rs`

**Interfaces:**
- Consumes: `ast::Expr::Closure`, `ast::ParamType::Fn` (Task 2).
- Produces: `infer_expr` handles `Expr::Closure`, returning `PlumType::TFun`. No other function signatures change — `check_fn`'s existing `ParamType` match (used to bind each param's type into the local env) needs one new arm for `ParamType::Fn` too, converting it to `PlumType::TFun`.

- [ ] **Step 1: Write failing tests**

Append to `plum-checker/tests/checker_tests.rs`:

```rust
#[test]
fn closure_literal_infers_as_a_function_type() {
    let src = "\
useClosure() -> Bool =
  cb = |v|
    True
  cb(5)
";
    let source = parse(src);
    let result = check_source(&source);
    assert!(result.is_ok(), "expected Ok, got {:?}", result.err());
}

#[test]
fn fn_value_typed_param_can_be_called() {
    let src = "\
each(cb: fn(Int) -> Bool) -> Bool =
  cb(5)
";
    let source = parse(src);
    let result = check_source(&source);
    assert!(result.is_ok(), "expected Ok, got {:?}", result.err());
}

#[test]
fn closure_passed_to_fn_value_typed_param_type_checks() {
    let src = "\
each(cb: fn(Int) -> Bool) -> Bool =
  cb(5)

use() -> Bool =
  each(|v|
    True)
";
    let source = parse(src);
    let result = check_source(&source);
    assert!(result.is_ok(), "expected Ok, got {:?}", result.err());
}
```

- [ ] **Step 2: Run to see them fail**

Run: `cargo test -p plum-checker --test checker_tests closure`
Expected: fails to compile initially (`infer_expr` doesn't exhaustively handle `Expr::Closure` — a non-exhaustive match compile error) or, once that's stubbed minimally, fails at runtime because `ParamType::Fn` isn't converted to a `PlumType` anywhere yet.

- [ ] **Step 3: Add `infer_expr`'s `Closure` arm**

In `plum-checker/src/lib.rs`, find `infer_expr`'s match (it has arms for `Expr::Int`, `Expr::Var`, etc., ending around the `Expr::Attribute` arm). Add:

```rust
        ast::Expr::Closure(cl) => {
            let mut closure_env = env.clone();
            let param_types: Vec<PlumType> = cl.params.iter().map(|p| {
                let t = PlumType::TVar(format!("_closure_{}", p));
                closure_env.insert(p.clone(), TypeScheme::mono(t.clone()));
                t
            }).collect();
            let body_ty = match &cl.body.stmts.last() {
                Some(ast::Stmt::Expr(e)) => infer_expr(e, &closure_env, ctx)?,
                Some(ast::Stmt::Return(Some(e))) => infer_expr(e, &closure_env, ctx)?,
                _ => PlumType::TUnit,
            };
            Ok(PlumType::TFun(param_types, Box::new(body_ty)))
        }
```

- [ ] **Step 4: Convert `ParamType::Fn` to a `PlumType` wherever `plum_type_from_ast`-style conversion happens for params**

Find every place in `plum-checker/src/lib.rs` that matches on `ast::ParamType::Type(t) => ...` / `ast::ParamType::Variadic(t) => ...` together (e.g. inside `build_global_tables`, `check_fn`) — there are a few such call sites. For each, add a third arm:

```rust
                    ast::ParamType::Fn(param_types, ret) => PlumType::TFun(
                        param_types.iter().map(plum_type_from_ast).collect(),
                        Box::new(ret.as_ref().map(|r| plum_type_from_ast(r)).unwrap_or(PlumType::TUnit)),
                    ),
```

(Match the exact surrounding style at each call site — some are inline closures passed to `.map()`, some are `match &p.ty { ... }` blocks. Search for `ast::ParamType::Variadic` to find every site that needs the parallel `Fn` arm — the compiler will also flag any non-exhaustive match as a hard error, which is the authoritative list.)

- [ ] **Step 5: Run checker tests**

Run: `cargo test -p plum-checker --test checker_tests closure`
Expected: all 3 pass.

- [ ] **Step 6: Run the full checker crate suite**

Run: `cargo test -p plum-checker`
Expected: green.

- [ ] **Step 7: Commit**

```bash
git add plum-checker/src/lib.rs plum-checker/tests/checker_tests.rs
git commit -m "feat(plum-checker): infer closure literal types; type-check fn(...)-typed params"
```

---

### Task 4: Codegen infrastructure — wasm function table + element section support

**Files:**
- Modify: `plum-wasm-codegen/src/lib.rs`
- Test: `plum-wasm-codegen/tests/codegen_tests.rs`

**Interfaces:**
- Consumes: nothing from other tasks — this is pure `WasmModule` infrastructure, independently testable without any closure-compiling logic existing yet.
- Produces: `WasmModule::add_table_element(&mut self, func_idx: u32) -> u32`, appending `func_idx` to a single funcref table and returning its table index. `WasmModule::finish()` now emits a `Table` section (only if any elements were added) between the existing Function and Memory sections, and an `Element` section between the existing Export and Code sections — both are the correct binary positions per the wasm module section order (Type, Import, Function, **Table**, Memory, Global, Export, **Element**, Code, Data).

- [ ] **Step 1: Write a failing test proving a table + element section round-trips through a real module**

Append to `plum-wasm-codegen/tests/codegen_tests.rs`:

```rust
#[test]
fn wasm_module_with_a_table_element_validates_and_call_indirect_works() {
    // Exercises WasmModule's new table/element support directly, independent of any
    // closure-compiling logic (which doesn't exist yet) — builds a tiny module by
    // hand: one function that returns 42, registered as table element 0, called via
    // `call_indirect` from `main` using a runtime-computed (not compile-time-constant)
    // table index, to prove the table/element wiring is real, not coincidentally
    // skipped by validation.
    let mut module = plum_wasm_codegen::WasmModule::new();
    let ret42_type = module.add_type(&[], &[wasm_encoder::ValType::I64]);
    let ret42_idx = module.add_function(ret42_type, &{
        let mut body = Vec::new();
        wasm_encoder::Instruction::I64Const(42).encode(&mut body);
        wasm_encoder::Instruction::End.encode(&mut body);
        body
    });
    let table_idx = module.add_table_element(ret42_idx);
    assert_eq!(table_idx, 0);

    let main_type = module.add_type(&[], &[wasm_encoder::ValType::I64]);
    let main_idx = module.add_function(main_type, &{
        let mut body = Vec::new();
        wasm_encoder::Instruction::I32Const(0).encode(&mut body); // table index operand
        wasm_encoder::Instruction::CallIndirect { type_index: ret42_type, table_index: 0 }.encode(&mut body);
        wasm_encoder::Instruction::End.encode(&mut body);
        body
    });
    module.add_export("main", wasm_encoder::ExportKind::Func, main_idx);

    let bytes = module.finish();
    let result = wasmparser::validate(&bytes);
    assert!(result.is_ok(), "wasm validation failed: {:?}", result.err());

    let engine = wasmtime::Engine::default();
    let wasm_module = wasmtime::Module::new(&engine, &bytes).expect("module should be loadable");
    let mut store = wasmtime::Store::new(&engine, ());
    let instance = wasmtime::Instance::new(&mut store, &wasm_module, &[]).expect("module should instantiate");
    let main = instance.get_typed_func::<(), i64>(&mut store, "main").expect("main should have signature () -> i64");
    assert_eq!(main.call(&mut store, ()).expect("main should not trap"), 42);
}
```

- [ ] **Step 2: Run to see it fail**

Run: `cargo test -p plum-wasm-codegen --test codegen_tests wasm_module_with_a_table`
Expected: fails to compile — `add_table_element` doesn't exist yet.

- [ ] **Step 3: Add table/element support to `WasmModule`**

In `plum-wasm-codegen/src/lib.rs`, find the `WasmModule` struct:

```rust
pub struct WasmModule {
    types: Vec<FuncType>,
    imports: Vec<(String, String, u32)>,
    functions: Vec<(u32, Vec<u8>)>,
    exports: Vec<(String, ExportKind, u32)>,
    memories: Vec<MemoryType>,
    globals: Vec<(ValType, bool, Vec<u8>)>,
    data_segments: Vec<(u32, Vec<u8>)>,
    pub func_import_count: u32,
    pub func_count: u32,
    global_count: u32,
}
```

Add a new field:

```rust
pub struct WasmModule {
    types: Vec<FuncType>,
    imports: Vec<(String, String, u32)>,
    functions: Vec<(u32, Vec<u8>)>,
    exports: Vec<(String, ExportKind, u32)>,
    memories: Vec<MemoryType>,
    globals: Vec<(ValType, bool, Vec<u8>)>,
    data_segments: Vec<(u32, Vec<u8>)>,
    /// Function indices, in table order — the single funcref table used for
    /// closure `call_indirect` dispatch. Index into this vec IS the table index.
    table_elements: Vec<u32>,
    pub func_import_count: u32,
    pub func_count: u32,
    global_count: u32,
}
```

Update `WasmModule::new()`'s struct literal to add `table_elements: Vec::new(),`.

Add a new method, near `add_memory`:

```rust
    /// Registers `func_idx` as the next slot in the single funcref table used for
    /// closure `call_indirect` dispatch, returning its table index.
    pub fn add_table_element(&mut self, func_idx: u32) -> u32 {
        let table_idx = self.table_elements.len() as u32;
        self.table_elements.push(func_idx);
        table_idx
    }
```

In `finish()`, find the Function section block, ending with `module.section(&funcs); }`, and insert a Table section right after it (before the Memory section block):

```rust
        // Table section
        if !self.table_elements.is_empty() {
            let mut tables = TableSection::new();
            tables.table(TableType {
                element_type: RefType::FUNCREF,
                minimum: self.table_elements.len() as u64,
                maximum: Some(self.table_elements.len() as u64),
                table64: false,
                shared: false,
            });
            module.section(&tables);
        }
```

Find the Export section block, ending with `module.section(&exports); }`, and insert an Element section right after it (before the Code section block):

```rust
        // Element section
        if !self.table_elements.is_empty() {
            let mut elements = ElementSection::new();
            let offset = ConstExpr::i32_const(0);
            elements.active(Some(0), &offset, Elements::Functions(std::borrow::Cow::Borrowed(&self.table_elements)));
            module.section(&elements);
        }
```

- [ ] **Step 4: Run the new test**

Run: `cargo test -p plum-wasm-codegen --test codegen_tests wasm_module_with_a_table`
Expected: passes.

- [ ] **Step 5: Run the full workspace suite**

Run: `cargo test --workspace`
Expected: green — confirm no existing test that builds a `WasmModule` and checks its exact byte output (if any) is affected by the new always-present-but-conditionally-emitted table field (it should be a pure no-op when `table_elements` is empty, since both new sections are gated on `!self.table_elements.is_empty()`).

- [ ] **Step 6: Commit**

```bash
git add plum-wasm-codegen/src/lib.rs plum-wasm-codegen/tests/codegen_tests.rs
git commit -m "feat(plum-wasm-codegen): add wasm function table + element section support"
```

---

### Task 5: Codegen — compile closure literals and closure calls

**Files:**
- Modify: `plum-wasm-codegen/src/lib.rs`
- Test: `plum-wasm-codegen/tests/codegen_tests.rs`

**This is the largest, most judgment-requiring task in this plan** — genuinely new algorithmic code (free-variable analysis, a program-wide discovery pre-pass, synthetic function generation), not a modification of existing tested logic. Treat the design below as your starting architecture, not verified-correct code to transcribe — validate every piece via the tests as you go, and stop and report BLOCKED/NEEDS_CONTEXT if a specific piece doesn't hold up under testing, per this plan's Global Constraints.

**Interfaces:**
- Consumes: `ast::Expr::Closure`, `ast::ParamType::Fn` (Task 2); `WasmModule::add_table_element` (Task 4); `PlumType::TFun` (already exists; Task 3 makes the checker produce it for closures).
- Produces: `compile_source` compiles every closure literal in the program to its own wasm function and table entry; `compile_expr` handles `Expr::Closure` (constructing the runtime closure value) and calling a closure-typed binding (`cb(x)` where `cb`'s inferred type is `TFun`, via `call_indirect` instead of a direct `Call`).

**Design:**

1. **Runtime representation.** A closure value is a single `i32` pointer to a heap-allocated pair `{table_index: i32, env_pointer: i32}` (8-byte stride per field, matching every other class-like struct already in this codegen — `table_index` at offset 0, `env_pointer` at offset 8). The captured-environment struct itself is a separate heap allocation: one 8-byte slot per captured variable, in a stable (e.g. alphabetical, or first-appearance) order.

2. **Discovery pre-pass.** Before compiling any function body, walk every `ast::Item::Fn`'s body (a new walker, since `Collector` runs per-function *after* registration and doesn't cross function boundaries) to find every `ast::Expr::Closure` node. For each one found as an argument to a call whose corresponding declared param type is `ast::ParamType::Fn(param_types, ret)` (already-monomorphized concrete types, since this pre-pass runs on the output of `monomorphize_source`), record:
   - a synthetic mangled name (e.g. `format!("closure${}", n)` with a simple incrementing counter);
   - the concrete wasm param/return `ValType`s (from `param_types`/`ret` via the existing `ast_type_to_wasm`);
   - its free variables — every `Var(name)` referenced in the closure's body that is NOT one of the closure's own params — resolved against the *enclosing* function's locals (both real params/assigned locals AND, if the closure is itself nested inside another closure, that outer closure's own free variables/params) — with each free variable's `PlumType` (via `infer_local_type` against the enclosing scope's type env, exactly as `Collector`/`compile_expr` already do elsewhere in this file).

   Store the result in a new `HashMap<usize, ClosureInfo>` (keyed by the closure `Expr`'s pointer identity, the same keying convention `classcall_scratch`/`match_scratch_index` already use), where:
   ```rust
   struct ClosureInfo {
       mangled_name: String,
       func_idx: u32,
       table_idx: u32,
       param_vts: Vec<ValType>,
       ret_vt: Option<ValType>,
       free_vars: Vec<(String, PlumType)>, // stable order
   }
   ```
   Add this map to `CompileCtx` as `pub closures: HashMap<usize, ClosureInfo>`.

3. **Registering each closure as a real function.** For each discovered closure, `module.add_type(...)` for `(env_ptr: I32, ...param_vts) -> ret_vt`, then reserve a function slot via `module.add_function(type_idx, &[])` (placeholder body, patched later — exactly how `compile_source` already pre-registers every ordinary `Fn`'s slot before compiling bodies) to get `func_idx`, then `module.add_table_element(func_idx)` to get `table_idx`.

4. **Compiling each closure's own body.** Build a `LocalCtx` for the closure much like `compile_fn_body` already does for an ordinary `Fn`, except: local 0 is the (unused-by-name) env pointer param; each free variable gets its own local slot, loaded from the env pointer at function entry (`LocalGet(env_ptr_local); {I64,F64,I32}Load(offset); LocalSet(free_var_local)`, offsets assigned in the same stable order used when constructing the env struct at the call site); each real closure param gets ordinary param-local treatment, offset by 1 (to account for the implicit env-pointer param at index 0) plus however many free-var locals precede it in your chosen local-numbering scheme. Compile the closure's `Block` body via the existing `compile_block_as_fn_body`/value-position machinery, exactly as an ordinary function's block body already is.

5. **Compiling a closure literal's construction site** (new `Expr::Closure` arm in `compile_expr`): reserve two scratch locals (extend the existing `classcall_scratch`-style pool, or add a parallel `closure_scratch: HashMap<usize, u32>` pool reserving 2 consecutive slots per closure literal — mirroring how `Collector`/`LocalCtx` already reserve scratch slots for `ClassCall`). Bump-allocate the env struct (size `8 * free_vars.len()`), store each free variable's *current* value (loaded via the existing `Expr::Var` local-lookup path in the *enclosing* function) into it; bump-allocate the 2-word closure struct, store `table_idx` (a compile-time `i32.const`) at offset 0 and the env struct's pointer at offset 8; leave the closure struct's pointer as the result.

6. **Compiling a closure call.** In `compile_expr`'s existing `ast::Expr::FnCall(call) => { ... }` arm, before the existing `ctx.func_ids.get(&call.name)` lookup: check whether `call.name` is a *local* (`ctx.locals.contains_key(&call.name)`) whose inferred type (`infer_local_type` on `ast::Expr::Var(call.name.clone())`) is `PlumType::TFun(..)`. If so, compile as a closure call instead of a direct `Call`: `LocalGet` the closure pointer, `I32Load` its `env_pointer` (offset 8) into a scratch, push that env pointer, then compile+push each real argument, then `LocalGet` the closure pointer again and `I32Load` its `table_index` (offset 0), then `Instruction::CallIndirect { type_index, table_index: 0 }` — `type_index` resolved from the call site's own already-known concrete arg/return types (build/reuse a `module.add_type` call keyed by that signature, or thread the specific `ClosureInfo`/type index through if the call site's closure binding can be traced back to a specific `ClosureInfo` — use your judgment on the cleanest way to get a consistent type index here, since multiple closures of the same signature can share one).

- [ ] **Step 1: Write failing tests**

Append to `plum-wasm-codegen/tests/codegen_tests.rs`:

```rust
#[test]
fn non_capturing_closure_passed_and_called_runs_correctly() {
    let src = "\
each(cb: fn(Int) -> Int) -> Int =
  cb(5)

main() -> Int =
  each(|v|
    v)
";
    let source = parse(src);
    let bytes = compile_source(&source).expect("compile failed");
    assert_eq!(run_main(&bytes), 5);
}

#[test]
fn capturing_closure_snapshots_value_at_creation_time_runs_correctly() {
    let src = "\
each(cb: fn(Int) -> Int) -> Int =
  cb(0)

useClosure() -> Int =
  x = 10
  cb = |v|
    x + v
  x = 999
  each(cb)

main() -> Int =
  useClosure()
";
    let source = parse(src);
    let bytes = compile_source(&source).expect("compile failed");
    // The closure must see x==10 (its value when the closure was created), not 999
    // (its value when `each(cb)` is actually called) — proving snapshot-by-value
    // capture, not a live/shared reference.
    assert_eq!(run_main(&bytes), 10);
}

#[test]
fn closure_passed_through_already_generic_higher_order_function_runs_correctly() {
    let src = "\
identity(value: a) -> a =
  value

each(cb: fn(Int) -> Int) -> Int =
  cb(identity(7))

main() -> Int =
  each(|v|
    v * 2)
";
    let source = parse(src);
    let bytes = compile_source(&source).expect("compile failed");
    assert_eq!(run_main(&bytes), 14);
}
```

- [ ] **Step 2: Run to see them fail**

Run: `cargo test -p plum-wasm-codegen --test codegen_tests capturing_closure non_capturing_closure closure_passed_through`
Expected: all fail — none of the compiling logic exists yet.

- [ ] **Step 3: Implement the design above**

Follow the 6-part design. Iterate test-by-test — get `non_capturing_closure_passed_and_called_runs_correctly` passing first (no free-variable analysis needed for that one, simplifying the first pass), then tackle capture, then the generics-interop test.

- [ ] **Step 4: Run all three new tests, then the full codegen suite**

```bash
cargo test -p plum-wasm-codegen --test codegen_tests
```

Expected: all pass, including every pre-existing test (confirming the new discovery pre-pass and `FnCall` arm change don't regress ordinary function calls).

- [ ] **Step 5: Run the full workspace and tree-sitter suites**

```bash
cargo test --workspace
cd tooling/tree-sitter-plum && npx --yes tree-sitter-cli test
```

Expected: fully green.

- [ ] **Step 6: Commit**

```bash
git add plum-wasm-codegen/src/lib.rs plum-wasm-codegen/tests/codegen_tests.rs
git commit -m "feat(plum-wasm-codegen): compile closure literals and closure calls via a function table"
```

---

### Task 6: Examples and docs

**Files:**
- Modify or create: `examples/closures.plum` (new file, following this repo's existing `examples/*.plum` convention — one file per feature area, exercised by both crates' `examples_test.rs`)
- Modify: `plum-checker/tests/examples_test.rs`, `plum-wasm-codegen/tests/examples_test.rs`
- Modify: `README.md`

**Interfaces:**
- Consumes: everything from Tasks 1-5.
- Produces: nothing further downstream — final integration/documentation task.

- [ ] **Step 1: Add `examples/closures.plum`**

```plum
each(cb: fn(Int) -> Int) -> Int =
  cb(5)

double(v: Int) -> Int =
  v * 2

useNamedFunctionAsValue() -> Int =
  each(double)

useCapturingClosure() -> Int =
  offset = 100
  each(|v|
    v + offset)

main() -> Int =
  each(|v|
    v * 3)
```

(If `each(double)` — passing a top-level *named function* wherever a `fn(...)`-typed param is expected, not just a closure literal — isn't yet supported by Task 5's design (it may need a small addition: wrapping a plain function reference in the same `{table_index, env_pointer}` shape with an empty/null env), either extend Task 5 minimally to support it, or drop `useNamedFunctionAsValue`/adjust this example to closures only and note the named-function-as-value gap in README's Known Gaps instead. Verify which via a quick test before deciding.)

- [ ] **Step 2: Extend both crates' `examples_test.rs`**

In `plum-checker/tests/examples_test.rs`, confirm `examples/closures.plum` is picked up automatically (it likely already iterates every `.plum` file in the directory — check `example_files()`'s implementation; if so, no change needed beyond adding the file).

In `plum-wasm-codegen/tests/examples_test.rs`, add:

```rust
#[test]
fn closures_example_compiles_and_runs_correctly() {
    let bytes = assert_compiles("closures.plum");
    let engine = wasmtime::Engine::default();
    let module = wasmtime::Module::new(&engine, &bytes).expect("module should be loadable");
    let mut store = wasmtime::Store::new(&engine, ());
    let instance = wasmtime::Instance::new(&mut store, &module, &[]).expect("module should instantiate");
    let main = instance.get_typed_func::<(), i64>(&mut store, "main").expect("main should have signature () -> i64");
    assert_eq!(main.call(&mut store, ()).expect("main should not trap"), 15);
}
```

(Adjust the expected result if Step 1's `main` body ends up different from what's shown above.)

- [ ] **Step 3: Run both examples suites**

```bash
cargo test -p plum-checker --test examples_test
cargo test -p plum-wasm-codegen --test examples_test
```

Expected: both green.

- [ ] **Step 4: Update README**

Add a new "Closures" section (find a sensible place — likely after the existing "Generics" section, before "`self`, field access, and methods"), documenting: the `|params| body` syntax, the `fn(...)` / `fn(...) -> T` type annotation syntax (positional types only), snapshot-by-value capture semantics, and a link to `examples/closures.plum`. Update the "Known gaps" list: remove the `closure` bullet (`` `closure` (`|params| body`) exists in `grammar.js` but isn't wired into any reachable rule yet, so it doesn't actually parse in context. ``) if it's still present verbatim — check the file's current end for this exact sentence — and add any residual gap actually discovered during Task 5 (e.g. named-function-as-closure-value, if that turned out to be unsupported per Task 6 Step 1's note).

- [ ] **Step 5: Run the full workspace and tree-sitter suites one final time**

```bash
cargo test --workspace
cd tooling/tree-sitter-plum && npx --yes tree-sitter-cli test
```

Expected: fully green, zero known failures.

- [ ] **Step 6: Commit**

```bash
git add examples/closures.plum plum-checker/tests/examples_test.rs plum-wasm-codegen/tests/examples_test.rs README.md
git commit -m "docs+test: closures complete; add example and update README"
```