Round-six κ development, definitions part 66

-- Round-six κ development (integer-multiplication-kappa), generated by skeleton subtraction.

import Definitions.Def_IntMul_Kappa_R6_Defs65
import Mathlib
import Definitions.Def_IntMul_HvdH_M_Defs8
import Definitions.Def_IntMul_MultitapeModel
import Definitions.Def_IntMul_HvdH_Numeric
import Definitions.Def_IntMul_HvdH_ExpApprox
import Definitions.Def_IntMul_HvdH_RealExp
import Definitions.Def_IntMul_HvdH_MachineRef
import Definitions.Def_IntMul_HvdH_Twiddle
import Definitions.Def_IntMul_HvdH_SApprox
import Definitions.Def_IntMul_HvdH_DiagApprox
import Definitions.Def_IntMul_HvdH_EApprox
import Definitions.Def_IntMul_HvdH_JApprox
import Definitions.Def_IntMul_HvdH_BApprox
import Definitions.Def_IntMul_HvdH_FPApprox
import Definitions.Def_IntMul_KappaFramework
import Definitions.Def_IntMul_KappaJainData
import Definitions.Def_IntMul_Kappa_K_Defs14
import Definitions.Def_IntMul_HvdH_M_Defs2
import Definitions.Def_IntMul_Kappa_K_Defs29
import Definitions.Def_IntMul_HvdH_M_Defs3
import Definitions.Def_IntMul_Kappa_R6_Defs3
import Definitions.Def_IntMul_HvdH_M_Defs1
import Definitions.Def_IntMul_Kappa_R6_Defs14
import Definitions.Def_IntMul_Kappa_R6_Defs1
import Definitions.Def_IntMul_Kappa_R6_Defs16
import Definitions.Def_IntMul_Kappa_R6_Defs13
import Definitions.Def_IntMul_Kappa_K_Defs10
import Definitions.Def_IntMul_Kappa_K_Defs27
import Definitions.Def_IntMul_Kappa_K_Defs28

set_option autoImplicit false
set_option relaxedAutoImplicit false
set_option maxHeartbeats 1000000

section
-- Source module: Solutions.J6SegKr7
/-!
# B3 (chunk kernel), part 7: the block loop

`loopB_runs`: the loop `while BC > 0: blkI` runs through the descriptor list `ds` (encoded on the
table stack by `dsEnc`), consuming the blocks of the row and pushing `gflat (chunkB …).reverse`
onto `OUT`.
-/

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

namespace J6Seg.KR

open IntMul NLogN NLogN.KIL NLogN.KFD NLogN.KFB NLogN.KFV

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

/-- All stage values of all blocks bounded. -/
def chunkOK (P : ℕ) (A : ℤ) (di gi gam dd : List ℤ) : List (ℕ × List ℤ × List ℤ) → List GI → Prop
  | [], _ => True
  | (n, am, ap) :: ds, z => blkOK P A (2 * n) di gi gam dd am ap (z.take n) ∧
      chunkOK P A di gi gam dd ds (z.drop n)

/-- Cost of the block loop. -/
def cLoop (pk b φD : ℕ) (φT : ℕ → ℕ) (B BH : ℕ) : List (ℕ × List ℤ × List ℤ) → ℕ
  | [] => 30
  | d :: ds => cBlk d.1 pk b φD (φT d.1) B BH + cLoop pk b φD φT B BH ds + 30

end J6Seg.KR

end

section
-- Source module: Solutions.J6SegKr8
/-!
# B3 (chunk kernel), part 8: the chunk kernel in IL

`chunkI pM`: with the chunk table `tabK …` on `dT` and the row `x` (length `λ`) on stack `0`, it
replaces `x` by `chunkG2 … x` and leaves the table in place; all its other stacks end empty.

The table (`tabK`): the header `[λ, 2^P, 2^(P+1), nb, |tab|, 2^pk, 2^b, 2(λ - Σn)]`, the four
stage tables reversed, then per block `blkHdr n pk b ++ am ++ ap`.
-/

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

namespace J6Seg.KR

open IntMul NLogN NLogN.KIL NLogN.KFD NLogN.KFB NLogN.KFV

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

/-- Header registers: `λ`, the table length, `2(λ - Σn)`. -/
abbrev RL : Fin HS := 30
abbrev RT : Fin HS := 34
abbrev R2 : Fin HS := 39

/-- **The kernel frame**: the twelve stacks the kernel keeps for the whole chunk. -/
def KF (vl vp vp1 vt vc vb vh ti tg ta td vd : List ℤ) : Fin HS → List ℤ :=
  Function.update (Function.update (Function.update (Function.update (Function.update (Function.update (Function.update (Function.update (Function.update (Function.update (Function.update (Function.update ((fun _ => [])) RL vl) rP vp) rP1 vp1) RT vt) kc vc) kB vb) R2 vh) TDI ti) TGI tg) TGA ta) TDD td) kD vd

section KFl
variable {vl vp vp1 vt vc vb vh ti tg ta td vd : List ℤ}

@[simp] theorem KF_vl : KF vl vp vp1 vt vc vb vh ti tg ta td vd RL = vl := by
  unfold KF; rw [Function.update_of_ne (show RL ≠ kD by decide), Function.update_of_ne (show RL ≠ TDD by decide), Function.update_of_ne (show RL ≠ TGA by decide), Function.update_of_ne (show RL ≠ TGI by decide), Function.update_of_ne (show RL ≠ TDI by decide), Function.update_of_ne (show RL ≠ R2 by decide), Function.update_of_ne (show RL ≠ kB by decide), Function.update_of_ne (show RL ≠ kc by decide), Function.update_of_ne (show RL ≠ RT by decide), Function.update_of_ne (show RL ≠ rP1 by decide), Function.update_of_ne (show RL ≠ rP by decide), Function.update_self]
@[simp] theorem KF_vp : KF vl vp vp1 vt vc vb vh ti tg ta td vd rP = vp := by
  unfold KF; rw [Function.update_of_ne (show rP ≠ kD by decide), Function.update_of_ne (show rP ≠ TDD by decide), Function.update_of_ne (show rP ≠ TGA by decide), Function.update_of_ne (show rP ≠ TGI by decide), Function.update_of_ne (show rP ≠ TDI by decide), Function.update_of_ne (show rP ≠ R2 by decide), Function.update_of_ne (show rP ≠ kB by decide), Function.update_of_ne (show rP ≠ kc by decide), Function.update_of_ne (show rP ≠ RT by decide), Function.update_of_ne (show rP ≠ rP1 by decide), Function.update_self]
@[simp] theorem KF_vp1 : KF vl vp vp1 vt vc vb vh ti tg ta td vd rP1 = vp1 := by
  unfold KF; rw [Function.update_of_ne (show rP1 ≠ kD by decide), Function.update_of_ne (show rP1 ≠ TDD by decide), Function.update_of_ne (show rP1 ≠ TGA by decide), Function.update_of_ne (show rP1 ≠ TGI by decide), Function.update_of_ne (show rP1 ≠ TDI by decide), Function.update_of_ne (show rP1 ≠ R2 by decide), Function.update_of_ne (show rP1 ≠ kB by decide), Function.update_of_ne (show rP1 ≠ kc by decide), Function.update_of_ne (show rP1 ≠ RT by decide), Function.update_self]
@[simp] theorem KF_vt : KF vl vp vp1 vt vc vb vh ti tg ta td vd RT = vt := by
  unfold KF; rw [Function.update_of_ne (show RT ≠ kD by decide), Function.update_of_ne (show RT ≠ TDD by decide), Function.update_of_ne (show RT ≠ TGA by decide), Function.update_of_ne (show RT ≠ TGI by decide), Function.update_of_ne (show RT ≠ TDI by decide), Function.update_of_ne (show RT ≠ R2 by decide), Function.update_of_ne (show RT ≠ kB by decide), Function.update_of_ne (show RT ≠ kc by decide), Function.update_self]
@[simp] theorem KF_vc : KF vl vp vp1 vt vc vb vh ti tg ta td vd kc = vc := by
  unfold KF; rw [Function.update_of_ne (show kc ≠ kD by decide), Function.update_of_ne (show kc ≠ TDD by decide), Function.update_of_ne (show kc ≠ TGA by decide), Function.update_of_ne (show kc ≠ TGI by decide), Function.update_of_ne (show kc ≠ TDI by decide), Function.update_of_ne (show kc ≠ R2 by decide), Function.update_of_ne (show kc ≠ kB by decide), Function.update_self]
@[simp] theorem KF_vb : KF vl vp vp1 vt vc vb vh ti tg ta td vd kB = vb := by
  unfold KF; rw [Function.update_of_ne (show kB ≠ kD by decide), Function.update_of_ne (show kB ≠ TDD by decide), Function.update_of_ne (show kB ≠ TGA by decide), Function.update_of_ne (show kB ≠ TGI by decide), Function.update_of_ne (show kB ≠ TDI by decide), Function.update_of_ne (show kB ≠ R2 by decide), Function.update_self]
@[simp] theorem KF_vh : KF vl vp vp1 vt vc vb vh ti tg ta td vd R2 = vh := by
  unfold KF; rw [Function.update_of_ne (show R2 ≠ kD by decide), Function.update_of_ne (show R2 ≠ TDD by decide), Function.update_of_ne (show R2 ≠ TGA by decide), Function.update_of_ne (show R2 ≠ TGI by decide), Function.update_of_ne (show R2 ≠ TDI by decide), Function.update_self]
@[simp] theorem KF_ti : KF vl vp vp1 vt vc vb vh ti tg ta td vd TDI = ti := by
  unfold KF; rw [Function.update_of_ne (show TDI ≠ kD by decide), Function.update_of_ne (show TDI ≠ TDD by decide), Function.update_of_ne (show TDI ≠ TGA by decide), Function.update_of_ne (show TDI ≠ TGI by decide), Function.update_self]
@[simp] theorem KF_tg : KF vl vp vp1 vt vc vb vh ti tg ta td vd TGI = tg := by
  unfold KF; rw [Function.update_of_ne (show TGI ≠ kD by decide), Function.update_of_ne (show TGI ≠ TDD by decide), Function.update_of_ne (show TGI ≠ TGA by decide), Function.update_self]
