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

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

import Definitions.Def_IntMul_Kappa_K_Defs26

set_option autoImplicit false
set_option relaxedAutoImplicit false
set_option maxHeartbeats 1000000

section
-- Source module: Solutions.SynthModel
/-!
# The concrete ring `R_r = ℂ[y]/(y^r + 1)` is a coefficient model

`Rr n = AdjoinRoot (X^{n+1} + 1)` over `ℂ` (so `r = n + 1`), with `y` its root (a unit,
`y^r = -1`) and coefficients in the power basis `1, y, …, y^{r-1}`.
`rrModel n : CoeffModel (yU n) (n + 1)` — so every statement of `SynthNorm` applies to `R_r`.
-/

set_option linter.unusedSectionVars false

namespace Synth

open Polynomial Finset

/-- `ℂ[y]/(y^{n+1} + 1)`. -/
abbrev Rr (n : ℕ) := AdjoinRoot (X ^ (n + 1) + 1 : ℂ[X])

lemma monic_rr (n : ℕ) : (X ^ (n + 1) + 1 : ℂ[X]).Monic := by
  have := monic_X_pow_add_C (1 : ℂ) (Nat.succ_ne_zero n); simpa using this

lemma natDegree_rr (n : ℕ) : (X ^ (n + 1) + 1 : ℂ[X]).natDegree = n + 1 := by
  have := natDegree_X_pow_add_C (n := n + 1) (r := (1 : ℂ)); simpa using this

/-- The root `y`. -/
noncomputable def yR (n : ℕ) : Rr n := AdjoinRoot.root _

lemma yR_pow (n : ℕ) : Synth.yR n ^ (n + 1) = -1 := by
  have h := AdjoinRoot.mk_self (f := (X ^ (n + 1) + 1 : ℂ[X]))
  rw [map_add, map_pow, AdjoinRoot.mk_X, map_one] at h
  exact eq_neg_of_add_eq_zero_left h

