Harvey–van der Hoeven machine development, part 6: definitions for the final O(n log n) assembly

-- Harvey–van der Hoeven multitape-TM development (NLogN), generated by skeleton subtraction.

import Definitions.Def_IntMul_HvdH_M_Defs5

set_option autoImplicit false
set_option relaxedAutoImplicit false
set_option maxHeartbeats 1000000

section
-- Source module: NLogN.HGl2
/-!
# Gaussian list primitives, part 2: stacks, encodings, cost lemmas, Gaussian zip products

Fixed stacks of the list pipeline (namespace `NLogN.GL`):

* `gD = 0` data, `gZA = 4` second operand of zips, `gZB = 5` reversed outputs, `gZT = 6` temp,
* counters `gC1 = 7` (element loops), `gCV = 8` (moves), `gCC = 9` (chunk loop),
* scratch registers `gR1 … gR4 = 10 … 13`, `gP1 … gP4 = 14 … 17`, dummies `gU1 = 18`, `gU2 = 19`,
* read-only parameters `gPP = 20` (`2^p`), `gPT = 21` (`2^(p+logT)`), `gPS = 22` (`2^(p-b)`),
  `gPQ = 23` (`S^2`), `gPC = 24` (`2^(2b+6Kα²+1)`), `gPR = 25` (`2^(p+1)`), `gPN = 26` (`T`),
  `gPM = 27` (`S`).

`zipP σ`: `gD : gints U ++ R₂`, `gZA : gints A ++ R₁` ↦ `gD : gints (zipWith (gzF σ p) A U) ++ R₂`,
`gZA : R₁` with `gzF σ p a z = ((z₁a₁ ± z₂a₂) tdiv 2^p, (±z₂a₁ ± z₁a₂) tdiv 2^p)`; four
multiplication calls per element.
-/

set_option linter.unusedSimpArgs false

namespace NLogN.GL

open IntMul NLogN

abbrev gD : Fin HS := 0
abbrev gZA : Fin HS := 4
abbrev gZB : Fin HS := 5
abbrev gZT : Fin HS := 6
abbrev gC1 : Fin HS := 7
abbrev gCV : Fin HS := 8
abbrev gCC : Fin HS := 9
abbrev gR1 : Fin HS := 10
abbrev gR2 : Fin HS := 11
abbrev gR3 : Fin HS := 12
abbrev gR4 : Fin HS := 13
abbrev gP1 : Fin HS := 14
abbrev gP2 : Fin HS := 15
abbrev gP3 : Fin HS := 16
abbrev gP4 : Fin HS := 17
abbrev gU1 : Fin HS := 18
abbrev gU2 : Fin HS := 19
abbrev gPP : Fin HS := 20
abbrev gPT : Fin HS := 21
abbrev gPS : Fin HS := 22
abbrev gPQ : Fin HS := 23
abbrev gPC : Fin HS := 24
abbrev gPR : Fin HS := 25
abbrev gPN : Fin HS := 26
abbrev gPM : Fin HS := 27

/-! ## Encodings -/

def genc (z : ℤ × ℤ) : List Item := [.int z.1, .int z.2]
def ienc (x : ℤ) : List Item := [.int x]

/-- Bounded Gaussian entries. -/
def GB (B : ℕ) (L : List (ℤ × ℤ)) : Prop := ∀ z ∈ L, zsz z.1 ≤ B ∧ zsz z.2 ≤ B

/-! ## Instruction cost lemmas -/

/-! ## The Gaussian zip product -/

def sgE (b : Bool) (e : Ex HS) : Ex HS := if b then .neg e else e
def sgZ (b : Bool) (x : ℤ) : ℤ := if b then -x else x

def reE (σ1 : Bool) : Ex HS := .quo (.add (.reg gP1) (sgE σ1 (.reg gP2))) (.reg gPP)
def imE (σ2 σ3 : Bool) : Ex HS := .quo (.add (sgE σ2 (.reg gP3)) (sgE σ3 (.reg gP4))) (.reg gPP)

/-- Signs `(σ1, σ2, σ3)`. -/
abbrev Sg := Bool × Bool × Bool

def zipBody (σ : Sg) : RM := RM.seqL [
  .rs (.push gR1 (.reg gZA)), .rs (.pop gZA), .rs (.push gR2 (.reg gZA)), .rs (.pop gZA),
  .rs (.push gR3 (.reg gD)), .rs (.pop gD), .rs (.push gR4 (.reg gD)), .rs (.pop gD),
  .mul gP1 gR1 gR3, .mul gP2 gR2 gR4, .mul gP3 gR1 gR4, .mul gP4 gR2 gR3,
  .rs (.push gZB (NLogN.GL.reE σ.1)), .rs (.push gZB (NLogN.GL.imE σ.2.1 σ.2.2)),
  .rs (.pop gP4), .rs (.pop gP3), .rs (.pop gP2), .rs (.pop gP1),
  .rs (.pop gR4), .rs (.pop gR3), .rs (.pop gR2), .rs (.pop gR1)]

/-- The element function of `zipBody σ`. -/
def gzF (σ : Sg) (p : ℕ) (a z : ℤ × ℤ) : ℤ × ℤ :=
  ((z.1 * a.1 + sgZ σ.1 (z.2 * a.2)).tdiv (2 ^ p),
    (sgZ σ.2.1 (z.2 * a.1) + sgZ σ.2.2 (z.1 * a.2)).tdiv (2 ^ p))

/-- `mulZ p (conjZ a) z`. -/
abbrev sgConjMul : Sg := (false, false, true)
/-- `conjZ (mulZ p a b)`. -/
abbrev sgMulConj : Sg := (true, true, true)

end NLogN.GL

end

section
-- Source module: NLogN.HGl3
/-!
# Gaussian list primitives, part 3: moves, transfers, the zip product program

* `mv_all`: `mvList` as an `RM` step (run / ok / cost in one statement).
* `xferP i j n`: move `n` integers from `i` to `j` keeping their order (via `gZT`).
* `zipP σ cN`: `gD := zipWith (gzF σ p) A U` (`A` consumed from `gZA`), `cN ∈ {gPN, gPM}` holds
  the length.
-/

set_option linter.unusedSimpArgs false

namespace NLogN.GL

open IntMul NLogN

/-- Order-preserving transfer through `gZT`. -/
def xferP (i j : Fin HS) (n : Ex HS) : RM :=
  .seq (.rs (mvList i gZT gCV n)) (.rs (mvList gZT j gCV n))

/-! ## Count expressions -/

def cnt2 (cN : Fin HS) : Ex HS := .add (.reg cN) (.reg cN)

/-- Count stacks are `gPN` or `gPM`. -/
def IsCnt (cN : Fin HS) : Prop := cN = gPN ∨ cN = gPM

/-! ## The zip product program -/

def zipP (σ : Sg) (cN : Fin HS) : RM :=
  .seq (.rep gC1 (.reg cN) (zipBody σ)) (.rs (mvList gZB gD gCV (cnt2 cN)))

end NLogN.GL

end

section
-- Source module: NLogN.HGl4
/-!
# Gaussian list primitives, part 4: scaling, `tdiv`, rounding, copying

* `scaleP cN`: `gD : ints u ++ R ↦ gints (u.map (a ↦ (a·2^k, 0))) ++ R` (`gPS = 2^k`).
* `tdivP cN`: `gD : gints W ++ R ↦ gints (W.map (z ↦ (z₁ tdiv 2^k, z₂ tdiv 2^k))) ++ R`
  (`gPT = 2^k`).
* `roundP cN`: `gD : gints Y ++ R ↦ ints (Y.map rnd) ++ R` with
  `rnd z = (S²·z₁·2^(c+1) + 2^p) / 2^(p+1)` (floor division; `gPQ = S²`, `gPC = 2^(c+1)`,
  `gPP = 2^p`, `gPR = 2^(p+1)`); one multiplication call per element.
* `copyP cN`: `gCH : gints A ++ R ↦ gints A ++ R` and `gZA : X ↦ gints A ++ X`.

Floor division and `tdiv` by powers of two use only `quo`/`rem`/`sel` (split-based) expressions.
-/

set_option linter.unusedSimpArgs false

namespace NLogN.GL

open IntMul NLogN

abbrev gCH : Fin HS := 28

/-! ## Restoring the order of a reversed output -/

/-! ## Scaling -/

def scaleBody : RM := RM.seqL [.rs (.push gZB (.shl (.reg gD) (.reg gPS))), .rs (.pop gD),
  .rs (.push gZB (.lit 0))]

def scaleF (k : ℕ) (a : ℤ) : ℤ × ℤ := (a * 2 ^ k, 0)

def scaleP (cN : Fin HS) : RM :=
  .seq (.rep gC1 (.reg cN) scaleBody) (.rs (mvList gZB gD gCV (cnt2 cN)))

/-! ## `tdiv` by a power of two -/

def tdivBody : RM := RM.seqL [.rs (.push gR1 (.reg gD)), .rs (.pop gD),
  .rs (.push gZB (.quo (.reg gR1) (.reg gPT))), .rs (.push gZB (.quo (.reg gD) (.reg gPT))),
  .rs (.pop gD), .rs (.pop gR1)]

def tdivF (k : ℕ) (z : ℤ × ℤ) : ℤ × ℤ := (z.1.tdiv (2 ^ k), z.2.tdiv (2 ^ k))

def tdivP (cN : Fin HS) : RM :=
  .seq (.rep gC1 (.reg cN) tdivBody) (.rs (mvList gZB gD gCV (cnt2 cN)))

