4ba0db3 | plum

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

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

4ba0db3 — Peter John 2026-07-20T22:22:15+05:30

feat(plum-wasm-codegen): compile closure literals and closure calls via a function table

Files changed (8) hide show

plum-core/src/parser.rs +14 -6
plum-wasm-codegen/src/lib.rs +789 -14
plum-wasm-codegen/tests/codegen_tests.rs +61 -0
tooling/tree-sitter-plum/grammar.js +4 -1
tooling/tree-sitter-plum/src/grammar.json +0 -0
tooling/tree-sitter-plum/src/node-types.json +0 -0
tooling/tree-sitter-plum/src/parser.c +0 -0
tooling/tree-sitter-plum/test/corpus/function.txt +27 -0

plum-core/src/parser.rs CHANGED Viewed

@@ -519,16 +519,24 @@ impl<'a> AstParser<'a> {
     }
     fn parse_closure(&self, node: Node) -> Closure {
-        // closure: "|" var_identifier,* "|" body
+        // closure: "|" var_identifier,* "|" (expression | body)
         let params: Vec<String> = self.children_of_kind(node, "var_identifier")
             .into_iter()
             .map(|n| self.text(n))
             .collect();
+        // The body is either an indented `body` block or a single inline expression
+        // (`|v| v`); normalize the inline form into a one-statement block so codegen and
+        // the checker only ever see a `Block`.
-        let body = self.children_of_kind(node, "body")
+        let body = match self.children_of_kind(node, "body").into_iter().next() {
-            .into_iter()
-            .next()
-            .map(|n| self.parse_block(n))
+            Some(block_node) => self.parse_block(block_node),
+            None => match node.child_by_field_name("body") {
+                Some(expr_node) => {
+                    let unwrapped = self.unwrap_expr_node(expr_node);
+                    Block { stmts: vec![Stmt::Expr(self.parse_expression(unwrapped))] }
+                }
-            .unwrap_or(Block { stmts: vec![] });
+                None => Block { stmts: vec![] },
+            },
+        };
         Closure { params, body }
     }

plum-wasm-codegen/src/lib.rs CHANGED Viewed

@@ -214,8 +214,34 @@ pub struct FuncSig {
     pub ret: Option<ValType>,
 }
+/// Everything codegen needs to know about one closure *literal* found in the program.
+/// wasm has no native closures: each literal `|v| body` becomes its own real wasm
+/// function (registered in the funcref table), and a closure *value* is a single i32
+/// pointer to a heap pair `{table_index: i32 @0, env_pointer: i32 @8}`. The env is a
+/// separate heap allocation: one 8-byte slot per captured (free) variable, in
+/// `free_vars` order.
+pub struct ClosureInfo {
+    /// Reserved wasm function index for this closure's compiled body.
+    pub func_idx: u32,
+    /// Index of `func_idx` in the funcref table (the `i32` stored at struct offset 0).
+    pub table_idx: u32,
+    /// Closure param val types (NOT including the implicit leading env pointer).
+    pub param_vts: Vec<ValType>,
+    /// Closure param plum types (for the closure body's own type env).
+    pub param_ptypes: Vec<PlumType>,
+    /// Closure return val type (`None` for a `Unit`-returning closure).
+    pub ret_vt: Option<ValType>,
+    /// Free variables captured by value, in a stable (first-appearance) order; the
+    /// index into this vec IS the variable's env-struct slot (offset = idx * 8).
+    pub free_vars: Vec<(String, PlumType)>,
+}
+/// Key for deduplicating `call_indirect` function-type indices: the full wasm
+/// signature (leading env-ptr param included) of a closure.
+type ClosureSigKey = (Vec<ValType>, Option<ValType>);
 /// Global, read-only lookup tables shared by every function body being compiled.
