Harvey–van der Hoeven machine development, part 3: transforms, top level, encoding and transposition

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

import Definitions.Def_IntMul_HvdH_M_Defs2

set_option autoImplicit false
set_option relaxedAutoImplicit false
set_option maxHeartbeats 1000000

section
-- Source module: NLogN.SSDif
/-!
# The forward FFT (decimation in frequency) over `ℤ/(2^n+1)`, as a recursive procedure
-/

namespace NLogN

open IntMul

abbrev DX : Fin NS := 8
abbrev DL : Fin NS := 9
abbrev DE : Fin NS := 10
abbrev U : Fin NS := 11
abbrev V : Fin NS := 12
abbrev P1 : Fin NS := 13
abbrev P2 : Fin NS := 14
abbrev CC : Fin NS := 15
abbrev W : Fin NS := 16
abbrev HH : Fin NS := 17

/-! ## Reference functions -/

/-- Sums of the butterflies. -/
def bflyT (F : ℤ) : List ℤ → List ℤ → List ℤ
  | u :: A, v :: B => ((v + u) % F) :: bflyT F A B
  | _, _ => []

/-- Twisted differences of the butterflies, starting at index `j`. -/
def bflyB (F : ℤ) (E j : ℕ) : List ℤ → List ℤ → List ℤ
  | u :: A, v :: B => (((-v + u) * 2 ^ (j * E)) % F) :: bflyB F E (j + 1) A B
  | _, _ => []

/-- Entries are canonical residues mod `F`. -/
def Canon (F : ℤ) (A : List ℤ) : Prop := ∀ a ∈ A, 0 ≤ a ∧ a < F

/-- The forward transform: input of length `2^j`, output in "bit-reversed" order. -/
def difZ (F : ℤ) : ℕ → ℕ → List ℤ → List ℤ
  | _, 0, A => A
  | E, j + 1, A => difZ F (2 * E) j (bflyT F (A.take (2 ^ j)) (A.drop (2 ^ j))) ++
      difZ F (2 * E) j (bflyB F E 0 (A.take (2 ^ j)) (A.drop (2 ^ j)))

/-! ## The butterfly step -/

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

def bflyStep : IStmt NS P :=
  .seq (.seq (.op (.dup DX T4)) (.seq (.op (.dup NLogN.V P1)) (.op (.add T4 P1 P1)))) (.seq (fred P1)
    (.seq (.seq (.op (.mv DX T4)) (.seq (.op (.neg T4)) (.seq (.op (.mv NLogN.V P2))
      (.seq (.op (.add T4 P2 P2)) (.op (.shl P2 W)))))) (.seq (fred P2)
        (.seq (.op (.shl W NLogN.DE)) (decC CC)))))

end NLogN

namespace NLogN

open IntMul

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

/-- Move `len` integers from `X` to `Y`, where `len` is on top of `HH`. -/
def mvH (X Y : Fin NS) : IStmt NS P := .seq (.op (.dup HH CC)) (.seq (mvN X Y CC) (.op (.drop CC)))

def difSegA : IStmt NS P :=
  .seq (.op (.dup DL HH)) (.seq (.op (.const T4 2)) (.seq (.op (.split HH T4 T3))
    (.seq (.op (.drop T3)) (.op (.drop T4)))))

def difSegD : IStmt NS P :=
  .seq (.op (.const W 1)) (.seq (.op (.dup HH CC)) (.seq (.loop CC (fun c => c != .zero) bflyStep)
    (.seq (.op (.drop CC)) (.op (.drop W)))))

def difSegE : IStmt NS P :=
  .seq (.op (.dup HH DL)) (.seq (.op (.dup NLogN.DE T4)) (.seq (.op (.shl T4 NLogN.DE)) (.op (.mv T4 NLogN.DE))))

def difSegH : IStmt NS P := .seq (.op (.drop NLogN.DE)) (.seq (.op (.drop DL)) (.op (.drop HH)))

def difMain (pD : Fin P) : IStmt NS P :=
  .seq difSegA (.seq (mvH DX U) (.seq (mvH U NLogN.V) (.seq difSegD (.seq difSegE
    (.seq (mvH P2 DX) (.seq (.call pD) (.seq (mvH P1 DX) (.seq (.call pD) difSegH))))))))

def difTest : IStmt NS P := .seq (.op (.dup DL CC)) (.seq (.op (.const T4 (-1))) (.op (.add CC T4 CC)))

def difBody (pD : Fin P) : IStmt NS P :=
  .seq difTest (.ite CC (fun c => c == .zero) (.op (.drop CC)) (.seq (.op (.drop CC)) (difMain pD)))

def difCost (n j : ℕ) : ℕ := 400000 * (n + 3) * (j + 1) * 2 ^ j

end NLogN

end

section
-- Source module: NLogN.SSDit
/-!
# The inverse FFT (decimation in time) over `ℤ/(2^n+1)`, as a recursive procedure
-/

namespace NLogN

open IntMul

/-! ## Reference functions -/

/-- `t_i - 2^(-iE) b_i`, written with `2^(n-(i+1)E)·2^E = 2^(n-iE)`. -/
def dcomX (F : ℤ) (n E i : ℕ) : List ℤ → List ℤ → List ℤ
  | t :: T, b :: B => ((t + -((b * 2 ^ (n - i * E)) % F)) % F) :: dcomX F n E (i + 1) T B
  | _, _ => []

def dcomY (F : ℤ) (n E i : ℕ) : List ℤ → List ℤ → List ℤ
  | t :: T, b :: B => ((t + (b * 2 ^ (n - i * E)) % F) % F) :: dcomY F n E (i + 1) T B
  | _, _ => []

/-- The inverse transform (input in the output order of `difZ`). -/
def ditZ (F : ℤ) (n : ℕ) : ℕ → ℕ → List ℤ → List ℤ
  | _, 0, B => B
  | E, j + 1, B =>
    dcomX F n E 0 (ditZ F n (2 * E) j (B.take (2 ^ j))) (ditZ F n (2 * E) j (B.drop (2 ^ j))) ++
      dcomY F n E 0 (ditZ F n (2 * E) j (B.take (2 ^ j))) (ditZ F n (2 * E) j (B.drop (2 ^ j)))

/-! ## The butterfly step -/

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

def ditStep : IStmt NS P :=
  .seq (.seq (.op (.mv DX T4)) (.op (.shl T4 W))) (.seq (fred T4)
    (.seq (.seq (.op (.dup T4 P1)) (.seq (.op (.neg P1)) (.seq (.op (.dup NLogN.V T3))
      (.op (.add T3 P1 P1))))) (.seq (fred P1)
    (.seq (.seq (.op (.mv T4 P2)) (.seq (.op (.mv NLogN.V T3)) (.op (.add T3 P2 P2)))) (.seq (fred P2)
      (.seq (.op (.split W NLogN.DE T3)) (.seq (.op (.drop T3)) (decC CC))))))))

def ditSegD : IStmt NS P :=
  .seq (.op (.dup RN W)) (.seq (.op (.dup HH CC)) (.seq (.loop CC (fun c => c != .zero) ditStep)
    (.seq (.op (.drop CC)) (.op (.drop W)))))

def ditSegF : IStmt NS P := .seq (.op (.drop NLogN.DE)) (.op (.drop DL))

def ditMain (pT : Fin P) : IStmt NS P :=
  .seq difSegA (.seq difSegE (.seq (.call pT) (.seq (mvH DX U) (.seq (.call pT) (.seq ditSegF
    (.seq (mvH U NLogN.V) (.seq ditSegD (.seq (mvH P2 DX) (.seq (mvH P1 DX) (.op (.drop HH)))))))))))

def ditBody (pT : Fin P) : IStmt NS P :=
  .seq difTest (.ite CC (fun c => c == .zero) (.op (.drop CC)) (.seq (.op (.drop CC)) (ditMain pT)))

def ditCost (n j : ℕ) : ℕ := 500000 * (n + 3) * (j + 1) * 2 ^ j

end NLogN

end

section
-- Source module: NLogN.SSUtil
/-!
# Helper procedures: bit length and powers of two
-/

namespace NLogN

open IntMul

abbrev ZZ : Fin NS := 26

abbrev T5 : Fin NS := 29
abbrev T6 : Fin NS := 30
abbrev T7 : Fin NS := 31

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

/-! ## Bit length -/

/-- Push the bit length of the (nonnegative) top of `X` onto `C`. -/
def blenS (X C : Fin NS) : IStmt NS P :=
  .seq (.seq (.op (.dup X T5)) (.seq (.op (.const C 0)) (.op (.const T7 2))))
    (.seq (.loop T5 (fun c => c != .zero) (.seq (.op (.split T5 T7 T6)) (.seq (.op (.drop T6))
      (.seq (.op (.const T6 1)) (.op (.add T6 C C))))))
    (.seq (.op (.drop T5)) (.op (.drop T7))))

end NLogN

end

section
-- Source module: NLogN.HFFT
/-!
# The Schönhage–Strassen FFT procedures on a larger machine

