tw_verify.ml

(* SPDX-License-Identifier: MPL-2.0 *)
(* SPDX-FileCopyrightText: 2024-2026 hyperpolymath *)

(** typed-wasm ownership verifier — Stage 7 (per-path analysis added Stage 9).

    Statically verifies typed-wasm Level 7 (aliasing safety) and Level 10
    (linearity) constraints on AffineScript's Wasm IR.

    The verifier runs after codegen on the in-memory [Wasm.wasm_module]
    together with the ownership annotations collected during codegen.
    It is a second line of defence: the QTT quantity checker (Stage 1) already
    enforces these rules at the AffineScript source level; this module
    re-checks them on the emitted Wasm IR to catch any codegen bugs.

    Level 10 — Linearity: a parameter annotated [Linear] (TyOwn) must be
    loaded exactly once on every execution path.  Zero uses = dropped;
    more than one use = possible duplication; non-zero use on some paths
    but zero on others = partial drop (dropped on the zero-count path).

    Level 7 — Aliasing safety: a parameter annotated [ExclBorrow] (TyMut)
    must be loaded at most once on any execution path.

    [SharedBorrow] (TyRef) and [Unrestricted] params are not checked.

    Branch semantics (per-path min/max): for [If] instructions we compute
    [(min_then, max_then)] and [(min_else, max_else)] independently and
    combine as [(min_then, max_then, min_else, max_else) →
    (min min_then min_else, max max_then max_else)].  A Linear param used
    exactly once in BOTH branches gives (1, 1) — OK.  A param used in the
    then-branch only gives (min=0, max=1) — [LinearDroppedOnSomePath].

    @see typed-wasm LEVEL-STATUS.md — Level 7 and Level 10 sections.
*)

open Codegen

(* ============================================================================
   Error types
   ============================================================================ *)

