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

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

import Definitions.Def_IntMul_Kappa_K_Defs11

set_option autoImplicit false
set_option relaxedAutoImplicit false
set_option maxHeartbeats 1000000

section
-- Source module: Solutions.KIL4Shift
/-!
# KIL4: the uniform Lemma 10 shift body

`shFracU k d nst neg` depends only on the fixed fraction `k / d`, the sign and the number `nst`
of offset-table stages.  It reads every size parameter from the parameter stacks:

* `22` `Q`, `23` `-Q`, `24` `vol A`, `25` `vol C`, `26` `vol A·(Q·vol C)`, `27` `Q·vol C`;
* the stage stacks `34 … 37` (radix, control, suffix volume, radix·suffix volume of every table
  stage, last stage on top), consumed by the table generator `genPU`.

`shFracU_spec`: for any stage list `ac` whose table is the table of `A.zip (csOf Λ c A)`, the body
computes `shiftArr Λ i (fracU Q k d neg) c 0 xs`, exactly as KIL3's `shFrac`.
-/

set_option linter.unusedSimpArgs false
set_option linter.unnecessarySeqFocus false

namespace NLogN.KIL

open IntMul RS AModel AModelL12

/-- The uniform `shP`: table from the stage stacks, then `rotP` with parameters from registers. -/
def shPU (nst : ℕ) : RS :=
  .seq (genPU nst) (.seq (.push 1 (.reg sPA)) (.seq (.push 2 (.reg sQQ)) (.seq (.push 4 (.reg sPC))
    (.seq rotP (.seq (.pop 1) (.seq (.pop 2) (.pop 4)))))))

/-- The uniform `shN`. -/
def shNU (nst : ℕ) : RS :=
  .seq (genPU nst) (.seq (.push 1 (.reg sPA)) (.seq (.push 2 (.reg sQQ)) (.seq (.push 4 (.reg sPC))
    (.seq NLogN.KIL.negP (.seq (.pop 1) (.seq (.pop 2) (.pop 4)))))))

/-- Cost bound of `shPU`/`shNU`. -/
def cShiftUU (nst V Bb : ℕ) : ℕ := 10 ^ 9 * (nst + 1) * (V + 1) * (Bb + 2) + 2 * 10 ^ 8 * (Bb + 2) ^ 2

/-- The stage stacks hold the stage list `ac`. -/
structure StgOK (t : IState HS) (Q : ℕ) (ac : List (ℕ × ℕ)) (LA LC LM LAM : List Item) : Prop where
  hA : t.G sSA = ints (ac.reverse.map fun p => (p.1 : ℤ)) ++ LA
  hC : t.G sSC = ints (ac.reverse.map fun p => (p.2 : ℤ)) ++ LC
  hM : t.G sSM = ints (stgM ac.reverse []) ++ LM
  hAM : t.G sSAM = ints (stgAM ac.reverse []) ++ LAM

section Body

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

/-- The post state of the uniform table shifts: stack `0` replaced, stage stacks consumed. -/
def postSt (t : IState HS) (ys : List ℤ) (R LA LC LM LAM : List Item) : IState HS :=
  ((((t.upd 0 (ints ys ++ R)).upd sSA LA).upd sSC LC).upd sSM LM).upd sSAM LAM

/-- **The uniform shift body** `y ↦ ±k·d⁻¹·y` controlled shift. -/
def shFracU (k d nst : ℕ) (neg : Bool) : RS :=
  .seq (.seq (fdivPU d) (fmulPU k)) (if neg then shNU nst else shPU nst)

def cShiftFU (k d nst V Bb : ℕ) : ℕ :=
  ((k + 1) ^ 2 + (d + 1) ^ 2) * 10 ^ 9 * (V + 1) * (Bb + 2) ^ 2 + cShiftUU nst V Bb

end Body

end NLogN.KIL

end

section
-- Source module: Solutions.KIL4Env
/-!
# KIL4: pushing computed parameters

* `Ex.congrG`: evaluation, side conditions, bounds and cost of an expression only depend on the
  stacks it reads.
* `pushes l` pushes the values of the expressions of `l` (each reading only stacks not pushed by
  `l`), `pops l` pops; `Sp.pushes`, `Sp.pops`.
* `powE`, `prodE`: products of parameters; `stgPush`: the pushes of the table stages
  (radix, control, suffix volume, radix·suffix volume), `stgPush_sp`: they establish `StgOK`.
-/

set_option linter.unusedSimpArgs false
set_option linter.unnecessarySeqFocus false

namespace NLogN.KIL

open IntMul RS AModel

/-! ## Pushes and pops -/

def pushes : List (Fin HS × Ex HS) → RS
  | [] => .skip
  | (i, e) :: l => .seq (.push i e) (pushes l)

/-- Pops in reverse order: `popsR [i₁, …, iₙ]` pops `iₙ, …, i₁`. -/
def popsR : List (Fin HS) → RS
  | [] => .skip
  | i :: l => .seq (popsR l) (.pop i)

/-- The state after pushing the values `l` (in order). -/
def pushSt (t : IState HS) : List (Fin HS × ℤ) → IState HS
  | [] => t
  | (i, v) :: l => pushSt (t.upd i (.int v :: t.G i)) l

def pushCost (Bb : ℕ) (l : List (Fin HS × Ex HS)) : ℕ :=
  (l.map fun p => (p.2.size + 1) * 10 ^ 7 * (Bb + 2) ^ 2).sum + 1

/-- A push list: every expression is clean, reads no pushed stack, is defined and bounded. -/
def PushOK (t : IState HS) (Bb : ℕ) (l : List (Fin HS × Ex HS)) : Prop :=
  ∀ p ∈ l, p.1.val < 56 ∧ p.2.Clean ∧ (∀ r ∈ p.2.regs, r ∉ l.map Prod.fst) ∧ p.2.ok t ∧
    p.2.Bnd t Bb

/-! ## Products of positive parameters -/

/-- Positive products of registers and literals. -/
def PosE (s : IState HS) : Ex HS → Prop
  | .reg i => s.IsInt i ∧ 1 ≤ s.top i
  | .lit c => 1 ≤ c
  | .mul a b => PosE s a ∧ PosE s b
  | _ => False

/-- Product of a list of expressions. -/
def prodE : List (Ex HS) → Ex HS
  | [] => .lit 1
  | e :: l => .mul e (prodE l)

/-- `e ^ n` as a product. -/
def powE (e : Ex HS) (n : ℕ) : Ex HS := prodE (List.replicate n e)

/-! ## The table stage pushes -/

/-- Pushes of the stages `l` (radix and control expressions), first stage first. -/
def stgPush : List (Ex HS × Ex HS) → List (Fin HS × Ex HS)
  | [] => []
  | (a, c) :: l => (sSA, a) :: (sSC, c) :: (sSM, prodE (l.map Prod.fst)) ::
      (sSAM, .mul a (prodE (l.map Prod.fst))) :: stgPush l

/-- Suffix volumes. -/
def sufV (suf : List (ℕ × ℕ)) : List (ℕ × ℕ) → List ℤ
  | [] => []
  | _ :: ac => (vol ((ac ++ suf).map Prod.fst) : ℤ) :: sufV suf ac

def sufAV (suf : List (ℕ × ℕ)) : List (ℕ × ℕ) → List ℤ
  | [] => []
  | x :: ac => ((x.1 * vol ((ac ++ suf).map Prod.fst) : ℕ) : ℤ) :: sufAV suf ac

/-- Evaluation of a push list. -/
def evalP (t : IState HS) (l : List (Fin HS × Ex HS)) : List (Fin HS × ℤ) :=
  l.map fun p => (p.1, p.2.eval t)

end NLogN.KIL

end

section
-- Source module: Solutions.KIL4Ctrl
/-!
# KIL4: run-time controls `p·q⁻¹ mod N`

For a fixed fraction `p / q` (`p ∈ ℤ`, `q ≥ 1`), `ctrlE p q` is an expression reading only the
modulus `N` (stack `22`) whose value is the representative in `[0, N)` of `p·q⁻¹` in `ZMod N`
(when `q` is prime to `N`): with `r = N mod q` and the compile-time table `t_r` (`q ∣ r·t_r + 1`),
`x = (N·t_r + 1) / q` is an inverse of `q` mod `N`, and the value is
`(p·x + |p|·x·N) mod N`.  Cost `O(q·(Bb + 2)²)`.
-/

set_option linter.unusedSimpArgs false
set_option linter.unnecessarySeqFocus false

namespace NLogN.KIL

open IntMul RS AModel

/-- The least `t < q` with `q ∣ r·t + 1` (`0` if none). -/
def invT (q r : ℕ) : ℕ := ((List.range q).find? (fun t => (r * t + 1) % q = 0)).getD 0

/-- Selection by the value of `e` among `vs` (indices from `k`). -/
def selL (e : Ex HS) : ℕ → List ℤ → Ex HS
  | _, [] => .lit 0
  | k, v :: vs => .sel (.add e (.lit (-((k : ℤ) + 1)))) (.lit v) (selL e (k + 1) vs)

