From fffd3d295ed200e1996e457fe8b784660bb32dc3 Mon Sep 17 00:00:00 2001 From: Samuel Giuseppe Schlesinger Date: Fri, 30 Oct 2026 18:49:27 -0400 Subject: [PATCH 1/4] feat(WordRAM/AKS): word-RAM subroutines for the AKS primality test (#893) Euclid's gcd, the integer k-th root by binary search with a capped power, the perfect-power test (AKS step 1), and multiplication modulo X^r - 1 and n on coefficient arrays, each with an explicit instruction bound that does not depend on the word size. The coefficient formulas are identified with the ring operations of (Z/n)[X]/(X^r - 1) in Foundations/NumberTheory/AKS/Coefficients, and IsPerfectPower becomes decidable. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_015DLg8kwPs8yi7ePPdtoNvG --- Complexity.lean | 6 + .../Algorithms/WordRAM/AKS/PerfectPower.lean | 278 +++++++++++++ .../Algorithms/WordRAM/AKS/PolyMul.lean | 364 ++++++++++++++++++ Complexity/Algorithms/WordRAM/NatGcd.lean | 96 +++++ Complexity/Algorithms/WordRAM/NatNthRoot.lean | 324 ++++++++++++++++ .../NumberTheory/AKS/Coefficients.lean | 176 +++++++++ .../NumberTheory/AKS/PerfectPower.lean | 64 +++ ComplexityTests/WordRAMAKS.lean | 109 ++++++ 8 files changed, 1417 insertions(+) create mode 100644 Complexity/Algorithms/WordRAM/AKS/PerfectPower.lean create mode 100644 Complexity/Algorithms/WordRAM/AKS/PolyMul.lean create mode 100644 Complexity/Algorithms/WordRAM/NatGcd.lean create mode 100644 Complexity/Algorithms/WordRAM/NatNthRoot.lean create mode 100644 Complexity/Foundations/NumberTheory/AKS/Coefficients.lean create mode 100644 Complexity/Foundations/NumberTheory/AKS/PerfectPower.lean create mode 100644 ComplexityTests/WordRAMAKS.lean diff --git a/Complexity.lean b/Complexity.lean index dba79c39b..c2e70310d 100644 --- a/Complexity.lean +++ b/Complexity.lean @@ -80,6 +80,8 @@ public import Complexity.Algorithms.ThreeSum.ChanHe.Sizes public import Complexity.Algorithms.ThreeSum.Convolution public import Complexity.Algorithms.ThreeSum.Defs public import Complexity.Algorithms.ThreeSum.ThreeSets +public import Complexity.Algorithms.WordRAM.AKS.PerfectPower +public import Complexity.Algorithms.WordRAM.AKS.PolyMul public import Complexity.Algorithms.WordRAM.Deduplicate.Count public import Complexity.Algorithms.WordRAM.Deduplicate.Model public import Complexity.Algorithms.WordRAM.Deduplicate.Program @@ -354,7 +356,9 @@ public import Complexity.Algorithms.WordRAM.MatrixMultiplication.Schonhage.Wante public import Complexity.Algorithms.WordRAM.MatrixMultiplication.Schonhage.Wanted.Setup public import Complexity.Algorithms.WordRAM.MatrixMultiplication.Schonhage.Wanted.Width public import Complexity.Algorithms.WordRAM.NatClog +public import Complexity.Algorithms.WordRAM.NatGcd public import Complexity.Algorithms.WordRAM.NatLog +public import Complexity.Algorithms.WordRAM.NatNthRoot public import Complexity.Algorithms.WordRAM.NatSqrt public import Complexity.Algorithms.WordRAM.PrefixSums.Dynamic public import Complexity.Algorithms.WordRAM.PrefixSums.Query @@ -2758,9 +2762,11 @@ public import Complexity.Foundations.LinearAlgebra.ReplacementDesign public import Complexity.Foundations.LinearAlgebra.VariableCopyPlacement public import Complexity.Foundations.Logic.Function.Iterate public import Complexity.Foundations.MeasureTheory.Constructions.BorelSpace.Order +public import Complexity.Foundations.NumberTheory.AKS.Coefficients public import Complexity.Foundations.NumberTheory.AKS.Correctness public import Complexity.Foundations.NumberTheory.AKS.Criterion public import Complexity.Foundations.NumberTheory.AKS.Introspective +public import Complexity.Foundations.NumberTheory.AKS.PerfectPower public import Complexity.Foundations.NumberTheory.AKS.SmallR public import Complexity.Foundations.NumberTheory.Harmonic public import Complexity.Foundations.NumberTheory.PrimeCounting diff --git a/Complexity/Algorithms/WordRAM/AKS/PerfectPower.lean b/Complexity/Algorithms/WordRAM/AKS/PerfectPower.lean new file mode 100644 index 000000000..da3c1d2f4 --- /dev/null +++ b/Complexity/Algorithms/WordRAM/AKS/PerfectPower.lean @@ -0,0 +1,278 @@ +/- +Copyright (c) 2026 Samuel Schlesinger. All rights reserved. +Released under Apache 2.0 license as described in the file LICENSE. +Authors: Samuel Schlesinger +-/ + +module +-- upstream: candidate + +public import Complexity.Algorithms.WordRAM.NatLog +public import Complexity.Algorithms.WordRAM.NatNthRoot +public import Complexity.Foundations.NumberTheory.AKS.PerfectPower + +/-! +# Perfect-power detection on the word RAM + +Step 1 of the AKS algorithm [Agrawal, Kayal, Saxena, Section 4][AKS2004] asks whether `n = a ^ b` +for some `b > 1`. For `n ≥ 2` such an `a` is at least `2`, so `b ≤ log₂ n`, and `n` is a `b`-th +power exactly when `Nat.nthRoot b n ^ b = n` (`Cslib.AKS.isPerfectPower_iff_nthRoot`). The program +computes `⌊log₂ n⌋`, and for each `b` from `2` to `⌊log₂ n⌋` takes the `b`-th root by binary search +(`Cslib.Algorithms.WordRAM.NatNthRoot.program`) and compares its capped `b`-th power with `n`. +Every intermediate value is at most `n² + 2n + 1`. + +## Main definitions + +* `Cslib.Algorithms.WordRAM.AKS.PerfectPower.program`: the perfect-power test. +* `Cslib.Algorithms.WordRAM.AKS.PerfectPower.time`: its instruction bound, `O(log³ n)`. + +## Main results + +* `Cslib.Algorithms.WordRAM.AKS.PerfectPower.exec`: the program decides `Cslib.AKS.IsPerfectPower` + for `n ≥ 2` within `time n` instructions. + +## References + +* [M. Agrawal, N. Kayal, N. Saxena, *PRIMES is in P*][AKS2004] +-/ + +@[expose] public section + +namespace Cslib.Algorithms.WordRAM.AKS.PerfectPower + +open Cslib.WordRAM + +/-- The candidate `n`, preserved. -/ +abbrev n : Fin 14 := 0 +/-- The current exponent `b`. -/ +abbrev k : Fin 14 := 1 +/-- The root found by the search. -/ +abbrev lo : Fin 14 := 2 +/-- Upper end of the search interval. -/ +abbrev hi : Fin 14 := 3 +/-- Probe of the search. -/ +abbrev mid : Fin 14 := 4 +/-- Capped power. -/ +abbrev p : Fin 14 := 5 +/-- Counter of the capped power. -/ +abbrev i : Fin 14 := 6 +/-- Constant one. -/ +abbrev one : Fin 14 := 7 +/-- Constant two. -/ +abbrev two : Fin 14 := 8 +/-- Loop test of the search. -/ +abbrev t : Fin 14 := 9 +/-- `⌊log₂ n⌋ + 1`, the end of the exponent loop. -/ +abbrev lim : Fin 14 := 10 +/-- The answer: `1` for a perfect power. -/ +abbrev flag : Fin 14 := 11 +/-- Argument of the logarithm, consumed. -/ +abbrev larg : Fin 14 := 12 +/-- `⌊log₂ n⌋`. -/ +abbrev lres : Fin 14 := 13 + +/-- The root search runs on the first ten registers. -/ +def rootMap : Fin 10 → Fin 14 := Fin.castLE (by norm_num) + +/-- The final capped power raises the root `lo` to the exponent `k`. -/ +def capMap : Fin 6 → Fin 14 := ![lo, k, n, p, i, one] + +/-- The logarithm runs on `larg`, `lres` and the shared constants. -/ +def logMap : Fin 4 → Fin 14 := ![larg, lres, one, two] + +theorem rootMap_injective : rootMap.Injective := Fin.castLE_injective _ +theorem capMap_injective : capMap.Injective := by decide +theorem logMap_injective : logMap.Injective := by decide + +/-- Test one exponent: take the root, raise it back, and record equality with `n`. -/ +def body : Code .standard 14 := + .seq (NatNthRoot.program.rename rootMap) (.seq (NatNthRoot.CapPow.program.rename capMap) + (ram[ if p = n then [ flag ← one ] else [ flag ← flag ]; k ← k + one ])) + +/-- Compute `⌊log₂ n⌋`, then test every exponent `2 ≤ b ≤ ⌊log₂ n⌋`. -/ +def program : Code .standard 14 := + .seq (ram[ larg ← n ]) (.seq (NatLog.program.rename logMap) + (.seq (ram[ lim ← lres + one; flag ← #0; k ← #2 ]) (.loop (.lt k lim) body))) + +/-- The instruction bound of `program`, which is `O(log³ n)`. -/ +def time (n : ℕ) : ℕ := + (Nat.log 2 n + 1) * (Nat.clog 2 (n + 1) * (6 * Nat.log 2 n + 13) + 6 * Nat.log 2 n + 22) + + 4 * Nat.log 2 n + 11 + +/-- **The perfect-power test.** For `n ≥ 2` and a word size holding `n² + 2n + 2`, the program +leaves `1` in `flag` if `n` is a perfect power and `0` otherwise, within `time n` instructions, +preserving `n` and memory. -/ +theorem exec {w : ℕ} (x : Data 14) (hn : 2 ≤ x.regs n) + (hfit : Fits w (x.regs n * x.regs n + 2 * x.regs n + 2)) : + ∃ y, ExecWithin w program x y (time (x.regs n)) ∧ + y.regs flag = (if Cslib.AKS.IsPerfectPower (x.regs n) then 1 else 0) ∧ + y.regs n = x.regs n ∧ y.mem = x.mem := by + classical + obtain ⟨N, hN⟩ : ∃ N, x.regs n = N := ⟨_, rfl⟩ + rw [hN] at hn hfit ⊢ + obtain ⟨L, hL⟩ : ∃ L, Nat.log 2 N = L := ⟨_, rfl⟩ + have hLN : L < N := hL ▸ Nat.log_lt_self 2 (by omega) + have hLpos : 1 ≤ L := hL ▸ Nat.le_log_of_pow_le one_lt_two (by omega) + have hfitN : Fits w N := hfit.of_le (by nlinarith) + have h2 : Fits w 2 := hfit.of_le (by nlinarith) + have h1 : Fits w 1 := hfit.of_le (by nlinarith) + let Good (b : ℕ) : Prop := ∃ c, 2 ≤ c ∧ c < b ∧ Nat.nthRoot c N ^ c = N + let I (y : Data 14) : Prop := y.regs n = N ∧ y.regs one = 1 ∧ y.regs two = 2 ∧ + y.regs lim = L + 1 ∧ 2 ≤ y.regs k ∧ y.regs k ≤ L + 1 ∧ + y.regs flag = (if Good (y.regs k) then 1 else 0) ∧ y.mem = x.mem + set T₀ := Nat.clog 2 (N + 1) * (6 * L + 13) + 6 * L + 18 with hT₀ + have hbody (y : Data 14) (hy : I y) (hc : (Cond.lt k lim).Holds y) : + ∃ z, ExecWithin w body y z T₀ ∧ I z ∧ L + 1 - z.regs k < L + 1 - y.regs k := by + obtain ⟨hyn, hyo, hyt, hyl, hk2, hkL, hyf, hym⟩ := hy + have hkL' : y.regs k ≤ L := by simp only [Cond.holds_lt, hyl] at hc; omega + obtain ⟨b, hb⟩ : ∃ b, y.regs k = b := ⟨_, rfl⟩ + rw [hb] at hk2 hkL hkL' hyf + have hbfit : Fits w b := hfitN.of_le (by omega) + -- The root search. + have hrn : rootMap NatNthRoot.n = n := rfl + have hrk : rootMap NatNthRoot.k = k := rfl + have hrlo : rootMap NatNthRoot.lo = lo := rfl + have hro : rootMap NatNthRoot.one = one := rfl + have hrt : rootMap NatNthRoot.two = two := rfl + obtain ⟨y₁, hrun₁, hy₁lo, hy₁n, hy₁k, hy₁o, hy₁t, hy₁m⟩ := + NatNthRoot.exec_rename rootMap rootMap_injective y + (by rw [hrk, hb]; omega) (by rw [hrk, hb]; exact hbfit) + (by rw [hrn, hyn]; exact hfit) + rw [hrn, hyn] at hy₁n hrun₁ + rw [hrk, hb] at hy₁k hrun₁ + rw [hrlo, hrk, hrn, hb, hyn] at hy₁lo + rw [hro] at hy₁o + rw [hrt] at hy₁t + have hpres₁ (j : Fin 14) (hj : ∀ a, rootMap a ≠ j) : y₁.regs j = y.regs j := + hrun₁.regs_eq (Code.preserves_rename_of_forall_ne _ _ _ hj) + set a := Nat.nthRoot b N with ha + have hb0 : b ≠ 0 := by omega + have hale : a ^ b ≤ N := (Nat.pow_nthRoot_le_iff).mpr (Or.inl hb0) + have ha1 : 1 ≤ a := (Nat.le_nthRoot_iff hb0).mpr (by rw [one_pow]; omega) + have haN : a ≤ N := (Nat.le_self_pow hb0 a).trans hale + -- The capped power of the root. + have hcm : capMap NatNthRoot.CapPow.m = lo := rfl + have hck : capMap NatNthRoot.CapPow.k = k := rfl + have hcn : capMap NatNthRoot.CapPow.n = n := rfl + have hcp : capMap NatNthRoot.CapPow.p = p := rfl + obtain ⟨y₂, hrun₂, -, hpeq, -, hy₂k, hy₂n, hy₂o, hy₂m⟩ := + NatNthRoot.CapPow.exec_rename capMap capMap_injective y₁ + (by rw [hcm, hy₁lo]; exact ha1) h1 (by rw [hck, hy₁k]; exact hbfit) + (by rw [hcn, hcm, hy₁n, hy₁lo]; exact hfit.of_le (by nlinarith)) + rw [hcm, hck, hcn, hcp, hy₁lo, hy₁k, hy₁n] at hpeq + rw [hck, hy₁k] at hy₂k hrun₂ + rw [hcn, hy₁n] at hy₂n + have hy₂p : y₂.regs p = a ^ b := hpeq hale + have hpres₂ (j : Fin 14) (hj : ∀ a, capMap a ≠ j) : y₂.regs j = y₁.regs j := + hrun₂.regs_eq (Code.preserves_rename_of_forall_ne _ _ _ hj) + have hco : capMap NatNthRoot.CapPow.one = one := rfl + have hy₂o' : y₂.regs one = 1 := by rwa [hco] at hy₂o + have hy₂f : y₂.regs flag = y.regs flag := + (hpres₂ flag (by decide)).trans (hpres₁ flag (by decide)) + -- Record the comparison and advance the exponent. + have hinc : Fits w (b + 1) := hfitN.of_le (by omega) + let v := if a ^ b = N then 1 else y.regs flag + let z := (y₂.set flag v).set k (b + 1) + have hrun₃ : ExecWithin (B := .standard) w + (ram[ if p = n then [ flag ← one ] else [ flag ← flag ]; k ← k + one ]) y₂ z 4 := by + by_cases heq : a ^ b = N + · have hc' : (Cond.eq p n).Holds y₂ := by + simp only [Cond.holds_eq, hy₂p, hy₂n, heq] + refine ((ExecWithin.branch_pos (T := 1) hc' ExecWithin.copy).seq (T₂ := 1) ?_).mono + (by norm_num) + ram_eval [z, v, heq, hy₂o', hy₂k, hinc] + · have hc' : ¬(Cond.eq p n).Holds y₂ := by + simp only [Cond.holds_eq, hy₂p, hy₂n]; exact heq + refine ((ExecWithin.branch_neg (T := 1) hc' ExecWithin.copy).seq (T₂ := 1) ?_).mono + (by norm_num) + ram_eval [z, v, heq, hy₂o', hy₂k, hy₂f, hinc] + have hgood : Good (b + 1) ↔ Good b ∨ a ^ b = N := by + constructor + · rintro ⟨c, hc2, hcb, hc⟩ + rcases Nat.lt_succ_iff_lt_or_eq.mp hcb with hcb | rfl + · exact Or.inl ⟨c, hc2, hcb, hc⟩ + · exact Or.inr hc + · rintro (⟨c, hc2, hcb, hc⟩ | h) + · exact ⟨c, hc2, by omega, hc⟩ + · exact ⟨b, hk2, by omega, h⟩ + refine ⟨z, ?_, ?_, ?_⟩ + · refine (hrun₁.seq (hrun₂.seq hrun₃)).mono ?_ + have : Nat.clog 2 (N + 1) * (6 * b + 13) ≤ Nat.clog 2 (N + 1) * (6 * L + 13) := + Nat.mul_le_mul_left _ (by omega) + omega + · refine ⟨?_, ?_, ?_, ?_, ?_, ?_, ?_, hy₂m.trans (hy₁m.trans hym)⟩ + · simp +decide [z, hy₂n] + · simp +decide [z, hy₂o'] + · have := (hpres₂ two (by decide)).trans hy₁t + simp +decide [z, this] + · have := (hpres₂ lim (by decide)).trans ((hpres₁ lim (by decide)).trans hyl) + simp +decide [z, this] + · simp +decide [z]; omega + · simp +decide [z]; omega + · simp only [z, v, Data.set_regs, Function.update_self] + rw [Function.update_of_ne (by decide), Function.update_self, hyf] + by_cases heq : a ^ b = N + · simp [heq, hgood] + · have : Good (b + 1) ↔ Good b := by simp [hgood, heq] + simp [heq, this] + · simp [z, hb]; omega + -- The logarithm and the initialization. + have hlarg : Fits w N := hfitN + let x₁ := x.set larg N + have hrun₀ : ExecWithin (B := .standard) w (ram[ larg ← n ]) x x₁ 1 := by + simpa [x₁, hN] using ExecWithin.copy (B := .standard) (w := w) (d := larg) (s := n) (x := x) + have hla : logMap NatLog.arg = larg := rfl + have hlr : logMap NatLog.result = lres := rfl + have hlo : logMap NatLog.one = one := rfl + have hlt : logMap NatLog.two = two := rfl + obtain ⟨y₁, hrun₁, hy₁r, hy₁m, hy₁o, hy₁t⟩ := + NatLog.exec_rename logMap logMap_injective x₁ h2 (by simpa [hla, x₁] using hlarg) + (by simp only [hla, x₁, Data.set_regs, Function.update_self, hL] + exact hfitN.of_le (by omega)) + have hx₁ : x₁.regs larg = N := by simp [x₁] + rw [hla, hx₁, hL] at hrun₁ hy₁r + rw [hlr] at hy₁r + rw [hlo] at hy₁o + rw [hlt] at hy₁t + have hpres₁ (j : Fin 14) (hj : ∀ a, logMap a ≠ j) : y₁.regs j = x₁.regs j := + hrun₁.regs_eq (Code.preserves_rename_of_forall_ne _ _ _ hj) + have hy₁n : y₁.regs n = N := by + rw [hpres₁ n (by decide)]; simp +decide [x₁, hN] + have hL1 : Fits w (L + 1) := hfitN.of_le (by omega) + let x₂ := ((y₁.set lim (L + 1)).set flag 0).set k 2 + have hrun₂ : ExecWithin (B := .standard) w (ram[ lim ← lres + one; flag ← #0; k ← #2 ]) + y₁ x₂ 3 := by + ram_eval [x₂, hy₁r, hy₁o, hL1, h2] + have hx₂ : I x₂ := by + have hgood : ¬Good 2 := fun ⟨c, hc2, hc, _⟩ => by omega + refine ⟨?_, ?_, ?_, ?_, ?_, ?_, ?_, hy₁m⟩ <;> + simp +decide [x₂, hy₁n, hy₁o, hy₁t, hLpos, hgood] + obtain ⟨y, hrun₃, hy, hc⟩ := ExecWithin.loop_invariant I (fun y => L + 1 - y.regs k) hbody hx₂ + obtain ⟨hyn, -, -, hyl, hk2, hkL, hyf, hym⟩ := hy + have hkend : L + 1 ≤ y.regs k := by simp only [Cond.holds_lt, hyl] at hc; omega + refine ⟨y, ?_, ?_, hyn, hym⟩ + · refine (hrun₀.seq (hrun₁.seq (hrun₂.seq hrun₃))).mono ?_ + have hx₂k : L + 1 - x₂.regs k ≤ L + 1 := by omega + have := Nat.mul_le_mul_right (T₀ + 2) hx₂k + simp only [time, hL, hT₀] at this ⊢ + nlinarith + · have hiff : Good (y.regs k) ↔ Cslib.AKS.IsPerfectPower N := by + rw [Cslib.AKS.isPerfectPower_iff_nthRoot hn] + constructor + · rintro ⟨c, hc2, hck, hc⟩; exact ⟨c, by omega, hc2, hc⟩ + · rintro ⟨c, hcL, hc2, hc⟩; exact ⟨c, hc2, by omega, hc⟩ + rw [hyf] + exact if_congr hiff rfl rfl + +/-- Run the perfect-power test on any injectively chosen registers. -/ +theorem exec_rename {r w : ℕ} (f : Fin 14 → Fin r) (hf : f.Injective) (x : Data r) + (hn : 2 ≤ x.regs (f n)) (hfit : Fits w (x.regs (f n) * x.regs (f n) + 2 * x.regs (f n) + 2)) : + ∃ y, ExecWithin w (program.rename f) x y (time (x.regs (f n))) ∧ + y.regs (f flag) = (if Cslib.AKS.IsPerfectPower (x.regs (f n)) then 1 else 0) ∧ + y.regs (f n) = x.regs (f n) ∧ y.mem = x.mem := by + obtain ⟨y, hy, hfl, hyn, hym⟩ := exec (x.comap f) hn hfit + refine ⟨y.extend f x, hy.rename hf, ?_, ?_, hym⟩ + · exact (hf.extend_apply _ _ flag).trans hfl + · exact (hf.extend_apply _ _ n).trans hyn + +end Cslib.Algorithms.WordRAM.AKS.PerfectPower diff --git a/Complexity/Algorithms/WordRAM/AKS/PolyMul.lean b/Complexity/Algorithms/WordRAM/AKS/PolyMul.lean new file mode 100644 index 000000000..453fff1b1 --- /dev/null +++ b/Complexity/Algorithms/WordRAM/AKS/PolyMul.lean @@ -0,0 +1,364 @@ +/- +Copyright (c) 2026 Samuel Schlesinger. All rights reserved. +Released under Apache 2.0 license as described in the file LICENSE. +Authors: Samuel Schlesinger +-/ + +module +-- upstream: candidate + +public import Complexity.Computability.WordRAM.Notation +public import Complexity.Foundations.NumberTheory.AKS.Coefficients +import Mathlib.Tactic.Linarith + +/-! +# Multiplication modulo `X^r - 1` and `n` on the word RAM + +Polynomials modulo `X^r - 1` with coefficients modulo `n` are stored as arrays of `r` words. The +program clears the destination array and then, for every pair `i, j < r`, adds `p i * q j` into +cell `(i + j) mod r`, reducing modulo `n` after each addition: schoolbook multiplication, with +`O(r²)` instructions. Every product and sum stays below `n² + n`. The result has the coefficients +`Cslib.AKS.mulCoeff n r p q`, so its class is the product of the classes of the factors +(`Cslib.AKS.toQuot_mulCoeff`). The factors may coincide, which is how a polynomial is squared; the +destination must be disjoint from both. + +## Main definitions + +* `Cslib.Algorithms.WordRAM.AKS.PolyMul.program`: the multiplication program. + +## Main results + +* `Cslib.Algorithms.WordRAM.AKS.PolyMul.exec`: the destination receives `Cslib.AKS.mulCoeff` + within `14 (r + 1)²` instructions, and no other memory cell changes. + +## References + +* [M. Agrawal, N. Kayal, N. Saxena, *PRIMES is in P*][AKS2004] +-/ + +@[expose] public section + +namespace Cslib.Algorithms.WordRAM.AKS.PolyMul + +open Cslib.WordRAM Finset + +/-- Base address of the first factor. -/ +abbrev pa : Fin 14 := 0 +/-- Base address of the second factor. -/ +abbrev pb : Fin 14 := 1 +/-- Base address of the product. -/ +abbrev pt : Fin 14 := 2 +/-- The number `r` of coefficients. -/ +abbrev rr : Fin 14 := 3 +/-- The modulus `n` of the coefficients. -/ +abbrev nn : Fin 14 := 4 +/-- Index into the first factor. -/ +abbrev i : Fin 14 := 5 +/-- Index into the second factor. -/ +abbrev j : Fin 14 := 6 +/-- Index `(i + j) mod r` into the product. -/ +abbrev k : Fin 14 := 7 +/-- Address scratch. -/ +abbrev addr : Fin 14 := 8 +/-- The coefficient `p i`. -/ +abbrev u : Fin 14 := 9 +/-- The product `p i * q j`. -/ +abbrev v : Fin 14 := 10 +/-- The updated product coefficient. -/ +abbrev s : Fin 14 := 11 +/-- Constant one. -/ +abbrev one : Fin 14 := 12 +/-- Constant zero. -/ +abbrev zero : Fin 14 := 13 + +/-- Clear one cell of the product. -/ +def clearBody : Code .standard 14 := ram[ addr ← pt + i; *addr ← zero; i ← i + one ] + +/-- Form `p i * q j` and the index `(i + j) mod r`. -/ +def innerProd : Code .standard 14 := ram[ + addr ← pb + j; v ← *addr; v ← u * v; k ← i + j; k ← k % rr ] + +/-- Add the product into its cell, reducing modulo `n`, and advance `j`. -/ +def innerAcc : Code .standard 14 := ram[ + addr ← pt + k; s ← *addr; s ← s + v; s ← s % nn; *addr ← s; j ← j + one ] + +/-- Add `p i * q j` into cell `(i + j) mod r` of the product. -/ +def innerBody : Code .standard 14 := .seq innerProd innerAcc + +/-- Add the products of `p i` with every coefficient of `q`. -/ +def outerBody : Code .standard 14 := + .seq (ram[ addr ← pa + i; u ← *addr; j ← #0 ]) + (.seq (.loop (.lt j rr) innerBody) (ram[ i ← i + one ])) + +/-- Clear the product, then accumulate every pair of coefficients. -/ +def program : Code .standard 14 := + .seq (ram[ one ← #1; zero ← #0; i ← #0 ]) (.seq (.loop (.lt i rr) clearBody) + (.seq (ram[ i ← #0 ]) (.loop (.lt i rr) outerBody))) + +/-- Memory holding `f` in the block `[t, t + len)` and agreeing with `m` elsewhere. -/ +private def view (m : ℕ → ℕ) (t len : ℕ) (f : ℕ → ℕ) (a : ℕ) : ℕ := + if t ≤ a ∧ a < t + len then f (a - t) else m a + +private theorem view_of_not_mem {m f : ℕ → ℕ} {t len a : ℕ} (h : ¬(t ≤ a ∧ a < t + len)) : + view m t len f a = m a := by + simp only [view, h, ↓reduceIte] + +private theorem view_add {m f : ℕ → ℕ} {t len k₀ : ℕ} (h : k₀ < len) : + view m t len f (t + k₀) = f k₀ := by + simp only [view] + split_ifs with h' + · rw [Nat.add_sub_cancel_left] + · omega + +private theorem view_update {m f : ℕ → ℕ} {t len k₀ : ℕ} (hk₀ : k₀ < len) (val : ℕ) : + Function.update (view m t len f) (t + k₀) val = view m t len (Function.update f k₀ val) := by + funext a + by_cases ha : a = t + k₀ + · subst ha; simp [view, hk₀] + · rw [Function.update_of_ne ha] + simp only [view] + split_ifs with h + · rw [Function.update_of_ne (by omega)] + · rfl + +/-- The partial row sum `∑_{j' < j, (i + j') mod r = k} p i q j'`. -/ +private def rowSum (p q : ℕ → ℕ) (r i j k : ℕ) : ℕ := + ∑ j' ∈ range j, if (i + j') % r = k then p i * q j' else 0 + +/-- The sum over all complete rows `i' < i` and the first `j` entries of row `i`. -/ +private def partSum (p q : ℕ → ℕ) (r i j k : ℕ) : ℕ := + ∑ i' ∈ range i, rowSum p q r i' r k + rowSum p q r i j k + +private theorem partSum_succ (p q : ℕ → ℕ) (r i j k : ℕ) : + partSum p q r i (j + 1) k = + partSum p q r i j k + if (i + j) % r = k then p i * q j else 0 := by + simp only [partSum, rowSum, sum_range_succ]; ring + +private theorem partSum_row (p q : ℕ → ℕ) (r i k : ℕ) : + partSum p q r i r k = partSum p q r (i + 1) 0 k := by + simp [partSum, rowSum, sum_range_succ] + +/-- The registers that every phase preserves. -/ +private def Regs (a b t r n : ℕ) (y : Data 14) : Prop := y.regs pa = a ∧ y.regs pb = b ∧ + y.regs pt = t ∧ y.regs rr = r ∧ y.regs nn = n ∧ y.regs one = 1 ∧ y.regs zero = 0 + +/-- The block hypotheses of `exec`, with `m` the original memory. -/ +private structure Hyp (w a b t r n : ℕ) (m : ℕ → ℕ) : Prop where + hn0 : 0 < n + hta : t + r ≤ a ∨ a + r ≤ t + htb : t + r ≤ b ∨ b + r ≤ t + hpa : ∀ i' < r, m (a + i') < n + hpb : ∀ j' < r, m (b + j') < n + hfa : Fits w (a + r) + hfb : Fits w (b + r) + hft : Fits w (t + r) + hf2 : Fits w (2 * r) + hfn : Fits w (n * n + n) + +private theorem clear_step {w a b t r n : ℕ} {m : ℕ → ℕ} (H : Hyp w a b t r n m) (y : Data 14) + (hy : Regs a b t r n y ∧ y.regs i ≤ r ∧ y.mem = view m t (y.regs i) (fun _ => 0)) + (hc : (Cond.lt i rr).Holds y) : + ∃ z, ExecWithin w clearBody y z 3 ∧ + (Regs a b t r n z ∧ z.regs i ≤ r ∧ z.mem = view m t (z.regs i) (fun _ => 0)) ∧ + r - z.regs i < r - y.regs i := by + obtain ⟨⟨hya, hyb, hyt, hyr, hyn, hyo, hyz⟩, hyi, hym⟩ := hy + obtain ⟨I, hI⟩ : ∃ I, y.regs i = I := ⟨_, rfl⟩ + have hlt : I < r := by simpa only [Cond.holds_lt, hyr, hI] using hc + have hfaddr : Fits w (t + I) := H.hft.of_le (by omega) + have hfinc : Fits w (I + 1) := H.hf2.of_le (by omega) + let z := ((y.set addr (t + I)).write (t + I) 0).set i (I + 1) + have hrun : ExecWithin (B := .standard) w clearBody y z 3 := by + ram_eval [clearBody, z, hyt, hI, hyo, hyz, hfaddr, hfinc] + refine ⟨z, hrun, ⟨⟨?_, ?_, ?_, ?_, ?_, ?_, ?_⟩, ?_, ?_⟩, ?_⟩ <;> + simp +decide only [z, Data.set_regs, Data.write_regs, Data.set_mem, Data.write_mem, + Function.update_apply, ↓reduceIte, hya, hyb, hyt, hyr, hyn, hyo, hyz, hI] + · omega + · rw [hym, hI] + funext a' + simp only [Function.update_apply, view] + split_ifs <;> first | rfl | omega + · omega + +private theorem inner_step {w a b t r n I : ℕ} {m : ℕ → ℕ} (H : Hyp w a b t r n m) (hIr : I < r) + (y : Data 14) + (hy : Regs a b t r n y ∧ y.regs i = I ∧ y.regs u = m (a + I) ∧ y.regs j ≤ r ∧ + y.mem = view m t r + (fun k' => partSum (fun i' => m (a + i')) (fun j' => m (b + j')) r I (y.regs j) k' % n)) + (hc : (Cond.lt j rr).Holds y) : + ∃ z, ExecWithin w innerBody y z 11 ∧ + (Regs a b t r n z ∧ z.regs i = I ∧ z.regs u = m (a + I) ∧ z.regs j ≤ r ∧ + z.mem = view m t r + (fun k' => partSum (fun i' => m (a + i')) (fun j' => m (b + j')) r I (z.regs j) k' % n)) ∧ + r - z.regs j < r - y.regs j := by + obtain ⟨⟨hya, hyb, hyt, hyr, hyn, hyo, hyz⟩, hyi, hyu, hyj, hym⟩ := hy + obtain ⟨J, hJ⟩ : ∃ J, y.regs j = J := ⟨_, rfl⟩ + have hlt : J < r := by simpa only [Cond.holds_lt, hyr, hJ] using hc + rw [hJ] at hym + have hr0 : 0 < r := by omega + obtain ⟨k₀, hk₀⟩ : ∃ k₀, (I + J) % r = k₀ := ⟨_, rfl⟩ + have hk₀r : k₀ < r := hk₀ ▸ Nat.mod_lt _ hr0 + obtain ⟨p, hp⟩ : ∃ p, m (a + I) = p := ⟨_, rfl⟩ + obtain ⟨q, hq⟩ : ∃ q, m (b + J) = q := ⟨_, rfl⟩ + rw [hp] at hyu + have hQJ : y.mem (b + J) = q := by + rw [hym, view_of_not_mem (by rcases H.htb with h | h <;> omega), hq] + obtain ⟨S, hS⟩ : ∃ S, partSum (fun i' => m (a + i')) (fun j' => m (b + j')) r I J k₀ = S := + ⟨_, rfl⟩ + have hSm : y.mem (t + k₀) = S % n := by rw [hym, view_add hk₀r, hS] + have hPn : p < n := hp ▸ H.hpa I hIr + have hQn : q < n := hq ▸ H.hpb J hlt + have hprod : p * q ≤ n * n := Nat.mul_le_mul hPn.le hQn.le + have hSn : S % n < n := Nat.mod_lt _ H.hn0 + have hf₁ : Fits w (b + J) := H.hfb.of_le (by omega) + have hf₂ : Fits w (p * q) := H.hfn.of_le (by omega) + have hf₃ : Fits w (I + J) := H.hf2.of_le (by omega) + have hf₄ : Fits w k₀ := hf₃.of_le (hk₀ ▸ Nat.mod_le _ _) + have hf₅ : Fits w (t + k₀) := H.hft.of_le (by omega) + have hf₆ : Fits w (S % n + p * q) := H.hfn.of_le (by omega) + have hf₇ : Fits w ((S % n + p * q) % n) := H.hfn.of_le + (by have := Nat.mod_lt (S % n + p * q) H.hn0; nlinarith) + have hf₈ : Fits w (J + 1) := H.hf2.of_le (by omega) + let y₁ := ((((y.set addr (b + J)).set v q).set v (p * q)).set k (I + J)).set k k₀ + have hrun₁ : ExecWithin (B := .standard) w innerProd y y₁ 5 := by + ram_eval [innerProd, y₁, hyb, hJ, hyu, hyi, hyr, hQJ, hk₀, hf₁, hf₂, hf₃, hf₄] + have hy₁ (j' : Fin 14) (hj' : j' ≠ addr ∧ j' ≠ v ∧ j' ≠ k) : y₁.regs j' = y.regs j' := by + simp [y₁, hj'.1, hj'.2.1, hj'.2.2] + have hy₁v : y₁.regs v = p * q := by simp +decide [y₁] + have hy₁k : y₁.regs k = k₀ := by simp [y₁] + have hy₁m : y₁.mem = y.mem := rfl + let z := (((((y₁.set addr (t + k₀)).set s (S % n)).set s (S % n + p * q)).set s + ((S % n + p * q) % n)).write (t + k₀) ((S % n + p * q) % n)).set j (J + 1) + have hrun₂ : ExecWithin (B := .standard) w innerAcc y₁ z 6 := by + ram_eval [innerAcc, z, hy₁ pt (by decide), hy₁ s (by decide), hy₁ nn (by decide), + hy₁ j (by decide), hy₁ one (by decide), hy₁v, hy₁k, hy₁m, hyt, hyn, hJ, hyo, hSm, + hf₅, hf₆, hf₇, hf₈] + have hz (j' : Fin 14) (hj' : j' ≠ addr ∧ j' ≠ v ∧ j' ≠ k ∧ j' ≠ s ∧ j' ≠ j) : + z.regs j' = y.regs j' := by + simp [z, hj'.1, hj'.2.1, hj'.2.2.1, hj'.2.2.2.1, hj'.2.2.2.2, y₁] + have hzj : z.regs j = J + 1 := by simp [z] + refine ⟨z, (hrun₁.seq hrun₂).mono (by omega), ⟨⟨?_, ?_, ?_, ?_, ?_, ?_, ?_⟩, ?_, ?_, ?_, ?_⟩, ?_⟩ + · rw [hz pa (by decide), hya] + · rw [hz pb (by decide), hyb] + · rw [hz pt (by decide), hyt] + · rw [hz rr (by decide), hyr] + · rw [hz nn (by decide), hyn] + · rw [hz one (by decide), hyo] + · rw [hz zero (by decide), hyz] + · rw [hz i (by decide), hyi] + · rw [hz u (by decide), hyu, hp] + · rw [hzj]; omega + · change Function.update y.mem (t + k₀) ((S % n + p * q) % n) = _ + rw [hzj, hym, view_update hk₀r] + congr 1 + funext k' + rw [partSum_succ, Function.update_apply, hp, hq] + by_cases h : k' = k₀ + · subst h; simp [hk₀, hS, Nat.mod_add_mod] + · have h' : ¬(I + J) % r = k' := fun h' => h (hk₀ ▸ h'.symm) + simp only [h, h', ↓reduceIte, Nat.add_zero] + · rw [hzj, hJ]; omega + +private theorem outer_step {w a b t r n : ℕ} {m : ℕ → ℕ} (H : Hyp w a b t r n m) (y : Data 14) + (hy : Regs a b t r n y ∧ y.regs i ≤ r ∧ y.mem = view m t r + (fun k' => partSum (fun i' => m (a + i')) (fun j' => m (b + j')) r (y.regs i) 0 k' % n)) + (hc : (Cond.lt i rr).Holds y) : + ∃ z, ExecWithin w outerBody y z (13 * r + 6) ∧ + (Regs a b t r n z ∧ z.regs i ≤ r ∧ z.mem = view m t r + (fun k' => partSum (fun i' => m (a + i')) (fun j' => m (b + j')) r (z.regs i) 0 k' % n)) ∧ + r - z.regs i < r - y.regs i := by + obtain ⟨⟨hya, hyb, hyt, hyr, hyn, hyo, hyz⟩, hyi, hym⟩ := hy + obtain ⟨I, hI⟩ : ∃ I, y.regs i = I := ⟨_, rfl⟩ + have hlt : I < r := by simpa only [Cond.holds_lt, hyr, hI] using hc + rw [hI] at hym + have hfaddr : Fits w (a + I) := H.hfa.of_le (by omega) + have hPI : y.mem (a + I) = m (a + I) := by + rw [hym, view_of_not_mem (by rcases H.hta with h | h <;> omega)] + let y₁ := ((y.set addr (a + I)).set u (m (a + I))).set j 0 + have hrun₁ : ExecWithin (B := .standard) w (ram[ addr ← pa + i; u ← *addr; j ← #0 ]) y y₁ 3 := by + ram_eval [y₁, hya, hI, hfaddr, hPI] + have hy₁ : Regs a b t r n y₁ ∧ y₁.regs i = I ∧ y₁.regs u = m (a + I) ∧ y₁.regs j ≤ r ∧ + y₁.mem = view m t r + (fun k' => partSum (fun i' => m (a + i')) (fun j' => m (b + j')) r I (y₁.regs j) k' % + n) := by + refine ⟨⟨?_, ?_, ?_, ?_, ?_, ?_, ?_⟩, ?_, ?_, ?_, ?mem⟩ + case mem => simpa [y₁] using hym + all_goals simp +decide [y₁, hya, hyb, hyt, hyr, hyn, hyo, hyz, hI] + obtain ⟨y₂, hrun₂, hy₂, hc₂⟩ := ExecWithin.loop_invariant _ (fun y => r - y.regs j) + (inner_step H hlt) hy₁ + obtain ⟨⟨hya₂, hyb₂, hyt₂, hyr₂, hyn₂, hyo₂, hyz₂⟩, hyi₂, -, hyj₂, hym₂⟩ := hy₂ + have hJr : y₂.regs j = r := by simp only [Cond.holds_lt, hyr₂] at hc₂; omega + have hf₉ : Fits w (I + 1) := H.hf2.of_le (by omega) + have hrun₃ : ExecWithin (B := .standard) w (ram[ i ← i + one ]) y₂ (y₂.set i (I + 1)) 1 := by + ram_eval [hyi₂, hyo₂, hf₉] + refine ⟨y₂.set i (I + 1), ?_, ?_, ?_⟩ + · refine (hrun₁.seq (hrun₂.seq hrun₃)).mono ?_ + have : y₁.regs j = 0 := by simp [y₁] + rw [this]; omega + · refine ⟨⟨?_, ?_, ?_, ?_, ?_, ?_, ?_⟩, ?le, ?mem⟩ + case le => simp only [Data.set_regs, Function.update_self]; omega + case mem => + simp only [Data.set_mem, Data.set_regs, Function.update_self] + rw [hym₂, hJr] + congr 1 + funext k' + rw [partSum_row] + all_goals simp +decide [hya₂, hyb₂, hyt₂, hyr₂, hyn₂, hyo₂, hyz₂] + · simp; omega + +/-- **Multiplication modulo `X^r - 1` and `n`.** If the product block `[t, t + r)` is disjoint +from both factor blocks, the factors' coefficients are below `n ≥ 1`, and the word size holds the +block ends, `2 r` and `n² + n`, then the product block receives the coefficients +`Cslib.AKS.mulCoeff n r p q` within `14 (r + 1)²` instructions, and no other cell changes. -/ +theorem exec {w : ℕ} (x : Data 14) {a b t r n : ℕ} (ha : x.regs pa = a) (hb : x.regs pb = b) + (ht : x.regs pt = t) (hr : x.regs rr = r) (hn : x.regs nn = n) (hn0 : 0 < n) + (hta : t + r ≤ a ∨ a + r ≤ t) (htb : t + r ≤ b ∨ b + r ≤ t) + (hpa : ∀ i' < r, x.mem (a + i') < n) (hpb : ∀ j' < r, x.mem (b + j') < n) + (hfa : Fits w (a + r)) (hfb : Fits w (b + r)) (hft : Fits w (t + r)) (hf2 : Fits w (2 * r)) + (hfn : Fits w (n * n + n)) : + ∃ y, ExecWithin w program x y (14 * (r + 1) * (r + 1)) ∧ + (∀ k' < r, y.mem (t + k') = + Cslib.AKS.mulCoeff n r (fun i' => x.mem (a + i')) (fun j' => x.mem (b + j')) k') ∧ + (∀ a', ¬(t ≤ a' ∧ a' < t + r) → y.mem a' = x.mem a') ∧ + y.regs pa = a ∧ y.regs pb = b ∧ y.regs pt = t ∧ y.regs rr = r ∧ y.regs nn = n := by + have H : Hyp w a b t r n x.mem := ⟨hn0, hta, htb, hpa, hpb, hfa, hfb, hft, hf2, hfn⟩ + have h1 : Fits w 1 := hfn.of_le (by nlinarith) + let x₀ := ((x.set one 1).set zero 0).set i 0 + have hrun₀ : ExecWithin (B := .standard) w (ram[ one ← #1; zero ← #0; i ← #0 ]) x x₀ 3 := by + ram_eval [x₀, h1] + have hx₀ : Regs a b t r n x₀ ∧ x₀.regs i ≤ r ∧ + x₀.mem = view x.mem t (x₀.regs i) (fun _ => 0) := by + refine ⟨⟨?_, ?_, ?_, ?_, ?_, ?_, ?_⟩, ?_, ?mem⟩ + case mem => funext a'; simp [x₀, view] + all_goals simp +decide [x₀, ha, hb, ht, hr, hn] + obtain ⟨y₁, hrun₁, hy₁, hc₁⟩ := ExecWithin.loop_invariant _ (fun y => r - y.regs i) + (clear_step H) hx₀ + obtain ⟨⟨hya, hyb, hyt, hyr, hyn, hyo, hyz⟩, hyi, hym⟩ := hy₁ + have hir : y₁.regs i = r := by simp only [Cond.holds_lt, hyr] at hc₁; omega + have hrun₂ : ExecWithin (B := .standard) w (ram[ i ← #0 ]) y₁ (y₁.set i 0) 1 := + ExecWithin.const (fits_zero w) + have hy₂ : Regs a b t r n (y₁.set i 0) ∧ (y₁.set i 0).regs i ≤ r ∧ (y₁.set i 0).mem = + view x.mem t r (fun k' => partSum (fun i' => x.mem (a + i')) (fun j' => x.mem (b + j')) r + ((y₁.set i 0).regs i) 0 k' % n) := by + refine ⟨⟨?_, ?_, ?_, ?_, ?_, ?_, ?_⟩, ?_, ?mem⟩ + case mem => + simp only [Data.set_mem, Data.set_regs, Function.update_self] + rw [hym, hir] + congr 1 + all_goals simp +decide [hya, hyb, hyt, hyr, hyn, hyo, hyz] + obtain ⟨y₃, hrun₃, hy₃, hc₃⟩ := ExecWithin.loop_invariant _ (fun y => r - y.regs i) + (outer_step H) hy₂ + obtain ⟨⟨hya₃, hyb₃, hyt₃, hyr₃, hyn₃, -, -⟩, hyi₃, hym₃⟩ := hy₃ + have hir₃ : y₃.regs i = r := by simp only [Cond.holds_lt, hyr₃] at hc₃; omega + refine ⟨y₃, ?_, ?_, ?_, hya₃, hyb₃, hyt₃, hyr₃, hyn₃⟩ + · refine (hrun₀.seq (hrun₁.seq (hrun₂.seq hrun₃))).mono ?_ + have h0 : x₀.regs i = 0 := by simp [x₀] + have h0' : (y₁.set i 0).regs i = 0 := by simp + rw [h0, h0', Nat.sub_zero] + nlinarith + · intro k' hk' + rw [hym₃, hir₃, view_add hk'] + simp only [Cslib.AKS.mulCoeff, partSum, rowSum, sum_range_zero, add_zero] + · intro a' ha' + rw [hym₃, view_of_not_mem ha'] + +end Cslib.Algorithms.WordRAM.AKS.PolyMul diff --git a/Complexity/Algorithms/WordRAM/NatGcd.lean b/Complexity/Algorithms/WordRAM/NatGcd.lean new file mode 100644 index 000000000..a2c666c6a --- /dev/null +++ b/Complexity/Algorithms/WordRAM/NatGcd.lean @@ -0,0 +1,96 @@ +/- +Copyright (c) 2026 Samuel Schlesinger. All rights reserved. +Released under Apache 2.0 license as described in the file LICENSE. +Authors: Samuel Schlesinger +-/ + +module +-- upstream: candidate + +public import Complexity.Computability.WordRAM.Notation +public import Mathlib.Data.Nat.GCD.Basic + +/-! +# A charged native greatest common divisor + +Euclid's algorithm with one native remainder per round. The second argument strictly decreases, +so the number of rounds is at most that argument; this crude bound suffices for callers whose +second argument is itself polynomially bounded. + +## Main definitions + +* `NatGcd.program`: Euclid's algorithm on two registers. + +## Main results + +* `NatGcd.exec`: `Nat.gcd a b` within `5 b + 3` instructions, preserving memory. +* `NatGcd.exec_rename`: the same on injectively chosen registers. +-/ + +@[expose] public section + +namespace Cslib.Algorithms.WordRAM.NatGcd + +open Cslib.WordRAM + +/-- First argument, and the output. -/ +abbrev a : Fin 4 := 0 +/-- Second argument; zero on exit. -/ +abbrev b : Fin 4 := 1 +/-- Remainder scratch. -/ +abbrev t : Fin 4 := 2 +/-- Constant zero. -/ +abbrev zero : Fin 4 := 3 + +/-- One round of Euclid's algorithm: `(a, b) ↦ (b, a mod b)`. -/ +def body : Code .standard 4 := ram[ t ← a % b; a ← b; b ← t ] + +/-- Euclid's algorithm, looping while the second argument is positive. -/ +def program : Code .standard 4 := ram[ zero ← #0; while zero < b do [ ⟪body⟫ ] ] + +/-- Compute `Nat.gcd a b` within `5 b + 3` instructions, preserving memory. Every register value +stays below `max a b`, so no fit hypothesis beyond the arguments is needed. -/ +theorem exec {w : ℕ} (x : Data 4) (ha : Fits w (x.regs a)) (hb : Fits w (x.regs b)) : + ∃ y, ExecWithin w program x y (5 * x.regs b + 3) ∧ + y.regs a = Nat.gcd (x.regs a) (x.regs b) ∧ y.mem = x.mem := by + let g := Nat.gcd (x.regs a) (x.regs b) + let I (y : Data 4) : Prop := y.regs zero = 0 ∧ Nat.gcd (y.regs a) (y.regs b) = g ∧ + Fits w (y.regs a) ∧ Fits w (y.regs b) ∧ y.mem = x.mem + have hbody (y : Data 4) (hy : I y) (hc : (Cond.lt zero b).Holds y) : + ∃ z, ExecWithin w body y z 3 ∧ I z ∧ z.regs b < y.regs b := by + obtain ⟨hz, hg, hfa, hfb, hm⟩ := hy + have hpos : 0 < y.regs b := by simpa only [Cond.holds_lt, hz] using hc + have hmod : Fits w (y.regs a % y.regs b) := hfb.of_le (Nat.mod_lt _ hpos).le + refine ⟨?z, ?run, ?inv, ?dec⟩ + case run => ram_eval [body, hmod] + case inv => + refine ⟨?_, ?_, ?_, ?_, hm⟩ + · simpa +decide [Function.update_apply] using hz + · simp +decide only [Data.set_regs, Function.update_apply, ↓reduceIte] + rw [← hg, Nat.gcd_comm (y.regs b), ← Nat.gcd_rec, Nat.gcd_comm] + · simpa +decide [Function.update_apply] using hfb + · simpa +decide [Function.update_apply] using hmod + case dec => + simp +decide only [Data.set_regs, Function.update_apply, ↓reduceIte] + exact Nat.mod_lt _ hpos + let x₀ := x.set zero 0 + have hx₀ : I x₀ := by + refine ⟨by simp [x₀], rfl, ?_, ?_, rfl⟩ <;> simpa +decide [x₀, Function.update_apply] + obtain ⟨y, he, hy, hc⟩ := ExecWithin.loop_invariant I (fun y => y.regs b) hbody hx₀ + obtain ⟨hz, hg, -, -, hm⟩ := hy + have hb0 : y.regs b = 0 := by simp only [Cond.holds_lt, hz] at hc; omega + refine ⟨y, ?_, by rw [hb0, Nat.gcd_zero_right] at hg; exact hg, hm⟩ + have hinit : ExecWithin (B := .standard) w (ram[ zero ← #0 ]) x x₀ 1 := + ExecWithin.const (fits_zero w) + have hb₀ : x₀.regs b = x.regs b := by simp +decide [x₀] + exact (hinit.seq he).mono (by rw [hb₀]; omega) + +/-- Run Euclid's algorithm on any injectively chosen registers. -/ +theorem exec_rename {r w : ℕ} (f : Fin 4 → Fin r) (hf : f.Injective) (x : Data r) + (ha : Fits w (x.regs (f a))) (hb : Fits w (x.regs (f b))) : + ∃ y, ExecWithin w (program.rename f) x y (5 * x.regs (f b) + 3) ∧ + y.regs (f a) = Nat.gcd (x.regs (f a)) (x.regs (f b)) ∧ y.mem = x.mem := by + obtain ⟨y, hy, hya, hym⟩ := exec (x.comap f) ha hb + exact ⟨y.extend f x, hy.rename hf, (hf.extend_apply _ _ a).trans hya, hym⟩ + +end Cslib.Algorithms.WordRAM.NatGcd diff --git a/Complexity/Algorithms/WordRAM/NatNthRoot.lean b/Complexity/Algorithms/WordRAM/NatNthRoot.lean new file mode 100644 index 000000000..943d9eb06 --- /dev/null +++ b/Complexity/Algorithms/WordRAM/NatNthRoot.lean @@ -0,0 +1,324 @@ +/- +Copyright (c) 2026 Samuel Schlesinger. All rights reserved. +Released under Apache 2.0 license as described in the file LICENSE. +Authors: Samuel Schlesinger +-/ + +module +-- upstream: candidate + +public import Complexity.Computability.WordRAM.Notation +public import Complexity.Computability.WordRAM.RegisterFrame +public import Mathlib.Analysis.SpecialFunctions.Pow.NthRootLemmas +public import Mathlib.Data.Nat.Log +import Mathlib.Tactic.Linarith + +/-! +# A charged native integer `k`-th root + +The `k`-th root `Nat.nthRoot k n`, the largest `a` with `a ^ k ≤ n`, is found by binary search on +`[0, n + 1)`. Each probe compares `m ^ k` with `n` by a *capped power*: repeated multiplication +that stops as soon as the partial product exceeds `n`, so every product stays below `n (n + 1)` +and fits a word of about twice the bit length of `n`. The search halves the interval, so it makes +at most `⌈log₂ (n + 1)⌉` probes. + +## Main definitions + +* `NatNthRoot.CapPow.program`: compare `m ^ k` with a bound `n` by capped multiplication. +* `NatNthRoot.program`: binary search for `Nat.nthRoot k n`. + +## Main results + +* `NatNthRoot.CapPow.exec`: within `6 k + 5` instructions, the capped power is at most `n` + exactly when `m ^ k ≤ n`, and then equals `m ^ k`. +* `NatNthRoot.exec`: `Nat.nthRoot k n` within `⌈log₂ (n + 1)⌉ (6 k + 13) + 7` instructions. +-/ + +@[expose] public section + +namespace Cslib.Algorithms.WordRAM.NatNthRoot + +open Cslib.WordRAM + +namespace CapPow + +/-- Base, preserved. -/ +abbrev m : Fin 6 := 0 +/-- Exponent, preserved. -/ +abbrev k : Fin 6 := 1 +/-- Bound, preserved. -/ +abbrev n : Fin 6 := 2 +/-- Capped power, the output. -/ +abbrev p : Fin 6 := 3 +/-- Counter. -/ +abbrev i : Fin 6 := 4 +/-- Constant one. -/ +abbrev one : Fin 6 := 5 + +/-- Stop once the partial product exceeds the bound; otherwise multiply once more. -/ +def body : Code .standard 6 := + ram[ if n < p then [ i ← k ] else [ p ← p * m; i ← i + one ] ] + +/-- Capped repeated multiplication. -/ +def program : Code .standard 6 := + .seq (ram[ p ← #1; i ← #0; one ← #1 ]) (.loop (.lt i k) body) + +/-- For a positive base, the capped power is at most the bound exactly when the true power is, +and then the two agree. The products stay below `n m`. -/ +theorem exec {w : ℕ} (x : Data 6) (hm : 1 ≤ x.regs m) (h1 : Fits w 1) + (hk : Fits w (x.regs k)) (hprod : Fits w (x.regs n * x.regs m)) : + ∃ y, ExecWithin w program x y (6 * x.regs k + 5) ∧ + (y.regs p ≤ x.regs n ↔ x.regs m ^ x.regs k ≤ x.regs n) ∧ + (x.regs m ^ x.regs k ≤ x.regs n → y.regs p = x.regs m ^ x.regs k) ∧ + y.regs m = x.regs m ∧ y.regs k = x.regs k ∧ y.regs n = x.regs n ∧ + y.regs one = 1 ∧ y.mem = x.mem := by + set b := x.regs m + set e := x.regs k + set N := x.regs n + let I (y : Data 6) : Prop := y.regs m = b ∧ y.regs k = e ∧ y.regs n = N ∧ y.regs one = 1 ∧ + y.regs i ≤ e ∧ (y.regs p = b ^ y.regs i ∨ (N < y.regs p ∧ N < b ^ e)) ∧ y.mem = x.mem + have hbody (y : Data 6) (hy : I y) (hc : (Cond.lt i k).Holds y) : + ∃ z, ExecWithin w body y z 4 ∧ I z ∧ e - z.regs i < e - y.regs i := by + obtain ⟨hym, hyk, hyn, hyo, hyi, hyp, hymem⟩ := hy + have hlt : y.regs i < e := by simpa only [Cond.holds_lt, hyk] using hc + by_cases hover : N < y.regs p + · have hc' : (Cond.lt n p).Holds y := by simpa only [Cond.holds_lt, hyn] using hover + have hcopy : ExecWithin (B := .standard) w (ram[ i ← k ]) y (y.set i e) 1 := by + simpa only [hyk] using ExecWithin.copy (B := .standard) (w := w) (d := i) (s := k) (x := y) + refine ⟨y.set i e, (ExecWithin.branch_pos hc' hcopy).mono (by omega), ?_, ?_⟩ + · refine ⟨?_, ?_, ?_, ?_, ?_, ?_, hymem⟩ + · simpa +decide [Function.update_apply] using hym + · simpa +decide [Function.update_apply] using hyk + · simpa +decide [Function.update_apply] using hyn + · simpa +decide [Function.update_apply] using hyo + · simp + · refine Or.inr ⟨by simpa +decide [Function.update_apply] using hover, ?_⟩ + rcases hyp with h | h + · exact hover.trans_le (h ▸ Nat.pow_le_pow_right hm hyi) + · exact h.2 + · simp only [Data.set_regs, Function.update_self]; omega + · have hc' : ¬(Cond.lt n p).Holds y := by simpa only [Cond.holds_lt, hyn] using hover + have hpow : y.regs p = b ^ y.regs i := hyp.resolve_right (fun h => hover h.1) + have hmul : Fits w (y.regs p * y.regs m) := by + rw [hym]; exact hprod.of_le (Nat.mul_le_mul_right _ (by omega)) + have hinc : Fits w (y.regs i + 1) := hk.of_le (by omega) + refine ⟨?z, ?run, ?inv, ?decrease⟩ + case run => + refine ExecWithin.branch_neg hc' ?_ + ram_eval [hmul, hinc, hyo] + case inv => + refine ⟨?_, ?_, ?_, ?_, ?_, ?_, hymem⟩ + · simpa +decide [Function.update_apply] using hym + · simpa +decide [Function.update_apply] using hyk + · simpa +decide [Function.update_apply] using hyn + · simpa +decide [Function.update_apply] using hyo + · simp +decide only [Data.set_regs, Function.update_apply, ↓reduceIte]; omega + · left + simp +decide only [Data.set_regs, Function.update_apply, ↓reduceIte, hpow, hym, + pow_succ] + case decrease => + simp +decide only [Data.set_regs, Function.update_apply, ↓reduceIte]; omega + let x₀ := ((x.set p 1).set i 0).set one 1 + have hx₀ : I x₀ := by + refine ⟨?_, ?_, ?_, ?_, ?_, ?_, rfl⟩ <;> simp +decide [x₀, b, e, N] + obtain ⟨y, he, hy, hc⟩ := ExecWithin.loop_invariant I (fun y => e - y.regs i) hbody hx₀ + obtain ⟨hym, hyk, hyn, hyo, hyi, hyp, hymem⟩ := hy + have hend : y.regs i = e := by simp only [Cond.holds_lt, hyk] at hc; omega + have hinit : ExecWithin (B := .standard) w (ram[ p ← #1; i ← #0; one ← #1 ]) x x₀ 3 := by + ram_eval [x₀, h1] + refine ⟨y, ?_, ?_, ?_, hym, hyk, hyn, hyo, hymem⟩ + · have hi₀ : x₀.regs i = 0 := by simp +decide [x₀] + exact (hinit.seq he).mono (by rw [hi₀]; omega) + · rw [hend] at hyp + rcases hyp with h | h + · rw [h] + · exact ⟨fun hle => absurd hle (not_le.mpr h.1), fun hle => absurd hle (not_le.mpr h.2)⟩ + · intro hle + rw [hend] at hyp + exact hyp.resolve_right fun h => absurd hle (not_le.mpr h.2) + +/-- Run the capped power on any injectively chosen registers. -/ +theorem exec_rename {r w : ℕ} (f : Fin 6 → Fin r) (hf : f.Injective) (x : Data r) + (hm : 1 ≤ x.regs (f m)) (h1 : Fits w 1) (hk : Fits w (x.regs (f k))) + (hprod : Fits w (x.regs (f n) * x.regs (f m))) : + ∃ y, ExecWithin w (program.rename f) x y (6 * x.regs (f k) + 5) ∧ + (y.regs (f p) ≤ x.regs (f n) ↔ x.regs (f m) ^ x.regs (f k) ≤ x.regs (f n)) ∧ + (x.regs (f m) ^ x.regs (f k) ≤ x.regs (f n) → + y.regs (f p) = x.regs (f m) ^ x.regs (f k)) ∧ + y.regs (f m) = x.regs (f m) ∧ y.regs (f k) = x.regs (f k) ∧ + y.regs (f n) = x.regs (f n) ∧ y.regs (f one) = 1 ∧ y.mem = x.mem := by + obtain ⟨y, hy, hle, heq, hym, hyk, hyn, hyo, hymem⟩ := exec (x.comap f) hm h1 hk hprod + have hext (j : Fin 6) : (y.extend f x).regs (f j) = y.regs j := hf.extend_apply _ _ j + refine ⟨y.extend f x, hy.rename hf, ?_, ?_, ?_, ?_, ?_, ?_, hymem⟩ <;> + simp only [hext] <;> assumption + +end CapPow + +/-- Bound `n`, preserved. -/ +abbrev n : Fin 10 := 0 +/-- Root degree `k`, preserved. -/ +abbrev k : Fin 10 := 1 +/-- Lower end of the search interval, and the output. -/ +abbrev lo : Fin 10 := 2 +/-- Upper end of the search interval. -/ +abbrev hi : Fin 10 := 3 +/-- Probe. -/ +abbrev mid : Fin 10 := 4 +/-- Capped power of the probe. -/ +abbrev p : Fin 10 := 5 +/-- Counter of the capped power. -/ +abbrev i : Fin 10 := 6 +/-- Constant one. -/ +abbrev one : Fin 10 := 7 +/-- Constant two. -/ +abbrev two : Fin 10 := 8 +/-- `lo + 1`, for the loop test. -/ +abbrev t : Fin 10 := 9 + +/-- The registers of the capped power inside the search. -/ +def capMap : Fin 6 → Fin 10 := ![mid, k, n, p, i, one] + +theorem capMap_injective : capMap.Injective := by decide + +/-- One probe: halve the interval according to the capped power of its midpoint. -/ +def body : Code .standard 10 := + .seq (ram[ mid ← lo + hi; mid ← mid / two ]) (.seq (CapPow.program.rename capMap) + (ram[ if n < p then [ hi ← mid ] else [ lo ← mid ]; t ← lo + one ])) + +/-- Binary search for the `k`-th root on `[0, n + 1)`. -/ +def program : Code .standard 10 := + .seq (ram[ lo ← #0; one ← #1; two ← #2; hi ← n + one; t ← lo + one ]) (.loop (.lt t hi) body) + +private theorem clog_halve {d : ℕ} (hd : 2 ≤ d) {d' : ℕ} (hd' : d' ≤ (d + 1) / 2) : + Nat.clog 2 d' < Nat.clog 2 d := by + rw [Nat.clog_of_two_le one_lt_two hd] + exact Nat.lt_succ_of_le (Nat.clog_mono_right 2 (by omega)) + +/-- Compute `Nat.nthRoot k n` for `k ≥ 1` within `⌈log₂ (n + 1)⌉ (6 k + 13) + 7` instructions, +preserving `n`, `k` and memory. Every intermediate value is at most `n² + 2 n + 1`. -/ +theorem exec {w : ℕ} (x : Data 10) (hk : 1 ≤ x.regs k) (hkfit : Fits w (x.regs k)) + (hfit : Fits w (x.regs n * x.regs n + 2 * x.regs n + 2)) : + ∃ y, ExecWithin w program x y (Nat.clog 2 (x.regs n + 1) * (6 * x.regs k + 13) + 7) ∧ + y.regs lo = Nat.nthRoot (x.regs k) (x.regs n) ∧ y.regs n = x.regs n ∧ + y.regs k = x.regs k ∧ y.regs one = 1 ∧ y.regs two = 2 ∧ y.mem = x.mem := by + set N := x.regs n + set e := x.regs k + have he0 : e ≠ 0 := by omega + have h1 : Fits w 1 := hfit.of_le (by omega) + have h2 : Fits w 2 := hfit.of_le (by omega) + let I (y : Data 10) : Prop := y.regs n = N ∧ y.regs k = e ∧ y.regs one = 1 ∧ y.regs two = 2 ∧ + y.regs t = y.regs lo + 1 ∧ y.regs lo < y.regs hi ∧ y.regs hi ≤ N + 1 ∧ + y.regs lo ^ e ≤ N ∧ N < y.regs hi ^ e ∧ y.mem = x.mem + have hbody (y : Data 10) (hy : I y) (hc : (Cond.lt t hi).Holds y) : + ∃ z, ExecWithin w body y z (6 * e + 11) ∧ I z ∧ + Nat.clog 2 (z.regs hi - z.regs lo) < Nat.clog 2 (y.regs hi - y.regs lo) := by + obtain ⟨hyn, hyk, hyo, hyt2, hyt, hlohi, hhi, hlo, hhip, hymem⟩ := hy + obtain ⟨l, hl⟩ : ∃ l, y.regs lo = l := ⟨_, rfl⟩ + obtain ⟨h, hh⟩ : ∃ h, y.regs hi = h := ⟨_, rfl⟩ + simp only [hl, hh] at hlohi hlo hhip hyt hhi + have hgap : l + 1 < h := by simpa only [Cond.holds_lt, hyt, hh] using hc + have hmd1 : l < (l + h) / 2 := by omega + have hmd2 : (l + h) / 2 < h := by omega + have hsum : Fits w (l + h) := hfit.of_le (by nlinarith) + have hdiv : Fits w ((l + h) / 2) := hsum.of_le (Nat.div_le_self _ _) + -- After the two midpoint instructions. + let y₁ := (y.set mid (l + h)).set mid ((l + h) / 2) + have hrun₁ : ExecWithin (B := .standard) w (ram[ mid ← lo + hi; mid ← mid / two ]) y y₁ 2 := by + ram_eval [y₁, hl, hh, hsum, hdiv, hyt2] + have hy₁ (j : Fin 10) (hj : j ≠ mid) : y₁.regs j = y.regs j := by + simp [y₁, hj] + have hy₁m : y₁.regs mid = (l + h) / 2 := by simp [y₁] + have hmdN : (l + h) / 2 ≤ N := by omega + have hcm : capMap CapPow.m = mid := rfl + have hck : capMap CapPow.k = k := rfl + have hcn : capMap CapPow.n = n := rfl + have hcp : capMap CapPow.p = p := rfl + have hco : capMap CapPow.one = one := rfl + obtain ⟨y₂, hrun₂, hple, -, hy₂m, hy₂k, hy₂n, hy₂o, hy₂mem⟩ := + CapPow.exec_rename capMap capMap_injective y₁ (by rw [hcm, hy₁m]; omega) h1 + (by rw [hck, hy₁ k (by decide), hyk]; exact hkfit) + (by + rw [hcn, hcm, hy₁m, hy₁ n (by decide), hyn] + exact hfit.of_le ((Nat.mul_le_mul_left N hmdN).trans (by omega))) + have hpres (j : Fin 10) (hj : ∀ a, capMap a ≠ j) : y₂.regs j = y₁.regs j := + hrun₂.regs_eq (Code.preserves_rename_of_forall_ne _ _ _ hj) + rw [hcp, hcn, hcm, hck, hy₁m, hy₁ n (by decide), hy₁ k (by decide), hyn, hyk] at hple + rw [hck, hy₁ k (by decide), hyk] at hy₂k + rw [hcn, hy₁ n (by decide), hyn] at hy₂n + rw [hco] at hy₂o + rw [hcm, hy₁m] at hy₂m + rw [hck, hy₁ k (by decide), hyk] at hrun₂ + have hy₂lo : y₂.regs lo = l := (hpres lo (by decide)).trans ((hy₁ lo (by decide)).trans hl) + have hy₂hi : y₂.regs hi = h := (hpres hi (by decide)).trans ((hy₁ hi (by decide)).trans hh) + have hy₂two : y₂.regs two = 2 := + (hpres two (by decide)).trans ((hy₁ two (by decide)).trans hyt2) + have hy₂mid : y₂.regs mid = (l + h) / 2 := hy₂m + by_cases hover : N < y₂.regs p + · -- The probe is too large: `hi ← mid`. + have hc' : (Cond.lt n p).Holds y₂ := by simpa only [Cond.holds_lt, hy₂n] using hover + have hpow : N < ((l + h) / 2) ^ e := by + by_contra hle + exact absurd (hple.mpr (not_lt.mp hle)) (not_le.mpr hover) + have hinc : Fits w (l + 1) := hfit.of_le (by nlinarith) + let z := (y₂.set hi ((l + h) / 2)).set t (l + 1) + have hrun₃ : ExecWithin (B := .standard) w + (ram[ if n < p then [ hi ← mid ] else [ lo ← mid ]; t ← lo + one ]) y₂ z 3 := by + refine ((ExecWithin.branch_pos (T := 1) hc' (ExecWithin.copy (d := hi) (s := mid))).seq + (T₂ := 1) ?_).mono (by norm_num) + ram_eval [z, hy₂lo, hy₂o, hy₂mid, hinc] + refine ⟨z, (hrun₁.seq (hrun₂.seq hrun₃)).mono (by omega), ?_, ?_⟩ + · refine ⟨?_, ?_, ?_, ?_, ?_, ?_, ?_, ?_, ?_, hy₂mem.trans hymem⟩ <;> + simp +decide [z, hy₂n, hy₂k, hy₂o, hy₂two, hy₂lo] <;> omega + · simp +decide only [z, Data.set_regs, Function.update_apply, ↓reduceIte, hy₂lo, hl, hh] + exact clog_halve (by omega) (by omega) + · have hc' : ¬(Cond.lt n p).Holds y₂ := by simpa only [Cond.holds_lt, hy₂n] using hover + have hpow : ((l + h) / 2) ^ e ≤ N := hple.mp (not_lt.mp hover) + have hinc : Fits w ((l + h) / 2 + 1) := hfit.of_le (by nlinarith) + let z := (y₂.set lo ((l + h) / 2)).set t ((l + h) / 2 + 1) + have hrun₃ : ExecWithin (B := .standard) w + (ram[ if n < p then [ hi ← mid ] else [ lo ← mid ]; t ← lo + one ]) y₂ z 4 := by + refine ((ExecWithin.branch_neg (T := 1) hc' (ExecWithin.copy (d := lo) (s := mid))).seq + (T₂ := 1) ?_).mono (by norm_num) + ram_eval [z, hy₂o, hy₂mid, hinc] + refine ⟨z, (hrun₁.seq (hrun₂.seq hrun₃)).mono (by omega), ?_, ?_⟩ + · refine ⟨?_, ?_, ?_, ?_, ?_, ?_, ?_, ?_, ?_, hy₂mem.trans hymem⟩ <;> + simp +decide [z, hy₂n, hy₂k, hy₂o, hy₂two, hy₂hi] <;> omega + · simp +decide only [z, Data.set_regs, Function.update_apply, ↓reduceIte, hy₂hi, hl, hh] + exact clog_halve (by omega) (by omega) + have hn1 : Fits w (x.regs n + 1) := hfit.of_le (by nlinarith) + let x₀ := ((((x.set lo 0).set one 1).set two 2).set hi (N + 1)).set t 1 + have hx₀ : I x₀ := by + have hpow : N < (N + 1) ^ e := Nat.lt_of_lt_of_le (Nat.lt_succ_self N) (Nat.le_self_pow he0 _) + refine ⟨?_, ?_, ?_, ?_, ?_, ?_, ?_, ?_, ?_, rfl⟩ <;> + simp +decide [x₀, N, e, he0, hpow] + obtain ⟨y, hrun, hy, hc⟩ := ExecWithin.loop_invariant I + (fun y => Nat.clog 2 (y.regs hi - y.regs lo)) hbody hx₀ + obtain ⟨hyn, hyk, hyo, hytwo, hyt, hlohi, -, hlo, hhip, hymem⟩ := hy + have hhi : y.regs hi = y.regs lo + 1 := by simp only [Cond.holds_lt, hyt] at hc; omega + have hinit : ExecWithin (B := .standard) w + (ram[ lo ← #0; one ← #1; two ← #2; hi ← n + one; t ← lo + one ]) x x₀ 5 := by + ram_eval [x₀, N, h1, h2, hn1] + refine ⟨y, ?_, ?_, hyn, hyk, hyo, hytwo, hymem⟩ + · have hμ : Nat.clog 2 (x₀.regs hi - x₀.regs lo) = Nat.clog 2 (N + 1) := by + simp +decide [x₀] + refine (hinit.seq hrun).mono ?_ + rw [hμ, show 6 * e + 11 + 2 = 6 * e + 13 by omega] + omega + · rw [hhi] at hhip + exact (Nat.nthRoot_eq_of_le_of_lt hlo hhip).symm + +/-- Run the binary search on any injectively chosen registers. -/ +theorem exec_rename {r w : ℕ} (f : Fin 10 → Fin r) (hf : f.Injective) (x : Data r) + (hk : 1 ≤ x.regs (f k)) (hkfit : Fits w (x.regs (f k))) + (hfit : Fits w (x.regs (f n) * x.regs (f n) + 2 * x.regs (f n) + 2)) : + ∃ y, ExecWithin w (program.rename f) x y + (Nat.clog 2 (x.regs (f n) + 1) * (6 * x.regs (f k) + 13) + 7) ∧ + y.regs (f lo) = Nat.nthRoot (x.regs (f k)) (x.regs (f n)) ∧ + y.regs (f n) = x.regs (f n) ∧ y.regs (f k) = x.regs (f k) ∧ + y.regs (f one) = 1 ∧ y.regs (f two) = 2 ∧ y.mem = x.mem := by + obtain ⟨y, hy, hlo, hyn, hyk, hyo, hyt, hymem⟩ := exec (x.comap f) hk hkfit hfit + have hext (j : Fin 10) : (y.extend f x).regs (f j) = y.regs j := hf.extend_apply _ _ j + refine ⟨y.extend f x, hy.rename hf, ?_, ?_, ?_, ?_, ?_, hymem⟩ <;> simp only [hext] <;> + assumption + +end Cslib.Algorithms.WordRAM.NatNthRoot diff --git a/Complexity/Foundations/NumberTheory/AKS/Coefficients.lean b/Complexity/Foundations/NumberTheory/AKS/Coefficients.lean new file mode 100644 index 000000000..0d4565e02 --- /dev/null +++ b/Complexity/Foundations/NumberTheory/AKS/Coefficients.lean @@ -0,0 +1,176 @@ +/- +Copyright (c) 2026 Samuel Schlesinger. All rights reserved. +Released under Apache 2.0 license as described in the file LICENSE. +Authors: Samuel Schlesinger +-/ + +module +-- upstream: candidate + +public import Complexity.Foundations.NumberTheory.AKS.Criterion + +/-! +# Coefficient arithmetic modulo `X^r - 1` and `n` + +An element of `(ℤ / n)[X] / (X^r - 1)` is represented by the `r` coefficients of its reduced +polynomial, as natural numbers below `n`. This file gives the coefficient formulas that a program +evaluates and identifies them with the ring operations: `Cslib.AKS.toQuot n r c` is the class of +`∑_{i < r} c i X^i`, products of classes have the coefficients `Cslib.AKS.mulCoeff` +(`Cslib.AKS.toQuot_mulCoeff`), the classes of `X^e + a` have the coefficients +`Cslib.AKS.powAddCoeff` (`Cslib.AKS.toQuot_powAddCoeff`), and reduced coefficient vectors +determine their classes (`Cslib.AKS.toQuot_eq_toQuot_iff`). In particular the AKS congruence +holds exactly when the `n`-th power of the class of `X + a` has the coefficients of `X^n + a` +(`Cslib.AKS.congruence_iff_toQuot`). + +## Main definitions + +* `Cslib.AKS.coeffPoly`, `Cslib.AKS.toQuot`: the polynomial with given coefficients and its class. +* `Cslib.AKS.mulCoeff`: the coefficients of a product. +* `Cslib.AKS.powAddCoeff`: the coefficients of `X^e + a`. + +## Main results + +* `Cslib.AKS.toQuot_mulCoeff`: coefficients of products. +* `Cslib.AKS.toQuot_powAddCoeff`: coefficients of `X^e + a`. +* `Cslib.AKS.toQuot_eq_toQuot_iff`: reduced coefficients determine their class. +* `Cslib.AKS.congruence_iff_toQuot`: the AKS congruence in coefficients. + +## References + +* [M. Agrawal, N. Kayal, N. Saxena, *PRIMES is in P*][AKS2004] +-/ + +@[expose] public section + +namespace Cslib.AKS + +open Polynomial Finset + +variable {n r : ℕ} + +/-- The polynomial `∑_{i < r} c i X^i` over `ℤ / n`. -/ +noncomputable def coeffPoly (n r : ℕ) (c : ℕ → ℕ) : (ZMod n)[X] := + ∑ i ∈ range r, C (c i : ZMod n) * X ^ i + +/-- The class of `∑_{i < r} c i X^i` in `(ℤ / n)[X] / (X^r - 1)`. -/ +noncomputable def toQuot (n r : ℕ) (c : ℕ → ℕ) : AdjoinRoot (X ^ r - 1 : (ZMod n)[X]) := + AdjoinRoot.mk _ (coeffPoly n r c) + +/-- The coefficients of the product of the classes of `p` and `q`, reduced modulo `n`: +`∑_{(i + j) mod r = k} p i q j mod n`. -/ +def mulCoeff (n r : ℕ) (p q : ℕ → ℕ) (k : ℕ) : ℕ := + (∑ i ∈ range r, ∑ j ∈ range r, if (i + j) % r = k then p i * q j else 0) % n + +/-- The coefficients of `X^e + a`, reduced modulo `X^r - 1` and `n`. -/ +def powAddCoeff (n r e a : ℕ) (k : ℕ) : ℕ := + ((if k = e % r then 1 else 0) + (if k = 0 then a else 0)) % n + +/-- The root of `X^r - 1` is an `r`-th root of unity. -/ +theorem root_pow_self (n r : ℕ) : AdjoinRoot.root (X ^ r - 1 : (ZMod n)[X]) ^ r = 1 := by + have h := AdjoinRoot.mk_self (f := (X ^ r - 1 : (ZMod n)[X])) + rw [map_sub, map_pow, AdjoinRoot.mk_X, map_one, sub_eq_zero] at h + exact h + +/-- Powers of the root depend only on the exponent modulo `r`. -/ +theorem root_pow_mod (n r m : ℕ) : + AdjoinRoot.root (X ^ r - 1 : (ZMod n)[X]) ^ (m % r) = + AdjoinRoot.root (X ^ r - 1 : (ZMod n)[X]) ^ m := by + conv_rhs => rw [← Nat.mod_add_div m r, pow_add, pow_mul, root_pow_self, one_pow, mul_one] + +/-- Natural numbers in `(ℤ / n)[X] / (X^r - 1)` depend only on their residue modulo `n`. -/ +theorem natCast_mod (n r m : ℕ) : + (((m % n : ℕ)) : AdjoinRoot (X ^ r - 1 : (ZMod n)[X])) = m := by + have h : ((n : ℕ) : AdjoinRoot (X ^ r - 1 : (ZMod n)[X])) = 0 := by + rw [← map_natCast (AdjoinRoot.of (X ^ r - 1 : (ZMod n)[X])), ZMod.natCast_self, map_zero] + conv_rhs => rw [← Nat.mod_add_div m n] + rw [Nat.cast_add, Nat.cast_mul, h, zero_mul, add_zero] + +/-- The class of `∑_{i < r} c i X^i` as a sum of powers of the root. -/ +theorem toQuot_eq_sum (n r : ℕ) (c : ℕ → ℕ) : + toQuot n r c = ∑ i ∈ range r, + (c i : AdjoinRoot (X ^ r - 1 : (ZMod n)[X])) * AdjoinRoot.root _ ^ i := by + simp only [toQuot, coeffPoly, map_sum, map_mul, map_pow, AdjoinRoot.mk_X, AdjoinRoot.mk_C] + rfl + +/-- **Coefficients of products.** -/ +theorem toQuot_mulCoeff (n r : ℕ) (p q : ℕ → ℕ) : + toQuot n r (mulCoeff n r p q) = toQuot n r p * toQuot n r q := by + set ξ := AdjoinRoot.root (X ^ r - 1 : (ZMod n)[X]) + rcases Nat.eq_zero_or_pos r with rfl | hr + · simp [toQuot_eq_sum] + rw [toQuot_eq_sum, toQuot_eq_sum, toQuot_eq_sum, sum_mul_sum] + simp only [mulCoeff, natCast_mod, Nat.cast_sum, Nat.cast_ite, Nat.cast_mul, Nat.cast_zero, + sum_mul, ite_mul, zero_mul] + rw [sum_comm] + refine sum_congr rfl fun i _ => ?_ + rw [sum_comm] + refine sum_congr rfl fun j _ => ?_ + rw [sum_ite_eq (range r) ((i + j) % r)] + simp only [mem_range, Nat.mod_lt _ hr, ↓reduceIte] + rw [root_pow_mod, pow_add] + ring + +/-- **Coefficients of `X^e + a`.** -/ +theorem toQuot_powAddCoeff (hr : 0 < r) (e a : ℕ) : + toQuot n r (powAddCoeff n r e a) = + AdjoinRoot.mk (X ^ r - 1 : (ZMod n)[X]) (X ^ e + C (a : ZMod n)) := by + rw [toQuot_eq_sum] + simp only [powAddCoeff, natCast_mod, Nat.cast_add, Nat.cast_ite, Nat.cast_one, Nat.cast_zero, + add_mul, ite_mul, one_mul, zero_mul, sum_add_distrib] + rw [sum_ite_eq' (range r) (e % r), sum_ite_eq' (range r) 0] + simp only [mem_range, Nat.mod_lt _ hr, hr, ↓reduceIte] + rw [pow_zero, mul_one, root_pow_mod, map_add, map_pow, AdjoinRoot.mk_X, AdjoinRoot.mk_C] + rfl + +/-- The coefficients of `∑_{i < r} c i X^i`. -/ +theorem coeff_coeffPoly (c : ℕ → ℕ) (k : ℕ) : + (coeffPoly n r c).coeff k = if k < r then (c k : ZMod n) else 0 := by + simp only [coeffPoly, finsetSum_coeff, coeff_C_mul_X_pow] + rw [sum_ite_eq (range r) k] + simp only [mem_range] + +/-- **Reduced coefficients determine their class**: for `n ≥ 1` and `r ≥ 1`, coefficient vectors +with entries below `n` have equal classes exactly when they agree below `r`. -/ +theorem toQuot_eq_toQuot_iff (hn : 0 < n) (hr : 0 < r) {c d : ℕ → ℕ} + (hc : ∀ i < r, c i < n) (hd : ∀ i < r, d i < n) : + toQuot n r c = toQuot n r d ↔ ∀ i < r, c i = d i := by + constructor + · intro h + rcases (show n = 1 ∨ 2 ≤ n by omega) with rfl | hn2 + · intro i hi; have := hc i hi; have := hd i hi; omega + have : Nontrivial (ZMod n) := ZMod.nontrivial_iff.mpr (by omega) + rw [toQuot, toQuot, AdjoinRoot.mk_eq_mk] at h + have hzero : coeffPoly n r c - coeffPoly n r d = 0 := by + by_contra h0 + refine (monic_X_pow_sub_one (R := ZMod n) hr.ne').not_dvd_of_degree_lt h0 ?_ h + rw [← C_1, degree_X_pow_sub_C hr] + refine (degree_lt_iff_coeff_zero _ _).mpr fun k hk => ?_ + have hk' : ¬k < r := not_lt.mpr hk + rw [coeff_sub, coeff_coeffPoly, coeff_coeffPoly] + simp only [hk', ↓reduceIte, sub_zero] + intro i hi + have hcoeff := congrArg (fun f => f.coeff i) hzero + simp only [coeff_sub, coeff_coeffPoly, hi, ↓reduceIte, coeff_zero, sub_eq_zero] at hcoeff + rw [ZMod.natCast_eq_natCast_iff', Nat.mod_eq_of_lt (hc i hi), Nat.mod_eq_of_lt (hd i hi)] + at hcoeff + exact hcoeff + · intro h + rw [toQuot_eq_sum, toQuot_eq_sum] + exact sum_congr rfl fun i hi => by rw [h i (mem_range.mp hi)] + +/-- **The AKS congruence in coefficients**: `(X + a)^n ≡ X^n + a (mod X^r - 1, n)` exactly when +the `n`-th power of the class with the coefficients of `X + a` is the class with the coefficients +of `X^n + a`. -/ +theorem congruence_iff_toQuot (hr : 0 < r) (a : ℕ) : + Congruence n r a ↔ toQuot n r (powAddCoeff n r 1 a) ^ n = toQuot n r (powAddCoeff n r n a) := by + rw [toQuot_powAddCoeff hr, toQuot_powAddCoeff hr, pow_one, ← map_pow, Congruence] + +/-- Coefficients of `X^e + a` are reduced modulo `n`. -/ +theorem powAddCoeff_lt (hn : 0 < n) (e a k : ℕ) : powAddCoeff n r e a k < n := + Nat.mod_lt _ hn + +/-- Coefficients of products are reduced modulo `n`. -/ +theorem mulCoeff_lt (hn : 0 < n) (p q : ℕ → ℕ) (k : ℕ) : mulCoeff n r p q k < n := + Nat.mod_lt _ hn + +end Cslib.AKS diff --git a/Complexity/Foundations/NumberTheory/AKS/PerfectPower.lean b/Complexity/Foundations/NumberTheory/AKS/PerfectPower.lean new file mode 100644 index 000000000..11e565268 --- /dev/null +++ b/Complexity/Foundations/NumberTheory/AKS/PerfectPower.lean @@ -0,0 +1,64 @@ +/- +Copyright (c) 2026 Samuel Schlesinger. All rights reserved. +Released under Apache 2.0 license as described in the file LICENSE. +Authors: Samuel Schlesinger +-/ + +module +-- upstream: candidate + +public import Complexity.Foundations.NumberTheory.AKS.Criterion +public import Mathlib.Analysis.SpecialFunctions.Pow.NthRootLemmas +public import Mathlib.Data.Nat.Log + +/-! +# Deciding perfect powers + +Step 1 of the AKS algorithm [Agrawal, Kayal, Saxena, Section 4][AKS2004] asks whether `n = a ^ b` +for some `b > 1`. For `n ≥ 2` such an `a` is at least `2`, so `b ≤ log₂ n`, and `n` is a `b`-th +power exactly when it is the `b`-th power of its integer `b`-th root `Nat.nthRoot b n`. This makes +`Cslib.AKS.IsPerfectPower` decidable by a bounded search. + +## Main results + +* `Cslib.AKS.isPerfectPower_iff_nthRoot`: perfect powers through integer roots. +* `Cslib.AKS.decidablePredIsPerfectPower`: the resulting decision procedure. + +## References + +* [M. Agrawal, N. Kayal, N. Saxena, *PRIMES is in P*][AKS2004] +-/ + +@[expose] public section + +namespace Cslib.AKS + +/-- For `n ≥ 2`, `n` is a perfect power exactly when it is the `b`-th power of its integer `b`-th +root for some `2 ≤ b ≤ log₂ n`. -/ +theorem isPerfectPower_iff_nthRoot {n : ℕ} (hn : 2 ≤ n) : + IsPerfectPower n ↔ ∃ b < Nat.log 2 n + 1, 2 ≤ b ∧ Nat.nthRoot b n ^ b = n := by + constructor + · rintro ⟨a, b, hb, rfl⟩ + have ha : 2 ≤ a := by + rcases a with _ | _ | a + · simp [zero_pow (by omega : b ≠ 0)] at hn + · simp at hn + · omega + refine ⟨b, Nat.lt_succ_of_le (Nat.le_log_of_pow_le one_lt_two (Nat.pow_le_pow_left ha b)), + hb, ?_⟩ + rw [Nat.nthRoot_pow (by omega)] + · rintro ⟨b, -, hb, h⟩ + exact ⟨_, b, hb, h⟩ + +/-- `0` and `1` are perfect powers: `0 = 0 ^ 2` and `1 = 1 ^ 2`. -/ +theorem isPerfectPower_of_lt_two {n : ℕ} (hn : n < 2) : IsPerfectPower n := by + interval_cases n + · exact ⟨0, 2, one_lt_two, rfl⟩ + · exact ⟨1, 2, one_lt_two, rfl⟩ + +/-- Perfect powers are decidable, by the bounded search of `isPerfectPower_iff_nthRoot`. -/ +instance decidablePredIsPerfectPower : DecidablePred IsPerfectPower := fun n => + if hn : 2 ≤ n then decidable_of_iff _ (isPerfectPower_iff_nthRoot hn).symm + else isTrue (isPerfectPower_of_lt_two (by omega)) + +end Cslib.AKS diff --git a/ComplexityTests/WordRAMAKS.lean b/ComplexityTests/WordRAMAKS.lean new file mode 100644 index 000000000..32e95b9dc --- /dev/null +++ b/ComplexityTests/WordRAMAKS.lean @@ -0,0 +1,109 @@ +/- +Copyright (c) 2026 Samuel Schlesinger. All rights reserved. +Released under Apache 2.0 license as described in the file LICENSE. +Authors: Samuel Schlesinger +-/ + +import Complexity.Algorithms.WordRAM.NatGcd +import Complexity.Algorithms.WordRAM.AKS.PerfectPower +import Complexity.Algorithms.WordRAM.AKS.PolyMul + +/-! +# Word-RAM subroutines of the AKS primality test + +Concrete runs of Euclid's algorithm, the integer `k`-th root, the perfect-power test and +multiplication modulo `X^r - 1` and `n`, and instances of their correctness theorems. +-/ + +namespace ComplexityTests.WordRAMAKS + +open Cslib.WordRAM Cslib.Algorithms.WordRAM + +/-! ## Greatest common divisors -/ + +private def gcdInput (a b : ℕ) : Data 4 := (Data.init []).set NatGcd.a a |>.set NatGcd.b b + +/-- Euclid's algorithm on `(12, 18)`, `(17, 5)` and `(7, 0)`. -/ +example : (Code.eval 8 8 NatGcd.program (gcdInput 12 18)).map (·.1.regs NatGcd.a) = some 6 ∧ + (Code.eval 8 8 NatGcd.program (gcdInput 17 5)).map (·.1.regs NatGcd.a) = some 1 ∧ + (Code.eval 8 8 NatGcd.program (gcdInput 7 0)).map (·.1.regs NatGcd.a) = some 7 := by + decide + +/-- The correctness theorem at width `8`: `gcd(12, 18) = 6` within `93` instructions. -/ +example : ∃ y, ExecWithin 8 NatGcd.program (gcdInput 12 18) y 93 ∧ y.regs NatGcd.a = 6 := by + obtain ⟨y, hy, ha, -⟩ := NatGcd.exec (w := 8) (gcdInput 12 18) (by decide) (by decide) + exact ⟨y, hy, ha.trans (by decide)⟩ + +/-! ## Integer roots -/ + +private def rootInput (n k : ℕ) : Data 10 := + (Data.init []).set NatNthRoot.n n |>.set NatNthRoot.k k + +/-- `⌊√30⌋ = 5`, `⌊∛64⌋ = 4` and `⌊∛63⌋ = 3`, by binary search. -/ +example : (Code.eval 16 16 NatNthRoot.program (rootInput 30 2)).map (·.1.regs NatNthRoot.lo) = + some 5 ∧ + (Code.eval 16 16 NatNthRoot.program (rootInput 64 3)).map (·.1.regs NatNthRoot.lo) = some 4 ∧ + (Code.eval 16 16 NatNthRoot.program (rootInput 63 3)).map (·.1.regs NatNthRoot.lo) = some 3 := by + decide + +/-- The correctness theorem: the search returns `Nat.nthRoot 3 63`. -/ +example : ∃ y, ExecWithin 16 NatNthRoot.program (rootInput 63 3) y 193 ∧ + y.regs NatNthRoot.lo = Nat.nthRoot 3 63 := by + obtain ⟨y, hy, hlo, -⟩ := NatNthRoot.exec (w := 16) (rootInput 63 3) (by decide) (by decide) + (by decide) + exact ⟨y, hy.mono (by decide), hlo⟩ + +/-! ## Perfect powers -/ + +private def ppInput (n : ℕ) : Data 14 := (Data.init []).set AKS.PerfectPower.n n + +/-- `8 = 2³` and `9 = 3²` are perfect powers; `10` and `12` are not. -/ +example : (Code.eval 16 16 AKS.PerfectPower.program (ppInput 8)).map + (·.1.regs AKS.PerfectPower.flag) = some 1 ∧ + (Code.eval 16 16 AKS.PerfectPower.program (ppInput 9)).map + (·.1.regs AKS.PerfectPower.flag) = some 1 ∧ + (Code.eval 16 16 AKS.PerfectPower.program (ppInput 10)).map + (·.1.regs AKS.PerfectPower.flag) = some 0 ∧ + (Code.eval 16 16 AKS.PerfectPower.program (ppInput 12)).map + (·.1.regs AKS.PerfectPower.flag) = some 0 := by + decide + +/-- The decision procedure for `IsPerfectPower` agrees. -/ +example : Cslib.AKS.IsPerfectPower 27 ∧ ¬Cslib.AKS.IsPerfectPower 12 := by decide + +/-- The correctness theorem on `27 = 3³`. -/ +example : ∃ y, ExecWithin 16 AKS.PerfectPower.program (ppInput 27) y + (AKS.PerfectPower.time 27) ∧ y.regs AKS.PerfectPower.flag = 1 := by + obtain ⟨y, hy, hf, -⟩ := AKS.PerfectPower.exec (w := 16) (ppInput 27) (by decide) (by decide) + exact ⟨y, hy, hf.trans (by decide)⟩ + +/-! ## Multiplication modulo `X^r - 1` and `n` -/ + +/-- Factors at `10` and `13`, product at `20`, `r = 3`, `n = 5`: `p = X + 1`, `q = 4X² + 3`. -/ +private def mulInput (pa pb : ℕ) : Data 14 := + ⟨fun j => if j = AKS.PolyMul.pa then pa else if j = AKS.PolyMul.pb then pb else + if j = AKS.PolyMul.pt then 20 else if j = AKS.PolyMul.rr then 3 else + if j = AKS.PolyMul.nn then 5 else 0, + fun a => if a = 10 ∨ a = 11 then 1 else if a = 13 then 3 else if a = 15 then 4 else + if a = 23 then 99 else 0⟩ + +/-- `(X + 1)² = X² + 2X + 1` and `(X + 1)(4X² + 3) = 4X³ + 4X² + 3X + 3 ≡ 4X² + 3X + 2` +modulo `X³ - 1` and `5`. -/ +example : (Code.eval 16 4 AKS.PolyMul.program (mulInput 10 10)).map + (fun yn => [yn.1.mem 20, yn.1.mem 21, yn.1.mem 22]) = some [1, 2, 1] ∧ + (Code.eval 16 4 AKS.PolyMul.program (mulInput 10 13)).map + (fun yn => [yn.1.mem 20, yn.1.mem 21, yn.1.mem 22]) = some [2, 3, 4] := by + decide +kernel + +/-- The correctness theorem computes the same square through `Cslib.AKS.mulCoeff`. -/ +example : ∃ y, ExecWithin 16 AKS.PolyMul.program (mulInput 10 10) y 224 ∧ y.mem 21 = 2 ∧ + y.mem 23 = 99 := by + obtain ⟨y, hy, hcoeff, hout, -⟩ := AKS.PolyMul.exec (w := 16) (mulInput 10 10) + (a := 10) (b := 10) (t := 20) (r := 3) (n := 5) rfl rfl rfl rfl rfl + (by decide) (by decide) (by decide) (by decide) (by decide) (by decide) (by decide) + (by decide) (by decide) (by decide) + refine ⟨y, hy, ?_, hout 23 (by decide)⟩ + rw [show (21 : ℕ) = 20 + 1 from rfl, hcoeff 1 (by decide)] + decide + +end ComplexityTests.WordRAMAKS From c743e327ce6f25ac9318bb4edeee92e1993947f1 Mon Sep 17 00:00:00 2001 From: Samuel Giuseppe Schlesinger Date: Fri, 30 Oct 2026 18:55:04 -0400 Subject: [PATCH 2/4] feat(WordRAM/AKS): write and compare coefficient arrays (#893) Fill writes the coefficients of X^e + a modulo X^r - 1 and n into a block, over arbitrary prior contents; Compare tests two blocks for equality. Both are linear in r. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_015DLg8kwPs8yi7ePPdtoNvG --- Complexity.lean | 1 + .../Algorithms/WordRAM/AKS/PolyArray.lean | 306 ++++++++++++++++++ ComplexityTests/WordRAMAKS.lean | 40 ++- 3 files changed, 345 insertions(+), 2 deletions(-) create mode 100644 Complexity/Algorithms/WordRAM/AKS/PolyArray.lean diff --git a/Complexity.lean b/Complexity.lean index c2e70310d..b44bd1eb8 100644 --- a/Complexity.lean +++ b/Complexity.lean @@ -81,6 +81,7 @@ public import Complexity.Algorithms.ThreeSum.Convolution public import Complexity.Algorithms.ThreeSum.Defs public import Complexity.Algorithms.ThreeSum.ThreeSets public import Complexity.Algorithms.WordRAM.AKS.PerfectPower +public import Complexity.Algorithms.WordRAM.AKS.PolyArray public import Complexity.Algorithms.WordRAM.AKS.PolyMul public import Complexity.Algorithms.WordRAM.Deduplicate.Count public import Complexity.Algorithms.WordRAM.Deduplicate.Model diff --git a/Complexity/Algorithms/WordRAM/AKS/PolyArray.lean b/Complexity/Algorithms/WordRAM/AKS/PolyArray.lean new file mode 100644 index 000000000..4128f38cf --- /dev/null +++ b/Complexity/Algorithms/WordRAM/AKS/PolyArray.lean @@ -0,0 +1,306 @@ +/- +Copyright (c) 2026 Samuel Schlesinger. All rights reserved. +Released under Apache 2.0 license as described in the file LICENSE. +Authors: Samuel Schlesinger +-/ + +module +-- upstream: candidate + +public import Complexity.Computability.WordRAM.Notation +public import Complexity.Foundations.NumberTheory.AKS.Coefficients + +/-! +# Writing and comparing coefficient arrays on the word RAM + +Two linear-time array routines for polynomials modulo `X^r - 1` and `n`, stored as `r` words: +`Fill` writes the coefficients `Cslib.AKS.powAddCoeff n r e a` of `X^e + a` (it clears the block +and then adds `a` at index `0` and `1` at index `e mod r`), and `Compare` tests two blocks for +equality. + +## Main definitions + +* `Cslib.Algorithms.WordRAM.AKS.Fill.program`: write the coefficients of `X^e + a`. +* `Cslib.Algorithms.WordRAM.AKS.Compare.program`: compare two blocks. + +## Main results + +* `Cslib.Algorithms.WordRAM.AKS.Fill.exec`: the block receives `Cslib.AKS.powAddCoeff` within + `5 r + 13` instructions, and no other cell changes. +* `Cslib.Algorithms.WordRAM.AKS.Compare.exec`: the flag is `1` exactly when the blocks agree, + within `10 r + 5` instructions, and memory is unchanged. +-/ + +@[expose] public section + +namespace Cslib.Algorithms.WordRAM.AKS + +open Cslib.WordRAM + +namespace Fill + +/-- Base address of the block. -/ +abbrev base : Fin 10 := 0 +/-- The number `r` of coefficients. -/ +abbrev rr : Fin 10 := 1 +/-- The modulus `n`. -/ +abbrev nn : Fin 10 := 2 +/-- The index `e mod r` of the monomial. -/ +abbrev er : Fin 10 := 3 +/-- The constant term `a`. -/ +abbrev aa : Fin 10 := 4 +/-- Loop index. -/ +abbrev i : Fin 10 := 5 +/-- Address scratch. -/ +abbrev addr : Fin 10 := 6 +/-- Value scratch. -/ +abbrev t : Fin 10 := 7 +/-- Constant one. -/ +abbrev one : Fin 10 := 8 +/-- Constant zero. -/ +abbrev zero : Fin 10 := 9 + +/-- Clear one cell. -/ +def clearBody : Code .standard 10 := ram[ addr ← base + i; *addr ← zero; i ← i + one ] + +/-- Store `a mod n` at index `0`. -/ +def storeConst : Code .standard 10 := ram[ t ← aa % nn; *base ← t ] + +/-- Add `1` at index `e mod r`, modulo `n`. -/ +def addMonomial : Code .standard 10 := ram[ + addr ← base + er; t ← *addr; t ← t + one; t ← t % nn; *addr ← t ] + +/-- Store `a mod n` at index `0`, then add `1` at index `e mod r`. -/ +def finish : Code .standard 10 := .seq storeConst addMonomial + +/-- Clear the block, then write the two nonzero coefficients. -/ +def program : Code .standard 10 := + .seq (ram[ one ← #1; zero ← #0; i ← #0 ]) (.seq (.loop (.lt i rr) clearBody) finish) + +private theorem clear_step {w b r : ℕ} {m : ℕ → ℕ} (hfb : Fits w (b + r)) (y : Data 10) + (hy : y.regs base = b ∧ y.regs rr = r ∧ y.regs one = 1 ∧ y.regs zero = 0 ∧ y.regs i ≤ r ∧ + ∀ a', y.mem a' = if b ≤ a' ∧ a' < b + y.regs i then 0 else m a') + (hc : (Cond.lt i rr).Holds y) : + ∃ z, ExecWithin w clearBody y z 3 ∧ + (z.regs base = b ∧ z.regs rr = r ∧ z.regs one = 1 ∧ z.regs zero = 0 ∧ z.regs i ≤ r ∧ + ∀ a', z.mem a' = if b ≤ a' ∧ a' < b + z.regs i then 0 else m a') ∧ + r - z.regs i < r - y.regs i := by + obtain ⟨hyb, hyr, hyo, hyz, hyi, hym⟩ := hy + obtain ⟨I, hI⟩ : ∃ I, y.regs i = I := ⟨_, rfl⟩ + have hlt : I < r := by simpa only [Cond.holds_lt, hyr, hI] using hc + have hf₁ : Fits w (b + I) := hfb.of_le (by omega) + have hf₂ : Fits w (I + 1) := hfb.of_le (by omega) + let z := ((y.set addr (b + I)).write (b + I) 0).set i (I + 1) + have hrun : ExecWithin (B := .standard) w clearBody y z 3 := by + ram_eval [clearBody, z, hyb, hI, hyo, hyz, hf₁, hf₂] + refine ⟨z, hrun, ⟨?_, ?_, ?_, ?_, ?_, ?_⟩, ?_⟩ + · simp +decide [z, hyb] + · simp +decide [z, hyr] + · simp +decide [z, hyo] + · simp +decide [z, hyz] + · simp only [z, Data.set_regs, Function.update_self]; omega + · intro a' + simp only [z, Data.set_mem, Data.write_mem, Data.set_regs, Function.update_self, + Function.update_apply, hym, hI] + split_ifs <;> first | rfl | omega + · simp only [z, Data.set_regs, Function.update_self, hI]; omega + +/-- **Writing `X^e + a`.** For `r ≥ 1` and `n ≥ 1`, with `e mod r` in `er` and the block end, +`a`, `n` and `n + 1` fitting a word, the block `[b, b + r)` receives the coefficients +`Cslib.AKS.powAddCoeff n r e a` within `5 r + 13` instructions, and no other cell changes. -/ +theorem exec {w : ℕ} (x : Data 10) {b r n e a : ℕ} (hb : x.regs base = b) (hr : x.regs rr = r) + (hn : x.regs nn = n) (her : x.regs er = e % r) (ha : x.regs aa = a) (hr0 : 0 < r) + (hn0 : 0 < n) (hfb : Fits w (b + r)) (hfa : Fits w a) (hfn : Fits w (n + 1)) : + ∃ y, ExecWithin w program x y (5 * r + 13) ∧ + (∀ k < r, y.mem (b + k) = Cslib.AKS.powAddCoeff n r e a k) ∧ + (∀ a', ¬(b ≤ a' ∧ a' < b + r) → y.mem a' = x.mem a') ∧ + y.regs base = b ∧ y.regs rr = r ∧ y.regs nn = n := by + have h1 : Fits w 1 := hfn.of_le (by omega) + let x₀ := ((x.set one 1).set zero 0).set i 0 + have hrun₀ : ExecWithin (B := .standard) w (ram[ one ← #1; zero ← #0; i ← #0 ]) x x₀ 3 := by + ram_eval [x₀, h1] + have hx₀ : x₀.regs base = b ∧ x₀.regs rr = r ∧ x₀.regs one = 1 ∧ x₀.regs zero = 0 ∧ + x₀.regs i ≤ r ∧ ∀ a', x₀.mem a' = if b ≤ a' ∧ a' < b + x₀.regs i then 0 else x.mem a' := by + refine ⟨?_, ?_, ?_, ?_, ?_, fun a' => ?_⟩ + · simp +decide [x₀, hb] + · simp +decide [x₀, hr] + · simp +decide [x₀] + · simp +decide [x₀] + · simp +decide [x₀] + · simp only [x₀, Data.set_mem, Data.set_regs, Function.update_self] + split_ifs <;> first | rfl | omega + obtain ⟨y, hrun₁, ⟨hyb, hyr, hyo, -, hyi, hym⟩, hc⟩ := ExecWithin.loop_invariant _ + (fun y => r - y.regs i) (clear_step hfb) hx₀ + have hir : y.regs i = r := by simp only [Cond.holds_lt, hyr] at hc; omega + rw [hir] at hym + -- The two nonzero coefficients. + have hyn : y.regs nn = n := by + rw [hrun₁.regs_eq (by decide), show x₀.regs nn = x.regs nn by simp +decide [x₀], hn] + have hya : y.regs aa = a := by + rw [hrun₁.regs_eq (by decide), show x₀.regs aa = x.regs aa by simp +decide [x₀], ha] + have hye : y.regs er = e % r := by + rw [hrun₁.regs_eq (by decide), show x₀.regs er = x.regs er by simp +decide [x₀], her] + have her' : e % r < r := Nat.mod_lt _ hr0 + have hamod : a % n < n := Nat.mod_lt _ hn0 + have hf₁ : Fits w (a % n) := hfa.of_le (Nat.mod_le _ _) + have hf₂ : Fits w (b + e % r) := hfb.of_le (by omega) + -- The cell at `e mod r` after the first store. + obtain ⟨c, hc⟩ : ∃ c, c = if e % r = 0 then a % n else 0 := ⟨_, rfl⟩ + obtain ⟨y', hy'⟩ : ∃ y', (y.set t (a % n)).write b (a % n) = y' := ⟨_, rfl⟩ + have hrun₂ : ExecWithin (B := .standard) w storeConst y y' 2 := by + rw [← hy']; ram_eval [storeConst, hya, hyn, hyb, hf₁] + have hy'b : y'.regs base = b := by rw [← hy']; simp +decide [hyb] + have hy'e : y'.regs er = e % r := by rw [← hy']; simp +decide [hye] + have hy'o : y'.regs one = 1 := by rw [← hy']; simp +decide [hyo] + have hy'n : y'.regs nn = n := by rw [← hy']; simp +decide [hyn] + have hy'r : y'.regs rr = r := by rw [← hy']; simp +decide [hyr] + have hy'm : y'.mem = Function.update y.mem b (a % n) := by rw [← hy']; rfl + have hcell : y'.mem (b + e % r) = c := by + rw [hy'm] + by_cases h : e % r = 0 + · simp [hc, h] + · rw [Function.update_of_ne (by omega), hym] + simp [hc, h, show b + e % r < b + r by omega] + have hcn : c < n := by rw [hc]; split_ifs <;> omega + have hf₃ : Fits w (c + 1) := hfn.of_le (by omega) + have hf₄ : Fits w ((c + 1) % n) := hfn.of_le (by have := Nat.mod_lt (c + 1) hn0; omega) + let z := ((((y'.set addr (b + e % r)).set t c).set t (c + 1)).set t ((c + 1) % n)).write + (b + e % r) ((c + 1) % n) + have hrun₃ : ExecWithin (B := .standard) w addMonomial y' z 5 := by + ram_eval [addMonomial, z, hy'b, hy'e, hy'o, hy'n, hcell, hf₂, hf₃, hf₄] + refine ⟨z, (hrun₀.seq (hrun₁.seq (hrun₂.seq hrun₃))).mono ?_, ?_, ?_, ?_, ?_, ?_⟩ + · have : x₀.regs i = 0 := by simp +decide [x₀] + rw [this]; omega + · intro k hk + simp only [z, Data.write_mem, Data.set_mem] + by_cases hke : k = e % r + · subst hke + rw [Function.update_self] + simp only [Cslib.AKS.powAddCoeff, ↓reduceIte, hc] + split_ifs with h0 + · rw [Nat.mod_add_mod, Nat.add_comm] + · simp + · rw [Function.update_of_ne (by omega), hy'm] + by_cases hk0 : k = 0 + · subst hk0 + simp only [Nat.add_zero, Function.update_self, Cslib.AKS.powAddCoeff, hke, ↓reduceIte, + Nat.zero_add] + · rw [Function.update_of_ne (by omega), hym] + simp [Cslib.AKS.powAddCoeff, hke, hk0, show b + k < b + r by omega] + · intro a' ha' + simp only [z, Data.write_mem, Data.set_mem] + rw [Function.update_of_ne (by omega), hy'm, Function.update_of_ne (by omega), hym] + simp [ha'] + · simp +decide [z, hy'b] + · simp +decide [z, hy'r] + · simp +decide [z, hy'n] + +end Fill + +namespace Compare + +/-- Base address of the first block. -/ +abbrev p : Fin 9 := 0 +/-- Base address of the second block. -/ +abbrev q : Fin 9 := 1 +/-- The number `r` of coefficients. -/ +abbrev rr : Fin 9 := 2 +/-- Loop index. -/ +abbrev k : Fin 9 := 3 +/-- Address scratch. -/ +abbrev addr : Fin 9 := 4 +/-- Coefficient of the first block. -/ +abbrev u : Fin 9 := 5 +/-- Coefficient of the second block. -/ +abbrev v : Fin 9 := 6 +/-- The answer: `1` when the blocks agree. -/ +abbrev flag : Fin 9 := 7 +/-- Constant one. -/ +abbrev one : Fin 9 := 8 + +/-- Compare one pair of coefficients. -/ +def body : Code .standard 9 := ram[ + addr ← p + k; u ← *addr; addr ← q + k; v ← *addr; + if u = v then [ flag ← flag ] else [ flag ← #0 ]; k ← k + one ] + +/-- Compare two blocks of `r` words. -/ +def program : Code .standard 9 := + .seq (ram[ flag ← #1; one ← #1; k ← #0 ]) (.loop (.lt k rr) body) + +private theorem step {w a b r : ℕ} {m : ℕ → ℕ} (hfa : Fits w (a + r)) (hfb : Fits w (b + r)) + (y : Data 9) + (hy : y.regs p = a ∧ y.regs q = b ∧ y.regs rr = r ∧ y.regs one = 1 ∧ y.regs k ≤ r ∧ + y.mem = m ∧ (y.regs flag = if ∀ k' < y.regs k, m (a + k') = m (b + k') then 1 else 0)) + (hc : (Cond.lt k rr).Holds y) : + ∃ z, ExecWithin w body y z 9 ∧ + (z.regs p = a ∧ z.regs q = b ∧ z.regs rr = r ∧ z.regs one = 1 ∧ z.regs k ≤ r ∧ + z.mem = m ∧ (z.regs flag = if ∀ k' < z.regs k, m (a + k') = m (b + k') then 1 else 0)) ∧ + r - z.regs k < r - y.regs k := by + classical + obtain ⟨hya, hyb, hyr, hyo, hyk, hym, hyf⟩ := hy + obtain ⟨K, hK⟩ : ∃ K, y.regs k = K := ⟨_, rfl⟩ + have hlt : K < r := by simpa only [Cond.holds_lt, hyr, hK] using hc + rw [hK] at hyf + have hf₁ : Fits w (a + K) := hfa.of_le (by omega) + have hf₂ : Fits w (b + K) := hfb.of_le (by omega) + have hf₃ : Fits w (K + 1) := hfa.of_le (by omega) + have hiff : (∀ k' < K + 1, m (a + k') = m (b + k')) ↔ + (∀ k' < K, m (a + k') = m (b + k')) ∧ m (a + K) = m (b + K) := by + constructor + · intro h; exact ⟨fun k' hk' => h k' (by omega), h K (by omega)⟩ + · rintro ⟨h, hK'⟩ k' hk' + rcases Nat.lt_succ_iff_lt_or_eq.mp hk' with hk' | rfl + · exact h k' hk' + · exact hK' + let f := if m (a + K) = m (b + K) then y.regs flag else 0 + let z := ((((((y.set addr (a + K)).set u (m (a + K))).set addr (b + K)).set v + (m (b + K))).set flag f).set k (K + 1)) + have hrun : ExecWithin (B := .standard) w body y z 9 := by + by_cases heq : m (a + K) = m (b + K) + · refine ExecWithin.mono (T := 8) ?_ (by omega) + ram_eval [body, z, f, hya, hyb, hK, hyo, hym, heq, hf₁, hf₂, hf₃] + · ram_eval [body, z, f, hya, hyb, hK, hyo, hym, heq, hf₁, hf₂, hf₃] + refine ⟨z, hrun, ⟨?_, ?_, ?_, ?_, ?_, ?_, ?_⟩, ?_⟩ + · simp +decide [z, hya] + · simp +decide [z, hyb] + · simp +decide [z, hyr] + · simp +decide [z, hyo] + · simp +decide [z]; omega + · simp +decide [z, hym] + · simp only [z, f, Data.set_regs, Function.update_self] + rw [Function.update_of_ne (by decide), Function.update_self] + simp only [hiff] + by_cases heq : m (a + K) = m (b + K) <;> simp [heq, hyf] + · simp +decide [z, hK]; omega + +/-- **Comparing blocks.** With both block ends fitting a word, the program leaves `1` in `flag` +exactly when the blocks `[a, a + r)` and `[b, b + r)` agree, and `0` otherwise, within +`11 r + 5` instructions, without changing memory. -/ +theorem exec {w : ℕ} (x : Data 9) {a b r : ℕ} (ha : x.regs p = a) (hb : x.regs q = b) + (hr : x.regs rr = r) (hfa : Fits w (a + r)) (hfb : Fits w (b + r)) (h1 : Fits w 1) : + ∃ y, ExecWithin w program x y (11 * r + 5) ∧ y.regs flag ≤ 1 ∧ + (y.regs flag = 1 ↔ ∀ k' < r, x.mem (a + k') = x.mem (b + k')) ∧ y.mem = x.mem := by + classical + let x₀ := ((x.set flag 1).set one 1).set k 0 + have hrun₀ : ExecWithin (B := .standard) w (ram[ flag ← #1; one ← #1; k ← #0 ]) x x₀ 3 := by + ram_eval [x₀, h1] + have hx₀ : x₀.regs p = a ∧ x₀.regs q = b ∧ x₀.regs rr = r ∧ x₀.regs one = 1 ∧ + x₀.regs k ≤ r ∧ x₀.mem = x.mem ∧ + (x₀.regs flag = if ∀ k' < x₀.regs k, x.mem (a + k') = x.mem (b + k') then 1 else 0) := by + refine ⟨?_, ?_, ?_, ?_, ?_, rfl, ?_⟩ <;> simp +decide [x₀, ha, hb, hr] + obtain ⟨y, hrun₁, ⟨-, -, hyr, -, hyk, hym, hyf⟩, hc⟩ := ExecWithin.loop_invariant _ + (fun y => r - y.regs k) (step hfa hfb) hx₀ + have hkr : y.regs k = r := by simp only [Cond.holds_lt, hyr] at hc; omega + rw [hkr] at hyf + refine ⟨y, (hrun₀.seq hrun₁).mono ?_, ?_, ?_, hym⟩ + · have : x₀.regs k = 0 := by simp +decide [x₀] + rw [this]; omega + · rw [hyf]; split_ifs <;> omega + · rw [hyf] + split_ifs with h + · exact ⟨fun _ => h, fun _ => rfl⟩ + · exact ⟨fun h' => by simp at h', fun h' => absurd h' h⟩ + +end Compare + +end Cslib.Algorithms.WordRAM.AKS diff --git a/ComplexityTests/WordRAMAKS.lean b/ComplexityTests/WordRAMAKS.lean index 32e95b9dc..f71da24c4 100644 --- a/ComplexityTests/WordRAMAKS.lean +++ b/ComplexityTests/WordRAMAKS.lean @@ -7,12 +7,14 @@ Authors: Samuel Schlesinger import Complexity.Algorithms.WordRAM.NatGcd import Complexity.Algorithms.WordRAM.AKS.PerfectPower import Complexity.Algorithms.WordRAM.AKS.PolyMul +import Complexity.Algorithms.WordRAM.AKS.PolyArray /-! # Word-RAM subroutines of the AKS primality test -Concrete runs of Euclid's algorithm, the integer `k`-th root, the perfect-power test and -multiplication modulo `X^r - 1` and `n`, and instances of their correctness theorems. +Concrete runs of Euclid's algorithm, the integer `k`-th root, the perfect-power test, +multiplication modulo `X^r - 1` and `n`, and writing and comparing coefficient arrays, and +instances of their correctness theorems. -/ namespace ComplexityTests.WordRAMAKS @@ -106,4 +108,38 @@ example : ∃ y, ExecWithin 16 AKS.PolyMul.program (mulInput 10 10) y 224 ∧ y. rw [show (21 : ℕ) = 20 + 1 from rfl, hcoeff 1 (by decide)] decide +/-! ## Writing and comparing coefficient arrays -/ + +private def fillInput (e a : ℕ) : Data 10 := + ⟨fun j => if j = AKS.Fill.base then 20 else if j = AKS.Fill.rr then 3 else + if j = AKS.Fill.nn then 5 else if j = AKS.Fill.er then e % 3 else + if j = AKS.Fill.aa then a else 0, + fun a => if a = 23 then 99 else 7⟩ + +/-- `X^4 + 7 ≡ X + 2` and `X^3 + 4 ≡ 0` modulo `X³ - 1` and `5`, over dirty memory. -/ +example : (Code.eval 16 4 AKS.Fill.program (fillInput 4 7)).map + (fun yn => [yn.1.mem 20, yn.1.mem 21, yn.1.mem 22, yn.1.mem 23]) = some [2, 1, 0, 99] ∧ + (Code.eval 16 4 AKS.Fill.program (fillInput 3 4)).map + (fun yn => [yn.1.mem 20, yn.1.mem 21, yn.1.mem 22]) = some [0, 0, 0] := by + decide +kernel + +/-- The correctness theorem agrees with `Cslib.AKS.powAddCoeff`. -/ +example : ∃ y, ExecWithin 16 AKS.Fill.program (fillInput 4 7) y 28 ∧ y.mem 20 = 2 := by + obtain ⟨y, hy, hc, -⟩ := AKS.Fill.exec (w := 16) (fillInput 4 7) (b := 20) (r := 3) (n := 5) + (e := 4) (a := 7) rfl rfl rfl rfl rfl (by decide) (by decide) (by decide) (by decide) + (by decide) + exact ⟨y, hy, (hc 0 (by decide)).trans (by decide)⟩ + +private def cmpInput : Data 9 := + ⟨fun j => if j = AKS.Compare.p then 10 else if j = AKS.Compare.q then 20 else + if j = AKS.Compare.rr then 3 else 0, + fun a => if a = 12 then 5 else if a = 22 then 6 else 1⟩ + +/-- Blocks agreeing on two cells but not the third compare unequal; shorter prefixes agree. -/ +example : (Code.eval 16 4 AKS.Compare.program cmpInput).map (·.1.regs AKS.Compare.flag) = + some 0 ∧ + (Code.eval 16 4 AKS.Compare.program (cmpInput.set AKS.Compare.rr 2)).map + (·.1.regs AKS.Compare.flag) = some 1 := by + decide + end ComplexityTests.WordRAMAKS From 99fd3bcaa03666baa3fdbfdc940d8d0eef3c6a55 Mon Sep 17 00:00:00 2001 From: Samuel Giuseppe Schlesinger Date: Fri, 30 Oct 2026 21:56:08 -0400 Subject: [PATCH 3/4] fix(WordRAM/AKS): wrap long test line and correct two docstrings (#893) Wrap the 101-character line in ComplexityTests/WordRAMAKS.lean that failed lake test --wfail. Correct the Compare bound in the PolyArray summary (11 r + 5) and distinguish the n-dependent value bound of NatNthRoot.exec from its k-dependent counters. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_015DLg8kwPs8yi7ePPdtoNvG --- Complexity/Algorithms/WordRAM/AKS/PolyArray.lean | 2 +- Complexity/Algorithms/WordRAM/NatNthRoot.lean | 3 ++- ComplexityTests/WordRAMAKS.lean | 3 ++- 3 files changed, 5 insertions(+), 3 deletions(-) diff --git a/Complexity/Algorithms/WordRAM/AKS/PolyArray.lean b/Complexity/Algorithms/WordRAM/AKS/PolyArray.lean index 4128f38cf..ff4e015b8 100644 --- a/Complexity/Algorithms/WordRAM/AKS/PolyArray.lean +++ b/Complexity/Algorithms/WordRAM/AKS/PolyArray.lean @@ -28,7 +28,7 @@ equality. * `Cslib.Algorithms.WordRAM.AKS.Fill.exec`: the block receives `Cslib.AKS.powAddCoeff` within `5 r + 13` instructions, and no other cell changes. * `Cslib.Algorithms.WordRAM.AKS.Compare.exec`: the flag is `1` exactly when the blocks agree, - within `10 r + 5` instructions, and memory is unchanged. + within `11 r + 5` instructions, and memory is unchanged. -/ @[expose] public section diff --git a/Complexity/Algorithms/WordRAM/NatNthRoot.lean b/Complexity/Algorithms/WordRAM/NatNthRoot.lean index 943d9eb06..e8c4f6505 100644 --- a/Complexity/Algorithms/WordRAM/NatNthRoot.lean +++ b/Complexity/Algorithms/WordRAM/NatNthRoot.lean @@ -195,7 +195,8 @@ private theorem clog_halve {d : ℕ} (hd : 2 ≤ d) {d' : ℕ} (hd' : d' ≤ (d exact Nat.lt_succ_of_le (Nat.clog_mono_right 2 (by omega)) /-- Compute `Nat.nthRoot k n` for `k ≥ 1` within `⌈log₂ (n + 1)⌉ (6 k + 13) + 7` instructions, -preserving `n`, `k` and memory. Every intermediate value is at most `n² + 2 n + 1`. -/ +preserving `n`, `k` and memory. Products and search bounds are at most `n² + 2 n + 1`, but the +probe counters run up to `k` whatever `n` is, hence the separate hypothesis `hkfit`. -/ theorem exec {w : ℕ} (x : Data 10) (hk : 1 ≤ x.regs k) (hkfit : Fits w (x.regs k)) (hfit : Fits w (x.regs n * x.regs n + 2 * x.regs n + 2)) : ∃ y, ExecWithin w program x y (Nat.clog 2 (x.regs n + 1) * (6 * x.regs k + 13) + 7) ∧ diff --git a/ComplexityTests/WordRAMAKS.lean b/ComplexityTests/WordRAMAKS.lean index f71da24c4..3988d2bf9 100644 --- a/ComplexityTests/WordRAMAKS.lean +++ b/ComplexityTests/WordRAMAKS.lean @@ -45,7 +45,8 @@ private def rootInput (n k : ℕ) : Data 10 := example : (Code.eval 16 16 NatNthRoot.program (rootInput 30 2)).map (·.1.regs NatNthRoot.lo) = some 5 ∧ (Code.eval 16 16 NatNthRoot.program (rootInput 64 3)).map (·.1.regs NatNthRoot.lo) = some 4 ∧ - (Code.eval 16 16 NatNthRoot.program (rootInput 63 3)).map (·.1.regs NatNthRoot.lo) = some 3 := by + (Code.eval 16 16 NatNthRoot.program (rootInput 63 3)).map (·.1.regs NatNthRoot.lo) = + some 3 := by decide /-- The correctness theorem: the search returns `Nat.nthRoot 3 63`. -/ From a3649ca57f9661fcc0ec1e131603eba355fe3b04 Mon Sep 17 00:00:00 2001 From: Samuel Giuseppe Schlesinger Date: Fri, 30 Oct 2026 22:14:14 -0400 Subject: [PATCH 4/4] feat(WordRAM/AKS): repeated squaring and the congruence test (#893) PolyPow raises a coefficient array modulo X^r - 1 and n to the power e by repeated squaring in Nat.size e rounds, reusing PolyMul and ArrayCopy. Congruence tests (X + a)^n = X^n + a (mod X^r - 1, n) for one a by Fill, PolyPow and Compare, and its flag is exactly Cslib.AKS.Congruence n r a. Both have instruction bounds independent of the word size. Adds Cslib.AKS.toQuot_congr. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_015DLg8kwPs8yi7ePPdtoNvG --- Complexity.lean | 2 + .../Algorithms/WordRAM/AKS/Congruence.lean | 259 ++++++++++++ .../Algorithms/WordRAM/AKS/PolyPow.lean | 383 ++++++++++++++++++ .../NumberTheory/AKS/Coefficients.lean | 5 + ComplexityTests/WordRAMAKSCongruence.lean | 71 ++++ 5 files changed, 720 insertions(+) create mode 100644 Complexity/Algorithms/WordRAM/AKS/Congruence.lean create mode 100644 Complexity/Algorithms/WordRAM/AKS/PolyPow.lean create mode 100644 ComplexityTests/WordRAMAKSCongruence.lean diff --git a/Complexity.lean b/Complexity.lean index b44bd1eb8..2af464aa6 100644 --- a/Complexity.lean +++ b/Complexity.lean @@ -80,9 +80,11 @@ public import Complexity.Algorithms.ThreeSum.ChanHe.Sizes public import Complexity.Algorithms.ThreeSum.Convolution public import Complexity.Algorithms.ThreeSum.Defs public import Complexity.Algorithms.ThreeSum.ThreeSets +public import Complexity.Algorithms.WordRAM.AKS.Congruence public import Complexity.Algorithms.WordRAM.AKS.PerfectPower public import Complexity.Algorithms.WordRAM.AKS.PolyArray public import Complexity.Algorithms.WordRAM.AKS.PolyMul +public import Complexity.Algorithms.WordRAM.AKS.PolyPow public import Complexity.Algorithms.WordRAM.Deduplicate.Count public import Complexity.Algorithms.WordRAM.Deduplicate.Model public import Complexity.Algorithms.WordRAM.Deduplicate.Program diff --git a/Complexity/Algorithms/WordRAM/AKS/Congruence.lean b/Complexity/Algorithms/WordRAM/AKS/Congruence.lean new file mode 100644 index 000000000..9882c7c40 --- /dev/null +++ b/Complexity/Algorithms/WordRAM/AKS/Congruence.lean @@ -0,0 +1,259 @@ +/- +Copyright (c) 2026 Samuel Schlesinger. All rights reserved. +Released under Apache 2.0 license as described in the file LICENSE. +Authors: Samuel Schlesinger +-/ + +module +-- upstream: candidate + +public import Complexity.Algorithms.WordRAM.AKS.PolyArray +public import Complexity.Algorithms.WordRAM.AKS.PolyPow + +/-! +# The AKS congruence for one `a` on the word RAM + +Step 5 of the AKS algorithm [Agrawal, Kayal, Saxena, Section 4][AKS2004] tests +`(X + a)^n ≡ X^n + a (mod X^r - 1, n)`. The program writes the coefficients of `X + a` into a base +block `B` and those of `1` into a result block `R`, raises by repeated squaring +(`Cslib.Algorithms.WordRAM.AKS.PolyPow.program`) so that `R` holds `(X + a)^n`, writes the +coefficients of `X^n + a` into the scratch block `T`, and compares `R` with `T`. Both blocks are +reduced, so they agree exactly when the classes agree (`Cslib.AKS.toQuot_eq_toQuot_iff`), which is +the congruence (`Cslib.AKS.congruence_iff_toQuot`). + +## Main definitions + +* `Cslib.Algorithms.WordRAM.AKS.Congruence.program`: the congruence test. +* `Cslib.Algorithms.WordRAM.AKS.Congruence.time`: its instruction bound. + +## Main results + +* `Cslib.Algorithms.WordRAM.AKS.Congruence.exec`: the flag is `1` exactly when + `Cslib.AKS.Congruence n r a` holds, within `time r n` instructions. + +## References + +* [M. Agrawal, N. Kayal, N. Saxena, *PRIMES is in P*][AKS2004] +-/ + +@[expose] public section + +namespace Cslib.Algorithms.WordRAM.AKS.Congruence + +open Cslib.WordRAM Cslib.AKS Polynomial + +/-- Base address of the result block. -/ +abbrev R : Fin 25 := 0 +/-- Base address of the base block. -/ +abbrev B : Fin 25 := 1 +/-- Base address of the scratch block. -/ +abbrev T : Fin 25 := 2 +/-- The number `r` of coefficients. -/ +abbrev rr : Fin 25 := 3 +/-- The modulus `n`, which is also the exponent. -/ +abbrev nn : Fin 25 := 4 +/-- The exponent of the power. -/ +abbrev e : Fin 25 := 5 +/-- Block argument of the subroutines. -/ +abbrev pa : Fin 25 := 9 +/-- Second block argument of the comparison. -/ +abbrev pb : Fin 25 := 10 +/-- Constant one. -/ +abbrev one : Fin 25 := 19 +/-- The tested value `a`, preserved. -/ +abbrev av : Fin 25 := 21 +/-- The exponent residue for `Fill`. -/ +abbrev er : Fin 25 := 22 +/-- The constant term for `Fill`. -/ +abbrev aa : Fin 25 := 23 +/-- The answer: `1` when the congruence holds. -/ +abbrev flag : Fin 25 := 24 + +/-- The power runs on the first `21` registers. -/ +def powMap : Fin 21 → Fin 25 := Fin.castLE (by norm_num) + +/-- `Fill` writes the block at `pa`, with scratch shared with the power. -/ +def fillMap : Fin 10 → Fin 25 := ![pa, rr, nn, er, aa, 12, 15, 16, one, 20] + +/-- `Compare` compares the blocks at `pa` and `pb`. -/ +def cmpMap : Fin 9 → Fin 25 := ![pa, pb, rr, 12, 15, 16, 17, flag, one] + +theorem powMap_injective : powMap.Injective := Fin.castLE_injective _ +theorem fillMap_injective : fillMap.Injective := by decide +theorem cmpMap_injective : cmpMap.Injective := by decide + +/-- Write `X^e + a` into the block at `pa`. -/ +def fill : Code .standard 25 := Fill.program.rename fillMap + +/-- Raise `R` by `B` to the power in `e`. -/ +def pow : Code .standard 25 := PolyPow.program.rename powMap + +/-- Compare the blocks at `pa` and `pb`. -/ +def cmp : Code .standard 25 := Compare.program.rename cmpMap + +/-- Write `X + a` into `B` and `1` into `R`, raise to the power `n`, write `X^n + a` into `T`, and +compare `R` with `T`. -/ +def program : Code .standard 25 := + .seq (ram[ one ← #1; er ← one % rr; aa ← av; pa ← B ]) (.seq fill + (.seq (ram[ er ← #0; aa ← #0; pa ← R ]) (.seq fill + (.seq (ram[ e ← nn ]) (.seq pow + (.seq (ram[ er ← nn % rr; aa ← av; pa ← T ]) (.seq fill + (.seq (ram[ pa ← R; pb ← T ]) cmp)))))))) + +/-- The instruction bound of `program`. -/ +def time (r n : ℕ) : ℕ := PolyPow.time r n + 26 * r + 57 + +/-- The registers that every subroutine call preserves. -/ +def Keep (j : Fin 25) : Prop := j = R ∨ j = B ∨ j = T ∨ j = rr ∨ j = nn ∨ j = av + +instance : DecidablePred Keep := fun _ => by unfold Keep; infer_instance + +private theorem fill_preserves : ∀ j, Keep j → fill.Preserves j := by decide +private theorem pow_preserves : ∀ j, Keep j → pow.Preserves j := by decide +private theorem cmp_preserves : ∀ j, Keep j → cmp.Preserves j := by decide + +private theorem fill_exec {w : ℕ} (x : Data 25) {bs r n e' A : ℕ} (hb : x.regs pa = bs) + (hr : x.regs rr = r) (hn : x.regs nn = n) (her : x.regs er = e' % r) (ha : x.regs aa = A) + (hr0 : 0 < r) (hn0 : 0 < n) (hfb : Fits w (bs + r)) (hfa : Fits w A) (hfn : Fits w (n + 1)) : + ∃ y, ExecWithin w fill x y (5 * r + 13) ∧ + (∀ k < r, y.mem (bs + k) = powAddCoeff n r e' A k) ∧ + (∀ a', ¬(bs ≤ a' ∧ a' < bs + r) → y.mem a' = x.mem a') ∧ + (∀ j, Keep j → y.regs j = x.regs j) := by + obtain ⟨y, hy, hblk, hframe, -⟩ := Fill.exec (x.comap fillMap) hb hr hn her ha hr0 hn0 hfb hfa hfn + have hy' := hy.rename fillMap_injective + exact ⟨_, hy', hblk, hframe, fun j hj => hy'.regs_eq (fill_preserves j hj)⟩ + +private theorem cmp_exec {w : ℕ} (x : Data 25) {a b r : ℕ} (ha : x.regs pa = a) + (hb : x.regs pb = b) (hr : x.regs rr = r) (hfa : Fits w (a + r)) (hfb : Fits w (b + r)) + (h1 : Fits w 1) : + ∃ y, ExecWithin w cmp x y (11 * r + 5) ∧ y.regs flag ≤ 1 ∧ + (y.regs flag = 1 ↔ ∀ k' < r, x.mem (a + k') = x.mem (b + k')) ∧ + (∀ j, Keep j → y.regs j = x.regs j) := by + obtain ⟨y, hy, hle, hiff, -⟩ := Compare.exec (x.comap cmpMap) ha hb hr hfa hfb h1 + have hy' := hy.rename cmpMap_injective + have hfl : (y.extend cmpMap x).regs flag = y.regs Compare.flag := + cmpMap_injective.extend_apply _ _ Compare.flag + exact ⟨_, hy', by rw [hfl]; exact hle, by rw [hfl]; exact hiff, + fun j hj => hy'.regs_eq (cmp_preserves j hj)⟩ + +private theorem pow_exec {w : ℕ} (x : Data 25) {a b t r n : ℕ} (ha : x.regs R = a) + (hb : x.regs B = b) (ht : x.regs T = t) (hr : x.regs rr = r) (hn : x.regs nn = n) + (hn0 : 0 < n) (hab : a + r ≤ b ∨ b + r ≤ a) (hta : t + r ≤ a ∨ a + r ≤ t) + (htb : t + r ≤ b ∨ b + r ≤ t) (hpa : ∀ i < r, x.mem (a + i) < n) + (hpb : ∀ i < r, x.mem (b + i) < n) (hfa : Fits w (a + r)) (hfb : Fits w (b + r)) + (hft : Fits w (t + r)) (hf2 : Fits w (2 * r)) (hfn : Fits w (n * n + n)) + (hfe : Fits w (x.regs e)) : + ∃ y, ExecWithin w pow x y (PolyPow.time r (x.regs e)) ∧ + toQuot n r (fun i => y.mem (a + i)) = + toQuot n r (fun i => x.mem (a + i)) * toQuot n r (fun i => x.mem (b + i)) ^ x.regs e ∧ + (∀ i < r, y.mem (a + i) < n) ∧ (∀ j, Keep j → y.regs j = x.regs j) := by + obtain ⟨y, hy, hq, hred, -⟩ := PolyPow.exec (x.comap powMap) ha hb ht hr hn hn0 hab hta htb hpa + hpb hfa hfb hft hf2 hfn hfe + have hy' := hy.rename powMap_injective + exact ⟨_, hy', hq, hred, fun j hj => hy'.regs_eq (pow_preserves j hj)⟩ + +/-- The class of `1`, written as `X^0 + 0`. -/ +private theorem toQuot_one {n r : ℕ} (hr : 0 < r) : toQuot n r (powAddCoeff n r 0 0) = 1 := by + rw [toQuot_powAddCoeff hr, pow_zero, Nat.cast_zero, C_0, add_zero, map_one] + +/-- **The AKS congruence for one `a`.** For `r ≥ 1` and `n ≥ 1`, with the result, base and scratch +blocks `[a, a + r)`, `[b, b + r)` and `[t, t + r)` pairwise disjoint, and the word size holding the +block ends, `2 r`, `n² + n` and the tested value `A`, the program leaves `1` in `flag` if +`(X + A)^n ≡ X^n + A (mod X^r - 1, n)` and `0` otherwise, within `time r n` instructions. -/ +theorem exec {w : ℕ} (x : Data 25) {a b t r n A : ℕ} (ha : x.regs R = a) (hb : x.regs B = b) + (ht : x.regs T = t) (hr : x.regs rr = r) (hn : x.regs nn = n) (hA : x.regs av = A) + (hr0 : 0 < r) (hn0 : 0 < n) (hab : a + r ≤ b ∨ b + r ≤ a) (hta : t + r ≤ a ∨ a + r ≤ t) + (htb : t + r ≤ b ∨ b + r ≤ t) (hfa : Fits w (a + r)) (hfb : Fits w (b + r)) + (hft : Fits w (t + r)) (hf2 : Fits w (2 * r)) (hfn : Fits w (n * n + n)) (hfA : Fits w A) : + ∃ y, ExecWithin w program x y (time r n) ∧ y.regs flag ≤ 1 ∧ + (y.regs flag = 1 ↔ Congruence n r A) ∧ + ∀ j, Keep j → y.regs j = x.regs j := by + have h1 : Fits w 1 := hfn.of_le (by nlinarith) + have hfn1 : Fits w (n + 1) := hfn.of_le (by nlinarith) + have hfnn : Fits w n := hfn.of_le (by nlinarith) + have hfr : Fits w r := hf2.of_le (by omega) + -- `B := X + A`. + let x₁ := (((x.set one 1).set er (1 % r)).set aa A).set pa b + have hrun₁ : ExecWithin (B := .standard) w (ram[ one ← #1; er ← one % rr; aa ← av; pa ← B ]) + x x₁ 4 := by + ram_eval [x₁, hr, hA, hb, h1, h1.of_le (Nat.mod_le 1 r)] + have hx₁ (j : Fin 25) (hj : Keep j) : x₁.regs j = x.regs j := by + rcases hj with rfl | rfl | rfl | rfl | rfl | rfl <;> simp +decide [x₁] + obtain ⟨y₁, hrun₂, hB₁, hfr₁, hy₁⟩ := fill_exec (w := w) x₁ (bs := b) (r := r) (n := n) (e' := 1) + (A := A) (by simp +decide [x₁]) (by rw [hx₁ rr (by simp [Keep]), hr]) + (by rw [hx₁ nn (by simp [Keep]), hn]) (by simp +decide [x₁]) (by simp +decide [x₁]) hr0 hn0 + hfb hfA hfn1 + have hy₁' (j : Fin 25) (hj : Keep j) : y₁.regs j = x.regs j := by rw [hy₁ j hj, hx₁ j hj] + -- `R := 1`. + let x₂ := ((y₁.set er 0).set aa 0).set pa a + have hrun₃ : ExecWithin (B := .standard) w (ram[ er ← #0; aa ← #0; pa ← R ]) y₁ x₂ 3 := by + ram_eval [x₂, hy₁' R (by simp [Keep]), ha] + have hx₂ (j : Fin 25) (hj : Keep j) : x₂.regs j = x.regs j := by + rw [← hy₁' j hj] + rcases hj with rfl | rfl | rfl | rfl | rfl | rfl <;> simp +decide [x₂] + obtain ⟨y₂, hrun₄, hR₂, hfr₂, hy₂⟩ := fill_exec (w := w) x₂ (bs := a) (r := r) (n := n) (e' := 0) + (A := 0) (by simp +decide [x₂]) (by rw [hx₂ rr (by simp [Keep]), hr]) + (by rw [hx₂ nn (by simp [Keep]), hn]) (by simp +decide [x₂]) (by simp +decide [x₂]) hr0 hn0 hfa + (fits_zero w) hfn1 + have hy₂' (j : Fin 25) (hj : Keep j) : y₂.regs j = x.regs j := by rw [hy₂ j hj, hx₂ j hj] + have hB₂ (k : ℕ) (hk : k < r) : y₂.mem (b + k) = powAddCoeff n r 1 A k := by + rw [hfr₂ _ (by rcases hab with h | h <;> omega)] + exact hB₁ k hk + -- `R := R · B^n`. + let x₃ := y₂.set e n + have hrun₅ : ExecWithin (B := .standard) w (ram[ e ← nn ]) y₂ x₃ 1 := by + ram_eval [x₃, hy₂' nn (by simp [Keep]), hn] + have hx₃ (j : Fin 25) (hj : Keep j) : x₃.regs j = x.regs j := by + rw [← hy₂' j hj] + rcases hj with rfl | rfl | rfl | rfl | rfl | rfl <;> simp +decide [x₃] + obtain ⟨y₃, hrun₆, hq₃, hred₃, hy₃⟩ := pow_exec (w := w) x₃ (a := a) (b := b) (t := t) (r := r) + (n := n) (by rw [hx₃ R (by simp [Keep]), ha]) (by rw [hx₃ B (by simp [Keep]), hb]) + (by rw [hx₃ T (by simp [Keep]), ht]) (by rw [hx₃ rr (by simp [Keep]), hr]) + (by rw [hx₃ nn (by simp [Keep]), hn]) hn0 hab hta htb + (fun i hi => by rw [show x₃.mem = y₂.mem from rfl, hR₂ i hi]; exact powAddCoeff_lt hn0 _ _ _) + (fun i hi => by rw [show x₃.mem = y₂.mem from rfl, hB₂ i hi]; exact powAddCoeff_lt hn0 _ _ _) + hfa hfb hft hf2 hfn (by simpa [x₃] using hfnn) + have hx₃e : x₃.regs e = n := by simp +decide [x₃] + rw [hx₃e] at hrun₆ hq₃ + have hy₃' (j : Fin 25) (hj : Keep j) : y₃.regs j = x.regs j := by rw [hy₃ j hj, hx₃ j hj] + -- `T := X^n + A`. + let x₄ := ((y₃.set er (n % r)).set aa A).set pa t + have hrun₇ : ExecWithin (B := .standard) w (ram[ er ← nn % rr; aa ← av; pa ← T ]) y₃ x₄ 3 := by + ram_eval [x₄, hy₃' nn (by simp [Keep]), hn, hy₃' rr (by simp [Keep]), hr, + hy₃' av (by simp [Keep]), hA, hy₃' T (by simp [Keep]), ht, + hfnn.of_le (Nat.mod_le _ _)] + have hx₄ (j : Fin 25) (hj : Keep j) : x₄.regs j = x.regs j := by + rw [← hy₃' j hj] + rcases hj with rfl | rfl | rfl | rfl | rfl | rfl <;> simp +decide [x₄] + obtain ⟨y₄, hrun₈, hT₄, hfr₄, hy₄⟩ := fill_exec (w := w) x₄ (bs := t) (r := r) (n := n) (e' := n) + (A := A) (by simp +decide [x₄]) (by rw [hx₄ rr (by simp [Keep]), hr]) + (by rw [hx₄ nn (by simp [Keep]), hn]) (by simp +decide [x₄]) (by simp +decide [x₄]) hr0 hn0 + hft hfA hfn1 + have hy₄' (j : Fin 25) (hj : Keep j) : y₄.regs j = x.regs j := by rw [hy₄ j hj, hx₄ j hj] + have hR₄ (k : ℕ) (hk : k < r) : y₄.mem (a + k) = y₃.mem (a + k) := by + rw [hfr₄ _ (by rcases hta with h | h <;> omega)] + rfl + -- Compare `R` with `T`. + let x₅ := (y₄.set pa a).set pb t + have hrun₉ : ExecWithin (B := .standard) w (ram[ pa ← R; pb ← T ]) y₄ x₅ 2 := by + ram_eval [x₅, hy₄' R (by simp [Keep]), ha, hy₄' T (by simp [Keep]), ht] + have hx₅ (j : Fin 25) (hj : Keep j) : x₅.regs j = x.regs j := by + rw [← hy₄' j hj] + rcases hj with rfl | rfl | rfl | rfl | rfl | rfl <;> simp +decide [x₅] + obtain ⟨y₅, hrun₁₀, hle, hiff, hy₅⟩ := cmp_exec (w := w) x₅ (a := a) (b := t) (r := r) + (by simp +decide [x₅]) (by simp +decide [x₅]) (by rw [hx₅ rr (by simp [Keep]), hr]) hfa hft h1 + refine ⟨y₅, ?_, hle, ?_, fun j hj => by rw [hy₅ j hj, hx₅ j hj]⟩ + · refine (hrun₁.seq (hrun₂.seq (hrun₃.seq (hrun₄.seq (hrun₅.seq (hrun₆.seq (hrun₇.seq + (hrun₈.seq (hrun₉.seq hrun₁₀))))))))).mono ?_ + unfold time + omega + · rw [hiff, congruence_iff_toQuot hr0] + change (∀ k' < r, y₄.mem (a + k') = y₄.mem (t + k')) ↔ _ + have hRred (k : ℕ) (hk : k < r) : y₄.mem (a + k) < n := by rw [hR₄ k hk]; exact hred₃ k hk + have hTred (k : ℕ) (hk : k < r) : y₄.mem (t + k) < n := by + rw [hT₄ k hk]; exact powAddCoeff_lt hn0 _ _ _ + rw [← toQuot_eq_toQuot_iff hn0 hr0 hRred hTred, toQuot_congr hR₄, hq₃, toQuot_congr hT₄, + toQuot_congr (c := fun i => x₃.mem (a + i)) (fun k hk => hR₂ k hk), + toQuot_congr (c := fun i => x₃.mem (b + i)) (fun k hk => hB₂ k hk), toQuot_one hr0, one_mul] + +end Cslib.Algorithms.WordRAM.AKS.Congruence diff --git a/Complexity/Algorithms/WordRAM/AKS/PolyPow.lean b/Complexity/Algorithms/WordRAM/AKS/PolyPow.lean new file mode 100644 index 000000000..4595ec96e --- /dev/null +++ b/Complexity/Algorithms/WordRAM/AKS/PolyPow.lean @@ -0,0 +1,383 @@ +/- +Copyright (c) 2026 Samuel Schlesinger. All rights reserved. +Released under Apache 2.0 license as described in the file LICENSE. +Authors: Samuel Schlesinger +-/ + +module +-- upstream: candidate + +public import Complexity.Algorithms.WordRAM.AKS.PolyMul +public import Complexity.Computability.WordRAM.Copy +public import Complexity.Computability.WordRAM.RegisterFrame +public import Mathlib.Data.Nat.Size + +/-! +# Powers modulo `X^r - 1` and `n` on the word RAM + +Raising a polynomial modulo `X^r - 1` and `n` to the power `e` by repeated squaring: with a result +block `R`, a base block `B` and a scratch block `T`, each round multiplies `R` by `B` into `T` and +copies it back when `e` is odd, squares `B` into `T` and copies it back, and halves `e`. The class +of `R · B^e` is invariant, so when `e` reaches `0` the result block holds the class of +`R · B^e` for the original `R`, `B` and `e`. The loop makes `Nat.size e` rounds, each of `O(r²)` +instructions (`Cslib.Algorithms.WordRAM.AKS.PolyMul.program` and +`Cslib.WordRAM.ArrayCopy.program`). This is the polynomial arithmetic of step 5 of the AKS +algorithm [Agrawal, Kayal, Saxena, Section 4][AKS2004]. + +## Main definitions + +* `Cslib.Algorithms.WordRAM.AKS.PolyPow.program`: the repeated-squaring loop. +* `Cslib.Algorithms.WordRAM.AKS.PolyPow.time`: its instruction bound. + +## Main results + +* `Cslib.Algorithms.WordRAM.AKS.PolyPow.exec`: the result block receives the coefficients of + `R · B^e` within `time r e` instructions. + +## References + +* [M. Agrawal, N. Kayal, N. Saxena, *PRIMES is in P*][AKS2004] +-/ + +@[expose] public section + +namespace Cslib.Algorithms.WordRAM.AKS.PolyPow + +open Cslib.WordRAM Cslib.AKS Polynomial + +/-- Base address of the result block. -/ +abbrev R : Fin 21 := 0 +/-- Base address of the base block. -/ +abbrev B : Fin 21 := 1 +/-- Base address of the scratch block. -/ +abbrev T : Fin 21 := 2 +/-- The number `r` of coefficients. -/ +abbrev rr : Fin 21 := 3 +/-- The modulus `n` of the coefficients. -/ +abbrev nn : Fin 21 := 4 +/-- The exponent, halved every round. -/ +abbrev e : Fin 21 := 5 +/-- The parity of the exponent. -/ +abbrev par : Fin 21 := 6 +/-- Constant two. -/ +abbrev two : Fin 21 := 7 +/-- Constant zero, for the loop test. -/ +abbrev z : Fin 21 := 8 +/-- First argument of the multiplication, and source of the copy. -/ +abbrev pa : Fin 21 := 9 +/-- Second argument of the multiplication, and destination of the copy. -/ +abbrev pb : Fin 21 := 10 +/-- Product block of the multiplication. -/ +abbrev pt : Fin 21 := 11 + +/-- The multiplication runs on `pa`, `pb`, `pt`, `rr`, `nn` and the scratch registers `12, …, 20`. +-/ +def mulMap : Fin 14 → Fin 21 := ![pa, pb, pt, rr, nn, 12, 13, 14, 15, 16, 17, 18, 19, 20] + +/-- The copy runs from `pa` to `pb` over `rr` words, with scratch among `12, …, 20`. -/ +def copyMap : Fin 7 → Fin 21 := ![pa, pb, rr, 12, 15, 16, 19] + +theorem mulMap_injective : mulMap.Injective := by decide +theorem copyMap_injective : copyMap.Injective := by decide + +/-- The multiplication touches no register below `9` other than `rr` and `nn`. -/ +private theorem mulMap_ne {j : Fin 21} (hj : j.val < 9) (hr : j ≠ rr) (hn : j ≠ nn) (a : Fin 14) : + mulMap a ≠ j := by + intro h + subst h + fin_cases a <;> simp_all [mulMap] + +/-- The copy touches no register below `9` other than `rr`. -/ +private theorem copyMap_ne {j : Fin 21} (hj : j.val < 9) (hr : j ≠ rr) (a : Fin 7) : + copyMap a ≠ j := by + intro h + subst h + fin_cases a <;> simp_all [copyMap] + +/-- Multiply the blocks at `x` and `y` into `T`, then copy `T` into the block at `d`. -/ +def mulInto (x y d : Fin 21) : Code .standard 21 := + .seq (ram[ pa ← x; pb ← y; pt ← T ]) (.seq (PolyMul.program.rename mulMap) + (.seq (ram[ pa ← T; pb ← d ]) (ArrayCopy.program.rename copyMap))) + +/-- One round: multiply `R` by `B` if `e` is odd, square `B`, and halve `e`. -/ +def body : Code .standard 21 := ram[ + par ← e % two; if z < par then [ ⟪mulInto R B R⟫ ] else [ par ← par ]; ⟪mulInto B B B⟫; + e ← e / two ] + +/-- Repeated squaring while the exponent is positive. -/ +def program : Code .standard 21 := .seq (ram[ two ← #2; z ← #0 ]) (.loop (.lt z e) body) + +/-- The instruction bound of one multiplication followed by a copy. -/ +def mulTime (r : ℕ) : ℕ := 14 * (r + 1) * (r + 1) + 7 * r + 9 + +/-- The instruction bound of `program`: `Nat.size e` rounds of two multiplications. -/ +def time (r e : ℕ) : ℕ := Nat.size e * (2 * mulTime r + 6) + 4 + +/-- Halving a positive number shortens its binary representation. -/ +theorem size_div_two_lt {m : ℕ} (hm : 0 < m) : Nat.size (m / 2) < Nat.size m := by + have hs : 0 < Nat.size m := Nat.size_pos.mpr hm + have hlt : m < 2 ^ Nat.size m := Nat.lt_size_self m + have : Nat.size (m / 2) ≤ Nat.size m - 1 := by + rw [Nat.size_le] + have h2 : 2 ^ Nat.size m = 2 * 2 ^ (Nat.size m - 1) := by + rw [← pow_succ']; congr 1; omega + omega + omega + +private theorem mulInto_exec {w : ℕ} (X : Data 21) {x y d : Fin 21} (hx : x.val < 9) + (hy : y.val < 9) (hd : d.val < 9) {a b t dd r n : ℕ} (ha : X.regs x = a) (hb : X.regs y = b) + (ht : X.regs T = t) (hdd : X.regs d = dd) (hr : X.regs rr = r) (hn : X.regs nn = n) + (hn0 : 0 < n) (hta : t + r ≤ a ∨ a + r ≤ t) (htb : t + r ≤ b ∨ b + r ≤ t) + (htd : t + r ≤ dd ∨ dd + r ≤ t) + (hpa : ∀ i < r, X.mem (a + i) < n) (hpb : ∀ j < r, X.mem (b + j) < n) + (hfa : Fits w (a + r)) (hfb : Fits w (b + r)) (hft : Fits w (t + r)) (hfd : Fits w (dd + r)) + (hf2 : Fits w (2 * r)) (hfn : Fits w (n * n + n)) : + ∃ Y, ExecWithin w (mulInto x y d) X Y (mulTime r) ∧ + (∀ k < r, Y.mem (dd + k) = + mulCoeff n r (fun i => X.mem (a + i)) (fun j => X.mem (b + j)) k) ∧ + (∀ a', ¬(t ≤ a' ∧ a' < t + r) → ¬(dd ≤ a' ∧ a' < dd + r) → Y.mem a' = X.mem a') ∧ + (∀ j : Fin 21, j.val < 9 → Y.regs j = X.regs j) := by + have hxa : x ≠ pa := fun h => by subst h; simp at hx + have hya : y ≠ pa := fun h => by subst h; simp at hy + have hyb : y ≠ pb := fun h => by subst h; simp at hy + have hda : d ≠ pa := fun h => by subst h; simp at hd + have hdb : d ≠ pb := fun h => by subst h; simp at hd + have h1 : Fits w 1 := hfn.of_le (by nlinarith) + -- Set the arguments of the multiplication. + let X₁ := ((X.set pa a).set pb b).set pt t + have hrun₁ : ExecWithin (B := .standard) w (ram[ pa ← x; pb ← y; pt ← T ]) X X₁ 3 := by + ram_eval [X₁, ha, hb, ht, hxa, hya, hyb] + have hX₁ (j : Fin 21) (hj : j.val < 9) : X₁.regs j = X.regs j := by + have h9 : j ≠ pa ∧ j ≠ pb ∧ j ≠ pt := by + refine ⟨?_, ?_, ?_⟩ <;> (intro h; subst h; simp at hj) + simp [X₁, h9.1, h9.2.1, h9.2.2] + -- Multiply. + obtain ⟨Y₁, hrun₂, hprod, hframe, -, -, -, hY₁r, hY₁n⟩ := + PolyMul.exec (w := w) (X₁.comap mulMap) (a := a) (b := b) (t := t) (r := r) (n := n) + (by change X₁.regs pa = a; simp +decide [X₁]) (by change X₁.regs pb = b; simp +decide [X₁]) + (by change X₁.regs pt = t; simp +decide [X₁]) + (by change X₁.regs rr = r; rw [hX₁ rr (by simp), hr]) + (by change X₁.regs nn = n; rw [hX₁ nn (by simp), hn]) + hn0 hta htb hpa hpb hfa hfb hft hf2 hfn + have hrun₂' := hrun₂.rename mulMap_injective + set Y₁' := Y₁.extend mulMap X₁ with hY₁' + have hY₁'m : Y₁'.mem = Y₁.mem := rfl + have hY₁'r (j : Fin 21) (hj : j.val < 9) : Y₁'.regs j = X.regs j := by + by_cases hjr : j = rr + · subst hjr + rw [hr, hY₁', show rr = mulMap PolyMul.rr from rfl] + exact (mulMap_injective.extend_apply _ _ _).trans hY₁r + by_cases hjn : j = nn + · subst hjn + rw [hn, hY₁', show nn = mulMap PolyMul.nn from rfl] + exact (mulMap_injective.extend_apply _ _ _).trans hY₁n + rw [hrun₂'.regs_eq (Code.preserves_rename_of_forall_ne _ _ _ (mulMap_ne hj hjr hjn)), + hX₁ j hj] + -- Set the arguments of the copy. + let X₂ := (Y₁'.set pa t).set pb dd + have hrun₃ : ExecWithin (B := .standard) w (ram[ pa ← T; pb ← d ]) Y₁' X₂ 2 := by + ram_eval [X₂, hY₁'r T (by simp), ht, hY₁'r d hd, hdd, hda] + have hX₂ (j : Fin 21) (hj : j.val < 9) : X₂.regs j = X.regs j := by + have h9 : j ≠ pa ∧ j ≠ pb := by + refine ⟨?_, ?_⟩ <;> (intro h; subst h; simp at hj) + simp [X₂, h9.1, h9.2, hY₁'r j hj] + -- Copy. + have hXr : X₂.regs rr = r := by rw [hX₂ rr (by simp), hr] + have hsrc : (X₂.comap copyMap).regs ArrayCopy.src = t := by + change X₂.regs pa = t; simp +decide [X₂] + have hdst : (X₂.comap copyMap).regs ArrayCopy.dst = dd := by + change X₂.regs pb = dd; simp +decide [X₂] + have hlen : (X₂.comap copyMap).regs ArrayCopy.len = r := hXr + obtain ⟨Y₂, hrun₄, hY₂m, -, -, hY₂len⟩ := ArrayCopy.exec_of_separated_or_eq (w := w) + (X₂.comap copyMap) h1 (by rw [hsrc, hlen]; exact hft) (by rw [hdst, hlen]; exact hfd) + (by rw [hsrc, hdst, hlen]; omega) + have hrun₄' := hrun₄.rename copyMap_injective + rw [hlen, hsrc, hdst] at hY₂m + rw [hlen] at hrun₄' + refine ⟨Y₂.extend copyMap X₂, ?_, ?_, ?_, ?_⟩ + · exact (hrun₁.seq (hrun₂'.seq (hrun₃.seq hrun₄'))).mono (by unfold mulTime; omega) + · intro k hk + change Y₂.mem (dd + k) = _ + rw [hY₂m, ArrayCopy.copyPrefix_read hk] + exact hprod k hk + · intro a' hat had + change Y₂.mem a' = _ + rw [hY₂m, ArrayCopy.copyPrefix_outside had] + exact hframe a' hat + · intro j hj + by_cases hjr : j = rr + · subst hjr + change Function.extend copyMap Y₂.regs X₂.regs (copyMap ArrayCopy.len) = X.regs rr + rw [copyMap_injective.extend_apply, hY₂len, hlen, hr] + rw [hrun₄'.regs_eq (Code.preserves_rename_of_forall_ne _ _ _ (copyMap_ne hj hjr)), hX₂ j hj] + +/-- The hypotheses of `exec` on the three blocks, the modulus and the word size. -/ +private structure Hyp (w a b t r n e₀ : ℕ) : Prop where + hn0 : 0 < n + hab : a + r ≤ b ∨ b + r ≤ a + hta : t + r ≤ a ∨ a + r ≤ t + htb : t + r ≤ b ∨ b + r ≤ t + hfa : Fits w (a + r) + hfb : Fits w (b + r) + hft : Fits w (t + r) + hf2 : Fits w (2 * r) + hfn : Fits w (n * n + n) + hfe : Fits w e₀ + +/-- The loop invariant: the registers, the reduced blocks and the class of `R · B^e`. -/ +private def Inv (a b t r n e₀ : ℕ) (target : AdjoinRoot (X ^ r - 1 : (ZMod n)[X])) + (y : Data 21) : Prop := + y.regs R = a ∧ y.regs B = b ∧ y.regs T = t ∧ y.regs rr = r ∧ y.regs nn = n ∧ y.regs two = 2 ∧ + y.regs z = 0 ∧ y.regs e ≤ e₀ ∧ (∀ i < r, y.mem (a + i) < n) ∧ (∀ i < r, y.mem (b + i) < n) ∧ + toQuot n r (fun i => y.mem (a + i)) * toQuot n r (fun i => y.mem (b + i)) ^ y.regs e = target + +private theorem body_step {w a b t r n e₀ : ℕ} (H : Hyp w a b t r n e₀) + {target : AdjoinRoot (X ^ r - 1 : (ZMod n)[X])} (y : Data 21) + (hy : Inv a b t r n e₀ target y) (hc : (Cond.lt z e).Holds y) : + ∃ y', ExecWithin w body y y' (2 * mulTime r + 4) ∧ Inv a b t r n e₀ target y' ∧ + Nat.size (y'.regs e) < Nat.size (y.regs e) := by + obtain ⟨hyR, hyB, hyT, hyr, hyn, hy2, hyz, hye, hRred, hBred, hq⟩ := hy + obtain ⟨E, hE⟩ : ∃ E, y.regs e = E := ⟨_, rfl⟩ + have hEpos : 0 < E := by simpa only [Cond.holds_lt, hyz, hE] using hc + rw [hE] at hye hq + have h2 : Fits w 2 := H.hfn.of_le (by have := H.hn0; nlinarith) + have hfE : Fits w E := H.hfe.of_le hye + -- The parity. + let y₁ := y.set par (E % 2) + have hrun₁ : ExecWithin (B := .standard) w (ram[ par ← e % two ]) y y₁ 1 := by + ram_eval [y₁, hE, hy2, h2.of_le (Nat.mod_lt E (by norm_num)).le] + have hy₁ (j : Fin 21) (hj : j ≠ par) : y₁.regs j = y.regs j := by simp [y₁, hj] + have hz₁ : y₁.regs z = 0 := by rw [hy₁ z (by decide), hyz] + have hp₁ : y₁.regs par = E % 2 := by simp [y₁] + -- Multiply `R` by `B` when the exponent is odd. + obtain ⟨y₂, hrun₂, hy₂r, hy₂R, hy₂Rred, hy₂B⟩ : ∃ y₂, ExecWithin w + (.branch (.lt z par) (mulInto R B R) (ram[ par ← par ])) y₁ y₂ (mulTime r + 2) ∧ + (∀ j : Fin 21, j.val < 9 → y₂.regs j = y₁.regs j) ∧ + toQuot n r (fun i => y₂.mem (a + i)) = + toQuot n r (fun i => y.mem (a + i)) * toQuot n r (fun i => y.mem (b + i)) ^ (E % 2) ∧ + (∀ i < r, y₂.mem (a + i) < n) ∧ (∀ i < r, y₂.mem (b + i) = y.mem (b + i)) := by + rcases Nat.mod_two_eq_zero_or_one E with hpar | hpar + · refine ⟨y₁.set par (y₁.regs par), ?_, ?_, ?_, ?_, ?_⟩ + · refine (ExecWithin.branch_neg ?_ ExecWithin.copy).mono (by unfold mulTime; omega) + simp only [Cond.holds_lt, hz₁, hp₁, hpar]; omega + · intro j _; simp + · simp [y₁, hpar] + · simpa [y₁] using hRred + · intro i _; rfl + · have hc' : (Cond.lt z par).Holds y₁ := by + simp only [Cond.holds_lt, hz₁, hp₁, hpar]; omega + obtain ⟨Y, hrun, hprod, hframe, hregs⟩ := mulInto_exec (w := w) y₁ (x := R) (y := B) (d := R) + (by decide) (by decide) (by decide) (a := a) (b := b) (t := t) (dd := a) (r := r) (n := n) + (by rw [hy₁ R (by decide), hyR]) (by rw [hy₁ B (by decide), hyB]) + (by rw [hy₁ T (by decide), hyT]) (by rw [hy₁ R (by decide), hyR]) + (by rw [hy₁ rr (by decide), hyr]) (by rw [hy₁ nn (by decide), hyn]) H.hn0 H.hta H.htb H.hta + (by simpa [y₁] using hRred) (by simpa [y₁] using hBred) H.hfa H.hfb H.hft H.hfa H.hf2 H.hfn + refine ⟨Y, (ExecWithin.branch_pos hc' hrun).mono (by omega), hregs, ?_, ?_, ?_⟩ + · rw [toQuot_congr hprod, toQuot_mulCoeff, hpar, pow_one] + rfl + · intro i hi + rw [hprod i hi] + exact mulCoeff_lt H.hn0 _ _ _ + · intro i hi + refine hframe _ ?_ ?_ + · rcases H.htb with h | h <;> omega + · rcases H.hab with h | h <;> omega + -- Square `B`. + have hy₂' (j : Fin 21) (hj : j.val < 9) (hjp : j ≠ par) : y₂.regs j = y.regs j := by + rw [hy₂r j hj, hy₁ j hjp] + obtain ⟨y₃, hrun₃, hprod₃, hframe₃, hregs₃⟩ := mulInto_exec (w := w) y₂ (x := B) (y := B) + (d := B) (by decide) (by decide) (by decide) (a := b) (b := b) (t := t) (dd := b) (r := r) + (n := n) (by rw [hy₂' B (by decide) (by decide), hyB]) + (by rw [hy₂' B (by decide) (by decide), hyB]) (by rw [hy₂' T (by decide) (by decide), hyT]) + (by rw [hy₂' B (by decide) (by decide), hyB]) (by rw [hy₂' rr (by decide) (by decide), hyr]) + (by rw [hy₂' nn (by decide) (by decide), hyn]) H.hn0 H.htb H.htb H.htb + (fun i hi => by rw [hy₂B i hi]; exact hBred i hi) + (fun i hi => by rw [hy₂B i hi]; exact hBred i hi) + H.hfb H.hfb H.hft H.hfb H.hf2 H.hfn + have hy₃R (i : ℕ) (hi : i < r) : y₃.mem (a + i) = y₂.mem (a + i) := by + refine hframe₃ _ ?_ ?_ + · rcases H.hta with h | h <;> omega + · rcases H.hab with h | h <;> omega + -- Halve the exponent. + have hy₃e : y₃.regs e = E := by rw [hregs₃ e (by decide), hy₂' e (by decide) (by decide), hE] + have hy₃2 : y₃.regs two = 2 := by + rw [hregs₃ two (by decide), hy₂' two (by decide) (by decide), hy2] + let y₄ := y₃.set e (E / 2) + have hrun₄ : ExecWithin (B := .standard) w (ram[ e ← e / two ]) y₃ y₄ 1 := by + ram_eval [y₄, hy₃e, hy₃2, hfE.of_le (Nat.div_le_self E 2)] + have hy₄ (j : Fin 21) (hj : j.val < 9) (hjp : j ≠ par) (hje : j ≠ e) : y₄.regs j = y.regs j := by + simp only [y₄, Data.set_regs, Function.update_of_ne hje] + rw [hregs₃ j hj, hy₂' j hj hjp] + refine ⟨y₄, ?_, ?_, ?_⟩ + · exact (hrun₁.seq (hrun₂.seq (hrun₃.seq hrun₄))).mono (by omega) + · refine ⟨?_, ?_, ?_, ?_, ?_, ?_, ?_, ?_, ?_, ?_, ?_⟩ + · rw [hy₄ R (by decide) (by decide) (by decide), hyR] + · rw [hy₄ B (by decide) (by decide) (by decide), hyB] + · rw [hy₄ T (by decide) (by decide) (by decide), hyT] + · rw [hy₄ rr (by decide) (by decide) (by decide), hyr] + · rw [hy₄ nn (by decide) (by decide) (by decide), hyn] + · rw [hy₄ two (by decide) (by decide) (by decide), hy2] + · rw [hy₄ z (by decide) (by decide) (by decide), hyz] + · simp only [y₄, Data.set_regs, Function.update_self] + exact (Nat.div_le_self E 2).trans hye + · intro i hi + change y₃.mem (a + i) < n + rw [hy₃R i hi] + exact hy₂Rred i hi + · intro i hi + change y₃.mem (b + i) < n + rw [hprod₃ i hi] + exact mulCoeff_lt H.hn0 _ _ _ + · simp only [y₄, Data.set_regs, Function.update_self] + change toQuot n r (fun i => y₃.mem (a + i)) * toQuot n r (fun i => y₃.mem (b + i)) ^ (E / 2) = + target + rw [toQuot_congr hy₃R, toQuot_congr hprod₃, toQuot_mulCoeff, toQuot_congr hy₂B, hy₂R, ← hq, + ← sq, ← pow_mul, mul_assoc, ← pow_add, Nat.mod_add_div] + · simp only [y₄, Data.set_regs, Function.update_self, hE] + exact size_div_two_lt hEpos + +/-- **Powers modulo `X^r - 1` and `n`.** If the result, base and scratch blocks `[a, a + r)`, +`[b, b + r)` and `[t, t + r)` are pairwise disjoint, the result and base coefficients are below +`n ≥ 1`, and the word size holds the block ends, `2 r`, `n² + n` and the exponent `e`, then within +`time r e` instructions the result block receives reduced coefficients whose class is that of +`R · B^e`. -/ +theorem exec {w : ℕ} (x : Data 21) {a b t r n : ℕ} (ha : x.regs R = a) (hb : x.regs B = b) + (ht : x.regs T = t) (hr : x.regs rr = r) (hn : x.regs nn = n) (hn0 : 0 < n) + (hab : a + r ≤ b ∨ b + r ≤ a) (hta : t + r ≤ a ∨ a + r ≤ t) (htb : t + r ≤ b ∨ b + r ≤ t) + (hpa : ∀ i < r, x.mem (a + i) < n) (hpb : ∀ i < r, x.mem (b + i) < n) + (hfa : Fits w (a + r)) (hfb : Fits w (b + r)) (hft : Fits w (t + r)) (hf2 : Fits w (2 * r)) + (hfn : Fits w (n * n + n)) (hfe : Fits w (x.regs e)) : + ∃ y, ExecWithin w program x y (time r (x.regs e)) ∧ + toQuot n r (fun i => y.mem (a + i)) = + toQuot n r (fun i => x.mem (a + i)) * toQuot n r (fun i => x.mem (b + i)) ^ x.regs e ∧ + (∀ i < r, y.mem (a + i) < n) ∧ y.regs R = a ∧ y.regs B = b ∧ y.regs T = t ∧ + y.regs rr = r ∧ y.regs nn = n := by + have H : Hyp w a b t r n (x.regs e) := ⟨hn0, hab, hta, htb, hfa, hfb, hft, hf2, hfn, hfe⟩ + have h2 : Fits w 2 := hfn.of_le (by nlinarith) + let x₀ := (x.set two 2).set z 0 + have hrun₀ : ExecWithin (B := .standard) w (ram[ two ← #2; z ← #0 ]) x x₀ 2 := by + ram_eval [x₀, h2] + have hx₀ (j : Fin 21) (h2 : j ≠ two) (hz : j ≠ z) : x₀.regs j = x.regs j := by + simp [x₀, h2, hz] + have hI : Inv a b t r n (x.regs e) (toQuot n r (fun i => x.mem (a + i)) * + toQuot n r (fun i => x.mem (b + i)) ^ x.regs e) x₀ := by + refine ⟨?_, ?_, ?_, ?_, ?_, by simp +decide [x₀], by simp [x₀], ?_, hpa, hpb, ?_⟩ + · rw [hx₀ R (by decide) (by decide), ha] + · rw [hx₀ B (by decide) (by decide), hb] + · rw [hx₀ T (by decide) (by decide), ht] + · rw [hx₀ rr (by decide) (by decide), hr] + · rw [hx₀ nn (by decide) (by decide), hn] + · rw [hx₀ e (by decide) (by decide)] + · rw [hx₀ e (by decide) (by decide)]; rfl + obtain ⟨y, hrun, hy, hc⟩ := ExecWithin.loop_invariant _ (fun y => Nat.size (y.regs e)) + (body_step H) hI + obtain ⟨hyR, hyB, hyT, hyr, hyn, -, hyz, -, hRred, -, hq⟩ := hy + have hy0 : y.regs e = 0 := by simp only [Cond.holds_lt, hyz] at hc; omega + refine ⟨y, ?_, ?_, hRred, hyR, hyB, hyT, hyr, hyn⟩ + · refine (hrun₀.seq hrun).mono ?_ + rw [hx₀ e (by decide) (by decide), show 2 * mulTime r + 4 + 2 = 2 * mulTime r + 6 by ring] + unfold time + omega + · rw [← hq, hy0, pow_zero, mul_one] + +end Cslib.Algorithms.WordRAM.AKS.PolyPow diff --git a/Complexity/Foundations/NumberTheory/AKS/Coefficients.lean b/Complexity/Foundations/NumberTheory/AKS/Coefficients.lean index 0d4565e02..daed87fd9 100644 --- a/Complexity/Foundations/NumberTheory/AKS/Coefficients.lean +++ b/Complexity/Foundations/NumberTheory/AKS/Coefficients.lean @@ -92,6 +92,11 @@ theorem toQuot_eq_sum (n r : ℕ) (c : ℕ → ℕ) : simp only [toQuot, coeffPoly, map_sum, map_mul, map_pow, AdjoinRoot.mk_X, AdjoinRoot.mk_C] rfl +/-- The class of `∑_{i < r} c i X^i` depends only on the coefficients below `r`. -/ +theorem toQuot_congr {c d : ℕ → ℕ} (h : ∀ i < r, c i = d i) : toQuot n r c = toQuot n r d := by + rw [toQuot_eq_sum, toQuot_eq_sum] + exact sum_congr rfl fun i hi => by rw [h i (mem_range.mp hi)] + /-- **Coefficients of products.** -/ theorem toQuot_mulCoeff (n r : ℕ) (p q : ℕ → ℕ) : toQuot n r (mulCoeff n r p q) = toQuot n r p * toQuot n r q := by diff --git a/ComplexityTests/WordRAMAKSCongruence.lean b/ComplexityTests/WordRAMAKSCongruence.lean new file mode 100644 index 000000000..6d7e61b34 --- /dev/null +++ b/ComplexityTests/WordRAMAKSCongruence.lean @@ -0,0 +1,71 @@ +/- +Copyright (c) 2026 Samuel Schlesinger. All rights reserved. +Released under Apache 2.0 license as described in the file LICENSE. +Authors: Samuel Schlesinger +-/ + +import Complexity.Algorithms.WordRAM.AKS.Congruence + +/-! +# Powers and the AKS congruence on the word RAM + +Concrete runs of repeated squaring modulo `X^r - 1` and `n` and of the congruence test, and +instances of their correctness theorems, including a congruence refuted by running the program. +-/ + +namespace ComplexityTests.WordRAMAKSCongruence + +open Cslib.WordRAM Cslib.Algorithms.WordRAM + +/-! ## Repeated squaring -/ + +/-- `R = 1` at `0`, `B = X + 1` at `2`, scratch at `4`, with `r = 2`, `n = 3` and exponent `e`. -/ +private def powInput (e : ℕ) : Data 21 := + ⟨fun j => if j = AKS.PolyPow.B then 2 else if j = AKS.PolyPow.T then 4 else + if j = AKS.PolyPow.rr then 2 else if j = AKS.PolyPow.nn then 3 else + if j = AKS.PolyPow.e then e else 0, + fun a => if a = 0 ∨ a = 2 ∨ a = 3 then 1 else if a < 4 then 0 else 5⟩ + +/-- `(X + 1)^2 = X² + 2X + 1 ≡ 2X + 2` modulo `X² - 1` and `3`. -/ +example : (Code.eval 16 4 AKS.PolyPow.program (powInput 2)).map + (fun yn => [yn.1.mem 0, yn.1.mem 1]) = some [2, 2] := by + decide +kernel + +/-- The correctness theorem at the same input: the result block is reduced modulo `7`. -/ +example : ∃ y, ExecWithin 16 AKS.PolyPow.program (powInput 5) y (AKS.PolyPow.time 2 5) ∧ + ∀ i < 2, y.mem i < 3 := by + obtain ⟨y, hy, -, hred, -⟩ := AKS.PolyPow.exec (w := 16) (powInput 5) (a := 0) (b := 2) (t := 4) + (r := 2) (n := 3) rfl rfl rfl rfl rfl (by decide) (by decide) (by decide) (by decide) + (by decide) (by decide) (by decide) (by decide) (by decide) (by decide) (by decide) + (by decide) + exact ⟨y, hy, fun i hi => by simpa using hred i hi⟩ + +/-! ## The congruence test -/ + +/-- Blocks at `0`, `r` and `2r`, with `r`, `n` and `a` in their registers, over dirty memory. -/ +private def congInput (r n a : ℕ) : Data 25 := + ⟨fun j => if j = AKS.Congruence.B then r else if j = AKS.Congruence.T then 2 * r else + if j = AKS.Congruence.rr then r else if j = AKS.Congruence.nn then n else + if j = AKS.Congruence.av then a else 0, + fun _ => 3⟩ + +/-- `(X + 1)^3 ≡ X^3 + 1` modulo `X - 1` and `3`: both sides evaluate to `2` at `X = 1`. -/ +example : (Code.eval 16 4 AKS.Congruence.program (congInput 1 3 1)).map + (·.1.regs AKS.Congruence.flag) = some 1 := by + decide +kernel + +/-- The correctness theorem and a run refute `(X + 1)^4 ≡ X^4 + 1 (mod X - 1, 4)`: at `X = 1` +the sides are `16 ≡ 0` and `2`. -/ +example : ¬ Cslib.AKS.Congruence 4 1 1 := by + obtain ⟨y, ⟨m, -, hy⟩, -, hiff, -⟩ := AKS.Congruence.exec (w := 16) (congInput 1 4 1) + (a := 0) (b := 1) (t := 2) (r := 1) (n := 4) (A := 1) rfl rfl rfl rfl rfl rfl (by decide) + (by decide) (by decide) (by decide) (by decide) (by decide) (by decide) (by decide) + (by decide) (by decide) (by decide) + have hrun : (Code.eval 16 4 AKS.Congruence.program (congInput 1 4 1)).map + (·.1.regs AKS.Congruence.flag) = some 0 := by decide +kernel + obtain ⟨⟨y', k⟩, heval, hflag⟩ := Option.map_eq_some_iff.mp hrun + obtain ⟨rfl, -⟩ := Exec.unique hy (Exec.of_eval heval) + rw [← hiff, hflag] + decide + +end ComplexityTests.WordRAMAKSCongruence