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

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

import Definitions.Def_IntMul_Kappa_K_Defs10

set_option autoImplicit false
set_option relaxedAutoImplicit false
set_option maxHeartbeats 1000000

section
-- Source module: Solutions.KIL2Unit
/-!
# KIL2: the Lemma 10 controlled shift with unit coefficient `u = ±c` (`2 ≤ c ≤ 20`)

`shiftArr Λ i u c` with `u = c` (resp. `u = -c`), `c` a small positive integer prime to `Q`, is
done as the paper's Lemma 4: first the fixed multiplicative permutation `y ↦ c·y (mod Q)` of the
target field (`mulPc`: split into `c` increasing streams, merge by residue), then the rotation
`shP` (resp. the reflection `shN`) by the offsets `Σ_j c_j·d_j`.

* `outMul_recR`: `mulPc` computes `pull [P, Q, B] [P, Q, B] (mulAddr Q c)`.
* `mulPc_spec`: stack form, linear cost `cMulPc_le`.
* `shiftArr_u_eq`: any coefficient `u`, `shiftArr` as a pull on `[vol A, Q, vol C]`.
* `shiftArr_mul_eq`, `shiftArr_negmul_eq`: `shiftArr u = (rotation/reflection) ∘ (y ↦ c·y)`.
* `shiftArrMul_il_spec`, `shiftArrNegMul_il_spec`: the IL programs, cost
  `cMulPc + cShift = O((|A| + 1) · (V + 1) · (Bb + 2) + (Bb + 2)^2)`.
-/

set_option linter.unusedSimpArgs false
set_option linter.unnecessarySeqFocus false

namespace NLogN.KIL

open IntMul RS AModel

/-- Output address `(p, z, b)` ↦ input address `(p, c⁻¹·z, b)`. -/
def mulAddr (Q c : ℕ) (a : List ℕ) : List ℕ := [a.getD 0 0, NLogN.KIL.yOf Q c (a.getD 1 0), a.getD 2 0]

/-! ## The shift with any coefficient as a pull on `[vol A, Q, vol C]` -/

section ShiftU

open AModelL12

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

/-- Output address `(p, y, b)` ↦ input address `(p, u⁻¹·(y − a_p), b)`. -/
def uAddr (Q : ℕ) (u : ZMod Q) (as : ℕ → ℕ) (a : List ℕ) : List ℕ :=
  [a.getD 0 0, (u⁻¹ * ((a.getD 1 0 : ZMod Q) - (as (a.getD 0 0) : ZMod Q))).val, a.getD 2 0]

/-- `u = k`: multiplicative permutation, then the self-contained rotation. -/
def shMul (Q k P B : ℕ) (ac : List (ℕ × ℕ)) : RS := .seq (mulPc P Q k B) (shP Q P B ac)

/-- `u = -k`: multiplicative permutation, then the self-contained reflection. -/
def shNMul (Q k P B : ℕ) (ac : List (ℕ × ℕ)) : RS := .seq (mulPc P Q k B) (shN Q P B ac)

/-- Cost bound of `shMul` / `shNMul`. -/
def cShiftU (nA V Bb : ℕ) : ℕ := 10 ^ 7 * (V + 1) * (Bb + 2) + cShift nA V Bb

end ShiftU

/-! ## Step runner with small unit coefficients -/

section RunU

open AModelL12

variable {sh : Shape} {n Q : ℕ} (Λ : Layout sh n Q) {Pp : ℕ} {I : List Sym}
  {iprog : Fin Pp → IStmt HS Pp}

/-- The coefficients handled: `±1`, and `±k` with `2 ≤ k ≤ 20`, `k` prime to `Q`. -/
def SmallUnit (Q : ℕ) (u : ZMod Q) : Prop :=
  u = 1 ∨ u = -1 ∨ (2 ≤ u.val ∧ u.val ≤ 20 ∧ Nat.Coprime Q u.val) ∨
    (2 ≤ (-u).val ∧ (-u).val ≤ 20 ∧ Nat.Coprime Q (-u).val)

/-- Shift code for small unit coefficients. -/
def scodeU : Step n Q → IStmt HS Pp
  | .shift i u c =>
    if u = 1 ∨ u = -1 then scodePM Λ (.shift i u c)
    else if 2 ≤ u.val ∧ u.val ≤ 20 ∧ Nat.Coprime Q u.val then
      (shMul Q u.val (AModel.vol (sh.take (Λ.pos i))) (AModel.vol (sh.drop (Λ.pos i + 1)))
        ((sh.take (Λ.pos i)).zip (csOf Λ c (sh.take (Λ.pos i))))).code
    else (shNMul Q (-u).val (AModel.vol (sh.take (Λ.pos i))) (AModel.vol (sh.drop (Λ.pos i + 1)))
        ((sh.take (Λ.pos i)).zip (csOf Λ c (sh.take (Λ.pos i))))).code
  | .swap _ _ => .skip

end RunU

end NLogN.KIL

end

section
-- Source module: Solutions.KIL2Div
/-!
# KIL2: multiplying the target field by `c⁻¹` (Lemma 4, the inverse direction)

The mirror of `KIL2Mul`: the records `y = 0, …, Q − 1` of a block are distributed by residue,
record `y` going to stream `σ(y mod c)` (a residue counter on stack `13` and a `c`-way branch),
and the streams `0, …, c − 1` are then concatenated.  Stream `j` holds the records
`y = c·z − j·Q` for `z ∈ [lo j, lo (j+1))` in increasing order (`filt_eq`), so the output position
`z` receives record `(c·z) mod Q`: the permutation is `y ↦ c⁻¹·y`.

* `filt_eq`, `outD_eq`: the combinatorial identity.
* `divPc_sp`, `divPc_spec`: the program, stack form, linear cost (factor `c ≤ 20`).
* `shiftArr_div_eq`, `shiftArr_negdiv_eq`; `shiftArrDiv_il_spec`, `shiftArrNegDiv_il_spec`:
  `u = k⁻¹` and `u = -k⁻¹`.
* `SmallUnit2`, `scodeU2`, `runCodeU2_spec`: the step runner with coefficients
  `±1, ±k, ±k⁻¹` (`2 ≤ k ≤ 20`, `k` prime to `Q`).
-/

set_option linter.unusedSimpArgs false
set_option linter.unnecessarySeqFocus false

namespace NLogN.KIL

open IntMul RS AModel

/-! ## The residue classes are the pieces -/

section Arith

variable {Q c : ℕ}

/-- Records of stream `j`, in order. -/
def flatD (rp : ℕ → List ℤ) (Q c j : ℕ) : List ℤ :=
  ((List.range Q).filter (fun y => sig Q c (y % c) = j)).flatMap rp

end Arith

/-! ## Distribution by residue -/

section DistD

variable (rp : ℕ → List ℤ) (Q c B : ℕ)

/-- Stream `j` at time `y`: the records `y' < y` of class `j`, reversed. -/
def accD (y j : ℕ) : List ℤ :=
  (((List.range y).filter (fun y' => sig Q c (y' % c) = j)).flatMap rp).reverse

/-- Branch `v`: move record `y` (residue `v`) from stack `0` to stream `σ(v)`. -/
def brD (v : ℕ) : RS := NLogN.mvK 0 (NLogN.KIL.sI (sig Q c v)) 33 (.lit B)

def bodyD : RS := .seq (NLogN.KIL.sel 13 (brD Q c B) 0 c) (.asg 13 (mAdd c (.reg 13) (.lit 1)))

def distD : RS := .rep 31 (.lit Q) (bodyD Q c B)

