κ-improvement development, part 29: item-language primitives

-- OpenAI-preprint / κ-improvement development (integer-multiplication-kappa), generated by skeleton subtraction.

import Definitions.Def_IntMul_Kappa_K_Defs28

set_option autoImplicit false
set_option relaxedAutoImplicit false
set_option maxHeartbeats 1000000

section
-- Source module: Solutions.KFV10Rec
/-!
# KFV §14: the uniform record product `recU`

`recU pM` is `recI pM b r Δ c KG` (`KFV6Rec`) with every size read at run time from six
parameter stacks (`kc = 50 … kD = 55`):

| stack | value |
|---|---|
| `kc = 50` | `c = 2^p` |
| `kr = 51` | `r` |
| `kK = 52` | `KG = kG b r c` |
| `kB = 53` | `2^b` |
| `kJ = 54` | `(2^b)^(2r)` |
| `kD = 55` | `2^Δ` |

Derived literals are expressions over them: `2r = r + r`, `-KG`, `H = 2^b / 2`, `-H`.
The pieces are related to those of `recI` by `ISim` (`KFV8Sim`/`KFV9Gen`).
-/

set_option linter.unusedSectionVars false
set_option linter.unusedVariables false
set_option linter.unusedSimpArgs false

namespace NLogN.KFV

open IntMul KIL KFD KFT KFB KFC RS

variable {Pp : ℕ} {I : List Sym} {iprog : Fin Pp → IStmt HS Pp}

/-! ## Parameters -/

abbrev kc : Fin HS := 50
abbrev kr : Fin HS := 51
abbrev kK : Fin HS := 52
abbrev kB : Fin HS := 53
abbrev kJ : Fin HS := 54
abbrev kD : Fin HS := 55

/-- The parameter map of a record product. -/
def πR (c r KG B BJ D2 : ℤ) : PMap :=
  pmOf [(kc, c), (kr, r), (kK, KG), (kB, B), (kJ, BJ), (kD, D2)]

/-! ## Literal expressions -/

/-- `2r = r + r`. -/
def exTwoR : Ex HS := .add (.reg kr) (.reg kr)
/-- `H = 2^b / 2`. -/
def exHalf : Ex HS := .quo (.reg kB) (.lit 2)

/-! ## Write conditions -/

/-! ## Uniform pieces -/

def unzBU (X SA SB : Fin HS) : RS :=
  .seq (.push SA (.add (.reg X) (.reg kc))) (.seq (.push SA (.lit 0)) (.seq (.pop X)
    (.seq (.push SB (.add (.reg X) (.reg kc))) (.seq (.push SB (.lit 0)) (.pop X)))))

def unzU (X SA SB C : Fin HS) : RS := .rep C (.reg kr) (unzBU X SA SB)

/-- `rpt` with the count read from stack `R`. -/
def rptR (R A C D : Fin HS) (b : IStmt HS Pp) : IStmt HS Pp :=
  .seq (.op (.dup R A)) (.seq (pro A C D) (.seq (uloop C D b) (.seq (epi C D) (.op (.drop A)))))

/-- `rpt` with the count `e`. -/
def rptE (e : Ex HS) (A C D : Fin HS) (b : IStmt HS Pp) : IStmt HS Pp :=
  .seq (pushE e A) (.seq (pro A C D) (.seq (uloop C D b) (.seq (epi C D) (.op (.drop A)))))

/-- `rpt` with count `0`. -/
def rptZ (A C D : Fin HS) (b : IStmt HS Pp) : IStmt HS Pp :=
  .seq (.op (.const A 0)) (.seq (pro A C D) (.seq (uloop C D b) (.seq (epi C D) (.op (.drop A)))))

/-- Uniform `carI b r`. -/
def carU : IStmt HS Pp :=
  .seq (.op (.dup kB tR)) (.seq (.op (.const tY 0)) (.seq (rptR kr tA tC NLogN.KFT.tD c1B)
    (.seq (.op (.const NLogN.KFT.tP 1)) (.seq (.op (.add tY NLogN.KFT.tP tY)) (.seq (.op (.const NLogN.KFT.tP 1))
    (.seq (rptR kr tA tC NLogN.KFT.tD c2B) (.seq (.op (.neg NLogN.KFT.tP)) (.seq (.op (.add tY NLogN.KFT.tP 0))
    (.op (.drop tR))))))))))

/-- Uniform `digI b 0 (2r) (2r)`. -/
def digU : IStmt HS Pp :=
  .seq (.op (.dup kB tR)) (.seq (rptE exTwoR tA tC NLogN.KFT.tD cutB) (.seq (.op (.drop tX))
    (.seq (rptZ tA tC NLogN.KFT.tD KFT.zB) (.seq (.op (.const NLogN.KFT.tE (2 ^ 0)))
    (.seq (rptE exTwoR tA tC NLogN.KFT.tD putB) (.seq (.op (.drop NLogN.KFT.tE)) (.op (.drop tR))))))))

/-- Boolean check: no stack of `50 … 55`. -/
def allLow (L : List (Fin HS)) : Bool := L.all fun i => decide (i.val < 50 ∨ 55 < i.val)

/-! ## Packing, extraction and interleaving -/

/-- Uniform `packI S P b r KG`. -/
def packU (S P : Fin HS) : IStmt HS Pp :=
  .seq (fcode (NLogN.mvK S 0 NLogN.KFV.vL exTwoR))
    (.seq carU (fcode (.seq (.push P (.add (.reg 0) (.neg (.reg kK)))) (.pop 0))))

def xHU : Ex HS := .add (.add (.reg 0) (.reg vRH)) exHalf

def carBU : RS :=
  .seq (.push NLogN.KFV.vT (.add (.rem xHU (.reg kB)) (.neg exHalf)))
    (.seq (.asg vRH (.quo xHU (.reg kB))) (.seq (.pop 0) (.pop 0)))

def carLU : RS := .seq (.push vRH (.lit 0)) (.seq (.rep NLogN.KFV.vL exTwoR carBU) (.pop vRH))

def fldBU (O : Fin HS) : RS :=
  .seq (.push O (.quo (.add (.reg NLogN.KFV.vT) (.neg (.reg vT4))) (.reg kD))) (.seq (.pop NLogN.KFV.vT) (.pop vT4))

def fldPU (O : Fin HS) : RS := .rep NLogN.KFV.vL (.reg kr) (fldBU O)

/-- Uniform `extI e O b r Δ`. -/
def extU (e : Ex HS) (O : Fin HS) : IStmt HS Pp :=
  .seq (fcode (.push vZ (.sel e (.add e (.reg kJ)) e))) <|
  .seq (.op (.mv vZ tX)) <|
  .seq digU <|
  .seq (fcode carLU) <|
  .seq (fcode (NLogN.mvK NLogN.KFV.vT vT2 NLogN.KFV.vL (.reg kr))) <|
  .seq (fcode (NLogN.mvK vT2 vT4 NLogN.KFV.vL (.reg kr))) <|
  fcode (fldPU O)

def ilvU : RS := .rep NLogN.KFV.vL (.reg kr) ilvB

/-- **The uniform record product.** -/
def recU (pM : Fin Pp) : IStmt HS Pp :=
  .seq (fcode (unzU NLogN.KFV.vX vA NLogN.KFV.vB NLogN.KFV.vL)) <| .seq (fcode (unzU NLogN.KFV.vY vC NLogN.KFV.vE NLogN.KFV.vL)) <|
  .seq (packU vA vPA) <| .seq (packU NLogN.KFV.vB vPB) <|
  .seq (packU vC vPC) <| .seq (packU NLogN.KFV.vE vPE) <|
  .seq (KFM.mulC pM vM1 vPA vPC) <| .seq (KFM.mulC pM vM2 vPB vPE) <|
  .seq (KFM.mulC pM vM3 vPA vPE) <| .seq (KFM.mulC pM vM4 vPB vPC) <|
  .seq (extU subE vOR) <| .seq (extU addE vOI) <|
  .seq (fcode ilvU) <| fcode dropsP

section Sims

variable {c KG : ℤ} {b r : ℕ} {BJ D2 : ℤ}

/-- The parameter map of a record product with the given sizes. -/
abbrev πr (c KG : ℤ) (b r : ℕ) (BJ D2 : ℤ) : PMap := πR c r KG ((2 : ℤ) ^ b) BJ D2

end Sims

/-! ## The run -/

/-- Parameter stacks are empty in the managed part. -/
def PZ (π : PMap) (g : Fin HS → List ℤ) : Prop := ∀ i v, π i = some v → g i = []

/-- Boolean check: a register program writes only stacks below `50`. -/
def lowW (p : RS) : Bool := p.writes.all fun j => decide (j.val < 50)

end NLogN.KFV

end

section
-- Source module: Solutions.KFV11Kron
/-!
# KFV §15: the uniform pointwise Kronecker stage `kronU`

`kronU pM` multiplies two arrays of `N` records (`r` Gaussian integers each) record by record:
stack `0` holds `gflat zf ++ gflat zg ++ l`; afterwards it holds `gflat (kmulN r p ρ N zf zg) ++ l`,
where record `j` of `kmulN` is `kmul r p ρ` of records `j` of `zf` and `zg` (`KFV6Math`).

The code is fixed: the six record-product parameters sit on the frame stacks `50 … 55`
(`KFV10Rec`), `N` on the frame of `vRH` and `2Nr` on the frame of `vT`; the record loop is a unary
loop on the frame of `vZ` (`toUL`). **`kronU_runs`**.
-/

set_option linter.unusedSectionVars false
set_option linter.unusedVariables false
set_option linter.unusedSimpArgs false

namespace NLogN.KFV

open IntMul KIL KFD KFT KFB KFC RS

variable {Pp : ℕ} {I : List Sym} {iprog : Fin Pp → IStmt HS Pp}

/-! ## Records -/

/-- Record `j` (of length `r`) of an array. -/
def blk (r j : ℕ) (z : List GI) : List GI := (z.drop (j * r)).take r