The recursive procedures `difBody` / `ditBody` (over the `NS = 32` stacks of the
Schönhage–Strassen machine) are renamed along an injective stack map `f : Fin NS → Fin S'` and
keep their correctness theorems `dif_runs` / `dit_runs`: stacks outside the image of `f` are
untouched.

For the HvdH machine (`HS = 64` stacks, stacks `55..63` reserved) we use `ssH`, which maps the SS
stack `i` to the HS stack `i` (so the SS stacks occupy `0..31`).

We also check that the reference functions of `Def_IntMul_HvdH_MachineRef` agree with the
functions of the SS machine.
-/

namespace NLogN

open IntMul

/-! ## The stack map into the HvdH machine -/

/-- SS stack `i` is HS stack `i` (`HS = 64`). -/
def ssH : Fin NS → Fin 64 := Fin.castLE (by norm_num)

@[simp] theorem ssH_val (i : Fin NS) : (ssH i).val = i.val := rfl

/-! ## Renamed transforms -/

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

/-! ## The HvdH machine (`HS = 64` stacks) -/

end NLogN

/-! ## The reference functions agree with the SS machine's -/

namespace IntMul.HvdH.Ref

end IntMul.HvdH.Ref

end

section
-- Source module: NLogN.HTopBlen
/-!
# Bit length as a procedure call

Procedure `4` of the HvdH machine computes the bit length: SS `blenS T1 ZZ` relocated by `ssH`.
In `RM` it is the call `blenRM`: the top of `ssH T1` (a nonnegative integer `x`) gives
`zsz x` pushed onto `ssH ZZ`.
-/

namespace NLogN.Top

open IntMul NLogN

/-- The body of the bit-length procedure (procedure `4` of a 5-procedure program). -/
def blenProc : IStmt HS 5 := (blenS T1 ZZ : IStmt NS 1).ren ssH (fun _ => (4 : Fin 5))

abbrev BX : Fin HS := ssH T1
abbrev BC : Fin HS := ssH ZZ

def blenCost (x : ℤ) : ℕ := 2000 * (zsz x + 1) * (zsz (zsz x : ℤ) + 2) + 2

/-- The `RM` call of the bit-length procedure (procedure number `3 + 1`). -/
def blenRM : RM :=
  .prc 3 (fun s => s.upd NLogN.Top.BC (.int (zsz (s.top NLogN.Top.BX)) :: s.G NLogN.Top.BC)) (fun s => blenCost (s.top NLogN.Top.BX))
    (fun s => ∃ (x : ℕ) (LX : List Item), s.G NLogN.Top.BX = .int x :: LX) [NLogN.Top.BC]
    (fun s j hj => by simp at hj; simp [IState.upd_G, hj])
    (fun s => by simp)

end NLogN.Top

end

section
-- Source module: NLogN.HTop
/-!
# The multiplication procedure of the HvdH machine

Procedure `1` multiplies the two integers on top of `ARGM = 55` (any signs).  It takes absolute
values, computes `N = max(bitlen |a|, bitlen |b|)`, multiplies directly (schoolbook expression)
when `N < n₀`, and otherwise runs one step `st` of the recursive algorithm (an `RM` program that
may call the multiplication procedure on smaller operands).

The step `st` is abstract here: it is specified by `StepOK`.
-/

set_option linter.unusedSimpArgs false

namespace NLogN.Top

open IntMul NLogN

abbrev XA : Fin HS := 44
abbrev XB : Fin HS := 45
abbrev SA : Fin HS := 46
abbrev SB : Fin HS := 47
abbrev XN : Fin HS := 48
abbrev XR : Fin HS := 49
abbrev TT : Fin HS := 50

def sgnE (i : Fin HS) : Ex HS := .sel (.reg i) (.lit (-1)) (.lit 1)
def absE (i : Fin HS) : Ex HS := .sel (.reg i) (.neg (.reg i)) (.reg i)
def maxE (i j : Fin HS) : Ex HS := .sel (.add (.reg i) (.neg (.reg j))) (.reg j) (.reg i)
def resE : Ex HS :=
  .sel (.reg NLogN.Top.SA) (.sel (.reg SB) (.reg XR) (.neg (.reg XR)))
    (.sel (.reg SB) (.neg (.reg XR)) (.reg XR))

/-- The schoolbook branch. -/
def smallRM : RM :=
  .rs (RS.seqL [.push XR (.mul (.reg XA) (.reg XB)), .pop XA, .pop XB])

def mulPre : RS :=
  RS.seqL [.push XA (.reg ARGM), .pop ARGM, .push XB (.reg ARGM), .pop ARGM,
    .push NLogN.Top.SA (sgnE XA), .push SB (sgnE XB), .asg XA (absE XA), .asg XB (absE XB),
    .push NLogN.Top.BX (.reg XA)]

def mulMid : RS :=
  RS.seqL [.pop NLogN.Top.BX, .push TT (.reg NLogN.Top.BC), .pop NLogN.Top.BC, .push NLogN.Top.XN (maxE TT NLogN.Top.BC), .pop NLogN.Top.BC, .pop TT]

def mulPost : RS :=
  RS.seqL [.pop NLogN.Top.XN, .push ARGM resE, .pop XR, .pop NLogN.Top.SA, .pop SB]

/-- The body of the multiplication procedure (as an `RM` program). -/
def mulPrefix : RM :=
  RM.seqL [.rs mulPre, blenRM, .rs (RS.seqL [.pop NLogN.Top.BX, .push NLogN.Top.BX (.reg XB)]), blenRM, .rs mulMid]

def mulRM (n₀ : ℕ) (st : RM) : RM :=
  .seq mulPrefix (.seq (.ite (.add (.reg NLogN.Top.XN) (.lit (-(n₀ : ℤ)))) smallRM st) (.rs mulPost))

/-- The procedure map: `prc 0` = DIF (procedure 2), `prc 1` = DIT (3), `prc 3` = bit length (4). -/
def pcH (k : ℕ) : Fin 5 := if k = 1 then 2 else if k = 2 then 3 else 4

/-- What the recursive step must satisfy.  `calls N` lists pairs `(count, size)` of the
multiplications it performs on `N`-bit inputs. -/
structure StepOK (st : RM) (n₀ : ℕ) (lin : ℕ → ℕ) (calls : ℕ → List (ℕ × ℕ)) : Prop where
  run : ∀ (s : IState HS) (a b N : ℕ) (LA LB LN : List Item), n₀ ≤ N → a < 2 ^ N → b < 2 ^ N →
    s.G XA = .int a :: LA → s.G XB = .int b :: LB → s.G NLogN.Top.XN = .int N :: LN →
    st.run s = ((s.upd XA LA).upd XB LB).upd XR (.int (a * b : ℕ) :: s.G XR)
  okc : ∀ (M : MulSpec) (Φ : ℕ → ℕ) (s : IState HS) (a b N : ℕ) (LA LB LN : List Item),
    n₀ ≤ N → a < 2 ^ N → b < 2 ^ N →
    s.G XA = .int a :: LA → s.G XB = .int b :: LB → s.G NLogN.Top.XN = .int N :: LN →
    (∀ q ∈ calls N, MulBnd M q.2 (Φ q.2)) →
    st.ok M s ∧ st.cost M s ≤ lin N + ((calls N).map fun q => q.1 * Φ q.2).sum
  smaller : ∀ N, n₀ ≤ N → ∀ q ∈ calls N, q.2 < N
  prc : ∀ (I : List Sym) (iprog : Fin 5 → IStmt HS 5), iprog 2 = (difBody 2).ren ssH id →
    iprog 3 = (ditBody 3).ren ssH id → iprog 4 = blenProc → RM.PrcOK I iprog pcH st

end NLogN.Top

namespace NLogN.Top

open IntMul NLogN

def sg1 (a : ℤ) : ℤ := if a < 0 then -1 else 1

/-- The state before the branch. -/
def st3 (s : IState HS) (a b : ℤ) (L : List Item) (N : ℕ) : IState HS :=
  (((((s.upd ARGM L).upd XA (.int (a.natAbs : ℤ) :: s.G XA)).upd XB
    (.int (b.natAbs : ℤ) :: s.G XB)).upd NLogN.Top.SA (.int (sg1 a) :: s.G NLogN.Top.SA)).upd SB
    (.int (sg1 b) :: s.G SB)).upd NLogN.Top.XN (.int (N : ℤ) :: s.G NLogN.Top.XN)

def nOf (a b : ℤ) : ℕ := max (zsz (a.natAbs : ℤ)) (zsz (b.natAbs : ℤ))

/-- The state after either branch. -/
def stB (s : IState HS) (a b : ℤ) (L : List Item) (N : ℕ) : IState HS :=
  (((st3 s a b L N).upd XA (s.G XA)).upd XB (s.G XB)).upd XR
    (.int ((a.natAbs * b.natAbs : ℕ) : ℤ) :: s.G XR)

end NLogN.Top

end

section
-- Source module: NLogN.HDigits
/-!
# Chunking an integer into `b`-bit digits

`chunkP X RB D C n`: with `x ≥ 0` on `X` and the ruler `2^b` on `RB`, cut `n` digits off `x`
(each `split` costs `O(b)`, independently of the size of the rest) and push them on `D`.
-/

