Prove2Me

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

Definition
IntMul_HvdH_M_Defs5

by xuanji · Oct 13, 2026 · Mathlib 0df444a (Lean v4.33.1)

harvey-van-der-hoeven-tminteger-multiplicationturing-machines

Definitions (Lean namespace NLogN) for a multitape-Turing-machine implementation of the Harvey–van der Hoeven O(nlog⁡n)O(n\log n) integer multiplication algorithm, continuing IntMul_HvdH_M_Defs1–IntMul_HvdH_M_Defs4 (part 5): the programs and cost functions of the recursive step (parameter, chirp, key, CRT and resampling phases, the tensor/axis combinators and the Gaussian resampling tables) and the final assembly constants.

Source modules (in order): NLogN.SSBase, NLogN.HTrans, NLogN.HPerm, NLogN.HCrMath, NLogN.HCrLin, NLogN.HCrKey, NLogN.HCrLoops, NLogN.HCrOut, NLogN.HCrIn, NLogN.HCrChirpMath, NLogN.HCrChirp, NLogN.HCrChirp2, NLogN.HCrChirp3, NLogN.HCrChirp4, NLogN.HCrChirpCost, NLogN.HCrPar6, NLogN.HNumAtan, NLogN.HExprDer, NLogN.HNumPi, NLogN.HNumExp, NLogN.HNumCall, NLogN.HNumExp2, NLogN.HNumExp3, NLogN.HNumExpPos, NLogN.HNumTw, NLogN.HNumTw2, NLogN.HNumScale, NLogN.HNumTwD, NLogN.HCrTw, NLogN.HGl1.

The file was generated by deleting from the development all theorems not needed by the definitions; it keeps all definitions plus the few lemmas the definitions themselves need (proof fields of structures, decreasing_by facts) and simple @[simp]/rfl lemmas.

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

import Definitions.Def_IntMul_HvdH_M_Defs4

set_option autoImplicit false
set_option relaxedAutoImplicit false
set_option maxHeartbeats 1000000

section
-- Source module: NLogN.SSBase
/-!
# Schönhage–Strassen in the item language: stack layout and Fermat-ring arithmetic

The program uses `NS` item stacks with fixed roles, so that stack (in)equalities are decided by
`simp`.  The current ring `ℤ/(2^n+1)` is given by the ruler `2^n` on top of `RN` and the modulus
`2^n+1` on top of `RF`.
-/

namespace NLogN

open IntMul

/-! ## Fermat reduction -/

variable {P : ℕ}

/-- The value computed by `fredA`. -/
def fredW (n : ℕ) (z : ℤ) : ℤ :=
  let q := sg z (z.natAbs / 2 ^ n)
  sg z (z.natAbs % 2 ^ n) + -sg q (q.natAbs % 2 ^ n) + sg q (q.natAbs / 2 ^ n)

/-- The final value after the four normalization steps. -/
def fnorm (F w : ℤ) : ℤ :=
  let w1 := if w < 0 then F + w else w
  let w2 := if w1 < 0 then F + w1 else w1
  let w3 := if ¬ (w2 + -F < 0) then w2 + -F else w2
  if ¬ (w3 + -F < 0) then w3 + -F else w3

end NLogN

namespace NLogN

open IntMul

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

end NLogN

end

section
-- Source module: NLogN.HTrans
/-!
# Transposition of row-major arrays by radix sort

The `n × m` row-major array of records `D` (each a list of `d` integers) on stack `0` is replaced
by its transpose: the element at `(i, j)` gets the key `j * n + i`, the keyed records are radix
sorted, and the keys are stripped.

Registers: data `0`, `n` on `1`, `m` on `2`, number of key bits `kb` on `3`.
Scratch: `40` (temporary stack), `41` (current key), `42` (row key `i`), `43`, `44` (loop
counters), `45` (item count of `40`), `54` (record count), plus those of `sortP` and `mvK`.
-/

set_option linter.unusedSimpArgs false

namespace NLogN

open IntMul RS

/-! ## Stripping the keys -/

/-! ## The whole transposition -/

/-! ## Costs -/

/-! ## Integer and Gaussian-pair arrays -/

/-- Gaussian pairs as flat integer lists `re, im, re, im, …`. -/
def pairInts (L : List (ℤ × ℤ)) : List ℤ := L.flatMap fun p => [p.1, p.2]

end NLogN

end

section
-- Source module: NLogN.HPerm
/-!
# Permuting a list of records by a list of target indices

Stack `0` holds `n` records of `d` integers, `n` is on `1`, the number of key bits `kb`
(`n ≤ 2^kb`) on `2`, and stack `3` holds the target indices `Kl` (a permutation of `0, …, n-1`).
`permP d` moves record `p` to position `Kl[p]`; the index list on `3` is restored.
Scratch: `40`, `41`, `43`, `45`, `51`, `54` and those of `sortP`.
-/

set_option linter.unusedSimpArgs false

namespace NLogN

open IntMul RS

def pRec (Kl : List ℕ) (D : List (List ℤ)) (p : ℕ) : List ℤ := (Kl.getD p 0 : ℤ) :: D.getD p []

def pflat (Kl : List ℕ) (D : List (List ℤ)) (p : ℕ) : List ℤ :=
  ((List.range p).map (pRec Kl D)).flatten

def natInts (Kl : List ℕ) : List Item := ints (List.map (Nat.cast : ℕ → ℤ) Kl)

def pBody (d : ℕ) : RS :=
  NLogN.RS.seq (push 40 (.reg 3)) (NLogN.RS.seq (mvK 0 40 51 (.lit d)) (NLogN.RS.seq (push 41 (.reg 3))
    (NLogN.RS.seq (pop 3) (asg 45 (.add (.reg 45) (.lit (d + 1)))))))

def pSt (s : IState HS) (Kl : List ℕ) (D : List (List ℤ)) (Rest RK : List Item) (d p : ℕ) :
    IState HS :=
  ((((s.upd 0 (ints (D.drop p).flatten ++ Rest)).upd 3 (natInts (Kl.drop p) ++ RK)).upd 40
    (ints (pflat Kl D p).reverse ++ s.G 40)).upd 41
    ((natInts (Kl.take p)).reverse ++ s.G 41)).upd 45
    (.int ((p * (d + 1) : ℕ) : ℤ) :: s.G 45)

def pLoop (d : ℕ) : RS := rep 43 (.reg 1) (pBody d)

def pAttachP (d : ℕ) : RS :=
  NLogN.RS.seq (push 45 (.lit 0)) (NLogN.RS.seq (pLoop d) (NLogN.RS.seq (mvK 40 0 51 (.reg 45))
    (NLogN.RS.seq (mvK 41 3 51 (.reg 1)) (pop 45))))

/-- The permuted list: position `q` holds the record whose target index is `q`. -/
def permL (Kl : List ℕ) (D : List (List ℤ)) : List (List ℤ) :=
  (List.range Kl.length).map fun q => D.getD (Kl.idxOf q) []

def permP (d : ℕ) : RS :=
  NLogN.RS.seq (pAttachP d) (NLogN.RS.seq (sortP (d + 1)) (NLogN.RS.seq (push 54 (.reg 1)) (NLogN.RS.seq (push 45 (.lit 0))
    (NLogN.RS.seq (stripP d) (NLogN.RS.seq (pop 45) (pop 54))))))

/-! ## Costs -/

end NLogN

end

section
-- Source module: NLogN.HCrMath
/-!
# Pure facts for chunking, CRT reindexing and the chirp table (namespace `NLogN.CR`)
-/

set_option linter.unusedSimpArgs false
set_option linter.unusedVariables false

namespace NLogN.CR

open IntMul IntMul.HvdH IntMul.HvdH.Ref

/-! ## Mixed-radix ranks -/

/-- The box index of row-major position `q` (inverse of `rmPos`). -/
def unrk : List ℕ → ℕ → ℕ → ℕ
  | [], _ => fun _ => 0
  | n :: ns, q => fun k => if k = 0 then q / ns.prod else unrk ns (q % ns.prod) (k - 1)

/-! ## The CRT keys -/

/-- The residues of `k`. -/
def resJ (s : List ℕ) (k : ℕ) : ℕ → ℕ := fun i => k % s.getD i 1

/-- Target position of digit `k` in `crtIn`: the row-major position of its residues. -/
def crtKey (s : List ℕ) (k : ℕ) : ℕ := rmPos s (resJ s k)

/-- The key list fed to `permP` for `crtIn`. -/
def keyL (s : List ℕ) : List ℕ := (List.range s.prod).map (crtKey s)

/-- The source lengths are positive and pairwise coprime. -/
def CrtOK (s : List ℕ) : Prop := s.Pairwise Nat.Coprime ∧ ∀ x ∈ s, 0 < x

/-- Zero padding of the digit list to length `n`. -/
def padL (c : List ℤ) (n : ℕ) : List ℤ := c ++ List.replicate (n - c.length) 0