/-- The first `k` record products. -/
def kmulN (r p ρ k : ℕ) (zf zg : List GI) : List GI :=
  (List.range k).flatMap fun j => kmul r p ρ (NLogN.KFV.blk r j zf) (NLogN.KFV.blk r j zg)

/-! ## Output bounds -/

/-! ## The stage -/

/-- **The uniform pointwise Kronecker stage.** -/
def kronU (pM : Fin Pp) : IStmt HS Pp :=
  .seq (fcode (NLogN.mvK 0 vA NLogN.KFV.vL (.reg NLogN.KFV.vT))) <|
  .seq (fcode (NLogN.mvK 0 NLogN.KFV.vB NLogN.KFV.vL (.reg NLogN.KFV.vT))) <|
  .seq (fcode (NLogN.mvK vA NLogN.KFV.vX NLogN.KFV.vL (.reg NLogN.KFV.vT))) <|
  .seq (fcode (NLogN.mvK NLogN.KFV.vB NLogN.KFV.vY NLogN.KFV.vL (.reg NLogN.KFV.vT))) <|
  .seq (KFC.toUL vRH vZ) <|
  .seq (.loop vZ isMk (.seq (.op (.unmark vZ)) (recU pM))) <|
  fcode (NLogN.mvK vO 0 NLogN.KFV.vL (.reg NLogN.KFV.vT))

/-- Its cost. -/
def cKron (N r p b ρ φ : ℕ) : ℕ :=
  4 * cLp (2 * (N * r)) (250 * (p + 4)) (zsz ((2 * (N * r) : ℕ) : ℤ)) + cUL N +
    (N * (1 + 2000 * cRec r p b (p + ρ) φ + 30) + 30) +
    cLp (2 * (N * r)) (250 * (b + 4)) (zsz ((2 * (N * r) : ℕ) : ℤ))

end NLogN.KFV

end

section
-- Source module: Solutions.KFV13Tab
/-!
# KFV §17: the uniform `J_E` exponent table

`tabP Q M N T` (`KFV3Tab`) evaluates one expression per record whose *shape* is the list of terms
`T` (it grows with `n`).  Here the terms are run-time data: three stacks `tbC, tbK, tbL` hold the
columns `c, k₁, k₂` of `T` (first term on top), and **`tabU`** loops over the records (count `N`
from `tbN`) and, for each record `j`, over the terms (count `|T|` from `tbM`), accumulating
`Σ c · bit_{k₁}(j) · bit_{k₂}(j) mod Q` (`Q` from `tbQ`); the terms are moved aside and moved back
after each record.  The entries pushed on stack `14` are those of `tabP`:

**`tabU_fsp`**: `tabU` pushes `(range N).map (svv Q T)` on stack `14` (entry `0` on top), leaves
everything else as it was, at cost `≤ cTabU N |T| B` with `B = zsz U`
(`O(N (|T| + 1) (B + 2)^2)`, as `tabP`).
-/

set_option linter.unusedSimpArgs false
set_option linter.unusedVariables false

namespace NLogN.KFV

open IntMul RS KIL KFD KFB AModel

abbrev tbV : Fin HS := 3
abbrev tbI : Fin HS := 5
abbrev tbQ : Fin HS := 6
abbrev tbN : Fin HS := 7
abbrev tbM : Fin HS := 8
abbrev tbC : Fin HS := 9
abbrev tbK : Fin HS := 10
abbrev tbL : Fin HS := 11
abbrev tbC2 : Fin HS := 12
abbrev tbK2 : Fin HS := 13
abbrev tbL2 : Fin HS := 15
abbrev tbJ : Fin HS := 16
abbrev tbA : Fin HS := 17
abbrev tbT : Fin HS := 18
abbrev tbR : Fin HS := 19

/-! ## 1. Expressions -/

/-- The record index `j = J - 1`. -/
def jU : Ex HS := .add (.reg tbJ) (.lit (-1))

/-- `bit_k(j)` with `k` on top of stack `x`. -/
def bitU (x : Fin HS) : Ex HS := .nmod (.ndiv jU (.pow2 (.reg x))) (.lit 2)

/-- The current term `c · bit_{k₁}(j) · bit_{k₂}(j)`. -/
def trmU : Ex HS := .mul (.mul (.reg tbC) (bitU tbK)) (bitU tbL)

/-- The accumulator update `(A + t) mod Q`. -/
def accU : Ex HS := .nmod (.add (.reg tbA) (.reg tbT)) (.reg tbQ)

/-- The entry `(Q - A) mod Q`. -/
def entU : Ex HS := .nmod (.add (.reg tbQ) (.neg (.reg tbA))) (.reg tbQ)

section Eval

variable {st : IState HS} {j k c : ℕ}

end Eval

/-! ## 2. The program -/

/-- One term: accumulate it and move it aside. -/
def inB : RS :=
  .seq (.push tbT trmU) (.seq (.push tbT accU) (.seq (.asg tbA (.reg tbT)) (.seq (.pop tbT)
  (.seq (.pop tbT) (.seq (.push tbC2 (.reg tbC)) (.seq (.pop tbC) (.seq (.push tbK2 (.reg tbK))
  (.seq (.pop tbK) (.seq (.push tbL2 (.reg tbL)) (.pop tbL))))))))))

/-- One record: the term loop, the terms moved back, the entry pushed on `14`. -/
def recB : RS :=
  .seq (.push tbA (.lit 0)) (.seq (.rep tbI (.reg tbM) inB)
  (.seq (NLogN.mvK tbC2 tbC tbV (.reg tbM)) (.seq (NLogN.mvK tbK2 tbK tbV (.reg tbM))
  (.seq (NLogN.mvK tbL2 tbL tbV (.reg tbM)) (.seq (.push 14 entU) (.seq (.pop tbA)
    (.asg tbJ (.add (.reg tbJ) (.lit (-1))))))))))

/-- **The uniform exponent table.** -/
def tabU : RS := .seq (.push tbJ (.reg tbN)) (.seq (.rep tbR (.reg tbN) NLogN.KFV.recB) (.pop tbJ))

/-! ## 3. The term loop -/

/-- The value of term `t` at record `j`. -/
def tmv (j : ℕ) (t : ℕ × ℕ × ℕ) : ℕ := t.1 * NLogN.KFV.bt t.2.1 j * NLogN.KFV.bt t.2.2 j

/-- The state after `k` terms of record `j`. -/
def inSt (g : Fin HS → List ℤ) (T : List (ℕ × ℕ × ℕ)) (Q j k : ℕ) : Fin HS → List ℤ :=
  Function.update (Function.update (Function.update (Function.update (Function.update
    (Function.update (Function.update g
    tbC ((T.drop k).map fun t => (t.1 : ℤ)))
    tbK ((T.drop k).map fun t => (t.2.1 : ℤ)))
    tbL ((T.drop k).map fun t => (t.2.2 : ℤ)))
    tbC2 ((T.take k).map fun t => (t.1 : ℤ)).reverse)
    tbK2 ((T.take k).map fun t => (t.2.1 : ℤ)).reverse)
    tbL2 ((T.take k).map fun t => (t.2.2 : ℤ)).reverse)
    tbA [((((T.take k).map (tmv j)).sum % Q : ℕ) : ℤ)]

/-- Cost of one term. -/
def cInB (B : ℕ) : ℕ := 40 * 10 ^ 7 * (B + 2) ^ 2

/-- Cost of one record. -/
def cRecB (m B : ℕ) : ℕ := 5 * 10 ^ 8 * (m + 1) * (B + 2) ^ 2

/-! ## 4. The record loop -/

/-- Cost of the table. -/
def cTabU (N m B : ℕ) : ℕ := 10 ^ 9 * (N + 1) * (m + 1) * (B + 2) ^ 2

/-- The registers of the table program other than the term columns and stack `14`. -/
def tbWork : List (Fin HS) := [tbC2, tbK2, tbL2, tbA, tbT, tbI, tbV, tbJ, tbR]

end NLogN.KFV

end

section
-- Source module: Solutions.KFV14Gen
/-!
# KFV §18: the terms of a round, generated at run time

The table program `tabU` (`KFV13Tab`) reads the terms `T` of round `h` from the stacks
`tbC, tbK, tbL`.  For the paper's layout (`Synth.tgtL`) the terms are
`(2^{ℓ-h+t}, n-1-pos(i,h), n-1-pos(i,t))` for the axes `i` of the round and `t < h`, where the slot
position `pos(i,t)` has the closed form `posF` (prefix slot `nL + (b-1-t) D + i`, chunk slot
`Lst + (⌊e/K⌋ D + i) K + e mod K`, `e = qK - 1 - (t-b)`).  The generator reads the layout
parameters from registers (`GRegs`), so it is one fixed program for every `n`:

* `posE`: the position expression; `posE_eval`/`posE_ok`/`posE_bnd`;
* `termB`/`tLoop`: the terms of one axis (counted loop over `t < h`).
-/

set_option linter.unusedSimpArgs false
set_option linter.unusedVariables false

namespace NLogN.KFV

open IntMul RS KIL KFD KFB AModel

/-- Layout parameters. -/
structure LP where
  D : ℕ
  ℓ : ℕ
  K : ℕ
  b : ℕ
  q : ℕ
  nL : ℕ
  Lst : ℕ
  n : ℕ

/-- `e = qK - 1 - (t - b)`. -/
def eF (L : NLogN.KFV.LP) (t : ℕ) : ℕ := L.q * L.K - 1 - (t - L.b)

/-- **The slot position** of `(i, t)` (`t ≤ ℓ - 2`) in the paper's layout. -/
def posF (L : NLogN.KFV.LP) (i t : ℕ) : ℕ :=
  if t < L.b then L.nL + (L.b - 1 - t) * L.D + i
  else L.Lst + (NLogN.KFV.eF L t / L.K * L.D + i) * L.K + NLogN.KFV.eF L t % L.K

/-- The admissibility of the parameters. -/
structure LP.OK (L : NLogN.KFV.LP) : Prop where
  hD : 1 ≤ L.D
  hK : 1 ≤ L.K
  hLst : L.Lst = L.nL + L.b * L.D
  hn : L.n = L.Lst + L.D * L.q * L.K
  hl : L.b + L.q * L.K + 1 = L.ℓ