namespace NLogN

open IntMul IntMul.HvdH.Ref

/-- Cut `n` digits of `b` bits off `X` onto `D`, counting on `C`. -/
def chunkP (X RB D C : Fin HS) (n : Ex HS) : RM := .rep C n (.op (.split X RB D))

/-- The state after `k` digits. -/
def chunkSt (X D : Fin HS) (s : IState HS) (b x : ℕ) (LX : List Item) (k : ℕ) : IState HS :=
  (s.upd X (.int ((x / 2 ^ (b * k) : ℕ) : ℤ) :: LX)).upd D
    (ints (chunksL b k x).reverse ++ s.G D)

end NLogN

end

section
-- Source module: NLogN.HCarry
/-!
# Carrying and assembling digits: `evalDigits`

Input: nonnegative coefficients `c₀, c₁, …, c_{M-1}` on stack `CA` (`c₀` on top), the ruler
`2^b` on `CR`.  Output on `CY`: `∑ cₖ 2^{bk}`.

Phase 1 normalises: `t := y + cₖ`, digit `t mod 2^b` (by `split`, cost `O(b)`), carry `t / 2^b`.
Phase 2 assembles from the top with `cat` (cost `O(b)` per digit) starting from the sentinel
`carry + 1 ≠ 0`, while `CP := CP · 2^b` tracks `2^{bM}`; finally the sentinel is removed.
-/

set_option linter.unusedSimpArgs false

namespace NLogN

open IntMul IntMul.HvdH IntMul.HvdH.Ref

abbrev CA : Fin HS := 20
abbrev CD : Fin HS := 21
abbrev CY : Fin HS := 22
abbrev CR : Fin HS := 23
abbrev CP : Fin HS := 24
abbrev CQ : Fin HS := 25

/-! ## Pure carrying -/

def carryL (b : ℕ) : List ℤ → ℤ → List ℤ × ℤ
  | [], y => ([], y)
  | c :: cs, y => ((y + c) % 2 ^ b :: (carryL b cs ((y + c) / 2 ^ b)).1,
      (carryL b cs ((y + c) / 2 ^ b)).2)

/-! ## Phase 1: normalisation -/

/-- Register holding the number of coefficients `M`. -/
abbrev CM : Fin HS := 26

def carryBody : RM :=
  .seq (.rs (.seq (.asg CY (.add (.reg CY) (.reg CA))) (.pop CA))) (.op (.split CY CR CD))

def carry1 : RM := .rep CQ (.reg CM) carryBody

/-- The state after `k` normalisation steps. -/
def c1St (s : IState HS) (b : ℕ) (c : List ℤ) (y0 : ℤ) (RA LY : List Item) (k : ℕ) :
    IState HS :=
  ((s.upd CA (ints (c.drop k) ++ RA)).upd CY (.int (carryL b (c.take k) y0).2 :: LY)).upd CD
    (ints (carryL b (c.take k) y0).1.reverse ++ s.G CD)

/-! ## Phase 2: assembly -/

def asmBody : RM := .seq (.op (.cat CY CR CD)) (.op (.shl CP CR))

def carry2 : RM :=
  .seq (.rs (.seq (.asg CY (.add (.reg CY) (.lit 1))) (.push CP (.lit 1))))
    (.seq (.rep CQ (.reg CM) asmBody) (.rs (.seq (.asg CY (.add (.reg CY) (.neg (.reg CP))))
      (.pop CP))))

/-- The state after `j` assembly steps (`M = dl.length` digits, top carry `y`). -/
def c2St (s1 : IState HS) (b : ℕ) (dl : List ℤ) (y : ℤ) (LY D0 P0 : List Item) (j : ℕ) :
    IState HS :=
  ((s1.upd CY (.int (kpack (2 ^ b) (dl.drop (dl.length - j)) + (y + 1) * 2 ^ (b * j)) :: LY)).upd
    CD (ints (dl.take (dl.length - j)).reverse ++ D0)).upd CP (.int (2 ^ (b * j)) :: P0)

def carryP : RM := .seq carry1 carry2

/-! ## Side conditions -/

/-! ## Cost -/

end NLogN

end

section
-- Source module: NLogN.HListOps
/-!
# Basic list moves between stacks

`mvList i j c n`: move `n` integers from stack `i` to stack `j` (reversing their order),
counting on `c`.
-/

set_option linter.unusedSimpArgs false

namespace NLogN

open IntMul

def mvBody (i j : Fin HS) : RS := .seq (.push j (.reg i)) (.pop i)

def mvList (i j c : Fin HS) (n : Ex HS) : RS := .rep c n (mvBody i j)

/-- The state after `k` moves. -/
def mvSt (i j : Fin HS) (s : IState HS) (l : List ℤ) (Ri : List Item) (k : ℕ) : IState HS :=
  (s.upd i (ints (l.drop k) ++ Ri)).upd j (ints (l.take k).reverse ++ s.G j)

end NLogN

end

section
-- Source module: NLogN.HFermMath
/-!
# Pure facts behind the Fermat encode/decode routines

* `kunpack B n W` (balanced digits) equals a carry pass `balConv` over the standard digits of
  any nonnegative `x ≡ W (mod B^n)`; so the machine can cut standard digits by `split` and fix
  them up with small numbers only.
* `kpack` of a signed digit list is `kpack` of the positive parts minus `kpack` of the negative
  parts.
-/

namespace NLogN

open IntMul IntMul.HvdH IntMul.HvdH.Ref

/-- Standard base-`B` digits of `x`, least significant first. -/
def stdDigits (B : ℤ) : ℕ → ℤ → List ℤ
  | 0, _ => []
  | n + 1, x => x % B :: stdDigits B n (x / B)

/-- Balanced conversion of standard digits with an incoming carry `c ∈ {0, 1}`. -/
def balConv (B : ℤ) : List ℤ → ℤ → List ℤ
  | [], _ => []
  | d :: ds, c => if B ≤ 2 * (d + c) then (d + c - B) :: balConv B ds 1
      else (d + c) :: balConv B ds 0

/-- Positive and negative parts of a signed digit. -/
def posP (z : ℤ) : ℤ := max z 0
def negP (z : ℤ) : ℤ := max (-z) 0

end NLogN

end

section
-- Source module: NLogN.HEnc
/-!
# Encoding a block of Gaussian integers as a Fermat residue

A Gaussian list is stored flattened, `gints [(a₀,b₀), (a₁,b₁), …] = a₀ :: b₀ :: a₁ :: …`.
`splitPN` pops a block of `r` pairs from `GU` and pushes the positive/negative parts of the real
and imaginary parts onto `E1 … E4` (reversed).
-/

set_option linter.unusedSimpArgs false

namespace NLogN

open IntMul IntMul.HvdH IntMul.HvdH.Ref

def gints (L : List (ℤ × ℤ)) : List Item := L.flatMap fun z => [.int z.1, .int z.2]

@[simp] theorem gints_nil : gints [] = [] := rfl
@[simp] theorem gints_cons (z : ℤ × ℤ) (L : List (ℤ × ℤ)) :
    gints (z :: L) = .int z.1 :: .int z.2 :: gints L := rfl

abbrev GU : Fin HS := 45
abbrev E1 : Fin HS := 40
abbrev E2 : Fin HS := 41
abbrev E3 : Fin HS := 42
abbrev E4 : Fin HS := 43
abbrev ET : Fin HS := 47
abbrev EC : Fin HS := 46

def posE (i : Fin HS) : Ex HS := .sel (.reg i) (.lit 0) (.reg i)
def negE (i : Fin HS) : Ex HS := .sel (.reg i) (.neg (.reg i)) (.lit 0)

def splitBody : RS :=
  RS.seqL [.push ET (.reg GU), .pop GU, .push E1 (posE ET), .push E2 (negE ET), .pop ET,
    .push ET (.reg GU), .pop GU, .push E3 (posE ET), .push NLogN.E4 (negE ET), .pop ET]

def splitPN (n : Ex HS) : RS := .rep EC n splitBody

def spSt (s : IState HS) (blk : List (ℤ × ℤ)) (R : List Item) (k : ℕ) : IState HS :=
  ((((s.upd GU (gints (blk.drop k) ++ R)).upd E1
    (ints ((blk.take k).map (posP ∘ Prod.fst)).reverse ++ s.G E1)).upd E2
    (ints ((blk.take k).map (negP ∘ Prod.fst)).reverse ++ s.G E2)).upd E3
    (ints ((blk.take k).map (posP ∘ Prod.snd)).reverse ++ s.G E3)).upd NLogN.E4
    (ints ((blk.take k).map (negP ∘ Prod.snd)).reverse ++ s.G NLogN.E4)

/-! ## Packing one nonnegative digit list -/

/-- Move the reversed list on `Ei` to `CA`, assemble it with `carryP`, push the value on `X`. -/
def packOne (Ei X : Fin HS) : RM :=
  RM.seqL [.rs (mvList Ei CA EC (.reg CM)), .rs (.push CY (.lit 0)), carryP,
    .rs (.push X (.reg CY)), .rs (.pop CY)]

