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

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

import Definitions.Def_IntMul_Kappa_K_Defs12

set_option autoImplicit false
set_option relaxedAutoImplicit false
set_option maxHeartbeats 1000000

section
-- Source module: Solutions.KIL6Reg
/-!
# KIL6: register files

`regState F v D`: the managed state `F` with the natural `v i` pushed on every register
`i ∈ D`.  Assignment lists `asgL L` (each an `NE` expression) act on the valuation `v` only
(`Sp.asgs`); counted loops of assignment lists iterate the valuation map (`Sp.repA`).  `pushZs` /
`popDs` create and remove the register file.
-/

set_option linter.unusedSimpArgs false
set_option linter.unusedVariables false

namespace NLogN.KIL

open IntMul RS AModel AModelL12

/-- The register file `D` holding the valuation `v` on top of `F`. -/
def regState (F : Fin HS → List ℤ) (v : Fin HS → ℕ) (D : List (Fin HS)) : Fin HS → List ℤ :=
  fun i => if i ∈ D then (v i : ℤ) :: F i else F i

/-- Registers readable from the file or holding the valuation below it. -/
def RegsIn (F : Fin HS → List ℤ) (v : Fin HS → ℕ) (D : List (Fin HS)) (e : NE) : Prop :=
  ∀ j ∈ e.regs, j.val < 56 ∧ (j ∈ D ∨ ∃ L, F j = (v j : ℤ) :: L)

/-! ## Assignment lists -/

def asgL : List (Fin HS × NE) → RS
  | [] => .skip
  | (i, e) :: L => .seq (.asg i e.toEx) (asgL L)

def evalA : List (Fin HS × NE) → (Fin HS → ℕ) → (Fin HS → ℕ)
  | [], v => v
  | (i, e) :: L, v => evalA L (Function.update v i (e.val v))

def szA : List (Fin HS × NE) → ℕ
  | [] => 1
  | (_, e) :: L => (e.size + 1) + szA L

/-- Side conditions of an assignment list along the way. -/
def OKA (F : Fin HS → List ℤ) (D : List (Fin HS)) (M : ℕ) : List (Fin HS × NE) → (Fin HS → ℕ) → Prop
  | [], _ => True
  | (i, e) :: L, v => i ∈ D ∧ i.val < 56 ∧ v i ≤ M ∧ RegsIn F v D e ∧ e.OK v M ∧
      OKA F D M L (Function.update v i (e.val v))

/-! ## Creating and removing the file -/

def pushZs : List (Fin HS) → RS
  | [] => .skip
  | i :: D => .seq (.push i (.lit 0)) (pushZs D)

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

end NLogN.KIL

end

section
-- Source module: Solutions.KIL6Desc
/-!
# KIL6: the descriptor computations

Closed forms of the counted assignment loops used for the run-time descriptor of Lemma 12:

* `mulL X c`: `X := X · c` (powers `2^u`, `q^e`, `q^r`, `W^Kp`);
* `clogL b Y K`: `if 1 < Y then (K := K + 1; Y := ⌈Y / b⌉)`; iterated `N ≥ ⌈log_b y⌉` times
  from `(y, 0)` it leaves `K = ⌈log_b y⌉ = Nat.clog b y`.
-/

set_option linter.unusedSimpArgs false
set_option linter.unusedVariables false

namespace NLogN.KIL

open IntMul RS AModel AModelL12

/-! ## Multiplication loops -/

def mulL (X : Fin HS) (c : ℕ) : List (Fin HS × NE) := [(X, .mul (.reg X) (.lit c))]

/-! ## Ceiling-logarithm loops -/

/-- One step of the ceiling logarithm: `(y, k) ↦ (⌈y / b⌉, k + 1)` while `y > 1`. -/
def cstep (b : ℕ) (p : ℕ × ℕ) : ℕ × ℕ :=
  if 1 < p.1 then ((p.1 + b - 1) / b, p.2 + 1) else p

def clogL (b : ℕ) (Y K : Fin HS) : List (Fin HS × NE) :=
  [(K, .lt (.lit 1) (.reg Y) (.add (.reg K) (.lit 1)) (.reg K)),
   (Y, .lt (.lit 1) (.reg Y) (.div (.add (.reg Y) (.lit (b - 1))) (.lit b)) (.reg Y))]

/-! ## Registers of the Lemma 12 program -/

abbrev aT2 : Fin HS := 25
abbrev aY : Fin HS := 26
abbrev aE : Fin HS := 27
abbrev aQE : Fin HS := 28
abbrev aK : Fin HS := 29
abbrev aN : Fin HS := 30
abbrev aQR : Fin HS := 31
abbrev aQN : Fin HS := 32
abbrev aWK : Fin HS := 33
abbrev aRp : Fin HS := 34
abbrev aGR : Fin HS := 35
abbrev aJ : Fin HS := 36
abbrev aNN : Fin HS := 37
abbrev aQO : Fin HS := 53
abbrev aQM : Fin HS := 54
abbrev aDg : Fin HS := 14
abbrev aTw : Fin HS := 13
abbrev aL : Fin HS := 12
/-- The input register holding `u`. -/
abbrev aU : Fin HS := 44

/-- The register file of the program. -/
def regD : List (Fin HS) :=
  [aT2, aY, NLogN.KIL.aE, aQE, aK, aN, aQR, aQN, aWK, aRp, aGR, kQt, kQr, kR, pX, pB, pZ, pT, aJ, aNN, aQO,
    aQM, aDg, aTw, aL]

/-! ## Stage 1: `2^u` and `e = ⌈log_q 2^u⌉` -/

def desc1 (q : ℕ) : RS :=
  .seq (asgL [(aT2, .lit 1)])
    (.seq (.rep 23 (NE.reg aU).toEx (asgL (NLogN.KIL.mulL aT2 2)))
      (.seq (asgL [(aY, .reg aT2), (NLogN.KIL.aE, .lit 0)])
        (.rep 23 (NE.reg aU).toEx (asgL (clogL q aY NLogN.KIL.aE)))))

/-! ## Stage 2: `q^e` and `Kp = ⌈log_m e⌉` -/

def desc2 (q m : ℕ) : RS :=
  .seq (asgL [(aQE, .lit 1)])
    (.seq (.rep 23 (NE.reg NLogN.KIL.aE).toEx (asgL (NLogN.KIL.mulL aQE q)))
      (.seq (asgL [(aY, .reg NLogN.KIL.aE), (aK, .lit 0)])
        (.rep 23 (NE.reg NLogN.KIL.aE).toEx (asgL (clogL m aY aK)))))

/-! ## Stage 3: `r = min (W Kp) e`, `n = e - r`, `q^r`, `q^n`, `|G| q^r`, swap registers -/

def desc3 (q W : ℕ) : RS :=
  .seq (asgL [(aY, .mul (.lit W) (.reg aK)), (kR, .lt (.reg aY) (.reg NLogN.KIL.aE) (.reg aY) (.reg NLogN.KIL.aE)),
      (aN, .sub (.reg NLogN.KIL.aE) (.reg kR)), (aQR, .lit 1)])
    (.seq (.rep 23 (NE.reg kR).toEx (asgL (NLogN.KIL.mulL aQR q)))
      (asgL [(aQN, .div (.reg aQE) (.reg aQR)), (aGR, .mul (.reg kG) (.reg aQR)), (kQt, .lit 1),
        (kQr, .div (.reg aQE) (.lit q))]))

end NLogN.KIL

end

section
-- Source module: Solutions.KIL6Pieces
/-!
# KIL6: the Prop. 11 calls on the pieces

* `IRuns.irepF`: counted item-language loops with iteration-dependent costs.
* `pushNL` / `popsR`: pushing a list of natural expressions on descriptor stacks outside the
  register file, and popping them again.
* `pieceI`: one piece call — push the child descriptor (`dJ = j`, `dL = R'`, `dA = q^o`,
  `dB = q^(n-o-m^j) · |G'| · q^o`, `dC = q^(n-o-m^j) · |Suf|`), call procedure `0`, pop, and
  advance `q^o ← q^o · q^(m^j)`, `q^(n-o) ← q^(n-o) / q^(m^j)` (`aTw` holds `q^(m^j)`).
  `piece_run`: it computes `pieceStep`, for any oracle code meeting a call specification.