(** An ownership violation found in a Wasm function body. *)
type ownership_error =
  | LinearNotUsed of { func_idx : int; param_idx : int }
  (** Level 10: Linear parameter was never loaded on any path — dropped without
      consumption.  The owned resource leaks unconditionally. *)
  | LinearDroppedOnSomePath of { func_idx : int; param_idx : int }
  (** Level 10: Linear parameter is consumed on some execution paths but
      silently dropped on others (per-path min_uses = 0, max_uses >= 1).
      The caller's ownership guarantee is satisfied only conditionally —
      the drop path leaks the resource. *)
  | LinearUsedMultiple of { func_idx : int; param_idx : int; count : int }
  (** Level 10: Linear parameter was loaded [count] times on some path —
      potential duplication.  Only one consumer is permitted. *)
  | ExclBorrowAliased of { func_idx : int; param_idx : int; count : int }
  (** Level 7: ExclBorrow parameter was loaded [count] times — aliasing violation.
      An exclusive borrow must not have multiple simultaneous references. *)
  | ModuleNotIsolated of { reason : string }
  (** Level 13 (module isolation — the typed-wasm contract widening,
      issue #234/#10). The emitted module owns its own linear memory
      yet ALSO imports a memory or table: a cross-module shared-state
      channel outside the declared function-import boundary.
      AffineScript codegen always emits a private memory and never
      imports one, so this is the negative L13 property and a violation
      is a codegen/isolation regression — exactly the class
      [tw_verify] exists to catch on emitted wasm. Additive: NO
      ownership-section wire-format change, so the multi-producer ABI
      (ephapax + typed-wasm) is untouched. *)

(* ============================================================================
   Per-path use-range analysis
   ============================================================================ *)

(** Compute [(min_uses, max_uses)] — the minimum and maximum number of times
    [local_idx] is loaded across all execution paths within [instrs].

    Sequential instructions sum their ranges: if instr A uses the param
    (1, 1) times and instr B uses it (0, 1) times, the sequence is (1, 2).

    For [If] nodes only one branch executes, so we take the component-wise
    min/max of the two branch ranges:
      - [min = min(then_min, else_min)]: may be zero if one branch doesn't use it
      - [max = max(then_max, else_max)]: the worst-case duplication count

    Example: [If { LocalGet 0 } { LocalGet 0 }] → [(1, 1)] OK for Linear.
    Example: [If { LocalGet 0 } { [] }]         → [(0, 1)] LinearDroppedOnSomePath.
    Example: [If { LocalGet 0; LocalGet 0 } { LocalGet 0 }] → [(1, 2)]
      → LinearUsedMultiple (max = 2).

    Nested [Block] and [Loop] bodies are descended into (no new paths). *)
let rec count_uses_range (local_idx : int) (instrs : Wasm.instr list) : int * int =
  List.fold_left (fun (a_min, a_max) instr ->
    let (i_min, i_max) = uses_range_in_instr local_idx instr in
    (a_min + i_min, a_max + i_max)
  ) (0, 0) instrs

and uses_range_in_instr (local_idx : int) (instr : Wasm.instr) : int * int =
  match instr with
  | Wasm.LocalGet n -> if n = local_idx then (1, 1) else (0, 0)
  | Wasm.Block (_, body) | Wasm.Loop (_, body) ->
    count_uses_range local_idx body
  | Wasm.If (_, then_, else_) ->
    let (t_min, t_max) = count_uses_range local_idx then_ in
    let (e_min, e_max) = count_uses_range local_idx else_ in
    (min t_min e_min, max t_max e_max)
  | _ -> (0, 0)

(* ============================================================================
   Per-function verification
   ============================================================================ *)

(** Verify ownership constraints for one function.

    [func] is the Wasm function body.
    [param_kinds] are the ownership annotations for each parameter (in order).
    [func_idx] is the global function index (for error reporting).
    Returns all violations found (empty list = clean).

    Uses per-path min/max analysis:
    - [LinearNotUsed]: max_uses = 0 (param dropped on every path)
    - [LinearDroppedOnSomePath]: min_uses = 0, max_uses >= 1 (dropped conditionally)
    - [LinearUsedMultiple]: max_uses > 1 (duplicated on some path)
    Multiple violations can be reported for a single param (e.g. both
    [LinearDroppedOnSomePath] and [LinearUsedMultiple] if min=0, max>1). *)
let verify_function
    (func     : Wasm.func)
    (param_kinds : ownership_kind list)
    (func_idx : int)
  : ownership_error list =
  List.concat_map (fun (param_idx, kind) ->
    let (min_uses, max_uses) = count_uses_range param_idx func.Wasm.f_body in
    match kind with
    | Linear ->
      (* Exactly once on every path: zero everywhere = dropped; zero on some
         path = partial drop; more than one on some path = may duplicate. *)
      let drop_errors =
        if max_uses = 0 then [LinearNotUsed { func_idx; param_idx }]
        else if min_uses = 0 then [LinearDroppedOnSomePath { func_idx; param_idx }]
        else []
      in
      let dup_errors =
        if max_uses > 1 then [LinearUsedMultiple { func_idx; param_idx; count = max_uses }]
        else []
      in
      drop_errors @ dup_errors
    | ExclBorrow ->
      (* At most once on any path: max > 1 creates simultaneous aliases. *)
      if max_uses > 1 then
        [ExclBorrowAliased { func_idx; param_idx; count = max_uses }]
      else
        []
    | Unrestricted | SharedBorrow ->
      (* No constraints on these ownership kinds. *)
      []
  ) (List.mapi (fun i k -> (i, k)) param_kinds)

(* ============================================================================
   Module-level verification
   ============================================================================ *)

(** Verify ownership constraints across an entire Wasm module.

    [wasm_mod] is the compiled module.
    [annots] is the list of [(func_idx, param_kinds, ret_kind)] annotations
    collected by codegen and stored in the [typedwasm.ownership] custom
    section (but here provided in structured form directly from [Codegen]).

    Imported functions have no body and are skipped.
    Returns all violations found. *)
let verify_module
    (wasm_mod : Wasm.wasm_module)
    (annots   : (int * ownership_kind list * ownership_kind) list)
  : ownership_error list =
  let import_count = List.length wasm_mod.Wasm.imports in
  List.concat_map (fun (func_idx, param_kinds, _ret_kind) ->
    let local_idx = func_idx - import_count in
    if local_idx < 0 || local_idx >= List.length wasm_mod.Wasm.funcs then
      []   (* Imported function: no IR to inspect *)
    else
      let func = List.nth wasm_mod.Wasm.funcs local_idx in
      verify_function func param_kinds func_idx
  ) annots

(* ============================================================================
   Pipeline integration — parse ownership section from the module
   ============================================================================ *)

(** Parse the [typedwasm.ownership] custom section payload into structured
    [(func_idx, param_kinds, ret_kind)] annotations.

    Supports BOTH v1 (legacy, unversioned) and v2 (ADR-020, accepted
    + ratified 2026-05-24) on-wire formats; v2 is identified by the
    [0xAF] sentinel byte followed by a version byte.  v1 readers
    fail cleanly on v2 sections (the sentinel pollutes the
    entry-count low byte and yields an implausible count); v2
    readers see the sentinel and dispatch.

    v1 layout (legacy, unversioned):
      u32le  count
      for each entry:
        u32le  func_idx
        u8     n_params
        u8[n]  param_kinds  (0=Unrestricted, 1=Linear, 2=SharedBorrow, 3=ExclBorrow)
        u8     ret_kind

    v2 layout (ADR-020):
      u8     0xAF         (* "AffineScript Format" sentinel *)
      u8     version      (* 0x02 for v2.0 *)
      u32le  count
      entry*              (* same entry shape as v1 *)

    Per ADR-021's coordinated landing protocol, this verifier ships
    parse-support FIRST.  AffineScript's emitter ([Codegen.build_ownership_
    section] / [Tw_section.Encode.ownership]) stays on v1 until
    [hyperpolymath/typed-wasm]'s Rust verifier + [hyperpolymath/ephapax]'s
    OCaml verifier also ship parse-support; then all producers flip to v2
    emit together.  See [docs/specs/TYPED-WASM-COORDINATION-LEDGER.adoc]
    Q-001 for the current state of the cross-repo rollout. *)