/-- State at time `y`. -/
def stD (K : NLogN.KIL.KS) (rest : List ℤ) (y : ℕ) (kk : List ℤ) : Fin HS → List ℤ :=
  strF c (toF {K with x := (List.range' y (Q - y)).flatMap rp ++ rest, qMB := [((y % c : ℕ) : ℤ)], k1 := kk})
    (accD rp Q c y)

variable {rp Q c B}

def cDist (Q c B Bb : ℕ) : ℕ :=
  (1 + 1) * 2000 * (Bb + 2) + Q * (cBodyM c B Bb + 8000 * (Bb + 2) + 30) + 32

end DistD

/-! ## Concatenation of the streams -/

section Cat

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

/-- Number of records of class `j`. -/
def lenD (Q c j : ℕ) : ℕ := ((List.range Q).filter (fun y => sig Q c (y % c) = j)).length

/-- Move streams `k − 1, …, 0` (whole) onto stack `10`. -/
def catD (Q c B : ℕ) : ℕ → RS
  | 0 => .skip
  | k + 1 => .seq (NLogN.mvK (NLogN.KIL.sI k) 10 33 (.lit ((lenD Q c k * B : ℕ) : ℤ))) (catD Q c B k)

end Cat

/-! ## One block and the whole program (multiplication by `c⁻¹`) -/

/-- The output block: position `z` receives record `c·z mod Q`. -/
def outDB (rp : ℕ → List ℤ) (Q c : ℕ) : List ℤ := (List.range Q).flatMap fun z => rp (c * z % Q)

section BlockD

variable (r : ℕ → ℕ → List ℤ) (P Q c B : ℕ)

def blkD : RS :=
  .seq (.push 13 (.lit 0)) (.seq (distD Q c B) (.seq (.pop 13) (.seq (catD Q c B c)
    (NLogN.mvK 10 12 33 (.lit ((Q * B : ℕ) : ℤ))))))

def divPc : RS :=
  .seq (.rep 30 (.lit (P : ℤ)) (blkD Q c B)) (NLogN.mvK 12 0 33 (.lit ((P * (Q * B) : ℕ) : ℤ)))

/-- All output blocks. -/
def outDiv : List ℤ := (List.range P).flatMap fun p => outDB (r p) Q c

/-- State at the start of block `p`. -/
def dSt (p : ℕ) (kk : List ℤ) : KS :=
  { x := (List.range' p (P - p)).flatMap (blk3 r Q),
    o := ((List.range p).flatMap fun p => outDB (r p) Q c).reverse, k0 := kk }

def cBlkD (Q c B Bb : ℕ) : ℕ :=
  (1 + 1) * 2000 * (Bb + 2) + cDist Q c B Bb + (2 * Bb + 2) +
    (c * ((1 + 1) * 2000 * (Bb + 2) + 32) + 1 + c * (Q * B) * (14000 * (Bb + 2))) +
    ((1 + 1) * 2000 * (Bb + 2) + Q * B * (14000 * (Bb + 2)) + 32)

variable {r P Q c B}

def cDivPc (P Q c B Bb : ℕ) : ℕ :=
  (1 + 1) * 2000 * (Bb + 2) + P * (cBlkD Q c B Bb + 8000 * (Bb + 2) + 30) + 32 +
    ((1 + 1) * 2000 * (Bb + 2) + P * (Q * B) * (14000 * (Bb + 2)) + 32)

end BlockD

/-- Output address `(p, z, b)` ↦ input address `(p, c·z mod Q, b)`. -/
def divAddr (Q c : ℕ) (a : List ℕ) : List ℕ := [a.getD 0 0, c * a.getD 1 0 % Q, a.getD 2 0]

/-! ## The shifts with coefficient `±k⁻¹` -/

section ShiftDiv

open AModelL12

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

/-- `u = k⁻¹`: inverse multiplicative permutation, then the self-contained rotation. -/
def shDiv (Q k P B : ℕ) (ac : List (ℕ × ℕ)) : RS := .seq (divPc P Q k B) (shP Q P B ac)

/-- `u = -k⁻¹`: inverse multiplicative permutation, then the self-contained reflection. -/
def shNDiv (Q k P B : ℕ) (ac : List (ℕ × ℕ)) : RS := .seq (divPc P Q k B) (shN Q P B ac)

end ShiftDiv

/-! ## Step runner with coefficients `±1, ±k, ±k⁻¹` -/

section RunU2

open AModelL12

variable {sh : Shape} {n Q : ℕ} (Λ : Layout sh n Q) {Pp : ℕ} {I : List Sym}
  {iprog : Fin Pp → IStmt HS Pp}

/-- The coefficients handled: `±1`, `±k` and `±k⁻¹` with `2 ≤ k ≤ 20`, `k` prime to `Q`. -/
def SmallUnit2 (Q : ℕ) (u : ZMod Q) : Prop :=
  SmallUnit Q u ∨
    (2 ≤ (u⁻¹).val ∧ (u⁻¹).val ≤ 20 ∧ Nat.Coprime Q (u⁻¹).val ∧ u * u⁻¹ = 1) ∨
    (2 ≤ ((-u)⁻¹).val ∧ ((-u)⁻¹).val ≤ 20 ∧ Nat.Coprime Q ((-u)⁻¹).val ∧ -u * (-u)⁻¹ = 1)

/-- Shift code for coefficients `±1, ±k, ±k⁻¹`. -/
def scodeU2 : Step n Q → IStmt HS Pp
  | .shift i u c =>
    if u = 1 ∨ u = -1 ∨ (2 ≤ u.val ∧ u.val ≤ 20 ∧ Nat.Coprime Q u.val) ∨
        (2 ≤ (-u).val ∧ (-u).val ≤ 20 ∧ Nat.Coprime Q (-u).val) then scodeU Λ (.shift i u c)
    else if 2 ≤ (u⁻¹).val ∧ (u⁻¹).val ≤ 20 ∧ Nat.Coprime Q (u⁻¹).val ∧ u * u⁻¹ = 1 then
      (shDiv Q (u⁻¹).val (AModel.vol (sh.take (Λ.pos i))) (AModel.vol (sh.drop (Λ.pos i + 1)))
        ((sh.take (Λ.pos i)).zip (csOf Λ c (sh.take (Λ.pos i))))).code
    else (shNDiv Q ((-u)⁻¹).val (AModel.vol (sh.take (Λ.pos i))) (AModel.vol (sh.drop (Λ.pos i + 1)))
        ((sh.take (Λ.pos i)).zip (csOf Λ c (sh.take (Λ.pos i))))).code
  | .swap _ _ => .skip

end RunU2

end NLogN.KIL

end

section
-- Source module: Solutions.KIL3Gather
/-!
# KIL3: multiplication of the target field by `c⁻¹` for any fixed `c` (passes over one stack)

`divPg P Q c B` computes the same array as KIL2's `divPc` (`pull … (divAddr Q c)`, output
position `z` reads record `c·z mod Q`), for every `c ≥ 2` prime to `Q`, with no bound on `c`:
instead of one stream stack per residue class, it makes `c` passes over the block.  Pass `j`
scans the records still on stack `0` with a residue counter (stack `13`): a record of class `j`
(`σ(y mod c) = j`) goes to the output stack `12`, a record of a later class goes to stack `10`,
a record of an earlier class was already taken (skip); then stack `10` is moved back to `0`.
After the `c` passes the output stack holds the block's classes `0, …, c − 1` in order, which
is `outDB` (`outD_eq`).  Cost `O(c² · Q · B · (Bb + 2))` per block (a `c`-way branch per
record and pass).
-/

set_option linter.unusedSimpArgs false

namespace NLogN.KIL

open IntMul RS AModel

section GatherDefs

variable (rp : ℕ → List ℤ) (Q c B : ℕ)

/-- Records `y' ∈ [y, Q)` whose class satisfies `p`. -/
def selR (p : ℕ → Bool) (y : ℕ) : List ℤ :=
  ((List.range' y (Q - y)).filter (fun y' => p (sig Q c (y' % c)))).flatMap rp

/-- Records `y' < y` whose class satisfies `p`. -/
def preR (p : ℕ → Bool) (y : ℕ) : List ℤ :=
  ((List.range y).filter (fun y' => p (sig Q c (y' % c)))).flatMap rp

/-- Branch of pass `j` on residue `v`. -/
def brG (j v : ℕ) : RS :=
  if sig Q c v < j then .skip
  else if sig Q c v = j then NLogN.mvK 0 12 33 (.lit B) else NLogN.mvK 0 10 33 (.lit B)

def bodyG (j : ℕ) : RS := .seq (NLogN.KIL.sel 13 (brG Q c B j) 0 c) (.asg 13 (mAdd c (.reg 13) (.lit 1)))

/-- Number of records of a class `> j`. -/
def cntG (j : ℕ) : ℕ := ((List.range Q).filter (fun y => decide (j < sig Q c (y % c)))).length

def passG (j : ℕ) : RS :=
  .seq (.push 13 (.lit 0)) (.seq (.rep 31 (.lit Q) (bodyG Q c B j)) (.seq (.pop 13)
    (NLogN.mvK 10 0 33 (.lit ((cntG Q c j * B : ℕ) : ℤ)))))

/-- Passes `0, …, k − 1`. -/
def passesG : ℕ → RS
  | 0 => .skip
  | k + 1 => .seq (passesG k) (passG Q c B k)

/-- State of pass `j` at record `y`. -/
def gSt (K : NLogN.KIL.KS) (rest O : List ℤ) (j y : ℕ) (q kk : List ℤ) : KS :=
  {K with x := selR rp Q c (fun k => decide (j ≤ k)) y ++ rest,
          a := (preR rp Q c (fun k => decide (j < k)) y).reverse,
          o := (preR rp Q c (fun k => decide (k = j)) y).reverse ++ O,
          qMB := q, k1 := kk}

end GatherDefs

section GatherLemmas

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

/-- Cost of one branch. -/
def cBr (B Bb : ℕ) : ℕ := (1 + 1) * 2000 * (Bb + 2) + B * (14000 * (Bb + 2)) + 32

def cBodyG (c B Bb : ℕ) : ℕ := c * cSel Bb + NLogN.KIL.cBr B Bb + (14 + 1) * 2000 * (Bb + 2)

def cPassG (Q c B Bb : ℕ) : ℕ :=
  (1 + 1) * 2000 * (Bb + 2) +
    ((1 + 1) * 2000 * (Bb + 2) + Q * (cBodyG c B Bb + 8000 * (Bb + 2) + 30) + 32) +
    (2 * Bb + 2) + ((1 + 1) * 2000 * (Bb + 2) + Q * B * (14000 * (Bb + 2)) + 32)

end GatherLemmas

/-! ## Blocks and the whole program -/

section GatherBlock

variable (r : ℕ → ℕ → List ℤ) (P Q c B : ℕ)

def divPg : RS :=
  .seq (.rep 30 (.lit (P : ℤ)) (passesG Q c B c)) (NLogN.mvK 12 0 33 (.lit ((P * (Q * B) : ℕ) : ℤ)))

def cBlkG (Q c B Bb : ℕ) : ℕ := 1 + c * cPassG Q c B Bb

def cDivPg (P Q c B Bb : ℕ) : ℕ :=
  (1 + 1) * 2000 * (Bb + 2) + P * (cBlkG Q c B Bb + 8000 * (Bb + 2) + 30) + 32 +
    ((1 + 1) * 2000 * (Bb + 2) + P * (Q * B) * (14000 * (Bb + 2)) + 32)

variable {r P Q c B}

end GatherBlock

end NLogN.KIL

end

section
-- Source module: Solutions.KIL3Scatter
/-!
# KIL3: multiplication of the target field by `c` for any fixed `c` (passes over one stack)

`mulPg P Q c B` computes the same array as KIL2's `mulPc` (`pull … (mulAddr Q c)`, output
position `z` reads record `c⁻¹·z`), for every `c ≥ 2` prime to `Q`, with no bound on `c`.
The block's records form `c` increasing pieces `j = ⌊c·y/Q⌋` (contiguous, piece `j` is
`[lo j, lo (j+1))`), and output position `z` takes the next record of piece `σ(z mod c)`.
Pass `j` merges piece `j` (top of stack `0`) into the partial output on stack `10` (positions of
the classes `< j`, in order): for `z = 0, …, Q − 1` (residue counter on `13`), a position of an
earlier class moves its record from `10` to `11`, a position of class `j` moves the next record of
stack `0` to `11`, a later class is skipped; then `11` is moved back to `10`.
-/

set_option linter.unusedSimpArgs false

namespace NLogN.KIL

open IntMul RS AModel

section ScatterDefs

variable (rp : ℕ → List ℤ) (Q c B : ℕ)

/-- Branch of merge pass `j` on residue `v`. -/
def brS (j v : ℕ) : RS :=
  if sig Q c v < j then NLogN.mvK 10 11 33 (.lit B)
  else if sig Q c v = j then NLogN.mvK 0 11 33 (.lit B) else .skip

def bodyS (j : ℕ) : RS := .seq (NLogN.KIL.sel 13 (brS Q c B j) 0 c) (.asg 13 (mAdd c (.reg 13) (.lit 1)))

/-- Number of output positions of a class `≤ j`. -/
def cntS (j : ℕ) : ℕ := ((List.range Q).filter (fun z => decide (sig Q c (z % c) ≤ j))).length

def passS (j : ℕ) : RS :=
  .seq (.push 13 (.lit 0)) (.seq (.rep 31 (.lit Q) (bodyS Q c B j)) (.seq (.pop 13)
    (NLogN.mvK 11 10 33 (.lit ((cntS Q c j * B : ℕ) : ℤ)))))

def passesS : ℕ → RS
  | 0 => .skip
  | k + 1 => .seq (passesS k) (passS Q c B k)

/-- State of merge pass `j` at output position `z`. -/
def sSt (K : NLogN.KIL.KS) (rest : List ℤ) (j z : ℕ) (q kk : List ℤ) : KS :=
  {K with x := (List.range' (nxt Q c j z) (Q - nxt Q c j z)).flatMap rp ++ rest,
          a := selR (fun y => rp (NLogN.KIL.yOf Q c y)) Q c (fun k => decide (k < j)) z,
          c := (preR (fun y => rp (NLogN.KIL.yOf Q c y)) Q c (fun k => decide (k ≤ j)) z).reverse,
          qMB := q, k1 := kk}

end ScatterDefs

section ScatterLemmas

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

end ScatterLemmas

/-! ## Blocks and the whole program -/

section ScatterBlock

variable (r : ℕ → ℕ → List ℤ) (P Q c B : ℕ)

def blkS : RS := .seq (passesS Q c B c) (NLogN.mvK 10 12 33 (.lit ((Q * B : ℕ) : ℤ)))

def mulPg : RS :=
  .seq (.rep 30 (.lit (P : ℤ)) (blkS Q c B)) (NLogN.mvK 12 0 33 (.lit ((P * (Q * B) : ℕ) : ℤ)))

def cBlkS (Q c B Bb : ℕ) : ℕ :=
  1 + c * cPassG Q c B Bb + ((1 + 1) * 2000 * (Bb + 2) + Q * B * (14000 * (Bb + 2)) + 32)

def cMulPg (P Q c B Bb : ℕ) : ℕ :=
  (1 + 1) * 2000 * (Bb + 2) + P * (cBlkS Q c B Bb + 8000 * (Bb + 2) + 30) + 32 +
    ((1 + 1) * 2000 * (Bb + 2) + P * (Q * B) * (14000 * (Bb + 2)) + 32)

variable {r P Q c B}

end ScatterBlock

end NLogN.KIL

end

section
-- Source module: Solutions.KIL3Unit
/-!
# KIL3: the Lemma 10 controlled shift with any fixed unit `u = ±n·d⁻¹`

For fixed naturals `n, d ≥ 1` prime to `Q`, `shiftArr Λ i (±n·d⁻¹) c` is
`(rotation / reflection) ∘ (y ↦ n·y) ∘ (y ↦ d⁻¹·y)` on `[vol A, Q, vol C]`
(`shiftArr_frac_eq`, `shiftArr_negfrac_eq`); the two multiplicative permutations are the
pass programs `divPg` (`KIL3Gather`) and `mulPg` (`KIL3Scatter`), skipped when the factor is `1`.

* `shiftArr_il_unit`: the IL program `shFrac`, cost
  `((n + 1)² + (d + 1)²) · 10⁶ · (V + 1) · (Bb + 2) + cShift |A| V Bb`.
-/

set_option linter.unusedSimpArgs false
set_option linter.unnecessarySeqFocus false

namespace NLogN.KIL

open IntMul RS AModel

/-! ## Addresses -/

section Addr

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

end Addr

/-! ## Costs of the pass programs -/

section Cost

end Cost

/-! ## The two multiplicative stages (skipped for the factor `1`) -/

section Stage

/-- Multiplication by `d⁻¹` (skip when `d = 1`). -/
def fdivP (P Q d B : ℕ) : RS := if d < 2 then .skip else divPg P Q d B

/-- Multiplication by `n` (skip when `n = 1`). -/
def fmulP (P Q k B : ℕ) : RS := if k < 2 then .skip else mulPg P Q k B

end Stage

/-! ## The controlled shift with coefficient `±n·d⁻¹` -/

section ShiftFrac

open AModelL12

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

/-- The IL program: `y ↦ d⁻¹·y`, then `y ↦ k·y`, then the rotation (`neg = false`) or the
reflection (`neg = true`). -/
def shFrac (Q k d P B : ℕ) (neg : Bool) (ac : List (ℕ × ℕ)) : RS :=
  .seq (.seq (fdivP P Q d B) (fmulP P Q k B)) (if neg then shN Q P B ac else shP Q P B ac)

/-- The coefficient `±k·d⁻¹`. -/
def fracU (Q k d : ℕ) (neg : Bool) : ZMod Q :=
  if neg then -((k : ZMod Q) * (d : ZMod Q)⁻¹) else (k : ZMod Q) * (d : ZMod Q)⁻¹

/-- Cost bound of `shFrac`. -/
def cShiftF (k d nA V Bb : ℕ) : ℕ :=
  ((k + 1) ^ 2 + (d + 1) ^ 2) * 10 ^ 6 * (V + 1) * (Bb + 2) + cShift nA V Bb

end ShiftFrac

end NLogN.KIL

end

section
-- Source module: Solutions.KIL3Ops
/-!
# KIL3: generic stack operations on managed states

Element-wise loops used by the one-stack network: `dupK` (copy `n` integers to two stacks),
`dropK` (discard `n` integers), `zeroK` (push `n` zeros), `zipK` (pointwise addition mod 2 of two
stacks).  All are stated on managed states `encK s f`.
-/

set_option linter.unusedSimpArgs false

namespace NLogN.KIL

open IntMul RS AModel

/-- Copy `n` integers from `X` onto both `Y` and `Z` (each reversed), consuming them on `X`. -/
def dupK (X Y Z C : Fin HS) (n : ℕ) : RS :=
  .rep C (.lit n) (.seq (.push Y (.reg X)) (.seq (.push Z (.reg X)) (.pop X)))

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

def dupSt (f : Fin HS → List ℤ) (X Y Z : Fin HS) (L l : List ℤ) (k : ℕ) : Fin HS → List ℤ :=
  Function.update (Function.update (Function.update f X (L.drop k ++ l)) Y
    ((L.take k).reverse ++ f Y)) Z ((L.take k).reverse ++ f Z)

/-! ## Dropping `n` integers -/

def dropK (X C : Fin HS) (n : ℕ) : RS := .rep C (.lit n) (.pop X)

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

/-! ## Pushing `n` zeros -/

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

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

/-! ## Pointwise addition mod 2 -/

/-- `x + y` reduced once by `2` (the value of `mAdd 2`). -/
def xr (x y : ℤ) : ℤ := if x + y + -((2 : ℕ) : ℤ) < 0 then x + y else x + y + -((2 : ℕ) : ℤ)

/-- Bits. -/
def Bits (l : List ℤ) : Prop := ∀ x ∈ l, x = 0 ∨ x = 1

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

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

def zipSt (f : Fin HS → List ℤ) (A Z O : Fin HS) (a b la lb : List ℤ) (k : ℕ) :
    Fin HS → List ℤ :=
  Function.update (Function.update (Function.update f A (a.drop k ++ la)) Z (b.drop k ++ lb)) O
    ((List.zipWith xr (a.take k) (b.take k)).reverse ++ f O)

/-! ## IStmt forms -/

end NLogN.KIL

end

section
-- Source module: Solutions.KIL3Tab
/-!
# KIL3: a table of `W` streams on one stack

The `W` streams of the network are kept on stack `0`, one after the other (stream `0` on top):
`tab W F = F 0 ++ F 1 ++ ⋯ ++ F (W − 1)`, every stream of the same length `v`.

Operations (all `O(W · v)` item moves, plus the cost of the applied code):

* `applyC v w code`: run `code` (any stack-`0` program) on stream `w` in place
  (`applyC_runs`): the streams before `w` are parked on stack `15` and moved back.
* `copyC v w`: push a reversed copy of stream `w` on stack `16` (`copyC_sp`).
* `buildC v l`: push on stack `20` the (reversed) pointwise sum mod 2 of the streams `l`.
* `rebuildC W v T`: replace the table by the streams `w ↦ Σ_{w' ∈ T w} F w'` (`rebuildC_sp`).

Scratch stacks: `15` (park), `16` (copy), `17` (temp), `19` (zip out), `20` (new table),
`21` (loop counter).
-/

set_option linter.unusedSimpArgs false
set_option linter.unnecessarySeqFocus false

namespace NLogN.KIL

open IntMul RS AModel

/-! ## Tables -/

/-- Streams `a, …, a + n − 1` concatenated. -/
def tabR (a n : ℕ) (F : ℕ → List ℤ) : List ℤ := (List.range' a n).flatMap F

/-- Uniform streams: length `v`, bits. -/
def Uni (W v : ℕ) (F : ℕ → List ℤ) : Prop := ∀ w < W, (F w).length = v ∧ NLogN.KIL.Bits (F w)

/-! ## Moves with literal counts -/

def cMvL (n Bb : ℕ) : ℕ := (1 + 1) * 2000 * (Bb + 2) + n * (14000 * (Bb + 2)) + 32

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

section Apply

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

/-- Run `code` on stream `w` (streams of length `v`). -/
def applyC (v w : ℕ) (code : IStmt HS Pp) : IStmt HS Pp :=
  .seq (NLogN.mvK 0 15 21 (.lit ((w * v : ℕ) : ℤ))).code
    (.seq code (NLogN.mvK 15 0 21 (.lit ((w * v : ℕ) : ℤ))).code)

end Apply

/-! ## Cost unit -/

/-- One pass over at most `V` items. -/
def cOp (V Bb : ℕ) : ℕ := 50000 * (V + 1) * (Bb + 2)

/-! ## Copying a stream -/

/-- Push a reversed copy of stream `w` on stack `16`. -/
def copyC (v w : ℕ) : RS :=
  .seq (NLogN.mvK 0 15 21 (.lit ((w * v : ℕ) : ℤ))) (.seq (NLogN.KIL.dupK 0 16 17 21 v)
    (.seq (NLogN.mvK 17 0 21 (.lit (v : ℤ))) (NLogN.mvK 15 0 21 (.lit ((w * v : ℕ) : ℤ)))))

/-- Size hypotheses for the table operations. -/
structure TabB (W v Bb : ℕ) : Prop where
  four : zsz ((2 * 2 : ℕ) : ℤ) ≤ Bb
  vol : zsz ((W * v : ℕ) : ℤ) ≤ Bb
  pos : 0 < W

/-! ## Accumulating a stream on stack `20` -/

/-- `acc ← acc ⊕ (stream w)ʳᵉᵛ` on the top `v` entries of stack `20`. -/
def accC (v w : ℕ) : RS :=
  .seq (copyC v w) (.seq (NLogN.KIL.zipK 20 16 19 21 v) (NLogN.mvK 19 20 21 (.lit (v : ℤ))))

/-! ## Building a summed stream, rebuilding the table -/

def accL (v : ℕ) : List ℕ → RS
  | [] => .skip
  | w :: l => .seq (accC v w) (accL v l)

/-- IL-order fold: `acc ⊕ (F w)ʳᵉᵛ` for each `w ∈ l`. -/
def accR (F : ℕ → List ℤ) (a : List ℤ) (l : List ℕ) : List ℤ :=
  l.foldl (fun a w => List.zipWith xr a (F w).reverse) a

def buildC (v : ℕ) (l : List ℕ) : RS := .seq (NLogN.KIL.zeroK 20 21 v) (NLogN.KIL.accL v l)

/-- Pointwise sum mod 2 of the streams `l` (zero stream of length `v` if `l = []`). -/
def sumS (F : ℕ → List ℤ) (v : ℕ) (l : List ℕ) : List ℤ :=
  l.foldl (fun a w => List.zipWith xr a (F w)) (List.replicate v 0)

def buildsC (v : ℕ) (T : ℕ → List ℕ) : ℕ → RS
  | 0 => .skip
  | n + 1 => .seq (buildsC v T n) (buildC v (T n))

/-- Source lists: indices in range, at most `W` of them. -/
def SrcOK (W : ℕ) (T : ℕ → List ℕ) : Prop := ∀ w < W, (∀ w' ∈ T w, w' < W) ∧ (T w).length ≤ W

/-- Replace the table by `w ↦ Σ_{w' ∈ T w} F w'`. -/
def rebuildC (W v : ℕ) (T : ℕ → List ℕ) : RS :=
  .seq (buildsC v T W) (.seq (NLogN.KIL.dropK 0 21 (W * v)) (NLogN.mvK 20 0 21 (.lit ((W * v : ℕ) : ℤ))))

def cRebuild (W V Bb : ℕ) : ℕ := (W * (6 * W + 2) + 3) * cOp V Bb

end NLogN.KIL

end

section
-- Source module: Solutions.KIL3Net
/-!
# KIL3: the network of one Prop. 11 level on one stack (bits)

The `W` wire streams (bits, `ZMod 2`, encoded as the integers `0, 1`) are kept on stack `0` as a
table (`KIL3Tab`).  A gate `g` with touched set `S`:

1. for every `w ∈ S` (in order), the frame change `Φ (M_g − lab w)` is run in place on stream `w`
   (`applyC`, the stream code is any stack-`0` program, e.g. with child calls);
2. the table is rebuilt: a touched stream becomes the sum mod 2 of the touched streams with
   coefficient `1` in row `w` of the gate matrix, an untouched stream is copied.

At the end the final frame changes `Φ (Mout w − lab w)` are run on every stream.
`netC_runs`: the code computes `netRunF` (`AModelProp11`) with `Φ` given by the stream codes,
cost `Σ_frames (cE + 2·cOp) + Σ_gates (cRebuild + 2) + 1` (`O(V)` per frame and per gate,
`W` and the gate list fixed).
-/

set_option linter.unusedSimpArgs false
set_option linter.unnecessarySeqFocus false

namespace NLogN.KIL

open IntMul RS AModel

/-! ## Bits as integers -/

/-- Encoding of a bit list as integers `0, 1`. -/
def enc2 (ys : List (ZMod 2)) : List ℤ := ys.map fun a => ((a.val : ℕ) : ℤ)

/-- Sum mod 2 of the streams `l`, over `ZMod 2`. -/
def sumZ (Y : ℕ → List (ZMod 2)) (a : List (ZMod 2)) (l : List ℕ) : List (ZMod 2) :=
  l.foldl (fun a w => List.zipWith (· + ·) a (Y w)) a

/-! ## Sequences of statements -/

section Seq

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

def seqL : List (IStmt HS Pp) → IStmt HS Pp
  | [] => .skip
  | c :: cs => .seq c (seqL cs)

/-- Lift a `Fin W`-indexed family to `ℕ`. -/
def liftF {W : ℕ} {α : Type*} (d : α) (Y : Fin W → α) (n : ℕ) : α :=
  if h : n < W then Y ⟨n, h⟩ else d

end Seq

/-! ## List sums over `Fin W` -/

/-- The touched wires, in order. -/
def touchL {W : ℕ} (S : Finset (Fin W)) : List (Fin W) := (List.finRange W).filter (· ∈ S)

/-! ## The network -/

section Net

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

/-- The function form of a list map on streams of length `v`. -/
def funΦ (v : ℕ) (Φ : M → List (ZMod 2) → List (ZMod 2)) (X : M) (f : Fin v → ZMod 2) :
    Fin v → ZMod 2 := fun i => (Φ X (List.ofFn f)).getD (i : ℕ) 0

/-- The table of stream functions. -/
def tabOf (Y : Fin W → Fin v → ZMod 2) : ℕ → List ℤ := liftF [] fun w => enc2 (List.ofFn (Y w))

/-- Sources of the rebuild after gate `g`. -/
def srcT (g : NGate (Fin W) (ZMod 2) M) (n : ℕ) : List ℕ :=
  if h : n < W then
    (if (⟨n, h⟩ : Fin W) ∈ g.2.1 then
      ((List.finRange W).filter (fun w' => w' ∈ g.2.1 ∧ g.1 ⟨n, h⟩ w' = 1)).map Fin.val
    else [n])
  else []

variable (E : M → IStmt HS Pp)

/-- Code of one gate (labels `lab` before it). -/
def gateC (v : ℕ) (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) => applyC v (w : ℕ) (E (g.2.2 - lab w))))
    (rebuildC W v (srcT g)).code

/-- Code of the network: the gates, then the final frame changes. -/
def netC (v : ℕ) (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) => applyC v (w : ℕ) (E (Mout w - lab w)))
  | lab, g :: gs => .seq (gateC E v lab g) (netC v Mout (labStep lab g) gs)

variable (cE : M → ℕ)

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

variable {E cE}

end Net

end NLogN.KIL

end

section
-- Source module: Solutions.KIL3Level
/-!
# KIL3: one Prop. 11 level (`levelR`) as an IL procedure, bit data

The level of `AModelProp11.levelR` over `ZMod 2` (the bit network), with the bits stored as the
integers `0, 1` (`enc2`).  The code is

1. `D ← D − H` (`subSteps`, the KIL2 runner, coefficients `±1`);
2. the cyclic row split into the `W` streams (`bwdP`, KILRow), the streams one after the other on
   stack `0`;
3. the network (`netC`, KIL3Net): per gate the frame changes `Φ_X` of the touched streams in
   place (the Lemma 10 steps of `X`, shifts by any implemented shift code, interchanges by calls of
   the child procedures), then the pointwise gate (sums mod 2); finally the frame changes to the
   output labels;
4. the routing: stream `w` ← stream `ρ w` (`rebuildC`);
5. the row merge (`fwdP`);
6. `D ← H − D` (`flipSteps`).

Main results: `levelR_run` (the level code computes `levelR` for any shift code correct on
integer arrays, cost `levelCI`), `levelCI_le` (`levelCI ≤ lvlK·(V+1)·(Bb+2)² + nShift·cS +
rankSum C·(c₀ + 2)`), `levelR_il_small` (KIL2 shift codes, units `±1, ±k, ±k⁻¹`, `k ≤ 20`).
The children are only required on the interchanges `H_i ↔ D_j` (`edgeSteps_swap_HD`).
-/

set_option linter.unusedSimpArgs false
set_option linter.unnecessarySeqFocus false

namespace NLogN.KIL

open IntMul RS AModel AModelL12 AModelFac

/-! ## Commuting the bit encoding with array maps -/

/-! ## The step runner on bit data -/

section Runner

variable {sh : Shape} {n Q : ℕ} (Λ : Layout sh n Q) {Pp : ℕ} {I : List Sym}
  {iprog : Fin Pp → IStmt HS Pp}

end Runner

/-! ## Row split and merge with their parameters -/

/-- Push the parameters `1, W, P L, B` on stacks `1 … 4`, run `p`, pop them. -/
def withPar (W PL B : ℕ) (p : RS) : RS :=
  .seq (.push 1 (.lit ((1 : ℕ) : ℤ))) (.seq (.push 2 (.lit (W : ℤ))) (.seq (.push 3 (.lit ((PL : ℕ) : ℤ)))
    (.seq (.push 4 (.lit (B : ℤ))) (.seq p (.seq (.pop 1) (.seq (.pop 2) (.seq (.pop 3) (.pop 4))))))))

def cPar (W PL B Bb : ℕ) : ℕ := 4 * ((1 + 1) * 2000 * (Bb + 2)) + cMove 1 W PL B Bb + 4 * (2 * Bb + 2)

/-! ## The level -/

section Level

variable {m W : ℕ} [NeZero W] (C : Circuit (ZMod 2) m W) (N : ℕ) [NeZero N] (P L' : ℕ)
  (A B Cs : Shape) {Pp : ℕ} {I : List Sym} {iprog : Fin Pp → IStmt HS Pp}

/-- The frame change `Φ_X` on a stream: the Lemma 10 steps of `X` (shifts by `scode`,
interchanges by the calls `jOf`). -/
def edgeC (scode : Step (m + m) N → IStmt HS Pp) (jOf : Fin (m + m) → Fin (m + m) → Fin Pp)
    (X : IMat m) : IStmt HS Pp :=
  runCode scode jOf (AModel.edgeSteps C N X)

/-- Routing sources: stream `w` ← stream `ρ w`. -/
def permT (n : ℕ) : List ℕ := if h : n < W then [(C.ρ ⟨n, h⟩ : ℕ)] else []

/-- **The level code.** -/
def levelC (scode : Step (m + m) N → IStmt HS Pp) (jOf : Fin (m + m) → Fin (m + m) → Fin Pp) :
    IStmt HS Pp :=
  .seq (runCode (scodeU2 ((hdLayout A B Cs m N).cons2 P (W * L'))) jOf (subSteps m))
  (.seq (withPar W (P * L') (vol (hdShape A B Cs m N)) bwdP).code
  (.seq (netC (edgeC C N scode jOf) (vol (P :: L' :: hdShape A B Cs m N)) C.Mout C.Min C.gates)
  (.seq (rebuildC W (vol (P :: L' :: hdShape A B Cs m N)) (NLogN.KIL.permT C)).code
  (.seq (withPar W (P * L') (vol (hdShape A B Cs m N)) fwdP).code
  (runCode (scodeU2 ((hdLayout A B Cs m N).cons2 P (W * L'))) jOf (flipSteps m))))))

/-- The exact cost of the level code. -/
def levelCI (cS c0 Bb : ℕ) : ℕ :=
  2 * (m * cShiftU (P :: W * L' :: hdShape A B Cs m N).length
    (vol (P :: W * L' :: hdShape A B Cs m N)) Bb + 1) +
  2 * cPar W (P * L') (vol (hdShape A B Cs m N)) Bb +
  netCI (fun X => runCostI cS (fun _ _ => c0) (AModel.edgeSteps C N X))
    (W * vol (P :: L' :: hdShape A B Cs m N)) Bb C.Mout C.Min C.gates +
  cRebuild W (W * vol (P :: L' :: hdShape A B Cs m N)) Bb

/-- The frame predicate: `Π` is `0/1` and every shift of the Lemma 10 steps is implemented. -/
def FrOK (ShOK : Step (m + m) N → Prop) (X : IMat m) : Prop :=
  PiOK C.fac X ∧ ∀ i u c, Step.shift i u c ∈ AModel.edgeSteps C N X → ShOK (.shift i u c)

end Level

/-! ## The cost bound -/

section FrameSumNat

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

end FrameSumNat

section LevelCost

variable {m W : ℕ} [NeZero W] (C : Circuit (ZMod 2) m W) (N : ℕ) [NeZero N] (P L' : ℕ)
  (A B Cs : Shape)

/-- Number of frame changes of the network. -/
def nFr : ℕ := frameSum (fun _ => 1) C.gates C.Min C.Mout

/-- Number of shifts of all frame changes (`Σ_X (4m + 2 r_X)`). -/
def nShift : ℕ := frameSum (fun X => 4 * m + 2 * pivRank C X) C.gates C.Min C.Mout

/-- The constant of the linear part of one level. -/
def lvlK (nB : ℕ) : ℕ :=
  2 * (m * (2 * 10 ^ 8 + 10 ^ 6 * (nB + 1)) + 1) + 2 * (10 ^ 9 * (W + 1)) +
    nFr C * (1 + 100000) + (C.gates.length + 1) * ((W * (6 * W + 2) + 3) * 50000 + 1) + 1

end LevelCost

end NLogN.KIL

end

section
-- Source module: Solutions.KIL3UnitLevel
/-!
# KIL3: one Prop. 11 level with arbitrary fixed unit coefficients

The level code of `KIL3Level` with the shift code `scodeF`: a shift step with coefficient
`u = ±k·d⁻¹` runs `shFrac` (`KIL3Unit`).  Every unit of `ZMod N` that is (the image of) a fixed
fraction `±k/d` with `k, d ≤ D` prime to `N` is handled (`FracUnit D N u`), so the constant
`K` depends only on the circuit and `D`.

* `shiftF_runs`: one shift step as an `IRuns` statement, cost `cShiftD D |sh| (vol sh) Bb`.
* `levelR_il_frac`: the level as IL code, cost `K · (V + 1) · (Bb + 2)² + rankSum C · (c₀ + 2)`.
-/

set_option linter.unusedSimpArgs false
set_option linter.unnecessarySeqFocus false

namespace NLogN.KIL

open IntMul RS AModel AModelL12 AModelFac

/-! ## Representations of units as fractions -/

/-- `u = ±k·d⁻¹` with `1 ≤ k, d ≤ D`, both prime to `Q`. -/
def FracUnit (D Q : ℕ) (u : ZMod Q) : Prop :=
  ∃ k d : ℕ, ∃ neg : Bool, 1 ≤ k ∧ k ≤ D ∧ 1 ≤ d ∧ d ≤ D ∧ Nat.Coprime Q k ∧
    Nat.Coprime Q d ∧ u = fracU Q k d neg

/-- A representation reduced below `Q`. -/
def RepOK (D Q : ℕ) (r : ℕ × ℕ × Bool) (u : ZMod Q) : Prop :=
  1 ≤ r.1 ∧ r.1 ≤ D ∧ r.1 ≤ Q ∧ 1 ≤ r.2.1 ∧ r.2.1 ≤ D ∧ r.2.1 ≤ Q ∧ Nat.Coprime Q r.1 ∧
    Nat.Coprime Q r.2.1 ∧ u = fracU Q r.1 r.2.1 r.2.2

/-- A choice of representations. -/
noncomputable def repOf (D Q : ℕ) (u : ZMod Q) : ℕ × ℕ × Bool :=
  @dite _ (∃ r, RepOK D Q r u) (Classical.propDecidable _) (fun h => Classical.choose h)
    (fun _ => (1, 1, false))

/-! ## The shift code -/

section RunF

variable {sh : Shape} {n Q : ℕ} (Λ : Layout sh n Q) {Pp : ℕ} {I : List Sym}
  {iprog : Fin Pp → IStmt HS Pp}

/-- Shift code for coefficients `±k·d⁻¹`, given representations `rep`. -/
def scodeF (rep : ZMod Q → ℕ × ℕ × Bool) : Step n Q → IStmt HS Pp
  | .shift i u c =>
    (shFrac Q (rep u).1 (rep u).2.1 (vol (sh.take (Λ.pos i))) (vol (sh.drop (Λ.pos i + 1)))
      (rep u).2.2 ((sh.take (Λ.pos i)).zip (csOf Λ c (sh.take (Λ.pos i))))).code
  | .swap _ _ => .skip

/-- Cost bound of one shift with `k, d ≤ D`. -/
def cShiftD (D nA V Bb : ℕ) : ℕ := 2 * (D + 1) ^ 2 * 10 ^ 6 * (V + 1) * (Bb + 2) + cShift nA V Bb

end RunF

/-! ## The level -/

section LevelF

variable {m W : ℕ} [NeZero W] (C : Circuit (ZMod 2) m W) (N : ℕ) [NeZero N] (A B Cs : Shape)

/-- Shifts handled by `scodeF`. -/
def ShFrac (D : ℕ) : Step (m + m) N → Prop
  | .shift i u c => FracUnit D N u ∧ ∀ j, i ≤ j → c j = 0
  | .swap _ _ => False

/-- The constant of one level with fraction units (`k, d ≤ D`), shape length `nB`. -/
def Kf (D nB : ℕ) : ℕ :=
  lvlK C nB + nShift C * (2 * (D + 1) ^ 2 * 10 ^ 6 + 2 * 10 ^ 8 + 10 ^ 6 * (nB + 1))

end LevelF

end NLogN.KIL

end

section
-- Source module: Solutions.KIL3Prog
/-!
# KIL3: the recursive interchange program `prog` as IL procedures

`AModelProp11.prog C q k P L A B Cs` (bit data, `ZMod 2`) is realised by a finite table of IL
procedures, one per configuration `(level, rows, A, B, Cs)` reachable from the root, each level
`j + 1` procedure being the level code (`levelC` with the unit shift code `scodeF`) whose
interchanges call the level-`j` procedures of the children, and each level-`0` procedure being the
digit interchange (two Lemma 3 field moves).

* `swapB_sp`: the level-`0` interchange `pull … (swapPos …)`.
* `levelF_run`: one level with explicit code and cost (refines `levelR_il_frac`).
* `prog_il`: the table computes `prog` (by induction on the level), with cost `ilCost`, a
  recursion in the style of `progCost`.
-/

set_option linter.unusedSimpArgs false
set_option linter.unnecessarySeqFocus false

namespace NLogN.KIL

open IntMul RS AModel AModelL12 AModelFac

/-! ## Parameters on stacks `1 … 4` -/

/-- Push `a, b, c, d` on stacks `1 … 4`, run `p`, pop them. -/
def withP4 (a b c d : ℕ) (p : RS) : RS :=
  .seq (.push 1 (.lit (a : ℤ))) (.seq (.push 2 (.lit (b : ℤ))) (.seq (.push 3 (.lit (c : ℤ)))
    (.seq (.push 4 (.lit (d : ℤ))) (.seq p (.seq (.pop 1) (.seq (.pop 2) (.seq (.pop 3) (.pop 4))))))))

/-! ## Level `0`: the digit interchange by two field moves -/

/-- The level-`0` code on `A0 ++ [q] ++ B ++ [q] ++ Cs`. -/
def swapB (A0 B Cs : Shape) (q : ℕ) : RS :=
  .seq (withP4 (vol A0) q (vol (B ++ [q])) (vol Cs) fwdP)
    (withP4 (vol A0) q (vol B) (vol ([q] ++ Cs)) bwdP)

def cSwapB (A0 B Cs : Shape) (q Bb : ℕ) : ℕ :=
  (4 * ((1 + 1) * 2000 * (Bb + 2)) + cMove (vol A0) q (vol (B ++ [q])) (vol Cs) Bb +
    4 * (2 * Bb + 2)) +
  (4 * ((1 + 1) * 2000 * (Bb + 2)) + cMove (vol A0) q (vol B) (vol ([q] ++ Cs)) Bb +
    4 * (2 * Bb + 2))

/-! ## Lengths -/

section Lengths

variable {m W : ℕ} [NeZero W] (C : Circuit (ZMod 2) m W)

end Lengths

/-! ## The configurations and the procedure table -/

/-- A configuration: level, row count `L`, and the shapes `A`, `B`, `Cs`. -/
abbrev Cfg := ℕ × ℕ × Shape × Shape × Shape

/-- The child configuration of `c` for the interchange of `a` and `b` (as in `prog`). -/
def childCfg (m W q : ℕ) (c : Cfg) (a b : Fin (m + m)) : Cfg :=
  (c.1 - 1, c.2.1 / W, c.2.2.1 ++ List.replicate a (q ^ m ^ (c.1 - 1)),
    List.replicate (m - 1 - a) (q ^ m ^ (c.1 - 1)) ++
      (c.2.2.2.1 ++ List.replicate (b - m) (q ^ m ^ (c.1 - 1))),
    List.replicate (m + m - 1 - b) (q ^ m ^ (c.1 - 1)) ++ c.2.2.2.2)

/-- All configurations reachable in `f` steps. -/
def allCfg (m W q : ℕ) : ℕ → Cfg → List Cfg
  | 0, c => [c]
  | f + 1, c => c :: (List.finRange (m + m)).flatMap fun a =>
      (List.finRange (m + m)).flatMap fun b => allCfg m W q f (childCfg m W q c a b)

theorem self_mem_allCfg (m W q : ℕ) : ∀ f (c : Cfg), c ∈ allCfg m W q f c
  | 0, _ => List.mem_singleton_self _
  | _ + 1, _ => List.mem_cons_self

theorem allCfg_pos (m W q f : ℕ) (c : Cfg) : 0 < (allCfg m W q f c).length :=
  List.length_pos_of_mem (self_mem_allCfg m W q f c)

section Table

variable {m W : ℕ} [NeZero W] (C : Circuit (ZMod 2) m W) (D q P k : ℕ) (root : Cfg)

/-- Number of procedures. -/
def progPp : ℕ := (allCfg m W q k root).length

/-- The procedure index of a configuration. -/
def cfgIdx (c : Cfg) : Fin (progPp (m := m) (W := W) q k root) :=
  if h : c ∈ allCfg m W q k root then ⟨_, List.idxOf_lt_length_of_mem h⟩
  else ⟨0, allCfg_pos m W q k root⟩

/-- The code of a configuration: the digit interchange at level `0`, the level code above. -/
noncomputable def cfgCode (c : Cfg) : IStmt HS (progPp (m := m) (W := W) q k root) :=
  if c.1 = 0 then (swapB (P :: c.2.1 :: c.2.2.1) c.2.2.2.1 c.2.2.2.2 q).code
  else levelC C (q ^ m ^ (c.1 - 1)) P (c.2.1 / W) c.2.2.1 c.2.2.2.1 c.2.2.2.2
    (scodeF ((hdLayout c.2.2.1 c.2.2.2.1 c.2.2.2.2 m (q ^ m ^ (c.1 - 1))).cons2 P (c.2.1 / W))
      (repOf D (q ^ m ^ (c.1 - 1))))
    (fun a b => cfgIdx q k root (childCfg m W q c a b))

/-- **The procedure table** of `prog C q k P L A B Cs`. -/
noncomputable def progTbl : Fin (progPp (m := m) (W := W) q k root) → IStmt HS (progPp (m := m) (W := W) q k root) :=
  fun i => cfgCode C D q P k root ((allCfg m W q k root)[(i : ℕ)]'i.2)

end Table

/-! ## Cost of the recursion -/

/-- Cost bound of the level-`j` procedure on volume `V` with `n = |A| + |B| + |Cs|`
(the analogue of `progCost`). -/
def ilCost {m W : ℕ} (C : Circuit (ZMod 2) m W) (D q Bb : ℕ) : ℕ → ℕ → ℕ → ℕ
  | 0, V, _ => 2 * (16000 * (Bb + 2) + 300000 * q * (V + 1) * (Bb + 2) +
      2 * 10 ^ 8 * (Bb + 2) ^ 2 + 8 * Bb + 8)
  | j + 1, V, n => Kf C D (n + 2 * m + 2) * ((V + 1) * (Bb + 2) ^ 2) +
      rankSum C * (ilCost C D q Bb j (V / W) (n + 2 * (m - 1)) + 2)

/-! ## The recursion -/

section Main

variable {m W : ℕ} [NeZero W] (C : Circuit (ZMod 2) m W)

/-- The root configuration. -/
abbrev rootCfg (k L : ℕ) (A B Cs : Shape) : Cfg := (k, L, A, B, Cs)

end Main

end NLogN.KIL

end

section
-- Source module: Solutions.KIL4Sim
/-!
# KIL4: literal-to-register simulation of register programs

The KIL1–KIL3 programs carry the sizes of a configuration (field sizes, volumes, the modulus,
pass counts, …) as literals `Ex.lit v`.  A uniform program reads them from parameter stacks
instead.  `RSim π p p'` says that `p'` is `p` with some literal occurrences `.lit v` replaced by
register reads `.reg i`, where the parameter map `π` says `π i = some v`.

`RSim.sim`: on a state where every parameter stack `i` holds `v` on top (`PV π s`) and is not
written by `p`, the program `p'` is executable whenever `p` is, ends in the same state, and costs at
most `200 ×` the cost of `p` (a register read costs at most `200 ×` a literal of the same value).
Hence every KIL specification `Sp p s t c` transfers to `Sp p' s t (200 * c)` (`Sp.sim`).
-/

set_option linter.unusedSimpArgs false
set_option linter.unnecessarySeqFocus false

namespace NLogN.KIL

open IntMul RS

/-- Parameter map: stack `i` holds the literal `v` when `π i = some v`. -/
abbrev PMap := Fin HS → Option ℤ

/-- `e'` is `e` with some literals replaced by parameter reads. -/
inductive ESim (π : PMap) : Ex HS → Ex HS → Prop
  | reg (i : Fin HS) : ESim π (.reg i) (.reg i)
  | lit (c : ℤ) : ESim π (.lit c) (.lit c)
  | par {i : Fin HS} {v : ℤ} : π i = some v → ESim π (.lit v) (.reg i)
  | add {a a' b b' : Ex HS} : ESim π a a' → ESim π b b' → ESim π (.add a b) (.add a' b')
  | neg {a a' : Ex HS} : ESim π a a' → ESim π (.neg a) (.neg a')
  | mul {a a' b b' : Ex HS} : ESim π a a' → ESim π b b' → ESim π (.mul a b) (.mul a' b')
  | shl {a a' b b' : Ex HS} : ESim π a a' → ESim π b b' → ESim π (.shl a b) (.shl a' b')
  | quo {a a' b b' : Ex HS} : ESim π a a' → ESim π b b' → ESim π (.quo a b) (.quo a' b')
  | rem {a a' b b' : Ex HS} : ESim π a a' → ESim π b b' → ESim π (.rem a b) (.rem a' b')
  | pow2 {a a' : Ex HS} : ESim π a a' → ESim π (.pow2 a) (.pow2 a')
  | sel {c c' a a' b b' : Ex HS} : ESim π c c' → ESim π a a' → ESim π b b' →
      ESim π (.sel c a b) (.sel c' a' b')
  | ndiv {a a' b b' : Ex HS} : ESim π a a' → ESim π b b' → ESim π (.ndiv a b) (.ndiv a' b')
  | nmod {a a' b b' : Ex HS} : ESim π a a' → ESim π b b' → ESim π (.nmod a b) (.nmod a' b')

/-- The parameter stacks hold their values. -/
def PV (π : PMap) (s : IState HS) : Prop := ∀ i v, π i = some v → ∃ L, s.G i = .int v :: L

/-- Parameter stacks are ordinary stacks (not expression scratch). -/
def PLow (π : PMap) : Prop := ∀ i v, π i = some v → i.val < 56

section ESimLemmas

variable {π : PMap} {s : IState HS}

end ESimLemmas

/-- `p'` is `p` with some literals replaced by parameter reads. -/
inductive RSim (π : PMap) : RS → RS → Prop
  | skip : RSim π .skip .skip
  | asg (i : Fin HS) {e e' : Ex HS} : ESim π e e' → RSim π (.asg i e) (.asg i e')
  | push (i : Fin HS) {e e' : Ex HS} : ESim π e e' → RSim π (.push i e) (.push i e')
  | pop (i : Fin HS) : RSim π (.pop i) (.pop i)
  | seq {a a' b b' : RS} : RSim π a a' → RSim π b b' → RSim π (.seq a b) (.seq a' b')
  | ite {c c' : Ex HS} {a a' b b' : RS} : ESim π c c' → RSim π a a' → RSim π b b' →
      RSim π (.ite c a b) (.ite c' a' b')
  | rep (i : Fin HS) {n n' : Ex HS} {b b' : RS} : ESim π n n' → RSim π b b' →
      RSim π (.rep i n b) (.rep i n' b')

/-! ## Parameter maps and states -/

/-- The parameter map of a list of (stack, value) pairs. -/
def pmOf : List (Fin HS × ℤ) → PMap
  | [] => fun _ => none
  | (i, v) :: l => fun j => if j = i then some v else pmOf l j

end NLogN.KIL

end

section
-- Source module: Solutions.KIL4Tbl
/-!
# KIL4: the offset table generator with run-time stage data

KIL2's `genP Q ac` has one stage per field of the prefix `A` of the target, every number (`Q`,
the radix `a`, the control `c`, the suffix volume `M`) a literal.  The uniform generator `genPU n`
reads the stages from four stage stacks (`34 … 37`: `a`, `c`, `M`, `a·M`, the last stage on top) and
`-Q` from the parameter stack `23`; its code depends only on the number `n` of stages.

* `gStageT`, `rsim_gStage`: the stage as a template; the literal stage simulates it.
* `stagesU_sp`, `genPU_sp`: the generator pushes `tbl Q ac` on stack `14`, pops the stage stacks.
* `tbl_zeros`, `tbl_app_congr`: the table only depends on the controls; runs of control-`0`
  fields collapse to one field of their volume (the "collapsed" stage lists of the uniform
  shifts have at most `2m + 3` stages whatever the shape).
-/

set_option linter.unusedSimpArgs false
set_option linter.unnecessarySeqFocus false

namespace NLogN.KIL

open IntMul RS AModel

/-! ## Parameter stacks of the uniform shift -/

/-- `Q`. -/
abbrev sQQ : Fin HS := 22
/-- `-Q`. -/
abbrev sNQ : Fin HS := 23
/-- Prefix volume `P`. -/
abbrev sPA : Fin HS := 24
/-- Suffix volume `B`. -/
abbrev sPC : Fin HS := 25
/-- `P·(Q·B)`. -/
abbrev sVV : Fin HS := 26
/-- `Q·B`. -/
abbrev sQB : Fin HS := 27
/-- Residue `Q mod c`. -/
abbrev sRR : Fin HS := 28
/-- Count of a pass. -/
abbrev sCN : Fin HS := 29
/-- Stage stacks. -/
abbrev sSA : Fin HS := 34
abbrev sSC : Fin HS := 35
abbrev sSM : Fin HS := 36
abbrev sSAM : Fin HS := 37

/-! ## Templates -/

/-- `mAdd` with `-Q` given by an expression. -/
def mAddT (eNQ x y : Ex HS) : Ex HS :=
  .sel (.add (.add x y) eNQ) (.add x y) (.add (.add x y) eNQ)

def gInnerT (eNQ : Ex HS) : RS :=
  .seq (.push 11 (mAddT eNQ (.reg 14) (.reg 12))) (.seq (.push 10 (.reg 14)) (.pop 14))

def gBodyT (eNQ ec eM : Ex HS) : RS :=
  .seq (.rep 31 eM (gInnerT eNQ)) (.seq (NLogN.mvK 10 14 33 eM) (.asg 12 (mAddT eNQ (.reg 12) ec)))

def gStageT (eNQ ea ec eM eaM : Ex HS) : RS :=
  .seq (.push 12 (.lit 0)) (.seq (.rep 30 ea (gBodyT eNQ ec eM)) (.seq (.pop 12)
    (.seq (.rep 31 eM (.pop 14)) (NLogN.mvK 11 14 33 eaM))))

/-! ## The literal stage in stack form -/

/-! ## One uniform stage -/

/-- The uniform stage: parameters from the stage stacks, then the stage data popped. -/
def stageU : RS :=
  .seq (gStageT (.reg sNQ) (.reg sSA) (.reg sSC) (.reg sSM) (.reg sSAM))
    (.seq (.pop sSA) (.seq (.pop sSC) (.seq (.pop sSM) (.pop sSAM))))

def cStageU (a M Bb : ℕ) : ℕ := 200 * cStage a M Bb + 4 * (2 * Bb + 2)

/-! ## The generator -/

/-- `n` uniform stages. -/
def stagesU : ℕ → RS
  | 0 => .skip
  | n + 1 => .seq stageU (stagesU n)

/-- The uniform generator for `n` stages. -/
def genPU (n : ℕ) : RS := .seq (.push 14 (.lit 0)) (stagesU n)

/-- Stage stack contents for the stages `rv` (top first; `rv` is the stage list reversed)
above the already generated suffix `suf`: the radix, the control, the suffix volume `M`, `a·M`. -/
def stgM : List (ℕ × ℕ) → List (ℕ × ℕ) → List ℤ
  | [], _ => []
  | x :: rv, suf => (vol (suf.map Prod.fst) : ℤ) :: stgM rv (x :: suf)

def stgAM : List (ℕ × ℕ) → List (ℕ × ℕ) → List ℤ
  | [], _ => []
  | x :: rv, suf => ((x.1 * vol (suf.map Prod.fst) : ℕ) : ℤ) :: stgAM rv (x :: suf)

/-- Cost of the uniform stages. -/
def stCost (Bb : ℕ) : List (ℕ × ℕ) → List (ℕ × ℕ) → ℕ
  | [], _ => 1
  | x :: rv, suf => cStageU x.1 (vol (suf.map Prod.fst)) Bb + stCost Bb rv (x :: suf)

/-! ## Cost of the generator -/

/-! ## Collapsing control-`0` fields of the table -/

end NLogN.KIL

end

section
-- Source module: Solutions.KIL4Pass
/-!
# KIL4: multiplication of the target field by `c` or `c⁻¹` with run-time modulus

KIL3's `divPg P Q c B` / `mulPg P Q c B` (the `c` gather / scatter passes) have `P, Q, B`, the
pass counts `cntG Q c j · B` and the class of every residue (`sig Q c v`, which depends on
`Q mod c`) compiled in.  The uniform versions `divPgU c` / `mulPgU c` (code depending only on the
fixed factor `c`):

* read `Q, P, B, P·(Q·B), Q·B` from the parameter stacks `22, 24, 25, 26, 27`;
* compute `r = Q mod c` and branch on it (`sel`, `c` branches, one per residue class: the class
  function `sig r c` is compiled per branch);
* compute the count of every pass at run time, `cntG Q c j = (Q / c)·A_j + B_j` with constants
  `A_j, B_j` of the branch (`cntG_eq`), pushed on stack `29` around the pass.

Every pass is the KIL3 pass with its literals replaced by register reads (`RSim`), so its
specification is transferred by `Sp.sim`.
-/

set_option linter.unusedSimpArgs false
set_option linter.unnecessarySeqFocus false

namespace NLogN.KIL

open IntMul RS AModel

/-! ## Generic helpers -/

/-! ## Counting over a period -/

/-- `#{y < n | p y}`. -/
def cntR (p : ℕ → Bool) (n : ℕ) : ℕ := ((List.range n).filter p).length

/-- Constants of the gather pass counts. -/
def cGA (r c j : ℕ) : ℕ := cntR (fun v => decide (j < sig r c v)) c
def cGB (r c j : ℕ) : ℕ := cntR (fun v => decide (j < sig r c v)) r

/-- Constants of the scatter pass counts. -/
def cSA (r c j : ℕ) : ℕ := cntR (fun v => decide (sig r c v ≤ j)) c
def cSB (r c j : ℕ) : ℕ := cntR (fun v => decide (sig r c v ≤ j)) r

/-! ## Pass templates -/

def brGT (r c j : ℕ) (eB : Ex HS) (v : ℕ) : RS :=
  if sig r c v < j then .skip
  else if sig r c v = j then NLogN.mvK 0 12 33 eB else NLogN.mvK 0 10 33 eB

def bodyGT (r c j : ℕ) (eB : Ex HS) : RS :=
  .seq (NLogN.KIL.sel 13 (brGT r c j eB) 0 c) (.asg 13 (mAdd c (.reg 13) (.lit 1)))

def passGT (r c j : ℕ) (eQ eB eCnt : Ex HS) : RS :=
  .seq (.push 13 (.lit 0)) (.seq (.rep 31 eQ (bodyGT r c j eB)) (.seq (.pop 13)
    (NLogN.mvK 10 0 33 eCnt)))

def brST (r c j : ℕ) (eB : Ex HS) (v : ℕ) : RS :=
  if sig r c v < j then NLogN.mvK 10 11 33 eB
  else if sig r c v = j then NLogN.mvK 0 11 33 eB else .skip

def bodyST (r c j : ℕ) (eB : Ex HS) : RS :=
  .seq (NLogN.KIL.sel 13 (brST r c j eB) 0 c) (.asg 13 (mAdd c (.reg 13) (.lit 1)))

def passST (r c j : ℕ) (eQ eB eCnt : Ex HS) : RS :=
  .seq (.push 13 (.lit 0)) (.seq (.rep 31 eQ (bodyST r c j eB)) (.seq (.pop 13)
    (NLogN.mvK 11 10 33 eCnt)))

end NLogN.KIL

end

section
-- Source module: Solutions.KIL4Div
/-!
# KIL4: the uniform `c` gather (`c⁻¹·`) with run-time `Q, P, B`

`divPgU c` depends only on the fixed factor `c`.  It reads `Q` (stack `22`), `P` (stack `24`),
`B` (stack `25`), `P·(Q·B)` (stack `26`) and computes the residue `Q mod c` (pushed on stack `28`)
and the count of every pass (pushed on stack `29` around the pass).  `divPgU_spec`: it computes
the same array as KIL3's `divPg P Q c B` (the KIL3 specification is transferred by `Sp.sim`).
-/

set_option linter.unusedSimpArgs false
set_option linter.unnecessarySeqFocus false

namespace NLogN.KIL

open IntMul RS AModel

/-- The run-time pass count `((Q / c)·A + B')·B`. -/
def cntE (c a b : ℕ) : Ex HS :=
  .mul (.add (.mul (.ndiv (.reg sQQ) (.lit c)) (.lit a)) (.lit b)) (.reg sPC)

/-! ## One gather pass with run-time `Q, B` and count -/

def passGU (r c j : ℕ) : RS :=
  .seq (.push sCN (cntE c (cGA r c j) (cGB r c j)))
    (.seq (passGT r c j (.reg sQQ) (.reg sPC) (.reg sCN)) (.pop sCN))

def cPassGU (Q c B Bb : ℕ) : ℕ := 10 * 10 ^ 7 * (Bb + 2) ^ 2 + (200 * cPassG Q c B Bb + (2 * Bb + 2))

/-- Passes `0, …, k − 1`. -/
def passesGU (r c : ℕ) : ℕ → RS
  | 0 => .skip
  | k + 1 => .seq (passesGU r c k) (passGU r c k)

def cBlkGU (Q c B Bb : ℕ) : ℕ := 1 + c * cPassGU Q c B Bb

/-- The gather with residue `r = Q mod c` compiled in, `P` and `P·(Q·B)` read at run time. -/
def divPgR (r c : ℕ) : RS :=
  .seq (.rep 30 (.reg sPA) (passesGU r c c)) (NLogN.mvK 12 0 33 (.reg sVV))

def cDivPgR (P Q c B Bb : ℕ) : ℕ :=
  (1 + 1) * 2000 * (Bb + 2) + P * (cBlkGU Q c B Bb + 8000 * (Bb + 2) + 30) + 32 +
    ((1 + 1) * 2000 * (Bb + 2) + P * (Q * B) * (14000 * (Bb + 2)) + 32)

/-- **The uniform gather** (multiplication of the target field by `c⁻¹`): branch on `Q mod c`. -/
def divPgU (c : ℕ) : RS :=
  .seq (.push sRR (.nmod (.reg sQQ) (.lit c)))
    (.seq (NLogN.KIL.sel sRR (fun r => divPgR r c) 0 c) (.pop sRR))

def cDivPgU (P Q c B Bb : ℕ) : ℕ :=
  4 * 10 ^ 7 * (Bb + 2) ^ 2 + (c * cSel Bb + cDivPgR P Q c B Bb + (2 * Bb + 2))

/-- The uniform `fdivP`. -/
def fdivPU (d : ℕ) : RS := if d < 2 then .skip else divPgU d

end NLogN.KIL

end

section
-- Source module: Solutions.KIL4Mul
/-!
# KIL4: the uniform `c` scatter (`c·`) with run-time `Q, P, B`

`mulPgU c` depends only on the fixed factor `c` (as `divPgU c`, KIL4Div): parameters on stacks
`22, 24, 25, 26, 27` (`Q, P, B, P·(Q·B), Q·B`), residue `Q mod c` on `28`, pass counts on `29`.
`mulPgU_spec`: it computes the array of KIL3's `mulPg P Q c B`.  Also the uniform stages
`fdivPU d`, `fmulPU k` and their specifications.
-/

set_option linter.unusedSimpArgs false
set_option linter.unnecessarySeqFocus false

namespace NLogN.KIL

open IntMul RS AModel

def passSU (r c j : ℕ) : RS :=
  .seq (.push sCN (cntE c (cSA r c j) (cSB r c j)))
    (.seq (passST r c j (.reg sQQ) (.reg sPC) (.reg sCN)) (.pop sCN))

def passesSU (r c : ℕ) : ℕ → RS
  | 0 => .skip
  | k + 1 => .seq (passesSU r c k) (passSU r c k)

def blkSU (r c : ℕ) : RS := .seq (passesSU r c c) (NLogN.mvK 10 12 33 (.reg sQB))

def mulPgR (r c : ℕ) : RS :=
  .seq (.rep 30 (.reg sPA) (blkSU r c)) (NLogN.mvK 12 0 33 (.reg sVV))

def cBlkSU (Q c B Bb : ℕ) : ℕ :=
  1 + c * cPassGU Q c B Bb + ((1 + 1) * 2000 * (Bb + 2) + Q * B * (14000 * (Bb + 2)) + 32)

def cMulPgR (P Q c B Bb : ℕ) : ℕ :=
  (1 + 1) * 2000 * (Bb + 2) + P * (cBlkSU Q c B Bb + 8000 * (Bb + 2) + 30) + 32 +
    ((1 + 1) * 2000 * (Bb + 2) + P * (Q * B) * (14000 * (Bb + 2)) + 32)

/-- Branch on `Q mod c` (residue pushed on stack `28` around the branch). -/
def resSel (c : ℕ) (br : ℕ → RS) : RS :=
  .seq (.push sRR (.nmod (.reg sQQ) (.lit c))) (.seq (NLogN.KIL.sel sRR br 0 c) (.pop sRR))

/-- **The uniform scatter** (multiplication of the target field by `c`). -/
def mulPgU (c : ℕ) : RS := NLogN.KIL.resSel c (fun r => mulPgR r c)

def cMulPgU (P Q c B Bb : ℕ) : ℕ :=
  4 * 10 ^ 7 * (Bb + 2) ^ 2 + (c * cSel Bb + cMulPgR P Q c B Bb + (2 * Bb + 2))

/-- The uniform `fmulP`. -/
def fmulPU (k : ℕ) : RS := if k < 2 then .skip else mulPgU k

/-- The parameter facts the uniform stages read. -/
structure StP (t : IState HS) (P Q B : ℕ) : Prop where
  hQ : ∃ L, t.G sQQ = .int Q :: L
  hB : ∃ L, t.G sPC = .int B :: L
  hA : ∃ L, t.G sPA = .int P :: L
  hV : ∃ L, t.G sVV = .int ((P * (Q * B) : ℕ) : ℤ) :: L
  hQB : ∃ L, t.G sQB = .int ((Q * B : ℕ) : ℤ) :: L

end NLogN.KIL

namespace NLogN.KIL

end NLogN.KIL

end