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

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

import Definitions.Def_IntMul_Kappa_K_Defs21

set_option autoImplicit false
set_option relaxedAutoImplicit false
set_option maxHeartbeats 1000000

section
-- Source module: Solutions.KFMU2
/-!
# KFM U2: uniform root and chirp tables

`UNIFORM.md`: the table codes `chirpG`, `twC`, `axG`, `twG` (KFMTabP) and `obB`, `oprodP`,
`levP`, `chirpP`, `chirpTabP` (KFMTabQ) carry `n`-dependent literals (counts `t, 2t, n, 2nt`,
the twiddle arguments `Δ, k, j`, the ruler `2^Δ`, and the axis recursion over `d`). The uniform
codes here read them from registers:

| stack | value |
|---|---|
| `rDl = 24` | `Δ` |
| `rJ1 = 25` | `2t − s` (axis level) / `k` (twist table) |
| `rJ2 = 32` | `t − s` (axis level) / `j` (twist table) |
| `rN = 33` | `n` = current array length |
| `rN2 = 34` | `2·n·t` |
| `nX = 30`, `nX2 = 31` | `t`, `2t` |
| `pR = 45` | `2^Δ` (as before) |

All stacks avoid KFT's `sK, sK2, uY, uB = 26–29` and twiddleP's `28, 29, 53`.
The chirp-array loop adds `aX = 20` (`d`), counters `aA, aC, aD = 21–23`, and the descriptor
stack `dS = 35` holding `descT s t` = the axis records `[t, 2t, 2t−s, t−s, n, 2nt]` then `2∏t`.

**Stack discipline (for kmain).** Every spec here is frame-preserving for arbitrary contents of
other stacks. None of these registers may be live across a Prop 14 / Prop 18 call: KFV's
`kronU`/`recU` use stacks 1–23 and 56 as scratch and `P18IL` requires every stack except `0` and
its descriptor stack to be empty, so kmain drops all of 20–35 before such a call and rebuilds
them afterwards.

* **`chirpPU_runs`** (list form, per-axis costs via `rptXV_runs`/`uloopV_runs`),
  **`chirpTabPU_runs`** (output `chirpArr Δ d s t` on `pW`, as `chirpTabP_runs`),
  **`cChirpTabU_le`**: `≤ d(AX + 4000) + 10^11 (∏t + 1)(φ + Δ + 3)` (geometric sum over the levels),
  **`cAxU_le`**, **`cTwGU_le`**: the old costs plus `O(Δ + t)`.

* **`rptX_of_rpt`**: any run of `rpt A C D n b` gives a run of `rptX X A C D b` from a state whose
  `X` holds `n` (the leading `const` is replaced by a `dup`). Every literal-count lemma of
  KFMTabP/Q converts this way.
* **`chirpGU_runs`**, **`twCU_runs`**, **`axGU_runs`**, **`twGU_runs`**.
-/

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

namespace NLogN.KFM

open IntMul KFC KFB KFT KIL

abbrev rDl : Fin HS := 24
abbrev rJ1 : Fin HS := 25
abbrev rJ2 : Fin HS := 32
abbrev rN : Fin HS := 33
abbrev rN2 : Fin HS := 34

/-! ## The chirp array over a list of axes -/

/-- `chirpArr` over a list of axes `(s, t)` (first axis outermost). -/
def chirpArrL (Δ : ℕ) (L : List (ℕ × ℕ)) : List GI :=
  L.foldl (fun A p => oprod Δ A (axTab Δ p.1 p.2)) [oneP Δ]

/-! ## The axis descriptor -/

abbrev aX : Fin HS := 20
abbrev aA : Fin HS := 21
abbrev aC : Fin HS := 22
abbrev aD : Fin HS := 23
abbrev dS : Fin HS := 35

/-- One axis record `[t, 2t, 2t − s, t − s, n, 2nt]` (top first). -/
def tupL (s t n : ℕ) : List ℤ :=
  [(t : ℤ), ((2 * t : ℕ) : ℤ), ((2 * t - s : ℕ) : ℤ), ((t - s : ℕ) : ℤ), (n : ℤ),
    ((2 * (n * t) : ℕ) : ℤ)]

/-- The axis records, `n` = the array length before each axis. -/
def descA : ℕ → List (ℕ × ℕ) → List ℤ
  | _, [] => []
  | n, p :: L => tupL p.1 p.2 n ++ descA (n * p.2) L

/-- The array length before axis `k`. -/
def nK (L : List (ℕ × ℕ)) (k : ℕ) : ℕ := ((L.take k).map Prod.snd).prod

section IL

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

/-! ## Repetition with per-iteration costs -/

/-! ## The root-table loop -/

/-- `chirpG` with the count on `nX`. -/
def chirpGU (pM : Fin P) : IStmt HS P :=
  .seq (rptX nX pA pC pD (chirpB pM)) (.seq (rptX nX pA pC pD (mvPB pV pW)) drop6B)

/-! ## Base roots by `twiddleP`, arguments from registers -/

/-- Copy `n, k, j` from `rn, rk, rj` onto `28, 29, 53`, run `twiddleP`, drop, and move the root
onto `X` (`re` on top). -/
def twCU (rn rk rj X : Fin HS) : IStmt HS P :=
  .seq (.op (.dup rn 28)) (.seq (.op (.dup rk 29)) (.seq (.op (.dup rj 53)) (.seq (RS.code twiddleP)
    (.seq (.op (.drop 53)) (.seq (.op (.drop 29)) (.seq (.op (.drop 28))
      (.seq (.op (.mv 52 X)) (.op (.mv 51 X)))))))))

def cTwU (n k j : ℕ) : ℕ := cTw n k j + 3 * cCnt (n + k + j)

/-! ## Axis and twist tables -/

/-- **Uniform axis table**: `Δ` on `rDl`, `2t` on `nX2`, `2t − s` on `rJ1`, `t` on `nX`, `t − s`
on `rJ2`, ruler `2^Δ` on `pR`. -/
def axGU (pM : Fin P) : IStmt HS P :=
  .seq (.seq (.op (.const 0 0)) (.op (.dup pR 0))) (.seq (twCU rDl nX2 rJ1 pH)
    (.seq (twCU rDl nX rJ2 pK) (chirpGU pM)))

def cAxU (Δ s t φ : ℕ) : ℕ :=
  cTwU Δ (2 * t) (2 * t - s) + cTwU Δ t (t - s) + (cChirpG t φ (Δ + 2) Δ + 2 * cCnt t) +
    (dupCost (Δ + 1) + 17) + 2

/-- **Uniform twist table** `twTab Δ k j n`: `Δ` on `rDl`, `k` on `rJ1`, `j` on `rJ2`, `n` on
`nX`, ruler on `pR`. -/
def twGU (pM : Fin P) : IStmt HS P :=
  .seq (.seq (.op (.const 0 0)) (.op (.dup pR 0))) (.seq (twCU rDl rJ1 rJ2 pH)
    (.seq (.seq (.op (.const pK 0)) (.op (.dup pR pK))) (chirpGU pM)))

def cTwGU (Δ k j n φ : ℕ) : ℕ :=
  cTwU Δ k j + (cChirpG n φ (Δ + 2) Δ + 2 * cCnt n) + 2 * (dupCost (Δ + 1) + 19)