/-- The inverse key list. -/
def invL (s : List ℕ) : List ℕ := (List.range s.prod).map fun q => crtIdx s (unrk s q)

end NLogN.CR

end

section
-- Source module: NLogN.HCrLin
/-!
# Linear cost bounds for straight-line register code (namespace `NLogN.CR`)

`LinB B p s`: every expression evaluated by the loop-free program `p` from `s` is an `Ex.Lin`
expression whose values have at most `B` bits.  Then `p.cost s ≤ lsz p · 3000 · (B + 2)`.
-/

set_option linter.unusedSimpArgs false
set_option linter.unusedVariables false

namespace NLogN.CR

open IntMul RS

def LinB (B : ℕ) : RS → IState HS → Prop
  | .skip, _ => True
  | .asg i e, s => e.Lin ∧ e.Bnd s B ∧ zsz (s.top i) ≤ B
  | .push _ e, s => e.Lin ∧ e.Bnd s B
  | .pop i, s => zsz (s.top i) ≤ B
  | .seq a b, s => LinB B a s ∧ LinB B b (a.run s)
  | .ite c a b, s => c.Lin ∧ c.Bnd s B ∧ (if c.eval s < 0 then LinB B a s else LinB B b s)
  | .rep _ _ _, _ => False

def lsz : RS → ℕ
  | .skip => 1
  | .asg _ e | .push _ e => e.size + 1
  | .pop _ => 1
  | .seq a b => lsz a + lsz b
  | .ite c a b => c.size + 1 + lsz a + lsz b
  | .rep _ _ _ => 0

end NLogN.CR

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

abbrev kRA : Fin HS := 4
abbrev kRB : Fin HS := 5
abbrev kX1 : Fin HS := 6
abbrev kX2 : Fin HS := 7
abbrev kX3 : Fin HS := 8
abbrev kX4 : Fin HS := 9
abbrev kPS : Fin HS := 10
abbrev kT : Fin HS := 11
abbrev kC1 : Fin HS := 12
abbrev kC2 : Fin HS := 13
abbrev kC3 : Fin HS := 14
abbrev kSR : Fin HS := 15
abbrev kKR : Fin HS := 16
abbrev kK4 : Fin HS := 17

/-- Advance one record `[r, s, w, sw]` from `kRA` to `kRB` and update the position. -/
def axBody : RS :=
  seqL [push kX1 (.reg kRA), pop kRA, push kX2 (.reg kRA), pop kRA,
    push kX3 (.reg kRA), pop kRA, push kX4 (.reg kRA), pop kRA,
    ite (.add (.add (.reg kX1) (.lit 1)) (.neg (.reg kX2)))
      (NLogN.RS.seq (push kRB (.add (.reg kX1) (.lit 1))) (asg kPS (.add (.reg kPS) (.reg kX3))))
      (NLogN.RS.seq (push kRB (.lit 0)) (asg kPS (.add (.add (.reg kPS) (.reg kX3)) (.neg (.reg kX4))))),
    push kRB (.reg kX2), push kRB (.reg kX3), push kRB (.reg kX4),
    pop kX1, pop kX2, pop kX3, pop kX4]

/-! ## Records -/

/-- Residues of `k + 1` on the axes `< m`, of `k` on the others. -/
def mixJ (s : List ℕ) (k m : ℕ) : ℕ → ℕ :=
  fun i => if i < m then (k + 1) % s.getD i 1 else k % s.getD i 1

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

def innerSt (t1 : IState HS) (s : List ℕ) (k : ℕ) (Rst RB0 LP : List Item) (m : ℕ) : IState HS :=
  ((t1.upd kRA (ints ((recL s k).drop m).flatten ++ Rst)).upd kRB
    (ints ((recL s (k + 1)).take m).flatten.reverse ++ RB0)).upd kPS
    (.int (rmPos s (mixJ s k m)) :: LP)

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

def outBody : RS :=
  NLogN.RS.seq (push kT (.reg kPS)) (NLogN.RS.seq (rep kC2 (.reg kKR) NLogN.CR.axBody) (mvK kRB kRA kC3 (.reg kK4)))

/-- The keys emitted so far. -/
def keyPre (s : List ℕ) (k : ℕ) : List ℤ := (List.range k).map fun k' => (crtKey s k' : ℤ)

def keySt (s0 : IState HS) (s : List ℕ) (Rst LP T0 : List Item) (k : ℕ) : IState HS :=
  ((s0.upd kRA (ints (recL s k).flatten ++ Rst)).upd kPS (.int (crtKey s k) :: LP)).upd kT
    (ints (keyPre s k).reverse ++ T0)

/-- **Key generation.** -/
def keyGen : RS :=
  NLogN.RS.seq (push kPS (.lit 0)) (NLogN.RS.seq (rep kC1 (.reg kSR) outBody) (NLogN.RS.seq (pop kPS) (mvK kT 3 kC3 (.reg kSR))))

end NLogN.CR

end

section
-- Source module: NLogN.HCrLoops
/-! # Small counted loops: popping, pushing zeros, pushing `0, …, n-1` (namespace `NLogN.CR`) -/

set_option linter.unusedSimpArgs false
set_option linter.unusedVariables false

namespace NLogN.CR

open IntMul RS

/-! ## Popping `n` integers -/

def dropN (i C : Fin HS) (n : Ex HS) : RS := rep C n (pop i)

/-! ## Pushing `n` zeros -/

def zerosP (i C : Fin HS) (n : Ex HS) : RS := rep C n (push i (.lit 0))

/-! ## Pushing `0, 1, …, n-1` (`0` on top) -/

def rangeP (i C : Fin HS) (n : Ex HS) : RS := rep C n (push i (.add (.reg C) (.lit (-1))))

def rangeZ (N k : ℕ) : List ℤ := (List.range k).map fun j => ((N - k + j : ℕ) : ℤ)

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

def crtOutP : RS :=
  NLogN.RS.seq (rangeP 0 kC1 (.reg 1)) (NLogN.RS.seq (permP 1) (NLogN.RS.seq (mvK 0 kT kC3 (.reg 1))
  (NLogN.RS.seq (mvK kT 3 kC3 (.reg 1)) (NLogN.RS.seq (permP 1) (NLogN.RS.seq (dropN 3 kC1 (.reg 1))
  (NLogN.RS.seq (mvK 0 kT kC3 (.reg cMR)) (NLogN.RS.seq (dropN 0 kC1 (.add (.reg 1) (.neg (.reg cMR))))
    (mvK kT 0 kC3 (.reg cMR)))))))))

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 cX : Fin HS := 30

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

/-- `x` on `cX` (consumed), the ruler `2^b` on `cRB`, `L` on `cLR`, `S ≥ L` on `1`:
pushes `padL (chunksL b L x) S` onto `0`. -/
def chunkPadP : RM :=
  .seq (.rs (NLogN.CR.zerosP 0 kC1 (.add (.reg 1) (.neg (.reg cLR)))))
  (.seq (chunkP cX cRB kT kC1 (.reg cLR))
  (.rs (NLogN.RS.seq (mvK kT 0 kC3 (.reg cLR)) (pop cX))))

/-- **`crtIn ∘ chunksL`**: chunk `x`, pad, and permute by the keys on `3`. -/
def crtInP : RM := .seq chunkPadP (.rs (permP 1))

/-! ## `evalDigits` -/

/-- Nonnegative digits `c` on `CA` (consumed, `c₀` on top), `2^b` on `CR`, `|c|` on `CM`:
pushes `evalDigits b c` onto `CY`. -/
def evalP : RM := .seq (.rs (push CY (.lit 0))) carryP

end NLogN.CR

end

section
-- Source module: NLogN.HCrChirpMath
/-!
# Pure facts for the chirp table (namespace `NLogN.CR`)

* `unrk_snoc`, `wsum_snoc`: peeling the last axis of a row-major index.
* `axIt`: one step of the per-axis accumulation loop; `axIt_iter`: `n` steps over `n` equal axes.
* `chirpKey`: the chirp index mod `2r` of row-major position `q`; `chirpL_eq`.
* `scanF`: the merge scan over the sorted table/request records; `scanF_sorted`.
-/

set_option linter.unusedSimpArgs false
set_option linter.unusedVariables false

namespace NLogN.CR

open IntMul IntMul.HvdH IntMul.HvdH.Ref

/-! ## Row-major indices: the last axis -/

/-- `∑ᵢ uᵢ · jᵢ²` for the box index `j` of row-major position `q`. -/
def wsum (ts : List ℕ) (u : ℕ → ℕ) (q : ℕ) : ℕ :=
  ((List.range ts.length).map fun i => u i * unrk ts q i ^ 2).sum

/-! ## The accumulation loop over equal axes -/

/-- One step: add `u · (x mod t)²` to the accumulator (mod `M`), divide `x` by `t`. -/
def axIt (t u M : ℕ) (z : ℕ × ℕ) : ℕ × ℕ := ((z.1 + u * (z.2 % t) ^ 2) % M, z.2 / t)

