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

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

import Definitions.Def_IntMul_Kappa_K_Defs25

set_option autoImplicit false
set_option relaxedAutoImplicit false
set_option maxHeartbeats 1000000

section
-- Source module: Solutions.KFY3Core
/-!
# KFY 3: generic pieces of the line passes

* **`coreQ ax`**: nlogn's `GL.axisP ax` without its transposition prelude `preP` (the
  `O(V log V)` radix transpose `transP 2`). The rotation is done before, by `TransNeed`.
  **`coreQ_all`**: on `gD = gints L ++ R` it applies the slice map to the consecutive slices
  of `L` (`coreOut`).
* **`padIn`** / **`padOut`**: a slice program for slices of length `s` turned into one for the
  padded slices of length `t ≥ s`: `padIn` drops the last `t − s` entries of the input slice
  (`f (x.take s)`), `padOut` appends `t − s` zeros to the output slice
  (`f x ++ replicate (t − s) 0`).
* **`pushL`** / **`popL`**: push a list of register values given by expressions, pop them again.
-/

set_option linter.unusedSimpArgs false
set_option linter.unusedVariables false
set_option linter.unusedSectionVars false

namespace NLogN.KFY

open IntMul NLogN IntMul.HvdH.Ref NLogN.GL NLogN.AX NLogN.Rsp

/-! ## Small expression facts -/

/-- `2 t − 2 s` from registers `rT`, `rS`. -/
def cdE (rT rS : Fin HS) : Ex HS := .add (cnt2 rT) (.neg (cnt2 rS))

/-- Size conditions of the moves of a padded axis (`n = n' = t`, `m` slices). -/
structure AxQA (t m Bi Bo Q : ℕ) : Prop where
  hm : 1 ≤ m
  bB : Bi ≤ Q
  bBo : Bo ≤ Q
  bt : zsz ((2 * t : ℕ) : ℤ) ≤ Q
  bm : zsz ((2 * m : ℕ) : ℤ) ≤ Q
  btot : zsz ((2 * (m * t) : ℕ) : ℤ) ≤ Q

/-! ## The axis core -/

section core

variable {M : MulSpec} {Ctx : IState HS → Prop} {B : ℕ}

/-- `GL.axisP` without the transposition prelude. -/
def coreQ (ax : GL.Axis M Ctx B) : RM :=
  .seq (.rep gCC (.reg ax.rm) (chunkBody ax)) (.rs (mvList gZB NLogN.GL.gD gCV (cnt2 ax.rtot)))

