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

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

import Definitions.Def_IntMul_Kappa_K_Defs13

set_option autoImplicit false
set_option relaxedAutoImplicit false
set_option maxHeartbeats 1000000

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

Per block: load the `M` groups (`q` records each) into the mark-delimited accumulator, then for
each row `y < q` peel the first record of every group to the output (`unzipA`) and move the
remaining groups back (`backB`).  All counting is unary; cost `O(q·V)` entry moves.
-/

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

namespace NLogN

open IntMul AModel

namespace KFC

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

/-! ## Unzip column: `mvB A O CB DB; mvW A T; unmark A; mark T` -/

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

def unzipA : IStmt HS Pp := uloop sCM sDM ucol

def urow : IStmt HS Pp := .seq unzipA backB

/-! ## Rows of one block -/

/-- Records `y = k, …, q-1` of column `m`. -/
def sufF (rr : ℕ → ℕ → List ℤ) (q k m : ℕ) : List ℤ := flF (fun y => rr y m) k (q - k)

def rowsB : IStmt HS Pp := uloop sCQ sDQ urow

/-! ## Loading a block into the accumulator -/

def mvQ : IStmt HS Pp := uloop sCQ sDQ (mvB sX NLogN.KFC.sT sCB sDB)

def ldcol : IStmt HS Pp := .seq mvQ (.op (.mark NLogN.KFC.sT))

def loadT : IStmt HS Pp := uloop sCM sDM ldcol

/-! ## Clearing the `M` empty groups -/

def clrA : IStmt HS Pp := uloop sCM sDM (.op (.unmark sA))

end KFC

end NLogN

end

section
-- Source module: Solutions.KFCBwdMain
/-!
# KFC: the self-contained fine-cost backward field move `bwdF`

`bwdF_spec`: on `t` with `ints xs ++ R` on stack `0` (shape `[P, M, q, B]`) and `P, q, M, B` on
stacks `1 … 4`, `bwdF` ends in `t.upd 0 (ints (pull [P,M,q,B] [P,q,M,B] gS 0 xs) ++ R)` with cost
`≤ cBwdF P q M B w = (2q+18)·kap w·V + cUL P + cUL q + cUL M + cUL B + 1000`.
-/

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

namespace NLogN

open IntMul AModel

namespace KFC

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

/-! ## One block -/

def blkB : IStmt HS Pp := .seq loadT (.seq backB (.seq rowsB clrA))

def cBlkB (q M B w : ℕ) : ℕ := (2 * q + 10) * kap w * (q * M * B)

/-! ## All blocks -/

def allB : IStmt HS Pp := uloop sCP sDP blkB

/-! ## Core and program -/

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

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

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

end KFC

end NLogN

end

section
-- Source module: Solutions.KFDLevel
/-!
# KFD: the fine body of one uniform level

`bodyF C tpl j0` is KIL5's `bodyU` with every per-element piece replaced by its fine form:
`fcode (subU …)`, the row split / merge by KFC's `bwdF` / `fwdF`, the fine network `netCF`
with fine frame changes `edgeCF` (fine level steps `fcode (stepU …)`), the fine routing
`fcode (rebuildCU …)`.  `bodyF_run`: the same denotation `levelR` as `bodyU_run`, with the fine
cost `bodyCF`.
-/

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

namespace NLogN

open IntMul RS

namespace KFD

open KIL AModel AModelL12

section Gen

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

/-- Push `l`, run the item program `p`, pop. -/
def withPushI (l : List (Fin HS × Ex HS)) (p : IStmt HS Pp) : IStmt HS Pp :=
  .seq (pushes l).code (.seq p (popsR (l.map Prod.fst)).code)

end Gen

section Code

variable {Pp : ℕ}

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

/-- The fine code of a frame change. -/
def edgeCF (m : ℕ) (j0 : Fin Pp) (Ts : List (TStep m)) : IStmt HS Pp :=
  KIL.seqL (Ts.map (tCodeF m j0))

end Code

section Edge

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

end Edge

section Body

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

/-- **The fine uniform level body.** -/
def bodyF (tpl : IMat m → List (TStep m)) (j0 : Fin Pp) : IStmt HS Pp :=
  .seq (fcode (subU m rXb)) (.seq (withPushI (splitL W) KFC.bwdF)
  (.seq (netCF (fun X => edgeCF m j0 (tpl X)) C.Mout C.Min C.gates)
  (.seq (fcode (rebuildCU W (NLogN.KIL.permT C))) (.seq (withPushI (splitL W) KFC.fwdF)
    (fcode (flipU m rXb))))))

/-- Fine cost of a shift template on streams of volume `v` (bit data). -/
def cShTF (v Bb : ℕ) (T : TSh m) : ℕ :=
  cStepGF T.k T.d (stgT rXs T).length v 1 Bb (paT rXs T) (pcT T) (stgT rXs T)