/-- `p·q⁻¹` in `ZMod N`. -/
def fracZ (N : ℕ) (p : ℤ) (q : ℕ) : ZMod N := (p : ZMod N) * (q : ZMod N)⁻¹

def tTab (q : ℕ) : List ℤ := (List.range q).map fun r => (invT q r : ℤ)

/-- An inverse of `q` modulo `N` (stack `22`). -/
def invE (q : ℕ) : Ex HS :=
  .ndiv (.add (.mul (.reg sQQ) (selL (.nmod (.reg sQQ) (.lit q)) 0 (tTab q))) (.lit 1)) (.lit q)

/-- The control `p·q⁻¹ mod N` (representative in `[0, N)`). -/
def ctrlE (p : ℤ) (q : ℕ) : Ex HS :=
  .nmod (.add (.mul (.lit p) (invE q)) (.mul (.mul (.lit (p.natAbs : ℤ)) (invE q)) (.reg sQQ)))
    (.reg sQQ)

/-- The size bound needed by `ctrlE p q` at modulus `N`. -/
def ctrlBnd (p : ℤ) (q N : ℕ) : ℕ := 2 * (p.natAbs + 1) * (q + 1) * (N + 1) ^ 2

end NLogN.KIL

end

section
-- Source module: Solutions.KIL4Step
/-!
# KIL4: one uniform Lemma 10 step with run-time parameters

`stepG k d neg pa pc stg` is a fixed program text built from *expressions*: `pa`, `pc` compute the
outer volumes `vol A`, `vol C`, and `stg` lists the (radix, control) expressions of the offset
table stages.  The expressions may read the modulus register `sQQ` and any descriptor register
(they must not read the parameter or stage stacks they are pushed onto).  The program

1. pushes the four size parameters and the stage list,
2. runs the uniform shift body `shFracU`,
3. pops the size parameters again (the stage stacks are consumed by the body).

`stepG_spec`: if the stages evaluate to a stage list `ac` whose table is the table of
`A.zip (csOf Λ c A)` and the volume expressions evaluate to `vol A`, `vol C`, the step replaces the
array on stack `0` by `shiftArr Λ i (fracU Q k d neg) c 0 xs` and leaves every other stack
unchanged.
-/

set_option linter.unusedSimpArgs false
set_option linter.unnecessarySeqFocus false

namespace NLogN.KIL

open IntMul RS AModel AModelL12

/-- The size-parameter pushes. -/
def prmL (pa pc : Ex HS) : List (Fin HS × Ex HS) :=
  [(sPA, pa), (sPC, pc), (sVV, .mul pa (.mul (.reg sQQ) pc)), (sQB, .mul (.reg sQQ) pc)]

/-- **The uniform step**: a fixed program in the fraction `k/d`, the sign and the expressions. -/
def stepG (k d : ℕ) (neg : Bool) (pa pc : Ex HS) (stg : List (Ex HS × Ex HS)) : RS :=
  .seq (pushes (prmL pa pc ++ stgPush stg))
    (.seq (shFracU k d stg.length neg) (popsR [sPA, sPC, sVV, sQB]))

def cStepG (k d nst V Bb : ℕ) (pa pc : Ex HS) (stg : List (Ex HS × Ex HS)) : ℕ :=
  pushCost Bb (prmL pa pc ++ stgPush stg) + cShiftFU k d nst V Bb + (4 * (2 * Bb + 2) + 1)

section Body

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

end Body

end NLogN.KIL

end

section
-- Source module: Solutions.KIL4Coll
/-!
# KIL4: collapsed offset tables of the level layout

For the level layout `(hdLayout A B Cs m N).cons2 P L` and a shift of field `H_a` (`a < m`) or
`D_b` (`m + b`), the offset table of the prefix only depends on

* `vol (P :: L :: A)` (one zero-control stage),
* the controls of the earlier `H` fields (radix `N` each),
* `vol B` (one zero-control stage, `D` case only),
* the controls of the earlier `D` fields.

So a fixed stage-expression list, reading `vol (P :: L :: A)`, `vol B` and `N` from registers,
produces the right table at every level.
-/

set_option linter.unusedSimpArgs false
set_option linter.unnecessarySeqFocus false

namespace NLogN.KIL

open IntMul RS AModel AModelL12

section Coll

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

end Coll

section Lvl

variable {m : ℕ}

/-- Control values as natural numbers, indexed by `ℕ`. -/
def cv {N : ℕ} (c : Fin (m + m) → ZMod N) (j : ℕ) : ℕ :=
  if h : j < m + m then (c ⟨j, h⟩).val else 0

end Lvl

end NLogN.KIL

end

section
-- Source module: Solutions.KIL4StepH
/-!
# KIL4: the uniform level step `stepU`

A *shift template* `TSh m` fixes the active field `i`, the fraction `k/d` of the unit, its sign and
the rational controls `cs j = (p_j, q_j)` (control `j` is `p_j · q_j⁻¹ mod N`).  The template does
not mention the modulus `N`, the level or any volume.

`stepU m rX T` is one fixed program.  It reads

* `N` from `sQQ` (and `-N` from `sNQ`),
* `vol (P :: L :: A)` from the register `rX` (`rXs` or `rXb`),
* `vol B` from `rVB` and `vol Cs` from `rVC`,

builds the size parameters and the collapsed offset-table stages, computes every control
`p_j q_j⁻¹ mod N` at run time (`ctrlE`), and runs the uniform shift body.

`stepU_spec`: on the level layout `(hdLayout A B Cs m N).cons2 P L` it computes
`(T.eval N).arr`, the array action of the shift step the template denotes.
-/

set_option linter.unusedSimpArgs false
set_option linter.unnecessarySeqFocus false

namespace NLogN.KIL

open IntMul RS AModel AModelL12

/-- Small descriptor registers: `vol (P :: L' :: A)`, `vol (P :: W·L' :: A)`, `vol B`, `vol Cs`. -/
abbrev rXs : Fin HS := 38
abbrev rXb : Fin HS := 39
abbrev rVB : Fin HS := 40
abbrev rVC : Fin HS := 41

/-! ## Push-list hygiene -/

/-- Registers an expression of the step may read. -/
def RdOK (r : Fin HS) : Prop := r = sQQ ∨ (38 ≤ r.val ∧ r.val < 56)

structure GoodE (t : IState HS) (Bb : ℕ) (e : Ex HS) : Prop where
  rd : ∀ r ∈ e.regs, RdOK r
  ok : e.ok t
  bnd : e.Bnd t Bb

/-- A stage triple is good: positive radix reading allowed registers, good control. -/
def StgE (t : IState HS) (Bb : ℕ) (p : Ex HS × Ex HS) : Prop :=
  PosE t p.1 ∧ (∀ r ∈ p.1.regs, RdOK r) ∧ GoodE t Bb p.2

/-! ## Templates -/

/-- A shift template: field, unit `±k/d` and rational controls. -/
structure TSh (m : ℕ) where
  i : Fin (m + m)
  k : ℕ
  d : ℕ
  neg : Bool
  cs : ℕ → ℤ × ℕ

variable {m : ℕ}

/-- The controls of a template modulo `N`. -/
def TSh.ctrl (T : TSh m) (N : ℕ) : Fin (m + m) → ZMod N :=
  fun j => if j < T.i then fracZ N (T.cs j).1 (T.cs j).2 else 0

/-- The step a template denotes modulo `N`. -/
def TSh.eval (T : TSh m) (N : ℕ) : Step (m + m) N :=
  .shift T.i (fracU N T.k T.d T.neg) (T.ctrl N)

/-- The control stage `j` (radix `N`, control computed at run time). -/
def ctE (T : TSh m) (j : ℕ) : Ex HS × Ex HS := (.reg sQQ, ctrlE (T.cs j).1 (T.cs j).2)

/-- Stage expressions of the collapsed table. -/
def stgT (rX : Fin HS) (T : TSh m) : List (Ex HS × Ex HS) :=
  if (T.i : ℕ) < m then (.reg rX, .lit 0) :: (List.range T.i).map (ctE T)
  else (.reg rX, .lit 0) :: ((List.range m).map (ctE T) ++
    (.reg rVB, .lit 0) :: (List.range (T.i - m)).map fun j => ctE T (m + j))

/-- Volume of the prefix before the active field. -/
def paT (rX : Fin HS) (T : TSh m) : Ex HS :=
  if (T.i : ℕ) < m then .mul (.reg rX) (powE (.reg sQQ) T.i)
  else .mul (.reg rX) (.mul (powE (.reg sQQ) m) (.mul (.reg rVB) (powE (.reg sQQ) (T.i - m))))

/-- Volume of the suffix after the active field. -/
def pcT (T : TSh m) : Ex HS :=
  if (T.i : ℕ) < m then
    .mul (powE (.reg sQQ) (m - 1 - T.i)) (.mul (.reg rVB) (.mul (powE (.reg sQQ) m) (.reg rVC)))
  else .mul (powE (.reg sQQ) (m + m - 1 - T.i)) (.reg rVC)