/-! ## 1. Registers -/

abbrev gA : Fin HS := 20
abbrev gT : Fin HS := 21
abbrev gI : Fin HS := 22
abbrev gP : Fin HS := 23
abbrev gT2 : Fin HS := 24
abbrev gV : Fin HS := 25
abbrev gD : Fin HS := 40
abbrev gLl : Fin HS := 41
abbrev gKk : Fin HS := 42
abbrev gBb : Fin HS := 43
abbrev gQq : Fin HS := 44
abbrev gNL : Fin HS := 45
abbrev gLst : Fin HS := 46
abbrev gNn : Fin HS := 47
abbrev gH : Fin HS := 48
abbrev gLG : Fin HS := 49
abbrev gLG2 : Fin HS := 39

/-- Register `r` holds `v` on top. -/
def RV (s : IState HS) (g : Fin HS → List ℤ) (r : Fin HS) (v : ℤ) : Prop :=
  ∃ L, (NLogN.KIL.encK s g).G r = .int v :: L

/-- The layout registers, and the round `h`. -/
structure GRegs (s : IState HS) (g : Fin HS → List ℤ) (L : NLogN.KFV.LP) (h : ℕ) : Prop where
  d : RV s g NLogN.KFV.gD L.D
  l : RV s g gLl L.ℓ
  k : RV s g gKk L.K
  b : RV s g gBb L.b
  q : RV s g gQq L.q
  nl : RV s g gNL L.nL
  st : RV s g gLst L.Lst
  n : RV s g gNn L.n
  h : RV s g gH h

def gIn : List (Fin HS) := [NLogN.KFV.gD, gLl, gKk, gBb, gQq, gNL, gLst, gNn, gH]

/-! ## 2. The position expression -/

/-- `e = qK - 1 - (t - b)` with `t` on `tR`. -/
def eE (tR : Fin HS) : Ex HS :=
  .add (.add (.mul (.reg gQq) (.reg gKk)) (.lit (-1))) (.neg (.add (.reg tR) (.neg (.reg gBb))))

/-- The prefix branch `nL + (b - 1 - t) D + i`. -/
def preE (iR tR : Fin HS) : Ex HS :=
  .add (.reg gNL) (.add (.mul (.add (.add (.reg gBb) (.lit (-1))) (.neg (.reg tR))) (.reg NLogN.KFV.gD)) (.reg iR))

/-- The chunk branch `Lst + (⌊e/K⌋ D + i) K + e mod K`. -/
def chkE (iR tR : Fin HS) : Ex HS :=
  .add (.reg gLst) (.add (.mul (.add (.mul (.ndiv (NLogN.KFV.eE tR) (.reg gKk)) (.reg NLogN.KFV.gD)) (.reg iR)) (.reg gKk))
    (.nmod (NLogN.KFV.eE tR) (.reg gKk)))

/-- **The slot position** `pos(i, t)` with `i` on `iR`, `t` on `tR`. -/
def posE (iR tR : Fin HS) : Ex HS := .sel (.add (.reg tR) (.neg (.reg gBb))) (preE iR tR) (chkE iR tR)

section PosE

variable {s : IState HS} {g : Fin HS → List ℤ} {L : LP} {h : ℕ} (H : GRegs s g L h) (hL : L.OK)
  {iR tR : Fin HS} {i t : ℕ} (hiR : RV s g iR i) (htR : RV s g tR t)
  (hi : i < L.D) (ht : t + 2 ≤ L.ℓ)
include H hL hiR htR hi ht

end PosE

/-! ## 3. The terms of one axis -/

/-- The weight of a term: `2^a` (forward) or `Q - 2^a` (opposite round). -/
def tF (op : Bool) (Q a : ℕ) : ℕ := if op then Q - 2 ^ a else 2 ^ a

/-- The sign register value: `-1` forward, `0` opposite. -/
def opV (op : Bool) : ℤ := if op then 0 else -1

/-- The terms of axis `i` (top slot at `tp`) in round `h`. -/
def axTr (L : NLogN.KFV.LP) (op : Bool) (Q h i tp : ℕ) : List (ℕ × ℕ × ℕ) :=
  (List.range h).map fun t => (NLogN.KFV.tF op Q (L.ℓ - h + t), L.n - 1 - tp, L.n - 1 - posF L i t)

/-- Push the terms `T` (in order) on the columns. -/
def pushT (g : Fin HS → List ℤ) (T : List (ℕ × ℕ × ℕ)) : Fin HS → List ℤ :=
  Function.update (Function.update (Function.update g
    tbC ((T.reverse.map fun t => (t.1 : ℤ)) ++ g tbC))
    tbK ((T.reverse.map fun t => (t.2.1 : ℤ)) ++ g tbK))
    tbL ((T.reverse.map fun t => (t.2.2 : ℤ)) ++ g tbL)

/-- One term (`t` on `gT2`): push the weight, `n-1-tp`, `n-1-pos(i,t)`; `t += 1`. -/
def tE : Ex HS := .pow2 (.add (.add (.reg gLl) (.neg (.reg gH))) (.reg gT2))

def termB : RS :=
  .seq (.push tbC (.sel (.reg NLogN.KFV.gV) NLogN.KFV.tE (.add (.reg tbQ) (.neg NLogN.KFV.tE))))
  (.seq (.push tbK (.add (.add (.reg gNn) (.lit (-1))) (.neg (.reg NLogN.KFV.gP))))
  (.seq (.push tbL (.add (.add (.reg gNn) (.lit (-1))) (.neg (posE NLogN.KFV.gI gT2))))
    (.asg gT2 (.add (.reg gT2) (.lit 1)))))

/-- Cost of one term. -/
def cTermB (B : ℕ) : ℕ := 10 ^ 9 * (B + 2) ^ 2

/-- The terms of one axis: `t` from `0` to `h-1` on `gT2` (the loop counter `gT` is not read). -/
def tLoop : RS := .seq (.push gT2 (.lit 0)) (.seq (.rep NLogN.KFV.gT (.reg gH) termB) (.pop gT2))

def cTLoop (h B : ℕ) : ℕ := cLp h (cTermB B) B + 10 ^ 4 * (B + 2)

/-! ## 4. One axis: position, terms, count -/

/-- One axis with top slot `e`: push `gP := e`, the terms, `tbM += h`. -/
def AX (e : Ex HS) : RS :=
  .seq (.push NLogN.KFV.gP e) (.seq NLogN.KFV.tLoop (.seq (.asg tbM (.add (.reg tbM) (.reg gH))) (.pop NLogN.KFV.gP)))

def cAX (sz h B : ℕ) : ℕ := (sz + 1) * 10 ^ 7 * (B + 2) ^ 2 + cTLoop h B + 10 ^ 4 * (B + 2)

/-! ## 5. The axes of a round -/

abbrev gRk : Fin HS := 38

/-- Flag of an axis on `gLG`: `-1` long, `0` short. -/
def fv (b : Bool) : ℤ := if b then -1 else 0

/-- The rank of axis `i` among the long axes. -/
def rk (fl : List Bool) (i : ℕ) : ℕ := (fl.take i).count true

/-- The terms of axis `i` in round `h` (none for a short axis in the top round). -/
def axG (L : NLogN.KFV.LP) (fl : List Bool) (op : Bool) (Q h i : ℕ) : List (ℕ × ℕ × ℕ) :=
  if h + 1 < L.ℓ then axTr L op Q h i (posF L i h)
  else if fl.getD i false then axTr L op Q h i (rk fl i) else []

/-- `e < 0` iff the round is not the top one. -/
def eTop : Ex HS := .add (.add (.reg gH) (.lit 1)) (.neg (.reg gLl))

/-- The branch of one axis. -/
def brB : RS := .ite eTop (AX (posE NLogN.KFV.gI gH)) (.ite (.reg gLG) (AX (.reg gRk)) .skip)

def cBr (h B : ℕ) : ℕ := 2 * cAX 100 h B + 10 ^ 5 * (B + 2)

/-- One axis (`i` on `gI`): its branch, the rank update, the flag moved to `gLG2`, `i += 1`. -/
def axB : RS :=
  .seq NLogN.KFV.brB
  (.seq (.asg gRk (.add (.reg gRk) (.neg (.reg gLG))))
  (.seq (.push gLG2 (.reg gLG)) (.seq (.pop gLG) (.asg NLogN.KFV.gI (.add (.reg NLogN.KFV.gI) (.lit 1))))))

/-- All axes (the loop counter `gA` is not read). -/
def axL : RS := .seq (.push NLogN.KFV.gI (.lit 0)) (.seq (.rep NLogN.KFV.gA (.reg NLogN.KFV.gD) NLogN.KFV.axB) (.pop NLogN.KFV.gI))

def cAxB (h B : ℕ) : ℕ := NLogN.KFV.cBr h B + 10 ^ 5 * (B + 2)

/-- The terms of the first `k` axes. -/
def genL (L : NLogN.KFV.LP) (fl : List Bool) (op : Bool) (Q h k : ℕ) : List (ℕ × ℕ × ℕ) :=
  (List.range k).flatMap (NLogN.KFV.axG L fl op Q h)

/-- The loop state after `k` axes. -/
def axF (g : Fin HS → List ℤ) (L : NLogN.KFV.LP) (fl : List Bool) (op : Bool) (Q h m0 : ℕ)
    (lm rG rG2 : List ℤ) (k : ℕ) : Fin HS → List ℤ :=
  Function.update (Function.update (Function.update (Function.update
    (pushT g (genL L fl op Q h k)) tbM (((m0 + (genL L fl op Q h k).length : ℕ) : ℤ) :: lm))
    gLG ((fl.drop k).map fv ++ rG)) gLG2 ((fl.take k).reverse.map fv ++ rG2)) gRk [(rk fl k : ℤ)]

def cAxL (D h B : ℕ) : ℕ := cLp D (cAxB h B) B + 10 ^ 4 * (B + 2)

end NLogN.KFV

end

