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

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

import Definitions.Def_IntMul_Kappa_K_Defs32

set_option autoImplicit false
set_option relaxedAutoImplicit false
set_option maxHeartbeats 1000000

section
-- Source module: Solutions.KFX1Sw
/-!
# KFX §1: the Lemma 12 swap wrapper with expression parameters

`swE q eP eG eS eU` pushes the four Lemma 12 parameters as the values of arbitrary expressions
(read from registers at run time), runs `l12F`, pops.  **`swE_spec`** is `KFL.swI_spec` with the
literals replaced by expressions (`KFL.ESp` facts), and **`swE_geo`** gives the three slot swaps
`rixT g (swT i)` on the 8-axis geometry.
-/

set_option linter.unusedSimpArgs false
set_option linter.unusedVariables false

namespace NLogN.KFX

open IntMul RS KIL KFD KFE KFB KFL AModel AModelL12 NetProp7

/-- The parameter pushes of the swap. -/
def swPl (eP eG eS eU : Ex HS) : List (Fin HS × Ex HS) :=
  [(kP, eP), (NLogN.KIL.kG, eG), (NLogN.KIL.kS, eS), (aU, eU)]

/-- **The swap wrapper with expression parameters.** -/
def swE (q : ℕ) (eP eG eS eU : Ex HS) : IStmt HS 1 := withPushI (swPl eP eG eS eU) (l12C q)

/-- Expressions that do not read the four Lemma 12 parameter stacks. -/
def NoK (e : Ex HS) : Prop := ∀ r ∈ e.regs, r ≠ kP ∧ r ≠ NLogN.KIL.kG ∧ r ≠ NLogN.KIL.kS ∧ r ≠ aU

/-! ## The three slot swaps on the 8-axis geometry -/

/-- The prefix volume of slot swap `i`. -/
def svP (g : Geo) : Fin 3 → ℕ
  | 0 => g.N0
  | 1 => vol [g.N0, g.Q, g.N1]
  | 2 => g.N0

/-- The middle volume of slot swap `i`. -/
def svG (g : Geo) : Fin 3 → ℕ
  | 0 => g.N1
  | 1 => g.N2
  | 2 => vol [g.N1, g.Q, g.N2]

/-- The suffix volume of slot swap `i`. -/
def svS (g : Geo) : Fin 3 → ℕ
  | 0 => vol [g.N2, g.Q, g.N3, g.B]
  | 1 => vol [g.N3, g.B]
  | 2 => vol [g.N3, g.B]

end NLogN.KFX

end

section
-- Source module: Solutions.KFX2Rot
/-!
# KFX §2: register-driven table rotations (the lines (22) and the last bit)

* `pushSt_G_mem`, `pushOK_of`: pushed parameter lists built from `KFL.ESp` facts.
* `tabRot_runs`: a generic `withPushI pl (fcode U ; rotF)` where `U` pushes the rotation table.
* `tabGenU sgn b`: `KFL.tabGen` with the literals `P, Q, d0, d1, 2^ρ, f-1, 2^K` read from the
  registers `1, 2, 3, 5, 6, 7, 8` (`tabGenU_sub`: the literal instance is `RS.sub`).
* **`rotU_geo`**: one line `rixT g (rot2 Q (amt K ρ f l))` with parameters from expressions.
* `xorGenU`, **`xorU_geo`**: the last-bit step `rixT g (XLc K ρ f)`.
-/

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

namespace NLogN.KFX

open IntMul RS KIL KFD KFE KFB KFL KFW AModel Butterfly

/-! ## Pushed parameter lists -/

/-! ## A table generator followed by `rotF` -/

/-! ## The uniform table generator of a line -/

/-- `((p / d) % Q) / 2^ρ` with `d, Q, 2^ρ` from the registers `rd, 2, 6`. -/
def uEU (rd : Fin HS) : Ex HS :=
  .ndiv (.nmod (.ndiv (.add (.reg KFL.rP) (.lit (-1))) (.reg rd)) (.reg 2)) (.reg 6)

/-- `(sgn · rA) mod Q` with `Q` from register `2`. -/
def amtEU (sgn : ℤ) : Ex HS := .nmod (.add (.mul (.lit sgn) (.reg NLogN.KFL.rA)) (.reg 2)) (.reg 2)

def segBodyU (T : Ex HS) : RS :=
  .seq (.asg NLogN.KFL.rA (.add (.reg NLogN.KFL.rA) (.mul (.reg NLogN.KFL.rW) T)))
    (.seq (.asg NLogN.KFL.rU (.quo (.reg NLogN.KFL.rU) (.reg 8)))
    (.seq (.asg KFL.rV (.quo (.reg KFL.rV) (.reg 8)))
      (.asg NLogN.KFL.rW (.mul (.reg NLogN.KFL.rW) (.reg 8)))))

def segLoopU (T : Ex HS) : RS := .rep rC (.reg 7) (segBodyU T)

def tabBodyU (sgn b : ℤ) : RS :=
  .seq (.push NLogN.KFL.rU (uEU 3)) (.seq (.push KFL.rV (uEU 5)) (.seq (.push NLogN.KFL.rW (.reg 6))
  (.seq (.push NLogN.KFL.rA (.lit 0)) (.seq (segLoopU (offE b))
  (.seq (.push 14 (amtEU sgn)) (.seq (.pop NLogN.KFL.rA) (.seq (.pop NLogN.KFL.rW)
  (.seq (.pop KFL.rV) (.seq (.pop NLogN.KFL.rU) (.asg KFL.rP (.add (.reg KFL.rP) (.lit (-1)))))))))))))