@[simp] theorem KF_ta : KF vl vp vp1 vt vc vb vh ti tg ta td vd TGA = ta := by
  unfold KF; rw [Function.update_of_ne (show TGA ≠ kD by decide), Function.update_of_ne (show TGA ≠ TDD by decide), Function.update_self]
@[simp] theorem KF_td : KF vl vp vp1 vt vc vb vh ti tg ta td vd TDD = td := by
  unfold KF; rw [Function.update_of_ne (show TDD ≠ kD by decide), Function.update_self]
@[simp] theorem KF_vd : KF vl vp vp1 vt vc vb vh ti tg ta td vd kD = vd := by
  unfold KF; rw [Function.update_self]

end KFl

/-- Read the header and the stage tables. -/
def hdrK : RS :=
  .seq (tkRS RL) (.seq (tkRS rP) (.seq (tkRS rP1) (.seq (tkRS BC) (.seq (tkRS RT) (.seq (tkRS kc)
  (.seq (tkRS kB) (.seq (tkRS R2) (.seq (dupR dT DU TDI CC RL) (.seq (dupR dT DU TGI CC RL)
  (.seq (dupR dT DU TGA CC RL) (.seq (dupR dT DU TDD CC RL) (.push kD (.lit 1)))))))))))))

/-- The zeros, the output, restoring the table, clearing. -/
def finK : RS :=
  .seq (.pop BC) (.seq (zeroR OUT CC R2) (.seq (dropR 0 CC R2) (.seq (mvP OUT 0 TP CC RL)
  (.seq (dropR TDI CC RL) (.seq (dropR TGI CC RL) (.seq (dropR TGA CC RL) (.seq (dropR TDD CC RL)
  (.seq (NLogN.mvK DU dT CC (.reg RT)) (.seq (.pop RL) (.seq (.pop rP) (.seq (.pop rP1)
  (.seq (.pop RT) (.seq (.pop kc) (.seq (.pop kB) (.seq (.pop R2) (.pop kD))))))))))))))))

end J6Seg.KR

end

section
-- Source module: Solutions.J6SegKs0
/-!
# B3′ (chunk kernel with in-kernel Kronecker constants), part 0: the table layout

The block header no longer carries `kG b (2n) 2^pk` and `(2^b)^(4n)`. It carries
`[n, 2n, j, 2nb]` with `j = clog₂(2n)`, and the kernel computes the constants by doubling
(`j` steps) and `pow2`. All table entries are then `O(pk + b + log λ)` bits.

* `blkHdr2 n b = [n, 2n, clog 2 (2n), 2 n b]`;
* `encD2 b d = blkHdr2 d.1 b ++ am ++ ap`, `dsEnc2 b ds`;
* `tabK2 P pk b λ di gi gam dd ds`: the same header as `tabK`
  (`[λ, 2^P, 2^(P+1), |ds|, |tab|, 2^pk, 2^b, 2(λ − Σn)]`), then `di.rev ++ gi.rev ++ gam.rev ++
  dd.rev ++ dsEnc2 b ds`.
-/

namespace J6Seg.KR

/-- The new block header. -/
def blkHdr2 (n b : ℕ) : List ℤ :=
  [(n : ℤ), ((2 * n : ℕ) : ℤ), ((Nat.clog 2 (2 * n) : ℕ) : ℤ), ((2 * n * b : ℕ) : ℤ)]

/-- One block descriptor. -/
def encD2 (b : ℕ) (d : ℕ × List ℤ × List ℤ) : List ℤ := blkHdr2 d.1 b ++ d.2.1 ++ d.2.2

/-- The block descriptors. -/
def dsEnc2 (b : ℕ) (ds : List (ℕ × List ℤ × List ℤ)) : List ℤ := ds.flatMap (encD2 b)

/-- **The chunk table (new layout).** -/
def tabK2 (P pk b lam : ℕ) (di gi gam dd : List ℤ) (ds : List (ℕ × List ℤ × List ℤ)) : List ℤ :=
  [(lam : ℤ), 2 ^ P, 2 ^ (P + 1), (ds.length : ℤ),
   ((8 + 4 * lam + (dsEnc2 b ds).length : ℕ) : ℤ), 2 ^ pk, 2 ^ b,
   ((2 * (lam - sumN ds) : ℕ) : ℤ)] ++ di.reverse ++ gi.reverse ++ gam.reverse ++ dd.reverse ++
    dsEnc2 b ds

end J6Seg.KR

end

section
-- Source module: Solutions.J6SegILJKer1
/-!
# B-line IL, the `J` kernel (B4), part 1: list model and arithmetic

The `J` kernel `jKer` (J6SegSw3) of a padded `G`-row `v` is

  `jKer v = W ++ 0^{G-m}`, `W_r = dotF (cc_r) ((zext K z).drop r)`,
  `z = bwdK bs (fwdK fs (v.take m) []) [] []`

for the coefficient streams `fs = fsOf …`, `bs = bsOf …`, `cc_r = ccL … r`. This file gives

* the generic kernel `jKerG P K G m fs bs cc v` (any streams) and the step forms of `fwdK`
  (`fA_succ`), `bwdK` (`bwdS_succ`) and `dotF` on prefixes (`dotF_take_succ`);
* the arithmetic of the IL: `ediv_of_tdiv` (floor division from `tdiv`/`tmod`), `rs_eq`
  (`rs P x = ⌊(2x + 2^P) / 2^{P+1}⌋`) and **`fdiv_recip`**: with `|2 a 2^P + pv| < 2^k` and
  `pv > 0`, `fdiv P a pv` is `q₀ + δ` where `q₀ = ⌊N iv / 2^k⌋` (`iv = ivOf k pv = ⌊2^k / 2pv⌋`)
  and the correction `δ ∈ {-1,0,1}` is read off the signs of `N - q₀ D` and `N - q₀ D - D`
  (`N = 2 a 2^P + pv`, `D = 2 pv`): two multiplications instead of a long division.
-/

set_option linter.unusedSectionVars false
set_option linter.unusedVariables false

namespace J6Seg.JK

open Finset

/-! ## Arithmetic -/

/-- The reciprocal of a pivot: `⌊2^k / 2pv⌋`. -/
def ivOf (k : ℕ) (pv : ℤ) : ℤ := 2 ^ k / (2 * pv)

/-- The correction of the reciprocal quotient. -/
def corr (r D : ℤ) : ℤ := if r < 0 then -1 else if r - D < 0 then 0 else 1

/-! ## `dotF` on prefixes -/

/-! ## The forward kernel by steps -/

/-- The forward accumulator after `r` rows. -/
def fA (P : ℕ) (fs : List (List ℤ)) (v : List ℤ) (r : ℕ) : List ℤ := fwdK P (fs.take r) (v.take r) []

/-! ## The backward kernel by steps -/

/-- One backward step on the record `(cb, cz, pv, fl)` and the value `bt`. -/
def stepB (P : ℕ) (rc : List ℤ × List ℤ × ℤ × ℤ) (bt : ℤ) (st : List ℤ × List ℤ) :
    List ℤ × List ℤ :=
  (J6Seg.fdiv P (bt - dotF P rc.1 st.1 - dotF P rc.2.1 st.2) rc.2.2.1 :: st.1,
    if rc.2.2.2 = 1 then J6Seg.fdiv P (bt - dotF P rc.1 st.1 - dotF P rc.2.1 st.2) rc.2.2.1 :: st.2
    else st.2)

/-- The fold of the backward steps. -/
def bwdF (P : ℕ) (l : List ((List ℤ × List ℤ × ℤ × ℤ) × ℤ)) (st : List ℤ × List ℤ) :
    List ℤ × List ℤ :=
  l.foldl (fun st x => stepB P x.1 x.2 st) st

/-- The backward state after `j` records. -/
def bwdS (P : ℕ) (bs : List (List ℤ × List ℤ × ℤ × ℤ)) (bts : List ℤ) (j : ℕ) : List ℤ × List ℤ :=
  bwdF P ((bs.take j).zip (bts.take j)) ([], [])

/-! ## The generic kernel -/

/-- **The generic `J` kernel** on one real row. -/
def jKerG (P K G m : ℕ) (fs : List (List ℤ)) (bs : List (List ℤ × List ℤ × ℤ × ℤ))
    (cc : List (List ℤ)) (v : List ℤ) : List ℤ :=
  (List.range m).map (fun r => dotF P (cc.getD r [])
    ((zext K (bwdK P bs (fwdK P fs (v.take m) []) [] [])).drop r)) ++ List.replicate (G - m) 0

end J6Seg.JK

end

section
-- Source module: Solutions.J6SegILJKer2
/-!
# B-line IL, the `J` kernel (B4), part 2: specifications of multiplying register programs

`MSp M p s t c`: the multiplying register program `p : RM` (products through the
multiplication procedure with specification `M`) is executable from `s`, ends in `t` and costs
at most `c`; the analogue of `KIL.Sp`. Here: composition (`MSp.seq`), plain register programs
(`MSp.rs`), a multiplication on a managed state (`MSp.mul`: `Φm + O(X)` for `X`-bit operands,
given `MulBnd M X Φm`), the branch (`MSp.ite`), the counted loop through managed states
(`MSp.rep`, overhead quadratic only in the bit size of the *count*) and the compilation
(`MSp.irun`). Also `floorDivE` / `rsE`: the rounding `rs P` as linear work.
-/

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

namespace J6Seg.JK

open IntMul NLogN NLogN.KIL Function

/-! ## Specifications -/

/-- Specification of a multiplying register program. -/
def MSp (M : MulSpec) (p : RM) (s t : IState HS) (c : ℕ) : Prop :=
  p.ok M s ∧ p.run s = t ∧ p.cost M s ≤ c

/-! ## The rounding as linear work -/

/-- `⌊e / d⌋` for a ruler `d = 2^m`. -/
def floorDivE (e d : Ex HS) : Ex HS := .add (.quo e d) (.sel (.rem e d) (.lit (-1)) (.lit 0))

end J6Seg.JK