/-! ## The whole block -/

abbrev X1 : Fin HS := 44
abbrev X2 : Fin HS := 48
abbrev X3 : Fin HS := 49
abbrev X4 : Fin HS := 50
abbrev EO : Fin HS := 51
abbrev RW : Fin HS := 52
abbrev FR : Fin HS := 53

/-- `X1 - X2 + (X3 - X4) · 2^{2rP}` -/
def encV : Ex HS :=
  .add (.add (.reg X1) (.neg (.reg X2))) (.shl (.add (.reg X3) (.neg (.reg X4))) (.reg RW))

def encE : Ex HS := .sel NLogN.encV (.add NLogN.encV (.reg FR)) NLogN.encV

def encP : RM :=
  RM.seqL [.rs (splitPN (.reg CM)), packOne E1 X1, packOne E2 X2, packOne E3 X3, packOne NLogN.E4 X4,
    .rs (.push EO encE), .rs (.pop X4), .rs (.pop X3), .rs (.pop X2), .rs (.pop X1)]

/-- The encoded value before reduction. -/
def encRaw (P : ℕ) (blk : List (ℤ × ℤ)) : ℤ :=
  kpack (2 ^ P) (blk.map (posP ∘ Prod.fst)) - kpack (2 ^ P) (blk.map (negP ∘ Prod.fst)) +
    (kpack (2 ^ P) (blk.map (posP ∘ Prod.snd)) - kpack (2 ^ P) (blk.map (negP ∘ Prod.snd))) *
      2 ^ (2 * blk.length * P)

/-! ## Side conditions -/

end NLogN

end

section
-- Source module: NLogN.HTopCost
/-!
# Side conditions and cost of the multiplication procedure
-/

set_option linter.unusedSimpArgs false

namespace NLogN.Top

open IntMul NLogN

end NLogN.Top

namespace NLogN.Top

open IntMul NLogN

def preCost (N : ℕ) : ℕ := 300000 * (N + 3) + 4000 * (N + 1) * (zsz (N : ℤ) + 2) + 4

def Q1 (s : IState HS) (a b : ℤ) (L : List Item) : IState HS :=
  (((((s.upd ARGM L).upd XA (.int |a| :: s.G XA)).upd XB (.int |b| :: s.G XB)).upd NLogN.Top.SA
    (.int (sg1 a) :: s.G NLogN.Top.SA)).upd SB (.int (sg1 b) :: s.G SB)).upd NLogN.Top.BX (.int |a| :: s.G NLogN.Top.BX)

def Q2 (s : IState HS) (a b : ℤ) (L : List Item) : IState HS :=
  (Q1 s a b L).upd NLogN.Top.BC (.int (zsz |a| : ℤ) :: s.G NLogN.Top.BC)

def Q3 (s : IState HS) (a b : ℤ) (L : List Item) : IState HS :=
  (Q2 s a b L).upd NLogN.Top.BX (.int |b| :: s.G NLogN.Top.BX)

def Q4 (s : IState HS) (a b : ℤ) (L : List Item) : IState HS :=
  (Q3 s a b L).upd NLogN.Top.BC (.int (zsz |b| : ℤ) :: .int (zsz |a| : ℤ) :: s.G NLogN.Top.BC)

/-- Cost of the schoolbook branch. -/
def smallCost (N : ℕ) : ℕ := 4 * 10 ^ 7 * (2 * N + 3) ^ 2 + 8 * N + 8

def topCost (n₀ N : ℕ) : ℕ :=
  preCost N + 20000 * (N + zsz (n₀ : ℤ) + 3) + 200000 * (2 * N + 3)

end NLogN.Top

end

section
-- Source module: NLogN.HTopInd
/-!
# Correctness and cost of the multiplication procedure, by induction on the operand size
-/

set_option linter.unusedSimpArgs false

namespace NLogN.Top

open IntMul NLogN

/-- Multiplication of operands of fewer than `n` bits, at cost `G (max size)`. -/
def Mlt (G : ℕ → ℕ) (n : ℕ) : MulSpec :=
  ⟨fun a b => G (max (zsz a) (zsz b)), fun a b => zsz a < n ∧ zsz b < n⟩

end NLogN.Top

namespace NLogN.Top

open IntMul NLogN

/-- The `N log N` shape used for the cost function. -/
def TT' (N : ℕ) : ℕ := (N + 1) * (zsz (N : ℤ) + 2)

/-- The constant for the small inputs. -/
def smallS (n₀ : ℕ) : ℕ := ∑ N ∈ Finset.range n₀, (topCost n₀ N + smallCost N + 2)

end NLogN.Top

end

section
-- Source module: NLogN.HTopMain
/-!
# The main procedure: read `x`, `y`, multiply, write `bin_{2n}`
-/

set_option linter.unusedSimpArgs false

namespace NLogN.Top

open IntMul NLogN

abbrev M1 : Fin HS := 10
abbrev M2 : Fin HS := 11

def mainH : IStmt HS 5 :=
  .seq (.op (.read ARGM M1 codeSep)) (.seq (.op (.read ARGM M2 codeEnd)) (.seq (.call 1)
    (.seq (.op (.shl M1 M2)) (.op (.write ARGM M1)))))

def mainHCost (n : ℕ) (Φ : ℕ) : ℕ := 2 * readRCost n + Φ + shlCost n + outRCost (2 * n)

end NLogN.Top

end

section
-- Source module: NLogN.HTopFinal
/-!
# The `O(n log n)` bound, conditional on a recursive step `st` satisfying `StepOK`
-/

set_option linter.unusedSimpArgs false

namespace NLogN.Top

open IntMul NLogN

/-- The program of the machine. -/
def hvProg (n₀ : ℕ) (st : RM) : Fin 5 → IStmt HS 5 :=
  ![mainH, (mulRM n₀ st).code 1 pcH, (difBody 2).ren ssH id, (ditBody 3).ren ssH id, blenProc]

def hvMachine (n₀ : ℕ) (st : RM) : MultitapeTM :=
  machineOf (fun j => compile (hvProg n₀ st j)) (by norm_num)

end NLogN.Top

end

section
-- Source module: NLogN.HAsym
/-!
# Asymptotics of the recursion: the Fermat-ring products

For `q = mkParams K N` the `3 · 2^{logOuter}` products of `fn + 1`-bit residues made by the three
`fermatW` calls cost at most a quarter of `TT' N` once `K ≥ 10^4` and `N` is large.
-/

set_option linter.unusedSimpArgs false

namespace NLogN.Asym

open IntMul NLogN NLogN.Top IntMul.HvdH IntMul.HvdH.Ref

/-- Facts about `mkParams K N` for huge `N`. -/
structure PF (K N : ℕ) : Prop where
  b_eq : (mkParams K N).b = Nat.size N
  zN : zsz (N : ℤ) = (mkParams K N).b
  b_big : 2 ^ (8 * K) < (mkParams K N).b
  N_low : 2 ^ ((mkParams K N).b - 1) ≤ N
  N_high : N < 2 ^ (mkParams K N).b
  Lb : (mkParams K N).L * (mkParams K N).b ≤ N + (mkParams K N).b
  NbL : N ≤ (mkParams K N).b * (mkParams K N).L
  L1 : 1 ≤ (mkParams K N).L
  E8 : 2 ^ (mkParams K N).logT ≤ 8 * (mkParams K N).L
  E4 : 4 * (mkParams K N).L ≤ 2 ^ (mkParams K N).logT
  Eb : (mkParams K N).logT ≤ (mkParams K N).b + 3
  Ke : K * (mkParams K N).e ≤ (mkParams K N).logT + K
  eK : (mkParams K N).logT ≤ K * (mkParams K N).e
  e1 : 2 ≤ (mkParams K N).e
  dh : (mkParams K N).dh < K
  lo : (mkParams K N).logOuter + (mkParams K N).e = (mkParams K N).logT
  P_eq : (mkParams K N).P = 2 * (mkParams K N).p + (mkParams K N).logT + 8
  p_eq : (mkParams K N).p = 2 * (mkParams K N).b + 2 * (mkParams K N).logT +
    6 * K * (mkParams K N).α ^ 2 + 8 * Nat.size (mkParams K N).b + 4 * K + 64
  X64 : 64 * (6 * K * (mkParams K N).α ^ 2) ≤ (mkParams K N).b
  B64 : 64 * (8 * Nat.size (mkParams K N).b + 4 * K + 64) ≤ (mkParams K N).b
  BK : 10 * (K * (Nat.size (mkParams K N).b + 9) + 2) ≤ (mkParams K N).b
  bB : (mkParams K N).b < 2 ^ Nat.size (mkParams K N).b
  bB' : 2 ^ (Nat.size (mkParams K N).b - 1) ≤ (mkParams K N).b
  B8 : 8 * K < Nat.size (mkParams K N).b
  α1 : 2 * K ≤ (mkParams K N).α
  α4 : (mkParams K N).α ^ 4 ≤ 81 * K ^ 4 * (64 * (mkParams K N).b)

end NLogN.Asym

end

section
-- Source module: NLogN.HDec
/-!
# Decoding a Fermat residue into a block of Gaussian integers