/-! ## The outer-product round, loop and one level -/

/-- `obB` with the counts on `nX` (`b`) and `nX2` (`2b`). -/
def obBU (pM : Fin P) : IStmt HS P :=
  .seq (mvPB pQ1 pH) (.seq (rptX nX pA pC pD (cpyB pH 0)) (.seq (rptX nX2 pA pC pD (.op (.mv pQ2 pU)))
    (.seq (rptX nX2 pA pC pD (dmB pU pW pQ2)) (.seq (gmulPU pM)
    (.seq (rptX nX2 pA pC pD (.op (.mv 0 pQ3))) (.seq (.op (.drop pH)) (.op (.drop pH))))))))

def cObU (b w Δ φ : ℕ) : ℕ := cOb b w Δ φ + 6 * cCnt (2 * b)

/-- `oprodP` with the count `n = |A|` on `rN`. -/
def oprodPU (pM : Fin P) : IStmt HS P := rptX rN NLogN.KFM.oA NLogN.KFM.oC NLogN.KFM.oD (obBU pM)

/-- **One uniform level**: registers as for `axGU`, plus `n = |A|` on `rN` and `2nt` on `rN2`. -/
def levPU (pM : Fin P) : IStmt HS P :=
  .seq (axGU pM) (.seq (rptX nX2 pA pC pD (.op (.mv pW pU)))
    (.seq (rptX nX2 pA pC pD (.op (.mv pU pQ2))) (.seq (oprodPU pM)
    (.seq (rptX rN2 pA pC pD (.op (.mv pQ3 pQ1)))
      (rptX nX2 pA pC pD (.op (.drop pQ2)))))))

def cLevU (Δ s t n φ : ℕ) : ℕ :=
  cAxU Δ s t φ + 2 * ((2 * t) * (110 * (Δ + 2 + 1) + 64) + 60 + cRpt (2 * t)) +
    (n * (cObU t (Δ + 2) Δ φ + 64) + 60 + cRpt n + cCnt n) +
    ((2 * (n * t)) * (110 * (Δ + 2 + 1) + 64) + 60 + cRpt (2 * (n * t))) +
    ((2 * t) * (2 * (Δ + 2) + 2 + 64) + 60 + cRpt (2 * t)) + 3 * cCnt (2 * t) + cCnt (2 * (n * t))

/-! ## One axis step: load the record, run the level, drop the registers -/

def ldB : IStmt HS P :=
  .seq (.op (.mv NLogN.KFM.dS nX)) (.seq (.op (.mv NLogN.KFM.dS nX2)) (.seq (.op (.mv NLogN.KFM.dS rJ1)) (.seq (.op (.mv NLogN.KFM.dS rJ2))
    (.seq (.op (.mv NLogN.KFM.dS rN)) (.op (.mv NLogN.KFM.dS rN2))))))

def dropB6 : IStmt HS P :=
  .seq (.op (.drop rN2)) (.seq (.op (.drop rN)) (.seq (.op (.drop rJ2)) (.seq (.op (.drop rJ1))
    (.seq (.op (.drop nX2)) (.op (.drop nX))))))

/-- The size of the largest record entry. -/
def zV (t n : ℕ) : ℕ := zsz ((2 * t + 2 * (n * t) + n : ℕ) : ℤ)

def cLdV (t n : ℕ) : ℕ := 700 * (zV t n + 1)

/-- One axis step of the uniform chirp loop. -/
def levB (pM : Fin P) : IStmt HS P := .seq ldB (.seq (levPU pM) dropB6)

/-! ## The uniform chirp-array loop -/

/-- **Uniform chirp array**: `d` on `aX`, the records `descA 1 L` on `dS`, `Δ` on `rDl`, ruler on
`pR`; counters `aA, aC, aD`. -/
def chirpPU (pM : Fin P) : IStmt HS P :=
  .seq (.seq (.op (.const pQ1 0)) (.op (.dup pR pQ1))) (rptX aX aA aC NLogN.KFM.aD (levB pM))

/-- The cost of axis step `k`. -/
def cStep (Δ φ : ℕ) (L : List (ℕ × ℕ)) (k : ℕ) : ℕ :=
  cLevU Δ (L.getD k (0, 0)).1 (L.getD k (0, 0)).2 (nK L k) φ + cLdV (L.getD k (0, 0)).2 (nK L k)

def cChirpPU (Δ φ : ℕ) (L : List (ℕ × ℕ)) : ℕ :=
  2 + (dupCost (Δ + 1) + 17) + ((∑ k ∈ Finset.range L.length, (cStep Δ φ L k + 64)) + 60 +
    cRpt L.length + cCnt L.length)

/-! ## The uniform chirp table -/

/-- The axis list of `(s, t)`. -/
def axL {d : ℕ} (s t : Fin d → ℕ) : List (ℕ × ℕ) := List.ofFn fun i => (s i, t i)

/-- The chirp-table descriptor: the axis records, then `2 ∏ t_i`. -/
def descT {d : ℕ} (s t : Fin d → ℕ) : List ℤ := descA 1 (axL s t) ++ [((2 * ∏ i, t i : ℕ) : ℤ)]

/-- **Uniform chirp table** onto `pW`; consumes the descriptor on `dS`. -/
def chirpTabPU (pM : Fin P) : IStmt HS P :=
  .seq (chirpPU pM) (.seq (rptX NLogN.KFM.dS pA pC pD (.op (.mv pQ1 pU)))
    (.seq (rptX NLogN.KFM.dS pA pC pD (.op (.mv pU pW))) (.op (.drop NLogN.KFM.dS))))

def cChirpTabU (Δ φ : ℕ) (L : List (ℕ × ℕ)) (N : ℕ) : ℕ :=
  cChirpPU Δ φ L + 2 * ((2 * N) * (110 * (Δ + 2 + 1) + 64) + 60 + cRpt (2 * N)) +
    2 * cCnt (2 * N) + (2 * zsz ((2 * N : ℕ) : ℤ) + 2)

/-! ## Cost bounds -/

end IL

end NLogN.KFM

end

section
-- Source module: Solutions.KFMU3
/-!
# KFM U3: uniform digits in and carries out

`UNIFORM.md`: KFT's `digI b e M S` and `carI b N` push the literals `2^b`, `2^e` and use the
literal counts `M`, `2(S − M)`, `N`. The uniform codes copy them from registers:

| stack | value |
|---|---|
| `uB2 = 24` | `2^b` |
| `uM = 25` | `M` (digits) / `N` (carries) |
| `uZ = 32` | `2(S − M)` |
| `uE = 33` | `2^e` |

(the same stacks as the table registers of KFMU2; the two are used in different phases).
The bodies (`cutB`, `zB`, `putB`, `c1B`, `c2B`) and the scratch stacks 56–62 are KFT's, unchanged.
The proofs are KFT's with each `rpt` converted by `rptX_of_rpt` and each `const` by a `dup`.

* **`digIU_spec`**: as `digI_spec`, cost `cDigU = cDig + O(b + e) + O(log M + log S)`.
* **`carIU_spec`**: as `carI_spec`, cost `cCarU = cCar + O(b + log N)`.
-/

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

namespace NLogN.KFM

open IntMul KFC KFB KFT

