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

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

import Definitions.Def_IntMul_Kappa_K_Defs27

set_option autoImplicit false
set_option relaxedAutoImplicit false
set_option maxHeartbeats 1000000

section
-- Source module: Solutions.KFUURev5
/-!
# KFUU §2e: the `SwapNeed` entry inside a table, and one swap step of phase 2

* `SwOKU`: the conclusion of `KFM.SwapNeed` for a table `iprog` and an entry `eSw`.
* `swOKU_run`: the same with the frame written as a plain state (`t.G 0`, `t.G dS`).
-/

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

section Gen

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

/-- The cost bound of one `SwapNeed` call. -/
noncomputable def cSwN (K : ℝ) (w u : ℕ) (sh : Shape) : ℝ :=
  K * (((w : ℝ) + 1) * ((vol sh : ℝ) + 1) * (u : ℝ) ^ KFS.τS +
    (((sh.length + 1) * (Nat.size (vol sh) + w + u + 1) ^ 2 : ℕ) : ℝ))

/-- What `KFM.SwapNeed` provides, in the table `iprog` at the entry `eSw`. -/
def SwOKU (I : List Sym) (iprog : Fin P → IStmt HS P) (eSw : Fin P) (dS : Fin HS) (K : ℝ) : Prop :=
  ∀ (u w : ℕ) (sh : Shape) (p q : ℕ), 1 ≤ u → (∀ a ∈ sh, 1 ≤ a) →
    p < q → q < sh.length → sh.getD p 0 = 2 ^ u → sh.getD q 0 = 2 ^ u →
  ∀ (s : IState HS) (f : Fin HS → List ℤ) (zs : List GI) (l : List ℤ),
    (∀ j, j ≠ 0 → j ≠ dS → f j = []) →
    f dS = KFM.descSw w u p q sh →
    f 0 = gflat zs ++ l → zs.length = vol sh → NLogN.KFB.GB w zs →
    ∃ c : ℕ, IRuns I iprog (.call eSw) (NLogN.KIL.encK s f)
        (NLogN.KIL.encK s (Function.update f 0 (gflat (pull sh sh (swapPos p q) 0 zs) ++ l))) c ∧
      (c : ℝ) ≤ cSwN K w u sh

/-! ## Moving and dropping a fixed number of items -/

def mvK (X Y : Fin HS) : ℕ → IStmt HS P
  | 0 => .skip
  | k + 1 => .seq (.op (.mv X Y)) (mvK X Y k)

def dropK (X : Fin HS) : ℕ → IStmt HS P
  | 0 => .skip
  | k + 1 => .seq (.op (.drop X)) (dropK X k)

/-! ## Pushing a list of register values and constants -/

def pushOp (D : Fin HS) : Fin HS ⊕ ℤ → IStmt HS P
  | .inl X => .op (.dup X D)
  | .inr c => .op (.const D c)

def pushL (D : Fin HS) : List (Fin HS ⊕ ℤ) → IStmt HS P
  | [] => .skip
  | o :: os => .seq (pushOp D o) (pushL D os)

/-- The source `o` provides the value `v` in the state `t` (pushing onto `D`). -/
def PushOK (t : IState HS) (D : Fin HS) : Fin HS ⊕ ℤ → ℤ → Prop
  | .inl X, v => X ≠ D ∧ ∃ L, t.G X = .int v :: L
  | .inr c, v => c = v

/-! ## One swap through the entry -/

abbrev rA : Fin HS := 41
abbrev rB : Fin HS := 15
abbrev rM : Fin HS := 16
abbrev rH : Fin HS := 17

/-- The sources of the swap descriptor `[w, u, 1, 3, 5, A, 2^u, M, 2^u, C]`. -/
def descSrc : List (Fin HS ⊕ ℤ) :=
  [.inl rW, .inl NLogN.KFUU.rU, .inr 1, .inr 3, .inr 5, .inl rA, .inl rU2, .inl NLogN.KFUU.rM, .inl rU2, .inl NLogN.KFUU.rB]

/-- Build the descriptor on `Dsc`, move it to `dS`, call the swap entry, drop the descriptor. -/
def swS (eSw : Fin P) (dS Dsc : Fin HS) : IStmt HS P :=
  .seq (NLogN.KFUU.pushL Dsc descSrc) (.seq (mvK Dsc dS 10) (.seq (.call eSw) (dropK dS 10)))

end Gen

end NLogN.KFUU

end

section
-- Source module: Solutions.KFUURev6
/-!
# KFUU §2f: the layouts of phase 2 (swaps of the main fields)

After phase 1 the layout is `lay1 bg n 0 = mains ++ bitsR ++ [2n]`, of field sizes
`2^u (n times) ++ ctB ++ [cz]`. Swap `k` exchanges the main fields `k` and `n-1-k`; on the data it
is the 5-field swap `[U^k, U, U^{n-2-2k}, U, U^k·vol ctB·cz]`, `U = 2^u`.
-/

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

/-- The bit fields after phase 1. -/
def ctB (u cz : ℕ) (bs : List Bool) : Shape :=
  (bitsR (bgB bs) (List.range bs.length)).map (asz (mzB u bs.length cz))
end NLogN.KFUU
end

section
-- Source module: Solutions.KFUURev7
/-!
# KFUU §2g: phase 2 in IL (one iteration)

`p2B` divides `M` by `U²`, calls the swap entry with the descriptor
`[w, u, 1, 3, 5, U^k, U, U^{n-2-2k}, U, U^k·Bh]`, and multiplies `A = U^k` and `C = U^k·Bh` by `U`.
-/

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

section Gen

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

/-- `i := i / 2^m` (an exact division), the remainder through `k`. -/
def divE (i j k : Fin HS) : IStmt HS P := .seq (.op (.split i j k)) (.op (.drop k))

/-- The cost of one `SwapNeed` call on a 5-field shape of volume `V`. -/
noncomputable def cSwV (K : ℝ) (w u V : ℕ) : ℝ :=
  K * (((w : ℝ) + 1) * ((V : ℝ) + 1) * (u : ℝ) ^ KFS.τS +
    (((5 + 1) * (Nat.size V + w + u + 1) ^ 2 : ℕ) : ℝ))

/-- One iteration of phase 2. -/
def p2B (eSw : Fin P) (dS Dsc : Fin HS) : IStmt HS P :=
  .seq (divE NLogN.KFUU.rM rU2 tq) (.seq (divE NLogN.KFUU.rM rU2 tq) (.seq (swS eSw dS Dsc)
    (.seq (.op (.shl rA rU2)) (.op (.shl NLogN.KFUU.rB rU2)))))

end Gen

end NLogN.KFUU

end

section
-- Source module: Solutions.KFUURev8
/-!
# KFUU §2h: the set-up of phase 2 (volume, powers)

* `volU pM pS`: push `vol S` on `rB` (the sizes `S` on `pS`, restored).
* `powL`: push `U^n` on `rM` (`U = 2^u` on `rU2`, `n` on `nR`).
* `rptXU_runs`: `rptX_runs` with a simple cost.
-/

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

section Gen

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

/-! ## The volume of a size list -/

def volB (pM : Fin P) (pS : Fin HS) : IStmt HS P :=
  .seq (mulC pM NLogN.KFUU.rB pS tq) (.seq (.op (.drop NLogN.KFUU.rB)) (.op (.mv tq NLogN.KFUU.rB)))

def volU (pM : Fin P) (pS : Fin HS) : IStmt HS P :=
  .seq (.op (.const NLogN.KFUU.rB 1)) (.seq (rptX nR NLogN.KFUU.lA lC lD (.seq (volB pM pS) (.op (.mv pS tS))))
    (mvL tS pS))

/-! ## Powers -/

def powL : IStmt HS P := .seq (.op (.const NLogN.KFUU.rM 1)) (rptX nR NLogN.KFUU.lA lC lD (.op (.shl NLogN.KFUU.rM rU2)))

/-! ## The phase-2 loop -/