end

section
-- Source module: Solutions.J6SegILJKer3
/-!
# B-line IL, the `J` kernel (B4), part 3: the dot loop

**`dotP X`**: with the count `c = |cs|` on top of `CNT`, the coefficients `cs ++ R` on the
table stack `dTJ` and the accumulator `acc` (`c ≤ |acc|`) on `X`, it adds `dotF P cs acc` to
the top of `SUM`, moves `cs` (reversed) onto the consumed-table stack `TT` and leaves `X` as it
was (`dotP_spec`). Every product goes through the multiplication procedure (`RM.mul`); the
rounding `rs P` is linear work (`rsE`: `⌊(2x + 2^P) / 2^{P+1}⌋` by `tdiv` / `tmod` and a sign
test, with the rulers `2^P`, `2^{P+1}` on `rP`, `rP1`). Cost: `(c + 1) · uC Φm Bx Bc`.
-/

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

namespace J6Seg.JK

open IntMul NLogN NLogN.KIL Function

/-! ## Stacks -/

/-- The table. -/
abbrev dTJ : Fin HS := 33
/-- The consumed table (reversed). -/
abbrev TT : Fin HS := 34
abbrev TMP : Fin HS := 35
abbrev SUM : Fin HS := 36
abbrev PX : Fin HS := 37
abbrev CD : Fin HS := 38
abbrev CM : Fin HS := 39
abbrev CNT : Fin HS := 40
/-- `2^P`, `2^{P+1}`, `2^k`. -/
abbrev rP : Fin HS := 41
abbrev rP1 : Fin HS := 42
abbrev rK : Fin HS := 43

/-! ## Expressions -/

/-- `rs P (top r) = ⌊(2x + 2^P) / 2^{P+1}⌋`. -/
def rsE (r : Fin HS) : Ex HS := floorDivE (.add (.add (.reg r) (.reg r)) (.reg rP)) (.reg rP1)

/-! ## The dot loop -/

/-- One step: `SUM += rs P (top dTJ · top X)`, move the coefficient to `TT` and the
accumulator entry to `TMP`. -/
def kD (X : Fin HS) : RM :=
  .seq (.mul PX X dTJ) (.rs (.seq (.asg SUM (.add (.reg SUM) (rsE PX))) (.seq (.pop PX)
    (.seq (.push TT (.reg dTJ)) (.seq (.pop dTJ) (.seq (.push TMP (.reg X)) (.pop X)))))))

/-- The dot loop: `CNT` steps, then the accumulator back from `TMP` to `X`. -/
def dotP (X : Fin HS) : RM := .seq (.rep CD (.reg CNT) (kD X)) (.rs (mvK TMP X CM (.reg CNT)))

/-- The unit cost of the kernel: one multiplication, linear work on `Bx` bits, loop
overhead on `Bc`-bit counters. -/
def uC (Φm Bx Bc : ℕ) : ℕ := Φm + 2 * 10 ^ 5 * (Bx + 2) + 10 ^ 8 * (Bc + 2) ^ 2

section dot
variable (P : ℕ) (F : Fin HS → List ℤ) (X : Fin HS) (cs acc R T0 : List ℤ) (S0 : ℤ) (Ls : List ℤ)

/-- The state of the dot loop after `k` steps. -/
def dSt (k : ℕ) : Fin HS → List ℤ :=
  update (update (update (update (update F dTJ (cs.drop k ++ R)) TT ((cs.take k).reverse ++ T0))
    X (acc.drop k)) TMP (acc.take k).reverse) SUM ((S0 + dotF P (cs.take k) acc) :: Ls)

/-- `X` is a free accumulator stack. -/
structure XOK (X : Fin HS) : Prop where
  lt : X.val < 55
  n1 : X ≠ dTJ
  n2 : X ≠ TT
  n3 : X ≠ TMP
  n4 : X ≠ SUM
  n5 : X ≠ PX
  n6 : X ≠ CD
  n7 : X ≠ CM
  n8 : X ≠ CNT
  n9 : X ≠ rP
  n10 : X ≠ rP1

variable {P F X cs acc R T0 S0 Ls}

end dot

end J6Seg.JK

end

section
-- Source module: Solutions.J6SegILJKer5
/-!
# B-line IL, the `J` kernel (B4), part 5: the forward loop

The forward table `flatT fs = ⋃_r (|cs_r| :: cs_r)` lies on `dTJ`; row `r` reads its count,
runs the dot loop on the accumulator `A = fA P fs v r` and pushes `v_r - dotF P cs_r (fA r)`
(`fA_succ`). After `m = |fs|` rows (`m` on `HM`), `A = fA P fs v m` and the table has moved
(reversed) onto `TT` (`fwdP_spec`). Cost `(∑_r (|cs_r| + 3) + 1) · uC`.
-/

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

namespace J6Seg.JK

open IntMul NLogN NLogN.KIL Function Finset

abbrev A : Fin HS := 44
abbrev V : Fin HS := 46
abbrev CR : Fin HS := 48
abbrev HM : Fin HS := 53

/-- A coefficient row with its count. -/
def rowT (cs : List ℤ) : List ℤ := (cs.length : ℤ) :: cs

/-- A table of rows. -/
def flatT (l : List (List ℤ)) : List ℤ := l.flatMap rowT

/-- One forward row. -/
def fwdB : RM :=
  .seq (.rs (.seq (.push CNT (.reg dTJ)) (.seq (.push TT (.reg dTJ)) (.seq (.pop dTJ)
    (.push SUM (.lit 0)))))) (.seq (dotP A) (.rs (.seq (.push A (.add (.reg V) (.neg (.reg SUM))))
      (.seq (.pop SUM) (.seq (.pop V) (.pop CNT))))))

/-- The forward loop: `top HM` rows. -/
def fwdP : RM := .rep CR (.reg HM) fwdB

section fwd
variable (P : ℕ) (F : Fin HS → List ℤ) (fs : List (List ℤ)) (v R T0 : List ℤ)

/-- The state after `r` forward rows. -/
def fSt (r : ℕ) : Fin HS → List ℤ :=
  update (update (update (update F dTJ (flatT (fs.drop r) ++ R)) TT
    ((flatT (fs.take r)).reverse ++ T0)) A (fA P fs v r)) V (v.drop r)

/-- The frame of the forward loop. -/
structure FwdF (lp lp1 : List ℤ) : Prop where
  cnt : F CNT = []
  sum : F SUM = []
  px : F PX = []
  cd : F CD = []
  tmp : F TMP = []
  cr : F CR = []
  rp : F rP = 2 ^ P :: lp
  rp1 : F rP1 = 2 ^ (P + 1) :: lp1

variable {P F fs v R T0}

end fwd

end J6Seg.JK

end

section
-- Source module: Solutions.J6SegILJKer6
/-!
# B-line IL, the `J` kernel (B4), part 6: the fixed-point division by the reciprocal

`recip` computes `fdiv P num pv` (`pv > 0`) from the numerator on `NUM`, the pivot on `PV` and
its reciprocal `iv = ⌊2^k / 2pv⌋` (from the table) on `IV` with two multiplications and a
sign correction (`fdiv_recip`), provided `|2 num 2^P + pv| < 2^k`. `N_natAbs_lt`: this holds for
`k ≥ 2 Xb + 2P + Bc + 6` when the numerator is `bt` minus two partial dot sums of `Xb`-bit data.
-/

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

namespace J6Seg.JK

open IntMul NLogN NLogN.KIL Function Finset

abbrev Q0 : Fin HS := 19
abbrev PV : Fin HS := 20
abbrev IV : Fin HS := 21
abbrev FL : Fin HS := 22
abbrev BZ : Fin HS := 45
abbrev NN : Fin HS := 49
abbrev QQ : Fin HS := 50
abbrev DD : Fin HS := 51
abbrev NUM : Fin HS := 52
abbrev Z : Fin HS := 54

/-! ## Bounds -/

/-! ## The program -/

/-- The corrected quotient `q₀ + corr (N - q₀ D) D`. -/
def zE : Ex HS := .add (.reg Q0) (.sel (.add (.reg NN) (.neg (.reg QQ))) (.lit (-1))
  (.sel (.add (.add (.reg NN) (.neg (.reg QQ))) (.neg (.reg DD))) (.lit 0) (.lit 1)))

/-- `N = 2 num 2^P + pv`. -/
def nE : Ex HS := .add (.shl (.add (.reg NUM) (.reg NUM)) (.reg rP)) (.reg PV)

/-- `q₀ = ⌊N iv / 2^k⌋`. -/
def q0E : Ex HS := floorDivE (.reg QQ) (.reg rK)

/-- The division by the reciprocal; pushes `fdiv P num pv` onto `Z`. -/
def recip : RM :=
  .seq (.rs (.push NN nE)) (.seq (.mul QQ IV NN) (.seq (.rs (.seq (.push Q0 q0E) (.seq (.pop QQ)
    (.push DD (.add (.reg PV) (.reg PV)))))) (.seq (.mul QQ DD Q0) (.rs (.seq (.push Z zE)
      (.seq (.pop QQ) (.seq (.pop DD) (.seq (.pop Q0) (.pop NN)))))))))

end J6Seg.JK

end

section
-- Source module: Solutions.J6SegILJKer7
/-!
# B-line IL, the `J` kernel (B4), part 7: list moves on managed states

Plain register programs (no products, linear cost) used around the kernel loops:

* `rdT X`: one table entry from `dTJ` onto `X`, recorded on the consumed-table stack `TT`;
* `rdH X`: one header entry from `dTJ` onto `X` (not recorded);
* `popL X C n`: drop `n` entries of `X`;
* `dupMv X Y T C n`: copy the top `n` entries of `X` (reversed) onto `Y` and onto `T`, removing
  them from `X`.
-/

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

namespace J6Seg.JK

open IntMul NLogN NLogN.KIL Function Finset

/-- One table entry onto `X`, recorded on `TT`. -/
def rdT (X : Fin HS) : RS := .seq (.push X (.reg dTJ)) (.seq (.push TT (.reg dTJ)) (.pop dTJ))

