Harvey–van der Hoeven machine development, part 7: definitions for the final O(n log n) assembly

-- Harvey–van der Hoeven multitape-TM development (NLogN), generated by skeleton subtraction.

import Definitions.Def_IntMul_HvdH_M_Defs6

set_option autoImplicit false
set_option relaxedAutoImplicit false
set_option maxHeartbeats 1000000

section
-- Source module: NLogN.HAx2
/-!
# One resampling axis, part 2: the `sapL` slice program as a `GL.LProg`

Stack layout of the `S̃` axis (all outside `GL.myW = 0 … 30`):

| stack | content                    | role                                  |
|-------|----------------------------|---------------------------------------|
| 31    | `sOff s w`                 | `Ka` of `sapP`                        |
| 32    | `s + (2w+1)`               | `KN`                                  |
| 33    | slice input                | `CU` (= `LProg.I`)                    |
| 34    | `s`                        | `Ks`, `TS`, axis `rn`                 |
| 35    | `2^(2p)`                   | `BR`                                  |
| 36    | `t`                        | `BN`, `TT`, axis `rn'`                |
| 37    | `2w+1`                     | `BNR`, `TRI`                          |
| 38    | `t (2w+1)`                 | `BNW`                                 |
| 39    | the weight table           | `BW`                                  |
| 40    | the step table             | `BD`                                  |
| 42/43/44 | `m`, `kb`, `m t`        | axis `rm`, `rkb`, `rtot`              |
| 45    | slice output               | `LProg.O`                             |
| 46    | `α² t²`                    | `TK2`                                 |
| 48/51/52 | `p`, `w`, `betaT p α`   | `TP`, `TM`, `TBe`                     |

Here `w = sWidth p α`.  `sapP` is relocated by `σS`, the table builders by `σB`.
-/

set_option linter.unusedSimpArgs false

namespace NLogN.AX

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

abbrev aKa : Fin HS := 31
abbrev aKN : Fin HS := 32
abbrev aI : Fin HS := 33
abbrev aS : Fin HS := 34
abbrev aR : Fin HS := 35
abbrev aT : Fin HS := 36
abbrev aNR : Fin HS := 37
abbrev aNW : Fin HS := 38
abbrev aW : Fin HS := 39
abbrev aD : Fin HS := 40
abbrev aM : Fin HS := 42
abbrev aKB : Fin HS := 43
abbrev aTot : Fin HS := 44
abbrev aO : Fin HS := 45
abbrev aK2 : Fin HS := 46
abbrev aP : Fin HS := 48
abbrev aWd : Fin HS := 51
abbrev aBe : Fin HS := 52

/-- Relocation of `sapP`. -/
def σS : Equiv.Perm (Fin HS) :=
  swapsPerm [(21, 33), (30, 34), (13, 35), (14, 36), (15, 37), (16, 38), (6, 39), (8, 40), (9, 41)]

/-- The `sapL` slice program: `sapP` in place on `aI`, then transfer to `aO`. -/
def sapSlP : RM := .seq (RM.ren σS sapP) (xferP aI aO (cnt2 aT))

/-- The slice context of the `S̃` axis. -/
def CtxS (p s t α : ℕ) (st : IState HS) : Prop :=
  (∃ L, st.G aS = .int (s : ℤ) :: L) ∧
  (∃ L, st.G aKa = .int (sOff s (HvdH.sWidth p α) : ℤ) :: L) ∧
  (∃ L, st.G aKN = .int ((s + (2 * HvdH.sWidth p α + 1) : ℕ) : ℤ) :: L) ∧
  (∃ L, st.G aR = .int (2 ^ (2 * p)) :: L) ∧
  (∃ L, st.G aT = .int (t : ℤ) :: L) ∧
  (∃ L, st.G aNR = .int ((2 * HvdH.sWidth p α + 1 : ℕ) : ℤ) :: L) ∧
  (∃ L, st.G aNW = .int ((t * (2 * HvdH.sWidth p α + 1) : ℕ) : ℤ) :: L) ∧
  (∃ R, st.G aW = ints (sTab p s t α).flatten ++ R) ∧
  (∃ R, st.G aD = ints ((sDel s t).map (Nat.cast : ℕ → ℤ)) ++ R)

/-! ## Entry bounds of the `S` table -/

/-! ## Output bounds -/

/-! ## The relocated `sapP` -/

/-! ## The slice program as a `GL.LProg` -/

/-- Arithmetic side conditions of the `S̃` slice (input bound `Bi`, output bound `Bo`,
multiplication operand bound `mm`, accumulator bound `Bacc`, list-move bound `Q`). -/
structure SapHyp (p s t α Bi Bo mm Bacc Q : ℕ) : Prop where
  hs : 0 < s
  hst : s ≤ t
  hBi : Bi ≤ mm
  hW : (4 * p + 35) * 2 ^ (2 * p) < 2 ^ mm
  hAcc : (2 * HvdH.sWidth p α + 2) * 2 ^ (2 * mm) ≤ 2 ^ Bacc
  hmB : 2 * mm ≤ Bacc
  hLB : zsz ((s + (2 * HvdH.sWidth p α + 1) : ℕ) : ℤ) ≤ Bacc
  hpB : 2 * p + 1 ≤ Bacc
  hnB : zsz ((t * (2 * HvdH.sWidth p α + 1) : ℕ) : ℤ) ≤ Bacc
  htB : zsz (t : ℤ) ≤ Bacc
  hc : BD.cS1 p α * 2 ^ Bi ≤ 2 ^ Bo
  hBo : Bo ≤ Q
  htQ : zsz ((2 * t : ℕ) : ℤ) ≤ Q

/-- Cost bound of one `S̃` slice. -/
def sapSlC (p s t α mm Φ Bacc Q : ℕ) : ℕ :=
  gresCost s (2 * HvdH.sWidth p α + 1) t mm Φ Bacc 1 + (16000 * (Q + 2) + t * 60000 * (Q + 2) + 64)

end NLogN.AX

end

section
-- Source module: NLogN.HAx3
/-!
# One resampling axis, part 3: tables of the `S̃` axis, and the axis program `sapAxP`

* `sTabR = sTabP.ren σB`, `dTabR = dTabP.ren σB` push the `S` weight table on `aW = 39` and the
  step table on `aD = 40` (registers: see `HAx2`).
* `popK` removes them again.
-/

set_option linter.unusedSimpArgs false

namespace NLogN.AX

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

/-- Relocation of the table builders. -/
def σB : Equiv.Perm (Fin HS) := swapsPerm [(49, 34), (50, 36), (23, 37), (22, 46), (6, 39), (8, 40)]

/-- Arithmetic side conditions of the table builders. -/
structure TabHyp (p s t α Bq : ℕ) : Prop where
  hBq : 4 * (s + t + 2 * HvdH.sWidth p α + 2) ^ 4 < 2 ^ Bq
  hEq : (4 * p + 35) * 2 ^ (2 * p) < 2 ^ Bq
  hK2 : α ^ 2 * t ^ 2 < 2 ^ Bq

/-- Registers read by the `S` table builder. -/
structure STabRegs (p s t α : ℕ) (st : IState HS) : Prop where
  rP : ∃ L, st.G aP = .int (p : ℤ) :: L
  rS : ∃ L, st.G aS = .int (s : ℤ) :: L
  rT : ∃ L, st.G aT = .int (t : ℤ) :: L
  rM : ∃ L, st.G aWd = .int (HvdH.sWidth p α : ℤ) :: L
  rB : ∃ L, st.G aBe = .int (HvdH.betaT p α) :: L
  rK : ∃ L, st.G aK2 = .int ((α ^ 2 * t ^ 2 : ℕ) : ℤ) :: L
  rR : ∃ L, st.G aNR = .int ((2 * HvdH.sWidth p α + 1 : ℕ) : ℤ) :: L

def sTabR : RS := RS.ren σB sTabP
def dTabR : RS := RS.ren σB dTabP

end NLogN.AX

end

section
-- Source module: NLogN.HAx4
/-!
# One resampling axis, part 4: the `S̃` axis program `sapAxP`

`sapAxP` on `gD = gints L ++ R` (with `L` of dims `s :: rest`, `m = rest.prod`) computes
`gints (tensorGo [(t, sapL p s t α)] (s :: rest) L) ++ R`: it builds the `S` weight table on
`aW = 39` and the step table on `aD = 40`, runs `GL.axisP` with the slice program `sapSl`, and
pops the two tables again.  Every stack except `gD` is restored.
-/

set_option linter.unusedSimpArgs false

namespace NLogN.AX

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

/-- The axis core (= `GL.axisP (sapAx …)`, written out). -/
def sapCoreP : RM :=
  .seq (RM.seqL [.rs (.push 1 (.reg aS)), .rs (.push 2 (.reg aM)), .rs (.push 3 (.reg aKB)),
      .rs (transP 2), .rs (.pop 1), .rs (.pop 2), .rs (.pop 3)])
    (.seq (.rep gCC (.reg aM) (RM.seqL [.rs (mvList gD gZT gCV (cnt2 aS)),
        .rs (mvList gZT aI gCV (cnt2 aS)), sapSlP, .rs (mvList aO gZB gCV (cnt2 aT))]))
      (.rs (mvList gZB gD gCV (cnt2 aTot))))

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

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