abbrev uB2 : Fin HS := 24
abbrev uM : Fin HS := 25
abbrev uZ : Fin HS := 32
abbrev uE : Fin HS := 33

section IL

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

/-- **Uniform digits in.** -/
def digIU : IStmt HS P :=
  .seq (.op (.dup uB2 tR)) (.seq (rptX uM tA tC tD cutB) (.seq (.op (.drop tX))
    (.seq (rptX uZ tA tC tD NLogN.KFT.zB) (.seq (.op (.dup uE KFT.tE))
    (.seq (rptX uM tA tC tD putB) (.seq (.op (.drop KFT.tE)) (.op (.drop tR))))))))

def cDigU (b e M S : ℕ) : ℕ :=
  cDig b e M S + (dupCost (b + 1) + 17) + (dupCost (e + 1) + 17) + 2 * cCnt M + cCnt (2 * (S - M))

/-- **Uniform carries out.** -/
def carIU : IStmt HS P :=
  .seq (.op (.dup uB2 tR)) (.seq (.op (.const tY 0)) (.seq (rptX uM tA tC 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 (rptX uM tA tC tD c2B) (.seq (.op (.neg NLogN.KFT.tP)) (.seq (.op (.add tY NLogN.KFT.tP 0))
    (.op (.drop tR))))))))))

def cCarU (b B N : ℕ) : ℕ := cCar b B N + (dupCost (b + 1) + 17) + 2 * cCnt N

end IL

end NLogN.KFM

end

section
-- Source module: Solutions.KFMU4
/-!
# KFM U4: `lg n` from the ruler `2^n` in linear time

Setup step 0 needs `b = lg n = max (clog 2 n) 1` from the input. The IL `read` leaves the ruler
`2^n` (not `n`), and a bit-length procedure on an `n`-bit number costs `Θ(n log n)`, which is over
budget. Instead `lgLoop` splits `2^n` by the rulers `2^{2^0}, 2^{2^1}, …` (each squared by one `shl`),
so the `k`-th step costs `O(2^k)` and the loop `O(n)`:

* state `k`: `gA = 2^{n+1−2^k}` (or `0` once `2^k > n+1`), `gR = 2^{2^k}`, `gP = 2^{2^{k−1}}`, `gJ = k`;
* it stops at `J = Nat.size (n+1)` (`2^{J−1} ≤ n+1 < 2^J`).

**`lgP_runs`**: `lgP` (init, loop, one exact test `n ≤ 2^{J−1}` by a split of a copy of `2^n`)
pushes `lg n` on `gJ` for `n ≥ 3`, everything else restored, cost `≤ 10^6 (n + 1)`.

Stacks (KFM range): `gA = 20`, `gR = 21`, `gP = 22`, `gJ = 23`, `gT = 34`, input copy on `gN = 35`.
-/

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

namespace NLogN.KFM

open IntMul KFC KFB KFT

abbrev gA : Fin HS := 20
abbrev gR : Fin HS := 21
abbrev gP : Fin HS := 22
abbrev gJ : Fin HS := 23
abbrev gT : Fin HS := 34
abbrev gN : Fin HS := 35

/-! ## Arithmetic -/

/-- The remaining power after `k` splits. -/
def aK (n k : ℕ) : ℕ := if 2 ^ k ≤ n + 1 then 2 ^ (n + 1 - 2 ^ k) else 0

/-- The previous ruler exponent's power. -/
def prvK (k : ℕ) : ℕ := if k = 0 then 0 else 2 ^ (2 ^ (k - 1))

/-! ## One step -/

section IL

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

def lgB : IStmt HS P :=
  .seq (.op (.drop gP)) (.seq (.op (.dup gR gP)) (.seq (.op (.split gA gR gT)) (.seq (.op (.drop gT))
    (.seq (.op (.dup gR gT)) (.seq (.op (.shl gR gT)) (.seq (.op (.drop gT)) (.seq (.op (.const gT 1))
      (.op (.add gJ gT gJ)))))))))

/-- The loop state. -/
def lgS (s : IState HS) (LA LR LP LJ : List Item) (n k : ℕ) : IState HS :=
  (((s.upd gA (.int (aK n k : ℤ) :: LA)).upd gR (.int ((2 : ℤ) ^ (2 ^ k)) :: LR)).upd gP
    (.int (prvK k : ℤ) :: LP)).upd gJ (.int (k : ℤ) :: LJ)

def cLgB (k : ℕ) : ℕ := 10 ^ 4 * (2 ^ k + k + 2)

/-! ## The loop -/

def isPosC (c : Cls) : Bool := c == .pos

def lgLoop : IStmt HS P := .loop gA isPosC lgB

/-! ## The exact test and the full procedure -/

def lgPost : IStmt HS P :=
  .seq (.seq (.op (.drop gA)) (.op (.drop gR))) (.seq (.seq (.op (.const gT 2)) (.seq (.op (.shl gP gT))
    (.op (.drop gT)))) (.seq (.op (.dup gN gA)) (.seq (.op (.split gA gP gT))
    (.seq (.op (.drop gT)) (.seq (.ite gA isPosC .skip (.seq (.op (.const gT (-1)))
      (.op (.add gJ gT gJ)))) (.seq (.op (.drop gA)) (.op (.drop gP))))))))

def lgInit : IStmt HS P :=
  .seq (.op (.dup gN gA)) (.seq (.op (.const gR 2)) (.seq (.op (.const gP 0)) (.op (.const gJ 0))))

/-- **`lg n` from the ruler `2^n` on `gN`.** -/
def lgP : IStmt HS P := .seq lgInit (.seq lgLoop lgPost)

/-- The value left by the test. -/
def lgV (n : ℕ) : ℕ := if 2 ^ (Nat.size (n + 1) - 1) + 1 ≤ n then Nat.size (n + 1) else Nat.size (n + 1) - 1

end IL

end NLogN.KFM

end

section
-- Source module: Solutions.KFMU5
/-!
# KFM U5: `n` in binary from the ruler `2^n`, in linear time

The IL `read` leaves `val x` and the ruler `2^n`, but not `n`. Setup needs `n` itself, for
`L = LP n` (`4n ≤ 2^L b`), the digit counts and the small-`n` test. Counting by `+1` costs
`Θ(n log n)`. `nP` works like `lgP` (KFMU4):

1. **phase 1** (`nLoop`): the rulers `R_k = 2^{2^k}` are squared up to `J = size(n+1)` while a
   copy of `2^n` is divided by them (`O(2^k)` per step). Every ruler is **kept** on `gS`, above a
   sentinel `0`;
2. **phase 2** (`nLoop2`): a second copy of `2^n` is split greedily by `R_{J-1}, …, R_0`. Each
   split by `R_k` gives bit `k` of `n` (quotient nonzero) and costs `O(2^k)`, which is independent of
   the size of the number. The bits are collected by Horner on `gJ`.

**`nP_runs`**: from `2^n` on `gN = 35`, `nP` pushes `n` on `gJ = 23`, everything else restored,
cost `≤ 10^6 (n + 1)`. Scratch `20–22`, `26`, `34`.
-/

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

namespace NLogN.KFM

open IntMul KFC KFB KFT

abbrev gS : Fin HS := 26