/-! ## Rounding -/

def rndX : Ex HS := .add (.shl (.reg gR1) (.reg gPC)) (.reg gPP)
def rndE : Ex HS := .add (.quo rndX (.reg gPR)) (.sel (.rem rndX (.reg gPR)) (.lit (-1)) (.lit 0))

def roundBody : RM := RM.seqL [.mul gR1 gD gPQ, .rs (.pop gD), .rs (.pop gD),
  .rs (.push gZB NLogN.GL.rndE), .rs (.pop gR1)]

/-- The rounding of one entry: `(S2·z₁·2^(c+1) + 2^p) / 2^(p+1)`. -/
def rndF (S2 : ℤ) (c p : ℕ) (z : ℤ × ℤ) : ℤ := (S2 * z.1 * 2 ^ (c + 1) + 2 ^ p) / 2 ^ (p + 1)

def roundP (cN : Fin HS) : RM :=
  .seq (.rep gC1 (.reg cN) NLogN.GL.roundBody) (.rs (mvList gZB gD gCV (.reg cN)))

/-! ## Copying a Gaussian list (keeping the original) -/

def copyBody : RM := RM.seqL [.rs (.push gCH (.reg gZT)), .rs (.push gZA (.reg gZT)),
  .rs (.pop gZT)]

def copyP (cN : Fin HS) : RM :=
  .seq (.rs (mvList gCH gZT gCV (cnt2 cN))) (.rep gC1 (cnt2 cN) copyBody)

end NLogN.GL

end

section
-- Source module: NLogN.HGl5
/-!
# The `tensorGo` combinator, part 1: list lemmas and abstract programs

* `chunkL` / `transposeL` facts (lengths, flattening, membership).
* `LProg M Ctx Pre f C`: an abstract RM program that consumes a Gaussian list `x` (with
  `Pre x`) from its input stack `I` and pushes `gints (f x)` on its output stack `O`, all other
  stacks unchanged, with `ok` and cost `≤ C`, whenever `Ctx` holds.
* `Stab Ctx W`: `Ctx` only depends on stacks outside `W` (and `ip`, `out`).
-/

set_option linter.unusedSimpArgs false

namespace NLogN.GL

open IntMul NLogN IntMul.HvdH.Ref

/-! ## List lemmas -/

/-! ## Abstract programs -/

/-- `Ctx` is insensitive to the stacks in `W` (and only depends on stacks). -/
def Stab (Ctx : IState HS → Prop) (W : List (Fin HS)) : Prop :=
  ∀ s t : IState HS, Ctx s → (∀ j, j ∉ W → t.G j = s.G j) → Ctx t

/-- An abstract list program: `I : gints x ++ R ↦ R`, `O : old ↦ gints (f x) ++ old`. -/
structure LProg (M : MulSpec) (Ctx : IState HS → Prop) (Pre : List (ℤ × ℤ) → Prop)
    (f : List (ℤ × ℤ) → List (ℤ × ℤ)) (C : ℕ) where
  p : RM
  I : Fin HS
  O : Fin HS
  hIO : I ≠ O
  hI : I.val < 55
  hO : O.val < 55
  run : ∀ (s : IState HS) (x : List (ℤ × ℤ)) (R : List Item), Ctx s → Pre x →
    s.G I = gints x ++ R → p.run s = (s.upd I R).upd O (gints (f x) ++ s.G O)
  ok : ∀ (s : IState HS) (x : List (ℤ × ℤ)) (R : List Item), Ctx s → Pre x →
    s.G I = gints x ++ R → p.ok M s
  cost : ∀ (s : IState HS) (x : List (ℤ × ℤ)) (R : List Item), Ctx s → Pre x →
    s.G I = gints x ++ R → p.cost M s ≤ C

/-- An abstract two-input list program (the shape of `fermatW`):
`I1 : gints x ++ R1 ↦ R1`, `I2 : gints y ++ R2 ↦ R2`, `O : old ↦ gints (f x y) ++ old`. -/
structure LProg2 (M : MulSpec) (Ctx : IState HS → Prop)
    (Pre : List (ℤ × ℤ) → List (ℤ × ℤ) → Prop)
    (f : List (ℤ × ℤ) → List (ℤ × ℤ) → List (ℤ × ℤ)) (C : ℕ) where
  p : RM
  I1 : Fin HS
  I2 : Fin HS
  O : Fin HS
  hd : [I1, I2, O].Nodup
  hI1 : I1.val < 55
  hI2 : I2.val < 55
  hO : O.val < 55
  run : ∀ (s : IState HS) (x y : List (ℤ × ℤ)) (R1 R2 : List Item), Ctx s → Pre x y →
    s.G I1 = gints x ++ R1 → s.G I2 = gints y ++ R2 →
    p.run s = ((s.upd I1 R1).upd I2 R2).upd O (gints (f x y) ++ s.G O)
  ok : ∀ (s : IState HS) (x y : List (ℤ × ℤ)) (R1 R2 : List Item), Ctx s → Pre x y →
    s.G I1 = gints x ++ R1 → s.G I2 = gints y ++ R2 → p.ok M s
  cost : ∀ (s : IState HS) (x y : List (ℤ × ℤ)) (R1 R2 : List Item), Ctx s → Pre x y →
    s.G I1 = gints x ++ R1 → s.G I2 = gints y ++ R2 → p.cost M s ≤ C

/-! ## General count expressions -/

end NLogN.GL

end

section
-- Source module: NLogN.HGl6
/-!
# The `tensorGo` combinator, part 2: one axis

An axis carries an abstract slice program `sl : LProg M Ctx (fun x => x.length = n ∧ GB B x) f C`
and five registers holding `n` (slice length), `n'` (output slice length), `m` (number of
slices), `kb` (key bits for the transpose, `n m ≤ 2^kb`) and `m n'` (output size).

`axisP ax` on `gD : gints L ++ Rest` (`L.length = n m`): transpose `n × m` (`transP 2` on stacks
`0 … 3`), apply `sl` to each of the `m` consecutive slices of length `n` (moving every slice into
`sl.I` and every result from `sl.O`), and restore the order:
`gD : gints ((chunkL n (transposeL n m L)).flatMap f) ++ Rest`.
-/

set_option linter.unusedSimpArgs false

namespace NLogN.GL

open IntMul NLogN IntMul.HvdH.Ref

/-- The stacks used by the list pipeline (`0 … 19`, `gCH`, `gSA = 29`, `gV = 30`). -/
def myW : List (Fin HS) := [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19,
  20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30]

structure Axis (M : MulSpec) (Ctx : IState HS → Prop) (B : ℕ) where
  n : ℕ
  n' : ℕ
  m : ℕ
  kb : ℕ
  C : ℕ
  f : List (ℤ × ℤ) → List (ℤ × ℤ)
  sl : LProg M Ctx (fun x => x.length = n ∧ GB B x) f C
  rn : Fin HS
  rn' : Fin HS
  rm : Fin HS
  rkb : Fin HS
  rtot : Fin HS

variable {M : MulSpec} {Ctx : IState HS → Prop} {B : ℕ}

