Schönhage–Strassen development, part 2: item-language control flow and the FM / FFT procedures

-- Schönhage–Strassen multitape-TM development (NLogN), generated by skeleton subtraction.

import Definitions.Def_IntMul_SS_Defs1

set_option autoImplicit false
set_option relaxedAutoImplicit false
set_option maxHeartbeats 1000000

section
-- Source module: NLogNP.IComp
/-!
# Control flow of the item language and the compiler correctness theorem
-/

namespace NLogN

open IntMul

variable {S P : ℕ}

/-- Classes of the top of an item stack, for tests. -/
inductive Cls where
  | none | mark | zero | pos | neg
  deriving DecidableEq

def cls : Option Item → Cls
  | none => .none
  | some .mark => .mark
  | some (.int z) => if z = 0 then .zero else if z < 0 then .neg else .pos

inductive IStmt (S P : ℕ) where
  | skip
  | op (o : IOp S)
  | seq (p q : IStmt S P)
  | ite (i : Fin S) (f : Cls → Bool) (p q : IStmt S P)
  | loop (i : Fin S) (f : Cls → Bool) (p : IStmt S P)
  | call (j : Fin P)

def IState.cl (s : IState S) (i : Fin S) : Cls := cls (s.G i).head?

inductive IExec (I : List Sym) (iprog : Fin P → IStmt S P) :
    IStmt S P → IState S → IState S → ℕ → Prop
  | skip (s) : IExec I iprog .skip s s 1
  | op {o s s' c} : o.Step I s s' c → IExec I iprog (.op o) s s' c
  | seq {p q s s₁ s₂ c₁ c₂} : IExec I iprog p s s₁ c₁ → IExec I iprog q s₁ s₂ c₂ →
      IExec I iprog (.seq p q) s s₂ (c₁ + c₂)
  | iteT {i f p q s s' c} : f (s.cl i) = true → IExec I iprog p s s' c →
      IExec I iprog (.ite i f p q) s s' (c + 25)
  | iteF {i f p q s s' c} : f (s.cl i) = false → IExec I iprog q s s' c →
      IExec I iprog (.ite i f p q) s s' (c + 25)
  | loopT {i f p s s₁ s₂ c₁ c₂} : f (s.cl i) = true → IExec I iprog p s s₁ c₁ →
      IExec I iprog (.loop i f p) s₁ s₂ c₂ → IExec I iprog (.loop i f p) s s₂ (c₁ + c₂ + 30)
  | loopF {i f p s} : f (s.cl i) = false → IExec I iprog (.loop i f p) s s 30
  | call {j s s' c} : IExec I iprog (iprog j) s s' c → IExec I iprog (.call j) s s' (c + 2)

/-! ## Compilation -/

def branchOn (i : Fin S) (k : Cls → Stmt (S + 4) P) : Stmt (S + 4) P :=
  caseTop (ι i) fun o =>
    if o = none then k .none
    else if o = some hdMark then k .mark
    else if o = some hdNeg then k .neg
    else .seq (.pop (ι i)) (.ite (ι i) isBitT (.seq (.push (ι i) hdPos) (k .pos))
      (.seq (.push (ι i) hdPos) (k .zero)))

def flagOf (i : Fin S) (f : Cls → Bool) : Stmt (S + 4) P :=
  branchOn i fun c => .push (s3 S) (if f c then 1 else 0)

def loopTest (o : Option Sym) : Bool := decide (o = some 1)

def compile : IStmt S P → Stmt (S + 4) P
  | .skip => .skip
  | .op o => o.stmt
  | .seq p q => .seq (compile p) (compile q)
  | .ite i f p q => branchOn i fun c => if f c then compile p else compile q
  | .loop i f p => .seq (flagOf i f) (.seq (.loop (s3 S) loopTest
      (.seq (.pop (s3 S)) (.seq (compile p) (flagOf i f)))) (.pop (s3 S)))
  | .call j => .call j

def inner (i : Fin S) (f : Cls → Bool) (p : IStmt S P) : Stmt (S + 4) P :=
  .loop (s3 S) loopTest (.seq (.pop (s3 S)) (.seq (compile p) (flagOf i f)))

variable {I : List Sym} {prog : Fin P → Stmt (S + 4) P}

def encF (i : Fin S) (f : Cls → Bool) (s : IState S) : St (S + 4) :=
  (encI s).push (s3 S) (if f (s.cl i) then 1 else 0)

end NLogN

namespace NLogN

open IntMul

variable {S P : ℕ}

end NLogN

end

section
-- Source module: NLogNP.IRuns
/-!
# Reasoning about item-language programs
-/

namespace NLogN

open IntMul

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

/-- `p` runs from `s` to `s'` with cost at most `c`. -/
def IRuns (I : List Sym) (iprog : Fin P → IStmt S P) (p : IStmt S P) (s s' : IState S) (c : ℕ) :
    Prop := ∃ c' ≤ c, IExec I iprog p s s' c'

namespace IRuns

end IRuns

/-! ## Operations as `IRuns` -/

namespace IRuns

end IRuns

/-! ## State lemmas -/

end NLogN

end

section
-- Source module: NLogNP.IEval
/-!
# A functional evaluator for item-language programs

`IOp.run` and `IStmt.run` compute the result state and cost of a program (with fuel for loops and
calls).  A successful evaluation gives an execution (`IStmt.run_exec`), so straight-line code can
be checked by `simp` evaluation instead of step-by-step reasoning.  Input and output operations
are not evaluated (`none`).
-/

namespace NLogN

open IntMul

variable {S P : ℕ}

/-- The exponent of a ruler. -/
def rulerExp (r : ℤ) : Option ℕ :=
  if r = 2 ^ (Nat.log2 r.toNat) then some (Nat.log2 r.toNat) else none

def IOp.run (s : IState S) : IOp S → Option (IState S × ℕ)
  | .add i j k =>
    match s.G i, s.G j with
    | .int a :: Li, .int b :: Lj =>
      if i = j then none else
        some ((((s.upd i Li).upd j Lj).upd k (.int (a + b) :: ((s.upd i Li).upd j Lj).G k)),
          addIntCost (zsz a + zsz b))
    | _, _ => none
  | .neg i =>
    match s.G i with
    | .int a :: L => some (s.upd i (.int (-a) :: L), 20)
    | _ => none
  | .dup i j =>
    match s.G i with
    | .int a :: _ => if i = j then none else some (s.upd j (.int a :: s.G j), dupCost (zsz a) + 17)
    | _ => none
  | .drop i =>
    match s.G i with
    | .int a :: L => some (s.upd i L, 2 * zsz a + 2)
    | _ => none
  | .mv i j =>
    match s.G i with
    | .int a :: L =>
      if i = j then none else some ((s.upd i L).upd j (.int a :: s.G j), 2 * moveCost (zsz a) + 19)
    | _ => none
  | .const i c => some (s.upd i (.int c :: s.G i), zsz c + 2)
  | .mark i => some (s.upd i (.mark :: s.G i), 1)
  | .unmark i =>
    match s.G i with
    | .mark :: L => some (s.upd i L, 1)
    | _ => none
  | .shl i j =>
    match s.G i, s.G j with
    | .int a :: Li, .int r :: _ =>
      match rulerExp r with
      | some m => if i = j then none else some (s.upd i (.int (a * 2 ^ m) :: Li), shlCost m)
      | none => none
    | _, _ => none
  | .split i j k =>
    match s.G i, s.G j with
    | .int a :: Li, .int r :: _ =>
      match rulerExp r with
      | some m =>
        if i = j ∨ i = k ∨ j = k then none else
          some ((s.upd i (.int (sg a (a.natAbs / 2 ^ m)) :: Li)).upd k
            (.int (sg a (a.natAbs % 2 ^ m)) :: s.G k), splitCost m + 17)
      | none => none
    | _, _ => none
  | .cat i j k =>
    match s.G i, s.G j, s.G k with
    | .int a :: Li, .int r :: _, .int b :: Lk =>
      match rulerExp r with
      | some m =>
        if i = j ∨ i = k ∨ j = k ∨ a = 0 ∨ b < 0 ∨ 2 ^ m ≤ b then none else
          some ((s.upd i (.int (sg a (a.natAbs * 2 ^ m + b.toNat)) :: Li)).upd k Lk, catCost m)
      | none => none
    | _, _, _ => none
  | .read _ _ _ => none
  | .write _ _ => none

/-- Evaluation with fuel (each loop iteration and call uses one unit). -/
def IStmt.run (iprog : Fin P → IStmt S P) : ℕ → IStmt S P → IState S → Option (IState S × ℕ)
  | _, .skip, s => some (s, 1)
  | _, .op o, s => o.run s
  | n, .seq p q, s =>
    match IStmt.run iprog n p s with
    | some (s₁, c₁) =>
      match IStmt.run iprog n q s₁ with
      | some (s₂, c₂) => some (s₂, c₁ + c₂)
      | none => none
    | none => none
  | n, .ite i f p q, s =>
    if f (s.cl i) then (IStmt.run iprog n p s).map (fun r => (r.1, r.2 + 25))
    else (IStmt.run iprog n q s).map (fun r => (r.1, r.2 + 25))
  | 0, .loop _ _ _, _ => none
  | n + 1, .loop i f p, s =>
    if f (s.cl i) then
      match IStmt.run iprog (n + 1) p s with
      | some (s₁, c₁) =>
        match IStmt.run iprog n (.loop i f p) s₁ with
        | some (s₂, c₂) => some (s₂, c₁ + c₂ + 30)
        | none => none
      | none => none
    else some (s, 30)
  | 0, .call _, _ => none
  | n + 1, .call j, s => (IStmt.run iprog n (iprog j) s).map (fun r => (r.1, r.2 + 2))

end NLogN

end

section
-- Source module: NLogNP.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

abbrev NS : ℕ := 32

abbrev ARG : Fin NS := 0
abbrev PRM : Fin NS := 1
abbrev RN : Fin NS := 2
abbrev RF : Fin NS := 3
abbrev T1 : Fin NS := 4
abbrev T2 : Fin NS := 5
abbrev T3 : Fin NS := 6
abbrev T4 : Fin NS := 7

/-- `o` evaluates to the state `t` with cost at most `c`. -/
def runsTo {S : ℕ} (o : Option (IState S × ℕ)) (t : IState S) (c : ℕ) : Prop :=
  match o with
  | some r => r.1 = t ∧ r.2 ≤ c
  | none => False

@[simp] theorem cls_none' : cls none = .none := rfl
@[simp] theorem cls_mark' : cls (some .mark) = .mark := rfl

@[simp] theorem IState.cl_def {S : ℕ} (s : IState S) (i : Fin S) : s.cl i = cls (s.G i).head? := rfl

/-! ## Fermat reduction -/

variable {P : ℕ}

/-- Reduce the integer on top of `X` modulo `2^n+1` (needs `|z| < 2^(2n+1)`). -/
def fredA (X : Fin NS) : IStmt NS P :=
  .seq (.op (.split X RN T1)) (.seq (.op (.split X RN T2)) (.seq (.op (.neg T2))
    (.seq (.op (.add T1 T2 T1)) (.op (.add T1 X X)))))

def fadd (X : Fin NS) : IStmt NS P :=
  .ite X (fun c => c == .neg) (.seq (.op (.dup RF T1)) (.op (.add T1 X X))) .skip

def fsub (X : Fin NS) : IStmt NS P :=
  .seq (.op (.dup RF T1)) (.seq (.op (.neg T1)) (.seq (.op (.dup X T2)) (.seq (.op (.add T2 T1 T1))
    (.ite T1 (fun c => c != .neg) (.seq (.op (.drop X)) (.op (.mv T1 X))) (.op (.drop T1))))))

def fred (X : Fin NS) : IStmt NS P :=
  .seq (fredA X) (.seq (fadd X) (.seq (fadd X) (.seq (fsub X) (fsub X))))

/-- 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: NLogNP.SSLoops
/-!
# Generic counted loops over item arrays
-/

namespace NLogN

open IntMul

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

/-- Decrement the counter on top of `C` (uses `T3`). -/
def decC (C : Fin NS) : IStmt NS P := .seq (.op (.const T3 (-1))) (.op (.add C T3 C))

/-- Move `k` integers from `X` to `Y` (reversing them); `k` is the counter on top of `C`. -/
def mvN (X Y C : Fin NS) : IStmt NS P :=
  .loop C (fun c => c != .zero) (.seq (.op (.mv X Y)) (decC C))

def ints (A : List ℤ) : List Item := A.map .int

@[simp] theorem ints_nil : ints [] = [] := rfl
@[simp] theorem ints_cons (a : ℤ) (A : List ℤ) : ints (a :: A) = .int a :: ints A := rfl

/-- All entries have at most `m` bits. -/
def Bnd (m : ℕ) (A : List ℤ) : Prop := ∀ a ∈ A, zsz a ≤ m

end NLogN

end

section
-- Source module: NLogNP.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 V P1)) (.op (.add T4 P1 P1)))) (.seq (fred P1)
    (.seq (.seq (.op (.mv DX T4)) (.seq (.op (.neg T4)) (.seq (.op (.mv V P2))
      (.seq (.op (.add T4 P2 P2)) (.op (.shl P2 W)))))) (.seq (fred P2)
        (.seq (.op (.shl W 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 DE T4)) (.seq (.op (.shl T4 DE)) (.op (.mv T4 DE))))

def difSegH : IStmt NS P := .seq (.op (.drop 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 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: NLogNP.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 V T3))
      (.op (.add T3 P1 P1))))) (.seq (fred P1)
    (.seq (.seq (.op (.mv T4 P2)) (.seq (.op (.mv V T3)) (.op (.add T3 P2 P2)))) (.seq (fred P2)
      (.seq (.op (.split W 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 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 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: NLogNP.SSAlg
/-!
# Algebra of the Fermat-ring FFT: inversion and the convolution property
-/

namespace NLogN

open IntMul

section ring

variable (n : ℕ)

/-! ## The forward transform evaluates at roots of unity -/

/-- `∑_{i<m} A_i x^i`. -/
def peval (A : List ℤ) (x : (ZMod (2 ^ n + 1))) (m : ℕ) : (ZMod (2 ^ n + 1)) := ∑ i ∈ Finset.range m, (A.getD i 0 : (ZMod (2 ^ n + 1))) * x ^ i

end ring

section ring2

variable (n : ℕ)

end ring2

/-! ## Cyclic convolution and pointwise products -/

/-- Cyclic convolution of length `K`, reduced mod `F`. -/
def cycZ (F : ℤ) (K : ℕ) (A B : List ℤ) : List ℤ :=
  (List.range K).map fun m => (∑ i ∈ Finset.range K, ∑ l ∈ Finset.range K,
    if (i + l) % K = m then A.getD i 0 * B.getD l 0 else 0) % F

/-- Pointwise products, reduced mod `F`. -/
def pmulZ (F : ℤ) : List ℤ → List ℤ → List ℤ
  | a :: A, b :: B => (a * b) % F :: pmulZ F A B
  | _, _ => []

section ring3

variable (n : ℕ)

end ring3

end NLogN

end

section
-- Source module: NLogNP.SSMath
/-!
# Arithmetic around the Fermat-ring FFT: pieces, twisting, recomposition and carries
-/

namespace NLogN

open IntMul

/-- `∑ A_i B^i`. -/
def zev (B : ℤ) : List ℤ → ℤ
  | [] => 0
  | a :: A => a + B * zev B A

/-- The `K` signed `M`-bit pieces of `x` (least significant first). -/
def pcs (M : ℕ) : ℕ → ℤ → List ℤ
  | 0, _ => []
  | K + 1, x => sg x (x.natAbs % 2 ^ M) :: pcs M K (sg x (x.natAbs / 2 ^ M))

/-- Twisting by powers of `2^Q`. -/
def wtZ (F : ℤ) (Q : ℕ) : ℕ → List ℤ → List ℤ
  | _, [] => []
  | i, a :: A => (a * 2 ^ (i * Q)) % F :: wtZ F Q (i + 1) A

/-- Negacyclic convolution of length `K`. -/
def ncZ (K : ℕ) (A B : List ℤ) : List ℤ :=
  (List.range K).map fun m => ∑ i ∈ Finset.range K, ∑ l ∈ Finset.range K,
    ((if i + l = m then A.getD i 0 * B.getD l 0 else 0) -
      (if i + l = m + K then A.getD i 0 * B.getD l 0 else 0))

section ring

variable (n : ℕ)

end ring

/-- Carry propagation into `M`-bit digits. -/
def carZ (M : ℕ) : List ℤ → ℤ → List ℤ × ℤ
  | [], c => ([], c)
  | e :: E, c => ((e + c) % 2 ^ M :: (carZ M E ((e + c) / 2 ^ M)).1,
      (carZ M E ((e + c) / 2 ^ M)).2)

end NLogN

end

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

namespace NLogN

open IntMul

abbrev RM : Fin NS := 18
abbrev RQ : Fin NS := 19
abbrev QS : Fin NS := 20
abbrev XA : Fin NS := 21
abbrev W2 : Fin NS := 22
abbrev HB : Fin NS := 23
abbrev CY : Fin NS := 24
abbrev DG : Fin NS := 25
abbrev ZZ : Fin NS := 26
abbrev PW : Fin NS := 27
abbrev PB : Fin NS := 28
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: NLogNP.SSPow
/-! # Powers of two: `PW : m :: L ↦ 2^m :: L` (recursive on the bits of `m`) -/

namespace NLogN

open IntMul

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

def pow2Base : IStmt NS P := .seq (.op (.drop PW)) (.op (.const PW 1))

def pow2A : IStmt NS P :=
  .seq (.op (.const T7 2)) (.seq (.op (.split PW T7 T6)) (.seq (.op (.drop T7)) (.op (.mv T6 PB))))

def pow2D : IStmt NS P :=
  .seq (.op (.dup PW T6)) (.seq (.op (.shl PW T6)) (.seq (.op (.drop T6)) (.op (.mv PB T6))))

def pow2E : IStmt NS P :=
  .ite T6 (fun c => c == .pos) (.seq (.op (.const T7 2)) (.seq (.op (.shl PW T7)) (.op (.drop T7))))
    .skip

def pow2Body (pP : Fin P) : IStmt NS P :=
  .ite PW (fun c => c == .zero) pow2Base
    (.seq pow2A (.seq (.call pP) (.seq pow2D (.seq pow2E (.op (.drop T6))))))

def pow2Cost (m : ℕ) : ℕ := 20000 * (m + 1)

end NLogN

end

section
-- Source module: NLogNP.ISchool
/-!
# Schoolbook multiplication in the item language (base case)
-/

namespace NLogN

open IntMul

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

def smulBody (X Y W Q R T : Fin S) : IStmt S P :=
  .seq (.op (.split Y Q R)) (.seq (.ite R (fun c => c == .zero) .skip
    (.ite R (fun c => c == .pos) (.seq (.op (.dup X T)) (.op (.add T W W)))
      (.seq (.op (.dup X T)) (.seq (.op (.neg T)) (.op (.add T W W))))))
    (.seq (.op (.drop R)) (.op (.shl X Q))))

/-- `A : x :: y :: L  ↦  x*y :: L`. -/
def smul (A X Y W Q R T : Fin S) : IStmt S P :=
  .seq (.op (.mv A X)) (.seq (.op (.mv A Y)) (.seq (.op (.const W 0)) (.seq (.op (.const Q 2))
    (.seq (.loop Y (fun c => c != .zero) (smulBody X Y W Q R T))
      (.seq (.op (.drop Y)) (.seq (.op (.drop X)) (.seq (.op (.drop Q)) (.op (.mv W A)))))))))

/-- The frame of the schoolbook loop. -/
def schoolSt (s : IState S) (A X Y W Q : Fin S) (L : List Item) (a b c : ℤ) : IState S :=
  ((((s.upd A L).upd X (.int a :: s.G X)).upd Y (.int b :: s.G Y)).upd W (.int c :: s.G W)).upd Q
    (.int 2 :: s.G Q)

def schoolIter (L : ℕ) : ℕ := 1500 * (L + 2)

end NLogN

namespace NLogN

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

def smulCost (N : ℕ) : ℕ := 2000 * (N + 2) ^ 2

end NLogN

end

section
-- Source module: NLogNP.SSProg
/-!
# The Schönhage–Strassen procedure `FM`: program text

`FM` multiplies modulo `2^N+1`: `ARG : a :: b :: L ↦ (a*b mod 2^N+1) :: L`, with `N` on top of
`PRM` and the ring rulers `2^N`, `2^N+1` on top of `RN`, `RF`.  Small `N` use schoolbook
multiplication; otherwise the inputs are cut into `K = 2^k` pieces of `M` bits, twisted and
transformed in the ring `ℤ/(2^N'+1)`, multiplied pointwise by recursive calls, transformed back,
untwisted, lifted to signed coefficients, carried into `M`-bit digits and recombined.
-/

namespace NLogN

open IntMul

variable {P : ℕ}

/-- `ARG : a :: b :: L ↦ a'' :: b'' :: L` with `x'' = (x mod 2^N) - ⌊x/2^N⌋`. -/
def fmPre : IStmt NS P :=
  .seq (.op (.split ARG RN T1)) (.seq (.op (.neg ARG)) (.seq (.op (.add T1 ARG ARG))
    (.seq (.op (.mv ARG T1)) (.seq (.op (.split ARG RN T2)) (.seq (.op (.neg ARG))
      (.seq (.op (.add T2 ARG ARG)) (.op (.mv T1 ARG))))))))

/-- Move `K` integers (`K` on top of `DL`) from `X` to `Y`. -/
def mvK (X Y : Fin NS) : IStmt NS P := .seq (.op (.dup DL CC)) (.seq (mvN X Y CC) (.op (.drop CC)))

/-- One piece: cut `M` bits off `XA`, twist by `W = 2^(iQ)`, reduce, push to `QS`. -/
def prepStep : IStmt NS P :=
  .seq (.op (.split XA RM T4)) (.seq (.op (.shl T4 W)) (.seq (fred T4) (.seq (.op (.mv T4 QS))
    (.seq (.op (.shl W RQ)) (decC CC)))))

/-- `ARG : x :: L ↦ L`, `DX : R ↦ (forward transform of the twisted pieces of x) ++ R`. -/
def fmPrep (pD : Fin P) : IStmt NS P :=
  .seq (.seq (.op (.mv ARG XA)) (.seq (.op (.const W 1)) (.op (.dup DL CC))))
    (.seq (.loop CC (fun c => c != .zero) prepStep)
      (.seq (.seq (.op (.drop CC)) (.seq (.op (.drop W)) (.op (.drop XA))))
        (.seq (mvK QS DX) (.call pD))))

/-- One pointwise product by a recursive call. -/
def pwStep (pF : Fin P) : IStmt NS P :=
  .seq (.op (.mv DX ARG)) (.seq (.op (.mv V ARG)) (.seq (.call pF) (.seq (.op (.mv ARG P1))
    (decC CC))))

def fmPw (pF : Fin P) : IStmt NS P :=
  .seq (.op (.dup DL CC)) (.seq (.loop CC (fun c => c != .zero) (pwStep pF)) (.op (.drop CC)))

/-- Fix a negative remainder of the signed split. -/
def fixS : IStmt NS P :=
  .seq (.op (.dup RM T6)) (.seq (.op (.add T6 T5 T5)) (.seq (.op (.const T6 (-1)))
    (.op (.add T6 T4 T4))))

/-- Untwist, lift to the symmetric range, add the carry and split off an `M`-bit digit. -/
def postStep : IStmt NS P :=
  .seq (.seq (.op (.mv DX T4)) (.op (.shl T4 W))) (.seq (fred T4)
    (.seq (.seq (.op (.shl T4 W2)) (.seq (.op (.dup HB T5)) (.op (.add T5 T4 T4)))) (.seq (fred T4)
      (.seq (.seq (.op (.dup HB T5)) (.seq (.op (.neg T5)) (.seq (.op (.add T5 T4 T4))
        (.seq (.op (.mv CY T5)) (.seq (.op (.add T5 T4 T4)) (.op (.split T4 RM T5)))))))
      (.seq (.ite T5 (fun c => c == .neg) fixS .skip)
        (.seq (.op (.mv T5 DG)) (.seq (.op (.mv T4 CY)) (.seq (.op (.split W RQ T5))
          (.seq (.op (.drop T5)) (decC CC))))))))))

def fmPost : IStmt NS P :=
  .seq (.seq (.op (.dup RN W)) (.seq (.op (.const CY 0)) (.op (.dup DL CC))))
    (.seq (.loop CC (fun c => c != .zero) postStep) (.seq (.op (.drop CC)) (.op (.drop W))))

def hornerStep : IStmt NS P := .seq (.op (.cat ZZ RM DG)) (decC CC)

def fmHorner : IStmt NS P :=
  .seq (.seq (.op (.const ZZ 1)) (.op (.dup DL CC)))
    (.seq (.loop CC (fun c => c != .zero) hornerStep) (.op (.drop CC)))

/-- `ZZ := (ZZ - 2^N - CY) mod (2^N+1)`, pushed to `ARG`. -/
def fmFinal : IStmt NS P :=
  .seq (.seq (.op (.dup RN T4)) (.seq (.op (.neg T4)) (.seq (.op (.add T4 ZZ ZZ))
    (.seq (.op (.mv CY T4)) (.seq (.op (.neg T4)) (.op (.add T4 ZZ ZZ)))))))
    (.seq (fred ZZ) (.op (.mv ZZ ARG)))

/-! ## Parameters -/

/-- `ZZ := ℓ = blen N`, `T1 := k = ℓ/2 - blen ℓ`, `DL := K = 2^k` (the latter via `pow2`). -/
def parA (pP : Fin P) : IStmt NS P :=
  .seq (.seq (blenS PRM ZZ) (.seq (blenS ZZ CY) (.seq (.op (.const T4 2)) (.seq (.op (.dup ZZ T1))
    (.seq (.op (.split T1 T4 T2)) (.seq (.op (.drop T2)) (.seq (.op (.drop T4))
      (.seq (.op (.mv CY T2)) (.seq (.op (.neg T2)) (.seq (.op (.add T2 T1 T1))
        (.op (.dup T1 PW))))))))))))
    (.seq (.call pP) (.op (.mv PW DL)))

/-- `T2 := M = N/K`, `RM := 2^M`. -/
def parB (pP : Fin P) : IStmt NS P :=
  .seq (.seq (.op (.dup PRM T2)) (.seq (.op (.split T2 DL T3)) (.seq (.op (.drop T3))
    (.op (.dup T2 PW)))))
    (.seq (.call pP) (.op (.mv PW RM)))

/-- `T2 := Q = (2M+2k+3)/K + 1`, `T3 := N' = Q*K`, `RQ := 2^Q`, `DE := 2^(2Q)`. -/
def parC (pP : Fin P) : IStmt NS P :=
  .seq (.seq (.op (.dup T2 T3)) (.seq (.op (.add T3 T2 T2)) (.seq (.op (.dup T1 T3))
    (.seq (.op (.dup T1 T4)) (.seq (.op (.add T3 T4 T4)) (.seq (.op (.add T4 T2 T2))
      (.seq (.op (.const T3 3)) (.seq (.op (.add T3 T2 T2)) (.seq (.op (.split T2 DL T3))
        (.seq (.op (.drop T3)) (.seq (.op (.const T3 1)) (.seq (.op (.add T3 T2 T2))
          (.seq (.op (.dup T2 T3)) (.seq (.op (.shl T3 DL)) (.op (.dup T2 PW))))))))))))))))
    (.seq (.call pP) (.seq (.op (.mv PW RQ)) (.seq (.op (.dup RQ T4)) (.seq (.op (.shl T4 RQ))
      (.op (.mv T4 DE))))))

/-- The ring `2^N'+1` on `RN`, `RF`. -/
def parD (pP : Fin P) : IStmt NS P :=
  .seq (.op (.dup T3 PW)) (.seq (.call pP) (.seq (.op (.dup PW T4)) (.seq (.op (.mv PW RN))
    (.seq (.op (.const T5 1)) (.seq (.op (.add T5 T4 T4)) (.op (.mv T4 RF)))))))

/-- `W2 := 2^(N'-k)`, `HB := 2^(N'-1)`, `PRM := N'`, and the temporaries are dropped. -/
def parE (pP : Fin P) : IStmt NS P :=
  .seq (.seq (.op (.dup T3 PW)) (.seq (.op (.dup T1 T4)) (.seq (.op (.neg T4))
    (.op (.add T4 PW PW)))))
    (.seq (.call pP) (.seq (.seq (.op (.mv PW W2)) (.seq (.op (.dup T3 PW))
      (.seq (.op (.const T4 (-1))) (.op (.add T4 PW PW)))))
    (.seq (.call pP) (.seq (.op (.mv PW HB)) (.seq (.op (.mv T3 PRM)) (.seq (.op (.drop T2))
      (.seq (.op (.drop T1)) (.op (.drop ZZ)))))))))

def fmParams (pP : Fin P) : IStmt NS P :=
  .seq (parA pP) (.seq (parB pP) (.seq (parC pP) (.seq (parD pP) (parE pP))))

/-- Drop the inner ring and the transform parameters (keeping `RM`, `DL`). -/
def fmUnpar : IStmt NS P :=
  .seq (.op (.drop RN)) (.seq (.op (.drop RF)) (.seq (.op (.drop PRM)) (.seq (.op (.drop W2))
    (.seq (.op (.drop HB)) (.seq (.op (.drop DE)) (.op (.drop RQ)))))))

def fmMain (pF pD pT pP : Fin P) : IStmt NS P :=
  .seq fmPre (.seq (fmParams pP) (.seq (fmPrep pD) (.seq (mvK DX U) (.seq (mvK U V)
    (.seq (fmPrep pD) (.seq (fmPw pF) (.seq (mvK P1 DX) (.seq (.call pT) (.seq fmPost
      (.seq fmUnpar (.seq fmHorner (.seq fmFinal (.seq (.op (.drop RM)) (.op (.drop DL)))))))))))))))

/-- Schoolbook multiplication and reduction for small `N`. -/
def fmBase : IStmt NS P := .seq (smul ARG XA QS W CC U V) (fred ARG)

/-- Threshold on the bit length of `N`. -/
def fmThr : ℕ := 2048

def fmBody (pF pD pT pP : Fin P) : IStmt NS P :=
  .seq (.seq (blenS PRM ZZ) (.seq (.op (.const T4 (-(fmThr : ℤ)))) (.op (.add T4 ZZ ZZ))))
    (.ite ZZ (fun c => c == .neg) (.seq (.op (.drop ZZ)) fmBase)
      (.seq (.op (.drop ZZ)) (fmMain pF pD pT pP)))

end NLogN

end

section
-- Source module: NLogNP.SSFM1
/-!
# `FM`: correctness of the preprocessing segments
-/

namespace NLogN

open IntMul

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

/-! ## Pieces -/

/-- The quotient after `i` cuts of `M` bits. -/
def qit (M i : ℕ) (x : ℤ) : ℤ := (fun y : ℤ => sg y (y.natAbs / 2 ^ M))^[i] x

def prepCost (N n k : ℕ) : ℕ := difCost n k + 40000 * 2 ^ k * (n + 3) + 1000 * (N + n + 10)

end NLogN

end

section
-- Source module: NLogNP.SSFM2
/-!
# `FM`: pointwise products, untwisting with carries, and recombination
-/

namespace NLogN

open IntMul

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

/-- The specification of a call of `FM` in the ring `2^n+1`, with cost at most `C`. -/
def FMSpec (pF : Fin P) (iprog : Fin P → IStmt NS P) (I : List Sym) (n C : ℕ) : Prop :=
  ∀ (s : IState NS) (a b : ℤ) (LA LP LN LF : List Item), 0 ≤ a ∧ a ≤ 2 ^ n → 0 ≤ b ∧ b ≤ 2 ^ n →
    s.G ARG = .int a :: .int b :: LA → s.G PRM = .int n :: LP → s.G RN = .int (2 ^ n) :: LN →
    s.G RF = .int (2 ^ n + 1) :: LF →
    IRuns I iprog (.call pF) s (s.upd ARG (.int ((a * b) % (2 ^ n + 1)) :: LA)) C

end NLogN

end

section
-- Source module: NLogNP.SSFM3
/-!
# `FM`: untwisting, lifting and carrying
-/

namespace NLogN

open IntMul

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

/-- Untwist by `2^e`, unweight by `2^(n-k)` and lift to the symmetric range. -/
def liftE (n k e : ℕ) (d : ℤ) : ℤ :=
  (((d * 2 ^ e) % (2 ^ n + 1)) * 2 ^ (n - k) + 2 ^ (n - 1)) % (2 ^ n + 1) - 2 ^ (n - 1)

/-- The lifted coefficients of the inverse transform `D` (with `K` entries). -/
def liftL (n k Q K : ℕ) (D : List ℤ) : List ℤ :=
  (List.range K).map fun i => liftE n k ((K - i) * Q) (D.getD i 0)

end NLogN

end

section
-- Source module: NLogNP.SSFM4
/-!
# `FM`: parameter computation
-/

namespace NLogN

open IntMul

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

/-- The parameters of `FM` at size `N`. -/
def fmL (N : ℕ) : ℕ := zsz (N : ℤ)
def fmB (N : ℕ) : ℕ := zsz (fmL N : ℤ)
def fmK (N : ℕ) : ℕ := fmL N / 2 - fmB N
def fmM (N : ℕ) : ℕ := N / 2 ^ fmK N
def fmQ (N : ℕ) : ℕ := (2 * fmM N + 2 * fmK N + 3) / 2 ^ fmK N + 1
def fmN' (N : ℕ) : ℕ := fmQ N * 2 ^ fmK N

end NLogN

end

section
-- Source module: NLogNP.SSCorrect
/-!
# Correctness of one Schönhage–Strassen level
-/

namespace NLogN

open IntMul

/-- The folded input `x mod 2^N - ⌊x / 2^N⌋`. -/
def fold (N : ℕ) (x : ℤ) : ℤ := x % 2 ^ N + -(x / 2 ^ N)

end NLogN

end

section
-- Source module: NLogNP.SSFM5
/-!
# One level of `FM`
-/

namespace NLogN

open IntMul

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

/-- The cost of one level of `FM` beyond the recursive calls. -/
def fmLevelCost (N : ℕ) : ℕ :=
  2 * difCost (fmN' N) (fmK N) + ditCost (fmN' N) (fmK N) +
    200000 * 2 ^ fmK N * (fmN' N + 3) + 2000 * (fmL N + 1) ^ 2 + 400000 * (N + 3)

end NLogN

end

section
-- Source module: NLogNP.SSParam
/-!
# Parameter inequalities for `FM`

For `ℓ = fmL N ≥ 2048` and `2^k ∣ N` (`k = fmK N`), the recursive size `fmN' N` is smaller,
keeps the divisibility invariant, and the potential `fmPhi` absorbs one level.
-/

namespace NLogN

/-- `λ = blen (ℓ^4)`. -/
def fmLam (N : ℕ) : ℕ := zsz (((fmL N) ^ 4 : ℕ) : ℤ)

/-- The constant of the potential. -/
def fmA : ℕ := 2 ^ 2100

/-- The potential `A · N · ℓ · λ`. -/
def fmPhi (N : ℕ) : ℕ := fmA * N * fmL N * fmLam N

/-! ## Arithmetic helpers -/

end NLogN

end

section
-- Source module: NLogNP.SSMain
/-!
# The top-level Schönhage–Strassen program
-/

namespace NLogN

open IntMul

section blen

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

end blen

/-- Procedure 0: read `x`, `y`; `N0 = 2^(blen(n+1)+1)`; call `FM`; write `bin_{2n}`. -/
def mainBody : IStmt NS 5 :=
  .seq (.op (.read ARG T1 codeSep)) (.seq (.op (.read ARG T2 codeEnd)) (.seq mainA
    (.seq (.call 1) (.seq (.op (.shl T1 T2)) (.op (.write ARG T1))))))
where
  mainA : IStmt NS 5 :=
    .seq (blenS T1 ZZ) (.seq (blenS ZZ PW) (.seq (.seq (.op (.const T4 1)) (.op (.add T4 PW PW)))
      (.seq (.call 4) (.seq (.seq (.op (.mv PW PRM)) (.op (.dup PRM PW))) (.seq (.call 4)
        (.seq (.op (.dup PW T4)) (.seq (.op (.mv PW RN)) (.seq (.op (.const T5 1))
          (.seq (.op (.add T5 T4 T4)) (.op (.mv T4 RF)))))))))))

def ssProg : Fin 5 → IStmt NS 5 :=
  ![mainBody, fmBody 1 2 3 4, difBody 2, ditBody 3, pow2Body 4]

/-- The FM size used for `n`-bit inputs. -/
def mainN0 (n : ℕ) : ℕ := 2 ^ (zsz ((n + 1 : ℕ) : ℤ) + 1)

def mainACost (n : ℕ) : ℕ :=
  2000 * (n + 2) * (zsz ((n + 1 : ℕ) : ℤ) + 2) + 2000 * (zsz ((n + 1 : ℕ) : ℤ) + 1) *
    (zsz (zsz ((n + 1 : ℕ) : ℤ) : ℤ) + 2) + 60000 * (mainN0 n + 3)

def mainCost (n : ℕ) : ℕ :=
  2 * readRCost n + mainACost n + fmPhi (mainN0 n) + shlCost n + outRCost (2 * n)

end NLogN

end

section
-- Source module: NLogNP.SSFinal
/-!
# The Schönhage–Strassen bound `O(n log n log log n)`
-/

namespace NLogN

open IntMul

attribute [local irreducible] fmA

/-- The Schönhage–Strassen machine. -/
def ssMachine : MultitapeTM := machineOf (fun j => compile (ssProg j)) (by norm_num)

end NLogN

namespace IntMul.TM

end IntMul.TM

end