`decP` pops `z` (`0 ≤ z < F = 2^(4rP)+1`) from `KZ` and pushes
`gints (decBlk F P r z).reverse` onto `KO`:

1. `x := if F < 2z then z - 1 else z` (so `x ≥ 0` and `x ≡ balMod F z (mod 2^(4rP))`);
2. cut `4r` standard `P`-bit digits off `x` with `split` (`chunkP`);
3. a carry pass `balConv` over the digits gives `kunpack (2^P) (4r) (balMod F z)`;
4. the quarters are folded: real part `d_c + d_{c+r}`, imaginary part `d_{2r+c} + d_{3r+c}`.
-/

set_option linter.unusedSimpArgs false

namespace NLogN

open IntMul IntMul.HvdH IntMul.HvdH.Ref

/-! ## Pure facts -/

/-- The carry after a `balConv` pass. -/
def bcar (B : ℤ) : List ℤ → ℤ → ℤ
  | [], c => c
  | d :: ds, c => bcar B ds (if B ≤ 2 * (d + c) then 1 else 0)

/-- The value cut into digits. -/
def decX (N z : ℤ) : ℤ := if N + 1 + -(z + z) < 0 then z + -1 else z

/-- The folded block from the four quarters. -/
def foldF (q₁ q₂ q₃ q₄ : List ℤ) (c : ℕ) : ℤ × ℤ :=
  (q₁.getD c 0 + q₂.getD c 0, q₃.getD c 0 + q₄.getD c 0)

/-! ## Registers -/

/-- input `z` -/
abbrev KZ : Fin HS := 30
/-- output -/
abbrev KO : Fin HS := 31
/-- `2^P` -/
abbrev KB : Fin HS := 32
/-- `N = 2^(4rP)` -/
abbrev KN : Fin HS := 33
/-- `r` -/
abbrev KR : Fin HS := 34
/-- loop counter -/
abbrev KC : Fin HS := 35
/-- scratch: digits (high on top) -/
abbrev KD : Fin HS := 36
/-- scratch: digits (low on top) -/
abbrev KE : Fin HS := 37
/-- scratch: balanced digits (high on top) -/
abbrev KG : Fin HS := 38
/-- scratch: carry -/
abbrev KY : Fin HS := 39
/-- scratch: quarters -/
abbrev KQ1 : Fin HS := 27
abbrev KQ2 : Fin HS := 28
abbrev KQ3 : Fin HS := 29
abbrev KQ4 : Fin HS := 54

/-! ## The program -/

/-- `x := if F < 2z then z - 1 else z` (with `F = N + 1`). -/
def decXE : Ex HS :=
  .sel (.add (.add (.reg KN) (.lit 1)) (.neg (.add (.reg KZ) (.reg KZ))))
    (.add (.reg KZ) (.lit (-1))) (.reg KZ)

/-- `4r` -/
def fourR : Ex HS := .shl (.reg KR) (.lit 4)

/-- `d + c` -/
def balT : Ex HS := .add (.reg KE) (.reg KY)
/-- `B - 2(d + c) - 1` (negative iff `B ≤ 2(d+c)`) -/
def balCnd : Ex HS := .add (.add (.reg KB) (.neg (.add balT balT))) (.lit (-1))

def balBody : RS :=
  RS.seqL [.push KG (.sel balCnd (.add balT (.neg (.reg KB))) balT),
    .asg KY (.sel balCnd (.lit 1) (.lit 0)), .pop KE]

def balPass : RS := RS.seqL [.push KY (.lit 0), .rep KC fourR balBody, .pop KY]

def foldBody : RS :=
  RS.seqL [.push KO (.add (.reg KQ3) (.reg KQ4)), .push KO (.add (.reg KQ1) (.reg KQ2)),
    .pop KQ1, .pop KQ2, .pop KQ3, .pop KQ4]

def foldPass : RS := .rep KC (.reg KR) foldBody

def decP : RM :=
  RM.seqL [.rs (.asg KZ decXE), chunkP KZ KB KD KC fourR, .rs (.pop KZ),
    .rs (mvList KD KE KC fourR), .rs balPass,
    .rs (mvList KG KQ4 KC (.reg KR)), .rs (mvList KG KQ3 KC (.reg KR)),
    .rs (mvList KG KQ2 KC (.reg KR)), .rs (mvList KG KQ1 KC (.reg KR)), .rs foldPass]

/-! ## The carry pass -/

/-- The state of the carry pass after `k` digits. -/
def balSt (s : IState HS) (B : ℤ) (dl : List ℤ) (RE : List Item) (k : ℕ) : IState HS :=
  ((s.upd KE (ints (dl.drop k) ++ RE)).upd KG
    (ints (balConv B (dl.take k) 0).reverse ++ s.G KG)).upd KY
    (.int (bcar B (dl.take k) 0) :: s.G KY)

/-! ## The fold -/

/-- The state of the fold after `k` output pairs. -/
def foldSt (s : IState HS) (q₁ q₂ q₃ q₄ : List ℤ) (R₁ R₂ R₃ R₄ : List Item) (k : ℕ) :
    IState HS :=
  ((((s.upd KQ1 (ints (q₁.drop k) ++ R₁)).upd KQ2 (ints (q₂.drop k) ++ R₂)).upd KQ3
    (ints (q₃.drop k) ++ R₃)).upd KQ4 (ints (q₄.drop k) ++ R₄)).upd KO
    (gints ((List.range k).map (foldF q₁ q₂ q₃ q₄)).reverse ++ s.G KO)

/-! ## Side conditions of the pieces -/

/-- Expressions built from registers, literals, `add`, `neg`, `sel`. -/
def Ex.Plain : Ex HS → Bool
  | .reg _ | .lit _ => true
  | .add a b => a.Plain && b.Plain
  | .neg a => a.Plain
  | .sel c x y => c.Plain && x.Plain && y.Plain
  | _ => false

/-! ## The whole decoder -/

/-! ## Costs -/

end NLogN

end

section
-- Source module: NLogN.HSort
/-!
# Stable radix sort of fixed-width records on a register stack

* `mvK X Y C n`: move `n` integers from `X` to `Y` (reversed); `mvK_spec`, `mvK_cost`.
* `radixPass k`, `radixSort kb`: the pure stable partition pass / LSD radix sort on the first
  component (the key) of each record; `radixSort_keys`: for keys forming a permutation of
  `0, …, N-1` (with `N ≤ 2^kb`) the record with key `q` ends up at position `q`.
* `passP w`, `sortP w`: the register programs (records of width `w ≥ 1` on stack `0`, record
  count on `1`, key bits `kb` on `2`; scratch `46 … 53`); `sortP_spec`, `sortP_cost`.
All per-element code uses `Ex.Lin` expressions only; costs are bounded with the linear lemmas.
-/

set_option linter.unusedSimpArgs false

namespace NLogN

open IntMul RS

/-- Move `n` integers from `X` to `Y` one at a time (reversing them), counter on `C`. -/
def mvK (X Y C : Fin HS) (n : Ex HS) : RS := rep C n (NLogN.RS.seq (push Y (.reg X)) (pop X))

def mvSt_s (X Y : Fin HS) (s : IState HS) (L : List ℤ) (R : List Item) (k : ℕ) : IState HS :=
  (s.upd X (ints (L.drop k) ++ R)).upd Y ((ints (L.take k)).reverse ++ s.G Y)

/-! ## Radix sort: pure functions -/

/-- Bit `k` of the key (the first component) of a record; for keys `≥ 0` this is
`Nat.testBit key k`. -/
def keyBit (k : ℕ) (r : List ℤ) : Bool := decide (0 < ((r.headD 0).tdiv (2 ^ k)).tmod 2)

/-- One stable partition pass on bit `k`: records with bit `0` first, then those with bit `1`. -/
def radixPass (k : ℕ) (R : List (List ℤ)) : List (List ℤ) :=
  R.filter (fun r => !keyBit k r) ++ R.filter (keyBit k)

/-- LSD radix sort on the key bits `0, …, kb-1`. -/
def radixSort (kb : ℕ) (R : List (List ℤ)) : List (List ℤ) :=
  (List.range kb).foldl (fun R k => radixPass k R) R

/-! ## One pass as a register program

Registers: data `0`, record count `1`; scratch `46` (ones), `47` (zeros), `48`, `49` (their
element counts), `50` (loop counter), `51` (counter of `mvK`), `52` (holds `2^k`). -/

/-- `-(bit k of the key on top of 0)` for the divisor `2^k` on `52`. -/
def bitE : Ex HS := .neg (.rem (.quo (.reg 0) (.reg 52)) (.lit 2))

def recBody (w : ℕ) : RS :=
  ite bitE (NLogN.RS.seq (mvK 0 46 51 (.lit w)) (asg 48 (.add (.reg 48) (.lit w))))
    (NLogN.RS.seq (mvK 0 47 51 (.lit w)) (asg 49 (.add (.reg 49) (.lit w))))

def passLoop (w : ℕ) : RS := rep 50 (.reg 1) (recBody w)