/-- The state before swap `k`. -/
def p2St (t : IState HS) (R : List Item) (u cz cr : ℕ) (bs : List Bool) (xs0 : List ℤ)
    (LA LM LB : List Item) (k : ℕ) : IState HS :=
  (((t.upd 0 (ints (Dl (mzB u bs.length cz) (bgB bs) (List.range bs.length ++ [bs.length])
    (rv bs.length k (lay1 (bgB bs) bs.length 0)) xs0 0) ++ R)).upd rA
    (.int (((2 ^ u) ^ k : ℕ) : ℤ) :: LA)).upd NLogN.KFUU.rM
    (.int (((2 ^ u) ^ (bs.length - 2 * k) : ℕ) : ℤ) :: LM)).upd NLogN.KFUU.rB
    (.int (((2 ^ u) ^ k * (vol (ctB u cz bs) * cr) : ℕ) : ℤ) :: LB)

def p2L (eSw : Fin P) (dS Dsc : Fin HS) : IStmt HS P := rptX NLogN.KFUU.rH NLogN.KFUU.lA lC lD (p2B eSw dS Dsc)

end Gen

end NLogN.KFUU

end

section
-- Source module: Solutions.KFUURot7
/-!
# KFUU §7g: the forward shifts, uniform

`rotsFU`: a register-counted loop (`oD = d` iterations) of `preF; rotCU false; postF`; it runs
KFU's forward shifts of targets `d-1, …, 0` (`rotsFU_runs`: the result is KFU's
`pull … (Gf sF (csF sF μ) d)`).  Registers before: `oI = d`, `oPi = P_d`, `oBv = 2`,
`oSP = [s_{d-1}, …, s_0]`, `oML = [μ_{d-1}, …, μ_0]`; after: `oI = 0`, `oPi = 1`,
`oBv = vol (shR d (sP d sF))`, `oSP`, `oML` consumed.
-/

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

namespace NLogN.KFUU

open IntMul KIL KFC KFT AModel AModelL12

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

/-- The loop count `d`. -/
abbrev oD : Fin HS := 21

/-- One forward step. -/
def fwB : IStmt HS Pp := .seq preF.code (.seq (NLogN.KFUU.rotCU false) postF.code)

/-- The forward shifts. -/
def rotsFU : IStmt HS Pp := rptX oD NLogN.KFUU.lA lC lD fwB

/-- The address map after `k` forward steps. -/
def Hf (d : ℕ) (sF : ℕ → ℕ) (cs : ℕ → List ℕ) : ℕ → List ℕ → List ℕ
  | 0 => id
  | k + 1 => Hf d sF cs k ∘ KFU.gRot sF cs (d - 1 - k)

/-- The volume behind the remaining targets after `k` forward steps. -/
def BvF (d : ℕ) (sF : ℕ → ℕ) (k : ℕ) : ℕ := vol (((List.range d).drop (d - k) ++ [d]).map (KFU.sP d sF))

/-- The loop state after `k` forward steps. -/
def fwSt (d : ℕ) (sF μ : ℕ → ℕ) (t : IState HS) (R RSP RML LI LP LB : List Item) (xs : List ℤ)
    (k : ℕ) : IState HS :=
  (((((t.upd 0 (ints (pull (KFU.shR d (KFU.sP d sF)) (KFU.shR d (KFU.sP d sF))
    (Hf d sF (KFU.csF sF μ) k) 0 xs) ++ R)).upd oI (.int (d - k : ℕ) :: LI)).upd oPi
    (.int (Pv sF (d - k)) :: LP)).upd oBv (.int (BvF d sF k) :: LB)).upd oSP
    (ints ((List.range (d - k)).map fun j => (sF j : ℤ)).reverse ++ RSP)).upd oML
    (ints ((List.range (d - k)).map fun j => (μ j : ℤ)).reverse ++ RML)

/-- Per-step cost bound. -/
def cStepU (d PM V w B Bb : ℕ) : ℕ := NLogN.KFUU.cTab d PM B Bb + 68 * kap w * V + 3 * 10 ^ 9 * (B + 2) ^ 2

end NLogN.KFUU

end

section
-- Source module: Solutions.KFUURot8
/-!
# KFUU §7h: the backward shifts, uniform

`rotsBU`: a register-counted loop (`oD = d` iterations) of `preB; rotCU true; postB`; it runs
KFU's inverse shifts of targets `0, …, d-1` (`rotsBU_runs`: the result is KFU's
`pull … (Gb sF (csB sF μ) d)`).  Registers before: `oI = 0`, `oPi = 1`,
`oBv = vol (shR d (sP d sF))`, `oSL = [s_0, …, s_{d-1}]`, `oML = [μ_0, …, μ_{d-1}]`; after:
`oI = d`, `oPi = P_d`, `oBv = 2`, `oSP = [s_{d-1}, …, s_0]` pushed, `oSL`, `oML` consumed.
-/

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

namespace NLogN.KFUU

open IntMul KIL KFC KFT AModel AModelL12

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

/-- One backward step. -/
def bwB : IStmt HS Pp := .seq preB.code (.seq (NLogN.KFUU.rotCU true) postB.code)

/-- The backward shifts. -/
def rotsBU : IStmt HS Pp := rptX oD NLogN.KFUU.lA lC lD bwB

/-- The volume from target `k` on. -/
def BvB (d : ℕ) (sF : ℕ → ℕ) (k : ℕ) : ℕ := vol (((List.range d).drop k ++ [d]).map (KFU.sP d sF))

/-- The loop state after `k` backward steps. -/
def bwSt (d : ℕ) (sF μ : ℕ → ℕ) (t : IState HS) (R RSL RSP RML LI LP LB : List Item) (xs : List ℤ)
    (k : ℕ) : IState HS :=
  ((((((t.upd 0 (ints (pull (KFU.shR d (KFU.sP d sF)) (KFU.shR d (KFU.sP d sF))
    (KFU.Gb sF (KFU.csB sF μ) k) 0 xs) ++ R)).upd oI (.int k :: LI)).upd oPi
    (.int (Pv sF k) :: LP)).upd oBv (.int (BvB d sF k) :: LB)).upd oSL
    (ints (((List.range d).drop k).map fun j => (sF j : ℤ)) ++ RSL)).upd oSP
    (ints ((List.range k).map fun j => (sF j : ℤ)).reverse ++ RSP)).upd oML
    (ints (((List.range d).drop k).map fun j => (μ j : ℤ)) ++ RML)

end NLogN.KFUU

end

section
-- Source module: Solutions.KFUUSet1
/-!
# KFUU §8a: copying descriptor entries (register code)

* `mv2 A H X`: pop the top of `X` and push it on `A` and on `H`;
* `mvL2 C n A H X`: `n` of them (counter `C`); `mvL2_sp`.
-/

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

namespace NLogN.KFUU

open IntMul KIL AModel

/-- Copy the top of `X` to `A` and `H`, then pop it. -/
def mv2 (A H X : Fin HS) : RS := .seq (.push A (.reg X)) (.seq (.push H (.reg X)) (.pop X))

/-- `n` copies. -/
def mvL2 (C : Fin HS) (n : Ex HS) (A H X : Fin HS) : RS := .rep C n (NLogN.KFUU.mv2 A H X)

/-- The managed lists after `k` copies. -/
def mvF (F0 : Fin HS → List ℤ) (A H X : Fin HS) (L r : List ℤ) (k : ℕ) : Fin HS → List ℤ :=
  Function.update (Function.update (Function.update F0 X (L.drop k ++ r))
    A ((L.take k).reverse ++ F0 A)) H ((L.take k).reverse ++ F0 H)

end NLogN.KFUU

end

section
-- Source module: Solutions.KFUUSet2
/-!
# KFUU §8b: reading the `Φ` descriptor — the first pieces

Stacks of the set-up: `dP = 44` (the `Φ` descriptor), `hD = 30` (the entries popped, for the
restoration), `xA, xB, bT` (copies), `cL` (loop counter).

* `setA_sp`: `w, h, d` into `rW, rU, oD`;
* `sLoop_sp`: the sizes `s_0, …, s_{d-1}` copied to `Ys`, `pSR`, `xA` (each `[s_{d-1}, …, s_0]`),
  their product into `oPi`.
-/

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

namespace NLogN.KFUU

open IntMul KIL AModel

