Round-six κ development, definitions part 64

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

import Definitions.Def_IntMul_Kappa_R6_Defs63
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_HvdH_M_Defs2
import Definitions.Def_IntMul_Kappa_K_Defs14
import Definitions.Def_IntMul_Kappa_R6_Defs3
import Definitions.Def_IntMul_HvdH_M_Defs3
import Definitions.Def_IntMul_Kappa_K_Defs21
import Definitions.Def_IntMul_Kappa_K_Defs28
import Definitions.Def_IntMul_Kappa_K_Defs29
import Definitions.Def_IntMul_Kappa_K_Defs12
import Definitions.Def_IntMul_HvdH_M_Defs1
import Definitions.Def_IntMul_Kappa_R6_Defs10
import Definitions.Def_IntMul_Kappa_K_Defs10
import Definitions.Def_IntMul_Kappa_K_Defs27

set_option autoImplicit false
set_option relaxedAutoImplicit false
set_option maxHeartbeats 1000000

section
-- Source module: Solutions.J6SegKr2
/-!
# B3 (chunk kernel), part 2: register loops on managed states

Counted register loops (`RS.rep` with the count on top of a register) on `encK` states:

* `mvP X Y T C R`: move `n` Gaussian pairs from `X` to `Y`, keeping each pair's orientation
  (`gflat z ++ lx` on `X` ↦ `gflat z.reverse` on top of `Y`);
* `rdP V O C R`: round pairs (`rs P` on both parts) from a reversed flat list on `V` to `O`
  (`gsw u ++ lv` ↦ `gflat ((u.reverse).map rsG)` on top of `O`);
* `rgP A Y C R`: a real table as Gaussian pairs (`a ++ la` on `A` ↦ `gflat (a.reverse.map (·, 0))`
  on top of `Y`);
* `zeroR`, `dropR`, `dupR`: push zeros, drop, duplicate-move with a register count.
-/

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

/-! ## Flat lists -/

/-- Pairs flattened with the parts swapped. -/
def gsw (u : List GI) : List ℤ := u.flatMap fun z => [z.2, z.1]

@[simp] lemma gsw_nil : gsw [] = [] := rfl
@[simp] lemma gsw_cons (z : GI) (u : List GI) : gsw (z :: u) = z.2 :: z.1 :: gsw u := rfl

@[simp] lemma length_gsw (u : List GI) : (gsw u).length = 2 * u.length := by
  induction u with
  | nil => rfl
  | cons z u ih => simp [ih]; ring

/-! ## Single steps -/

/-! ## Moving pairs -/

/-- One pair from `X` to `Y` (through `T`). -/
def mvPB (X Y T : Fin HS) : RS :=
  .seq (.push T (.reg X)) (.seq (.pop X) (.seq (.push Y (.reg X)) (.seq (.pop X)
    (.seq (.push Y (.reg T)) (.pop T)))))

/-- `n` pairs from `X` to `Y` (count on `R`, loop counter `C`). -/
def mvP (X Y T C R : Fin HS) : RS := .rep C (.reg R) (mvPB X Y T)

/-- Cost of `mvP`. -/
def cMvP (n B : ℕ) : ℕ := (n + 1) * 40000 * (B + 2)

/-! ## Zeros, drops, duplicating moves -/

/-- Push `n` zeros on `Y`. -/
def zeroR (Y C R : Fin HS) : RS := .rep C (.reg R) (.push Y (.lit 0))

/-- Drop `n` integers from `X`. -/
def dropR (X C R : Fin HS) : RS := .rep C (.reg R) (.pop X)

/-- Move `n` integers from `X` to both `Y` and `Z` (each reversed). -/
def dupR (X Y Z C R : Fin HS) : RS :=
  .rep C (.reg R) (.seq (.push Y (.reg X)) (.seq (.push Z (.reg X)) (.pop X)))

/-- A real list `a` (reversed on `A`) as Gaussian pairs `(a_i, 0)` on top of `Y`. -/
def rgP (A Y C R : Fin HS) : RS :=
  .rep C (.reg R) (.seq (.push Y (.lit 0)) (.seq (.push Y (.reg A)) (.pop A)))

/-! ## Rounding pairs -/

/-- `2^P`. -/
abbrev rP : Fin HS := 31
/-- `2^{P+1}`. -/
abbrev rP1 : Fin HS := 32

/-- `rs P (top V)`: `⌊(2x + 2^P) / 2^{P+1}⌋` by `tdiv`/`tmod` and a sign test. -/
def rsE (V : Fin HS) : Ex HS :=
  .add (.quo (.add (.add (.reg V) (.reg V)) (.reg rP)) (.reg rP1))
    (.sel (.rem (.add (.add (.reg V) (.reg V)) (.reg rP)) (.reg rP1)) (.lit (-1)) (.lit 0))

/-- Pairs of `rs P` values. -/
def rsG (P : ℕ) (w : GI) : GI := (J6Seg.rs P w.1, J6Seg.rs P w.2)

/-- Round the pairs of a swapped flat list `gsw u` from `V` onto `O`. -/
def rdP (V O C R : Fin HS) : RS :=
  .rep C (.reg R) (.seq (.push O (rsE V)) (.seq (.pop V) (.seq (.push O (rsE V)) (.pop V))))

end J6Seg.KR

end

section
-- Source module: Solutions.J6SegKr3
/-!
# B3 (chunk kernel), part 3: the diagonal stage

`dgI pM X T U Y CN CC N`: with `n` on `N`, the row `gflat z ++ lx` on `X` (`|z| = n`) and the
table `c ++ lt` on `T` (`n ≤ |c|`), it pushes `gflat (dgG P c z).reverse` onto `Y`, consumes the
row and leaves the table in place (the used entries go through `U` and back). Each item costs two
calls of the multiplication procedure (`KFM.mulC`) and two roundings `rs P` (`rsE`).
-/

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 loop test: the counter is positive. -/
def isPos (c : Cls) : Bool := decide (c = .pos)

abbrev dM1 : Fin HS := 12
abbrev dM2 : Fin HS := 13

/-- The register part of one item. -/
def dgRS (X T U Y CN : Fin HS) : RS :=
  .seq (.pop X) (.seq (.push Y (rsE dM2)) (.seq (.push Y (rsE dM1)) (.seq (.pop dM1)
    (.seq (.pop dM2) (.seq (.push U (.reg T)) (.seq (.pop T)
      (.asg CN (.add (.reg CN) (.lit (-1))))))))))

/-- One item. -/
def dgBody (pM : Fin Pp) (X T U Y CN : Fin HS) : IStmt HS Pp :=
  .seq (KFM.mulC pM dM1 X T) (.seq (RS.pop X).code (.seq (KFM.mulC pM dM2 X T)
    (dgRS X T U Y CN).code))

/-- **The diagonal stage.** -/
def dgI (pM : Fin Pp) (X T U Y CN CC N : Fin HS) : IStmt HS Pp :=
  .seq (RS.push CN (.reg N)).code (.seq (.loop CN isPos (dgBody pM X T U Y CN))
    (.seq (RS.pop CN).code (NLogN.mvK U T CC (.reg N)).code))

/-- Cost of the diagonal stage on `n` items. -/
def cDG (n φ B : ℕ) : ℕ := (n + 1) * (2 * φ + 400000 * (B + 2))

end J6Seg.KR

end

section
-- Source module: Solutions.J6SegKr4
/-!
# B3 (chunk kernel), part 4: the Toeplitz stage

`tsI pM IN OUTS`: with `n` on `bN`, `2n` on `bN2`, a Gaussian block `t` (reversed, `gflat t.reverse`)
on `IN`, the real table `am` (length `n`) on top of the table stack `dT`, and the record-product
parameters for records of length `r = 2n` on the parameter stacks `50 … 55` (managed), it leaves
`gflat (lwG P (2n) am t)` on `OUTS`, consumes `am` from `dT` onto the stack `DU` of consumed
descriptors, and uses one uniform Kronecker record product (`recU`, `Δ = 0`).

The parameter stacks are managed lists here; around `recU` they are moved into the base state
(`encK_split`).
-/

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 register file of the chunk kernel -/

/-- The table stack. -/
abbrev dT : Fin HS := 49
/-- Block length `n` and `2n`. -/
abbrev bN : Fin HS := 37
abbrev bN2 : Fin HS := 38
/-- Consumed descriptors (reversed). -/
abbrev DU : Fin HS := 45
/-- Work stacks. -/
abbrev W1 : Fin HS := 47
abbrev W2 : Fin HS := 48
/-- Loop counter of the register loops (`vL`), pair-move scratch (`vT`), table scratch (`vT2`). -/
abbrev CC : Fin HS := 23
abbrev TP : Fin HS := 18
abbrev TA : Fin HS := 19

/-! ## Moving the parameters into the base state -/

def pmask (F : Fin HS → List ℤ) : Fin HS → List ℤ := fun i => if 50 ≤ i.val ∧ i.val ≤ 55 then F i else []
def rmask (F : Fin HS → List ℤ) : Fin HS → List ℤ := fun i => if 50 ≤ i.val ∧ i.val ≤ 55 then [] else F i

/-! ## Entry bound of a record product -/

/-! ## The stage -/

/-- **The Toeplitz stage.** -/
def tsI (pM : Fin Pp) (IN OUTS : Fin HS) : IStmt HS Pp :=
  .seq (zeroR vX CC bN2).code <|
  .seq (mvP IN vX TP CC bN).code <|
  .seq (zeroR vY CC bN2).code <|
  .seq (dupR dT DU TA CC bN).code <|
  .seq (rgP TA vY CC bN).code <|
  .seq (recU pM) <|
  .seq (dropR vO CC bN2).code <|
  (rdP vO OUTS CC bN).code

/-- Cost of the Toeplitz stage. -/
def cTS (n pk b φ B : ℕ) : ℕ :=
  6 * cMvP (2 * n) B + (n + 1) * 200000 * (B + 2) + 2000 * cRec (2 * n) pk b 0 φ

end J6Seg.KR

end

section
-- Source module: Solutions.J6SegKr5
/-!
# B3 (chunk kernel), part 5: one block

`blkI pM`: with a block header `[n, 2n, kG b (2n) 2^pk, (2^b)^(4n)]` and the block tables `am, ap`
on top of the table stack `dT`, the four stage tables on `TDI, TGI, TGA, TDD`, and the block
input `z` (length `n`) on top of the row stack `0`, it pushes `gflat (blkG …).reverse` onto `OUT`,
consumes the header and tables from `dT` onto the consumed-descriptor stack `DU`, and decrements
the block counter `BC`.

The stacks the block body changes are kept in the frame `BF` (twelve stacks over a base `G`).
-/

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}