/-- Fine cost of the body (`cC`: cost of one child call wrapper). -/
def bodyCF (tpl : IMat m → List (TStep m)) (N P L' : ℕ) (A B Cs : Shape) (cC Bb : ℕ) : ℕ :=
  cStepsUF rXb (vol (P :: W * L' :: hdShape A B Cs m N)) 1 Bb (subT m) +
  (cSplit W Bb + KFC.cBwdF 1 W (P * L') (vol (hdShape A B Cs m N)) 1) +
  (netCIF (fun X => edgeCost (cShTF (vol (P :: L' :: hdShape A B Cs m N)) Bb) cC (tpl X))
    (W * vol (P :: L' :: hdShape A B Cs m N)) Bb C.Mout C.Min C.gates +
  (cRebuildF W (W * vol (P :: L' :: hdShape A B Cs m N)) Bb +
  ((cSplit W Bb + KFC.cFwdF 1 W (P * L') (vol (hdShape A B Cs m N)) 1) +
  cStepsUF rXb (vol (P :: W * L' :: hdShape A B Cs m N)) 1 Bb (flipT m))))

end Body

end KFD

end NLogN

end

section
-- Source module: Solutions.KFDProg
/-!
# KFD: the fine uniform Prop. 11 procedure

* `swapEF`, `lev0F`: the level-`0` digit interchange by KFC's fine field moves (`fwdF`/`bwdF`),
  cost `≤ (4q+136)·kap w·V + 10^9 (Bb+2)²` (`cLev0F_le`).
* `levF`, `progF`: KIL5's `levU`/`progU` with the fine body `bodyF` and `lev0F`.
* `levelCostF_le`: one level costs `KlevF·(V+1) + KbF·(Bb+2)² + SWn·cC` — linear in the
  volume, plus a per-call `poly(Bb)` overhead.
* `progF_run`: the same denotation `prog C q k P L A B Cs` as `progU_run`, cost `ilCostF`.
-/

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

namespace NLogN

open IntMul RS

namespace KFD

open KIL AModel AModelL12

section Swap

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

/-- The fine level-`0` interchange with parameters `a = vol A0`, `q`, `b = vol B`, `c = vol Cs`. -/
def swapEF (a q b c : Ex HS) : IStmt HS Pp :=
  .seq (withPushI [(1, a), (2, q), (3, .mul b q), (4, c)] KFC.fwdF)
    (withPushI [(1, a), (2, q), (3, b), (4, .mul q c)] KFC.bwdF)

/-- Fine cost of `swapEF` (`w`: bit size of the entries). -/
def cSwapEF (a q b c : Ex HS) (qn V Bb w : ℕ) : ℕ :=
  pushCost Bb [(1, a), (2, q), (3, .mul b q), (4, c)] + pushCost Bb [(1, a), (2, q), (3, b), (4, .mul q c)] +
    ((2 * qn + 67) * KFC.kap w * V + (2 * qn + 69) * KFC.kap w * V) +
    2 * (4 * (2 * Bb + 2) + 1)

/-- The fine level-`0` code. -/
def lev0F : IStmt HS Pp :=
  swapEF (.mul (.reg NLogN.KIL.dP) (.mul (.reg dL) (.reg dA))) (.reg dT) (.reg dB) (.reg dC)

def cLev0F (q V Bb w : ℕ) : ℕ :=
  cSwapEF (.mul (.reg NLogN.KIL.dP) (.mul (.reg dL) (.reg dA))) (.reg dT) (.reg dB) (.reg dC) q V Bb w

end Swap

section Costs

variable {m : ℕ}

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

/-- The per-shift fine cost unit on bit data, `v` the stream volume. -/
def zU (v Bb : ℕ) : ℕ := 10 ^ 13 * ((v + 1) * (1 + 2)) + 10 ^ 10 * (Bb + 2) ^ 2

end Costs

section LevCost

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

/-- Fine template constants of all shifts of all frame changes of one level. -/
def KShF : ℕ := frameSum (fun X => shKF (tpl X)) C.gates C.Min C.Mout

/-- Number of table units (frame changes and gates) of one level. -/
def nTab : ℕ := nFr C + (C.gates.length + 1) * (W * (2 * W + 2) + 202)

/-- The fine level constants: linear part and per-call `(Bb+2)²` part. -/
def KlevF : ℕ :=
  3 * 10 ^ 13 * (tKFs rXb (subT m) + tKFs rXb (flipT m) + KShF C tpl) +
    2 * (2 * W + 69) * 276 + 10 ^ 9 * nTab C + 100

def KbF : ℕ :=
  10 ^ 10 * (tKFs rXb (subT m) + tKFs rXb (flipT m) + KShF C tpl) +
    2 * ((pushSz (splitL W) + 1) * 10 ^ 7 + 10) + 10 ^ 8 * nTab C +
    (pushSz (entL W m) + 5) * 10 ^ 7 + nTab C + 100

end LevCost

/-- **The fine cost recursion** (`K`: linear level constant, `Kb`: per-call `(Bb+2)²` constant,
`Kd`: per-child-call descriptor constant, `S`: number of child calls per level). -/
def ilCostF (W K Kb Kd S q Bb : ℕ) : ℕ → ℕ → ℕ
  | 0, V => K * (q + 1) * (V + 1) + Kb * (Bb + 2) ^ 2
  | j + 1, V => K * (V + 1) + Kb * (Bb + 2) ^ 2 +
      S * (ilCostF W K Kb Kd S q Bb j (V / W) + Kd * (Bb + 2) ^ 2)

section Main

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

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

/-- **The fine uniform procedure table** (one procedure). -/
def progF : Fin 1 → IStmt HS 1 :=
  fun _ => .ite dJ (fun c => c == .zero) lev0F (levF C tpl)

end Main

end KFD

end NLogN

end

section
-- Source module: Solutions.KFEPad
/-!
# KFE: fine padding loops

`padP` / `unpadP` (KIL6Pad) run through `KFD.fcode`: every counted loop is a unary loop, so
the per-element cost is `O(1)` for bit data and the per-block loop overhead is
`O(zsz (F'·Y))`, bounded by the block length.  Result: fine cost
`cPadF V Bb = 2000 (Bb + 2) + 20000 V` (`V` the padded volume) in place of
`cPad V Bb = 10^5 (V+1)(Bb+2)`.

* `FSp.zeroR`, `FSp.dropR`: fine `zeroKE` / `dropKE` with a register count;
* `padP_fsp`, `unpadP_fsp`: the fine specifications (same denotations as `padP_sp`,
  `unpadP_sp`), for bit data (`Bnd 1`);
* `padS_fsp`, `unpadS_fsp`: the stages with their register assignments.
-/

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

namespace NLogN.KFE

open IntMul RS AModel AModelL12 KIL KFD

/-! ## Fine `zeroKE` / `dropKE` with a register count -/

/-! ## The fine padding loops -/

/-- Fine cost of a padding / unpadding pass over a padded volume `V` (bit data). -/
def cPadF (V Bb : ℕ) : ℕ := 2000 * (Bb + 2) + 20000 * V

/-! ## The stages -/

end NLogN.KFE

end

section
-- Source module: Solutions.KFESwap
/-!
# KFE: fine binary embedding / extraction and fine high-digit swaps

* `preB_fsp`, `postB_fsp`: KIL6's two padding passes of the binary embedding / extraction with
  fine cost `2 (10^9 (B+2)^2 + cPadF V B)` (were `cPad`).
* `swapsF q`: the `r` high-digit interchanges as an item-language loop (`irep`) whose body is
  KFD's fine interchange `swapEF` (KFC `fwdF`/`bwdF`) plus the two register updates;
  `swapsF_runs`: same denotation as `swapsP_sp`, cost
  `(r + 1) (cSw1F q V B + 10^4 (B+2))` with `cSw1F q V B = (4q+136)·276·V + 10^9 (B+2)^2`
  (no `(B+2)^2` factor on the volume).
-/

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

namespace NLogN.KFE

open IntMul RS AModel AModelL12 KIL KFD

/-! ## Fine high-digit swaps -/

/-- The high-digit interchanges, fine version. -/
def swapsF {Pp : ℕ} (q : ℕ) : IStmt HS Pp :=
  irep 24 (.reg kR) (.seq (swapEF NLogN.KIL.swA (.lit (q : ℤ)) swB swC6)
    (RS.seq (.asg kQt (.mul (.reg kQt) (.lit (q : ℤ)))) (.asg kQr (.ndiv (.reg kQr) (.lit (q : ℤ))))).code)

/-- Fine cost bound of one digit interchange (bit data). -/
def cSw1F (q V Bb : ℕ) : ℕ := (4 * q + 136) * 276 * V + 10 ^ 9 * (Bb + 2) ^ 2

end NLogN.KFE

end

section
-- Source module: Solutions.KFEMain
/-!
# KFE: the fine Lemma 12 program

`l12F q m W` is KIL6's `l12I q m W` with every per-element pass made fine:

* the binary embedding / extraction (`preB`, `postB`) and the row padding of the middle part
  (`padR`, `unpadR`) are compiled by `KFD.fcode` (unary loop counters), cost `cPadF`;
* the high-digit interchanges are `swapsF` (KFC `fwdF`/`bwdF`), cost `cSw1F`;
* the descriptor (`preA`, `desc4`), the piece loop `piecesI` and the register file are unchanged
  (they cost `poly(Bb)` per step and touch no data).

`l12F_run`: for any oracle code meeting `CallSpec` (in particular the fine Prop. 11 code
`KFD.progF`), the program computes the variant denotation `l12v` (same as `l12_run`), at cost
`cL12F`.
-/

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

namespace NLogN.KFE

open IntMul RS AModel AModelL12 KIL KFD

/-- The middle part, fine version (row padding by `fcode`). -/
def midF (q m W : ℕ) : IStmt HS 1 :=
  .seq (desc4 W).code (.seq (fcode padR) (.seq (piecesI q m) (fcode unpadR)))

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

/-- **The fine Lemma 12 item-language program** (`q`, `m`, `W` are literals). -/
def l12F (q m W : ℕ) : IStmt HS 1 :=
  .seq (preA q m W).code (.seq (fcode preB) (.seq (NLogN.KFE.swapsF q)
    (.seq (.ite aN (fun c => c == Cls.zero) .skip (midF q m W)) (.seq (fcode postB) (popDs regD).code))))

/-- Fine cost of the Lemma 12 program. -/
def cL12F (u e r K n Vq V' m q : ℕ) (cO : ℕ → ℕ) (B : ℕ) : ℕ :=
  cPreA u e B + 4 * (10 ^ 9 * (B + 2) ^ 2 + cPadF Vq B) +
    (r + 1) * (cSw1F q Vq B + 10 ^ 4 * (B + 2)) + ((if n = 0 then 0 else cMidF K V' m n cO B) + 26) +
    25 * (2 * B + 2) + 1

end NLogN.KFE

end

section
-- Source module: Solutions.KIL6Inst
/-!
# KIL6: the Lemma 12 program with the Prop. 11 code as oracle

The oracle code is the uniform Prop. 11 program of KIL5 (`prog_il_uniform_rat`, i.e.
`progU (Qc.toCircuit 1) (edgeTplQ Qc)`).  Its denotation at the levels used is
`prog (Qc.toCircuit (q ^ m ^ Kp)) q`.  The variant denotation `l12v` is taken with the
length-normalised oracle `normO` (equal to the oracle on every call of the procedure, see
`normO_eq`; it only serves to state the uniform length invariant).

* `l12_il_uniform_rat`: the item-language program `l12I q m W` (with the oracle procedure table)
  computes `l12v (normO (prog …)) …`, with cost `cL12 … (ilCostU5 …)`, for every prime `q`, under
  the KIL5 level side conditions at the levels `j < Kp`;
* `l12_il_rat_correct`: if moreover the oracle is correct (`OracleSpec`, which holds for all large
  primes by `prop11_ratLabels`), the output is the interchange of the two `[2^u]` fields.
-/

set_option linter.unusedSimpArgs false
set_option linter.unusedVariables false

namespace NLogN.KIL

open IntMul RS AModel AModelL12 AModelRat AModelFac

section Norm

variable {α : Type*}

/-- The oracle made length-preserving (it is returned unchanged where it preserves length). -/
def normO (child : Oracle α) : Oracle α :=
  fun j P L A B Cs ys => if (child j P L A B Cs ys).length = ys.length then child j P L A B Cs ys
    else ys

end Norm

section Rat

variable {m W : ℕ} [NeZero W]

end Rat

end NLogN.KIL

end

section
-- Source module: Solutions.KIL6Cost
/-!
# KIL6: closed cost bounds

* `ilCostU5_le_coeff`: `ilCostU5 j V ≤ (Bb+2)² (V + W^j) G_j` with the coefficient sequence
  `G_{j+1} = (S/W) G_j + C`; hence (`power_width_recurrence`) `ilCostU5 j V ≤ K (Bb+2)² (V + W^j)
  (m^j)^τ` when `S / W < m^τ`.
* `cL12_le`: the cost of the Lemma 12 program is `≤ K · (vol (binShape … 2^u) + 1) · (Bb+2)² · u^τ`.
-/

set_option linter.unusedSimpArgs false
set_option linter.unusedVariables false

namespace NLogN.KIL

open IntMul RS AModel AModelL12 Finset

/-! ## The oracle cost -/

/-! ## The linear part -/

/-- The constant of the linear part. -/
def Nc (q m W : ℕ) : ℕ := 10 ^ 11 * (q + 1) ^ 3 * (m + 3) * (W + 1)

/-! ## The closed bound -/

end NLogN.KIL

end

section
-- Source module: Solutions.KFECost
/-!
# KFE: closed fine cost bound of the Lemma 12 program

* `cL12F_lin`: `cL12F ≤ Av·(Vb+1)(Kp+1) + Ao·(Bb+2)²(u+1) + Σ_j digit_j · cO j` — the data passes
  are linear in the bit volume `Vb` (no `(Bb+2)` factor); the descriptor work costs
  `O((u+1)(Bb+2)²)`.
* `calls_le`: the per-call overhead of the oracle, `Σ_{j<L} digit_j (S+1)^j ≤ m·n^(c+1)` with
  `c = clog_m (S+1)`: the number of Prop. 11 calls is polynomial in `n ≤ u`.
* `cL12F_le`: with the fine oracle cost `ilCostF` and `S/W < m^τ`,
  `cL12F ≤ K·((Vb + 1)·u^τ + (u+1)^(c+1)·(Bb+2)²)`, `c = clog_m (S+1)`.
-/

set_option linter.unusedSimpArgs false
set_option linter.unusedVariables false

namespace NLogN.KFE

open IntMul RS AModel AModelL12 Finset KIL KFD

/-- Volume constant of the fine Lemma 12 cost. -/
def Av (q W : ℕ) : ℕ := 10 ^ 6 * (q + 1) ^ 3 * (W + 1)

/-- Descriptor-overhead constant of the fine Lemma 12 cost. -/
def Ao (m W : ℕ) : ℕ := 10 ^ 11 * (m + 3) * (W + 1)

/-! ## The number of oracle calls -/

/-! ## The closed bound -/

end NLogN.KFE

end

section
-- Source module: Solutions.KIL6Tok
/-!
# KIL6: the template side conditions for large primes

* `tok_edgeTpl`: the `TOK` side conditions of the templates `edgeTplQ Qc X` hold at every modulus
  `q ^ e` (`e ≥ 1`) for every prime `q` beyond a constant of the frame matrix, once `Bb` holds
  `tokBnd A · (q^e + 1)²` (this is the open `tok_edgeTplQ` of KIL5_WIP, in a slightly different
  form: the constant is explicit).
* `framesOK_toCircuit`: a property of the rational frame changes transfers to the integer frame
  changes of every `Qc.toCircuit Q`.
* `subBnd_of`: `SubBnd N Bb` from `zsz (8 (N+1)²) ≤ Bb`.
* `SWn_toCircuit_eq`: the number of child calls per level of the uniform Prop. 11 code is the
  rank sum of the rational circuit.
-/

set_option linter.unusedSimpArgs false
set_option linter.unusedVariables false

namespace NLogN.KIL

open IntMul RS AModel AModelL12 AModelFac AModelRat Finset

variable {m : ℕ}

/-! ## TOK -/

/-- The size of the constants of a rational step. -/
def qsB : QStep m → ℕ
  | .sh _ u c => u.num.natAbs + u.den + ∑ j : Fin (m + m), 2 * ((c j).num.natAbs + 1) * ((c j).den + 1)
  | .sw _ _ => 0

/-- The template constant of a frame change. -/
noncomputable def tokBnd (A : QMat m) : ℕ :=
  ((shearQ (ratFacQ A).1 (piQ A) (ratFacQ A).2.2).map qsB).sum

/-! ## Frames of the integer circuits -/

/-! ## SubBnd -/

/-! ## The number of child calls -/

namespace QStep

def isSw : QStep m → Bool
  | .sh _ _ _ => false
  | .sw _ _ => true

end QStep

end NLogN.KIL

end

section
-- Source module: Solutions.KFBGauss
/-!
# KFB: Gaussian-integer data on item stacks, linear-combination expressions

* `GI = ℤ × ℤ` (Gaussian integer `a + b i`), stored on a stack as the two items `a, b`
  (`gflat`).  `toC` is the embedding into `ℂ`; `gmul` the product.
* `lcE l`: the register expression `Σ c_t · top(r_t)` for coefficients `c_t ∈ {-1, 0, 1}`, built
  only from `add`, `neg`, `reg`, `lit` (a linear expression: cost `O(size · (B + 2))`).
  `lcE_eval`, `lcE_bnd` (all partial sums of `n ≤ 2^e` terms of size `≤ w` have size `≤ w + e`).
* `FSp.mvN`: the fine counted move with the count read from a register and its loop overhead
  bounded by the count itself (`≤ N (250 w + 1600) + 2000`, no `Bb`).
-/

set_option linter.unusedSimpArgs false
set_option linter.unusedVariables false

namespace NLogN.KFB

open IntMul RS AModel KIL KFD Complex

/-! ## Gaussian integers -/

abbrev GI := ℤ × ℤ

/-- Items of a Gaussian array: real and imaginary part of each entry, in order. -/
def gflat (zs : List GI) : List ℤ := zs.flatMap fun z => [z.1, z.2]

@[simp] theorem gflat_nil : gflat [] = [] := rfl
@[simp] theorem gflat_cons (z : GI) (zs : List GI) : gflat (z :: zs) = z.1 :: z.2 :: gflat zs := rfl

/-- Bit-size bound of a Gaussian array. -/
def GB (w : ℕ) (zs : List GI) : Prop := ∀ z ∈ zs, zsz z.1 ≤ w ∧ zsz z.2 ≤ w

/-- Gaussian product `(a + b i)(x + y i)`. -/
def gmul (c z : GI) : GI := (c.1 * z.1 - c.2 * z.2, c.2 * z.1 + c.1 * z.2)

/-- The embedding `ℤ[i] → ℂ`. -/
noncomputable def toC (z : GI) : ℂ := (z.1 : ℂ) + (z.2 : ℂ) * I

/-- Coefficients with both components in `{-1, 0, 1}`. -/
def Sm (c : GI) : Prop := c.1.natAbs ≤ 1 ∧ c.2.natAbs ≤ 1

/-! ## Linear-combination expressions -/

/-- `c · top(r)` for `c ∈ {-1, 0, 1}`. -/
def tmE (c : ℤ) (r : Fin HS) : Ex HS :=
  if c = 1 then .reg r else if c = -1 then .neg (.reg r) else .lit 0

/-- `Σ c_t · top(r_t)`. -/
def lcE : List (ℤ × Fin HS) → Ex HS
  | [] => .lit 0
  | (c, r) :: l => .add (tmE c r) (lcE l)

/-! ## Fine moves with the loop overhead bounded by the count -/

end NLogN.KFB

end

section
-- Source module: Solutions.KFBQop
/-!
# KFB: the two-stream pass `qop` (fine cost)

On a Gaussian array viewed as consecutive blocks `[u ‖ v]` of two streams of `B1` entries each,
`qop Ms` applies the small `2 × 2` Gaussian matrix `M_p = Ms[p mod L]` (`L = |Ms|`) to block `p`:
`(u, v) ↦ (M₀₀ u + M₀₁ v, M₁₀ u + M₁₁ v)` entrywise.  With a single matrix this is a
selected-bit kernel on a field `[2]` (the `C` kernel (25) scaled by 2, a phase on the selected
half, a sign); with the period `[1, X]` it is a CNOT whose control is the field just before.

Per block: one move of the block to a scratch stack (reversal), the two streams distributed to
four register stacks (real/imaginary parts), one zip loop computing the outputs with linear
register expressions (`lcE`), the second output stream moved back in order.  Every entry is
touched `O(1)` times; loop counters are unary (`fcode`) with overhead linear in the counts.

Stacks: data `0`; scratch `qT = 33`, output accumulator `qA = 34`, `qO = 35`; stream registers
`36 … 39`; counters `30, 31`; counts `48 … 52` (`G`, `B1`, `2B1`, `4B1`, total items).
-/

set_option linter.unusedSimpArgs false
set_option linter.unusedVariables false

namespace NLogN.KFB

open IntMul RS AModel KIL KFD

abbrev qT : Fin HS := 33
abbrev qA : Fin HS := 34
abbrev qO : Fin HS := 35
abbrev rUr : Fin HS := 36
abbrev rUi : Fin HS := 37
abbrev rVr : Fin HS := 38
abbrev rVi : Fin HS := 39
abbrev cG : Fin HS := 30
abbrev c1 : Fin HS := 31
abbrev rG : Fin HS := 48
abbrev rB : Fin HS := 49
abbrev rB2 : Fin HS := 50
abbrev rB4 : Fin HS := 51
abbrev rN : Fin HS := 52

/-- Per-element cost of a distribution step. -/
def cD (w : ℕ) : ℕ := 10000 * (w + 2)

/-- Distribute a reversed Gaussian stream from `qT` to the register stacks `Xr`, `Xi`. -/
def distB (Xr Xi : Fin HS) : RS :=
  .rep NLogN.KFB.c1 (.reg NLogN.KFB.rB) (.seq (.push Xi (.reg qT)) (.seq (.pop qT) (.seq (.push Xr (.reg qT)) (.pop qT))))

/-! ## Pushing a linear combination -/

/-! ## The zip loop -/

/-- A `2 × 2` Gaussian matrix `[[a, b], [c, d]]`. -/
structure M2G where
  a : GI
  b : GI
  c : GI
  d : GI

def M2G.Ok (M : M2G) : Prop := KFB.Sm M.a ∧ KFB.Sm M.b ∧ KFB.Sm M.c ∧ KFB.Sm M.d

/-- `p x + q y`. -/
def lin2 (p q x y : GI) : GI := gmul p x + gmul q y

def reT (p q : GI) : List (ℤ × Fin HS) := [(p.1, rUr), (-p.2, rUi), (q.1, rVr), (-q.2, rVi)]
def imT (p q : GI) : List (ℤ × Fin HS) := [(p.2, rUr), (p.1, rUi), (q.2, rVr), (q.1, rVi)]

/-- One output entry of each stream: `qA ← M₀₀ u + M₀₁ v`, `qO ← M₁₀ u + M₁₁ v`. -/
def zb (M : M2G) : RS :=
  .seq (.push qA (lcE (reT M.a M.b))) (.seq (.push qA (lcE (imT M.a M.b)))
    (.seq (.push qO (lcE (reT M.c M.d))) (.seq (.push qO (lcE (imT M.c M.d)))
      (.seq (.pop rUr) (.seq (.pop rUi) (.seq (.pop rVr) (.pop rVi)))))))

def zipR (M : M2G) : RS := .rep NLogN.KFB.c1 (.reg NLogN.KFB.rB) (zb M)

def cZ (w : ℕ) : ℕ := 113000 * (w + 4)

/-! ## One block -/

/-- Output of one block `[u ‖ v]`. -/
def pairOut (M : M2G) (u v : List GI) : List GI :=
  List.zipWith (lin2 M.a M.b) u v ++ List.zipWith (lin2 M.c M.d) u v

/-- The per-block program. -/
def blk (M : M2G) : RS :=
  .seq (NLogN.mvK 0 qT NLogN.KFB.c1 (.reg rB4)) (.seq (distB rVr rVi) (.seq (distB rUr rUi)
    (.seq (zipR M) (.seq (NLogN.mvK qO qT NLogN.KFB.c1 (.reg rB2)) (NLogN.mvK qT qA NLogN.KFB.c1 (.reg rB2))))))

def cBlk (w B1 : ℕ) : ℕ := B1 * (135000 * w + 530000) + 12000

/-- Scratch stacks of `qop`. -/
structure Scr (f : Fin HS → List ℤ) : Prop where
  t : f qT = []
  o : f qO = []
  ur : f rUr = []
  ui : f rUi = []
  vr : f rVr = []
  vi : f rVi = []
  c : f NLogN.KFB.c1 = []

/-- Counts of `qop` (block level). -/
structure Cnt (s : IState HS) (f : Fin HS → List ℤ) (B1 : ℕ) : Prop where
  b : ∃ L, (encK s f).G NLogN.KFB.rB = .int (B1 : ℤ) :: L
  b2 : ∃ L, (encK s f).G rB2 = .int ((2 * B1 : ℕ) : ℤ) :: L
  b4 : ∃ L, (encK s f).G rB4 = .int ((4 * B1 : ℕ) : ℤ) :: L

/-! ## A group of blocks, and the whole pass -/

/-- Outputs of consecutive blocks with the matrices `Ms`. -/
def grpOut : List M2G → ℕ → List GI → List GI
  | [], _, _ => []
  | M :: Ms, B1, zs => pairOut M (zs.take B1) ((zs.drop B1).take B1) ++ grpOut Ms B1 (zs.drop (2 * B1))

def blks : List M2G → RS
  | [] => .skip
  | M :: Ms => .seq (NLogN.KFB.blk M) (blks Ms)

/-- **The two-stream pass**: `G` groups of `|Ms|` blocks, block `p` transformed by
`Ms[p mod |Ms|]`, then the outputs moved back to stack `0` in order. -/
def qop (Ms : List M2G) : RS :=
  .seq (.rep cG (.reg rG) (blks Ms)) (NLogN.mvK qA 0 NLogN.KFB.c1 (.reg rN))

/-- Denotation of `qop`. -/
def qopL (Ms : List M2G) (B1 G : ℕ) (zs : List GI) : List GI :=
  (List.range G).flatMap fun g => grpOut Ms B1 (zs.drop (g * (2 * B1 * Ms.length)))

/-- Cost of `qop`. -/
def cQop (w B1 L G : ℕ) : ℕ :=
  G * (L * NLogN.KFB.cBlk w B1 + 2101) + 2000 + (2 * (G * (2 * B1 * L))) * (250 * (w + 2) + 2600) + 2000

/-! ## Fine compilability -/

/-- Stacks touched by the pass. -/
def qT_all : List (Fin HS) := [NLogN.KFB.c1, NLogN.KFB.rB, rB2, rB4, qT, qO, qA, 0, rUr, rUi, rVr, rVi]

end NLogN.KFB

end

section
-- Source module: Solutions.KFBSeg
/-!
# KFB: a self-contained segment wrapper for `qop`, and its index semantics

`qopW Ms B1 G off` pushes the counts, parks the first `off` Gaussian entries of stack `0`,
runs `qop Ms` on the following `G · 2B1|Ms|` entries, and moves the parked prefix back:
on a clean managed state it replaces `pre ++ seg ++ rest` by `pre ++ qopL Ms B1 G seg ++ rest`.

Index semantics (`getD_qopL_diag`): with one diagonal matrix `diag(α, β)`, entry `j` of the
segment is multiplied by `α` or `β` according to bit `j / B1 mod 2`.
-/

set_option linter.unusedSimpArgs false
set_option linter.unusedVariables false

namespace NLogN.KFB

open IntMul RS AModel KIL KFD

abbrev qP : Fin HS := 40
abbrev rP : Fin HS := 53

/-- The wrapper: counts as literals, park, pass, unpark, pop the counts. -/
def qopW (Ms : List M2G) (B1 G off : ℕ) : RS :=
  .seq (.push NLogN.KFB.rB (.lit (B1 : ℤ))) (.seq (.push rB2 (.lit ((2 * B1 : ℕ) : ℤ)))
  (.seq (.push rB4 (.lit ((4 * B1 : ℕ) : ℤ))) (.seq (.push rG (.lit (G : ℤ)))
  (.seq (.push rN (.lit ((2 * (G * (2 * B1 * Ms.length)) : ℕ) : ℤ)))
  (.seq (.push NLogN.KFB.rP (.lit ((2 * off : ℕ) : ℤ)))
  (.seq (NLogN.mvK 0 qP NLogN.KFB.c1 (.reg NLogN.KFB.rP)) (.seq (qop Ms) (.seq (NLogN.mvK qP 0 NLogN.KFB.c1 (.reg NLogN.KFB.rP))
  (.seq (.pop NLogN.KFB.rB) (.seq (.pop rB2) (.seq (.pop rB4) (.seq (.pop rG) (.seq (.pop rN)
    (.pop NLogN.KFB.rP))))))))))))))

/-- Clean scratch for the wrapper. -/
structure Cl (f : Fin HS → List ℤ) : Prop where
  scr : Scr f
  a : f qA = []
  g : f cG = []
  b : f NLogN.KFB.rB = []
  b2 : f rB2 = []
  b4 : f rB4 = []
  rg : f rG = []
  rn : f rN = []
  rp : f NLogN.KFB.rP = []
  p : f qP = []

/-- Cost of the wrapper. -/
def cQopW (w B1 L G off Z : ℕ) : ℕ :=
  cQop w B1 L G + 2 * (2 * off * (250 * w + 2600) + 2000) + 6 * (4000 * (Z + 2) + 2 * Z + 2)

/-! ## Index semantics -/

/-- A diagonal `2 × 2` Gaussian matrix. -/
def dg (α β : GI) : M2G := ⟨α, 0, 0, β⟩

/-! ## Index semantics, general periodic matrices -/

/-- Default matrix. -/
def M0 : M2G := ⟨0, 0, 0, 0⟩

/-- Output entry `j` of a pass with the matrix `M` on blocks of `2·B1`. -/
def pent (M : M2G) (B1 : ℕ) (zs : List GI) (j : ℕ) : GI :=
  if j % (2 * B1) < B1 then lin2 M.a M.b (zs.getD j 0) (zs.getD (j + B1) 0)
  else lin2 M.c M.d (zs.getD (j - B1) 0) (zs.getD j 0)

end NLogN.KFB

end

section
-- Source module: Solutions.KFBPhase
/-!
# KFB: the `phase` step as a fine-cost item-language program

The state `x : Fin W × Addr m (Fin n) → ℂ` of the restricted interface is encoded role-major
as a Gaussian array: role `w`, address `b` at index `w · 2^{mn} + aidx b`, where
`aidx b = Σ_{(r,c)} b(r,c) · 2^{c + n r}` (`finProdFinEquiv`).

`phaseP` = one scalar pass `i^{k₀ + k₁ n}` and one diagonal pass per address bit `(r, c)`
(multiplying by `i^{a_r}` where the bit is set), each a `qopW` on the segment of role `w`.

* `segApply` / `getD_segApply`: index semantics of one segment pass.
* `rep_phase`: if `zs` encodes `x`, the result encodes `(RStep.den T (phase w k₀ k₁ a)) *ᵥ x`.
* `phaseP_fsp`: the program, its result and its fine cost
  `(mn + 1) · O((W + 1) 2^{mn} (w + 1))`.
-/

set_option linter.unusedSimpArgs false
set_option linter.unusedVariables false

namespace NLogN.KFB

open IntMul RS AModel KIL KFD Complex

/-! ## Powers of `i` -/

/-- `i^k` as a Gaussian integer. -/
def ipw (k : ℕ) : GI :=
  if k % 4 = 0 then (1, 0) else if k % 4 = 1 then (0, 1) else if k % 4 = 2 then (-1, 0) else (0, -1)

/-! ## Binary digits -/

/-- Address index: bit `(r, c)` at position `c + n r`. -/
def aidx {m n : ℕ} (b : Fin m × Fin n → Fin 2) : ℕ :=
  ∑ p : Fin m × Fin n, (b p).val * 2 ^ (finProdFinEquiv p).val

/-- The digit function of an address. -/
def adig {m n : ℕ} (b : Fin m × Fin n → Fin 2) (i : ℕ) : ℕ :=
  if h : i < m * n then (b (finProdFinEquiv.symm ⟨i, h⟩)).val else 0

/-! ## One segment pass on the list level -/

/-- The list produced by `qopW Ms B1 G off`. -/
def segApply (off : ℕ) (Ms : List M2G) (B1 G : ℕ) (zs : List GI) : List GI :=
  zs.take off ++ (qopL Ms B1 G (zs.drop off) ++ (zs.drop off).drop (G * (2 * B1 * Ms.length)))

/-! ## The phase passes on the list level -/

/-- After the scalar pass and `t` bit passes. -/
def phL (N off u : ℕ) (aE : ℕ → ℕ) (zs : List GI) : ℕ → List GI
  | 0 => segApply off [dg (ipw u) (ipw u)] 1 (N / 2) zs
  | t + 1 => segApply off [dg (ipw 0) (ipw (aE t))] (2 ^ t) (N / 2 ^ (t + 1)) (phL N off u aE zs t)

/-! ## The array encoding and the `phase` step -/

/-- `zs` encodes the state `x` (role-major, address index `aidx`). -/
def Rep {W m n : ℕ} (zs : List GI) (x : Fin W × (Fin m × Fin n → Fin 2) → ℂ) : Prop :=
  zs.length = W * 2 ^ (m * n) ∧ ∀ w b, NLogN.KFB.toC (zs.getD (w.val * 2 ^ (m * n) + aidx b) 0) = x (w, b)

/-- Bit exponents of the phase vector `a`. -/
def phE {m n : ℕ} (a : Fin m → ℕ) (i : ℕ) : ℕ :=
  if h : i < m * n then a (finProdFinEquiv.symm ⟨i, h⟩).1 else 0

/-! ## The phase program -/

/-- The scalar pass and `t` bit passes, on the segment at `off`. -/
def phP (N off u : ℕ) (aE : ℕ → ℕ) : ℕ → RS
  | 0 => qopW [dg (ipw u) (ipw u)] 1 (N / 2) off
  | t + 1 => .seq (phP N off u aE t) (qopW [dg (ipw 0) (ipw (aE t))] (2 ^ t) (N / 2 ^ (t + 1)) off)

/-- Cost of one pass over a segment of `N` entries at offset `off`. -/
def cPass (w N off Z : ℕ) : ℕ :=
  N * (136000 * w + 551000) + 4 * off * (250 * w + 2600) + 8000 + 6 * (4000 * (Z + 2) + 2 * Z + 2)

/-- The program of the step `phase w k₀ k₁ a` (roles `Fin W`, `m` slots, `n` columns). -/
def phaseProg (m n : ℕ) (w k₀ k₁ : ℕ) (a : Fin m → ℕ) : RS :=
  phP (2 ^ (m * n)) (w * 2 ^ (m * n)) (k₀ + k₁ * n) (phE (n := n) a) (m * n)

end NLogN.KFB

end

section
-- Source module: Solutions.KFBGate
/-!
# KFB: role gates on adjacent roles and role permutations

* `segApply` with any matrix list (`getD_segApply`).
* `pass_fsp`: one `qopW` pass on a clean managed state, as `segApply`.
* **`gate2_iruns`**: the step `gate (g2Mat r M)` (the `2 × 2` Gaussian matrix `M` with
  `{-1,0,1}` components on the adjacent roles `r, r+1`, identity elsewhere) as one pass
  (`B1 = 2^{mn}`, one block), on the encoding `Rep`.
* **`perm_iruns`**: every role permutation `σ` as a product of adjacent swaps, each one pass;
  the program and the number of passes depend only on `σ`.
-/

set_option linter.unusedSimpArgs false
set_option linter.unusedVariables false

namespace NLogN.KFB

open IntMul RS AModel KIL KFD Complex

/-! ## Gates on two adjacent roles -/

/-- The role matrix: `M` on roles `r, r+1`, identity elsewhere. -/
noncomputable def g2Mat (W r : ℕ) (M : M2G) : Matrix (Fin W) (Fin W) ℂ := fun p q =>
  if p.val = r ∧ q.val = r then NLogN.KFB.toC M.a
  else if p.val = r ∧ q.val = r + 1 then NLogN.KFB.toC M.b
  else if p.val = r + 1 ∧ q.val = r then NLogN.KFB.toC M.c
  else if p.val = r + 1 ∧ q.val = r + 1 then NLogN.KFB.toC M.d
  else if p = q then 1 else 0

/-- The list after the gate pass on roles `r, r+1`. -/
def g2L (N r : ℕ) (M : M2G) (zs : List GI) : List GI := segApply (r * N) [M] N 1 zs

/-! ## Bounds -/

/-! ## Role permutations -/

/-- The swap of roles `r, r+1` as a `2 × 2` matrix. -/
def swM : M2G := ⟨0, ipw 0, ipw 0, 0⟩

/-- Adjacent swaps `r ∈ L` (the last one first). -/
def permP (N : ℕ) : List ℕ → RS
  | [] => .skip
  | r :: L => .seq (permP N L) (qopW [swM] N 1 (r * N))

def permL (N : ℕ) : List ℕ → List GI → List GI
  | [], zs => zs
  | r :: L, zs => g2L N r swM (permL N L zs)

/-- The adjacent swap with index `i : Fin k` on `Fin (k + 1)`. -/
def asw (k : ℕ) (i : Fin k) : Equiv.Perm (Fin (k + 1)) := Equiv.swap i.castSucc i.succ

end NLogN.KFB

end

section
-- Source module: Solutions.KFBSwap
/-!
# KFB: swapping two address bits of a Gaussian array (KFC field moves)

With `q = 2`, KFD's fine field interchange `swapEF` (one `fwdF` and one `bwdF`, linear cost)
swaps two binary fields of the integer array `gflat zs`.  Viewing the Gaussian array of
`W · 2^K` entries as the shape `[W 2^{K-Q-1}, 2, 2^{Q-P-1}, 2, 2^{P+1}]` of integers, the two
binary fields are the index bits `Q > P`.

* `swbit P Q k`: `k` with bits `P` and `Q` interchanged.
* `gflat_gswL`: `gflat (gswL P Q zs) = pull … (swapPos 1 3) … (gflat zs)`.
* `rep_gsw`: on the encoding, the swap of the address bits at positions `P`, `Q` (all roles).
* `gsw_iruns`: the fine IL run.
-/

set_option linter.unusedSimpArgs false
set_option linter.unusedVariables false

namespace NLogN.KFB

open IntMul RS AModel KIL KFD Complex

/-- `k` with bits `P` and `Q` interchanged. -/
def swbit (P Q k : ℕ) : ℕ :=
  if k / 2 ^ P % 2 = k / 2 ^ Q % 2 then k
  else if k / 2 ^ P % 2 = 1 then k - 2 ^ P + 2 ^ Q else k + 2 ^ P - 2 ^ Q

/-- The integer shape of a Gaussian array of `W·2^K` entries, split at bits `Q > P`. -/
def swSh (W K P Q : ℕ) : Shape := [W * 2 ^ (K - Q - 1), 2, 2 ^ (Q - P - 1), 2, 2 ^ (P + 1)]

/-! ## The swapped Gaussian array -/

/-- Entry `k` read from `swbit P Q k`. -/
def gswL (P Q : ℕ) (zs : List GI) : List GI :=
  (List.range zs.length).map fun k => zs.getD (swbit P Q k) 0

/-! ## On addresses -/

/-- The address with the bits at positions `P`, `Q` interchanged. -/
def swA {m n : ℕ} (P Q : Fin (m * n)) (b : Fin m × Fin n → Fin 2) : Fin m × Fin n → Fin 2 :=
  fun p => b (finProdFinEquiv.symm (Equiv.swap P Q (finProdFinEquiv p)))

/-! ## The fine IL run -/

/-- The bit-swap program: `swapEF` with literal parameters. -/
def gswP {Pp : ℕ} (W K P Q : ℕ) : IStmt HS Pp :=
  swapEF (.lit ((W * 2 ^ (K - Q - 1) : ℕ) : ℤ)) (.lit 2) (.lit ((2 ^ (Q - P - 1) : ℕ) : ℤ))
    (.lit ((2 ^ (P + 1) : ℕ) : ℤ))

/-- Its cost (`V = 2 W 2^K` integers of `≤ wb` bits). -/
def cGsw (W K P Q wb : ℕ) : ℕ :=
  cSwapEF (.lit ((W * 2 ^ (K - Q - 1) : ℕ) : ℤ)) (.lit 2) (.lit ((2 ^ (Q - P - 1) : ℕ) : ℤ))
    (.lit ((2 ^ (P + 1) : ℕ) : ℤ)) 2 (2 * (W * 2 ^ K)) (zsz (((2 * (W * 2 ^ K)) : ℕ) : ℤ)) wb

end NLogN.KFB

end

section
-- Source module: Solutions.KFBCnot
/-!
# KFB: controlled address-bit flips (control above target) on one role

`qopW (cnMs d) 2^t G (w·2^K)`: the periodic matrix list `cnMs d` (`I` or the swap `X` according
to bit `d` of the block number) applied to blocks `[u ‖ v]` of `2^t` entries flips address bit
`t` exactly where bit `c = t + d + 1` is set, on the segment of role `w` only.

* `getD_cnL`: index semantics.
* `cnA`: the address map `b(tgt) += b(ctl)`; `rep_cnot`: the pass on the encoding.
* `cnot_fsp`: the fine run, cost `cPass`.
-/

set_option linter.unusedSimpArgs false
set_option linter.unusedVariables false

namespace NLogN.KFB

open IntMul RS AModel KIL KFD Complex

/-- The identity `2 × 2` matrix. -/
def idM : M2G := ⟨ipw 0, 0, 0, ipw 0⟩

/-- `X` on blocks whose number has bit `d` set, `I` elsewhere (period `2^{d+1}`). -/
def cnMs (d : ℕ) : List M2G :=
  (List.range (2 ^ (d + 1))).map fun p => if p / 2 ^ d % 2 = 1 then swM else idM

/-- The CNOT pass on the list level. -/
def cnL (K t d off : ℕ) (zs : List GI) : List GI :=
  segApply off (cnMs d) (2 ^ t) (2 ^ K / 2 ^ (t + d + 2)) zs

/-! ## On addresses -/

/-- `b(tgt) += b(ctl)` (positions in `Fin (m n)`). -/
def cnA {m n : ℕ} (ctl tgt : Fin (m * n)) (b : Fin m × Fin n → Fin 2) : Fin m × Fin n → Fin 2 :=
  fun p => if finProdFinEquiv p = tgt then b p + b (finProdFinEquiv.symm ctl) else b p

/-! ## On the encoding -/

end NLogN.KFB

end

section
-- Source module: Solutions.KFBRow
/-!
# KFB: the `rowAdd` step as a fine IL program

`rowAdd w i j`: `x(j, c) += x(i, c)` for every column `c`, on the address bits of role `w`.
Column `c` is one controlled flip of bit `pos(j,c)` controlled by bit `pos(i,c)`
(`pos(r,c) = c + n r`):

* `i > j`: the control is above the target: one `cnot` pass;
* `i < j`: swap the two bits (all roles, `gswP`), the pass with control and target exchanged,
  swap back.

* `den_rowAdd_mulVec`: the denotation on states.
* `rep_col`, `rep_rowP`: the passes on the encoding; **`rowAdd_iruns`**: the fine IL run.
-/

set_option linter.unusedSimpArgs false
set_option linter.unusedVariables false

namespace NLogN.KFB

open IntMul RS AModel KIL KFD Complex

/-! ## Address maps -/

/-- Position of slot `r`, column `c`. -/
def pos {m n : ℕ} (r : Fin m) (c : Fin n) : Fin (m * n) := finProdFinEquiv (r, c)

/-- Column `C` of the row addition. -/
def colA {m n : ℕ} (i j : Fin m) (C : ℕ) (b : Fin m × Fin n → Fin 2) : Fin m × Fin n → Fin 2 :=
  fun p => if p.1 = j ∧ p.2.val = C then b p + b (i, p.2) else b p

/-- Columns `< C` of the row addition. -/
def raC {m n : ℕ} (i j : Fin m) (C : ℕ) (b : Fin m × Fin n → Fin 2) : Fin m × Fin n → Fin 2 :=
  fun p => if p.1 = j ∧ p.2.val < C then b p + b (i, p.2) else b p

/-! ## One column -/

/-- The list after column `c`. -/
def colL (W m n w : ℕ) (i j : Fin m) (c : Fin n) (zs : List GI) : List GI :=
  if j.val < i.val then
    cnL (m * n) (NLogN.KFB.pos j c).val (n * (i.val - j.val) - 1) (w * 2 ^ (m * n)) zs
  else
    gswL (NLogN.KFB.pos i c).val (NLogN.KFB.pos j c).val
      (cnL (m * n) (NLogN.KFB.pos i c).val (n * (j.val - i.val) - 1) (w * 2 ^ (m * n))
        (gswL (NLogN.KFB.pos i c).val (NLogN.KFB.pos j c).val zs))

/-- The program of column `c`. -/
def colP {Pp : ℕ} (W m n w : ℕ) (i j : Fin m) (c : Fin n) : IStmt HS Pp :=
  if j.val < i.val then
    fcode (qopW (cnMs (n * (i.val - j.val) - 1)) (2 ^ (NLogN.KFB.pos j c).val)
      (2 ^ (m * n) / 2 ^ ((NLogN.KFB.pos j c).val + (n * (i.val - j.val) - 1) + 2)) (w * 2 ^ (m * n)))
  else
    .seq (gswP W (m * n) (NLogN.KFB.pos i c).val (NLogN.KFB.pos j c).val)
      (.seq (fcode (qopW (cnMs (n * (j.val - i.val) - 1)) (2 ^ (NLogN.KFB.pos i c).val)
        (2 ^ (m * n) / 2 ^ ((NLogN.KFB.pos i c).val + (n * (j.val - i.val) - 1) + 2)) (w * 2 ^ (m * n))))
        (gswP W (m * n) (NLogN.KFB.pos i c).val (NLogN.KFB.pos j c).val))

/-- Uniform cost bound of one column. -/
def cColU (W K wb : ℕ) : ℕ :=
  2 * cGsw W K 0 0 wb + NLogN.KFB.cPass wb (2 ^ K) (W * 2 ^ K) (zsz (((4 * (3 * 2 ^ K + W * 2 ^ K)) : ℕ) : ℤ))

/-! ## All columns -/

def rowL (W m n w : ℕ) (i j : Fin m) : ℕ → List GI → List GI
  | 0, zs => zs
  | C + 1, zs => if h : C < n then colL W m n w i j ⟨C, h⟩ (rowL W m n w i j C zs)
      else rowL W m n w i j C zs

def rowP {Pp : ℕ} (W m n w : ℕ) (i j : Fin m) : ℕ → IStmt HS Pp
  | 0 => .skip
  | C + 1 => if h : C < n then .seq (rowP W m n w i j C) (colP W m n w i j ⟨C, h⟩)
      else rowP W m n w i j C

end NLogN.KFB

end

section
-- Source module: Solutions.KFBChild
/-!
# KFB: the `child` step (a recursive call on one slot of one role)

Role `w`'s segment, as Gaussian indices `(hi, s, lo)` (`s` = the `n` bits of slot `h`), is
rearranged to `(hi, lo, s)` by KFC's `fwdF` (`q = 2^n`, the real/imaginary pair kept as the
last field), so that slot `h` becomes `2^{n(m-1)}` contiguous blocks of `2^n` entries; a program
`PT` computing an operator `T` blockwise runs on the segment; `bwdF` moves the field back.

* arithmetic: `sum_bits_window`, `sidx`, `aidx_slot`, `aidx_upd`;
* `mvL`, `gflat_mvL`: the two field moves as Gaussian index maps (= KFC `pull` on `gflat`);
* `den_child_mulVec`: the child matrix on vectors;
* `rep_child`: the denotation `RStep.den T (child w h fwd)` on the encoding, given the
  blockwise semantics of the called program.
-/

set_option linter.unusedSimpArgs false
set_option linter.unusedVariables false

namespace NLogN.KFB

open IntMul RS AModel KIL KFD Complex

/-! ## Bit windows -/

/-- The index of a slot value. -/
def sidx {n : ℕ} (σ : Fin n → Fin 2) : ℕ := ∑ c : Fin n, (σ c).val * 2 ^ c.val

/-- `b` with slot `h` replaced by `σ`. -/
def updS {m n : ℕ} (b : Fin m × Fin n → Fin 2) (h : Fin m) (σ : Fin n → Fin 2) :
    Fin m × Fin n → Fin 2 := fun p => if p.1 = h then σ p.2 else b p

/-! ## Field moves as Gaussian index maps -/

/-- Input index read by output index `k` of the move `[A, Y, X] → [A, X, Y]`. -/
def g4 (X Y k : ℕ) : ℕ := (k / (X * Y) * Y + k % Y) * X + k / Y % X

/-- The Gaussian array after the move `[A, Y, X] → [A, X, Y]`. -/
def mvL (X Y : ℕ) (zs : List GI) : List GI :=
  (List.range zs.length).map fun k => zs.getD (g4 X Y k) 0

/-! ## The child matrix on vectors -/

/-! ## The child step on the encoding -/

/-- Blockwise semantics of a called program `F` on Gaussian arrays: on every array of length a
multiple of `2^n`, each block of `2^n` consecutive entries is mapped by `U`. -/
def BlkT {n : ℕ} (U : Butterfly.Op (Fin n)) (F : List GI → List GI) : Prop :=
  ∀ ys : List GI, 2 ^ n ∣ ys.length → (F ys).length = ys.length ∧
    ∀ β (σ : Fin n → Fin 2), β * 2 ^ n + 2 ^ n ≤ ys.length →
      NLogN.KFB.toC ((F ys).getD (β * 2 ^ n + sidx σ) 0) =
        ∑ σ', U σ σ' * NLogN.KFB.toC (ys.getD (β * 2 ^ n + sidx σ') 0)

/-- The child step: move slot `h` of the segment `[off, off+N)` innermost, apply `F`, move back. -/
def chL (N off q M : ℕ) (F : List GI → List GI) (zs : List GI) : List GI :=
  zs.take off ++ (NLogN.KFB.mvL q M (F (NLogN.KFB.mvL M q ((zs.drop off).take N))) ++ zs.drop (off + N))

end NLogN.KFB

end

section
-- Source module: Solutions.KFCRot
/-!
# KFC: fine-cost block rotation (paper Lemma 4, `u = 1`)

`rotF` computes the same denotation as `KIL.rotP`: on `[P, Q, B]` with the offset table
`a_0, …, a_{P-1}` (`a_p ≤ Q`, binary, one entry per block) on stack `14` (consumed), it produces
`pull [P,Q,B] [P,Q,B] (rotAddr Q as) 0 xs`.

Per block: `c = Q - a_p` is computed in binary (`O(zsz Q)`), both `a_p` and `c` are converted to
unary counters with the linear `toUL` (`O(Q)`), and the block is rotated by counted record moves
(`mvR`): the first `c` records go to `T`, then to `A` (restoring their order), the last `a_p`
records go to the output, then the first `c`.  Every entry is moved `≤ 4` times.
-/

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

namespace NLogN

open IntMul AModel

namespace KFC

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

/-! ## Moving `k` records of `B` entries -/

def mvR (X Y C D : Fin HS) : IStmt HS Pp := uloop C D (mvB X Y sCB sDB)

/-! ## `c = Q - a` in binary -/

abbrev rV : Fin HS := 31
abbrev rW : Fin HS := 32
abbrev sTb : Fin HS := 14
abbrev sQ : Fin HS := 2

def cmpC : IStmt HS Pp :=
  .seq (.op (.dup sQ NLogN.KFC.rV)) (.seq (.op (.dup sTb NLogN.KFC.rW)) (.seq (.op (.neg NLogN.KFC.rW)) (.op (.add NLogN.KFC.rV NLogN.KFC.rW NLogN.KFC.rV))))

/-! ## One block, in three parts -/

def rpre : IStmt HS Pp :=
  .seq cmpC (.seq (pro sTb sCQ sDQ) (.seq (pro NLogN.KFC.rV sCM sDM) (.op (.drop NLogN.KFC.rV))))

def rmid : IStmt HS Pp :=
  .seq (NLogN.KFC.mvR sX NLogN.KFC.sT sCM sDM) (.seq (NLogN.KFC.mvR NLogN.KFC.sT sA sCM sDM) (.seq (NLogN.KFC.mvR sX NLogN.KFC.sO sCQ sDQ)
    (NLogN.KFC.mvR sA NLogN.KFC.sO sCM sDM)))

def rpost : IStmt HS Pp := .seq (epi sCQ sDQ) (.seq (epi sCM sDM) (.op (.drop sTb)))

def rblk : IStmt HS Pp := .seq rpre (.seq rmid rpost)

def cRB (Q B w : ℕ) : ℕ := 3 * Q * (B * kap w + 124) + 1100 * Q + 2200 * zsz (Q : ℤ) + 4000

end KFC

end NLogN

end

section
-- Source module: Solutions.KFCRotMain
/-!
# KFC: the self-contained fine-cost block rotation `rotF`

`rotF_spec`: the same result as `KIL.rotP_spec`
(`(t.upd 0 (ints (pull [P,Q,B] [P,Q,B] (rotAddr Q as) 0 xs) ++ R)).upd 14 L14`), with cost
`cRotF P Q B w ≤ 68·kap w·(P·Q·B)` (`cRotF_le`), `w` a bound on the bit size of the entries.
-/

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

namespace NLogN

open IntMul AModel

namespace KFC

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

def allR : IStmt HS Pp := uloop sCP sDP rblk

/-! ## Core, wrapper, program -/

def coreR : IStmt HS Pp :=
  .seq (.op (.mark NLogN.KFC.sO)) (.seq allR (.seq (mvW NLogN.KFC.sO sX) (.op (.unmark NLogN.KFC.sO))))

def rotF : IStmt HS Pp := NLogN.KFC.wrap 4 sCB sDB (NLogN.KFC.wrap 1 sCP sDP coreR)

def cRotF (P Q B w : ℕ) : ℕ :=
  P * (cRB Q B w + 64) + (P * Q * B) * kap w + cUL P + cUL B + 31 * (P + B) + 300

end KFC

end NLogN

end

section
-- Source module: Solutions.KFCShift
/-!
# KFC: the Lemma 10 controlled shift (`u = 1`) with the fine-cost rotation

`rotF_shift_spec`: `rotF` with the offset table `offs Λ c A` on stack `14` computes
`shiftArr Λ i 1 c 0 xs` (same denotation as `KIL.rotP_shift_spec`), cost `cRotF`.

`shiftArrF_spec`: the self-contained program `shF` = (`genP` table generator + literal pushes,
register code) ; `rotF` ; (pops).  Its cost is
`cRotF (vol A) Q (vol C) w + 400000·(|A|+1)·vol A·(Bb+2) + 13000·(Bb+2)`:
the per-element part is `O(kap w · V)` with no `Bb`, and `Bb` (a bound on the bit size of the
descriptor numbers `vol A, Q, vol C, 2Q` only) enters with the number of blocks `vol A`.
-/

set_option linter.unusedSimpArgs false

namespace NLogN

open IntMul AModel

namespace KFC

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

section

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

/-- Table generation and parameter pushes (register code). -/
def shPre (Q P B : ℕ) (ac : List (ℕ × ℕ)) : RS :=
  .seq (KIL.genP Q ac) (.seq (.push 1 (.lit P)) (.seq (.push 2 (.lit Q)) (.push 4 (.lit B))))

def shPost : RS := .seq (.pop 1) (.seq (.pop 2) (.pop 4))

/-- The self-contained fine-cost Lemma 10 shift. -/
def shF (Q P B : ℕ) (ac : List (ℕ × ℕ)) : IStmt HS Pp :=
  .seq (shPre Q P B ac).code (.seq rotF shPost.code)

end

end KFC

end NLogN

end

section
-- Source module: Solutions.KFGLin
/-!
# KFG: scaled Gaussian coefficients and scaled linear-combination expressions

Dyadic Gaussian data are stored as Gaussian integers with a per-role scale `2^{-e}`.  Aligning
two roles of different scales multiplies one stream by `2^d`, which costs `O(w + d)` per entry
through the ruler shift `Ex.shl` (a linear expression).

* `SC`: a scaled coefficient `c · 2^d` with `c ∈ {-1,0,1}²` (`SC.Ok`), acting by `smulS`.
* `tsE c d r = (c · top r) ≪ 2^d`, `lsE l = Σ tsE`: linear register expressions
  (`lsE_eval`, `lsE_bnd`, `lsE_lin`, `lsE_size`).
* `push_lsE`: one push of such a combination, at cost `(5 |l| + 2) · 2000 · (w + D + e + 2)`.
* `M2S`: a `2 × 2` matrix of scaled coefficients; `lin2S p q x y = p·x + q·y`; `lin2S_bnd`.
-/

set_option linter.unusedSimpArgs false
set_option linter.unusedVariables false

namespace NLogN.KFG

open IntMul RS AModel KIL KFD KFB Complex

/-! ## Scaled coefficients -/

/-- A scaled Gaussian coefficient `c · 2^d`. -/
structure SC where
  c : GI
  d : ℕ

def SC.Ok (p : NLogN.KFG.SC) : Prop := Sm p.c

/-- The Gaussian integer `c · 2^d`. -/
def SC.val (p : NLogN.KFG.SC) : GI := (p.c.1 * 2 ^ p.d, p.c.2 * 2 ^ p.d)

/-- Action of a scaled coefficient. -/
def smulS (p : NLogN.KFG.SC) (z : GI) : GI := gmul p.val z

/-- `p x + q y`. -/
def lin2S (p q : NLogN.KFG.SC) (x y : GI) : GI := smulS p x + smulS q y

/-- A `2 × 2` matrix of scaled coefficients `[[a, b], [c, d]]`. -/
structure M2S where
  a : NLogN.KFG.SC
  b : NLogN.KFG.SC
  c : NLogN.KFG.SC
  d : NLogN.KFG.SC

def M2S.Ok (M : M2S) : Prop := M.a.Ok ∧ M.b.Ok ∧ M.c.Ok ∧ M.d.Ok

/-- All shifts of `M` are at most `D`. -/
def M2S.Sh (M : M2S) (D : ℕ) : Prop := M.a.d ≤ D ∧ M.b.d ≤ D ∧ M.c.d ≤ D ∧ M.d.d ≤ D

/-! ## Scaled terms -/

/-- `c · 2^d · top(r)` for `c ∈ {-1, 0, 1}`. -/
def tsE (c : ℤ) (d : ℕ) (r : Fin HS) : Ex HS := .shl (tmE c r) (.lit ((2 : ℤ) ^ d))

/-- `Σ c_t · 2^{d_t} · top(r_t)`. -/
def lsE : List (ℤ × ℕ × Fin HS) → Ex HS
  | [] => .lit 0
  | (c, d, r) :: l => .add (tsE c d r) (lsE l)

/-- The value of a scaled combination. -/
def lsV (s : IState HS) (l : List (ℤ × ℕ × Fin HS)) : ℤ :=
  (l.map fun t => t.1 * 2 ^ t.2.1 * s.top t.2.2).sum

/-! ## Real and imaginary parts of `p u + q v` -/

def reS (p q : NLogN.KFG.SC) : List (ℤ × ℕ × Fin HS) :=
  [(p.c.1, p.d, rUr), (-p.c.2, p.d, rUi), (q.c.1, q.d, rVr), (-q.c.2, q.d, rVi)]
def imS (p q : NLogN.KFG.SC) : List (ℤ × ℕ × Fin HS) :=
  [(p.c.2, p.d, rUr), (p.c.1, p.d, rUi), (q.c.2, q.d, rVr), (q.c.1, q.d, rVi)]

/-! ## Bit bounds -/

end NLogN.KFG

end

section
-- Source module: Solutions.KFGZip
/-!
# KFG: the scaled zip loop

`zipRS M` is KFB's `zipR` with scaled coefficients: from the four register streams
`rUr, rUi, rVr, rVi` (real and imaginary parts of `u` and `v`) it pushes `M₀₀ u + M₀₁ v` on `qA`
and `M₁₀ u + M₁₁ v` on `qO`, entrywise, where every coefficient is `c · 2^d`
(`c ∈ {-1,0,1}²`, `d ≤ D`).  Cost `O(w + D)` per entry (`cZS`).
-/

set_option linter.unusedSimpArgs false
set_option linter.unusedVariables false

namespace NLogN.KFG

open IntMul RS AModel KIL KFD KFB Complex

/-- One output entry of each stream. -/
def zbS (M : M2S) : RS :=
  .seq (.push qA (lsE (reS M.a M.b))) (.seq (.push qA (lsE (imS M.a M.b)))
    (.seq (.push qO (lsE (reS M.c M.d))) (.seq (.push qO (lsE (imS M.c M.d)))
      (.seq (.pop rUr) (.seq (.pop rUi) (.seq (.pop rVr) (.pop rVi)))))))

def zipRS (M : M2S) : RS := .rep NLogN.KFB.c1 (.reg NLogN.KFB.rB) (zbS M)

/-- Per-entry cost of the scaled zip loop. -/
def cZS (w D : ℕ) : ℕ := 180000 * (w + D + 4)

end NLogN.KFG

end