/-- Size conditions of the list moves / transposition (`Q` bounds every moved entry). -/
structure AxQ (s t m kb Bi Bo Q : ℕ) : Prop where
  hm : 1 ≤ m
  hkb : s * m ≤ 2 ^ kb
  bB : Bi ≤ Q
  bBo : Bo ≤ Q
  bn : zsz ((2 * s : ℕ) : ℤ) ≤ Q
  bn' : zsz ((2 * t : ℕ) : ℤ) ≤ Q
  bm : zsz ((2 * m : ℕ) : ℤ) ≤ Q
  bkb : zsz (kb : ℤ) ≤ Q
  btot : zsz ((2 * (m * t) : ℕ) : ℤ) ≤ Q

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

end NLogN.AX

end

section
-- Source module: NLogN.HAx5
/-!
# One resampling axis, part 5: the `bapL` slice program as a `GL.LProg`

Stack layout of the `B̃` axis (all outside `GL.myW = 0 … 30`); `e = eWidth p α`,
`J = jIter p s t α`, `u = ((s : ZMod t)⁻¹).val`:

| stack | content (at entry)          | role                                        |
|-------|-----------------------------|---------------------------------------------|
| 31    | `sOff s e`                  | `Ka`                                        |
| 32    | `s + (2e+1)`                | `KN`                                        |
| 33    | (`D̃'` table, built)          | `W3`                                        |
| 34    | `t`                         | stack `1`; `TT`; axis `rn`                  |
| 35    | `kb`                        | stack `2`; axis `rkb`                       |
| 36    | (`P̃_t` keys, built)          | stack `3`                                   |
| 37    | (`E` table, built)          | `BW`                                        |
| 38    | (`E` steps, built)          | `BD`                                        |
| 39    | `2^p`                       | `BR`                                        |
| 40    | `s`                         | `BN`                                        |
| 41    | `J`                         | `JN`                                        |
| 42    | `2e+1`                      | `BNR`; `TRI`                                |
| 43    | `2s`                        | `JL`                                        |
| 44    | `s`                         | `PS`; `TS`; axis `rn'`                      |
| 45    | (`C̃` steps, built)           | `D1`                                        |
| 46    | (`C̃` table, built)           | `W1`                                        |
| 47    | `u` (then `P̃_s⁻¹` keys)      | `TU` for the key builder, then `KS`         |
| 48    | `s (2e+1)`                  | `BNW`                                       |
| 49    | `α²`                        | `TBe` for the builders; slice input `CU`    |
| 50    | `2^p`                       | `PR`                                        |
| 51    | `s`                         | `Ks`                                        |
| 52    | `p`                         | `TP` for the builders; slice output         |
| 53/54 | `m`, `m s`                  | axis `rm`, `rtot`                           |
-/

set_option linter.unusedSimpArgs false

namespace NLogN.AX

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

abbrev bKa : Fin HS := 31
abbrev bKN : Fin HS := 32
abbrev bW3 : Fin HS := 33
abbrev bT : Fin HS := 34
abbrev bKB : Fin HS := 35
abbrev bK3 : Fin HS := 36
abbrev bW : Fin HS := 37
abbrev bD : Fin HS := 38
abbrev bR : Fin HS := 39
abbrev bN : Fin HS := 40
abbrev bJN : Fin HS := 41
abbrev bNR : Fin HS := 42
abbrev bJL : Fin HS := 43
abbrev bS : Fin HS := 44
abbrev bD1 : Fin HS := 45
abbrev bW1 : Fin HS := 46
abbrev bKS : Fin HS := 47
abbrev bNW : Fin HS := 48
abbrev bI : Fin HS := 49
abbrev bPR : Fin HS := 50
abbrev bKs : Fin HS := 51
abbrev bO : Fin HS := 52
abbrev bM : Fin HS := 53
abbrev bTot : Fin HS := 54

/-- Relocation of `bapP`. -/
def σP : Equiv.Perm (Fin HS) :=
  swapsPerm [(0, 33), (1, 34), (2, 35), (3, 36), (6, 37), (8, 38), (13, 39), (14, 40), (15, 42),
    (16, 48), (21, 49), (24, 50), (30, 51), (9, 22)]

/-- The `bapL` slice program: `bapP` in place on `bI`, then transfer to `bO`. -/
def bapSlP : RM := .seq (RM.ren σP bapP) (xferP bI bO (cnt2 bS))

/-- The slice context of the `B̃` axis. -/
structure CtxB (p s t α kb : ℕ) (st : IState HS) : Prop where
  t1 : ∃ L, st.G bT = .int (t : ℤ) :: L
  kb2 : ∃ L, st.G bKB = .int (kb : ℤ) :: L
  k3 : ∃ R, st.G bK3 = natInts (ptKeys s t) ++ R
  kS : ∃ R, st.G bKS = natInts (psKeys s t) ++ R
  ps : ∃ L, st.G bS = .int (s : ℤ) :: L
  pr : ∃ L, st.G bPR = .int (2 ^ p) :: L
  ks : ∃ L, st.G bKs = .int (s : ℤ) :: L
  ka : ∃ L, st.G bKa = .int (sOff s (HvdH.eWidth p α) : ℤ) :: L
  kn : ∃ L, st.G bKN = .int ((s + (2 * HvdH.eWidth p α + 1) : ℕ) : ℤ) :: L
  br : ∃ L, st.G bR = .int (2 ^ p) :: L
  bn : ∃ L, st.G NLogN.AX.bN = .int (s : ℤ) :: L
  bnr : ∃ L, st.G bNR = .int ((2 * HvdH.eWidth p α + 1 : ℕ) : ℤ) :: L
  bnw : ∃ L, st.G bNW = .int ((s * (2 * HvdH.eWidth p α + 1) : ℕ) : ℤ) :: L
  bw : ∃ R, st.G bW = ints (eTab p s t α).flatten ++ R
  bd : ∃ R, st.G bD = ints ((eDel s).map (Nat.cast : ℕ → ℤ)) ++ R
  jl : ∃ L, st.G bJL = .int ((2 * s : ℕ) : ℤ) :: L
  jn : ∃ L, st.G bJN = .int (HvdH.jIter p s t α : ℤ) :: L
  w1 : ∃ R, st.G bW1 = ints (cTab s).flatten ++ R
  d1 : ∃ R, st.G bD1 = ints ((cDel s t).map (Nat.cast : ℕ → ℤ)) ++ R
  w3 : ∃ R, st.G bW3 = ints (dTab p s t α).flatten ++ R

/-! ## Magnitudes of the intermediate lists -/

/-- Arithmetic side conditions of the `B̃` slice (input bound `Bi`, output bound `Bo`,
multiplication operand bound `mm`, accumulator bound `Bacc`, list-move bound `Q`). -/
structure BapHyp (p s t α kb Bi Bo mm Bacc Q : ℕ) : Prop where
  hs : 0 < s
  hst : s ≤ t
  hco : Nat.Coprime s t
  hkb : t ≤ 2 ^ kb
  hBi : Bi ≤ mm
  hmm : 1 ≤ mm
  hW : (4 * p + 35) * 2 ^ (2 * p) < 2 ^ mm
  hJc : (HvdH.jIter p s t α + 1) * BD.cE1 p α ^ HvdH.jIter p s t α * 2 ^ Bi ≤ 2 ^ mm
  hmB : 2 * mm ≤ Bacc
  hAcc : (2 * HvdH.eWidth p α + 2) * 2 ^ (2 * mm) ≤ 2 ^ Bacc
  hpB : p + 1 ≤ Bacc
  hT : zsz ((t + 1 : ℕ) : ℤ) ≤ Bacc
  hSW : zsz (((s + 1) * (2 * HvdH.eWidth p α + 1) : ℕ) : ℤ) ≤ Bacc
  hJ : zsz (HvdH.jIter p s t α : ℤ) ≤ Bacc
  hc : BD.cB1 p α (HvdH.jIter p s t α) * 2 ^ Bi ≤ 2 ^ Bo
  hBo : Bo ≤ Q
  hsQ : zsz ((2 * s : ℕ) : ℤ) ≤ Q

/-! ## The slice program as a `GL.LProg` -/

/-- Cost bound of one `B̃` slice. -/
def bapSlC (p s t α kb mm Φ Bacc Q : ℕ) : ℕ :=
  bapCost p s t α kb mm Φ Bacc + (16000 * (Q + 2) + s * 60000 * (Q + 2) + 64)

end NLogN.AX

end

section
-- Source module: NLogN.HRspDTab
/-!
# Building the `D̃'` table `dTab p s t α` (one row `dTilde p s t α l`, `l < s`) on `W3`

Scalars: `TP = p`, `TS = s`, `TT = t`, `TK2 = s²`, `TBe = α²`, `TSg = 2α² - 2`.
-/

set_option linter.unusedSimpArgs false