/-- **The uniform level step.**  A fixed program in `m`, the register `rX` and the template. -/
def stepU (m : ℕ) (rX : Fin HS) (T : TSh m) : RS :=
  stepG T.k T.d T.neg (paT rX T) (pcT T) (stgT rX T)

/-- Template side conditions (independent of the level up to the coprimality with `N`). -/
structure TOK (T : TSh m) (N Bb : ℕ) : Prop where
  k1 : 1 ≤ T.k
  kN : T.k ≤ N
  d1 : 1 ≤ T.d
  dN : T.d ≤ N
  ck : Nat.Coprime N T.k
  cd : Nat.Coprime N T.d
  cs : ∀ j < (T.i : ℕ), 1 ≤ (T.cs j).2 ∧ Nat.Coprime N (T.cs j).2 ∧
    zsz (ctrlBnd (T.cs j).1 (T.cs j).2 N) ≤ Bb

end NLogN.KIL

end

section
-- Source module: Solutions.KIL4Steps
/-!
# KIL4: uniform step lists and the uniform `sub` / `flip` phases

`stepsU m rX Ts` runs a list of shift templates one after another (a fixed program for a fixed
template list).  `stepsU_spec`: on the level layout it computes `runArr` of the denoted steps.

The `D ← D − H` and `D ← H − D` phases of every level (`subSteps`, `flipSteps`) are denoted by
fixed template lists `subT m`, `flipT m` at **every** modulus `N` (`subT_eval`, `flipT_eval`).
Hence `subU m rX`, `flipU m rX` are single programs serving all levels (`subU_spec`, `flipU_spec`).
-/

set_option linter.unusedSimpArgs false
set_option linter.unnecessarySeqFocus false

namespace NLogN.KIL

open IntMul RS AModel AModelL12

variable {m : ℕ}

/-- Run templates in order. -/
def stepsU (m : ℕ) (rX : Fin HS) : List (TSh m) → RS
  | [] => .skip
  | T :: Ts => .seq (stepU m rX T) (stepsU m rX Ts)

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

/-! ## The `sub` and `flip` templates -/

/-- Template of `dhStep u c j` with `u = ±1` (`neg`) and `c ∈ ℤ`. -/
def dhT (neg : Bool) (cc : ℤ) (j : Fin m) : TSh m where
  i := Fin.natAdd m j
  k := 1
  d := 1
  neg := neg
  cs := fun l => if l = (j : ℕ) then (cc, 1) else (0, 1)

/-- Templates of `subSteps`. -/
def subT (m : ℕ) : List (TSh m) := (List.finRange m).map (dhT false (-1))

/-- Templates of `flipSteps`. -/
def flipT (m : ℕ) : List (TSh m) := (List.finRange m).map (dhT true 1)

/-- The uniform `sub` phase (one program for every level). -/
def subU (m : ℕ) (rX : Fin HS) : RS := stepsU m rX (subT m)

/-- The uniform `flip` phase. -/
def flipU (m : ℕ) (rX : Fin HS) : RS := stepsU m rX (flipT m)

/-- Size condition for the `±1` controls. -/
def SubBnd (N Bb : ℕ) : Prop :=
  zsz ((ctrlBnd 1 1 N : ℕ) : ℤ) ≤ Bb ∧ zsz ((ctrlBnd (-1) 1 N : ℕ) : ℤ) ≤ Bb ∧
    zsz ((ctrlBnd 0 1 N : ℕ) : ℤ) ≤ Bb

/-! ## Per-step cost: a template constant times `(V + 1)(Bb + 2)²` -/

/-- Total expression size of a push list. -/
def pushSz (l : List (Fin HS × Ex HS)) : ℕ := (l.map fun p => p.2.size + 1).sum

/-- The constant of a template step (depends only on the template and `m`). -/
def tK (rX : Fin HS) (T : TSh m) : ℕ :=
  pushSz (prmL (paT rX T) (pcT T) ++ stgPush (stgT rX T)) + (T.k + 1) ^ 2 + (T.d + 1) ^ 2 +
    (stgT rX T).length + 3

end NLogN.KIL

end

section
-- Source module: Solutions.KIL5Swap
/-!
# KIL5: the level-`0` interchange with run-time sizes, and the descriptor stacks

`prog C q 0 P L A B Cs` is the interchange of the two digits (fields of size `q`) of
`P :: L :: swShape A B Cs q`.  KIL3 realised it by `swapB`, two Lemma 3 field moves whose size
parameters are *literals*.  Here the size parameters are *expressions*, evaluated at run time:

* `withPush l p`: push the expressions `l`, run `p`, pop them (`Sp.withPush`);
* `swapE a q b c`: the two field moves with parameters read from expressions (`swapE_sp`):
  on `A0 ++ q :: (B ++ q :: Cs)` it computes the digit interchange when `a, q, b, c` evaluate to
  `vol A0, q, vol B, vol Cs`;
* the descriptor stacks `dJ … dT` (level, `P`, `L`, `vol A`, `vol B`, `vol Cs`, tower of moduli),
  `DescIn`, and `lev0U`, the uniform level-`0` code reading the descriptor (`lev0U_sp`).
-/

set_option linter.unusedSimpArgs false
set_option linter.unnecessarySeqFocus false

namespace NLogN.KIL

open IntMul RS AModel AModelL12

/-! ## Push, run, pop -/

/-- Push the expressions `l` (in order), run `p`, pop them again. -/
def withPush (l : List (Fin HS × Ex HS)) (p : RS) : RS :=
  .seq (pushes l) (.seq p (popsR (l.map Prod.fst)))

/-! ## Good expressions -/

/-! ## The two field moves with run-time parameters -/

/-- The level-`0` interchange with parameters `a = vol A0`, `q`, `b = vol B`, `c = vol Cs`. -/
def swapE (a q b c : Ex HS) : RS :=
  .seq (withPush [(1, a), (2, q), (3, .mul b q), (4, c)] fwdP)
    (withPush [(1, a), (2, q), (3, b), (4, .mul q c)] bwdP)

/-- Cost bound of `swapE`. -/
def cSwapE (a q b c : Ex HS) (qn V Bb : ℕ) : ℕ :=
  pushCost Bb [(1, a), (2, q), (3, .mul b q), (4, c)] + pushCost Bb [(1, a), (2, q), (3, b), (4, .mul q c)] +
    2 * (300000 * qn * (V + 1) * (Bb + 2) + 2 * 10 ^ 8 * (Bb + 2) ^ 2) +
    2 * (4 * (2 * Bb + 2) + 1)

/-! ## The descriptor stacks -/

/-- Descriptor stacks: level, `P`, `L`, `vol A`, `vol B`, `vol Cs`, and the tower of moduli
`q ^ m ^ k, …, q ^ m ^ 0` (top first). -/
abbrev dJ : Fin HS := 46
abbrev dP : Fin HS := 47
abbrev dL : Fin HS := 48
abbrev dA : Fin HS := 49
abbrev dB : Fin HS := 50
abbrev dC : Fin HS := 51
abbrev dT : Fin HS := 52

/-- The tower `q ^ m ^ k, …, q ^ m ^ 0`, top first. -/
def towerL (q m k : ℕ) : List ℤ := (List.range (k + 1)).reverse.map fun j => ((q ^ m ^ j : ℕ) : ℤ)

/-- The input descriptor of a call at level `k` on `P :: L :: swShape A B Cs (q ^ m ^ k)`. -/
structure DescIn (t : IState HS) (q m k P L : ℕ) (A B Cs : Shape) : Prop where
  j : ∃ R, t.G dJ = .int (k : ℤ) :: R
  p : ∃ R, t.G NLogN.KIL.dP = .int (P : ℤ) :: R
  l : ∃ R, t.G dL = .int (L : ℤ) :: R
  a : ∃ R, t.G dA = .int ((vol A : ℕ) : ℤ) :: R
  b : ∃ R, t.G dB = .int ((vol B : ℕ) : ℤ) :: R
  c : ∃ R, t.G dC = .int ((vol Cs : ℕ) : ℤ) :: R
  tw : ∃ R, t.G dT = ints (towerL q m k) ++ R

/-! ## The uniform level-`0` code -/

/-- The level-`0` interchange reading the descriptor: `vol (P :: L :: A)`, `q`, `vol B`,
`vol Cs`. -/
def lev0U : RS :=
  swapE (.mul (.reg NLogN.KIL.dP) (.mul (.reg dL) (.reg dA))) (.reg dT) (.reg dB) (.reg dC)

/-- Cost of `lev0U` (`q` is the field size, `V` the volume). -/
def cLev0 (q V Bb : ℕ) : ℕ :=
  cSwapE (.mul (.reg NLogN.KIL.dP) (.mul (.reg dL) (.reg dA))) (.reg dT) (.reg dB) (.reg dC) q V Bb