-pub struct CompileCtx {
+pub struct CompileCtx<'a> {
     pub func_ids: HashMap<String, u32>,
     pub func_sigs: HashMap<String, FuncSig>,
     pub classes: ClassEnv,
@@ -223,6 +249,13 @@ pub struct CompileCtx {
     pub enum_variants: EnumVariants,
     pub global_env: TypeEnv,
     pub bump_global: u32,
+    /// Closure literal (keyed by `&Expr::Closure` pointer identity) -> its `ClosureInfo`.
+    pub closures: HashMap<usize, ClosureInfo>,
+    /// The AST of each discovered closure literal, keyed the same way, so its body can
+    /// be compiled in a second pass after all closures are registered.
+    pub closure_asts: HashMap<usize, &'a ast::Closure>,
+    /// Closure wasm signature -> function-type index, for `call_indirect` at call sites.
+    pub closure_call_types: HashMap<ClosureSigKey, u32>,
 }
 /// Per-module state that accumulates as function bodies are compiled: the running
@@ -242,8 +275,15 @@ struct LocalCtx<'a> {
     match_scratch_base: u32,
     /// `Match` stmt identity (pointer address) -> scratch slot offset.
     match_scratch_index: HashMap<usize, u32>,
+    /// First local index reserved for closure-construction scratch temporaries; each
+    /// closure literal uses two consecutive slots (env struct ptr, closure struct ptr).
+    closure_scratch_base: u32,
+    /// `Closure` expr identity (pointer address) -> offset of its first (of two) slots.
+    closure_scratch: HashMap<usize, u32>,
     func_ids: &'a HashMap<String, u32>,
     func_sigs: &'a HashMap<String, FuncSig>,
+    closures: &'a HashMap<usize, ClosureInfo>,
+    closure_call_types: &'a HashMap<ClosureSigKey, u32>,
     classes: &'a ClassEnv,
     methods: &'a MethodEnv,
     enum_variants: &'a EnumVariants,
@@ -286,9 +326,24 @@ fn plum_type_to_valtype(t: &PlumType) -> ValType {
     match t {
         PlumType::TInt => ValType::I64,
         PlumType::TFloat => ValType::F64,
+        // A closure value is a single i32 pointer to its heap-allocated
+        // `{table_index, env_pointer}` pair, so `TFun` is an i32 like every other
+        // heap reference (`TBool`/`TStr`/`TNamed`).
-        PlumType::TBool | PlumType::TStr | PlumType::TNamed(_) => ValType::I32,
+        PlumType::TBool | PlumType::TStr | PlumType::TNamed(_) | PlumType::TFun(_, _) => ValType::I32,
-        PlumType::TVar(_) | PlumType::TUnit | PlumType::TFun(_, _) => ValType::I64,
+        PlumType::TVar(_) | PlumType::TUnit => ValType::I64,
+    }
+}
+/// Maps an `ast::ParamType::Fn(params, ret)` to the wasm signature of the *closure
+/// function* it compiles to: an implicit leading `env_ptr: I32`, then one param per
+/// declared param type, returning `ret`'s val type (or nothing for `Unit`).
+fn fn_param_type_to_wasm_sig(params: &[ast::Type], ret: &Option<Box<ast::Type>>) -> (Vec<ValType>, Option<ValType>) {
+    let mut vts = vec![ValType::I32]; // env pointer
+    for p in params {
+        vts.push(ast_type_to_wasm(&p.name).unwrap_or(ValType::I32));
     }
+    let ret_vt = ret.as_ref().and_then(|t| ast_type_to_wasm(&t.name));
+    (vts, ret_vt)
 }
 fn block_type_for(result_vt: Option<ValType>) -> BlockType {
@@ -341,6 +396,9 @@ pub fn compile_source(source: &ast::Source) -> Result<Vec<u8>, String> {
     let mut func_ids: HashMap<String, u32> = HashMap::new();
     let mut func_sigs: HashMap<String, FuncSig> = HashMap::new();
+    // Closure wasm signature -> function-type index, deduped so every closure/call site
+    // of the same shape shares one `call_indirect` type.
+    let mut closure_call_types: HashMap<ClosureSigKey, u32> = HashMap::new();
     // Register every function AND method signature up front (methods get an implicit
     // leading `self: pointer` param and are keyed as "Receiver::method").
@@ -360,16 +418,89 @@ pub fn compile_source(source: &ast::Source) -> Result<Vec<u8>, String> {
             let key = fn_key(f);
             func_ids.insert(key.clone(), func_idx);
             func_sigs.insert(key, FuncSig { params: param_types, ret });
+            // Any `fn(...) -> ...`-typed param is callable via `call_indirect`; register
+            // its wasm signature (leading env-ptr param included) so call sites can
+            // resolve a consistent type index even if no matching closure literal exists.
+            for p in &f.params {
+                if let ast::ParamType::Fn(params, ret) = &p.ty {
+                    let (sig_params, ret_vt) = fn_param_type_to_wasm_sig(params, ret);
+                    let sig_key: ClosureSigKey = (sig_params.clone(), ret_vt);
+                    if !closure_call_types.contains_key(&sig_key) {
+                        let results: Vec<ValType> = ret_vt.into_iter().collect();
+                        let tidx = module.add_type(&sig_params, &results);
+                        closure_call_types.insert(sig_key, tidx);
+                    }
+                }
+            }
         }
     }
-    let ctx = CompileCtx { func_ids, func_sigs, classes, methods, enum_variants, global_env, bump_global };
     let fns: Vec<&ast::Fn> = source.items.iter().filter_map(|item| match item {
         ast::Item::Fn(f) => Some(f),
         _ => None,
     }).collect();
+    // ---- Discovery pre-pass: find every closure literal in every function body. ----
+    // (Runs on the monomorphized source, so any generic types in a closure's context
+    // are already concrete.) Registers each closure as its own wasm function + table
+    // element and records the free variables it must capture.
+    let fn_decls: HashMap<String, &ast::Fn> = fns.iter().map(|f| (f.name.clone(), *f)).collect();
+    let mut raw_closures: Vec<RawClosure> = Vec::new();
+    for f in &fns {
+        let mut env = global_env.clone();
+        if let Some(recv) = &f.type_param {
+            env.insert("self".to_string(), TypeScheme::mono(PlumType::TNamed(recv.clone())));
+        }
+        for p in &f.params {
+            env.insert(p.name.clone(), TypeScheme::mono(param_plum_type(&p.ty)));
+        }
+        let mut walker = ClosureWalker {
+            env,
+            cctx: check_ctx_of(&classes, &methods, &enum_variants),
+            fn_decls: &fn_decls,
+            found: Vec::new(),
+        };
+        match &f.body {
+            ast::FnBody::Block(block) => walker.walk_block(block),
+            ast::FnBody::Expr(e) => walker.walk_expr(e, None),
+        }
+        raw_closures.extend(walker.found);
+    }
+    let mut closures: HashMap<usize, ClosureInfo> = HashMap::new();
+    let mut closure_asts: HashMap<usize, &ast::Closure> = HashMap::new();
+    for rc in &raw_closures {
+        let mut sig_params = vec![ValType::I32]; // env pointer
+        sig_params.extend(rc.param_vts.iter().copied());
+        let sig_key: ClosureSigKey = (sig_params.clone(), rc.ret_vt);
+        let type_idx = match closure_call_types.get(&sig_key) {
+            Some(t) => *t,
+            None => {
+                let results: Vec<ValType> = rc.ret_vt.into_iter().collect();
+                let t = module.add_type(&sig_params, &results);
+                closure_call_types.insert(sig_key, t);
+                t
+            }
+        };
+        let func_idx = module.add_function(type_idx, &[]);
+        let table_idx = module.add_table_element(func_idx);
+        closures.insert(rc.ptr, ClosureInfo {
+            func_idx,
+            table_idx,
+            param_vts: rc.param_vts.clone(),
+            param_ptypes: rc.param_ptypes.clone(),
+            ret_vt: rc.ret_vt,
+            free_vars: rc.free_vars.clone(),
+        });
+        closure_asts.insert(rc.ptr, rc.closure);
+    }
+    let ctx = CompileCtx {
+        func_ids, func_sigs, classes, methods, enum_variants, global_env, bump_global,
+        closures, closure_asts, closure_call_types,
+    };
     let mut state = ModuleState { next_string_offset: STRING_BASE, data_segments: Vec::new() };
     let mut compiled_bodies: Vec<(String, Vec<u8>)> = Vec::new();
@@ -385,6 +516,18 @@ pub fn compile_source(source: &ast::Source) -> Result<Vec<u8>, String> {
         module.functions[slot].1 = body.clone();
     }
+    // Compile each closure literal's own body into its reserved function slot.
+    let mut closure_bodies: Vec<(u32, Vec<u8>)> = Vec::new();
+    for (ptr, info) in &ctx.closures {
+        let cl = ctx.closure_asts.get(ptr).expect("every registered closure has its AST recorded");
+        let body = compile_closure_body(cl, info, &ctx, &mut state)?;
+        closure_bodies.push((info.func_idx, body));
+    }
+    for (func_idx, body) in &closure_bodies {
+        let slot = (*func_idx - module.func_import_count) as usize;
+        module.functions[slot].1 = body.clone();
+    }
     for (offset, bytes) in &state.data_segments {
         module.add_data_segment(*offset, bytes);
     }
@@ -396,6 +539,352 @@ pub fn compile_source(source: &ast::Source) -> Result<Vec<u8>, String> {
     Ok(module.finish())
 }
+/// The `PlumType` of a declared parameter, including `fn(...)`-typed params as `TFun`.
+fn param_plum_type(pt: &ast::ParamType) -> PlumType {
+    match pt {
+        ast::ParamType::Type(t) => plum_checker::plum_type_from_ast(t),
+        ast::ParamType::Variadic(t) => plum_checker::plum_type_from_ast(t),
+        ast::ParamType::Fn(params, ret) => {
+            let param_types = params.iter().map(plum_checker::plum_type_from_ast).collect();
+            let ret_ty = ret.as_ref().map(|r| plum_checker::plum_type_from_ast(r)).unwrap_or(PlumType::TUnit);
+            PlumType::TFun(param_types, Box::new(ret_ty))
+        }
+    }
+}
+/// One closure literal discovered by the pre-pass, before it is registered as a wasm
+/// function. `ptr` is the closure's `&ast::Closure` pointer identity (its stable key).
+struct RawClosure<'a> {
+    ptr: usize,
+    closure: &'a ast::Closure,
+    param_vts: Vec<ValType>,
+    param_ptypes: Vec<PlumType>,
+    ret_vt: Option<ValType>,
+    free_vars: Vec<(String, PlumType)>,
+}
+/// Walks a function body (maintaining an evolving type env, exactly like `Collector`)
+/// to find every closure literal and determine its concrete signature and captured
+/// free variables. A closure passed directly as a `fn(...)`-typed call argument takes
+/// its signature from that declared param type; any other closure (e.g. one assigned to
+/// a local) falls back to the checker's inference of the closure expression itself.
+struct ClosureWalker<'a, 'c> {
+    env: TypeEnv,
+    cctx: plum_checker::CheckCtx<'c>,
+    fn_decls: &'a HashMap<String, &'a ast::Fn>,
+    found: Vec<RawClosure<'a>>,
+}
+impl<'a, 'c> ClosureWalker<'a, 'c> {
+    fn walk_block(&mut self, block: &'a ast::Block) {
+        for s in &block.stmts {
+            self.walk_stmt(s);
+        }
+    }
+    fn walk_stmt(&mut self, stmt: &'a ast::Stmt) {
+        match stmt {
+            ast::Stmt::Assign(a) => {
+                for (target, value) in a.targets.iter().zip(a.values.iter()) {
+                    self.walk_expr(value, None);
+                    let ty = plum_checker::infer_expr(value, &self.env, &self.cctx).unwrap_or(PlumType::TInt);
+                    self.env.insert(target.clone(), TypeScheme::mono(ty));
+                }
+            }
+            ast::Stmt::Return(Some(e)) => self.walk_expr(e, None),
+            ast::Stmt::Return(None) => {}
+            ast::Stmt::If(i) => {
+                self.walk_expr(&i.condition, None);
+                self.walk_block(&i.body);
+                for ei in &i.else_ifs {
+                    self.walk_expr(&ei.condition, None);
+                    self.walk_block(&ei.body);
+                }
+                if let Some(e) = &i.else_ {
+                    self.walk_block(e);
+                }
+            }
+            ast::Stmt::While(w) => {
+                self.walk_expr(&w.condition, None);
+                self.walk_block(&w.body);
+            }
+            ast::Stmt::For(f) => {
+                self.walk_expr(&f.iter, None);
+                for v in &f.vars {
+                    self.env.insert(v.clone(), TypeScheme::mono(PlumType::TInt));
+                }
+                self.walk_block(&f.body);
+            }
+            ast::Stmt::Expr(e) => self.walk_expr(e, None),
+            ast::Stmt::Assert(e) => self.walk_expr(e, None),
+            ast::Stmt::Match(m) => {
+                for s in &m.subjects {
+                    self.walk_expr(s, None);
+                }
+                for case in &m.cases {
+                    self.walk_block(&case.body);
+                }
+            }
+            ast::Stmt::Break | ast::Stmt::Continue | ast::Stmt::Todo => {}
+        }
+    }
+    /// `expected_fn` carries the declared `fn(params) -> ret` type when this expression is
+    /// a call argument in a `fn`-typed parameter position, giving a closure literal its
+    /// concrete signature.
+    fn walk_expr(&mut self, expr: &'a ast::Expr, expected_fn: Option<&'a ast::ParamType>) {
+        match expr {
+            ast::Expr::Closure(cl) => self.record_closure(cl, expected_fn),
+            ast::Expr::Binary(b) => { self.walk_expr(&b.left, None); self.walk_expr(&b.right, None); }
+            ast::Expr::Bool(b) => { self.walk_expr(&b.left, None); self.walk_expr(&b.right, None); }
+            ast::Expr::Compare(c) => { self.walk_expr(&c.left, None); self.walk_expr(&c.right, None); }
+            ast::Expr::Not(inner) => self.walk_expr(inner, None),
+            ast::Expr::Unary(u) => self.walk_expr(&u.operand, None),
+            ast::Expr::Paren(inner) => self.walk_expr(inner, None),
+            ast::Expr::Ternary(t) => {
+                self.walk_expr(&t.condition, None);
+                self.walk_expr(&t.then, None);
+                self.walk_expr(&t.else_, None);
+            }
+            ast::Expr::FnCall(call) => {
+                let callee = self.fn_decls.get(&call.name).copied();
+                for (i, arg) in call.args.iter().enumerate() {
+                    let arg_expr = arg_expr_of(arg);
+                    let expected = callee.and_then(|f| f.params.get(i)).map(|p| &p.ty)
+                        .filter(|pt| matches!(pt, ast::ParamType::Fn(_, _)));
+                    self.walk_expr(arg_expr, expected);
+                }
+            }
+            ast::Expr::ClassCall(call) => {
+                for fa in &call.fields {
+                    self.walk_expr(&fa.value, None);
+                }
+            }
+            ast::Expr::Attribute(a) => {
+                self.walk_expr(&a.object, None);
+                if let ast::AttrKind::Method(call) = &a.attr {
+                    for arg in &call.args {
+                        self.walk_expr(arg_expr_of(arg), None);
+                    }
+                }
+            }
+            ast::Expr::Int(_)
+            | ast::Expr::Float(_)
+            | ast::Expr::String(_)
+            | ast::Expr::Self_
+            | ast::Expr::Var(_)
+            | ast::Expr::TypeName(_) => {}
+        }
+    }
+    fn record_closure(&mut self, cl: &'a ast::Closure, expected_fn: Option<&'a ast::ParamType>) {
+        let (param_vts, param_ptypes, ret_vt) = match expected_fn {
+            Some(ast::ParamType::Fn(params, ret)) => {
+                let param_vts = params.iter().map(|t| ast_type_to_wasm(&t.name).unwrap_or(ValType::I32)).collect();
+                let param_ptypes = params.iter().map(plum_checker::plum_type_from_ast).collect();
+                let ret_vt = ret.as_ref().and_then(|t| ast_type_to_wasm(&t.name));
+                (param_vts, param_ptypes, ret_vt)
+            }
+            _ => {
+                // Not a direct `fn`-typed call argument: infer the closure's own type.
+                match plum_checker::infer_expr(&ast::Expr::Closure(Box::new(cl.clone())), &self.env, &self.cctx) {
+                    Ok(PlumType::TFun(ptypes, ret)) => {
+                        let param_vts = ptypes.iter().map(plum_type_to_valtype).collect();
+                        let ret_vt = match *ret {
+                            PlumType::TUnit => None,
+                            other => Some(plum_type_to_valtype(&other)),
+                        };
+                        (param_vts, ptypes, ret_vt)
+                    }
+                    _ => {
+                        let param_vts = cl.params.iter().map(|_| ValType::I64).collect();
+                        let param_ptypes = cl.params.iter().map(|_| PlumType::TInt).collect();
+                        (param_vts, param_ptypes, Some(ValType::I64))
+                    }
+                }
+            }
+        };
+        let free_vars = collect_free_vars(cl, &self.env);
+        self.found.push(RawClosure {
+            ptr: cl as *const ast::Closure as usize,
+            closure: cl,
+            param_vts,
+            param_ptypes,
+            ret_vt,
+            free_vars,
+        });
+    }
+}
+fn arg_expr_of(arg: &ast::Arg) -> &ast::Expr {
+    match arg {
+        ast::Arg::Positional(e) => e,
+        ast::Arg::Keyword { value, .. } => value,
+        ast::Arg::Pair { value, .. } => value,
+    }
+}
+/// Determines a closure's captured free variables: every `Var` referenced in its body
+/// that is neither one of the closure's own params nor assigned locally inside the body,
+/// in first-appearance order. Each free variable's type is looked up in the *enclosing*
+/// scope's type env.
+fn collect_free_vars(cl: &ast::Closure, env: &TypeEnv) -> Vec<(String, PlumType)> {
+    let mut bound: std::collections::HashSet<String> = cl.params.iter().cloned().collect();
+    fv_collect_bound_block(&cl.body, &mut bound);
+    let mut seen: std::collections::HashSet<String> = std::collections::HashSet::new();
+    let mut free: Vec<(String, PlumType)> = Vec::new();
+    fv_collect_refs_block(&cl.body, &bound, &mut seen, &mut free, env);
+    free
+}
+fn fv_collect_bound_block(block: &ast::Block, bound: &mut std::collections::HashSet<String>) {
+    for s in &block.stmts {
+        match s {
+            ast::Stmt::Assign(a) => {
+                for t in &a.targets {
+                    bound.insert(t.clone());
+                }
+            }
+            ast::Stmt::For(f) => {
+                for v in &f.vars {
+                    bound.insert(v.clone());
+                }
+                fv_collect_bound_block(&f.body, bound);
+            }
+            ast::Stmt::If(i) => {
+                fv_collect_bound_block(&i.body, bound);
+                for ei in &i.else_ifs {
+                    fv_collect_bound_block(&ei.body, bound);
+                }
+                if let Some(e) = &i.else_ {
+                    fv_collect_bound_block(e, bound);
+                }
+            }
+            ast::Stmt::While(w) => fv_collect_bound_block(&w.body, bound),
+            ast::Stmt::Match(m) => {
+                for case in &m.cases {
+                    for p in &case.patterns {
+                        fv_collect_pattern_bindings(p, bound);
+                    }
+                    fv_collect_bound_block(&case.body, bound);
+                }
+            }
+            _ => {}
+        }
+    }
+}
+fn fv_collect_pattern_bindings(pat: &ast::CasePattern, bound: &mut std::collections::HashSet<String>) {
+    match pat {
+        ast::CasePattern::Name(n) => { bound.insert(n.clone()); }
+        ast::CasePattern::Class { fields, .. } => {
+            for f in fields {
+                fv_collect_pattern_bindings(f, bound);
+            }
+        }
+        _ => {}
+    }
+}
+fn fv_collect_refs_block(
+    block: &ast::Block,
+    bound: &std::collections::HashSet<String>,
+    seen: &mut std::collections::HashSet<String>,
+    free: &mut Vec<(String, PlumType)>,
+    env: &TypeEnv,
+) {
+    for s in &block.stmts {
+        match s {
+            ast::Stmt::Assign(a) => {
+                for v in &a.values {
+                    fv_collect_refs_expr(v, bound, seen, free, env);
+                }
+            }
+            ast::Stmt::Return(Some(e)) | ast::Stmt::Expr(e) | ast::Stmt::Assert(e) => {
+                fv_collect_refs_expr(e, bound, seen, free, env);
+            }
+            ast::Stmt::If(i) => {
+                fv_collect_refs_expr(&i.condition, bound, seen, free, env);
+                fv_collect_refs_block(&i.body, bound, seen, free, env);
+                for ei in &i.else_ifs {
+                    fv_collect_refs_expr(&ei.condition, bound, seen, free, env);
+                    fv_collect_refs_block(&ei.body, bound, seen, free, env);
+                }
+                if let Some(e) = &i.else_ {
+                    fv_collect_refs_block(e, bound, seen, free, env);
+                }
+            }
+            ast::Stmt::While(w) => {
+                fv_collect_refs_expr(&w.condition, bound, seen, free, env);
+                fv_collect_refs_block(&w.body, bound, seen, free, env);
+            }
+            ast::Stmt::For(f) => {
+                fv_collect_refs_expr(&f.iter, bound, seen, free, env);
+                fv_collect_refs_block(&f.body, bound, seen, free, env);
+            }
+            ast::Stmt::Match(m) => {
+                for subj in &m.subjects {
+                    fv_collect_refs_expr(subj, bound, seen, free, env);
+                }
+                for case in &m.cases {
+                    fv_collect_refs_block(&case.body, bound, seen, free, env);
+                }
+            }
+            _ => {}
+        }
+    }
+}
+fn fv_collect_refs_expr(
+    expr: &ast::Expr,
+    bound: &std::collections::HashSet<String>,
+    seen: &mut std::collections::HashSet<String>,
+    free: &mut Vec<(String, PlumType)>,
+    env: &TypeEnv,
+) {
+    match expr {
+        ast::Expr::Var(name) => {
+            if !bound.contains(name) && seen.insert(name.clone()) {
+                let ty = plum_checker::lookup(env, name).unwrap_or(PlumType::TInt);
+                free.push((name.clone(), ty));
+            }
+        }
+        ast::Expr::Binary(b) => { fv_collect_refs_expr(&b.left, bound, seen, free, env); fv_collect_refs_expr(&b.right, bound, seen, free, env); }
+        ast::Expr::Bool(b) => { fv_collect_refs_expr(&b.left, bound, seen, free, env); fv_collect_refs_expr(&b.right, bound, seen, free, env); }
+        ast::Expr::Compare(c) => { fv_collect_refs_expr(&c.left, bound, seen, free, env); fv_collect_refs_expr(&c.right, bound, seen, free, env); }
+        ast::Expr::Not(inner) => fv_collect_refs_expr(inner, bound, seen, free, env),
+        ast::Expr::Unary(u) => fv_collect_refs_expr(&u.operand, bound, seen, free, env),
+        ast::Expr::Paren(inner) => fv_collect_refs_expr(inner, bound, seen, free, env),
+        ast::Expr::Ternary(t) => {
+            fv_collect_refs_expr(&t.condition, bound, seen, free, env);
+            fv_collect_refs_expr(&t.then, bound, seen, free, env);
+            fv_collect_refs_expr(&t.else_, bound, seen, free, env);
+        }
+        ast::Expr::FnCall(call) => {
+            for arg in &call.args {
+                fv_collect_refs_expr(arg_expr_of(arg), bound, seen, free, env);
+            }
+        }
+        ast::Expr::ClassCall(call) => {
+            for fa in &call.fields {
+                fv_collect_refs_expr(&fa.value, bound, seen, free, env);
+            }
+        }
+        ast::Expr::Attribute(a) => {
+            fv_collect_refs_expr(&a.object, bound, seen, free, env);
+            if let ast::AttrKind::Method(call) = &a.attr {
+                for arg in &call.args {
+                    fv_collect_refs_expr(arg_expr_of(arg), bound, seen, free, env);
+                }
+            }
+        }
+        // A nested closure isn't supported by the discovery pre-pass; its own free-var
+        // references are not hoisted here.
+        _ => {}
+    }
+}
 /// Walks a function body once to determine: (1) every locally-assigned/bound name and
 /// its inferred type, (2) how many `ClassCall` scratch temporaries it needs, and (3)
 /// the subject type for every `match` statement (for its own scratch temporary).
@@ -407,6 +896,9 @@ struct Collector<'a> {
     classcall_scratch: HashMap<usize, u32>,
     match_scratch: HashMap<usize, PlumType>,
     next_classcall_slot: u32,
+    /// `Closure` expr identity -> offset of its first (of two) construction scratch slots.
+    closure_scratch: HashMap<usize, u32>,
+    next_closure_slot: u32,
 }
 impl<'a> Collector<'a> {
@@ -558,9 +1050,15 @@ impl<'a> Collector<'a> {
             | ast::Expr::Self_
             | ast::Expr::Var(_)
             | ast::Expr::TypeName(_) => {}
+            // A closure literal needs two construction scratch slots in the *enclosing*
+            // function (env struct ptr, closure struct ptr). Its body's own locals are
-            // TODO: closures aren't compiled yet; codegen for closure bodies lands
+            // NOT this function's — they belong to the separate closure function — so we
-            // in a later task.
+            // don't recurse into the body here.
-            ast::Expr::Closure(_) => {}
+            ast::Expr::Closure(cl) => {
+                let base = self.next_closure_slot;
+                self.next_closure_slot += 2;
+                self.closure_scratch.insert(cl.as_ref() as *const ast::Closure as usize, base);
+            }
         }
     }