abbrev dP : Fin HS := 44
abbrev hD : Fin HS := 30
abbrev xA : Fin HS := 31
abbrev xB : Fin HS := 33
abbrev bT : Fin HS := 11
abbrev cL : Fin HS := 10
abbrev pSK : Fin HS := 5
abbrev pTK : Fin HS := 6
abbrev pTR : Fin HS := 8
abbrev pSR : Fin HS := 43
/-- Backward temporaries. -/
abbrev xS : Fin HS := 45
abbrev xM : Fin HS := 9

/-- `w, h, d` into registers. -/
def setA : RS := .seq (NLogN.KFUU.mv2 rW NLogN.KFUU.hD dP) (.seq (NLogN.KFUU.mv2 NLogN.KFUU.rU NLogN.KFUU.hD dP) (NLogN.KFUU.mv2 oD NLogN.KFUU.hD dP))

/-- One step of the size loop. -/
def sB (Ys : Fin HS) : RS :=
  .seq (.push Ys (.reg dP)) (.seq (.push pSR (.reg dP))
    (.seq (.asg oPi (.mul (.reg oPi) (.reg dP))) (NLogN.KFUU.mv2 xA NLogN.KFUU.hD dP)))

/-- The size loop. -/
def sLoop (Ys : Fin HS) : RS := .rep NLogN.KFUU.cL (.reg oD) (NLogN.KFUU.sB Ys)

/-- The sizes `s_0, …, s_{k-1}` as integers. -/
def sZ (sF : ℕ → ℕ) (k : ℕ) : List ℤ := (List.range k).map fun j => (sF j : ℤ)

/-- The managed lists after `k` steps. -/
def sF_k (F0 : Fin HS → List ℤ) (Ys : Fin HS) (sF : ℕ → ℕ) (d : ℕ) (r : List ℤ) (k : ℕ) :
    Fin HS → List ℤ :=
  Function.update (Function.update (Function.update (Function.update (Function.update
    (Function.update F0 dP ((sZ sF d).drop k ++ r)) Ys ((sZ sF k).reverse ++ F0 Ys))
    pSR ((sZ sF k).reverse ++ F0 pSR)) oPi [(Pv sF k : ℤ)]) xA ((sZ sF k).reverse ++ F0 xA))
    NLogN.KFUU.hD ((sZ sF k).reverse ++ F0 NLogN.KFUU.hD)

/-- The padded sizes `t_j = 2^h (1 + b_j)`. -/
def tE : Ex HS := .mul (.reg rU2) (.add (.lit 1) (.reg NLogN.KFUU.bT))

/-- One step of the padded-size loop (bits from `bT`, top `b_{d-1}`). -/
def tB : RS := .seq (.push pTK NLogN.KFUU.tE) (.seq (.push xB NLogN.KFUU.tE) (.pop NLogN.KFUU.bT))

def tLoop : RS := .rep NLogN.KFUU.cL (.reg oD) NLogN.KFUU.tB

/-- The padded sizes `t_0, …, t_{k-1}`. -/
def tZ (h : ℕ) (bL : ℕ → Bool) (k : ℕ) : List ℤ :=
  (List.range k).map fun j => ((2 ^ h * (if bL j then 2 else 1) : ℕ) : ℤ)

/-- The managed lists after `k` steps. -/
def tF_k (F0 : Fin HS → List ℤ) (h d : ℕ) (bL : ℕ → Bool) (k : ℕ) : Fin HS → List ℤ :=
  Function.update (Function.update (Function.update F0 NLogN.KFUU.bT ((KFM.bitsZ d bL).reverse.drop k))
    pTK (((tZ h bL d).reverse.take k).reverse ++ F0 pTK))
    xB (((tZ h bL d).reverse.take k).reverse ++ F0 xB)

end NLogN.KFUU

end

section
-- Source module: Solutions.KFUUSet3
/-!
# KFUU §8c: reading the `Φ` descriptor

`readD Ys Ym` (register code) reads `descPhi w h d bL sF μ` from `dP` and restores it:
`rW = w`, `rU = h`, `oD = d`, `oPi = P_d`; the bits reversed on `bT`; the sizes
`[s_{d-1}, …, s_0]` on `Ys`, `pSR`, `xA`; the multipliers `[μ_{d-1}, …, μ_0]` on `Ym`.
-/

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

namespace NLogN.KFUU

open IntMul KIL AModel

/-- `3 + 3 d`. -/
def n3d : Ex HS := .add (.add (.add (.lit 3) (.reg oD)) (.reg oD)) (.reg oD)

def readD (Ys Ym : Fin HS) : RS :=
  .seq setA (.seq (mvL2 NLogN.KFUU.cL (.reg oD) NLogN.KFUU.bT NLogN.KFUU.hD dP) (.seq (.push oPi (.lit 1)) (.seq (sLoop Ys)
    (.seq (mvL2 NLogN.KFUU.cL (.reg oD) Ym NLogN.KFUU.hD dP) (NLogN.mvK NLogN.KFUU.hD dP NLogN.KFUU.cL n3d)))))

/-- The multipliers as integers. -/
def μZ (μ : ℕ → ℕ) (k : ℕ) : List ℤ := (List.range k).map fun j => (μ j : ℤ)

/-- The managed lists after `readD`. -/
def rdF (F0 : Fin HS → List ℤ) (Ys Ym : Fin HS) (w h d : ℕ) (bL : ℕ → Bool) (sF μ : ℕ → ℕ) :
    Fin HS → List ℤ :=
  Function.update (Function.update (Function.update (Function.update (Function.update
    (Function.update (Function.update (Function.update (Function.update F0
      rW [(w : ℤ)]) NLogN.KFUU.rU [(h : ℤ)]) oD [(d : ℤ)]) oPi [(Pv sF d : ℤ)]) Ys (sZ sF d).reverse)
      pSR (sZ sF d).reverse) xA (sZ sF d).reverse) NLogN.KFUU.bT (KFM.bitsZ d bL).reverse) Ym (μZ μ d).reverse

def cRead (d B : ℕ) : ℕ := 10 ^ 10 * (d + 1) * (B + 2) ^ 2

end NLogN.KFUU

end

section
-- Source module: Solutions.KFUUSet4
/-!
# KFUU §8d: building the shape lists

After `readD`, `mkL` (register code) builds the lists that the passes read:
`rU2 = 2^h`; `pSK = [2, s_0, …, s_{d-1}]` (`(zl SK).reverse`); `pTK = [2, t_0, …]`
(`(zl TK).reverse`); `pTR = [t_{d-1}, …, t_0]` (`zl (shP h bs)` for the reversal);
`pSR = [2, s_{d-1}, …, s_0]` (`(zl SR).reverse`); `rCZ = 2`; `nR = [d+1, d]`; `oI = d`; `oBv = 2`.
-/

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

namespace NLogN.KFUU

open IntMul KIL AModel

def mkL : RS :=
  .seq (.push rU2 (.pow2 (.reg NLogN.KFUU.rU)))
  (.seq (NLogN.mvK xA pSK NLogN.KFUU.cL (.reg oD))
  (.seq (.push pSK (.lit 2))
  (.seq NLogN.KFUU.tLoop
  (.seq (.push pTK (.lit 2))
  (.seq (NLogN.mvK xB pTR NLogN.KFUU.cL (.reg oD))
  (.seq (.push pSR (.lit 2))
  (.seq (.push rCZ (.lit 2))
  (.seq (.push nR (.reg oD))
  (.seq (.push nR (.add (.reg oD) (.lit 1)))
  (.seq (.push oI (.reg oD)) (.push oBv (.lit 2))))))))))))

/-- The managed lists after `mkL`. -/
def mkF (F1 : Fin HS → List ℤ) (h d : ℕ) (bL : ℕ → Bool) (sF : ℕ → ℕ) : Fin HS → List ℤ :=
  Function.update (Function.update (Function.update (Function.update (Function.update
  (Function.update (Function.update (Function.update (Function.update (Function.update
  (Function.update (Function.update F1
    rU2 [((2 ^ h : ℕ) : ℤ)]) xA []) pSK (2 :: sZ sF d)) NLogN.KFUU.bT []) pTK (2 :: tZ h bL d)) xB [])
    pTR (tZ h bL d).reverse) pSR (2 :: (sZ sF d).reverse)) rCZ [2]) nR [(d : ℤ) + 1, (d : ℤ)])
    oI [(d : ℤ)]) oBv [2]