/-- The kept rulers `prvK (k-1), …, prvK 0` (top first; `prvK 0 = 0` is the sentinel). -/
def rl (k : ℕ) : List ℤ := (List.range k).reverse.map fun j => (prvK j : ℤ)

section IL

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

/-! ## Phase 1 -/

def nB : IStmt HS P :=
  .seq (.op (.mv gP gS)) (.seq (.op (.dup gR gP)) (.seq (.op (.split gA gR gT)) (.seq (.op (.drop gT))
    (.seq (.op (.dup gR gT)) (.seq (.op (.shl gR gT)) (.op (.drop gT)))))))

def nS1 (s : IState HS) (LA LR LP LS : List Item) (n k : ℕ) : IState HS :=
  (((s.upd gA (.int (aK n k : ℤ) :: LA)).upd gR (.int ((2 : ℤ) ^ (2 ^ k)) :: LR)).upd gP
    (.int (prvK k : ℤ) :: LP)).upd gS (ints (rl k) ++ LS)

def cNB (k : ℕ) : ℕ := 10 ^ 4 * (2 ^ k + 2)

def nLoop : IStmt HS P := .loop gA isPosC nB

/-! ## Phase 2 -/

/-- Phase-2 state at level `k` (rulers `R_{k-1}, …, R_0` and the sentinel left on `gS`):
`gA = 2^{n mod 2^k}`, `gJ = ⌊n / 2^k⌋`. -/
def nT (s : IState HS) (LA LS LJ : List Item) (n k : ℕ) : IState HS :=
  ((s.upd gA (.int ((2 : ℤ) ^ (n % 2 ^ k)) :: LA)).upd gS (ints (rl (k + 1)) ++ LS)).upd gJ
    (.int ((n / 2 ^ k : ℕ) : ℤ) :: LJ)

def nB2 : IStmt HS P :=
  .seq (.op (.mv gS gR)) (.seq (.op (.const gT 2)) (.seq (.op (.shl gJ gT)) (.seq (.op (.drop gT))
    (.seq (.op (.split gA gR gT))
      (.seq (.ite gA isPosC (.seq (.op (.drop gT)) (.seq (.op (.const gT 1)) (.op (.add gJ gT gJ))))
        (.seq (.op (.drop gA)) (.op (.mv gT gA))))
      (.op (.drop gR)))))))

def cNB2 (k J : ℕ) : ℕ := 10 ^ 4 * (2 ^ k + J + 2)

def nLoop2 : IStmt HS P := .loop gS isPosC nB2

/-! ## Init, middle, post, and the whole procedure -/

def nInit : IStmt HS P := .seq (.op (.dup gN gA)) (.seq (.op (.const gR 2)) (.op (.const gP 0)))

def nMid : IStmt HS P :=
  .seq (.op (.drop gA)) (.seq (.op (.drop gR)) (.seq (.op (.mv gP gS)) (.seq (.op (.dup gN gA))
    (.op (.const gJ 0)))))

def nPost : IStmt HS P := .seq (.op (.drop gS)) (.op (.drop gA))

/-- **`n` from the ruler `2^n` on `gN`.** -/
def nP : IStmt HS P := .seq nInit (.seq nLoop (.seq nMid (.seq nLoop2 nPost)))

end IL

end NLogN.KFM

end

section
-- Source module: Solutions.KFMS0
/-!
# KFM S0: generic counted loops of register programs (`RS.rep`) for the setup arithmetic

The setup parameters (§9.2: `d, K, α, L, ℓ, r, s_i, μ_i, …`) are small numbers (`O(log n)` bits)
computed by register programs `RS` (nlogn `HReg`; `RS.code_runs` gives the IL run). This file
proves one generic loop lemma, so that every parameter program is a short instance:

* **`genLoop_run`**: `rep i n body` with `n = N`, where on the iteration-`k` state `lpSt … k c`
  (counter `N − k` on `i`, accumulator `c` on `a`) `body` only replaces the accumulator by
  `G k c`. The result is `accSeq G a0 N` on `a`; all else unchanged.
* **`genLoop_cost`**: cost `≤ (|n|+1)·10^7(B+2)² + N·(C + 4·10^7(B+2)² + 30) + 32` when each
  body run costs `≤ C`.
* Instances: **`sumSeq`** (`G k c = c + g (N − k)`, so the result is `a0 + Σ_{x=1}^{N} g x`,
  **`accSeq_sum`**) and **`satSq`** (`G k c = min(c², cap)`, `m` squarings give
  `min(x^{2^m}, cap)`, **`accSeq_satSq`**).
-/

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

namespace NLogN.KFM

open IntMul RS

/-- The accumulator sequence `a0, G 0 a0, G 1 (G 0 a0), …`. -/
def accSeq (G : ℕ → ℤ → ℤ) (a0 : ℤ) : ℕ → ℤ
  | 0 => a0
  | k + 1 => G k (accSeq G a0 k)

/-- The state at the start of iteration `k`: counter `N − k` on `i`, accumulator `c` on `a`. -/
def lpSt (s : IState HS) (i a : Fin HS) (La : List Item) (N k : ℕ) (c : ℤ) : IState HS :=
  (s.upd i (.int ((N : ℤ) - k) :: s.G i)).upd a (.int c :: La)

/-! ## Instances -/

/-- Sums: `G k c = c + g (N - k)`. -/
def sumG (g : ℕ → ℤ) (N : ℕ) : ℕ → ℤ → ℤ := fun k c => c + g (N - k)

/-- Saturated squaring `c ↦ min(c², cap)` (as `sel (cap - c²) cap c²`). -/
def satG (cap : ℤ) : ℕ → ℤ → ℤ := fun _ c => if cap - c * c < 0 then cap else c * c

end NLogN.KFM

end

section
-- Source module: Solutions.KFMS1
/-!
# KFM S1: integer roots by counting (setup of `d = ⌊b^{2^{-75}}⌋` and `K = ⌊d^{2^{-56}}⌋`)

`rootP m rv rd` pushes onto `rd` the number of `x ∈ [1, v]` with `x^{2^m} ≤ v`, where `v ≥ 1` is
the top of `rv`. Each test makes `m` saturated squarings `y ↦ min(y², v+1)` (`m` is a literal, a
true constant: `75` or `56`), so every value has `O(log v)` bits. The predicate is downward closed,
so the count is the largest such `x` (**`count_eq_findGreatest`**), i.e. the integer
`2^m`-th root of `v`.

* **`rootP_run`**: the run (`RS`), all stacks but `rd` restored.
* **`count_root_spec`**: `c^{2^m} ≤ v < (c+1)^{2^m}`.
Scratch: `sI = 20` (outer counter), `sY = 21`, `sJ = 22` (inner counter).
-/

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

namespace NLogN.KFM

open IntMul RS

abbrev sI : Fin HS := 20
abbrev sY : Fin HS := 21
abbrev sJ : Fin HS := 22

/-! ## Counting a downward-closed predicate -/

/-! ## The program -/

/-- `min(y², v+1)` as an expression (`v` on `rv`, `y` on `sY`). -/
def satE (rv : Fin HS) : Ex HS :=
  .sel (.add (.add (.reg rv) (.lit 1)) (.neg (.mul (.reg sY) (.reg sY)))) (.add (.reg rv) (.lit 1))
    (.mul (.reg sY) (.reg sY))