end NLogN.KIL

end

section
-- Source module: Solutions.KIL5Tab
/-!
# KIL5: the stream table with run-time stream length

KIL3's table operations (`applyC`, `copyC`, `accC`, `buildC`, `rebuildC`, KIL3Tab) carry the
stream length `v` and the offsets `w · v` as literals.  Here the counts are read from registers:

* `rV` (`42`) holds `v`, `rWV` (`43`) holds `W · v`;
* the offset `w · v` (`w < W` a fixed wire) is computed into the temporary register `rT` (`44`)
  by one multiplication (`push rT (w · rV)`) and popped afterwards.

The element loops `dupK`, `zipK`, `zeroK`, `dropK` with register counts are obtained from the
literal ones by the literal-to-register simulation (`Sp.sim`, KIL4Sim).

`applyCU_runs` is relative to an invariant: the stream code only has to be correct on states
`t` with `Inv G0 t` (every protected stack `r ≥ 22`, `r ≠ 44`, holds `G0 r`).  This is what the
uniform procedure needs: its child calls read the descriptor and the level registers from these
stacks.
-/

set_option linter.unusedSimpArgs false
set_option linter.unnecessarySeqFocus false

namespace NLogN.KIL

open IntMul RS AModel

/-! ## Registers and the invariant -/

/-- The stream length `v`. -/
abbrev rV : Fin HS := 42
/-- `W · v`. -/
abbrev rWV : Fin HS := 43
/-- Temporary `w · v`. -/
abbrev rT : Fin HS := 44

/-- Protected stacks: the level registers, the descriptor and the KIL4 parameter stacks. -/
def Prot (r : Fin HS) : Prop := 22 ≤ r.val ∧ r.val ≠ 44

instance (r : Fin HS) : Decidable (Prot r) := by unfold Prot; infer_instance

/-- The invariant: every protected stack holds its value `G0 r`. -/
def Inv (G0 : Fin HS → List Item) (t : IState HS) : Prop := ∀ r, Prot r → t.G r = G0 r

/-- Facts about the register stacks of a managed state. -/
structure RegS (s : IState HS) (f : Fin HS → List ℤ) (W v : ℕ) : Prop where
  fV : f NLogN.KIL.rV = []
  fWV : f rWV = []
  fT : f NLogN.KIL.rT = []
  sV : ∃ L, s.G NLogN.KIL.rV = .int (v : ℤ) :: L
  sWV : ∃ L, s.G rWV = .int ((W * v : ℕ) : ℤ) :: L

/-! ## Moves and loops with register counts -/

def dupKE (X Y Z C : Fin HS) (n : Ex HS) : RS :=
  .rep C n (.seq (.push Y (.reg X)) (.seq (.push Z (.reg X)) (.pop X)))

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

def zeroKE (Y C : Fin HS) (n : Ex HS) : RS := .rep C n (.push Y (.lit 0))

def zipKE (A Z O C : Fin HS) (n : Ex HS) : RS :=
  .rep C n (.seq (.push O (mAdd 2 (.reg A) (.reg Z))) (.seq (.pop A) (.pop Z)))

/-! ## Applying a stream code in place -/

section Apply

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

/-- Park the streams before `w` (`w · v` items, counted in `rT`), run `code`, move them back. -/
def applyCU (w : ℕ) (code : IStmt HS Pp) : IStmt HS Pp :=
  .seq (RS.seq (.push NLogN.KIL.rT (.mul (.lit (w : ℤ)) (.reg NLogN.KIL.rV))) (NLogN.mvK 0 15 21 (.reg NLogN.KIL.rT))).code
    (.seq code (RS.seq (NLogN.mvK 15 0 21 (.reg NLogN.KIL.rT)) (.pop NLogN.KIL.rT)).code)

/-- The push of `w · v`. -/
def cPushT (Bb : ℕ) : ℕ := 4 * 10 ^ 7 * (Bb + 2) ^ 2

end Apply

/-! ## Copying, accumulating and rebuilding with register counts -/

/-- Cost unit of one register-count table operation. -/
def cCU (V Bb : ℕ) : ℕ := 300 * cOp V Bb + cPushT Bb + (2 * Bb + 2)

/-- Push a reversed copy of stream `w` on stack `16`. -/
def copyCU (w : ℕ) : RS :=
  .seq (.push NLogN.KIL.rT (.mul (.lit (w : ℤ)) (.reg NLogN.KIL.rV))) (.seq (NLogN.mvK 0 15 21 (.reg NLogN.KIL.rT))
    (.seq (dupKE 0 16 17 21 (.reg NLogN.KIL.rV)) (.seq (NLogN.mvK 17 0 21 (.reg NLogN.KIL.rV))
      (.seq (NLogN.mvK 15 0 21 (.reg NLogN.KIL.rT)) (.pop NLogN.KIL.rT)))))

/-- `acc ← acc ⊕ (stream w)ʳᵉᵛ` on the top `v` entries of stack `20`. -/
def accCU (w : ℕ) : RS :=
  .seq (copyCU w) (.seq (zipKE 20 16 19 21 (.reg NLogN.KIL.rV)) (NLogN.mvK 19 20 21 (.reg NLogN.KIL.rV)))

def accLU : List ℕ → RS
  | [] => .skip
  | w :: l => .seq (accCU w) (accLU l)

def buildCU (l : List ℕ) : RS := .seq (zeroKE 20 21 (.reg NLogN.KIL.rV)) (accLU l)

def buildsCU (T : ℕ → List ℕ) : ℕ → RS
  | 0 => .skip
  | n + 1 => .seq (buildsCU T n) (buildCU (T n))

/-- Replace the table by `w ↦ Σ_{w' ∈ T w} F w'` (`W` fixed, `v` and `W · v` from registers). -/
def rebuildCU (W : ℕ) (T : ℕ → List ℕ) : RS :=
  .seq (buildsCU T W) (.seq (dropKE 0 21 (.reg rWV)) (NLogN.mvK 20 0 21 (.reg rWV)))

def cRebuildU (W V Bb : ℕ) : ℕ := (W * (2 * W + 2) + 202) * cCU V Bb

end NLogN.KIL

end

section
-- Source module: Solutions.KIL5Net
/-!
# KIL5: the network of one level with run-time stream length

KIL3Net's `frames_runs`, `gateC_runs`, `netC_runs` re-proved for the register-count table code
(`applyCU`, `rebuildCU`, KIL5Tab).  The stream codes only have to be correct on states satisfying
the invariant `Inv G0` (the protected stacks `≥ 22`, `≠ 44`, hold `G0`): the network keeps these
stacks fixed, so a code that reads the level registers / the descriptor (e.g. a child call) may
rely on them.

Cost: every frame change costs `cE X + cCU (W·v) Bb`, every gate `cRebuildU W (W·v) Bb + 1`.
-/

set_option linter.unusedSimpArgs false
set_option linter.unnecessarySeqFocus false

namespace NLogN.KIL

open IntMul RS AModel

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 gateCU (lab : Fin W → M) (g : NGate (Fin W) (ZMod 2) M) : IStmt HS Pp :=
  .seq (seqL ((touchL g.2.1).map fun (w : Fin W) => applyCU (w : ℕ) (E (g.2.2 - lab w))))
    (rebuildCU W (srcT g)).code

/-- 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 netCU (Mout : Fin W → M) : (Fin W → M) → List (NGate (Fin W) (ZMod 2) M) → IStmt HS Pp
  | lab, [] => seqL ((List.finRange W).map fun (w : Fin W) => applyCU (w : ℕ) (E (Mout w - lab w)))
  | lab, g :: gs => .seq (gateCU E lab g) (netCU Mout (labStep lab g) gs)

variable (cE : M → ℕ)

/-- Cost of the register-count network code. -/
def netCIU (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) + cCU V Bb).sum + 1
  | lab, g :: gs => ((touchL g.2.1).map fun (w : Fin W) => cE (g.2.2 - lab w) + cCU V Bb).sum +
      1 + cRebuildU W V Bb + netCIU V Bb Mout (labStep lab g) gs

variable {E cE}

end Net

end NLogN.KIL

end

section
-- Source module: Solutions.KIL5Tpl
/-!
# KIL5: the edge steps of the rational circuits are template-denoted

`TStep m` is a step template: a shift template (`TSh`) or a field transposition.  The rational
circuit `QCircuit.toCircuit Q` (`AModelRat`) takes as Lemma 10 data of a frame change `X` the
integer lifts mod `Q` of the rational Bruhat factors `ratFacQ A` of `A = D⁻¹ • X`
(`facD Q D`).  Reduced mod `N ∣ Q`, these lifts are the reductions `redQ N` of the fixed
rationals, so every step of `edgeSteps` is the evaluation at `N` of a template whose unit is
`±|num|/den` and whose controls are `num/den` of fixed rationals.