section
-- Source module: Solutions.KFV15Tw
/-!
# KFV §19: the uniform twiddle pass `twCU`

`KFV.twC N r = dblI N (2r) ; rotC N (2r) ; unpadI N (2r) (4r)` carries the literals `N` and `2r`.
`twCU` reads them from the parameter stacks `qX = N`, `qB = 2r`, `qZ = 2r` (`KFUU.π3`):
`dblIU` (register-counted doubling), `rotCU` (`const 1 N`, `const 2 Q` become `dup`s), and
KFUU's `unpadIU`.  **`twCU_runs`**: `twC_spec` transported by `KFV.ISim.runs`.
-/

set_option linter.unusedSimpArgs false
set_option linter.unusedVariables false
set_option maxRecDepth 20000

namespace NLogN.KFV

open IntMul KIL KFC KFT KFU KFB KFH KFUU AModel

variable {Pp : ℕ} {I : List Sym} {iprog : Fin Pp → IStmt HS Pp}

def dblBlkU : IStmt HS Pp :=
  .seq (rptX NLogN.KFUU.qB zA zC NLogN.KFU.zD dnStep) (.seq (rptX NLogN.KFUU.qB zA zC NLogN.KFU.zD (.op (.mv zX zY)))
    (rptX NLogN.KFUU.qB zA zC NLogN.KFU.zD (.op (.mv zY zO))))

/-- Uniform doubling scan. -/
def dblIU : IStmt HS Pp :=
  .seq (.op (.mark zO)) (.seq (rptX NLogN.KFUU.qX oA oC NLogN.KFU.oD dblBlkU) (.seq (mvW zO 0) (.op (.unmark zO))))

/-- Uniform rotation (parameters `N`, `Q` duplicated from `qX`, `qB`). -/
def rotCU : IStmt HS Pp :=
  .seq (.op (.dup NLogN.KFUU.qX 1)) (.seq (.op (.dup NLogN.KFUU.qB 2)) (.seq (.op (.const 4 2))
    (.seq rotF (.seq (.op (.drop 1)) (.seq (.op (.drop 2)) (.op (.drop 4)))))))

/-- **The uniform twiddle pass.** -/
def twCU : IStmt HS Pp := .seq dblIU (.seq NLogN.KFV.rotCU unpadIU)

/-- Its parameters. -/
def πT (N r : ℕ) : PMap := π3 N (2 * r) (4 * r - 2 * r)

end NLogN.KFV

end

section
-- Source module: Solutions.KFV16Twist
/-!
# KFV §20: the uniform twist `J_E` of a round

`twistU op` (`op = false`: `J_E` of a forward round, `op = true`: `J_{-E}` of an opposite round):

* `genT op`: `Q = 2^{ℓ+1}`, `N = 2^n` on `tbQ`, `tbN`; the terms of the round (`axL`) on the
  columns, their number on `tbM` (**`genT_fsp`**);
* `tabU` (`KFV13Tab`): the exponent table on stack `14`;
* `clrT`: the term columns are emptied; the parameters `N, 2r, 2r` of `twCU` are pushed;
* `twCU` (`KFV15Tw`), then the parameters are popped.

Everything `n`-dependent is read from the layout registers `GRegs` and the flag list on `gLG`.
-/

set_option linter.unusedSimpArgs false
set_option linter.unusedVariables false

namespace NLogN.KFV

open IntMul RS KIL KFD KFB AModel

/-- The term generation of a round. -/
def genT (op : Bool) : RS :=
  .seq (.push NLogN.KFV.gV (.lit (opV op)))
  (.seq (.push tbQ (.pow2 (.add (.reg gLl) (.lit 1))))
  (.seq (.push tbN (.pow2 (.reg gNn)))
  (.seq (.push tbM (.lit 0))
  (.seq (.push gRk (.lit 0))
  (.seq axL
  (.seq (.pop gRk)
  (.seq (NLogN.mvK gLG2 gLG NLogN.KFV.gA (.reg NLogN.KFV.gD)) (.pop NLogN.KFV.gV))))))))

def cGenT (D h B : ℕ) : ℕ := cAxL D h B + 10 ^ 9 * (B + 2) ^ 2 + 10 ^ 4 * D * (B + 2)

/-- The generated terms of round `h`. -/
def genG (L : LP) (fl : List Bool) (op : Bool) (h : ℕ) : List (ℕ × ℕ × ℕ) :=
  genL L fl op (2 ^ (L.ℓ + 1)) h L.D

/-! ## Emptying the term columns -/

def clrB : RS := .seq (.pop tbC) (.seq (.pop tbK) (.pop tbL))

/-- Drop `|T|` entries from each term column. -/
def clrT : RS := .rep NLogN.KFV.gA (.reg tbM) clrB

def colSt (g : Fin HS → List ℤ) (T : List (ℕ × ℕ × ℕ)) (k : ℕ) : Fin HS → List ℤ :=
  Function.update (Function.update (Function.update g
    tbC (((T.map fun t => (t.1 : ℤ))).drop k ++ g tbC))
    tbK (((T.map fun t => (t.2.1 : ℤ))).drop k ++ g tbK))
    tbL (((T.map fun t => (t.2.2 : ℤ))).drop k ++ g tbL)

/-! ## The register part of the twist -/

abbrev qX : Fin HS := KFUU.qX
abbrev qB : Fin HS := KFUU.qB
abbrev qZ : Fin HS := KFUU.qZ

/-- Terms, table, clean-up, the parameters of `twCU`. -/
def twA (op : Bool) : RS :=
  .seq (genT op) (.seq NLogN.KFV.tabU (.seq clrT (.seq (.pop tbM) (.seq (.pop tbN) (.seq (.pop tbQ)
    (.seq (.push qX (.pow2 (.reg gNn))) (.seq (.push NLogN.KFV.qB (.pow2 (.add (.reg gLl) (.lit 1))))
      (.push qZ (.pow2 (.add (.reg gLl) (.lit 1)))))))))))

/-- Pop the parameters of `twCU`. -/
def twB : RS := .seq (.pop qX) (.seq (.pop NLogN.KFV.qB) (.pop qZ))

def cTwA (N D h m B : ℕ) : ℕ := cGenT D h B + cTabU N m B + cLp m (6 * B + 6) B + 10 ^ 9 * (B + 2) ^ 2

/-- The table of a round. -/
def tabG (L : LP) (fl : List Bool) (op : Bool) (h : ℕ) : List ℤ :=
  (List.range (2 ^ L.n)).map (svv (2 ^ (L.ℓ + 1)) (genG L fl op h).reverse)

/-! ## The twist pass -/

section Twist

variable {Pp : ℕ} {I : List Sym} {iprog : Fin Pp → IStmt HS Pp}

/-- **The uniform twist** of a round (`op = true`: the opposite twist). -/
def twistU (op : Bool) : IStmt HS Pp := .seq (fcode (twA op)) (.seq twCU (fcode twB))

def cTwistU (n ℓ D h m W B : ℕ) : ℕ :=
  cTwA (2 ^ n) D h m B + 2000 * cTw (2 ^ n) (2 ^ ℓ) W + 10 ^ 4 * (B + 2)

end Twist

end NLogN.KFV

end

section
-- Source module: Solutions.KFV17Link
/-!
# KFV §21: the generated terms are the paper's exponent table

For the paper's layout (`pLP`: `D, ℓ, K, b = bb, q = qq, nL = nLong, Lst = |lstL|, n = |tgtL|`)
the closed-form positions are the slot positions of `tgtL`:

* `spos_eq_posF`: `spos (i, t) = posF i t` for `t ≤ ℓ - 2` (`pre_slot`, `tgt_slot`);
* `spos_top`: the top slot of a long axis sits at its rank among the long axes;

hence (**`genG_perm`**) the terms generated at run time (`genG`) are a permutation of
`tabT h = trips tgtL ℓ h (partAxes al h)` (forward) resp. of `negTr (2r) (tabT h)` (opposite),
and the exponent tables agree (**`svv_genG`**).
-/

set_option linter.unusedSectionVars false
set_option linter.unusedSimpArgs false
set_option linter.unusedVariables false

namespace NLogN.KFV

open Finset Synth GlueLayout GlueP14 GlueError

/-! ## List lemmas -/

/-! ## The paper's layout parameters -/

section Paper

variable (D ℓ K : ℕ) [NeZero D] (long : Fin D → Bool)

/-- The flags of the axes. -/
def flOf : List Bool := (List.finRange D).map long

/-- The layout parameters of `tgtL`. -/
def pLP : LP :=
  ⟨D, ℓ, K, bb ℓ K, qq ℓ K, nLong D long, (lstL D ℓ K long).length, (tgtL D ℓ K long).length⟩

end Paper

end NLogN.KFV

end

section
-- Source module: Solutions.KFV18Ite
/-!
# KFV §22: a sign test with arbitrary (call-containing) branches

`iteI e a b` is the shape of the compiled `RS.ite` with arbitrary item-language branches:
evaluate `e`, branch on its sign (`a` if `e < 0`). **`iteI_runsT`**, **`iteI_runsF`**.
-/

set_option linter.unusedSimpArgs false
set_option linter.unusedVariables false

namespace NLogN.KFV

open IntMul RS KIL KFD

variable {Pp : ℕ} {I : List Sym} {iprog : Fin Pp → IStmt HS Pp}

/-- Branch on the sign of `e`. -/
def iteI (e : Ex HS) (a b : IStmt HS Pp) : IStmt HS Pp :=
  .seq (e.code NLogN.zX) (.seq (.op (.mv NLogN.zX.V NLogN.zX.A))
    (.ite NLogN.zX.A (fun k => k == .neg) (.seq (.op (.drop NLogN.zX.A)) a) (.seq (.op (.drop NLogN.zX.A)) b)))

/-- Overhead of the test. -/
def cIte (e : Ex HS) (s : IState HS) : ℕ :=
  e.cost s + (2 * moveCost (zsz (e.eval s)) + 19) + (2 * zsz (e.eval s) + 2) + 25

end NLogN.KFV

end