/-- **The uniform table generator.** -/
def tabGenU (sgn b : ℤ) : RS :=
  .seq (.push KFL.rP (.reg 1)) (.seq (.rep rq (.reg 1) (tabBodyU sgn b)) (.pop KFL.rP))

/-- The register assignment of `tabGenU`. -/
def σT (P Q d0 d1 r2 fm k2 : ℤ) : Fin HS → Option ℤ := fun i =>
  if i = 1 then some P else if i = 2 then some Q else if i = 3 then some d0 else
  if i = 5 then some d1 else if i = 6 then some r2 else if i = 7 then some fm else
  if i = 8 then some k2 else none

/-! ## One line, register-driven -/

/-- The parameter list of a line: `P, Q, B` for `rotF` on `1, 2, 4`, and `d0, d1, 2^ρ, f-1, 2^K`
for the table on `3, 5, 6, 7, 8`. -/
def rotPl (eP eQ eB ed0 ed1 er ef eK : Ex HS) : List (Fin HS × Ex HS) :=
  [(1, eP), (2, eQ), (4, eB), (3, ed0), (5, ed1), (6, er), (7, ef), (8, eK)]

/-- **One line with register parameters.** -/
def rotU {Pp : ℕ} (sgn b : ℤ) (eP eQ eB ed0 ed1 er ef eK : Ex HS) : IStmt HS Pp :=
  withPushI (rotPl eP eQ eB ed0 ed1 er ef eK) (.seq (fcode (tabGenU sgn b)) KFC.rotF)

/-- The expressions of a parameter list read none of the low stacks `1 … 12`. -/
def NoLow (e : Ex HS) : Prop := ∀ r ∈ e.regs, 13 ≤ r.val

/-- The cost of one register-driven line. -/
def cRotU (K f : ℕ) (b : ℤ) (P Q B w Bb : ℕ) : ℕ :=
  400 * cTabGen K f b P Bb + KFC.cRotF P Q B w + 10 ^ 10 * (Bb + 2) ^ 2

/-! ## The last-bit step, register-driven -/

/-- `(((p / d) % Q) / 2^j) % 2` with `d, Q, 2^j` from the registers `3, 5, 6`. -/
def bEU : Ex HS :=
  .nmod (.ndiv (.nmod (.ndiv (.add (.reg KFL.rP) (.lit (-1))) (.reg 3)) (.reg 5)) (.reg 6)) (.lit 2)

def xorBodyU : RS := .seq (.push 14 bEU) (.asg KFL.rP (.add (.reg KFL.rP) (.lit (-1))))

/-- **The uniform bit table.** -/
def xorGenU : RS := .seq (.push KFL.rP (.reg 1)) (.seq (.rep rq (.reg 1) xorBodyU) (.pop KFL.rP))

/-- The register assignment of `xorGenU`. -/
def σX (P d Q r2 : ℤ) : Fin HS → Option ℤ := fun i =>
  if i = 1 then some P else if i = 3 then some d else if i = 5 then some Q else
  if i = 6 then some r2 else none

/-- The parameter list of the last-bit step: `P, 2, B` for `rotF`, `d, Q, 2^j` for the table. -/
def xorPl (eP eB ed eQ ej : Ex HS) : List (Fin HS × Ex HS) :=
  [(1, eP), (2, .lit 2), (4, eB), (3, ed), (5, eQ), (6, ej)]

/-- **The last-bit step with register parameters.** -/
def xorU {Pp : ℕ} (eP eB ed eQ ej : Ex HS) : IStmt HS Pp :=
  withPushI (xorPl eP eB ed eQ ej) (.seq (fcode xorGenU) KFC.rotF)

def cXorU (P B w Bb : ℕ) : ℕ := 400 * cXorGen P Bb + KFC.cRotF P 2 B w + 10 ^ 10 * (Bb + 2) ^ 2

end NLogN.KFX

end

section
-- Source module: Solutions.KFX3Corr
/-!
# KFX §3: the guard-set correction with register parameters

`corrRSU` is `KFL.corrRS K ρ f g A` with the size-dependent literals read from the registers
`1 … 10`: `recG.S0, Q, recG.S1, recG.S2, B, 2^ρ, f-1, 2^K, M = recG.V, A`.  `corrRSU_sub`: the
literal instance is `RS.sub`.  **`corrU_spec`**: `withPushI plC (fcode corrRSU)` computes
`rixT g (corrc K ρ f)` at the cost of `KFL.corrI_spec` (times `400`, plus the pushes).
-/

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

namespace NLogN.KFX

open IntMul RS KIL KFD KFE KFB KFL KFW AModel Butterfly

/-! ## The uniform program -/

def gcEU (rs : Fin HS) : Ex HS := .nmod (.ndiv (.reg NLogN.KFL.rm) (.reg rs)) (.reg 2)

def keyEU : Ex HS :=
  .add (.reg NLogN.KFL.rm) (.add (.mul (.add (.reg rY) (.neg (.reg rY0))) (.reg 3))
    (.mul (.add (.reg NLogN.KFL.rZ) (.neg (.reg rZ0))) (.reg 4)))

def segStepU (eu ev ea T : Ex HS) (tgt : Fin HS) (efin : Ex HS) : RS :=
  .seq (.push NLogN.KFL.rU eu) (.seq (.push KFL.rV ev) (.seq (.push NLogN.KFL.rW (.reg 6)) (.seq (.push NLogN.KFL.rA ea)
  (.seq (segLoopU T) (.seq (.asg tgt efin) (.seq (.pop NLogN.KFL.rA) (.seq (.pop NLogN.KFL.rW)
  (.seq (.pop KFL.rV) (.pop NLogN.KFL.rU)))))))))