/-! ## The chirp index -/

/-- The weights `r / tᵢ`. -/
def uW (q : Params) (i : ℕ) : ℕ := q.r / 2 ^ q.te i

/-- The key of the chirp entry at row-major position `x`: the chirp index mod `2r`. -/
def chirpKey (q : Params) (x : ℕ) : ℕ := wsum q.tdims (uW q) x % (2 * q.r)

/-! ## The merge scan -/

/-- Table record `[2c, re, im]`. -/
def tabRec (v : ℕ → ℤ × ℤ) (c : ℕ) : List ℤ := [2 * (c : ℤ), (v c).1, (v c).2]

/-- Request record `[2κ(x)+1, x, 0]`. -/
def reqRec (κ : ℕ → ℕ) (x : ℕ) : List ℤ := [2 * (κ x : ℤ) + 1, x, 0]

/-- The scan: a table record sets the current value, a request emits `(x, current)`. -/
def scanF (cur : ℤ × ℤ) : List (List ℤ) → List (ℤ × (ℤ × ℤ))
  | [] => []
  | r :: Z => if r.headD 0 % 2 = 1 then (r.getD 1 0, cur) :: scanF cur Z
      else scanF (r.getD 1 0, r.getD 2 0) Z

/-- The current value after scanning. -/
def curF (cur : ℤ × ℤ) : List (List ℤ) → ℤ × ℤ
  | [] => cur
  | r :: Z => if r.headD 0 % 2 = 1 then curF cur Z else curF (r.getD 1 0, r.getD 2 0) Z

/-- The request indices of `Z`, in order. -/
def reqXs (Z : List (List ℤ)) : List ℕ :=
  (Z.filter fun r => r.headD 0 % 2 = 1).map fun r => (r.getD 1 0).toNat

def IsRec (v : ℕ → ℤ × ℤ) (κ : ℕ → ℕ) (r : List ℤ) : Prop :=
  (∃ c, r = tabRec v c) ∨ (∃ x, r = reqRec κ x)

/-- Every request finds its table entry in the rest, or it is the current value and every
remaining table entry comes later. -/
def Good (v : ℕ → ℤ × ℤ) (κ : ℕ → ℕ) (cur : ℤ × ℤ) (Z : List (List ℤ)) : Prop :=
  ∀ x, reqRec κ x ∈ Z → tabRec v (κ x) ∈ Z ∨ (cur = v (κ x) ∧ ∀ c, tabRec v c ∈ Z → κ x < c)

/-! ## Gathering through the sorted merge -/

/-- All table records `c < R` followed by all requests `x < T`. -/
def mergeL (v : ℕ → ℤ × ℤ) (κ : ℕ → ℕ) (R T : ℕ) : List (List ℤ) :=
  (List.range R).map (tabRec v) ++ (List.range T).map (reqRec κ)

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 hI : Fin HS := 11
abbrev hXq : Fin HS := 12
abbrev hAc : Fin HS := 13

/-! ## One axis step -/

def fE (t : Fin HS) : Ex HS := .rem (.reg hXq) (.reg t)

def axE (t : Fin HS) (u : ℤ) : Ex HS :=
  .rem (.add (.reg hAc) (.mul (.lit u) (.mul (fE t) (fE t)))) (.reg hM2)

def axStep (t : Fin HS) (u : ℤ) : RS :=
  NLogN.RS.seq (asg hAc (axE t u)) (asg hXq (.quo (.reg hXq) (.reg t)))

def axLoop (t : Fin HS) (u : ℤ) (cnt : Fin HS) : RS := rep NLogN.CR.hI (.reg cnt) (axStep t u)

def axSt (s : IState HS) (LA LX : List Item) (z : ℕ × ℕ) : IState HS :=
  (s.upd hAc (.int z.1 :: LA)).upd hXq (.int z.2 :: LX)

/-! ## One request record -/

def reqBody : RS :=
  NLogN.RS.seq (push hXq (.add (.reg hQ) (.lit (-1))))
  (NLogN.RS.seq (push hAc (.lit 0))
  (NLogN.RS.seq (axLoop hRr 1 hNa)
  (NLogN.RS.seq (axLoop hRh 2 hNb)
  (NLogN.RS.seq (push 0 (.lit 0))
  (NLogN.RS.seq (push 0 (.add (.reg hQ) (.lit (-1))))
  (NLogN.RS.seq (push 0 (.add (.add (.reg hAc) (.reg hAc)) (.lit 1)))
  (NLogN.RS.seq (pop hAc) (pop hXq))))))))

/-- The register contents the chirp programs read. -/
structure ChirpRegs (q : Params) (w : IState HS) : Prop where
  m2 : ∃ L, w.G hM2 = .int ((2 ^ (q.e + 1) : ℕ) : ℤ) :: L
  rr : ∃ L, w.G hRr = .int ((2 ^ q.e : ℕ) : ℤ) :: L
  rh : ∃ L, w.G hRh = .int ((2 ^ (q.e - 1) : ℕ) : ℤ) :: L
  na : ∃ L, w.G hNa = .int ((q.K - q.dh : ℕ) : ℤ) :: L
  nb : ∃ L, w.G hNb = .int (q.dh : ℤ) :: L
  tt : ∃ L, w.G hTT = .int (q.tdims.prod : ℤ) :: L

/-! ## All request records -/

def reqGen : RS := rep hQ (.reg hTT) reqBody

/-- The last `k` request records (record `T - k` first). -/
def reqTail (κ : ℕ → ℕ) (T k : ℕ) : List (List ℤ) :=
  (List.range k).map fun j => reqRec κ (T - k + j)

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 hC : Fin HS := 14
abbrev hVA : Fin HS := 15
abbrev hVB : Fin HS := 16

/-- The abstract twiddle program (relocated by the caller): it reads `c` on `hC`, `p` on `hPp`
and `R` on `hM2`, pushes `twD p R c` onto `hVA`, `hVB`, does not write the loop counter `hQ`,
and costs at most `Ctw`. -/
def TwSpec (M : MulSpec) (tw : RM) (p R Ctw : ℕ) : Prop :=
  hQ ∉ tw.writes ∧ ∀ (t : IState HS) (c : ℕ) (L LP LM : List Item), t.G hC = .int c :: L →
    t.G hPp = .int p :: LP → t.G hM2 = .int R :: LM → c < R →
    tw.ok M t ∧ tw.run t = (t.upd hVA (.int (twD p R c).1 :: t.G hVA)).upd hVB
      (.int (twD p R c).2 :: t.G hVB) ∧ tw.cost M t ≤ Ctw

def tabPush : RS :=
  NLogN.RS.seq (push 0 (.reg hVB)) (NLogN.RS.seq (push 0 (.reg hVA)) (NLogN.RS.seq (push 0 (.add (.reg hC) (.reg hC)))
    (NLogN.RS.seq (pop hVA) (NLogN.RS.seq (pop hVB) (pop hC)))))

def tabBody (tw : RM) : RM :=
  .seq (.rs (push hC (.add (.reg hQ) (.lit (-1))))) (.seq tw (.rs tabPush))

def tabGen (tw : RM) : RM := .rep hQ (.reg hM2) (tabBody tw)

/-- The last `k` table records (record `R - k` first). -/
def tabTail (v : ℕ → ℤ × ℤ) (R k : ℕ) : List (List ℤ) :=
  (List.range k).map fun j => tabRec v (R - k + j)

end NLogN.CR

end

section
-- Source module: NLogN.HCrChirp3
/-!
# The chirp table program, part 3: the merge scan (namespace `NLogN.CR`)

`scanLoop` consumes `n` (on stack `1`) records `[k, a, b]` from stack `0`: a record with even
key sets the current value `(a, b)` (registers `hCa`, `hCb`); a record with odd key emits the
current value onto `hOut` and `a` onto `hKs`.
-/

set_option linter.unusedSimpArgs false
set_option linter.unusedVariables false

namespace NLogN.CR

open IntMul RS IntMul.HvdH IntMul.HvdH.Ref

abbrev hK : Fin HS := 18
abbrev hXx : Fin HS := 19
abbrev hYy : Fin HS := 21
abbrev hCa : Fin HS := 22
abbrev hCb : Fin HS := 23
abbrev hOut : Fin HS := 24
abbrev hKs : Fin HS := 25

def parE : Ex HS := .neg (.rem (.reg NLogN.CR.hK) (.lit 2))

def emitP : RS := NLogN.RS.seq (push hOut (.reg hCa)) (NLogN.RS.seq (push hOut (.reg hCb)) (push hKs (.reg hXx)))

def setP : RS := NLogN.RS.seq (asg hCa (.reg hXx)) (asg hCb (.reg hYy))

def scanPre : RS :=
  NLogN.RS.seq (push NLogN.CR.hK (.reg 0)) (NLogN.RS.seq (pop 0) (NLogN.RS.seq (push hXx (.reg 0)) (NLogN.RS.seq (pop 0)
  (NLogN.RS.seq (push hYy (.reg 0)) (pop 0)))))