@@ -601,9 +1099,15 @@ fn compile_fn_body(f: &ast::Fn, ctx: &CompileCtx, state: &mut ModuleState) -> Re
         classcall_scratch: HashMap::new(),
         match_scratch: HashMap::new(),
         next_classcall_slot: 0,
+        closure_scratch: HashMap::new(),
+        next_closure_slot: 0,
     };
     if let ast::FnBody::Block(block) = &f.body {
         collector.walk_block(block);
+    } else if let ast::FnBody::Expr(e) = &f.body {
+        // An expression-bodied fn can still contain a closure literal (e.g.
+        // `main() -> Int = each(|v| v)`), which needs construction scratch slots.
+        collector.walk_expr(e);
     }
     // ---- assign local indices: [self?][params][named...][classcall scratch...][match scratch...] ----
@@ -641,6 +1145,12 @@ fn compile_fn_body(f: &ast::Fn, ctx: &CompileCtx, state: &mut ModuleState) -> Re
         idx += 1;
     }
+    let closure_scratch_base = idx;
+    let closure_scratch_count = collector.closure_scratch.values().copied().max().map(|m| m + 2).unwrap_or(0);
+    for _ in 0..closure_scratch_count {
+        groups.push(ValType::I32);
+    }
     if groups.is_empty() {
         body.push(0);
     } else {
@@ -657,8 +1167,12 @@ fn compile_fn_body(f: &ast::Fn, ctx: &CompileCtx, state: &mut ModuleState) -> Re
         classcall_scratch: collector.classcall_scratch,
         match_scratch_base,
         match_scratch_index,
+        closure_scratch_base,
+        closure_scratch: collector.closure_scratch,
         func_ids: &ctx.func_ids,
         func_sigs: &ctx.func_sigs,
+        closures: &ctx.closures,
+        closure_call_types: &ctx.closure_call_types,
         classes: &ctx.classes,
         methods: &ctx.methods,
         enum_variants: &ctx.enum_variants,
@@ -905,7 +1419,14 @@ fn compile_stmt(stmt: &ast::Stmt, body: &mut Vec<u8>, ctx: &LocalCtx, state: &mu
 /// Returns true if the expression leaves a value on the wasm stack.
 fn expr_has_result(expr: &ast::Expr, ctx: &LocalCtx) -> bool {
     match expr {
+        ast::Expr::FnCall(call) => {
+            if ctx.locals.contains_key(&call.name) {
+                if let PlumType::TFun(_, ret) = infer_local_type(&ast::Expr::Var(call.name.clone()), ctx) {
+                    return !matches!(*ret, PlumType::TUnit);
+                }
+            }
-        ast::Expr::FnCall(call) => ctx.func_sigs.get(&call.name).map(|s| s.ret.is_some()).unwrap_or(true),
+            ctx.func_sigs.get(&call.name).map(|s| s.ret.is_some()).unwrap_or(true)
+        }
         ast::Expr::Attribute(attr) => match &attr.attr {
             ast::AttrKind::Method(call) => {
                 if let PlumType::TNamed(class_name) = infer_local_type(&attr.object, ctx) {
@@ -1293,7 +1814,13 @@ fn compile_expr(expr: &ast::Expr, body: &mut Vec<u8>, ctx: &LocalCtx, state: &mu
             Instruction::End.encode(body);
         }
         ast::Expr::FnCall(call) => {
+            // A call whose callee name is a *local* of function type is a closure call,
+            // dispatched via `call_indirect` — not a direct `Call` to a named function.
+            let is_closure_call = ctx.locals.contains_key(&call.name)
+                && matches!(infer_local_type(&ast::Expr::Var(call.name.clone()), ctx), PlumType::TFun(_, _));
+            if is_closure_call {
+                compile_closure_call(call, body, ctx, state)?;
-            if let Some(info) = ctx.enum_variants.get(&call.name) {
+            } else if let Some(info) = ctx.enum_variants.get(&call.name) {
                 compile_variant_construction(info, call, expr, body, ctx, state)?;
             } else {
                 for arg in &call.args {
@@ -1434,10 +1961,8 @@ fn compile_expr(expr: &ast::Expr, body: &mut Vec<u8>, ctx: &LocalCtx, state: &mu
             state.data_segments.push((offset, data));
             Instruction::I32Const(offset as i32).encode(body);
         }
-        // TODO: closure codegen (capturing free variables, emitting an indirect-callable
-        // value) lands in a later task.
-        ast::Expr::Closure(_) => {
+        ast::Expr::Closure(cl) => {
-            return Err("codegen: closures are not yet supported".to_string());
+            compile_closure_literal(cl, body, ctx, state)?;
         }
     }
     Ok(())
@@ -1503,3 +2028,253 @@ fn compile_variant_construction(
     Instruction::LocalGet(scratch_local).encode(body);
     Ok(())
 }
+/// Emits a width-appropriate store for a heap slot: `[addr, value]` -> memory.
+fn emit_store(vt: ValType, offset: u64, body: &mut Vec<u8>) {
+    match vt {
+        ValType::I64 => Instruction::I64Store(MemArg { offset, align: 3, memory_index: 0 }),
+        ValType::F64 => Instruction::F64Store(MemArg { offset, align: 3, memory_index: 0 }),
+        _ => Instruction::I32Store(MemArg { offset, align: 2, memory_index: 0 }),
+    }.encode(body);
+}
+/// Emits a width-appropriate load for a heap slot: `[addr]` -> value.
+fn emit_load(vt: ValType, offset: u64, body: &mut Vec<u8>) {
+    match vt {
+        ValType::I64 => Instruction::I64Load(MemArg { offset, align: 3, memory_index: 0 }),
+        ValType::F64 => Instruction::F64Load(MemArg { offset, align: 3, memory_index: 0 }),
+        _ => Instruction::I32Load(MemArg { offset, align: 2, memory_index: 0 }),
+    }.encode(body);
+}
+/// Compiles a closure *literal* at its construction site. Snapshots each captured free
+/// variable's CURRENT value into a fresh heap env struct, then builds the 2-word closure
+/// struct `{table_index @0, env_pointer @8}` and leaves its pointer on the stack.
+fn compile_closure_literal(
+    cl: &ast::Closure,
+    body: &mut Vec<u8>,
+    ctx: &LocalCtx,
+    _state: &mut ModuleState,
+) -> Result<(), String> {
+    let key = cl as *const ast::Closure as usize;
+    let info = ctx
+        .closures
+        .get(&key)
+        .ok_or_else(|| "codegen: closure literal was not discovered (nested closures are not supported)".to_string())?;
+    let scratch_off = *ctx
+        .closure_scratch
+        .get(&key)
+        .ok_or_else(|| "internal codegen error: missing closure scratch slot".to_string())?;
+    let env_scratch = ctx.closure_scratch_base + scratch_off;
+    let closure_scratch = env_scratch + 1;
+    // Bump-allocate the env struct: one 8-byte slot per captured variable.
+    let env_size = (info.free_vars.len() as i32) * 8;
+    Instruction::GlobalGet(ctx.bump_global).encode(body);
+    Instruction::LocalSet(env_scratch).encode(body);
+    Instruction::GlobalGet(ctx.bump_global).encode(body);
+    Instruction::I32Const(env_size).encode(body);
+    Instruction::I32Add.encode(body);
+    Instruction::GlobalSet(ctx.bump_global).encode(body);
+    // Snapshot each free variable's current value from the enclosing function's local.
+    for (i, (name, ty)) in info.free_vars.iter().enumerate() {
+        let local_idx = *ctx
+            .locals
+            .get(name)
+            .ok_or_else(|| format!("codegen: captured variable '{}' is not a local in the enclosing scope", name))?;
+        Instruction::LocalGet(env_scratch).encode(body);
+        Instruction::LocalGet(local_idx).encode(body);
+        emit_store(plum_type_to_valtype(ty), (i as u64) * 8, body);
+    }
+    // Bump-allocate the 2-word closure struct.
+    Instruction::GlobalGet(ctx.bump_global).encode(body);
+    Instruction::LocalSet(closure_scratch).encode(body);
+    Instruction::GlobalGet(ctx.bump_global).encode(body);
+    Instruction::I32Const(16).encode(body);
+    Instruction::I32Add.encode(body);
+    Instruction::GlobalSet(ctx.bump_global).encode(body);
+    // table_index @ offset 0
+    Instruction::LocalGet(closure_scratch).encode(body);
+    Instruction::I32Const(info.table_idx as i32).encode(body);
+    Instruction::I32Store(MemArg { offset: 0, align: 2, memory_index: 0 }).encode(body);
+    // env_pointer @ offset 8
+    Instruction::LocalGet(closure_scratch).encode(body);
+    Instruction::LocalGet(env_scratch).encode(body);
+    Instruction::I32Store(MemArg { offset: 8, align: 2, memory_index: 0 }).encode(body);
+    // The closure value is its struct pointer.
+    Instruction::LocalGet(closure_scratch).encode(body);
+    Ok(())
+}
+/// Compiles a call to a closure-typed local via `call_indirect`. Stack order matches the
+/// closure function's signature `(env_ptr, ...args)`: push the env pointer, then each
+/// argument, then the table index (the `call_indirect` operand).
+fn compile_closure_call(
+    call: &ast::FnCall,
+    body: &mut Vec<u8>,
+    ctx: &LocalCtx,
+    state: &mut ModuleState,
+) -> Result<(), String> {
+    let closure_local = *ctx
+        .locals
+        .get(&call.name)
+        .ok_or_else(|| format!("codegen: closure '{}' is not a local", call.name))?;
+    let (param_ptypes, ret_ptype) = match infer_local_type(&ast::Expr::Var(call.name.clone()), ctx) {
+        PlumType::TFun(p, r) => (p, *r),
+        other => return Err(format!("codegen: '{}' is not callable (type {:?})", call.name, other)),
+    };
+    let mut sig_params = vec![ValType::I32]; // env pointer
+    for p in &param_ptypes {
+        sig_params.push(plum_type_to_valtype(p));
+    }
+    let ret_vt = match ret_ptype {
+        PlumType::TUnit => None,
+        other => Some(plum_type_to_valtype(&other)),
+    };
+    let sig_key: ClosureSigKey = (sig_params, ret_vt);
+    let type_index = *ctx
+        .closure_call_types
+        .get(&sig_key)
+        .ok_or_else(|| format!("internal codegen error: no call_indirect type for closure '{}'", call.name))?;
+    // env pointer (struct offset 8)
+    Instruction::LocalGet(closure_local).encode(body);
+    Instruction::I32Load(MemArg { offset: 8, align: 2, memory_index: 0 }).encode(body);
+    // real arguments
+    for arg in &call.args {
+        compile_expr(arg_expr_of(arg), body, ctx, state)?;
+    }
+    // table index (struct offset 0) — the call_indirect operand
+    Instruction::LocalGet(closure_local).encode(body);
+    Instruction::I32Load(MemArg { offset: 0, align: 2, memory_index: 0 }).encode(body);
+    Instruction::CallIndirect { type_index, table_index: 0 }.encode(body);
+    Ok(())
+}
+/// Compiles a closure literal's own body into a standalone wasm function. Local 0 is the
+/// implicit env pointer; the closure's params follow; then each captured free variable
+/// gets a local loaded from the env struct at function entry (restoring the snapshot).
+fn compile_closure_body(
+    cl: &ast::Closure,
+    info: &ClosureInfo,
+    ctx: &CompileCtx,
+    state: &mut ModuleState,
+) -> Result<Vec<u8>, String> {
+    let mut body = Vec::new();
+    // Base type env: globals + captured free vars + closure params.
+    let mut base_env = ctx.global_env.clone();
+    for (name, ty) in &info.free_vars {
+        base_env.insert(name.clone(), TypeScheme::mono(ty.clone()));
+    }
+    for (name, pty) in cl.params.iter().zip(info.param_ptypes.iter()) {
+        base_env.insert(name.clone(), TypeScheme::mono(pty.clone()));
+    }
+    let mut collector = Collector {
+        env: base_env.clone(),
+        cctx: check_ctx_of(&ctx.classes, &ctx.methods, &ctx.enum_variants),
+        named: Vec::new(),
+        named_set: Default::default(),
+        classcall_scratch: HashMap::new(),
+        match_scratch: HashMap::new(),
+        next_classcall_slot: 0,
+        closure_scratch: HashMap::new(),
+        next_closure_slot: 0,
+    };
+    collector.walk_block(&cl.body);
+    // ---- local index layout ----
+    // [env_ptr][closure params][free-var locals][named...][classcall][match][closure scratch]
+    let mut locals: HashMap<String, u32> = HashMap::new();
+    let mut groups: Vec<ValType> = Vec::new();
+    let mut idx = 0u32;
+    idx += 1; // local 0 = env pointer (a param, so not declared below)
+    for name in &cl.params {
+        locals.insert(name.clone(), idx);
+        idx += 1;
+    }
+    for (name, ty) in &info.free_vars {
+        locals.insert(name.clone(), idx);
+        groups.push(plum_type_to_valtype(ty));
+        idx += 1;
+    }
+    for (name, ty) in &collector.named {
+        if locals.contains_key(name) {
+            continue;
+        }
+        locals.insert(name.clone(), idx);
+        groups.push(plum_type_to_valtype(ty));
+        idx += 1;
+    }
+    let classcall_scratch_base = idx;
+    let classcall_count = collector.classcall_scratch.values().copied().max().map(|m| m + 1).unwrap_or(0);
+    for _ in 0..classcall_count {
+        groups.push(ValType::I32);
+        idx += 1;
+    }
+    let match_scratch_base = idx;
+    let mut match_scratch_index: HashMap<usize, u32> = HashMap::new();
+    for (ptr, ty) in collector.match_scratch.iter() {
+        match_scratch_index.insert(*ptr, idx - match_scratch_base);
+        groups.push(plum_type_to_valtype(ty));
+        idx += 1;
+    }
+    let closure_scratch_base = idx;
+    let closure_scratch_count = collector.closure_scratch.values().copied().max().map(|m| m + 2).unwrap_or(0);
+    for _ in 0..closure_scratch_count {
+        groups.push(ValType::I32);
+    }
+    if groups.is_empty() {
+        body.push(0);
+    } else {
+        body.extend(encode_leb128_u32(groups.len() as u32));
+        for g in &groups {
+            body.extend(encode_leb128_u32(1));
+            g.encode(&mut body);
+        }
+    }
+    // Restore each captured free variable from the env struct (local 0) at entry.
+    for (i, (name, ty)) in info.free_vars.iter().enumerate() {
+        let local_idx = *locals.get(name).expect("free var local was assigned above");
+        Instruction::LocalGet(0).encode(&mut body); // env pointer
+        emit_load(plum_type_to_valtype(ty), (i as u64) * 8, &mut body);
+        Instruction::LocalSet(local_idx).encode(&mut body);
+    }
+    let local_ctx = LocalCtx {
+        locals,
+        classcall_scratch_base,
+        classcall_scratch: collector.classcall_scratch,
+        match_scratch_base,
+        match_scratch_index,
+        closure_scratch_base,
+        closure_scratch: collector.closure_scratch,
+        func_ids: &ctx.func_ids,
+        func_sigs: &ctx.func_sigs,
+        closures: &ctx.closures,
+        closure_call_types: &ctx.closure_call_types,
+        classes: &ctx.classes,
+        methods: &ctx.methods,
+        enum_variants: &ctx.enum_variants,
+        type_env: RefCell::new(base_env),
+        bump_global: ctx.bump_global,
+    };
+    compile_block_as_fn_body(&cl.body, &mut body, &local_ctx, state, info.ret_vt)?;
+    Instruction::End.encode(&mut body);
+    Ok(body)
+}

plum-wasm-codegen/tests/codegen_tests.rs CHANGED Viewed

@@ -900,3 +900,64 @@ fn wasm_module_with_a_table_element_validates_and_call_indirect_works() {
     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);
 }
+#[test]
+fn non_capturing_closure_passed_and_called_runs_correctly() {
+    // The brief's plan scopes closures as "passable as an ordinary call argument
+    // (`each(|v| ...)`)"; the inline single-expression body is that supported form.
+    // (A multi-line indented closure body immediately followed by `)` on the same line
+    // needs the external scanner to emit a dedent at a closing bracket, which is a
+    // parser/scanner concern outside this codegen task — see task-5-report.md.)
+    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);
+}

tooling/tree-sitter-plum/grammar.js CHANGED Viewed

@@ -372,7 +372,10 @@ module.exports = grammar({
         "|",
         field("parameters", optional(commaSep1($.var_identifier))),
         "|",
+        // A closure body is either a single inline expression (`|v| v`, usable as an
+        // ordinary call argument like `each(|v| v)`) or an indented block — the same
+        // `choice($.expression, $.body)` shape already used by `fn` and `case` bodies.
-        field("body", $.body),
+        field("body", choice($.expression, $.body)),
       ),
     // The member name always lexes as `fn_identifier` (a superset of `var_identifier`,

tooling/tree-sitter-plum/src/grammar.json CHANGED Viewed

Binary file

tooling/tree-sitter-plum/src/node-types.json CHANGED Viewed

Binary file

tooling/tree-sitter-plum/src/parser.c CHANGED Viewed

Binary file

tooling/tree-sitter-plum/test/corpus/function.txt CHANGED Viewed

@@ -356,3 +356,30 @@ each(cb: fn(a) -> b) -> Bool =
       (expression
         (primary_expression
           (type_identifier))))))
+================================================================================
+function - inline closure literal as a call argument
+================================================================================
+main() -> Int =
+  each(|v| v)
+--------------------------------------------------------------------------------
+(source
+  (fn
+    name: (fn_identifier)
+    returns: (return_type
+      (type_identifier))
+    body: (body
+      (expression
+        (primary_expression
+          (fn_call
+            function: (var_identifier)
+            arguments: (fn_argument_list
+              (expression
+                (closure
+                  parameters: (var_identifier)
+                  body: (expression
+                    (primary_expression
+                      (var_identifier))))))))))))