/-- One header entry onto `X`. -/
def rdH (X : Fin HS) : RS := .seq (.push X (.reg dTJ)) (.pop dTJ)

/-- Drop `n` entries of `X`. -/
def popL (X C : Fin HS) (n : Ex HS) : RS := .rep C n (.pop X)

/-- Copy the top `n` entries of `X` onto `Y` and `T`. -/
def dupMv (X Y T C : Fin HS) (n : Ex HS) : RS :=
  .rep C n (.seq (.push Y (.reg X)) (.seq (.push T (.reg X)) (.pop X)))

end J6Seg.JK

end

section
-- Source module: Solutions.J6SegILJKer8
/-!
# B-line IL, the `J` kernel (B4), part 8: one backward record

The backward table `flatB k bs = ⋃_j (|cb| :: cb ++ |cz| :: cz ++ [pv, iv, fl])` (with the
reciprocal `iv = ivOf k pv`) lies on `dTJ`. Record `j` computes both dot products against
`Z = (bwdS j).1` and `BZ = (bwdS j).2`, the numerator `bt_j - d₁ - d₂`, its fixed-point quotient
by `pv` (`recip`, or `0` for `pv = 0`), pushes it onto `Z` and, for `fl = 1`, onto `BZ`
(`stepB`): `bwd_body`.
-/

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

namespace J6Seg.JK

open IntMul NLogN NLogN.KIL Function Finset

/-- A backward record with its counts and reciprocal. -/
def recT (k : ℕ) (rc : List ℤ × List ℤ × ℤ × ℤ) : List ℤ :=
  rowT rc.1 ++ (rowT rc.2.1 ++ [rc.2.2.1, ivOf k rc.2.2.1, rc.2.2.2])

/-- The backward table. -/
def flatB (k : ℕ) (bs : List (List ℤ × List ℤ × ℤ × ℤ)) : List ℤ := bs.flatMap (recT k)

/-- The quotient: by the reciprocal for `pv > 0`, else `0`. -/
def divB : RM := .ite (.neg (.reg PV)) recip (.rs (.push Z (.lit 0)))

/-- The border copy. -/
def flB : RM := .ite (.add (.reg FL) (.lit (-1))) (.rs .skip) (.rs (.push BZ (.reg Z)))

/-- One backward record. -/
def bwdB : RM :=
  .seq (.rs (.seq (rdT CNT) (.push SUM (.lit 0)))) (.seq (dotP Z) (.seq (.rs (.seq (.pop CNT)
    (rdT CNT))) (.seq (dotP BZ) (.seq (.rs (.seq (.pop CNT) (.seq (.push NUM (.add (.reg A)
      (.neg (.reg SUM)))) (.seq (.pop SUM) (.seq (rdT PV) (.seq (rdT IV) (rdT FL)))))))
        (.seq divB (.seq flB (.rs (.seq (.pop FL) (.seq (.pop IV) (.seq (.pop PV)
          (.seq (.pop NUM) (.pop A))))))))))))

section bwd
variable (P k : ℕ) (F : Fin HS → List ℤ) (bs : List (List ℤ × List ℤ × ℤ × ℤ)) (bts R T0 : List ℤ)

/-- The state after `j` backward records. -/
def bSt (j : ℕ) : Fin HS → List ℤ :=
  update (update (update (update (update F dTJ (flatB k (bs.drop j) ++ R)) TT
    ((flatB k (bs.take j)).reverse ++ T0)) A (bts.drop j)) Z (bwdS P bs bts j).1) BZ
      (bwdS P bs bts j).2

/-- The frame of the backward loop. -/
structure BwdF (lp lp1 lk : List ℤ) : Prop where
  cnt : F CNT = []
  sum : F SUM = []
  px : F PX = []
  cd : F CD = []
  tmp : F TMP = []
  cr : F CR = []
  num : F NUM = []
  pv : F PV = []
  iv : F IV = []
  fl : F FL = []
  nn : F NN = []
  qq : F QQ = []
  dd : F DD = []
  q0 : F Q0 = []
  rp : F rP = 2 ^ P :: lp
  rp1 : F rP1 = 2 ^ (P + 1) :: lp1
  rk : F rK = 2 ^ k :: lk

/-- The data of the backward loop: shapes and a-priori bounds. -/
structure BwdD (Xb Bc : ℕ) : Prop where
  len : bs.length = bts.length
  wb : ∀ j (h : j < bs.length), bs[j].1.length ≤ j
  wz : ∀ j (h : j < bs.length), bs[j].2.1.length ≤ (bwdS P bs bts j).2.length
  pv0 : ∀ j (h : j < bs.length), 0 ≤ bs[j].2.2.1
  fl01 : ∀ j (h : j < bs.length), bs[j].2.2.2 = 0 ∨ bs[j].2.2.2 = 1
  bcb : ∀ j (h : j < bs.length), Bnd Xb bs[j].1
  bcz : ∀ j (h : j < bs.length), Bnd Xb bs[j].2.1
  bpv : ∀ j (h : j < bs.length), zsz bs[j].2.2.1 ≤ Xb
  bbt : Bnd Xb bts
  bz : ∀ j ≤ bs.length, Bnd Xb (bwdS P bs bts j).1
  bb : ∀ j ≤ bs.length, Bnd Xb (bwdS P bs bts j).2
  mB : bs.length < 2 ^ Bc

variable {P k F bs bts R T0}

end bwd

end J6Seg.JK

end

section
-- Source module: Solutions.J6SegILJKer9
/-!
# B-line IL, the `J` kernel (B4), part 9: the backward loop

`bwdP` runs `bwdB` for the `m = top HM` records: from `A = bts`, `Z = BZ = []` and the table
`flatB k bs` on `dTJ` it leaves `Z = (bwdS P bs bts m).1 = bwdK P bs bts [] []`
(`bwdK_eq_bwdS`), `BZ = (bwdS m).2`, `A = []`, the table moved onto `TT` (`bwdP_spec`).
-/

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

namespace J6Seg.JK

open IntMul NLogN NLogN.KIL Function Finset

/-- The backward loop. -/
def bwdP : RM := .rep CR (.reg HM) bwdB

/-- The record lengths. -/
def recL (bs : List (List ℤ × List ℤ × ℤ × ℤ)) (j : ℕ) : ℕ :=
  (bs.getD j ([], [], 0, 0)).1.length + (bs.getD j ([], [], 0, 0)).2.1.length

end J6Seg.JK

end

section
-- Source module: Solutions.J6SegILJKer10
/-!
# B-line IL, the `J` kernel (B4), part 10: the cyclic extension

`zextP` builds `E = zext K zl = zl.drop (m - K) ++ zl ++ zl.take K` from `Z = zl` (`m = |zl|`
on `HM`, `K ≤ m` on `HK`), leaving `Z` as it was: the reversed extension is assembled on `ER` by
three copies (`dupMv`) and moved onto `E` (`zextP_spec`). Linear cost.
-/

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

namespace J6Seg.JK

open IntMul NLogN NLogN.KIL Function Finset

abbrev TM2 : Fin HS := 29
abbrev ER : Fin HS := 30
abbrev E : Fin HS := 31
abbrev HK : Fin HS := 25

/-- The counts. -/
def cK : Ex HS := .reg HK
def cMK : Ex HS := .add (.reg HM) (.neg (.reg HK))
def cM : Ex HS := .reg HM
def cM2K : Ex HS := .add (.reg HM) (.add (.reg HK) (.reg HK))

/-- The cyclic extension of `Z` onto `E`. -/
def zextP : RS :=
  .seq (mvK Z TMP CM cMK) (.seq (dupMv Z ER TM2 CM cK) (.seq (mvK TM2 Z CM cK)
    (.seq (mvK TMP Z CM cMK) (.seq (dupMv Z ER TMP CM cM) (.seq (mvK TMP Z CM cM)
      (.seq (dupMv Z ER TMP CM cK) (.seq (mvK TMP Z CM cK) (mvK ER E CM cM2K))))))))

/-- The states of `zextP`. -/
def zS (G : Fin HS → List ℤ) (zv tm tm2 er : List ℤ) : Fin HS → List ℤ :=
  update (update (update (update G Z zv) TMP tm) TM2 tm2) ER er

end J6Seg.JK

end

section
-- Source module: Solutions.J6SegILJKer11
/-!
# B-line IL, the `J` kernel (B4), part 11: the coupling window

Row `r` of `cplP` reads `|cc_r|`, runs the dot loop on the window `E = ze.drop r` and pushes
`W_r = dotF P cc_r (ze.drop r)` onto `W`, then drops the head of `E` (`cpl_body`). After
`m = top HM` rows `W = (W_0, …, W_{m-1}).reverse` (`cplP_spec`).
-/

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

namespace J6Seg.JK

open IntMul NLogN NLogN.KIL Function Finset

abbrev W : Fin HS := 47

/-- One coupling row. -/
def cplB : RM :=
  .seq (.rs (.seq (rdT CNT) (.push SUM (.lit 0)))) (.seq (dotP E) (.rs (.seq (.push W (.reg SUM))
    (.seq (.pop SUM) (.seq (.pop CNT) (.pop E))))))

/-- The coupling loop. -/
def cplP : RM := .rep CR (.reg HM) cplB

/-- The window products. -/
def wF (P : ℕ) (cc : List (List ℤ)) (ze : List ℤ) (r : ℕ) : ℤ := dotF P (cc.getD r []) (ze.drop r)

section cpl
variable (P : ℕ) (F : Fin HS → List ℤ) (cc : List (List ℤ)) (ze R T0 : List ℤ)

/-- The state after `r` coupling rows. -/
def cSt (r : ℕ) : Fin HS → List ℤ :=
  update (update (update (update F dTJ (flatT (cc.drop r) ++ R)) TT
    ((flatT (cc.take r)).reverse ++ T0)) E (ze.drop r)) W ((List.range r).map (wF P cc ze)).reverse

/-- The frame of the coupling loop. -/
structure CplF (lp lp1 : List ℤ) : Prop where
  cnt : F CNT = []
  sum : F SUM = []
  px : F PX = []
  cd : F CD = []
  tmp : F TMP = []
  cr : F CR = []
  rp : F rP = 2 ^ P :: lp
  rp1 : F rP1 = 2 ^ (P + 1) :: lp1