/-- Output accumulator, block counter, diagonal-stage counter, table scratch. -/
abbrev OUT : Fin HS := 46
abbrev BC : Fin HS := 25
abbrev CN : Fin HS := 24
abbrev TU : Fin HS := 44
/-- The four stage tables. -/
abbrev TDI : Fin HS := 40
abbrev TGI : Fin HS := 41
abbrev TGA : Fin HS := 42
abbrev TDD : Fin HS := 43

/-- **The block frame**: the twelve stacks the block body changes, over a base `G`. -/
def BF (G : Fin HS → List ℤ) (r0 w1 w2 o dt du bn bn2 vr vk vj bc : 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 G 0 r0) W1 w1) W2 w2) OUT o) dT dt) DU du) bN bn) bN2 bn2) kr vr) kK vk) kJ vj) BC bc

section BFl
variable {G : Fin HS → List ℤ} {r0 w1 w2 o dt du bn bn2 vr vk vj bc : List ℤ}

@[simp] theorem BF_r0 : BF G r0 w1 w2 o dt du bn bn2 vr vk vj bc 0 = r0 := by
  unfold BF; rw [Function.update_of_ne (show 0 ≠ BC by decide), Function.update_of_ne (show 0 ≠ kJ by decide), Function.update_of_ne (show 0 ≠ kK by decide), Function.update_of_ne (show 0 ≠ kr by decide), Function.update_of_ne (show 0 ≠ bN2 by decide), Function.update_of_ne (show 0 ≠ bN by decide), Function.update_of_ne (show 0 ≠ DU by decide), Function.update_of_ne (show 0 ≠ dT by decide), Function.update_of_ne (show 0 ≠ OUT by decide), Function.update_of_ne (show 0 ≠ W2 by decide), Function.update_of_ne (show 0 ≠ W1 by decide), Function.update_self]
@[simp] theorem BF_w1 : BF G r0 w1 w2 o dt du bn bn2 vr vk vj bc W1 = w1 := by
  unfold BF; rw [Function.update_of_ne (show W1 ≠ BC by decide), Function.update_of_ne (show W1 ≠ kJ by decide), Function.update_of_ne (show W1 ≠ kK by decide), Function.update_of_ne (show W1 ≠ kr by decide), Function.update_of_ne (show W1 ≠ bN2 by decide), Function.update_of_ne (show W1 ≠ bN by decide), Function.update_of_ne (show W1 ≠ DU by decide), Function.update_of_ne (show W1 ≠ dT by decide), Function.update_of_ne (show W1 ≠ OUT by decide), Function.update_of_ne (show W1 ≠ W2 by decide), Function.update_self]
@[simp] theorem BF_w2 : BF G r0 w1 w2 o dt du bn bn2 vr vk vj bc W2 = w2 := by
  unfold BF; rw [Function.update_of_ne (show W2 ≠ BC by decide), Function.update_of_ne (show W2 ≠ kJ by decide), Function.update_of_ne (show W2 ≠ kK by decide), Function.update_of_ne (show W2 ≠ kr by decide), Function.update_of_ne (show W2 ≠ bN2 by decide), Function.update_of_ne (show W2 ≠ bN by decide), Function.update_of_ne (show W2 ≠ DU by decide), Function.update_of_ne (show W2 ≠ dT by decide), Function.update_of_ne (show W2 ≠ OUT by decide), Function.update_self]
@[simp] theorem BF_o : BF G r0 w1 w2 o dt du bn bn2 vr vk vj bc OUT = o := by
  unfold BF; rw [Function.update_of_ne (show OUT ≠ BC by decide), Function.update_of_ne (show OUT ≠ kJ by decide), Function.update_of_ne (show OUT ≠ kK by decide), Function.update_of_ne (show OUT ≠ kr by decide), Function.update_of_ne (show OUT ≠ bN2 by decide), Function.update_of_ne (show OUT ≠ bN by decide), Function.update_of_ne (show OUT ≠ DU by decide), Function.update_of_ne (show OUT ≠ dT by decide), Function.update_self]
@[simp] theorem BF_dt : BF G r0 w1 w2 o dt du bn bn2 vr vk vj bc dT = dt := by
  unfold BF; rw [Function.update_of_ne (show dT ≠ BC by decide), Function.update_of_ne (show dT ≠ kJ by decide), Function.update_of_ne (show dT ≠ kK by decide), Function.update_of_ne (show dT ≠ kr by decide), Function.update_of_ne (show dT ≠ bN2 by decide), Function.update_of_ne (show dT ≠ bN by decide), Function.update_of_ne (show dT ≠ DU by decide), Function.update_self]
@[simp] theorem BF_du : BF G r0 w1 w2 o dt du bn bn2 vr vk vj bc DU = du := by
  unfold BF; rw [Function.update_of_ne (show DU ≠ BC by decide), Function.update_of_ne (show DU ≠ kJ by decide), Function.update_of_ne (show DU ≠ kK by decide), Function.update_of_ne (show DU ≠ kr by decide), Function.update_of_ne (show DU ≠ bN2 by decide), Function.update_of_ne (show DU ≠ bN by decide), Function.update_self]
@[simp] theorem BF_bn : BF G r0 w1 w2 o dt du bn bn2 vr vk vj bc bN = bn := by
  unfold BF; rw [Function.update_of_ne (show bN ≠ BC by decide), Function.update_of_ne (show bN ≠ kJ by decide), Function.update_of_ne (show bN ≠ kK by decide), Function.update_of_ne (show bN ≠ kr by decide), Function.update_of_ne (show bN ≠ bN2 by decide), Function.update_self]
@[simp] theorem BF_bn2 : BF G r0 w1 w2 o dt du bn bn2 vr vk vj bc bN2 = bn2 := by
  unfold BF; rw [Function.update_of_ne (show bN2 ≠ BC by decide), Function.update_of_ne (show bN2 ≠ kJ by decide), Function.update_of_ne (show bN2 ≠ kK by decide), Function.update_of_ne (show bN2 ≠ kr by decide), Function.update_self]
@[simp] theorem BF_vr : BF G r0 w1 w2 o dt du bn bn2 vr vk vj bc kr = vr := by
  unfold BF; rw [Function.update_of_ne (show kr ≠ BC by decide), Function.update_of_ne (show kr ≠ kJ by decide), Function.update_of_ne (show kr ≠ kK by decide), Function.update_self]
@[simp] theorem BF_vk : BF G r0 w1 w2 o dt du bn bn2 vr vk vj bc kK = vk := by
  unfold BF; rw [Function.update_of_ne (show kK ≠ BC by decide), Function.update_of_ne (show kK ≠ kJ by decide), Function.update_self]
@[simp] theorem BF_vj : BF G r0 w1 w2 o dt du bn bn2 vr vk vj bc kJ = vj := by
  unfold BF; rw [Function.update_of_ne (show kJ ≠ BC by decide), Function.update_self]
@[simp] theorem BF_bc : BF G r0 w1 w2 o dt du bn bn2 vr vk vj bc BC = bc := by
  unfold BF; rw [Function.update_self]

end BFl

/-! ## Taking header values off the table stack -/

/-- Copy the top of `dT` to `R` and to `DU`, and pop it from `dT`. -/
def tkRS (R : Fin HS) : RS := .seq (.push R (.reg dT)) (.seq (.push DU (.reg dT)) (.pop dT))

/-! ## The block -/

/-- Read the block header. -/
def hdrRS : RS := .seq (tkRS bN) (.seq (.push bN2 (.reg dT)) (.seq (tkRS kr) (.seq (tkRS kK)
  (tkRS kJ))))

/-- Clear the block registers and count the block. -/
def popRS : RS := .seq (.pop bN) (.seq (.pop bN2) (.seq (.pop kr) (.seq (.pop kK) (.seq (.pop kJ)
  (.asg BC (.add (.reg BC) (.lit (-1))))))))

/-- Cost of one block. -/
def cBlk (n pk b φD φT B BH : ℕ) : ℕ :=
  80000 * (BH + 2) + 5 * cDG n φD B + 4 * cMvP n B + 2 * cTS n pk b φT B

/-- Entry bounds. -/
def AB (A : ℤ) (u : List GI) : Prop := ∀ w ∈ u, |w.1| ≤ A ∧ |w.2| ≤ A
def TB (A : ℤ) (c : List ℤ) : Prop := ∀ a ∈ c, |a| ≤ A

/-- All stage values of a block bounded by `A`. -/
def blkOK (P : ℕ) (A : ℤ) (r : ℕ) (di gi gam dd am ap : List ℤ) (z : List GI) : Prop :=
  AB A z ∧ AB A (dgG P di z) ∧ AB A (dgG P gi (dgG P di z)) ∧
  AB A (lwG P r am (dgG P gi (dgG P di z))) ∧
  AB A (dgG P gam (lwG P r am (dgG P gi (dgG P di z)))) ∧
  AB A (upG P r ap (dgG P gam (lwG P r am (dgG P gi (dgG P di z))))) ∧
  AB A (dgG P gi (upG P r ap (dgG P gam (lwG P r am (dgG P gi (dgG P di z)))))) ∧
  AB A (blkG P r di gi gam dd am ap z)

end J6Seg.KR

end

section
-- Source module: Solutions.J6AIL3
/-!
# A lines, IL side 3: the kernel frame and the header moves

The `A` kernel (`J6AIL4`) reads its parameters from the table stack `dT`: a header of register
values on top of the tables. Here:

* `AF G r0 w1 w2 o dt du`: the six stacks the kernel body changes (`0, W1, W2, OUT, dT, DU`)
  over a base `G` (the registers live in `G`);
* `tkL Rs`: the header values `vs` on top of `dT` into the registers `Rs` (`tkL_sp`:
  `G ↦ pushAll G Rs vs`);
* `rtLR Rs`: back from the registers onto `dT`, the registers emptied (`rtLR_sp`).
-/

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

namespace J6Seg.KR.AK

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

/-- **The kernel frame**: the six stacks the kernel body changes, over a base `G`. -/
def AF (G : Fin HS → List ℤ) (r0 w1 w2 o dt du : List ℤ) : Fin HS → List ℤ :=
  Function.update (Function.update (Function.update (Function.update (Function.update (Function.update (G) (0 : Fin HS) r0) W1 w1) W2 w2) OUT o) dT dt) DU du

section AFl
variable {G : Fin HS → List ℤ} {r0 w1 w2 o dt du : List ℤ}

@[simp] theorem AF_r0 : AF G r0 w1 w2 o dt du (0 : Fin HS) = r0 := by
  unfold AF; rw [Function.update_of_ne (show (0 : Fin HS) ≠ DU by decide), Function.update_of_ne (show (0 : Fin HS) ≠ dT by decide), Function.update_of_ne (show (0 : Fin HS) ≠ OUT by decide), Function.update_of_ne (show (0 : Fin HS) ≠ W2 by decide), Function.update_of_ne (show (0 : Fin HS) ≠ W1 by decide), Function.update_self]
