feat(linear-arrows): enforce -[q]-> quantity annotations on lambda params

Three coordinated fixes:

**typecheck.ml — lambda synth:**
Replace the hardcoded QOmega in ExprLambda synthesis with the actual
p_quantity from each parameter declaration. A lambda `|@linear x: T| e`
now synthesises type `T -[1]-> U` instead of `T -[ω]-> U`, so the
arrow type faithfully records the linearity contract.

**typecheck.ml — lambda check mode:**
When checking a lambda against an expected `TArrow (_, q, _, _)`, verify
that any explicit param quantity annotation matches the expected arrow
quantity. Unannotated params silently inherit the expected quantity
(correct for context-driven inference).

**quantity.ml — lambda body param checking:**
Add an `errors` accumulator to `env` (since `infer_usage_expr` is unit).
In `ExprLambda`, declare annotated params via `env_declare` so `env_use`
tracks their usage counts. After walking the body, check each annotated
param with `check_quantity` and push violations to `env.errors`. Drain
accumulated errors at the end of `check_function_quantities` (step 4).

Also fixes: save/restore `env.quantities` entries that lambda params
shadow, so nested scopes cannot leak stale quantity entries.

**codegen.ml / wasm.ml / wasm_encode.ml:**
Add `custom_sections` field to `wasm_module` and emit it in the encoder
(Wasm section ID 0). Wire `ownership_kind` annotations collected during
codegen into the `affinescript.ownership` custom section for typed-wasm
Level 7/10 verification. Fix `codegen.ml` record literal to include the
new field.

**Tests:**
- `test/e2e/fixtures/linear_arrow.affine` — valid single-use @linear lambda
- `test/e2e/fixtures/linear_arrow_violation.affine` — double-use @linear lambda
- `E2E Linear Arrows` suite: 2 cases (valid accepted, violation rejected)
- 75 tests total, all passing.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

Author: hyperpolymath

lib/codegen.ml

@@ -10,6 +10,14 @@
 open Ast
 open Wasm
 