variable {P F cc ze R T0}

end cpl

end J6Seg.JK

end

section
-- Source module: Solutions.J6SegILJKer12
/-!
# B-line IL, the `J` kernel (B4), part 12: the Gaussian row in and out

* `unzP`: the row `gflat x` on stack `0` is split into its real and imaginary parts
  (`RE = (x.map fst).reverse`, `IM = (x.map snd).reverse`);
* `zerP`: `G - m` zero pairs onto stack `0`;
* `outP`: the outputs `WR = a.reverse`, `WI = b.reverse` (`|a| = |b| = m`) interleaved onto
  stack `0`: `gflat (a.zip b)`.
-/

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

namespace J6Seg.JK

open IntMul NLogN NLogN.KIL NLogN.KFB Function Finset

abbrev HG : Fin HS := 24
abbrev WI : Fin HS := 15
abbrev WR : Fin HS := 16
abbrev IM : Fin HS := 17
abbrev RE : Fin HS := 18

/-- Split the row. -/
def unzP : RS := .rep CR (.reg HG) (.seq (.push RE (.reg 0)) (.seq (.pop 0)
  (.seq (.push IM (.reg 0)) (.pop 0))))

/-- `G - m` zero pairs. -/
def zerP : RS := .rep CR (.add (.reg HG) (.neg (.reg HM))) (.seq (.push 0 (.lit 0)) (.push 0 (.lit 0)))

/-- Interleave the outputs. -/
def outP : RS := .rep CR (.reg HM) (.seq (.push 0 (.reg WI)) (.seq (.pop WI)
  (.seq (.push 0 (.reg WR)) (.pop WR))))

section io
variable {s : IState HS} {F : Fin HS → List ℤ}

/-- A list of Gaussian integers is bounded. -/
def GBnd (B : ℕ) (x : List GI) : Prop := Bnd B (gflat x)

end io

end J6Seg.JK

end

section
-- Source module: Solutions.J6SegILJKer13
/-!
# B-line IL, the `J` kernel (B4), part 13: one real part

`partP`: from the row `v` (length `G`) on `V` and the table body
`flatT fs ++ flatB k bs ++ flatT cc` on `dTJ`, the forward loop, the backward loop, the cyclic
extension and the coupling window leave `W = Wl.reverse` with
`Wl = (range m).map (wF P cc (zext K (bwdS P bs (fA P fs v m) m).1))`; every other stack is as
before (`V` emptied), the table is restored from `TT` (`partP_spec`).
-/

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

namespace J6Seg.JK

open IntMul NLogN NLogN.KIL NLogN.KFB Function Finset

abbrev HNB : Fin HS := 23
abbrev HL : Fin HS := 14

def cGM : Ex HS := .add (.reg HG) (.neg (.reg HM))
def cK2 : Ex HS := .add (.reg HK) (.reg HK)

/-- One real part of the kernel. -/
def partP : RM :=
  .seq fwdP (.seq (.rs (popL V CM cGM)) (.seq bwdP (.seq (.rs zextP) (.seq cplP
    (.rs (.seq (popL E CM cK2) (.seq (popL Z CM cM) (.seq (popL BZ CM (.reg HNB))
      (mvK TT dTJ CM (.reg HL))))))))))

/-- The table body. -/
def jBody (k : ℕ) (fs : List (List ℤ)) (bs : List (List ℤ × List ℤ × ℤ × ℤ))
    (cc : List (List ℤ)) : List ℤ := flatT fs ++ (flatB k bs ++ flatT cc)

/-- The number of border records. -/
def nBd (bs : List (List ℤ × List ℤ × ℤ × ℤ)) : ℕ := (bs.filter (fun rc => rc.2.2.2 = 1)).length

/-- The output of one part. -/
def wOut (P k K : ℕ) (fs : List (List ℤ)) (bs : List (List ℤ × List ℤ × ℤ × ℤ))
    (cc : List (List ℤ)) (v : List ℤ) : List ℤ :=
  (List.range fs.length).map (wF P cc (zext K (bwdS P bs (fA P fs v fs.length) fs.length).1))

/-- The frame of a part: scratch empty, the header registers. -/
structure PartF (P k K G m nb L : ℕ) (F : Fin HS → List ℤ) : Prop where
  cnt : F CNT = []
  sum : F SUM = []
  px : F PX = []
  cd : F CD = []
  tmp : F TMP = []
  cr : F CR = []
  cm : F CM = []
  num : F NUM = []
  pv : F PV = []
  iv : F IV = []
  fl : F FL = []
  nn : F NN = []
  qq : F QQ = []
  dd : F DD = []
  q0 : F Q0 = []
  tm2 : F TM2 = []
  er : F ER = []
  a : F A = []
  z : F Z = []
  bz : F BZ = []
  e : F E = []
  w : F W = []
  tt : F TT = []
  rp : ∃ l, F rP = 2 ^ P :: l
  rp1 : ∃ l, F rP1 = 2 ^ (P + 1) :: l
  rk : ∃ l, F rK = 2 ^ k :: l
  hm : ∃ l, F HM = (m : ℤ) :: l
  hk : ∃ l, F HK = (K : ℤ) :: l
  hg : ∃ l, F HG = (G : ℤ) :: l
  hnb : ∃ l, F HNB = (nb : ℤ) :: l
  hl : ∃ l, F HL = (L : ℤ) :: l

/-- The shapes and bounds of a part. -/
structure PartD (P k K G Xb Bc : ℕ) (fs : List (List ℤ)) (bs : List (List ℤ × List ℤ × ℤ × ℤ))
    (cc : List (List ℤ)) (v : List ℤ) : Prop where
  lv : v.length = G
  lb : bs.length = fs.length
  lc : cc.length = fs.length
  mG : fs.length ≤ G
  Km : K ≤ fs.length
  wf : ∀ r (h : r < fs.length), fs[r].length ≤ r
  wc : ∀ r (h : r < cc.length), cc[r].length + r ≤ fs.length + 2 * K
  bf : ∀ r (h : r < fs.length), Bnd Xb fs[r]
  bc : ∀ r (h : r < cc.length), Bnd Xb cc[r]
  bv : Bnd Xb v
  bA : ∀ r ≤ fs.length, Bnd Xb (fA P fs v r)
  bwd : BwdD P bs (fA P fs v fs.length) Xb Bc
  cnts : G + 2 * fs.length + 2 * K + (jBody k fs bs cc).length < 2 ^ Bc
  kB : 2 * Xb + 2 * P + Bc + 6 ≤ k

/-- The cost of one part. -/
def partC (k K G : ℕ) (fs : List (List ℤ)) (bs : List (List ℤ × List ℤ × ℤ × ℤ))
    (cc : List (List ℤ)) (u : ℕ) : ℕ :=
  (∑ r ∈ range fs.length, ((fs.getD r []).length + 3) + ∑ j ∈ range bs.length, (recL bs j + 6) +
    ∑ r ∈ range cc.length, ((cc.getD r []).length + 3) + G + 3 * fs.length + 3 * K +
      (jBody k fs bs cc).length + 13) * u

end J6Seg.JK

end

section
-- Source module: Solutions.J6SegILJKer14
/-!
# B-line IL, the `J` kernel (B4), part 14: the kernel states

`kS B z0 re im v w tmp wr wi`: the managed state `B` (registers loaded, table body on `dTJ`)
with the stacks `0, RE, IM, V, W, TMP, WR, WI` replaced; lookup and update lemmas.
The header moves `hdrIn` / `hdrOut`.
-/

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

namespace J6Seg.JK

open IntMul NLogN NLogN.KIL NLogN.KFB Function Finset

/-- The kernel states. -/
def kS (B : Fin HS → List ℤ) (z0 re im v w tmp wr wi : List ℤ) : Fin HS → List ℤ :=
  update (update (update (update (update (update (update (update B 0 z0) RE re) IM im) V v) W w)
    TMP tmp) WR wr) WI wi

section ks
variable {B : Fin HS → List ℤ} {z0 re im v w tmp wr wi : List ℤ}

@[simp] theorem kS_0 : kS B z0 re im v w tmp wr wi 0 = z0 := by simp +decide [kS, update_apply]
@[simp] theorem kS_RE : kS B z0 re im v w tmp wr wi RE = re := by simp +decide [kS, update_apply]
@[simp] theorem kS_IM : kS B z0 re im v w tmp wr wi IM = im := by simp +decide [kS, update_apply]
@[simp] theorem kS_V : kS B z0 re im v w tmp wr wi V = v := by simp +decide [kS, update_apply]
@[simp] theorem kS_W : kS B z0 re im v w tmp wr wi W = w := by simp +decide [kS, update_apply]
@[simp] theorem kS_TMP : kS B z0 re im v w tmp wr wi TMP = tmp := by simp +decide [kS, update_apply]
@[simp] theorem kS_WR : kS B z0 re im v w tmp wr wi WR = wr := by simp +decide [kS, update_apply]
@[simp] theorem kS_WI : kS B z0 re im v w tmp wr wi WI = wi := by simp [kS]

end ks

/-! ## The header -/

/-- The header registers, in table order. -/
def hdrRegs : List (Fin HS) := [rP, rP1, rK, HM, HK, HG, HNB, HL]

/-- Read the header. -/
def hdrIn : RS := .seq (rdH rP) (.seq (rdH rP1) (.seq (rdH rK) (.seq (rdH HM) (.seq (rdH HK)
  (.seq (rdH HG) (.seq (rdH HNB) (rdH HL)))))))

/-- Write one header entry back. -/
def wrH (X : Fin HS) : RS := .seq (.push dTJ (.reg X)) (.pop X)

/-- Write the header back. -/
def hdrOut : RS := .seq (wrH HL) (.seq (wrH HNB) (.seq (wrH HG) (.seq (wrH HK) (.seq (wrH HM)
  (.seq (wrH rK) (.seq (wrH rP1) (wrH rP)))))))