def scanPost : RS := NLogN.RS.seq (pop NLogN.CR.hK) (NLogN.RS.seq (pop hXx) (pop hYy))

def scanBody : RS := NLogN.RS.seq scanPre (NLogN.RS.seq (ite parE emitP setP) scanPost)

def scanLoop : RS := rep hQ (.reg 1) scanBody

/-! ## The scan loop -/

def outRev (E : List (ℤ × (ℤ × ℤ))) : List ℤ := ((E.map fun e => [e.2.1, e.2.2]).flatten).reverse

def keysRev (E : List (ℤ × (ℤ × ℤ))) : List ℤ := (E.map Prod.fst).reverse

/-- A record `[k, a, b]` with `k ≥ 0`. -/
def Rec3 (r : List ℤ) : Prop := ∃ k a b : ℤ, r = [k, a, b] ∧ 0 ≤ k

def scSt (s : IState HS) (Z : List (List ℤ)) (Rst : List Item) (cur0 : ℤ × ℤ) (LA LB : List Item)
    (k : ℕ) : IState HS :=
  ((((s.upd 0 (ints (Z.drop k).flatten ++ Rst)).upd hCa (.int (curF cur0 (Z.take k)).1 :: LA)).upd
    hCb (.int (curF cur0 (Z.take k)).2 :: LB)).upd hOut
    (ints (outRev (scanF cur0 (Z.take k))) ++ s.G hOut)).upd hKs
    (ints (keysRev (scanF cur0 (Z.take k))) ++ s.G hKs)

end NLogN.CR

end

section
-- Source module: NLogN.HCrChirp4
/-!
# The chirp table program, part 4: sort, scan, gather (namespace `NLogN.CR`)
-/

set_option linter.unusedSimpArgs false
set_option linter.unusedVariables false

namespace NLogN.CR

open IntMul RS IntMul.HvdH IntMul.HvdH.Ref

/-- Sort the merged records, scan them, and gather the emitted values by the request indices. -/
def chirpTail : RS :=
  NLogN.RS.seq (push 1 (.add (.reg hM2) (.reg hTT)))
  (NLogN.RS.seq (sortP 3)
  (NLogN.RS.seq (push hCa (.lit 0)) (NLogN.RS.seq (push hCb (.lit 0))
  (NLogN.RS.seq scanLoop
  (NLogN.RS.seq (pop hCa) (NLogN.RS.seq (pop hCb) (NLogN.RS.seq (pop 1)
  (NLogN.RS.seq (mvK hOut 0 NLogN.CR.hI (.add (.reg hTT) (.reg hTT)))
  (NLogN.RS.seq (mvK hKs 3 NLogN.CR.hI (.reg hTT))
  (NLogN.RS.seq (push 1 (.reg hTT))
  (NLogN.RS.seq (permP 2)
  (NLogN.RS.seq (pop 1)
  (dropN 3 NLogN.CR.hI (.reg hTT))))))))))))))

/-! ## The whole chirp table -/

/-- **The chirp table.**  Requests, twiddle table, sort, scan, gather. -/
def chirpP (tw : RM) : RM := .seq (.rs reqGen) (.seq (tabGen tw) (.rs chirpTail))

end NLogN.CR

end

section
-- Source module: NLogN.HCrChirpCost
/-!
# Costs of the chirp table program (namespace `NLogN.CR`)

`QB B p s`: every expression of the loop-free program `p` evaluated from `s` has all
intermediate values of at most `B` bits; then `p.cost s ≤ qsz p · 10^8 · (B + 2)^2`.
-/

set_option linter.unusedSimpArgs false
set_option linter.unusedVariables false

namespace NLogN.CR

open IntMul RS IntMul.HvdH IntMul.HvdH.Ref

def QB (B : ℕ) : RS → IState HS → Prop
  | .skip, _ => True
  | .asg i e, s => e.Bnd s B ∧ zsz (s.top i) ≤ B
  | .push _ e, s => e.Bnd s B
  | .pop i, s => zsz (s.top i) ≤ B
  | .seq a b, s => QB B a s ∧ QB B b (a.run s)
  | .ite c a b, s => c.Bnd s B ∧ (if c.eval s < 0 then QB B a s else QB B b s)
  | .rep _ _ _, _ => False

def qsz : RS → ℕ
  | .skip => 1
  | .asg _ e | .push _ e => e.size + 1
  | .pop _ => 1
  | .seq a b => qsz a + qsz b
  | .ite c a b => c.size + 1 + qsz a + qsz b
  | .rep _ _ _ => 0

/-! ## The axis step and loop -/

/-! ## One request record -/

def reqPost : RS := NLogN.RS.seq (push 0 (.lit 0)) (NLogN.RS.seq (push 0 (.add (.reg hQ) (.lit (-1))))
  (NLogN.RS.seq (push 0 (.add (.add (.reg hAc) (.reg hAc)) (.lit 1))) (NLogN.RS.seq (pop hAc) (pop hXq))))

/-- The bit bound used for all register values of the chirp programs. -/
def ChirpB (q : Params) (B : ℕ) : Prop := 4 * (q.r * q.r) + 4 * q.r + q.tdims.prod + q.K + 8 < 2 ^ B

/-! ## The twiddle table -/

/-- All table values have at most `B` bits. -/
def TwB (p R B : ℕ) : Prop := ∀ c < R, zsz (twD p R c).1 ≤ B ∧ zsz (twD p R c).2 ≤ B

/-! ## The scan -/

/-! ## Sort, scan, gather -/

/-! ## The whole chirp table -/

end NLogN.CR

end

section
-- Source module: NLogN.HCrPar6
/-!
# Parameter computation, part 6: cost of stage D and of `parP` (namespace `NLogN.CR`)
-/

set_option linter.unusedSimpArgs false
set_option linter.unusedVariables false

namespace NLogN.CR

open IntMul RS IntMul.HvdH IntMul.HvdH.Ref

def bndA (K N nb : ℕ) : ℕ := 16 * N + 4 * nb + 4 * K + 16
def bndB (K N : ℕ) : ℕ :=
  2 ^ ((mkParams K N).b + 1) + 4 * 2 ^ (mkParams K N).e * (mkParams K N).P +
    (mkParams K N).b * (mkParams K N).b + 2 * (mkParams K N).b + 1
def bndC (K N : ℕ) : ℕ :=
  2 * (facN K).factorial * facN K + 8 * K * copQ K (2 ^ (mkParams K N).e) +
    4 * K * 2 ^ (mkParams K N).e + 8 * K + 8
def bndD (K N : ℕ) : ℕ :=
  K + (mkParams K N).dh + m1V K (2 ^ (mkParams K N).e) + m2V K (2 ^ (mkParams K N).e) +
    copQ K (2 ^ (mkParams K N).e) * m1V K (2 ^ (mkParams K N).e) +
    copQ K (2 ^ (mkParams K N).e) * m2V K (2 ^ (mkParams K N).e) +
    copQ K (2 ^ (mkParams K N).e) + 3 * (copS K (mkParams K N).dh (2 ^ (mkParams K N).e)).prod + 2

/-- The bit bound for `parP`: every value computed has at most `B` bits when `parBnd K N nb < 2^B`.
With `q = mkParams K N`, `r = 2^q.e`: `parBnd = O(N + 2^e·P + (256K²)!·K²·(r+1) + q.S)`. -/
def parBnd (K N nb : ℕ) : ℕ := bndA K N nb + bndB K N + bndC K N + bndD K N

/-! ## Register contents after `parP` -/

end NLogN.CR

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 loop invariant, relative to the state `s0` before the pushes. -/
def atanInv' (x N K : ℕ) (σ : ℤ) (s0 : IState HS) (k : ℕ) (t : IState HS) : Prop :=
  t.G 40 = .int (atanSum x N σ k) :: s0.G 40 ∧
  t.G 41 = .int ((x : ℤ) ^ (2 * k + 1)) :: s0.G 41 ∧
  t.G 42 = .int k :: s0.G 42 ∧
  t.G 43 = .int (σ ^ k) :: s0.G 43 ∧
  t.G 44 = .int (2 ^ N) :: s0.G 44 ∧
  t.G 45 = .int ((K : ℤ) - k) :: s0.G 45 ∧
  ∀ j : Fin HS, j.val < 40 ∨ 45 < j.val → t.G j = s0.G j

def atanBody (x : ℕ) (σ : ℤ) : RS :=
  NLogN.RS.seqL [
      asg 40 (.add (.reg 40) (.mul (.reg 43) (.ndiv (.reg 44)
        (.mul (.add (.add (.reg 42) (.reg 42)) (.lit 1)) (.reg 41))))),
      asg 41 (.mul (.reg 41) (.lit ((x : ℤ) * x))),
      asg 42 (.add (.reg 42) (.lit 1)),
      asg 43 (.mul (.reg 43) (.lit σ))]