namespace NLogN.Rsp

open IntMul NLogN RS

/-- `2α² - 2` (the `σ` argument of `expPos`). -/
abbrev TSg : Fin HS := 25

/-- `t l - s r_l`. -/
def dDE : Ex HS := .add (.mul (.reg NLogN.Rsp.TT) kE) (.neg (.mul (.reg TS) cRE))
/-- `α² (t l - s r_l)²`. -/
def XDE : Ex HS := .mul (.reg TBe) (.mul dDE dDE)

/-- The `D̃'` weight program. -/
def dTWP : RS :=
  NLogN.RS.seq (push 16 (.reg TP)) (NLogN.RS.seq (push 17 (.reg TK2)) (NLogN.RS.seq (push 18 XDE) (NLogN.RS.seq (push 27 (.reg TSg))
    (NLogN.RS.seq expPosP (NLogN.RS.seq (push W3 (.reg 19)) (NLogN.RS.seq (pop 19) (NLogN.RS.seq (pop 27) (NLogN.RS.seq (pop 18)
      (NLogN.RS.seq (pop 17) (pop 16))))))))))

def dTWf (p s t A K2 g : ℕ) (_i k : ℕ) : ℤ := HvdH.expPos p K2 (A * HvdH.kOf s t k) g

def dTWCost (p g Bq : ℕ) : ℕ := 100 * 10 ^ 7 * (Bq + 2) ^ 2 + expPosCost p Bq Bq g

/-- The `D̃'` table builder: pushes `ints (dTab p s t α).flatten` on `W3`. -/
def dTabBP : RS := row1P W3 TS dTWP

end NLogN.Rsp

end

section
-- Source module: NLogN.HRspETab
/-!
# Building the `E` weight table `eTab p s t α` on `BW`

Row `i < 2m+1` (`h = i - m`), column `l < s`: `eW p s t α i l`.
Scalars: `TP = p`, `TS = s`, `TT = t`, `TM = m = eWidth p α`, `TK2 = s²`, `TBe = α²`,
`TRI = 2m+1`; counters `TCI`, `TCK`.
The floor division in `bOf s t (l+h)` (numerator possibly negative) is computed as
`(N + 2s·tm) / (2s) - tm` with a non-negative shifted numerator.
-/

set_option linter.unusedSimpArgs false

namespace NLogN.Rsp

open IntMul NLogN RS

/-- `h = i - m`. -/
def hE : Ex HS := .add iE (.neg (.reg TM))
/-- `l + h`. -/
def xE : Ex HS := .add kE hE
/-- `2t(l+h) + s + 2s·t·m ≥ 0`. -/
def nE : Ex HS :=
  .add (.add (.mul (.add (.reg NLogN.Rsp.TT) (.reg NLogN.Rsp.TT)) NLogN.Rsp.xE) (.reg TS))
    (.mul (.mul (.add (.reg TS) (.reg TS)) (.reg NLogN.Rsp.TT)) (.reg TM))
/-- `⌊(2t(l+h) + s) / (2s)⌋`. -/
def flE : Ex HS := .add (.ndiv nE (.add (.reg TS) (.reg TS))) (.neg (.mul (.reg NLogN.Rsp.TT) (.reg TM)))
/-- `b_{l+h}`. -/
def bxE : Ex HS := .add (.mul (.reg NLogN.Rsp.TT) NLogN.Rsp.xE) (.neg (.mul (.reg TS) flE))
/-- `t h + b_l`. -/
def aE' : Ex HS := .add (.mul (.reg NLogN.Rsp.TT) hE) dDE
/-- `K_{l,h}`. -/
def kzE : Ex HS := .add (.mul aE' aE') (.neg (.mul bxE bxE))
/-- `α² · K_{l,h}.toNat`. -/
def XEE : Ex HS := .mul (.reg TBe) (.sel kzE (.lit 0) kzE)

/-! ## a generic bound for polynomial expressions -/

/-- A degree-like measure: `|e| ≤ D ^ deg e` when all leaves are `≤ D` (`D ≥ 2`). -/
def deg : Ex HS → ℕ
  | .reg _ | .lit _ => 1
  | .add a b => max (deg a) (deg b) + 1
  | .mul a b => deg a + deg b
  | .neg a => deg a
  | .ndiv a b | .nmod a b => max (deg a) (deg b)
  | .sel c x y => max (deg c) (max (deg x) (deg y))
  | _ => 0

/-- Leaves bounded by `D`, divisions of non-negatives by positives, no shifts/powers. -/
def PolyOK (σ : IState HS) (D : ℤ) : Ex HS → Prop
  | .reg i => |σ.top i| ≤ D
  | .lit c => |c| ≤ D
  | .add a b | .mul a b => PolyOK σ D a ∧ PolyOK σ D b
  | .neg a => PolyOK σ D a
  | .ndiv a b | .nmod a b => PolyOK σ D a ∧ PolyOK σ D b ∧ 0 ≤ a.eval σ ∧ 0 < b.eval σ
  | .sel c x y => PolyOK σ D c ∧ PolyOK σ D x ∧ PolyOK σ D y
  | _ => False

end NLogN.Rsp

end

section
-- Source module: NLogN.HRspETab2
/-!
# The `E` weight program and table builder `eTabBP` (pushes `ints (eTab p s t α).flatten` on `BW`)
-/

set_option linter.unusedSimpArgs false

namespace NLogN.Rsp

open IntMul NLogN RS

/-- The weight: `0` for `h = 0`, else the `expNeg` value on `19`. -/
def wEE : Ex HS := .sel (.neg (.mul hE hE)) (.reg 19) (.lit 0)

/-- The `E` weight program. -/
def eWP : RS :=
  NLogN.RS.seq (push 16 (.reg TP)) (NLogN.RS.seq (push 17 (.reg TK2)) (NLogN.RS.seq (push 18 XEE) (NLogN.RS.seq expNegP
    (NLogN.RS.seq (push NLogN.Rsp.BW wEE) (NLogN.RS.seq (pop 19) (NLogN.RS.seq (pop 18) (NLogN.RS.seq (pop 17) (pop 16))))))))

def eWf (p s t m A K2 : ℕ) (i l : ℕ) : ℤ :=
  if i = m then 0 else HvdH.expNeg p K2 (A * (HvdH.eKz s t l ((i : ℤ) - m)).toNat)

def eWCost (p Bq : ℕ) : ℕ := 400 * 10 ^ 7 * (Bq + 2) ^ 2 + expNegCost p Bq Bq

/-! ## The full `E` table -/

/-- The `E` weight table builder: pushes `ints (eTab p s t α).flatten` on `BW`
(needs `TRI = 2m+1`, `TS = s` and the scalar parameters). -/
def eTabBP : RS := tabP NLogN.Rsp.BW TCI TCK TRI TS eWP

end NLogN.Rsp

end

section
-- Source module: NLogN.HAx6
/-!
# One resampling axis, part 6: the tables of the `B̃` axis

Each `Rsp` table builder is relocated by its own permutation so that it reads its scalar
parameters from the `B̃` layout (see `HAx5`) and pushes its table on the stack the relocated
`bapP` reads.  Three derived scalars are pushed on the temporaries `10` (`e`), `11` (`s²`),
`12` (`2α² - 2`) for the duration of the table building.
-/

set_option linter.unusedSimpArgs false

namespace NLogN.AX

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

abbrev xE : Fin HS := 10
abbrev xK : Fin HS := 11
abbrev xG : Fin HS := 12

def πE : Equiv.Perm (Fin HS) :=
  swapsPerm [(6, 37), (23, 42), (22, 11), (48, 49), (51, 10), (50, 34), (49, 44), (48, 52)]
def πED : Equiv.Perm (Fin HS) := swapsPerm [(8, 38), (49, 44)]
def πC : Equiv.Perm (Fin HS) := swapsPerm [(49, 44)]
def πST : Equiv.Perm (Fin HS) := swapsPerm [(50, 34), (49, 44)]
def πD : Equiv.Perm (Fin HS) :=
  swapsPerm [(0, 33), (25, 12), (22, 11), (48, 49), (50, 34), (49, 44), (48, 52)]
def πK : Equiv.Perm (Fin HS) := swapsPerm [(3, 36), (29, 47), (50, 34)]

def eTabR : RS := RS.ren πE eTabBP
def eDelR : RS := RS.ren πED eDelP
def cTabR : RS := RS.ren πC cTabP
def cDelR : RS := RS.ren πST cDelP
def dTabR' : RS := RS.ren πD dTabBP
def ptKR : RS := RS.ren πK ptKeysP
def psKR : RS := RS.ren πST psKeysP

end NLogN.AX

end

section
-- Source module: NLogN.HAx7
/-!
# One resampling axis, part 7: derived scalars and table removal for the `B̃` axis
-/

set_option linter.unusedSimpArgs false

namespace NLogN.AX

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