* `rsh i u c`: the shift template of the rational step `(i, u, c)`; `rsh_eval`.
* `QStep`: rational steps; `QStep.red N` reduces mod `N`; `QStep.tpl` is the template;
  `tpl_eval` (`(s.tpl).eval N = s.red N`).
* `shearQ E₁ Π E₂`: the rational `shearSteps`; `shearQ_red` (reduction commutes).
* `edgeTplQ C X`: the template list of a frame change; **`edgeSteps_toCircuit`**:
  `edgeSteps (C.toCircuit Q) N X = (edgeTplQ C X).map (TStep.eval N)` for every `N ∣ Q = q ^ b`
  with `1 < N`, whenever the prime `q` is good for `A` (`GoodAtQ`).
-/

set_option linter.unusedSimpArgs false
set_option linter.unnecessarySeqFocus false

namespace NLogN.KIL

open AModel AModelL12 AModelFac AModelRat

variable {m : ℕ}

/-! ## Step templates -/

/-- A step template: a shift template or a field transposition. -/
inductive TStep (m : ℕ) where
  | sh (T : TSh m)
  | sw (i j : Fin (m + m))

def TStep.eval {m : ℕ} (N : ℕ) : TStep m → Step (m + m) N
  | .sh T => T.eval N
  | .sw i j => .swap i j

/-! ## Reduction facts -/

/-! ## Rational shift templates -/

/-- The template of the rational shift `x_i ↦ u x_i + Σ_{j < i} c_j x_j`. -/
def rsh (i : Fin (m + m)) (u : ℚ) (c : Fin (m + m) → ℚ) : TSh m where
  i := i
  k := u.num.natAbs
  d := u.den
  neg := decide (u.num < 0)
  cs := fun j => if h : j < m + m then ((c ⟨j, h⟩).num, (c ⟨j, h⟩).den) else (0, 1)

/-- The reduction mod `N` of a rational shift. -/
def redSh (N : ℕ) (i : Fin (m + m)) (u : ℚ) (c : Fin (m + m) → ℚ) : Step (m + m) N :=
  .shift i (redQ N u) (fun j => if j < i then redQ N (c j) else 0)

/-! ## Rational steps -/

/-- A rational step: a shift `(i, u, c)` or a transposition. -/
inductive QStep (m : ℕ) where
  | sh (i : Fin (m + m)) (u : ℚ) (c : Fin (m + m) → ℚ)
  | sw (i j : Fin (m + m))

namespace QStep

def red (N : ℕ) : QStep m → Step (m + m) N
  | .sh i u c => redSh N i u c
  | .sw i j => .swap i j

def tpl : QStep m → TStep m
  | .sh i u c => .sh (rsh i u c)
  | .sw i j => .sw i j

/-- The inverse step: shift by `u⁻¹` with controls `-u⁻¹ c`. -/
def rev : QStep m → QStep m
  | .sh i u c => .sh i u⁻¹ (fun j => -(u⁻¹ * c j))
  | .sw i j => .sw i j

/-- Denominators (and the unit numerator) coprime to `N`. -/
def Good (N : ℕ) : QStep m → Prop
  | .sh i u c => u ≠ 0 ∧ Nat.Coprime u.den N ∧ Nat.Coprime u.num.natAbs N ∧
      ∀ j, j < i → Nat.Coprime (c j).den N
  | .sw _ _ => True

end QStep

/-- Inverse of a rational step list. -/
def revQ (l : List (QStep m)) : List (QStep m) := (l.map QStep.rev).reverse

/-! ## The rational Lemma 10 steps -/

/-- Row `a` of `E` on the `H` block. -/
def hcQ (E : QMat m) (a : Fin m) : Fin (m + m) → ℚ := Fin.addCases (fun b => E a b) (fun _ => 0)

/-- Row `a` of `E` on the `D` block. -/
def dcQ (E : QMat m) (a : Fin m) : Fin (m + m) → ℚ := Fin.addCases (fun _ => 0) (fun b => E a b)

/-- `H ← E H`. -/
def hQ (E : QMat m) : List (QStep m) :=
  (List.finRange m).reverse.map fun a => .sh (Fin.castAdd m a) (E a a) (hcQ E a)

/-- `D ← E D`. -/
def dQ (E : QMat m) : List (QStep m) :=
  (List.finRange m).reverse.map fun a => .sh (Fin.natAdd m a) (E a a) (dcQ E a)

/-- The pivot gadget. -/
def pivotQ (i j : Fin m) : List (QStep m) :=
  [.sh (Fin.natAdd m j) 1 (Pi.single (Fin.castAdd m i) 1), .sw (Fin.castAdd m i) (Fin.natAdd m j),
   .sh (Fin.natAdd m j) (-1) (Pi.single (Fin.castAdd m i) 1)]

/-- Pivots of an integer `0/1` matrix. -/
def pivotsZ (Pi : IMat m) : List (Fin m × Fin m) :=
  (List.finRange m ×ˢ List.finRange m).filter fun p => Pi p.1 p.2 = 1

def addQ (ps : List (Fin m × Fin m)) : List (QStep m) := ps.flatMap fun p => pivotQ p.1 p.2

/-- **The rational `shearSteps`.** -/
def shearQ (E₁ : QMat m) (Pi : IMat m) (E₂ : QMat m) : List (QStep m) :=
  dQ E₂ ++ revQ (hQ E₁) ++ addQ (pivotsZ Pi) ++ hQ E₁ ++ revQ (dQ E₂)

/-- Good factors mod `N`: denominators and diagonal numerators coprime to `N`. -/
def FGood (N : ℕ) (E : QMat m) : Prop :=
  (∀ i j, Nat.Coprime (E i j).den N) ∧ (∀ i, E i i ≠ 0) ∧ (∀ i, Nat.Coprime (E i i).num.natAbs N)

/-! ## Reduction of the rational steps -/

/-! ## The edge steps of the rational circuit -/

/-- The template list of a frame change `X` of the rational circuit (independent of `Q`). -/
noncomputable def edgeTplQ {R : Type*} [CommRing R] {W : ℕ} (C : QCircuit R m W) (X : IMat m) :
    List (TStep m) :=
  let A := ((C.den : ℚ)⁻¹ • ratM X)
  (shearQ (ratFacQ A).1 (piQ A) (ratFacQ A).2.2).map QStep.tpl

end NLogN.KIL

end

section
-- Source module: Solutions.KIL5Edge
/-!
# KIL5: the uniform frame change (template lists, child calls through the descriptor)

The frame change `Φ_X` of a level is the run of a template list `Ts` (KIL5Tpl):

* a shift template `T` is the fixed program `stepU m rXs T` (KIL4StepH), reading the modulus
  `N` and the volumes from the level registers;
* an interchange template `sw a b` is `swC m j0 a b`: push the child descriptor (level `k - 1`,
  `L / W`, `vol A · N^a`, `N^(m-1-a) · vol B · N^(b-m)`, `N^(m+m-1-b) · vol Cs`) computed from
  the level registers, call the uniform procedure `j0`, pop the descriptor.

Both codes depend only on `m`, `j0` and the template, not on the level, the modulus or the
sizes.  `edgeCU_runs`: the code of `Ts` computes `runC` of the denoted steps (any child function
whose call code is correct); `swC_runs`: the call wrapper realises a child given by a recursive
specification of the procedure `j0` on descriptor states.
-/

set_option linter.unusedSimpArgs false
set_option linter.unnecessarySeqFocus false

namespace NLogN.KIL

open IntMul RS AModel

/-! ## Level registers used by the call wrapper -/

/-- `L' = L / W`. -/
abbrev rLp : Fin HS := 45
/-- The saved top of the tower (the field size of the level). -/
abbrev rS : Fin HS := 53
/-- `vol A`. -/
abbrev rVA : Fin HS := 54
/-- The child level `k - 1`. -/
abbrev rJ1 : Fin HS := 55

/-- The child descriptor of the interchange of fields `a < m ≤ b`. -/
def chL (m a b : ℕ) : List (Fin HS × Ex HS) :=
  [(dJ, .reg rJ1), (dL, .reg rLp), (dA, .mul (.reg rVA) (powE (.reg sQQ) a)),
   (dB, .mul (powE (.reg sQQ) (m - 1 - a)) (.mul (.reg rVB) (powE (.reg sQQ) (b - m)))),
   (dC, .mul (powE (.reg sQQ) (m + m - 1 - b)) (.reg rVC))]

section Code

variable {Pp : ℕ}

/-- The child call of an interchange template. -/
def swC (m : ℕ) (j0 : Fin Pp) (a b : ℕ) : IStmt HS Pp :=
  .seq (pushes (NLogN.KIL.chL m a b)).code (.seq (.call j0) (popsR [dJ, dL, dA, dB, dC]).code)

/-- The code of one template. -/
def tCode (m : ℕ) (j0 : Fin Pp) : TStep m → IStmt HS Pp
  | .sh T => (stepU m rXs T).code
  | .sw a b => NLogN.KIL.swC m j0 (a : ℕ) (b : ℕ)