@[simp] theorem AF_w1 : AF G r0 w1 w2 o dt du W1 = w1 := by
  unfold AF; rw [Function.update_of_ne (show W1 ≠ DU by decide), Function.update_of_ne (show W1 ≠ dT by decide), Function.update_of_ne (show W1 ≠ OUT by decide), Function.update_of_ne (show W1 ≠ W2 by decide), Function.update_self]
@[simp] theorem AF_w2 : AF G r0 w1 w2 o dt du W2 = w2 := by
  unfold AF; rw [Function.update_of_ne (show W2 ≠ DU by decide), Function.update_of_ne (show W2 ≠ dT by decide), Function.update_of_ne (show W2 ≠ OUT by decide), Function.update_self]
@[simp] theorem AF_o : AF G r0 w1 w2 o dt du OUT = o := by
  unfold AF; rw [Function.update_of_ne (show OUT ≠ DU by decide), Function.update_of_ne (show OUT ≠ dT by decide), Function.update_self]
@[simp] theorem AF_dt : AF G r0 w1 w2 o dt du dT = dt := by
  unfold AF; rw [Function.update_of_ne (show dT ≠ DU by decide), Function.update_self]
@[simp] theorem AF_du : AF G r0 w1 w2 o dt du DU = du := by
  unfold AF; rw [Function.update_self]

end AFl

/-! ## Header moves -/

/-- The registers outside the frame. -/
def FFree (R : Fin HS) : Prop :=
  R.val < 56 ∧ ∀ k ∈ ([(0 : Fin HS), W1, W2, OUT, dT, DU] : List (Fin HS)), R ≠ k

/-- One header value from `dT` into `R`. -/
def tkH (R : Fin HS) : RS := .seq (.push R (.reg dT)) (.pop dT)

/-- One register value back onto `dT`, the register popped. -/
def rtH (R : Fin HS) : RS := .seq (.push dT (.reg R)) (.pop R)

/-- The header into the registers, in order. -/
def tkL : List (Fin HS) → RS
  | [] => .skip
  | R :: Rs => .seq (tkH R) (tkL Rs)

/-- The registers back onto `dT`, last first. -/
def rtLR : List (Fin HS) → RS
  | [] => .skip
  | R :: Rs => .seq (rtLR Rs) (rtH R)

/-- The registers after reading the header. -/
def pushAll (G : Fin HS → List ℤ) : List (Fin HS) → List ℤ → Fin HS → List ℤ
  | R :: Rs, v :: vs => pushAll (Function.update G R (v :: G R)) Rs vs
  | _, _ => G

end J6Seg.KR.AK

end

section
-- Source module: Solutions.J6SegKr6
/-!
# B3 (chunk kernel), part 6: the chunk on pairs with per-block records, and growth bounds

* `chunkB`, `chunkG2`: the chunk kernel on pairs, each block of length `n` with its own record
  length `r = 2n`; `liftG_chunkL2`: `liftG (chunkL …) = chunkG2 …`.
* growth: if the row entries are bounded by `X` and the tables by `2^P·Kt`, every stage value of a
  block of length `n ≤ λ` is bounded by `Xs ((λ+1)·Kt+1) X 7`, where `Xs K X i` iterates `X ↦ K·X+1`.
-/

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 chunk on pairs, per-block records -/

/-- Total length of the blocks. -/
def sumN : List (ℕ × List ℤ × List ℤ) → ℕ
  | [] => 0
  | d :: ds => d.1 + sumN ds

section chunk
variable (P : ℕ) (di gi gam dd : List ℤ)

/-- The blocks part of the chunk kernel. -/
def chunkB : List (ℕ × List ℤ × List ℤ) → List GI → List GI
  | [], _ => []
  | (n, am, ap) :: ds, z => blkG P (2 * n) di gi gam dd am ap (z.take n) ++ chunkB ds (z.drop n)

/-- The chunk kernel on pairs (zeros after the last block). -/
def chunkG2 (ds : List (ℕ × List ℤ × List ℤ)) (z : List GI) : List GI :=
  chunkB P di gi gam dd ds z ++ (z.drop (sumN ds)).map fun _ => ((0 : ℤ), (0 : ℤ))

end chunk

/-! ## Growth of the stage values -/

/-- Iterates of `X ↦ K·X + 1`. -/
def Xs (K X : ℤ) : ℕ → ℤ
  | 0 => X
  | i + 1 => K * Xs K X i + 1

end J6Seg.KR

end

section
-- Source module: Solutions.J6AIL4
/-!
# A lines, IL side 4: the per-block `A` kernel as an IL program

`aKprog pM` computes `J6A.aKerL w2 bk win Q aTab rTab cTab` on the row on stack `0`. It reads
every parameter from its table on `dT`: the table `aKTab …` is the header `aHdr …` (register
values) followed by `aTab ++ rTab ++ cTab`.

1. `tkL aRegs`: the header into the registers (`R1 … R7`, `bN`, `bN2`, `rP`, `rP1`, and the
   record-product parameters `kc … kD`);
2. `dgI` (scale by `aTab`, rounding `Q`, exact because `aTab = 2^Q a`), drop the rest of the row,
   pad with `bk` zero pairs;
3. `aTab` off `dT` (onto `DU`); `tsI`: the Toeplitz stage (one record product, `Δ = 0`), rounding
   `Q`, exact because `rTab = 2^Q r`; it consumes `rTab` onto `DU`;
4. drop the first `w2` outputs; `dgI` with `cTab` (rounding `Q`) onto `OUT`;
5. `win - bk` zero pairs, the outputs back onto stack `0`;
6. restore the table (`DU → dT`) and the header (`rtLR aRegs`).

**`aK_runs`**: the run on the managed state of `RowBody6` (every stack but `0`, `dT` empty).
-/

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

namespace J6Seg.KR.AK

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

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

/-- Counts: `w2`, `2(win - w2)`, `2bk`, `2w2`, `bk`, `2(win - bk)`, `2w2 + bk`. -/
abbrev R1 : Fin HS := 26
abbrev R2 : Fin HS := 27
abbrev R3 : Fin HS := 28
abbrev R4 : Fin HS := 29
abbrev R5 : Fin HS := 30
abbrev R6 : Fin HS := 33
abbrev R7 : Fin HS := 34

/-- The header registers, in table order. -/
def aRegs : List (Fin HS) := [R7, R1, R2, R3, R4, R5, R6, bN, bN2, rP, rP1, kc, kr, kK, kB, kJ, kD]

/-- The header values. -/
def aHdr (w2 bk win Q pk b : ℕ) : List ℤ :=
  [((w2 + bk + w2 : ℕ) : ℤ), (w2 : ℤ), ((2 * (win - w2) : ℕ) : ℤ), ((2 * bk : ℕ) : ℤ),
   ((2 * w2 : ℕ) : ℤ), (bk : ℤ), ((2 * (win - bk) : ℕ) : ℤ), ((w2 + bk : ℕ) : ℤ),
   ((2 * (w2 + bk) : ℕ) : ℤ), 2 ^ Q, 2 ^ (Q + 1), 2 ^ pk, ((2 * (w2 + bk) : ℕ) : ℤ),
   kG b (2 * (w2 + bk)) (2 ^ pk), 2 ^ b, (2 ^ b) ^ (2 * (2 * (w2 + bk))), 1]

/-- Cost of the `A` kernel. -/
def cAK (w2 bk win pk b φ1 φ2 B BH : ℕ) : ℕ :=
  (17 * (5000 * (BH + 2)) + 1) + cDG w2 φ1 B + cMvP (2 * (win - w2)) B + cMvP (2 * bk) B +
  (2 * 2000 * (B + 2) + w2 * (14000 * (B + 2)) + 32) + cTS (w2 + bk) pk b φ2 B +
  cMvP (2 * w2) B + cDG bk φ1 B + cMvP (2 * (win - bk)) B + cMvP bk B +
  (2 * 2000 * (B + 2) + (w2 + bk + w2) * (14000 * (B + 2)) + 32) + (17 * (5000 * (BH + 2)) + 1)

/-- The registers after the header. -/
def aG (w2 bk win Q pk b : ℕ) : Fin HS → List ℤ :=
  pushAll (fun _ => []) aRegs (aHdr w2 bk win Q pk b)

end J6Seg.KR.AK

end

section
-- Source module: Solutions.J6SegILKer2
/-!
# B-line IL, B2 (kernel-pass wrapper) §2: table moves and the row loop

* `tLoad_sp`: pop the length `|T|` of the next table from the store `ρ 10`, push it on `U`
  (`ρ 15`) and keep it in `ρ 16`, move `T` to `dT` in order (through `T1 = ρ 14`).
* `tUnload_sp`: move `T` from `dT` onto `U` (reversed), drop the length.
* `rows_run`: the loop over the `post` rows of one block: call the kernel `e` on the row on top
  of stack `0` (tables on `dT`), move its output (`2m` integers) onto `Acc` (`ρ 17`).
-/

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

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

theorem RegOK.ne {ρ : Fin 19 → Fin HS} {dT dF : Fin HS} (hR : RegOK ρ dT dF) {a b : Fin 19}
    (h : a ≠ b) : ρ a ≠ ρ b := hR.inj.ne h

/-! ## Loading and unloading a table -/

/-- Load the next table onto `dT`. -/
def tLoad (ρ : Fin 19 → Fin HS) (dT : Fin HS) : RS :=
  .seq (.push (ρ 16) (.reg (ρ 10))) (.seq (.pop (ρ 10)) (.seq (.push (ρ 15) (.reg (ρ 16)))
    (.seq (mvK (ρ 10) (ρ 14) (ρ 18) (.reg (ρ 16))) (mvK (ρ 14) dT (ρ 18) (.reg (ρ 16))))))

/-- Put the table on `dT` back (reversed) onto `U`. -/
def tUnload (ρ : Fin 19 → Fin HS) (dT : Fin HS) : RS :=
  .seq (mvK dT (ρ 15) (ρ 18) (.reg (ρ 16))) (.pop (ρ 16))

/-- The cost of one table move. -/
def cTab (B n : ℕ) : ℕ := 40000 * (B + 2) * (n + 1)

/-! ## The row loop -/

/-- One row: the kernel, then the output onto `Acc`. -/
def rowB {Pp : ℕ} (ρ : Fin 19 → Fin HS) (e : Fin Pp) : IStmt HS Pp :=
  .seq (.call e) (mvK 0 (ρ 17) (ρ 18) (.reg (ρ 7))).code

/-- The rows of one block. -/
def rowsL {Pp : ℕ} (ρ : Fin 19 → Fin HS) (e : Fin Pp) : IStmt HS Pp :=
  irep (ρ 13) (.reg (ρ 4)) (rowB ρ e)

/-- The cost of the rows of one block. -/
def cRows (Ckn B m post : ℕ) : ℕ := post * (Ckn + 40000 * (B + 2) * (m + 1)) + 10000 * (B + 2)