section
-- Source module: Solutions.KFV19Layer
/-!
# KFV §23: uniform layers in a context

`LayerU Ctx code T N p W cost`: on any state `encK s (o0 (gflat zs ++ l))` (only stack `0` managed)
whose base `s` satisfies `Ctx`, `code` replaces the array by one representing `T xs`. The round's
context (layout registers, flags) lives in the base, so a layer may call procedures that need
every other managed stack empty (Proposition 14).

* `RecRep_unique`: the record representation determines the integers, so the output of a layer
  does not depend on the base (used for loops).
* `encK_encK`: a managed frame inside the base can be moved into the managed part.
* `RCtx L fl h`: the base is `encK s1 fG` with the layout registers (`GRegs`), the flags of the
  long axes on `gLG`, and the twist scratch empty in `fG`.
* **`layerU_twist`**: `twistU op` is the layer `twLc r sh Ex` for every `Ex` whose exponents
  modulo `2r` are the generated table's (`r = 2^ℓ`).
-/

set_option linter.unusedSectionVars false
set_option linter.unusedSimpArgs false
set_option linter.unusedVariables false

namespace NLogN.KFV

open IntMul RS KIL KFD KFB KFH AModel

variable {Pp : ℕ} {I : List Sym} {iprog : Fin Pp → IStmt HS Pp}

/-! ## Basic facts -/

/-- Only stack `0` managed. -/
def o0 (v : List ℤ) : Fin HS → List ℤ := Function.update (fun _ => []) 0 v

/-! ## Layers in a context -/

section Layer

variable {r : ℕ}