/-- `e = ((2e+1) - 1) / 2`. -/
def eE : Ex HS := .quo (.add (.reg bNR) (.lit (-1))) (.lit 2)
/-- `s²`. -/
def s2E : Ex HS := .mul (.reg bS) (.reg bS)
/-- `2α² - 2`. -/
def gE : Ex HS := .add (.add (.reg bI) (.reg bI)) (.lit (-2))

/-- Push the derived scalars `e`, `s²`, `2α²-2` on `xE`, `xK`, `xG`. -/
def tmpP : RS := .seq (.push xE eE) (.seq (.push xK s2E) (.push xG gE))
/-- Pop them again. -/
def tmpQ : RS := .seq (.pop xG) (.seq (.pop xK) (.pop xE))

/-- The state after `tmpP`. -/
def tmpSt (e s α : ℕ) (st : IState HS) : IState HS :=
  ((st.upd xE (.int (e : ℤ) :: st.G xE)).upd xK (.int ((s ^ 2 : ℕ) : ℤ) :: st.G xK)).upd xG
    (.int ((2 * α ^ 2 - 2 : ℕ) : ℤ) :: st.G xG)

/-! ## Removing a table -/

/-! ## Arithmetic facts for the `B̃` tables -/

/-- Size conditions of the `B̃` table builders. -/
structure TabHypB (p s t α Bq : ℕ) : Prop where
  h24 : (α ^ 2 + s + t + 2 * HvdH.eWidth p α + 2) ^ 24 < 2 ^ Bq
  hEq : (4 * p + 35) * 2 ^ (2 * p) < 2 ^ Bq

end NLogN.AX

end

section
-- Source module: NLogN.HAx8
/-!
# One resampling axis, part 8: the `B̃` axis program `bapAxP`

`bapAxP` on `gD = gints L ++ R` (with `L` of dims `t :: rest`, `m = rest.prod`) computes
`gints (tensorGo [(s, bapL p s t α)] (t :: rest) L) ++ R`: it pushes the derived scalars
`e, s², 2α²-2`, builds the seven tables of `bapP` (layout: `HAx5`), pops the derived scalars,
runs `GL.axisP` with the slice program `bapSl`, and pops the seven tables again.  Every stack
except `gD` is restored.
-/

set_option linter.unusedSimpArgs false

namespace NLogN.AX

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

/-- The axis core (= `GL.axisP (bapAx …)`, written out). -/
def bapCoreP : RM :=
  .seq (RM.seqL [.rs (.push 1 (.reg bT)), .rs (.push 2 (.reg bM)), .rs (.push 3 (.reg bKB)),
      .rs (transP 2), .rs (.pop 1), .rs (.pop 2), .rs (.pop 3)])
    (.seq (.rep gCC (.reg bM) (RM.seqL [.rs (mvList gD gZT gCV (cnt2 bT)),
        .rs (mvList gZT bI gCV (cnt2 bT)), bapSlP, .rs (mvList bO gZB gCV (cnt2 bS))]))
      (.rs (mvList gZB gD gCV (cnt2 bTot))))

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

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

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

end NLogN.AX

end

section
-- Source module: NLogN.HFwW1
/-!
# Wrapping the Fermat convolution `fwP` as an abstract `LProg2` (part 1)

The frame (stacks outside `GL.myW`, values on top, arbitrary contents below):

* `KB = 32`: `2^P`, `KN = 33`: `2^fn`, `KR = 34`: `r`, `RW = 52`: `2^(2rP)`, `FR = 53`: `2^fn + 1`,
* `wLO = 35`: `2^logOuter`, `wKB = 36`: `kb`, `wKM = 37`: `K - 1`, `wSH = 38`: `2^(fn - logOuter)`,
* `wDS = 39`: `ints ((descs q.outerAxes 1).flatMap enc8) ++ _`.

Inputs `wI1 = GU = 45` (`gints U`), `wI2 = 40` (`gints V`), output `wO = 41`.

`fwWrap`: push copies of the registers `fwP` reads inside `myW`, move the descriptors from `wDS` to
`AXS` and `V` from `wI2` to `VI` (order preserving), run `fwP`, move the result from `OUT` to `wO`
and the descriptors back (order preserving), pop the copies.
-/

set_option linter.unusedSimpArgs false
set_option linter.unusedVariables false

namespace NLogN.FWW

open IntMul IntMul.HvdH NLogN NLogN.FW

abbrev wLO : Fin HS := 35
abbrev wKB : Fin HS := 36
abbrev wKM : Fin HS := 37
abbrev wSH : Fin HS := 38
abbrev wDS : Fin HS := 39
abbrev wI1 : Fin HS := 45
abbrev wI2 : Fin HS := 40
abbrev wO : Fin HS := 41

/-- The frame the wrapper needs (only stacks outside `myW ++ [45, 40, 41]`). -/
structure FwCtx (q : Ref.Params) (kb : ℕ) (s : IState HS) : Prop where
  vKB : ∃ L, s.G NLogN.KB = .int (2 ^ q.P) :: L
  vKN : ∃ L, s.G NLogN.KN = .int (2 ^ q.fn) :: L
  vKR : ∃ L, s.G KR = .int (q.r : ℤ) :: L
  vRW : ∃ L, s.G RW = .int (2 ^ (2 * q.r * q.P)) :: L
  vFR : ∃ L, s.G FR = .int (2 ^ q.fn + 1) :: L
  vLO : ∃ L, s.G wLO = .int ((2 ^ q.logOuter : ℕ) : ℤ) :: L
  vKBR : ∃ L, s.G wKB = .int (kb : ℤ) :: L
  vKM : ∃ L, s.G wKM = .int ((q.K - 1 : ℕ) : ℤ) :: L
  vSH : ∃ L, s.G wSH = .int (2 ^ (q.fn - q.logOuter)) :: L
  vDS : ∃ R, s.G wDS = ints ((descs q.outerAxes 1).flatMap enc8) ++ R

/-- The precondition on the two inputs. -/
def FwPre (q : Ref.Params) (U V : List (ℤ × ℤ)) : Prop :=
  U.length = 2 ^ q.logOuter * q.r ∧ V.length = 2 ^ q.logOuter * q.r ∧
    (∀ z ∈ U, |z.1| < 2 ^ q.P ∧ |z.2| < 2 ^ q.P) ∧ (∀ z ∈ V, |z.1| < 2 ^ q.P ∧ |z.2| < 2 ^ q.P)

/-! ## Count expressions -/

/-- `2 · 2^logOuter · r`: the number of integers of a Gaussian list of length `T`. -/
def cT : Ex HS := .shl (.add (.reg wLO) (.reg wLO)) (.reg KR)

/-- `8 (K - 1)`: the number of descriptor integers. -/
def cD : Ex HS := .shl (.reg wKM) (.lit 8)

/-! ## Order-preserving transfers as `Spec`s -/

/-! ## Facts about `fermatW` -/

/-! ## The program -/

/-- Copies of the frame values into the registers `fwP` reads. -/
def pushL : List RM :=
  [.rs (.push CR (.reg NLogN.KB)), .rs (.push CM (.reg KR)), .rs (.push NLogN.FW.XN (.reg NLogN.KN)),
    .rs (.push XF (.reg FR)), .rs (.push NN (.reg wLO)), .rs (.push KBR (.reg wKB)),
    .rs (.push NAX (.reg wKM)), .rs (.push SH (.reg wSH))]

/-- Remove the copies. -/
def popL : List RM :=
  [.rs (.pop CR), .rs (.pop CM), .rs (.pop NLogN.FW.XN), .rs (.pop XF), .rs (.pop NN), .rs (.pop KBR),
    .rs (.pop NAX), .rs (.pop SH)]

/-- The core: descriptors to `AXS`, `V` to `VI`, `fwP`, result to `wO`, descriptors back. -/
def coreL : List RM :=
  [GL.xferP wDS AXS cD, GL.xferP wI2 VI cT, fwP, GL.xferP OUT wO cT, GL.xferP AXS wDS cD]

def fwWrap : RM := .seq (RM.seqL pushL) (.seq (RM.seqL coreL) (RM.seqL popL))

/-- The cost of `fwWrap`. -/
def fwWCost (q : Ref.Params) (kb Φm : ℕ) : ℕ :=
  fwCost q kb Φm + 10 ^ 6 * (2 ^ q.logOuter * q.r + 1) * (fwW2 q kb + 2) +
    10 ^ 6 * (q.K + 1) * (fwW q kb + 2)

/-! ## The main specification -/

/-! ## `run`, `ok`, `cost` -/

section corollaries

variable (M : MulSpec) {Φm : ℕ} (q : Ref.Params) (kb : ℕ) (hMB : MulBnd M (q.fn + 1) Φm)
    (hP : 1 ≤ q.P) (hK : q.K ≤ q.fn) (hlo : q.logOuter ≤ q.fn) (hkb : q.logOuter ≤ kb)
    (s : IState HS) (hC : FwCtx q kb s) (U V : List (ℤ × ℤ)) (R1 R2 : List Item)
    (hPre : FwPre q U V) (hU : s.G wI1 = gints U ++ R1) (hV : s.G wI2 = gints V ++ R2)

include hMB hP hK hlo hkb hC hPre hU hV

