cubical-transport-hott-lean4/Topolei/RenderProbe.lean

/-
  Topolei.Render.Probe
  ====================
  Empirical check for `render_faithful` — the axiom in `Topolei.GPU.Spec`
  that claims the GPU's runtime output equals the compiled shader's
  `ShaderSemantic` at every pixel and uniform setting.

  The Rust FFI `topolei_canvas_render_probe_path_pixel` accepts a Lean
  `EMLPath` inductive structurally, builds the fragment shader as a
  `naga::Module` directly via `emit_naga::build_probe_module` (no GLSL
  text intermediary), submits the module to wgpu's SPIR-V backend, and
  reads back a single pixel from an `Rgba32Float` offscreen target.
  This module binds that FFI and exercises it on the demo paths
  `plotExp` / `plotLn` / `plotTransp`.

  ## Tolerance

  IEEE 754 rounding differs slightly between the CPU (Lean-side
  `ShaderSemantic` evaluation via `Float.exp`, `Float.log`, `Float.cos`)
  and the GPU (wgpu → Vulkan → driver → shader unit).  The comparison
  therefore uses an absolute tolerance rather than bit-exact equality.
  `5e-3` is generous; a bit-exact probe on known IEEE 754 drivers
  (e.g. SwiftShader) might tighten it to 1e-6 or better.

  ## Running

  The probe requires a Vulkan/Metal/DX12 adapter (or a software
  rasterizer like SwiftShader) on the host.  When no adapter is
  available the FFI returns a sentinel `RGBA { r := -1, g := -1, ... }`
  and the test skips that comparison rather than asserting failure.

  Run with:  `lake build probe-test && ./.lake/build/bin/probe-test`.
-/

import Topolei.EML
import Topolei.EML.Path
import Topolei.GPU.Spec

namespace TopoleiProbe

/-- Rendered-pixel RGBA as four Lean `Float` values.  Layout matches
    the `topolei_canvas_shim_io_ok_rgba` helper in `canvas-rs/shim.c`:
    tag 0, 4 boxed-Float fields. -/
structure RGBA where
  r : Float
  g : Float
  b : Float
  a : Float
deriving Repr, Inhabited

/-- **Structured probe.**  `renderProbePath path width height time
    pathParam x y` — passes the `EMLPath` inductive *structurally* to
    the Rust side, which walks the object tree
    (`eml::emlpath_from_lean`), builds a `naga::Module` directly via
    `emit_naga::build_probe_module`, and runs the offscreen probe.

    No shader-source string crosses the FFI boundary.  See
    `NAGA_IR_PLAN.md` for the construction; `compileEMLPath_correct`
    in `Topolei.GPU.Spec` is the contract the IR builder must satisfy. -/
@[extern "topolei_canvas_render_probe_path_pixel"]
opaque renderProbePath
    (path : @& EMLPath)
    (width height : UInt32)
    (time pathParam : Float)
    (x y : UInt32) : IO RGBA

/-- True when `renderProbePath` returned its GPU-unavailable sentinel.
    A channel of `-1.0` is outside the valid RGBA range `[0, 1]`, so
    tests can use this as a reliable skip signal. -/
def RGBA.isSentinel (c : RGBA) : Bool :=
  c.r < 0.0 && c.g < 0.0 && c.b < 0.0 && c.a < 0.0

/-- Bridge to `GPU/Spec.lean`: the CPU-side `PixelColor` that the GPU
    is expected to match.  One-to-one translation. -/
def RGBA.ofPixelColor (c : PixelColor) : RGBA :=
  { r := c.r, g := c.g, b := c.b, a := c.a }

/-- Absolute-tolerance comparison for `Float`.  Both arguments must
    differ by no more than `tol`. -/
def floatClose (a b tol : Float) : Bool :=
  let d := a - b
  let ad := if d < 0.0 then -d else d
  ad <= tol

/-- Compare two RGBA values channel-wise under `tol`. -/
def rgbaClose (expected actual : RGBA) (tol : Float) : Bool :=
  floatClose expected.r actual.r tol &&
  floatClose expected.g actual.g tol &&
  floatClose expected.b actual.b tol &&
  floatClose expected.a actual.a tol

