Harvey–van der Hoeven machine development, part 4: Fermat convolution program, bounds, parameter program

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

import Definitions.Def_IntMul_HvdH_M_Defs3

set_option autoImplicit false
set_option relaxedAutoImplicit false
set_option maxHeartbeats 1000000

section
-- Source module: NLogN.HFW8
/-!
# Fermat-ring convolution machine, part 8: scaling and decoding

`dcL`: for each residue `z` on `XD`: push `(z · 2^{2n-l}) mod F` on `KZ` (linear operations,
ruler `2^{n-l}` on `SH = 28`), drop `z`, run `decP` (pushes the reversed block on `KO`).
`mvPairs`: move the `N·r` Gaussian pairs from `KO` to `OUT = 0` keeping each pair's order.
-/

set_option linter.unusedSimpArgs false
set_option linter.unusedVariables false

namespace NLogN.FW

open IntMul IntMul.HvdH NLogN

abbrev SH : Fin HS := 28
abbrev OUT : Fin HS := 0

/-! ## Bounds on decoded blocks -/

/-! ## The scaling expression -/

def wE : Ex HS := .shl (.reg XD) (.reg SH)
def dsE : Ex HS := .add (.quo wE (.reg NLogN.FW.XN)) (.neg (.rem wE (.reg NLogN.FW.XN)))
/-- `(z · 2^{2n-l}) mod (2^n + 1)` for `z` on `XD`, `2^{n-l}` on `SH`. -/
def scE : Ex HS := .sel dsE (.add dsE (.reg XF)) dsE

/-- The scaled residue. -/
def scl (n l : ℕ) (z : ℤ) : ℤ := (z * 2 ^ (2 * n - l)) % (2 ^ n + 1)

/-! ## The decoding loop -/

def dcBody : RM := RM.seqL [.rs (.push KZ scE), .rs (.pop XD), decP]

def dcL : RM := .rep C0 (.reg NN) dcBody

def dSt (P r l : ℕ) (C : List ℤ) (s : IState HS) (RD : List Item) (k : ℕ) : IState HS :=
  (s.upd XD (ints (C.drop k) ++ RD)).upd KO
    (gints (((C.take k).map (scl (4 * r * P) l)).flatMap
      (Ref.decBlk (2 ^ (4 * r * P) + 1) P r)).reverse ++ s.G KO)

/-- Cost of decoding `N` residues. -/
def dcCost (N r P Q W : ℕ) : ℕ :=
  N * (1000000 * (r + 1) * (Q + 5) + 200000 * (r * P + 1) + 100000 * (W + 2)) + 4000 * (W + 2) + 32

/-! ## Moving the Gaussian pairs -/

/-- `N · r` -/
def cntE : Ex HS := .shl (.reg NN) (.reg KR)

def mpBody : RM :=
  RM.seqL [.rs (.push TMP (.reg KO)), .rs (.pop KO), .rs (.push OUT (.reg KO)), .rs (.pop KO),
    .rs (.push OUT (.reg TMP)), .rs (.pop TMP)]

def mvPairs : RM := .rep C0 cntE mpBody

def mpSt (Z : List (ℤ × ℤ)) (s : IState HS) (RK : List Item) (k : ℕ) : IState HS :=
  (s.upd KO (gints (Z.drop k) ++ RK)).upd OUT (gints (Z.take k).reverse ++ s.G OUT)

end NLogN.FW

end

section
-- Source module: NLogN.HFW9
/-!
# Fermat-ring convolution machine, part 9: facts about the parameters and the levels
-/

set_option linter.unusedSimpArgs false
set_option linter.unusedVariables false

namespace NLogN.FW

open IntMul IntMul.HvdH NLogN

namespace Prm

variable (q : Ref.Params)

end Prm

/-! ## The two kinds of passes -/

def Gd (n : ℕ) : ℕ → ℕ → List ℤ → List ℤ := fun E j => difZ (2 ^ n + 1) E j
def Gt (n : ℕ) : ℕ → ℕ → List ℤ → List ℤ := fun E j => ditZ (2 ^ n + 1) n E j

/-- Conditions on the level list, independent of the transform. -/
def LvlOK (n kb W N J : ℕ) (D : List Desc) : Prop :=
  ∀ d ∈ D, d.1 * (2 ^ d.2.1 * d.2.2.2) = N ∧ 1 ≤ d.1 ∧ 1 ≤ d.2.2.2 ∧
    2 ^ d.2.1 * d.2.2.1 = 2 * n ∧ d.2.1 ≤ J

end NLogN.FW

end