def atanPre : RS :=
  NLogN.RS.seqL [push 40 (.lit 0), push 41 (.lit 0), push 42 (.lit 0), push 43 (.lit 1),
    push 44 (.pow2 (.reg 0))]

def atanPost : RS := NLogN.RS.seqL [push 2 (.reg 40), pop 44, pop 43, pop 42, pop 41, pop 40]

/-- `Σ_{k<K} σ^k ⌊2^N / ((2k+1) x^{2k+1})⌋`, for `N` on register `0` and `K` on register `1`,
pushed on register `2`. -/
def atanP (x : ℕ) (σ : ℤ) : RS :=
  NLogN.RS.seq (NLogN.RS.seq atanPre (asg 41 (.lit x))) (NLogN.RS.seq (NLogN.RS.rep 45 (.reg 1) (atanBody x σ)) atanPost)

end NLogN

namespace NLogN
open IntMul RS

/-- The size parameter of `atanP`. -/
def atanB (x N K : ℕ) : ℕ := N + (2 * K + 3) * (x + 1) + 2

def atanM (x N K : ℕ) : ℕ := 2 ^ N * (2 * K + 3) * x ^ (2 * K + 3)

end NLogN

end

section
-- Source module: NLogN.HExprDer
/-!
# Derived expressions

Floor division and modulus for an arbitrary numerator and a positive divisor, built from `sel`
and `ndiv`/`nmod`.
-/

namespace NLogN

open IntMul

namespace Ex

/-- `a / b` (floor) for `b > 0`. -/
def fdiv (a b : Ex HS) : Ex HS :=
  .sel a (.neg (.ndiv (.add (.neg a) (.add b (.lit (-1)))) b)) (.ndiv a b)

end Ex

end NLogN

end

section
-- Source module: NLogN.HNumPi
/-!
# `piBits` and `log2Bits` as register programs

Input `n` on register `3`; the result is pushed on register `4`.
Uses registers `0, 1, 2, 40 … 47`.
-/

set_option linter.unusedSimpArgs false

namespace NLogN

open IntMul RS

/-- `M = 2n + 9` on registers `0` and `1`. -/
def mE : Ex HS := .add (.add (.reg 3) (.reg 3)) (.lit 9)

/-- Final rounding: `(c₁·A + c₂·B + 2^{n+8}) / 2^{n+9}` with `A`, `B` on registers `46`, `47`. -/
def roundE (c₁ c₂ : ℤ) : Ex HS :=
  Ex.fdiv (.add (.add (.mul (.lit c₁) (.reg 46)) (.mul (.lit c₂) (.reg 47)))
    (.pow2 (.add (.reg 3) (.lit 8)))) (.pow2 (.add (.reg 3) (.lit 9)))

/-- Two arctan-type series at `M = 2n+9`, combined as `(c₁·A + c₂·B + 2^{n+8}) / 2^{n+9}`. -/
def twoAtanP (x₁ x₂ : ℕ) (σ : ℤ) (c₁ c₂ : ℤ) : RS :=
  NLogN.RS.seqL [push 0 NLogN.mE, push 1 NLogN.mE,
    atanP x₁ σ, push 46 (.reg 2), pop 2,
    atanP x₂ σ, push 47 (.reg 2), pop 2,
    pop 1, pop 0, push 4 (roundE c₁ c₂), pop 47, pop 46]

def piP : RS := twoAtanP 5 239 (-1) 16 (-4)

/-- `log2Bits` uses one series; the second (with `x = 3`, coefficient `0`) is harmless. -/
def log2P : RS := twoAtanP 3 3 1 2 0

end NLogN

namespace NLogN

open IntMul RS

/-- Cost bound of `piP` / `log2P`. -/
def piCost (x₁ x₂ n : ℕ) : ℕ :=
  10 ^ 11 * (2 * n + 11) * (atanB x₁ (2 * n + 9) (2 * n + 9) + atanB x₂ (2 * n + 9) (2 * n + 9) + 6) ^ 2

end NLogN

end

section
-- Source module: NLogN.HNumExp
/-!
# `expBits` as a register program

Inputs `n` on register `5`, `x` on `6`, `y` on `7`; the two components of `expBits n x y` are
pushed on registers `8` and `9`.  Scratch registers `30 … 39`.
-/

set_option linter.unusedSimpArgs false

namespace NLogN

open IntMul RS IntMul.HvdH

/-! ## Arithmetic facts -/

/-! ## The program -/

/-- `x · 2^(n+9)` -/
def eXE : Ex HS := .mul (.reg 6) (.pow2 (.add (.reg 5) (.lit 9)))
/-- `y · 2^(n+9)` -/
def eYE : Ex HS := .mul (.reg 7) (.pow2 (.add (.reg 5) (.lit 9)))
/-- `2^(2n+9)` -/
def ePE : Ex HS := .pow2 (.add (.add (.reg 5) (.reg 5)) (.lit 9))
/-- `2n+13` -/
def eKE : Ex HS := .add (.add (.reg 5) (.reg 5)) (.lit 13)
/-- `(j+1) 2^N` -/
def dvE : Ex HS := .mul (.add (.reg 33) (.lit 1)) (.reg 32)
def nu1E : Ex HS := .add (.mul (.reg 34) (.reg 30)) (.neg (.mul (.reg 35) (.reg 31)))
def nu2E : Ex HS := .add (.mul (.reg 34) (.reg 31)) (.mul (.reg 35) (.reg 30))
/-- rounding `(S + 2^(n+8)) / 2^(n+9)` of register `i` -/
def rndE (i : Fin HS) : Ex HS :=
  Ex.fdiv (.add (.reg i) (.pow2 (.add (.reg 5) (.lit 8)))) (.pow2 (.add (.reg 5) (.lit 9)))

def expBody : RS :=
  NLogN.RS.seqL [asg 36 (.add (.reg 36) (.reg 34)), asg 37 (.add (.reg 37) (.reg 35)),
    asg 38 (Ex.fdiv nu1E dvE), asg 35 (Ex.fdiv nu2E dvE), asg 34 (.reg 38),
    asg 33 (.add (.reg 33) (.lit 1))]

def expPre : RS :=
  NLogN.RS.seqL [push 30 eXE, push 31 eYE, push 32 ePE, push 33 (.lit 0), push 34 (.reg 32),
    push 35 (.lit 0), push 36 (.lit 0), push 37 (.lit 0), push 38 (.reg 32)]

def expPost : RS :=
  NLogN.RS.seqL [push 8 (NLogN.rndE 36), push 9 (NLogN.rndE 37), pop 38, pop 37, pop 36, pop 35, pop 34, pop 33,
    pop 32, pop 31, pop 30]

def expP : RS := NLogN.RS.seq expPre (NLogN.RS.seq (NLogN.RS.rep 39 eKE expBody) expPost)

/-- The loop invariant, relative to the initial state `s0`. -/
def expInv (N K : ℕ) (X Y : ℤ) (s0 : IState HS) (k : ℕ) (t : IState HS) : Prop :=
  t.G 30 = .int X :: s0.G 30 ∧
  t.G 31 = .int Y :: s0.G 31 ∧
  t.G 32 = .int (2 ^ N) :: s0.G 32 ∧
  t.G 33 = .int k :: s0.G 33 ∧
  t.G 34 = .int (expTerm N X Y k).1 :: s0.G 34 ∧
  t.G 35 = .int (expTerm N X Y k).2 :: s0.G 35 ∧
  t.G 36 = .int (expApprox N k X Y).1 :: s0.G 36 ∧
  t.G 37 = .int (expApprox N k X Y).2 :: s0.G 37 ∧
  t.G 38 = .int (expTerm N X Y k).1 :: s0.G 38 ∧
  t.G 39 = .int ((K : ℤ) - k) :: s0.G 39 ∧
  ∀ j : Fin HS, j.val < 30 ∨ 39 < j.val → t.G j = s0.G j

/-! ## Cost -/

/-- Bit-size parameter: every value the program handles has fewer than `expB n m` bits. -/
def expB (n m : ℕ) : ℕ := (2 * n + 13) * (m + n + 11) + 2 * n + 10

/-- The cost bound of `expP`, polynomial in `n` and the bit sizes of `x`, `y`. -/
def expCost (n zx zy : ℕ) : ℕ := 10 ^ 12 * (n + zx + zy + 20) ^ 5

end NLogN

end

section
-- Source module: NLogN.HNumCall
/-!
# Tools for composing numeric register programs

* general integer division `Ex.ediv`, remainder `Ex.emod`, and `Ex.max0` (`toNat`),
* "top-based" wrappers of `piP`, `log2P`, `expP` usable as conditional `simp` lemmas,
* crude size bounds on `piBits`, `log2Bits`, `expBits`.
-/

set_option linter.unusedSimpArgs false

namespace NLogN

open IntMul RS IntMul.HvdH

namespace Ex

/-- Euclidean (Lean) integer division `a / b` for any `b` (`a / 0 = 0`). -/
def ediv (a b : Ex HS) : Ex HS :=
  .sel (.neg b) (fdiv a b) (.sel b (.neg (fdiv a (.neg b))) (.lit 0))