+(** Ownership kind for typed-wasm schema annotations.
+    Maps AffineScript ownership qualifiers to typed-wasm Level 7/10 verification. *)
+type ownership_kind =
+  | Unrestricted  (** Plain value, no ownership constraint (Wasm i32/f64 etc.) *)
+  | Linear        (** TyOwn / own — consumed exactly once (typed-wasm Level 10 linearity) *)
+  | SharedBorrow  (** TyRef / ref — read-only aliasing safety (typed-wasm Level 7) *)
+  | ExclBorrow    (** TyMut / mut — exclusive mutable aliasing safety (typed-wasm Level 7) *)
+
 (** Code generation context *)
 type context = {
   types : func_type list;            (** type section *)
@@ -29,6 +37,10 @@ type context = {
   string_data : (string * int) list; (** string content -> memory offset *)
   next_string_offset : int;          (** next available offset for string data *)
   datas : data list;                 (** data segments *)
+  ownership_annots : (int * ownership_kind list * ownership_kind) list;
+  (** Collected ownership annotations: (func_index, param_kinds, return_kind).
+      Emitted as the [affinescript.ownership] Wasm custom section for typed-wasm
+      Level 7/10 verification. Kind encoding: 0=Unrestricted, 1=Linear, 2=SharedBorrow, 3=ExclBorrow. *)
 }
 
 (** Code generation error *)
@@ -70,8 +82,63 @@ let create_context () : context = {
   string_data = [];
   next_string_offset = 2048;  (* Start strings after heap at offset 2048 *)
   datas = [];
+  ownership_annots = [];
 }
 
+(** Extract ownership kind from a parameter declaration.
+    Checks p_ownership first; falls back to the shape of p_ty. *)
+let ownership_kind_of_param (p : param) : ownership_kind =
+  match p.p_ownership with
+  | Some Own -> Linear
+  | Some Ref -> SharedBorrow
+  | Some Mut -> ExclBorrow
+  | None ->
+    match p.p_ty with
+    | TyOwn _ -> Linear
+    | TyRef _ -> SharedBorrow
+    | TyMut _ -> ExclBorrow
+    | _ -> Unrestricted
+
+(** Extract ownership kind from an optional return type expression *)
+let ownership_kind_of_ret (ret : type_expr option) : ownership_kind =
+  match ret with
+  | Some (TyOwn _) -> Linear
+  | Some (TyRef _) -> SharedBorrow
+  | Some (TyMut _) -> ExclBorrow
+  | _ -> Unrestricted
+
+(** Encode an ownership_kind as a single byte (0–3) *)
+let ownership_kind_byte = function
+  | Unrestricted -> 0 | Linear -> 1 | SharedBorrow -> 2 | ExclBorrow -> 3
+
+(** Build the payload for the [affinescript.ownership] Wasm custom section.
+    Encoding (all little-endian):
+      u32  entry_count
+      per entry:
+        u32  func_index
+        u8   param_count
+        u8*  param_kind  (one per param, see kind encoding above)
+        u8   return_kind *)
+let build_ownership_section (annots : (int * ownership_kind list * ownership_kind) list) : bytes =
+  if annots = [] then Bytes.empty
+  else
+    let buf = Buffer.create 64 in
+    let write_u32_le n =
+      Buffer.add_char buf (Char.chr  (n         land 0xff));
+      Buffer.add_char buf (Char.chr ((n lsr  8) land 0xff));
+      Buffer.add_char buf (Char.chr ((n lsr 16) land 0xff));
+      Buffer.add_char buf (Char.chr ((n lsr 24) land 0xff))
+    in
+    let write_u8 n = Buffer.add_char buf (Char.chr (n land 0xff)) in
+    write_u32_le (List.length annots);
+    List.iter (fun (func_idx, param_kinds, ret_kind) ->
+      write_u32_le func_idx;
+      write_u8 (List.length param_kinds);
+      List.iter (fun k -> write_u8 (ownership_kind_byte k)) param_kinds;
+      write_u8 (ownership_kind_byte ret_kind)
+    ) annots;
+    Buffer.to_bytes buf
+
 (** Map AffineScript type to WASM value type *)
 let type_to_wasm (ty : type_expr) : value_type result =
   match ty with
@@ -1617,9 +1684,14 @@ let gen_decl (ctx : context) (decl : top_level) : context result =
     (* Determine function index before generating *)
     let func_idx = import_func_count ctx_with_type + List.length ctx_with_type.funcs in
 
-    (* Register function name to index mapping *)
+    (* Stage 2: Extract ownership annotations for typed-wasm [affinescript.ownership] section *)
+    let param_kinds = List.map ownership_kind_of_param fd.fd_params in
+    let ret_kind = ownership_kind_of_ret fd.fd_ret_ty in
+
+    (* Register function name to index mapping and record ownership annotations *)
     let ctx_with_func_idx = { ctx_with_type with
-      func_indices = ctx_with_type.func_indices @ [(fd.fd_name.name, func_idx)]
+      func_indices = ctx_with_type.func_indices @ [(fd.fd_name.name, func_idx)];
+      ownership_annots = ctx_with_type.ownership_annots @ [(func_idx, param_kinds, ret_kind)];
     } in
 
     (* Generate function with correct type index *)
@@ -1722,6 +1794,14 @@ let generate_module (prog : program) : wasm_module result =
   (* Add memory export *)
   let exports_with_mem = { e_name = "memory"; e_desc = ExportMemory 0 } :: ctx'.exports in
 
+  (* Stage 2: Build [affinescript.ownership] custom section from collected annotations *)
+  let ownership_payload = build_ownership_section ctx'.ownership_annots in
+  let custom_sections = if Bytes.length ownership_payload > 0 then
+    [("affinescript.ownership", ownership_payload)]
+  else
+    []
+  in
+
   Ok {
     types = ctx'.types;
     funcs = all_funcs;
@@ -1733,4 +1813,5 @@ let generate_module (prog : program) : wasm_module result =
     elems = elems;
     datas = List.rev ctx'.datas;  (* Reverse to get original order *)
     start = None;
+    custom_sections;
   }

lib/quantity.ml

@@ -198,13 +198,18 @@ type env = {
     (** Declared quantity for each tracked variable. *)
   usages : (string, usage) Hashtbl.t;
     (** Current observed usage count for each tracked variable. *)
+  errors : (quantity_error * Span.t) list ref;
+    (** Quantity errors accumulated by nested constructs (lambda params).
+        [infer_usage_expr] returns [unit], so nested violations are
+        pushed here and drained at the end of [check_function_quantities]. *)
 }
 
 (** Create a fresh quantity environment. *)
 let create_env () : env =
   {
     quantities = Hashtbl.create 16;
     usages = Hashtbl.create 16;
+    errors = ref [];
   }
 
 (** Declare a variable with its quantity annotation in the environment.
@@ -299,23 +304,55 @@ let rec infer_usage_expr (env : env) (expr : expr) : unit =
 
        Consequence: a lambda that captures an @linear variable raises
        LinearVariableUsedMultiple, because scale_usage QOmega UOne = UMany.
-       This closes BUG-004. *)
+       This closes BUG-004.
+
+       Additionally: lambda params with explicit quantity annotations are
+       declared in the env so env_use tracks them, and their usage is
+       verified against their declared quantity after the body walk.
+       Violations are pushed to env.errors (since infer_usage_expr is unit)
+       and drained by check_function_quantities at the top level. *)
     let param_names =
       List.map (fun (p : param) -> p.p_name.name) lam.elam_params
     in
     let before_snapshot = env_snapshot env in