def cMk (d B : ℕ) : ℕ := 10 ^ 10 * (d + 1) * (B + 2) ^ 2

end NLogN.KFUU

end

section
-- Source module: Solutions.KFUURev9
/-!
# KFUU §2i: phase 2 in IL

`p2U` computes `vol ctB · cr` (from `vol (shP u bs)`, `U^n` and `cz = 2 cr`) and `⌊n/2⌋`, then
runs the `⌊n/2⌋` swaps. The main fields are reversed: the layout `lay1 bg n 0` becomes
`rv n (n/2) (lay1 bg n 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 KFB KFC KFT KFU AModel AModelL12

section Gen

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

/-- Phase 2. -/
def p2U (pM : Fin P) (pS : Fin HS) (eSw : Fin P) (dS Dsc : Fin HS) : IStmt HS P :=
  .seq (volU pM pS) (.seq powL (.seq (divE NLogN.KFUU.rB NLogN.KFUU.rM tq) (.seq (mulC pM rCZ NLogN.KFUU.rB tq)
    (.seq (.op (.drop NLogN.KFUU.rB)) (.seq (.op (.mv tq NLogN.KFUU.rB)) (.seq (.op (.const tq 2)) (.seq (divE NLogN.KFUU.rB tq tq2)
    (.seq (.op (.dup nR NLogN.KFUU.rH)) (.seq (divE NLogN.KFUU.rH tq tq2) (.seq (.op (.drop tq)) (.seq (.op (.const rA 1))
    (.seq (p2L eSw dS Dsc) (.seq (.op (.drop rA)) (.seq (.op (.drop NLogN.KFUU.rB)) (.seq (.op (.drop NLogN.KFUU.rM))
      (.op (.drop NLogN.KFUU.rH)))))))))))))))))

end Gen

end NLogN.KFUU

end

section
-- Source module: Solutions.KFUURev10
/-!
# KFUU §2j: the layouts of phase 3

With `M = (range n).reverse`, phase 3 treats the axes `M[i] = n-1-i` in order:
`lay3 bg n i = layA (M.take i) ++ mains (M.drop i) ++ bitsF (M.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 KFB KFC KFT KFU AModel AModelL12

def lay3 (bg : ℕ → Bool) (n i : ℕ) : List ℕ :=
  layA bg ((List.range n).reverse.take i) ++ mains ((List.range n).reverse.drop i) ++
    bitsF bg ((List.range n).reverse.drop i) ++ [2 * n]

/-! ## Volumes -/

end NLogN.KFUU

end

section
-- Source module: Solutions.KFUURev11
/-!
# KFUU §2k: one step of phase 3 in IL

Axis `j = n-1-i`: `rM := U^j`, `qX := U·G` (`G = vol (S.drop (j+1))` on `pXr`), `qB := U^j`,
`qZ := cz·P/U^j` (`P = vol (S.take j)` on `pV`), then the bit moves back (`mvBU`) when `S_j > U`,
and `G := S_j·G`.
-/

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

section Gen

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

/-- `pXr := S · pXr` (`S` on `tS`). -/
def accG (pM : Fin P) : IStmt HS P :=
  .seq (mulC pM pXr tS tq) (.seq (.op (.drop pXr)) (.op (.mv tq pXr)))

/-- One step of phase 3. -/
def p3B (pM : Fin P) : IStmt HS P :=
  .seq (divE NLogN.KFUU.rM rU2 tq) (.seq (mulC pM pXr rU2 NLogN.KFUU.qX) (.seq (.op (.dup NLogN.KFUU.rM NLogN.KFUU.qB))
    (.seq (mulC pM pV rCZ NLogN.KFUU.qZ) (.seq (divE NLogN.KFUU.qZ NLogN.KFUU.rM tq) (.seq bitT (.seq (.ite tq NLogN.KFUU.isPos mvBU .skip)
    (.seq (.op (.drop tq)) (.seq (accG pM) (.seq (.op (.drop NLogN.KFUU.qX)) (.seq (.op (.drop NLogN.KFUU.qB))
    (.seq (.op (.drop NLogN.KFUU.qZ)) (.op (.drop pV)))))))))))))

end Gen

end NLogN.KFUU

end

section
-- Source module: Solutions.KFUURev12
/-!
# KFUU §2l: phase 3 in IL (the loop)
-/

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

section Gen

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

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

/-- Phase 3. -/
def p3U (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 pXr 1))
    (.seq powL (.seq (p3L pM pS) (.seq (.op (.drop pXr)) (.op (.drop NLogN.KFUU.rM))))))))

end Gen

end NLogN.KFUU

end

section
-- Source module: Solutions.KFUURev13
/-!
# KFUU §2m: the uniform reversal `revU`

`revU` reverses the `n = |bs|` outer axes of an array of shape `shP u bs ++ [cz]`
(`shP u bs = [2^{u+b_0}, …]`, `cz = 2 cr` the fixed inner block), with every size read from
registers (`2^u`, `u`, `w`, `cz`, `n`) and the size list `zl (shP u bs)` on `pS`. Phase 2 calls
the `SwapNeed` entry (`SwOKU`).
-/

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

section Gen

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

/-- The uniform reversal of the outer axes. -/
def revU (pM : Fin P) (pS : Fin HS) (eSw : Fin P) (dS Dsc : Fin HS) : IStmt HS P :=
  .seq (p1U pM pS) (.seq (p2U pM pS eSw dS Dsc) (p3U pM pS))

/-- The cost of `revU`. -/
noncomputable def cRevU (K : ℝ) (u w cz cr n Z : ℕ) (V : ℕ) : ℕ :=
  (n + 1) * (⌈cSwV K w u (V * cr)⌉₊ + 2 * KFU.cMv (V * cz) w + 160 * cU Z) + 30 * cU Z

end Gen

end NLogN.KFUU

end

section
-- Source module: Solutions.KFUUGlue
/-!
# KFUU §9a: the passes on managed states

The pass lemmas (`padAllU_runs`, `unpadAllU_runs`, `revU_runs`, `rotsFU_runs`, `rotsBU_runs`)
restated on managed states `encK s F` (data `xs ++ l` on stack `0`).
-/

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

namespace NLogN.KFUU

open IntMul KIL KFB KFC KFT KFU AModel AModelL12

section Gen

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

end Gen

/-! ## The shape lists -/

/-- The bits `bL 0, …, bL (d-1)`. -/
def bsL (d : ℕ) (bL : ℕ → Bool) : List Bool := (List.range d).map bL

end NLogN.KFUU

end

section
-- Source module: Solutions.KFUUClean
/-!
# KFUU §8e: clearing registers and lists

* `dropR X C n`: pop `n` (an expression) entries of `X`;
* `popsL L`: pop one entry of every register of `L`;
* `clrF F L`: `F` with the stacks of `L` emptied.
-/

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

namespace NLogN.KFUU

open IntMul KIL AModel

def dropR (X C : Fin HS) (n : Ex HS) : RS := .rep C n (.pop X)

def popsL : List (Fin HS) → RS
  | [] => .skip
  | i :: L => .seq (.pop i) (popsL L)

def clrF (F : Fin HS → List ℤ) (L : List (Fin HS)) : Fin HS → List ℤ :=
  fun j => if j ∈ L then [] else F j

/-- `oD + 1`. -/
def nD1 : Ex HS := .add (.reg oD) (.lit 1)

/-- The registers holding one value at the end of `Φ`. -/
def Lp : List (Fin HS) := [rW, NLogN.KFUU.rU, oD, oPi, rU2, rCZ, nR, oI, oBv]

/-- The forward cleanup. -/
def cleanF : RS :=
  .seq (dropR pSR NLogN.KFUU.cL nD1) (.seq (dropR pSK NLogN.KFUU.cL nD1) (.seq (dropR pTK NLogN.KFUU.cL nD1) (.seq (dropR pTR NLogN.KFUU.cL nD1)
    (.seq (.pop nR) (NLogN.KFUU.popsL NLogN.KFUU.Lp)))))