section
-- Source module: NLogN.HFW10
/-!
# Fermat-ring convolution machine, part 10: the program `fwP` computing `Ref.fermatW`

Stack layout of `fwP` (all stacks hold the value on top, below it arbitrary contents):

* inputs: `GU = 45`: `gints U`, `VI = 12`: `gints V` (both consumed);
* encoder registers: `CR = 23`: `2^P`, `CM = 26`: `r`, `RW = 52`: `2^(2rP)`, `FR = 53`: `2^fn + 1`;
* decoder registers: `KB = 32`: `2^P`, `KN = 33`: `2^fn`, `KR = 34`: `r`;
* transform registers: `XN = 2`: `2^fn`, `XF = 3`: `2^fn + 1`;
* sizes: `NN = 19`: `N = 2^logOuter`, `KBR = 27`: `kb` (with `N ≤ 2^kb`), `NAX = 18`: `K - 1`,
  `SH = 28`: `2^(fn - logOuter)`;
* level descriptors: `AXS = 15`: `ints ((descs q.outerAxes 1).flatMap enc8)` (restored at the end);
* output: `OUT = 0`: `gints (fermatW q U V)` pushed on top of the old contents.

Everything else (`XD = 8`, `YC = 4`, `SA = 20`, `AXR = 16`, `KO = 31`, counters, scratch) is restored.
-/

set_option linter.unusedSimpArgs false
set_option linter.unusedVariables false

namespace NLogN.FW

open IntMul IntMul.HvdH NLogN

abbrev VI : Fin HS := 12
abbrev SA : Fin HS := 20

def σV : Equiv.Perm (Fin HS) := swapsPerm [(GU, VI)]

/-- The encoder reading from `VI` instead of `GU`. -/
def encV : RM := encL.ren σV

/-! ## The program -/

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

def fwP : RM :=
  RM.seqL [encL, passP AXS AXR difPrc, .rs (mvList XD NLogN.FW.SA C2 (.reg NN)), encV,
    .rs (mvList AXR AXS C2 c8E), passP AXS AXR difPrc, .rs (mvList NLogN.FW.SA YC C2 (.reg NN)), pmL,
    passP AXR AXS ditPrc, dcL, mvPairs]

/-- The bit bound used throughout. -/
def fwW (q : Ref.Params) (kb : ℕ) : ℕ := 2 * q.fn + kb + 3

/-- The bit bound of the decoded pairs and their count. -/
def fwW2 (q : Ref.Params) (kb : ℕ) : ℕ := q.P + kb + q.e + 3

/-- The bit bound of the encoder/decoder scalars. -/
def fwQ (q : Ref.Params) : ℕ := q.P + q.e + 1

/-- The explicit cost of `fwP`. -/
def fwCost (q : Ref.Params) (kb Φm : ℕ) : ℕ :=
  2 * encCost (2 ^ q.logOuter) q.r q.P (fwQ q) (fwW q kb) +
    3 * passCost (q.K - 1) q.fn kb (fwW q kb) (2 ^ q.logOuter) q.e +
    2 * (2 ^ q.logOuter * (14000 * (fwW q kb + 2) + 30) + 4000 * (fwW q kb + 2) + 32) +
    (8000 * (fwW q kb + 2) + 8 * (q.K - 1) * (14000 * (fwW q kb + 2) + 30) + 32) +
    pmCost (2 ^ q.logOuter) (fwW q kb) Φm + dcCost (2 ^ q.logOuter) q.r q.P (fwQ q) (fwW q kb) +
    (8000 * (fwW2 q kb + 2) + 2 ^ q.logOuter * 2 ^ q.e * (40000 * (fwW2 q kb + 2)) + 32)

/-! ## The math -/

/-! ## The main theorem -/

/-! ## A simplified cost bound -/