end corollaries

/-! ## The abstract program and its `FwOK` conditions -/

/-! ## A simplified cost bound -/

end NLogN.FWW

end

section
-- Source module: NLogN.HEx1
/-!
# The full `fwdL` / `exactConvL` programs, part 1: generic tools (namespace `NLogN.EX`)

* `tensorGo_snoc`: `tensorGo` over `As ++ [a]` is one more single-axis step after `As`
  (`tgD As dims` = the axis list after `As`).
* `pushes l` / `pops rs`: sequences of register pushes and pops; `pushV s vs` is the state with
  the values `vs` pushed; `PushOK` is the per-push side condition (each expression evaluated in
  the state after the previous pushes).
* `invP X C Sr Tr`: pushes `((s : ZMod t)⁻¹).val` on `X` by summing `x · [s x ≡ 1 (mod t)]`
  over `x < t` (a counted loop on `C`).
-/

set_option linter.unusedSimpArgs false
set_option linter.unusedVariables false

namespace NLogN.EX

open IntMul NLogN IntMul.HvdH IntMul.HvdH.Ref

/-! ## `tensorGo` one axis at a time -/

/-- The axis list after applying the maps `As` (as in `tensorGo`). -/
def tgD {α : Type*} : List (ℕ × (List α → List α)) → List ℕ → List ℕ
  | [], d => d
  | (n', _) :: As, _ :: rest => tgD As (rest ++ [n'])
  | _ :: _, [] => []

/-! ## Sequences of pushes and pops -/

def pushes : List (Fin HS × Ex HS) → RS
  | [] => .skip
  | (r, e) :: l => .seq (.push r e) (pushes l)

/-- Pops in reverse order: `pops [r₁, …, rₙ]` pops `rₙ` first. -/
def pops : List (Fin HS) → RS
  | [] => .skip
  | r :: l => .seq (pops l) (.pop r)

/-- The state with the values `vs` pushed (in order). -/
def pushV (s : IState HS) : List (Fin HS × ℤ) → IState HS
  | [] => s
  | (r, v) :: l => pushV (s.upd r (.int v :: s.G r)) l