def cClean (d B : ℕ) : ℕ := 10 ^ 6 * (d + 2) * (B + 2)

end NLogN.KFUU

end

section
-- Source module: Solutions.KFUUTrP
/-!
# KFUU §11b: the uniform axis rotation `transU`

`readT` reads `descTr w h bs = [w, h, k, b_0, …]` from `dP` (restored); `mkT` builds `2^h`,
the sizes `t_j = 2^{h+b_j}` (on `pTK` and `xB`), `nR = k`, `rCZ = 2`; the first `revU`
reverses all `k` axes; `midT` sets `rCZ = 2 t_0`, the sizes `t_{k-1}, …, t_1` on `pTR`,
`nR = k - 1`; the second `revU` reverses the first `k - 1` axes (inner block `2 t_0`);
`cleanT` clears. **`transU_runs`**: the result is `gflat (transposeL t_0 (V / t_0) zs)`.
-/

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

namespace NLogN.KFUU

open IntMul KIL KFB KFC KFT KFU AModel AModelL12

instance decLS_T (i : Fin HS) : Decidable (LS i) := by unfold LS; infer_instance

/-- The bits as a function. -/
def bLs (bs : List Bool) : ℕ → Bool := fun j => bs.getD j false

/-- `3 + k`. -/
def n3k : Ex HS := .add (.lit 3) (.reg oD)

def readT : RS :=
  .seq setA (.seq (mvL2 NLogN.KFUU.cL (.reg oD) NLogN.KFUU.bT NLogN.KFUU.hD dP) (NLogN.mvK NLogN.KFUU.hD dP NLogN.KFUU.cL n3k))

/-- The managed lists after `readT`. -/
def rdT (F0 : Fin HS → List ℤ) (w h : ℕ) (bs : List Bool) : Fin HS → List ℤ :=
  Function.update (Function.update (Function.update (Function.update F0
    rW [(w : ℤ)]) NLogN.KFUU.rU [(h : ℤ)]) oD [(bs.length : ℤ)]) NLogN.KFUU.bT (KFM.bitsZ bs.length (bLs bs)).reverse

/-- `2^h`, the sizes, `nR = k`, `rCZ = 2`. -/
def mkT : RS :=
  .seq (.push rU2 (.pow2 (.reg NLogN.KFUU.rU))) (.seq NLogN.KFUU.tLoop (.seq (.push nR (.reg oD)) (.push rCZ (.lit 2))))

def mkTF (F1 : Fin HS → List ℤ) (h k : ℕ) (bL : ℕ → Bool) : Fin HS → List ℤ :=
  Function.update (Function.update (Function.update (Function.update (Function.update
  (Function.update F1 rU2 [((2 ^ h : ℕ) : ℤ)]) NLogN.KFUU.bT []) pTK (tZ h bL k)) xB (tZ h bL k))
    nR [(k : ℤ)]) rCZ [2]

/-- `k - 1`. -/
def nK1 : Ex HS := .add (.reg oD) (.lit (-1))

/-- Between the reversals. -/
def midT : RS :=
  .seq (.asg rCZ (.add (.reg xB) (.reg xB))) (.seq (.pop xB)
    (.seq (NLogN.mvK xB pTR NLogN.KFUU.cL nK1) (.asg nR nK1)))

/-- The clearing. -/
def cleanT : RS :=
  .seq (dropR pTK NLogN.KFUU.cL (.reg oD)) (.seq (dropR pTR NLogN.KFUU.cL (.reg nR)) (NLogN.KFUU.popsL [rW, NLogN.KFUU.rU, oD, rU2, rCZ, nR]))

section Gen

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

/-- The uniform axis rotation. -/
def transU (pM eSw : Fin P) (sdS Dsc : Fin HS) : IStmt HS P :=
  .seq (RS.seq readT mkT).code
  (.seq (revU pM pTK eSw sdS Dsc)
  (.seq midT.code
  (.seq (revU pM pTR eSw sdS Dsc) cleanT.code)))

/-- The cost of `transU`. -/
noncomputable def cTr (K : ℝ) (w h : ℕ) (bs : List Bool) (B Z : ℕ) : ℕ :=
  10 ^ 9 * (bs.length + 1) * (B + 2) ^ 2 + 10 ^ 10 * (bs.length + 1) * (B + 2) ^ 2 +
    cRevU K h w 2 1 bs.length Z (vol (shP h bs)) +
    10 ^ 6 * (bs.length + 1) * (B + 2) ^ 2 +
    cRevU K h w (2 * 2 ^ (h + if bs.headD false then 1 else 0)) (2 ^ (h + if bs.headD false then 1 else 0))
      (bs.length - 1) Z (vol (shP h bs.tail.reverse)) +
    10 ^ 6 * (bs.length + 2) * (B + 2)

end Gen

end NLogN.KFUU

end

section
-- Source module: Solutions.KFUUPhiF
/-!
# KFUU §9b: the uniform `Φ`

`phiFU`: read the descriptor (`readD`), build the shape lists (`mkL`), pad `S_k → T_k`, reverse
(`revU`, the `SwapNeed` entry), unpad `T_r → S_r`, the `d` forward shifts (`rotsFU`), pad
`S_r → T_r`, clear. **`phiFU_runs`**: on the managed state of a `PhiNeed` call it computes
`KFU.PhiD`.
-/

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

namespace NLogN.KFUU

open IntMul KIL KFB KFC KFT KFU AModel AModelL12

instance decLS_F (i : Fin HS) : Decidable (LS i) := by unfold LS; infer_instance

section Gen

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

/-- The uniform `Φ`. -/
def phiFU (pM eSw : Fin P) (sdS Dsc : Fin HS) : IStmt HS P :=
  .seq (RS.seq (readD oSP oML) mkL).code
  (.seq (padAllU pM pSK pTK)
  (.seq (RS.pop nR).code
  (.seq (revU pM pTR eSw sdS Dsc)
  (.seq (RS.seq (.push pTR (.lit 2)) (.push nR nD1)).code
  (.seq (unpadAllU pM pSR pTR)
  (.seq rotsFU
  (.seq (padAllU pM pSR pTR) cleanF.code)))))))

/-- The cost of `phiFU`. -/
noncomputable def cFwd (K : ℝ) (w h d : ℕ) (bL : ℕ → Bool) (sF : ℕ → ℕ) (B Z Br Bb : ℕ) : ℕ :=
  cRead d B + cMk d B + 3 * ((d + 2) * (2000 * cScan (vol (NLogN.KFU.shR d (NLogN.KFU.tP d h bL))) w + 25 * cU Z)) +
    (2 * B + 2) + cRevU K h w 2 1 d Z (vol (shP h (bsL d bL).reverse)) + 6 * 10 ^ 7 * (B + 2) ^ 2 +
    (d + 1) * (cStepU d (vol ((List.range (d - 1)).map sF)) (vol (NLogN.KFU.shR d (NLogN.KFU.sP d sF))) w Br Bb + 3000) +
    cClean d B

end Gen

end NLogN.KFUU

end

section
-- Source module: Solutions.KFUUPhiB
/-!
# KFUU §9c: the uniform `Φ⁻¹`

`phiBU`: read the descriptor (`readD` onto `xS`, `xM`), build the lists (`mkL`, the sizes and
multipliers in forward order on `oSL`, `oML`), unpad `T_r → S_r`, the `d` backward shifts
(`rotsBU`), pad `S_r → T_r`, reverse back (`revU`), unpad `T_k → S_k`, clear.
**`phiBU_runs`**: it computes `KFU.PhiInvD`.
-/

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

namespace NLogN.KFUU

open IntMul KIL KFB KFC KFT KFU AModel AModelL12

/-- The extra setup of `Φ⁻¹`. -/
def setB : RS :=
  .seq (NLogN.mvK NLogN.KFUU.xS oSL NLogN.KFUU.cL (.reg oD)) (.seq (NLogN.mvK xM oML NLogN.KFUU.cL (.reg oD))
    (.seq (.push pTR (.lit 2)) (.seq (.asg oBv (.add (.reg oPi) (.reg oPi)))
      (.seq (.asg oPi (.lit 1)) (.asg oI (.lit 0))))))