/-- The root body: test one candidate `x` (the counter). -/
def rootB (m : ℕ) (rv rd : Fin HS) : RS :=
  .seq (.push sY (.reg NLogN.KFM.sI)) (.seq (.rep sJ (.lit m) (.asg sY (satE rv)))
    (.seq (.asg rd (.add (.reg rd) (.sel (.add (.reg rv) (.neg (.reg sY))) (.lit 0) (.lit 1))))
      (.pop sY)))

/-- **`rootP m rv rd`**: push `#{x ∈ [1, v] : x^{2^m} ≤ v}` onto `rd`. -/
def rootP (m : ℕ) (rv rd : Fin HS) : RS :=
  .seq (.push rd (.lit 0)) (.rep NLogN.KFM.sI (.reg rv) (rootB m rv rd))

/-- The indicator of the test. -/
def rootInd (m v : ℕ) (x : ℕ) : ℤ := if x ^ (2 ^ m) ≤ v then 1 else 0

/-- Register hypotheses of `rootP`. -/
structure RootRegs (rv rd : Fin HS) : Prop where
  v56 : rv.val < 56
  d56 : rd.val < 56
  vz : rv.val < 56
  vI : rv ≠ NLogN.KFM.sI
  vY : rv ≠ sY
  vJ : rv ≠ sJ
  dI : rd ≠ NLogN.KFM.sI
  dY : rd ≠ sY
  dJ : rd ≠ sJ
  vd : rv ≠ rd

/-! ## Runs -/

/-! ## Cost -/

/-- The bit bound of the root computation for the input `v`. -/
def Bv (v : ℕ) : ℕ := 2 * zsz ((v + 1 : ℕ) : ℤ) + 2

end NLogN.KFM

end

section
-- Source module: Solutions.KFMS2
/-!
# KFM S2: counting loops with an expression test, and products

* **`cntP n tst rd`**: push onto `rd` the number of `x ∈ [1, N]` (`N = n`) for which the test
  expression `tst` (read with the candidate `x` on top of `sI = 20`) is negative.
  **`cntP_run`**, **`cntP_cost`**.
* **`prodP n rd`**: push `∏_{x=1}^{N} x = N!` onto `rd`. **`prodP_run`**, **`prodP_cost`**.
* `Ex.Bnd_upd`: bounds are unchanged by updating a stack the expression does not read.
-/

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

namespace NLogN.KFM

open IntMul RS

/-! ## Counting with an expression test -/

/-- The count body. -/
def cntB (tst : Ex HS) (rd : Fin HS) : RS :=
  .asg rd (.add (.reg rd) (.sel tst (.lit 1) (.lit 0)))

def cntP (n tst : Ex HS) (rd : Fin HS) : RS :=
  .seq (.push rd (.lit 0)) (.rep NLogN.KFM.sI n (cntB tst rd))

/-- The state at candidate `x`. -/
def xSt (s : IState HS) (x : ℕ) : IState HS := s.upd NLogN.KFM.sI (.int (x : ℤ) :: s.G NLogN.KFM.sI)

/-! ## Factorials -/

def prodG (N : ℕ) : ℕ → ℤ → ℤ := fun k c => c * ((N : ℤ) - k)

def prodB (rd : Fin HS) : RS := .asg rd (.mul (.reg rd) (.reg NLogN.KFM.sI))

/-- **`prodP n rd`**: push `N!` onto `rd` (`N = n`). -/
def prodP (n : Ex HS) (rd : Fin HS) : RS := .seq (.push rd (.lit 1)) (.rep NLogN.KFM.sI n (prodB rd))

end NLogN.KFM

end

section
-- Source module: Solutions.KFMS3
/-!
# KFM S3: the scalar setup parameters as counts (math side)

For `n ≥ n₀` the §9.2 parameters are the counts computed by `rootP` / `cntP`:

* **`d_count`**: `dP n = findGreatest (x^{2^75} ≤ b) b`;
* **`K_count`**: `KP n = findGreatest (x^{2^56} ≤ d) d`;
* **`α_count`**: `αP n = 1 + findGreatest (x^4 < X) X`, `X = 12 d² b`;
* **`L_count`**: `LP n = findGreatest (2^{x-1} b < 4n) b`.
-/

set_option linter.unusedSimpArgs false

namespace NLogN.KFM

open IntMul IntMul.OpenAI.Glue

/-! ## The setup register program -/

open RS

abbrev zN : Fin HS := 23
abbrev zB : Fin HS := 24
abbrev zD : Fin HS := 25
abbrev zK : Fin HS := 26
abbrev zAl : Fin HS := 27
abbrev zL : Fin HS := 28
abbrev zEl : Fin HS := 29
abbrev zR : Fin HS := 30
abbrev zDp : Fin HS := 31
abbrev zQ : Fin HS := 32
abbrev zMh : Fin HS := 33
abbrev zMf : Fin HS := 34

/-- `X = 12 d² b`. -/
def XE : Ex HS := .mul (.mul (.lit 12) (.mul (.reg NLogN.KFM.zD) (.reg NLogN.KFM.zD))) (.reg zB)
/-- `x^4 − X` (`x` on `sI`). -/
def tstA : Ex HS := .add (.mul (.mul (.reg NLogN.KFM.sI) (.reg NLogN.KFM.sI)) (.mul (.reg NLogN.KFM.sI) (.reg NLogN.KFM.sI))) (.neg NLogN.KFM.XE)
/-- `2^{x−1} b − 4n`. -/
def tstL : Ex HS := .add (.mul (.pow2 (.add (.reg NLogN.KFM.sI) (.lit (-1)))) (.reg zB)) (.neg (.mul (.lit 4) (.reg zN)))
/-- `ℓ = ⌊(L + d − 1)/d⌋`. -/
def elE : Ex HS := .ndiv (.add (.add (.reg zL) (.reg NLogN.KFM.zD)) (.lit (-1))) (.reg NLogN.KFM.zD)
/-- `d' = dℓ − L`. -/
def dpE : Ex HS := .add (.mul (.reg NLogN.KFM.zD) (.reg zEl)) (.neg (.reg zL))
/-- `D = 256 d²`. -/
def DE : Ex HS := .mul (.lit 256) (.mul (.reg NLogN.KFM.zD) (.reg NLogN.KFM.zD))
/-- `Q = 2F(⌊r/(2FD)⌋ + 1)` (`F` on `zQ`). -/
def QE : Ex HS :=
  .mul (.mul (.lit 2) (.reg zQ)) (.add (.ndiv (.reg zR) (.mul (.mul (.lit 2) (.reg zQ)) DE)) (.lit 1))
/-- `⌊((4d−1)t − 4d)/(4dQ)⌋` for `t` given by `tE`. -/
def ME (tE : Ex HS) : Ex HS :=
  .ndiv (.add (.mul (.add (.mul (.lit 4) (.reg NLogN.KFM.zD)) (.lit (-1))) tE) (.neg (.mul (.lit 4) (.reg NLogN.KFM.zD))))
    (.mul (.mul (.lit 4) (.reg NLogN.KFM.zD)) (.reg zQ))