/-- `a % b` (Lean `Int.emod`). -/
def emod (a b : Ex HS) : Ex HS := .add a (.neg (.mul b (ediv a b)))

/-- `max 0 a`, i.e. `a.toNat`. -/
def max0 (a : Ex HS) : Ex HS := .sel a (.lit 0) a

end Ex

/-! ## Top-based wrappers -/

/-! ## Size bounds of the constants -/

end NLogN

end

section
-- Source module: NLogN.HNumExp2
/-!
# `expNegFix`, `expNeg`, `expPos` as register programs

`expNegFixP`: inputs `n` on `10`, `M` on `11`, `X` on `12`; result pushed on `13`.
Uses (with pushes popped) registers `3 … 9` and the scratch registers of `log2P`
(`0, 1, 2, 40 … 47`) and `expP` (`30 … 39`).
-/

set_option linter.unusedSimpArgs false

namespace NLogN

open IntMul RS IntMul.HvdH

/-- `τ = (-X) / L` with `X` on `12`, `L` on `4`. -/
def tauE : Ex HS := Ex.ediv (.neg (.reg 12)) (.reg 4)

/-- `s = (M - n - 1).toNat + τ.toNat`, with `τ` on `14`. -/
def sE : Ex HS := .add (Ex.max0 (.add (.reg 11) (.neg (.add (.reg 10) (.lit 1))))) (Ex.max0 (.reg 14))

def expNegFixP : RS :=
  NLogN.RS.seqL [push 3 (.reg 11), log2P, push 14 tauE, push 15 NLogN.sE, push 5 (.reg 11),
    push 6 (.add (.reg 12) (.mul (.reg 14) (.reg 4))), push 7 (.lit 0), expP,
    push 13 (Ex.fdiv (.add (.reg 8) (.pow2 (.reg 15))) (.pow2 (.add (.reg 15) (.lit 1)))),
    pop 9, pop 8, pop 7, pop 6, pop 5, pop 15, pop 14, pop 4, pop 3]

/-- The size parameter of `expNegFixP`. -/
def expNegB (n M zX T : ℕ) : ℕ := expB M (3 * M + 14) + n + M + T + zX + 10

/-- Cost bound of `expNegFixP`, in terms of `n`, `M`, `zsz X` and `τ = ((-X) / log2Bits M).toNat`. -/
def expNegFixCost (n M zX T : ℕ) : ℕ := 10 ^ 20 * (n + M + zX + T + 20) ^ 5

/-! Top-based wrappers. -/

end NLogN

end

section
-- Source module: NLogN.HNumExp3
/-!
# `expNeg` as a register program

`expNegP`: inputs `n` on `16`, `k` on `17`, `j` on `18`; `expNeg n k j` pushed on `19`.
It uses (with pushes popped) registers `3 … 15` and the scratch registers of
`piP`/`log2P` (`0, 1, 2, 40 … 47`) and `expP` (`30 … 39`).
-/

set_option linter.unusedSimpArgs false

namespace NLogN

open IntMul RS IntMul.HvdH

/-- `2n + 10` with `n` on `16`. -/
def m10E : Ex HS := .add (.add (.reg 16) (.reg 16)) (.lit 10)

/-- `-(P j / k)` with `P` on `4`, `j` on `18`, `k` on `17`. -/
def xNegE : Ex HS := .neg (Ex.ediv (.mul (.reg 4) (.reg 18)) (.reg 17))

def expNegBody : RS :=
  NLogN.RS.seqL [push 3 m10E, piP, push 10 (.reg 16), push 11 m10E, push 12 xNegE, expNegFixP,
    push 19 (.reg 13), pop 13, pop 12, pop 11, pop 10, pop 4, pop 3]

def expNegP : RS :=
  NLogN.RS.ite (.add (.mul (.reg 17) (.reg 16)) (.neg (.reg 18))) (push 19 (.lit 0)) expNegBody

/-- The size parameter of `expNegBody`. -/
def expNegBB (n zX T zj zk : ℕ) : ℕ := expNegB n (2 * n + 10) zX T + zj + zk + 12 * n + 100

/-- Cost bound of `expNegP`, in terms of `n` and the bit sizes of `k`, `j`. -/
def expNegCost (n zk zj : ℕ) : ℕ := 10 ^ 35 * (n + zk + zj + 30) ^ 10

end NLogN

end

section
-- Source module: NLogN.HNumExpPos
/-!
# `expPos` as a register program

`expPosP`: inputs `n` on `16`, `k` on `17`, `j` on `18`, `σ` on `27`; `expPos n k j σ` pushed
on `19`. It uses (with pushes popped) registers `3 … 15` and the scratch registers of
`piP`/`log2P` (`0, 1, 2, 40 … 47`) and `expP` (`30 … 39`).
-/

set_option linter.unusedSimpArgs false

namespace NLogN

open IntMul RS IntMul.HvdH

/-- `min 0 (P j / k - σ L)` with `P` on `14`, `j` on `18`, `k` on `17`, `σ` on `27`, `L` on `4`. -/
def xPosE : Ex HS :=
  .neg (Ex.max0 (.neg (.add (Ex.ediv (.mul (.reg 14) (.reg 18)) (.reg 17))
    (.neg (.mul (.reg 27) (.reg 4))))))

def expPosP : RS :=
  NLogN.RS.seqL [push 3 m10E, piP, push 14 (.reg 4), log2P, push 12 xPosE, pop 4, pop 14, pop 4,
    push 10 (.reg 16), push 11 m10E, expNegFixP,
    push 19 (.reg 13), pop 13, pop 12, pop 11, pop 10, pop 3]

/-- The size parameter of `expPosP`. -/
def expPosBB (n zX T zj zk zσ : ℕ) : ℕ :=
  expNegB n (2 * n + 10) zX T + 2 * zj + zk + 2 * zσ + 30 * n + 200

/-- Cost bound of `expPosP`, in terms of `n`, the bit sizes of `k`, `j`, and `σ`. -/
def expPosCost (n zk zj σ : ℕ) : ℕ := 10 ^ 35 * (n + zk + zj + σ + 30) ^ 10

end NLogN

end

section
-- Source module: NLogN.HNumTw
/-!
# `octExp` and `twiddle` as register programs

* `octExpP`: inputs `n` on `48`, `k` on `49`, `a` on `50`; pushes `(octExp n k a).1` on `51`
  and `(octExp n k a).2` on `52`. Uses (pushed and popped) registers `3 … 9` and the scratch
  registers of `piP` (`0, 1, 2, 40 … 47`) and `expP` (`30 … 39`).
* `twiddleP`: inputs `n` on `28`, `k` on `29`, `j` on `53`; pushes `(twiddle n k j).1` on `51`
  and `(twiddle n k j).2` on `52`. Uses in addition registers `10 … 14` and `48 … 50`.
-/

set_option linter.unusedSimpArgs false

namespace NLogN

open IntMul RS IntMul.HvdH

/-- `n + 4` with `n` on `48`. -/
def n4E : Ex HS := .add (.reg 48) (.lit 4)

/-- `P a / (4 k)` with `P` on `4`, `a` on `50`, `k` on `49`. -/
def phiE : Ex HS := Ex.ediv (.mul (.reg 4) (.reg 50)) (.mul (.lit 4) (.reg 49))

/-- `(u + 2^3) / 2^4` with `u` on `r`. -/
def rnd4E (r : Fin HS) : Ex HS := Ex.fdiv (.add (.reg r) (.pow2 (.lit 3))) (.pow2 (.lit 4))

def octExpP : RS :=
  NLogN.RS.seqL [push 3 n4E, piP, push 5 n4E, push 6 (.lit 0), push 7 phiE, expP,
    push 51 (rnd4E 8), push 52 (rnd4E 9), pop 9, pop 8, pop 7, pop 6, pop 5, pop 4, pop 3]

/-- The size parameter of `octExpP`. -/
def octBB (n za zk : ℕ) : ℕ := expB (n + 4) (3 * (n + 4) + 14 + za) + 2 * za + 2 * zk + 20 * n + 100

/-- Cost bound of `octExpP`, in terms of `n` and the bit sizes of `a`, `k`. -/
def octExpCost (n za zk : ℕ) : ℕ := 10 ^ 20 * (n + za + zk + 30) ^ 5

end NLogN

end

section
-- Source module: NLogN.HNumTw2
/-!
# `twiddle` as a register program

`twiddleP`: inputs `n` on `28`, `k` on `29`, `j` on `53`; pushes `(twiddle n k j).1` on `51`
and `(twiddle n k j).2` on `52`. Uses (pushed and popped) registers `10 … 14`, `48 … 50` and
those of `octExpP` (`3 … 9`, `0, 1, 2, 30 … 47`).
-/

set_option linter.unusedSimpArgs false

namespace NLogN

open IntMul RS IntMul.HvdH