/-- `y` as a unit: `y · (-y^n) = 1`. -/
noncomputable def yU (n : ℕ) : (Rr n)ˣ :=
  Units.mkOfMulEqOne (Synth.yR n) (-(Synth.yR n) ^ n) (by rw [mul_neg, ← pow_succ', yR_pow, neg_neg])

/-- The power basis `y^k`, indexed by `Fin (n+1)`. -/
noncomputable def rrBasis (n : ℕ) : Module.Basis (Fin (n + 1)) ℂ (Rr n) :=
  (AdjoinRoot.powerBasis' (monic_rr n)).basis.reindex (finCongr (by
    rw [AdjoinRoot.powerBasis'_dim, natDegree_rr]))

lemma rrBasis_apply (n : ℕ) (k : Fin (n + 1)) : rrBasis n k = Synth.yR n ^ (k : ℕ) := by
  simp [rrBasis, Module.Basis.reindex_apply, PowerBasis.basis_eq_pow, AdjoinRoot.powerBasis'_gen,
    Synth.yR]

/-- **`R_r` is a coefficient model.** -/
noncomputable def rrModel (n : ℕ) : CoeffModel (Synth.yU n) (n + 1) where
  pos := Nat.succ_pos n
  co := (rrBasis n).equivFun.toLinearMap
  inj := (rrBasis n).equivFun.injective
  co_one := by
    funext j
    have : (1 : Rr n) = rrBasis n 0 := by rw [rrBasis_apply]; simp
    simp only [LinearEquiv.coe_coe]
    rw [this, Module.Basis.equivFun_self, Pi.single_apply]
    by_cases hj : j = 0
    · subst hj; simp
    · rw [if_neg (Ne.symm hj), if_neg (fun h => hj (by rw [h]; rfl))]
  co_y := by
    intro z
    apply (rrBasis n).equivFun.symm.injective
    simp only [LinearEquiv.coe_coe, LinearEquiv.symm_apply_apply]
    rw [Module.Basis.equivFun_symm_apply]
    set c := (rrBasis n).equivFun z
    have hz : z = ∑ k, c k • rrBasis n k := by
      rw [← Module.Basis.equivFun_symm_apply, LinearEquiv.symm_apply_apply]
    show (Synth.yU n : Rr n) * z = _
    conv_lhs => rw [hz]
    simp only [rrBasis_apply, mul_sum, mul_smul_comm, ← pow_succ']
    have hyU : ((Synth.yU n : (Rr n)ˣ) : Rr n) = Synth.yR n := rfl
    rw [hyU]
    rw [Fin.sum_univ_castSucc, Fin.sum_univ_succ]
    simp only [Fin.val_last, yR_pow, Fin.val_castSucc, Fin.val_succ, Fin.val_zero, pow_zero]
    have h0 : cyc c 0 = -c (Fin.last n) := by simp [cyc, Fin.last]
    have hs : ∀ i : Fin n, cyc c i.succ = c i.castSucc := by
      intro i; simp only [cyc, Fin.val_succ]; rw [dif_neg (by omega)]; rfl
    simp only [h0, hs]
    rw [← pow_succ', yR_pow, add_comm]
    simp only [← pow_succ', smul_neg, neg_smul, one_smul]

end Synth

end

section
-- Source module: Solutions.SynthKron
/-!
# §7 Lemma 17: signed packing, centered extraction, negacyclic fold, truncation

All the arithmetic content of Lemma 17, stated over `ℤ` (and over a coefficient model for the
ring-level identification).

* **Packing** (`gam`, `dig`): with carry `γ₀ = 0`, `d_j = (u_j + γ_j) mod B`,
  `γ_{j+1} = (u_j + γ_j - d_j)/B`; `pack_digit`, `pack_carry` (`γ_j ∈ {-1,0}` when `|u_j| < B`),
  `pack_telescope` (`∑ u_j B^j = ∑ d_j B^j + γ_r B^r`), `pack_sign` (`γ_r = -1 ↔ A_u < 0`).
* `pack_bound`: `|u_j| ≤ c`, `2c < B` give `2|A_u| < B^r`.
* `cnv` (ordinary product coefficients), `eval_mul_cnv` (`A_a A_c = ∑_{j<2r} h_j B^j`),
  `abs_cnv_le` and `coeff_bound31` (`|h_j| ≤ r 2^{2p} < 2^{3p}`, `2·2^{3p} < B₀ = 2^{4p}`).
* `twos_complement_exists`: a fixed-width signed block representation `Z = ∑ q_j B^j + ε B^J`.
* **Extraction** (`cent`, `rho`): `extract` proves the residual identity, `ρ_j ∈ {0,1}`, exact
  recovery `ĥ_j = h_j`, and `ρ_J + ε = 0`; `kron_recover` is the instance of Lemma 17.
* **Negacyclic fold**: `fold_formula`, `co_mul` and `co_mul_gauss` (the Gaussian numerators of
  `f g` are `(ac - be) + i(ae + bc)` folded modulo `y^r + 1`).
* **Truncation**: `truncNum_div` (`trunc₀(H / 2^{p+ℓ})` is `H.tdiv 2^{p+ℓ}`), `Qp_normalized`.
* **Scaling by appending zeros**: `append_zeros`.
-/

set_option linter.unusedSectionVars false

namespace Synth

open Finset Butterfly

/-! ## 1. Signed packing -/

section Pack

variable (B : ℤ) (u : ℕ → ℤ)

variable {B}

end Pack

/-! ## 2. Ordinary product coefficients -/

/-- `h_j = ∑_{i + k = j} a_i c_k` for polynomials of degree `< r`. -/
def cnv (r : ℕ) (a c : ℕ → ℤ) (j : ℕ) : ℤ :=
  ∑ i ∈ range r, ∑ k ∈ range r, if i + k = j then a i * c k else 0

/-! ## 3. Fixed-width two's complement -/

/-! ## 4. Centered extraction -/

/-- The centered residue in `[-H, H)` modulo `2H`. -/
def cent (H x : ℤ) : ℤ := (x + H) % (2 * H) - H

/-- The extraction carry `ρ_j`. -/
def rho (H : ℤ) (q : ℕ → ℤ) : ℕ → ℤ
  | 0 => 0
  | j + 1 => (q j + rho H q j - cent H (q j + rho H q j)) / (2 * H)

/-- The residual `∑_{k=j}^{J-1} f_k B^{k-j}`. -/
def tail (B : ℤ) (f : ℕ → ℤ) (J j : ℕ) : ℤ := ∑ k ∈ range (J - j), f (j + k) * B ^ k

/-! ## 5. The negacyclic fold and the ring product -/

section Ring

variable {R : Type*} [CommRing R] [Algebra ℂ R] {y : Rˣ} {r : ℕ} (C : CoeffModel y r)

end Ring

/-! ## 6. Truncation of the normalized product -/

/-! ## 7. Scaling by `M = 2^m`: appending zero bits -/

end Synth

end

section
-- Source module: Solutions.KFV1Coef
/-!
# KFV §1: Lemma 16's computed procedures on coefficient vectors

The IL works on record coefficients (`RecRep r F zs (xs.map C.co)`), so the computed procedures
`procQ` / `procQI` of `GlueError` / `GlueProp18` (over a coefficient model `C` of `R`) are
transported to lists of coefficient vectors `Fin r → ℂ`:

* `nrotS s v`: the negacyclic rotation by `s` (`(y^s z)` in coefficients), read off the doubled
  record `[v, -v]` (`dbl`): `nrotS s v q = dbl v ((q + 2r - s mod 2r) mod 2r)`;
  `co_ypow`, `co_yp_mul`: `co (y^e z) = nrotS (e mod 2r) (co z)`.
* `twLc`, `layerQc`, `compRc`, `compRIc`, `procQc`, `procQIc`: the coefficient versions;
  **`co_procQ`**, **`co_procQI`**: `(procQ … xs).map co = procQc … (xs.map co)` (and `procQI`).
-/

set_option linter.unusedSectionVars false
set_option linter.unusedSimpArgs false

namespace NLogN.KFV

open Finset Butterfly AModel Synth GlueLayout GlueP14 GlueError GlueProp18

/-! ## 1. Negacyclic rotation -/

section Rot

variable {r : ℕ}

/-- The doubled record `[v, -v]` (zero beyond `2r`). -/
def dbl (v : Fin r → ℂ) (k : ℕ) : ℂ :=
  if h : k < r then v ⟨k, h⟩ else if h2 : k < 2 * r then -v ⟨k - r, by omega⟩ else 0

/-- Negacyclic rotation by `s`. -/
def nrotS (s : ℕ) (v : Fin r → ℂ) : Fin r → ℂ :=
  fun q => NLogN.KFV.dbl v ((q + 2 * r - s % (2 * r)) % (2 * r))

end Rot

/-! ## 2. Monomials in coefficients -/

section Mono

variable {R : Type*} [CommRing R] [Algebra ℂ R] {y : Rˣ} {r : ℕ} (C : CoeffModel y r)

/-- The exponent `e mod 2r` as a natural number. -/
def sOf (r : ℕ) (e : ℤ) : ℕ := (e % (2 * r : ℤ)).toNat

end Mono

/-! ## 3. Rounds on coefficients -/

section Rounds

variable {R : Type*} [CommRing R] [Algebra ℂ R] {y : Rˣ} {r : ℕ} (C : CoeffModel y r)
variable {D : ℕ} {sh : Shape} {ι : Type} [Fintype ι] [DecidableEq ι]

/-- `J_E` on coefficient vectors. -/
noncomputable def twLc (r : ℕ) (sh : Shape) (Ex : ℕ → ℤ) (xs : List (Fin r → ℂ)) :
    List (Fin r → ℂ) :=
  (List.range (AModel.vol sh)).map fun j => nrotS (sOf r (Ex j)) (xs.getD j 0)

/-- The computed layer on coefficients: `Q_p ∘ H0^{⊗}`. -/
noncomputable def layerQc (p : ℕ) (Sb : SelBits sh ι) (xs : List (Fin r → ℂ)) :
    List (Fin r → ℂ) :=
  QpArr p (selOp Sb (Butterfly.tens H0) xs)

/-- The computed forward round on coefficients. -/
noncomputable def compRc (p : ℕ) (Sb : SelBits sh ι) (Ex : ℕ → ℤ) (xs : List (Fin r → ℂ)) :
    List (Fin r → ℂ) :=
  twLc r sh Ex (layerQc p Sb xs)

/-- The computed opposite round on coefficients. -/
noncomputable def compRIc (p : ℕ) (Sb : SelBits sh ι) (Ex : ℕ → ℤ) (xs : List (Fin r → ℂ)) :
    List (Fin r → ℂ) :=
  layerQc p Sb (twLc r sh (fun j => -Ex j) xs)

/-- The computed forward procedure on coefficients (rounds `n-1, …, 0`). -/
noncomputable def procQc (p : ℕ) {a : Fin D → ℕ} (L : CoordLayout sh a) (ℓ : ℕ)
    (P : ℕ → Finset (Fin D)) (Sb : (h : ℕ) → SelBits sh {i // i ∈ P h}) :
    ℕ → List (Fin r → ℂ) → List (Fin r → ℂ) :=
  runDown fun h => compRc p (Sb h) (roundEx L ℓ h (P h))

/-- The computed opposite procedure on coefficients (rounds `0, …, n-1`). -/
noncomputable def procQIc (p : ℕ) {a : Fin D → ℕ} (L : CoordLayout sh a) (ℓ : ℕ)
    (P : ℕ → Finset (Fin D)) (Sb : (h : ℕ) → SelBits sh {i // i ∈ P h}) (n : ℕ)
    (xs : List (Fin r → ℂ)) : List (Fin r → ℂ) :=
  (List.range n).foldl (fun u k => compRIc p (Sb k) (roundEx L ℓ k (P k)) u) xs

end Rounds

end NLogN.KFV

end

section
-- Source module: Solutions.KFV2Tw
/-!
# KFV §3: the twiddle pass `J_E` (per-record negacyclic rotation by a table)

Records of `r` Gaussian coefficients (`2r` integers) on stack `0`, a table of rotation amounts
`s_j < 2r` (one per record) on stack `14`.

`twC N r = dblI N (2r); rotC N (2r); unpadI N (2r) (4r)`:
double every record to `[z, -z]`, rotate each doubled record (shape `[N, 2r, 2]`) by its table
entry with KFC `rotF` (the table is consumed), keep the first half.

* `twZ r N sv zs`: the result on Gaussian lists; **`twC_spec`** (IL run, cost
  `≤ 10^4 (w+1)(12 N r + 2N + 2) + 70·kap w·4Nr + O(log)`), **`twZ_rep`**: if `zs` represents
  `xs` (`RecRep r F`) then `twZ … zs` represents `twLc r sh Ex xs` when `s_j = sOf r (Ex j)`
  (`J_E` in coefficients, `KFV1Coef`).
-/

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

namespace NLogN.KFV

open IntMul KFC KFT KFU KFB KFH AModel AModelL12

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

/-! ## 1. Index arithmetic and closed forms -/

/-- A component of a Gaussian integer. -/
def gc (z : GI) (c : ℕ) : ℤ := if c = 0 then z.1 else z.2

/-! ## 2. The Gaussian result -/

/-- Entry `i < 2r` of the doubled record `j`: `[z, -z]`. -/
def dblGI (zs : List GI) (r j i : ℕ) : GI :=
  if i < r then zs.getD (j * r + i) 0 else -zs.getD (j * r + (i - r)) 0

/-- The twiddled array: record `j` rotated negacyclically by `sv j`. -/
def twZ (r N : ℕ) (sv : ℕ → ℕ) (zs : List GI) : List GI :=
  (List.range (N * r)).map fun k => dblGI zs r (k / r) ((k % r + 2 * r - sv (k / r)) % (2 * r))

/-! ## 3. Records -/

/-! ## 4. The program -/

/-- `rotF` on `[N, Q, 2]` with pushed parameters (the table on stack `14` is consumed). -/
def rotC (N Q : ℕ) : IStmt HS Pp :=
  .seq (.op (.const 1 N)) (.seq (.op (.const 2 Q)) (.seq (.op (.const 4 2))
    (.seq rotF (.seq (.op (.drop 1)) (.seq (.op (.drop 2)) (.op (.drop 4)))))))

def cRotC (N Q w : ℕ) : ℕ :=
  cRotF N Q 2 w + 3 * (zsz (N : ℤ) + zsz (Q : ℤ) + zsz (2 : ℤ)) + 12

/-- **The twiddle pass.** -/
def twC (N r : ℕ) : IStmt HS Pp :=
  .seq (dblI N (2 * r)) (.seq (rotC N (2 * r)) (unpadI N (2 * r) (4 * r)))

def cTw (N r w : ℕ) : ℕ := cDblI N (2 * r) w + cRotC N (2 * r) w + cUnpadI N (2 * r) (4 * r) w

end NLogN.KFV

end

section
-- Source module: Solutions.KFV3Tab
/-!
# KFV §4: the `J_E` exponent table, one record at a time

Round `h` of Lemma 16 twists record `j` by `y^{E_j}` with
`E_j = roundEx L ℓ h P j = Σ_{i ∈ P} twE ℓ h (x_i)`, `x_i` the coordinate of axis `i` at address
`j`. In a slot layout every `x_i` is a bit sum of the binary address, so
`E_j = -Σ_{(c, k₁, k₂) ∈ T} c · bit_{k₁}(j) · bit_{k₂}(j)`: the mask bit (slot `h` of axis `i`)
times the linear form of the lower slots (`c = 2^{ℓ-h+t}`). This is a masked linear table.

* `svE Q M T`: one expression for `(-Σ c·bit·bit) mod Q` of the index in register `rN`
  (`bitE k = (j / 2^k) mod 2`); **`tabP`** loops over the records with a down-counter on `rN` and
  pushes the entries on stack `14`, ending with entry `0` on top (the layout `rotF` consumes).
* **`tabP_fsp`**: the table `(range N).map (fun j => (-tv T j) mod Q)`, cost
  `cLp N ((18|T| + 7)·10^7 (B+2)^2 + …) B` with `B = zsz U`, `U` bounding the index, the moduli,
  the coefficients and the sums. That is `O(|T| B²)` **per record**: `|T| ≤ D ℓ`, `B = O(log N + ℓ)`,
  so the bound is polylogarithmic per record and does not grow with the record length `r`.
* **`roundEx_slot`**: in any slot layout (`slotLayout hS`, in particular the paper's
  `paperLayout`) `roundEx … j = -tv (trips …) j`, so `sOf r (roundEx …) = tabEntry` with
  `Q = 2r`.
-/

set_option linter.unusedSimpArgs false
set_option linter.unusedVariables false

namespace NLogN.KFV

open IntMul RS KIL KFD KFB AModel

/-! ## 1. The expression -/

/-- Bit `k` of `j`. -/
def bt (k j : ℕ) : ℕ := j / 2 ^ k % 2

/-- The value of the sum at index `j`. -/
def tv (T : List (ℕ × ℕ × ℕ)) (j : ℕ) : ℕ := (T.map fun t => t.1 * NLogN.KFV.bt t.2.1 j * NLogN.KFV.bt t.2.2 j).sum

/-! ## 2. Evaluation, side conditions, bounds -/

section Eval

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

variable {U : ℕ}

end Eval

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

/-- The entries. -/
def svv (Q : ℕ) (T : List (ℕ × ℕ × ℕ)) (j : ℕ) : ℤ := (-(tv T j : ℤ)) % (Q : ℤ)

/-! ## 4. The table of a round in a slot layout -/

section Slot

open Finset Synth GlueLayout GlueError Butterfly

variable {D : ℕ} [NeZero D] {a : Fin D → ℕ} {arr : List ℕ}

/-- The terms of round `h` for axis `i`: `2^{ℓ-h+t} · bit(slot h) · bit(slot t)`, `t < h`. -/
def tripsAx (arr : List ℕ) (ℓ h : ℕ) (i : Fin D) : List (ℕ × ℕ × ℕ) :=
  (List.range h).map fun t =>
    (2 ^ (ℓ - h + t), arr.length - 1 - spos D arr i h, arr.length - 1 - spos D arr i t)

/-- The terms of round `h`. -/
noncomputable def trips (arr : List ℕ) (ℓ h : ℕ) (P : Finset (Fin D)) : List (ℕ × ℕ × ℕ) :=
  (roundAxes P).flatMap (tripsAx arr ℓ h)

/-! ### The opposite round: the negated exponent -/

/-- Coefficients `Q - c`: the same table then holds `(+Σ) mod Q`. -/
def negTr (Q : ℕ) (T : List (ℕ × ℕ × ℕ)) : List (ℕ × ℕ × ℕ) := T.map fun t => (Q - t.1, t.2)

end Slot

end NLogN.KFV

end

section
-- Source module: Solutions.KFV3Lay
/-!
# KFV §5: the prefix rounds and the top round are Proposition 14 layers with `K = 1`

In the paper's layout `tgtL = Lst ++ chunks`, `Lst` starts with the top slots `(i, ℓ-1)` of the
`m` long axes, then the slots `(i, h)` for `h = b-1, …, 0`, `D` consecutive positions per `h`
(axis order). So:

* **prefix round** `h < b`: the selected bits are the `D` consecutive positions
  `preP h + i` (`preP h = m + (b-1-h) D`); `pre_round_layer`: the round's selected-bit operator
  equals the one of `layerSel (2^{preP h}) 1 D (2^{sufP h}) 0` (one-bit "chunks", `ρ = 0`);
* **top round** `h = ℓ-1`: only the long axes take part, at positions `< m`;
  `top_round_layer`: the same with `layerSel 1 1 m (2^{n-m}) 0`, `m = |partAxes al (ℓ-1)|`.

Hence (`pre_layerQc`, `top_layerQc`) the computed layer `Q_p ∘ H0^{⊗}` of those rounds is
`prop14Proc … (tens Cm)` on that layer, which KFH's `p14Leaf` (individual kernels: `3D + 1`
kernel passes and one truncation pass) computes.
-/

set_option linter.unusedSectionVars false
set_option linter.unusedSimpArgs false

namespace NLogN.KFV

open Finset Butterfly AModel Synth GlueLayout GlueP14 GlueError

section Paper

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

/-- Number of long axes. -/
def nLong : ℕ := ((List.finRange D).filter fun i => long i).length

/-- Bits before the block of prefix round `h`. -/
def preP (h : ℕ) : ℕ := nLong D long + (bb ℓ K - 1 - h) * D

/-- Bits after the block of prefix round `h`. -/
def sufP (h : ℕ) : ℕ := (tgtL D ℓ K long).length - NLogN.KFV.preP D ℓ K long h - D

end Paper

end NLogN.KFV

end

section
-- Source module: Solutions.KFV4Round
/-!
# KFV §6: one round of Lemma 16 in IL

A round of the computed forward procedure on coefficients is
`compRc p Sb Ex = twLc Ex ∘ layerQc p Sb` (`KFV1Coef`), the opposite round is
`compRIc p Sb Ex = layerQc p Sb ∘ twLc (-Ex)`.

* `LayerOK code T N r p W cost`: on every clean state holding a record array `zs` (scale `p`,
  width `≤ W`, `N` records in the unit disk) that represents `xs`, `code` leaves an array
  representing `T xs` (scale `p`, width `≤ W`, `N` records in the disk), at cost `≤ cost`.
* **`twist_run`**: the `J_E` pass of a round: the exponent table (`tabP`, `KFV3Tab`) and the
  per-record negacyclic rotation (`twC`, `KFV2Tw`), cost `cTab + cTw`.
* `LayerOK.twist` / `LayerOK.twistPre`: a layer followed / preceded by a twist pass;
  `LayerOK.comp`: sequential composition.
-/

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

namespace NLogN.KFV

open IntMul KFC KFT KFU KFB KFH KFD KIL AModel

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

/-! ## 1. Width and disk through the twist -/

/-- Records in the unit disk. -/
def Disk {r : ℕ} (xs : List (Fin r → ℂ)) : Prop := ∀ j (h : j < xs.length), ‖xs[j]‖ ≤ 1

/-! ## 2. The twist pass -/

/-- The rotation amounts read off a table. -/
def svN (Q : ℕ) (T : List (ℕ × ℕ × ℕ)) (j : ℕ) : ℕ := (svv Q T j).toNat

/-! ## 3. Layers and their composition -/

section Layer

variable {r : ℕ}

end Layer

end NLogN.KFV

end

section
-- Source module: Solutions.KFV4Lay
/-!
# KFV §7: the layers of the rounds of Lemma 16 as `LayerOK` codes

* `layerOK_leaf`: KFH's `p14Leaf` (individual kernels) computes `prop14Proc … (tens Cm)` on any
  layer layout; cost `cP14 ≤ 10^6 (D+1)(W + 8D + 8) N r + 10^5 (D+1)(Z+2)`.
* `layerOK_p14`: any run of KFM's `p14Ph` (the statement of `KFM.prop14Ph_kTable`) is a
  `LayerOK` for `prop14Proc`.
* `layerOK_trunc0`: the top round without long axes is `Q_p` alone (`truncP 0`).
* On the paper's layout (`KFV3Lay`): `layerOK_pre` (prefix rounds `h < b`, `p14Leaf` with
  one-bit chunks), `layerOK_top` (round `ℓ - 1`), and `layerOK_chunk` (chunk rounds, from any
  Proposition 14 implementation on `layerSel (2^A) K D (2^B) ρ`).
-/

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

namespace NLogN.KFV

open IntMul KFC KFT KFU KFB KFH KFD KIL AModel Butterfly

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

/-! ## 1. Generic layers -/

/-! ## 2. The rounds of the paper's layout -/

section Paper

open Synth GlueLayout GlueP14 GlueError

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

/-- Number of address bits. -/
abbrev nB : ℕ := (tgtL D ℓ K long).length

end Paper

end NLogN.KFV

end

section
-- Source module: Solutions.KFV4Proc
/-!
# KFV §8: Lemma 16 forward and opposite as IL programs (paper layout)

* `fwdP Rc n = Rc (n-1); …; Rc 0` and `oppP Rc n = Rc 0; …; Rc (n-1)`;
  `fwd_ok` / `opp_ok`: by induction on the rounds, they compute `runDown F n` / the `foldl` of
  the rounds (`procQc` / `procQIc` of `KFV1Coef`).
* `layerCode h`: prefix rounds `h < b` by `p14Leaf` with one-bit chunks (`dK` row: individual
  kernels), the top round by `topCode`, chunk rounds by any Proposition 14 code `Cch h`
  (in the merged table: `KFM.prop14Ph_kTable`, packaged by `layerOK_chunk`).
* `fwdRound h = layerCode h ; J_E` (exponent table `trips … h` and `twC`),
  `oppRound h = J_{-E} ; layerCode h` (table `negTr (2r) (trips … h)`).
* **`lemma16_fwd_IL`**: `fwdP fwdRound ℓ` computes `procQc p (paperLayout) ℓ (partAxes al)
  (paperSel) ℓ`, i.e. (`co_procQ`) the coefficients of `GlueError.procQ`, the computed forward
  procedure of Lemma 16 / `GlueProp18`;
  **`lemma16_opp_IL`**: `oppP oppRound ℓ` computes `procQIc …` (`co_procQI`: `GlueProp18.procQI`).
  Cost: the sum of the round costs `layerCost h + cTab + cTw`.
-/

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

namespace NLogN.KFV

open IntMul KFC KFT KFU KFB KFH KFD KIL AModel Butterfly

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

/-! ## 1. Round induction -/

section Ind

variable {r : ℕ}

end Ind

/-! ## 2. The rounds of the paper's layout -/

section Paper

open Synth GlueLayout GlueP14 GlueError

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

/-- The exponent table of round `h`. -/
noncomputable def tabT (h : ℕ) : List (ℕ × ℕ × ℕ) :=
  trips (tgtL D ℓ K long) ℓ h (partAxes (al D ℓ long) h)

end Paper

end NLogN.KFV

end

section
-- Source module: Solutions.KFV7Need
/-!
# KFV §11: the shape of the uniform Proposition 14 that Prop 18 (KFV) needs

Uniform conventions (`UNIFORM.md`): one fixed procedure `e14` in the table; every
`n`-dependent parameter is run-time data on a descriptor stack `dS14`.

**`desc14 P K D S ρ r p w Z`** is the descriptor `[P, K, D, S, ρ, r, p, w, Z]` (the layer shape
`layerShape P K D S`, the root exponent `ρ < K`, the record length `r`, the scale `p`, the
input width `w`, the counter size `Z`).

**`P14Need Isym tbl e14 dS14 A C₀`**: for every parameter set satisfying the side conditions of
`KFM.prop14Ph_kTable` and `KFM.round_cost_closed`, a call of `e14` with `desc14` on `dS14` and
the array on stack `0`

* computes `prop14Proc … (tens Cm)` exactly as `P14Run` (`KFV4Lay`) does (same `RecRep`,
  `GB (p+1)`, and the `√2·2^{-p}` distance to `layerH`);
* leaves every other stack unchanged (the descriptor stays on `dS14`);
* needs empty only stacks `1 … 49` and `55 … 62` (other than `dS14`); stacks `50 … 54` and `63`
  hold the caller's (KFV's) data and are **arbitrary and preserved**;
* costs `≤ C₀ · vol · r · (w + 1) · D^{1 - 2^{-54}}` (the closed bound of
  `KFM.round_cost_closed`; any per-call descriptor overhead is absorbed using
  `2 (2Z+4)^3 + 2 ≤ r`).

`p14Run_of_need` checks that this is enough for KFV: on states whose free stacks are empty it
gives a `P14Run` (the hypothesis of `layerOK_p14`) with a constant code, after the descriptor is in
place.
-/

set_option linter.unusedSectionVars false
set_option linter.unusedVariables false

namespace NLogN.KFV

open IntMul KFC KFT KFU KFB KFH KFD KIL AModel Butterfly

/-- The descriptor of a uniform Proposition 14 call. -/
def desc14 (P K D S ρ r p w Z : ℕ) : List ℤ :=
  [(P : ℤ), (K : ℤ), (D : ℤ), (S : ℤ), (ρ : ℤ), (r : ℤ), (p : ℤ), (w : ℤ), (Z : ℤ)]

/-- The stacks a uniform Proposition 14 call may require to be empty (all but `0`, the
descriptor stack, and KFV's live stacks `50 … 54`, `63`). -/
def p14Free (dS14 j : Fin HS) : Prop :=
  j ≠ 0 ∧ j ≠ dS14 ∧ (j.val < 50 ∨ (55 ≤ j.val ∧ j.val ≤ 62))

/-- **The uniform Proposition 14 KFV needs** (to be supplied by the KFW worker as `P14U`). -/
def P14Need {Pp : ℕ} (Isym : List Sym) (tbl : Fin Pp → IStmt HS Pp) (e14 : Fin Pp)
    (dS14 : Fin HS) (A : ℕ) (C₀ : ℝ) : Prop :=
  ∀ {K r Z P D S ρ p w : ℕ} (hρ : ρ < K),
    6 ≤ K → 1 ≤ r → 80 * Z * (Z + 2) ≤ 2 ^ K → 1 ≤ S → 1 ≤ P → 1 ≤ D →
    zsz ((32 * (vol (layerShape P K D S) * r) : ℕ) : ℤ) ≤ Z → zsz (p : ℤ) ≤ Z →
    (K : ℝ) ≤ 2 * (D : ℝ) ^ ((2 : ℝ)⁻¹ ^ 56) → KFM.EwB A D ≤ w + 1 → w + KFM.EwB A D + 2 < 2 ^ Z →
    2 * (2 * Z + 4) ^ 3 + 2 ≤ r → (K * D + P + 10) * (Z + 2) ≤ vol (layerShape P K D S) * r →
    ∀ (st : IState HS) (f : Fin HS → List ℤ) (zs : List GI) (l : List ℤ)
      (xs : List (Fin r → ℂ)), RecRep r p zs xs →
      xs.length = vol (layerShape P K D S) → f 0 = gflat zs ++ l →
      f dS14 = desc14 P K D S ρ r p w Z → (∀ j, p14Free dS14 j → f j = []) →
      NLogN.KFB.GB w zs → (∀ j (h : j < xs.length), ‖xs[j]‖ ≤ 1) →
      ∃ (zs' : List GI) (c : ℕ),
        RecRep r p zs' (prop14Proc (P := P) (S := S) hρ p (Butterfly.tens Cm : Butterfly.Op (Fin D)) xs) ∧
        NLogN.KFB.GB (p + 1) zs' ∧
        (∀ j (h : j < (prop14Proc (P := P) (S := S) hρ p (Butterfly.tens Cm : Butterfly.Op (Fin D)) xs).length)
          (h' : j < (layerH P K D S ρ hρ xs).length), ∀ c,
          ‖(prop14Proc (P := P) (S := S) hρ p (Butterfly.tens Cm : Butterfly.Op (Fin D)) xs)[j] c -
            (layerH P K D S ρ hρ xs)[j] c‖ < Real.sqrt 2 * 2⁻¹ ^ p) ∧
        IRuns Isym tbl (.call e14) (NLogN.KIL.encK st f) (NLogN.KIL.encK st (Function.update f 0 (gflat zs' ++ l))) c ∧
        (c : ℝ) ≤ C₀ * ((vol (layerShape P K D S) * r : ℕ) : ℝ) * (w + 1) *
          (D : ℝ) ^ (1 - (2 : ℝ)⁻¹ ^ 54)

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 -/

def allHi (L : List (Fin HS)) : Bool := L.all fun j => decide (56 ≤ j.val)

def gW (W : List (Fin HS)) (i : Fin HS) : Bool :=
  decide (i ∈ W) || decide (56 ≤ i.val) || decide (i.val = 26 ∨ i.val = 27 ∨ i.val = 28 ∨ i.val = 29)

/-- The `toUL` scratch. -/
def ulS (i : Fin HS) : Prop := i.val = 26 ∨ i.val = 27 ∨ i.val = 28 ∨ i.val = 29

/-! ## Generalized substitution -/

/-- `e'` evaluates to `v` at a cost comparable to the literal. -/
def LitOK (π : PMap) (v : ℤ) (e' : Ex HS) : Prop :=
  e'.Clean ∧ ∀ s : IState HS, PV π s → e'.ok s ∧ e'.eval s = v ∧
    e'.cost s + (2 * moveCost (zsz v) + 19) ≤ 2000 * (zsz v + 2)

/-- `e'` is `e` with some literals replaced by expressions over the parameters. -/
inductive EG (π : PMap) : Ex HS → Ex HS → Prop
  | refl (e : Ex HS) : EG π e e
  | lit {v : ℤ} {e' : Ex HS} : LitOK π v e' → EG π (.lit v) e'
  | add {a a' b b' : Ex HS} : EG π a a' → EG π b b' → EG π (.add a b) (.add a' b')
  | neg {a a' : Ex HS} : EG π a a' → EG π (.neg a) (.neg a')
  | mul {a a' b b' : Ex HS} : EG π a a' → EG π b b' → EG π (.mul a b) (.mul a' b')
  | shl {a a' b b' : Ex HS} : EG π a a' → EG π b b' → EG π (.shl a b) (.shl a' b')
  | quo {a a' b b' : Ex HS} : EG π a a' → EG π b b' → EG π (.quo a b) (.quo a' b')
  | rem {a a' b b' : Ex HS} : EG π a a' → EG π b b' → EG π (.rem a b) (.rem a' b')
  | pow2 {a a' : Ex HS} : EG π a a' → EG π (.pow2 a) (.pow2 a')
  | sel {c c' a a' b b' : Ex HS} : EG π c c' → EG π a a' → EG π b b' →
      EG π (.sel c a b) (.sel c' a' b')
  | ndiv {a a' b b' : Ex HS} : EG π a a' → EG π b b' → EG π (.ndiv a b) (.ndiv a' b')
  | nmod {a a' b b' : Ex HS} : EG π a a' → EG π b b' → EG π (.nmod a b) (.nmod a' b')

/-- Register programs with generalized literal substitution. -/
inductive GSim (π : PMap) : RS → RS → Prop
  | skip : GSim π .skip .skip
  | asg (i : Fin HS) {e e' : Ex HS} : EG π e e' → GSim π (.asg i e) (.asg i e')
  | push (i : Fin HS) {e e' : Ex HS} : EG π e e' → GSim π (.push i e) (.push i e')
  | pop (i : Fin HS) : GSim π (.pop i) (.pop i)
  | seq {a a' b b' : RS} : GSim π a a' → GSim π b b' → GSim π (.seq a b) (.seq a' b')
  | ite {c c' : Ex HS} {a a' b b' : RS} : EG π c c' → GSim π a a' → GSim π b b' →
      GSim π (.ite c a b) (.ite c' a' b')
  | rep (i : Fin HS) {n n' : Ex HS} {b b' : RS} : EG π n n' → GSim π b b' →
      GSim π (.rep i n b) (.rep i n' b')

/-- Push the value of an expression on any stack `j ≠ 56`. -/
def pushE (e : Ex HS) (j : Fin HS) : IStmt HS Pp := .seq (e.code NLogN.zX) (.op (.mv 56 j))

end NLogN.KFV

end

section
-- Source module: Solutions.KFUUBase
/-!
# KFUU §0: tools for the uniform `Φ` / `Φ⁻¹`

* `MS`: the `MulOK` spec of `kTable` (`KFM.mulCostH (zsz a + zsz b)`, total domain).
  `mulC pM a b r` pushes `top b * top a` on `r` with one call (`mulC_runs`).
* `ISim.uloop`, `ISim.rpt`: literal-count repetitions `rpt` become register-count `rptX`.
* `scan_runs`: a register-counted loop that moves a list from `X` to `Y` one entry per
  iteration and runs a body before each move.
-/

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.KFUU

open IntMul KIL KFC KFT AModel AModelL12

/-- The multiplication spec of `kTable`. -/
def MS : MulSpec := ⟨fun a b => KFM.mulCostH (zsz a + zsz b), fun _ _ => True⟩

section Gen

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

/-! ## One multiplication -/

/-- Push `top b * top a` on `r` (via `ARGM = 55`). -/
def mulC (pM : Fin P) (a b r : Fin HS) : IStmt HS P :=
  .seq (.op (.dup a ARGM)) (.seq (.op (.dup b ARGM)) (.seq (.call pM) (.op (.mv ARGM r))))

def cMul (x y : ℤ) : ℕ :=
  (dupCost (zsz x) + 17) + (dupCost (zsz y) + 17) + KFM.mulCostH (zsz y + zsz x) +
    (2 * moveCost (zsz (y * x)) + 19)

/-! ## Literal counts to register counts -/

/-! ## Scanning a list -/

/-- The scan state: `U k` with `X` holding the rest and `Y` the moved entries. -/
def scS (U : ℕ → IState HS) (X Y : Fin HS) (L : List ℤ) (RX RY : List Item) (k : ℕ) : IState HS :=
  ((U k).upd X (ints (L.drop k) ++ RX)).upd Y (ints (L.take k).reverse ++ RY)

/-- The state after the body of iteration `k` (before the move). -/
def scM (U : ℕ → IState HS) (X Y : Fin HS) (L : List ℤ) (RX RY : List Item) (k : ℕ) : IState HS :=
  ((U (k + 1)).upd X (ints (L.drop k) ++ RX)).upd Y (ints (L.take k).reverse ++ RY)

/-! ## The cost unit -/

/-- One arithmetic step (a multiplication or a small stack operation) on numbers of size `≤ Z`. -/
def cU (Z : ℕ) : ℕ := 10 ^ 4 * (Z + 1) + KFM.fmB (2 * Z)

end Gen

end NLogN.KFUU

end

section
-- Source module: Solutions.KFUURot1
/-!
# KFUU §7a: register programs for the uniform shifts — helpers and the stage lists

The uniform prefix-controlled shift of target `i` (sizes `s_j = sF j`, coefficient list
`cs = cF neg sF μ Q`) computes the stage data of KIL4's uniform table generator at run time:

* `Pv sF j = ∏_{k<j} s_k`;
* `stgM_rev` / `stgAM_rev`: closed forms of the stage stacks `stgM`, `stgAM` of the
  stage list `(A.zip cs).reverse` (`A = [s_0, …, s_{i-1}]`): `P_i / P_{j+1}` and `P_i / P_j`;
* `encK_lift`: a stack prefix of the unmanaged part becomes managed.
-/

set_option linter.unusedSimpArgs false
set_option linter.unusedVariables false
set_option linter.unnecessarySeqFocus false

namespace NLogN.KFUU

open IntMul KIL AModel

/-! ## Managed-state helpers -/

/-! ## Prefix products and the stage lists -/

/-- `P_j = ∏_{k<j} s_k`. -/
def Pv (sF : ℕ → ℕ) (j : ℕ) : ℕ := vol ((List.range j).map sF)

end NLogN.KFUU

end

section
-- Source module: Solutions.KFUURot2
/-!
# KFUU §7b: the coefficient loop of a uniform shift

Registers of a shift of target `i` (all values natural numbers):
`oI = i`, `oQ = Q = s_i`, `oMu = μ_i`, `oPi = P_i`, `oBv = B` (volume behind the target),
`oSP`: the sizes `s_{i-1}, …, s_0` (top first).

`coefB neg` is one step `j` of the coefficient loop: it pushes the stage data of field `j`
(`c_j`, `P_i / P_j`, `P_i / P_{j+1}`, `s_j`) on the stage stacks of KIL4 and advances
`oPr = P_j` to `P_{j+1}`, moving `s_j` from `oSQ` to `oSP`.  `c_j = μ P_j mod Q`
(forward, `neg = false`) or `(Q - μ P_j mod Q) mod Q` (backward).
-/

set_option linter.unusedSimpArgs false
set_option linter.unusedVariables false
set_option linter.unnecessarySeqFocus false

namespace NLogN.KFUU

open IntMul KIL AModel

abbrev oI : Fin HS := 13
abbrev oQ : Fin HS := 18
abbrev oMu : Fin HS := 19
abbrev oPi : Fin HS := 20
abbrev oPr : Fin HS := 22
abbrev oBv : Fin HS := 24
abbrev oSP : Fin HS := 25
abbrev oSQ : Fin HS := 26
abbrev oCC : Fin HS := 27
abbrev oCS : Fin HS := 28
abbrev oCM : Fin HS := 29

/-- `μ · P_j mod Q`. -/
def cMx : Ex HS := .nmod (.mul (.reg oMu) (.reg oPr)) (.reg oQ)

/-- The control of the current field. -/
def cE : Bool → Ex HS
  | false => cMx
  | true => .nmod (.add (.reg oQ) (.neg cMx)) (.reg oQ)

/-- Its value. -/
def cF (neg : Bool) (sF : ℕ → ℕ) (μ Q j : ℕ) : ℕ :=
  if neg then (Q - μ * Pv sF j % Q) % Q else μ * Pv sF j % Q

/-- One step of the coefficient loop. -/
def coefB (neg : Bool) : RS :=
  .seq (.push sSC (NLogN.KFUU.cE neg)) (.seq (.push sSAM (.ndiv (.reg oPi) (.reg oPr)))
    (.seq (.push sSM (.ndiv (.reg oPi) (.mul (.reg oPr) (.reg oSQ))))
      (.seq (.push NLogN.KIL.sSA (.reg oSQ)) (.seq (.asg oPr (.mul (.reg oPr) (.reg oSQ)))
        (.seq (.push oSP (.reg oSQ)) (.pop oSQ))))))

/-- The managed lists after `k` steps. -/
def cfF (f0 : Fin HS → List ℤ) (neg : Bool) (sF : ℕ → ℕ) (μ Q i k : ℕ) : Fin HS → List ℤ :=
  Function.update (Function.update (Function.update (Function.update (Function.update
    (Function.update (Function.update f0
      oSQ (((List.range i).map fun j => (sF j : ℤ)).drop k))
      oSP (((List.range k).map fun j => (sF j : ℤ)).reverse))
      oPr [(Pv sF k : ℤ)])
      sSC (((List.range k).map fun j => (cF neg sF μ Q j : ℤ)).reverse))
      sSAM (((List.range k).map fun j => ((Pv sF i / Pv sF j : ℕ) : ℤ)).reverse))
      sSM (((List.range k).map fun j => ((Pv sF i / Pv sF (j + 1) : ℕ) : ℤ)).reverse))
      NLogN.KIL.sSA (((List.range k).map fun j => (sF j : ℤ)).reverse)

end NLogN.KFUU

end

section
-- Source module: Solutions.KFUURot3
/-!
# KFUU §7c: the loops of the uniform table set-up

* `coefL_sp`: `i` steps of `coefB` (counter `oCC`, count `oI = i`);
* `stgL_sp`: `i` uniform stages `stageU` of KIL4 (counter `oCS`), which turn the stage stacks
  of `ac` into the offset table `tbl Q ac` on stack `14`.
-/

set_option linter.unusedSimpArgs false
set_option linter.unusedVariables false
set_option linter.unnecessarySeqFocus false

namespace NLogN.KFUU

open IntMul KIL AModel

/-- The coefficient loop. -/
def coefL (neg : Bool) : RS := .rep oCC (.reg oI) (coefB neg)

/-! ## The stage loop -/

/-- The stage list of target `i`. -/
def acL (neg : Bool) (sF : ℕ → ℕ) (μ Q i : ℕ) : List (ℕ × ℕ) :=
  ((List.range i).map sF).zip ((List.range i).map (cF neg sF μ Q))

/-- The managed lists after `k` stages (`rv = ac.reverse`). -/
def stF (g : Fin HS → List ℤ) (rv : List (ℕ × ℕ)) (Q k : ℕ) : Fin HS → List ℤ :=
  Function.update (Function.update (Function.update (Function.update (Function.update g
    14 (tbl Q (rv.take k).reverse))
    NLogN.KIL.sSA ((rv.drop k).map fun p => (p.1 : ℤ)))
    sSC ((rv.drop k).map fun p => (p.2 : ℤ)))
    sSM (stgM (rv.drop k) (rv.take k).reverse))
    sSAM (stgAM (rv.drop k) (rv.take k).reverse)

end NLogN.KFUU

end

section
-- Source module: Solutions.KFUURot4
/-!
# KFUU §7d: the uniform table set-up of one shift

`tabU neg` (register code): from the registers `oI = i`, `oQ = Q = s_i`, `oMu = μ`,
`oPi = P_i`, `oBv = B` and the sizes `s_{i-1}, …, s_0` on `oSP`, it pushes the offset table
`tbl Q (acL neg sF μ Q i)` on stack `14`, `P_i, Q, B` on stacks `1, 2, 4` (the parameters of
KFC `rotF`) and `-Q` on `sNQ`; `oSP` is restored.  `tabU_sp`, cost `cTab`.
-/

set_option linter.unusedSimpArgs false
set_option linter.unusedVariables false
set_option linter.unnecessarySeqFocus false

namespace NLogN.KFUU

open IntMul KIL AModel

/-- The table set-up. -/
def tabU (neg : Bool) : RS :=
  .seq (.push oPr (.lit 1)) (.seq (NLogN.mvK oSP oSQ oCM (.reg oI)) (.seq (coefL neg)
    (.seq (.pop oPr) (.seq (.push sNQ (.neg (.reg oQ))) (.seq (.push 14 (.lit 0))
      (.seq (.rep oCS (.reg oI) stageU) (.seq (.push 1 (.reg oPi))
        (.seq (.push 2 (.reg oQ)) (.push 4 (.reg oBv))))))))))

/-- Its cost bound. -/
def cTab (i P B Bb : ℕ) : ℕ := 10 ^ 10 * (i + 1) * ((B + 2) ^ 2 + (P + 1) * (Bb + 2))

/-- The managed lists at the end of the set-up. -/
def tabEnd (neg : Bool) (sF : ℕ → ℕ) (μ Q i Pi Bv : ℕ) : Fin HS → List ℤ :=
  Function.update (Function.update (Function.update (Function.update (Function.update
    (Function.update (fun _ => []) oSP ((List.range i).map fun j => (sF j : ℤ)).reverse)
    14 (tbl Q (acL neg sF μ Q i))) 1 [(Pi : ℤ)]) 2 [(Q : ℤ)]) 4 [(Bv : ℤ)]) sNQ [-(Q : ℤ)]

end NLogN.KFUU

end

section
-- Source module: Solutions.KFUURot5
/-!
# KFUU §7e: one uniform prefix-controlled shift

`rotCU neg`: the table set-up `tabU neg` (register code), KFC `rotF`, and the pops of the
parameters.  `rotCU_runs`: for `A = [s_0, …, s_{i-1}]` and any suffix `C` with
`vol C = Bv`, the data `xs` of shape `A ++ Q :: C` become KFU's controlled shift
`pull … (rotG |A| Q (fun a => lf cs a % Q))` with `cs = [cF neg sF μ Q j | j < i]`; the code does
not depend on `i, Q, μ, C`.
-/

set_option linter.unusedSimpArgs false
set_option linter.unusedVariables false
set_option linter.unnecessarySeqFocus false

namespace NLogN.KFUU

open IntMul KIL KFC AModel AModelL12

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

/-- The pops after the shift. -/
def postU : RS := .seq (.pop 1) (.seq (.pop 2) (.seq (.pop 4) (.pop sNQ)))

/-- One uniform shift. -/
def rotCU (neg : Bool) : IStmt HS Pp := .seq (NLogN.KFUU.tabU neg).code (.seq rotF postU.code)

/-- The coefficients of target `i`. -/
def csU (neg : Bool) (sF : ℕ → ℕ) (μ Q i : ℕ) : List ℕ := (List.range i).map (cF neg sF μ Q)

/-- Its cost. -/
def cRotU (i P Q Bv w B Bb : ℕ) : ℕ := NLogN.KFUU.cTab i P B Bb + cRotF P Q Bv w + 8 * B + 8

end NLogN.KFUU

end

section
-- Source module: Solutions.KFUURot6
/-!
# KFUU §7f: register updates around the uniform shifts

* `rotIU_runs`: `rotCU` realises KFU's shift `gRot` of target `i` on `shR d (sP d sF)`
  (`neg = false`: the forward coefficients `csF`, `neg = true`: the backward `csB`).
* `preF`/`postF` (forward, targets `d-1, …, 0`) and `preB`/`postB` (backward, `0, …, d-1`):
  pop `s_i`, `μ_i` into `oQ`, `oMu` and update `oPi = P_i`, `oI`, `oBv` (register code).
-/

set_option linter.unusedSimpArgs false
set_option linter.unusedVariables false
set_option linter.unnecessarySeqFocus false

namespace NLogN.KFUU

open IntMul KIL KFC AModel AModelL12

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

/-- The list of remaining sizes (backward). -/
abbrev oSL : Fin HS := 3
/-- The list of multipliers. -/
abbrev oML : Fin HS := 32

/-- The volume behind target `i`. -/
def Cv (d : ℕ) (sF : ℕ → ℕ) (i : ℕ) : ℕ := vol (((List.range d).drop (i + 1) ++ [d]).map (KFU.sP d sF))

/-! ## Register updates -/

/-- Forward: take target `i` (from `oI = i+1`, `oPi = P_{i+1}`). -/
def preF : RS :=
  .seq (.push oQ (.reg oSP)) (.seq (.pop oSP) (.seq (.push oMu (.reg oML)) (.seq (.pop oML)
    (.seq (.asg oPi (.ndiv (.reg oPi) (.reg oQ))) (.asg oI (.add (.reg oI) (.lit (-1))))))))

/-- Forward: after the shift (`oBv := oBv · Q`, pop `Q`, `μ`). -/
def postF : RS :=
  .seq (.asg oBv (.mul (.reg oBv) (.reg oQ))) (.seq (.pop oQ) (.pop oMu))

/-- Backward: take target `i` (from `oSL`, `oML`; `oBv := oBv / Q`). -/
def preB : RS :=
  .seq (.push oQ (.reg oSL)) (.seq (.pop oSL) (.seq (.push oMu (.reg oML)) (.seq (.pop oML)
    (.asg oBv (.ndiv (.reg oBv) (.reg oQ))))))

/-- Backward: after the shift (push `Q` on `oSP`, `oPi := oPi · Q`, `oI := oI + 1`, pops). -/
def postB : RS :=
  .seq (.push oSP (.reg oQ)) (.seq (.asg oPi (.mul (.reg oPi) (.reg oQ)))
    (.seq (.asg oI (.add (.reg oI) (.lit 1))) (.seq (.pop oQ) (.pop oMu))))

end NLogN.KFUU

end

section
-- Source module: Solutions.KFUUPad
/-!
# KFUU §1: uniform padding and deletion scans

`padIU` / `unpadIU` are KFU's `padI X B Z` / `unpadI X B C` with the three literal counts read
from the parameter stacks `qX, qB, qZ = 50, 51, 52` (for `unpadIU`, `qZ` holds `C - B`).
`padIU_runs` / `unpadIU_runs` follow from `padI_spec` / `unpadI_spec` by `KFV.ISim.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.KFUU

open IntMul KIL KFC KFT KFU AModel AModelL12

abbrev qX : Fin HS := 50
abbrev qB : Fin HS := 51
abbrev qZ : Fin HS := 52

section Gen

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

def padBlkU : IStmt HS P :=
  .seq (rptX NLogN.KFUU.qB zA zC NLogN.KFU.zD (.op (.mv 0 zO))) (rptX NLogN.KFUU.qZ zA zC NLogN.KFU.zD (.op (.const zO 0)))

/-- Uniform padding scan. -/
def padIU : IStmt HS P :=
  .seq (.op (.mark zO)) (.seq (rptX NLogN.KFUU.qX oA oC NLogN.KFU.oD padBlkU) (.seq (mvW zO 0) (.op (.unmark zO))))

def unpBlkU : IStmt HS P :=
  .seq (rptX NLogN.KFUU.qB zA zC NLogN.KFU.zD (.op (.mv 0 zO))) (rptX NLogN.KFUU.qZ zA zC NLogN.KFU.zD (.op (.drop 0)))

/-- Uniform deletion scan. -/
def unpadIU : IStmt HS P :=
  .seq (.op (.mark zO)) (.seq (rptX NLogN.KFUU.qX oA oC NLogN.KFU.oD unpBlkU) (.seq (mvW zO 0) (.op (.unmark zO))))

/-- The parameter map of one scan. -/
def π3 (X B Z : ℕ) : PMap := pmOf [(NLogN.KFUU.qX, (X : ℤ)), (NLogN.KFUU.qB, (B : ℤ)), (NLogN.KFUU.qZ, (Z : ℤ))]

/-! ## The field loop of `embed` -/

abbrev tS : Fin HS := 47
abbrev tT : Fin HS := 48
abbrev pV : Fin HS := 49
abbrev pXr : Fin HS := 53
abbrev tq : Fin HS := 54
abbrev lA : Fin HS := 34
abbrev lC : Fin HS := 35
abbrev lD : Fin HS := 36

/-- Parameters of one padding scan: `qX := X`, `qB := V_k`, `qZ := T_k V_{k+1} - V_k`. -/
def prmB (pM : Fin P) : IStmt HS P :=
  .seq (.op (.dup pXr NLogN.KFUU.qX)) (.seq (.op (.mv pV NLogN.KFUU.qB)) (.seq (mulC pM pV NLogN.KFUU.tT NLogN.KFUU.qZ)
    (.seq (.op (.dup NLogN.KFUU.qB tq)) (.seq (.op (.neg tq)) (.op (.add NLogN.KFUU.qZ tq NLogN.KFUU.qZ))))))

def cPrm (X Vk V1 Tk : ℤ) : ℕ :=
  (dupCost (zsz X) + 17) + (2 * moveCost (zsz Vk) + 19) + cMul V1 Tk + (dupCost (zsz Vk) + 17) + 20 +
    addIntCost (zsz (Tk * V1) + zsz (-Vk))

/-- Stacks that may hold the input lists of the generic passes. -/
def LS (i : Fin HS) : Prop :=
  i.val = 5 ∨ i.val = 6 ∨ i.val = 7 ∨ i.val = 8 ∨ i.val = 9 ∨ i.val = 43 ∨ i.val = 44 ∨ i.val = 45

/-- The count register of the generic passes. -/
abbrev nR : Fin HS := 46

/-- After one padding scan: drop the parameters, `X ← T_k X`, move `T_k` back to `pT`. -/
def finB (pM : Fin P) (pT : Fin HS) : IStmt HS P :=
  .seq (.op (.drop NLogN.KFUU.qX)) (.seq (.op (.drop NLogN.KFUU.qB)) (.seq (.op (.drop NLogN.KFUU.qZ)) (.seq (mulC pM pXr NLogN.KFUU.tT tq)
    (.seq (.op (.drop pXr)) (.seq (.op (.mv tq pXr)) (.op (.mv NLogN.KFUU.tT pT)))))))

def cFin (a b c X Tk : ℤ) : ℕ :=
  (2 * zsz a + 2) + (2 * zsz b + 2) + (2 * zsz c + 2) + cMul X Tk + (2 * zsz X + 2) +
    (2 * moveCost (zsz (Tk * X)) + 19) + (2 * moveCost (zsz Tk) + 19)

/-- One field of the padding loop. -/
def padMB (pM : Fin P) (pT : Fin HS) : IStmt HS P := .seq (prmB pM) (.seq padIU (finB pM pT))

def cMB (X Vk V1 Tk w : ℕ) : ℕ :=
  cPrm X Vk V1 Tk + 2000 * cPadI X Vk (Tk * V1 - Vk) w +
    cFin X Vk ((Tk * V1 - Vk : ℕ) : ℤ) X Tk

/-! ## The volume lists -/

/-- A shape as a list of integers. -/
def zl (S : Shape) : List ℤ := List.map (fun a : ℕ => (a : ℤ)) S

/-- The suffix volumes `vol (S.drop i)` for `j ≤ i ≤ |S|`. -/
def sufV (S : Shape) (j : ℕ) : List ℤ :=
  ((List.range (S.length + 1)).drop j).map (fun i => (vol (S.drop i) : ℤ))

/-- Phase 1 of `padAllU`: the suffix volumes of `S` (given reversed on `pS`) on `pV`, while `S`
moves to `tS` (in order). -/
def sufP (pM : Fin P) (pS : Fin HS) : IStmt HS P :=
  rptX nR NLogN.KFUU.lA lC lD (.seq (mulC pM pV pS pV) (.op (.mv pS tS)))

/-- Move a list of `n` entries (`n` on `nR`) from `X` to `Y` (reversing it). -/
def mvL (X Y : Fin HS) : IStmt HS P := rptX nR NLogN.KFUU.lA lC lD (.seq .skip (.op (.mv X Y)))

/-- The main loop of `padAllU`. -/
def padLp (pM : Fin P) (pS pT : Fin HS) : IStmt HS P :=
  rptX nR NLogN.KFUU.lA lC lD (.seq (padMB pM pT) (.op (.mv tS pS)))

/-- **Uniform padding** `S ↦ T` (`S.reverse` on `pS`, `T.reverse` on `pT`, `|S|` on `nR`). -/
def padAllU (pM : Fin P) (pS pT : Fin HS) : IStmt HS P :=
  .seq (.op (.const pV 1)) (.seq (sufP pM pS) (.seq (mvL pT NLogN.KFUU.tT) (.seq (.op (.const pXr 1))
    (.seq (padLp pM pS pT) (.seq (.op (.drop pXr)) (.op (.drop pV)))))))

end Gen

end NLogN.KFUU

end

section
-- Source module: Solutions.KFUUUnpad
/-!
# KFUU §1b: uniform deletion of the padding

`unpadAllU pM pS pT`: `ys` on the padded shape `T` becomes `extract S T 0 ys` (KFU
`unpadAll_spec` with the per-field counts computed at run time from `S.reverse` on `pS`,
`T.reverse` on `pT` and `|S|` on `nR`).
-/

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.KFUU

open IntMul KIL KFC KFT KFU AModel AModelL12

section Gen

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

/-- Parameters of one deletion scan: `qX := X`, `qB := S_k W_{k+1}`, `qZ := W_k - S_k W_{k+1}`. -/
def prmU (pM : Fin P) : IStmt HS P :=
  .seq (.op (.dup pXr NLogN.KFUU.qX)) (.seq (.op (.mv pV NLogN.KFUU.qZ)) (.seq (mulC pM pV NLogN.KFUU.tT NLogN.KFUU.qB)
    (.seq (.op (.dup NLogN.KFUU.qB tq)) (.seq (.op (.neg tq)) (.op (.add NLogN.KFUU.qZ tq NLogN.KFUU.qZ))))))

def cPrmU (X Wk W1 Sk : ℤ) : ℕ :=
  (dupCost (zsz X) + 17) + (2 * moveCost (zsz Wk) + 19) + cMul W1 Sk +
    (dupCost (zsz (Sk * W1)) + 17) + 20 + addIntCost (zsz Wk + zsz (-(Sk * W1)))

/-- One field of the deletion loop. -/
def unpMB (pM : Fin P) (pS : Fin HS) : IStmt HS P := .seq (prmU pM) (.seq unpadIU (finB pM pS))

def cUMB (X Wk W1 Sk w : ℕ) : ℕ :=
  cPrmU X Wk W1 Sk + 2000 * cUnpadI X (Sk * W1) Wk w +
    cFin X ((Sk * W1 : ℕ) : ℤ) ((Wk - Sk * W1 : ℕ) : ℤ) X Sk

/-- The main loop of `unpadAllU`. -/
def unpLp (pM : Fin P) (pS pT : Fin HS) : IStmt HS P :=
  rptX nR NLogN.KFUU.lA lC lD (.seq (unpMB pM pS) (.op (.mv tS pT)))

/-- **Uniform deletion of the padding** `T ↦ S` (`S.reverse` on `pS`, `T.reverse` on `pT`,
`|S|` on `nR`). -/
def unpadAllU (pM : Fin P) (pS pT : Fin HS) : IStmt HS P :=
  .seq (.op (.const pV 1)) (.seq (sufP pM pT) (.seq (mvL pS NLogN.KFUU.tT) (.seq (.op (.const pXr 1))
    (.seq (unpLp pM pS pT) (.seq (.op (.drop pXr)) (.op (.drop pV)))))))

end Gen

end NLogN.KFUU

end

section
-- Source module: Solutions.KFUURev1
/-!
# KFUU §2a: register-driven one-bit moves on layouts

`mvFU` / `mvBU` are KFU's `mvFC a 2 g c` / `mvBC a g 2 c` with `a, g, c` read from the
registers `qX, qB, qZ` (`mvFU_lay`, `mvBU_lay`, the `Dl`-layout forms of `mvF_lay` / `mvB_lay`).
-/

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.KFUU

open IntMul KIL KFC KFT KFU AModel AModelL12

section Gen

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

/-- `fwdF` with `vol A, vol G, vol C` from `qX, qB, qZ` and `vol F = 2`. -/
def mvFU : IStmt HS P :=
  .seq (.op (.dup NLogN.KFUU.qX 1)) (.seq (.op (.const 2 2)) (.seq (.op (.dup NLogN.KFUU.qB 3)) (.seq (.op (.dup NLogN.KFUU.qZ 4))
    (.seq fwdF (.seq (.op (.drop 1)) (.seq (.op (.drop 2)) (.seq (.op (.drop 3)) (.op (.drop 4)))))))))

/-- `bwdF` with `vol A, vol G, vol C` from `qX, qB, qZ` and `vol F = 2`. -/
def mvBU : IStmt HS P :=
  .seq (.op (.dup NLogN.KFUU.qX 1)) (.seq (.op (.const 2 2)) (.seq (.op (.dup NLogN.KFUU.qB 3)) (.seq (.op (.dup NLogN.KFUU.qZ 4))
    (.seq bwdF (.seq (.op (.drop 1)) (.seq (.op (.drop 2)) (.seq (.op (.drop 3)) (.op (.drop 4)))))))))

/-- The register overhead of one move. -/
def cK3 (a g c : ℕ) : ℕ :=
  (dupCost (zsz (a : ℤ)) + 17) + (zsz (2 : ℤ) + 2) + (dupCost (zsz (g : ℤ)) + 17) +
    (dupCost (zsz (c : ℤ)) + 17) + (2 * zsz (a : ℤ) + 2) + (2 * zsz (2 : ℤ) + 2) +
    (2 * zsz (g : ℤ) + 2) + (2 * zsz (c : ℤ) + 2)

/-! ## Layout forms -/

end Gen

end NLogN.KFUU

end

section
-- Source module: Solutions.KFUURev2
/-!
# KFUU §2b: the layouts of the uniform reversal

The axes are `L = range n` with sizes `2^(u + b_j)` (`bs = [b_0, …]`), the inner block is the
atom `cid = n` of size `cz`: `mzB u n cz`, `bgB bs`. Phase 1 runs through `lay1 i`
(`i = n, …, 0`): `layA (L.take i) ++ mains (L.drop i) ++ bitsR (L.drop i) ++ [2n]`.
-/

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.KFUU

open IntMul KIL KFC KFT KFU AModel AModelL12

/-- Main sizes `2^u` on the axes `j < n`, `cz` for the inner block. -/
def mzB (u n cz : ℕ) (j : ℕ) : ℕ := if j < n then 2 ^ u else cz

/-- The extra bits. -/
def bgB (bs : List Bool) (j : ℕ) : Bool := bs.getD j false

/-- The size list. -/
abbrev shP (u : ℕ) (bs : List Bool) : Shape := KFM.shPad u bs

/-! ## Phase 1 layouts -/

/-- The layout after phase 1 has treated the axes `i, …, n-1`. -/
def lay1 (bg : ℕ → Bool) (n i : ℕ) : List ℕ :=
  layA bg ((List.range n).take i) ++ mains ((List.range n).drop i) ++
    bitsR bg ((List.range n).drop i) ++ [2 * n]

/-! ## Volumes -/

/-! ## One step of phase 1 -/

abbrev rU2 : Fin HS := 37
abbrev tq2 : Fin HS := 42

section Gen

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

/-- `tq := S_i - 2^u` (zero, or positive for an axis with an extra bit). -/
def bitT : IStmt HS P :=
  .seq (.op (.dup tS tq)) (.seq (.op (.dup rU2 tq2)) (.seq (.op (.neg tq2)) (.op (.add tq tq2 tq))))

/-- `qB := S_i · qB`. -/
def accB (pM : Fin P) : IStmt HS P :=
  .seq (mulC pM NLogN.KFUU.qB tS tq) (.seq (.op (.drop NLogN.KFUU.qB)) (.op (.mv tq NLogN.KFUU.qB)))

def isPos (c : Cls) : Bool := decide (c = .pos)

/-- One step of phase 1. -/
def p1B (pM : Fin P) : IStmt HS P :=
  .seq (mulC pM pV rU2 NLogN.KFUU.qX) (.seq bitT (.seq (.ite tq NLogN.KFUU.isPos mvFU .skip) (.seq (.op (.drop tq))
    (.seq (accB pM) (.seq (.op (.drop NLogN.KFUU.qX)) (.op (.drop pV)))))))

end Gen

theorem len_shP (u : ℕ) (bs : List Bool) : (shP u bs).length = bs.length := by simp [KFM.shPad]

section Gen

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

end Gen

end NLogN.KFUU

end

section
-- Source module: Solutions.KFUURev3
/-!
# KFUU §2c: phase 1 of the uniform reversal

`p1U pM pS` (sizes `S = shP u bs` on `pS` in order, `n = |bs|` on `nR`, `2^u` on `rU2`, `cz` on
`rCZ`) moves every extra bit behind the main atoms: `Dl (lay1 n) ↦ Dl (lay1 0)`.
-/

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.KFUU

open IntMul KIL KFC KFT KFU AModel AModelL12

abbrev rU : Fin HS := 38
abbrev rCZ : Fin HS := 39
abbrev rW : Fin HS := 40

section Gen

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

/-- Phase 1. -/
def p1U (pM : Fin P) (pS : Fin HS) : IStmt HS P :=
  .seq (.op (.const pV 1)) (.seq (sufP pM pS) (.seq (.op (.drop pV)) (.seq (.op (.const NLogN.KFUU.qB 1))
    (.seq (.op (.dup rCZ NLogN.KFUU.qZ)) (.seq (rptX nR NLogN.KFUU.lA lC lD (.seq (p1B pM) (.op (.mv tS pS))))
      (.seq (.op (.drop NLogN.KFUU.qB)) (.op (.drop NLogN.KFUU.qZ))))))))

/-- The phase-1 loop alone. -/
def p1L (pM : Fin P) (pS : Fin HS) : IStmt HS P :=
  rptX nR NLogN.KFUU.lA lC lD (.seq (p1B pM) (.op (.mv tS pS)))

end Gen

end NLogN.KFUU

end

section
-- Source module: Solutions.KFUURev4
/-!
# KFUU §2d: one Lemma 12 swap through the `SwapNeed` entry

* `pull_swap5`: exchanging two equal fields of `A ++ F :: M ++ F :: C` only depends on
  `vol A, vol M, vol C` (the 5-field shape `[vol A, F, vol M, F, vol C]`).
* `gpair`: pairs of a flat integer list (`gflat (gpair xs) = xs`).
-/

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.KFUU

open IntMul KIL KFB KFC KFT KFU AModel AModelL12

/-! ## Pairs -/

def gpair : List ℤ → List GI
  | a :: b :: r => (a, b) :: gpair r
  | _ => []

end NLogN.KFUU

end