end J6Seg

end

section
-- Source module: Solutions.J6SegILKer3
/-!
# B-line IL, B2 (kernel-pass wrapper) §3: the block loops

`segT tabOf q = |tabOf q| :: tabOf q`; the table store holds `tabsT tabOf nc ++ R0`
(`tabsT = ⋃_{q < nc} segT q`).

* `qLoop_run`: for one outer index `i`, the loop over the chunks `q < nc`: load table `q`, run the
  `post` rows of block `i·nc + q` (kernel `g q`), unload; the store is consumed and `U` (`ρ 15`)
  collects the used tables, reversed.
* `iLoop_run`: the loop over `i < pre`: the chunk loop, then the store is restored from `U`.

The rows come from an abstract stream (`RowStream`): row `r` is `xr r`, on top of stack `0`
before it is `zA r`, the outputs so far are `aA r` (reversed), the kernel of row `r` is
`g (r / post % nc)`.
-/

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

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

/-- One stored table. -/
def segT (tabOf : ℕ → List ℤ) (q : ℕ) : List ℤ := ((tabOf q).length : ℤ) :: tabOf q

/-- The table store. -/
def tabsT (tabOf : ℕ → List ℤ) (nc : ℕ) : List ℤ := (List.range nc).flatMap (segT tabOf)

/-- The row stream. -/
def RowStream (Pre : List GI → Prop) (g : ℕ → List GI → List GI) (B m post nc R : ℕ)
    (zA aA : ℕ → List ℤ) (xr : ℕ → List GI) : Prop :=
  ∀ r < R, zA r = gflat (xr r) ++ zA (r + 1) ∧ Pre (xr r) ∧
    (g (r / post % nc) (xr r)).length = m ∧ Bnd B (gflat (g (r / post % nc) (xr r))) ∧
    aA (r + 1) = (gflat (g (r / post % nc) (xr r))).reverse ++ aA r

/-- The chunk body. -/
def qBody {Pp : ℕ} (ρ : Fin 19 → Fin HS) (dT : Fin HS) (e : Fin Pp) : IStmt HS Pp :=
  .seq (tLoad ρ dT).code (.seq (rowsL ρ e) (tUnload ρ dT).code)

/-- The chunk loop. -/
def qLoop {Pp : ℕ} (ρ : Fin 19 → Fin HS) (dT : Fin HS) (e : Fin Pp) : IStmt HS Pp :=
  irep (ρ 12) (.reg (ρ 6)) (qBody ρ dT e)

/-- The outer body: the chunk loop, then restore the store. -/
def iBody {Pp : ℕ} (ρ : Fin 19 → Fin HS) (dT : Fin HS) (e : Fin Pp) : IStmt HS Pp :=
  .seq (qLoop ρ dT e) (mvK (ρ 15) (ρ 10) (ρ 18) (.reg (ρ 9))).code

/-- The outer loop. -/
def iLoop {Pp : ℕ} (ρ : Fin 19 → Fin HS) (dT : Fin HS) (e : Fin Pp) : IStmt HS Pp :=
  irep (ρ 11) (.reg (ρ 5)) (iBody ρ dT e)

/-- Costs. -/
def cQb (Ckn B m post Lmax : ℕ) : ℕ := 2 * cTab B Lmax + cRows Ckn B m post
def cIb (Ckn B m post Lmax nc Tot : ℕ) : ℕ :=
  nc * (cQb Ckn B m post Lmax + 10000 * (B + 2)) + 20000 * (B + 2) * (Tot + 2)
def cIL (Ckn B m post Lmax nc Tot pre : ℕ) : ℕ :=
  pre * (cIb Ckn B m post Lmax nc Tot + 10000 * (B + 2)) + 10000 * (B + 2)

end J6Seg

end

section
-- Source module: Solutions.J6SegILKer4
/-!
# B-line IL, B2: the kernel-pass wrapper `kerW` (`kerPass` on stack `0`)

`kerW ρ dT dF eF eB e`:

1. push `descFine w f A post` on `dF` (by `dup` from `ρ 0, 2, 3, 4`), call `eF`, drop it:
   stack `0` becomes the exposure `X = blkT (2^f) post A L` (`A = pre·nc`);
2. `iLoop`: for `i < pre`, `q < nc`: load table `q` on `dT`, run the kernel `e` on the `post`
   rows of block `i·nc + q` (each output moved onto `Acc = ρ 17`), unload; restore the store;
3. move `Acc` back onto stack `0` (`ρ 8 = 2·A·post·2^f` integers);
4. push `descFine w' f A post`, call `eB`, drop it.

**`kerW_run`**: on the managed state `F` (registers as in `J6SegILKer1`, `F dT = []`) with
`F 0 = gflat L ++ l`, `LRep A (2^f) post L φ` and the kernel hypotheses on the rows `φ b c`,
the wrapper leaves `gflat (kerPass (2^f) post A (fun b => g (b % nc)) L) ++ l` on stack `0` and
everything else unchanged, at cost
`C₀ · cFine w f A post + C₀ · cFine w' f A post + kerC …` (`kerC`: the rows at `⌊Ck⌋₊ +
O((B+2)·2^f)` each, the table moves `O((B+2)·Lmax)` per block, `O((B+2)·|tabs|)` per `i`,
`O((B+2)·N)` for the final move).
-/

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

namespace J6Seg
open IntMul NLogN NLogN.KIL NLogN.KFB IntMul.Kappa.R6 Function
open IntMul.HvdH IntMul.HvdH.Ref NLogN.KFM NLogN.KFY

/-- Push the fine descriptor `[w, f, A, post]` (`w` from `rw`) on `dF`. -/
def descI {Pp : ℕ} (ρ : Fin 19 → Fin HS) (dF rw : Fin HS) : IStmt HS Pp :=
  .seq (.op (.dup (ρ 4) dF)) (.seq (.op (.dup (ρ 3) dF)) (.seq (.op (.dup (ρ 2) dF))
    (.op (.dup rw dF))))

/-- Drop it. -/
def dropI {Pp : ℕ} (dF : Fin HS) : IStmt HS Pp :=
  .seq (.op (.drop dF)) (.seq (.op (.drop dF)) (.seq (.op (.drop dF)) (.op (.drop dF))))

/-- **The kernel-pass wrapper.** -/
def kerW {Pp : ℕ} (ρ : Fin 19 → Fin HS) (dT dF : Fin HS) (eF eB e : Fin Pp) : IStmt HS Pp :=
  .seq (descI ρ dF (ρ 0)) (.seq (.call eF) (.seq (dropI dF) (.seq (iLoop ρ dT e)
    (.seq (mvK (ρ 17) 0 (ρ 18) (.reg (ρ 8))).code (.seq (descI ρ dF (ρ 1))
      (.seq (.call eB) (dropI dF)))))))

/-- The cost of the wrapper outside the two fine exposures. -/
def kerC (Ckn B m post Lmax nc Tot pre N : ℕ) : ℕ :=
  cIL Ckn B m post Lmax nc Tot pre + 20000 * (B + 2) * (N + 1) + 4000 * (B + 1)

/-! ## The exposed rows -/

end J6Seg

end

section
-- Source module: Solutions.J6AIL11
/-!
# A lines, IL side 11: the copy step `dupL 32`

* `Sp.dupKE`: `dupK` with a general count expression;
* `cpBody T T2 C r`: from `T = G.reverse`, push `G` onto stack `0` (via `T2`) and restore `T`;
  `rpt k b`: `k` copies of `b` in sequence;
* **`dupP`**: read `pre, post, t` of the descriptor on `dA` into the register file, compute
  `2·pre·t·post`, move the array to `T`, push it back `32` times, drop `T`, leave the register file.
-/

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

namespace J6A.DP

open IntMul NLogN NLogN.KIL RS J6Seg.PJ IntMul.HvdH J6A.PA

/-! ## One copy -/

abbrev T1 : Fin HS := 16
abbrev T2 : Fin HS := 17
abbrev CC : Fin HS := 18

@[local simp] theorem neF0 : ((0 : Fin HS) = T1) = False := eq_false (by decide)
@[local simp] theorem neF1 : ((0 : Fin HS) = T2) = False := eq_false (by decide)
@[local simp] theorem neF2 : ((0 : Fin HS) = CC) = False := eq_false (by decide)
@[local simp] theorem neF3 : ((0 : Fin HS) = (20 : Fin HS)) = False := eq_false (by decide)
@[local simp] theorem neF4 : (T1 = (0 : Fin HS)) = False := eq_false (by decide)
@[local simp] theorem neF5 : (T1 = T2) = False := eq_false (by decide)
@[local simp] theorem neF6 : (T1 = CC) = False := eq_false (by decide)
@[local simp] theorem neF7 : (T1 = (20 : Fin HS)) = False := eq_false (by decide)
@[local simp] theorem neF8 : (T2 = (0 : Fin HS)) = False := eq_false (by decide)
@[local simp] theorem neF9 : (T2 = T1) = False := eq_false (by decide)
@[local simp] theorem neF10 : (T2 = CC) = False := eq_false (by decide)
@[local simp] theorem neF11 : (T2 = (20 : Fin HS)) = False := eq_false (by decide)
@[local simp] theorem neF12 : (CC = (0 : Fin HS)) = False := eq_false (by decide)
@[local simp] theorem neF13 : (CC = T1) = False := eq_false (by decide)
@[local simp] theorem neF14 : (CC = T2) = False := eq_false (by decide)
@[local simp] theorem neF15 : (CC = (20 : Fin HS)) = False := eq_false (by decide)
@[local simp] theorem neF16 : ((20 : Fin HS) = (0 : Fin HS)) = False := eq_false (by decide)
@[local simp] theorem neF17 : ((20 : Fin HS) = T1) = False := eq_false (by decide)
@[local simp] theorem neF18 : ((20 : Fin HS) = T2) = False := eq_false (by decide)
@[local simp] theorem neF19 : ((20 : Fin HS) = CC) = False := eq_false (by decide)

/-- From `T1 = G.reverse` (count `|G|` on `20`) push `G` onto `0` and restore `T1`. -/
def cpBody : RS := .seq (dupKE T1 0 T2 CC (.reg 20)) (mvK T2 T1 CC (.reg 20))

/-- `k` copies of a program in sequence. -/
def rpt : ℕ → RS → RS
  | 0, _ => .skip
  | k + 1, b => .seq (rpt k b) b

/-- The cost of one copy. -/
def cB (n Bb : ℕ) : ℕ :=
  ((1 + 1) * 2000 * (Bb + 2) + n * ((1 + 1) * 2000 * (Bb + 2) +
    (1 + 1) * 2000 * (Bb + 2) + (2 * Bb + 2) + 8000 * (Bb + 2) + 30) + 32) +
  ((1 + 1) * 2000 * (Bb + 2) + n * (14000 * (Bb + 2)) + 32)