/-- The registers loaded with the header `h` (`|h| = 8`), the body on `dTJ`. -/
def hdrSt (F : Fin HS → List ℤ) (h body : List ℤ) : Fin HS → List ℤ :=
  update (update (update (update (update (update (update (update (update F rP [h.getD 0 0])
    rP1 [h.getD 1 0]) rK [h.getD 2 0]) HM [h.getD 3 0]) HK [h.getD 4 0]) HG [h.getD 5 0])
      HNB [h.getD 6 0]) HL [h.getD 7 0]) dTJ body

end J6Seg.JK

end

section
-- Source module: Solutions.J6SegILJKer15
/-!
# B-line IL, the `J` kernel (B4), part 15: the generic kernel as a register program

`jKprog`: read the header, split the Gaussian row, run `partP` on the real and on the imaginary
part, interleave the outputs with the `G - m` zero pairs, write the header back.
`jKprog_spec`: on the managed state with the table `jTab` on `dTJ` and the row `gflat x ++ l` on
stack `0` (every other managed stack empty) it leaves `gflat (liftG (jKerG P K G m fs bs cc) x) ++ l`
on stack `0` and everything else as it was, at cost `jC`.
-/

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

namespace J6Seg.JK

open IntMul NLogN NLogN.KIL NLogN.KFB Function Finset

/-- The header. -/
def jHdr (P k K G : ℕ) (fs : List (List ℤ)) (bs : List (List ℤ × List ℤ × ℤ × ℤ))
    (cc : List (List ℤ)) : List ℤ :=
  [2 ^ P, 2 ^ (P + 1), 2 ^ k, (fs.length : ℤ), (K : ℤ), (G : ℤ), (nBd bs : ℤ),
    ((jBody k fs bs cc).length : ℤ)]

/-- **The `J`-kernel table.** -/
def jTab (P k K G : ℕ) (fs : List (List ℤ)) (bs : List (List ℤ × List ℤ × ℤ × ℤ))
    (cc : List (List ℤ)) : List ℤ := jHdr P k K G fs bs cc ++ jBody k fs bs cc

/-- Move the output of a part onto `O`. -/
def outMv (O : Fin HS) : RS := .seq (mvK W TMP CM (.reg HM)) (mvK TMP O CM (.reg HM))

/-- **The `J` kernel.** -/
def jKprog : RM :=
  .seq (.rs hdrIn) (.seq (.rs unzP) (.seq (.rs (mvK RE V CM (.reg HG))) (.seq partP
    (.seq (.rs (outMv WR)) (.seq (.rs (mvK IM V CM (.reg HG))) (.seq partP
      (.seq (.rs (outMv WI)) (.rs (.seq zerP (.seq outP hdrOut))))))))))

/-- The cost of the kernel. -/
def jC (k K G : ℕ) (fs : List (List ℤ)) (bs : List (List ℤ × List ℤ × ℤ × ℤ))
    (cc : List (List ℤ)) (u : ℕ) : ℕ :=
  2 * partC k K G fs bs cc u + (3 * G + 4 * fs.length + 40) * u

/-- The precondition on a row: its length, and the shapes and a-priori bounds of both parts. -/
def JPre (P k K G Xb Bc : ℕ) (fs : List (List ℤ)) (bs : List (List ℤ × List ℤ × ℤ × ℤ))
    (cc : List (List ℤ)) (x : List GI) : Prop :=
  x.length = G ∧ PartD P k K G Xb Bc fs bs cc (x.map Prod.fst) ∧
    PartD P k K G Xb Bc fs bs cc (x.map Prod.snd)

/-! ## The list side -/

end J6Seg.JK

end

section
-- Source module: Solutions.J6SegILJKer16
/-!
# B-line IL, the `J` kernel (B4), part 16: the kernel run

`jKprog_spec`: the run of `jKprog` on a managed state (table on `dTJ`, row on `0`, everything
else empty). `partF_kS`: the frame of a part inside the kernel.
-/

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

namespace J6Seg.JK

open IntMul NLogN NLogN.KIL NLogN.KFB Function Finset

section run
variable {P k K G Xb Bc : ℕ} {fs : List (List ℤ)} {bs : List (List ℤ × List ℤ × ℤ × ℤ)}
  {cc : List (List ℤ)}

/-- The base state: header registers loaded, body on `dTJ`. -/
abbrev jB (F0 : Fin HS → List ℤ) := hdrSt F0 (jHdr P k K G fs bs cc) (jBody k fs bs cc)

end run

end J6Seg.JK

end

section
-- Source module: Solutions.J6SegILJKer18
/-!
# B-line IL, the `J` kernel (B4), part 18: the `RowBody6` statements

`jK_row`: if procedure `e` of the table is the compiled `jKprog`, it is a `RowBody6` row kernel
for the generic `J` map `jKerG` with the table `jTab` on `dTJ`, at cost `jC + 2`.

`jKer_row`: the same for the `J` kernel of the segmented line (`jKer P R s t λ w α K G`), with
the streams `jFs`, `jBs`, `jCc` of the line.
-/

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

namespace J6Seg.JK

open IntMul NLogN NLogN.KIL NLogN.KFB Function Finset

section line
variable (P R s t lam w α K : ℕ)

/-- The forward streams of the line. -/
def jFs : List (List ℤ) :=
  fsOf P (LJ s t lam w) (lam + w) (s - lam - w) (capI P (J6Seg.TT P R s t α) (BB s t lam w))

/-- The backward records of the line. -/
def jBs : List (List ℤ × List ℤ × ℤ × ℤ) :=
  bsOf P (LJ s t lam w) (lam + w) (s - lam - w) (capI P (J6Seg.TT P R s t α) (BB s t lam w))

/-- The coupling windows of the line. -/
def jCc : List (List ℤ) := (List.range (LJ s t lam w).length).map (ccL P R s t lam w α K)

end line

end J6Seg.JK

end

section
-- Source module: Solutions.J6SegPreIface
/-!
# B5 contract (`PreU6`): the kernel-pass set-up from the axis descriptor