/-- Part A: `d`, `K`, `α`. -/
def setupA : RS :=
  .seq (rootP 75 zB NLogN.KFM.zD) (.seq (rootP 56 NLogN.KFM.zD zK) (.seq (cntP NLogN.KFM.XE tstA zAl)
    (.asg zAl (.add (.reg zAl) (.lit 1)))))

/-- Part B: `L`, `ℓ`, `r`, `d'`. -/
def setupB : RS :=
  .seq (cntP (.reg zB) tstL zL) (.seq (.push zEl elE) (.seq (.push zR (.pow2 (.reg zEl)))
    (.push zDp dpE)))

/-- Part C: `Q` (through `F = D!`), `M_{r/2}`, `M_r`. -/
def setupC : RS :=
  .seq (prodP DE zQ) (.seq (.asg zQ QE) (.seq (.push zMh (ME (.ndiv (.reg zR) (.lit 2))))
    (.push zMf (ME (.reg zR)))))

/-- **The scalar setup program** (inputs `n` on `zN`, `b` on `zB`). -/
def setupS : RS := .seq setupA (.seq setupB setupC)

end NLogN.KFM

end

section
-- Source module: Solutions.KFMSm
/-!
# KFM Sm: bit bounds for register expressions from bounds on their atoms

`exSm A s e`: every register, literal and `2^a` that `e` reads is at most `A` in absolute value,
and every divisor is positive. `exDg e` is a degree (sums add one, products add). Then
`|e.eval s| ≤ A^{exDg e}` (**`exSm_abs`**), and `A^{exDg e} < 2^B` gives `e.Bnd s B`
(**`exSm_bnd`**), which is the hypothesis of the RS cost lemmas.
-/

set_option linter.unusedSimpArgs false

namespace NLogN.KFM

open IntMul RS

/-- Atom bounds (see the module doc). -/
def exSm (A : ℕ) (s : IState HS) : Ex HS → Prop
  | .reg i => |s.top i| ≤ (A : ℤ)
  | .lit c => |c| ≤ (A : ℤ)
  | .add a b | .mul a b => exSm A s a ∧ exSm A s b
  | .neg a => exSm A s a
  | .ndiv a b | .nmod a b => exSm A s a ∧ exSm A s b ∧ 0 < b.eval s
  | .pow2 a => exSm A s a ∧ 0 ≤ a.eval s ∧ 2 ^ (a.eval s).toNat ≤ A
  | _ => False

/-- The degree. -/
def exDg : Ex HS → ℕ
  | .reg _ | .lit _ => 1
  | .add a b => max (exDg a) (exDg b) + 1
  | .mul a b => exDg a + exDg b
  | .neg a | .pow2 a => exDg a
  | .ndiv a b | .nmod a b => max (exDg a) (exDg b)
  | _ => 1

end NLogN.KFM

end

section
-- Source module: Solutions.KFMS4
/-!
# KFM S4: the cost of the scalar setup program

For `n ≥ n₀`, with `b = lg n`, every atom the setup expressions read is at most `A = 2^{b+2}`.
The expressions have degree at most `8`, so all values have fewer than `B = 8b+17` bits
(`KFMSm`). This gives

* **`setupA_cost`**, **`setupB_cost`**, **`setupC_cost`**: each at most `10^15 (b+1)^4`.
-/

set_option linter.unusedSimpArgs false
set_option linter.unusedVariables false

namespace NLogN.KFM

open IntMul IntMul.OpenAI.Glue IntMul.OpenAI.RP RS

/-- The state after `setupS`: `d, K, α, L, ℓ, r, d', Q, M_{r/2}, M_r` pushed on `zD … zMf`. -/
noncomputable def setupSt (n : ℕ) (s : IState HS) : IState HS :=
  let p := fun (t : IState HS) (i : Fin HS) (v : ℕ) => t.upd i (.int (v : ℤ) :: s.G i)
  p (p (p (p (p (p (p (p (p (p s NLogN.KFM.zD (dP n)) zK (KP n)) zAl (αP n)) zL (LP n)) zEl (ℓP n)) zR (IntMul.OpenAI.Glue.rP n))
    zDp (mcnt (LP n) (dP n))) zQ (srcQ (dP n) (IntMul.OpenAI.Glue.rP n))) zMh (srcM (dP n) (IntMul.OpenAI.Glue.rP n) (IntMul.OpenAI.Glue.rP n / 2)))
    zMf (srcM (dP n) (IntMul.OpenAI.Glue.rP n) (IntMul.OpenAI.Glue.rP n))

end NLogN.KFM

end

section
-- Source module: Solutions.KFV8Sim
/-!
# KFV §12: literal-to-register simulation of item-language code

`ISim π c c'`: `c'` is `c` with some literal pushes `.const j v` replaced by copies `.dup i j` of a
parameter stack `i` holding `v` (`π i = some v`). It is the `IStmt`-level analogue of
`KIL.RSim` and covers compiled register programs (`Ex.code` compiles `.lit v` to `.const` and
`.reg i` to `.dup`) as well as hand-written codes such as KFT's `carI`/`digI`.

**`ISim.runs`**: if `c` does not write any parameter stack and makes no calls (`NW π c`), every
run of `c` from a state with the parameters on top (`PV π s`) is matched by a run of `c'` with
the same final state and at most `2000 ×` the cost. This turns old codes whose *shape* does not
depend on `n` into uniform codes reading their sizes from stacks.
-/

set_option linter.unusedSectionVars false
set_option linter.unusedVariables false

namespace NLogN.KFV

open IntMul KIL

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

/-- The stacks an operation may change. -/
def opW : IOp HS → List (Fin HS)
  | .add i j k => [i, j, k]
  | .neg i => [i]
  | .dup _ j => [j]
  | .drop i => [i]
  | .mv i j => [i, j]
  | .const i _ => [i]
  | .mark i => [i]
  | .unmark i => [i]
  | .shl i _ => [i]
  | .split i _ k => [i, k]
  | .read i j _ => [i, j]
  | .write i j => [i, j]
  | .cat i _ k => [i, k]

/-- `c` writes no parameter stack and makes no calls. -/
def NW (π : PMap) : IStmt HS Pp → Prop
  | .skip => True
  | .op o => ∀ i ∈ opW o, π i = none
  | .seq a b => NW π a ∧ NW π b
  | .ite _ _ a b => NW π a ∧ NW π b
  | .loop _ _ a => NW π a
  | .call _ => False