-/

set_option linter.unusedSimpArgs false
set_option linter.unusedVariables false

namespace NLogN.KIL

open IntMul RS AModel AModelL12

section IRepF

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

end IRepF

/-! ## Pushing natural expressions on stacks outside the register file -/

/-- The managed state after pushing the values of `L` (evaluated in `v`). -/
def pushF (F : Fin HS → List ℤ) (v : Fin HS → ℕ) : List (Fin HS × NE) → Fin HS → List ℤ
  | [] => F
  | (i, e) :: L => pushF (Function.update F i ((e.val v : ℤ) :: F i)) v L

def pushNL (L : List (Fin HS × NE)) : RS := pushes (L.map fun p => (p.1, p.2.toEx))

/-- Side conditions of a push list. -/
def PushNOK (F : Fin HS → List ℤ) (v : Fin HS → ℕ) (D : List (Fin HS)) (M : ℕ)
    (L : List (Fin HS × NE)) : Prop :=
  ∀ p ∈ L, p.1.val < 56 ∧ p.1 ∉ D ∧ RegsIn F v D p.2 ∧ p.2.OK v M ∧
    ∀ j ∈ p.2.regs, j ∉ L.map Prod.fst

/-! ## One piece call -/

/-- The child descriptor of a piece. -/
def pcL : List (Fin HS × NE) :=
  [(dJ, .reg aJ), (dL, .reg aRp), (dA, .reg aQO),
   (dB, .mul (.mul (.div (.reg aQM) (.reg aTw)) (.reg aGR)) (.reg aQO)),
   (dC, .mul (.div (.reg aQM) (.reg aTw)) (.reg NLogN.KIL.kS))]

/-- Advance the offset registers by one piece. -/
def advL : List (Fin HS × NE) :=
  [(aQO, .mul (.reg aQO) (.reg aTw)), (aQM, .div (.reg aQM) (.reg aTw))]

/-- One piece: push the child descriptor, call procedure `0`, pop, advance. -/
def pieceI : IStmt HS 1 :=
  .seq (pushNL pcL).code (.seq (.call 0) (.seq (NLogN.KIL.popsR (pcL.map Prod.fst)).code
    (asgL advL).code))