/-- Convert framebuffer pixel `(x, y)` at resolution `(w, h)` into the
    `uv ∈ [0, 1]²` coordinates the probe shader sees at that pixel's
    fragment center.  Framebuffer y is top-down (wgpu origin) while
    the vertex shader's uv tracks NDC, so we flip: `uv.y = 1 -
    (y+0.5)/h`.  `uv.x = (x+0.5)/w`.  The `+0.5` lands us on the pixel
    center where the rasterizer samples. -/
def pixelUV (x y width height : UInt32) : PixelCoord :=
  let w := width.toNat.toFloat
  let h := height.toNat.toFloat
  let px := (x.toNat.toFloat + 0.5) / w
  let py := 1.0 - (y.toNat.toFloat + 0.5) / h
  { x := px, y := py }

/-- Run a single probe: pass the `EMLPath` directly to the Rust
    offscreen probe; compare the GPU-returned pixel `(x, y)` against
    Lean's `EMLPath.toColor path (pixelUV ...) u`. -/
def runProbe
    (path : EMLPath) (label : String)
    (width height : UInt32)
    (time pathParam : Float)
    (x y : UInt32)
    (tol : Float := 5e-3) : IO Bool := do
  let actual ← renderProbePath path width height time pathParam x y
  if actual.isSentinel then
    IO.println s!"  ⏭  SKIP [{label}] at ({x}, {y}) — no GPU adapter"
    return true
  let p : PixelCoord := pixelUV x y width height
  let u : FrameUniforms :=
    { time, resWidth := width.toNat.toFloat, resHeight := height.toNat.toFloat,
      pathParam }
  let expectedColor := EMLPath.toColor path p u
  let expected := RGBA.ofPixelColor expectedColor
  if rgbaClose expected actual tol then
    IO.println s!"  ✅ [{label}] pixel ({x}, {y}): GPU matches Lean (r={expected.r} g={expected.g} b={expected.b})"
    return true
  else
    IO.println s!"  ❌ [{label}] pixel ({x}, {y}) divergent"
    IO.println s!"       expected: r={expected.r} g={expected.g} b={expected.b} a={expected.a}"
    IO.println s!"       actual:   r={actual.r} g={actual.g} b={actual.b} a={actual.a}"
    return false

/-- Full probe battery.  Returns number of failures.

    Each entry asserts `EMLPath.toColor path p u ≈ GPU_pixel(p, u)`
    within tolerance — i.e. the inhabited form of `render_faithful`
    on the naga-IR pipeline. -/
def runProbes : IO UInt32 := do
  IO.println "── Topolei render_faithful probes ──"
  IO.println "  Naga-IR pixel probes (EMLPath → naga::Module → SPIR-V → GPU):"
  -- Constant path (dim variable absent from body): `pathParam` is
  -- irrelevant and the two endpoints coincide.
  let n1 ← runProbe plotExp.toEMLPath "plotExp" 128 128 0.0 0.0 64 64
  let n2 ← runProbe plotExp.toEMLPath "plotExp" 128 128 0.0 1.0 64 64
  let n3 ← runProbe plotLn.toEMLPath  "plotLn"  128 128 0.0 0.5 32 64
  -- Parametric path: `pathParam` changes the color — sample across
  -- values to cover the sweep.
  let n4 ← runProbe plotTransp.toEMLPath "plotTransp" 128 128 0.0 0.0 64 64
  let n5 ← runProbe plotTransp.toEMLPath "plotTransp" 128 128 0.0 0.5 64 64
  let n6 ← runProbe plotTransp.toEMLPath "plotTransp" 128 128 0.0 1.0 64 64
  let n7 ← runProbe plotTransp.toEMLPath "plotTransp" 128 128 0.0 0.5 32 96
  let mut fails : UInt32 := 0
  for ok in [n1, n2, n3, n4, n5, n6, n7] do
    if !ok then fails := fails + 1
  IO.println s!"── {7 - fails.toNat} / 7 probes passed ──"
  return fails

end TopoleiProbe