def lineRU (l : Line) : RS :=
  segStepU (.ndiv (.reg NLogN.KFL.rX) (.reg 6)) (.ndiv (.reg (othR l)) (.reg 6)) (.lit 0) (offE l.b) (tgtR l)
    (.nmod (.add (.add (.reg (tgtR l)) (.reg 2)) (.neg (.mul (.lit l.sgn) (.reg NLogN.KFL.rA)))) (.reg 2))

def t0RU : RS :=
  segStepU (.ndiv (.reg NLogN.KFL.rX) (.reg 6)) (.ndiv (.reg rY) (.reg 6)) (.reg rY) t0E rY (.reg NLogN.KFL.rA)

def progLU : List Line → RS
  | [] => .skip
  | l :: ls => .seq (progLU ls) (lineRU l)

def keyProgU : RS :=
  .seq (.push rY0 (gcEU 3)) (.seq (.push rZ0 (gcEU 4)) (.seq (.push NLogN.KFL.rX (gcEU 1))
  (.seq (.push rY (gcEU 3)) (.seq (.push NLogN.KFL.rZ (gcEU 4)) (.seq (progLU lines)
  (.seq t0RU (.seq (.push NLogN.KFL.rK keyEU) (.seq (.pop NLogN.KFL.rX) (.seq (.pop rY)
  (.seq (.pop NLogN.KFL.rZ) (.seq (.pop rY0) (.pop rZ0))))))))))))

def badBU : RS :=
  .seq (NLogN.mvK 0 sKB NLogN.KFL.cMv (.reg 5)) (.seq (.push sKK (.reg NLogN.KFL.rK))
    (.seq (.push sFL (.lit 1)) (.asg rNB (.add (.reg rNB) (.lit 1)))))

def goodBU : RS := .seq (NLogN.mvK 0 NLogN.KFL.sT NLogN.KFL.cMv (.reg 5)) (.push sFL (.lit 0))

def brBU : RS := .ite e1E badBU (.ite e2E badBU goodBU)

def p1BU : RS := .seq keyProgU (.seq brBU (.seq (.pop NLogN.KFL.rK) (.asg NLogN.KFL.rm (.add (.reg NLogN.KFL.rm) (.lit 1)))))

def p1ProgU : RS := .rep KFL.c1 (.reg 9) p1BU

def bucketU (ZK ZB n nB : Fin HS) : RS :=
  .seq (.push ZK (.quo (.reg sKK) (.lit 2))) (.seq (.pop sKK) (.seq (NLogN.mvK sKB ZB NLogN.KFL.cMv (.reg 5))
    (.seq (.asg n (.add (.reg n) (.lit 1))) (.asg nB (.add (.reg nB) (.reg 5))))))

def entBU : RS := .ite condE (bucketU sZ1K sZ1B rn1 rn1B) (bucketU sZ0K sZ0B rn0 rn0B)

def entLoopU : RS := .rep c3 (.reg rNB) entBU

def passBU : RS :=
  .seq (.push rn0 (.lit 0)) (.seq (.push rn1 (.lit 0)) (.seq (.push rn0B (.lit 0))
  (.seq (.push rn1B (.lit 0)) (.seq entLoopU
  (.seq (NLogN.mvK sZ0B sKB NLogN.KFL.cMv (.reg rn0B)) (.seq (NLogN.mvK sZ0K sKK NLogN.KFL.cMv (.reg rn0))
  (.seq (NLogN.mvK sZ1B sKB NLogN.KFL.cMv (.reg rn1B)) (.seq (NLogN.mvK sZ1K sKK NLogN.KFL.cMv (.reg rn1))
  (.seq (.pop rn0) (.seq (.pop rn1) (.seq (.pop rn0B) (.pop rn1B))))))))))))

def radixProgU : RS := .rep c2 (.reg 10) passBU

def p3BU : RS :=
  .seq (.ite (.neg (.reg sFL)) (.seq (NLogN.mvK sKB 0 NLogN.KFL.cMv (.reg 5)) (.pop sKK))
    (NLogN.mvK NLogN.KFL.sT 0 NLogN.KFL.cMv (.reg 5))) (.pop sFL)

def p3ProgU : RS := .rep c4 (.reg 9) p3BU

/-- **The uniform correction.** -/
def corrRSU : RS :=
  .seq (.push NLogN.KFL.rm (.lit 0)) (.seq (.push rNB (.lit 0)) (.seq p1ProgU (.seq (.pop NLogN.KFL.rm)
    (.seq radixProgU (.seq p3ProgU (.pop rNB))))))

/-- The register assignment of `corrRSU`. -/
def σC (s0 q s1 s2 b r2 fm k2 m a : ℤ) : Fin HS → Option ℤ := fun i =>
  if i = 1 then some s0 else if i = 2 then some q else if i = 3 then some s1 else
  if i = 4 then some s2 else if i = 5 then some b else if i = 6 then some r2 else
  if i = 7 then some fm else if i = 8 then some k2 else if i = 9 then some m else
  if i = 10 then some a else none

/-- The register values of the correction on the geometry `g`. -/
def σCg (K ρ f : ℕ) (g : Geo) (A : ℕ) : Fin HS → Option ℤ :=
  σC (g.recG.S0 : ℤ) ((2 : ℤ) ^ (f * K)) (g.recG.S1 : ℤ) (g.recG.S2 : ℤ) (g.B : ℤ) ((2 : ℤ) ^ ρ)
    ((f - 1 : ℕ) : ℤ) ((2 : ℤ) ^ K) (g.recG.V : ℤ) (A : ℤ)