def passP (w : ℕ) : RS :=
  NLogN.RS.seq (push 48 (.lit 0)) (NLogN.RS.seq (push 49 (.lit 0)) (NLogN.RS.seq (passLoop w)
    (NLogN.RS.seq (mvK 46 0 51 (.reg 48)) (NLogN.RS.seq (mvK 47 0 51 (.reg 49)) (NLogN.RS.seq (pop 49) (pop 48))))))

/-- The state after `j` records of the pass loop (relative to the state `s` before the pass). -/
def passSt (k : ℕ) (s : IState HS) (R : List (List ℤ)) (Rest : List Item) (j : ℕ) : IState HS :=
  ((((s.upd 0 (ints (R.drop j).flatten ++ Rest)).upd 46
    (ints (((R.take j).filter (keyBit k)).flatten.reverse) ++ s.G 46)).upd 47
    (ints (((R.take j).filter (fun r => !keyBit k r)).flatten.reverse) ++ s.G 47)).upd 48
    (.int (((R.take j).filter (keyBit k)).flatten.length) :: s.G 48)).upd 49
    (.int (((R.take j).filter (fun r => !keyBit k r)).flatten.length) :: s.G 49)

section
variable (k : ℕ) (s : IState HS) (R : List (List ℤ)) (Rest : List Item) (j : ℕ)

end

/-! ## The radix sort program

Data on `0`, record count on `1`, number of key bits on `2`; scratch `46 … 53`. -/

def sortBody (w : ℕ) : RS := NLogN.RS.seq (NLogN.passP w) (asg 52 (.shl (.reg 52) (.lit 2)))

def sortLoop (w : ℕ) : RS := rep 53 (.reg 2) (sortBody w)

/-- Radix sort of the records of width `w` on stack `0` by their first components. -/
def sortP (w : ℕ) : RS := NLogN.RS.seq (push 52 (.lit 1)) (NLogN.RS.seq (sortLoop w) (pop 52))

def sortSt (s : IState HS) (R : List (List ℤ)) (Rest : List Item) (k : ℕ) : IState HS :=
  (s.upd 0 (ints (radixSort k R).flatten ++ Rest)).upd 52 (.int (2 ^ k) :: s.G 52)

/-! ## Costs -/

/-! ## Correctness of radix sort for distinct keys -/

/-- The key of a record. -/
def key (r : List ℤ) : ℤ := r.headD 0

/-- Sorted by the key modulo `2^k`. -/
def SortedMod (k : ℕ) (L : List (List ℤ)) : Prop :=
  L.Pairwise (fun r r' => key r % 2 ^ k ≤ key r' % 2 ^ k)

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

/-- The keyed record for position `p = i * m + j`: key `j * n + i`. -/
def keyRec (n m : ℕ) (D : List (List ℤ)) (p : ℕ) : List ℤ :=
  (((p % m) * n + p / m : ℕ) : ℤ) :: D.getD p []

def kflat (n m : ℕ) (D : List (List ℤ)) (p : ℕ) : List ℤ :=
  ((List.range p).map (keyRec n m D)).flatten

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

/-- Inner loop states: `j` elements of row `i` done. -/
def innerSt (s u : IState HS) (D : List (List ℤ)) (Rest : List Item) (n m d i j : ℕ) :
    IState HS :=
  ((((u.upd 0 (ints (D.drop (i * m + j)).flatten ++ Rest)).upd 40
    (ints (kflat n m D (i * m + j)).reverse ++ s.G 40)).upd 41
    (.int ((j * n + i : ℕ) : ℤ) :: s.G 41)).upd 45
    (.int (((i * m + j) * (d + 1) : ℕ) : ℤ) :: s.G 45)).upd 54
    (.int ((i * m + j : ℕ) : ℤ) :: s.G 54)

def innerLoop (d : ℕ) : RS := rep 44 (.reg 2) (elemBody d)

def rowBody (d : ℕ) : RS :=
  NLogN.RS.seq (push 41 (.reg 42)) (NLogN.RS.seq (innerLoop d) (NLogN.RS.seq (pop 41) (asg 42 (.add (.reg 42) (.lit 1)))))

/-- Outer loop states: `i` rows done. -/
def outerSt (s : IState HS) (D : List (List ℤ)) (Rest : List Item) (n m d i : ℕ) : IState HS :=
  ((((s.upd 0 (ints (D.drop (i * m)).flatten ++ Rest)).upd 40
    (ints (kflat n m D (i * m)).reverse ++ s.G 40)).upd 42
    (.int (i : ℤ) :: s.G 42)).upd 45
    (.int ((i * m * (d + 1) : ℕ) : ℤ) :: s.G 45)).upd 54
    (.int ((i * m : ℕ) : ℤ) :: s.G 54)

def outerLoop (d : ℕ) : RS := rep 43 (.reg 1) (rowBody d)

/-- Attach the keys: the keyed records end up on `0`. -/
def attachP (d : ℕ) : RS :=
  NLogN.RS.seq (push 42 (.lit 0)) (NLogN.RS.seq (push 45 (.lit 0)) (NLogN.RS.seq (push 54 (.lit 0))
    (NLogN.RS.seq (outerLoop d) (mvK 40 0 51 (.reg 45)))))

/-! ## Stripping the keys -/

def stripBody (d : ℕ) : RS :=
  NLogN.RS.seq (pop 0) (NLogN.RS.seq (mvK 0 40 51 (.lit d)) (asg 45 (.add (.reg 45) (.lit d))))

def stripLoop (d : ℕ) : RS := rep 43 (.reg 54) (stripBody d)

def stripP (d : ℕ) : RS := NLogN.RS.seq (asg 45 (.lit 0)) (NLogN.RS.seq (stripLoop d) (mvK 40 0 51 (.reg 45)))

def stripSt (s : IState HS) (S : List (List ℤ)) (Rest : List Item) (d q : ℕ) : IState HS :=
  ((s.upd 0 (ints (S.drop q).flatten ++ Rest)).upd 40
    (ints ((S.take q).map List.tail).flatten.reverse ++ s.G 40)).upd 45
    (.int ((q * d : ℕ) : ℤ) :: (s.G 45).tail)

/-! ## The whole transposition -/

/-- Radix sort with the record count on `54` and the key bits on `3`. -/
def sortW (w : ℕ) : RS :=
  NLogN.RS.seq (push 1 (.reg 54)) (NLogN.RS.seq (push 2 (.reg 3)) (NLogN.RS.seq (sortP w) (NLogN.RS.seq (pop 2) (pop 1))))

def transP (d : ℕ) : RS :=
  NLogN.RS.seq (attachP d) (NLogN.RS.seq (sortW (d + 1)) (NLogN.RS.seq (stripP d) (NLogN.RS.seq (pop 54) (NLogN.RS.seq (pop 45) (pop 42)))))

/-! ## Costs -/

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

end NLogN

end

section
-- Source module: NLogN.HFW1
/-!
# Fermat-ring convolution machine, part 1: infrastructure

* `FW.Spec M p s s' c`: `p` may run from `s` (`ok`), ends in `s'` and costs at most `c`.
  Sequencing, register pushes/pops, list moves and counted loops in this form.
* The relocated transposition `trP` (data on stack `8`, `n, m, kb` on `1, 2, 3`).
-/

set_option linter.unusedSimpArgs false
set_option linter.unusedVariables false

namespace NLogN.FW

open IntMul NLogN

/-! ## Specifications -/