let parse_ownership_section_payload
    (payload : bytes)
  : (int * ownership_kind list * ownership_kind) list =
  let kind_of_byte = function
    | 1 -> Linear
    | 2 -> SharedBorrow
    | 3 -> ExclBorrow
    | _ -> Unrestricted
  in
  let pos = ref 0 in
  let len = Bytes.length payload in
  (* v2 sentinel dispatch.  Two bytes minimum: [0xAF, version].  The
     sentinel is byte-distinct from any v1 entry-count low byte for
     a module with <= 0xAE entries (i.e. effectively every module);
     ADR-020 catalogues this trade-off explicitly. *)
  let v2_detected =
    len >= 2
    && Char.code (Bytes.get payload 0) = 0xAF
  in
  if v2_detected then begin
    let v2_version = Char.code (Bytes.get payload 1) in
    pos := 2;
    (* v2.0 is the only known version today.  An unknown v2.X is a
       future-version section we cannot interpret; per ADR-021
       conservative-disposition (sound fallback), treat as empty.
       The verifier therefore enforces nothing on the section —
       loud-correct rather than silent-wrong. *)
    if v2_version <> 0x02 then []
    else
      let read_u32_le () =
        if !pos + 4 > len then 0
        else begin
          let b0 = Char.code (Bytes.get payload  !pos)      in
          let b1 = Char.code (Bytes.get payload (!pos + 1)) in
          let b2 = Char.code (Bytes.get payload (!pos + 2)) in
          let b3 = Char.code (Bytes.get payload (!pos + 3)) in
          pos := !pos + 4;
          b0 lor (b1 lsl 8) lor (b2 lsl 16) lor (b3 lsl 24)
        end
      in
      let read_u8 () =
        if !pos >= len then 0
        else begin
          let b = Char.code (Bytes.get payload !pos) in
          pos := !pos + 1;
          b
        end
      in
      let count = read_u32_le () in
      List.init count (fun _ ->
        let func_idx    = read_u32_le () in
        let n_params    = read_u8 ()     in
        let param_kinds = List.init n_params (fun _ -> kind_of_byte (read_u8 ())) in
        let ret_kind    = kind_of_byte (read_u8 ()) in
        (func_idx, param_kinds, ret_kind))
  end else begin
    (* v1 (legacy) path.  Identical to the pre-ADR-020 behaviour;
       preserved verbatim for backward compatibility with already-
       emitted modules and with the in-tree fixtures. *)
    let read_u32_le () =
      if !pos + 4 > len then 0
      else begin
        let b0 = Char.code (Bytes.get payload  !pos)      in
        let b1 = Char.code (Bytes.get payload (!pos + 1)) in
        let b2 = Char.code (Bytes.get payload (!pos + 2)) in
        let b3 = Char.code (Bytes.get payload (!pos + 3)) in
        pos := !pos + 4;
        b0 lor (b1 lsl 8) lor (b2 lsl 16) lor (b3 lsl 24)
      end
    in
    let read_u8 () =
      if !pos >= len then 0
      else begin
        let b = Char.code (Bytes.get payload !pos) in
        pos := !pos + 1;
        b
      end
    in
    let count = read_u32_le () in
    List.init count (fun _ ->
      let func_idx    = read_u32_le () in
      let n_params    = read_u8 ()     in
      let param_kinds = List.init n_params (fun _ -> kind_of_byte (read_u8 ())) in
      let ret_kind    = kind_of_byte (read_u8 ()) in
      (func_idx, param_kinds, ret_kind))
  end