/-! ## The register-driven correction -/

/-- The parameter list of the correction (stacks `1 … 10`). -/
def corrPl (e1 e2 e3 e4 e5 e6 e7 e8 e9 e10 : Ex HS) : List (Fin HS × Ex HS) :=
  [(1, e1), (2, e2), (3, e3), (4, e4), (5, e5), (6, e6), (7, e7), (8, e8), (9, e9), (10, e10)]

/-- **The register-driven correction.** -/
def corrU (e1 e2 e3 e4 e5 e6 e7 e8 e9 e10 : Ex HS) : IStmt HS 1 :=
  withPushI (corrPl e1 e2 e3 e4 e5 e6 e7 e8 e9 e10) (fcode corrRSU)

/-- The cost bound of `KFL.corrI_spec`. -/
noncomputable def corrB (K f : ℕ) (g : Geo) : ℝ :=
  ((g.V / g.B : ℕ) : ℝ) * ((zsz (g.V : ℤ) : ℝ) + 1) ^ 3 + (g.V : ℝ) +
    min 1 (80 * ((f : ℝ) - 1) * 2⁻¹ ^ K) * ((g.V / g.B : ℕ) : ℝ) *
      ((zsz (g.V : ℤ) : ℝ) + 1) * ((g.B : ℝ) + (zsz (g.V : ℤ) : ℝ) + 1)

end NLogN.KFX

end

section
-- Source module: Solutions.KFX4Core
/-!
# KFX §4: the register-driven Lemma 13 core

The fixed code `l13U q` is `KFL.l13code` with every size-dependent literal read from the
*master registers* (`MReg`): `w, 2·M·R` on `1, 6` (for `toBitsU`/`fromBitsU`), the swap flags on
`11, 12`, and `Q = 2^{fK}, B = R·2(w+1), 2^ρ, gx, N0, N1, N2, N3, A = zsz V, 2^K` on
`13, 15, 16, 17, 18, 19, 46, 47, 48, 49`, `f, K` on `36, 37`.  The slot swaps `preL px py` are
realized by two optional swaps (`optA`: swap `0` or `2` by the flag on `11`, `optB`: swap `1` by
the flag on `12`).  **`l13U_runs`** is `KFL.l13_runs` for this code.
-/

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

namespace NLogN.KFX

open IntMul RS KIL KFD KFE KFB KFL KFW KFH AModel Butterfly

/-! ## Master registers -/

abbrev mW : Fin HS := 1
abbrev mNt : Fin HS := 6
abbrev mFa : Fin HS := 11
abbrev mFb : Fin HS := 12
abbrev mQ : Fin HS := 13
abbrev mB : Fin HS := 15
abbrev mR2 : Fin HS := 16
abbrev mGx : Fin HS := 17
abbrev mN0 : Fin HS := 18
abbrev mN1 : Fin HS := 19
abbrev mN2 : Fin HS := 46
abbrev mN3 : Fin HS := 47
abbrev mA : Fin HS := 48
abbrev mK2 : Fin HS := 49
abbrev mF : Fin HS := 36
abbrev mK : Fin HS := 37

/-- A register assignment of the sixteen master registers. -/
def σM (vW vNt vFa vFb vQ vB vR2 vGx vN0 vN1 vN2 vN3 vA vK2 vF vK : ℤ) : Fin HS → Option ℤ :=
  fun i =>
  if i = mW then some vW else if i = mNt then some vNt else if i = mFa then some vFa else
  if i = mFb then some vFb else if i = mQ then some vQ else if i = mB then some vB else
  if i = mR2 then some vR2 else if i = mGx then some vGx else if i = mN0 then some vN0 else
  if i = mN1 then some vN1 else if i = mN2 then some vN2 else if i = mN3 then some vN3 else
  if i = mA then some vA else if i = mK2 then some vK2 else if i = mF then some vF else
  if i = mK then some vK else none

/-- The flag of swap `A` (`px ≠ 0`). -/
def flA (px : Fin 3) : ℕ := if px = 0 then 0 else 1

/-- The flag of swap `B` (`py = 2`, or `px = 2 ∧ py = 0`). -/
def flB (px py : Fin 3) : ℕ := (if py = 2 then 1 else 0) + (if px = 2 ∧ py = 0 then 1 else 0)

/-- The geometry of the bits. -/
abbrev gB (N0 N1 N2 N3 K f R w : ℕ) : Geo := geoB N0 N1 N2 N3 (f * K) R w

/-- The middle factor of swap `A`: `1` (swap `0`) or `Q·N2` (swap `2`). -/
def gxV (N2 K f : ℕ) (px : Fin 3) : ℕ := if px = 2 then 2 ^ (f * K) * N2 else 1

/-- **The master registers** of the row addition with parameters `N0 … w, px, py`. -/
def σMR (N0 N1 N2 N3 K f ρ R w : ℕ) (px py : Fin 3) : Fin HS → Option ℤ :=
  σM (w : ℕ) ((2 * (vol (rect N0 N1 N2 N3 (f * K)) * R) : ℕ) : ℤ) (flA px : ℕ) (flB px py : ℕ)
    ((2 ^ (f * K) : ℕ) : ℤ) ((R * (2 * (w + 1)) : ℕ) : ℤ) ((2 ^ ρ : ℕ) : ℤ) (gxV N2 K f px : ℕ)
    (N0 : ℕ) (N1 : ℕ) (N2 : ℕ) (N3 : ℕ) (zsz ((NLogN.KFX.gB N0 N1 N2 N3 K f R w).V : ℤ) : ℕ)
    ((2 ^ K : ℕ) : ℤ) (f : ℕ) (K : ℕ)