/-- `code` implements the layer `T` on bases satisfying `Ctx`. -/
def LayerU (I : List Sym) (iprog : Fin Pp → IStmt HS Pp) (Ctx : IState HS → Prop)
    (code : IStmt HS Pp) (T : List (Fin r → ℂ) → List (Fin r → ℂ)) (N p W cost : ℕ) : Prop :=
  ∀ s, Ctx s → ∀ (zs : List GI) (l : List ℤ) (xs : List (Fin r → ℂ)),
    RecRep r p zs xs → xs.length = N → NLogN.KFB.GB W zs → Disk xs →
    ∃ zs', RecRep r p zs' (T xs) ∧ NLogN.KFB.GB W zs' ∧ (T xs).length = N ∧ Disk (T xs) ∧
      IRuns I iprog code (NLogN.KIL.encK s (o0 (gflat zs ++ l))) (NLogN.KIL.encK s (o0 (gflat zs' ++ l))) cost

variable {Ctx Ctx' : IState HS → Prop}

end Layer

/-! ## The round context -/

/-- Stacks the twist needs empty in the frame. -/
def twFree : List (Fin HS) :=
  [0, NLogN.KFV.gT, gT2, NLogN.KFV.gV, tbQ, tbN, tbM, gRk, NLogN.KFV.gA, NLogN.KFV.gI, tbC, tbK, tbL, qX, NLogN.KFV.qB, qZ] ++ tbWork

/-- The base of a round: a frame `fG` with the layout registers, the flags, and the twist
scratch empty, over an arbitrary state. -/
def RCtx (L : LP) (fl : List Bool) (h : ℕ) (s : IState HS) : Prop :=
  ∃ s1 fG rG, s = NLogN.KIL.encK s1 fG ∧ GRegs s1 fG L h ∧ fG gLG = fl.map fv ++ rG ∧
    ∀ j ∈ twFree, fG j = []

/-! ## The twist as a layer -/

/-! ## The paper's rounds -/

section Paper

open Synth GlueLayout GlueP14 GlueError

variable (D ℓ K : ℕ) [NeZero D] (long : Fin D → Bool)

end Paper

end NLogN.KFV

end

section
-- Source module: Solutions.KFV12Leaf
/-!
# KFV §16: the uniform individual-kernel leaf `leafU` (prefix and top rounds of Lemma 16)

`leafU` is `KFH.p14Leaf P 1 D S 0 r` with every size read from registers (`LRegs`): `P, D`,
`S r`, `P 2^{D-1}`, `2^{D-1} S r`, `2^{dScale D}`, `2·vol·r`, `D mod 8` on the stacks `6 … 13` (top
of the managed list or the background: the specs only read `(encK s f).G`).  The kernel passes
are one counted loop per tensor power (`kersU`, counter `5`), whose body is `kerCore M`
(literally `qopW [M] B1 G 0` after substituting the registers `1 … 4, 12`, `KFW.FSp.of_sub`)
followed by the register update `B1 ← 2 B1`, `G ← G / 2`; the scalar kernel `MD D` is selected by
`D mod 8` (`gD_mod`), and the truncation `truncU` is `truncP` with `N, 2^Δ` in registers.
Every stack `1 … 13` it uses besides the inputs is empty on entry and exit.
Main results: `leafU_fsp`, `leafU_iruns`, `cLeafU_le`.
-/

set_option linter.unusedSectionVars false
set_option linter.unusedVariables false
set_option linter.unusedSimpArgs false

namespace NLogN.KFV

open IntMul RS AModel KIL KFD KFB KFH

abbrev yB : Fin HS := 1
abbrev yB2 : Fin HS := 2
abbrev yB4 : Fin HS := 3
abbrev yG : Fin HS := 4
abbrev yC : Fin HS := 5
abbrev yP : Fin HS := 6
abbrev yD : Fin HS := 7
abbrev yQ : Fin HS := 8
abbrev yG0 : Fin HS := 9
abbrev yQD : Fin HS := 10
abbrev yW : Fin HS := 11
abbrev yN : Fin HS := 12
abbrev yM : Fin HS := 13

/-- The kernel pass with its sizes in registers. -/
def kerCore (M : M2G) : RS :=
  .seq (.push NLogN.KFB.rB (.reg NLogN.KFV.yB)) (.seq (.push NLogN.KFB.rB2 (.reg yB2))
  (.seq (.push rB4 (.reg yB4)) (.seq (.push NLogN.KFB.rG (.reg yG))
  (.seq (.push NLogN.KFB.rN (.reg yN))
  (.seq (.push NLogN.KFB.rP (.lit ((2 * 0 : ℕ) : ℤ)))
  (.seq (NLogN.mvK 0 qP NLogN.KFB.c1 (.reg NLogN.KFB.rP)) (.seq (qop [M]) (.seq (NLogN.mvK qP 0 NLogN.KFB.c1 (.reg NLogN.KFB.rP))
  (.seq (.pop NLogN.KFB.rB) (.seq (.pop NLogN.KFB.rB2) (.seq (.pop rB4) (.seq (.pop NLogN.KFB.rG) (.seq (.pop NLogN.KFB.rN)
    (.pop NLogN.KFB.rP))))))))))))))

/-- The substitution of a kernel pass. -/
def sgK (B1 G N : ℕ) : Fin HS → Option ℤ := fun i =>
  if i = NLogN.KFV.yB then some (B1 : ℤ) else if i = yB2 then some ((2 * B1 : ℕ) : ℤ)
  else if i = yB4 then some ((4 * B1 : ℕ) : ℤ) else if i = yG then some (G : ℤ)
  else if i = yN then some (N : ℤ) else none

/-- The leaf's own registers. -/
def yAll : List (Fin HS) := [NLogN.KFV.yB, yB2, yB4, yG, yC, NLogN.KFV.yP, yD, yQ, yG0, yQD, NLogN.KFV.yW, yN, NLogN.KFV.yM]

/-- The kernel registers on the managed lists. -/
def regF (B1 G : ℕ) (f : Fin HS → List ℤ) : Fin HS → List ℤ :=
  Function.update (Function.update (Function.update (Function.update f NLogN.KFV.yB [(B1 : ℤ)]) yB2
    [((2 * B1 : ℕ) : ℤ)]) yB4 [((4 * B1 : ℕ) : ℤ)]) yG [(G : ℤ)]

/-- The register update between kernel passes: `B1 ← 2 B1`, `G ← G / 2`. -/
def kUpd : RS :=
  .seq (.asg NLogN.KFV.yB (.reg yB2)) (.seq (.asg yB2 (.reg yB4))
    (.seq (.asg yB4 (.add (.reg yB4) (.reg yB4))) (.asg yG (.quo (.reg yG) (.lit 2)))))

/-! ## The kernel loop -/

/-- The uniform kernel loop: `D` passes over the bits `D-1, …, 0` of `[P] × [2]^D × [S]`. -/
def kersU (M : M2G) : RS := .rep yC (.reg yD) (.seq (kerCore M) kUpd)

/-- The arrays after `k` passes. -/
def zsK (M : M2G) (B0 G0 : ℕ) (zs : List GI) : ℕ → List GI
  | 0 => zs
  | k + 1 => qopL [M] (2 ^ k * B0) (G0 / 2 ^ k) (zsK M B0 G0 zs k)

/-- The operator after `k` passes. -/
noncomputable def PK (A : Butterfly.M2) (D k : ℕ) : Butterfly.Op (Fin D) :=
  ((((List.finRange D).drop (D - k)).map fun i => Butterfly.kerAt i A).map Butterfly.tensF).prod

/-- Cost of the kernel loop. -/
def cKers (D N Z w : ℕ) : ℕ :=
  200 * (zsz (D : ℤ) + 2) + cRep D + D * (400 * ck N Z (w + 2 * D) + 40000 * (Z + 2))

/-! ## Loading and clearing the kernel registers -/

/-- Load `B1 ← a`, `2 B1`, `4 B1`, `G ← b`. -/
def kPush (a b : Fin HS) : RS :=
  .seq (.push NLogN.KFV.yB (.reg a)) (.seq (.push yB2 (.add (.reg NLogN.KFV.yB) (.reg NLogN.KFV.yB)))
    (.seq (.push yB4 (.add (.reg yB2) (.reg yB2))) (.push yG (.reg b))))

/-- Clear the kernel registers. -/
def kPop : RS := .seq (.pop yG) (.seq (.pop yB4) (.seq (.pop yB2) (.pop NLogN.KFV.yB)))

/-- The kernel registers are free. -/
def KFree (f : Fin HS → List ℤ) : Prop := f NLogN.KFV.yB = [] ∧ f yB2 = [] ∧ f yB4 = [] ∧ f yG = []

/-! ## The leaf's register file and layer implementations -/

/-- The leaf's inputs: `P, D`, `S r`, `P 2^{D-1}`, `2^{D-1} S r`, `2^{dScale D}`, `2 · vol · r`,
`D mod 8`. -/
structure LRegs (s : IState HS) (f : Fin HS → List ℤ) (P D S r : ℕ) : Prop where
  p : ∃ L, (NLogN.KIL.encK s f).G NLogN.KFV.yP = .int (P : ℤ) :: L
  d : ∃ L, (NLogN.KIL.encK s f).G yD = .int (D : ℤ) :: L
  q : ∃ L, (NLogN.KIL.encK s f).G yQ = .int ((S * r : ℕ) : ℤ) :: L
  g0 : ∃ L, (NLogN.KIL.encK s f).G yG0 = .int ((P * 2 ^ (D - 1) : ℕ) : ℤ) :: L
  qd : ∃ L, (NLogN.KIL.encK s f).G yQD = .int ((2 ^ (D - 1) * (S * r) : ℕ) : ℤ) :: L
  w : ∃ L, (NLogN.KIL.encK s f).G NLogN.KFV.yW = .int ((2 ^ dScale D : ℕ) : ℤ) :: L
  n : ∃ L, (NLogN.KIL.encK s f).G yN = .int ((2 * (vol (Butterfly.layerShape P 1 D S) * r) : ℕ) : ℤ) :: L
  m : ∃ L, (NLogN.KIL.encK s f).G NLogN.KFV.yM = .int ((D % 8 : ℕ) : ℤ) :: L

/-- The input registers. -/
def yIn : List (Fin HS) := [NLogN.KFV.yP, yD, yQ, yG0, yQD, NLogN.KFV.yW, yN, NLogN.KFV.yM]

/-- `code` implements `T` on the layer `[P] × [2]^D × [S]` whose sizes are in the leaf registers. -/
def LU (P D S r : ℕ) (T : Butterfly.Op (Fin D)) (code : RS) (e dw : ℕ) (cost : ℕ → ℕ) : Prop :=
  ∀ (s : IState HS) (f : Fin HS → List ℤ) (zs : List GI) (l : List ℤ) (xs : List (Fin r → ℂ))
    (F w : ℕ), Cl f → KFree f → f yC = [] → NLogN.KFV.LRegs s f P D S r → f 0 = gflat zs ++ l →
    RecRep r F zs xs → xs.length = vol (Butterfly.layerShape P 1 D S) → NLogN.KFB.GB w zs →
    ∃ zs', RecRep r (F + e) zs' (Butterfly.selOp (Butterfly.layerSel P 1 D S 0 Nat.zero_lt_one) T xs) ∧
      NLogN.KFB.GB (w + dw) zs' ∧
      FSp code (NLogN.KIL.encK s f) (NLogN.KIL.encK s (Function.update f 0 (gflat zs' ++ l))) (cost w)

/-- One loop of kernels with the register load and clear. -/
def kLoop (M : M2G) : RS := .seq (kPush yQ yG0) (.seq (kersU M) kPop)

/-! ## The scalar kernel `MD D`, chosen by `D mod 8` -/

/-- Cost of the scalar kernel stage. -/
def kdU (N Z w : ℕ) : ℕ := 24000 * (Z + 2) + 8 * cSel Z + 400 * ck N Z w + 4 * (2 * Z + 2)

/-- The scalar kernel on bit `0`. -/
def kerDU : RS :=
  .seq (kPush yQD NLogN.KFV.yP) (.seq (NLogN.KIL.sel NLogN.KFV.yM (fun c => kerCore (MD c)) 0 8) kPop)

/-! ## The truncation with its sizes in registers -/

/-- `truncP` reading `N` from `yN` and `2^Δ` from `yW`. -/
def truncU : RS :=
  .seq (.push NLogN.KFB.rN (.reg yN)) (.seq (.rep NLogN.KFB.c1 (.reg NLogN.KFB.rN) (.seq (.push qO (.quo (.reg 0) (.reg NLogN.KFV.yW))) (.pop 0)))
    (.seq (NLogN.mvK qO 0 NLogN.KFB.c1 (.reg NLogN.KFB.rN)) (.pop NLogN.KFB.rN)))

/-- Its substitution. -/
def sgT (Δ N : ℕ) : Fin HS → Option ℤ := fun i =>
  if i = yN then some (N : ℤ) else if i = NLogN.KFV.yW then some ((2 : ℤ) ^ Δ) else none

/-! ## The uniform leaf -/

/-- **The uniform individual-kernel leaf**: `p14Leaf P 1 D S 0 r` with all sizes in registers. -/
def leafU : RS := .seq (.seq (.seq (.seq (kLoop NLogN.KFH.MS) (kLoop MC)) (kLoop NLogN.KFH.MS)) kerDU) NLogN.KFV.truncU

/-- Cost of one kernel loop stage. -/
def kcU (D N Z w : ℕ) : ℕ := 24000 * (Z + 2) + (cKers D N Z w + 4 * (2 * Z + 2))

/-- Cost of the leaf. -/
def cLeafU (D N Z w : ℕ) : ℕ :=
  kcU D N Z w + kcU D N Z (w + 2 * D) + kcU D N Z (w + (2 * D + 2 * D)) +
    kdU N Z (w + (2 * D + 2 * D + 2 * D)) + 400 * cTrunc (w + 6 * D + 2) (dScale D) (2 * N) Z

/-! ## Fine compilability, cost bound and the IL run -/

end NLogN.KFV

end

section
-- Source module: Solutions.KFV20Leaf
/-!
# KFV §24: the leaf rounds with register-computed parameters

`leafR eP eD eS` pushes the eight leaf registers (`LRegs`) of the layer
`[2^a] × [2]^d × [2^c]` (records of `2^ℓ` coefficients), computed at run time from expressions
`eP, eD, eS` over the layout registers (`a, d, c` their values, `ℓ` on `gLl`), runs the uniform
leaf `leafU` (`KFV12Leaf`), and pops the registers. **`leafR_runs`**.

* `EV st B e v`: `e` is evaluable on `st` with value `v`, all intermediate values `B`-bit;
  closure lemmas for the constructors used.
-/

set_option linter.unusedSectionVars false
set_option linter.unusedSimpArgs false
set_option linter.unusedVariables false

namespace NLogN.KFV

open IntMul RS KIL KFD KFB KFH AModel

variable {Pp : ℕ} {I : List Sym} {iprog : Fin Pp → IStmt HS Pp}

/-! ## Evaluated expressions -/

/-- `e` evaluates on `st` to `v`, with every intermediate value `B`-bit. -/
def EV (st : IState HS) (B : ℕ) (e : Ex HS) (v : ℤ) : Prop :=
  e.ok st ∧ e.eval st = v ∧ e.Bnd st B

namespace EV

variable {st : IState HS} {B : ℕ}

end EV

/-! ## The leaf register file -/

/-- The register pushes. -/
def lfPush (eP eD eS : Ex HS) : RS :=
  .seq (.push NLogN.KFV.yP (.pow2 eP))
  (.seq (.push yD eD)
  (.seq (.push yQ (.pow2 (.add eS (.reg gLl))))
  (.seq (.push yG0 (.pow2 (.add (.add eP eD) (.lit (-1)))))
  (.seq (.push yQD (.pow2 (.add (.add (.add eD (.lit (-1))) eS) (.reg gLl))))
  (.seq (.push NLogN.KFV.yW (.pow2 (.add (.add eD (.ndiv eD (.lit 2))) (.nmod eD (.lit 2)))))
  (.seq (.push yN (.pow2 (.add (.add (.add (.add eP eD) eS) (.reg gLl)) (.lit 1))))
    (.push NLogN.KFV.yM (.nmod eD (.lit 8)))))))))

/-- The register pops. -/
def lfPop : RS :=
  .seq (.pop NLogN.KFV.yM) (.seq (.pop yN) (.seq (.pop NLogN.KFV.yW) (.seq (.pop yQD) (.seq (.pop yG0)
    (.seq (.pop yQ) (.seq (.pop yD) (.pop NLogN.KFV.yP)))))))

/-- **The leaf with register-computed parameters.** -/
def leafR (eP eD eS : Ex HS) : RS := .seq (lfPush eP eD eS) (.seq leafU lfPop)

/-- The register values (`P = 2^a`, `S = 2^c`, `r = 2^ℓ`). -/
def lfF (v : List ℤ) (a d c ℓ : ℕ) : Fin HS → List ℤ :=
  Function.update (Function.update (Function.update (Function.update (Function.update
    (Function.update (Function.update (Function.update (o0 v) NLogN.KFV.yP [((2 ^ a : ℕ) : ℤ)])
    yD [(d : ℤ)]) yQ [((2 ^ c * 2 ^ ℓ : ℕ) : ℤ)]) yG0 [((2 ^ a * 2 ^ (d - 1) : ℕ) : ℤ)])
    yQD [((2 ^ (d - 1) * (2 ^ c * 2 ^ ℓ) : ℕ) : ℤ)]) NLogN.KFV.yW [((2 ^ dScale d : ℕ) : ℤ)])
    yN [((2 * (vol (Butterfly.layerShape (2 ^ a) 1 d (2 ^ c)) * 2 ^ ℓ) : ℕ) : ℤ)])
    NLogN.KFV.yM [((d % 8 : ℕ) : ℤ)]

/-- Cost of the pushes. -/
def cLfPush (sz B : ℕ) : ℕ := 8 * (3 * sz + 12) * 10 ^ 7 * (B + 2) ^ 2

/-- All registers of `e` are layout registers. -/
def RG (e : Ex HS) : Prop := ∀ r ∈ e.regs, r ∈ NLogN.KFV.gIn

/-- Cost of a register leaf. -/
def cLeafR (sz a d c ℓ w B : ℕ) : ℕ :=
  cLfPush sz B + cLeafU d (vol (Butterfly.layerShape (2 ^ a) 1 d (2 ^ c)) * 2 ^ ℓ)
    (zsz ((12 * (vol (Butterfly.layerShape (2 ^ a) 1 d (2 ^ c)) * 2 ^ ℓ) : ℕ) : ℤ)) w +
    8 * (2 * B + 2)

end NLogN.KFV

end

section
-- Source module: Solutions.KFV21Rounds
/-!
# KFV §25: the prefix and top rounds of the paper's layout, uniform

* **`layerU_pre`**: for `h < b`, `preC = leafR ePre gD eSpre` (with `ePre = nL + (b-1-h) D`,
  `eSpre = n - ePre - D` read from the layout registers) is the computed layer of round `h`.
* **`layerU_top`**: `topC` tests `nL` at run time: without long axes it is `Q_p` alone
  (`truncTop`), else the leaf on `[1] × [2]^{nL} × [2^{n-nL}]`.
-/

set_option linter.unusedSectionVars false
set_option linter.unusedSimpArgs false
set_option linter.unusedVariables false

namespace NLogN.KFV

open IntMul RS KIL KFD KFB KFH AModel Butterfly

variable {Pp : ℕ} {I : List Sym} {iprog : Fin Pp → IStmt HS Pp}

/-! ## Prefix rounds -/

/-- `nL + (b - 1 - h) D`. -/
def ePre : Ex HS :=
  .add (.reg gNL) (.mul (.add (.add (.reg gBb) (.lit (-1))) (.neg (.reg gH))) (.reg NLogN.KFV.gD))

/-- `n - ePre - D`. -/
def eSpre : Ex HS := .add (.add (.reg gNn) (.neg ePre)) (.neg (.reg NLogN.KFV.gD))

/-- The prefix round code. -/
def preC : RS := leafR ePre (.reg NLogN.KFV.gD) eSpre

section Paper

open Synth GlueLayout GlueP14 GlueError

variable (D ℓ K : ℕ) [NeZero D] (long : Fin D → Bool)

/-- The leaf register bound of a round. -/
def bL : ℕ := (tgtL D ℓ K long).length + D + ℓ + 3

end Paper

/-! ## The sign test with a register expression -/

/-- Overhead bound of `iteI`. -/
def cIteB (sz B : ℕ) : ℕ := (sz + 1) * 10 ^ 7 * (B + 2) ^ 2 + 2 * B + 27

/-! ## The top round -/

/-- `Q_p` alone: `N = 2^{n+ℓ+1}` integers truncated by `2^0`. -/
def truncTop : RS :=
  .seq (.push yN (.pow2 (.add (.add (.reg gNn) (.reg gLl)) (.lit 1))))
  (.seq (.push NLogN.KFV.yW (.lit 1)) (.seq NLogN.KFV.truncU (.seq (.pop NLogN.KFV.yW) (.pop yN))))

/-- `nL - 1` (negative iff no long axes). -/
def eTopT : Ex HS := .add (.reg gNL) (.lit (-1))

/-- The top leaf: `[1] × [2]^{nL} × [2^{n - nL}]`. -/
def topLeaf : RS := leafR (.lit 0) (.reg gNL) (.add (.reg gNn) (.neg (.reg gNL)))

/-- The top round code. -/
def topC : IStmt HS Pp := iteI eTopT (fcode truncTop) (fcode topLeaf)

def cTruncTop (n ℓ W B : ℕ) : ℕ :=
  10 * 10 ^ 7 * (B + 2) ^ 2 + 400 * cTrunc W 0 (2 * (2 ^ n * 2 ^ ℓ)) B + 2 * (2 * B + 2)

section Paper

open Synth GlueLayout GlueP14 GlueError

variable (D ℓ K : ℕ) [NeZero D] (long : Fin D → Bool)

/-- Cost of the top round. -/
def cTopC (sz W B : ℕ) : ℕ :=
  cIteB 3 B + cTruncTop (tgtL D ℓ K long).length ℓ W B +
    cLeafR sz 0 (nLong D long) ((tgtL D ℓ K long).length - nLong D long) ℓ W B

end Paper

end NLogN.KFV

end

section
-- Source module: Solutions.KFV25Moves
/-!
# KFV §29: an explicit single-bit move sequence from the row-major layout to the paper's layout

`initL` (axis order, every axis most significant bit first: the row-major order of the records)
is taken to `tgtL` (Lemma 16's layout) by explicit single-bit moves `LOp.mv c p` whose positions
have closed forms in the layout parameters, so a fixed program can generate them at run time.

* `bmOps c δ K`: the `K` moves `mv (c+t) (c+t+δ)`; **`applyN_bm`**: they take
  `A ++ B ++ X ++ E` to `A ++ X ++ B ++ E` (`|A| = c`, `|B| = δ`, `|X| = K`).
* `genOps`: one pass over a list of triples `(U, T, V)`: every `T` is moved to the growing
  prefix; **`applyN_gen`**.
* Phases 1–2 (tops of the long axes; the low bits `h = b-1, …, 0`, at most `D (b+1)` moves)
  and **`applyN_moves12`**: `applyN (moves12 …) initL = Lst ++ (chunks in axis-major order)`.
  The chunk phase uses Lemma 12 exchanges (`KFV27Chunks`), not single-bit moves (cost).
-/

set_option linter.unusedSectionVars false
set_option linter.unusedSimpArgs false

namespace NLogN.KFV

open Synth AModel AModelL12

/-! ## Block moves -/

/-- `K` single moves bringing the block at `[c+δ, c+δ+K)` to `c`. -/
def bmOps (c δ K : ℕ) : List LOp := (List.range K).map fun t => LOp.mv (c + t) (c + t + δ)

/-! ## Passes -/

/-- One pass: every middle part `T` goes to the prefix. -/
def genOps : ℕ → ℕ → List (List ℕ × List ℕ × List ℕ) → List LOp
  | _, _, [] => []
  | c, d, (U, T, V) :: xs =>
    bmOps c (d + U.length) T.length ++ genOps (c + T.length) (d + U.length + V.length) xs

def trip (x : List ℕ × List ℕ × List ℕ) : List ℕ := x.1 ++ (x.2.1 ++ x.2.2)

def tOf (x : List ℕ × List ℕ × List ℕ) : List ℕ := x.2.1

def uvOf (x : List ℕ × List ℕ × List ℕ) : List ℕ := x.1 ++ x.2.2

/-! ## The paper's layout -/

section Paper

variable (D ℓ K : ℕ) (long : Fin D → Bool)

/-- The top slot of a long axis. -/
def topS (i : Fin D) : List ℕ := if long i then [nm D i (ℓ - 1)] else []

/-- The low slots `h = b-1-u, …, 0`. -/
def loS (u : ℕ) (i : Fin D) : List ℕ := (List.range (bb ℓ K - u)).reverse.map (nm D i)

/-- Phase 1: the tops of the long axes. -/
def xs1 : List (List ℕ × List ℕ × List ℕ) :=
  (List.finRange D).map fun i => ([], topS D ℓ long i, restA D ℓ K i ++ loS D ℓ K 0 i)

/-- Phase 2, pass `u`: the low slot `h = b-1-u` of every axis. -/
def xs2 (u : ℕ) : List (List ℕ × List ℕ × List ℕ) :=
  (List.finRange D).map fun i => (restA D ℓ K i, [nm D i (bb ℓ K - 1 - u)], loS D ℓ K (u + 1) i)

/-- The prefix after `u` passes of phase 2. -/
def pre2 (u : ℕ) : List ℕ :=
  ((List.finRange D).filter fun i => long i).map (fun i => nm D i (ℓ - 1)) ++
    (List.range u).flatMap fun u' => (List.finRange D).map fun i => nm D i (bb ℓ K - 1 - u')

/-- **The single-bit moves of phases 1–2.** -/
def moves12 : List LOp :=
  genOps 0 0 (xs1 D ℓ K long) ++
    ((List.range (bb ℓ K)).flatMap fun u => genOps (NLogN.KFV.pre2 D ℓ K long u).length 0 (xs2 D ℓ K u))

/-- The arrangement during phase 2. -/
def arr2 (u : ℕ) : List ℕ := (List.finRange D).flatMap fun i => restA D ℓ K i ++ loS D ℓ K u i

end Paper

end NLogN.KFV

end

section
-- Source module: Solutions.KFV27Chunks
/-!
# KFV §31: the chunk phase of the layout, a selection sort by Lemma 12 exchanges

After the single-bit moves (`KFV25Moves`) the `N = D q` chunks of `K` bits are in axis-major
order; Lemma 16's layout wants them chunk-major. Position `t` of the target holds the chunk
originally at `src t = (t mod D) q + ⌊t / D⌋`. Selection sort: at step `t` the wanted chunk sits
at the *chase* position `x_t`: start at `src t` and apply `src` while the value is `< t`
(`chF`, fuel `N`); exchange positions `t` and `x_t` (nothing if equal). At most `N` exchanges,
and the positions are computed by a fixed arithmetic loop, so a fixed program can do it.

* `FH t y z`: `z` is the first value `≥ t` on the `src`-orbit of `y`; **`chF_FH`** (the fuelled
  chase finds it), **`FH_bound`** (within `N` steps when `src` is a permutation of `[0, N)`).
* `piS t`: the arrangement after `t` steps (position ↦ original index); **`piS_final`**:
  `piS N s = src s` for `s < N`.
* `arrC t`: the slot list after `t` steps; **`arrC_step`**: one step is the bit-level exchange
  `LOp.xch` (or nothing); **`applyN_xchOps`**: all steps give `Lst ++ chunkMajor.flatten`.
-/

set_option linter.unusedSectionVars false
set_option linter.unusedSimpArgs false

namespace NLogN.KFV

open Synth AModel AModelL12

section Chase

variable (src : ℕ → ℕ)

/-- The fuelled chase: `m` times `y ← src y` while `y < t`. -/
def chF (t : ℕ) : ℕ → ℕ → ℕ
  | 0, y => y
  | m + 1, y => chF t m (if y < t then src y else y)

/-- `z` is the first value `≥ t` on the orbit of `y`. -/
def FH (t y z : ℕ) : Prop := ∃ k, src^[k] y = z ∧ t ≤ z ∧ ∀ i < k, src^[i] y < t

variable {N : ℕ} (hsrc : ∀ x < N, src x < N)
  (hinj : ∀ x < N, ∀ y < N, src x = src y → x = y)
include hsrc

include hinj

end Chase

/-! ## The selection sort -/

section SelSort

variable (src : ℕ → ℕ) (N : ℕ)

/-- The exchange partner of step `t`. -/
def xT (t : ℕ) : ℕ := chF src t N (src t)

/-- A transposition of positions. -/
def swF (a b s : ℕ) : ℕ := if s = a then b else if s = b then a else s

/-- The arrangement after `t` steps: position ↦ original index. -/
def piS : ℕ → ℕ → ℕ
  | 0 => id
  | t + 1 => fun s => piS t (NLogN.KFV.swF t (xT src N t) s)

/-- The sort invariant. -/
def SInv (t : ℕ) : Prop :=
  (∀ s < t, NLogN.KFV.piS src N t s = src s) ∧
    ∀ u, t ≤ u → u < N → ∃ z, FH src t (src u) z ∧ NLogN.KFV.piS src N t z = src u

variable {src N} (hsrc : ∀ x < N, src x < N)
  (hinj : ∀ x < N, ∀ y < N, src x = src y → x = y)
include hsrc hinj

end SelSort

/-! ## The slot lists -/

section Lists

variable (src : ℕ → ℕ) (N : ℕ) (am : List (List ℕ)) (K : ℕ) (Lst : List ℕ)

/-- The slot list after `t` steps. -/
def arrC (t : ℕ) : List ℕ := Lst ++ ((List.range N).map fun s => am.getD (NLogN.KFV.piS src N t s) []).flatten

/-- The exchange of step `t` (nothing if the chunk is in place). -/
def opT (t : ℕ) : List LOp :=
  if xT src N t = t then [] else [LOp.xch (Lst.length + t * K) (Lst.length + xT src N t * K) K]

/-- All exchanges. -/
def xchOps : List LOp := (List.range N).flatMap (opT src N K Lst)

variable {src N} (hsrc : ∀ x < N, src x < N)
  (hinj : ∀ x < N, ∀ y < N, src x = src y → x = y)

variable {am K Lst} (ham : ∀ c ∈ am, c.length = K) (hamN : am.length = N)
include hsrc hinj ham hamN

end Lists

/-! ## The paper's chunks -/

section PaperChunks

variable (D ℓ K : ℕ) (long : Fin D → Bool)

/-- The source of target position `t`: chunk `t mod D` (axis) `⌊t/D⌋` (index). -/
def srcT (q t : ℕ) : ℕ := (t % D) * q + t / D

/-- The whole layout schedule: single-bit moves, then the chunk sort. -/
def layOps : List LOp :=
  moves12 D ℓ K long ++ xchOps (NLogN.KFV.srcT D (Synth.qq ℓ K)) (D * Synth.qq ℓ K) K (lstL D ℓ K long)

end PaperChunks

end NLogN.KFV

end

section
-- Source module: Solutions.KFV22Chunk
/-!
# KFV §26: the chunk rounds through the uniform Proposition 14

A chunk round `b ≤ h < ℓ - 1` pushes the descriptor
`desc14 (2^A) K D (2^B) ρ r p W Z` on stack `36` (`A = |Lst| + j₀ D K`, `B = (q - 1 - j₀) D K`,
`j₀ = (qK - 1 - (h - b)) / K`, `ρ = (h - b) mod K`, `r = 2^ℓ`; `p, W, Z` from the registers
`gPp, gWw, gZz` of the base), calls the Proposition 14 entry `e14`, and pops the descriptor.
**`layerU_chunk`**: given `P14Need` for `e14` (the uniform Proposition 14 of KFW), this is the
computed layer of round `h`.
-/

set_option linter.unusedSectionVars false
set_option linter.unusedSimpArgs false
set_option linter.unusedVariables false

namespace NLogN.KFV

open IntMul RS KIL KFD KFB KFH AModel Butterfly

variable {Pp : ℕ} {I : List Sym} {iprog : Fin Pp → IStmt HS Pp}

abbrev gPp : Fin HS := 30
abbrev gWw : Fin HS := 31
abbrev gZz : Fin HS := 32
abbrev dD : Fin HS := 36

/-- The precision, width and size bound of the chunk rounds, in the base. -/
def XR (s : IState HS) (p W Z : ℕ) : Prop :=
  (∃ L, s.G gPp = .int (p : ℤ) :: L) ∧ (∃ L, s.G gWw = .int (W : ℤ) :: L) ∧
    (∃ L, s.G gZz = .int (Z : ℤ) :: L)

/-! ## The descriptor -/

/-- `qK - 1 - h + b`. -/
def eEi : Ex HS :=
  .add (.add (.mul (.reg gQq) (.reg gKk)) (.lit (-1))) (.add (.neg (.reg gH)) (.reg gBb))

/-- `j₀ = e / K`. -/
def eJg : Ex HS := .ndiv eEi (.reg gKk)

/-- `A = |Lst| + j₀ D K`. -/
def ePa : Ex HS := .add (.reg gLst) (.mul (.mul eJg (.reg NLogN.KFV.gD)) (.reg gKk))

/-- `B = (q - 1 - j₀) D K`. -/
def eSb : Ex HS := .mul (.mul (.add (.add (.reg gQq) (.lit (-1))) (.neg eJg)) (.reg NLogN.KFV.gD)) (.reg gKk)

/-- `ρ = (h - b) mod K`. -/
def eRo : Ex HS := .nmod (.add (.reg gH) (.neg (.reg gBb))) (.reg gKk)

/-- The descriptor pushes (`Z` first, `P` last on top). -/
def dPush : RS :=
  .seq (.push dD (.reg gZz)) (.seq (.push dD (.reg gWw)) (.seq (.push dD (.reg gPp))
  (.seq (.push dD (.pow2 (.reg gLl))) (.seq (.push dD eRo) (.seq (.push dD (.pow2 eSb))
  (.seq (.push dD (.reg NLogN.KFV.gD)) (.seq (.push dD (.reg gKk)) (.push dD (.pow2 ePa)))))))))

/-- The descriptor pops. -/
def dPop : RS :=
  .seq (.pop dD) (.seq (.pop dD) (.seq (.pop dD) (.seq (.pop dD) (.seq (.pop dD)
  (.seq (.pop dD) (.seq (.pop dD) (.seq (.pop dD) (.pop dD))))))))

/-- **The chunk round code.** -/
def chunkC (e14 : Fin Pp) : IStmt HS Pp := .seq (fcode dPush) (.seq (.call e14) (fcode dPop))

def cDPush (B : ℕ) : ℕ := 9 * (60 * 10 ^ 7 * (B + 2) ^ 2)

section Paper

open Synth GlueLayout GlueP14 GlueError

variable (D ℓ K : ℕ) [NeZero D] (long : Fin D → Bool)

/-- Cost bound of a chunk round. -/
noncomputable def cChunk (W B : ℕ) (C₀ : ℝ) : ℕ :=
  cDPush B + ⌈C₀ * ((2 ^ NLogN.KFV.nB D ℓ K long * 2 ^ ℓ : ℕ) : ℝ) * (W + 1) *
    (D : ℝ) ^ (1 - (2 : ℝ)⁻¹ ^ 54)⌉₊ + 9 * (2 * B + 2)

end Paper

end NLogN.KFV

end

section
-- Source module: Solutions.KFV23Loop
/-!
# KFV §27: unary round loops over layers

**`rptX_layers`**: a unary loop `rptX X A C D b` (`KFM.rptXV_runs`) whose body maps the array by
`F k` in iteration `k`, on bases `Bk k s1` that move with the iteration (the round register) and
commute with the loop's own stacks `A, C, D`, computes `iterL F n`. The data of each iteration
are fixed by `RecRep_unique`, so the body may be specified for every base.

* `iterL F n = F (n-1) ∘ … ∘ F 0`; `runDown_iterL`, `foldl_iterL`: the forward procedure
  (`runDown`, rounds `n-1, …, 0`) and the opposite one (`foldl`, rounds `0, …, n-1`).
-/

set_option linter.unusedSectionVars false
set_option linter.unusedSimpArgs false
set_option linter.unusedVariables false

namespace NLogN.KFV

open IntMul RS KIL KFD KFB KFH KFC KFT AModel

variable {Pp : ℕ} {I : List Sym} {iprog : Fin Pp → IStmt HS Pp}

/-- `F (n-1) ∘ … ∘ F 0`. -/
def iterL {α : Type*} (F : ℕ → α → α) : ℕ → α → α
  | 0, x => x
  | n + 1, x => F n (iterL F n x)

section Loop

variable {r : ℕ}

/-- The body of a layer loop: iteration `k` maps the array by `F k` and the base by `Bk`. -/
def LoopBody (I : List Sym) (iprog : Fin Pp → IStmt HS Pp) (b : IStmt HS Pp)
    (Q : IState HS → Prop) (Bk : ℕ → IState HS → IState HS)
    (F : ℕ → List (Fin r → ℂ) → List (Fin r → ℂ)) (N p W : ℕ) (cs : ℕ → ℕ) (n : ℕ) : Prop :=
  ∀ k < n, ∀ (s1 : IState HS), Q s1 → ∀ (zs : List GI) (l : List ℤ) (xs : List (Fin r → ℂ)),
    RecRep r p zs xs → xs.length = N → NLogN.KFB.GB W zs → Disk xs →
    ∃ zs', RecRep r p zs' (F k xs) ∧ NLogN.KFB.GB W zs' ∧ (F k xs).length = N ∧ Disk (F k xs) ∧
      IRuns I iprog b (encK (Bk k s1) (o0 (gflat zs ++ l)))
        (encK (Bk (k + 1) s1) (o0 (gflat zs' ++ l))) (cs k)

end Loop

end NLogN.KFV

end