The axis program (`AxisU6`, S2) runs the two kernel passes of `axisPipe` with the wrapper
`kerW ρ dT dF eF eB e` (`J6SegILKer4`, run lemma `kerW_run`). Before each such pass the registers
`ρ 0 … ρ 10` must hold (`kerW_run`'s `h0 … h10`)

    ρ0 = [w]  ρ1 = [w']  ρ2 = [f]  ρ3 = [pre·nc]  ρ4 = [post]  ρ5 = [pre]  ρ6 = [nc]
    ρ7 = [2·2^f]  ρ8 = [2·(pre·nc·post·2^f)]  ρ9 = [|T|]  ρ10 = T,  T = tabsT tabOf nc

(`ρ 11 … ρ 18` empty). `PreU6` asks for procedures (uniform in the table, as `FineU6`) that,
**from the axis descriptor `descAx6 P n src D k` on `dA`** and nothing else, build exactly this
register file, and procedures that clear it again:

* the **chunk pass** (`axisPipe` passes 2 and 6): `f = log₂ λ`, `pre = ∏_{j<k} t_j`,
  `nc = t_k / λ`, `post = ∏_{j>k} t_j`, `tabOf = chunkTab6 … q` (the table `tabK` of `J6SegKr8`
  read by the chunk kernel `chunkI`), widths `w = w' = wd6`;
* the **`J` pass** (pass 4): `f = log₂ G` (`G = Gj6`), the same `pre`, `post`, `nc = 1`,
  `tabOf = fun _ => jTab6 …` (the table `jTab` of `J6SegILJKer15` read by `jKprog`), widths
  `w = w' = wd6`.

All line parameters are functions of the descriptor entries `s = src k`, `t = t_k`, `p`, `α`,
`d` (and the fixed `P : Prm`):

    Q = segPF p d s     (the Kronecker / fixed-point precision, `P` of `chunkK`, `jKer`)
    R = p / 7,  λ = lamP d,  w = K = wF d,  nc = t / λ,  G = 2^⌈log₂ |LJ s t λ w|⌉

and the **plan** `pl : Plan6` fixes the free precisions as affine functions of `Q`:
`pk = pkA·Q + pkC` (stage bound `2^pk`), `b = bA·Q + bC` (Kronecker digit width),
`kJ = kA·Q + kC` (the `J` kernel's reciprocal precision), `wd = wA·Q + wC` (array width).

The slab control lists (passes 1, 3, 5, 7) are **not** part of this contract (B1).

The table-size facts the assembly needs are separate pure statements: `ChunkTabBnd6`
(all chunk-table entries `≤ 2^(Q+R+2)` in modulus).
-/

set_option linter.unusedSectionVars false
set_option linter.unusedVariables false

noncomputable section

namespace J6Seg
open IntMul NLogN NLogN.KIL NLogN.KFB IntMul.Kappa.R6 Function

/-! ## The plan and the line parameters -/

/-- The free precisions, as affine functions `a·Q + c` of `Q = segPF p d s`. -/
structure Plan6 where
  pkA : ℕ
  pkC : ℕ
  bA : ℕ
  bC : ℕ
  kA : ℕ
  kC : ℕ
  wA : ℕ
  wC : ℕ

section par
variable (p d s t : ℕ)

/-- The precision `Q` of the line kernels. -/
def Q6 : ℕ := segPF p d s
/-- `R = p / 7`. -/
def R6 : ℕ := p / 7
/-- The padded length of the near-boundary line (a power of two). -/
def Gj6 : ℕ := 2 ^ Nat.clog 2 (LJ s t (lamP d) (wF d)).length

def pk6 (pl : Plan6) : ℕ := pl.pkA * Q6 p d s + pl.pkC
def b6 (pl : Plan6) : ℕ := pl.bA * Q6 p d s + pl.bC
def kJ6 (pl : Plan6) : ℕ := pl.kA * Q6 p d s + pl.kC
def wd6 (pl : Plan6) : ℕ := pl.wA * Q6 p d s + pl.wC

variable (α : ℕ) (pl : Plan6)

/-- **The tables of chunk `q`** (read by `chunkI2` on `dT = 49`; layout `KR.tabK2`, `J6SegKs0`). -/
def chunkTab6 (q : ℕ) : List ℤ :=
  KR.tabK2 (Q6 p d s) (pk6 p d s pl) (b6 p d s pl) (lamP d)
    (diL (Q6 p d s) (R6 p) s t (lamP d) α) (giL (Q6 p d s) (R6 p) s t (lamP d) α)
    (gamL (Q6 p d s) (R6 p) s t (lamP d) α) (ddL (Q6 p d s) (R6 p) s t (lamP d) α)
    (descsOf s t (lamP d) (TT (Q6 p d s) (R6 p) s t α) q)

/-- **The table of the `J` kernel** (read by `jKprog` on `dTJ = 33`). -/
def jTab6 : List ℤ :=
  JK.jTab (Q6 p d s) (kJ6 p d s pl) (wF d) (Gj6 d s t)
    (JK.jFs (Q6 p d s) (R6 p) s t (lamP d) (wF d) α)
    (JK.jBs (Q6 p d s) (R6 p) s t (lamP d) (wF d) α)
    (JK.jCc (Q6 p d s) (R6 p) s t (lamP d) (wF d) α (wF d))

end par

/-! ## The register file of a kernel pass -/

/-- The register file `ρ 0 … ρ 18` of `kerW` for a pass with widths `w, w'`, rows `2^f`,
`pre, nc, post` and tables `T` (`ρ 11 … ρ 18` empty). -/
def kerRegs (w w' f pre nc post : ℕ) (T : List ℤ) : Fin 19 → List ℤ := fun i =>
  if i.val = 0 then [(w : ℤ)] else if i.val = 1 then [(w' : ℤ)] else
  if i.val = 2 then [(f : ℤ)] else if i.val = 3 then [((pre * nc : ℕ) : ℤ)] else
  if i.val = 4 then [(post : ℤ)] else if i.val = 5 then [(pre : ℤ)] else
  if i.val = 6 then [(nc : ℤ)] else if i.val = 7 then [((2 * 2 ^ f : ℕ) : ℤ)] else
  if i.val = 8 then [((2 * (pre * nc * post * 2 ^ f) : ℕ) : ℤ)] else
  if i.val = 9 then [(T.length : ℤ)] else if i.val = 10 then T else []

section regs
variable (P : Prm) (n : ℕ) (src : ℕ → ℕ) (D k : ℕ) (pl : Plan6)

/-- The registers of the **chunk pass** of axis `k`. -/
def chunkRegs6 : Fin 19 → List ℤ :=
  let p := pP6 P n; let d := dP6 P n; let s := src k; let t := tF6 P n k
  kerRegs (wd6 p d s pl) (wd6 p d s pl) (Nat.log 2 (lamF d)) (preT P n D k) (t / lamP d)
    (postT P n D k) (tabsT (chunkTab6 p d s t (αN6 P n) pl) (t / lamP d))

/-- The registers of the **`J` pass** of axis `k`. -/
def jRegs6 : Fin 19 → List ℤ :=
  let p := pP6 P n; let d := dP6 P n; let s := src k; let t := tF6 P n k
  kerRegs (wd6 p d s pl) (wd6 p d s pl) (Nat.clog 2 (LJ s t (lamP d) (wF d)).length)
    (preT P n D k) 1 (postT P n D k) (tabsT (fun _ => jTab6 p d s t (αN6 P n) pl) 1)

end regs

/-! ## The contract -/

/-- One set-up / tear-down pair: procedure `eUp` builds the register file `V` from the
descriptor on `dA` (registers empty before), `eDn` clears it; every other stack is unchanged,
the base state `s` is preserved, each at cost `≤ C`. -/
def SetupBody6 {Pp : ℕ} (Isym : List Sym) (tbl : Fin Pp → IStmt HS Pp) (eUp eDn : Fin Pp)
    (ρ : Fin 19 → Fin HS) (dA : Fin HS) (desc : List ℤ) (V : Fin 19 → List ℤ) (C : ℝ) : Prop :=
  ∀ (s : IState HS) (F : Fin HS → List ℤ),
    (∀ j, j ≠ dA → (∀ i, ρ i ≠ j) → F j = []) → (∀ i, F (ρ i) = []) → F dA = desc →
    (∃ c : ℕ, IRuns Isym tbl (.call eUp) (encK s F) (encK s (extend ρ V F)) c ∧ (c : ℝ) ≤ C) ∧
    (∃ c : ℕ, IRuns Isym tbl (.call eDn) (encK s (extend ρ V F)) (encK s F) c ∧ (c : ℝ) ≤ C)

/-- **`PreU6 P pl tab` (B5).** For every admissible register map `ρ` and descriptor stack `dA`
there is a family `U` of procedures, uniform in the table (as `FineU6`), with two set-up /
tear-down pairs: `(eC, eCd)` for the chunk pass and `(eJ, eJd)` for the `J` pass of axis `k`,
reading only the descriptor `descAx6 P n src D k` on `dA`, each at cost
`Cp · cU6 P n A / (D + 1)`. -/
def PreU6 (P : Prm) (pl : Plan6) (tab : Fin 1 → IStmt HS 1) : Prop :=
  ∀ (ρ : Fin 19 → Fin HS) (dA : Fin HS), Injective ρ → (∀ i, ρ i ≠ dA) → (∀ i, ρ i ≠ 0) →
    dA ≠ 0 → (∀ i, (ρ i).val < 56) → dA.val < 56 →
  ∃ (nU : ℕ) (U : ∀ E : ℕ, (Fin nU → Fin (7 + E)) → Fin nU → IStmt HS (7 + E))
    (eC eCd eJ eJd : Fin nU) (Cp : ℝ) (A n₁ : ℕ), 0 ≤ Cp ∧
  ∀ n, n₁ ≤ n → ∀ src : ℕ → ℕ, SrcOK6 P n src → ∀ D : ℕ, dP6 P n = D + 1 → ∀ k ≤ D,
  ∀ (E : ℕ) (ext : Fin E → IStmt HS (7 + E)) (idx : Fin nU → Fin (7 + E)),
    (∀ i, IL.kTableP tab ext (idx i) = U E idx i) →
    ∀ Isym : List Sym,
      SetupBody6 Isym (IL.kTableP tab ext) (idx eC) (idx eCd) ρ dA (descAx6 P n src D k)
        (chunkRegs6 P n src D k pl) (Cp * cU6 P n A / (D + 1)) ∧
      SetupBody6 Isym (IL.kTableP tab ext) (idx eJ) (idx eJd) ρ dA (descAx6 P n src D k)
        (jRegs6 P n src D k pl) (Cp * cU6 P n A / (D + 1))

/-- What B5 proves (in the round-six regime, from the table facts `LineNeed6` provides). -/
def PreNeed6 (P : Prm) (pl : Plan6) : Prop :=
  ∀ tab : Fin 1 → IStmt HS 1, IL.SwapNeedτ tab P.τ → IL2.RowUCP 576 P.τ tab → FineU6 tab →
    PreU6 P pl tab

/-! ## The table-size facts (pure) -/

/-- **The chunk-table entries are `≤ 2^(Q+R+2)`** in modulus (`Kt = 2^(R+2)` in
`J6SegKr9.chunk_rowBody6`), for every axis of the round-six regime. -/
def ChunkTabBnd6 (P : Prm) : Prop :=
  ∃ n₁ : ℕ, ∀ n, n₁ ≤ n → ∀ src : ℕ → ℕ, SrcOK6 P n src → ∀ D : ℕ, dP6 P n = D + 1 →
  ∀ k ≤ D,
    let p := pP6 P n; let d := dP6 P n; let s := src k; let t := tF6 P n k
    let T := TT (Q6 p d s) (R6 p) s t (αN6 P n)
    let A : ℤ := 2 ^ (Q6 p d s + R6 p + 2)
    ∀ u < lamP d, |T.di u| ≤ A ∧ |T.gi u| ≤ A ∧ |T.gam u| ≤ A ∧ |T.dd u| ≤ A ∧
      ∀ c, |T.am c u| ≤ A ∧ |T.ap c u| ≤ A

end J6Seg

end

end

section
-- Source module: Solutions.J6SegAxis2
/-!
# B-line axis program (`AxisU6`), part 2: a kernel pass as a `PassOK` step

`kpSt eUp eDn ρ dT dF eF eB e = call eUp; kerW ρ dT dF eF eB e; call eDn`: the set-up of the
register file `kerRegs …` from the descriptor (`SetupBody6`, B5), the kernel-pass wrapper
(`kerW_run`, B2) with a row kernel (`RowBody6`), the tear-down.

**`kpass_ok`**: under the hypotheses of `kerW_run` (the row facts given by `lineAt`, the rows of
the array), the step is `PassOK` for `kerPass (2^f) post (pre·nc) (fun b => g (b % nc))`, at
cost `2 Cp + (the kerW_run bound)`.
-/

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

noncomputable section

namespace J6Seg.Ax
open IntMul NLogN NLogN.KIL NLogN.KFB IntMul.Kappa.R6 Function

/-- The kernel-pass step. -/
def kpSt {Pp : ℕ} (eUp eDn : Fin Pp) (ρ : Fin 19 → Fin HS) (dT dF : Fin HS) (eF eB e : Fin Pp) :
    IStmt HS Pp :=
  .seq (.call eUp) (.seq (kerW ρ dT dF eF eB e) (.call eDn))

end J6Seg.Ax

end

end

section
-- Source module: Solutions.J6SegAxis3
/-!
# B-line axis program (`AxisU6`), part 3: the slab passes as `PassOK` steps

The slab passes (B1, `J6SegILSlab8`) read the skeleton registers
`rC = flatC cs`, `rM = [2 post]`, `rNC = [|cs|]`, `rNB = [pre]`, `rTOT = [pre · n' · 2 post]`
(and `rK1 = [2^(P-p)]`, `rK2 = [2^p]` for the output pass). They are built from the descriptor
by a set-up / tear-down pair:

* `SetupS Isym tbl stUp stDn ρ dA desc V C`: from any managed state in which only `0`, `sZ`,
  `dA` and the registers `ρ` may be nonempty (registers empty, `desc` on `dA`), `stUp` writes the
  register file `V` and `stDn` clears it, everything else unchanged, each at cost `≤ C`
  (B1's source-table producers and `ctlB`; `J6SegILSrc*`, `J6SegILCtl*`).
* `slabV cs pre post n' K1 K2`: the register file on `ρS = (rC, rM, rNC, rNB, rTOT, rK1, rK2)`.

`slabId_ok`, `slabSc_ok` (`PassOK`) and `slabOut_ok` (`Pass2OK`): set-up; the B1 pass; tear-down.
-/

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

noncomputable section

namespace J6Seg.Ax
open IntMul NLogN NLogN.KIL NLogN.KFB IntMul.Kappa.R6 Function
open J6Seg.Slab
open NLogN.KFUU (lA lC lD)

/-- **A set-up / tear-down pair** of statements for the register file `V` on `ρ`. -/
def SetupS {Pp m : ℕ} (Isym : List Sym) (tbl : Fin Pp → IStmt HS Pp) (stUp stDn : IStmt HS Pp)
    (ρ : Fin m → Fin HS) (dA : Fin HS) (desc : List ℤ) (V : Fin m → List ℤ) (C : ℝ) : Prop :=
  ∀ (s : IState HS) (F : Fin HS → List ℤ),
    (∀ j, j ≠ 0 → j ≠ sZ → j ≠ dA → (∀ i, ρ i ≠ j) → F j = []) → (∀ i, F (ρ i) = []) →
    F dA = desc →
    (∃ c : ℕ, IRuns Isym tbl stUp (encK s F) (encK s (extend ρ V F)) c ∧ (c : ℝ) ≤ C) ∧
    (∃ c : ℕ, IRuns Isym tbl stDn (encK s (extend ρ V F)) (encK s F) c ∧ (c : ℝ) ≤ C)

/-- The slab registers. -/
def ρS : Fin 7 → Fin HS := ![rC, rM, rNC, rNB, rTOT, rK1, rK2]

/-- The slab register file. -/
def slabV (cs : List (ℤ × ℤ)) (pre post n' : ℕ) (K1 K2 : List ℤ) : Fin 7 → List ℤ :=
  ![flatC cs, [((2 * post : ℕ) : ℤ)], [(cs.length : ℤ)], [(pre : ℤ)],
    [((pre * (n' * (2 * post)) : ℕ) : ℤ)], K1, K2]

/-- The stacks a slab pass touches besides the registers (must be free of data). -/
def SlabFree (dA : Fin HS) : Prop :=
  dA ≠ 0 ∧ dA ≠ sZ ∧ (∀ i, ρS i ≠ dA) ∧ dA ≠ lA ∧ dA ≠ lC ∧ dA ≠ lD ∧ dA ≠ rACC

/-- The slab-pass step. -/
def slSt {Pp : ℕ} (stUp stDn body : IStmt HS Pp) : IStmt HS Pp := .seq stUp (.seq body stDn)

end J6Seg.Ax

end

end

section
-- Source module: Solutions.J6SegKs1
/-!
# B3′, part 1: computing the Kronecker constants of a block in the kernel

`kgRS` (with `j` on `kK`, `m = b·2n` on `kJ`, `2^pk` on `kc`, `2^b` on `kB`, scratch `vX, vY, CN`):
`acc := 2^pk`, `ρ := 2^b`; `j` times `acc := acc + acc·ρ`, `ρ := ρ·ρ` (so `acc = KFV.kG b 2^k 2^pk`,
`ρ = 2^(b 2^k)`); then `kJ := 2^m`, `kK := acc mod 2^m = KFV.kG b (2n) 2^pk`, `kJ := 2^m · 2^m`.
Cost `O(j · b 2^j + pow2Cost m)` — linear in `n b` up to the factor `j = clog₂ 2n`.
-/

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

namespace J6Seg.KR

open IntMul NLogN NLogN.KIL NLogN.KFD NLogN.KFB NLogN.KFV

/-! ## `kG` arithmetic -/

/-! ## `2^m` by `pow2` -/

/-- Cost bound of `i := 2^i`. -/
def cP2 (m B : ℕ) : ℕ :=
  2000 * (B + 2) + 2 * moveCost B + pow2Cost m + 2 * B + 100 + (2 * B + 2) + 2 * moveCost B + 19

/-! ## The constants of one block -/

/-- The doubling body. -/
def kgBody : RS :=
  .seq (.asg vX (.add (.reg vX) (.shl (.reg vX) (.reg vY)))) (.asg vY (.shl (.reg vY) (.reg vY)))

/-- **The block constants.** -/
def kgRS : RS :=
  .seq (.push vX (.reg kc)) <| .seq (.push vY (.reg kB)) <| .seq (.rep CN (.reg kK) kgBody) <|
  .seq (.asg kJ (.pow2 (.reg kJ))) <| .seq (.asg kK (.rem (.reg vX) (.reg kJ))) <|
  .seq (.asg kJ (.shl (.reg kJ) (.reg kJ))) <| .seq (.pop vX) (.pop vY)

/-- Cost of `kgRS`. -/
def cKG (j m B : ℕ) : ℕ :=
  2 * (4000 * (B + 2)) + (2 * 2000 * (B + 2) + j * (20000 * (B + 2) + 8000 * (B + 2) + 30) + 32) +
    cP2 m B + 2 * (4 * 2000 * (B + 2)) + 2 * (2 * B + 2)

end J6Seg.KR

end

section
-- Source module: Solutions.J6SegKs2
/-!
# B3′, part 2: the new block header and block clearing

* `hdrRS2 = hdrRS; kgRS`: reads `[n, 2n, j, 2nb]` and computes `kG b (2n) 2^pk`, `(2^b)^(4n)`;
* `hdrF2`, `popF2`: the header and the clearing on the block frame `BF`, with the small header
  values at `BH` and the two block constants at the per-block bound `BkN b n = 8 n b + 4`.
-/

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

namespace J6Seg.KR

open IntMul NLogN NLogN.KIL NLogN.KFD NLogN.KFB NLogN.KFV

/-- The new header program. -/
def hdrRS2 : RS := .seq hdrRS kgRS

/-- The per-block bound of the block constants. -/
def BkN (b n : ℕ) : ℕ := 4 * (b * (2 * n)) + 4

end J6Seg.KR

end

section
-- Source module: Solutions.J6SegKs3
/-!
# B3′, part 3: one block in IL with in-kernel constants

`blkI2 pM` is `blkI pM` with the header `hdrRS2` (reads `blkHdr2 n b`, computes the constants).
`blkI2_runs`: as `blkI_runs`, with `blkHdr2` and the cost `cBlk2`.
-/

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

namespace J6Seg.KR

open IntMul NLogN NLogN.KIL NLogN.KFD NLogN.KFB NLogN.KFV

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

/-- **One block** (new header). -/
def blkI2 (pM : Fin Pp) : IStmt HS Pp :=
  .seq hdrRS2.code <|
  .seq (dgI pM 0 TDI TU W1 CN CC bN) <|
  .seq (mvP W1 W2 TP CC bN).code <|
  .seq (dgI pM W2 TGI TU W1 CN CC bN) <|
  .seq (tsI pM W1 W2) <|
  .seq (dgI pM W2 TGA TU W1 CN CC bN) <|
  .seq (mvP W1 W2 TP CC bN).code <|
  .seq (tsI pM W2 W1) <|
  .seq (mvP W1 W2 TP CC bN).code <|
  .seq (dgI pM W2 TGI TU W1 CN CC bN) <|
  .seq (mvP W1 W2 TP CC bN).code <|
  .seq (dgI pM W2 TDD TU OUT CN CC bN) <|
  popRS.code

/-- Cost of one block (new header). -/
def cBlk2 (n pk b φD φT B BH : ℕ) : ℕ :=
  cBlk n pk b φD φT B BH + cKG (Nat.clog 2 (2 * n)) (b * (2 * n)) (BkN b n) + 2 * (2 * BkN b n + 2)

end J6Seg.KR

end

section
-- Source module: Solutions.J6SegKs4
/-!
# B3′, part 4: the block loop with in-kernel constants (`loopB2_runs`, cost `cLoop2`).
-/

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

namespace J6Seg.KR

open IntMul NLogN NLogN.KIL NLogN.KFD NLogN.KFB NLogN.KFV

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

/-- Cost of the block loop (new header). -/
def cLoop2 (pk b φD : ℕ) (φT : ℕ → ℕ) (B BH : ℕ) : List (ℕ × List ℤ × List ℤ) → ℕ
  | [] => 30
  | d :: ds => cBlk2 d.1 pk b φD (φT d.1) B BH + cLoop2 pk b φD φT B BH ds + 30

end J6Seg.KR

end

section
-- Source module: Solutions.J6SegKs5
/-!
# B3′, part 5: the chunk kernel with in-kernel constants

`chunkI2 pM` (`hdrK`, the block loop of `blkI2`, `finK`) on the table `tabK2 …` (`J6SegKs0`).
`chunkI2_runs` is `chunkI_runs` for the new layout; the cost `cKer2` uses `cLoop2`.
-/

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

namespace J6Seg.KR

open IntMul NLogN NLogN.KIL NLogN.KFD NLogN.KFB NLogN.KFV

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

/-- **The chunk kernel.** -/
def chunkI2 (pM : Fin Pp) : IStmt HS Pp :=
  .seq hdrK.code (.seq (.loop BC isPos (blkI2 pM)) finK.code)

/-- Cost of the chunk kernel. -/
def cKer2 (pk b lam φD : ℕ) (φT : ℕ → ℕ) (B BH : ℕ) (ds : List (ℕ × List ℤ × List ℤ)) : ℕ :=
  8 * 12000 * (BH + 2) + 4 * cMvP lam BH + 4000 * (BH + 2) + cLoop2 pk b φD φT B BH ds +
    (2 * BH + 2) + 2 * cMvP (2 * (lam - sumN ds)) BH + cMvP lam B + 4 * cMvP lam BH +
    (2 * 2000 * (BH + 2) + (8 + 4 * lam + (dsEnc2 b ds).length) * (14000 * (BH + 2)) + 32) +
    8 * (2 * BH + 2)

end J6Seg.KR

end