/-- Multiplication by `i^t` for `t ∈ {0, 1, 2, 3}` (and `t ≥ 3` treated as `3`). -/
def rot4 (t : ℤ) (p : ℤ × ℤ) : ℤ × ℤ :=
  if t < 1 then p else if t < 2 then (-p.2, p.1) else if t < 3 then (-p.1, -p.2) else (p.2, -p.1)

/-! ## Top-based wrappers for `octExpP` -/

/-! ## The program -/

/-- `j mod k` with `j` on `53`, `k` on `29`. -/
def jkE : Ex HS := Ex.emod (.reg 53) (.reg 29)
/-- `q = 8 (j mod k) / k`. -/
def qE : Ex HS := Ex.ediv (.mul (.lit 8) jkE) (.reg 29)
/-- `r = 8 (j mod k) mod k`. -/
def rE : Ex HS := Ex.emod (.mul (.lit 8) jkE) (.reg 29)
/-- `q mod 2` with `q` on `10`. -/
def cE : Ex HS := Ex.emod (.reg 10) (.lit 2)
/-- `((q + 1) / 2) mod 4` with `q` on `10`. -/
def tE : Ex HS := Ex.emod (Ex.ediv (.add (.reg 10) (.lit 1)) (.lit 2)) (.lit 4)
/-- The octant argument: `r` if `c = 0`, else `k - r` (`r` on `11`, `c` on `12`). -/
def aE : Ex HS := .sel (.add (.reg 12) (.lit (-1))) (Ex.max0 (.reg 11))
  (Ex.max0 (.add (.reg 29) (.neg (.reg 11))))
/-- The imaginary part, conjugated if `c = 1` (`c` on `12`, value on `52`). -/
def yE : Ex HS := .sel (.add (.reg 12) (.lit (-1))) (.reg 52) (.neg (.reg 52))
/-- Real part of `i^t (x, y)` with `t` on `13`, `x` on `51`, `y` on `14`. -/
def reE : Ex HS := .sel (.add (.reg 13) (.lit (-1))) (.reg 51)
  (.sel (.add (.reg 13) (.lit (-2))) (.neg (.reg 14))
    (.sel (.add (.reg 13) (.lit (-3))) (.neg (.reg 51)) (.reg 14)))
/-- Imaginary part of `i^t (x, y)`. -/
def imE : Ex HS := .sel (.add (.reg 13) (.lit (-1))) (.reg 14)
  (.sel (.add (.reg 13) (.lit (-2))) (.reg 51)
    (.sel (.add (.reg 13) (.lit (-3))) (.neg (.reg 14)) (.neg (.reg 51))))

/-- Set-up: `q, r, c, t` on `10 … 13` and the arguments of `octExpP` on `48 … 50`. -/
def twA : RS :=
  NLogN.RS.seqL [push 10 qE, push 11 NLogN.rE, push 12 NLogN.cE, push 13 NLogN.tE, push 48 (.reg 28), push 49 (.reg 29),
    push 50 NLogN.aE]

/-- Conjugation/rotation of the result of `octExpP`, then clean-up. -/
def twB : RS :=
  NLogN.RS.seqL [push 14 yE, push 50 NLogN.reE, push 49 NLogN.imE, pop 51, pop 52,
    push 51 (.reg 50), push 52 (.reg 49), pop 49, pop 50,
    pop 14, pop 50, pop 49, pop 48, pop 13, pop 12, pop 11, pop 10]

def twiddleP : RS := NLogN.RS.seqL [twA, octExpP, twB]

/-! ## Cost -/

/-- The size parameter of `twiddleP`. -/
def twBB (n zj zk : ℕ) : ℕ :=
  2 * expB (n + 4) (3 * (n + 4) + 14 + zj + zk + 4) + 4 * zj + 4 * zk + 20 * n + 200

/-- The state after `twA`. -/
def twAState (n k j : ℕ) (s : IState HS) : IState HS :=
  ((((((s.upd 10 (.int (8 * ((j : ℤ) % k) / k) :: s.G 10)).upd 11 (.int (8 * ((j : ℤ) % k) % k) :: s.G 11)).upd 12
    (.int (8 * ((j : ℤ) % k) / k % 2) :: s.G 12)).upd 13 (.int ((8 * ((j : ℤ) % k) / k + 1) / 2 % 4) :: s.G 13)).upd 48
    (.int n :: s.G 48)).upd 49 (.int k :: s.G 49)).upd 50 (.int (((if 8 * ((j : ℤ) % k) / k % 2 < 1 then (8 * ((j : ℤ) % k) % k).toNat else ((k : ℤ) - 8 * ((j : ℤ) % k) % k).toNat) : ℕ) : ℤ) :: s.G 50)

/-- Cost bound of `twiddleP`, in terms of `n` and the bit sizes of `j`, `k`. -/
def twiddleCost (n zj zk : ℕ) : ℕ := 10 ^ 22 * (n + zj + zk + 30) ^ 5

end NLogN

end

section
-- Source module: NLogN.HNumScale
/-!
# `ceilSqrt`, `scaleZ` and `twD` as register programs

* `isqrtP`: inputs `v` on `41`, a bit bound `B` on `42` (`0 ≤ v < 4^B`); pushes `⌊√v⌋` on `43`.
  Binary search with the loop counter on `44`.
* `ceilSqrtP`: same inputs; pushes `ceilSqrt v` on `43`.
-/

set_option linter.unusedSimpArgs false

namespace NLogN

open IntMul RS IntMul.HvdH

/-! ## Binary-search square root -/

/-- The value of the search variable after `k` of the `B` steps. -/
def isq (v : ℤ) (B : ℕ) : ℕ → ℤ
  | 0 => 0
  | k + 1 => if v - (isq v B k + 2 ^ (B - k - 1)) ^ 2 < 0 then isq v B k
      else isq v B k + 2 ^ (B - k - 1)

/-- `r + 2^(c-1)` with `r` on `43`, `c` on `44`. -/
def isqHE : Ex HS := .add (.reg 43) (.pow2 (.add (.reg 44) (.lit (-1))))
/-- One binary-search step. -/
def isqBE : Ex HS := .sel (.add (.reg 41) (.neg (.mul isqHE isqHE))) (.reg 43) isqHE

def isqrtP : RS := NLogN.RS.seq (push 43 (.lit 0)) (NLogN.RS.rep 44 (.reg 42) (asg 43 isqBE))

/-- The loop state after `k` steps. -/
def isqSt (s : IState HS) (v : ℤ) (B k : ℕ) : IState HS :=
  (s.upd 43 (.int (isq v B k) :: s.G 43)).upd 44 (.int ((B : ℤ) - k) :: s.G 44)

/-- Cost bound of `isqrtP`. -/
def isqrtCost (zv B : ℕ) : ℕ := 10 ^ 10 * (B + 1) * (zv + 2 * B + 12) ^ 2

/-- `r + 1` if `r² < v`, else `r` (`r` on `43`, `v` on `41`). -/
def csE : Ex HS := .sel (.add (.mul (.reg 43) (.reg 43)) (.neg (.reg 41))) (.add (.reg 43) (.lit 1))
  (.reg 43)

def ceilSqrtP : RS := NLogN.RS.seq isqrtP (asg 43 NLogN.csE)

/-- Cost bound of `ceilSqrtP`. -/
def ceilSqrtCost (zv B : ℕ) : ℕ := 10 ^ 11 * (B + 1) * (zv + 2 * B + 12) ^ 2

/-! ## `scaleZ` -/

/-- Truncated division `a.tdiv m`, valid for `m > 0`. -/
def tdivE (a m : Ex HS) : Ex HS := .sel a (.neg (Ex.fdiv (.neg a) m)) (Ex.fdiv a m)

/-- `2^k x` with `k` on `11`, `x` on `12`. -/
def w1E : Ex HS := .mul (.pow2 (.reg 11)) (.reg 12)
/-- `2^k y` with `k` on `11`, `y` on `13`. -/
def w2E : Ex HS := .mul (.pow2 (.reg 11)) (.reg 13)
/-- `w₁² + w₂²` with `w₁, w₂` on `17, 18`. -/
def vE : Ex HS := .add (.mul (.reg 17) (.reg 17)) (.mul (.reg 18) (.reg 18))
/-- `4^p - v` with `p` on `10`, `v` on `19`. -/
def scCondE : Ex HS := .add (.pow2 (.add (.reg 10) (.reg 10))) (.neg (.reg 19))
/-- The bit bound `k + B_z + 1` for the square root. -/
def scBE : Ex HS := .add (.add (.reg 11) (.reg 14)) (.lit 1)

def scaleThen : RS :=
  NLogN.RS.seqL [push 41 (.reg 19), push 42 scBE, ceilSqrtP,
    push 15 (tdivE (.mul (.reg 17) (.pow2 (.reg 10))) (.reg 43)),
    push 16 (tdivE (.mul (.reg 18) (.pow2 (.reg 10))) (.reg 43)), pop 43, pop 42, pop 41]