/-! ## `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) (U V : List (ℤ × ℤ))
    (RU RV Rs LCR LCM LRW LFR LKB LKN LKR LXN LXF LNN LKBR LNAX LSH : List Item)
    (hUl : U.length = 2 ^ q.logOuter * q.r) (hVl : V.length = 2 ^ q.logOuter * q.r)
    (hUb : ∀ z ∈ U, |z.1| < 2 ^ q.P ∧ |z.2| < 2 ^ q.P)
    (hVb : ∀ z ∈ V, |z.1| < 2 ^ q.P ∧ |z.2| < 2 ^ q.P)
    (hU : s.G GU = gints U ++ RU) (hV : s.G VI = gints V ++ RV)
    (hCR : s.G CR = .int (2 ^ q.P) :: LCR) (hCM : s.G CM = .int (q.r : ℤ) :: LCM)
    (hRW : s.G RW = .int (2 ^ (2 * q.r * q.P)) :: LRW) (hFR : s.G FR = .int (2 ^ q.fn + 1) :: LFR)
    (hKB : s.G KB = .int (2 ^ q.P) :: LKB) (hKN : s.G KN = .int (2 ^ q.fn) :: LKN)
    (hKR : s.G KR = .int (q.r : ℤ) :: LKR)
    (hXN : s.G XN = .int (2 ^ q.fn) :: LXN) (hXF : s.G XF = .int (2 ^ q.fn + 1) :: LXF)
    (hNN : s.G NN = .int ((2 ^ q.logOuter : ℕ) : ℤ) :: LNN) (hKBR : s.G KBR = .int kb :: LKBR)
    (hNAX : s.G NAX = .int ((q.K - 1 : ℕ) : ℤ) :: LNAX)
    (hAXS : s.G AXS = ints ((descs q.outerAxes 1).flatMap enc8) ++ Rs)
    (hSH : s.G SH = .int (2 ^ (q.fn - q.logOuter)) :: LSH)

include hMB hP hK hlo hkb hUl hVl hUb hVb hU hV hCR hCM hRW hFR hKB hKN hKR hXN hXF hNN hKBR hNAX
  hAXS hSH

end corollaries

/-! ## The procedures called by `fwP` -/

section prc

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

end prc

end NLogN.FW

end

section
-- Source module: NLogN.HNumAtan
/-!
# `atanInv` / `atanhInv` as a register program

Inputs `N` on register `0`, `K` on register `1`; the result is pushed on register `2`.
Scratch registers `40 … 45`.
-/

set_option linter.unusedSimpArgs false

namespace NLogN

open IntMul RS

/-- The partial sums. -/
def atanSum (x N : ℕ) (σ : ℤ) (k : ℕ) : ℤ :=
  ∑ m ∈ Finset.range k, σ ^ m * ((2 ^ N / ((2 * m + 1) * x ^ (2 * m + 1)) : ℕ) : ℤ)

end NLogN

namespace NLogN
open IntMul RS

end NLogN

end

section
-- Source module: NLogN.HBd1
/-!
# Crude magnitude bounds, part 1: scalars and the one-dimensional maps

All bounds are on `Int.natAbs` of each component.  `ZB z B` says both components of the
Gaussian integer `z` are at most `B` in absolute value; `LB L B` says this for every entry of a
list, `FB f B` for every value of a function `ℕ → ℤ × ℤ`.

Main results:
* `expNegFix_natAbs_le`: `|expNegFix n M X| ≤ (2M+15) 2^n` for **every** `X`;
  hence `expNeg_natAbs_le`, `expPos_natAbs_le`, `dTilde_natAbs_le`: `≤ (4p+35) 2^p`.