section Gen

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

/-- The uniform `Φ⁻¹`. -/
def phiBU (pM eSw : Fin P) (sdS Dsc : Fin HS) : IStmt HS P :=
  .seq (RS.seq (readD NLogN.KFUU.xS xM) (RS.seq mkL setB)).code
  (.seq (unpadAllU pM pSR pTR)
  (.seq rotsBU
  (.seq (padAllU pM pSR pTR)
  (.seq (RS.seq (.pop nR) (.pop pTK)).code
  (.seq (revU pM pTK eSw sdS Dsc)
  (.seq (RS.seq (.push pTK (.lit 2)) (.push nR nD1)).code
  (.seq (unpadAllU pM pSK pTK)
    (RS.seq (dropR oSP NLogN.KFUU.cL (.reg oD)) cleanF).code)))))))

/-- The cost of `phiBU`. -/
noncomputable def cBwd (K : ℝ) (w h d : ℕ) (bL : ℕ → Bool) (sF : ℕ → ℕ) (B Z Br Bb : ℕ) : ℕ :=
  cRead d B + cMk d B + 10 ^ 8 * (d + 1) * (B + 2) ^ 2 +
    3 * ((d + 2) * (2000 * cScan (vol (NLogN.KFU.shR d (NLogN.KFU.tP d h bL))) w + 25 * cU Z)) +
    (d + 1) * (cStepU d (vol ((List.range (d - 1)).map sF)) (vol (NLogN.KFU.shR d (NLogN.KFU.sP d sF))) w Br Bb + 3000) +
    (4 * B + 4) + cRevU K h w 2 1 d Z (vol (shP h (bsL d bL))) + 6 * 10 ^ 7 * (B + 2) ^ 2 +
    10 ^ 5 * (d + 1) * (B + 2) + cClean d B

end Gen

end NLogN.KFUU

end

section
-- Source module: Solutions.KFUUAbs
/-!
# KFUU §9d: absorbing the costs of `Φ` into `cPhiU`

With `L = size V_T` (`V_T ≥ 2^{d+1}`, so `d + 2 ≤ L` and `L^4 ≤ 64 V_T`), every bit-size
quantity of the setup, the table and the multiplications is polylogarithmic in `V_T`, or is
covered by the descriptor term `(d+1)(h + w + size V_S + 1)^2`.

* `pow4_le`: `x^4 ≤ 32 · 2^x`; `fmB_le`: `fmB n ≤ 1280 K₀ (n+1)^2`; `cU_le`.
* `natPart_le`: all `ℕ` costs `≤ C₁ ((d+1) cScan V_T w + D + d·cRb)`.
-/

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

namespace NLogN.KFUU

open IntMul KIL KFB KFC KFT KFU AModel AModelL12

/-- The constant of `fmB`. -/
def K0 : ℕ := 8 * fmA + 10 ^ 7 + 10 ^ 5

/-- The constant of `cU`. -/
def CU : ℕ := 10 ^ 4 + 5120 * NLogN.KFUU.K0

/-! ## The sizes -/

/-! ## The ℕ part -/

/-- The `ℕ` cost pieces: setup (`B`), passes (`Z`), table (`Br`, `Bb`) and the `ℕ` part of a
reversal. -/
def natCost (w d VT VS PM B Z Br Bb : ℕ) : ℕ :=
  3 * 10 ^ 10 * (d + 2) * (B + 2) ^ 2 + 3 * ((d + 2) * (2000 * cScan VT w + 25 * cU Z)) +
    (d + 1) * (cStepU d PM VS w Br Bb + 3000) +
    ((d + 1) * (1 + 2 * KFU.cMv VT w + 160 * cU Z) + 30 * cU Z) + 2

/-- The constant `C₁`. -/
def C1 : ℕ := 10 ^ 16 * NLogN.KFUU.CU

end NLogN.KFUU

end

section
-- Source module: Solutions.KFUUNeed
/-!
# KFUU §10: `SwapNeed → PhiNeed`

The table is KFM's merged table `mTable ext₁ ∅ ∅ own` (`KFM.mExt`): the `SwapNeed` extensions
relabelled (`IRuns.relabel` via `KFM.mTable_g1`), then `own = [phiFU, phiBU]`. The swap entry
of the reversals is the relabelled `SwapNeed` entry. Costs: `cFwd`, `cBwd ≤ natCost + (d+1)·cSwV`
and `natCost_le`, `swq_le` absorb everything into `(C₁ + 781 K) · cPhiU`.

**`phiNeed_of_swapNeed : KFM.SwapNeed → KFM.PhiNeed`**.
-/

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

namespace NLogN.KFUU

open IntMul KIL KFB KFC KFT KFU AModel AModelL12

/-! ## The ℕ / ℝ cost comparison -/

/-! ## The theorem -/

/-- The empty extension list. -/
def ext0 : Fin 0 → IStmt HS (7 + 0) := fun k => k.elim0

/-- The reversal scratch stack (`18`, or `19` if the swap descriptor sits on `18`). -/
def dscOf (sdS : Fin HS) : Fin HS := if sdS = 18 then 19 else 18

/-- Our two procedures, in the merged table of size `7 + (E₁ + 0 + 0 + 2)`. -/
noncomputable def ownPhi (E₁ : ℕ) (entry : Fin (7 + E₁)) (sdS : Fin HS) :
    Fin 2 → IStmt HS (7 + (E₁ + 0 + 0 + 2)) := fun k =>
  if k.val = 0 then
    phiFU (KFM.pMul _) (KFM.gI E₁ (E₁ + 0 + 0 + 2) 0 (by omega) entry) sdS (dscOf sdS)
  else phiBU (KFM.pMul _) (KFM.gI E₁ (E₁ + 0 + 0 + 2) 0 (by omega) entry) sdS (dscOf sdS)

end NLogN.KFUU

end

section
-- Source module: Solutions.KFUUTrNeed
/-!
# KFUU §11c: `SwapNeed → TransNeed`, and `SwapNeed → PhiNeed ∧ TransNeed`

* `trNat`: the `ℕ` cost of `transU` (with `B = w + h + k + 5`, `Z = zsz (2V) + w + h + 3`) is
  `≤ C₂ (N + D) + (2k+1) ⌈cSwV K w h V⌉` with `N = (k+1)(w+1)(V+1)`, `D = (k+1)(h+w+k+1)^2`.
* `trAbsorb`: so the cost is `≤ (C₂ + 2 + 1536 K) ((k+1)(w+1)(V+1)(h^τ+1) + D)`.
* **`transNeed_of_swapNeed`**, **`phiTrans_of_swapNeed`**.
-/

set_option linter.unusedSimpArgs false
set_option linter.unusedVariables false

namespace NLogN.KFUU

open IntMul KIL KFB KFC KFT KFU AModel AModelL12

/-- The constant `C₂`. -/
def C2 : ℕ := 10 ^ 13 + 10 ^ 5 * NLogN.KFUU.CU

/-! ## The theorem -/

/-- Our procedure, in the merged table of size `7 + (E₁ + 0 + 0 + 1)`. -/
noncomputable def ownTr (E₁ : ℕ) (entry : Fin (7 + E₁)) (sdS : Fin HS) :
    Fin 1 → IStmt HS (7 + (E₁ + 0 + 0 + 1)) := fun _ =>
  transU (KFM.pMul _) (KFM.gI E₁ (E₁ + 0 + 0 + 1) 0 (by omega) entry) sdS (dscOf sdS)

end NLogN.KFUU

end

section
-- Source module: Solutions.KFV6Math
/-!
# KFV §10: the arithmetic of one Kronecker record product

The pointwise step of Proposition 18 replaces records `x, z` by `Q_p(x z / r)` (`QQ ∘ PP` of
`GlueProp18`). On coefficient vectors:

* `ncv u v`: the negacyclic convolution (`co (X Z) = ncv (co X) (co Z)`, `KFV.co_mul`),
  `ppq r p u v = Q_p(ncv u v / r)`; **`co_ppq`**: `co (Q_p(r⁻¹ (X Z))) = ppq r p (co X) (co Z)`.