def scaleElse : RS := NLogN.RS.seqL [push 15 (.reg 17), push 16 (.reg 18)]

def scaleZP : RS :=
  NLogN.RS.seqL [push 17 w1E, push 18 w2E, push 19 vE, NLogN.RS.ite scCondE scaleThen scaleElse, pop 19, pop 18,
    pop 17]

/-- The state after the first three pushes of `scaleZP`. -/
def scSt (s : IState HS) (k : ℕ) (x y : ℤ) : IState HS :=
  ((s.upd 17 (.int (2 ^ k * x) :: s.G 17)).upd 18 (.int (2 ^ k * y) :: s.G 18)).upd 19
    (.int (2 ^ k * x * (2 ^ k * x) + 2 ^ k * y * (2 ^ k * y)) :: s.G 19)

/-- The state of `scaleThen` after `ceilSqrtP`. -/
def scSt2 (s : IState HS) (k Bz : ℕ) (x y : ℤ) : IState HS :=
  (((NLogN.scSt s k x y).upd 41 (.int (2 ^ k * x * (2 ^ k * x) + 2 ^ k * y * (2 ^ k * y)) ::
    (NLogN.scSt s k x y).G 41)).upd 42 (.int ((k : ℤ) + Bz + 1) :: (NLogN.scSt s k x y).G 42)).upd 43
    (.int (ceilSqrt (2 ^ k * x * (2 ^ k * x) + 2 ^ k * y * (2 ^ k * y)).toNat) ::
      (NLogN.scSt s k x y).G 43)

/-- Cost bound of `scaleZP`. -/
def scaleCost (p k Bz : ℕ) : ℕ := 10 ^ 14 * (p + k + Bz + 20) ^ 3

/-! ## `twD` -/

end NLogN

end

section
-- Source module: NLogN.HNumTwD
/-!
# `twD` as a register program

`twDP`: inputs `p` on `28`, `n` on `29`, `j` on `53`; pushes `(twD p n j).1` on `15` and
`(twD p n j).2` on `16`. It runs `twiddleP` and then `scaleZP` with `k = 0` and the bit bound
`twBz p` (an upper bound for the size of the twiddle factors). Uses (pushed and popped)
registers `10 … 14`, `17 … 19`, `41 … 44`, `51`, `52` and those of `twiddleP`.
-/

set_option linter.unusedSimpArgs false

namespace NLogN

open IntMul RS IntMul.HvdH

/-- The bit bound used for the twiddle factors. -/
def twBz (p : ℕ) : ℕ := expB (p + 4) (3 * (p + 4) + 14) + 7

/-- `twBz p` with `p` on `28`. -/
def bzE : Ex HS := .add (.mul (.add (.add (.reg 28) (.reg 28)) (.lit 21))
  (.add (.mul (.lit 4) (.reg 28)) (.lit 41))) (.add (.add (.reg 28) (.reg 28)) (.lit 25))

def twDP : RS :=
  NLogN.RS.seqL [twiddleP, push 10 (.reg 28), push 11 (.lit 0), push 12 (.reg 51), push 13 (.reg 52),
    push 14 bzE, scaleZP, pop 14, pop 13, pop 12, pop 11, pop 10, pop 52, pop 51]

/-- The state after `twiddleP`. -/
def twT1 (p n j : ℕ) (s : IState HS) : IState HS :=
  (s.upd 51 (.int (twiddle p n j).1 :: s.G 51)).upd 52 (.int (twiddle p n j).2 :: s.G 52)

/-- The state before `scaleZP`. -/
def twT2 (p n j : ℕ) (s : IState HS) : IState HS :=
  (push 14 bzE).run ((push 13 (.reg 52)).run ((push 12 (.reg 51)).run ((push 11 (.lit 0)).run
    ((push 10 (.reg 28)).run (twT1 p n j s)))))

/-- Cost bound of `twDP`, in terms of `p` and the bit sizes of `n`, `j`. -/
def twDCost (p zn zj : ℕ) : ℕ := 10 ^ 25 * (p + zn + zj + 30) ^ 6

end NLogN

end

section
-- Source module: NLogN.HCrTw
/-!
# The twiddle program for the chirp table (namespace `NLogN.CR`)

`twR` instantiates the abstract hypothesis `TwSpec` of `chirpP_spec`/`chirpP_cost` with the
numerics program `twDP` (HNumTwD.lean): a wrapper copies `p` (`hPp`), `R` (`hM2`), `c` (`hC`)
onto `28, 29, 53`, runs `twDP` and restores; the whole is relocated by the swap `10 ↔ 54` so
that it does not write the loop counter `hQ = 10`.
-/

set_option linter.unusedSimpArgs false
set_option linter.unusedVariables false

namespace NLogN.CR

open IntMul RS IntMul.HvdH IntMul.HvdH.Ref

/-- Copy `p, R, c` to the inputs of `twDP`, run it, restore. -/
def twW : RS :=
  seqL [push 28 (.reg hPp), push 29 (.reg hM2), push 53 (.reg hC), twDP, pop 53, pop 29, pop 28]

def twσ : Equiv.Perm (Fin HS) := Equiv.swap 10 54

/-- **The twiddle program** used for the chirp table. -/
def twR : RM := RM.ren twσ (.rs twW)

/-- The cost constant of `twR` for `p` on `hPp` and `R` on `hM2`. -/
def twCost (p R : ℕ) : ℕ :=
  9 * 10 ^ 7 * (zsz (p : ℤ) + zsz (R : ℤ) + 3) ^ 2 + twDCost p (zsz (R : ℤ)) (zsz (R : ℤ))

/-! ## The chirp table with the concrete twiddle program -/

end NLogN.CR

end

section
-- Source module: NLogN.HGl1
/-!
# Generic element-stream loops (Gaussian list primitives, part 1)

`RM.rep c n b` where the body `b` pops one element from each of the input stacks `A1`, `A2`
and pushes the (reversed) encodings of two results onto `B1`, `B2`.  After `N` rounds the
outputs are the reversed flattened encodings of the `zipWith` lists.  The wrappers `loop11`,
`loop21`, `loop12` are the one/two-input, one/two-output special cases.

Outputs are pushed in reversed item order, so `B = (enc out).reverse ++ old`; one `mvList` of
all items restores the natural order (`mvList_run` reverses).
-/

set_option linter.unusedSimpArgs false

namespace NLogN.GL

open IntMul NLogN

section Generic

variable {α1 α2 β1 β2 : Type*}

/-- The state after `k` rounds of the generic stream loop. -/
def lSt (A1 A2 B1 B2 : Fin HS) (e1 : α1 → List Item) (e2 : α2 → List Item)
    (d1 : β1 → List Item) (d2 : β2 → List Item) (f1 : α1 → α2 → β1) (f2 : α1 → α2 → β2)
    (s : IState HS) (L1 : List α1) (L2 : List α2) (R1 R2 : List Item) (k : ℕ) : IState HS :=
  (((s.upd A1 ((L1.drop k).flatMap e1 ++ R1)).upd A2 ((L2.drop k).flatMap e2 ++ R2)).upd B1
    (((List.zipWith f1 (L1.take k) (L2.take k)).flatMap d1).reverse ++ s.G B1)).upd B2
    (((List.zipWith f2 (L1.take k) (L2.take k)).flatMap d2).reverse ++ s.G B2)

/-- Body hypothesis of the generic loop. -/
def BodyOK (M : MulSpec) (b : RM) (A1 A2 B1 B2 c : Fin HS) (e1 : α1 → List Item)
    (e2 : α2 → List Item) (d1 : β1 → List Item) (d2 : β2 → List Item) (f1 : α1 → α2 → β1)
    (f2 : α1 → α2 → β2) (s : IState HS) (L1 : List α1) (L2 : List α2) (C : ℕ) : Prop :=
  ∀ x ∈ L1, ∀ y ∈ L2, ∀ (t : IState HS) (r1 r2 o1 o2 : List Item),
    (∀ j, j ∉ [A1, A2, B1, B2, c] → t.G j = s.G j) → t.G A1 = e1 x ++ r1 →
    t.G A2 = e2 y ++ r2 → t.G B1 = o1 ++ s.G B1 → t.G B2 = o2 ++ s.G B2 →
    b.run t = (((t.upd A1 r1).upd A2 r2).upd B1 ((d1 (f1 x y)).reverse ++ t.G B1)).upd B2
      ((d2 (f2 x y)).reverse ++ t.G B2) ∧ b.ok M t ∧ b.cost M t ≤ C

end Generic

/-! ## Special cases -/

section Wrappers

variable {α α2 β β2 : Type*}

end Wrappers

end NLogN.GL

end
Source
D. Harvey, J. van der Hoeven, Integer multiplication in time O(n log n), Ann. of Math. 193 (2021), 563–617 (Theorem 1.1: multitape Turing machine multiplication in O(n log n)). Definitions of the Lean development NLogN, written for the prove2.me integer-multiplication campaign.

View graph