/-- The oracle call specification used by the pieces (discharged by `progU_run`). -/
def CallSpec (I : List Sym) (iprog : Fin 1 → IStmt HS 1) (child : Oracle (ZMod 2))
    (q m j P R' : ℕ) (V : ℕ) (cO : ℕ) : Prop :=
  ∀ (A B Cs : Shape) (t : IState HS) (R : List Item) (ys : List (ZMod 2)),
    DescIn t q m j P R' A B Cs → (∀ a ∈ A, 1 ≤ a) → (∀ a ∈ B, 1 ≤ a) → (∀ a ∈ Cs, 1 ≤ a) →
    ys.length = vol (P :: R' :: swShape A B Cs (q ^ m ^ j)) →
    vol (P :: R' :: swShape A B Cs (q ^ m ^ j)) = V →
    IRuns I iprog (iprog 0) (t.upd 0 (ints (enc2 ys) ++ R))
      (t.upd 0 (ints (enc2 (child j P R' A B Cs ys)) ++ R)) cO

end NLogN.KIL

end

section
-- Source module: Solutions.KIL6Levels
/-!
# KIL6: the levels of piece calls

The pieces of `pieceLevels m n` are enumerated level by level: level `j` (width `m^j`) has
`n / m^j % m` pieces.  `lvlI`: one level (count into `aDg`, `irep` of `pieceI`,
`aNN ← aNN / m`, `aDg ← 0`).  `towerUp`: `aTw ← aTw^m`, push it on the tower stack `dT`,
`aJ ← aJ + 1`.  The register valuation at the start / end of a level is `vLev`.
-/

set_option linter.unusedSimpArgs false
set_option linter.unusedVariables false

namespace NLogN.KIL

open IntMul RS AModel AModelL12

/-! ## Arithmetic of the levels -/

/-- The first `J` levels of pieces. -/
def levels (m n J : ℕ) : List ℕ := (List.range J).flatMap fun j => List.replicate (n / m ^ j % m) j

/-! ## Counting the digits -/

/-- `(y, L) ↦ (y / m, L + 1)` while `y > 0`. -/
def dstep (m : ℕ) (p : ℕ × ℕ) : ℕ × ℕ := if 0 < p.1 then (p.1 / m, p.2 + 1) else p

def dlenL (m : ℕ) (Y L : Fin HS) : List (Fin HS × NE) :=
  [(L, .lt (.lit 0) (.reg Y) (.add (.reg L) (.lit 1)) (.reg L)),
   (Y, .lt (.lit 0) (.reg Y) (.div (.reg Y) (.lit m)) (.reg Y))]

/-! ## Powers by a register -/

def mulR (X Y : Fin HS) : List (Fin HS × NE) := [(X, .mul (.reg X) (.reg Y))]

/-! ## The level valuation -/

/-- The registers at a level boundary. -/
def vLev (vB : Fin HS → ℕ) (j a b c d : ℕ) : Fin HS → ℕ :=
  Function.update (Function.update (Function.update (Function.update (Function.update vB aJ j)
    aNN a) aQO b) aQM c) aTw d

/-! ## One level -/

def dgL (m : ℕ) : List (Fin HS × NE) := [(aDg, .mod (.reg aNN) (.lit m))]
def nnL (m : ℕ) : List (Fin HS × NE) := [(aNN, .div (.reg aNN) (.lit m)), (aDg, .lit 0)]

/-- One level of pieces. -/
def lvlI (m : ℕ) : IStmt HS 1 :=
  .seq (asgL (dgL m)).code (.seq (irep 10 (NE.reg aDg).toEx pieceI) (asgL (nnL m)).code)

/-- Cost of a level with `d` pieces of oracle cost `cO`. -/
def cLvl (d cO B : ℕ) : ℕ := d * (cO + 2 * 10 ^ 9 * (B + 2) ^ 2) + 10 ^ 9 * (B + 2) ^ 2

/-! ## The tower step -/

def tw1L : List (Fin HS × NE) := [(aY, .reg aTw)]
def tw2L : List (Fin HS × NE) := [(aJ, .add (.reg aJ) (.lit 1)), (aY, .lit 0)]

/-- `aTw ← aTw ^ m`, push it on the tower, `aJ ← aJ + 1`. -/
def towerUp (m : ℕ) : RS :=
  .seq (asgL tw1L) (.seq (.rep 9 (NE.lit (m - 1)).toEx (asgL (mulR aTw aY)))
    (.seq (pushNL [(dT, .reg aTw)]) (asgL tw2L)))

end NLogN.KIL

end

section
-- Source module: Solutions.KIL6PcMain
/-!
# KIL6: all piece calls

`piecesI q m`: initialise the level registers (`aJ = 0`, `aNN = n`, `q^0`, `q^n`, `aTw = q`),
count the base-`m` digits of `n` (`aL`), push `dP = P` and the tower `dT = [q]`; run level `0`,
then `aL - 1` times (`towerUp`; next level); pop the tower and `dP`.

`pieces_run`: on the padded row array it computes
`runPieces (pieceStep child P R' G' Suf q m n) (m^·) 0 (pieceLevels m n)`, given the call
specification of the oracle code at every level `j` with `m^j ≤ n`.
-/

set_option linter.unusedSimpArgs false
set_option linter.unusedVariables false

namespace NLogN.KIL

open IntMul RS AModel AModelL12

/-! ## Initialisation -/

def inL (q : ℕ) : List (Fin HS × NE) :=
  [(aJ, .lit 0), (aNN, .reg aN), (aQO, .lit 1), (aQM, .reg aQN), (aTw, .lit q), (aY, .reg aN),
   (aL, .lit 0), (aDg, .lit 0)]

def y0L : List (Fin HS × NE) := [(aY, .lit 0)]

def pcInit (q m : ℕ) : RS :=
  .seq (asgL (inL q)) (.seq (.rep 23 (NE.reg aN).toEx (asgL (dlenL m aY aL)))
    (.seq (asgL y0L) (pushNL [(NLogN.KIL.dP, .reg kP), (dT, .reg aTw)])))

/-- The registers written by the piece phase. -/
def pcW : List (Fin HS) := [aJ, aNN, aQO, aQM, aTw, aY, aL, aDg]

/-! ## Removing the tower -/

def pcFin : RS := .seq (.rep 8 (NE.reg aL).toEx (.pop dT)) (.pop NLogN.KIL.dP)

/-! ## All pieces -/

def piecesI (q m : ℕ) : IStmt HS 1 :=
  .seq (pcInit q m).code (.seq (lvlI m)
    (.seq (irep 11 (NE.sub (.reg aL) (.lit 1)).toEx (.seq (towerUp m).code (lvlI m))) pcFin.code))

/-- Cost of all piece calls (`cO j`: the oracle cost at level `j`). -/
def cPcs (m n : ℕ) (cO : ℕ → ℕ) (B : ℕ) : ℕ :=
  ∑ j ∈ Finset.range (Nat.digits m n).length,
      (cLvl (n / m ^ j % m) (cO j) B + (m + 2) * 10 ^ 9 * (B + 2) ^ 2) +
    2 * 10 ^ 9 * (n + 2) * (B + 2) ^ 2

end NLogN.KIL

end

section
-- Source module: Solutions.KIL6Stage
/-!
# KIL6: padding stages

`padS x b z t` sets the four padding registers from `NE` expressions (`X`, `F·Y`, `(F'-F)·Y`,
and the total) and runs `padP`; `unpadS` likewise runs `unpadP`.  Their specifications on the
register file (`padS_sp`, `unpadS_sp`) reduce every padding of the main program to arithmetic.
-/

set_option linter.unusedSimpArgs false
set_option linter.unusedVariables false

namespace NLogN.KIL

open IntMul RS AModel AModelL12

/-! ## Padding stages -/

def padAsg (x b z t : NE) : List (Fin HS × NE) := [(pX, x), (pB, b), (pZ, z), (pT, t)]

def padS (x b z t : NE) : RS := .seq (asgL (padAsg x b z t)) padP

def unpadS (x b z t : NE) : RS := .seq (asgL (padAsg x b z t)) unpadP

/-- Side conditions of one padding-register expression. -/
def PadE (F : Fin HS → List ℤ) (v : Fin HS → ℕ) (M : ℕ) (e : NE) : Prop :=
  RegsIn F v regD e ∧ e.OK v M ∧ pX ∉ e.regs ∧ pB ∉ e.regs ∧ pZ ∉ e.regs ∧ pT ∉ e.regs

/-- The padding-register values after `padAsg`. -/
def padV (v : Fin HS → ℕ) (a b c d : ℕ) : Fin HS → ℕ :=
  Function.update (Function.update (Function.update (Function.update v pX a) pB b) pZ c) pT d

end NLogN.KIL

end

section
-- Source module: Solutions.KIL6PreA
/-!
# KIL6: the run-time descriptor

`preA q m W`: create the register file and compute `2^u`, `e = ⌈log_q 2^u⌉`, `q^e`,
`Kp = ⌈log_m e⌉`, `r = min (W·Kp) e`, `n = e - r`, `q^r`, `q^n`, `vol G · q^r` and the swap
registers.
-/

set_option linter.unusedSimpArgs false
set_option linter.unusedVariables false

namespace NLogN.KIL

open IntMul RS AModel AModelL12

def preA (q m W : ℕ) : RS :=
  .seq (pushZs regD) (.seq (desc1 q) (.seq (desc2 q m) (desc3 q W)))

/-- Cost of the descriptor phase. -/
def cPreA (u e B : ℕ) : ℕ := 10 ^ 9 * (u + 2 * e + 4) * (B + 2) ^ 2

/-- The input valuation: `kP, kG, kS, aU` hold `P, g, sv, u`; the file holds `0`. -/
def v0 (P g sv u : ℕ) : Fin HS → ℕ :=
  Function.update (Function.update (Function.update (Function.update (fun _ => 0) kP P) kG g)
    NLogN.KIL.kS sv) aU u

end NLogN.KIL

end

section
-- Source module: Solutions.KIL6PreB
/-!
# KIL6: binary embedding, high-digit swaps, binary extraction

`preB`: the two padding passes realising `embed (binShape … 2^u) (binShape … q^e)`;
`swapsP q`: the `r` primitive digit interchanges; `postB`: the two unpadding passes realising
`extract`.  All on the bit array `enc2 xs`.
-/

set_option linter.unusedSimpArgs false
set_option linter.unusedVariables false

namespace NLogN.KIL

open IntMul RS AModel AModelL12

/-! ## The padding expressions -/

def xA : NE := .mul (.reg kP) (.mul (.reg aT2) (.reg kG))
def bA : NE := .mul (.reg aT2) (.reg NLogN.KIL.kS)
def zA : NE := .mul (.sub (.reg aQE) (.reg aT2)) (.reg NLogN.KIL.kS)
def tA : NE := .mul NLogN.KIL.xA (.mul (.reg aQE) (.reg NLogN.KIL.kS))
def yB : NE := .mul (.reg kG) (.mul (.reg aQE) (.reg NLogN.KIL.kS))
def bB : NE := .mul (.reg aT2) yB
def zB : NE := .mul (.sub (.reg aQE) (.reg aT2)) yB
def tB : NE := .mul (.reg kP) (.mul (.reg aQE) yB)
def tUB : NE := .mul (.reg kP) (.mul (.reg aT2) yB)
def tUA : NE := .mul NLogN.KIL.xA (.mul (.reg aT2) (.reg NLogN.KIL.kS))

def preB : RS := .seq (padS NLogN.KIL.xA bA NLogN.KIL.zA NLogN.KIL.tA) (padS (.reg kP) NLogN.KIL.bB NLogN.KIL.zB tB)
def postB : RS := .seq (unpadS (.reg kP) NLogN.KIL.bB NLogN.KIL.zB tUB) (unpadS NLogN.KIL.xA bA NLogN.KIL.zA tUA)

/-! ## The high-digit swaps on the file -/

end NLogN.KIL

end

section
-- Source module: Solutions.KIL6Mid
/-!
# KIL6: the middle part

When digits remain (`n = e - r ≥ 1`): compute `W^Kp` and the padded row count
`R' = padRows W Kp (q^r)` (`desc4`), pad the rows (`padR`), run all piece calls (`piecesI`),
and remove the row padding (`unpadR`).  `mid_run`: this realises `l12mid`.
-/

set_option linter.unusedSimpArgs false
set_option linter.unusedVariables false

namespace NLogN.KIL

open IntMul RS AModel AModelL12

def rpE : NE := .mul (.reg aWK) (.div (.sub (.add (.reg aQR) (.reg aWK)) (.lit 1)) (.reg aWK))

def desc4 (W : ℕ) : RS :=
  .seq (asgL [(aWK, .lit 1)]) (.seq (.rep 23 (NE.reg aK).toEx (asgL (NLogN.KIL.mulL aWK W)))
    (asgL [(aRp, rpE)]))

/-! ## Row padding -/

def yR : NE := .mul (.reg aQN) (.mul (.reg aGR) (.mul (.reg aQN) (.reg NLogN.KIL.kS)))
def bR : NE := .mul (.reg aQR) yR
def zR : NE := .mul (.sub (.reg aRp) (.reg aQR)) yR
def tR : NE := .mul (.reg kP) (.mul (.reg aRp) yR)
def tUR : NE := .mul (.reg kP) (.mul (.reg aQR) yR)

def padR : RS := padS (.reg kP) NLogN.KIL.bR zR NLogN.KIL.tR
def unpadR : RS := unpadS (.reg kP) NLogN.KIL.bR zR tUR

def midI (q m W : ℕ) : IStmt HS 1 :=
  .seq (RS.seq (desc4 W) padR).code (.seq (piecesI q m) unpadR.code)

/-- Cost of the middle part. -/
def cMid (K V m n : ℕ) (cO : ℕ → ℕ) (B : ℕ) : ℕ :=
  10 ^ 9 * (K + 3) * (B + 2) ^ 2 + 2 * cPad V B + cPcs m n cO B

end NLogN.KIL

end

section
-- Source module: Solutions.KIL6Main
/-!
# KIL6: the Lemma 12 program

`l12I q m W`: descriptor (`preA`), binary embedding (`preB`), `min (W·Kp) e` high-digit swaps
(`swapsP q`), the middle part if digits remain (`midI`, a branch on `n = 0`), binary extraction
(`postB`), removal of the register file.

`l12_run`: on the bit array `enc2 xs` of shape `binShape Pre G Suf (2^u)`, given the run
specification `CallSpec` of the Prop. 11 code at every level used, the program computes
`l12v child W m q Kp (W·Kp) e u Pre G Suf 0 xs` (with `e = ⌈log_q 2^u⌉`, `Kp = ⌈log_m e⌉`).
-/

set_option linter.unusedSimpArgs false
set_option linter.unusedVariables false

namespace NLogN.KIL

open IntMul RS AModel AModelL12

/-- **The Lemma 12 item-language program** (`q`, `m`, `W` are literals). -/
def l12I (q m W : ℕ) : IStmt HS 1 :=
  .seq (RS.seq (preA q m W) (RS.seq preB (NLogN.KIL.swapsP q))).code
    (.seq (.ite aN (fun c => c == Cls.zero) .skip (midI q m W)) (RS.seq postB (popDs regD)).code)

/-- Cost of the Lemma 12 program. -/
def cL12 (u e r K n Vq V' m q : ℕ) (cO : ℕ → ℕ) (B : ℕ) : ℕ :=
  cPreA u e B + 4 * (10 ^ 9 * (B + 2) ^ 2 + cPad Vq B) +
    (r + 1) * (cSw1 q Vq B + 10 ^ 4 * (B + 2)) + ((if n = 0 then 0 else cMid K V' m n cO B) + 26) +
    25 * (2 * B + 2) + 1

end NLogN.KIL

end

section
-- Source module: Solutions.KFCBase
/-!
# KFC base: fine-cost item-language loops

Item-language (`IStmt`) building blocks whose per-element cost is `O(1 + zsz a)` for the moved
entry `a`, with **no** binary loop counter per element:

* `mvW X Y`   : move integers from `X` to `Y` until the top of `X` is not an integer
                (delimiter-driven, cost `(76 w + 131)` per entry of bit size `≤ w`);
* `uloop C D b`: run `b` once per mark on the *unary counter* `C` (the marks are moved to `D`
                and then moved back, so the counter is restored); overhead `64` per iteration;
* `mvB X Y C D`: move exactly `n` integers, `n` = number of marks on `C`;
* `toU A X K K2`: push `n` marks on `X`, where `n ≥ 0` is the integer on top of `A`
                (binary counter, cost `O((n+1)(zsz n + 2))`; superseded by the linear
                `KFCUnary.toUL`, which the programs use);
* `clrM X`    : drop all marks on top of `X`.
-/

set_option linter.unusedSimpArgs false

namespace NLogN

open IntMul

namespace KFC

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

/-! ## Classes and item lists -/

def isMk (c : Cls) : Bool := c == .mark
def isIn (c : Cls) : Bool := c == .zero || c == .pos || c == .neg

def marks (n : ℕ) : List Item := List.replicate n .mark

@[simp] theorem marks_zero : marks 0 = [] := rfl

/-- The head of `L` is not a mark. -/
def NM (L : List Item) : Prop := L.head? ≠ some .mark
/-- The head of `L` is not an integer. -/
def NI (L : List Item) : Prop := ∀ a : ℤ, L.head? ≠ some (.int a)

/-- Mark-terminated encoding of a list of groups. -/
def enc (gs : List (List ℤ)) : List Item := gs.flatMap fun g => ints g ++ [.mark]

@[simp] theorem enc_nil : enc [] = [] := rfl

/-! ## `mvW`: move integers until a delimiter -/

def mvW (X Y : Fin HS) : IStmt HS Pp := .loop X isIn (.op (.mv X Y))

/-- Per-entry move cost bound for entries of bit size `≤ w`. -/
def kap (w : ℕ) : ℕ := 76 * w + 200

/-! ## Unary counted loops -/

/-- One iteration per mark on `C`; the marks are moved to `D` and then back. -/
def uloop (C D : Fin HS) (b : IStmt HS Pp) : IStmt HS Pp :=
  .seq (.loop C isMk (.seq (.op (.unmark C)) (.seq (.op (.mark D)) b)))
    (.loop D isMk (.seq (.op (.unmark D)) (.op (.mark C))))

/-- The counter frame: `T` with `C`, `D` overwritten. -/
def cf (T : IState HS) (C D : Fin HS) (LC LD : List Item) : IState HS := (T.upd C LC).upd D LD

/-! ## `mvB`: move a counted block -/

def mvB (X Y C D : Fin HS) : IStmt HS Pp := uloop C D (.op (.mv X Y))

/-! ## `clrM`: drop the marks on top of a stack -/

def clrM (X : Fin HS) : IStmt HS Pp := .loop X isMk (.op (.unmark X))

/-! ## `toU`: binary to unary (prologue only) -/

end KFC

end NLogN

end

section
-- Source module: Solutions.KFCUnary
/-!
# KFC: linear-time binary → unary conversion

`toUL A X` pushes `n` marks on `X` (`n ≥ 0` the integer on top of `A`) in time
`O(n + zsz n)`: the bits of `n` are peeled off with `split` by the ruler `2` (cost `O(1)` each,
since `split` costs `O(m)` for the ruler `2^m`), then the unary count is rebuilt from the most
significant bit by doubling (`dbl`, cost `O(c)` for count `c`; the doublings sum to `O(n)`).

Scratch stacks `sK = 26`, `sK2 = 27`, `sY = 28`, `sBt = 29` are restored.
-/

set_option linter.unusedSimpArgs false
set_option linter.unnecessarySeqFocus false

namespace NLogN

open IntMul

namespace KFC

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

abbrev uK : Fin HS := 26
abbrev uR : Fin HS := 27
abbrev uY : Fin HS := 28
abbrev uB : Fin HS := 29

/-! ## Doubling a unary count -/

def dbl (X Y : Fin HS) : IStmt HS Pp :=
  .seq (.loop X isMk (.seq (.op (.unmark X)) (.seq (.op (.mark Y)) (.op (.mark Y)))))
    (.loop Y isMk (.seq (.op (.unmark Y)) (.op (.mark X))))

/-! ## Peeling the bits -/

def isNZ' (c : Cls) : Bool := c != .zero

def splitL : IStmt HS Pp := .loop uK isNZ' (.op (.split uK uR uB))

/-! ## Rebuilding from the most significant bit -/

/-- Value of a most-significant-first bit list. -/
def valM : List ℕ → ℕ
  | [] => 0
  | d :: D => d * 2 ^ D.length + valM D

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

def bld (X : Fin HS) : IStmt HS Pp :=
  .loop uB isIn (.seq (dbl X uY) (.seq (.ite uB isPos (.op (.mark X)) .skip) (.op (.drop uB))))

/-! ## The conversion -/

def toUL (A X : Fin HS) : IStmt HS Pp :=
  .seq (.op (.dup A uK)) (.seq (.op (.const uR 2)) (.seq (.op (.mark uB)) (.seq NLogN.KFC.splitL
    (.seq (.op (.drop uK)) (.seq (.op (.drop uR)) (.seq (.op (.const uY 0)) (.seq (bld X)
      (.seq (.op (.drop uY)) (.op (.unmark uB))))))))))

/-- Cost of the linear conversion. -/
def cUL (n : ℕ) : ℕ := 1000 * (n + 1) + 300 * zsz (n : ℤ)

end KFC

end NLogN

end

section
-- Source module: Solutions.KFDBase
/-!
# KFD base: fine compilation of register programs

The register programs `RS` of KIL1–KIL6 are compiled by `RS.code`, whose counted loop
`rep i n b` keeps a *binary* counter on `i` and decrements it in every iteration (cost
`Θ(zsz n)` per iteration).  `fcode` compiles the same program with a *unary* counter instead:

* `rep i n b ↦ push i n; toUL i i; loop i isMk (unmark i; fcode b); drop i`

(the value `n` stays below the `n` marks as a sentinel; `toUL` is the linear binary → unary
conversion of `KFCUnary`).  Everything else is compiled as by `RS.code`.

`fcode_runs` (**simulation**): if `p.ok s` and every loop body leaves its counter alone
(`FOK p`: the counter is not *touched* — read or written — by the body, and avoids the
`toUL` scratch stacks `26 … 29`), then `fcode p` realises `p.run` with the cost `fcost p s`.
The run is proved relative to a set of *ghost* stacks (the counters of the enclosing loops,
which hold marks in the fine run and the binary count in the `RS` semantics).

`fcost_le_cost`: `fcost p s ≤ 100 · p.cost s` (coarse transfer of every `KIL` cost bound).
The point of `fcost` is the loop: per iteration it charges `fcost b + 31` plus the amortised
`toUL` cost, i.e. `O(1)` overhead per iteration, independent of the size of the count.

`FSp p s t c` is the fine analogue of `KIL.Sp`.
-/

set_option linter.unusedSimpArgs false
set_option linter.unnecessarySeqFocus false

namespace NLogN

open IntMul RS

namespace KFD

open KFC (marks marks_zero isMk NM toUL cUL uK uR uY uB)

/-! ## Expression congruence -/

/-! ## Touched stacks, fine compilability -/

/-- Stacks read or written by a program. -/
def touch : RS → List (Fin HS)
  | .skip => []
  | .asg i e => i :: e.regs
  | .push i e => i :: e.regs
  | .pop i => [i]
  | .seq a b => touch a ++ touch b
  | .ite c a b => c.regs ++ (touch a ++ touch b)
  | .rep i n b => i :: (n.regs ++ touch b)

/-- The `toUL` scratch stacks. -/
def scr : List (Fin HS) := [uK, uR, uY, uB]

/-- Fine compilability (boolean): loop bodies do not touch their counters; counters avoid the
`toUL` scratch stacks. -/
def fok : RS → Bool
  | .seq a b => fok a && fok b
  | .ite _ a b => fok a && fok b
  | .rep i _ b => !((touch b).contains i) && !(NLogN.KFD.scr.contains i) && fok b
  | _ => true

/-! ## The fine compilation and its cost -/

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

def fcode : RS → IStmt HS Pp
  | .skip => .skip
  | .asg i e => (RS.asg i e).code
  | .push i e => (RS.push i e).code
  | .pop i => .op (.drop i)
  | .seq a b => .seq (fcode a) (fcode b)
  | .ite c a b => .seq (c.code zX) (.seq (.op (.mv zX.V zX.A))
      (.ite zX.A (fun k => k == .neg) (.seq (.op (.drop zX.A)) (fcode a))
        (.seq (.op (.drop zX.A)) (fcode b))))
  | .rep i n b => .seq (RS.push i n).code (.seq (toUL i i)
      (.seq (.loop i isMk (.seq (.op (.unmark i)) (fcode b))) (.op (.drop i))))

def fcost : RS → IState HS → ℕ
  | .skip, _ => 1
  | .asg i e, s => asgCost i e s
  | .push _ e, s => e.cost s + (2 * moveCost (zsz (e.eval s)) + 19)
  | .pop i, s => 2 * zsz (s.top i) + 2
  | .seq a b, s => fcost a s + fcost b (a.run s)
  | .ite c a b, s => c.cost s + (2 * moveCost (zsz (c.eval s)) + 19) + (2 * zsz (c.eval s) + 2) +
      25 + (if c.eval s < 0 then fcost a s else fcost b s)
  | .rep i n b, s => n.cost s + (2 * moveCost (zsz (n.eval s)) + 19) + cUL (n.eval s).toNat +
      ((∑ k ∈ Finset.range (n.eval s).toNat, (fcost b (repSt i n b s k) + 31)) + 30) +
      (2 * zsz (n.eval s) + 2)

/-! ## Ghost stacks -/

/-- `t` with the stacks `gs` taken from `s'`. -/
def ovr (gs : List (Fin HS)) (s' t : IState HS) : IState HS :=
  ⟨fun j => if j ∈ gs then s'.G j else t.G j, t.ip, t.out⟩

@[simp] theorem ovr_G (gs : List (Fin HS)) (s' t : IState HS) (j : Fin HS) :
    (ovr gs s' t).G j = if j ∈ gs then s'.G j else t.G j := rfl
@[simp] theorem ovr_ip (gs : List (Fin HS)) (s' t : IState HS) : (ovr gs s' t).ip = t.ip := rfl
@[simp] theorem ovr_out (gs : List (Fin HS)) (s' t : IState HS) : (ovr gs s' t).out = t.out := rfl

/-- `s'` agrees with `s` off the ghost stacks. -/
def Agr (gs : List (Fin HS)) (s s' : IState HS) : Prop :=
  (∀ j, j ∉ gs → s'.G j = s.G j) ∧ s'.ip = s.ip ∧ s'.out = s.out

/-! ## The simulation -/

/-! ## Top level -/

/-! ## Coarse transfer: `fcost ≤ 100 · cost` -/

/-! ## Fine specifications -/

/-- The register program `p` runs from `s` to `t` with fine cost at most `c`. -/
def FSp (p : RS) (s t : IState HS) (c : ℕ) : Prop := p.ok s ∧ p.run s = t ∧ fcost p s ≤ c

/-- Programs without loops. -/
def norep : RS → Bool
  | .seq a b => norep a && norep b
  | .ite _ a b => norep a && norep b
  | .rep _ _ _ => false
  | _ => true

/-! ## Loops -/

/-- Loop overhead of a fine loop with count `N`. -/
def cRep (N : ℕ) : ℕ := 400 * zsz (N : ℤ) + 1100 * N + 1200

/-! ## Literal-to-register simulation -/

end KFD

end NLogN

end

section
-- Source module: Solutions.KFDOps
/-!
# KFD: fine element loops

Fine (`FSp`) versions of the element loops of KIL3Ops (`dupK`, `dropK`, `zeroK`, `zipK`), of
`mvK` with a register count, and of the conditional / `c`-way branch (`sel`).  The per-element
cost depends only on a *data bound* `w` (bits: `w = 1`); the count enters only through the
loop overhead `cRep N = O(N + zsz N)`.
-/

set_option linter.unusedSimpArgs false
set_option linter.unnecessarySeqFocus false

namespace NLogN

open IntMul RS

namespace KFD

open KIL

/-! ## Expression costs -/

/-- Fine cost of an element loop: `N` iterations of cost `c`, count of size `≤ Bb`. -/
def cLp (N c Bb : ℕ) : ℕ := 1000 * (Bb + 2) + N * (c + 1100)

/-! ## Moves with a register count -/

/-! ## `dupK`, `dropK`, `zeroK`, `zipK` -/

/-! ## Conditionals and the `c`-way branch -/

end KFD

end NLogN

end

section
-- Source module: Solutions.KFDGen
/-!
# KFD: the fine offset-table generator

Fine specifications of KIL2's `genP` stages and KIL4's uniform generator `genPU` (programs
unchanged).  All entries are residues `< Q`, so every element operation costs `O(b)` with the
*local* bound `zsz (2Q) ≤ b`; loop counts enter only linearly (`cLp N c N`, since
`zsz N ≤ N`).  `genPU_fsp`: cost `≤ 10^9 · (n + 1) · (vol A + 1) · (b + 2)`, where `vol A` is
the number of table entries (blocks).  No global size bound is needed.
-/

set_option linter.unusedSimpArgs false
set_option linter.unnecessarySeqFocus false

namespace NLogN

open IntMul RS

namespace KFD

open KIL AModel

section Proofs

variable (s : IState HS) {Q b : ℕ}

def cBodyF (M b : ℕ) : ℕ :=
  cLp M (cInner b) M + cLp M (250 * (b + 2)) M + (14 + 1) * 2000 * (b + 2)

def cOuterF (a M b : ℕ) : ℕ := cLp a (cBodyF M b) a

def cStageF (a M b : ℕ) : ℕ :=
  (1 + 1) * 2000 * (b + 2) + cOuterF a M b + (2 * b + 2) + cLp M (2 * b + 2) M +
    cLp (a * M) (250 * (b + 2)) (a * M)

end Proofs

/-! ## The uniform stage and generator -/

def cStageUF (a M b : ℕ) : ℕ := 200 * cStageF a M b + 2 * (2 * a + 2) + 2 * (2 * b + 2) +
  2 * (2 * M + 2) + 2 * (2 * (a * M) + 2)

/-- Fine cost of the uniform stages. -/
def stCostF (b : ℕ) : List (ℕ × ℕ) → List (ℕ × ℕ) → ℕ
  | [], _ => 1
  | x :: rv, suf => cStageUF x.1 (vol (suf.map Prod.fst)) b + stCostF b rv (x :: suf)

end KFD

end NLogN

end

section
-- Source module: Solutions.KFDRot
/-!
# KFD: the fine block rotation `rotP`

Fine specification of KILRot's `rotP` (program unchanged).  Every record move is a fine `mvK`
(`O(B·(w+2))` for a record of `B` entries of size `≤ w`), every per-block number (`Q`, `a_p`,
`Q - a_p`) is bounded by `Q` itself (`zsz Q ≤ Q`), so a block of `Q` records costs
`O(Q·B·(w+2))`; one multiplication `P·Q·B` per call costs `O((Bb+2)^2)`.

`rotP_fspec`: `rotP` at cost `≤ 10^6 (P·Q·B + 1)(w + 2) + 10^8 (Bb + 2)^2`.
-/

set_option linter.unusedSimpArgs false
set_option linter.unnecessarySeqFocus false

namespace NLogN

open IntMul RS

namespace KFD

open KIL AModel

/-- Fine cost of moving one record of `B` entries of size `≤ w`. -/
def cRecF (B w : ℕ) : ℕ := 3000 * (B + 1) * (w + 2)

/-- Fine cost bound of one rotated block. -/
def cRBlkF (Q B w : ℕ) : ℕ :=
  4 * (8000 * (Q + 2) + (1500 * Q + 1200) + Q * cRecF B w) + (2 * Q + 2)

section Proof

variable (s : IState HS) (r : ℕ → ℕ → List ℤ) (P Q B w : ℕ) (as : ℕ → ℕ)

/-- Fine cost bound of `rotP`. -/
def cRotF (P Q B w Bb : ℕ) : ℕ :=
  (5 + 1) * 10 ^ 7 * (Bb + 2) ^ 2 + (1000 * (Bb + 2) + P * (cRBlkF Q B w + 1100)) +
    (200 * (Bb + 2) + (1500 * (P * (Q * B)) + 1200) + P * (Q * B) * (250 * (w + 2))) + (2 * Bb + 2)

end Proof

end KFD

end NLogN

end

section
-- Source module: Solutions.KFDNeg
/-!
# KFD: the fine block reflection `negP`

Fine specification of KIL2Neg's `negP` (program unchanged), as `rotP_fspec` (KFDRot):
`negP_fspec` at cost `≤ 10^6 (P·Q·B + 1)(w + 2) + 10^8 (Bb + 2)^2`.
-/

set_option linter.unusedSimpArgs false
set_option linter.unnecessarySeqFocus false

namespace NLogN

open IntMul RS

namespace KFD

open KIL AModel

/-- Fine cost bound of one reflected block. -/
def cNBlkF (Q B w : ℕ) : ℕ :=
  2 * (10000 * (Q + 2) + (1500 * Q + 1200) + Q * (2 * cRecF B w)) +
  2 * (12000 * (Q + 2) + (1500 * Q + 1200) + Q * cRecF B w) + (2 * Q + 2)

section Proof

variable (s : IState HS) (r : ℕ → ℕ → List ℤ) (P Q B w : ℕ) (as : ℕ → ℕ)

/-- Fine cost bound of `negP`. -/
def cNegF (P Q B w Bb : ℕ) : ℕ :=
  (5 + 1) * 10 ^ 7 * (Bb + 2) ^ 2 + (1000 * (Bb + 2) + P * (cNBlkF Q B w + 1100)) +
    (200 * (Bb + 2) + (1500 * (P * (Q * B)) + 1200) + P * (Q * B) * (250 * (w + 2))) + (2 * Bb + 2)

end Proof

end KFD

end NLogN

end

section
-- Source module: Solutions.KFDPass
/-!
# KFD: the fine gather `divPgU` (multiplication of the target field by `c⁻¹`)

Fine specifications of KIL3Gather / KIL4Div (programs unchanged).  Three size bounds:

* `w` bounds the data entries (per-entry cost `O(w + 2)`);
* `b ≥ zsz (2c)` bounds the residue counter and the `c`-way branch (per-record cost
  `O(c · (b + 2))`, a constant for fixed `c`);
* `bL` bounds the per-block numbers `Q`, `B`, `Q·B` of the run-time pass counts (`cntE`, one
  multiplication per pass and block: `O((bL + 2)^2)` per block, `= O(Q·B)` for `bL ≈ zsz (Q·B)`);
* `Bb` bounds the per-call numbers `P`, `P·Q·B` (`O((Bb + 2)^2)` per call).
-/

set_option linter.unusedSimpArgs false
set_option linter.unnecessarySeqFocus false

namespace NLogN

open IntMul RS

namespace KFD

open KIL AModel

/-! ## Literal passes -/

def cBrF (B w : ℕ) : ℕ := cLp B (250 * (w + 2)) B

section GatherLemmas

variable {rp : ℕ → List ℤ} {Q c B : ℕ}

def cBodyGF (c B w b : ℕ) : ℕ := c * cSel b + cBrF B w + (14 + 1) * 2000 * (b + 2)

def cPassGF (Q c B w b : ℕ) : ℕ :=
  (1 + 1) * 2000 * (b + 2) + cLp Q (cBodyGF c B w b) Q + (2 * b + 2) +
    cLp (Q * B) (250 * (w + 2)) (Q * B)

/-! ## Uniform passes -/

def cPassGUF (Q c B w b bL : ℕ) : ℕ :=
  10 * 10 ^ 7 * (bL + 2) ^ 2 + (200 * cPassGF Q c B w b + (2 * bL + 2))

def cBlkGUF (Q c B w b bL : ℕ) : ℕ := 1 + c * cPassGUF Q c B w b bL

def cDivPgRF (P Q c B w b bL Bb : ℕ) : ℕ :=
  cLp P (cBlkGUF Q c B w b bL) Bb + cLp (P * (Q * B)) (250 * (w + 2)) Bb

def cDivPgUF (P Q c B w b bL Bb : ℕ) : ℕ :=
  4 * 10 ^ 7 * (Bb + 2) ^ 2 + (c * cSel Bb + cDivPgRF P Q c B w b bL Bb + (2 * Bb + 2))

end GatherLemmas

end KFD

end NLogN

end

section
-- Source module: Solutions.KFDScat
/-!
# KFD: the fine scatter `mulPgU` (multiplication of the target field by `c`)

Fine specifications of KIL3Scatter / KIL4Mul (programs unchanged), with the bounds `w`, `b`,
`bL`, `Bb` of KFDPass.
-/

set_option linter.unusedSimpArgs false
set_option linter.unnecessarySeqFocus false

namespace NLogN

open IntMul RS

namespace KFD

open KIL AModel

section ScatterLemmas

variable {rp : ℕ → List ℤ} {Q c B : ℕ}

def cBlkSUF (Q c B w b bL : ℕ) : ℕ :=
  1 + c * cPassGUF Q c B w b bL + cLp (Q * B) (250 * (w + 2)) bL

def cMulPgRF (P Q c B w b bL Bb : ℕ) : ℕ :=
  cLp P (cBlkSUF Q c B w b bL) Bb + cLp (P * (Q * B)) (250 * (w + 2)) Bb

def cMulPgUF (P Q c B w b bL Bb : ℕ) : ℕ :=
  4 * 10 ^ 7 * (Bb + 2) ^ 2 + (c * cSel Bb + cMulPgRF P Q c B w b bL Bb + (2 * Bb + 2))

end ScatterLemmas

end KFD

end NLogN

end

section
-- Source module: Solutions.KFDShift
/-!
# KFD: closed bounds for the fine gather / scatter, fine shifts

* `sq_zsz_le`: `(zsz x + 2)² ≤ 18 x` for `x ≥ 1` (local per-block numbers cost `O(block)`).
* closed bounds of `cDivPgUF`, `cMulPgUF` with the local choices `b = 2c`, `bL = zsz (c·Q·B)`:
  `≤ 10^10 (c+1)^3 (V+1)(w+2) + 10^8 (c+1)(Bb+2)²`.
-/

set_option linter.unusedSimpArgs false
set_option linter.unnecessarySeqFocus false

namespace NLogN

open IntMul RS

namespace KFD

open KIL AModel

/-! ## Closed bounds of the gather / scatter pieces -/

section Bounds

variable {Q c B w b bL : ℕ}

end Bounds

section Bounds2

variable {P Q c B w b bL Bb : ℕ}

end Bounds2

/-! ## Fine `fdivPU` / `fmulPU` -/

/-- Closed cost of the fine `fdivPU d` / `fmulPU d`. -/
def cFrF (d V w Bb : ℕ) : ℕ :=
  10 ^ 10 * ((d + 1) ^ 3 * ((V + 1) * (w + 2))) + 10 ^ 8 * ((d + 1) * (Bb + 2) ^ 2)

/-! ## Fine table shifts -/

/-- Closed cost of the fine `shPU` / `shNU`. -/
def cShUF (nst V w Bb : ℕ) : ℕ :=
  2 * 10 ^ 12 * ((nst + 1) * ((V + 1) * (w + 2))) + 2 * 10 ^ 8 * (Bb + 2) ^ 2

section Body

variable {sh : Shape} {n Q : ℕ} (Λ : Layout sh n Q)

/-- Closed cost of the fine `shFracU`. -/
def cShFF (k d nst V w Bb : ℕ) : ℕ := cFrF d V w Bb + cFrF k V w Bb + cShUF nst V w Bb

end Body

end KFD

end NLogN

end

section
-- Source module: Solutions.KFDStep
/-!
# KFD: the fine uniform level step

`stepG_fspec` / `stepU_fspec`: the KIL4 step programs (unchanged) with fine cost
`cStepGF`; the data bound `w` and the global bound `Bb` are separate.
`cStepGF_le`: `cStepGF ≤ tKF·(10^13 (V+1)(w+2) + 10^10 (Bb+2)²)`.
-/

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

namespace NLogN

open IntMul RS

namespace KFD

open KIL AModel AModelL12

/-- Fine cost of `stepG`. -/
def cStepGF (k d nst V w Bb : ℕ) (pa pc : Ex HS) (stg : List (Ex HS × Ex HS)) : ℕ :=
  100 * pushCost Bb (prmL pa pc ++ stgPush stg) + cShFF k d nst V w Bb + 100 * (4 * (2 * Bb + 2) + 1)

section Body

variable {sh : Shape} {n Q : ℕ} (Λ : Layout sh n Q)

end Body

variable {m : ℕ}

/-- Fine step weight of a template. -/
def tKF (rX : Fin HS) (T : TSh m) : ℕ :=
  pushSz (prmL (paT rX T) (pcT T) ++ stgPush (stgT rX T)) + (T.k + 1) ^ 3 + (T.d + 1) ^ 3 +
    (stgT rX T).length + 3

def cStepsUF (rX : Fin HS) (V w Bb : ℕ) : List (TSh m) → ℕ
  | [] => 1
  | T :: Ts => cStepGF T.k T.d (stgT rX T).length V w Bb (paT rX T) (pcT T) (stgT rX T) +
      cStepsUF rX V w Bb Ts

def tKFs (rX : Fin HS) (Ts : List (TSh m)) : ℕ := (Ts.map (tKF rX)).sum

end KFD

end NLogN

end

section
-- Source module: Solutions.KFDTab
/-!
# KFD: fine stream-table operations

The register programs `copyCU`, `accCU`, `accLU`, `buildCU`, `buildsCU`, `rebuildCU` of KIL5Tab
are unchanged; here they get *fine* specifications (`FSp`, i.e. compiled by `fcode`, unary loop
counters).  The cost unit is `cCF V Bb = O(V) + O((Bb + 2)^2)`: linear in the table volume
`V = W · v`, plus one multiplication (`w · v`) per operation.  `applyCF` is the fine `applyCU`.
-/

set_option linter.unusedSimpArgs false
set_option linter.unnecessarySeqFocus false

namespace NLogN

open IntMul RS

namespace KFD

open KIL

/-- Fine cost unit of an element loop over at most `V` bits. -/
def cOF (V Bb : ℕ) : ℕ := 10 ^ 6 * (V + 1) + 1000 * (Bb + 2)

/-- Fine cost unit of one table operation. -/
def cCF (V Bb : ℕ) : ℕ := 1000 * cOF V Bb + cPushT Bb + (2 * Bb + 2)

/-- Fine cost of a rebuild. -/
def cRebuildF (W V Bb : ℕ) : ℕ := (W * (2 * W + 2) + 202) * cCF V Bb

/-! ## Fine compilability of the table programs -/

/-! ## Applying a stream code in place (fine) -/

section Apply

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

/-- The fine `applyCU`. -/
def applyCF (w : ℕ) (code : IStmt HS Pp) : IStmt HS Pp :=
  .seq (fcode (RS.seq (.push NLogN.KIL.rT (.mul (.lit (w : ℤ)) (.reg NLogN.KIL.rV))) (NLogN.mvK 0 15 21 (.reg NLogN.KIL.rT))))
    (.seq code (fcode (RS.seq (NLogN.mvK 15 0 21 (.reg NLogN.KIL.rT)) (.pop NLogN.KIL.rT))))

end Apply

end KFD

end NLogN

end

section
-- Source module: Solutions.KFDNet
/-!
# KFD: the fine network of one level

KIL5Net's `framesU_runs`, `gateCU_runs`, `netCU_runs` for the fine table code (`applyCF`,
`fcode rebuildCU`, KFDTab).  Every frame change costs `cE X + cCF (W·v) Bb`, every gate
`cRebuildF W (W·v) Bb + 1`; `cCF V Bb = O(V) + O((Bb+2)^2)`.
-/

set_option linter.unusedSimpArgs false
set_option linter.unnecessarySeqFocus false

namespace NLogN.KFD

open IntMul RS AModel KIL

section Seq

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

end Seq

section Net

variable {Pp : ℕ} {I : List Sym} {iprog : Fin Pp → IStmt HS Pp}
variable {W v : ℕ} {M : Type*} [AddGroup M]

variable (E : M → IStmt HS Pp)

/-- Code of one gate (labels `lab` before it), register counts. -/
def gateCF (lab : Fin W → M) (g : NGate (Fin W) (ZMod 2) M) : IStmt HS Pp :=
  .seq (KIL.seqL ((touchL g.2.1).map fun (w : Fin W) => applyCF (w : ℕ) (E (g.2.2 - lab w))))
    (fcode (rebuildCU W (srcT g)))

/-- Code of the network, register counts: the gates, then the final frame changes.  It depends
only on `E`, the gates and the labels, not on the stream length. -/
def netCF (Mout : Fin W → M) : (Fin W → M) → List (NGate (Fin W) (ZMod 2) M) → IStmt HS Pp
  | lab, [] => KIL.seqL ((List.finRange W).map fun (w : Fin W) => applyCF (w : ℕ) (E (Mout w - lab w)))
  | lab, g :: gs => .seq (gateCF E lab g) (netCF Mout (labStep lab g) gs)

variable (cE : M → ℕ)

/-- Cost of the register-count network code. -/
def netCIF (V Bb : ℕ) (Mout : Fin W → M) : (Fin W → M) → List (NGate (Fin W) (ZMod 2) M) → ℕ
  | lab, [] => ((List.finRange W).map fun (w : Fin W) => cE (Mout w - lab w) + cCF V Bb).sum + 1
  | lab, g :: gs => ((touchL g.2.1).map fun (w : Fin W) => cE (g.2.2 - lab w) + cCF V Bb).sum +
      1 + cRebuildF W V Bb + netCIF V Bb Mout (labStep lab g) gs

variable {E cE}

end Net

end NLogN.KFD

end

section
-- Source module: Solutions.KFCPass
/-!
# KFC passes over a mark-delimited accumulator

Stack roles (all below 56, disjoint from the size parameters on `1 … 4`):
`sA` accumulator, `sT` temporary, `sO` output, `(sCB,sDB)` unary counter `B`,
`(sCM,sDM)` counter `M`, `(sCQ,sDQ)` counter `q`, `(sCP,sDP)` counter `P`, `sK, sK2` scratch.

The accumulator holds `encF g 0 n = g 0 ++ ▮ ++ g 1 ++ ▮ ++ …` (groups terminated by marks).
-/

set_option linter.unusedSimpArgs false
set_option linter.unnecessarySeqFocus false

namespace NLogN

open IntMul

namespace KFC

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

abbrev sX : Fin HS := 0
abbrev sA : Fin HS := 15
abbrev sT : Fin HS := 16
abbrev sO : Fin HS := 17
abbrev sCB : Fin HS := 18
abbrev sDB : Fin HS := 19
abbrev sCM : Fin HS := 20
abbrev sDM : Fin HS := 21
abbrev sCQ : Fin HS := 22
abbrev sDQ : Fin HS := 23
abbrev sCP : Fin HS := 24
abbrev sDP : Fin HS := 25
abbrev sK : Fin HS := 26
abbrev sK2 : Fin HS := 27

/-! ## Group lists given by functions -/

def encF (f : ℕ → List ℤ) (a n : ℕ) : List Item := enc ((List.range' a n).map f)
def flF (f : ℕ → List ℤ) (a n : ℕ) : List ℤ := (List.range' a n).flatMap f

@[simp] theorem encF_zero (f : ℕ → List ℤ) (a : ℕ) : encF f a 0 = [] := rfl
@[simp] theorem flF_zero (f : ℕ → List ℤ) (a : ℕ) : flF f a 0 = [] := rfl

/-! ## Forward zip column: `mvW A T; unmark A; mvB X T CB DB; mark T` -/

def zcol : IStmt HS Pp :=
  .seq (mvW sA NLogN.KFC.sT) (.seq (.op (.unmark sA)) (.seq (mvB sX NLogN.KFC.sT sCB sDB) (.op (.mark NLogN.KFC.sT))))

/-- Forward zip pass over `n` columns. -/
def zipA : IStmt HS Pp := uloop sCM sDM zcol

/-! ## Moving the groups back: `unmark T; mark A; mvW T A` -/

def bcol : IStmt HS Pp := .seq (.op (.unmark NLogN.KFC.sT)) (.seq (.op (.mark sA)) (mvW NLogN.KFC.sT sA))

/-- Move the `n` groups back from `T` to `A`. -/
def backB : IStmt HS Pp := uloop sCM sDM bcol

/-! ## Initialising and flushing the accumulator -/

def initA : IStmt HS Pp := uloop sCM sDM (.op (.mark sA))

def fcol : IStmt HS Pp := .seq (mvW sA NLogN.KFC.sO) (.op (.unmark sA))

def flushA : IStmt HS Pp := uloop sCM sDM fcol

end KFC

end NLogN

end

section
-- Source module: Solutions.KFCFwd
/-!
# KFC: fine-cost paper Lemma 3, forward direction `[P, q, M, B] → [P, M, q, B]`

`fwdF` computes the same denotation as `KIL.fwdP` (`pull [P,q,M,B] [P,M,q,B] gS 0 xs`), but every
per-element step is an item move with `O(1)` overhead: rows/columns/records are counted with unary
counters (marks), groups in the accumulator are delimited by marks.  The cost is
`(2q+16)·kap w·V + cUL P + cUL q + cUL M + cUL B + 1000 ≤ (2q+67)·kap w·V` (see `KFCFwdMain`)
where `w` bounds the bit size of the **data** entries only and `kap w = 76 w + 200`;
`cUL n = 1000(n+1) + 300·zsz n` is the one-time linear binary-to-unary conversion
(`KFCUnary.toUL`) of the four size parameters.
-/

set_option linter.unusedSimpArgs false
set_option linter.unnecessarySeqFocus false

namespace NLogN

open IntMul

namespace KFC

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

/-! ## One row: zip into the accumulator, then move the groups back -/

def zrow : IStmt HS Pp := .seq zipA backB

/-! ## The rows of one block -/

/-- Group `m` of the accumulator after `k` rows (records `rr y m`, `y < k`). -/
def grpF (rr : ℕ → ℕ → List ℤ) (k m : ℕ) : List ℤ := (List.range k).flatMap fun y => rr y m

def rowsF : IStmt HS Pp := uloop sCQ sDQ zrow

/-- The block body in row order. -/
def rowsOf (rr : ℕ → ℕ → List ℤ) (M : ℕ) (y : ℕ) : List ℤ := flF (rr y) 0 M

/-! ## One block: init, rows, flush -/

def blkF : IStmt HS Pp := .seq initA (.seq rowsF flushA)

def cBlk (q M B w : ℕ) : ℕ := (2 * q + 8) * kap w * (q * M * B)

/-! ## All blocks -/

def allF : IStmt HS Pp := uloop sCP sDP blkF

end KFC

end NLogN

end

section
-- Source module: Solutions.KFCFwdMain
/-!
# KFC: the self-contained fine-cost forward field move `fwdF`

`fwdF_spec`: on `t` with `ints xs ++ R` on stack `0` and `P, q, M, B` on stacks `1 … 4`,
`fwdF` ends in `t.upd 0 (ints (pull [P,q,M,B] [P,M,q,B] gS 0 xs) ++ R)` (all scratch stacks
restored) with cost `≤ cFwdF P q M B w = (2q+16)·kap w·V + cUL P + cUL q + cUL M + cUL B + 1000`,
`V = P·q·M·B`, `w` a bound on the bit size of the entries of `xs` only.
-/

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

namespace NLogN

open IntMul AModel

namespace KFC

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

/-! ## Prologue / epilogue of a unary counter pair -/

/-- Push sentinels `0` on `C`, `D`, then `n` marks on `C` (`n` = top of `A`). -/
def pro (A C D : Fin HS) : IStmt HS Pp :=
  .seq (.op (.const C 0)) (.seq (.op (.const D 0)) (toUL A C))

/-- Drop the marks and the sentinels of a counter pair. -/
def epi (C D : Fin HS) : IStmt HS Pp := .seq (clrM C) (.seq (.op (.drop C)) (.op (.drop D)))

/-- Wrap a body with the counter pair `(C, D)` holding `n` = top of `A`. -/
def wrap (A C D : Fin HS) (b : IStmt HS Pp) : IStmt HS Pp := .seq (pro A C D) (.seq b (epi C D))

/-! ## The core: blocks, then copy the output back -/

def coreF : IStmt HS Pp :=
  .seq (.op (.mark NLogN.KFC.sO)) (.seq (.op (.mark NLogN.KFC.sT)) (.seq allF (.seq (mvW NLogN.KFC.sO sX)
    (.seq (.op (.unmark NLogN.KFC.sO)) (.op (.unmark NLogN.KFC.sT))))))

/-! ## The program -/

def fwdF : IStmt HS Pp :=
  NLogN.KFC.wrap 4 sCB sDB (NLogN.KFC.wrap 3 sCM sDM (NLogN.KFC.wrap 2 sCQ sDQ (NLogN.KFC.wrap 1 sCP sDP coreF)))

def cFwdF (P q M B w : ℕ) : ℕ :=
  (2 * q + 16) * kap w * (P * q * M * B) + cUL P + cUL q + cUL M + cUL B + 1000

end KFC

end NLogN

end