/-! ## Expression specifications from the masters -/

/-! ## The expressions of the pieces -/

def eFK : Ex HS := .mul (.reg mF) (.reg mK)
def eFm1 : Ex HS := .add (.reg mF) (.lit (-1))
/-- `N3 · B` (`= S2`). -/
def eS2 : Ex HS := .mul (.reg mN3) (.reg mB)
/-- `N0 · Q · N1`. -/
def eP1 : Ex HS := .mul (.mul (.reg mN0) (.reg mQ)) (.reg mN1)
/-- `N1 · gx`. -/
def eGA : Ex HS := .mul (.reg mN1) (.reg mGx)
/-- `N2 · Q · N3 · B / gx`. -/
def eSA : Ex HS := .ndiv (.mul (.mul (.mul (.reg mN2) (.reg mQ)) (.reg mN3)) (.reg mB)) (.reg mGx)
/-- `N0 · Q · N1 · Q · N2`. -/
def eRP : Ex HS := .mul (.mul eP1 (.reg mQ)) (.reg mN2)
/-- `N2 · Q · N1`. -/
def eD0 : Ex HS := .mul (.mul (.reg mN2) (.reg mQ)) (.reg mN1)
/-- `2^K / (2 · 2^ρ) = 2^{K-ρ-1}`. -/
def eE : Ex HS := .ndiv (.reg mK2) (.mul (.lit 2) (.reg mR2))
/-- `Q / (2 · 2^{K-ρ-1}) = 2^{ρ + (f-1) K}`. -/
def eD : Ex HS := .ndiv (.reg mQ) (.mul (.lit 2) NLogN.KFX.eE)
def eXP : Ex HS := .mul eP1 NLogN.KFX.eE
/-- `2^j · S1`, `S1 = Q · S2 · N2`. -/
def eXB : Ex HS := .mul eD (.mul (.mul (.reg mQ) eS2) (.reg mN2))
def eXd : Ex HS := .mul (.reg mN1) NLogN.KFX.eE
/-- `recG.S1 = Q · N3 · N2`. -/
def eRS1 : Ex HS := .mul (.mul (.reg mQ) (.mul (.reg mN3) (.lit 1))) (.reg mN2)
def eRS0 : Ex HS := .mul (.mul (.reg mQ) eRS1) (.reg mN1)
def eRV : Ex HS := .mul (.mul (.reg mQ) eRS0) (.reg mN0)

section Vals

variable {N0 N1 N2 N3 K f ρ R w : ℕ} {px py : Fin 3} {t : IState HS}

/-- Positivity data of a parameter set. -/
structure PPos (N0 N1 N2 N3 K f ρ R : ℕ) : Prop where
  n0 : 1 ≤ N0
  n1 : 1 ≤ N1
  n2 : 1 ≤ N2
  n3 : 1 ≤ N3
  r : 1 ≤ R
  k : 1 ≤ K
  f : 1 ≤ f
  ρ : ρ < K

variable (hp : PPos N0 N1 N2 N3 K f ρ R) (hA : Agree (σMR N0 N1 N2 N3 K f ρ R w px py) t)
include hp hA

end Vals

end NLogN.KFX

/-! ## The code -/

namespace NLogN.KFX

open IntMul RS KIL KFD KFE KFB KFL KFW KFH AModel Butterfly

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

/-- Slot swap `1` (`y ↔ z`). -/
def swBU (q : ℕ) : IStmt HS 1 := swE q eP1 (.reg mN2) eS2 eFK
/-- Slot swap `0` or `2` (by `gx`). -/
def swAU (q : ℕ) : IStmt HS 1 := swE q (.reg mN0) eGA eSA eFK
def optA (q : ℕ) : IStmt HS 1 := .ite mFa NLogN.KFX.isPos (swAU q) .skip
def optB (q : ℕ) : IStmt HS 1 := .ite mFb NLogN.KFX.isPos (swBU q) .skip

def rotLU (l : Line) : IStmt HS 1 :=
  rotU l.sgn l.b eRP (.reg mQ) eS2 eD0 (.reg mN2) (.reg mR2) eFm1 (.reg mK2)

def lineU (q : ℕ) (l : Line) : IStmt HS 1 :=
  if l.tgtY then .seq (swBU q) (.seq (rotLU l) (swBU q)) else rotLU l

def linesU (q : ℕ) : List Line → IStmt HS 1
  | [] => .skip
  | l :: ls => .seq (lineU q l) (linesU q ls)

def corrLU : IStmt HS 1 :=
  corrU eRS0 (.reg mQ) eRS1 (.reg mN3) (.reg mB) (.reg mR2) eFm1 (.reg mK2) eRV (.reg mA)

def xorLU : IStmt HS 1 := xorU eXP eXB eXd (.reg mQ) eD

def coreBU (q : ℕ) : IStmt HS 1 := .seq (linesU q lines) (.seq corrLU xorLU)

/-- **The register-driven Lemma 13 code.** -/
def l13U (q : ℕ) : IStmt HS 1 :=
  .seq toBitsU (.seq (optA q) (.seq (optB q) (.seq (coreBU q) (.seq (optB q) (.seq (optA q)
    fromBitsU)))))

/-! ## Swap lists -/

def lA (px : Fin 3) : List (Fin 3) := if px = 0 then [] else if px = 1 then [0] else [2]
def lB (px py : Fin 3) : List (Fin 3) := if flB px py = 1 then [1] else []