-    (* Shadow lambda params so their usage inside the body doesn't leak
-       into the outer quantity environment. *)
-    List.iter (fun name ->
-      if Hashtbl.mem env.usages name then
-        Hashtbl.replace env.usages name UZero
-      (* Params not yet in env are fine — they are lambda-local. *)
-    ) param_names;
+    (* Save any quantity entries that the param names may shadow, so we
+       can restore them after the lambda body is analysed. *)
+    let old_quantities = List.map (fun name ->
+      (name, Hashtbl.find_opt env.quantities name)
+    ) param_names in
+    (* Declare each lambda param in the env so env_use can track it.
+       Annotated params get their declared quantity; unannotated params
+       get QOmega (unrestricted — no constraint to verify). *)
+    List.iter (fun (p : param) ->
+      let q = Option.value p.p_quantity ~default:QOmega in
+      env_declare env p.p_name.name q
+    ) lam.elam_params;
     infer_usage_expr env lam.elam_body;
     let after_snapshot = env_snapshot env in
+    (* Check annotated params against their declared quantities.
+       Only params with an explicit annotation are verified — unannotated
+       params default to QOmega which accepts any usage count. *)
+    List.iter (fun (p : param) ->
+      match p.p_quantity with
+      | None -> ()
+      | Some declared_q ->
+        let actual =
+          Hashtbl.find_opt after_snapshot p.p_name.name
+          |> Option.value ~default:UZero
+        in
+        (match check_quantity p.p_name declared_q actual with
+         | Ok () -> ()
+         | Error e -> env.errors := e :: !(env.errors))
+    ) lam.elam_params;
     (* Restore outer env, then re-apply QOmega-scaled deltas for captured
        outer variables only (exclude lambda params). *)
     env_restore env before_snapshot;