/-- The cost bound of one axis. -/
def axCost (ax : Axis M Ctx B) (Q : ℕ) : ℕ :=
  3 * 10 ^ 6 * (ax.kb + 1) * (ax.n * ax.m + 1) * (B + ax.kb + 6) +
    ax.m * (ax.C + 100000 * (ax.n + ax.n' + 1) * (Q + 2)) + 100000 * (ax.m * ax.n' + 1) * (Q + 2)

def chunkBody (ax : Axis M Ctx B) : RM := RM.seqL [.rs (mvList gD gZT gCV (cnt2 ax.rn)),
  .rs (mvList gZT ax.sl.I gCV (cnt2 ax.rn)), ax.sl.p, .rs (mvList ax.sl.O gZB gCV (cnt2 ax.rn'))]

def preP (ax : Axis M Ctx B) : RM := RM.seqL [.rs (.push 1 (.reg ax.rn)),
  .rs (.push 2 (.reg ax.rm)), .rs (.push 3 (.reg ax.rkb)), .rs (transP 2), .rs (.pop 1),
  .rs (.pop 2), .rs (.pop 3)]

def axisP (ax : Axis M Ctx B) : RM :=
  .seq (preP ax) (.seq (.rep gCC (.reg ax.rm) (chunkBody ax))
    (.rs (mvList gZB gD gCV (cnt2 ax.rtot))))

/-- The output of one axis (the step of `tensorGo`). -/
def axOut (ax : Axis M Ctx B) (L : List (ℤ × ℤ)) : List (ℤ × ℤ) :=
  (chunkL ax.n (transposeL ax.n ax.m L)).flatMap ax.f

/-! ## All axes -/

end NLogN.GL

end

section
-- Source module: NLogN.HGl7
/-!
# `convL` and `FtL` programs (parametric in the `fermatW` program)

* `convP fw`: `gD : gints U ++ R`, `gZA : gints A ++ R₁` ↦ `gD : gints (convL q U A) ++ R`,
  `gZA : R₁` (transfers to the abstract `fermatW` program `fw`, then `tdiv 2^(p+logT)`).
* `ftlP fw`: `gD : gints U ++ R` ↦ `gD : gints (FtL q U) ++ R` with the chirp table
  `chirpL q` kept on `gCH`.
-/

set_option linter.unusedSimpArgs false

namespace NLogN.GL

open IntMul NLogN IntMul.HvdH IntMul.HvdH.Ref

abbrev gSA : Fin HS := 29
abbrev gV : Fin HS := 30

/-! ## `convL` -/

section
variable {M : MulSpec} {Ctx : IState HS → Prop} {PreW : List (ℤ × ℤ) → List (ℤ × ℤ) → Prop}
  {q : Params} {CW : ℕ}

def convP (fw : LProg2 M Ctx PreW (fermatW q) CW) : RM :=
  RM.seqL [xferP gD fw.I1 (cnt2 gPN), xferP gZA fw.I2 (cnt2 gPN), fw.p,
    xferP fw.O gD (cnt2 gPN), tdivP gPN]

/-- The conditions on the `fermatW` program used here. -/
structure FwOK (fw : LProg2 M Ctx PreW (fermatW q) CW) : Prop where
  i1 : fw.I1 ∉ myW
  i2 : fw.I2 ∉ myW
  o : fw.O ∉ myW
  stab : Stab Ctx (myW ++ [fw.I1, fw.I2, fw.O])

end

/-! ## `FtL` -/

section
variable {M : MulSpec} {Ctx : IState HS → Prop} {PreW : List (ℤ × ℤ) → List (ℤ × ℤ) → Prop}
  {q : Params} {CW : ℕ}

def ftlP (fw : LProg2 M Ctx PreW (fermatW q) CW) : RM :=
  RM.seqL [copyP gPN, zipP sgConjMul gPN, copyP gPN, convP fw, copyP gPN, zipP sgConjMul gPN]

/-- The Bluestein input `ã̄ · u` handed to `convL`. -/
def ftlV (q : Params) (U : List (ℤ × ℤ)) : List (ℤ × ℤ) :=
  List.zipWith (fun a u => mulZ q.p (conjZ a) u) (chirpL q) U

end

end NLogN.GL

end

section
-- Source module: NLogN.HGl8
/-!
# `fwdL` and `exactConvL` programs

* `fwdP axsS axsB fw`: `gD : gints U ++ R` ↦ `gD : gints (fwdL q U) ++ R`
  (`S̃` along all axes via `tgP axsS`, `FtL` via `ftlP fw`, `B̃` along all axes via `tgP axsB`).
* `exactP axsS axsB fw`: `gD : ints u ++ R`, `gV : ints v ++ RV` ↦
  `gD : ints (exactConvL q u v) ++ R`, `gV : RV`.
-/

set_option linter.unusedSimpArgs false

namespace NLogN.GL

open IntMul NLogN IntMul.HvdH IntMul.HvdH.Ref

/-- The `S̃` maps of `fwdL`. -/
def fwdAxS (q : Params) : List (ℕ × (List (ℤ × ℤ) → List (ℤ × ℤ))) :=
  (List.range q.K).map fun i => (2 ^ q.te i, sapL q.p (q.s.getD i 1) (2 ^ q.te i) q.α)

/-- The `B̃` maps of `fwdL`. -/
def fwdAxB (q : Params) : List (ℕ × (List (ℤ × ℤ) → List (ℤ × ℤ))) :=
  (List.range q.K).map fun i => (q.s.getD i 1, bapL q.p (q.s.getD i 1) (2 ^ q.te i) q.α)

/-- The input of the `FtL` stage of `fwdL q U`. -/
def fwdX (q : Params) (U : List (ℤ × ℤ)) : List (ℤ × ℤ) := tensorGo (fwdAxS q) q.sdims U

section
variable {M : MulSpec} {Ctx : IState HS → Prop} {PreW : List (ℤ × ℤ) → List (ℤ × ℤ) → Prop}
  {q : Params} {CW B : ℕ}

end

/-! ## `exactConvL` -/

/-- The scaling exponent `c = 2b + 6Kα²` of the final rounding. -/
def exC (q : Params) : ℕ := 2 * q.b + 6 * q.K * q.α ^ 2

/-- The scaled Gaussian input `u ↦ (u·2^(p-b), 0)`. -/
def exU (q : Params) (u : List ℤ) : List (ℤ × ℤ) := u.map (scaleF (q.p - q.b))

/-- The pointwise product handed to the third `fwdL`. -/
def exW (q : Params) (u v : List ℤ) : List (ℤ × ℤ) :=
  List.zipWith (fun a b => conjZ (mulZ q.p a b)) (fwdL q (exU q u)) (fwdL q (exU q v))

section
variable {M : MulSpec} {Ctx : IState HS → Prop} {PreW : List (ℤ × ℤ) → List (ℤ × ℤ) → Prop}
  {q : Params} {CW B : ℕ}

end

end NLogN.GL

end

section
-- Source module: NLogN.HAx1
/-!
# One resampling axis, part 1: `axisP` with separate input / output bounds

`GL.axisP_all` needs the slice map to preserve one bound `GB B`.  The resampling maps grow
their entries, so here `AxReady2 ax Bo s Q` asks instead that the slice map sends `GB B`
inputs (of length `n`) to `GB Bo` outputs (of length `n'`).  The program is unchanged
(`GL.axisP ax`).
-/

set_option linter.unusedSimpArgs false

namespace NLogN.AX

open IntMul NLogN IntMul.HvdH.Ref NLogN.GL

variable {M : MulSpec} {Ctx : IState HS → Prop} {B : ℕ}

/-- `GL.AxReady` with an output bound `Bo` for the slice map. -/
structure AxReady2 (ax : GL.Axis M Ctx B) (Bo : ℕ) (s : IState HS) (Q : ℕ) : Prop where
  hn : 1 ≤ ax.n
  hm : 1 ≤ ax.m
  hkb : ax.n * ax.m ≤ 2 ^ ax.kb
  hf : ∀ x, x.length = ax.n → GB B x → (ax.f x).length = ax.n' ∧ GB Bo (ax.f x)
  sI : ax.sl.I ∉ myW
  sO : ax.sl.O ∉ myW
  sW : gCC ∉ ax.sl.p.writes
  stab : Stab Ctx (myW ++ [ax.sl.I, ax.sl.O])
  regs : [ax.rn, ax.rn', ax.rm, ax.rkb, ax.rtot].all (fun r => decide (r ∉ myW) &&
    decide (r ≠ ax.sl.I) && decide (r ≠ ax.sl.O) && decide (r.val < 55)) = true
  vn : ∃ L, s.G ax.rn = .int ax.n :: L
  vn' : ∃ L, s.G ax.rn' = .int ax.n' :: L
  vm : ∃ L, s.G ax.rm = .int ax.m :: L
  vkb : ∃ L, s.G ax.rkb = .int ax.kb :: L
  vtot : ∃ L, s.G ax.rtot = .int (ax.m * ax.n') :: L
  bB : B ≤ Q
  bBo : Bo ≤ Q
  bn : zsz ((2 * ax.n : ℕ) : ℤ) ≤ Q
  bn' : zsz ((2 * ax.n' : ℕ) : ℤ) ≤ Q
  bm : zsz ((2 * ax.m : ℕ) : ℤ) ≤ Q
  bkb : zsz (ax.kb : ℤ) ≤ Q
  btot : zsz ((2 * (ax.m * ax.n') : ℕ) : ℤ) ≤ Q

end NLogN.AX

end

section
-- Source module: NLogN.HRspBand
/-!
# Banded multiply–accumulate on integer lists (core of Gaussian resampling)

`bandP` computes, for an input list `X` (stack `BX`), a delta list `Dl` (stack `BD`, the
gather positions are the prefix sums `c_k = Dl[0] + ⋯ + Dl[k]`), a weight table
`Ws = [W_0, …, W_{NR-1}]` (stack `BW`, rows concatenated) and an accumulator `acc` (stack `BA`):

  `acc[k] + ∑_{i < NR} (W_i[k] * X[c_k + i]).tdiv R`     (`R = 2^sh` on `BR`)

replacing `acc` on `BA`; `X` is replaced by `X.drop NR`; the tables are restored.
Every product goes through `mul`; everything else is linear list work.
-/

set_option linter.unusedSimpArgs false

namespace NLogN.Rsp

open IntMul NLogN RS

/-! ## Stacks -/

abbrev BX : Fin HS := 4
abbrev BY : Fin HS := 5
abbrev BW : Fin HS := 6
abbrev BWT : Fin HS := 7
abbrev BD : Fin HS := 8
abbrev BDT : Fin HS := 9
abbrev BA : Fin HS := 10
abbrev BAT : Fin HS := 11
abbrev BP : Fin HS := 12
abbrev BR : Fin HS := 13
abbrev BN : Fin HS := 14
abbrev BNR : Fin HS := 15
abbrev BNW : Fin HS := 16
abbrev BCR : Fin HS := 17
abbrev BCK : Fin HS := 18
abbrev BCM : Fin HS := 19
abbrev BCC : Fin HS := 20

/-! ## The list function -/

/-- Prefix sums of the deltas: `pre Dl k = Dl[0] + ⋯ + Dl[k-1]`. -/
def pre (Dl : List ℕ) (k : ℕ) : ℕ := (Dl.take k).sum

/-- One round: `acc[k] + (W[k] * X[pre (k+1)]).tdiv R`. -/
def bandRound (R : ℤ) (X : List ℤ) (Dl : List ℕ) (Wi acc : List ℤ) : List ℤ :=
  (List.range acc.length).map fun k =>
    acc.getD k 0 + (Wi.getD k 0 * X.getD (NLogN.Rsp.pre Dl (k + 1)) 0).tdiv R

/-- All rounds; round `i` reads `X.drop i`. -/
def bandF (R : ℤ) (Dl : List ℕ) : List (List ℤ) → List ℤ → List ℤ → List ℤ
  | [], _, acc => acc
  | Wi :: Ws, X, acc => bandF R Dl Ws (X.drop 1) (bandRound R X Dl Wi acc)

/-! ## Programs -/

/-- One output position: move `Dl[k]` entries of `X` to `Y`, multiply, accumulate. -/
def kBody : RM :=
  .seq (.rs (mvK NLogN.Rsp.BX BY BCM (.reg BD)))
    (.seq (.rs (asg BCC (.add (.reg BCC) (.reg BD))))
      (.seq (.mul BP NLogN.Rsp.BX NLogN.Rsp.BW)
        (.rs (seqL [push BAT (.add (.reg BA) (.quo (.reg BP) (.reg BR))), pop BP, pop BA,
          push BWT (.reg NLogN.Rsp.BW), pop NLogN.Rsp.BW, push BDT (.reg BD), pop BD]))))

/-- One round. -/
def roundBody : RM :=
  .seq (.rs (push BCC (.lit 0)))
    (.seq (.rep BCK (.reg BN) kBody)
      (.rs (seqL [mvK BY NLogN.Rsp.BX BCK (.reg BCC), pop BCC, pop NLogN.Rsp.BX, mvK BAT BA BCK (.reg BN),
        mvK BDT BD BCK (.reg BN)])))

/-- The whole program; `BNW` holds `n · NR`. -/
def bandP : RM :=
  .seq (.rep BCR (.reg BNR) NLogN.Rsp.roundBody) (.rs (mvK BWT NLogN.Rsp.BW BCK (.reg BNW)))

/-! ## One output position -/

/-- The state after `k` positions of a round started in `u`. -/
def kSt (u : IState HS) (R : ℤ) (Xi : List ℤ) (Dl : List ℕ) (Wi acc : List ℤ)
    (RX RW RD RA LCC : List Item) (k : ℕ) : IState HS :=
  ((((((((u.upd NLogN.Rsp.BX (ints (Xi.drop (NLogN.Rsp.pre Dl k)) ++ RX)).upd BY
    (ints (Xi.take (NLogN.Rsp.pre Dl k)).reverse ++ u.G BY)).upd BCC (.int (NLogN.Rsp.pre Dl k : ℤ) :: LCC)).upd NLogN.Rsp.BW
    (ints (Wi.drop k) ++ RW)).upd BWT (ints (Wi.take k).reverse ++ u.G BWT)).upd BD
    (ints ((Dl.map (Nat.cast : ℕ → ℤ)).drop k) ++ RD)).upd BDT
    (ints ((Dl.map (Nat.cast : ℕ → ℤ)).take k).reverse ++ u.G BDT)).upd BA (ints (acc.drop k) ++ RA)).upd BAT
    (ints ((bandRound R Xi Dl Wi acc).take k).reverse ++ u.G BAT)

/-! ## One round -/

end NLogN.Rsp

end

section
-- Source module: NLogN.HRspBand2
/-!
# Banded multiply–accumulate: rounds and the whole program (ok and cost)
-/

set_option linter.unusedSimpArgs false

namespace NLogN.Rsp

open IntMul NLogN RS

/-- The tail of a round: restore `X`, drop its head, restore `A` and `D`. -/
def postL : RS := seqL [mvK BY NLogN.Rsp.BX BCK (.reg BCC), pop BCC, pop NLogN.Rsp.BX, mvK BAT BA BCK (.reg BN),
  mvK BDT BD BCK (.reg BN)]

/-- Cost bound of one round. -/
def roundCost (n m Φ B dmax xl : ℕ) : ℕ :=
  25000 * (B + 2) + n * (Φ + 1000 * (m + 1) + 100000 * (B + 2) + dmax * (14000 * (B + 2))) +
    xl * (14000 * (B + 2))

/-! ## All rounds -/

/-- Outer loop states: `j` rounds done. -/
def rSt (s : IState HS) (R : ℤ) (Xl : List ℤ) (Dl : List ℕ) (Ws : List (List ℤ)) (acc : List ℤ)
    (RX RW RA : List Item) (j : ℕ) : IState HS :=
  (((s.upd NLogN.Rsp.BX (ints (Xl.drop j) ++ RX)).upd NLogN.Rsp.BW (ints (Ws.drop j).flatten ++ RW)).upd BWT
    (ints (Ws.take j).flatten.reverse ++ s.G BWT)).upd BA (ints (bandF R Dl (Ws.take j) Xl acc) ++ RA)

/-- Cost of `bandP`. -/
def bandCost (n NR m Φ B dmax xl : ℕ) : ℕ :=
  NR * (roundCost n m Φ B dmax xl + 10000 * (B + 2)) + (n * NR) * (14000 * (B + 2)) +
    10000 * (B + 2)

/-! ## Size of the accumulators -/

end NLogN.Rsp

end

section
-- Source module: NLogN.HRspCyc
/-!
# Cyclic padding of an integer list

`cycP` pops a list `u` (length `s ≥ 1`) from `CU` and pushes
`cycL u a N = [u[(a + j) mod s] | j < N]` onto `BX` (the input of `bandP`).
Parameters: `s` on `Ks`, `a < s` on `Ka`, `N` on `KN`.  The cost is linear in `N + s`.
-/

set_option linter.unusedSimpArgs false

namespace NLogN.Rsp

open IntMul NLogN RS

abbrev CU : Fin HS := 21
abbrev CZ : Fin HS := 25
abbrev CO : Fin HS := 26
abbrev CRm : Fin HS := 27
abbrev CC1 : Fin HS := 28
abbrev CC2 : Fin HS := 29
abbrev Ks : Fin HS := 30
abbrev Ka : Fin HS := 31
abbrev KN : Fin HS := 32

/-! ## Popping a list -/

/-- Pop `n` entries from `X`, counter on `C`. -/
def popK (X C : Fin HS) (n : Ex HS) : RS := rep C n (pop X)

/-- The cyclic padding. -/
def cycL (u : List ℤ) (a N : ℕ) : List ℤ := (List.range N).map fun j => u.getD ((a + j) % u.length) 0

def genBody : RS :=
  seqL [push CO (.reg CU), push CZ (.reg CU), pop CU, asg CRm (.add (.reg CRm) (.lit (-1))),
    ite (.add (.reg CRm) (.lit (-1))) (NLogN.RS.seq (mvK CZ CU CC2 (.reg Ks)) (asg CRm (.reg Ks))) skip]

def cycP : RS :=
  seqL [mvK CU CZ CC1 (.reg Ka), push CRm (.add (.reg Ks) (.neg (.reg Ka))),
    rep CC1 (.reg KN) genBody, mvK CO NLogN.Rsp.BX CC1 (.reg KN), mvK CU CZ CC1 (.reg CRm),
    rep CC1 (.reg Ks) (pop CZ), pop CRm]

/-- Generation states. -/
def gSt (s1 : IState HS) (u : List ℤ) (a : ℕ) (RU RZ LRm : List Item) (j : ℕ) : IState HS :=
  (((s1.upd CU (ints (u.drop ((a + j) % u.length)) ++ RU)).upd CZ
    (ints (u.take ((a + j) % u.length)).reverse ++ RZ)).upd CRm
    (.int ((u.length - (a + j) % u.length : ℕ) : ℤ) :: LRm)).upd CO
    (ints (cycL u a j).reverse ++ s1.G CO)

/-! ## Costs -/

/-- Cost bound of `cycP`. -/
def cycCost (s N B : ℕ) : ℕ := 100000 * (N + 2 * s + 1) * (B + 2)

end NLogN.Rsp

end

section
-- Source module: NLogN.HRspRes
/-!
# One banded resampling pass on an integer list

`resP` pops `u` from `CU`, pads it cyclically (`cycP`), runs `bandP` against the tables on
`BW`, `BD` with a zero accumulator, and pushes the result back on `CU`.
-/

set_option linter.unusedSimpArgs false

namespace NLogN.Rsp

open IntMul NLogN RS

/-- Push `n` zeros on `X`, counter on `C`. -/
def zerosP (X C : Fin HS) (n : Ex HS) : RS := rep C n (push X (.lit 0))

/-- The banded pass on an integer list. -/
def resP : RM :=
  RM.seqL [.rs cycP, .rs (NLogN.Rsp.zerosP BA CC1 (.reg BN)), bandP, .rs (popK NLogN.Rsp.BX CC1 (.reg Ks)),
    .rs (mvK BA CO CC1 (.reg BN)), .rs (mvK CO CU CC1 (.reg BN))]

/-- The list computed by `resP`. -/
def resL (R : ℤ) (Dl : List ℕ) (Ws : List (List ℤ)) (u : List ℤ) (a n : ℕ) : List ℤ :=
  bandF R Dl Ws (cycL u a (u.length + Ws.length)) (List.replicate n 0)

end NLogN.Rsp

namespace NLogN.Rsp

open IntMul NLogN RS

/-- Cost bound of `resP`. -/
def resCost (len NR n m Φ B dmax : ℕ) : ℕ :=
  cycCost len (len + NR) B + (4000 * (B + 2) + n * (12000 * (B + 2) + 200) + 32) +
    bandCost n NR m Φ B dmax (len + NR) + (4000 * (B + 2) + len * (12000 * (B + 2)) + 32) +
    2 * (4000 * (B + 2) + n * (14000 * (B + 2)) + 32)

end NLogN.Rsp

end

section
-- Source module: NLogN.HRspTab
/-!
# Generic table builder

`tabP T CI CK RI RK wP` runs the weight program `wP` for `i = NI-1, …, 0` (outer counter `CI`,
holding `i+1`) and `k = NK-1, …, 0` (inner counter `CK`, holding `k+1`); `wP` pushes `wf i k` on
`T`.  The result is the row-major table `tabL wf NI NK` on `T` (entry `(0,0)` on top).
-/

set_option linter.unusedSimpArgs false

namespace NLogN.Rsp

open IntMul NLogN RS

def rowL (wf : ℕ → ℕ → ℤ) (NK i : ℕ) : List ℤ := (List.range NK).map (wf i)

def tabL (wf : ℕ → ℕ → ℤ) (NI NK : ℕ) : List ℤ := ((List.range NI).map (rowL wf NK)).flatten

def tabP (T CI CK RI RK : Fin HS) (wP : RS) : RS := rep CI (.reg RI) (rep CK (.reg RK) wP)

/-- Abstract specification of a weight program. -/
structure WSpec (wP : RS) (T CI CK : Fin HS) (s0 : IState HS) (NI NK : ℕ) (wf : ℕ → ℕ → ℤ)
    (Cw : ℕ) : Prop where
  wI : CI ∉ wP.writes
  wK : CK ∉ wP.writes
  spec : ∀ (s : IState HS) (i k : ℕ) (Li Lk : List Item), i < NI → k < NK →
    s.G CI = .int ((i + 1 : ℕ) : ℤ) :: Li → s.G CK = .int ((k + 1 : ℕ) : ℤ) :: Lk →
    (∀ x, x ≠ CI → x ≠ CK → x ≠ T → s.G x = s0.G x) → s.ip = s0.ip → s.out = s0.out →
    wP.ok s ∧ wP.run s = s.upd T (.int (wf i k) :: s.G T) ∧ wP.cost s ≤ Cw

structure TDist (T CI CK RI RK : Fin HS) : Prop where
  tI : T ≠ CI
  tK : T ≠ CK
  iK : CI ≠ CK
  rI1 : RI ≠ CI
  rI2 : RI ≠ CK
  rI3 : RI ≠ T
  rK1 : RK ≠ CI
  rK2 : RK ≠ CK
  rK3 : RK ≠ T
  i5 : CI.val < 56
  k5 : CK.val < 56
  ri5 : RI.val < 56
  rk5 : RK.val < 56

def oSt (T : Fin HS) (s0 : IState HS) (wf : ℕ → ℕ → ℤ) (NI NK j : ℕ) : IState HS :=
  s0.upd T (ints (((List.range NI).map (rowL wf NK)).drop (NI - j)).flatten ++ s0.G T)

def iSt (T : Fin HS) (t0 : IState HS) (wf : ℕ → ℕ → ℤ) (NK i j : ℕ) : IState HS :=
  t0.upd T (ints ((rowL wf NK i).drop (NK - j)) ++ t0.G T)

/-- Cost of the table builder. -/
def tabCost (NI NK Cw : ℕ) : ℕ :=
  2 * 2000 * (zsz (NI : ℤ) + 2) +
    NI * ((2 * 2000 * (zsz (NK : ℤ) + 2) + NK * (Cw + 8000 * (zsz (NK : ℤ) + 2) + 30) + 32) +
      8000 * (zsz (NI : ℤ) + 2) + 30) + 32

end NLogN.Rsp

end

section
-- Source module: NLogN.HRspSapM
/-!
# The banded pass `resL` computes the real / imaginary parts of `sapL`
-/

set_option linter.unusedSimpArgs false

namespace NLogN.Rsp

open IntMul NLogN

/-- `c_k = ⌈s k / t⌉`. -/
def cS (s t k : ℕ) : ℕ := (s * k + t - 1) / t

/-- Steps of `c`: `[c_0, c_1 - c_0, …, c_{t-1} - c_{t-2}]` (`c_0 = 0`). -/
def sDel (s t : ℕ) : List ℕ := (List.range t).map fun k => NLogN.Rsp.cS s t k - NLogN.Rsp.cS s t (k - 1)

/-- Cyclic offset `a ≡ -m (mod s)`. -/
def sOff (s m : ℕ) : ℕ := (s - m % s) % s

/-- The `sTerm` weight of offset `i` for output `k`. -/
def sW (p s t α i k : ℕ) : ℤ :=
  HvdH.betaT p α * HvdH.expNeg p (α ^ 2 * t ^ 2)
    (((t : ℤ) * (((NLogN.Rsp.cS s t k : ℕ) : ℤ) - ((HvdH.sWidth p α : ℕ) : ℤ) + ((i : ℕ) : ℤ)) -
      (s : ℤ) * (k : ℤ)).natAbs ^ 2)

/-- The weight table: row `i < 2m+1`, column `k < t`. -/
def sTab (p s t α : ℕ) : List (List ℤ) :=
  (List.range (2 * HvdH.sWidth p α + 1)).map fun i => (List.range t).map fun k => NLogN.Rsp.sW p s t α i k

end NLogN.Rsp

end

section
-- Source module: NLogN.HRspSTab
/-!
# Building the `S` weight table and the step table

`sTabP` pushes `ints (sTab p s t α).flatten` on `BW`; `sDelP` pushes `ints (sDel s t)` on `BD`.
Scalar parameters (registers): `TP = p`, `TS = s`, `TT = t`, `TM = m = sWidth p α`,
`TBe = betaT p α`, `TK2 = α² t²`, `TRI = 2m+1`; counters `TCI`, `TCK`.
-/

set_option linter.unusedSimpArgs false

namespace NLogN.Rsp

open IntMul NLogN RS

abbrev TP : Fin HS := 48
abbrev TS : Fin HS := 49
abbrev TT : Fin HS := 50
abbrev TM : Fin HS := 51
abbrev TBe : Fin HS := 52
abbrev TK2 : Fin HS := 22
abbrev TRI : Fin HS := 23
abbrev TCI : Fin HS := 54
abbrev TCK : Fin HS := 53

/-- `k` (from the counter `k+1`). -/
def kE : Ex HS := .add (.reg TCK) (.lit (-1))
/-- `i` (from the counter `i+1`). -/
def iE : Ex HS := .add (.reg TCI) (.lit (-1))
/-- `c_k = (s k + t - 1) / t`. -/
def cE : Ex HS := .ndiv (.add (.mul (.reg TS) kE) (.add (.reg NLogN.Rsp.TT) (.lit (-1)))) (.reg NLogN.Rsp.TT)
/-- `d = t (c_k - m + i) - s k`. -/
def dE : Ex HS := .add (.mul (.reg NLogN.Rsp.TT) (.add (.add NLogN.Rsp.cE (.neg (.reg TM))) iE)) (.neg (.mul (.reg TS) kE))
/-- `d²`. -/
def XE : Ex HS := .mul NLogN.Rsp.dE NLogN.Rsp.dE

/-- The `S` weight program. -/
def sWP : RS :=
  NLogN.RS.seq (push 16 (.reg TP)) (NLogN.RS.seq (push 17 (.reg TK2)) (NLogN.RS.seq (push 18 NLogN.Rsp.XE) (NLogN.RS.seq expNegP
    (NLogN.RS.seq (push NLogN.Rsp.BW (.mul (.reg TBe) (.reg 19))) (NLogN.RS.seq (pop 19) (NLogN.RS.seq (pop 18) (NLogN.RS.seq (pop 17) (pop 16))))))))

/-- The weights with explicit scalar parameters. -/
def sWf (p s t m : ℕ) (beta : ℤ) (K2 : ℕ) (i k : ℕ) : ℤ :=
  beta * HvdH.expNeg p K2
    (((t : ℤ) * (((NLogN.Rsp.cS s t k : ℕ) : ℤ) - (m : ℤ) + (i : ℤ)) - (s : ℤ) * (k : ℤ)).natAbs ^ 2)

/-- Cost bound of one weight. -/
def sWCost (p Bq : ℕ) : ℕ := 100 * 10 ^ 7 * (Bq + 2) ^ 2 + expNegCost p Bq Bq

/-! ## The full `S` table -/

/-- The `S` weight table builder: pushes `ints (sTab p s t α).flatten` on `BW`
(needs `TRI = 2m+1`, `TT = t` and the scalar parameters). -/
def sTabP : RS := tabP NLogN.Rsp.BW TCI TCK TRI NLogN.Rsp.TT sWP

/-! ## The step table `sDel` -/

/-- `c_{k-1}` (for `k ≥ 1`). -/
def cE1 : Ex HS :=
  .ndiv (.add (.mul (.reg TS) (.add (.reg TCK) (.lit (-2)))) (.add (.reg NLogN.Rsp.TT) (.lit (-1)))) (.reg NLogN.Rsp.TT)

/-- `c_k - c_{k-1}` (`0` for `k = 0`). -/
def dlE : Ex HS := .sel (.add (.reg TCK) (.lit (-2))) (.lit 0) (.add NLogN.Rsp.cE (.neg NLogN.Rsp.cE1))

def dWP : RS := push BD dlE

def dWf (s t : ℕ) (_i k : ℕ) : ℤ := ((NLogN.Rsp.cS s t k - NLogN.Rsp.cS s t (k - 1) : ℕ) : ℤ)

/-- The step-table builder: pushes `ints ((sDel s t).map cast)` on `BD` (needs `TS = s`, `TT = t`). -/
def dTabP : RS := NLogN.RS.seq (push TRI (.lit 1)) (NLogN.RS.seq (tabP BD TCI TCK TRI NLogN.Rsp.TT dWP) (pop TRI))

end NLogN.Rsp

end

section
-- Source module: NLogN.HRspSplit
/-!
# Splitting a Gaussian list into real / imaginary integer lists and merging back
-/

set_option linter.unusedSimpArgs false

namespace NLogN.Rsp

open IntMul NLogN RS

/-- Move one item from `X` to `Y`. -/
def mv1 (X Y : Fin HS) : RS := NLogN.RS.seq (push Y (.reg X)) (pop X)

/-! ## split -/

def splitP (X A B C : Fin HS) (n : Ex HS) : RS := rep C n (NLogN.RS.seq (NLogN.Rsp.mv1 X A) (NLogN.Rsp.mv1 X B))

def spSt (X A B : Fin HS) (s : IState HS) (U : List (ℤ × ℤ)) (R : List Item) (j : ℕ) :
    IState HS :=
  ((s.upd X (gints (U.drop j) ++ R)).upd A (ints ((U.take j).map Prod.fst).reverse ++ s.G A)).upd B
    (ints ((U.take j).map Prod.snd).reverse ++ s.G B)

structure Dist4 (X A B C : Fin HS) : Prop where
  xa : X ≠ A
  xb : X ≠ B
  xc : X ≠ C
  ab : A ≠ B
  ac : A ≠ C
  bc : B ≠ C
  x5 : X.val < 56
  a5 : A.val < 56
  b5 : B.val < 56
  c5 : C.val < 56

/-! ## merge -/

def mergeP (X A B C : Fin HS) (n : Ex HS) : RS := rep C n (NLogN.RS.seq (NLogN.Rsp.mv1 B X) (NLogN.Rsp.mv1 A X))

def mgSt (X A B : Fin HS) (s : IState HS) (LA LB : List ℤ) (RA RB : List Item) (j : ℕ) :
    IState HS :=
  ((s.upd A (ints (LA.drop j) ++ RA)).upd B (ints (LB.drop j) ++ RB)).upd X
    (gints ((LA.zip LB).take j).reverse ++ s.G X)

end NLogN.Rsp

end

section
-- Source module: NLogN.HRspGres
/-!
# The banded pass on a Gaussian list (both components)
-/

set_option linter.unusedSimpArgs false

namespace NLogN.Rsp

open IntMul NLogN RS

abbrev SA : Fin HS := 33
abbrev SB : Fin HS := 34
abbrev SC : Fin HS := 35

/-- Componentwise banded pass on the Gaussian list on `CU` (length on `Ks`, output length on `BN`). -/
def gresP : RM :=
  RM.seqL [.rs (splitP CU NLogN.Rsp.SA NLogN.Rsp.SB SC (.reg Ks)), .rs (mvK NLogN.Rsp.SA CU SC (.reg Ks)), resP,
    .rs (mvK CU NLogN.Rsp.SA SC (.reg BN)), .rs (mvK NLogN.Rsp.SB CU SC (.reg Ks)), resP,
    .rs (mvK CU NLogN.Rsp.SB SC (.reg BN)), .rs (mergeP CU NLogN.Rsp.SA NLogN.Rsp.SB SC (.reg BN))]

/-- The list computed by `gresP`. -/
def gresL (R : ℤ) (Dl : List ℕ) (Ws : List (List ℤ)) (U : List (ℤ × ℤ)) (a n : ℕ) :
    List (ℤ × ℤ) :=
  (resL R Dl Ws (U.map Prod.fst) a n).zip (resL R Dl Ws (U.map Prod.snd) a n)

/-- Cost bound of `gresP`. -/
def gresCost (len NR n m Φ B dmax : ℕ) : ℕ :=
  2 * resCost len NR n m Φ B dmax + 3 * (4000 * (B + 2) + len * (20000 * (B + 2)) + 32) +
    3 * (4000 * (B + 2) + n * (20000 * (B + 2)) + 32)

end NLogN.Rsp

end

section
-- Source module: NLogN.HRspSap
/-!
# The `sapL` slice program
-/

set_option linter.unusedSimpArgs false

namespace NLogN.Rsp

open IntMul NLogN RS

/-- The `sapL` slice program (= `gresP` run against the `S` tables). -/
abbrev sapP : RM := gresP

end NLogN.Rsp

end

section
-- Source module: NLogN.HRspEM
/-!
# The banded pass `resL` computes `Eapprox` and `Dapprox` (componentwise)
-/

set_option linter.unusedSimpArgs false

namespace NLogN.Rsp

open IntMul NLogN

/-- Steps `[0, 1, …, 1]` (length `s`): output `l` reads from position `l`. -/
def eDel (s : ℕ) : List ℕ := (List.range s).map fun l => if l = 0 then 0 else 1

/-- The `eTerm` weight of offset `i` (`h = i - m`) for output `l` (`0` for `h = 0`). -/
def eW (p s t α i l : ℕ) : ℤ :=
  if i = HvdH.eWidth p α then 0 else
    HvdH.expNeg p (s ^ 2) (α ^ 2 * (HvdH.eKz s t l ((i : ℤ) - (HvdH.eWidth p α : ℤ))).toNat)

def eTab (p s t α : ℕ) : List (List ℤ) :=
  (List.range (2 * HvdH.eWidth p α + 1)).map fun i => (List.range s).map fun l => NLogN.Rsp.eW p s t α i l

/-- The list form of `Eapprox` on a slice of length `s`. -/
def eapL (p s t α : ℕ) (u : List (ℤ × ℤ)) : List (ℤ × ℤ) :=
  (List.range s).map (HvdH.Eapprox p s t α fun j => u.getD j (0, 0))

/-! ## `Dapprox` -/

/-- The single-row table of `dTilde`. -/
def dTab (p s t α : ℕ) : List (List ℤ) := [(List.range s).map (HvdH.dTilde p s t α)]

def dapL (p s t α : ℕ) (u : List (ℤ × ℤ)) : List (ℤ × ℤ) :=
  (List.range s).map (HvdH.Dapprox p s t α fun j => u.getD j (0, 0))

end NLogN.Rsp

end

section
-- Source module: NLogN.HRspJ1
/-!
# Pointwise loops on integer lists: halving, signed accumulation
-/

set_option linter.unusedSimpArgs false

namespace NLogN.Rsp

open IntMul NLogN RS

/-! ## halving -/

def hBody (X T : Fin HS) : RS := NLogN.RS.seq (push T (.quo (.reg X) (.lit 2))) (pop X)

def halfLoop (X T C : Fin HS) (n : Ex HS) : RS := rep C n (hBody X T)

def hSt (X T : Fin HS) (s : IState HS) (L : List ℤ) (R : List Item) (j : ℕ) : IState HS :=
  (s.upd X (ints (L.drop j) ++ R)).upd T (ints ((L.take j).map fun x => x.tdiv 2).reverse ++ s.G T)

/-- Halve every entry (rounding towards zero) of the `n`-list on `X` (`n` on `K`). -/
def halfP (X T C K : Fin HS) : RS := NLogN.RS.seq (halfLoop X T C (.reg K)) (mvK T X C (.reg K))

/-! ## signed accumulation `A += ±X` -/

def sgF (ng : Bool) (a w : ℤ) : ℤ := a + (if ng then -w else w)

def sgV (ng : Bool) : ℤ := if ng then -1 else 1

def aE (X A Sg : Fin HS) : Ex HS := .add (.reg A) (.sel (.reg Sg) (.neg (.reg X)) (.reg X))

def aBody (X A Sg T W : Fin HS) : RS :=
  NLogN.RS.seq (push T (NLogN.Rsp.aE X A Sg)) (NLogN.RS.seq (pop A) (NLogN.RS.seq (push W (.reg X)) (pop X)))

structure Dist6 (X A Sg T W C : Fin HS) : Prop where
  d : [X, A, Sg, T, W, C].Nodup
  lt : ∀ i ∈ [X, A, Sg, T, W, C], i.val < 56

def aSt (X A T W : Fin HS) (s : IState HS) (ng : Bool) (Lw La : List ℤ) (RX RA : List Item)
    (j : ℕ) : IState HS :=
  (((s.upd X (ints (Lw.drop j) ++ RX)).upd A (ints (La.drop j) ++ RA)).upd T
    (ints ((List.zipWith (sgF ng) La Lw).take j).reverse ++ s.G T)).upd W
    (ints (Lw.take j).reverse ++ s.G W)

/-- `A += ±X` pointwise (sign from `Sg`: `-1` subtracts), lists of length `n` on `K`;
`X` is restored. -/
def addP (X A Sg T W C K : Fin HS) : RS :=
  NLogN.RS.seq (rep C (.reg K) (aBody X A Sg T W)) (NLogN.RS.seq (mvK T A C (.reg K)) (mvK W X C (.reg K)))

/-- Cost bound of `addP`. -/
def addCost (n B : ℕ) : ℕ := 3 * (4000 * (B + 2) + n * (40000 * (B + 2)) + 32)

end NLogN.Rsp

end

section
-- Source module: NLogN.HRspJM
/-!
# List form of `Japprox`: iterates `vL`, alternating partial sums `accL`
-/

set_option linter.unusedSimpArgs false

namespace NLogN.Rsp

open IntMul NLogN

def halfL (u : List (ℤ × ℤ)) : List (ℤ × ℤ) := u.map fun z => (z.1.tdiv 2, z.2.tdiv 2)

/-- Pointwise `a ± w` on Gaussian lists. -/
def gsgF (ng : Bool) (a w : ℤ × ℤ) : ℤ × ℤ := (sgF ng a.1 w.1, sgF ng a.2 w.2)

/-- The iterates `ṽ⁽ʲ⁾` as lists of length `s`. -/
def vL (p s t α : ℕ) (u : List (ℤ × ℤ)) : ℕ → List (ℤ × ℤ)
  | 0 => halfL u
  | j + 1 => eapL p s t α (vL p s t α u j)

/-- Alternating partial sums. -/
def accL (p s t α : ℕ) (u : List (ℤ × ℤ)) (n : ℕ) : List (ℤ × ℤ) :=
  (List.range s).map fun l =>
    (∑ j ∈ Finset.range n, (-1) ^ j * ((vL p s t α u j).getD l (0, 0)).1,
     ∑ j ∈ Finset.range n, (-1) ^ j * ((vL p s t α u j).getD l (0, 0)).2)

/-! ## agreement with `vIter` / `Japprox` -/

/-- The list form of `Japprox`. -/
def japL (p s t α : ℕ) (u : List (ℤ × ℤ)) : List (ℤ × ℤ) :=
  (List.range s).map (HvdH.Japprox p s t α fun j => u.getD j (0, 0))

end NLogN.Rsp

end

section
-- Source module: NLogN.HRspJ
/-!
# `Japprox` on a slice: the `E` pass (`gresP` against the `E` tables) iterated with signed
accumulation
-/

set_option linter.unusedSimpArgs false

namespace NLogN.Rsp

open IntMul NLogN RS

/-- The registers / tables that `gresP` reads, set for the `E` pass. -/
structure EParams (σ : IState HS) (p s t α : ℕ) : Prop where
  ks : ∃ L, σ.G Ks = .int (s : ℤ) :: L
  ka : ∃ L, σ.G Ka = .int (sOff s (HvdH.eWidth p α) : ℤ) :: L
  kn : ∃ L, σ.G KN = .int ((s + (2 * HvdH.eWidth p α + 1) : ℕ) : ℤ) :: L
  br : ∃ L, σ.G BR = .int (2 ^ p) :: L
  bn : ∃ L, σ.G BN = .int (s : ℤ) :: L
  bnr : ∃ L, σ.G BNR = .int ((2 * HvdH.eWidth p α + 1 : ℕ) : ℤ) :: L
  bnw : ∃ L, σ.G BNW = .int ((s * (2 * HvdH.eWidth p α + 1) : ℕ) : ℤ) :: L
  bw : ∃ R, σ.G NLogN.Rsp.BW = ints (eTab p s t α).flatten ++ R
  bd : ∃ R, σ.G BD = ints ((eDel s).map (Nat.cast : ℕ → ℤ)) ++ R

/-- The stacks `EParams` depends on. -/
def eRegs : List (Fin HS) := [Ks, Ka, KN, BR, BN, BNR, BNW, NLogN.Rsp.BW, BD]

/-! ## The `J` loop -/

abbrev JA : Fin HS := 36
abbrev JT : Fin HS := 37
abbrev JW : Fin HS := 38
abbrev JS : Fin HS := 39
abbrev JC : Fin HS := 40
abbrev JN : Fin HS := 41
abbrev JK : Fin HS := 42
abbrev JL : Fin HS := 43

def jBody : RM :=
  RM.seqL [.rs (addP CU JA JS JT JW JK JL), gresP, .rs (asg JS (.neg (.reg JS)))]

def jSt (s1 : IState HS) (p s t α : ℕ) (u : List (ℤ × ℤ)) (RU RA LS : List Item) (j : ℕ) :
    IState HS :=
  ((s1.upd CU (gints (vL p s t α u j) ++ RU)).upd JA (ints (pairInts (accL p s t α u j)) ++ RA)).upd
    JS (.int (sgV (decide (j % 2 = 1))) :: LS)

/-- Cost of one round of the `J` loop. -/
def jRoundCost (s NR mm Φ B : ℕ) : ℕ :=
  addCost (2 * s) B + (gresCost s NR s mm Φ B 1 + 3 * 2000 * (B + 2))

/-- `Japprox` on the Gaussian slice on `CU` (length `s`; `JL = 2s`, `JN = n` rounds, `E`
parameters set). -/
def JP : RM :=
  RM.seqL [.rs (halfP CU JT JK JL), .rs (push JS (.lit 1)), .rs (NLogN.Rsp.zerosP JA JK (.reg JL)),
    .rep JC (.reg JN) jBody, .rs (popK CU JK (.reg JL)), .rs (mvK JA JT JK (.reg JL)),
    .rs (mvK JT CU JK (.reg JL)), .rs (pop JS)]

/-- Cost bound of `JP`. -/
def jCost (s NR n mm Φ B : ℕ) : ℕ :=
  10 * (4000 * (B + 2) + (2 * s) * (30000 * (B + 2)) + 32) +
    n * (jRoundCost s NR mm Φ B + 8000 * (B + 2) + 30)

end NLogN.Rsp

end

section
-- Source module: NLogN.HRspStage
/-!
# A `gresP` stage with its scalar parameters pushed from expressions (and popped afterwards)
-/

set_option linter.unusedSimpArgs false

namespace NLogN.Rsp

open IntMul NLogN RS

/-! ## congruence of expressions -/

/-! ## the stage -/

/-- Expressions for the seven scalar parameters of `gresP`. -/
structure SArgs where
  eKs : Ex HS
  eKa : Ex HS
  eKN : Ex HS
  eBR : Ex HS
  eBN : Ex HS
  eBNR : Ex HS
  eBNW : Ex HS

def SArgs.l (A : SArgs) : List (Ex HS) := [A.eKs, A.eKa, A.eKN, A.eBR, A.eBN, A.eBNR, A.eBNW]

def stTg : List (Fin HS) := [Ks, Ka, KN, BR, BN, BNR, BNW]

def pushesP (A : SArgs) : RS :=
  NLogN.RS.seq (push Ks A.eKs) (NLogN.RS.seq (push Ka A.eKa) (NLogN.RS.seq (push KN A.eKN) (NLogN.RS.seq (push BR A.eBR)
    (NLogN.RS.seq (push BN A.eBN) (NLogN.RS.seq (push BNR A.eBNR) (push BNW A.eBNW))))))

def popsP : RS :=
  NLogN.RS.seq (pop BNW) (NLogN.RS.seq (pop BNR) (NLogN.RS.seq (pop BN) (NLogN.RS.seq (pop BR) (NLogN.RS.seq (pop KN) (NLogN.RS.seq (pop Ka) (pop Ks))))))

/-- `gresP` with its parameters pushed from `A` (and popped afterwards). -/
def stageP (A : SArgs) : RM := .seq (.rs (pushesP A)) (.seq gresP (.rs popsP))

/-- Expressions read no target register. -/
def SArgs.Good (A : SArgs) : Prop := ∀ e ∈ A.l, ∀ i ∈ e.regs, i ∉ stTg ∧ i.val < 56

/-- The state after the pushes. -/
def pSt (s0 : IState HS) (v : Fin 7 → ℤ) : IState HS :=
  ((((((s0.upd Ks (.int (v 0) :: s0.G Ks)).upd Ka (.int (v 1) :: s0.G Ka)).upd KN
    (.int (v 2) :: s0.G KN)).upd BR (.int (v 3) :: s0.G BR)).upd BN (.int (v 4) :: s0.G BN)).upd
    BNR (.int (v 5) :: s0.G BNR)).upd BNW (.int (v 6) :: s0.G BNW)

def SArgs.v (A : SArgs) (s0 : IState HS) : Fin 7 → ℤ :=
  ![A.eKs.eval s0, A.eKa.eval s0, A.eKN.eval s0, A.eBR.eval s0, A.eBN.eval s0, A.eBNR.eval s0,
    A.eBNW.eval s0]

/-- The parameter expressions are linear, clean, defined and `B`-bounded at `s0`. -/
def SArgs.Fine (A : SArgs) (s0 : IState HS) (B : ℕ) : Prop :=
  ∀ e ∈ A.l, e.Clean ∧ e.Lin ∧ e.ok s0 ∧ e.Bnd s0 B ∧ e.size ≤ 3

end NLogN.Rsp

end

section
-- Source module: NLogN.HRspBM
/-!
# List form of `bapL`: `ptL`, `cselL`, `japL`, `dapL`, `psL`
-/

set_option linter.unusedSimpArgs false

namespace NLogN.Rsp

open IntMul NLogN

/-! ## `C̃`: the gather `r_l = round(t l / s)` -/

/-- `r_l = ⌊(2 t l + s) / (2 s)⌋`. -/
def cR (s t l : ℕ) : ℕ := (2 * t * l + s) / (2 * s)

/-- Steps of `r`: `[r_0, r_1 - r_0, …]` (`r_0 = 0`). -/
def cDel (s t : ℕ) : List ℕ := (List.range s).map fun l => cR s t l - cR s t (l - 1)

/-- The (single-row) weight table of `C̃`: all ones. -/
def cTab (s : ℕ) : List (List ℤ) := [List.replicate s 1]

/-- List form of `C̃` (input length `t`, output length `s`). -/
def cselL (s t : ℕ) (v : List (ℤ × ℤ)) : List (ℤ × ℤ) :=
  (List.range s).map fun l => v.getD (cR s t l) (0, 0)

/-! ## permutations as key lists for `permP` -/

def gRec (z : ℤ × ℤ) : List ℤ := [z.1, z.2]

/-- `P̃_t` as a list: `(P̃_t V)_k = V_{-sk mod t}`. -/
def ptL (s t : ℕ) (V : List (ℤ × ℤ)) : List (ℤ × ℤ) :=
  (List.range t).map fun (k : ℕ) => V.getD (-((s : ZMod t) * (k : ZMod t))).val (0, 0)

/-- Target index of record `j` for `P̃_t`: `-s⁻¹ j mod t`. -/
def ptKeys (s t : ℕ) : List ℕ :=
  (List.range t).map fun (j : ℕ) => (-((s : ZMod t)⁻¹ * (j : ZMod t))).val

/-- `P̃_s⁻¹` as a list: `(P̃_s⁻¹ W)_j = W_{t⁻¹ j mod s}`. -/
def psL (s t : ℕ) (W : List (ℤ × ℤ)) : List (ℤ × ℤ) :=
  (List.range s).map fun (j : ℕ) => W.getD ((t : ZMod s)⁻¹ * (j : ZMod s)).val (0, 0)

/-- Target index of record `i` for `P̃_s⁻¹`: `t i mod s`. -/
def psKeys (s t : ℕ) : List ℕ := (List.range s).map fun i => t * i % s

/-! ## the composition -/

end NLogN.Rsp

end

section
-- Source module: NLogN.HRspBap
/-!
# The `bapL` slice program

`bapP = P̃_s⁻¹ ∘ D̃' ∘ J̃' ∘ C̃ ∘ P̃_t` on the slice stack `CU`:
* `P̃_t`, `P̃_s⁻¹`: `permP 2` relocated (stack `0 ↔ CU`, key stack `3 ↔ KS`);
* `C̃`: a `gresP` stage with one all-ones weight row (tables relocated `BW ↔ W1`, `BD ↔ D1`);
* `J̃'`: `JP` (the `E` parameters / tables are the native ones);
* `D̃'`: a `gresP` stage with the diagonal table (relocated `BW ↔ W3`).
-/

set_option linter.unusedSimpArgs false

namespace NLogN.Rsp

open IntMul NLogN RS

/-! ## relocated stages -/

/-- `σ` fixes `CU`, the parameter stacks and every register read by `A`. -/
def StFix (σ : Equiv.Perm (Fin HS)) (A : SArgs) : Prop :=
  σ CU = CU ∧ (∀ i ∈ stTg, σ i = i) ∧ ∀ e ∈ A.l, ∀ i ∈ e.regs, σ i = i

/-! ## registers and the program -/

/-- `s` (slice output length). -/
abbrev PS : Fin HS := 44
/-- `2^p`. -/
abbrev PR : Fin HS := 24
/-- The `P̃_s⁻¹` key list. -/
abbrev KS : Fin HS := 47
/-- The `C̃` weight table (all ones). -/
abbrev W1 : Fin HS := 46
/-- The `C̃` step table. -/
abbrev D1 : Fin HS := 45
/-- The `D̃'` weight table. -/
abbrev W3 : Fin HS := 0

/-- Parameters of the `C̃` stage: input length `t` (stack `1`), offset `0`, one row, `R = 1`. -/
def Ac : SArgs := ⟨.reg 1, .lit 0, .add (.reg 1) (.lit 1), .lit 1, .reg PS, .lit 1, .reg PS⟩

/-- Parameters of the `D̃'` stage. -/
def Ad : SArgs := ⟨.reg PS, .lit 0, .add (.reg PS) (.lit 1), .reg NLogN.Rsp.PR, .reg PS, .lit 1, .reg PS⟩

def σt : Equiv.Perm (Fin HS) := swapsPerm [(0, CU)]
def σs : Equiv.Perm (Fin HS) := swapsPerm [(0, CU), (3, KS)]
def σc : Equiv.Perm (Fin HS) := swapsPerm [(NLogN.Rsp.BW, W1), (BD, D1)]
def σd : Equiv.Perm (Fin HS) := swapsPerm [(NLogN.Rsp.BW, W3)]

/-- `P̃_t`: `permP` with the records on `CU`, `n = t` on `1`, `kb` on `2`, keys on `3`. -/
def ptP : RS := (permP 2).ren σt

/-- `P̃_s⁻¹`: `permP` with the records on `CU`, `n = s` pushed on `1`, `kb` on `2`, keys on `KS`. -/
def psP : RS := NLogN.RS.seq (push 1 (.reg PS)) (NLogN.RS.seq ((permP 2).ren σs) (pop 1))

/-- **The `bapL` slice program.** -/
def bapP : RM := RM.seqL [.rs ptP, (stageP Ac).ren σc, JP, (stageP Ad).ren σd, .rs psP]

/-! ## the precondition -/

/-- Parameter stacks and tables of `bapP` (all slice-independent, left intact). -/
structure BParams (σ : IState HS) (p s t α kb : ℕ) : Prop where
  pt : ∃ L, σ.G 1 = .int (t : ℤ) :: L
  pk : ∃ L, σ.G 2 = .int (kb : ℤ) :: L
  kt : ∃ R, σ.G 3 = natInts (ptKeys s t) ++ R
  ks : ∃ R, σ.G KS = natInts (psKeys s t) ++ R
  ps : ∃ L, σ.G PS = .int (s : ℤ) :: L
  pr : ∃ L, σ.G NLogN.Rsp.PR = .int (2 ^ p) :: L
  e : EParams σ p s t α
  jl : ∃ L, σ.G JL = .int ((2 * s : ℕ) : ℤ) :: L
  jn : ∃ L, σ.G JN = .int (HvdH.jIter p s t α : ℤ) :: L
  w1 : ∃ R, σ.G W1 = ints (cTab s).flatten ++ R
  d1 : ∃ R, σ.G D1 = ints ((cDel s t).map (Nat.cast : ℕ → ℤ)) ++ R
  w3 : ∃ R, σ.G W3 = ints (dTab p s t α).flatten ++ R

/-! ## run -/

/-! ## ok and cost -/

/-- Cost of one relocated `gresP` stage (pushes + `gresP` + pops). -/
def stCost (len NR n mm Φ B dmax : ℕ) : ℕ :=
  7 * (4 * 2000 * (B + 2)) + gresCost len NR n mm Φ B dmax + 7 * (2 * B + 2)

/-- Cost bound of `bapP`: two permutations, the `C̃` and `D̃'` stages and the `J̃'` loop. -/
def bapCost (p s t α kb mm Φ B : ℕ) : ℕ :=
  10 ^ 6 * (kb + 1) * (t + 1) * (2 + 1) * (B + kb + 2 + 4) +
  stCost t 1 s mm Φ B (t / s + 1) +
  jCost s (2 * HvdH.eWidth p α + 1) (HvdH.jIter p s t α) mm Φ B +
  stCost s 1 s mm Φ B 1 +
  (2 * 2000 * (B + 2) + 10 ^ 6 * (kb + 1) * (s + 1) * (2 + 1) * (B + kb + 2 + 4) + (2 * B + 2))

end NLogN.Rsp

end

section
-- Source module: NLogN.HRspBTab
/-!
# Table builders for `bapP`

One-row tables built by `row1P T RK wP` (`TRI := 1` pushed, `tabP`, popped), with a single-push
weight program `push T e`:
* `cTabP` : `ints (cTab s).flatten` (all ones) on `W1`;
* `eDelP` : `ints (eDel s)` on `BD`;
* `cDelP` : `ints (cDel s t)` on `D1`;
* `psKeysP` : `natInts (psKeys s t)` on `KS`;
* `ptKeysP` : `natInts (ptKeys s t)` on `3` (needs `TU = (s⁻¹ mod t)`).
Scalars: `TS = s`, `TT = t`, `TU`; counters `TCI`, `TCK`.
-/

set_option linter.unusedSimpArgs false

namespace NLogN.Rsp

open IntMul NLogN RS

/-- `s⁻¹ mod t` (for the `P̃_t` keys). -/
abbrev TU : Fin HS := 29

/-! ## generic one-row builder -/

def row1P (T RK : Fin HS) (wP : RS) : RS :=
  NLogN.RS.seq (push TRI (.lit 1)) (NLogN.RS.seq (tabP T TCI TCK TRI RK wP) (pop TRI))

/-! ## `cTab`: all ones -/

def cTabP : RS := row1P W1 TS (push W1 (.lit 1))

/-! ## `eDel` -/

def eDelE : Ex HS := .sel (.add (.reg TCK) (.lit (-2))) (.lit 0) (.lit 1)

def eDelP : RS := row1P BD TS (push BD eDelE)

/-! ## `cDel` -/

def kE1 : Ex HS := .add (.reg TCK) (.lit (-2))
def cRE : Ex HS := .ndiv (.add (.mul (.add (.reg NLogN.Rsp.TT) (.reg NLogN.Rsp.TT)) kE) (.reg TS)) (.add (.reg TS) (.reg TS))
def cRE1 : Ex HS :=
  .ndiv (.add (.mul (.add (.reg NLogN.Rsp.TT) (.reg NLogN.Rsp.TT)) kE1) (.reg TS)) (.add (.reg TS) (.reg TS))
def cDelE : Ex HS := .sel kE1 (.lit 0) (.add cRE (.neg cRE1))

def cDelP : RS := row1P D1 TS (push D1 cDelE)

/-! ## the permutation keys -/

def psKE : Ex HS := .nmod (.mul (.reg NLogN.Rsp.TT) kE) (.reg TS)

def psKeysP : RS := row1P KS TS (push KS psKE)

def ptKE : Ex HS := .nmod (.add (.reg NLogN.Rsp.TT) (.neg (.nmod (.mul (.reg TU) kE) (.reg NLogN.Rsp.TT)))) (.reg NLogN.Rsp.TT)

def ptKeysP : RS := row1P 3 NLogN.Rsp.TT (push 3 ptKE)

end NLogN.Rsp

end