/-- The swap property of `swE_geo` for one prime. -/
def SwEOK (q : ℕ) (Ksw : ℝ) : Prop :=
  ∀ (i : Fin 3) (g : Geo), g.Pos → ∀ u : ℕ, 1 ≤ u → g.Q = 2 ^ u →
    ∀ (I : List Sym) (t : IState HS) (zs : List ℤ) (R : List Item),
    KFL.Bits zs → t.G 0 = ints zs ++ R → zs.length = g.V →
    ∀ (eP eG eS eU : Ex HS) (Bb : ℕ), Bb ≤ zsz (g.V : ℤ) + 3 →
    eP.size + eG.size + eS.size + eU.size ≤ 50 →
    NoK eP → NoK eG → NoK eS → NoK eU →
    ESp eP t Bb (svP g i : ℤ) → ESp eG t Bb (svG g i : ℤ) → ESp eS t Bb (svS g i : ℤ) →
    ESp eU t Bb (u : ℤ) →
    ∃ c : ℕ, IRuns I l12prog (swE q eP eG eS eU) t
        (t.upd 0 (ints (rixT g (swT i) zs) ++ R)) c ∧
        (c : ℝ) ≤ Ksw * ((g.V : ℝ) + 1) * (u : ℝ) ^ (1 - (2 : ℝ)⁻¹ ^ 50)

section Runs

variable {N0 N1 N2 N3 K f ρ R w : ℕ} {px py : Fin 3} {t : IState HS}

/-- The correction property of `corrU_spec`. -/
def CorrUOK (Kc : ℝ) : Prop :=
  ∀ (K ρ f : ℕ), 6 ≤ K → ρ < K → 1 ≤ f →
    ∀ (g : Geo), g.Pos → g.Q = 2 ^ (f * K) →
    ∀ (I : List Sym) (iprog : Fin 1 → IStmt HS 1) (t : IState HS) (zs : List ℤ) (R : List Item),
    KFL.Bits zs → t.G 0 = ints zs ++ R → zs.length = g.V →
    ∀ (e1 e2 e3 e4 e5 e6 e7 e8 e9 e10 : Ex HS) (Bb : ℕ),
    ((corrPl e1 e2 e3 e4 e5 e6 e7 e8 e9 e10).map fun p => p.2.size + 1).sum ≤ 100 →
    NoLow e1 → NoLow e2 → NoLow e3 → NoLow e4 → NoLow e5 → NoLow e6 → NoLow e7 → NoLow e8 →
    NoLow e9 → NoLow e10 →
    ESp e1 t Bb (g.recG.S0 : ℤ) → ESp e2 t Bb ((2 : ℤ) ^ (f * K)) →
    ESp e3 t Bb (g.recG.S1 : ℤ) → ESp e4 t Bb (g.recG.S2 : ℤ) → ESp e5 t Bb (g.B : ℤ) →
    ESp e6 t Bb ((2 : ℤ) ^ ρ) → ESp e7 t Bb ((f - 1 : ℕ) : ℤ) → ESp e8 t Bb ((2 : ℤ) ^ K) →
    ESp e9 t Bb (g.recG.V : ℤ) → ESp e10 t Bb (zsz (g.V : ℤ) : ℤ) →
    ∃ c : ℕ, IRuns I iprog (corrU e1 e2 e3 e4 e5 e6 e7 e8 e9 e10) t
        (t.upd 0 (ints (rixT g (corrc K ρ f) zs) ++ R)) c ∧
      (c : ℝ) ≤ Kc * corrB K f g + 10 ^ 10 * ((Bb : ℝ) + 2) ^ 2

end Runs

end NLogN.KFX

end

section
-- Source module: Solutions.KFX5Head
/-!
# KFX §5: the header of the uniform row addition

From the row-addition parameters on the background stacks `36 … 45`
(`[mC^k, K, e, t, S0, B, r, H, i, j]`) the header computes the master registers `σMR` of
`l13U`: the third row `h ∉ {i, j}`, the sorted rows `s0 < s1 < s2`, the slots `px, py`, the
factors `N0 … N3` of the shape `shR`, `Q = 2^{fK}`, `2^K`, `2^ρ` (`ρ = e % K`), the swap flags,
`gx`, and (by a halving loop) `A = zsz V`.
-/

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

namespace NLogN.KFX

open IntMul RS KIL KFD KFE KFB KFL KFW KFH AModel Butterfly

/-! ## The sorted rows -/

/-- A third row: the least of `0, 1, 2` different from `i` and `j`. -/
def hR (i j : ℕ) : ℕ := if i ≠ 0 ∧ j ≠ 0 then 0 else if i ≠ 1 ∧ j ≠ 1 then 1 else 2
def s0R (i j : ℕ) : ℕ := min (NLogN.KFX.hR i j) (min i j)
def s2R (i j : ℕ) : ℕ := max (NLogN.KFX.hR i j) (max i j)
def s1R (i j : ℕ) : ℕ := min i j + max i j + NLogN.KFX.hR i j - s0R i j - s2R i j
def pxR (i j : ℕ) : Fin 3 := if i = s0R i j then 0 else if i = s2R i j then 2 else 1
def pyR (i j : ℕ) : Fin 3 := if j = s0R i j then 0 else if j = s2R i j then 2 else 1
def σR (i j : ℕ) (x : Fin 3) : ℕ := if x = 0 then s0R i j else if x = 1 then s1R i j else s2R i j

section Rows

variable {i j : ℕ} (hij : i ≠ j)
include hij

end Rows

/-! ## Expression combinators -/

section Comb