* `ncv_gauss`: with `u = s (a + i b)`, `v = s (c + i e)`, the Gaussian numerators are the folded
  integer convolutions `(H^{ac} - H^{be}) + i (H^{ae} + H^{bc})` (as `co_mul_gauss`).
* `kmul r p ρ zx zy`: the integer output record, `((hRe_j - hRe_{j+r}).tdiv 2^{p+ρ}, …)` with
  `hRe = cnv(re, re) - cnv(im, im)`, `hIm = cnv(re, im) + cnv(im, re)`; **`toC_kmul`**: for
  `r = 2^ρ` it is the numerator record (scale `2^p`) of `ppq r p u v`.
* **`extract_digits`**: for `Z = ∑_{j<J} h_j B^j` with `B = 2^b`, `h_j ∈ [-B/2, B/2)`, the
  shifted value `zfix B J Z = Z + [Z < 0] B^J` lies in `[0, B^J)`, and the centered carry pass
  over its base-`B` digits (`qd`) recovers every `h_j` (Lemma 17's extraction, `Synth.extract`).
* `rho_succ_eq`: the carry step is `ρ_{j+1} = (q_j + ρ_j + H) / (2H)` (one ruler `quo`).
-/

set_option linter.unusedSectionVars false
set_option linter.unusedSimpArgs false

namespace NLogN.KFV

open Finset Butterfly Synth GlueError GlueProp18 KFB

/-! ## 1. Coefficients of the normalized truncated product -/

section NC

variable {r : ℕ}

/-- The negacyclic convolution of two coefficient vectors. -/
def ncv (u v : Fin r → ℂ) (j : Fin r) : ℂ :=
  ∑ i : Fin r, u i * (if h : (i : ℕ) ≤ j then v ⟨j - i, by omega⟩
    else - v ⟨j + r - i, by have := j.isLt; omega⟩)

/-- `Q_p (u v / r)` on coefficient vectors. -/
noncomputable def ppq (r p : ℕ) (u v : Fin r → ℂ) : Fin r → ℂ :=
  fun j => Qp p ((r : ℂ)⁻¹ * ncv u v j)

end NC

/-! ## 2. The integer output record -/

section Rec

/-- Real parts of a record. -/
def reF (zs : List GI) (i : ℕ) : ℤ := (zs.getD i 0).1

/-- Imaginary parts of a record. -/
def imF (zs : List GI) (i : ℕ) : ℤ := (zs.getD i 0).2

/-- Unfolded real numerators `cnv(a, c) - cnv(b, e)`. -/
def hRe (r : ℕ) (zx zy : List GI) (j : ℕ) : ℤ :=
  cnv r (reF zx) (reF zy) j - cnv r (imF zx) (imF zy) j

/-- Unfolded imaginary numerators `cnv(a, e) + cnv(b, c)`. -/
def hIm (r : ℕ) (zx zy : List GI) (j : ℕ) : ℤ :=
  cnv r (reF zx) (imF zy) j + cnv r (imF zx) (reF zy) j

/-- The output record: fold modulo `y^r + 1`, then truncate by `2^{p+ρ}`. -/
def kmul (r p ρ : ℕ) (zx zy : List GI) : List GI :=
  (List.range r).map fun j => ((hRe r zx zy j - hRe r zx zy (j + r)).tdiv (2 ^ (p + ρ)),
    (hIm r zx zy j - hIm r zx zy (j + r)).tdiv (2 ^ (p + ρ)))

end Rec

/-! ## 3. Extraction from the base-`2^b` digits -/

section Ext

/-- `Z + [Z < 0] B^J`. -/
def zfix (B : ℤ) (J : ℕ) (Z : ℤ) : ℤ := if Z < 0 then Z + B ^ J else Z

/-- Digit `j` of `x` in base `2^b`. -/
def qd (b x : ℕ) (j : ℕ) : ℤ := ((x / 2 ^ (b * j) % 2 ^ b : ℕ) : ℤ)

end Ext

end NLogN.KFV

end

section
-- Source module: Solutions.KFV6Loop
/-!
# KFV §10: the four register loops of one Kronecker record product

All on managed states `encK s f`, with the loop counter `vL`:

* **`unzP_fsp`**: split a record `gflat zs` (from stack `X`) into the offset real parts
  `(z.1 + c, 0)` (on `SA`) and offset imaginary parts `(z.2 + c, 0)` (on `SB`), each reversed, ready
  for KFT's packer `carI` after one `mvK`.
* **`carL_fsp`**: the centered carry pass over base-`2^b` digits `(q_j, 0)` on stack `0`:
  pushes `cent H (q_j + ρ_j)` on `vT` (`H = 2^{b-1}`), the carry `ρ` on `vRH` (two ruler ops).
* **`fldP_fsp`**: `(x - y).tdiv 2^Δ` of two equal-length stacks (the fold `h_j - h_{j+r}` and the
  truncation of Lemma 17 / `Q_p`).
* **`ilvP_fsp`**: interleave two stacks into Gaussian pairs.

Every loop step costs `O(w)` for `w`-bit values: the cost is linear in the record length.
-/

set_option linter.unusedSimpArgs false
set_option linter.unusedVariables false
set_option linter.unusedSectionVars false
set_option maxRecDepth 200000

namespace NLogN.KFV

open IntMul RS KIL KFD KFB Synth

/-! ## Stacks of the product phase (all below `26`, outside the KFN scratch) -/

abbrev vX : Fin HS := 1
abbrev vY : Fin HS := 2
abbrev vO : Fin HS := 3
abbrev vA : Fin HS := 4
abbrev vB : Fin HS := 5
abbrev vC : Fin HS := 6
abbrev vE : Fin HS := 7
abbrev vPA : Fin HS := 8
abbrev vPB : Fin HS := 9
abbrev vPC : Fin HS := 10
abbrev vPE : Fin HS := 11
abbrev vM1 : Fin HS := 12
abbrev vM2 : Fin HS := 13
abbrev vM3 : Fin HS := 14
abbrev vM4 : Fin HS := 15
abbrev vZ : Fin HS := 16
abbrev vRH : Fin HS := 17
abbrev vT : Fin HS := 18
abbrev vT2 : Fin HS := 19
abbrev vT4 : Fin HS := 20
abbrev vOR : Fin HS := 21
abbrev vOI : Fin HS := 22
abbrev vL : Fin HS := 23

/-! ## 1. Splitting a record into offset real and imaginary parts -/

/-- `(z.1 + c, 0)`. -/
def offG (c : ℤ) (z : GI) : GI := (z.1 + c, 0)

def unzB (X SA SB : Fin HS) (c : ℤ) : RS :=
  .seq (.push SA (.add (.reg X) (.lit c))) (.seq (.push SA (.lit 0)) (.seq (.pop X)
    (.seq (.push SB (.add (.reg X) (.lit c))) (.seq (.push SB (.lit 0)) (.pop X)))))

def unzP (X SA SB C : Fin HS) (c : ℤ) (n : ℕ) : RS := .rep C (.lit n) (unzB X SA SB c)

/-! ## 2. The centered carry pass -/

def xHE (H : ℤ) : Ex HS := .add (.add (.reg 0) (.reg vRH)) (.lit H)

def carB (H B : ℤ) : RS :=
  .seq (.push NLogN.KFV.vT (.add (.rem (xHE H) (.lit B)) (.lit (-H))))
    (.seq (.asg vRH (.quo (xHE H) (.lit B))) (.seq (.pop 0) (.pop 0)))

def carL (H B : ℤ) (J : ℕ) : RS :=
  .seq (.push vRH (.lit 0)) (.seq (.rep NLogN.KFV.vL (.lit J) (carB H B)) (.pop vRH))

/-- The extracted (centered) digit `j`. -/
def cdig (H : ℤ) (q : ℕ → ℤ) (j : ℕ) : ℤ := cent H (q j + rho H q j)

/-- The digit list as Gaussian pairs `(q, 0)`. -/
def dpair (qs : List ℤ) : List GI := qs.map fun x => (x, 0)

/-! ## 3. Fold and truncate: `(x - y).tdiv 2^Δ` of two stacks -/

def fldB (O : Fin HS) (Δ : ℕ) : RS :=
  .seq (.push O (.quo (.add (.reg NLogN.KFV.vT) (.neg (.reg vT4))) (.lit (2 ^ Δ)))) (.seq (.pop NLogN.KFV.vT) (.pop vT4))

def fldP (O : Fin HS) (Δ n : ℕ) : RS := .rep NLogN.KFV.vL (.lit n) (fldB O Δ)

/-- The fold-and-truncate of one pair. -/
def fdt (Δ : ℕ) (x y : ℤ) : ℤ := (x - y).tdiv (2 ^ Δ)

/-! ## 4. Interleaving two stacks into Gaussian pairs -/

def ilvB : RS := .seq (.push vO (.reg vOR)) (.seq (.push vO (.reg vOI)) (.seq (.pop vOR) (.pop vOI)))

def ilvP (n : ℕ) : RS := .rep NLogN.KFV.vL (.lit n) ilvB

end NLogN.KFV

end

section
-- Source module: Solutions.KFV6Pack
/-!
# KFV §10: packing, the four ring products, and the extraction of one record product

* `carE`, `digE`, `mvE`, `mulE`: KFT's packer `carI` and digit cutter `digI`, an `mv`, and one
  `MulOK` product (`KFM.mulC`) on managed states `encK s f`.
* **`packI_runs`**: `mvK` a reversed offset record onto stack `0`, `carI`, subtract the offset
  constant: the signed Kronecker integer `∑_h a_h 2^{b h}` on stack `P` (Lemma 17's packing).
* **`extI_runs`**: from two products `m₁, m₂` with `m₁ ± m₂ = ∑_{j<2r} h_j 2^{b j}`
  (`h_j ∈ [-2^{b-1}, 2^{b-1})`): `zfix`, `digI`, the centered carry pass `carL`, two `mvK`, and the
  fold-and-truncate `fldP`; pushes `((h_j - h_{j+r}).tdiv 2^Δ)_{j<r}` on stack `O`
  (`KFV6Math.extract_digits`).
-/

set_option linter.unusedSimpArgs false
set_option linter.unusedVariables false
set_option linter.unusedSectionVars false
set_option maxRecDepth 10000

namespace NLogN.KFV

open IntMul RS KIL KFD KFB KFT Synth

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

/-! ## 1. IState-level pieces on managed states -/

/-! ## 2. Packing one offset part -/

/-- The packed value of an offset part, minus the offset constant. -/
def packI (S P : Fin HS) (b r : ℕ) (KG : ℤ) : IStmt HS Pp :=
  .seq (fcode (NLogN.mvK S 0 NLogN.KFV.vL (.lit ((2 * r : ℕ) : ℤ))))
    (.seq (carI b r) (fcode (.seq (.push P (.add (.reg 0) (.lit (-KG)))) (.pop 0))))

def cPack (r w b Bp : ℕ) : ℕ :=
  cLp (2 * r) (250 * (w + 2)) Bp + cCar b w r + (8000 * (Bp + 2) + (2 * Bp + 2))

/-! ## 3. Extraction of one Gaussian component -/

/-- `Z + [Z < 0] B^J` as an expression. -/
def zfE (e : Ex HS) (BJ : ℤ) : Ex HS := .sel e (.add e (.lit BJ)) e

/-- Extract one component: `zfix`, digits, centered carries, fold, truncate. -/
def extI (e : Ex HS) (O : Fin HS) (b r Δ : ℕ) : IStmt HS Pp :=
  .seq (fcode (.push vZ (zfE e ((2 ^ b) ^ (2 * r))))) <|
  .seq (.op (.mv vZ tX)) <|
  .seq (digI b 0 (2 * r) (2 * r)) <|
  .seq (fcode (carL (2 ^ (b - 1)) (2 ^ b) (2 * r))) <|
  .seq (fcode (NLogN.mvK NLogN.KFV.vT vT2 NLogN.KFV.vL (.lit (r : ℤ)))) <|
  .seq (fcode (NLogN.mvK vT2 vT4 NLogN.KFV.vL (.lit (r : ℤ)))) <|
  fcode (fldP O Δ r)

def cExt (esz r b Δ Bz Bb : ℕ) : ℕ :=
  (3 * esz + 4) * 2000 * (Bz + 2) + 110 * (Bz + 1) + cDig b 0 (2 * r) (2 * r) +
    (4000 * (b + 4) + cLp (2 * r) (40000 * (b + 4)) Bb + (2 * (b + 2) + 2)) +
    cLp r (250 * (b + 4)) Bb + cLp r (250 * (b + 4)) Bb + cLp r (20000 * (b + 2 + Δ + 3)) Bb

end NLogN.KFV

end

section
-- Source module: Solutions.KFV6Rec
/-!
# KFV §10: one Kronecker record product in IL

**`recI pM b r Δ c KG`** consumes record `x` (stack `vX`) and record `z` (stack `vY`) and pushes the
reversed output record on `vO`:

1. `unzP` twice: offset real / imaginary parts (`+ c`, `c = 2^p`);
2. `packI` four times: `A = ∑ a_h B^h`, `Bi = ∑ b_h B^h`, `C`, `E` (`B = 2^b`, Lemma 17 packing);
3. four `MulOK` products `CA, EB, EA, CB` (FM in `kTable`);
4. `extI` on `CA - EB` and on `EA + CB`: the folded, truncated numerators;
5. `ilvP`, then drop the scratch integers.

**`recI_runs`**: the result is `kmul r p ρ zx zy` (`KFV6Math`), whose `toC` is `2^p` times
`Q_p(x z / r)` when `r = 2^ρ` (`toC_kmul`). The cost is `O(r (b + p))` plus four products of
`(b r + 2)`-bit integers.
-/

set_option linter.unusedSimpArgs false
set_option linter.unusedVariables false
set_option linter.unusedSectionVars false
set_option maxRecDepth 10000

namespace NLogN.KFV

open IntMul RS KIL KFD KFB KFT Synth

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

/-- The scratch stacks of the record product are empty. -/
structure Tmp (f : Fin HS → List ℤ) : Prop where
  a : f vA = []
  b : f NLogN.KFV.vB = []
  c : f vC = []
  e : f NLogN.KFV.vE = []
  pa : f vPA = []
  pb : f vPB = []
  pc : f vPC = []
  pe : f vPE = []
  m1 : f vM1 = []
  m2 : f vM2 = []
  m3 : f vM3 = []
  m4 : f vM4 = []
  z : f vZ = []
  x : f tX = []
  rh : f vRH = []
  t : f NLogN.KFV.vT = []
  t2 : f vT2 = []
  t4 : f vT4 = []
  or : f vOR = []
  oi : f vOI = []
  l : f NLogN.KFV.vL = []

def dropsP : RS := .seq (.pop vPA) (.seq (.pop vPB) (.seq (.pop vPC) (.seq (.pop vPE)
  (.seq (.pop vM1) (.seq (.pop vM2) (.seq (.pop vM3) (.pop vM4)))))))

def subE : Ex HS := .add (.reg vM1) (.neg (.reg vM2))
def addE : Ex HS := .add (.reg vM3) (.reg vM4)

/-- The offset constant `c ∑_{h<r} B^h`. -/
def kG (b r : ℕ) (c : ℤ) : ℤ := ∑ h ∈ Finset.range r, c * (2 ^ b : ℤ) ^ h

/-- Cost of one record product (`φ` bounds one product of `(b r + 2)`-bit factors). -/
def cRec (r p b Δ φ : ℕ) : ℕ :=
  2 * cLp r (30000 * (p + 4)) (zsz ((2 * r : ℕ) : ℤ)) + 4 * cPack r (p + 2) b (b * r + 2) +
    4 * (φ + 1000 * (b * r + 3)) +
    cExt 4 r b Δ (b * (2 * r) + 2) (zsz ((2 * r : ℕ) : ℤ)) +
    cExt 3 r b Δ (b * (2 * r) + 2) (zsz ((2 * r : ℕ) : ℤ)) +
    cLp r (10000 * (b + 4)) (zsz ((2 * r : ℕ) : ℤ)) + 8 * (2 * (b * (2 * r) + 2) + 2)

end NLogN.KFV

end