* `twD_ZB`: `twD p n j` has both components `≤ 2^p` (it is `scaleZ`'d into the disc).
* `Sapprox_FB`, `Eapprox_FB`, `vIter_FB`, `Japprox_FB`, `Dapprox_FB`, `Bapprox_FB`,
  `sapL_LB`, `bapL_LB` (and the lengths `sapL_length`, `bapL_length`).
-/

set_option linter.unusedSimpArgs false

namespace NLogN.BD

open IntMul.HvdH IntMul.HvdH.Ref

/-! ## Bound predicates -/

/-- The sup norm of a Gaussian integer. -/
def gB (z : ℤ × ℤ) : ℕ := max z.1.natAbs z.2.natAbs

/-- Both components of `z` are at most `B` in absolute value. -/
def ZB (z : ℤ × ℤ) (B : ℕ) : Prop := z.1.natAbs ≤ B ∧ z.2.natAbs ≤ B

/-- Every entry of `L` is bounded by `B`. -/
def LB (L : List (ℤ × ℤ)) (B : ℕ) : Prop := ∀ z ∈ L, z.1.natAbs ≤ B ∧ z.2.natAbs ≤ B

/-- Every entry of an integer list is bounded by `B`. -/
def LBz (L : List ℤ) (B : ℕ) : Prop := ∀ a ∈ L, a.natAbs ≤ B

/-- Every value of `f` is bounded by `B`. -/
def FB (f : ℕ → ℤ × ℤ) (B : ℕ) : Prop := ∀ j, (f j).1.natAbs ≤ B ∧ (f j).2.natAbs ≤ B

/-! ## Arithmetic helpers -/

/-! ## `log2Bits M ≤ 2^M` -/

/-! ## `expNegFix` -/

/-! ## Roots of unity -/

/-! ## Products -/

/-! ## The one-dimensional maps -/

/-- The factor of `S̃`: `(2m+1)(4p+35)`, `m = sWidth p α`. -/
def cS1 (p α : ℕ) : ℕ := (2 * sWidth p α + 1) * (4 * p + 35)

/-- The factor of `Ẽ`: `(2m+1)(4p+35)`, `m = eWidth p α`. -/
def cE1 (p α : ℕ) : ℕ := (2 * eWidth p α + 1) * (4 * p + 35)

/-- The factor of `B̃` for a given number of Neumann terms `J`. -/
def cB1 (p α J : ℕ) : ℕ := (4 * p + 35) * (J * cE1 p α ^ J)

/-! ## The list forms -/

end NLogN.BD

end

section
-- Source module: NLogN.HBd2
/-!
# Crude magnitude bounds, part 2: list plumbing, the Fermat convolution and the pipeline

* `chunkL_LB`, `transposeL_LB`, `flatMap_LB`, `tensorGo_LB` (general bound function `g`),
  `tensorGo_LB_mul` (factor `c` per axis).
* `kunpack_natAbs_le`, `decBlk_LB`, `fermatW_LB`: every output entry of `fermatW` has
  components `≤ 2^P`, unconditionally; `convL_LB`, `FtL_LB` (unconditional as well).
* `fwdL_stages`, `fwdL_LB`, `exactConvL` intermediates, `crtIn`/`crtOut`/`chunksL`.
* `refStep_bounds`: all intermediate lists of `refStep q x y` for a general `q` satisfying
  `HJ q` (every `jIter` is at most `Jm q`).
-/

set_option linter.unusedSimpArgs false

namespace NLogN.BD

open IntMul.HvdH IntMul.HvdH.Ref

/-! ## List plumbing -/

/-! ## Integer lists -/

/-! ## The Fermat convolution -/

/-- The bound `2 · 2^P / 2^(p+logT)` of the DFT oracle (unconditional). -/
def BC (q : Params) : ℕ := 2 * (2 ^ q.P / 2 ^ (q.p + q.logT))

/-! ## `fwdL` -/

/-- The `S̃` axes of `fwdL`. -/
def axS (q : Params) : List (ℕ × (List (ℤ × ℤ) → List (ℤ × ℤ))) :=
  (List.range q.K).map fun i => (2 ^ q.te i, sapL q.p (q.s.getD i 1) (2 ^ q.te i) q.α)

/-- The `B̃` axes of `fwdL`. -/
def axB (q : Params) : List (ℕ × (List (ℤ × ℤ) → List (ℤ × ℤ))) :=
  (List.range q.K).map fun i => (q.s.getD i 1, bapL q.p (q.s.getD i 1) (2 ^ q.te i) q.α)

/-- The FtL input inside `fwdL`. -/
def ftIn (q : Params) (U : List (ℤ × ℤ)) : List (ℤ × ℤ) :=
  List.zipWith (fun a u => mulZ q.p (conjZ a) u) (chirpL q) (tensorGo (axS q) q.sdims U)

/-- Uniform bound on the number of Neumann terms: `4Kp/α² + 1`. -/
def Jm (q : Params) : ℕ := 4 * q.K * q.p / q.α ^ 2 + 1

/-- Hypothesis: every axis has `jIter ≤ Jm q`. -/
def HJ (q : Params) : Prop := ∀ i < q.K, jIter q.p (q.s.getD i 1) (2 ^ q.te i) q.α ≤ Jm q

/-- Per-axis factors. -/
def cSq (q : Params) : ℕ := cS1 q.p q.α
def cBq (q : Params) : ℕ := cB1 q.p q.α (Jm q)

/-- `fwdL` output bound (independent of the input). -/
def BF (q : Params) : ℕ := cBq q ^ q.K * NLogN.BD.BC q

/-! ## `exactConvL` -/

/-- The scaled input `U` of `exactConvL`. -/
def ecU (q : Params) (u : List ℤ) : List (ℤ × ℤ) := u.map fun a => (a * 2 ^ (q.p - q.b), (0 : ℤ))

/-- The pointwise product `W` of `exactConvL`. -/
def ecW (q : Params) (u v : List ℤ) : List (ℤ × ℤ) :=
  List.zipWith (fun a b => conjZ (mulZ q.p a b)) (fwdL q (ecU q u)) (fwdL q (ecU q v))

/-- `Y = fwdL W`. -/
def ecY (q : Params) (u v : List ℤ) : List (ℤ × ℤ) := fwdL q (ecW q u v)

/-- The bound of `W`. -/
def BW (q : Params) : ℕ := 2 * BF q * BF q / 2 ^ q.p

/-- The output bound of `exactConvL`. -/
def BX (q : Params) : ℕ := BF q * q.S ^ 2 * 2 ^ (2 * q.b + 6 * q.K * q.α ^ 2) / 2 ^ q.p + 2

/-! ## The whole step -/

/-- The input digits on the source box. -/
def rU (q : Params) (x : ℕ) : List ℤ := crtIn q.s (chunksL q.b q.L x)

end NLogN.BD

end

section
-- Source module: NLogN.HBd3
/-!
# Crude magnitude bounds, part 3: the parameters `mkParams K N`

* `copS_upper`: `4K sᵢ ≤ (4K-1) tᵢ` for the source lengths of `mkParams` (direct from `copS`).
* `jIter_le`: then `jIter ≤ 4Kp/α² + 1 = Jm q`, i.e. `HJ (mkParams K N)`.
* `Bq q = 2 · max(2^p, BW q) · cSq q ^ K`, and for `K ≥ 20`, `N ≥ 2^(2^(8K)+1)`:
  `Bq_size`: `Nat.size (Bq q) ≤ p + p/8` and `Bq_lt`: `Bq q < 2^P`.
* `refStep_LB_mkParams`: every intermediate list of `refStep (mkParams K N) x y` is bounded by
  `Bq`, except the raw `fermatW` outputs which are bounded by `2^P` (structurally).
-/

set_option linter.unusedSimpArgs false

namespace NLogN.BD

open IntMul.HvdH IntMul.HvdH.Ref NLogN.Asym

/-! ## The source lengths -/

/-! ## The global bound -/

/-- The global bound: twice the largest `fwdL` input bound times `cSq^K`. -/
def Bq (q : Params) : ℕ := 2 * max (2 ^ q.p) (BW q) * cSq q ^ q.K

/-- Explicit exponent bound: `Bq ≤ 2^(p + X)`. -/
def Xq (q : Params) (B : ℕ) : ℕ :=
  20 + 2 * q.K * (B + 6) + q.K * (3 * B + 10) + q.K * Jm q * (4 * B + 22)

/-! ## Comparing the individual bounds with `Bq` -/

/-! ## The main statement -/

end NLogN.BD

end

section
-- Source module: NLogN.HCrKey
/-!
# Generating the CRT keys `keyL s` incrementally (namespace `NLogN.CR`)

Records `[k mod sᵢ, sᵢ, wᵢ, sᵢ wᵢ]` (`wᵢ = ∏_{i' > i} s_{i'}`) on stack `kRA`; for each `k < S`
the current key `rmPos s (k mod sᵢ)ᵢ` (kept on `kPS`) is emitted and every record is advanced.
-/

set_option linter.unusedSimpArgs false
set_option linter.unusedVariables false

namespace NLogN.CR

open IntMul IntMul.HvdH IntMul.HvdH.Ref RS

/-! ## Records -/

/-- `wᵢ = ∏_{i' > i} s_{i'}` -/
def wgt (s : List ℕ) (i : ℕ) : ℕ := (s.drop (i + 1)).prod

def recI (s : List ℕ) (k i : ℕ) : List ℤ :=
  [((k % s.getD i 1 : ℕ) : ℤ), (s.getD i 1 : ℤ), (wgt s i : ℤ), ((s.getD i 1 * wgt s i : ℕ) : ℤ)]

/-- The record list at step `k`. -/
def recL (s : List ℕ) (k : ℕ) : List (List ℤ) := (List.range s.length).map (recI s k)

/-! ## The inner loop over the axes -/

/-! ## The outer loop over `k < S` -/

end NLogN.CR

end

section
-- Source module: NLogN.HCrOut
/-!
# `crtOut` as a register program (namespace `NLogN.CR`)

Input: `w` (length `S = ∏ sᵢ`) on stack `0`, `S` on `1`, `kb` (`S ≤ 2^kb`) on `2`, the keys
`natInts (keyL s)` on `3`, `M ≤ S` on `cMR`.  Output: `crtOut s M w` on `0`; all else restored.
-/

set_option linter.unusedSimpArgs false
set_option linter.unusedVariables false

namespace NLogN.CR

open IntMul IntMul.HvdH IntMul.HvdH.Ref RS

abbrev cMR : Fin HS := 33

end NLogN.CR

end

section
-- Source module: NLogN.HCrIn
/-!
# Chunking, `crtIn` and `evalDigits` as register programs (namespace `NLogN.CR`)
-/

set_option linter.unusedSimpArgs false
set_option linter.unusedVariables false

namespace NLogN.CR

open IntMul IntMul.HvdH IntMul.HvdH.Ref RS

abbrev cRB : Fin HS := 31
abbrev cLR : Fin HS := 32

/-! ## Chunking and zero padding onto stack `0` -/

/-! ## `evalDigits` -/

end NLogN.CR

end

section
-- Source module: NLogN.HCrChirp
/-!
# The chirp table program, part 1: the request keys (namespace `NLogN.CR`)

Registers (inputs, read only): `hM2` = `2r`, `hRr` = `r`, `hRh` = `r/2`, `hNa` = `K - dh`,
`hNb` = `dh`, `hTT` = `T`.  Scratch: `hQ`, `hI`, `hXq`, `hAc`.

`reqGen` pushes onto stack `0` the request records `[2·chirpKey x + 1, x, 0]` for `x < T`
(record `0` on top).
-/

set_option linter.unusedSimpArgs false
set_option linter.unusedVariables false

namespace NLogN.CR

open IntMul RS IntMul.HvdH IntMul.HvdH.Ref

abbrev hM2 : Fin HS := 4
abbrev hRr : Fin HS := 5
abbrev hRh : Fin HS := 6
abbrev hNa : Fin HS := 7
abbrev hNb : Fin HS := 8
abbrev hTT : Fin HS := 9
abbrev hQ : Fin HS := 10

/-! ## One axis step -/

/-! ## One request record -/

/-! ## All request records -/

end NLogN.CR

end

section
-- Source module: NLogN.HCrChirp2
/-!
# The chirp table program, part 2: the twiddle table (namespace `NLogN.CR`)

The twiddle program `tw` is abstract (`TwSpec`): from `c < R` on `hC`, `p` on `hPp` and `R` on
`hM2` it pushes `(twD p R c).1` onto `hVA` and `(twD p R c).2` onto `hVB`, and nothing else.
`tabGen tw` pushes the table records `[2c, re, im]` for `c < R` (record `0` on top).
-/

set_option linter.unusedSimpArgs false
set_option linter.unusedVariables false

namespace NLogN.CR

open IntMul RS IntMul.HvdH IntMul.HvdH.Ref

abbrev hPp : Fin HS := 17

end NLogN.CR

end

section
-- Source module: NLogN.HCrPar
/-!
# Parameter computation, part 1: search loops and the factorial loop (namespace `NLogN.CR`)

* `srch m C n c`: `m := 0`, then `n` times `m := m + (c < 0 ? 0 : 1)`.  If `c < 0` exactly when
  `m ≥ μ` (and `μ ≤ n`), the result is `μ` (`srch_spec`).  Instances: `Nat.size`, `Nat.clog 2`,
  `ceilSqrt`.
* `facP f C n`: `f := 1`, then `f := f · C` for `C = n, …, 1`: `n!`.
-/

set_option linter.unusedSimpArgs false
set_option linter.unusedVariables false

namespace NLogN.CR

open IntMul RS IntMul.HvdH IntMul.HvdH.Ref

/-! ## Search loops -/

def srchBody (m : Fin HS) (c : Ex HS) : RS := asg m (.add (.reg m) (.sel c (.lit 0) (.lit 1)))

def srch (m C : Fin HS) (n c : Ex HS) : RS := NLogN.RS.seq (push m (.lit 0)) (rep C n (srchBody m c))

/-! ## The factorial loop -/

def facP (f C : Fin HS) (n : Ex HS) : RS :=
  NLogN.RS.seq (push f (.lit 1)) (rep C n (asg f (.mul (.reg f) (.reg C))))

end NLogN.CR

end

section
-- Source module: NLogN.HCrPar2
/-!
# Parameter computation, part 2: facts about `mkParams` (namespace `NLogN.CR`)
-/

set_option linter.unusedSimpArgs false
set_option linter.unusedVariables false

namespace NLogN.CR

open IntMul RS IntMul.HvdH IntMul.HvdH.Ref

/-- The exponent `E = ⌈log₂ 4L⌉` of `mkParams`. -/
def parEx (K N : ℕ) : ℕ := Nat.clog 2 (4 * (mkParams K N).L)

end NLogN.CR

end

section
-- Source module: NLogN.HCrPar3
/-!
# Parameter computation, part 3: the scalar stage (namespace `NLogN.CR`)

Input: `N ≥ 1` on `rN`, `nb` with `N < 2^nb` on `rNb`.
Stage A pushes `b, L, E, e, dh, K-dh, r, r/2, 2r, T, logOuter, 2^b, 2L-1`.
-/

set_option linter.unusedSimpArgs false
set_option linter.unusedVariables false

namespace NLogN.CR

open IntMul RS IntMul.HvdH IntMul.HvdH.Ref

abbrev rFn : Fin HS := 26
abbrev rS : Fin HS := 28
abbrev rN : Fin HS := 34
abbrev rNb : Fin HS := 35
abbrev rB : Fin HS := 36
abbrev rE : Fin HS := 37
abbrev re : Fin HS := 38
abbrev rLO : Fin HS := 39
abbrev rRec : Fin HS := 40
abbrev rAl : Fin HS := 42
abbrev rPP : Fin HS := 44

/-- Push the integer `v` onto stack `r`. -/
def pushV (s : IState HS) (r : Fin HS) (v : ℤ) : IState HS := s.upd r (.int v :: s.G r)

@[simp] theorem pushV_ip (s : IState HS) (r : Fin HS) (v : ℤ) : (pushV s r v).ip = s.ip := rfl
@[simp] theorem pushV_out (s : IState HS) (r : Fin HS) (v : ℤ) : (pushV s r v).out = s.out := rfl

/-- `max x 0` as an expression. -/
def max0 (x : Ex HS) : Ex HS := .sel x (.lit 0) x

/-! ## Stage A -/

def stA (K : ℕ) : RS :=
  seqL [srch rB hQ (.reg rNb) (.add (.reg rN) (.neg (.pow2 (.reg rB)))),
    push cLR (.ndiv (.add (.add (.reg rN) (.reg rB)) (.lit (-1))) (.reg rB)),
    srch rE hQ (.add (.reg rNb) (.lit 2))
      (.add (.add (.mul (.lit 4) (.reg cLR)) (.lit (-1))) (.neg (.pow2 (.reg rE)))),
    push re (.ndiv (.add (.reg rE) (.lit ((K : ℤ) - 1))) (.lit K)),
    push hNb (.add (.mul (.lit K) (.reg re)) (.neg (.reg rE))),
    push hNa (.add (.lit K) (.neg (.reg hNb))),
    push hRr (.pow2 (.reg re)),
    push hRh (.pow2 (.add (.reg re) (.lit (-1)))),
    push hM2 (.pow2 (.add (.reg re) (.lit 1))),
    push hTT (.pow2 (.reg rE)),
    push rLO (.add (.reg rE) (.neg (.reg re))),
    push cRB (.pow2 (.reg rB)),
    push cMR (.add (.mul (.lit 2) (.reg cLR)) (.lit (-1)))]

/-- The state after stage A. -/
def outA (s : IState HS) (K N : ℕ) : IState HS :=
  let q := mkParams K N
  let E := parEx K N
  pushV (pushV (pushV (pushV (pushV (pushV (pushV (pushV (pushV (pushV (pushV (pushV (pushV s
    rB q.b) cLR q.L) rE E) re q.e) hNb q.dh) hNa ((q.K - q.dh : ℕ) : ℤ)) hRr ((2 ^ q.e : ℕ) : ℤ))
    hRh ((2 ^ (q.e - 1) : ℕ) : ℤ)) hM2 ((2 ^ (q.e + 1) : ℕ) : ℤ)) hTT (q.tdims.prod : ℤ))
    rLO (q.logOuter : ℤ)) cRB ((2 ^ q.b : ℕ) : ℤ)) cMR ((2 * q.L - 1 : ℕ) : ℤ)

/-! ## Stage B: `α, p, P, fn` -/

abbrev pT1 : Fin HS := 18
abbrev pT2 : Fin HS := 19
abbrev pT3 : Fin HS := 21

/-- `x - 1 - m·m < 0 ↔ ceilSqrt x ≤ m`. -/
def csE (x m : Fin HS) : Ex HS := .add (.add (.reg x) (.lit (-1))) (.neg (.mul (.reg m) (.reg m)))

def pE (K : ℕ) : Ex HS :=
  .add (.add (.add (.add (.mul (.lit 2) (.reg rB)) (.mul (.lit 2) (.reg rE)))
    (.mul (.lit (6 * K)) (.mul (.reg rAl) (.reg rAl)))) (.mul (.lit 8) (.reg pT3))) (.lit (4 * K + 64))

def stB (K : ℕ) : RS :=
  seqL [srch pT1 hQ (.reg rB) (csE rB pT1), srch pT2 hQ (.reg pT1) (csE pT1 pT2),
    push rAl (.mul (.lit K) (.add (.reg pT2) (.lit 2))),
    srch pT3 hQ (.reg rB) (.add (.reg rB) (.neg (.pow2 (.reg pT3)))),
    push hPp (pE K), push rPP (.add (.add (.mul (.lit 2) (.reg hPp)) (.reg rE)) (.lit 8)),
    push rFn (.mul (.mul (.lit 4) (.reg hRr)) (.reg rPP)),
    pop pT3, pop pT2, pop pT1]

def alV (K b : ℕ) : ℕ := K * (ceilSqrt (ceilSqrt b) + 2)
def pV (K b E : ℕ) : ℕ := 2 * b + 2 * E + 6 * K * alV K b ^ 2 + 8 * Nat.size b + 4 * K + 64
def PV (K b E : ℕ) : ℕ := 2 * pV K b E + E + 8

/-! ## Stage C: `Q, m1, m2` (the factorial `(256K²)!` by a loop) -/

abbrev pT4 : Fin HS := 22
abbrev pT5 : Fin HS := 23
abbrev pT6 : Fin HS := 24
abbrev pTQ : Fin HS := 25

def facN (K : ℕ) : ℕ := 256 * K ^ 2

def QE (K : ℕ) : Ex HS :=
  .mul (.mul (.lit 2) (.reg pT4))
    (.add (.ndiv (.reg hRr) (.mul (.mul (.lit 2) (.reg pT4)) (.lit (facN K)))) (.lit 1))

def mE (K : ℕ) (x : Fin HS) : Ex HS :=
  .ndiv (max0 (.add (.mul (.lit ((4 * K - 1 : ℕ) : ℤ)) (.reg x)) (.lit (-(4 * K : ℤ)))))
    (.mul (.lit (4 * K)) (.reg pTQ))

def stC (K : ℕ) : RS :=
  seqL [facP pT4 hQ (.lit (facN K)), push pTQ (QE K), push pT5 (mE K hRh), push pT6 (mE K hRr),
    pop pT4]

def m1V (K r : ℕ) : ℕ := ((4 * K - 1) * (r / 2) - 4 * K) / (4 * K * copQ K r)
def m2V (K r : ℕ) : ℕ := ((4 * K - 1) * r - 4 * K) / (4 * K * copQ K r)

/-! ## Stage D: the source lengths `sᵢ`, their records and `S = ∏ sᵢ` -/

abbrev pTS : Fin HS := 20

def iEx : Ex HS := .add (.reg hQ) (.lit (-1))

def sE : Ex HS :=
  .add (.lit 1) (.mul (.reg pTQ) (.sel (.add iEx (.neg (.reg hNb)))
    (max0 (.add (.reg pT5) (.neg iEx))) (max0 (.add (.reg pT6) (.neg (.add iEx (.neg (.reg hNb))))))))

def sBody : RS :=
  seqL [push pTS sE, push rRec (.mul (.reg pTS) (.reg rS)), push rRec (.reg rS), push rRec (.reg pTS),
    push rRec (.lit 0), asg rS (.mul (.reg pTS) (.reg rS)), pop pTS]

/-- The records of the axes `K - k, …, K - 1` (the first on top). -/
def recsF (sl : List ℕ) (K k : ℕ) : List ℤ :=
  ((List.range k).map fun j => recI sl 0 (K - k + j)).flatten

def Wk (sl : List ℕ) (K k : ℕ) : ℕ := (sl.drop (K - k)).prod

def dSt (t : IState HS) (sl : List ℕ) (K k : ℕ) : IState HS :=
  pushV (t.upd rRec (ints (recsF sl K k) ++ t.G rRec)) rS (Wk sl K k)

def stD (K : ℕ) : RS := seqL [push rS (.lit 1), rep hQ (.lit K) sBody, pop pT6, pop pT5, pop pTQ]

/-- The state after stages C and D: records `[0, sᵢ, wᵢ, sᵢwᵢ]` on `rRec` and `S` on `rS`. -/
def outD (t : IState HS) (sl : List ℕ) : IState HS :=
  pushV (t.upd rRec (ints (recL sl 0).flatten ++ t.G rRec)) rS (sl.prod)

/-! ## The whole parameter program -/

/-- **`mkParams` as a register program.** -/
def parP (K : ℕ) : RS := seqL [stA K, NLogN.CR.stB K, stC K, stD K]

/-- The state after `parP`. -/
def parOut (s : IState HS) (K N : ℕ) : IState HS :=
  let q := mkParams K N
  outD (pushV (pushV (pushV (pushV (outA s K N) rAl q.α) hPp q.p) rPP q.P) rFn q.fn) q.s

end NLogN.CR

end