variable {u : IState HS} {Bb : ℕ}

def eSub (a b : Ex HS) : Ex HS := .add a (.neg b)
def eMin (a b : Ex HS) : Ex HS := .sel (eSub a b) a b
def eMax (a b : Ex HS) : Ex HS := .sel (eSub a b) b a
def eSq (a : Ex HS) : Ex HS := .mul a a

end Comb

/-! ## The parameter expressions -/

/-- Parameter register `q` (stack `36 + q`). -/
abbrev pr (q : Fin 10) : Ex HS := .reg (KFW.rowR q)

def eHr : Ex HS := .sel (.neg (.mul (NLogN.KFX.pr 8) (NLogN.KFX.pr 9))) (.lit 0)
  (.sel (.neg (.mul (eSq (eSub (NLogN.KFX.pr 8) (.lit 1))) (eSq (eSub (NLogN.KFX.pr 9) (.lit 1))))) (.lit 1) (.lit 2))
def eAm : Ex HS := eMin (NLogN.KFX.pr 8) (NLogN.KFX.pr 9)
def eBm : Ex HS := eMax (NLogN.KFX.pr 8) (NLogN.KFX.pr 9)
def eT0 : Ex HS := eMin eHr eAm
def eT2 : Ex HS := eMax eHr eBm
def eT1 : Ex HS := eSub (eSub (.add (.add eAm eBm) eHr) eT0) NLogN.KFX.eT2
def eRho : Ex HS := .nmod (NLogN.KFX.pr 2) (NLogN.KFX.pr 1)
def eFKp : Ex HS := .mul (NLogN.KFX.pr 0) (NLogN.KFX.pr 1)
/-- `sR K f e s = e - K (f (s+1) - 1) - e % K`. -/
def eSRe (es : Ex HS) : Ex HS :=
  eSub (eSub (NLogN.KFX.pr 2) (.mul (NLogN.KFX.pr 1) (eSub (.mul (NLogN.KFX.pr 0) (.add es (.lit 1))) (.lit 1)))) eRho
def eN0 : Ex HS := .mul (NLogN.KFX.pr 5) (.pow2 (eSub (eSub (NLogN.KFX.pr 3) (eSRe eT0)) eFKp))
def eN1 : Ex HS := .pow2 (eSub (eSub (eSRe eT0) (eSRe eT1)) eFKp)
def eN2 : Ex HS := .pow2 (eSub (eSub (eSRe eT1) (eSRe NLogN.KFX.eT2)) eFKp)
def eN3 : Ex HS := .mul (.pow2 (eSRe NLogN.KFX.eT2)) (NLogN.KFX.pr 4)
def eQp : Ex HS := .pow2 eFKp
def eK2p : Ex HS := .pow2 (NLogN.KFX.pr 1)
def eR2p : Ex HS := .pow2 eRho
def eNtp : Ex HS := .mul (.lit 2) (.mul (.mul (.mul (NLogN.KFX.pr 5) (.pow2 (NLogN.KFX.pr 3))) (NLogN.KFX.pr 4)) (NLogN.KFX.pr 6))
def eBp : Ex HS := .mul (NLogN.KFX.pr 6) (.mul (.lit 2) (.add (NLogN.KFX.pr 7) (.lit 1)))
def eFa : Ex HS := .sel (.neg (eSq (eSub (NLogN.KFX.pr 8) eT0))) (.lit 1) (.lit 0)
def eFb : Ex HS := .add (.sel (.neg (eSq (eSub (NLogN.KFX.pr 9) NLogN.KFX.eT2))) (.lit 0) (.lit 1))
  (.mul (.sel (.neg (eSq (eSub (NLogN.KFX.pr 8) NLogN.KFX.eT2))) (.lit 0) (.lit 1))
    (.sel (.neg (eSq (eSub (NLogN.KFX.pr 9) eT0))) (.lit 0) (.lit 1)))
def eGx : Ex HS := .sel (.neg (eSq (eSub (NLogN.KFX.pr 8) NLogN.KFX.eT2))) (.lit 1) (.mul eQp eN2)

/-! ## Arithmetic of the rows -/

/-! ## The shape factors -/

/-- `sR` of row `s` at level `k`. -/
abbrev srw (K e k s : ℕ) : ℕ := KFR.sR K (mC ^ k) e s

def N0R (k K e T B i j : ℕ) : ℕ := B * 2 ^ (T - srw K e k (s0R i j) - mC ^ k * K)
def N1R (k K e i j : ℕ) : ℕ := 2 ^ (srw K e k (s0R i j) - srw K e k (s1R i j) - mC ^ k * K)
def N2R (k K e i j : ℕ) : ℕ := 2 ^ (srw K e k (s1R i j) - srw K e k (s2R i j) - mC ^ k * K)
def N3R (k K e S0 i j : ℕ) : ℕ := 2 ^ srw K e k (s2R i j) * S0

/-! ## Specifications of the parameter expressions -/

/-- The volume bound `V = 2 n r (H+1)` of a row addition. -/
def VbR (T S0 B r H : ℕ) : ℕ := 2 * (B * 2 ^ T * S0) * r * (H + 1)

/-- The admissibility facts of a row-addition parameter set. -/
structure PF (k K e T S0 B r : ℕ) : Prop where
  k1 : 1 ≤ K
  he : K * (mC ^ (k + 1) - 1) ≤ e
  ht : e - e % K + K ≤ T
  s0 : 1 ≤ S0
  b : 1 ≤ B
  r : 1 ≤ r

section PFacts

variable {k K e T S0 B r H : ℕ} (hp : PF k K e T S0 B r)
include hp