(** Level 13 — module isolation (negative form). A well-isolated
    module owns its private linear memory and reaches other modules
    ONLY through the declared function-import boundary. If the module
    has its own memory yet also imports a memory or a table, that
    imported entity is a shared-state channel outside the boundary —
    not isolated. (A module with no own memory that imports one is a
    pure consumer shim, not in scope here; AffineScript codegen never
    emits that shape.) Carrier-free: read only the standard wasm
    import/memory sections, so no multi-producer ABI change. *)
let verify_module_isolation
    (wasm_mod : Wasm.wasm_module)
  : ownership_error list =
  if wasm_mod.Wasm.mems = [] then []
  else
    List.filter_map
      (fun (im : Wasm.import) ->
        match im.Wasm.i_desc with
        | Wasm.ImportMemory ->
          Some (ModuleNotIsolated {
            reason =
              Printf.sprintf
                "module owns linear memory yet imports memory '%s.%s' \
                 (cross-module shared memory breaks L13 isolation)"
                im.Wasm.i_module im.Wasm.i_name })
        | Wasm.ImportTable ->
          Some (ModuleNotIsolated {
            reason =
              Printf.sprintf
                "module owns linear memory yet imports table '%s.%s' \
                 (externally-backed table breaks L13 isolation)"
                im.Wasm.i_module im.Wasm.i_name })
        | Wasm.ImportFunc _ -> None)
      wasm_mod.Wasm.imports

(** Verify a Wasm module using the embedded [typedwasm.ownership] custom
    section.  This is the primary entry point for the pipeline and the CLI.

    Returns [Ok ()] if no violations are found, [Error errs] otherwise. *)
let verify_from_module
    (wasm_mod : Wasm.wasm_module)
  : (unit, ownership_error list) Result.t =
  match List.assoc_opt "typedwasm.ownership" wasm_mod.Wasm.custom_sections with
  | None ->
    (* No ownership section: nothing to verify.  This is not an error —
       modules compiled without ownership qualifiers have no constraints. *)
    Ok ()
  | Some payload ->
    let annots = parse_ownership_section_payload payload in
    (* L7/L10 (ownership) + L13 (module isolation). The isolation check
       is gated behind the same ownership-section presence so the
       "no section ⇒ Ok" contract is preserved. *)
    let errors =
      verify_module wasm_mod annots @ verify_module_isolation wasm_mod
    in
    if errors = [] then Ok () else Error errors

(* ============================================================================
   Pretty-printing
   ============================================================================ *)

(** Format a single ownership error for human-readable output. *)
let pp_error (fmt : Format.formatter) (err : ownership_error) : unit =
  match err with
  | LinearNotUsed { func_idx; param_idx } ->
    Format.fprintf fmt
      "Level 10 violation: function %d, param %d — Linear (own) param dropped \
       on all paths (must be consumed exactly once)"
      func_idx param_idx
  | LinearDroppedOnSomePath { func_idx; param_idx } ->
    Format.fprintf fmt
      "Level 10 violation: function %d, param %d — Linear (own) param dropped \
       on some paths (per-path min uses = 0; must be consumed on every path)"
      func_idx param_idx
  | LinearUsedMultiple { func_idx; param_idx; count } ->
    Format.fprintf fmt
      "Level 10 violation: function %d, param %d — Linear (own) param loaded \
       %d times on some path (exactly 1 required; possible duplication)"
      func_idx param_idx count
  | ExclBorrowAliased { func_idx; param_idx; count } ->
    Format.fprintf fmt
      "Level 7 violation: function %d, param %d — ExclBorrow (mut) param \
       aliased (%d simultaneous references; at most 1 permitted)"
      func_idx param_idx count
  | ModuleNotIsolated { reason } ->
    Format.fprintf fmt "Level 13 violation: %s" reason

(** Format a full verification report. Prints "OK" if no errors. *)
let pp_report (fmt : Format.formatter) (errs : ownership_error list) : unit =
  match errs with
  | [] ->
    Format.fprintf fmt "typed-wasm ownership verification: OK@."
  | _ ->
    Format.fprintf fmt "typed-wasm ownership verification: %d violation(s)@."
      (List.length errs);
    List.iter (fun e ->
      Format.fprintf fmt "  %a@." pp_error e
    ) errs