/-- Side conditions of `pushes l` producing the values `vs`. -/
def PushOK (B : ℕ) : IState HS → List (Fin HS × Ex HS) → List (Fin HS × ℤ) → Prop
  | _, [], [] => True
  | s, (r, e) :: l, (r', v) :: vs => r' = r ∧ r.val < 56 ∧ e.Clean ∧ e.ok s ∧ e.Bnd s B ∧
      e.eval s = v ∧ PushOK B (s.upd r (.int v :: s.G r)) l vs
  | _, _, _ => False

/-- Cost weight of a push list. -/
def pwt (l : List (Fin HS × Ex HS)) : ℕ := (l.map fun x => x.2.size + 1).sum

/-! ## The inverse modulo `t` by a counted sum -/

/-- `s x mod t = 1` test (negative iff true), `x = ctr - 1`. -/
def invX (C : Fin HS) : Ex HS := .add (.reg C) (.lit (-1))

def invY (C Sr Tr : Fin HS) : Ex HS := .nmod (.mul (.reg Sr) (invX C)) (.reg Tr)

def invC (C Sr Tr : Fin HS) : Ex HS :=
  .add (.mul (.add (invY C Sr Tr) (.lit (-1))) (.add (invY C Sr Tr) (.lit (-1)))) (.lit (-1))

def invB (X C Sr Tr : Fin HS) : RS :=
  .asg X (.add (.reg X) (.sel (invC C Sr Tr) (invX C) (.lit 0)))

/-- `X : R ↦ ((s : ZMod t)⁻¹).val :: R` (`s` on `Sr`, `t` on `Tr`). -/
def invP (X C Sr Tr : Fin HS) : RS := .seq (.push X (.lit 0)) (.rep C (.reg Tr) (invB X C Sr Tr))

/-- The summand. -/
def invF (s t x : ℕ) : ℕ := if s * x % t = 1 then x else 0

/-- The partial sums (candidates `t-1, t-2, …`). -/
def invAcc (s t k : ℕ) : ℕ := ∑ j ∈ Finset.range k, invF s t (t - 1 - j)

section inv

variable (X C Sr Tr : Fin HS)

/-- The states of the inverse loop. -/
def invSt (w : IState HS) (s t k : ℕ) : IState HS := w.upd X (.int (invAcc s t k) :: w.G X)

end inv

end NLogN.EX

end

section
-- Source module: NLogN.HEx2
/-!
# The full `fwdL` / `exactConvL` programs, part 2: stacks, parameters, stage lists

Stacks used by the axis loops (all inside `GL.myW`, away from `gD`, `gZA`, the `GL` constants
`20 … 28`, `gSA`, `gV`; every program below restores them):

| stack | content                                  |
|-------|------------------------------------------|
| 1  `xS` | `ints q.s ++ _` (walked by the loops)  |
| 2  `xU` | used part of `q.s` (reversed)          |
| 3  `xM` | the current array size (loop only)     |
| 5  `gp` | `p`        | 6 `ga2` | `α²`          |
| 7  `gw` | `w = sWidth p α` | 8 `gw1` | `2w+1`  |
| 9  `ge` | `ε = eWidth p α` | 10 `ge1` | `2ε+1` |
| 11 `gbe` | `betaT p α` | 12 `g2p2` | `2^(2p)` |
| 13 `g2p` | `2^p`  | 14 `gkb` | `kb` (`T ≤ 2^kb`) |
| 15 `gK` | `K` | 16 `gKd` | `K - dh`           |
| 17 `grh` | `2^(e-1)` | 18 `gr` | `2^e`       |
| 19 `cMv` | counter of the restoring move      |
| 55 `cL` | the axis loop counter (no `< 55` stack is free of the axis programs' writes) |
-/

set_option linter.unusedSimpArgs false
set_option linter.unusedVariables false

namespace NLogN.EX

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

abbrev xS : Fin HS := 1
abbrev xU : Fin HS := 2
abbrev xM : Fin HS := 3
abbrev gp : Fin HS := 5
abbrev ga2 : Fin HS := 6
abbrev gw : Fin HS := 7
abbrev gw1 : Fin HS := 8
abbrev ge : Fin HS := 9
abbrev ge1 : Fin HS := 10
abbrev gbe : Fin HS := 11
abbrev g2p2 : Fin HS := 12
abbrev g2p : Fin HS := 13
abbrev gkb : Fin HS := 14
abbrev gK : Fin HS := 15
abbrev gKd : Fin HS := 16
abbrev grh : Fin HS := 17
abbrev gr : Fin HS := 18
abbrev cMv : Fin HS := 19
abbrev cL : Fin HS := 55

/-- The scalar parameters the axis loops read (values on top, anything below). -/
structure Glob (q : Params) (kb : ℕ) (s : IState HS) : Prop where
  p : ∃ L, s.G gp = .int (q.p : ℤ) :: L
  a2 : ∃ L, s.G ga2 = .int ((q.α ^ 2 : ℕ) : ℤ) :: L
  w : ∃ L, s.G gw = .int (sWidth q.p q.α : ℤ) :: L
  w1 : ∃ L, s.G gw1 = .int ((2 * sWidth q.p q.α + 1 : ℕ) : ℤ) :: L
  e : ∃ L, s.G ge = .int (eWidth q.p q.α : ℤ) :: L
  e1 : ∃ L, s.G ge1 = .int ((2 * eWidth q.p q.α + 1 : ℕ) : ℤ) :: L
  be : ∃ L, s.G gbe = .int (betaT q.p q.α) :: L
  p2 : ∃ L, s.G g2p2 = .int (2 ^ (2 * q.p)) :: L
  p1 : ∃ L, s.G g2p = .int (2 ^ q.p) :: L
  kb : ∃ L, s.G gkb = .int (kb : ℤ) :: L
  K : ∃ L, s.G gK = .int (q.K : ℤ) :: L
  Kd : ∃ L, s.G gKd = .int ((q.K - q.dh : ℕ) : ℤ) :: L
  rh : ∃ L, s.G grh = .int ((2 ^ (q.e - 1) : ℕ) : ℤ) :: L
  r : ∃ L, s.G gr = .int ((2 ^ q.e : ℕ) : ℤ) :: L

/-! ## Per-axis quantities -/

/-- Source length of axis `k`. -/
def sK (q : Params) (k : ℕ) : ℕ := q.s.getD k 1

/-- Target length of axis `k`. -/
def tK (q : Params) (k : ℕ) : ℕ := 2 ^ q.te k

/-- `T = ∏ tᵢ`. -/
def TT (q : Params) : ℕ := q.tdims.prod

/-- Axis list before the `S̃` step `k`. -/
def dS (q : Params) (k : ℕ) : List ℕ := q.s.drop k ++ q.tdims.take k

/-- Number of slices of the `S̃` step `k`. -/
def mS (q : Params) (k : ℕ) : ℕ := (q.s.drop (k + 1) ++ q.tdims.take k).prod

/-- Array size before the `S̃` step `k`. -/
def NS (q : Params) (k : ℕ) : ℕ := (dS q k).prod

/-- The array after `k` `S̃` steps. -/
def XS (q : Params) (U : List (ℤ × ℤ)) (k : ℕ) : List (ℤ × ℤ) :=
  tensorGo ((fwdAxS q).take k) q.sdims U

/-- Axis list before the `B̃` step `k`. -/
def dB (q : Params) (k : ℕ) : List ℕ := q.tdims.drop k ++ q.s.take k

/-- Number of slices of the `B̃` step `k`. -/
def mB (q : Params) (k : ℕ) : ℕ := (q.tdims.drop (k + 1) ++ q.s.take k).prod

/-- Array size before the `B̃` step `k`. -/
def NB (q : Params) (k : ℕ) : ℕ := (NLogN.EX.dB q k).prod

/-- The array after `k` `B̃` steps (from the `FtL` output `Y`). -/
def XB (q : Params) (Y : List (ℤ × ℤ)) (k : ℕ) : List (ℤ × ℤ) :=
  tensorGo ((fwdAxB q).take k) q.tdims Y

end NLogN.EX

end

section
-- Source module: NLogN.HEx3
/-!
# The full `fwdL` / `exactConvL` programs, part 3: the scalar expressions of the axis loops

In iteration `k` of an axis loop the counter `cL` holds `K - k`, `xS` holds `sₖ` on top and
`xM` the current array size `N`.  The pushed scalars are computed by small expressions over these
and the globals (`Glob`): `tₖ = (K - dh) - (K - k) < 0 ? 2^(e-1) : 2^e`, `sOff`, `N / sₖ`, …
-/

set_option linter.unusedSimpArgs false
set_option linter.unusedVariables false

namespace NLogN.EX

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

/-- `tₖ` from the counter. -/
def tE : Ex HS := .sel (.add (.reg gKd) (.neg (.reg cL))) (.reg grh) (.reg gr)

/-- `sOff s w = (s - w mod s) mod s`. -/
def sOffE (sr wr : Fin HS) : Ex HS :=
  .nmod (.add (.reg sr) (.neg (.nmod (.reg wr) (.reg sr)))) (.reg sr)

/-- `N / s`. -/
def mE (Nr sr : Fin HS) : Ex HS := .ndiv (.reg Nr) (.reg sr)

/-- The bit-size budget of all scalars. -/
def Gx (q : Params) (kb : ℕ) : ℕ :=
  NLogN.EX.TT q + 2 ^ q.e + q.p + q.α ^ 2 + sWidth q.p q.α + eWidth q.p q.α + q.K + kb + 2

def Vx (q : Params) (kb : ℕ) : ℕ := 2 ^ (2 * q.p) + 4 * Gx q kb ^ 3

/-! ## Evaluation bundles -/

/-- `e` evaluates to `v` in `w`, with its side conditions and bit bound `B`. -/
def EA (w : IState HS) (B : ℕ) (e : Ex HS) (v : ℤ) : Prop :=
  e.eval w = v ∧ e.ok w ∧ e.Bnd w B ∧ e.Clean

/-! ## `tE` -/

/-! ## `sOffE` and `mE` -/

end NLogN.EX

namespace NLogN.EX

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

end NLogN.EX

end

section
-- Source module: NLogN.HEx4
/-!
# The full `fwdL` / `exactConvL` programs, part 4: the body of the `S̃` axis loop

`sBody` (iteration `k`): push the registers of `sapAxP` (layout `HAx2`) computed from `sₖ`
(top of `xS`), the counter, the size `N` (top of `xM`) and the globals; run `sapAxP`; pop the
registers; then `xM := (N / sₖ) tₖ`, move `sₖ` from `xS` to `xU`.
-/

set_option linter.unusedSimpArgs false
set_option linter.unusedVariables false

namespace NLogN.EX

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

/-- The pushed registers of the `S̃` step. -/
def sPush : List (Fin HS × Ex HS) :=
  [(aS, .reg xS), (aT, NLogN.EX.tE), (aKa, sOffE xS gw), (aKN, .add (.reg xS) (.reg gw1)),
    (aR, .reg g2p2), (aNR, .reg gw1), (aNW, .mul NLogN.EX.tE (.reg gw1)), (aM, NLogN.EX.mE xM xS),
    (aKB, .reg gkb), (aK2, .mul (.reg ga2) (.mul NLogN.EX.tE NLogN.EX.tE)), (aP, .reg gp), (aWd, .reg gw),
    (aBe, .reg gbe), (aTot, .mul (NLogN.EX.mE xM xS) NLogN.EX.tE)]

/-- Their values in iteration `k`. -/
def sVals (q : Params) (kb k : ℕ) : List (Fin HS × ℤ) :=
  [(aS, ((sK q k : ℕ) : ℤ)), (aT, ((tK q k : ℕ) : ℤ)),
    (aKa, ((sOff (sK q k) (sWidth q.p q.α) : ℕ) : ℤ)),
    (aKN, ((sK q k + (2 * sWidth q.p q.α + 1) : ℕ) : ℤ)), (aR, 2 ^ (2 * q.p)),
    (aNR, ((2 * sWidth q.p q.α + 1 : ℕ) : ℤ)), (aNW, ((tK q k * (2 * sWidth q.p q.α + 1) : ℕ) : ℤ)),
    (aM, ((mS q k : ℕ) : ℤ)), (aKB, (kb : ℤ)), (aK2, ((q.α ^ 2 * tK q k ^ 2 : ℕ) : ℤ)),
    (aP, (q.p : ℤ)), (aWd, (sWidth q.p q.α : ℤ)), (aBe, betaT q.p q.α),
    (aTot, ((mS q k * tK q k : ℕ) : ℤ))]

/-- Loop bookkeeping: `xM := (N / s) t`, `s` from `xS` to `xU`. -/
def sUpd : RS := .seq (.asg xM (.mul (NLogN.EX.mE xM xS) NLogN.EX.tE)) (.seq (.push xU (.reg xS)) (.pop xS))

def sBody : RM := RM.seqL [.rs (pushes sPush), sapAxP, .rs (pops (sPush.map Prod.fst)), .rs sUpd]

/-- The per-axis hypotheses of the `S̃` step `k`. -/
structure SAx (q : Params) (kb mm Bacc Q Bq : ℕ) (BS : ℕ → ℕ) (k : ℕ) : Prop where
  sap : SapHyp q.p (sK q k) (tK q k) q.α (BS k) (BS (k + 1)) mm Bacc Q
  tab : TabHyp q.p (sK q k) (tK q k) q.α Bq
  axq : AxQ (sK q k) (tK q k) (mS q k) kb (BS k) (BS (k + 1)) Q

/-- Overhead of one iteration (besides `sapAxP`). -/
def sOver (B : ℕ) : ℕ := 10 ^ 9 * (B + 2) ^ 2

end NLogN.EX

end

section
-- Source module: NLogN.HEx5
/-!
# The full `fwdL` / `exactConvL` programs, part 5: the `S̃` axis loop

`sLoop = rep cL K sBody`: iteration `k` turns `gints (XS q U k)` into `gints (XS q U (k+1))`,
walking `q.s` from `xS` to `xU` and keeping the current size on `xM`.
-/

set_option linter.unusedSimpArgs false
set_option linter.unusedVariables false

namespace NLogN.EX

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

/-- `q.s` as integers. -/
def sl (q : Params) : List ℤ := q.s.map (fun x : ℕ => (x : ℤ))

/-- The `S̃` loop. -/
def sLoop : RM := .rep cL (.reg gK) NLogN.EX.sBody

/-- State after `k` iterations. -/
def sSt (q : Params) (U : List (ℤ × ℤ)) (w : IState HS) (R RS LM : List Item) (k : ℕ) :
    IState HS :=
  (((w.upd gD (gints (XS q U k) ++ R)).upd xM (.int (NS q k : ℤ) :: LM)).upd xU
    (ints ((NLogN.EX.sl q).take k).reverse ++ w.G xU)).upd xS (ints ((NLogN.EX.sl q).drop k) ++ RS)

/-- Overhead of the loop besides the `sapAxP` calls (uniform bound `CS` on them). -/
def sLoopCost (K CS Bx : ℕ) : ℕ := 2 * 2000 * (Bx + 2) + K * (CS + sOver Bx + 8000 * (Bx + 2) + 30) + 32

end NLogN.EX

end

section
-- Source module: NLogN.HEx6
/-!
# The full `fwdL` / `exactConvL` programs, part 6: the body of the `B̃` axis loop

`bBody` (iteration `k`, `s = sₖ` on top of `xS`, `N = tₖ m` on top of `xM`): push the registers
of `bapAxP` (layout `HAx5`), compute `u = (s⁻¹ mod t).val` onto `bKS` with `invP`, run `bapAxP`,
pop everything, then `xM := (N / t) s` and move `s` from `xS` to `xU`.
-/

set_option linter.unusedSimpArgs false
set_option linter.unusedVariables false

namespace NLogN.EX

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

/-- `(p s + a (t - s) - 1) / (a (t - s))` from expressions for `p, s, t, a`. -/
def jEx (pe se te ae : Ex HS) : Ex HS :=
  .ndiv (.add (.mul pe se) (.add (.mul ae (.add te (.neg se))) (.lit (-1))))
    (.mul ae (.add te (.neg se)))

/-- `N / t` (the rest size of the `B̃` step). -/
def mBE : Ex HS := .ndiv (.reg xM) NLogN.EX.tE

def jE : Ex HS := jEx (.reg gp) (.reg xS) NLogN.EX.tE (.reg ga2)

/-- The pushed registers of the `B̃` step (`bKS` is computed afterwards by `invP`). -/
def bPush : List (Fin HS × Ex HS) :=
  [(bKa, sOffE xS ge), (bKN, .add (.reg xS) (.reg ge1)), (bT, NLogN.EX.tE), (bKB, .reg gkb),
    (bR, .reg g2p), (NLogN.AX.bN, .reg xS), (bJN, jE), (bNR, .reg ge1), (bJL, .add (.reg xS) (.reg xS)),
    (bS, .reg xS), (bNW, .mul (.reg xS) (.reg ge1)), (bI, .reg ga2), (bPR, .reg g2p),
    (bKs, .reg xS), (bO, .reg gp), (bM, mBE), (bTot, .mul mBE (.reg xS))]

/-- Their values in iteration `k`. -/
def bVals (q : Params) (kb k : ℕ) : List (Fin HS × ℤ) :=
  [(bKa, ((sOff (sK q k) (eWidth q.p q.α) : ℕ) : ℤ)),
    (bKN, ((sK q k + (2 * eWidth q.p q.α + 1) : ℕ) : ℤ)), (bT, ((tK q k : ℕ) : ℤ)),
    (bKB, (kb : ℤ)), (bR, 2 ^ q.p), (NLogN.AX.bN, ((sK q k : ℕ) : ℤ)),
    (bJN, (jIter q.p (sK q k) (tK q k) q.α : ℤ)), (bNR, ((2 * eWidth q.p q.α + 1 : ℕ) : ℤ)),
    (bJL, ((2 * sK q k : ℕ) : ℤ)), (bS, ((sK q k : ℕ) : ℤ)),
    (bNW, ((sK q k * (2 * eWidth q.p q.α + 1) : ℕ) : ℤ)), (bI, ((q.α ^ 2 : ℕ) : ℤ)),
    (bPR, 2 ^ q.p), (bKs, ((sK q k : ℕ) : ℤ)), (bO, (q.p : ℤ)), (bM, ((mB q k : ℕ) : ℤ)),
    (bTot, ((mB q k * sK q k : ℕ) : ℤ))]

/-- The inverse register. -/
def uVal (q : Params) (k : ℕ) : ℤ := ((((sK q k : ℕ) : ZMod (tK q k))⁻¹).val : ℕ)

def bUpd : RS := .seq (.asg xM (.mul mBE (.reg xS))) (.seq (.push xU (.reg xS)) (.pop xS))

def bBody : RM :=
  RM.seqL [.rs (pushes bPush), .rs (invP bKS cMv bS bT), bapAxP,
    .rs (pops ((bPush.map Prod.fst) ++ [bKS])), .rs bUpd]

/-- The per-axis hypotheses of the `B̃` step `k`. -/
structure BAx (q : Params) (kb mm Bacc Q Bq : ℕ) (BB : ℕ → ℕ) (k : ℕ) : Prop where
  bap : BapHyp q.p (sK q k) (tK q k) q.α kb (BB k) (BB (k + 1)) mm Bacc Q
  tab : TabHypB q.p (sK q k) (tK q k) q.α Bq
  axq : AxQ (tK q k) (sK q k) (mB q k) kb (BB k) (BB (k + 1)) Q
  lt : sK q k < tK q k
  t2 : 2 ≤ tK q k

/-- Overhead of one `B̃` iteration (besides `bapAxP`), including the inverse loop. -/
def bOver (q : Params) (B : ℕ) : ℕ := (NLogN.EX.TT q + 3) * 10 ^ 10 * (B + 2) ^ 2

end NLogN.EX

namespace NLogN.EX

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

end NLogN.EX

end

section
-- Source module: NLogN.HEx7
/-!
# The full `fwdL` / `exactConvL` programs, part 7: the `B̃` axis loop
-/

set_option linter.unusedSimpArgs false
set_option linter.unusedVariables false

namespace NLogN.EX

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

/-- The `B̃` loop. -/
def bLoop : RM := .rep cL (.reg gK) bBody

/-- State after `k` iterations. -/
def bSt (q : Params) (U : List (ℤ × ℤ)) (w : IState HS) (R RS LM : List Item) (k : ℕ) :
    IState HS :=
  (((w.upd gD (gints (XB q U k) ++ R)).upd xM (.int (NB q k : ℤ) :: LM)).upd xU
    (ints ((NLogN.EX.sl q).take k).reverse ++ w.G xU)).upd xS (ints ((NLogN.EX.sl q).drop k) ++ RS)

/-- Overhead of the loop besides the `bapAxP` calls (uniform bound `CS` on them). -/
def bLoopCost (q : Params) (CS Bx : ℕ) : ℕ :=
  2 * 2000 * (Bx + 2) + q.K * (CS + bOver q Bx + 8000 * (Bx + 2) + 30) + 32

end NLogN.EX

end

section
-- Source module: NLogN.HEx8
/-!
# The full `fwdL` / `exactConvL` programs, part 8: the two axis phases and `fwdLP`

A phase pushes the initial array size on `xM`, runs the axis loop, pops `xM` and moves `q.s`
back from `xU` to `xS`; it changes nothing but `gD`.
`fwdLP fw = S̃ phase; ftlP fw; B̃ phase` computes `Ref.fwdL`.
-/

set_option linter.unusedSimpArgs false
set_option linter.unusedVariables false

namespace NLogN.EX

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

/-- A phase around an axis loop; `cst` holds the initial size. -/
def phaseP (cst : Fin HS) (loop : RM) : RM :=
  RM.seqL [.rs (.push xM (.reg cst)), loop, .rs (.pop xM), .rs (mvList xU xS cMv (.reg gK))]

/-- Overhead of a phase besides its loop. -/
def phOver (K Bx : ℕ) : ℕ := (K + 1) * 20000 * (Bx + 2)

end NLogN.EX

namespace NLogN.EX

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

/-- The `S̃` phase (initial size `S` read from `gPM`). -/
def sPhase : RM := phaseP gPM sLoop

/-- The `B̃` phase (initial size `T` read from `gPN`). -/
def bPhase : RM := phaseP gPN bLoop

end NLogN.EX

end

section
-- Source module: NLogN.HEx9
/-!
# The full `fwdL` / `exactConvL` programs, part 9: `fwdLP`

`fwdLP fw = S̃ phase (loop over the axes, `sapAxP` per axis); ftlP fw; B̃ phase (loop over the
axes, `bapAxP` per axis)`, generic in the `fermatW` program `fw` (any `GL.LProg2` with
`GL.FwOK`, e.g. `FWW.fwL`). Its context `Ctx` (e.g. the `FwCtx` frame on stacks
`32 … 39, 52, 53`) is a precondition that survives: every axis program restores every stack
but `gD`.
-/

set_option linter.unusedSimpArgs false
set_option linter.unusedVariables false

namespace NLogN.EX

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

section
variable {M : MulSpec} {Ctx : IState HS → Prop} {PreW : List (ℤ × ℤ) → List (ℤ × ℤ) → Prop}
  {q : Params} {CW : ℕ}

def fwdLP (fw : LProg2 M Ctx PreW (fermatW q) CW) : RM := RM.seqL [sPhase, ftlP fw, bPhase]

/-- Static facts (parameters, per-axis hypotheses, size budgets) for `fwdLP`. -/
structure FwdCfgX (M : MulSpec) (q : Params) (kb mm Φ Bacc Qa Bq Bx CS CB m Q : ℕ)
    (BS BB : ℕ → ℕ) : Prop where
  hM : MulBnd M mm Φ
  hsl : q.s.length = q.K
  hst : ∀ i < q.K, sK q i ≤ tK q i
  hα : 1 ≤ q.α
  hBx : Vx q kb < 2 ^ Bx
  sax : ∀ k < q.K, SAx q kb mm Bacc Qa Bq BS k
  bax : ∀ k < q.K, BAx q kb mm Bacc Qa Bq BB k
  cS : ∀ k < q.K,
    sapAxCost q.p (sK q k) (tK q k) q.α (mS q k) kb (BS k) mm Φ Bacc Qa Bq ≤ CS
  cB : ∀ k < q.K,
    bapAxCost q.p (sK q k) (tK q k) q.α (mB q k) kb (BB k) mm Φ Bacc Qa Bq ≤ CB
  chl : (chirpL q).length = NLogN.EX.TT q
  chb : GB m (chirpL q)
  Sm : BS q.K ≤ m
  Bm : BB q.K ≤ m
  mQ : 2 * m + 1 ≤ Q
  k : q.p + q.logT + 1 ≤ Q
  NQ : zsz ((2 * NLogN.EX.TT q : ℕ) : ℤ) ≤ Q

/-- What `fwdLP` needs in the machine state (besides the data on `gD`). -/
structure FwdEnvX (q : Params) (kb : ℕ) (Ctx : IState HS → Prop) (s : IState HS) : Prop where
  ctx : Ctx s
  glob : Glob q kb s
  xs : ∃ RS, s.G xS = ints (NLogN.EX.sl q) ++ RS
  ch : ∃ RC, s.G gCH = gints (chirpL q) ++ RC
  pn : ∃ L, s.G gPN = .int (NLogN.EX.TT q : ℤ) :: L
  pm : ∃ L, s.G gPM = .int (q.S : ℤ) :: L
  pt : ∃ L, s.G gPT = .int (2 ^ (q.p + q.logT)) :: L
  pp : ∃ L, s.G gPP = .int (2 ^ q.p) :: L

/-- What `fwdLP` needs about the data `U`. -/
structure FwdInX (q : Params) (PreW : List (ℤ × ℤ) → List (ℤ × ℤ) → Prop) (BS BB : ℕ → ℕ)
    (m Q : ℕ) (U : List (ℤ × ℤ)) : Prop where
  len : U.length = q.S
  b : GB (BS 0) U
  pre : PreW (ftlV q (fwdX q U)) (chirpL q)
  wl : (fermatW q (ftlV q (fwdX q U)) (chirpL q)).length = NLogN.EX.TT q
  wb : GB Q (fermatW q (ftlV q (fwdX q U)) (chirpL q))
  cb : GB m (convL q (ftlV q (fwdX q U)) (chirpL q))
  yb : GB (BB 0) (FtL q (fwdX q U))

/-- The cost bound of `fwdLP`: two axis phases (`K` iterations each, every iteration bounded
by `CS` resp. `CB` plus explicit overhead) and `FtL` (`CW` for `fermatW`, `8 T` products of
size `m`, linear rest). -/
def fwdLCost (q : Params) (CS CB Bx CW Q Φm : ℕ) : ℕ :=
  (sLoopCost q.K CS Bx + phOver q.K Bx) +
    (CW + 8 * NLogN.EX.TT q * Φm + 1000000 * (NLogN.EX.TT q + 1) * (Q + 2)) +
    (bLoopCost q CB Bx + phOver q.K Bx)

end

end NLogN.EX

end

section
-- Source module: NLogN.HEx10
/-!
# The full `fwdL` / `exactConvL` programs, part 10: `exactLP`

`exactLP fw` follows `GL.exactP` (scale `u`, `fwdL`, park, scale `v`, `fwdL`, pointwise
`mulConj`, `fwdL`, round) with `fwdLP fw` in place of `GL.fwdP`.
-/

set_option linter.unusedSimpArgs false
set_option linter.unusedVariables false

namespace NLogN.EX

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

section
variable {M : MulSpec} {Ctx : IState HS → Prop} {PreW : List (ℤ × ℤ) → List (ℤ × ℤ) → Prop}
  {q : Params} {CW : ℕ}

def exactLP (fw : LProg2 M Ctx PreW (fermatW q) CW) : RM :=
  RM.seqL [scaleP gPM, fwdLP fw, xferP gD gSA (cnt2 gPM), xferP gV gD (.reg gPM),
    scaleP gPM, fwdLP fw, xferP gSA gZA (cnt2 gPM), zipP sgMulConj gPM,
    fwdLP fw, roundP gPM]

end

end NLogN.EX

end

section
-- Source module: NLogN.HSt1
/-!
# The recursive step, part 1: state frames, list loops (namespace `NLogN.ST`)

* `W s L`: the state `s` with the lists `L i` pushed on top of every stack `i`.
* `slLoop`: extracts the source lengths `sᵢ` from the CRT records (consuming them).
* `dropN` wrappers.
-/

set_option linter.unusedSimpArgs false
set_option linter.unusedVariables false

namespace NLogN.ST

open IntMul NLogN

/-! ## Frames -/

/-- `s` with `L i` pushed on top of each stack `i`. -/
def W (s : IState HS) (L : Fin HS → List Item) : IState HS :=
  ⟨fun i => L i ++ s.G i, s.ip, s.out⟩

@[simp] theorem W_G (s : IState HS) (L : Fin HS → List Item) (i : Fin HS) :
    (W s L).G i = L i ++ s.G i := rfl
@[simp] theorem W_ip (s : IState HS) (L : Fin HS → List Item) : (W s L).ip = s.ip := rfl
@[simp] theorem W_out (s : IState HS) (L : Fin HS → List Item) : (W s L).out = s.out := rfl

/-! ## Spec helpers -/

open NLogN.FW (Spec)

/-! ## Extracting the source lengths from the CRT records -/

abbrev sR : Fin HS := 4
abbrev sT : Fin HS := 41
abbrev sC : Fin HS := 43

def slBody : RM :=
  RM.seqL [.rs (.pop sR), .rs (.push NLogN.ST.sT (.reg sR)), .rs (.pop sR), .rs (.pop sR), .rs (.pop sR)]

/-- Consume `K` records `[a, x, b, c]` from stack `4`, collecting the `x`s on `41` (reversed). -/
def slLoop (K : ℕ) : RM := .rep sC (.lit K) slBody

/-- Second entries of 4-records. -/
def snd4 (r : List ℤ) : ℤ := r.getD 1 0

def Rec4 (r : List ℤ) : Prop := ∃ a x b c : ℤ, r = [a, x, b, c]

/-- The family of states of `slLoop`. -/
def slSt (w : IState HS) (rs : List (List ℤ)) (R4 : List Item) (k : ℕ) : IState HS :=
  (w.upd sR (ints (rs.drop k).flatten ++ R4)).upd NLogN.ST.sT (ints ((rs.take k).map snd4).reverse ++ w.G NLogN.ST.sT)

end NLogN.ST

end

section
-- Source module: NLogN.HSt2
/-!
# The recursive step, part 2: the Fermat level descriptors (namespace `NLogN.ST`)

`descLoop K` pushes `ints ((descs q.outerAxes 1).flatMap enc8)` onto stack `39`, reading
`r` (`47`), `r/2` (`48`), `K - dh` (`49`), `logOuter` (`51`), `fn` (`54`); scratch `41`, `19`,
counter `52`.
-/

set_option linter.unusedSimpArgs false
set_option linter.unusedVariables false

namespace NLogN.ST

open IntMul NLogN IntMul.HvdH IntMul.HvdH.Ref
open NLogN.FW (Spec Desc descs enc8)

/-! ## The descriptors as an explicit list -/

/-- `Σ_{a ≤ i < a + n} te i`. -/
def jsum (q : Params) (a n : ℕ) : ℕ := ((List.range' a n).map q.te).sum

/-- The descriptor of level `i` (outer axes `0 … n-1`, starting at `a`). -/
def dI (q : Params) (a n i : ℕ) : Desc :=
  (2 ^ jsum q a (i - a), q.te i, 2 * q.fn / 2 ^ q.te i, 2 ^ jsum q (i + 1) (a + n - (i + 1)))

/-! ## The loop -/

abbrev tR : Fin HS := 47
abbrev tRh : Fin HS := 48
abbrev tKd : Fin HS := 49
abbrev tLo : Fin HS := 51
abbrev tFn : Fin HS := 54
abbrev dX : Fin HS := 41
abbrev dB : Fin HS := 19
abbrev dC : Fin HS := 52
abbrev dO : Fin HS := 39

/-- `i < dh` test for `i = ctr - 1`: `ctr - 1 - K + (K - dh) < 0`. -/
def tcE (K : ℕ) : Ex HS := .add (.add (.reg dC) (.reg tKd)) (.lit (-1 - (K : ℤ)))

/-- `2^{te i}`. -/
def tE (K : ℕ) : Ex HS := .sel (tcE K) (.reg tRh) (.reg tR)

def preE : Ex HS := .ndiv (.pow2 (.reg tLo)) (.mul (.reg dX) (.reg NLogN.ST.dB))
def expE : Ex HS := .pow2 (.ndiv (.add (.reg tFn) (.reg tFn)) (.reg dX))

def dPush : List (Fin HS × Ex HS) :=
  [(dO, preE), (dO, .reg dX), (dO, expE), (dO, .reg NLogN.ST.dB), (dO, .reg NLogN.ST.dB), (dO, expE), (dO, .reg dX),
    (dO, preE)]

def dBody (K : ℕ) : RM :=
  RM.seqL [.rs (.push dX (tE K)), .rs (EX.pushes dPush), .rs (.asg NLogN.ST.dB (.mul (.reg NLogN.ST.dB) (.reg dX))),
    .rs (.pop dX)]

def descLoop (K : ℕ) : RM :=
  RM.seqL [.rs (.push NLogN.ST.dB (.lit 1)), .rep dC (.lit ((K : ℤ) - 1)) (dBody K), .rs (.pop NLogN.ST.dB)]

end NLogN.ST

end