+    (* Restore quantity entries that lambda params may have shadowed. *)
+    List.iter (fun (name, old_q) ->
+      match old_q with
+      | Some q -> Hashtbl.replace env.quantities name q
+      | None   -> Hashtbl.remove  env.quantities name
+    ) old_quantities;
     Hashtbl.iter (fun name after_u ->
       if List.mem name param_names then ()
       else begin
@@ -562,18 +599,29 @@ let check_function_quantities (fd : fn_decl) : unit result =
   | FnBlock blk -> infer_usage_block env blk
   | FnExpr e -> infer_usage_expr env e
   end;
-  (* Step 3: check each parameter *)
-  List.fold_left (fun acc (param : param) ->
-    match acc with
-    | Error _ -> acc  (* Stop at first error *)
-    | Ok () ->
-      let declared = Option.value param.p_quantity ~default:QOmega in
-      let actual =
-        Hashtbl.find_opt env.usages param.p_name.name
-        |> Option.value ~default:UZero
-      in
-      check_quantity param.p_name declared actual
-  ) (Ok ()) fd.fd_params
+  (* Step 3: check each top-level parameter *)
+  let param_result =
+    List.fold_left (fun acc (param : param) ->
+      match acc with
+      | Error _ -> acc  (* Stop at first error *)
+      | Ok () ->
+        let declared = Option.value param.p_quantity ~default:QOmega in
+        let actual =
+          Hashtbl.find_opt env.usages param.p_name.name
+          |> Option.value ~default:UZero
+        in
+        check_quantity param.p_name declared actual
+    ) (Ok ()) fd.fd_params
+  in
+  (* Step 4: drain errors accumulated by nested lambda param checks.
+     infer_usage_expr is unit, so nested violations are pushed to
+     env.errors during the body walk above. Report the first one. *)
+  match param_result with
+  | Error _ -> param_result
+  | Ok () ->
+    match !(env.errors) with
+    | [] -> Ok ()
+    | (err, span) :: _ -> Error (err, span)
 
 (** Check quantities for all functions in a program.
 

lib/typecheck.ml

@@ -686,11 +686,21 @@ let rec synth (ctx : context) (expr : expr) : ty result =
       | Some sc -> Hashtbl.replace ctx.name_types n sc
       | None -> Hashtbl.remove ctx.name_types n
     ) old;
-    (* Build the arrow type: curried for multi-param *)
+    (* Build the arrow type: curried for multi-param.
+       Each arrow carries the declared quantity of the corresponding parameter.
+       If the parameter has no explicit annotation, we default to QOmega
+       (unrestricted), matching the convention used by check_fn_decl. *)
     let eff = fresh_effvar ctx.level in
-    let ty = List.fold_right (fun param_ty acc ->
-      TArrow (param_ty, QOmega, acc, eff)
-    ) param_tys body_ty in
+    let param_qty_pairs = List.map2 (fun (p : param) param_ty ->
+      let q = match p.p_quantity with
+        | Some q -> lower_quantity q
+        | None   -> QOmega
+      in
+      (q, param_ty)
+    ) elam_params param_tys in
+    let ty = List.fold_right (fun (q, param_ty) acc ->
+      TArrow (param_ty, q, acc, eff)
+    ) param_qty_pairs body_ty in
     Ok ty
 
   (* Function application *)
@@ -954,13 +964,30 @@ and check_stmt (ctx : context) (stmt : stmt) : unit result =
 (** Check that an expression has the expected type. *)
 and check (ctx : context) (expr : expr) (expected : ty) : unit result =
   match expr with
-  (* Lambda against arrow type: check mode is more precise *)
+  (* Lambda against arrow type: check mode is more precise.
+     We peel the expected arrow type one param at a time.  For each param:
+     - If the param has an explicit quantity annotation, we verify it is
+       consistent with the arrow quantity from the expected type.
+     - The arrow quantity from the expected type is used as the definitive
+       quantity for binding (so an unannotated lambda param correctly inherits
+       the quantity from its context, e.g. a @linear annotation on the let). *)
   | ExprLambda { elam_params; elam_body; elam_ret_ty = _ }
     when (match repr expected with TArrow _ -> true | _ -> false) ->
     let rec peel_arrows ty params =
       match params, repr ty with
       | [], _ -> Ok ()
-      | p :: rest, TArrow (param_ty, _q, ret_ty, _eff) ->
+      | p :: rest, TArrow (param_ty, q, ret_ty, _eff) ->
+        (* Validate explicit quantity annotation against the expected arrow. *)
+        let* () = match p.p_quantity with
+          | Some pq ->
+            let pq' = lower_quantity pq in
+            if pq' = q then Ok ()
+            else Error (TypeMismatch {
+              expected = TArrow (param_ty, q,   ret_ty, EPure);
+              got      = TArrow (param_ty, pq', ret_ty, EPure);
+            })
+          | None -> Ok ()
+        in
         bind_var ctx p.p_name.name param_ty;
         peel_arrows ret_ty rest
       | _ -> synth_and_unify ctx expr expected

lib/wasm.ml

@@ -242,6 +242,11 @@ type wasm_module = {
   elems : elem list;   (** element segments for table initialization *)
   datas : data list;   (** data segments for memory initialization *)
   start : int option;  (** optional start function index *)
+  custom_sections : (string * bytes) list;
+  (** Named custom sections (Wasm section ID 0).
+      Used for [affinescript.ownership] — carries ownership annotations
+      (TyOwn/TyRef/TyMut) that survive to the binary for typed-wasm
+      Level 7/10 verification. *)
 }
 [@@deriving show, eq]
 
@@ -257,4 +262,5 @@ let empty_module () : wasm_module = {
   elems = [];
   datas = [];
   start = None;
+  custom_sections = [];
 }

lib/wasm_encode.ml

@@ -360,6 +360,14 @@ let write_module_to_file path (m : wasm_module) : unit =
   add_section buf 10 (fun b -> add_vec b m.funcs add_code);
   add_section buf 11 (fun b -> add_vec b m.datas add_data);
 
+  (* Emit custom sections (Wasm section ID 0) — includes [affinescript.ownership] typed-wasm schema *)
+  List.iter (fun (name, payload) ->
+    add_section buf 0 (fun b ->
+      add_string b name;
+      add_bytes b payload
+    )
+  ) m.custom_sections;
+
   let oc = open_out_bin path in
   output_bytes oc (Bytes.of_string (Buffer.contents buf));
   close_out oc

test/e2e/fixtures/linear_arrow.affine

@@ -0,0 +1,15 @@
+// SPDX-License-Identifier: PMPL-1.0-or-later
+// End-to-end test: linear arrow quantity enforcement in lambda expressions
+// Tests:
+//   - Lambda with @linear param produces T -[1]-> U type (not T -[ω]-> U)
+//   - Lambda with @linear param used once passes quantity check
+//   - Named function with @linear param passes quantity check
+
+fn consume_once(@linear x: Int) -> Int = x + 1;
+
+fn wrap_linear() -> Int {
+  let f = |@linear x: Int| x + 1;
+  f(41)
+}
+
+fn hof_linear(g: (Int) -> Int, n: Int) -> Int = g(n);

test/e2e/fixtures/linear_arrow_violation.affine

@@ -0,0 +1,10 @@
+// SPDX-License-Identifier: PMPL-1.0-or-later
+// Regression: @linear param in a lambda must not be used twice inside the body.
+// The quantity checker must reject this program because `x` is used
+// two times inside the lambda body, violating the @linear (must use once)
+// contract declared on the lambda parameter.
+
+fn trigger_double_use() -> Int {
+  let f = |@linear x: Int| x + x;
+  f(21)
+}