/-- Cost bound of `coreQ`. -/
def coreCost (ax : GL.Axis M Ctx B) (Q : ℕ) : ℕ :=
  ax.m * (ax.C + 100000 * (ax.n + ax.n' + 1) * (Q + 2)) + 100000 * (ax.m * ax.n' + 1) * (Q + 2)

/-- The output of the core: the slice map on the consecutive slices. -/
def coreOut (ax : GL.Axis M Ctx B) (L : List (ℤ × ℤ)) : List (ℤ × ℤ) :=
  (chunkL ax.n L).flatMap ax.f

end core

/-! ## Pushing zeros -/

/-- Push `n` zero items onto `X` (counter `C`). -/
def pushZ (X C : Fin HS) (n : Ex HS) : RS := .rep C n (.push X (.lit 0))

/-! ## Padded slice programs -/

section pad

variable {M : MulSpec} {Ctx : IState HS → Prop}

/-- Drop the last `t − s` entries of the slice on `I` (via `Z`, counter `K`), then run `p`. -/
def padInP (I Z K rS rT : Fin HS) (p : RM) : RM :=
  RM.seqL [.rs (mvList I Z gCV (cnt2 rS)), .rs (popK I K (cdE rT rS)), .rs (mvList Z I gCV (cnt2 rS)),
    p]

/-- The side conditions of `padIn`. -/
structure PadInOk (I O Z K rS rT : Fin HS) : Prop where
  hIZ : I ≠ Z
  hIV : I ≠ gCV
  hZV : Z ≠ gCV
  hI55 : I.val < 55
  hZ55 : Z.val < 55
  hIK : I ≠ K
  hK56 : K.val < 56
  hS56 : rS.val < 56
  hT56 : rT.val < 56
  hSI : rS ≠ I
  hSZ : rS ≠ Z
  hTI : rT ≠ I
  hTZ : rT ≠ Z

/-- Push `t − s` zero entries on `O` (counter `K`), then run `p`. -/
def padOutP (O K rS rT : Fin HS) (p : RM) : RM := .seq (.rs (pushZ O K (cdE rT rS))) p

end pad

/-! ## Pushing and popping register lists -/

/-- Push the value of `e` onto `r`, for each `(r, e)` in order. -/
def pushL : List (Fin HS × Ex HS) → RS
  | [] => .skip
  | (r, e) :: l => .seq (.push r e) (pushL l)

/-- Pop every register of the list. -/
def popL : List (Fin HS) → RS
  | [] => .skip
  | r :: l => .seq (.pop r) (popL l)

/-- The well-formedness of a push list at `s`: distinct targets `< 56`, clean expressions
defined at `s` that read no target. -/
structure PushOk (l : List (Fin HS × Ex HS)) (s : IState HS) : Prop where
  nd : (l.map Prod.fst).Nodup
  lt : ∀ p ∈ l, p.1.val < 56
  cl : ∀ p ∈ l, p.2.Clean
  ok : ∀ p ∈ l, p.2.ok s
  rd : ∀ p ∈ l, ∀ q ∈ l, q.1 ∉ p.2.regs

/-- The registers of a pop list hold integers of size `≤ B` on top. -/
def PopOk (l : List (Fin HS)) (s : IState HS) (B : ℕ) : Prop :=
  l.Nodup ∧ ∀ r ∈ l, r.val < 56 ∧ ∃ (v : ℤ) (L : List Item), s.G r = .int v :: L ∧ zsz v ≤ B

end NLogN.KFY

end

section
-- Source module: Solutions.KFY1Pad
/-!
# KFY 1: the padded line passes (math side)

The machine keeps the array in the padded box `∏ t_i` on every axis (so that the axis rotations of
`TransNeed` always see the shapes `shPad h bs`). Axis by axis it applies

* **`A`**: a `t_i`-line `X` (only its first `s_i` entries are read) ↦ `sapL p s_i t_i α X`
  (`t_i` outputs);
* **`B₀`**: a `t_i`-line `X` ↦ `(B̃₀ X)_{0..s_i-1}` followed by `t_i − s_i` zeros (`b0Pad`).

As list pipelines these are `tensorGo (mapsA …)` and `tensorGo (mapsB …)` on the dims
`List.ofFn t`. This file proves that they meet the representation and error clauses of
`KFM.LineNeed`:

* **`passA`**: `RepPad zs u`, `‖u‖ ≤ 1/2` ↦ `RepFull L (wOf L)`, `GB (p+2) L` and
  `‖wOf L − tensCLM (kerA M.A) u‖ ≤ εA 2^{−p}` (base `tensor_fp`).
* **`passB`**: `RepFull ys y`, `‖y‖ ≤ 1` ↦ `RepPad L (vOf L)`, `GB (p+2) L` and
  `‖vOf L − tensCLM (kerB M.B) y‖ ≤ εB 2^{−p}`. The padded entries are exactly zero because
  `B̃₀ 0 = 0` (`B0P_zero`) and the slice maps are local (`sliceCompZ_local`).
* `discL_stepA/B`: every axis step keeps the entries in the disc `‖·‖ ≤ 2^p`, so all
  intermediate arrays are `GB (p+2)` (what the rotations and nlogn's slice programs need).
* `getD_listOf'`, `rmPos_digL`, `listOf_self`: `listOf` vs the digits `digL` of `KFMBridge`.
-/

set_option linter.unusedSimpArgs false
set_option linter.unusedVariables false
set_option linter.unusedSectionVars false

namespace NLogN.KFY

open IntMul IntMul.HvdH IntMul.HvdH.Ref IntMul.HvdH.RefPf NLogN.KFM IntMul.OpenAI.RT Complex
open KFB (GI gflat)

/-! ## Row-major digits and `listOf` -/

/-! ## Box points as natural-number index functions -/

variable {d : ℕ}

/-- Coordinates of a box point, extended by `0` beyond `d`. -/
def extF (j : Fin d → ℕ) : ℕ → ℕ := fun i => if h : i < d then j ⟨i, h⟩ else 0

/-- The input of a pass as a box function. -/
def gIn (dims : List ℕ) (zs : List GI) : (Fin d → ℕ) → GI :=
  fun j => zs.getD (rmPos dims (extF j)) 0

/-! ## Discs and bit bounds -/

/-- Every entry in the disc of radius `2^p`. -/
def DiscL (p : ℕ) (L : List GI) : Prop := ∀ z ∈ L, ‖gC z‖ ≤ 2 ^ p

/-! ## The line maps -/

/-- The padded `B̃₀` line map: `B̃₀` on `[0, s)`, zero on `[s, t)`. -/
def PB (p s t α : ℕ) (U : ℕ → ℤ × ℤ) (l : ℕ) : ℤ × ℤ := if l < s then B0P p s t α U l else (0, 0)

/-- The padded `B̃₀` list map. -/
def b0Pad (p s t α : ℕ) (X : List GI) : List GI :=
  (List.range s).map (B0P p s t α fun j => X.getD j (0, 0)) ++ List.replicate (t - s) (0, 0)

/-! ## Locality of slice maps -/

section padded

variable (p α : ℕ) (s t : Fin d → ℕ)

/-- The padded `B̃₀` maps of the axes. -/
def PBf : Fin d → (ℕ → ℤ × ℤ) → (ℕ → ℤ × ℤ) := fun i => PB p (s i) (t i) α

/-- The plain `B̃₀` maps of the axes. -/
def B0f : Fin d → (ℕ → ℤ × ℤ) → (ℕ → ℤ × ℤ) := fun i => B0P p (s i) (t i) α

end padded

/-! ## The `A` pass -/

section passA

variable (p α : ℕ) (s t : Fin d → ℕ) [∀ i, NeZero (s i)] [∀ i, NeZero (t i)]

/-- The `A` maps for `tensorGo` (axis `i` of the padded box, all `t_i` long). -/
def mapsA : List (ℕ × (List GI → List GI)) :=
  (List.range d).map fun i => ((List.ofFn t).getD i 1,
    sapL p ((List.ofFn s).getD i 1) ((List.ofFn t).getD i 1) α)

/-- The `B₀` maps for `tensorGo` (padded back to `t_i`). -/
def mapsB : List (ℕ × (List GI → List GI)) :=
  (List.range d).map fun i => ((List.ofFn t).getD i 1,
    b0Pad p ((List.ofFn s).getD i 1) ((List.ofFn t).getD i 1) α)

/-- The output of the `A` pass as a function on the `t`-box (scale `2^{-p}`). -/
noncomputable def wOf (L : List GI) : ((i : Fin d) → ZMod (t i)) → ℂ :=
  fun pt => gC (L.getD (rmPos (List.ofFn t) (extF fun i => (pt i).val)) 0) / 2 ^ p

/-- The output of the `B₀` pass as a function on the `s`-box (scale `2^{-p}`). -/
noncomputable def vOf (L : List GI) : ((i : Fin d) → ZMod (s i)) → ℂ :=
  fun pt => gC (L.getD (rmPos (List.ofFn t) (extF fun i => (pt i).val)) 0) / 2 ^ p

/-! ## The `B₀` pass -/

/-! ## Each axis step keeps the disc -/

end passA

end NLogN.KFY

end

section
-- Source module: Solutions.KFY2Trans
/-!
# KFY 2: the axis rotation of `TransNeed` as an `RM` procedure node

The line passes are register programs (`RM`, nlogn `HRM`). The axis exposure is `TransNeed`'s
rotation, an item-language procedure with its descriptor on an unknown stack `dT ≠ 0`. It enters
the `RM` program as one `RM.prc` node (`trPrc`) whose code calls a small wrapper procedure:

`trW X Y dT jT = xfer X Y dT; call jT; unxfer dT`

The caller builds the descriptor `descTr w h bs` on its own stack `X` above a mark (any `X, Y` with
`X, Y ≠ dT`). The wrapper moves it to `dT` (KFMSC `xfer`), calls the rotation, and drops it
again.

* `trF` / `trPre` / `trC`: the functional specification, the precondition and the cost of the
  node. `trF` decodes the descriptor and the array from the state, so it is an honest function
  of the state. The cost is an upper bound (`IRuns` bounds costs from above), so the node does not
  depend on the input tape `Isym`.
* **`trW_runs`** / **`trPrc_ok`** (`PrcOK`): the wrapper meets the node's specification, from the
  `TransNeed` run (stated here for any table that contains the wrapper and the rotation).
* **`trF_eq`**, **`trC_le`**: on a state with explicit `w h bs zs l R`, `trF` is the rotation
  `transposeL t (V/t)` on stack `0`, and its cost is the `TransNeed` bound plus `O(k B)`.
-/

set_option linter.unusedSimpArgs false
set_option linter.unusedVariables false
set_option linter.unusedSectionVars false

namespace NLogN.KFY

open IntMul IntMul.HvdH IntMul.HvdH.Ref NLogN.KFM KIL KFB AModel

/-! ## Decoding -/

/-- The integer of an item (`0` for a mark). -/
def itZ : Item → ℤ
  | .int z => z
  | .mark => 0

/-- The items above the first mark. -/
def blk (L : List Item) : List Item := L.takeWhile fun it => it != .mark

/-- The items below the first mark. -/
def aft (L : List Item) : List Item := (L.dropWhile fun it => it != .mark).tail

/-- `h` from a rotation descriptor `[w, h, k, b…]`. -/
def hOf (D : List Item) : ℕ := (itZ (D.getD 1 .mark)).toNat

/-- `w` from a rotation descriptor. -/
def wOfD (D : List Item) : ℕ := (itZ (D.getD 0 .mark)).toNat

/-- The bits from a rotation descriptor. -/
def bsOf (D : List Item) : List Bool := (D.drop 3).map fun it => decide (it = .int 1)

/-- The first `n` Gaussian pairs of an item list. -/
def pairsOf : ℕ → List Item → List GI
  | 0, _ => []
  | n + 1, a :: b :: rest => (itZ a, itZ b) :: pairsOf n rest
  | _ + 1, _ => []

/-! ## The node -/

section node

variable (X dT : Fin HS) (Kc : ℝ)

/-- Shape of the descriptor on `X`. -/
def trSh (u : IState HS) : Shape := shPad (hOf (NLogN.KFY.blk (u.G X))) (bsOf (NLogN.KFY.blk (u.G X)))

/-- **The functional specification**: drop the descriptor block on `X`, rotate stack `0`. -/
def trF (u : IState HS) : IState HS :=
  (u.upd X (aft (u.G X))).upd 0
    (ints (gflat (transposeL ((trSh X u).headD 1) (vol (trSh X u).tail)
      (pairsOf (vol (trSh X u)) (u.G 0)))) ++ (u.G 0).drop (2 * vol (trSh X u)))

/-- The precondition: a rotation descriptor over a mark on `X`, a matching array on `0`. -/
def trPre (u : IState HS) : Prop :=
  ∃ (w h : ℕ) (bs : List Bool) (zs : List GI) (l R : List Item),
    u.G X = ints (descTr w h bs) ++ .mark :: R ∧ u.G 0 = ints (gflat zs) ++ l ∧
      1 ≤ h ∧ bs ≠ [] ∧ zs.length = vol (shPad h bs) ∧ NLogN.KFB.GB w zs

/-- The `TransNeed` cost bound for `w h k V`. -/
noncomputable def trBound (w h k V : ℕ) : ℝ :=
  Kc * (((k + 1 : ℕ) : ℝ) * ((w : ℝ) + 1) * ((V : ℝ) + 1) * ((h : ℝ) ^ KFS.τS + 1) +
    (((k + 1) * (h + w + k + 1) ^ 2 : ℕ) : ℝ))

/-- Bit size of the descriptor entries. -/
def trB (w h k : ℕ) : ℕ := w + h + k + 1

/-- The overhead of the wrapper (`xfer`, `unxfer`, calls). -/
def trOver (w h k : ℕ) : ℕ := 3 * (k + 3) * NLogN.KFM.cMv (trB w h k) + 110

/-- The cost of the node (an upper bound). -/
noncomputable def trC (u : IState HS) : ℕ :=
  ⌈trBound Kc (wOfD (NLogN.KFY.blk (u.G X))) (hOf (NLogN.KFY.blk (u.G X))) (bsOf (NLogN.KFY.blk (u.G X))).length
      (vol (trSh X u))⌉₊ +
    trOver (wOfD (NLogN.KFY.blk (u.G X))) (hOf (NLogN.KFY.blk (u.G X))) (bsOf (NLogN.KFY.blk (u.G X))).length

theorem trF_frame (u : IState HS) (j : Fin HS) (hj : j ∉ [X, 0]) : (trF X u).G j = u.G j := by
  simp only [List.mem_cons, List.not_mem_nil, or_false, not_or] at hj
  simp [trF, hj.1, hj.2]

theorem trF_io (u : IState HS) : (trF X u).ip = u.ip ∧ (trF X u).out = u.out := by
  simp [trF]

/-- **The rotation node.** Procedure `pc 1` is the wrapper. -/
noncomputable def trPrc : RM :=
  .prc 0 (trF X) (trC X Kc) (trPre X) [X, 0] (trF_frame X) (trF_io X)

end node

/-! ## The wrapper procedure -/

/-- The wrapper body: move the descriptor to `dT`, rotate, drop the descriptor. -/
def trW {P : ℕ} (X Y dT : Fin HS) (jT : Fin P) : IStmt HS P :=
  .seq (xfer X Y dT) (.seq (.call jT) (unxfer dT))

/-- The rotation, as an `IRuns` hypothesis on a table (the `TransNeed` clause, relabelled). -/
def TrHyp {P : ℕ} (I : List Sym) (iprog : Fin P → IStmt HS P) (jT : Fin P) (dT : Fin HS)
    (Kc : ℝ) : Prop :=
  ∀ (w h : ℕ) (bs : List Bool), 1 ≤ h → bs ≠ [] →
  ∀ (s : IState HS) (f : Fin HS → List ℤ) (zs : List GI) (l : List ℤ),
    (∀ j, j ≠ 0 → j ≠ dT → f j = []) →
    f dT = descTr w h bs →
    f 0 = gflat zs ++ l → zs.length = vol (shPad h bs) → NLogN.KFB.GB w zs →
    ∃ c : ℕ, IRuns I iprog (.call jT) (NLogN.KIL.encK s f)
        (NLogN.KIL.encK s (Function.update f 0
          (gflat (transposeL ((shPad h bs).headD 1) (vol (shPad h bs).tail) zs) ++ l))) c ∧
      (c : ℝ) ≤ trBound Kc w h bs.length (vol (shPad h bs))

end NLogN.KFY

end

section
-- Source module: Solutions.KFY6Gen
/-!
# KFY 6: generic combinators for the line passes

* **`brk l P`**: push the register list `l`, run `P`, pop the registers again. If `P` changes
  only one stack `j ∉ l` (`P.run t' = t'.upd j Z`), so does the bracket (**`brk_all`**).
* **`trBlk X Kc rw rm`**: the axis rotation of `TransNeed` as an `RM` block: mark `X`, build the
  rotation descriptor (KFMSE `descTrP`, `w` from `rw`, rotation count from `rm`), run the
  procedure node `trPrc` (KFY2) (**`trBlk_all`**).
-/

set_option linter.unusedSimpArgs false
set_option linter.unusedVariables false
set_option linter.unusedSectionVars false

namespace NLogN.KFY

open IntMul NLogN

/-! ## Brackets -/

/-- Push `l`, run `P`, pop `l`. -/
def brk (l : List (Fin HS × Ex HS)) (P : RM) : RM :=
  .seq (.rs (NLogN.KFY.pushL l)) (.seq P (.rs (popL (l.map Prod.fst))))

/-- The cost of a bracket around a program of cost `C`. -/
def brkC (n K B C : ℕ) : ℕ := n * ((K + 1) * 10 ^ 7 * (B + 2) ^ 2) + 1 + C + (n * (2 * B + 2) + 1)

/-! ## The rotation block -/

open NLogN.KFM in
/-- Mark `X`, build the rotation descriptor on `X`, rotate. -/
noncomputable def trBlk (X : Fin HS) (Kc : ℝ) (rw rm : Fin HS) : RM :=
  .seq (.op (.mark X)) (.seq (.rs (descTrP X rw rm)) (trPrc X Kc))

end NLogN.KFY

end

section
-- Source module: Solutions.KFY7Pro
/-!
# KFY 7: the prologue of the line passes

The line-map descriptor `[p, α, h, d, b_0..b_{d−1}, s_0..s_{d−1}]` (top first) lies on the
descriptor stack `dS = 26`. The prologue `proP` takes `p, α, ℓ = h + 1, d` into registers
(`yP = 11`, `yA = 12`, KFM's `zEl = 29`, `zD = 25`), counts the zero bits into `d'` (KFM's
`zDp = 31`, the bits are `[d' ≤ j]`), and pushes `L = dℓ − d'` (`yL = 13`, the base-2 logarithm
of the array length). The sources stay on `dS` (**`proP_all`**).

All registers are pushed above the unknown contents of their stacks.
-/

set_option linter.unusedSimpArgs false
set_option linter.unusedVariables false
set_option linter.unusedSectionVars false

namespace NLogN.KFY

open IntMul NLogN

abbrev dS : Fin HS := 26
abbrev yP : Fin HS := 11
abbrev yA : Fin HS := 12
abbrev yL : Fin HS := 13
abbrev yW : Fin HS := 14
abbrev yS : Fin HS := 15

/-! ## Generic facts -/

/-- `run`, `ok` and a cost bound of an `RS` program. -/
def RSAll (P : RS) (s t : IState HS) (C : ℕ) : Prop := P.ok s ∧ P.run s = t ∧ P.cost s ≤ C

/-! ## Taking a descriptor entry -/

/-- Push `e` on `r`, pop the descriptor stack. -/
def takeD (r : Fin HS) (e : Ex HS) : RS := .seq (.push r e) (.pop NLogN.KFY.dS)

/-! ## Counting the zero bits -/

/-- The bits `[d' ≤ j]`, `j < d`. -/
def bitL (d d' : ℕ) : List ℤ := (List.range d).map fun j => if d' ≤ j then 1 else 0

/-- One bit: `d' += 1 − b`, drop `b`. -/
def dpB : RS :=
  .seq (.asg KFM.zDp (.add (.reg KFM.zDp) (.add (.lit 1) (.neg (.reg NLogN.KFY.dS))))) (.pop NLogN.KFY.dS)

def dpLoop : RS := .rep KFM.mOc (.reg KFM.zD) dpB

/-- The state of the bit loop after `k` bits. -/
def dpSt (t : IState HS) (d d' : ℕ) (srcs : List ℤ) (R LZ : List Item) (k : ℕ) : IState HS :=
  (t.upd NLogN.KFY.dS (ints ((bitL d d').drop k ++ srcs) ++ R)).upd KFM.zDp (.int (min k d' : ℕ) :: LZ)

/-! ## The prologue -/

/-- `L = dℓ − d'`. -/
def lE : Ex HS := .add (.mul (.reg KFM.zD) (.reg KFM.zEl)) (.neg (.reg KFM.zDp))

def proP : RS :=
  .seq (takeD NLogN.KFY.yP (.reg NLogN.KFY.dS)) (.seq (takeD yA (.reg NLogN.KFY.dS)) (.seq (takeD KFM.zEl (.add (.reg NLogN.KFY.dS) (.lit 1)))
    (.seq (takeD KFM.zD (.reg NLogN.KFY.dS)) (.seq (.push KFM.zDp (.lit 0)) (.seq dpLoop (.push yL lE))))))

/-- The state after the prologue. -/
def proSt (t : IState HS) (p α h d d' : ℕ) (srcs : List ℤ) (R : List Item) : IState HS :=
  ((((((t.upd NLogN.KFY.dS (ints srcs ++ R)).upd NLogN.KFY.yP (.int p :: t.G NLogN.KFY.yP)).upd yA (.int α :: t.G yA)).upd
    KFM.zEl (.int ((h + 1 : ℕ) : ℤ) :: t.G KFM.zEl)).upd KFM.zD (.int d :: t.G KFM.zD)).upd
    KFM.zDp (.int d' :: t.G KFM.zDp)).upd yL (.int ((d * (h + 1) - d' : ℕ) : ℤ) :: t.G yL)

/-- The cost bound of the prologue. -/
def proC (d B : ℕ) : ℕ := (d + 20) * (20 * 10 ^ 7 * (B + 2) ^ 2)

end NLogN.KFY

end

section
-- Source module: Solutions.KFY8Wid
/-!
# KFY 8: the widths `sWidth p α = ⌈√p⌉ α` and `eWidth p α = ⌈(⌈√p⌉ + α − 1)/α⌉ + 1`

`⌈√p⌉ = #{x ∈ [1, p] : (x − 1)² < p}` (KFMS2's counting loop `cntP`, counter `sI = 20`),
then one assignment. The result is pushed on `yW = 14` (**`wA_all`**, **`wB_all`**).
-/

set_option linter.unusedSimpArgs false
set_option linter.unusedVariables false
set_option linter.unusedSectionVars false

namespace NLogN.KFY

open IntMul NLogN
open IntMul.HvdH (ceilSqrt)

/-- The test `(x − 1)² − p` for `x` on `sI`. -/
def sqT : Ex HS :=
  .add (.mul (.add (.reg KFM.sI) (.lit (-1))) (.add (.reg KFM.sI) (.lit (-1)))) (.neg (.reg NLogN.KFY.yP))

/-- `⌈√p⌉` on `yW`. -/
def csP : RS := KFM.cntP (.reg NLogN.KFY.yP) sqT NLogN.KFY.yW

/-- `sWidth`: `⌈√p⌉ α`. -/
def wA : RS := .seq csP (.asg NLogN.KFY.yW (.mul (.reg NLogN.KFY.yW) (.reg yA)))

/-- `eWidth`: `(⌈√p⌉ + α − 1)/α + 1`. -/
def wB : RS :=
  .seq csP (.asg NLogN.KFY.yW (.add (.ndiv (.add (.reg NLogN.KFY.yW) (.add (.reg yA) (.lit (-1)))) (.reg yA)) (.lit 1)))

end NLogN.KFY

end

section
-- Source module: Solutions.KFY9Epi
/-!
# KFY 9: the epilogue of the line passes

After the axis loop the sources lie reversed on `yS`; the epilogue moves them back to `dS`,
pushes the bits `[d' ≤ j]` (KFMSB `bitP`), `d`, `h = ℓ − 1`, `α`, `p` on `dS` (restoring the
line-map descriptor) and pops the registers `yW, yL, zDp, zD, zEl, yA, yP` (**`epiP_all`**).
-/

set_option linter.unusedSimpArgs false
set_option linter.unusedVariables false
set_option linter.unusedSectionVars false

namespace NLogN.KFY

open IntMul NLogN

/-- The registers of the passes (popped by the epilogue). -/
def regsP : List (Fin HS) := [NLogN.KFY.yW, yL, KFM.zDp, KFM.zD, KFM.zEl, yA, NLogN.KFY.yP]

def epiP : RS :=
  .seq (mvList yS NLogN.KFY.dS KFM.mOc (.reg KFM.zD)) (.seq (KFM.bitP NLogN.KFY.dS) (.seq (.push NLogN.KFY.dS (.reg KFM.zD))
    (.seq (.push NLogN.KFY.dS KFM.hE) (.seq (.push NLogN.KFY.dS (.reg yA)) (.seq (.push NLogN.KFY.dS (.reg NLogN.KFY.yP))
      (popL regsP))))))

/-- The descriptor `[p, α, h, d, bits, sources]`. -/
def descL (p α h d d' : ℕ) (srcs : List ℤ) : List ℤ :=
  (p : ℤ) :: (α : ℤ) :: (h : ℤ) :: (d : ℤ) :: (bitL d d' ++ srcs)

/-- The state after the epilogue. -/
def epiSt (v : IState HS) (LS : List Item) (D : List ℤ) (R : List Item) : IState HS :=
  { v with G := fun j => if j ∈ regsP then (v.G j).tail else if j = yS then LS else
      if j = NLogN.KFY.dS then ints D ++ R else v.G j }

/-- The cost bound of the epilogue. -/
def epiC (d B : ℕ) : ℕ := (d + 10) * (30 * 10 ^ 7 * (B + 2) ^ 2)

end NLogN.KFY

end

section
-- Source module: Solutions.KFY4Sap
/-!
# KFY 4: the `A` line pass on one axis (`sapAxQ`)

nlogn's `sapAxP` (`HAx4`) with two changes:

* no transposition prelude (`coreQ` instead of `GL.axisP`): the slices are the consecutive
  slices of the array on `gD` (the rotation is done before, by `TransNeed`);
* padded slices: every slice has length `t` (the padded axis) and only its first `s` entries
  are read (`padIn`, through stack `41`, counter `47`); the slice map is
  `x ↦ sapL p s t α (x.take s)` (length `t`).

**`sapAxQ_all`**: on `gD = gints L ++ R`, `L.length = m t`, it computes
`gints ((chunkL t L).flatMap (sapL p s t α ∘ take s)) ++ R`; every other stack is restored.
-/

set_option linter.unusedSimpArgs false
set_option linter.unusedVariables false
set_option linter.unusedSectionVars false

namespace NLogN.KFY

open IntMul NLogN IntMul.HvdH.Ref NLogN.GL NLogN.AX NLogN.Rsp

/-- The axis core, written out. -/
def sapCoreQ : RM :=
  .seq (.rep gCC (.reg aM) (RM.seqL [.rs (mvList NLogN.GL.gD gZT gCV (cnt2 aT)),
      .rs (mvList gZT aI gCV (cnt2 aT)), padInP aI 41 47 NLogN.AX.aS aT sapSlP,
      .rs (mvList aO gZB gCV (cnt2 aT))]))
    (.rs (mvList gZB NLogN.GL.gD gCV (cnt2 aTot)))

/-- **The padded `S̃` axis program.** -/
def sapAxQ : RM :=
  RM.seqL [.rs sTabR, .rs dTabR, sapCoreQ, .rs (popK aW gCC (.reg aNW)),
    .rs (popK NLogN.AX.aD gCC (.reg aT))]

/-- The scalar registers read by `sapAxQ`. -/
structure SapRegsQ (p s t α m : ℕ) (st : IState HS) : Prop where
  tab : STabRegs p s t α st
  rKa : ∃ L, st.G aKa = .int (sOff s (HvdH.sWidth p α) : ℤ) :: L
  rKN : ∃ L, st.G aKN = .int ((s + (2 * HvdH.sWidth p α + 1) : ℕ) : ℤ) :: L
  rR : ∃ L, st.G aR = .int (2 ^ (2 * p)) :: L
  rNW : ∃ L, st.G aNW = .int ((t * (2 * HvdH.sWidth p α + 1) : ℕ) : ℤ) :: L
  rM : ∃ L, st.G aM = .int (m : ℤ) :: L
  rTot : ∃ L, st.G aTot = .int ((m * t : ℕ) : ℤ) :: L

/-- Cost bound of `sapAxQ`. -/
def sapAxCostQ (p s t α m mm Φ Bacc Q Bq : ℕ) : ℕ :=
  tabCost (2 * HvdH.sWidth p α + 1) t (sWCost p Bq) +
    (tabCost 1 t (40 * 10 ^ 7 * (Bq + 2) ^ 2) + 10 ^ 9) +
    (m * (sapSlC p s t α mm Φ Bacc Q + 100000 * (t + 1) * (Q + 2) +
      100000 * (t + t + 1) * (Q + 2)) + 100000 * (m * t + 1) * (Q + 2)) +
    (2 * 2000 * (Bq + 2) + t * (2 * HvdH.sWidth p α + 1) * (12000 * (Bq + 2)) + 32) +
    (2 * 2000 * (Bq + 2) + t * (12000 * (Bq + 2)) + 32)

/-- The output of the padded `S̃` axis. -/
def sapOutQ (p s t α : ℕ) (L : List (ℤ × ℤ)) : List (ℤ × ℤ) :=
  (chunkL t L).flatMap fun x => sapL p s t α (x.take s)

end NLogN.KFY

end

section
-- Source module: Solutions.KFY10Iter
/-!
# KFY 10: `tensorGo` as an iteration over rotated shapes

`stepL R f L`: transpose the head axis of the shape `R` to the end (`transposeL`), apply `f` to
the consecutive lines. `iterL dims fs k` applies the steps for the rotations `dims.rotate j`,
`j < k`. **`tensorGo_iterL`**: `tensorGo` of the list `(dims_i, fs i)_i` is `iterL dims fs d`.

Also the rotated padded shapes: **`shPad_bsN`**: `shPad (ℓ − 1) (bsN n k)` is the box
`tB n` rotated by `k` (the shape `TransNeed` sees at step `k`).
-/

set_option linter.unusedSimpArgs false
set_option linter.unusedVariables false

namespace NLogN.KFY

open IntMul IntMul.HvdH IntMul.HvdH.Ref NLogN.KFM IntMul.OpenAI.Glue IntMul.OpenAI.RP
open KFB (GI gflat)

/-- One step of the pass over the shape `R`. -/
def stepL (R : List ℕ) (f : List GI → List GI) (L : List GI) : List GI :=
  (chunkL (R.headD 1) (transposeL (R.headD 1) R.tail.prod L)).flatMap f

/-- The first `k` steps over the rotations of `dims`. -/
def iterL (dims : List ℕ) (fs : ℕ → List GI → List GI) : ℕ → List GI → List GI
  | 0, L => L
  | k + 1, L => stepL (dims.rotate k) (fs k) (iterL dims fs k L)

/-- The map list of a pass. -/
def mapsL (dims : List ℕ) (fs : ℕ → List GI → List GI) : List (ℕ × (List GI → List GI)) :=
  (List.range dims.length).map fun i => (dims.getD i 1, fs i)

/-! ## Lengths and bounds along the iteration -/

/-! ## The rotated padded shapes -/

end NLogN.KFY

end

section
-- Source module: Solutions.KFY11Facts
/-!
# KFY 11: size facts of the padded box

* `2 ≤ ℓ`, `d' < d`, `dℓ − d' = L`;
* **`prod_tB`**: `∏ t_i = 2^L` (`= TP n`);
* the head and the tail product of the rotated box (`t_k`, `2^{L − log t_k}`).
-/

set_option linter.unusedSimpArgs false
set_option linter.unusedVariables false

namespace NLogN.KFY

open IntMul IntMul.OpenAI.Glue IntMul.OpenAI.RP NLogN.KFM

/-- The bit of axis `j`. -/
noncomputable def bN (n j : ℕ) : ℕ := if dpN n ≤ j then 1 else 0

end NLogN.KFY

end

section
-- Source module: Solutions.KFY12AxA
/-!
# KFY 12: one axis of the `A` pass

`axA X Kc` (one iteration of the axis loop, counter `cX = 55` holding `d − k`):

1. `takeD yS (reg dS)`: move `s_k` from the descriptor stack to `yS`;
2. push `b_k = [d' ≤ k]` (`yB`), then `t_k = 2^{ℓ−1+b_k}` (`yT`), `m_k = 2^{L−ℓ+1−b_k}` (`yM`);
3. push `w = p + 2` (`zw`) and `k` (`zm`), rotate (`trBlk`: `TransNeed`'s axis rotation of the
   box `t` rotated by `k`), pop;
4. push the registers of `sapAxQ` (`regsA`), run `sapAxQ` (the padded `S̃` on the `m_k`
   consecutive lines of length `t_k`), pop.

**`axA_all`**: the array `X` on stack `0` becomes `stepL R_k (sapL p s_k t_k α) X` with
`R_k = (ofFn t).rotate k`; `s_k` moves from `dS` to `yS`; nothing else changes.
-/

set_option linter.unusedSimpArgs false
set_option linter.unusedVariables false
set_option linter.unusedSectionVars false

namespace NLogN.KFY

open IntMul NLogN IntMul.HvdH.Ref NLogN.GL NLogN.AX NLogN.Rsp
open NLogN.KFM (shPad bsN)

abbrev yB : Fin HS := 16
abbrev yT : Fin HS := 18
abbrev yM : Fin HS := 19
abbrev zw : Fin HS := 22
abbrev zm : Fin HS := 23
/-- The axis counter. -/
abbrev cX : Fin HS := 55

/-- `k = d − (d − k)`. -/
def kE : Ex HS := .add (.reg KFM.zD) (.neg (.reg cX))
/-- `b_k = [d' ≤ k]`. -/
def bitE : Ex HS := .sel (.add kE (.neg (.reg KFM.zDp))) (.lit 0) (.lit 1)
/-- `t_k = 2^{ℓ − 1 + b_k}`. -/
def tE : Ex HS := .pow2 (.add (.add (.reg KFM.zEl) (.lit (-1))) (.reg NLogN.KFY.yB))
/-- `m_k = 2^{L − ℓ + 1 − b_k}`. -/
def mE : Ex HS :=
  .pow2 (.add (.add (.reg yL) (.neg (.reg KFM.zEl))) (.add (.lit 1) (.neg (.reg NLogN.KFY.yB))))
/-- `w = p + 2`. -/
def wE : Ex HS := .add (.reg NLogN.KFY.yP) (.lit 2)

/-- `2 W + 1` from `yW`. -/
def nrE : Ex HS := .add (.add (.reg NLogN.KFY.yW) (.reg NLogN.KFY.yW)) (.lit 1)

/-- The registers of `sapAxQ`. -/
def regsA : List (Fin HS × Ex HS) :=
  [(aP, .reg NLogN.KFY.yP), (NLogN.AX.aS, .reg yS), (aT, .reg yT), (aWd, .reg NLogN.KFY.yW),
    (aBe, .ndiv (.pow2 (.reg NLogN.KFY.yP)) (.add (.reg yA) (.reg yA))),
    (aK2, .mul (.mul (.reg yA) (.reg yA)) (.mul (.reg yT) (.reg yT))),
    (aNR, nrE),
    (aKa, .nmod (.add (.reg yS) (.neg (.nmod (.reg NLogN.KFY.yW) (.reg yS)))) (.reg yS)),
    (aKN, .add (.reg yS) nrE),
    (aR, .pow2 (.add (.reg NLogN.KFY.yP) (.reg NLogN.KFY.yP))),
    (aNW, .mul (.reg yT) nrE),
    (aM, .reg NLogN.KFY.yM), (aTot, .mul (.reg NLogN.KFY.yM) (.reg yT))]

/-- **One axis of the `A` pass.** -/
noncomputable def axA (X : Fin HS) (Kc : ℝ) : RM :=
  .seq (.rs (takeD yS (.reg NLogN.KFY.dS))) (brk [(NLogN.KFY.yB, bitE)] (brk [(yT, NLogN.KFY.tE), (NLogN.KFY.yM, mE)]
    (.seq (brk [(zw, wE), (zm, kE)] (trBlk X Kc zw zm)) (brk regsA sapAxQ))))

/-! ## Small facts -/

/-- The cost bound of one `A` axis. -/
noncomputable def axAC (Kc : ℝ) (n k B K Ctr Csap : ℕ) : ℕ :=
  (2 * 10 ^ 7 * (B + 2) ^ 2 + (2 * B + 2)) +
    brkC 1 K B (brkC 2 K B (brkC 2 K B (1 + (10 ^ 12 * (IntMul.OpenAI.Glue.bP n + 1) ^ 3 + Ctr)) +
      brkC 13 K B Csap))

/-- Size bounds for one axis (`A` pass). -/
structure AxBnd (n A B : ℕ) : Prop where
  env : KFM.CostEnv A B
  pA : 2 ^ (2 * IntMul.OpenAI.Glue.pP n) ≤ A
  αA : IntMul.OpenAI.Glue.αP n ^ 2 * IntMul.OpenAI.Glue.rP n ^ 2 ≤ A
  wA : IntMul.OpenAI.Glue.rP n * (2 * HvdH.sWidth (IntMul.OpenAI.Glue.pP n) (IntMul.OpenAI.Glue.αP n) + 1) + IntMul.OpenAI.Glue.rP n ≤ A
  LA : 2 ^ IntMul.OpenAI.Glue.LP n ≤ A
  dA : IntMul.OpenAI.Glue.dP n ≤ A
  pB : 2 * IntMul.OpenAI.Glue.pP n + 1 ≤ B
  LB : IntMul.OpenAI.Glue.LP n + 2 ≤ B

end NLogN.KFY

end

section
-- Source module: Solutions.KFY14Loop
/-!
# KFY 14: the axis loop

`RM.rep cX (reg zD) ax`: `d` iterations of an axis body `ax`. **`loop_all`**: if every iteration
`k` moves `s_k` from `dS` to `yS` and replaces the array `X` on stack `0` by
`stepL (dims.rotate k) (fs k) X` (for arrays satisfying an invariant `P` that the steps keep),
the loop ends in `loopSt … d`: the array is `iterL dims fs d zs`, the sources lie reversed on
`yS`, nothing else changed.
-/

set_option linter.unusedSimpArgs false
set_option linter.unusedVariables false
set_option linter.unusedSectionVars false

namespace NLogN.KFY

open IntMul NLogN IntMul.HvdH.Ref

/-- The loop state after `k` axes. -/
def loopSt (u0 : IState HS) (srcs : List ℤ) (R R0 : List Item) (dims : List ℕ)
    (fs : ℕ → List KFB.GI → List KFB.GI) (zs : List KFB.GI) (k : ℕ) : IState HS :=
  ((u0.upd yS (ints (srcs.take k).reverse ++ u0.G yS)).upd NLogN.KFY.dS (ints (srcs.drop k) ++ R)).upd 0
    (ints (KFB.gflat (NLogN.KFY.iterL dims fs k zs)) ++ R0)

end NLogN.KFY

end

section
-- Source module: Solutions.KFY15Pass
/-!
# KFY 15: a whole line pass (generic in the width program and the axis body)

`passP wP ax = rs (proP; wP) ; rep cX (reg zD) ax ; rs epiP`.

**`pass_all`**: on the descriptor `[p, α, ℓ − 1, d, bits, s_0..s_{d−1}]` on `dS` and the array
`zs` on stack `0`, if `wP` pushes the width `w` on `yW` and the axis body `ax` performs axis
step `k` (in the register layout of the prologue), the pass leaves everything unchanged except
stack `0`, which holds `iterL (ofFn t) fs d zs`.
-/

set_option linter.unusedSimpArgs false
set_option linter.unusedVariables false
set_option linter.unusedSectionVars false

namespace NLogN.KFY

open IntMul NLogN IntMul.HvdH.Ref NLogN.GL NLogN.AX NLogN.Rsp

/-- A line pass. -/
def passP (wP : RS) (ax : RM) : RM :=
  .seq (.rs (.seq proP wP)) (.seq (.rep cX (.reg KFM.zD) ax) (.rs epiP))

/-- The sources `s_0..s_{d−1}`. -/
noncomputable def srcsN (n : ℕ) : List ℤ := (List.range (IntMul.OpenAI.Glue.dP n)).map fun j => (KFM.srcN n j : ℤ)

/-- The cost bound of a pass. -/
def passC (d B Cw C : ℕ) : ℕ :=
  (proC d B + Cw) + ((2 * 10 ^ 7 * (B + 2) ^ 2 + d * (C + 4 * 10 ^ 7 * (B + 2) ^ 2 + 30) + 32) +
    epiC d B)

end NLogN.KFY

end

section
-- Source module: Solutions.KFY5Bap
/-!
# KFY 5: the `B₀` line pass on one axis (`bap0AxQ`)

`KFMBap0Ax.bap0AxP` with two changes:

* no transposition prelude (`coreQ` instead of `GL.axisP`): the slices are the consecutive
  slices of the array on `gD` (the rotation is done before, by `TransNeed`);
* padded output: every output slice has length `t`: the `s` values of `b0L` followed by `t − s`
  zeros (`padOut`, counter `10`). The slice map is `x ↦ b0L p s t α x ++ replicate (t − s) 0`.

**`bap0AxQ_all`**: on `gD = gints L ++ R`, `L.length = m t`, it computes
`gints (b0OutQ p s t α L) ++ R`; every other stack is restored. The proof is
`bap0AxP_all` with the axis lemma replaced by `coreQ_all`.
-/

set_option linter.unusedSimpArgs false
set_option linter.unusedVariables false
set_option linter.unusedSectionVars false

namespace NLogN.KFY

open IntMul NLogN IntMul.HvdH.Ref NLogN.GL NLogN.AX NLogN.Rsp NLogN.KFM

/-- The axis core, written out. -/
def bap0CoreQ : RM :=
  .seq (.rep gCC (.reg bM) (RM.seqL [.rs (mvList NLogN.GL.gD gZT gCV (cnt2 NLogN.AX.bT)),
      .rs (mvList gZT bI gCV (cnt2 NLogN.AX.bT)), padOutP bO 10 bS NLogN.AX.bT bap0SlP,
      .rs (mvList bO gZB gCV (cnt2 NLogN.AX.bT))]))
    (.rs (mvList gZB NLogN.GL.gD gCV (cnt2 bTot)))

/-- **The padded `B̃₀` axis program.** -/
def bap0AxQ : RM :=
  RM.seqL [.rs tmpP, .rs eTabR, .rs eDelR, .rs cTabR, .rs cDelR, .rs dTabR', .rs ptKR, .rs psKR,
    .rs tmpQ, bap0CoreQ, .rs (popK bW gCC (.reg bNW)), .rs (popK bD gCC (.reg bS)),
    .rs (popK bW1 gCC (.reg bS)), .rs (popK bD1 gCC (.reg bS)), .rs (popK bW3 gCC (.reg bS)),
    .rs (popK bK3 gCC (.reg NLogN.AX.bT)), .rs (popK bKS gCC (.reg bS))]

/-- The scalar registers read by `bap0AxQ` (`BapRegs` with the output size `m t`). -/
structure BapRegsQ (p s t α m kb : ℕ) (st : IState HS) : Prop where
  t1 : ∃ L, st.G NLogN.AX.bT = .int (t : ℤ) :: L
  kb2 : ∃ L, st.G bKB = .int (kb : ℤ) :: L
  ps : ∃ L, st.G bS = .int (s : ℤ) :: L
  pr : ∃ L, st.G bPR = .int (2 ^ p) :: L
  ks : ∃ L, st.G bKs = .int (s : ℤ) :: L
  ka : ∃ L, st.G bKa = .int (sOff s (HvdH.eWidth p α) : ℤ) :: L
  kn : ∃ L, st.G bKN = .int ((s + (2 * HvdH.eWidth p α + 1) : ℕ) : ℤ) :: L
  br : ∃ L, st.G NLogN.AX.bR = .int (2 ^ p) :: L
  bn : ∃ L, st.G NLogN.AX.bN = .int (s : ℤ) :: L
  bnr : ∃ L, st.G bNR = .int ((2 * HvdH.eWidth p α + 1 : ℕ) : ℤ) :: L
  bnw : ∃ L, st.G bNW = .int ((s * (2 * HvdH.eWidth p α + 1) : ℕ) : ℤ) :: L
  jl : ∃ L, st.G bJL = .int ((2 * s : ℕ) : ℤ) :: L
  jn : ∃ L, st.G bJN = .int (HvdH.jIter p s t α : ℤ) :: L
  rp : ∃ L, st.G bO = .int (p : ℤ) :: L
  ra : ∃ L, st.G bI = .int ((α ^ 2 : ℕ) : ℤ) :: L
  ru : ∃ L, st.G bKS = .int (((s : ZMod t)⁻¹).val : ℤ) :: L
  rm : ∃ L, st.G bM = .int (m : ℤ) :: L
  rtot : ∃ L, st.G bTot = .int ((m * t : ℕ) : ℤ) :: L

/-- Cost bound of `bap0AxQ`. -/
def bap0AxCostQ (p s t α m mm Φ Bacc Q Bq : ℕ) : ℕ :=
  16 * 10 ^ 7 * (Bq + 2) ^ 2 +
  (tabCost (2 * HvdH.eWidth p α + 1) s (eWCost p Bq) +
  ((tabCost 1 s ((eDelE.size + 1) * 10 ^ 7 * (s + 4) ^ 2) + 10 ^ 9) +
  ((tabCost 1 s (2 * 10 ^ 7 * (1 + 2) ^ 2) + 10 ^ 9) +
  ((tabCost 1 s ((cDelE.size + 1) * 10 ^ 7 * (Bq + 2) ^ 2) + 10 ^ 9) +
  ((tabCost 1 s (dTWCost p (2 * α ^ 2 - 2) Bq) + 10 ^ 9) +
  ((tabCost 1 t ((ptKE.size + 1) * 10 ^ 7 * (Bq + 2) ^ 2) + 10 ^ 9) +
  ((tabCost 1 s ((psKE.size + 1) * 10 ^ 7 * (Bq + 2) ^ 2) + 10 ^ 9) +
  ((6 * Bq + 6) +
  ((m * (bap0SlC p s t α mm Φ Bacc Q + 100000 * (t + 1) * (Q + 2) +
      100000 * (t + t + 1) * (Q + 2)) + 100000 * (m * t + 1) * (Q + 2)) +
  ((2 * 2000 * (Bq + 2) + s * (2 * HvdH.eWidth p α + 1) * (12000 * (Bq + 2)) + 32) +
  ((2 * 2000 * (Bq + 2) + s * (12000 * (Bq + 2)) + 32) +
  ((2 * 2000 * (Bq + 2) + s * (12000 * (Bq + 2)) + 32) +
  ((2 * 2000 * (Bq + 2) + s * (12000 * (Bq + 2)) + 32) +
  ((2 * 2000 * (Bq + 2) + s * (12000 * (Bq + 2)) + 32) +
  ((2 * 2000 * (Bq + 2) + t * (12000 * (Bq + 2)) + 32) +
  (2 * 2000 * (Bq + 2) + s * (12000 * (Bq + 2)) + 32))))))))))))))))

/-- The output of the padded `B̃₀` axis. -/
def b0OutQ (p s t α : ℕ) (L : List (ℤ × ℤ)) : List (ℤ × ℤ) :=
  (chunkL t L).flatMap fun x => b0L p s t α x ++ List.replicate (t - s) (0, 0)

end NLogN.KFY

end

section
-- Source module: Solutions.KFY13AxB
/-!
# KFY 13: one axis of the `B₀` pass

`axB X Kc` is `axA X Kc` with two changes:

* after `t_k`, `m_k` it computes `u_k = (s_k mod t_k)^{-1}` into `yU = 17` (`invP`: KFM's
  `sumLoop`, `u = ∑_{y ∈ [1, t]} y·[y s ≡ 1 (mod t)]`, **`inv_sum`**), popped at the end;
* the line program is the padded `B̃₀` (`bap0AxQ`, registers `regsB`).

**`axB_all`**: the array `X` on stack `0` becomes `stepL R_k (b0Pad' p s_k t_k α) X`, with the
line map `x ↦ b0L p s t α x ++ replicate (t − s) 0`; `s_k` moves from `dS` to `yS`.
-/

set_option linter.unusedSimpArgs false
set_option linter.unusedVariables false
set_option linter.unusedSectionVars false

namespace NLogN.KFY

open IntMul NLogN IntMul.HvdH.Ref NLogN.GL NLogN.AX NLogN.Rsp
open NLogN.KFM (shPad bsN)

abbrev yU : Fin HS := 17

/-! ## The modular inverse -/

/-- `u = s^{-1} mod t` into `yU` (`s` on `yS`, `t` on `yT`). -/
def invP : RS :=
  .seq (.push KFM.mS (.reg yT)) (.seq (.push KFM.mP (.reg yS)) (.seq (.push KFM.mM (.lit 0))
    (.seq (.rep KFM.mIc (.reg KFM.mS) KFM.sumB) (.seq (.push NLogN.KFY.yU (.reg KFM.mM))
      (.seq (.pop KFM.mM) (.seq (.pop KFM.mP) (.pop KFM.mS)))))))

/-- The cost bound of `invP`. -/
def invC (t B : ℕ) : ℕ :=
  (t + 1) * 40 * (10 ^ 7 * (B + 2) ^ 2) + 8 * (10 ^ 7 * (B + 2) ^ 2) + 6 * B + 6

/-! ## The `B̃₀` registers -/

/-- `α² (t − s)`. -/
def dJE : Ex HS := .mul (.mul (.reg yA) (.reg yA)) (.add (.reg yT) (.neg (.reg yS)))
/-- `J = ⌈p s / (α² (t − s))⌉`. -/
def jE : Ex HS := .ndiv (.add (.add (.mul (.reg NLogN.KFY.yP) (.reg yS)) dJE) (.lit (-1))) dJE

/-- The registers of `bap0AxQ`. -/
def regsB : List (Fin HS × Ex HS) :=
  [(NLogN.AX.bT, .reg yT), (bKB, .reg yT), (bS, .reg yS), (bPR, .pow2 (.reg NLogN.KFY.yP)), (bKs, .reg yS),
    (bKa, .nmod (.add (.reg yS) (.neg (.nmod (.reg NLogN.KFY.yW) (.reg yS)))) (.reg yS)),
    (bKN, .add (.reg yS) nrE), (bR, .pow2 (.reg NLogN.KFY.yP)), (AX.bN, .reg yS), (bNR, nrE),
    (bNW, .mul (.reg yS) nrE), (bJL, .add (.reg yS) (.reg yS)), (bJN, jE), (bO, .reg NLogN.KFY.yP),
    (bI, .mul (.reg yA) (.reg yA)), (bKS, .reg NLogN.KFY.yU), (bM, .reg NLogN.KFY.yM),
    (bTot, .mul (.reg NLogN.KFY.yM) (.reg yT))]

/-- **One axis of the `B₀` pass.** -/
noncomputable def axB (X : Fin HS) (Kc : ℝ) : RM :=
  .seq (.rs (takeD yS (.reg NLogN.KFY.dS))) (brk [(NLogN.KFY.yB, bitE)] (brk [(yT, NLogN.KFY.tE), (NLogN.KFY.yM, mE)]
    (.seq (.rs invP) (.seq (.seq (brk [(zw, wE), (zm, kE)] (trBlk X Kc zw zm))
      (brk regsB bap0AxQ)) (.rs (.pop NLogN.KFY.yU))))))

/-- The cost bound of one `B₀` axis. -/
noncomputable def axBC (Kc : ℝ) (n t B K Ctr Cbap : ℕ) : ℕ :=
  (2 * 10 ^ 7 * (B + 2) ^ 2 + (2 * B + 2)) +
    brkC 1 K B (brkC 2 K B (invC t B + ((brkC 2 K B (1 + (10 ^ 12 * (IntMul.OpenAI.Glue.bP n + 1) ^ 3 + Ctr)) +
      brkC 18 K B Cbap) + (2 * B + 2))))

/-- Size bounds for one axis (`B₀` pass). -/
structure AxBndB (n A B : ℕ) : Prop where
  env : KFM.CostEnv A B
  pA : 2 ^ (2 * IntMul.OpenAI.Glue.pP n) ≤ A
  αA : IntMul.OpenAI.Glue.αP n ^ 2 * IntMul.OpenAI.Glue.rP n ^ 2 ≤ A
  jA : IntMul.OpenAI.Glue.pP n * IntMul.OpenAI.Glue.rP n + IntMul.OpenAI.Glue.αP n ^ 2 * IntMul.OpenAI.Glue.rP n ≤ A
  wA : IntMul.OpenAI.Glue.rP n * (2 * HvdH.eWidth (IntMul.OpenAI.Glue.pP n) (IntMul.OpenAI.Glue.αP n) + 1) + IntMul.OpenAI.Glue.rP n ≤ A
  LA : 2 ^ IntMul.OpenAI.Glue.LP n ≤ A
  dA : IntMul.OpenAI.Glue.dP n ≤ A
  pB : 2 * IntMul.OpenAI.Glue.pP n + 1 ≤ B
  LB : IntMul.OpenAI.Glue.LP n + 2 ≤ B

/-- The padded `B̃₀` line map. -/
def b0Q (p s t α : ℕ) (x : List KFB.GI) : List KFB.GI :=
  KFM.b0L p s t α x ++ List.replicate (t - s) (0, 0)

end NLogN.KFY

end

section
-- Source module: Solutions.KFY16PassAB
/-!
# KFY 16: the `A` pass and the `B₀` pass

* `progA X Kc = passP wA (axA X Kc)`, **`passA_all`**: the array becomes
  `tensorGo (mapsA p α s t) (ofFn t) zs` (the list pipeline of `KFY1Pad.passA`);
* `progB X Kc = passP wB (axB X Kc)`, **`passB_all`**: `tensorGo (mapsB p α s t) (ofFn t) ys`.

The arrays keep the invariant `length = ∏ t_i`, `DiscL p` (which gives `GB (p + 2)`), using an
`M : AxisMaps` (`disc_sapL`, `disc_b0Pad`).
-/

set_option linter.unusedSimpArgs false
set_option linter.unusedVariables false
set_option linter.unusedSectionVars false

namespace NLogN.KFY

open IntMul NLogN IntMul.HvdH.Ref NLogN.GL NLogN.AX NLogN.Rsp

noncomputable def progA (X : Fin HS) (Kc : ℝ) : RM := passP wA (axA X Kc)
noncomputable def progB (X : Fin HS) (Kc : ℝ) : RM := passP NLogN.KFY.wB (NLogN.KFY.axB X Kc)

/-- `t_k`. -/
noncomputable def tN (n k : ℕ) : ℕ := 2 ^ (IntMul.OpenAI.Glue.ℓP n - 1 + bN n k)
/-- `m_k = ∏_{i ≠ k} t_i`. -/
noncomputable def mN (n k : ℕ) : ℕ := 2 ^ (IntMul.OpenAI.Glue.LP n - (IntMul.OpenAI.Glue.ℓP n - 1 + bN n k))

/-- The `A` line maps. -/
noncomputable def fsA (n : ℕ) (k : ℕ) : List KFB.GI → List KFB.GI :=
  sapL (IntMul.OpenAI.Glue.pP n) ((List.ofFn (KFM.sB n)).getD k 1) ((List.ofFn (KFM.tB n)).getD k 1) (IntMul.OpenAI.Glue.αP n)

/-- The `B₀` line maps. -/
noncomputable def fsB (n : ℕ) (k : ℕ) : List KFB.GI → List KFB.GI :=
  b0Pad (IntMul.OpenAI.Glue.pP n) ((List.ofFn (KFM.sB n)).getD k 1) ((List.ofFn (KFM.tB n)).getD k 1) (IntMul.OpenAI.Glue.αP n)

/-- The array invariant. -/
def ArrInv (n : ℕ) (X : List KFB.GI) : Prop :=
  X.length = (List.ofFn (KFM.tB n)).prod ∧ DiscL (IntMul.OpenAI.Glue.pP n) X

end NLogN.KFY

end

section
-- Source module: Solutions.KFY17Par
/-!
# KFY 17: the numeric parameters of the passes

For `n` the bit sizes of the line programs (uniform over the axes):

* `lq = ⌊log₂(4p + 35)⌋ + 1`, `Jb = ⌊4dp/α²⌋ + 1` (a bound of every `jIter p s t α`, from the
  source window `4d s ≤ (4d − 1) t`), `mm = 3p + 2lq + 4 + (2lq + 1) Jb` (the multiplier width),
  `Bacc`, `Bo`, `Q`, `Bq`, the register environment `A = 2^E`, `B = 8E + 1`;
* **`sapHypN`**, **`tabHypN`**, **`axQAN`**, **`bapHypN`**, **`tabHypBN`**, **`axBndN`**,
  **`axBndBN`**, **`mulBndN`**: the hypotheses of `passA_all` / `passB_all`.

The facts of `n` used (`NF n`) come from KFQ `chainSizes` and KFMSetup `srcS_ok_n`.
-/

set_option linter.unusedSimpArgs false
set_option linter.unusedVariables false
set_option linter.unusedSectionVars false

namespace NLogN.KFY

open IntMul NLogN IntMul.OpenAI.Glue NLogN.GL NLogN.AX NLogN.Rsp

noncomputable section

def lqN (n : ℕ) : ℕ := Nat.log 2 (4 * pP n + 35) + 1
def JbN (n : ℕ) : ℕ := 4 * IntMul.OpenAI.Glue.dP n * pP n / αP n ^ 2 + 1
def mmN (n : ℕ) : ℕ := 3 * pP n + 2 * lqN n + 4 + (2 * lqN n + 1) * JbN n
def BaccN (n : ℕ) : ℕ := 2 * mmN n + 4 * lqN n + ℓP n + 2 * pP n + 8
def BoN (n : ℕ) : ℕ := pP n + 2 + mmN n + 3 * lqN n
def QN (n : ℕ) : ℕ := BoN n + IntMul.OpenAI.Glue.LP n + ℓP n + 4
def BqN (n : ℕ) : ℕ := 24 * (ℓP n + 2 * lqN n + 4) + 2 * pP n + lqN n + 2
def EN (n : ℕ) : ℕ := 2 * pP n + IntMul.OpenAI.Glue.LP n + 2 * ℓP n + 2 * lqN n + 8
def AN (n : ℕ) : ℕ := 2 ^ EN n
def BN (n : ℕ) : ℕ := 8 * EN n + 1

/-- The closed bound of `mulCostH` (KFMMul `mulCostH_le`). -/
def ΦF (m : ℕ) : ℕ :=
  (8 * fmA + 10 ^ 7 + 10 ^ 5) * (m + 1) * (zsz ((m + 1 : ℕ) : ℤ) + 2) *
    (zsz (((zsz ((m + 1 : ℕ) : ℤ) + 2) ^ 4 : ℕ) : ℤ) + 1)
def ΦN (n : ℕ) : ℕ := ΦF (2 * mmN n)

/-- The multiplication spec of the merged table. -/
def MS : MulSpec := ⟨fun a b => KFM.mulCostH (zsz a + zsz b), fun _ _ => True⟩

/-! ## The facts of `n` -/

/-- The facts of `n` used by the parameter proofs. -/
structure NF (n : ℕ) : Prop where
  hn : IntMul.OpenAI.Glue.n0 ≤ n
  α2 : 2 ≤ αP n
  αp : αP n ^ 2 ≤ pP n
  d2 : 2 ≤ IntMul.OpenAI.Glue.dP n
  win : ∀ k < IntMul.OpenAI.Glue.dP n, 2 ≤ KFM.srcN n k ∧ Odd (KFM.srcN n k) ∧
    4 * IntMul.OpenAI.Glue.dP n * KFM.srcN n k ≤ (4 * IntMul.OpenAI.Glue.dP n - 1) * tN n k

/-! ## Per-axis consequences -/

/-- The per-axis facts. -/
structure AxF (n k : ℕ) : Prop where
  s0 : 0 < KFM.srcN n k
  st : KFM.srcN n k < tN n k
  co : Nat.Coprime (KFM.srcN n k) (tN n k)
  sw : KFM.srcN n k ≤ 4 * IntMul.OpenAI.Glue.dP n * (tN n k - KFM.srcN n k)
  tr : tN n k ≤ 2 ^ ℓP n
  t2 : 2 ≤ tN n k
  J : HvdH.jIter (pP n) (KFM.srcN n k) (tN n k) (αP n) ≤ JbN n

/-! ## Width bounds -/

/-! ## The hypotheses -/

end

end NLogN.KFY

end

section
-- Source module: Solutions.KFY18CostL
/-!
# KFY 18: polynomial shapes of nlogn's resampling cost bounds

Every cost function of the per-axis programs (`roundCost`, `bandCost`, `resCost`, `gresCost`,
`stCost`, `jCost`, `tabCost`, …) is bounded by a product of a few linear factors:

* the slice costs are `O(NR · (n + len + NR) · (Φ + mm + B))` (`NR` rounds over `n` outputs),
* the table builders are `O(NI · NK · (Cw + log))`.

The per-axis bounds **`sapAx_le`** and **`bapAx_le`** split the cost of one axis into the main
term (`m` slices, linear in `m · t`) and a table term (polynomial in `t` and the bit sizes).

`polyle` proves `P ≤ Q` for polynomials over `ℕ` whose difference has nonnegative coefficients.
-/

set_option linter.unusedSimpArgs false
set_option linter.unusedVariables false
set_option linter.unusedSectionVars false
set_option maxRecDepth 100000

namespace NLogN.KFY

open IntMul NLogN NLogN.Rsp NLogN.AX NLogN.KFM

/-! ## One axis -/

/-- The constant sizes of the `B₀` table expressions. -/
def ksB : ℕ := eDelE.size + cDelE.size + ptKE.size + psKE.size + 1

end NLogN.KFY

end

section
-- Source module: Solutions.KFY19N
/-!
# KFY 19: the cost of the two passes for `n` (natural-number part)

For `n` past a threshold (`NG n`):

* every per-axis bound splits as `Ctr + main + overN n`, with `main` linear in `2^L` and
  **`overN n = 10^80 (p+1)^12 (2^ℓ + 2)^3`** collecting the tables and constant overheads
  (`axA_le`, `axB_le`);
* **`over_le`**: `(3d + 1) · overN n + 2 + d (d+1) (ℓ + p + d + 2)² + d ≤ 2^L`, since
  `L ≥ 4ℓ − 4` (`d ≥ 4`) and `2^ℓ` beats every fixed power of `ℓ` (`bigpow`);
* **`ZN_le`**: the per-element factor `Z = Φ + mm + Bacc + Q + 2 ≤ cZ (p + 1)(lq + 1)^6`;
* **`JE_le`**: `(Jb + 1)(2E + 2) ≤ 24 ⌈√p⌉ α` (the `B₀` rounds), from `α⁴ ≥ 12 d² b`.
-/

set_option linter.unusedSimpArgs false
set_option linter.unusedVariables false
set_option linter.unusedSectionVars false
set_option maxRecDepth 100000
set_option exponentiation.threshold 2048

namespace NLogN.KFY

open IntMul NLogN IntMul.OpenAI.Glue NLogN.GL NLogN.AX NLogN.Rsp

noncomputable section

/-! ## Growth -/

/-! ## The facts of `n` -/

/-- The facts of `n` used by the cost bounds. -/
structure NG (n : ℕ) : Prop where
  nf : NF n
  α4 : 12 * IntMul.OpenAI.Glue.dP n ^ 2 * bP n ≤ αP n ^ 4
  bl : bP n ≤ 2 * IntMul.OpenAI.Glue.dP n * ℓP n
  d5 : 2 ^ 18 * IntMul.OpenAI.Glue.dP n ^ 5 ≤ bP n
  dbig : 2 ^ 12 ≤ IntMul.OpenAI.Glue.dP n
  Lb : IntMul.OpenAI.Glue.LP n < bP n
  win2 : ∀ k < IntMul.OpenAI.Glue.dP n, (4 * IntMul.OpenAI.Glue.dP n - 2) * tN n k < 4 * IntMul.OpenAI.Glue.dP n * KFM.srcN n k

section facts

variable {n : ℕ} (h : NG n)
include h

/-- `⌈√p⌉`. -/
abbrev SN (n : ℕ) : ℕ := HvdH.ceilSqrt (pP n)

end facts

/-! ## Per-axis facts -/

section axis

variable {n : ℕ} (h : NG n) {k : ℕ} (hk : k < dP n)
include h hk

end axis

/-! ## The size parameters -/

section sizes

variable {n : ℕ} (h : NG n)
include h

/-- The constant of the per-element factor. -/
def cZ : ℕ := (8 * fmA + 10 ^ 7 + 10 ^ 5 + 1) * 10 ^ 6

/-- The per-element factor. -/
def ZN (n : ℕ) : ℕ := ΦN n + mmN n + BaccN n + QN n + 2

end sizes

/-! ## The per-axis bounds -/

/-- The overhead collector. -/
def overN (n : ℕ) : ℕ := 10 ^ 80 * (pP n + 1) ^ 12 * (2 ^ ℓP n + 2) ^ 3

/-- The main term of an `A` axis. -/
def mainA (n : ℕ) : ℕ := 4 * 10 ^ 11 * (2 * HvdH.sWidth (pP n) (αP n) + 2) * 2 ^ IntMul.OpenAI.Glue.LP n * ZN n

/-- The main term of a `B₀` axis. -/
def mainB (n : ℕ) : ℕ :=
  9 * 10 ^ 14 * (JbN n + 1) * (2 * HvdH.eWidth (pP n) (αP n) + 2) * 2 ^ IntMul.OpenAI.Glue.LP n * ZN n

section axb

variable {n : ℕ} (h : NG n)
include h

end axb

section axis2

variable {n : ℕ} (h : NG n) {k : ℕ} (hk : k < dP n)
include h hk

end axis2

/-! ## The pass and wrapper overheads -/

section pass

variable {n : ℕ} (h : NG n)
include h

end pass

/-! ## The overheads fit in `2^L` -/

end

end NLogN.KFY

end

section
-- Source module: Solutions.KFY21Main
/-!
# KFY 21: `TransNeed → LineNeed`

**`lineNeed_of_transNeed`**. The table merges `TransNeed`'s extensions (relabelled by
`gI E₁ _ 0`) with three own procedures:

* `own 0 = trW X Y dT jT`: the wrapper around `TransNeed`'s entry (descriptor moved from `X` to
  `TransNeed`'s descriptor stack `dT`), called by the `prc` node of every axis;
* `own 1 = (progA X K).code`, `own 2 = (progB X K).code`: the two passes (`eA`, `eB`).

The descriptor stack is `dS = 26`; `C₀ = 17 K + 1 + 216·10^14 · cZ · 4 · 58^6`, `A = 6`.
-/

set_option linter.unusedSimpArgs false
set_option linter.unusedVariables false
set_option linter.unusedSectionVars false
set_option maxRecDepth 100000

namespace NLogN.KFY

open IntMul NLogN IntMul.OpenAI.Glue NLogN.GL NLogN.AX NLogN.Rsp IntMul.HvdH.Ref KIL KFB

noncomputable section

/-! ## Procedure-free programs -/

section basic

variable {P : ℕ} {I : List Sym} {iprog : Fin P → IStmt HS P} (pc : ℕ → Fin P)

end basic

/-- No `prc` node. -/
def noPrc : RM → Bool
  | .rs _ | .mul _ _ _ | .op _ => true
  | .seq a b | .ite _ a b => noPrc a && noPrc b
  | .rep _ _ b => noPrc b
  | .prc .. => false

section prc

variable {P : ℕ} {I : List Sym} {iprog : Fin P → IStmt HS P} {jT : Fin P}
  {X Y dT : Fin HS} {Kc : ℝ} (hT : TrHyp I iprog jT dT Kc) (hdT : dT ≠ 0) (hXY : X ≠ Y)
  (hXD : X ≠ dT) (hYD : Y ≠ dT) (hX0 : X ≠ 0) (pc : ℕ → Fin P)
  (hW : iprog (pc 1) = trW X Y dT jT)
include hT hdT hXY hXD hYD hX0 hW

end prc

/-! ## The descriptor and the arrays -/

/-! ## The axis maps exist -/

/-! ## The writes of the axes -/

/-! ## The assembly -/

/-- The constant of `LineNeed`. -/
def C0 (K : ℝ) : ℝ := 17 * K + 1 + ((216 * 10 ^ 14 : ℕ) : ℝ) * NLogN.KFY.cZ * 4 * 58 ^ 6

end

end NLogN.KFY

end

section
-- Source module: Solutions.KFV2Dbl
/-!
# KFV §2: the doubling scan (every block followed by its negation)

`dblI X m`: stack `0` holds `X` blocks of `m` integers; afterwards it holds
`dblPre m xs X` (block `b` replaced by `b ++ (-b)`).  Per block: `dup` to `zX`, `neg`, move to
the buffer `zO`; the negated copy goes `zX → zY → zO`; finally the whole buffer goes back.
Cost `≤ 10^4 (w+1)(2Xm + X + 1)` (`cDblI_le`).  Stacks: `0`, buffers `zX = 53`, `zY = 54`,
`zO = 63`, counters `57–62` (all restored).
-/

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

namespace NLogN.KFV

open IntMul KFC KFT KFU AModel AModelL12

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

abbrev zX : Fin HS := 53
abbrev zY : Fin HS := 54

/-- Negate every entry. -/
def negL (l : List ℤ) : List ℤ := l.map (fun a => -a)

/-- One step: `dup 0 zX; neg zX; mv 0 zO`. -/
def dnStep : IStmt HS Pp := .seq (.op (.dup 0 zX)) (.seq (.op (.neg zX)) (.op (.mv 0 zO)))

/-- The doubled prefix after `k` blocks. -/
def dblPre (m : ℕ) (xs : List ℤ) (k : ℕ) : List ℤ :=
  (List.range k).flatMap fun x => (xs.drop (x * m)).take m ++ NLogN.KFV.negL ((xs.drop (x * m)).take m)

/-- One block. -/
def dblBlk (m : ℕ) : IStmt HS Pp :=
  .seq (rpt NLogN.KFU.zA zC NLogN.KFU.zD m dnStep) (.seq (rpt NLogN.KFU.zA zC NLogN.KFU.zD m (.op (.mv zX zY)))
    (rpt NLogN.KFU.zA zC NLogN.KFU.zD m (.op (.mv zY zO))))

def cDBlk (m w : ℕ) : ℕ :=
  (m * (500 * (w + 1) + 64) + 60 + cRpt m) + 2 * (m * (76 * w + 101 + 64) + 60 + cRpt m)

/-- **The doubling scan.** -/
def dblI (X m : ℕ) : IStmt HS Pp :=
  .seq (.op (.mark zO)) (.seq (rpt NLogN.KFU.oA NLogN.KFU.oC NLogN.KFU.oD X (dblBlk m)) (.seq (mvW zO 0) (.op (.unmark zO))))

def cDblI (X m w : ℕ) : ℕ :=
  1 + ((X * (cDBlk m w + 64) + 60 + cRpt X) + (X * (2 * m) * kap w + 30 + 1))

end NLogN.KFV

end

section
-- Source module: Solutions.GlueLayout
/-!
# Glue 2a: the paper's layout as a `CoordLayout`, its selected bits, and the bridge to `layerSel`

* §0 bit arithmetic: `getBit_bitsum`, `bitsum_update` (binary digit sums), `getD_decode`
  (a digit of `decode`), `getD_decode_binSh` (bits of a binary address).
* §1 `PosData`/`posSel`: selected bits of a binary address given by bit positions;
  `selOp_transport`: two selected-bit structures related by an address map preserving the
  linear index give the same `selOp` for every tensor kernel.
* §2 `slotLayout`: a slot arrangement `arr` (a list of slot names `nm i h` covering every slot
  `h < a i` exactly once) is a `CoordLayout (binSh |arr|) a`; `slotSel` (round `h`, axes `P`)
  is `SelCompat` with it (`slot_compat`).
* §3 `chunk_bridge`: if the selected positions are `A + iK + (K-1-ρ)` in a binary address of
  length `A + DK + B`, the selected-bit operator *is* `layerSel (2^A) K D (2^B) ρ` of §5.
* §4 the paper's layout `tgtL` (`SynthLayoutPaper`): `paperLayout`, `paperSel`,
  **`synth_forward_paper`** (§6.1 on the concrete layout: the rounds implement `B F^-`), and
  **`paper_round_layer`** (for every slot `b ≤ h < b + qK` the round's simultaneous layer is
  exactly Proposition 14's `layerH (2^A) K D (2^B) ρ` with `ρ = (h-b) mod K`, via `tgt_slot`).
-/

set_option linter.unusedSectionVars false
set_option linter.unusedSimpArgs false

namespace GlueLayout

open Finset Butterfly AModel Synth

/-! ## 0. Bit arithmetic -/

/-! ## 1. Selected bits given by positions; transport -/

variable {n : ℕ} {ι : Type*} [Fintype ι] [DecidableEq ι]

/-! ## 2. Slot arrangements are coordinate layouts -/

section Slots

variable (D : ℕ) [NeZero D]

variable (a : Fin D → ℕ) (arr : List ℕ)

variable {D a arr}

variable (D a arr)

variable {D a arr}

variable {R : Type*} [CommRing R] [Algebra ℂ R] (y : Rˣ)

end Slots

/-! ## 3. Chunk positions are Proposition 14's `layerSel` -/

section Chunk

variable {A K D B ρ : ℕ}

end Chunk

/-! ## 4. The paper's layout `tgtL` -/

section Paper

variable (D ℓ K : ℕ) [NeZero D] (long : Fin D → Bool)

variable {R : Type*} [CommRing R] [Algebra ℂ R] (y : Rˣ)

/-! ### The chunk rounds are Proposition 14's layer -/

/-- `e = qK - 1 - (h - b)`. -/
def eIdx (h : ℕ) : ℕ := Synth.qq ℓ K * K - 1 - (h - bb ℓ K)
/-- The chunk group `j₀ = ⌊e/K⌋` of round `h`. -/
def jGrp (h : ℕ) : ℕ := eIdx ℓ K h / K
/-- The common offset `ρ = (h - b) mod K`. -/
def rhoOf (h : ℕ) : ℕ := (h - bb ℓ K) % K
/-- Prefix bits before chunk group `j₀`: `A = |Lst| + j₀ D K`. -/
def preA (h : ℕ) : ℕ := (lstL D ℓ K long).length + jGrp ℓ K h * D * K
/-- Suffix bits after chunk group `j₀`: `B = (q - 1 - j₀) D K`. -/
def sufB (h : ℕ) : ℕ := (Synth.qq ℓ K - 1 - jGrp ℓ K h) * D * K

lemma rhoOf_lt (hK : 0 < K) (h : ℕ) : rhoOf ℓ K h < K := Nat.mod_lt _ hK

end Paper

end GlueLayout

end