/-- The code of a frame change given by a template list. -/
def edgeCU (m : ℕ) (j0 : Fin Pp) (Ts : List (TStep m)) : IStmt HS Pp :=
  seqL (Ts.map (tCode m j0))

end Code

/-- Cost of one template step. -/
def tCost {m : ℕ} (cS : TSh m → ℕ) (cC : ℕ) : TStep m → ℕ
  | .sh T => cS T
  | .sw _ _ => cC

/-- Cost of a frame change. -/
def edgeCost {m : ℕ} (cS : TSh m → ℕ) (cC : ℕ) (Ts : List (TStep m)) : ℕ :=
  (Ts.map (tCost cS cC)).sum + 1

/-! ## The frame change -/

section Edge

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

end Edge

/-! ## The child call -/

section Call

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

/-- Cost of the call wrapper around a child of cost `c0`. -/
def cSwC (m a b Bb c0 : ℕ) : ℕ := pushCost Bb (NLogN.KIL.chL m a b) + (c0 + 2) + (5 * (2 * Bb + 2) + 1)

end Call

end NLogN.KIL

end

section
-- Source module: Solutions.KIL5Level
/-!
# KIL5: the body of one uniform level

The level body is one fixed code `bodyU C tpl j0` (it depends on the circuit, the template
function `tpl` and the procedure index `j0` only):

1. `D ← D − H`: `subU m rXb` (KIL4Steps);
2. the row split `bwdP` with the parameters `1, W, P · L', vol hd` pushed from the registers
   (`withPush (splitL W)`);
3. the network `netCU` (KIL5Net) whose frame changes are the template codes `edgeCU m j0 (tpl X)`
   (KIL5Edge);
4. the routing `rebuildCU W (permT C)`;
5. the row merge `fwdP`;
6. `D ← H − D`: `flipU m rXb`.

All sizes are read from the level registers (stacks `≥ 22`), whose contents are fixed by the
invariant `Inv G0` and described by `LevG G0 …`.  `bodyU_run`: the body computes `levelR`.
-/

set_option linter.unusedSimpArgs false
set_option linter.unnecessarySeqFocus false

namespace NLogN.KIL

open IntMul RS AModel AModelL12

/-! ## The row split / merge parameters from the registers -/

/-- Push `1, W, P · L', v / (P · L')` on stacks `1 … 4`. -/
def splitL (W : ℕ) : List (Fin HS × Ex HS) :=
  [(1, .lit 1), (2, .lit (W : ℤ)), (3, .mul (.reg NLogN.KIL.dP) (.reg rLp)),
   (4, .ndiv (.reg NLogN.KIL.rV) (.mul (.reg NLogN.KIL.dP) (.reg rLp)))]

def cSplit (W Bb : ℕ) : ℕ := pushCost Bb (NLogN.KIL.splitL W) + (4 * (2 * Bb + 2) + 1)

/-! ## The level registers -/

