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

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

import Definitions.Def_IntMul_Kappa_K_Defs9
import Definitions.Def_IntMul_HvdH_M_Defs8
import Definitions.Def_IntMul_HvdH_BApprox
import Definitions.Def_IntMul_HvdH_DiagApprox
import Definitions.Def_IntMul_HvdH_EApprox
import Definitions.Def_IntMul_HvdH_FPApprox
import Definitions.Def_IntMul_HvdH_JApprox
import Definitions.Def_IntMul_HvdH_SApprox

set_option autoImplicit false
set_option relaxedAutoImplicit false
set_option maxHeartbeats 1000000

section
-- Source module: Solutions.AModel
/-!
# Array model: mixed-radix arrays, pull maps, and linear address maps via Bruhat factors

Arrays are lists in row-major lexicographic order over fields (first field most significant).
All denotations are plain `List` functions.
-/

namespace AModel

open Matrix

/-! ## 1. Mixed-radix basics -/

/-! ## 2. Pull maps -/

section Pull
variable {α : Type*}

end Pull

/-! ## 3. Active fields -/

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

namespace Layout

end Layout

/-! ## 4. Controlled shift -/

/-! ## 5. Field transposition -/

section Shift

/-- Array action of the controlled shift: `pull` with the explicit inverse address map. -/
def shiftArr {α : Type*} (i : Fin n) (u : ZMod Q) (c : Fin n → ZMod Q) (d : α) (xs : List α) :
    List α :=
  AModel.pull sh sh (shiftInv Λ i u c) d xs

end Shift

section Swap

end Swap

/-! ## Primitive steps and their composition -/

namespace Step

variable [NeZero Q]

end Step

section Run
variable [NeZero Q]

end Run

/-! ## 6. Linear address maps through Bruhat factors -/

section Bruhat

variable {E : Matrix (Fin n) (Fin n) (ZMod Q)}

variable [NeZero Q]

end Bruhat

end AModel

end

section
-- Source module: Solutions.ButterflyPacked
/-!
# §5.2 Lemma 13: packed selected-bit addition on one rectangle

1. (20)/(21): eight controlled `±1` updates of `(u, w)` give `(u + c(1 - 2(u mod 2)), w)`.
2. The packed lines (22) on `L = fK`-bit chunks `y, z` (mod `2^L`), with the `f - 1` segments
   at `j_i = ρ + iK` and controls `c_i = x_{j_i}`: outside the guard set `B` (23), their
   composite `S0` equals the ideal map `T0` (flip `y_{j_i}` where `c_i = 1`, fix `x, z`).
3. Bijections agreeing off `B` with `T0(B) = B` satisfy `S0(B) = B` (the repair permutes `B`).
4. (24): the guard count `|B| · 2^K ≤ 80 (f-1) · 2^{3L}`.
-/

namespace Butterfly

open Finset

/-! ## 1. The eight scalar updates (20) -/

/-- One line of (20): `tgtY` = update `u` testing `w` (else update `w` testing `u`),
sign `sgn`, tested parity `b`. -/
structure Line where
  tgtY : Bool
  sgn : ℤ
  b : ℤ

/-- The eight lines of (20) in chronological order. -/
def lines : List Line :=
  [⟨true, -1, 0⟩, ⟨false, -1, 0⟩, ⟨true, 1, 0⟩, ⟨false, 1, 0⟩,
   ⟨false, 1, 1⟩, ⟨true, 1, 1⟩, ⟨false, -1, 1⟩, ⟨true, -1, 1⟩]

/-- A line on one segment pair with Boolean control `c`. -/
def Line.apply (l : Line) (c : Bool) (p : ℤ × ℤ) : ℤ × ℤ :=
  if l.tgtY then (p.1 + l.sgn * (if c ∧ p.2 % 2 = l.b then 1 else 0), p.2)
  else (p.1, p.2 + l.sgn * (if c ∧ p.1 % 2 = l.b then 1 else 0))

/-- The composite of the eight lines. -/
def run8 (c : Bool) (p : ℤ × ℤ) : ℤ × ℤ := lines.foldl (fun p l => l.apply c p) p

/-! ## 2. Segment packing -/

variable (K : ℕ)

/-- Pack `n` segments `u 0, …, u (n-1)` of `K` bits above a high part `h`:
`u 0 + 2^K (u 1 + 2^K (… + 2^K h))`. -/
def pk : ℕ → (ℕ → ℤ) → ℤ → ℤ
  | 0, _, h => h
  | n + 1, u, h => u 0 + 2 ^ K * pk n (fun i => u (i + 1)) h

/-- Bit `j` of an integer. -/
def bitAt (t : ℤ) (j : ℕ) : ℤ := (t / 2 ^ j) % 2

/-- A packed number `lo + 2^ρ · pk(u, h)`. -/
def packed (ρ n : ℕ) (lo : ℤ) (u : ℕ → ℤ) (h : ℤ) : ℤ := lo + 2 ^ ρ * Butterfly.pk K n u h

/-! ## 3. The packed lines (22) -/

variable (ρ f : ℕ)

/-- The offset of a packed line: `∑_{i<f-1} 2^{j_i} [x_{j_i} = 1][t_{j_i} = b]`. -/
def offset (x t b : ℤ) : ℤ :=
  ∑ i ∈ range (f - 1), (2 : ℤ) ^ (ρ + i * K) *
    (if bitAt x (ρ + i * K) = 1 ∧ bitAt t (ρ + i * K) = b then 1 else 0)

/-- A packed line (22) on the chunk pair `(y, z)` of width `L = fK`, modulo `2^L`. -/
def pline (x : ℤ) (l : Line) (p : ℤ × ℤ) : ℤ × ℤ :=
  if l.tgtY then ((p.1 + l.sgn * Butterfly.offset K ρ f x p.2 l.b) % 2 ^ (f * K), p.2)
  else (p.1, (p.2 + l.sgn * Butterfly.offset K ρ f x p.1 l.b) % 2 ^ (f * K))

/-- `S0`: the eight packed lines in order. -/
def S0 (x : ℤ) (p : ℤ × ℤ) : ℤ × ℤ := lines.foldl (fun p l => pline K ρ f x l p) p

/-- The control of segment `i`. -/
def ctl (x : ℤ) (i : ℕ) : Bool := decide (bitAt x (ρ + i * K) = 1)

/-- The ideal map `T0`: flip bit `j_i` of `y` for every `i < f-1` with control `1`. -/
def T0y (x y : ℤ) : ℤ :=
  y + ∑ i ∈ range (f - 1), (2 : ℤ) ^ (ρ + i * K) *
    ((if Butterfly.ctl K ρ x i then 1 else 0) * (1 - 2 * bitAt y (ρ + i * K)))

variable {K ρ f}

/-! ## 4. Address form: guards and the ideal map -/

variable (K ρ f)

/-- The guard value of segment `i` of chunk `t`: `g_i = ⌊t / 2^{j_i+1}⌋ mod 2^{K-1}`. -/
def guard (t : ℤ) (i : ℕ) : ℤ := (t / 2 ^ (ρ + i * K + 1)) % 2 ^ (K - 1)

/-- (23): every used guard value lies in `[10, 2^{K-1} - 11]`. -/
def GuardOK (t : ℤ) : Prop := ∀ i < f - 1, 10 ≤ guard K ρ t i ∧ guard K ρ t i ≤ 2 ^ (K - 1) - 11

variable {K ρ f}

/-! ## 5. The repair permutes `B` -/

/-! ## 6. The guard count (24) -/

/-- Guard `i` of a natural chunk value. -/
def guardN (K ρ : ℕ) (n i : ℕ) : ℕ := (n / 2 ^ (ρ + i * K + 1)) % 2 ^ (K - 1)

end Butterfly

end

section
-- Source module: Solutions.KIL6Den
/-!
# KIL6: Lemma 12 as a uniform IL program — the list denotations

The IL program of KIL6 realises the Lemma 12 interchange with the following data flow (array
model, `AModelL12`):

1. **binary padding** `binShape … (2^u) → binShape … (q^e)`: two one-field paddings
   (`padL`, `embed_binShape`);
2. **high swaps**: `r = min ρ e` single-digit interchanges (`highSwaps`);
3. if `n = e - r > 0`, the **middle part** `l12mid`: the row field is the prefix
   `Pre ++ [q]^r` (volume `vol Pre · q^r`), padded to `padRows W Kp (q^r)` rows (`padL`); the
   Prop. 11 oracle runs on the pieces (`runPieces … pieceLevels`), with the gap
   `G ++ [q]^r` (the already interchanged high digits of the second chunk belong to the gap);
   the padding is scanned away (`unpadL`);