/-! ## The copy step -/

/-- `2 · pre · (t · post)`. -/
def nE : Ex HS := .mul (.lit 2) (.mul (.reg 1) (.mul (.reg 4) (.reg 2)))

/-- Read `pre, post, t` and compute the flat length. -/
def rdN (X : Fin HS) : RS :=
  .seq (rdK X 1 1) (.seq (rdK X 2 2) (.seq (rdK X 4 4) (.asg 20 nE)))

/-- **The copy step.** -/
def dupP (X : Fin HS) : RS :=
  .seq (initP rg) (.seq (rdN X) (.seq (mvK 0 T1 CC (.reg 20)) (.seq (rpt 31 cpBody)
    (.seq (mvK T1 0 CC (.reg 20)) (finiP rg)))))

/-- The cost of the copy step. -/
def cDupP (W n Bb : ℕ) : ℕ :=
  (21 * 8001 + 1) + (cRdK 1 W + (cRdK 2 W + (cRdK 4 W + J6Seg.PJ.cT nE W))) +
    ((1 + 1) * 2000 * (Bb + 2) + n * (14000 * (Bb + 2)) + 32) + (31 * cB n Bb + 1) +
    ((1 + 1) * 2000 * (Bb + 2) + n * (14000 * (Bb + 2)) + 32) + (21 * (2 * (W + 1) + 2) + 1)

end J6A.DP

end

section
-- Source module: Solutions.J6AIL12
/-!
# A lines, IL side 12: the parameter block on any descriptor stack

`parAX X` is `parA` reading the descriptor from the stack `X` (`parA = parAX 31`), and
**`Rn.parAX`** is `Rn.parA` for every `X ∉ rg`, `X < 56`, `X ≠ 15, 47`.
-/

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

namespace J6A.PA

open IntMul NLogN NLogN.KIL RS J6Seg.PJ IntMul.HvdH

/-- **The parameter block on the stack `X`.** -/
def parAX (X : Fin HS) : RS := seqL [rdK X 1 1, rdK X 2 2, rdK X 3 3, rdK X 4 4, rdK X 5 5, rdK X 6 6,
  .asg 20 (.lit 0), .asg 21 (.lit 0), .rep cL (.reg 5) sqB, .asg 7 mE,
  .asg 20 (.lit 1), .asg 8 (.lit 1), .asg 22 (.lit 0), .rep cL (.reg 7) lgB,
  .asg 9 (.add (.add (.reg 8) (.reg 7)) (.lit 3)),
  .asg 11 (.mul (.lit 8) (.reg 8)),
  .asg 12 (.add (.reg 22) (.lit 3)),
  .asg 10 (.ndiv (.reg 4) (.mul (.lit 16) (.reg 8)))]

end J6A.PA

end

section
-- Source module: Solutions.J6AIL13
/-!
# A lines, IL side 13: the fine-descriptor set-ups

For a fine descriptor `descFine w f a y` with `w = p + 2`, `f` a literal and `a`, `y` expressions in
the registers of `parAX X`:

* `fsR X ea ey fl Q`: enter the register file, `parAX X`, push `w, f, a, y` on four temporaries
  `Q 0 … Q 3` (outside `rg`, away from `dF`), leave the register file;
* `fsUp`: then copy `y, a, f, w` onto `dF` and drop the temporaries; `fsDn`: drop four from `dF`.

**`fsUp_runs`** / **`fsDn_runs`**.
-/

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

namespace J6A.FS

open IntMul NLogN NLogN.KIL RS J6Seg.PJ IntMul.HvdH Function J6A.PA IntMul.Kappa.R6

/-- The base of the temporaries. -/
def qb (dF : Fin HS) : ℕ := if 16 ≤ dF.val ∧ dF.val ≤ 19 then 32 else 16

theorem qb_cases (dF : Fin HS) : (qb dF = 32 ∧ 16 ≤ dF.val ∧ dF.val ≤ 19) ∨
    (qb dF = 16 ∧ ¬ (16 ≤ dF.val ∧ dF.val ≤ 19)) := by
  unfold qb; split_ifs with h
  · exact Or.inl ⟨rfl, h⟩
  · exact Or.inr ⟨rfl, h⟩

/-- The temporaries. -/
def Q (dF : Fin HS) (i : Fin 4) : Fin HS :=
  ⟨qb dF + i.val, by rcases qb_cases dF with h | h <;> (rw [h.1]; have := i.isLt; unfold HS; omega)⟩

/-- `w = p + 2`. -/
def eW : Ex HS := .add (.reg 5) (.lit 2)

/-- The register part. -/
def fsR (X : Fin HS) (ea ey : Ex HS) (fl : ℕ) (dF : Fin HS) : RS :=
  .seq (initP rg) (.seq (parAX X) (.seq (.push (Q dF 0) eW) (.seq (.push (Q dF 1) (.lit fl))
    (.seq (.push (Q dF 2) ea) (.seq (.push (Q dF 3) ey) (finiP rg))))))

/-- **The set-up**: registers, then copy onto `dF`, then drop the temporaries. -/
def fsUp {Pp : ℕ} (X : Fin HS) (ea ey : Ex HS) (fl : ℕ) (dF : Fin HS) : IStmt HS Pp :=
  .seq (fsR X ea ey fl dF).code (.seq (.op (.dup (Q dF 3) dF)) (.seq (.op (.dup (Q dF 2) dF))
    (.seq (.op (.dup (Q dF 1) dF)) (.seq (.op (.dup (Q dF 0) dF))
    (.seq (.op (.drop (Q dF 0))) (.seq (.op (.drop (Q dF 1))) (.seq (.op (.drop (Q dF 2)))
      (.op (.drop (Q dF 3))))))))))

/-- **The tear-down.** -/
def fsDn {Pp : ℕ} (dF : Fin HS) : IStmt HS Pp :=
  .seq (.op (.drop dF)) (.seq (.op (.drop dF)) (.seq (.op (.drop dF)) (.op (.drop dF))))

/-- The register file after the four pushes, back on the stacks. -/
def fsG (F : Fin HS → List ℤ) (dF : Fin HS) (w f a y : ℤ) : Fin HS → List ℤ :=
  update (update (update (update F (Q dF 0) [w]) (Q dF 1) [f]) (Q dF 2) [a]) (Q dF 3) [y]

/-- The cost of the register part. -/
def cFsR (W p m : ℕ) (ea ey : Ex HS) (fl : ℕ) : ℕ :=
  (21 * 8001 + 1) + (cParA W p m + (J6Seg.PJ.cT eW W + (J6Seg.PJ.cT (.lit fl : Ex HS) W +
    (J6Seg.PJ.cT ea W + J6Seg.PJ.cT ey W)))) + (21 * (2 * (W + 1) + 2) + 1)

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

/-- The cost of the stack part. -/
def cFsI (w f a y : ℤ) : ℕ :=
  (dupCost (zsz y) + 17) + (dupCost (zsz a) + 17) + (dupCost (zsz f) + 17) +
    (dupCost (zsz w) + 17) + (2 * zsz w + 2) + (2 * zsz f + 2) + (2 * zsz a + 2) + (2 * zsz y + 2)

end il

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

/-- The cost of the tear-down. -/
def cFsD (w f a y : ℕ) : ℕ :=
  (2 * zsz (w : ℤ) + 2) + (2 * zsz (f : ℤ) + 2) + (2 * zsz (a : ℤ) + 2) + (2 * zsz (y : ℤ) + 2)

end runs

end J6A.FS

end

section
-- Source module: Solutions.KFUURev7
/-!
# KFUU §2g: phase 2 in IL (one iteration)

`p2B` divides `M` by `U²`, calls the swap entry with the descriptor
`[w, u, 1, 3, 5, U^k, U, U^{n-2-2k}, U, U^k·Bh]`, and multiplies `A = U^k` and `C = U^k·Bh` by `U`.
-/

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

namespace NLogN.KFUU

open IntMul KIL KFB KFC KFT KFU AModel AModelL12

section Gen

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