/-- The level registers (in the protected stacks `G0`) of a level with modulus `N` on
`P :: (W · L') :: hdShape A B Cs m N`. -/
structure LevG (G0 : Fin HS → List Item) (m W N P L' : ℕ) (A B Cs : Shape) : Prop where
  qq : ∃ L, G0 sQQ = .int (N : ℤ) :: L
  nq : ∃ L, G0 sNQ = .int (-(N : ℤ)) :: L
  xs : ∃ L, G0 rXs = .int ((vol (P :: L' :: A) : ℕ) : ℤ) :: L
  xb : ∃ L, G0 rXb = .int ((vol (P :: W * L' :: A) : ℕ) : ℤ) :: L
  vb : ∃ L, G0 rVB = .int ((vol B : ℕ) : ℤ) :: L
  vc : ∃ L, G0 rVC = .int ((vol Cs : ℕ) : ℤ) :: L
  v : ∃ L, G0 NLogN.KIL.rV = .int ((vol (P :: L' :: hdShape A B Cs m N) : ℕ) : ℤ) :: L
  wv : ∃ L, G0 rWV = .int ((W * vol (P :: L' :: hdShape A B Cs m N) : ℕ) : ℤ) :: L
  p : ∃ L, G0 NLogN.KIL.dP = .int (P : ℤ) :: L
  lp : ∃ L, G0 rLp = .int (L' : ℤ) :: L

/-! ## The body -/

section Body

variable {m W : ℕ} [NeZero W] (C : Circuit (ZMod 2) m W) {Pp : ℕ} {I : List Sym}
  {iprog : Fin Pp → IStmt HS Pp}

/-- **The uniform level body.** -/
def bodyU (tpl : IMat m → List (TStep m)) (j0 : Fin Pp) : IStmt HS Pp :=
  .seq (subU m rXb).code (.seq (withPush (NLogN.KIL.splitL W) bwdP).code
  (.seq (netCU (fun X => edgeCU m j0 (tpl X)) C.Mout C.Min C.gates)
  (.seq (rebuildCU W (NLogN.KIL.permT C)).code (.seq (withPush (NLogN.KIL.splitL W) fwdP).code (flipU m rXb).code))))

/-- Cost of a shift template on streams of volume `v`. -/
def cShT (v Bb : ℕ) (T : TSh m) : ℕ :=
  cStepG T.k T.d (stgT rXs T).length v Bb (paT rXs T) (pcT T) (stgT rXs T)

/-- Cost of the body (`cC`: cost of one child call wrapper). -/
def bodyCI (tpl : IMat m → List (TStep m)) (N P L' : ℕ) (A B Cs : Shape) (cC Bb : ℕ) : ℕ :=
  cStepsU rXb (vol (P :: W * L' :: hdShape A B Cs m N)) Bb (subT m) +
  2 * (cSplit W Bb + cMove 1 W (P * L') (vol (hdShape A B Cs m N)) Bb) +
  netCIU (fun X => edgeCost (cShT (vol (P :: L' :: hdShape A B Cs m N)) Bb) cC (tpl X))
    (W * vol (P :: L' :: hdShape A B Cs m N)) Bb C.Mout C.Min C.gates +
  cRebuildU W (W * vol (P :: L' :: hdShape A B Cs m N)) Bb +
  cStepsU rXb (vol (P :: W * L' :: hdShape A B Cs m N)) Bb (flipT m)

end Body

end NLogN.KIL

end

section
-- Source module: Solutions.KIL5Prog
/-!
# KIL5: the uniform procedure (one fixed program for every level)

`progU C tpl : Fin 1 → IStmt HS 1` is a single procedure.  It reads the descriptor
(`dJ` level, `dP`, `dL`, `dA`, `dB`, `dC`, the tower `dT`, KIL5Swap):

* level `0` (`dJ = 0`): the digit interchange `lev0U` (KIL5Swap);
* level `k + 1`: `entryU` (save the tower top in `rS`, pop it, push the level registers
  computed from the descriptor), the level body `bodyU C tpl 0` (KIL5Level, children by the call
  wrappers `swC`, KIL5Edge), `exitU` (pop the registers, restore the tower top).

The program depends only on the circuit `C` and the template function `tpl`.
-/

set_option linter.unusedSimpArgs false
set_option linter.unnecessarySeqFocus false

namespace NLogN.KIL

open IntMul RS AModel AModelL12

/-! ## Pushed states -/

/-! ## Entry and exit -/

/-- `vol hd` from the descriptor: `vol A · F · vol B · F · vol Cs`, `F = N ^ m`. -/
def hdE (m : ℕ) : Ex HS :=
  .mul (.reg dA) (.mul (powE (.reg dT) m) (.mul (.reg dB) (.mul (powE (.reg dT) m) (.reg dC))))

/-- The level registers computed from the descriptor (tower top already popped). -/
def entL (W m : ℕ) : List (Fin HS × Ex HS) :=
  [(sQQ, .reg dT), (sNQ, .neg (.reg dT)), (rLp, .ndiv (.reg dL) (.lit (W : ℤ))),
   (rJ1, .add (.reg dJ) (.lit (-1))), (rVA, .reg dA), (rVB, .reg dB), (rVC, .reg dC),
   (rXs, .ndiv (.mul (.reg NLogN.KIL.dP) (.mul (.reg dL) (.reg dA))) (.lit (W : ℤ))),
   (rXb, .mul (.reg NLogN.KIL.dP) (.mul (.reg dL) (.reg dA))),
   (NLogN.KIL.rV, .ndiv (.mul (.reg NLogN.KIL.dP) (.mul (.reg dL) (hdE m))) (.lit (W : ℤ))),
   (rWV, .mul (.reg NLogN.KIL.dP) (.mul (.reg dL) (hdE m)))]

def entryU (W m : ℕ) : RS := .seq (.push NLogN.KIL.rS (.reg dT)) (.seq (.pop dT) (pushes (entL W m)))

def exitU (W m : ℕ) : RS :=
  .seq (popsR ((entL W m).map Prod.fst)) (.seq (.push dT (.reg NLogN.KIL.rS)) (.pop NLogN.KIL.rS))

/-- The values pushed by the entry (level `k + 1`, modulus `N = q ^ m ^ k`). -/
def entV (W m N P L k : ℕ) (A B Cs : Shape) : List (Fin HS × ℤ) :=
  [(sQQ, (N : ℤ)), (sNQ, -(N : ℤ)), (rLp, ((L / W : ℕ) : ℤ)), (rJ1, (k : ℤ)),
   (rVA, ((vol A : ℕ) : ℤ)), (rVB, ((vol B : ℕ) : ℤ)), (rVC, ((vol Cs : ℕ) : ℤ)),
   (rXs, ((vol (P :: L / W :: A) : ℕ) : ℤ)), (rXb, ((vol (P :: W * (L / W) :: A) : ℕ) : ℤ)),
   (NLogN.KIL.rV, ((vol (P :: L / W :: hdShape A B Cs m N) : ℕ) : ℤ)),
   (rWV, ((W * vol (P :: L / W :: hdShape A B Cs m N) : ℕ) : ℤ))]

/-- The state after saving and popping the tower top. -/
def entS (t : IState HS) (F : ℕ) (T : List Item) : IState HS :=
  (t.upd NLogN.KIL.rS (.int (F : ℤ) :: t.G NLogN.KIL.rS)).upd dT T

def cEntry (W m Bb : ℕ) : ℕ := 2 * 10 ^ 7 * (Bb + 2) ^ 2 + (2 * Bb + 2) + pushCost Bb (entL W m)

def cExit (_W _m Bb : ℕ) : ℕ :=
  (11 * (2 * Bb + 2) + 1) + 2 * 10 ^ 7 * (Bb + 2) ^ 2 + (2 * Bb + 2)

/-! ## Cost bounds -/

section FrameLe

variable {Wr R M : Type*} [Fintype Wr] [DecidableEq Wr] [AddGroup M]

end FrameLe

section Costs

variable {m : ℕ}

/-- Sum of the template constants of the shifts of a template list. -/
def shK : List (TStep m) → ℕ
  | [] => 0
  | .sh T :: Ts => NLogN.KIL.tK rXs T + shK Ts
  | .sw _ _ :: Ts => shK Ts

/-- Number of interchanges (child calls) of a template list. -/
def swN : List (TStep m) → ℕ
  | [] => 0
  | .sh _ :: Ts => swN Ts
  | .sw _ _ :: Ts => swN Ts + 1

/-- Sum of the template constants of a list of shift templates. -/
def tKs (rX : Fin HS) (Ts : List (TSh m)) : ℕ := (Ts.map (NLogN.KIL.tK rX)).sum

end Costs

section LevCost

variable {m W : ℕ} [NeZero W] (C : Circuit (ZMod 2) m W) (tpl : IMat m → List (TStep m))

/-- Template constants of all shifts of all frame changes of one level. -/
def KSh : ℕ := frameSum (fun X => NLogN.KIL.shK (tpl X)) C.gates C.Min C.Mout

/-- Number of child calls of one level. -/
def SWn : ℕ := frameSum (fun X => swN (tpl X)) C.gates C.Min C.Mout

/-- The level constant. -/
def Klev : ℕ :=
  tKs rXb (subT m) * 10 ^ 9 + tKs rXb (flipT m) * 10 ^ 9 +
    2 * ((pushSz (NLogN.KIL.splitL W) + 1) * 10 ^ 7 + 10 + (300000 * W + 2 * 10 ^ 8)) +
    KSh C tpl * 10 ^ 9 + nFr C * (1 + 10 ^ 8) +
    C.gates.length * ((W * (2 * W + 2) + 202) * 10 ^ 8 + 1) +
    (W * (2 * W + 2) + 202) * 10 ^ 8 + (pushSz (entL W m) + 5) * 10 ^ 7 + 100

end LevCost

/-! ## The uniform procedure -/

/-- Bound of the child descriptor push sizes. -/
def chS (m : ℕ) : ℕ :=
  ((List.range (m + m)).map fun a => ((List.range (m + m)).map fun b => pushSz (NLogN.KIL.chL m a b)).sum).sum

/-- The per-call descriptor constant. -/
def Kd (m : ℕ) : ℕ := 10 ^ 7 * (chS m + 1)

/-- The cost of the uniform procedure (`K`: level constant, `Kd`: per-call constant, `S`: number
of child calls per level). The level-`0` term has the factor `q + 1` (the digit interchange of
`KILLemma3` moves fields of size `q`). -/
def ilCostU5 (W K Kd S q Bb : ℕ) : ℕ → ℕ → ℕ
  | 0, V => K * (q + 1) * ((V + 1) * (Bb + 2) ^ 2)
  | j + 1, V => K * ((V + 1) * (Bb + 2) ^ 2) +
      S * (ilCostU5 W K Kd S q Bb j (V / W) + Kd * (Bb + 2) ^ 2)

section Main

variable {m W : ℕ} [NeZero W] (C : Circuit (ZMod 2) m W) (tpl : IMat m → List (TStep m))

/-- One level `≥ 1`: entry, body, exit. -/
def levU : IStmt HS 1 := .seq (entryU W m).code (.seq (bodyU C tpl (0 : Fin 1)) (exitU W m).code)

/-- **The uniform procedure table** (one procedure). -/
def progU : Fin 1 → IStmt HS 1 :=
  fun _ => .ite dJ (fun c => c == .zero) lev0U.code (levU C tpl)

/-- The level side conditions at modulus `N`: on every frame change of the circuit the Lemma 10
steps are denoted by the templates and the shift templates are implementable; the `±1` controls
fit. -/
structure LvOK (N Bb : ℕ) : Prop where
  frames : FramesOK (fun X => AModel.edgeSteps C N X = (tpl X).map (TStep.eval N) ∧
    ∀ T, TStep.sh T ∈ tpl X → TOK T N Bb) C.gates C.Min C.Mout
  sub : SubBnd N Bb

end Main

end NLogN.KIL

end

section
-- Source module: Solutions.KIL6Pad
/-!
# KIL6: one-field padding and extraction as IL loops

* `padP`: on stack `0` an array `xs` of `X · (F · Y)` entries (`[X, F, Y]`, row-major); it
  produces `padL X F F' Y 0 xs` (every block of `F · Y` entries followed by `(F' - F) · Y`
  zeros).  Per block: move `F · Y` entries to the buffer stack `15` (`mvK`), push the zeros;
  finally move everything back (the two reversals cancel).
* `unpadP`: the inverse scan `unpadL X F F' Y` (move `F · Y`, drop `(F' - F) · Y`).

The four counts are read from the parameter registers `16 … 19` (pushed by the caller); the
loop counters are `20`, `21`.  Both cost `O((V' + 1)(Bb + 2))`, `V'` the padded volume
(`padP_sp`, `unpadP_sp`).
-/

set_option linter.unusedSimpArgs false
set_option linter.unnecessarySeqFocus false

namespace NLogN.KIL

open IntMul RS AModel AModelL12

/-! ## Loops with register counts -/

def cZeroR (n Bb : ℕ) : ℕ :=
  (1 + 1) * 2000 * (Bb + 2) + n * ((1 + 1) * 2000 * (Bb + 2) + 8000 * (Bb + 2) + 30) + 32

def cDropR (n Bb : ℕ) : ℕ :=
  (1 + 1) * 2000 * (Bb + 2) + n * ((2 * Bb + 2) + 8000 * (Bb + 2) + 30) + 32

/-! ## The padding program -/

/-- Parameter registers of the padding loops. -/
abbrev pX : Fin HS := 16
abbrev pB : Fin HS := 17
abbrev pZ : Fin HS := 18
abbrev pT : Fin HS := 19

/-- `padL`: per block move `F·Y` entries (`pB`) to stack `15` and push `(F'-F)·Y` zeros (`pZ`);
move the `X·F'·Y` entries (`pT`) back. -/
def padP : RS :=
  .seq (.rep 20 (.reg pX) (.seq (NLogN.mvK 0 15 21 (.reg pB)) (zeroKE 15 21 (.reg pZ))))
    (NLogN.mvK 15 0 21 (.reg pT))

/-- `unpadL`: per block move `F·Y` entries (`pB`) to stack `15` and drop `(F'-F)·Y` (`pZ`);
move the `X·F·Y` entries (`pT`) back. -/
def unpadP : RS :=
  .seq (.rep 20 (.reg pX) (.seq (NLogN.mvK 0 15 21 (.reg pB)) (dropKE 0 21 (.reg pZ))))
    (NLogN.mvK 15 0 21 (.reg pT))

/-- Cost bound of the padding loops (`V` bounds the padded volume `X · F' · Y`). -/
def cPad (V Bb : ℕ) : ℕ := 10 ^ 5 * (V + 1) * (Bb + 2)

/-- The managed state of the padding loop after `k` blocks. -/
def padF (f : Fin HS → List ℤ) (_X F F' Y : ℕ) (xs rest : List ℤ) (k : ℕ) : Fin HS → List ℤ :=
  Function.update (Function.update f 0 (xs.drop (k * (F * Y)) ++ rest)) 15
    (padL k F F' Y 0 xs).reverse

/-- The managed state of the extraction loop after `k` blocks. -/
def unpadF (f : Fin HS → List ℤ) (F F' Y : ℕ) (ys rest : List ℤ) (k : ℕ) : Fin HS → List ℤ :=
  Function.update (Function.update f 0 (ys.drop (k * (F' * Y)) ++ rest)) 15
    (unpadL k F F' Y ys).reverse

end NLogN.KIL

end

section
-- Source module: Solutions.KIL6Loop
/-!
# KIL6: counted loops whose body is item-language code (with procedure calls)

`irep C n body` is the item-language counterpart of `RS.rep C n b` for an arbitrary
`IStmt` body (which may contain `.call`s): push the count `n` on the counter stack `C`, run
`body; C := C - 1` while the top of `C` is nonzero, pop `C`.  `IRuns.irep` is the analogue of
`Sp.rep` on managed states.

Also: `Sp.asgQ` (assignment of a general expression, quadratic cost) and the conversion of a
stack-`0` specification on `t.upd 0 …` states to managed states (`Sp.top0`).
-/

set_option linter.unusedSimpArgs false
set_option linter.unnecessarySeqFocus false

namespace NLogN.KIL

open IntMul RS AModel AModelL12

section IRep

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

/-- The counted loop with an item-language body. -/
def irep (C : Fin HS) (n : Ex HS) (body : IStmt HS Pp) : IStmt HS Pp :=
  .seq (RS.push C n).code
    (.seq (.loop C (fun c => c != .zero) (.seq body (RS.asg C (decE C)).code)) (RS.pop C).code)

end IRep

end NLogN.KIL

end

section
-- Source module: Solutions.KIL6Swap
/-!
# KIL6: the high-digit interchanges as an IL loop

`swapsP q`: `r` iterations (count in register `kR`) of the Lemma 3 single-digit interchange
`swapE` (KIL5Swap) of the digit pair `Pre.length + t ↔ Pre.length + e + |G| + t` of
`qShape Pre G Suf q e`, `t = 0, …, r - 1`.  The volumes `vol Pre · q^t`, `q^(e-1-t) · vol G · q^t`,
`q^(e-1-t) · vol Suf` are products of the registers `kP, kG, kS` (the collapsed shape) and the
running powers `kQt = q^t`, `kQr = q^(e-1-t)` (updated by one multiplication / division per
iteration).

`swapsP_sp`: the result is `highSwaps (qShape Pre G Suf q e) |Pre| (|Pre| + e + |G|) r 0 xs`,
the registers end with `q^r` and `q^(e-r) / q`; cost `≤ (r+1) · 2·10⁹ (q+1)(V+1)(Bb+2)²`.
-/

set_option linter.unusedSimpArgs false
set_option linter.unnecessarySeqFocus false

namespace NLogN.KIL

open IntMul RS AModel AModelL12

/-! ## Registers of the Lemma 12 program (readable by `GoodE`: stacks `38 … 55`) -/

abbrev kP : Fin HS := 38
abbrev kG : Fin HS := 39
abbrev kS : Fin HS := 40
abbrev kQt : Fin HS := 41
abbrev kQr : Fin HS := 42
abbrev kR : Fin HS := 43

/-! ## The loop -/

/-- The four volume expressions of the digit interchange at offset `t`. -/
def swA : Ex HS := .mul (.reg kP) (.reg kQt)
def swB : Ex HS := .mul (.mul (.reg kQr) (.reg kG)) (.reg kQt)
def swC6 : Ex HS := .mul (.reg kQr) (.reg NLogN.KIL.kS)

/-- The high-digit interchanges. -/
def swapsP (q : ℕ) : RS :=
  .rep 24 (.reg kR) (.seq (swapE NLogN.KIL.swA (.lit (q : ℤ)) swB swC6)
    (.seq (.asg kQt (.mul (.reg kQt) (.lit (q : ℤ)))) (.asg kQr (.ndiv (.reg kQr) (.lit (q : ℤ))))))

/-- Cost bound of one digit interchange. -/
def cSw1 (q V Bb : ℕ) : ℕ := 2 * 10 ^ 9 * (q + 1) * ((V + 1) * (Bb + 2) ^ 2)

/-- Managed state after `k` interchanges. -/
def swF (f : Fin HS → List ℤ) (sh : Shape) (i j q e : ℕ) (xs rest : List ℤ) (l1 l2 : List ℤ)
    (k : ℕ) : Fin HS → List ℤ :=
  Function.update (Function.update (Function.update f 0 (highSwaps sh i j k 0 xs ++ rest)) kQt
    (((q ^ k : ℕ) : ℤ) :: l1)) kQr (((q ^ (e - k) / q : ℕ) : ℤ) :: l2)

end NLogN.KIL

end

section
-- Source module: Solutions.KIL6NE
/-!
# KIL6: natural-number expressions

`NE` is a small language of natural-number expressions over registers (`+`, `-` (truncation
excluded by hypothesis), `*`, `/`, `%`, `if a < b then x else y`), compiled to `Ex HS`.
`NE.good`: on a state whose registers hold the naturals `v i`, if every intermediate value is
`≤ M` and `zsz M ≤ B`, then the compiled expression is clean, well defined, `B`-bounded and
evaluates to the natural value.  This turns every descriptor computation into arithmetic.
-/

set_option linter.unusedSimpArgs false
set_option linter.unusedVariables false

namespace NLogN.KIL

open IntMul RS AModel AModelL12

inductive NE where
  | reg (i : Fin HS)
  | lit (c : ℕ)
  | add (a b : NE)
  | sub (a b : NE)
  | mul (a b : NE)
  | div (a b : NE)
  | mod (a b : NE)
  | lt (a b x y : NE)

namespace NE

def toEx : NE → Ex HS
  | reg i => .reg i
  | lit c => .lit (c : ℤ)
  | add a b => .add a.toEx b.toEx
  | sub a b => .add a.toEx (.neg b.toEx)
  | mul a b => .mul a.toEx b.toEx
  | div a b => .ndiv a.toEx b.toEx
  | mod a b => .nmod a.toEx b.toEx
  | lt a b x y => .sel (.add a.toEx (.neg b.toEx)) x.toEx y.toEx

def val (v : Fin HS → ℕ) : NE → ℕ
  | reg i => v i
  | lit c => c
  | add a b => a.val v + b.val v
  | sub a b => a.val v - b.val v
  | mul a b => a.val v * b.val v
  | div a b => a.val v / b.val v
  | mod a b => a.val v % b.val v
  | lt a b x y => if a.val v < b.val v then x.val v else y.val v

def regs : NE → List (Fin HS)
  | reg i => [i]
  | lit _ => []
  | add a b | sub a b | mul a b | div a b | mod a b => a.regs ++ b.regs
  | lt a b x y => a.regs ++ b.regs ++ x.regs ++ y.regs

/-- All intermediate values are `≤ M`, subtractions do not truncate, divisors are positive. -/
def OK (v : Fin HS → ℕ) (M : ℕ) : NE → Prop
  | reg i => v i ≤ M
  | lit c => c ≤ M
  | add a b => a.OK v M ∧ b.OK v M ∧ a.val v + b.val v ≤ M
  | sub a b => a.OK v M ∧ b.OK v M ∧ b.val v ≤ a.val v
  | mul a b => a.OK v M ∧ b.OK v M ∧ a.val v * b.val v ≤ M
  | div a b => a.OK v M ∧ b.OK v M ∧ 0 < b.val v
  | mod a b => a.OK v M ∧ b.OK v M ∧ 0 < b.val v
  | lt a b x y => a.OK v M ∧ b.OK v M ∧ (if a.val v < b.val v then x.OK v M else y.OK v M)

/-- No multiplication or division (the compiled expression is linear). -/
def IsLin : NE → Prop
  | reg _ | lit _ => True
  | add a b | sub a b => a.IsLin ∧ b.IsLin
  | mul _ _ | div _ _ | mod _ _ => False
  | lt a b x y => a.IsLin ∧ b.IsLin ∧ x.IsLin ∧ y.IsLin

def size : NE → ℕ
  | reg _ | lit _ => 1
  | add a b | mul a b | div a b | mod a b => a.size + b.size + 1
  | sub a b => a.size + b.size + 2
  | lt a b x y => a.size + b.size + x.size + y.size + 3

/-- The registers of `e` hold the naturals `v i` in `t`. -/
def Holds (t : IState HS) (v : Fin HS → ℕ) (e : NE) : Prop :=
  ∀ i ∈ e.regs, ∃ L, t.G i = .int (v i : ℤ) :: L

/-- Registers below `56` give clean expressions. -/
def RegsLt (e : NE) : Prop := ∀ i ∈ e.regs, i.val < 56

end NE

/-! ## Managed-state statements -/

end NLogN.KIL

end