4. **binary extraction** (two `unpadL`).

Compared with `AModelL12.l12q` this variant needs no joining rotation: the interchanged high
digits `d⁺` simply stay inside the gap.  `l12v_correct` proves that it is the interchange
`swapPos` of the two `[2^u]` fields, from the same oracle specification `OracleSpec`.

The list functions `padL` / `unpadL` are the one-field embeddings / extractions in the chunk
form the IL loops produce (`embed_eq_padL`, `extract_eq_unpadL`).
-/

set_option linter.unusedSimpArgs false
set_option linter.unnecessarySeqFocus false

namespace NLogN.KIL

open AModel AModelL12

section Chunks

variable {α : Type*}

end Chunks

/-! ## One-field padding and extraction in chunk form -/

section Pad

variable {α : Type*}

/-- Pad the field `F ↦ F'` of `[X, F, Y]` (row-major): every block of `F · Y` entries is followed
by `(F' - F) · Y` default entries. -/
def padL (X F F' Y : ℕ) (d : α) (xs : List α) : List α :=
  (List.range X).flatMap fun x =>
    (xs.drop (x * (F * Y))).take (F * Y) ++ List.replicate ((F' - F) * Y) d

/-- Extract the field `F` from `[X, F', Y]`: keep the first `F · Y` entries of every block of
`F' · Y`. -/
def unpadL (X F F' Y : ℕ) (ys : List α) : List α :=
  (List.range X).flatMap fun x => (ys.drop (x * (F' * Y))).take (F * Y)

/-! ## Composing embeddings and extractions -/

end Pad

/-! ## The binary and row paddings in chunk form -/

section Shapes

variable {α : Type*}

end Shapes

/-! ## The variant procedure and its correctness -/

section Variant

variable {α : Type*}

/-- The middle part: rows `Pre ++ [q]^r` (volume `vol Pre · q^r`), padded to
`padRows W Kp (q^r)`; gap `G ++ [q]^r`; the Prop. 11 oracle on the pieces of the remaining
`n` digits. -/
noncomputable def l12mid (child : Oracle α) (W m q Kp r : ℕ) (Pre G Suf : Shape) (n : ℕ) (d : α)
    (ys : List α) : List α :=
  extract (rowShape (AModel.vol Pre) (q ^ r) (G ++ List.replicate r q) Suf q n)
    (rowShape (AModel.vol Pre) (padRows W Kp (q ^ r)) (G ++ List.replicate r q) Suf q n) d
    (runPieces (pieceStep child (AModel.vol Pre) (padRows W Kp (q ^ r)) (G ++ List.replicate r q) Suf q m n)
      (fun j => m ^ j) 0 (pieceLevels m n)
      (AModelL12.embed (rowShape (AModel.vol Pre) (q ^ r) (G ++ List.replicate r q) Suf q n)
        (rowShape (AModel.vol Pre) (padRows W Kp (q ^ r)) (G ++ List.replicate r q) Suf q n) d ys))

/-- **The Lemma 12 procedure realised by the KIL6 program**: binary padding, `min ρ e` high
digit swaps, the middle part (if digits remain), binary extraction. -/
noncomputable def l12v (child : Oracle α) (W m q Kp ρ e u : ℕ) (Pre G Suf : Shape) (d : α)
    (xs : List α) : List α :=
  extract (binShape Pre G Suf (2 ^ u)) (binShape Pre G Suf (q ^ e)) d
    (if e - min ρ e = 0 then
      highSwaps (qShape Pre G Suf q e) Pre.length (Pre.length + e + G.length) (min ρ e) d
        (AModelL12.embed (binShape Pre G Suf (2 ^ u)) (binShape Pre G Suf (q ^ e)) d xs)
    else l12mid child W m q Kp (min ρ e) Pre G Suf (e - min ρ e) d
      (highSwaps (qShape Pre G Suf q e) Pre.length (Pre.length + e + G.length) (min ρ e) d
        (AModelL12.embed (binShape Pre G Suf (2 ^ u)) (binShape Pre G Suf (q ^ e)) d xs)))

end Variant

/-! ## Auxiliary facts: pieces, digits, maps -/

section Aux

variable {α β : Type*}

end Aux

end NLogN.KIL

end

section
-- Source module: NLogN.SSBase
/-!
# Schönhage–Strassen in the item language: stack layout and Fermat-ring arithmetic

The program uses `NS` item stacks with fixed roles, so that stack (in)equalities are decided by
`simp`.  The current ring `ℤ/(2^n+1)` is given by the ruler `2^n` on top of `RN` and the modulus
`2^n+1` on top of `RF`.
-/

namespace NLogN

open IntMul

abbrev ARG : Fin NS := 0
abbrev PRM : Fin NS := 1

/-! ## Fermat reduction -/

variable {P : ℕ}

end NLogN

namespace NLogN

open IntMul

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

end NLogN

end

section
-- Source module: Solutions.KILBase
/-!
# KIL base: register programs on "managed" stacks

`encK s f` is the state `s` with the list of integers `f j` pushed on top of every stack `j`
(`f j = []` for unmanaged stacks).  All programs of the KIL files are specified on such states;
updates are `Function.update`s of `f`, which keeps the bookkeeping small.

`Sp p s t c` : the register program `p` is executable from `s` (`RS.ok`), ends in `t`, and costs
at most `c`.  `Sp.irun` turns it into an `IRuns` of the compiled item-language code.
-/

set_option linter.unusedSimpArgs false

namespace NLogN

open IntMul RS

namespace KIL

/-- `s` with `f j` on top of each stack `j`. -/
def encK (s : IState HS) (f : Fin HS → List ℤ) : IState HS :=
  { s with G := fun j => ints (f j) ++ s.G j }

@[simp] theorem encK_G (s : IState HS) (f : Fin HS → List ℤ) (j : Fin HS) :
    (NLogN.KIL.encK s f).G j = ints (f j) ++ s.G j := rfl
@[simp] theorem encK_ip (s : IState HS) (f : Fin HS → List ℤ) : (NLogN.KIL.encK s f).ip = s.ip := rfl
@[simp] theorem encK_out (s : IState HS) (f : Fin HS → List ℤ) : (NLogN.KIL.encK s f).out = s.out := rfl

/-- Specification of a register program. -/
def Sp (p : RS) (s t : IState HS) (c : ℕ) : Prop := p.ok s ∧ p.run s = t ∧ p.cost s ≤ c

/-! ## Elementary statements -/

section
variable (s : IState HS) (f : Fin HS → List ℤ)

end

/-! ## Moving integers between stacks -/

/-! ## Counted loops through a family of managed states -/

/-! ## Expression side conditions on managed states -/

end KIL

end NLogN

end

section
-- Source module: Solutions.KILState
/-!
# KIL managed states

`KS` names the contents of the managed stacks used by the KIL programs; `toF` places them on
their stacks.  `upd_*` rewrite a `Function.update` of a managed state into a record update and
`toF_*` read a field.
-/

namespace NLogN

namespace KIL

/-- Contents of the managed stacks (each a list of integers, top first). -/
structure KS where
  x : List ℤ := []
  P : List ℤ := []
  q : List ℤ := []
  M : List ℤ := []
  B : List ℤ := []
  MB : List ℤ := []
  u : List ℤ := []
  z : List ℤ := []
  V : List ℤ := []
  qB : List ℤ := []
  a : List ℤ := []
  c : List ℤ := []
  o : List ℤ := []
  qMB : List ℤ := []
  w : List ℤ := []
  k0 : List ℤ := []
  k1 : List ℤ := []
  k2 : List ℤ := []
  k3 : List ℤ := []

/-- The managed stack contents as a function on stacks. -/
def toF (k : NLogN.KIL.KS) (j : Fin HS) : List ℤ :=
  match j.val with
  | 0 => k.x
  | 1 => k.P
  | 2 => k.q
  | 3 => k.M
  | 4 => k.B
  | 5 => k.MB
  | 6 => k.u
  | 7 => k.z
  | 8 => k.V
  | 9 => k.qB
  | 10 => k.a
  | 11 => k.c
  | 12 => k.o
  | 13 => k.qMB
  | 14 => k.w
  | 30 => k.k0
  | 31 => k.k1
  | 32 => k.k2
  | 33 => k.k3
  | _ => []

@[simp] theorem toF_x (k : NLogN.KIL.KS) : toF k 0 = k.x := rfl

@[simp] theorem toF_P (k : NLogN.KIL.KS) : toF k 1 = k.P := rfl

@[simp] theorem toF_q (k : NLogN.KIL.KS) : toF k 2 = k.q := rfl

@[simp] theorem toF_M (k : NLogN.KIL.KS) : toF k 3 = k.M := rfl

@[simp] theorem toF_B (k : NLogN.KIL.KS) : toF k 4 = k.B := rfl

@[simp] theorem toF_MB (k : NLogN.KIL.KS) : toF k 5 = k.MB := rfl

@[simp] theorem toF_u (k : NLogN.KIL.KS) : toF k 6 = k.u := rfl

@[simp] theorem toF_z (k : NLogN.KIL.KS) : toF k 7 = k.z := rfl

@[simp] theorem toF_V (k : NLogN.KIL.KS) : toF k 8 = k.V := rfl

@[simp] theorem toF_qB (k : NLogN.KIL.KS) : toF k 9 = k.qB := rfl

@[simp] theorem toF_a (k : NLogN.KIL.KS) : toF k 10 = k.a := rfl

@[simp] theorem toF_c (k : NLogN.KIL.KS) : toF k 11 = k.c := rfl

@[simp] theorem toF_o (k : NLogN.KIL.KS) : toF k 12 = k.o := rfl

@[simp] theorem toF_qMB (k : NLogN.KIL.KS) : toF k 13 = k.qMB := rfl

@[simp] theorem toF_w (k : NLogN.KIL.KS) : toF k 14 = k.w := rfl

@[simp] theorem toF_k0 (k : NLogN.KIL.KS) : toF k 30 = k.k0 := rfl

@[simp] theorem toF_k1 (k : NLogN.KIL.KS) : toF k 31 = k.k1 := rfl

@[simp] theorem toF_k2 (k : NLogN.KIL.KS) : toF k 32 = k.k2 := rfl

@[simp] theorem toF_k3 (k : NLogN.KIL.KS) : toF k 33 = k.k3 := rfl

end KIL

end NLogN

end

section
-- Source module: Solutions.KILList
/-!
# KIL list layer: block transposes and their `pull` denotations

Arrays on `[P, q, M, B]` are described by a record function `r p y m : List ℤ` (each of length
`B`).  `inp r P q M` lists the records in the order `(p, y, m)`, `outF r P q M` in the order
`(p, m, y)`.  Moving the field `q` behind the field `M` maps `inp` to `outF` (and back).
-/

namespace NLogN.KIL

open AModel

/-- One block in input order `(y, m)`. -/
def blkIn (r : ℕ → ℕ → ℕ → List ℤ) (q M p : ℕ) : List ℤ :=
  (List.range q).flatMap fun y => (List.range M).flatMap fun m => r p y m

/-- One block in output order `(m, y)`. -/
def blkOut (r : ℕ → ℕ → ℕ → List ℤ) (q M p : ℕ) : List ℤ :=
  (List.range M).flatMap fun m => (List.range q).flatMap fun y => r p y m

def inp (r : ℕ → ℕ → ℕ → List ℤ) (P q M : ℕ) : List ℤ := (List.range P).flatMap (blkIn r q M)

def outF (r : ℕ → ℕ → ℕ → List ℤ) (P q M : ℕ) : List ℤ := (List.range P).flatMap (blkOut r q M)

/-- Records of an array stored on `[P, q, M, B]`. -/
def recOf (xs : List ℤ) (q M B p y m : ℕ) : List ℤ :=
  (List.range B).map fun b => xs.getD (((p * q + y) * M + m) * B + b) 0

/-- Records of an array stored on `[P, M, q, B]`. -/
def recOf' (xs : List ℤ) (q M B p y m : ℕ) : List ℤ :=
  (List.range B).map fun b => xs.getD (((p * M + m) * q + y) * B + b) 0

/-- Exchange of the address digits `1` and `2`. -/
def gS (a : List ℕ) : List ℕ := [a.getD 0 0, a.getD 2 0, a.getD 1 0, a.getD 3 0]

/-! ## Generic list lemmas -/

/-! ## The two orders of the records -/

end NLogN.KIL

end

section
-- Source module: Solutions.KILMove
/-!
# Paper Lemma 3 (field move), forward direction: `[P, q, M, B] → [P, M, q, B]`

`fwdP` reads the array `inp r P q M` (records `r p y m` of length `B`, order `(p, y, m)`) from
stack `0` and leaves `outF r P q M` (order `(p, m, y)`), with `P, q, M, B` on stacks `1 … 4`.
Per block `p` it zips the rows `y = 0, …, q-1` one at a time into an accumulator of `M`
groups (stacks `10`, `11`), so the cost is `O(q · V)` item moves.
-/

set_option linter.unusedSimpArgs false

namespace NLogN

namespace KIL

open IntMul RS

/-! ## The program -/

def fBody : RS := .seq (NLogN.mvK 10 11 33 (.reg 6)) (NLogN.mvK 0 11 33 (.reg 4))

def fStep : RS :=
  .seq (.rep 32 (.reg 3) fBody) (.seq (.asg 6 (.add (.reg 6) (.reg 4)))
    (.seq (.asg 7 (.add (.reg 7) (.reg 5))) (NLogN.mvK 11 10 33 (.reg 7))))

def fBlock : RS :=
  .seq (.rep 31 (.reg 2) fStep) (.seq (NLogN.mvK 10 12 33 (.reg 7))
    (.seq (.asg 6 (.lit 0)) (.asg 7 (.lit 0))))

def fwdP : RS :=
  .seq (.push 5 (.mul (.reg 3) (.reg 4))) (.seq (.push 8 (.mul (.reg 1) (.mul (.reg 2) (.reg 5))))
    (.seq (.push 6 (.lit 0)) (.seq (.push 7 (.lit 0)) (.seq (.rep 30 (.reg 1) fBlock)
      (.seq (NLogN.mvK 12 0 33 (.reg 8)) (.seq (.pop 7) (.seq (.pop 6) (.seq (.pop 8) (.pop 5)))))))))

/-! ## Lists describing the intermediate states -/

section Lists
variable (r : ℕ → ℕ → ℕ → List ℤ)

/-- Rest of row `k` of block `p`, from column `m0`. -/
def rowR (p k m0 M : ℕ) : List ℤ := (List.range' m0 (M - m0)).flatMap (r p k)
/-- Group `m` of the accumulator after `k` rows. -/
def grp (p k m : ℕ) : List ℤ := (List.range k).flatMap fun y => r p y m
/-- Accumulator groups from `m0` on. -/
def Tp (p k m0 M : ℕ) : List ℤ := (List.range' m0 (M - m0)).flatMap (grp r p k)
/-- Reversed new accumulator groups `< m0`. -/
def Cp (p k m0 : ℕ) : List ℤ := ((List.range m0).flatMap (grp r p (k + 1))).reverse
/-- Rows `k, …, q-1` of block `p`. -/
def rowsR (p k q M : ℕ) : List ℤ := (List.range' k (q - k)).flatMap fun y =>
  (List.range M).flatMap (r p y)
/-- Blocks `p, …, P-1`. -/
def blksR (p P q M : ℕ) : List ℤ := (List.range' p (P - p)).flatMap (blkIn r q M)
/-- Reversed output of blocks `< p`. -/
def Ob (q M p : ℕ) : List ℤ := ((List.range p).flatMap (blkOut r q M)).reverse

end Lists

/-! ## States -/

/-- The managed state of `fwdP` inside the main loop. -/
def St (P q M B : ℕ) (x a c o : List ℤ) (u z : ℤ) (k0 k1 k2 : List ℤ) : KS :=
  { x := x, P := [(P : ℤ)], q := [(q : ℤ)], M := [(M : ℤ)], B := [(B : ℤ)],
    MB := [(M : ℤ) * B], V := [(P : ℤ) * ((q : ℤ) * ((M : ℤ) * B))], u := [u], z := [z],
    a := a, c := c, o := o, k0 := k0, k1 := k1, k2 := k2 }

section Proof

variable (s : IState HS) (r : ℕ → ℕ → ℕ → List ℤ) (P q M B Bb : ℕ)

/-- Inner loop state: column `m` of step `k` of block `p`. -/
def innSt (p k m : ℕ) (k2 : List ℤ) : KS :=
  St P q M B (rowR r p k m M ++ (rowsR r p (k + 1) q M ++ blksR r (p + 1) P q M))
    (Tp r p k m M) (Cp r p k m) (Ob r q M p) ((k * B : ℕ) : ℤ) ((k * (M * B) : ℕ) : ℤ)
    [(P : ℤ) - p] [(q : ℤ) - k] k2

/-- State before step `k` of block `p`. -/
def stepSt (p k : ℕ) (k1 : List ℤ) : KS :=
  St P q M B (rowsR r p k q M ++ blksR r (p + 1) P q M) (Tp r p k 0 M) [] (Ob r q M p)
    ((k * B : ℕ) : ℤ) ((k * (M * B) : ℕ) : ℤ) [(P : ℤ) - p] k1 []

/-- Cost bound of one step. -/
def cStep (q M B Bb : ℕ) : ℕ :=
  2 * 2000 * (Bb + 2) + M * (2 * (2 * 2000 * (Bb + 2) + 32) + (q * B) * (14000 * (Bb + 2)) +
    8000 * (Bb + 2) + 30) + 32 + 2 * (4 * 2000 * (Bb + 2)) +
    (2 * 2000 * (Bb + 2) + q * (M * B) * (14000 * (Bb + 2)) + 32)

/-- Cost bound of one block. -/
def cBlk (q M B Bb : ℕ) : ℕ :=
  (1 + 1) * 2000 * (Bb + 2) + q * (NLogN.KIL.cStep q M B Bb + 8000 * (Bb + 2) + 30) + 32 +
    ((1 + 1) * 2000 * (Bb + 2) + q * (M * B) * (14000 * (Bb + 2)) + 32 +
      ((1 + 1) * 2000 * (Bb + 2) + (1 + 1) * 2000 * (Bb + 2)))

/-- State at the start of block `p`. -/
def blkSt (p : ℕ) (k0 : List ℤ) : KS :=
  St P q M B (blksR r p P q M) [] [] (Ob r q M p) 0 0 k0 [] []

/-- Cost bound of the whole forward move. -/
def cFwd (P q M B Bb : ℕ) : ℕ :=
  (3 + 1) * 10 ^ 7 * (Bb + 2) ^ 2 + (5 + 1) * 10 ^ 7 * (Bb + 2) ^ 2 + (1 + 1) * 2000 * (Bb + 2) +
    (1 + 1) * 2000 * (Bb + 2) +
    ((1 + 1) * 2000 * (Bb + 2) + P * (NLogN.KIL.cBlk q M B Bb + 8000 * (Bb + 2) + 30) + 32) +
    ((1 + 1) * 2000 * (Bb + 2) + P * (q * (M * B)) * (14000 * (Bb + 2)) + 32) +
    4 * (2 * Bb + 2)

/-- The input record of `fwdP`: the array on `0`, `P, q, M, B` on `1 … 4`. -/
def fIn (x : List ℤ) (P q M B : ℕ) : KS :=
  { x := x, P := [(P : ℤ)], q := [(q : ℤ)], M := [(M : ℤ)], B := [(B : ℤ)] }

/-- Prologue states of `fwdP`. -/
def pre (x : List ℤ) (P q M B : ℕ) (V u z : List ℤ) : KS :=
  { x := x, P := [(P : ℤ)], q := [(q : ℤ)], M := [(M : ℤ)], B := [(B : ℤ)], MB := [(M : ℤ) * B],
    V := V, u := u, z := z }

end Proof

end KIL

end NLogN

namespace NLogN

namespace KIL

end KIL

end NLogN

end

section
-- Source module: Solutions.KILMoveB
/-!
# Paper Lemma 3 (field move), backward direction: `[P, M, q, B] → [P, q, M, B]`

`bwdP` reads `outF r P q M` (order `(p, m, y)`) from stack `0` and leaves `inp r P q M` (order
`(p, y, m)`), with `P, q, M, B` on stacks `1 … 4`.  Per block it peels, `q` times, the first
record of every group onto the output, so the cost is `O(q · V)` item moves.
-/

set_option linter.unusedSimpArgs false

namespace NLogN

namespace KIL

open IntMul RS

/-! ## The program -/

def bBody : RS := .seq (NLogN.mvK 10 12 33 (.reg 4)) (NLogN.mvK 10 11 33 (.add (.reg 6) (.neg (.reg 4))))

def bStep : RS :=
  .seq (.rep 32 (.reg 3) bBody) (.seq (.asg 6 (.add (.reg 6) (.neg (.reg 4))))
    (.seq (.asg 7 (.add (.reg 7) (.neg (.reg 5)))) (NLogN.mvK 11 10 33 (.reg 7))))

def bBlock : RS :=
  .seq (NLogN.mvK 0 11 33 (.reg 13)) (.seq (NLogN.mvK 11 10 33 (.reg 13)) (.seq (.asg 6 (.reg 9))
    (.seq (.asg 7 (.reg 13)) (.rep 31 (.reg 2) bStep))))

def bwdP : RS :=
  .seq (.push 5 (.mul (.reg 3) (.reg 4))) (.seq (.push 9 (.mul (.reg 2) (.reg 4)))
    (.seq (.push 13 (.mul (.reg 2) (.reg 5))) (.seq (.push 8 (.mul (.reg 1) (.reg 13)))
    (.seq (.push 6 (.lit 0)) (.seq (.push 7 (.lit 0)) (.seq (.rep 30 (.reg 1) bBlock)
      (.seq (NLogN.mvK 12 0 33 (.reg 8)) (.seq (.pop 7) (.seq (.pop 6) (.seq (.pop 8)
        (.seq (.pop 13) (.seq (.pop 9) (.pop 5)))))))))))))

/-! ## Lists -/

section Lists
variable (r : ℕ → ℕ → ℕ → List ℤ)

/-- Remaining records `y ≥ k` of group `m` of block `p`. -/
def grpR (q p k m : ℕ) : List ℤ := (List.range' k (q - k)).flatMap fun y => r p y m
/-- Groups `m ≥ m0` of the remaining records. -/
def Ap (q p k m0 M : ℕ) : List ℤ := (List.range' m0 (M - m0)).flatMap (grpR r q p k)
/-- Reversed groups `< m0` of the next step. -/
def Cq (q p k m0 : ℕ) : List ℤ := ((List.range m0).flatMap (grpR r q p (k + 1))).reverse
/-- Blocks `p, …, P-1` in output order. -/
def blksO (p P q M : ℕ) : List ℤ := (List.range' p (P - p)).flatMap (blkOut r q M)
/-- The reversed output so far. -/
def Oi (q M p k m : ℕ) : List ℤ :=
  ((List.range p).flatMap (blkIn r q M) ++ (List.range k).flatMap (fun y =>
    (List.range M).flatMap (r p y)) ++ (List.range m).flatMap (r p k)).reverse

end Lists

/-- The managed state of `bwdP` inside the main loop. -/
def St2 (P q M B : ℕ) (x a c o : List ℤ) (u z : ℤ) (k0 k1 k2 : List ℤ) : KS :=
  { x := x, P := [(P : ℤ)], q := [(q : ℤ)], M := [(M : ℤ)], B := [(B : ℤ)],
    MB := [(M : ℤ) * B], qB := [(q : ℤ) * B], qMB := [(q : ℤ) * ((M : ℤ) * B)],
    V := [(P : ℤ) * ((q : ℤ) * ((M : ℤ) * B))], u := [u], z := [z],
    a := a, c := c, o := o, k0 := k0, k1 := k1, k2 := k2 }

section Proof

variable (s : IState HS) (r : ℕ → ℕ → ℕ → List ℤ) (P q M B Bb : ℕ)

def bInn (p k m : ℕ) (k2 : List ℤ) : KS :=
  St2 P q M B (blksO r (p + 1) P q M) (Ap r q p k m M) (Cq r q p k m) (Oi r q M p k m)
    (((q - k) * B : ℕ) : ℤ) (((q - k) * (M * B) : ℕ) : ℤ) [(P : ℤ) - p] [(q : ℤ) - k] k2

def bStepSt (p k : ℕ) (k1 : List ℤ) : KS :=
  St2 P q M B (blksO r (p + 1) P q M) (Ap r q p k 0 M) [] (Oi r q M p k 0)
    (((q - k) * B : ℕ) : ℤ) (((q - k) * (M * B) : ℕ) : ℤ) [(P : ℤ) - p] k1 []

/-- Cost bound of one backward step. -/
def bcStep (q M B Bb : ℕ) : ℕ :=
  (1 + 1) * 2000 * (Bb + 2) + M * ((1 + 1) * 2000 * (Bb + 2) + B * (14000 * (Bb + 2)) + 32 +
    ((4 + 1) * 2000 * (Bb + 2) + (q * B) * (14000 * (Bb + 2)) + 32) + 8000 * (Bb + 2) + 30) + 32 +
    ((4 + 1) * 2000 * (Bb + 2) + (4 + 1) * 2000 * (Bb + 2)) +
    ((1 + 1) * 2000 * (Bb + 2) + q * (M * B) * (14000 * (Bb + 2)) + 32)

/-- Cost bound of one backward block. -/
def bcBlk (q M B Bb : ℕ) : ℕ :=
  ((1 + 1) * 2000 * (Bb + 2) + q * (M * B) * (14000 * (Bb + 2)) + 32) +
  ((1 + 1) * 2000 * (Bb + 2) + q * (M * B) * (14000 * (Bb + 2)) + 32) +
  ((1 + 1) * 2000 * (Bb + 2) + (1 + 1) * 2000 * (Bb + 2)) +
  ((1 + 1) * 2000 * (Bb + 2) + q * (bcStep q M B Bb + 8000 * (Bb + 2) + 30) + 32)

def bBlkSt (p : ℕ) (k0 : List ℤ) : KS :=
  St2 P q M B (blksO r p P q M) [] [] (Oi r q M p 0 0) 0 0 k0 [] []

/-- Cost bound of the whole backward move. -/
def cBwd (P q M B Bb : ℕ) : ℕ :=
  4 * ((3 + 1) * 10 ^ 7 * (Bb + 2) ^ 2) + (1 + 1) * 2000 * (Bb + 2) + (1 + 1) * 2000 * (Bb + 2) +
    ((1 + 1) * 2000 * (Bb + 2) + P * (bcBlk q M B Bb + 8000 * (Bb + 2) + 30) + 32) +
    ((1 + 1) * 2000 * (Bb + 2) + P * (q * (M * B)) * (14000 * (Bb + 2)) + 32) +
    6 * (2 * Bb + 2)

/-- Prologue states of `bwdP`. -/
def pre2 (x : List ℤ) (P q M B : ℕ) (qB qMB V u z : List ℤ) : KS :=
  { x := x, P := [(P : ℤ)], q := [(q : ℤ)], M := [(M : ℤ)], B := [(B : ℤ)], MB := [(M : ℤ) * B],
    qB := qB, qMB := qMB, V := V, u := u, z := z }

end Proof

end KIL

end NLogN

namespace NLogN

namespace KIL

end KIL

end NLogN

end

section
-- Source module: Solutions.KILShape
/-!
# Collapsing shapes: a swap of two adjacent field blocks is a 4-field transpose

`pull_swapMid`: on `A ++ F ++ G ++ C`, exchanging the field blocks `F` and `G` is the `pull`
by `gS` on the collapsed shapes `[vol A, vol F, vol G, vol C] → [vol A, vol G, vol F, vol C]`.
With `F = [q]` this is "move the field `q` behind the fields `G`" (Lemma 3, item 2); with
`G = [q]` it is the move of `q` in front of `F`.
-/

namespace NLogN.KIL

open AModel AModelL12

/-- Output address `(xA, xG, xF, xC)` ↦ input address `(xA, xF, xG, xC)`. -/
def swapMid (nA nF nG : ℕ) (a : List ℕ) : List ℕ :=
  a.take nA ++ (a.drop (nA + nG)).take nF ++ (a.drop nA).take nG ++ a.drop (nA + nG + nF)

end NLogN.KIL

end

section
-- Source module: Solutions.KILLemma3
/-!
# Paper Lemma 3 (single-field move) in the framework's stack-disciplined form

* `fwdP_spec`: for a state `t` with the array `xs` (shape `[P, q, M, B]`, row-major) on stack `0`
  and `P, q, M, B` on stacks `1 … 4`, `fwdP` is executable, ends in
  `t.upd 0 (ints (pull [P, q, M, B] [P, M, q, B] gS 0 xs) ++ R)` (all other stacks restored), and
  costs at most `2·10^5 · q · (V + 1) · (Bb + 2) + 10^8 · (Bb + 2)^2` (`V = P q M B`, `Bb` bounds
  the bit size of every entry and of `V`).
* `bwdP_spec`: the inverse move `[P, M, q, B] → [P, q, M, B]`, same cost shape.
* `fwdP_iruns`, `bwdP_iruns`: the same as `IRuns` of the compiled item-language code.
-/

set_option linter.unusedSimpArgs false

namespace NLogN

namespace KIL

open IntMul RS AModel

/-- Base state below the arguments. -/
def baseOf (t : IState HS) (R L1 L2 L3 L4 : List Item) : IState HS :=
  ((((t.upd 0 R).upd 1 L1).upd 2 L2).upd 3 L3).upd 4 L4

/-- Cost bound of the move programs. -/
def cMove (P q M B Bb : ℕ) : ℕ :=
  300000 * q * (P * (q * (M * B)) + 1) * (Bb + 2) + 2 * 10 ^ 8 * (Bb + 2) ^ 2

/-! ## General shapes: exchanging two adjacent field blocks

On the shape `A ++ F ++ G ++ C`, exchanging the blocks `F` and `G` (in particular moving one
field `F = [q]` behind the fields `G`, or `G = [q]` in front of `F`) is computed by `fwdP` with
cost `O(vol F · V)` and by `bwdP` with cost `O(vol G · V)`. -/

end KIL

end NLogN

end

section
-- Source module: Solutions.KILRot
/-!
# Primitive (d): controlled shift as a rotation of blocks

`rotP` reads an array on `[P, Q, B]` (stack `0`) and a list of offsets `a_p < Q` (stack `14`,
consumed) and rotates the `Q` records of block `p` cyclically by `a_p`:
record `y` of block `p` moves to `(y + a_p) mod Q`.  Cost `O(V + P · Q)` item operations.

* `Sp.repMv`: a counted loop moving one record of `B` integers per iteration.
* `rotP_sp`, `rotP_spec`: the program.
* `rotOut_eq_pull`: the output is `pull [P, Q, B] [P, Q, B] (rotAddr Q a) 0 xs`.
-/

set_option linter.unusedSimpArgs false

namespace NLogN.KIL

open IntMul RS AModel

/-! ## Moving records one at a time -/

end NLogN.KIL

namespace NLogN.KIL

open IntMul RS AModel

/-! ## The rotation program -/

/-- `Q - a` for `Q` on stack `2` and the offset `a` on top of stack `14`. -/
def qma : Ex HS := .add (.reg 2) (.neg (.reg 14))

def rBlk : RS :=
  .seq (.rep 31 qma (NLogN.mvK 0 11 33 (.reg 4))) (.seq (.rep 31 (.reg 14) (NLogN.mvK 0 12 33 (.reg 4)))
    (.seq (.rep 31 qma (NLogN.mvK 11 10 33 (.reg 4))) (.seq (.rep 31 qma (NLogN.mvK 10 12 33 (.reg 4)))
      (.pop 14))))

def rotP : RS :=
  .seq (.push 8 (.mul (.reg 1) (.mul (.reg 2) (.reg 4)))) (.seq (.rep 30 (.reg 1) rBlk)
    (.seq (NLogN.mvK 12 0 33 (.reg 8)) (.pop 8)))

section Lists
variable (r : ℕ → ℕ → List ℤ) (Q : ℕ) (as : ℕ → ℕ)

/-- Records `0, …, Q - a_p - 1` of block `p`. -/
def firstR (p : ℕ) : List (List ℤ) := (List.range (Q - as p)).map (r p)
/-- Records `Q - a_p, …, Q - 1` of block `p`. -/
def secondR (p : ℕ) : List (List ℤ) := (List.range' (Q - as p) (as p)).map (r p)
/-- Block `p` in input order. -/
def blk3 (p : ℕ) : List ℤ := (List.range Q).flatMap (r p)
/-- Block `p` rotated by `a_p`. -/
def rotBlk (p : ℕ) : List ℤ := (secondR r Q as p).flatten ++ (firstR r Q as p).flatten
def inR (P : ℕ) : List ℤ := (List.range P).flatMap (blk3 r Q)
def outR (P : ℕ) : List ℤ := (List.range P).flatMap (rotBlk r Q as)

end Lists

/-- State at the start of block `p`. -/
def rSt (r : ℕ → ℕ → List ℤ) (P Q B : ℕ) (as : ℕ → ℕ) (p : ℕ) (k0 : List ℤ) : KS :=
  { x := (List.range' p (P - p)).flatMap (blk3 r Q), P := [(P : ℤ)], q := [(Q : ℤ)],
    B := [(B : ℤ)], V := [(P : ℤ) * ((Q : ℤ) * B)],
    w := (List.range' p (P - p)).map (fun p => (as p : ℤ)),
    o := ((List.range p).flatMap (rotBlk r Q as)).reverse, k0 := k0 }

/-- Cost bound of one rotated block. -/
def cRBlk (Q B Bb : ℕ) : ℕ :=
  4 * ((4 + 1) * 2000 * (Bb + 2) + Q * ((1 + 1) * 2000 * (Bb + 2) + B * (14000 * (Bb + 2)) + 32 +
    8000 * (Bb + 2) + 30) + 32) + (2 * Bb + 2)

section Proof

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

/-- Managed input/output states of `rotP`. -/
def rIO (x w : List ℤ) (P Q B : ℕ) : KS :=
  { x := x, P := [(P : ℤ)], q := [(Q : ℤ)], B := [(B : ℤ)], w := w }

/-- Cost bound of `rotP`. -/
def cRot (P Q B Bb : ℕ) : ℕ :=
  (5 + 1) * 10 ^ 7 * (Bb + 2) ^ 2 + ((1 + 1) * 2000 * (Bb + 2) + P * (cRBlk Q B Bb + 8000 * (Bb + 2) + 30) + 32) +
    ((1 + 1) * 2000 * (Bb + 2) + P * (Q * B) * (14000 * (Bb + 2)) + 32) + (2 * Bb + 2)

end Proof

end NLogN.KIL

namespace NLogN.KIL

open AModel

/-! ## Denotation of the rotation -/

/-- Records of an array stored on `[P, Q, B]`. -/
def recR (xs : List ℤ) (Q B p y : ℕ) : List ℤ :=
  (List.range B).map fun b => xs.getD ((p * Q + y) * B + b) 0

/-- Output address `(p, y, b)` ↦ input address `(p, (y - a_p) mod Q, b)`. -/
def rotAddr (Q : ℕ) (as : ℕ → ℕ) (a : List ℕ) : List ℕ :=
  [a.getD 0 0, (a.getD 1 0 + Q - as (a.getD 0 0)) % Q, a.getD 2 0]

end NLogN.KIL

namespace NLogN.KIL

open IntMul RS AModel

/-- Base state below the arguments of `rotP`. -/
def rBase (t : IState HS) (R L1 L2 L4 L14 : List Item) : IState HS :=
  ((((t.upd 0 R).upd 1 L1).upd 2 L2).upd 4 L4).upd 14 L14

end NLogN.KIL

end

section
-- Source module: Solutions.KILShift
/-!
# The controlled shift of Lemma 10 as a block rotation

For a layout `Λ` on `sh = A ++ Q :: C` with the target field `i` at position `|A|`, unit
coefficient `u = 1` and controls only before the target (`c j = 0` for `j ≥ i`), the controlled
shift `shiftArr Λ i 1 c` is the block rotation `pull [vol A, Q, vol C] [vol A, Q, vol C]
(rotAddr Q (offs Λ c A))` where block `p` (the prefix digits) is rotated by
`offs Λ c A p = (Σ_j c_j · digit_j(decode A p)).val`.

* `shiftArr_eq_rot`: the list identity.
* `rotP_shift_spec`: `rotP` computes `shiftArr Λ i 1 c` given the offset table on stack `14`.
-/

namespace NLogN.KIL

open IntMul RS AModel AModelL12

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

/-- Rotation offset of block `p` (prefix digits `decode A p`). -/
def offs (c : Fin n → ZMod Q) (A : Shape) (p : ℕ) : ℕ :=
  (∑ j, c j * (((decode A p).getD (Λ.pos j) 0 : ℕ) : ZMod Q)).val

end NLogN.KIL

end

section
-- Source module: Solutions.KIL2Table
/-!
# KIL2: generating the offset table of a controlled shift

`genP Q ac` pushes onto stack `14` the table `tbl Q ac` of the offsets
`(Σ_k c_k · d_k) mod Q`, `p = (d_0, …, d_{r-1})` ranging over the mixed-radix shape
`ac.map Prod.fst` in increasing order (first field most significant), where `ac = [(a_k, c_k)]`.
It is built field by field from the last one: `T ← concat_{d < a} (T + c·d mod Q)`.
Everything (`a_k`, `c_k`, the suffix volumes) is a literal of the code.

* `genP_spec`: the generator, linear cost `O(|A| · (vol A + 1) · (Bb + 2))`.
* `tbl_eq_offs`: `tbl` is the table `(range (vol A)).map (offs Λ c A)` of `KILShift`.
-/

set_option linter.unusedSimpArgs false

namespace NLogN.KIL

open IntMul RS AModel

/-! ## Expressions -/

/-- `(x + y) mod Q` for `0 ≤ x, y < Q`, by one conditional subtraction. -/
def mAdd (Q : ℕ) (x y : Ex HS) : Ex HS :=
  .sel (.add (.add x y) (.lit (-(Q : ℤ)))) (.add x y) (.add (.add x y) (.lit (-(Q : ℤ))))

/-! ## The program -/

/-- Move one integer from `14` to `10`, pushing its shifted copy (by `s` on `12`) on `11`. -/
def gInner (Q : ℕ) : RS :=
  .seq (.push 11 (mAdd Q (.reg 14) (.reg 12))) (.seq (.push 10 (.reg 14)) (.pop 14))

/-- One copy `T + s` of the table, restoring `T`, then `s ← s + c mod Q`. -/
def gBody (Q c M : ℕ) : RS :=
  .seq (.rep 31 (.lit M) (gInner Q))
    (.seq (NLogN.mvK 10 14 33 (.lit M)) (.asg 12 (mAdd Q (.reg 12) (.lit c))))

/-- One field `(a, c)` of the table: `T ← concat_{d < a} (T + c·d mod Q)`, `|T| = M`. -/
def gStage (Q a c M : ℕ) : RS :=
  .seq (.push 12 (.lit 0)) (.seq (.rep 30 (.lit a) (gBody Q c M)) (.seq (.pop 12)
    (.seq (.rep 31 (.lit M) (.pop 14)) (NLogN.mvK 11 14 33 (.lit (a * M))))))

/-- The generator for the fields `ac = [(a_k, c_k)]`. -/
def genP (Q : ℕ) : List (ℕ × ℕ) → RS
  | [] => .push 14 (.lit 0)
  | (a, c) :: ac => .seq (genP Q ac) (gStage Q a c (AModel.vol (ac.map Prod.fst)))

/-! ## Denotation -/

/-- `T + s mod Q`. -/
def shT (Q : ℕ) (s : ℤ) (T : List ℤ) : List ℤ := T.map fun v => (v + s) % Q

/-- One stage. -/
def stg (Q a c : ℕ) (T : List ℤ) : List ℤ :=
  (List.range a).flatMap fun d : ℕ => shT Q (((c : ℤ) * (d : ℤ)) % Q) T

/-- The table. -/
def tbl (Q : ℕ) : List (ℕ × ℕ) → List ℤ
  | [] => [0]
  | (a, c) :: ac => stg Q a c (tbl Q ac)

/-! ## Proofs -/

/-- Values in `[0, Q)`. -/
def InR (Q : ℕ) (T : List ℤ) : Prop := ∀ v ∈ T, 0 ≤ v ∧ v < Q

def cInner (Bb : ℕ) : ℕ := (14 + 1) * 2000 * (Bb + 2) + ((1 + 1) * 2000 * (Bb + 2) + (2 * Bb + 2))

section Proofs

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

def cInnerL (M Bb : ℕ) : ℕ :=
  (1 + 1) * 2000 * (Bb + 2) + M * (cInner Bb + 8000 * (Bb + 2) + 30) + 32

def cBody (M Bb : ℕ) : ℕ :=
  cInnerL M Bb + ((1 + 1) * 2000 * (Bb + 2) + M * (14000 * (Bb + 2)) + 32) +
    (14 + 1) * 2000 * (Bb + 2)

/-- The first `d` copies of a stage. -/
def stgP (Q c : ℕ) (T : List ℤ) (d : ℕ) : List ℤ :=
  (List.range d).flatMap fun d' : ℕ => shT Q (((c : ℤ) * (d' : ℤ)) % Q) T

def cOuter (a M Bb : ℕ) : ℕ :=
  (1 + 1) * 2000 * (Bb + 2) + a * (cBody M Bb + 8000 * (Bb + 2) + 30) + 32

def cStage (a M Bb : ℕ) : ℕ :=
  (1 + 1) * 2000 * (Bb + 2) + cOuter a M Bb + (2 * Bb + 2) +
    ((1 + 1) * 2000 * (Bb + 2) + M * ((2 * Bb + 2) + 8000 * (Bb + 2) + 30) + 32) +
    ((1 + 1) * 2000 * (Bb + 2) + a * M * (14000 * (Bb + 2)) + 32)

/-- Cost of the generator. -/
def genCost (Bb : ℕ) : List (ℕ × ℕ) → ℕ
  | [] => (1 + 1) * 2000 * (Bb + 2)
  | (a, _) :: ac => genCost Bb ac + cStage a (AModel.vol (ac.map Prod.fst)) Bb

end Proofs

/-! ## Stack form -/

/-! ## Denotation of the table -/

/-- `Σ_k c_k d_k`. -/
def lf : List ℕ → List ℕ → ℕ
  | c :: cs, d :: D => c * d + lf cs D
  | _, _ => 0

section Offs

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

/-- Coefficient of position `k`. -/
def cfOf (c : Fin n → ZMod Q) (k : ℕ) : ZMod Q := ∑ j, if (Λ.pos j : ℕ) = k then c j else 0

/-- The per-position coefficients of the prefix `A`. -/
def csOf (c : Fin n → ZMod Q) (A : Shape) : List ℕ :=
  (List.range A.length).map fun k => (cfOf Λ c k).val

end Offs

end NLogN.KIL

end

section
-- Source module: Solutions.KIL2Shift
/-!
# KIL2: the Lemma 10 controlled shift, self-contained

`shP` = offset-table generator `genP` (stack `14`), then the size parameters
`vol A, Q, vol C` pushed on stacks `1, 2, 4`, then `rotP`, then the parameters popped.

* `Sp.pushLit`, `Sp.popI`: pushes of literals and pops on arbitrary states.
* `shiftArr_il_spec`: `shP` computes `shiftArr Λ i 1 c 0 xs` (Lemma 10 case: unit coefficient
  `1`, controls before the target), stack disciplined, cost
  `O((|A| + 1) · (V + 1) · (Bb + 2) + (Bb + 2)^2)`.
-/

set_option linter.unusedSimpArgs false

namespace NLogN.KIL

open IntMul RS AModel

/-- The self-contained Lemma 10 shift on `[P, Q, B]` with table fields `ac`. -/
def shP (Q P B : ℕ) (ac : List (ℕ × ℕ)) : RS :=
  .seq (genP Q ac) (.seq (.push 1 (.lit P)) (.seq (.push 2 (.lit Q)) (.seq (.push 4 (.lit B))
    (.seq rotP (.seq (.pop 1) (.seq (.pop 2) (.pop 4)))))))

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

section

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

end

end NLogN.KIL

end

section
-- Source module: Solutions.KIL2Neg
/-!
# KIL2: controlled shift with coefficient `-1` (block reflection)

For coefficient `u = -1` the controlled shift is `y ↦ a_p − y (mod Q)` on the target field of
block `p`: the inverse address map reads record `(a_p − y) mod Q`, so each block is cut after
record `a_p` and both pieces are output in reversed record order (records themselves intact).

`negP` does this with the same skeleton as `rotP` (offsets `a_p < Q` on stack `14`):
a record-reversing double move `0 → 10 → 11` (`dmv`), then a plain move `11 → 12`.

* `Sp.repMv2`: the counted double-move loop (reverses the record order, keeps records).
* `negP_spec`: `negP` computes `pull [P, Q, B] [P, Q, B] (refAddr Q as) 0 xs`, linear cost.
* `shiftInv_split_u`: the controlled-shift inverse on `A ++ Q :: C` for any coefficient.
* `shiftArr_neg_eq_ref`, `shN`, `shiftArrNeg_il_spec`: the self-contained `u = -1` shift
  (offset table generated by `genP`).
-/

set_option linter.unusedSimpArgs false

namespace NLogN.KIL

open IntMul RS AModel

/-! ## The double-move loop -/

/-- Move one record `0 → 10 → 11` (record kept in order, pushed on `11`). -/
def dmv : RS := .seq (NLogN.mvK 0 10 33 (.reg 4)) (NLogN.mvK 10 11 33 (.reg 4))

def cMv (B Bb : ℕ) : ℕ := (1 + 1) * 2000 * (Bb + 2) + B * (14000 * (Bb + 2)) + 32

/-! ## The reflection program -/

/-- `a + 1` for the offset `a` on stack `14`. -/
def a1E : Ex HS := .add (.reg 14) (.lit 1)
/-- `Q − a − 1`. -/
def qa1E : Ex HS := .add (.add (.reg 2) (.neg (.reg 14))) (.lit (-1))

def nBlk : RS :=
  .seq (.rep 31 a1E dmv) (.seq (.rep 31 a1E (NLogN.mvK 11 12 33 (.reg 4)))
    (.seq (.rep 31 qa1E dmv) (.seq (.rep 31 qa1E (NLogN.mvK 11 12 33 (.reg 4))) (.pop 14))))

def negP : RS :=
  .seq (.push 8 (.mul (.reg 1) (.mul (.reg 2) (.reg 4)))) (.seq (.rep 30 (.reg 1) nBlk)
    (.seq (NLogN.mvK 12 0 33 (.reg 8)) (.pop 8)))

section Lists
variable (r : ℕ → ℕ → List ℤ) (Q : ℕ) (as : ℕ → ℕ)

/-- Records `0, …, a_p` of block `p`. -/
def FR (p : ℕ) : List (List ℤ) := (List.range (as p + 1)).map (r p)
/-- Records `a_p + 1, …, Q − 1` of block `p`. -/
def SR (p : ℕ) : List (List ℤ) := (List.range' (as p + 1) (Q - (as p + 1))).map (r p)
/-- Block `p` reflected: both pieces in reversed record order. -/
def refBlk (p : ℕ) : List ℤ := (NLogN.KIL.FR r as p).reverse.flatten ++ (SR r Q as p).reverse.flatten
def outN (P : ℕ) : List ℤ := (List.range P).flatMap (refBlk r Q as)

end Lists

/-- State at the start of block `p`. -/
def nSt (r : ℕ → ℕ → List ℤ) (P Q B : ℕ) (as : ℕ → ℕ) (p : ℕ) (k0 : List ℤ) : KS :=
  { x := (List.range' p (P - p)).flatMap (blk3 r Q), P := [(P : ℤ)], q := [(Q : ℤ)],
    B := [(B : ℤ)], V := [(P : ℤ) * ((Q : ℤ) * B)],
    w := (List.range' p (P - p)).map (fun p => (as p : ℤ)),
    o := ((List.range p).flatMap (refBlk r Q as)).reverse, k0 := k0 }

def cNBlk (Q B Bb : ℕ) : ℕ :=
  2 * ((6 + 1) * 2000 * (Bb + 2) + Q * (NLogN.KIL.cMv B Bb + NLogN.KIL.cMv B Bb + 8000 * (Bb + 2) + 30) + 32) +
  2 * ((6 + 1) * 2000 * (Bb + 2) + Q * ((1 + 1) * 2000 * (Bb + 2) + B * (14000 * (Bb + 2)) + 32 +
    8000 * (Bb + 2) + 30) + 32) + (2 * Bb + 2)

section Proof

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

end Proof

/-- Cost bound of `negP`. -/
def cNeg (P Q B Bb : ℕ) : ℕ :=
  (5 + 1) * 10 ^ 7 * (Bb + 2) ^ 2 + ((1 + 1) * 2000 * (Bb + 2) +
    P * (cNBlk Q B Bb + 8000 * (Bb + 2) + 30) + 32) +
    ((1 + 1) * 2000 * (Bb + 2) + P * (Q * B) * (14000 * (Bb + 2)) + 32) + (2 * Bb + 2)

/-! ## Denotation of the reflection -/

/-- Output address `(p, y, b)` ↦ input address `(p, (a_p − y) mod Q, b)`. -/
def refAddr (Q : ℕ) (as : ℕ → ℕ) (a : List ℕ) : List ℕ :=
  [a.getD 0 0, (as (a.getD 0 0) + Q - a.getD 1 0) % Q, a.getD 2 0]

/-! ## The controlled shift with coefficient `-1` -/

section ShiftNeg

open AModelL12

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

/-- The self-contained reflection shift: table, parameters, `negP`, parameters popped. -/
def shN (Q P B : ℕ) (ac : List (ℕ × ℕ)) : RS :=
  .seq (genP Q ac) (.seq (.push 1 (.lit P)) (.seq (.push 2 (.lit Q)) (.seq (.push 4 (.lit B))
    (.seq NLogN.KIL.negP (.seq (.pop 1) (.seq (.pop 2) (.pop 4)))))))

end ShiftNeg

end NLogN.KIL

end

section
-- Source module: Solutions.KIL2Mul
/-!
# KIL2: multiplying the target field by a fixed unit (Lemma 4, `y ↦ c·y` and `y ↦ c⁻¹·y`)

For a fixed positive integer `c` prime to `Q` (`c ≤ 20`, one stack per stream), the permutation
`y ↦ c·y (mod Q)` of the `Q` records of every block is done as in the paper: the block is split into
the `c` increasing pieces `j = ⌊c·y / Q⌋` (stacks `35 + j`), and the output `z = 0, …, Q − 1` is
read from stream `σ(z mod c)`, the unique `j < c` with `j·Q + z ≡ 0 (mod c)` (a residue counter on
stack `13` and a `c`-way branch).  The inverse `y ↦ c⁻¹·y` swaps the roles: records are
distributed by residue and the streams are concatenated.

Number theory (this part):
* `nxt Q c j z = ⌈(j·Q + z)/c⌉`, `lo Q c j = nxt Q c j 0`; `nxt_hit`, `nxt_miss`.
* `sig Q c v`: the stream of residue `v`; `sig_spec`, `sig_unique`.
-/

set_option linter.unusedSimpArgs false

namespace NLogN.KIL

open IntMul RS AModel

/-! ## Ceiling arithmetic -/

section Arith

variable (Q c : ℕ)

/-- `⌈(j·Q + z) / c⌉`: the first record of stream `j` not yet output at time `z`. -/
def nxt (j z : ℕ) : ℕ := (j * Q + z + c - 1) / c

/-- First record of piece `j`. -/
def lo (j : ℕ) : ℕ := nxt Q c j 0

variable {Q c}

/-! ## The stream of a residue -/

variable (Q c)

/-- The unique `j < c` with `j·Q + v ≡ 0 (mod c)`. -/
def sig (v : ℕ) : ℕ := (-(v : ZMod c) * (Q : ZMod c)⁻¹).val

variable {Q c}

/-- The input record read at output position `z` (multiplication by `c`). -/
def yOf (Q c z : ℕ) : ℕ := (sig Q c (z % c) * Q + z) / c

end Arith

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

/-- Branch `br v` for the value `v ∈ [k, k + m)` on top of stack `r`. -/
def sel (r : Fin HS) (br : ℕ → RS) : ℕ → ℕ → RS
  | _, 0 => .skip
  | k, 1 => br k
  | k, m + 2 => .ite (.add (.reg r) (.lit (-((k : ℤ) + 1)))) (br k) (sel r br (k + 1) (m + 1))

/-- Cost of one level of the branch. -/
def cSel (Bb : ℕ) : ℕ := (3 + 1) * 2000 * (Bb + 2) + (2 * Bb + 2) + 25

/-! ## Stream stacks `35, …, 35 + c − 1` -/

/-- Stack of stream `j`. -/
def sI (j : ℕ) : Fin HS := ⟨(35 + j) % 64, Nat.mod_lt _ (by decide)⟩

/-- The managed stacks `f`, with stream `j < c` holding `g j`. -/
def strF (c : ℕ) (f : Fin HS → List ℤ) (g : ℕ → List ℤ) : Fin HS → List ℤ :=
  fun i => if 35 ≤ i.val ∧ i.val < 35 + c then g (i.val - 35) else f i

/-! ## The merge by residue (multiplication by `c`) -/

section MergeM

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

/-- Remaining records of stream `j` at time `z`. -/
def contM (z j : ℕ) : List ℤ :=
  ((List.range' (nxt Q c j z) (lo Q c (j + 1) - nxt Q c j z)).map rp).flatten

/-- Output of the first `z` positions. -/
def outM (z : ℕ) : List ℤ := (List.range z).flatMap fun z' => rp (NLogN.KIL.yOf Q c z')

/-- Branch `v`: move one record from stream `σ(v)` to the output stack `12`. -/
def brM (v : ℕ) : RS := NLogN.mvK (NLogN.KIL.sI (sig Q c v)) 12 33 (.lit B)

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

def mergeM : RS := .rep 31 (.lit Q) (bodyM Q c B)

/-- State at time `z`. -/
def stM (K : NLogN.KIL.KS) (o0 : List ℤ) (z : ℕ) (kk : List ℤ) : Fin HS → List ℤ :=
  strF c (toF {K with o := (outM rp Q c z).reverse ++ o0, qMB := [((z % c : ℕ) : ℤ)], k1 := kk})
    (contM rp Q c z)

def cBodyM (c B Bb : ℕ) : ℕ :=
  c * cSel Bb + ((1 + 1) * 2000 * (Bb + 2) + B * (14000 * (Bb + 2)) + 32) +
    (14 + 1) * 2000 * (Bb + 2)

variable {rp Q c B}

end MergeM

section MergeLoop

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

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

end MergeLoop

/-! ## Distribution by contiguous pieces -/

section Dist

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

/-- Move pieces `k − 1, …, 0` from stack `10` (block reversed) to their stream stacks. -/
def dSeq (Q c B : ℕ) : ℕ → RS
  | 0 => .skip
  | k + 1 => .seq (NLogN.mvK 10 (NLogN.KIL.sI k) 33 (.lit (((lo Q c (k + 1) - lo Q c k) * B : ℕ) : ℤ)))
      (dSeq Q c B k)

end Dist

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

section Block

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

def blkMc : RS :=
  .seq (NLogN.mvK 0 10 33 (.lit ((Q * B : ℕ) : ℤ))) (.seq (dSeq Q c B c) (.seq (.push 13 (.lit 0))
    (.seq (mergeM Q c B) (.pop 13))))

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

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

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

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

variable {r P Q c B}

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

end Block

end NLogN.KIL

end

section
-- Source module: Solutions.KIL2Run
/-!
# KIL2: running a list of steps, interchanges delegated to child procedures

`runCode scode jOf steps` is item-language code (`IStmt`, so it may `call` procedures):
a shift step runs the code `scode s`, an interchange `swap a b` is the call `.call (jOf a b)`.

* `runCode_spec`: given IRuns specs of the shift codes and of the children (stack `0` holds the
  array, every other stack is restored), the code computes `runC Λ child 0 steps xs`, the
  `AModelProp11` denotation, with cost `Σ_steps (shift cost | child cost + 2) + 1`.
* `runCode_runArr`: with correct children this is `runArr Λ 0 steps xs`.
* `scode1`, `runCode1_spec`: the concrete instance where every shift has unit coefficient `1`
  (Lemma 10 case), implemented by `shP` (`KIL2Shift`); cost `cShift` per shift.
-/

set_option linter.unusedSimpArgs false

namespace NLogN.KIL

open IntMul AModel

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

/-- Code of one step. -/
def stepCode (scode : Step n Q → IStmt HS Pp) (jOf : Fin n → Fin n → Fin Pp) :
    Step n Q → IStmt HS Pp
  | .shift i u c => scode (.shift i u c)
  | .swap a b => .call (jOf a b)

/-- Code of a list of steps. -/
def runCode (scode : Step n Q → IStmt HS Pp) (jOf : Fin n → Fin n → Fin Pp) :
    List (Step n Q) → IStmt HS Pp
  | [] => .skip
  | s :: rest => .seq (stepCode scode jOf s) (runCode scode jOf rest)

/-- Cost of one step: `cS` for a shift, the child cost plus the call overhead for a swap. -/
def stepCost (cS : ℕ) (cC : Fin n → Fin n → ℕ) : Step n Q → ℕ
  | .shift _ _ _ => cS
  | .swap a b => cC a b + 2

/-- Cost of a list of steps. -/
def runCostI (cS : ℕ) (cC : Fin n → Fin n → ℕ) (steps : List (Step n Q)) : ℕ :=
  (steps.map (stepCost cS cC)).sum + 1

/-- Data invariant of the array on stack `0`. -/
def DOK (sh : Shape) (Bb : ℕ) (ys : List ℤ) : Prop := ys.length = AModel.vol sh ∧ Bnd Bb ys

/-! ## The unit-coefficient instance -/

/-- Shift code for unit coefficient `1`: `shP` on `[vol A, Q, vol C]`, `A`/`C` the fields before /
after the target. -/
def scode1 : Step n Q → IStmt HS Pp
  | .shift i _ c => (shP Q (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

/-! ## Coefficients `±1` -/

/-- Shift code for coefficients `1` (`shP`, rotation) and `-1` (`shN`, reflection). -/
def scodePM : Step n Q → IStmt HS Pp
  | .shift i u c =>
    if u = 1 then (shP Q (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 (shN Q (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 NLogN.KIL

end