/-- `p` runs from `s` to `s'` with cost at most `c`. -/
def Spec (M : MulSpec) (p : RM) (s s' : IState HS) (c : ℕ) : Prop :=
  p.ok M s ∧ p.run s = s' ∧ p.cost M s ≤ c

/-! ## Counted loops -/

/-! ## List moves -/

/-! ## Chunks -/

open IntMul.HvdH

end NLogN.FW

end

section
-- Source module: NLogN.HFW2
/-!
# Fermat-ring convolution machine, part 2: multi-axis transforms as levels

`mdif F axes` (and `mdit`) are recursive over the axes; the machine applies one *level* per
axis to the whole array: the array is cut into `pre` blocks of size `2^j · B` and `colMap` is
applied to every block.  `descs axes pre` lists the levels `(pre, j, E, B)`.
-/

set_option linter.unusedSimpArgs false
set_option linter.unusedVariables false

namespace NLogN.FW

open IntMul IntMul.HvdH NLogN

/-- One level `(pre, j, E, B)`: `pre` blocks, `2^j` rows of `B` columns, root exponent `E`. -/
abbrev Desc := ℕ × ℕ × ℕ × ℕ

def descs : List (ℕ × ℕ) → ℕ → List Desc
  | [], _ => []
  | (j, E) :: rest, pre => (pre, j, E, Ref.blockSize rest) :: descs rest (pre * 2 ^ j)

/-- Apply `g E j` to the columns of every block of a level. -/
def lvl (g : ℕ → ℕ → List ℤ → List ℤ) (d : Desc) (A : List ℤ) : List ℤ :=
  (Ref.chunkL (2 ^ d.2.1 * d.2.2.2) A).flatMap (Ref.colMap d.2.1 d.2.2.2 (g d.2.2.1 d.2.1))

def lvls (g : ℕ → ℕ → List ℤ → List ℤ) (ds : List Desc) (A : List ℤ) : List ℤ :=
  ds.foldl (fun A d => lvl g d A) A

/-! ## Chunks of concatenations -/

/-! ## Lengths -/

/-! ## The multi-axis transforms as levels -/

/-! ## Canonical residues -/

end NLogN.FW

end

section
-- Source module: NLogN.HFW3
/-!
# Fermat-ring convolution machine, part 3: procedure calls and the relocated transposition

* `difPrc` / `ditPrc`: `RM.prc` calls of procedure number `0` (DIF, `pc 1`) and `1` (DIT, `pc 2`),
  with total functional specifications reading the SS registers
  `DX = 8` (data), `DL = 9` (`2^j`), `DE = 10` (`2^E`), `RN = 2` (`2^n`), `RF = 3` (`2^n+1`).
* `trP`: the transposition `transP 1` relocated so that its data stack is `8`.
-/

set_option linter.unusedSimpArgs false
set_option linter.unusedVariables false

namespace NLogN.FW

open IntMul IntMul.HvdH NLogN

/-- The data stack of the transforms (`ssH DX`). -/
abbrev XD : Fin HS := 8
/-- `2^j` (`ssH DL`) -/
abbrev XL : Fin HS := 9
/-- `2^E` (`ssH DE`) -/
abbrev XE : Fin HS := 10
/-- `2^n` (`ssH RN`) -/
abbrev XN : Fin HS := 2
/-- `2^n + 1` (`ssH RF`) -/
abbrev XF : Fin HS := 3

def itInt : Item → ℤ
  | .int a => a
  | _ => 0

/-! ## Specifications of the transforms -/

/-- Number of entries transformed: `2^j` read from `XL`. -/
def tlen (s : IState HS) : ℕ := (s.top XL).toNat

/-- The transformed data read from the state. -/
def tdat (s : IState HS) : List ℤ := ((s.G XD).take (tlen s)).map itInt

def difF (s : IState HS) : IState HS :=
  s.upd XD (ints (difZ (s.top XF) (Nat.log2 (s.top XE).toNat) (Nat.log2 (tlen s)) (tdat s)) ++
    (s.G XD).drop (tlen s))

def ditF (s : IState HS) : IState HS :=
  s.upd XD (ints (ditZ (s.top XF) (Nat.log2 (s.top NLogN.FW.XN).toNat) (Nat.log2 (s.top XE).toNat)
    (Nat.log2 (tlen s)) (tdat s)) ++ (s.G XD).drop (tlen s))

def difC (s : IState HS) : ℕ := difCost (Nat.log2 (s.top NLogN.FW.XN).toNat) (Nat.log2 (tlen s))
def ditC (s : IState HS) : ℕ := ditCost (Nat.log2 (s.top NLogN.FW.XN).toNat) (Nat.log2 (tlen s))

/-- Precondition of the SS transforms. -/
def preT (s : IState HS) : Prop :=
  ∃ (n j E : ℕ) (A : List ℤ) (R LL LE LN LF : List Item), 1 ≤ n ∧ A.length = 2 ^ j ∧
    (j = 0 ∨ 2 ^ j * E = 2 * n) ∧ Canon (2 ^ n + 1) A ∧ s.G XD = ints A ++ R ∧
    s.G XL = .int (2 ^ j) :: LL ∧ s.G XE = .int (2 ^ E) :: LE ∧ s.G NLogN.FW.XN = .int (2 ^ n) :: LN ∧
    s.G XF = .int (2 ^ n + 1) :: LF

theorem difF_frame (s : IState HS) (j : Fin HS) (hj : j ∉ [XD]) : (difF s).G j = s.G j := by
  simp only [List.mem_singleton] at hj; simp [difF, IState.upd_G, hj]

theorem ditF_frame (s : IState HS) (j : Fin HS) (hj : j ∉ [XD]) : (ditF s).G j = s.G j := by
  simp only [List.mem_singleton] at hj; simp [ditF, IState.upd_G, hj]

/-- DIF on the top `2^j` entries of `XD` (procedure `pc 1`). -/
def difPrc : RM := .prc 0 difF difC preT [XD] difF_frame (fun s => ⟨rfl, rfl⟩)

/-- DIT on the top `2^j` entries of `XD` (procedure `pc 2`). -/
def ditPrc : RM := .prc 1 ditF ditC preT [XD] ditF_frame (fun s => ⟨rfl, rfl⟩)

/-! ## The procedures meet their specifications -/

section prc

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

end prc

/-! ## The relocated transposition -/

/-- Swap the data stack `0` of `transP` with `XD = 8`. -/
def σT : Equiv.Perm (Fin HS) := swapsPerm [(0, 8)]

/-- Transpose the `n × m` integer array on top of `XD`; `n, m, kb` on `1, 2, 3`. -/
def trP : RM := (RM.rs (transP 1)).ren σT

end NLogN.FW

end

section
-- Source module: NLogN.HFW4
/-!
# Fermat-ring convolution machine, part 4: one level of a multi-axis transform

Registers (all levels):
* `XD = 8` the array, `XL = 9` `2^j`, `XE = 10` `2^E`, `XN = 2` `2^n`, `XF = 3` `2^n+1`;
* `PR = 13` number of blocks, `BS = 14` number of columns `B`, `KBR = 27` `kb`
  (`2^j B ≤ 2^kb`), `NN = 19` the array length;
* scratch: `YC = 4` (columns), `ZB = 5` (blocks), counters `C0 = 6`, `C1 = 7`, `C2 = 11`.
-/

set_option linter.unusedSimpArgs false
set_option linter.unusedVariables false

namespace NLogN.FW

open IntMul IntMul.HvdH NLogN

abbrev YC : Fin HS := 4
abbrev ZB : Fin HS := 5
abbrev C0 : Fin HS := 6
abbrev C1 : Fin HS := 7
abbrev C2 : Fin HS := 11
abbrev PR : Fin HS := 13
abbrev BS : Fin HS := 14
abbrev NN : Fin HS := 19
abbrev KBR : Fin HS := 27

/-! ## Transposition with parameters from registers -/

/-- Transpose the `top a × top b` array on `XD`. -/
def trWith (a b : Fin HS) : RM :=
  RM.seqL [.rs (.push 1 (.reg a)), .rs (.push 2 (.reg b)), .rs (.push 3 (.reg KBR)), trP,
    .rs (.pop 1), .rs (.pop 2), .rs (.pop 3)]

/-! ## List helpers -/

/-! ## The column loop -/

/-- The SS registers hold `2^j`, `2^E`, `2^n`, `2^n + 1`. -/
def RegsOK (n j E : ℕ) (t : IState HS) : Prop :=
  ∃ LL LE LN LF : List Item, t.G XL = .int (2 ^ j) :: LL ∧ t.G XE = .int (2 ^ E) :: LE ∧
    t.G NLogN.FW.XN = .int (2 ^ n) :: LN ∧ t.G XF = .int (2 ^ n + 1) :: LF

/-- A transform program on the top `2^j` entries of `XD`. -/
def TrOK (M : MulSpec) (T : RM) (g : List ℤ → List ℤ) (n j E cT : ℕ) : Prop :=
  ∀ (t : IState HS) (A : List ℤ) (R : List Item), A.length = 2 ^ j → Canon (2 ^ n + 1) A →
    t.G XD = ints A ++ R → RegsOK n j E t → Spec M T t (t.upd XD (ints (g A) ++ R)) cT

/-- Transforms keep length and canonical form. -/
def GOK (g : List ℤ → List ℤ) (n j : ℕ) : Prop :=
  ∀ A : List ℤ, A.length = 2 ^ j → Canon (2 ^ n + 1) A →
    (g A).length = 2 ^ j ∧ Canon (2 ^ n + 1) (g A)

def colBody (T : RM) : RM := .seq T (.rs (mvList XD YC C2 (.reg XL)))

def colLoop (T : RM) : RM := .rep C1 (.reg BS) (colBody T)

def colSt (g : List ℤ → List ℤ) (s : IState HS) (Cs : List (List ℤ)) (RD : List Item) (k : ℕ) :
    IState HS :=
  (s.upd XD (ints (Cs.drop k).flatten ++ RD)).upd YC (ints ((Cs.take k).flatMap g).reverse ++ s.G YC)

/-! ## One block -/

/-- `B · 2^j` -/
def cbE : Ex HS := .shl (.reg BS) (.reg XL)

def blkBody (T : RM) : RM :=
  RM.seqL [trWith XL BS, colLoop T, .rs (mvList YC XD C2 cbE), trWith BS XL,
    .rs (mvList XD NLogN.FW.ZB C2 cbE)]

/-- Cost of one block (`Cb = 2^j B` entries). -/
def blkCost (j B kb W cT : ℕ) : ℕ :=
  4 * 10 ^ 6 * (kb + 1) * (2 ^ j * B + 1) * (W + kb + 5) + 40000 * (W + 2) +
    (B * (cT + 2 ^ j * (14000 * (W + 2) + 30) + 12000 * (W + 2) + 62) + 4000 * (W + 2) + 32) +
    2 * (8000 * (W + 2) + 2 ^ j * B * (14000 * (W + 2) + 30) + 32)

def blkSt (g : List ℤ → List ℤ) (j B : ℕ) (s : IState HS) (Bl : List (List ℤ)) (RD : List Item)
    (k : ℕ) : IState HS :=
  (s.upd XD (ints (Bl.drop k).flatten ++ RD)).upd NLogN.FW.ZB
    (ints ((Bl.take k).flatMap (Ref.colMap j B g)).reverse ++ s.G NLogN.FW.ZB)

/-! ## One level -/

def lvlP (T : RM) : RM :=
  .seq (.rep C0 (.reg PR) (blkBody T)) (.rs (mvList NLogN.FW.ZB XD C2 (.reg NN)))

/-- Cost of one level with `pre` blocks. -/
def lvlCost (pre j B kb W cT : ℕ) : ℕ :=
  pre * (blkCost j B kb W cT + 8000 * (W + 2) + 30) + pre * (2 ^ j * B) * (14000 * (W + 2) + 30) +
    8000 * (W + 2) + 64

end NLogN.FW

end

section
-- Source module: NLogN.HFW5
/-!
# Fermat-ring convolution machine, part 5: passes over all axes

A level descriptor `(pre, j, E, B)` is stored as the four integers `pre, 2^j, 2^E, B`
(`encF`).  The forward pass reads the descriptors from `AXS = 15` (moving them to `AXR = 16`,
where they end up reversed); the backward pass reads them from `AXR` in reverse order and moves
them back.  Counter `CAX = 17`, number of levels on `NAX = 18`.
-/

set_option linter.unusedSimpArgs false
set_option linter.unusedVariables false

namespace NLogN.FW

open IntMul IntMul.HvdH NLogN

abbrev AXS : Fin HS := 15
abbrev AXR : Fin HS := 16
abbrev CAX : Fin HS := 17
abbrev NAX : Fin HS := 18

/-! ## Cost of a level -/

/-! ## Reading descriptors -/

/-- Copy the top of `src` to `r` and move it to `dst`. -/
def rd (src dst r : Fin HS) : RM :=
  RM.seqL [.rs (.push r (.reg src)), .rs (.push dst (.reg src)), .rs (.pop src)]

/-- Move the top of `src` to `dst`. -/
def mv1 (src dst : Fin HS) : RM := RM.seqL [.rs (.push dst (.reg src)), .rs (.pop src)]

/-- A descriptor as a palindrome of eight integers. -/
def enc8 (d : Desc) : List ℤ :=
  [(d.1 : ℤ), 2 ^ d.2.1, 2 ^ d.2.2.1, (d.2.2.2 : ℤ), (d.2.2.2 : ℤ), 2 ^ d.2.2.1, 2 ^ d.2.1, (d.1 : ℤ)]

/-- Read a descriptor into `PR, XL, XE, BS`, moving its eight integers from `src` to `dst`. -/
def rd8 (src dst : Fin HS) : RM :=
  RM.seqL [rd src dst PR, rd src dst XL, rd src dst XE, rd src dst BS, mv1 src dst, mv1 src dst,
    mv1 src dst, mv1 src dst]

/-! ## One axis -/

def axBody (src dst : Fin HS) (T : RM) : RM :=
  RM.seqL [rd8 src dst, lvlP T, .rs (.pop PR), .rs (.pop XL), .rs (.pop XE), .rs (.pop BS)]

/-- Side conditions of a level descriptor for an array of length `N`. -/
def DescOK (n kb W N : ℕ) (d : Desc) : Prop :=
  d.1 * (2 ^ d.2.1 * d.2.2.2) = N ∧ 1 ≤ d.1 ∧ 1 ≤ d.2.2.2 ∧ N ≤ 2 ^ kb ∧ n + 1 ≤ W ∧
    d.2.1 + 1 ≤ W ∧ d.2.2.1 + 1 ≤ W ∧ kb + 1 ≤ W

/-! ## A pass over all levels -/

def passP (src dst : Fin HS) (T : RM) : RM := .rep CAX (.reg NAX) (axBody src dst T)

/-- All conditions on the levels of a pass. -/
def PassOK (M : MulSpec) (T : RM) (G : ℕ → ℕ → List ℤ → List ℤ) (cT : ℕ → ℕ) (n kb W N J : ℕ)
    (D : List Desc) : Prop :=
  ∀ d ∈ D, DescOK n kb W N d ∧ TrOK M T (G d.2.2.1 d.2.1) n d.2.1 d.2.2.1 (cT d.2.1) ∧
    GOK (G d.2.2.1 d.2.1) n d.2.1 ∧ d.2.1 ≤ J

def pSt (G : ℕ → ℕ → List ℤ → List ℤ) (src dst : Fin HS) (s : IState HS) (D : List Desc)
    (A : List ℤ) (RD Rs : List Item) (k : ℕ) : IState HS :=
  ((s.upd XD (ints (lvls G (D.take k) A) ++ RD)).upd src (ints ((D.drop k).flatMap enc8) ++ Rs)).upd
    dst (ints ((D.take k).flatMap enc8).reverse ++ s.G dst)

/-- Cost of a pass over `m` levels. -/
def passCost (m n kb W N J : ℕ) : ℕ :=
  m * (10 ^ 7 * (kb + 1) * N * (W + kb + 5) + 500000 * (n + 3) * (J + 1) * N + 98000 * (W + 2) + 30) +
    4000 * (W + 2) + 32

end NLogN.FW

end

section
-- Source module: NLogN.HFW6
/-!
# Fermat-ring convolution machine, part 6: encoding all blocks

`encL`: the `N` blocks of `r` Gaussian integers on `GU = 45` are encoded by `encP` (residues on
`EO = 51`, reversed) and moved to `XD` in order.  Counter `C0`, count on `NN`.
-/

set_option linter.unusedSimpArgs false
set_option linter.unusedVariables false

namespace NLogN.FW

open IntMul IntMul.HvdH NLogN

/-! ## Relocation -/

/-! ## Gaussian lists -/

/-! ## The encoding loop -/

def encL : RM := .seq (.rep C0 (.reg NN) encP) (.rs (mvList EO XD C2 (.reg NN)))

/-- The residue of one block. -/
def encRes (P r : ℕ) (blk : List (ℤ × ℤ)) : ℤ := Ref.encBlk P r blk % (2 ^ (4 * r * P) + 1)

def eSt (P r : ℕ) (s : IState HS) (Bl : List (List (ℤ × ℤ))) (RU : List Item) (k : ℕ) :
    IState HS :=
  (s.upd GU (gints (Bl.drop k).flatten ++ RU)).upd EO
    (ints ((Bl.take k).map (encRes P r)).reverse ++ s.G EO)

/-- Cost of encoding `N` blocks of `r` entries. -/
def encCost (N r P Q W : ℕ) : ℕ :=
  N * (2000000 * (r + 1) * (Q + 3) + 1000000 * (r * P + 1) + 8000 * (W + 2) + 30) +
    N * (14000 * (W + 2) + 30) + 8000 * (W + 2) + 64

end NLogN.FW

end

section
-- Source module: NLogN.HFW7
/-!
# Fermat-ring convolution machine, part 7: pointwise products and scaling

Reduction modulo `F = 2^n + 1` with linear operations: `x = lo + 2^n hi ≡ lo - hi`.
-/

set_option linter.unusedSimpArgs false
set_option linter.unusedVariables false

namespace NLogN.FW

open IntMul IntMul.HvdH NLogN

/-! ## Pointwise products -/

abbrev TMP : Fin HS := 24

/-- `lo - hi` for `x = lo + 2^n hi` on `TMP`. -/
def dE : Ex HS := .add (.rem (.reg TMP) (.reg NLogN.FW.XN)) (.neg (.quo (.reg TMP) (.reg NLogN.FW.XN)))

/-- `x mod (2^n + 1)` for `0 ≤ x ≤ 2^{2n}`. -/
def redE : Ex HS := .sel dE (.add dE (.reg XF)) dE

def pmBody : RM :=
  RM.seqL [.mul TMP XD YC, .rs (.push NLogN.FW.ZB redE), .rs (.pop TMP), .rs (.pop XD), .rs (.pop YC)]

def pmL : RM := .seq (.rep C0 (.reg NN) pmBody) (.rs (mvList NLogN.FW.ZB XD C2 (.reg NN)))

def pmSt (n : ℕ) (A B : List ℤ) (s : IState HS) (RD RY : List Item) (k : ℕ) : IState HS :=
  ((s.upd XD (ints (B.drop k) ++ RD)).upd YC (ints (A.drop k) ++ RY)).upd NLogN.FW.ZB
    (ints (Ref.pmulZ (2 ^ n + 1) (A.take k) (B.take k)).reverse ++ s.G NLogN.FW.ZB)

/-- Cost of `N` pointwise products. -/
def pmCost (N W Φm : ℕ) : ℕ :=
  N * (Φm + 110000 * (W + 2)) + N * (14000 * (W + 2) + 30) + 8000 * (W + 2) + 64

end NLogN.FW

end