theorem st5 (t : IState HS) (a b c d e : List Item) :
    ((((t.upd rM a).upd rM b).upd 0 c).upd rA d).upd rB e =
      (((t.upd 0 c).upd rA d).upd rM b).upd rB e := by
  ((apply IState.ext' <;> (try intros) <;> (simp +decide only [IState.upd_G, IState.upd_ip, IState.upd_out, ↓reduceIte]) <;> (try split_ifs) <;> (try subst_vars) <;> (try simp_all +decide)))

end Gen

end NLogN.KFUU

end

section
-- Source module: Solutions.J6SegILSlab2
/-!
# B-line IL, B1 (streaming slab pass) §2: registers, lanes, the unit loops

**Register convention** (fixed stacks, all below `56`, none `0`; the caller's data elsewhere is
preserved, since every scratch stack is restored):

| stack | role |
|---|---|
| `rK1 = 41`, `rK2 = 42` | constants of the item programs (inputs, preserved) |
| `0` | lane 1 input (the stream; the rest `l` below it is kept), and the output |
| `rZ = 53` | lane 2 input (the second array of the zipped output pass) |
| `rT = 50`, `rT2 = 51` | lane temps (scratch) |
| `rC = 45` | the control list `flatC cs` (input, preserved) |
| `rCU = 47`, `rOP = 48`, `rARG = 49`, `rACC = 52` | scratch: used commands, current command, output |
| `rNB = 37`, `rNC = 38`, `rM = 39`, `rTOT = 40` | inputs (one value each, preserved): `pre`, `|cs|`, `m`, the output length |
| `43, 46, 54` | scratch of the item programs (§5) |
| `34, 35, 36` | the `rptX` counters (`lA lC lD`) |

`55` (`ARGM`) is the multiplication argument stack. A *lane layout* `lin inp tmp G` puts the
remaining input units `inp` and the moved units `tmp` on the lane stacks over the managed state
`G`; `LaneOK lin` says it leaves the skeleton registers alone. The unit programs (`rd`, `bk`, `dr`,
`it`) are abstract with their specifications (`RdOK`, `BkOK`, `DrOK`, `ItOK`).

Here: the managed-state op lemmas, `cf3_encK` (the `rptX` counters over a managed state), and the
loops over one slab: `dropL_runs` (drop `m` units), `zeroL_runs` (emit `m` zeros), `copyL_runs`
(move `m` units to the temps while emitting `g a u`), `backL_runs` (move them back).
-/

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

namespace J6Seg.Slab

open IntMul NLogN NLogN.KIL Function
open NLogN.KFUU (lA lC lD)
open NLogN.KFT (cf3 rptX)

abbrev rC : Fin HS := 45
abbrev rCU : Fin HS := 47
abbrev rOP : Fin HS := 48
abbrev rARG : Fin HS := 49
abbrev rT : Fin HS := 50
abbrev rT2 : Fin HS := 51
abbrev rACC : Fin HS := 52
abbrev rZ : Fin HS := 53
abbrev rNB : Fin HS := 37
abbrev rNC : Fin HS := 38
abbrev rM : Fin HS := 39
abbrev rTOT : Fin HS := 40
abbrev rK1 : Fin HS := 41
abbrev rK2 : Fin HS := 42

/-- The skeleton registers (never lane stacks). -/
def skel : List (Fin HS) := [rC, rCU, rOP, rARG, rACC, rNB, rNC, rM, rTOT, rK1, rK2, lA, lC, lD]

/-! ## Managed-state operations -/

/-! ## Lanes -/

section lanes
variable {U : Type*}

/-- The lane layout leaves the skeleton registers alone and is the identity when empty. -/
structure LaneOK (lin : List U → List U → (Fin HS → List ℤ) → (Fin HS → List ℤ)) : Prop where
  rd : ∀ j ∈ skel, ∀ inp tmp G, lin inp tmp G j = G j
  up : ∀ j ∈ skel, ∀ inp tmp G v, lin inp tmp (update G j v) = update (lin inp tmp G) j v
  nil : ∀ G, lin [] [] G = G

variable {P : ℕ} {I : List Sym} {iprog : Fin P → IStmt HS P}
  (lin : List U → List U → (Fin HS → List ℤ) → (Fin HS → List ℤ)) (OkU : U → Prop)

/-- Move one unit from the inputs to the temps. -/
def RdOK (rd : IStmt HS P) (c : ℕ) : Prop :=
  ∀ s G u inp tmp, OkU u →
    IRuns I iprog rd (encK s (lin (u :: inp) tmp G)) (encK s (lin inp (u :: tmp) G)) c

/-- Move one unit back from the temps to the inputs. -/
def BkOK (bk : IStmt HS P) (c : ℕ) : Prop :=
  ∀ s G u inp tmp, OkU u →
    IRuns I iprog bk (encK s (lin inp (u :: tmp) G)) (encK s (lin (u :: inp) tmp G)) c

/-- Drop one input unit. -/
def DrOK (dr : IStmt HS P) (c : ℕ) : Prop :=
  ∀ s G u inp tmp, OkU u →
    IRuns I iprog dr (encK s (lin (u :: inp) tmp G)) (encK s (lin inp tmp G)) c

/-- Emit `g a u` for the unit `u` on top of the temps (`a` on `rARG`; `KG`: the constants). -/
def ItOK (OkA : ℤ → Prop) (KG : (Fin HS → List ℤ) → Prop) (g : ℤ → U → ℤ) (it : IStmt HS P)
    (c : ℕ) : Prop :=
  ∀ s G u inp tmp a rest, OkU u → OkA a → G rARG = a :: rest → KG G →
    IRuns I iprog it (encK s (lin inp (u :: tmp) G))
      (encK s (lin inp (u :: tmp) (update G rACC (g a u :: G rACC)))) c

end lanes

/-! ## The slab loops -/

end J6Seg.Slab

end

section
-- Source module: Solutions.J6SegILSlab3
/-!
# B-line IL, B1 (streaming slab pass) §3: one command, the command loop, one block

The skeleton contents are kept in the form `skS G C CU OP ARG ACC` (control `rC`, used
commands `rCU`, current op `rOP`, current argument `rARG`, output `rACC`, over `G`).

* `cmdP`: pop `(op, arg)` from `rC` to `rOP`, `rARG`; branch on the sign of `op`:
  drop a slab (`op < 0`), emit `m` zeros (`op = 0`), emit the slab mapped by `g arg`, reading
  the units to the temps and moving them back (`op > 0`); push `op`, `arg` to `rCU`.
  **`cmd_runs`**.
* `cmdLoop`: `rptX rNC cmdP` — the whole command list; **`cmdLoop_runs`**: the output is
  `runC m g cs inp` (reversed onto `rACC`), the input loses `nDrop cs` slabs.
* `restLoop`: `rptX rNC (mv rCU rC; mv rCU rC)` restores the control list.
  **`restLoop_runs`**.
* `blockP = cmdLoop; restLoop`: **`block_runs`**.
-/

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

namespace J6Seg.Slab

open IntMul NLogN NLogN.KIL Function
open NLogN.KFUU (lA lC lD)
open NLogN.KFT (cf3 rptX)

/-! ## The skeleton contents -/

/-- The skeleton contents over `G`. -/
def skS (G : Fin HS → List ℤ) (C CU OP ARG ACC : List ℤ) : Fin HS → List ℤ :=
  update (update (update (update (update G rC C) rCU CU) rOP OP) rARG ARG) rACC ACC

section sk
variable (G : Fin HS → List ℤ) (C CU OP ARG ACC : List ℤ)

@[simp] theorem skS_C : skS G C CU OP ARG ACC rC = C := by
  simp (config := { decide := true }) [skS, update_apply]
@[simp] theorem skS_CU : skS G C CU OP ARG ACC rCU = CU := by
  simp (config := { decide := true }) [skS, update_apply]
@[simp] theorem skS_OP : skS G C CU OP ARG ACC rOP = OP := by
  simp (config := { decide := true }) [skS, update_apply]
@[simp] theorem skS_ARG : skS G C CU OP ARG ACC rARG = ARG := by
  simp (config := { decide := true }) [skS, update_apply]
@[simp] theorem skS_ACC : skS G C CU OP ARG ACC rACC = ACC := by simp [skS]

theorem skS_ne {j : Fin HS} (h1 : j ≠ rC) (h2 : j ≠ rCU) (h3 : j ≠ rOP) (h4 : j ≠ rARG)
    (h5 : j ≠ rACC) : skS G C CU OP ARG ACC j = G j := by
  simp [skS, h1, h2, h3, h4, h5]

@[simp] theorem skS_lA : skS G C CU OP ARG ACC lA = G lA := skS_ne _ _ _ _ _ _ (by decide)
  (by decide) (by decide) (by decide) (by decide)
@[simp] theorem skS_lC : skS G C CU OP ARG ACC lC = G lC := skS_ne _ _ _ _ _ _ (by decide)
  (by decide) (by decide) (by decide) (by decide)
@[simp] theorem skS_lD : skS G C CU OP ARG ACC lD = G lD := skS_ne _ _ _ _ _ _ (by decide)
  (by decide) (by decide) (by decide) (by decide)
@[simp] theorem skS_rM : skS G C CU OP ARG ACC rM = G rM := skS_ne _ _ _ _ _ _ (by decide)
  (by decide) (by decide) (by decide) (by decide)
@[simp] theorem skS_rNC : skS G C CU OP ARG ACC rNC = G rNC := skS_ne _ _ _ _ _ _ (by decide)
  (by decide) (by decide) (by decide) (by decide)
@[simp] theorem skS_rNB : skS G C CU OP ARG ACC rNB = G rNB := skS_ne _ _ _ _ _ _ (by decide)
  (by decide) (by decide) (by decide) (by decide)
@[simp] theorem skS_rTOT : skS G C CU OP ARG ACC rTOT = G rTOT := skS_ne _ _ _ _ _ _ (by decide)
  (by decide) (by decide) (by decide) (by decide)

end sk

/-! ## Lane-state operations -/

/-! ## One command -/

section cmd
variable {U : Type*} {P : ℕ} {I : List Sym} {iprog : Fin P → IStmt HS P}

/-- The unit programs and their specifications. -/
structure UnitsOK (lin : List U → List U → (Fin HS → List ℤ) → (Fin HS → List ℤ))
    (OkU : U → Prop) (OkA : ℤ → Prop) (KG : (Fin HS → List ℤ) → Prop) (g : ℤ → U → ℤ)
    (dr rd it bk : IStmt HS P) (cdr crd cit cbk : ℕ) : Prop where
  hL : LaneOK lin
  hdr : DrOK (I := I) (iprog := iprog) lin OkU dr cdr
  hrd : RdOK (I := I) (iprog := iprog) lin OkU rd crd
  hit : ItOK (I := I) (iprog := iprog) lin OkU OkA KG g it cit
  hbk : BkOK (I := I) (iprog := iprog) lin OkU bk cbk

/-- Emit a slab: read it to the temps while emitting, then move it back. -/
def emitP (rd it bk : IStmt HS P) : IStmt HS P :=
  .seq (rptX rM lA lC lD (.seq rd it)) (rptX rM lA lC lD bk)

/-- **One command.** -/
def cmdP (dr rd it bk : IStmt HS P) : IStmt HS P :=
  .seq (.op (.mv rC rOP)) (.seq (.op (.mv rC rARG)) (.seq
    (.ite rOP (fun c => c == .neg) (rptX rM lA lC lD dr)
      (.ite rOP (fun c => c == .zero) (rptX rM lA lC lD (.op (.const rACC 0))) (emitP rd it bk)))
    (.seq (.op (.mv rOP rCU)) (.op (.mv rARG rCU)))))

/-- The cost bound of one command (`cu = cdr + crd + cit + cbk`). -/
def cmdC (m Ba cu : ℕ) : ℕ := 4 * (2 * moveCost Ba + 19) + 50 + (m + 1) * (cu + 12002)

variable {lin : List U → List U → (Fin HS → List ℤ) → (Fin HS → List ℤ)} {OkU : U → Prop}
  {OkA : ℤ → Prop} {KG : (Fin HS → List ℤ) → Prop} {g : ℤ → U → ℤ} {dr rd it bk : IStmt HS P}
  {cdr crd cit cbk : ℕ}

end cmd

/-! ## The command loop -/

/-- The control list of a command list. -/
def flatC (cs : List (ℤ × ℤ)) : List ℤ := cs.flatMap fun c => [c.1, c.2]

@[simp] theorem flatC_nil : flatC [] = [] := rfl

section loop
variable {U : Type*} {P : ℕ} {I : List Sym} {iprog : Fin P → IStmt HS P}
  {lin : List U → List U → (Fin HS → List ℤ) → (Fin HS → List ℤ)} {OkU : U → Prop}
  {OkA : ℤ → Prop} {KG : (Fin HS → List ℤ) → Prop} {g : ℤ → U → ℤ} {dr rd it bk : IStmt HS P}
  {cdr crd cit cbk : ℕ}

/-- The command loop. -/
def cmdLoop (dr rd it bk : IStmt HS P) : IStmt HS P := rptX rNC lA lC lD (cmdP dr rd it bk)

/-- Admissible commands: bounded sizes, admissible arguments for the emits. -/
def CmdsOK (OkA : ℤ → Prop) (Ba : ℕ) (cs : List (ℤ × ℤ)) : Prop :=
  ∀ c ∈ cs, zsz c.1 ≤ Ba ∧ zsz c.2 ≤ Ba ∧ (0 < c.1 → OkA c.2)

/-- The restore loop body. -/
def restP : IStmt HS P := .seq (.op (.mv rCU rC)) (.op (.mv rCU rC))

/-- One block: the command loop, then the restore loop. -/
def blockP (dr rd it bk : IStmt HS P) : IStmt HS P := .seq (cmdLoop dr rd it bk) (rptX rNC lA lC lD restP)

/-- The cost bound of one block. -/
def blockC (L m Ba cu : ℕ) : ℕ :=
  (L + 1) * (cmdC m Ba cu + 3000) + (L + 1) * (2 * (2 * moveCost Ba + 19) + 3000)

end loop

end J6Seg.Slab

end

section
-- Source module: Solutions.J6SegILSlab4
/-!
# B-line IL, B1 (streaming slab pass) §4: the whole pass (generic lanes)

`passP = rptX rNB blockP; rptX rTOT (mv rACC 0)`: the blocks, then the output (collected
reversed on `rACC`) is moved onto stack `0`. **`pass_runs`**: from the lane state with the input
`xs` (`pre` blocks of `n` slabs of `m` units) the output `passOut … ++ G 0` lands on stack `0`; all
skeleton registers are restored (`rACC` empty again). **`passOut_eq`**: for the control list
`ctlOf n n' src argOf` of a monotone source, `passOut` is `slabPass` (block form).
-/

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

namespace J6Seg.Slab

open IntMul NLogN NLogN.KIL Function
open NLogN.KFUU (lA lC lD)
open NLogN.KFT (cf3 rptX)

/-- The output of the pass. -/
def passOut {U : Type*} (m : ℕ) (g : ℤ → U → ℤ) (cs : List (ℤ × ℤ)) (n pre : ℕ) (xs : List U) :
    List ℤ :=
  (List.range pre).flatMap fun i => runC m g cs (xs.drop (i * (n * m)))

/-- **The pass program.** -/
def passP {P : ℕ} (dr rd it bk : IStmt HS P) : IStmt HS P :=
  .seq (rptX rNB lA lC lD (blockP dr rd it bk)) (rptX rTOT lA lC lD (.op (.mv rACC 0)))

/-- The cost bound of the pass. -/
def passC (pre L m Ba cu T Bo : ℕ) : ℕ :=
  (pre + 1) * (blockC L m Ba cu + 3000) + (T + 1) * (2 * moveCost Bo + 19 + 3000)

/-! ## The pass output is the slab pass -/

end J6Seg.Slab

end

section
-- Source module: Solutions.J6SegILSlab5
/-!
# B-line IL, B1 (streaming slab pass) §5: the integer lane, identity and scale items

**The integer lane** `linZ inp tmp G`: the remaining input units on stack `0`, the moved units on
`rT = 50` (both over `G`). Units: `rdZ = mv 0 rT`, `bkZ = mv rT 0`, `drZ = drop 0`; `OkU u` is
`zsz u ≤ Bu`.

**Items** (`KG` trivial):
* identity (passes 3, 5): `itI = dup rT rACC`, `g a u = u`;
* scale (pass 1): `itS = dup rT rACC; shl rACC rARG`, `g a u = u * a` for `a = 2^e`, `e ≤ E`
  (the argument of every emit command is `2^e`).

**`unitsI`**, **`unitsS`**: the `UnitsOK` bundles, with the costs `O(Bu + E + 1)` per unit.
-/

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

namespace J6Seg.Slab

open IntMul NLogN NLogN.KIL Function
open NLogN.KFUU (lA lC lD)
open NLogN.KFT (cf3 rptX)

/-! ## Managed `shl`, `neg`, `add`, `split` -/

/-! ## The integer lane -/

/-- The integer lane. -/
def linZ (inp tmp : List ℤ) (G : Fin HS → List ℤ) : Fin HS → List ℤ :=
  update (update G 0 (inp ++ G 0)) rT (tmp ++ G rT)

section units
variable {P : ℕ} {I : List Sym} {iprog : Fin P → IStmt HS P}

/-- The identity item. -/
def itI : IStmt HS P := .op (.dup rT rACC)

/-- The scale item. -/
def itS : IStmt HS P := .seq (.op (.dup rT rACC)) (.op (.shl rACC rARG))

/-- The scale arguments: powers of two. -/
def PowA (E : ℕ) (a : ℤ) : Prop := ∃ e ≤ E, a = 2 ^ e

end units

end J6Seg.Slab

end

section
-- Source module: Solutions.J6SegILSlab6
/-!
# B-line IL, B1 (streaming slab pass) §6: the pair lane and the output item (pass 7)

**The pair lane** `linP inp tmp G` (units `(y, y₂) : ℤ × ℤ`, the zipped input): first components
on stack `0` / temps `rT = 50`, second components on `rZ = 53` / temps `rT2 = 51`. Units:
`rdP = mv 0 rT; mv rZ rT2`, `bkP` (reverse), `drP = drop 0; drop rZ`.

**The output item** `itO pM` (constants `2^k₁` on `rK1 = 41`, `2^k₂` on `rK2 = 42`, scratch
`43, 46, 54`, multiplication via `mulC pM` / `ARGM = 55`):
`g a (y, y₂) = (a * ((y - y₂) / 2^k₁)).tdiv (2^k₂)` (`outG`), i.e. `outE` with `a = d̃_l`,
`k₁ = P - p`, `k₂ = p`. The floor division is the truncating `split` corrected by the sign of
the remainder; the cost is `O(cU Z)` with `Z = Bu + Ba + 2` (one multiplication).
-/

set_option linter.unusedSectionVars false
set_option linter.unusedVariables false
set_option linter.unusedSimpArgs false
set_option maxRecDepth 20000
set_option maxHeartbeats 1000000

namespace J6Seg.Slab

open IntMul NLogN NLogN.KIL Function
open NLogN.KFUU (lA lC lD cU)
open NLogN.KFT (cf3 rptX)

/-! ## Managed `neg`, `add`, `split`, `mulC` -/

/-! ## The pair lane -/

/-- The pair lane. -/
def linP (inp tmp : List (ℤ × ℤ)) (G : Fin HS → List ℤ) : Fin HS → List ℤ :=
  update (update (update (update G 0 (inp.map Prod.fst ++ G 0)) rZ (inp.map Prod.snd ++ G rZ))
    rT (tmp.map Prod.fst ++ G rT)) rT2 (tmp.map Prod.snd ++ G rT2)

/-- Both components bounded. -/
def OkP (Bu : ℕ) (u : ℤ × ℤ) : Prop := zsz u.1 ≤ Bu ∧ zsz u.2 ≤ Bu

section units
variable {P : ℕ} {I : List Sym} {iprog : Fin P → IStmt HS P}

def rdP : IStmt HS P := .seq (.op (.mv 0 rT)) (.op (.mv rZ rT2))
def bkP : IStmt HS P := .seq (.op (.mv rT 0)) (.op (.mv rT2 rZ))
def drP : IStmt HS P := .seq (.op (.drop 0)) (.op (.drop rZ))

end units

/-! ## The output item -/

/-- The output map of an entry pair. -/
def outG (k₁ k₂ : ℕ) (a : ℤ) (u : ℤ × ℤ) : ℤ := (a * ((u.1 - u.2) / 2 ^ k₁)).tdiv (2 ^ k₂)

/-- The constants of the output item. -/
def KO (k₁ k₂ : ℕ) (G : Fin HS → List ℤ) : Prop :=
  (∃ l, G rK1 = (2 : ℤ) ^ k₁ :: l) ∧ (∃ l, G rK2 = (2 : ℤ) ^ k₂ :: l)

section item
variable {P : ℕ} {I : List Sym} {iprog : Fin P → IStmt HS P}

/-- The floor correction. -/
def fixP : IStmt HS P :=
  .ite 43 (fun c => c == .neg) (.seq (.op (.const 46 (-1))) (.op (.add 54 46 54))) .skip

/-- **The output item.** -/
def itO (pM : Fin P) : IStmt HS P :=
  .seq (.op (.dup rT 43)) (.seq (.op (.dup rT2 46)) (.seq (.op (.neg 46)) (.seq (.op (.add 43 46 54))
  (.seq (.op (.split 54 rK1 43)) (.seq fixP (.seq (.op (.drop 43))
  (.seq (NLogN.KFUU.mulC pM 54 rARG 43) (.seq (.op (.drop 54)) (.seq (.op (.split 43 rK2 46))
  (.seq (.op (.drop 46)) (.op (.mv 43 rACC))))))))))))

/-- The cost bound of the output item. -/
def cItO (Z k₁ k₂ : ℕ) : ℕ := 12 * cU Z + 200 * (k₁ + 1) + 200 * (k₂ + 1) + 100

end item

end J6Seg.Slab

end

section
-- Source module: Solutions.J6SegILSlab8
/-!
# B-line IL, B1 (streaming slab pass) §8: the slab passes of `axisPipe` on `GI` arrays

The three IL slab passes, on `gflat` arrays (`post` `GI` entries = `m = 2 post` integers per slab):

* **`slabId_runs`** (passes 3, 5: `h = fun _ z => z`; control `ctlOf n n' src (fun _ => 0)`),
* **`slabSc_runs`** (pass 1: `h = selH p P`; control `ctlOf n n' src (fun _ => 2^(P-p-1))`),
* **`slabOut_runs`** (pass 7: `h = outH p P s t α` on `zip Y Y₂`, source `some l`; control
  `ctlOf N N some argO` with `argO l = d̃_l` for `l < s`, else `0`; the two arrays on stacks
  `0` and `rZ`; constants `2^(P-p)`, `2^p` on `rK1`, `rK2`; multiplication `MulOK pM`).

Each takes the skeleton inputs (`rC` = `flatC cs ++ RC`, `rACC = []`, `rNB = pre`, `rNC = |cs|`,
`rM = 2 post`, `rTOT = pre · n' · 2 post`, the `rptX` counters empty) and leaves
`gflat (slabPass …) ++ l` on stack `0`, everything else unchanged. The cost is `passC …`;
with `passC_le` and `|cs| ≤ n + n'` (`length_ctlOf_le`) it is, per block,
`O((n + n' + 1) (post + 1) (cu + Ba + 1))` with `cu = O(Bu + 1)` (identity),
`O(Bu + P + 1)` (scale), `O(cU (Bu + Ba + 2) + P)` (output: one multiplication per entry).
-/

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

namespace J6Seg.Slab

open IntMul NLogN NLogN.KIL Function
open NLogN.KFB (GI gflat)
open NLogN.KFUU (lA lC lD cU)
open NLogN.KFT (cf3 rptX)

/-! ## The commands of `ctlOf` -/

/-! ## Flattening -/

section inst
variable {P : ℕ} {I : List Sym} {iprog : Fin P → IStmt HS P}

/-- The output arguments. -/
def argO (p s t α : ℕ) (l : ℕ) : ℤ := if l < s then IntMul.HvdH.dTilde p s t α l else 0

end inst

end J6Seg.Slab

end

section
-- Source module: Solutions.J6SegILCtl1
/-!
# B-line IL, B1 follow-up: building the control list, §1 the list level

The control builder streams a **source table** `srcTab src argOf n'` (two integers per output
slab `l < n'`: `code (src l)` (`-1` for `none`, `k` for `some k`) and `argOf l`) and produces
`ctlOf n n' src argOf`. The list facts:

* `ctlAux_snoc`: `ctlAux` on `ls ++ [l]` is `ctlAux` on `ls` followed by the step of `l`
  (`stepC`: `none ↦ [(0,0)]`, `some k ↦ (k - cur)` drops then `(1, argOf l)`; `stepK`: the new
  position).
* `flatC_snoc`, `srcTab_succ`: the table and the control list grow one entry at a time.
-/

set_option linter.unusedSectionVars false
set_option linter.unusedVariables false

namespace J6Seg.Slab

/-- The code of a source slab. -/
def code : Option ℕ → ℤ
  | none => -1
  | some k => k

/-- **The source table** of a slab pass: `[code (src l), argOf l]` for `l < n'`. -/
def srcTab (src : ℕ → Option ℕ) (argOf : ℕ → ℤ) (n' : ℕ) : List ℤ :=
  (List.range n').flatMap fun l => [code (src l), argOf l]

/-- The commands of one output slab from position `cur`. -/
def stepC (src : ℕ → Option ℕ) (argOf : ℕ → ℤ) (l cur : ℕ) : List (ℤ × ℤ) :=
  match src l with
  | none => [(0, 0)]
  | some k => List.replicate (k - cur) J6Seg.dropC ++ [(1, argOf l)]

/-- The position after one output slab. -/
def stepK (src : ℕ → Option ℕ) (l cur : ℕ) : ℕ :=
  match src l with
  | none => cur
  | some k => k

end J6Seg.Slab

end

section
-- Source module: Solutions.J6SegILCtl2
/-!
# B-line IL, B1 follow-up: building the control list, §2 one table entry

**Register convention of the control builder** (all below `56`, none `0`, none in `26–29`;
`lA lC lD = 34–36` are the `rptX` counters, empty in the managed state):

| stack | role |
|---|---|
| `bX = 44` | the source table `srcTab src argOf n'` (input, restored) |
| `bN' = 20`, `bN = 21` | `n'`, `n` (inputs, one value each, preserved) |
| `bSv = 25` | the consumed table entries (scratch, emptied by the restore) |
| `bCur = 30` | the input position (scratch) |
| `bT = 33` | the reversed control list being built (scratch, emptied by the transfer) |
| `bCd = 24`, `bA = 23`, `bW = 22`, `bD = 32` | scratch of one entry |
| `rC = 45`, `rNC = 38` | outputs: `flatC cs` pushed on `rC`, `|cs|` pushed on `rNC` |

Here: `loopMV` (a variable-cost `rptX` over managed states) and **`entry_runs`** (one table
entry: the commands of one output slab pushed reversed onto `bT`, the position updated, the entry
saved on `bSv`).
-/

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

namespace J6Seg.Slab

open IntMul NLogN NLogN.KIL Function
open NLogN.KFUU (lA lC lD cU)
open NLogN.KFT (cf3 rptX)

abbrev bX : Fin HS := 44
abbrev bN' : Fin HS := 20
abbrev bN : Fin HS := 21
abbrev bW : Fin HS := 22
abbrev bA : Fin HS := 23
abbrev bCd : Fin HS := 24
abbrev bSv : Fin HS := 25
abbrev bCur : Fin HS := 30
abbrev bD : Fin HS := 32
abbrev bT : Fin HS := 33

/-! ## List level: one step -/

/-- The commands of one output slab with source `o` from position `cur`. -/
def stO (o : Option ℕ) (a : ℤ) (cur : ℕ) : List (ℤ × ℤ) :=
  match o with
  | none => [(0, 0)]
  | some k => List.replicate (k - cur) J6Seg.dropC ++ [(1, a)]

/-- The position after it. -/
def kO (o : Option ℕ) (cur : ℕ) : ℕ :=
  match o with
  | none => cur
  | some k => k

/-! ## Managed-state helpers -/

/-! ## The managed state of the builder -/

/-- The builder's registers over `G`. -/
def cS (G : Fin HS → List ℤ) (X Sv Cur T Cd A D W : List ℤ) : Fin HS → List ℤ :=
  update (update (update (update (update (update (update (update G bX X) bSv Sv) bCur Cur) bT T) bCd Cd) bA A) bD D) bW W

section cs
variable (G : Fin HS → List ℤ) (X Sv Cur T Cd A D W : List ℤ) (v : List ℤ)
@[simp] theorem cS_X : cS G X Sv Cur T Cd A D W bX = X := by simp +decide [cS]
@[simp] theorem cS_Sv : cS G X Sv Cur T Cd A D W bSv = Sv := by simp +decide [cS]
@[simp] theorem cS_Cur : cS G X Sv Cur T Cd A D W bCur = Cur := by simp +decide [cS]
@[simp] theorem cS_T : cS G X Sv Cur T Cd A D W bT = T := by simp +decide [cS]
@[simp] theorem cS_Cd : cS G X Sv Cur T Cd A D W bCd = Cd := by simp +decide [cS]
@[simp] theorem cS_A : cS G X Sv Cur T Cd A D W bA = A := by simp +decide [cS]
@[simp] theorem cS_D : cS G X Sv Cur T Cd A D W bD = D := by simp +decide [cS]
@[simp] theorem cS_W : cS G X Sv Cur T Cd A D W bW = W := by simp +decide [cS]
@[simp] theorem cS_lA : cS G X Sv Cur T Cd A D W lA = G lA := by simp +decide [cS]
@[simp] theorem cS_lC : cS G X Sv Cur T Cd A D W lC = G lC := by simp +decide [cS]
@[simp] theorem cS_lD : cS G X Sv Cur T Cd A D W lD = G lD := by simp +decide [cS]
theorem cS_upd_X : update (cS G X Sv Cur T Cd A D W) bX v = cS G v Sv Cur T Cd A D W := by
  funext x; by_cases hx : x = bX
  · subst hx; simp +decide [cS]
  · simp [cS, update_apply, hx]
theorem cS_upd_Sv : update (cS G X Sv Cur T Cd A D W) bSv v = cS G X v Cur T Cd A D W := by
  funext x; by_cases hx : x = bSv
  · subst hx; simp +decide [cS]
  · simp [cS, update_apply, hx]
theorem cS_upd_Cur : update (cS G X Sv Cur T Cd A D W) bCur v = cS G X Sv v T Cd A D W := by
  funext x; by_cases hx : x = bCur
  · subst hx; simp +decide [cS]
  · simp [cS, update_apply, hx]
theorem cS_upd_T : update (cS G X Sv Cur T Cd A D W) bT v = cS G X Sv Cur v Cd A D W := by
  funext x; by_cases hx : x = bT
  · subst hx; simp +decide [cS]
  · simp [cS, update_apply, hx]
theorem cS_upd_Cd : update (cS G X Sv Cur T Cd A D W) bCd v = cS G X Sv Cur T v A D W := by
  funext x; by_cases hx : x = bCd
  · subst hx; simp +decide [cS]
  · simp [cS, update_apply, hx]
theorem cS_upd_A : update (cS G X Sv Cur T Cd A D W) bA v = cS G X Sv Cur T Cd v D W := by
  funext x; by_cases hx : x = bA
  · subst hx; simp +decide [cS]
  · simp [cS, update_apply, hx]
theorem cS_upd_D : update (cS G X Sv Cur T Cd A D W) bD v = cS G X Sv Cur T Cd A v W := by
  funext x; by_cases hx : x = bD
  · subst hx; simp +decide [cS]
  · simp [cS, update_apply, hx]
theorem cS_upd_W : update (cS G X Sv Cur T Cd A D W) bW v = cS G X Sv Cur T Cd A D v := by
  funext x; by_cases hx : x = bW
  · subst hx; simp +decide [cS]
  · simp [cS, update_apply, hx]
end cs

/-! ## One entry -/

section entry
variable {P : ℕ} {I : List Sym} {iprog : Fin P → IStmt HS P}

/-- One drop command, pushed reversed. -/
def dropE : IStmt HS P := .seq (.op (.const bT (-1))) (.op (.const bT 0))

/-- The `none` branch: the zero command. -/
def zB : IStmt HS P := .seq (.op (.const bT 0)) (.op (.const bT 0))

/-- The `some k` branch: `k - cur` drops, the emit command, the new position `k`. -/
def eB : IStmt HS P :=
  .seq (.op (.dup bCd bW)) (.seq (.op (.dup bCur bD)) (.seq (.op (.neg bD)) (.seq
    (.op (.add bW bD bD)) (.seq (rptX bD lA lC lD dropE) (.seq (.op (.drop bD)) (.seq
    (.op (.const bT 1)) (.seq (.op (.dup bA bT)) (.seq (.op (.drop bCur))
    (.op (.dup bCd bCur))))))))))

/-- **One table entry.** -/
def entP : IStmt HS P :=
  .seq (.op (.mv bX bCd)) (.seq (.op (.mv bX bA)) (.seq
    (.ite bCd (fun c => c == .neg) zB eB) (.seq (.op (.mv bCd bSv)) (.op (.mv bA bSv)))))

/-- The cost bound of one entry with `d` drops. -/
def cEnt (Z d : ℕ) : ℕ := 21 * cU Z + 4000 * (d + 1)

end entry

end J6Seg.Slab

end

section
-- Source module: Solutions.J6SegILCtl3
/-!
# B-line IL, B1 follow-up: building the control list, §3 list facts and the pair transfer

* `pos src argOf j`: the input position after the first `j` output slabs; `pos_succ`, `pos_mono`,
  `pos_le_n`, `sum_tele` (the drops of the entries telescope to `pos`).
* `length_ctlOf`: `|ctlOf n n' src argOf| = n + n'` for a monotone source.
* `srcTab_drop`: the table seen from entry `j`.
* `pairMove_runs`: `rptX X (mv i j; mv i j)` moves a reversed list of `2L` entries back in order.
* `cS_upd_G`, `cS_self`: the builder state over other registers.
-/

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

namespace J6Seg.Slab

open IntMul NLogN NLogN.KIL Function
open NLogN.KFUU (lA lC lD cU)
open NLogN.KFT (cf3 rptX)

/-! ## Positions -/

/-- The input position after the first `j` output slabs. -/
def pos (src : ℕ → Option ℕ) (argOf : ℕ → ℤ) (j : ℕ) : ℕ :=
  (J6Seg.ctlAux src argOf (List.range j) 0).2

variable {n n' : ℕ} {src : ℕ → Option ℕ} (argOf : ℕ → ℤ)

/-! ## The pair transfer -/

/-! ## The builder state over other registers -/

end J6Seg.Slab

end