/-- `c'` is `c` with some literal pushes replaced by parameter copies (`par`), or by any code
computing the literal from the parameters at most `2000 ×` as expensively (`parC`). -/
inductive ISim (I : List Sym) (iprog : Fin Pp → IStmt HS Pp) (π : PMap) :
    IStmt HS Pp → IStmt HS Pp → Prop
  | skip : ISim I iprog π .skip .skip
  | op (o : IOp HS) : ISim I iprog π (.op o) (.op o)
  | call (j : Fin Pp) : ISim I iprog π (.call j) (.call j)
  | par {i j : Fin HS} {v : ℤ} : π i = some v → i ≠ j →
      ISim I iprog π (.op (.const j v)) (.op (.dup i j))
  | parC {j : Fin HS} {v : ℤ} {d : IStmt HS Pp} :
      (∀ s : IState HS, PV π s → IRuns I iprog d s (s.upd j (.int v :: s.G j)) (2000 * (zsz v + 2))) →
      ISim I iprog π (.op (.const j v)) d
  | seq {a a' b b' : IStmt HS Pp} : ISim I iprog π a a' → ISim I iprog π b b' →
      ISim I iprog π (.seq a b) (.seq a' b')
  | ite {i : Fin HS} {f : Cls → Bool} {a a' b b' : IStmt HS Pp} : ISim I iprog π a a' →
      ISim I iprog π b b' → ISim I iprog π (.ite i f a b) (.ite i f a' b')
  | loop {i : Fin HS} {f : Cls → Bool} {a a' : IStmt HS Pp} : ISim I iprog π a a' →
      ISim I iprog π (.loop i f a) (.loop i f a')

/-! ## Compiled register programs -/

end NLogN.KFV

end

section
-- Source module: Solutions.KFV9Gen
/-!
# KFV §13: generalized literal substitution and the write analysis

* `wl c` (stacks written, calls ignored) and `nc c` (no calls): `NW_of` gives `NW π c` when no
  parameter stack is written. `wl_ecode`: compiled expressions write only the scratch `56…63`;
  `wl_fcode`: compiled register programs write `p.writes`, the scratch, and the `toUL` stacks.
* `EG π e e'`: `e'` is `e` with literals `v` replaced by any clean expression over the parameters
  that evaluates to `v` at cost `≤ 2000 (zsz v + 2)` (e.g. `2r = r + r`, `-K`, `B / 2`).
  `GSim` is the register-program version. **`ISim.fcodeG`**: `GSim π p p'` gives
  `ISim π (fcode p) (fcode p')`. **`ISim.pushE`**: a literal push `const j v` on any stack
  `j ≠ 56` may be replaced by `e.code; mv 56 j`.
-/

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

namespace NLogN.KFV

open IntMul KIL KFD

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

/-! ## Write analysis -/

/-- Stacks written by a code (calls ignored). -/
def wl : IStmt HS Pp → List (Fin HS)
  | .skip => []
  | .op o => opW o
  | .seq a b => wl a ++ wl b
  | .ite _ _ a b => wl a ++ wl b
  | .loop _ _ a => wl a
  | .call _ => []

/-- No calls. -/
def nc : IStmt HS Pp → Bool
  | .call _ => false
  | .seq a b => nc a && nc b
  | .ite _ _ a b => nc a && nc b
  | .loop _ _ a => nc a
  | _ => true

/-! ## Generalized substitution -/

end NLogN.KFV

end

section
-- Source module: Solutions.KFWSwap1
/-!
# KFW: uniform Lemma 12 on integer arrays (parameters in registers)

The integer-array swap `KFS.swW q w Pre G Suf u` (to bits, Lemma 12, back to integers) with its
size-dependent literals read from registers:

* `toBitsU` / `fromBitsU`: `KFL.toBitsI N w` / `fromBitsI N w` with `N`, `w` copied from the
  registers `sN`, `sW` (`ISim`, at most `2000 ×` the cost);
* `swIU q`: the Lemma 12 wrapper `KFL.swI` with the four parameters pushed from the registers
  `sP, sG, sS, sU` (**`swIU_spec`**);
* `swWU q`: the composite (**`swWU_spec`**).
-/

set_option linter.unusedSimpArgs false
set_option linter.unusedVariables false

namespace NLogN.KFW

open IntMul RS KIL KFD KFE KFB KFL KFS KFV AModel AModelL12 NetProp7

/-- Registers: `w`, `u`, `vol Pre`, `vol G`, `vol Suf · (w+1)`, `N`. -/
abbrev sW : Fin HS := 1
abbrev sU : Fin HS := 2
abbrev sP : Fin HS := 3
abbrev sG : Fin HS := 4
abbrev sS : Fin HS := 5
abbrev sN : Fin HS := 6

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

/-! ## Conversions -/

def toBitsU : IStmt HS Pp :=
  .seq (.op (.dup sN cPN)) (.seq (.op (.dup NLogN.KFW.sW cPW)) (.seq (.op (.const cR2 2))
    (.seq (KFC.wrap cPN cCN cDN (KFC.wrap cPW cCw cDw toCore))
    (.seq (.op (.drop cR2)) (.seq (.op (.drop cPW)) (.op (.drop cPN)))))))

def fromBitsU : IStmt HS Pp :=
  .seq (.op (.dup sN cPN)) (.seq (.op (.dup NLogN.KFW.sW cPW)) (.seq (.op (.const cR2 2))
    (.seq (KFC.wrap cPN cCN cDN (KFC.wrap cPW cCw cDw fromCore))
    (.seq (.op (.drop cR2)) (.seq (.op (.drop cPW)) (.op (.drop cPN)))))))

/-- The parameter map of the conversions. -/
def πNW (N w : ℕ) : PMap := fun i => if i = sN then some (N : ℤ) else if i = NLogN.KFW.sW then some (w : ℤ) else none

/-! ## The Lemma 12 wrapper with register parameters -/

/-- The parameter pushes, from the registers. -/
def swLU : List (Fin HS × Ex HS) :=
  [(KIL.kP, .reg NLogN.KFW.sP), (KIL.kG, .reg sG), (KIL.kS, .reg sS), (aU, .reg sU)]

/-- **The uniform swap wrapper.** -/
def swIU (q : ℕ) : IStmt HS 1 := withPushI swLU (l12C q)

/-! ## The integer-array swap -/

/-- **The uniform integer-array swap**: to bits, Lemma 12, back to integers. -/
def swWU (q : ℕ) : IStmt HS 1 := .seq toBitsU (.seq (swIU q) fromBitsU)

end NLogN.KFW

end

section
-- Source module: Solutions.KFWSwap2
/-!
# KFW: reading the swap descriptor

The descriptor `descSw w u p q sh = [w, u, p, q, |sh|, sh_0, …]` sits on the managed stack `dS`.
`descP` reads it into the registers of `swWU` (`w`, `u`, `vol Pre`, `vol G`, `vol Suf · (w+1)`,
`N` for the Gaussian array of shape `sh ++ [2]`) by moving it item by item to a temporary stack
and back, with unary-counted loops (`rep`); **`descP_fsp`**.
-/

set_option linter.unusedSimpArgs false
set_option linter.unusedVariables false

namespace NLogN.KFW

open IntMul RS KIL KFD KFB KFS AModel

/-- Descriptor and scratch stacks. -/
abbrev dS : Fin HS := 7
abbrev sT : Fin HS := 8
abbrev xp : Fin HS := 9
abbrev xq : Fin HS := 10
abbrev xL : Fin HS := 11
abbrev aS : Fin HS := 12
abbrev cC : Fin HS := 13

/-- Move one item from `dS` to `sT`. -/
def mv1 : RS := .seq (.push sT (.reg NLogN.KFW.dS)) (.pop NLogN.KFW.dS)

/-- Move one item from `sT` back to `dS`. -/
def mvb1 : RS := .seq (.push NLogN.KFW.dS (.reg sT)) (.pop sT)

/-- Quadratic cost unit. -/
def cq (B : ℕ) : ℕ := 10 ^ 8 * (B + 2) ^ 2

/-- Read the top of `dS` into `X` and move it to `sT`. -/
def hdr (X : Fin HS) : RS := .seq (.push X (.reg NLogN.KFW.dS)) NLogN.KFW.mv1

/-- One step of a product phase: `A *= top dS`, move the top to `sT`. -/
def phB (A : Fin HS) : RS := .seq (.asg A (.mul (.reg A) (.reg NLogN.KFW.dS))) NLogN.KFW.mv1

/-- A product phase of `n` steps. -/
def phR (A : Fin HS) (n : Ex HS) : RS := .rep cC n (phB A)

/-- The state after `k` steps of a product phase. -/
def phS (f : Fin HS → List ℤ) (A : Fin HS) (a0 : ℕ) (ys : List ℕ) (rest T0 : List ℤ) (k : ℕ) :
    Fin HS → List ℤ :=
  Function.update (Function.update (Function.update f A [((a0 * (ys.take k).prod : ℕ) : ℤ)]) sT
    (((ys.take k).reverse.map fun y : ℕ => (y : ℤ)) ++ T0)) NLogN.KFW.dS
    ((ys.drop k).map (fun y : ℕ => (y : ℤ)) ++ rest)

/-- Move `n` items back from `sT` to `dS`. -/
def bkR (n : Ex HS) : RS := .rep cC n mvb1

def bkS (f : Fin HS → List ℤ) (zs T1 D0 : List ℤ) (k : ℕ) : Fin HS → List ℤ :=
  Function.update (Function.update f NLogN.KFW.dS ((zs.take k).reverse ++ D0)) sT (zs.drop k ++ T1)

/-! ## The descriptor program -/

/-- **Reading the descriptor** into the registers of `swWU`. -/
def descP : RS :=
  .seq (hdr NLogN.KFW.sW)
  (.seq (hdr sU)
  (.seq (hdr xp)
  (.seq (hdr xq)
  (.seq (hdr xL)
  (.seq (.push NLogN.KFW.sP (.lit 1))
  (.seq (.push sG (.lit 1))
  (.seq (.push NLogN.KFW.aS (.lit 1))
  (.seq (phR NLogN.KFW.sP (.reg xp))
  (.seq (NLogN.KFW.mv1)
  (.seq (phR sG (.add (.add (.reg xq) (.neg (.reg xp))) (.lit (-1))))
  (.seq (NLogN.KFW.mv1)
  (.seq (phR NLogN.KFW.aS (.add (.add (.reg xL) (.neg (.reg xq))) (.lit (-1))))
  (.seq (bkR (.reg xL))
  (.seq (mvb1)
  (.seq (mvb1)
  (.seq (mvb1)
  (.seq (mvb1)
  (.seq (mvb1)
  (.seq (.push sS (.mul (.mul (.reg NLogN.KFW.aS) (.lit 2)) (.add (.reg NLogN.KFW.sW) (.lit 1))))
  (.seq (.push sN (.mul (.mul (.mul (.mul (.mul (.lit 2) (.reg NLogN.KFW.sP)) (.reg sG)) (.reg NLogN.KFW.aS)) (.pow2 (.reg sU))) (.pow2 (.reg sU))))
  (.seq (.pop NLogN.KFW.aS)
  (.seq (.pop xL)
  (.seq (.pop xq)
  ((.pop xp)))))))))))))))))))))))))

/-- The managed stacks of `descP` as a record. -/
def stD (f : Fin HS → List ℤ) (D T W U P G S N X Y Lr A : List ℤ) : Fin HS → List ℤ := fun j =>
  if j = NLogN.KFW.dS then D else if j = sT then T else if j = NLogN.KFW.sW then W else if j = sU then U else
  if j = NLogN.KFW.sP then P else if j = sG then G else if j = sS then S else if j = sN then N else
  if j = xp then X else if j = xq then Y else if j = xL then Lr else if j = NLogN.KFW.aS then A else f j

/-- The registers `descP` produces. -/
def regsSw (f : Fin HS → List ℤ) (w u p q : ℕ) (sh : List ℕ) : Fin HS → List ℤ :=
  stD f (f NLogN.KFW.dS) [] [(w : ℤ)] [(u : ℤ)] [((preS sh p).prod : ℤ)] [((midS sh p q).prod : ℤ)]
    [(((sufS sh q).prod * 2 * (w + 1) : ℕ) : ℤ)] [((2 * sh.prod : ℕ) : ℤ)] [] [] [] []

/-! ### Record lemmas for `stD` -/

section StD
variable (f : Fin HS → List ℤ) (D T W U P G S N X Y Lr A : List ℤ) (v : List ℤ)

@[simp] theorem stD_dS : stD f D T W U P G S N X Y Lr A NLogN.KFW.dS = D := rfl

@[simp] theorem stD_sT : stD f D T W U P G S N X Y Lr A sT = T := rfl

@[simp] theorem stD_sW : stD f D T W U P G S N X Y Lr A NLogN.KFW.sW = W := rfl

@[simp] theorem stD_sU : stD f D T W U P G S N X Y Lr A sU = U := rfl

@[simp] theorem stD_sP : stD f D T W U P G S N X Y Lr A NLogN.KFW.sP = P := rfl

@[simp] theorem stD_sG : stD f D T W U P G S N X Y Lr A sG = G := rfl

@[simp] theorem stD_sS : stD f D T W U P G S N X Y Lr A sS = S := rfl

@[simp] theorem stD_sN : stD f D T W U P G S N X Y Lr A sN = N := rfl

@[simp] theorem stD_xp : stD f D T W U P G S N X Y Lr A xp = X := rfl

@[simp] theorem stD_xq : stD f D T W U P G S N X Y Lr A xq = Y := rfl

@[simp] theorem stD_xL : stD f D T W U P G S N X Y Lr A xL = Lr := rfl

@[simp] theorem stD_aS : stD f D T W U P G S N X Y Lr A NLogN.KFW.aS = A := rfl

end StD

/-- Per-piece cost budget of `descP`. -/
def xD (L B BL : ℕ) : ℕ := 2 * cq BL + 2 * L * cq B + cRep L

end NLogN.KFW

end

section
-- Source module: Solutions.KFWSwapG
/-!
# KFW: uniform Lemma 12 on Gaussian arrays (`SwapNeed`)

`swGU q = descP ; swWU q ; popP`: read the descriptor `descSw w u p q sh` on `dS` into the
registers of `swWU` (`descP`), run the register-driven integer swap on the shape `sh ++ [2]`,
and pop the registers again.  **`swGU_spec`**, and **`swapNeed : KFM.SwapNeed`** with the table
`kTable (fun _ => (swGU q).ren id (gL 1))` (one fixed prime `q`).
-/

set_option linter.unusedSimpArgs false
set_option linter.unusedVariables false

namespace NLogN.KFW

open IntMul RS KIL KFD KFB KFS KFL AModel AModelL12

/-- Pop the six registers of `swWU`. -/
def popP : RS :=
  .seq (.pop sN) (.seq (.pop sS) (.seq (.pop sG) (.seq (.pop sP) (.seq (.pop sU) (.pop NLogN.KFW.sW)))))

/-- **The uniform Gaussian swap.** -/
def swGU (q : ℕ) : IStmt HS 1 := .seq (fcode descP) (.seq (swWU q) (fcode popP))

end NLogN.KFW

end