end PFacts

section SFacts

variable {k K e T S0 B r : ℕ} (hp : PF k K e T S0 B r)
include hp

end SFacts

section Par

variable {k K e T S0 B r H : ℕ} {i j : Fin mC} {u : IState HS} {Bb : ℕ}
  (hp : PF k K e T S0 B r) (hA : Agree (rowσ k K e T S0 B r H i j) u)
  (hBb : zsz ((VbR T S0 B r H : ℕ) : ℤ) ≤ Bb)

include hp hA hBb

end Par

end NLogN.KFX

end

section
-- Source module: Solutions.KFX6Run
/-!
# KFX §6: the bit-length loop and the run of the uniform row addition

`loopZ` computes `A = zsz V` on stack `48` by halving `V` (pushed on stack `14`) until it is `0`.
`rowCodeU q = withPushI hdrL (zB (l13U q))` pushes the master registers computed from the
parameter registers (`hdrL`), runs `loopZ`, then `l13U q`.
-/

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

namespace NLogN.KFX

open IntMul RS KIL KFD KFE KFB KFL KFW KFH AModel Butterfly

/-! ## The bit-length loop -/

abbrev rX : Fin HS := 14

def eHalf : Ex HS := .ndiv (.reg rX) (.lit 2)
def eInc : Ex HS := .add (.reg mA) (.lit 1)
def bodyZ : RS := .seq (.asg rX eHalf) (.asg mA eInc)
def loopZ {Pp : ℕ} : IStmt HS Pp := .loop rX (fun c => c != .zero) bodyZ.code

/-- Loop state `m`: `A = m`, `X = V / 2^m`. -/
def ssZ (t : IState HS) (V m : ℕ) : IState HS :=
  (t.upd mA (.int (m : ℤ) :: t.G mA)).upd rX (.int ((V / 2 ^ m : ℕ) : ℤ) :: t.G rX)

def cZ (Bb : ℕ) : ℕ := 2 * (4 * 10 ^ 7 * (Bb + 2) ^ 2)

/-! ## The Lemma 13 core as a hypothesis -/

/-- The statement of `l13U_body` for a fixed `q` and constant. -/
def L13U (q : ℕ) (Kc : ℝ) : Prop :=
  ∀ {I : List Sym} (N0 N1 N2 N3 K f ρ R w : ℕ) (px py : Fin 3), px ≠ py → 6 ≤ K → ρ < K →
    1 ≤ f → 1 ≤ N0 → 1 ≤ N1 → 1 ≤ N2 → 1 ≤ N3 → 1 ≤ R →
    ∀ (t : IState HS) (zs : List GI) (R0 : List Item),
    Agree (σMR N0 N1 N2 N3 K f ρ R w px py) t →
    t.G 0 = ints (gflat zs) ++ R0 → zs.length = vol (rect N0 N1 N2 N3 (f * K)) * R → NLogN.KFB.GB w zs →
    ∃ c, IRuns I l12prog (l13U q) t
        (t.upd 0 (ints (gflat (permRecs (rect N0 N1 N2 N3 (f * K))
          (rowAddAddr K ρ f px py) R zs)) ++ R0)) c ∧
      (c : ℝ) ≤ Kc * ((zs.length : ℝ) * (w + 1) * (((f * K : ℕ) : ℝ) ^ (1 - (2 : ℝ)⁻¹ ^ 50) + 1) +
        (vol (rect N0 N1 N2 N3 (f * K)) : ℝ) * ((zsz (2 * zs.length * (w + 1) : ℤ) : ℝ) + 1) ^ 3 +
        min 1 (80 * ((f : ℝ) - 1) * 2⁻¹ ^ K) * (vol (rect N0 N1 N2 N3 (f * K)) : ℝ) *
          ((zsz (2 * zs.length * (w + 1) : ℤ) : ℝ) + 1) * ((R : ℝ) * (w + 1) +
            (zsz (2 * zs.length * (w + 1) : ℤ) : ℝ) + 1))

/-! ## The code -/

def eVx : Ex HS := .mul (.reg mNt) (.add (.reg mW) (.lit 1))

def hdrL : List (Fin HS × Ex HS) :=
  [(mW, NLogN.KFX.pr 7), (mNt, eNtp), (mFa, eFa), (mFb, eFb), (mQ, eQp), (mB, NLogN.KFX.eBp), (mR2, eR2p),
   (mGx, eGx), (mN0, eN0), (mN1, eN1), (mN2, eN2), (mN3, eN3), (mK2, eK2p)]

def zL : List (Fin HS × Ex HS) := [(mA, .lit 0), (rX, eVx)]

def zB (p : IStmt HS 1) : IStmt HS 1 :=
  .seq (NLogN.KIL.pushes NLogN.KFX.zL).code (.seq loopZ (.seq p (.seq (.op (.drop rX)) (.op (.drop mA)))))

/-- **The uniform row-addition code** (over `l12prog`). -/
def rowCodeU (q : ℕ) : IStmt HS 1 := withPushI hdrL (NLogN.KFX.zB (l13U q))

/-! ## The masters agree -/

/-! ## The geometry of a row addition -/

/-- The sorted rows as a map into `Fin mC`. -/
def σF (i j : Fin mC) : Fin 3 → Fin mC := fun x =>
  ⟨σR i j x % mC, Nat.mod_lt _ (by norm_num [mC])⟩

/-! ## Header cost -/

def cHd (Bb : ℕ) : ℕ := 10 ^ 12 * (Bb + 2) ^ 3

/-! ## The run -/

end NLogN.KFX

end