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/8] 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/8] 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/8] 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/8] 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 From 0f592f24122ab239e824bd5c0c28a8206726b130 Mon Sep 17 00:00:00 2001 From: Samuel Giuseppe Schlesinger Date: Fri, 30 Oct 2026 22:23:56 -0400 Subject: [PATCH 5/8] feat(AKS): integer thresholds, the order search and step 3 (#893) IntegerVariant states the AKS tests with c = ceil(log2 n): PassesUpTo n r A checks the congruence for 1 <= a <= A, and prime_iff_passesUpTo proves n.Prime <-> PassesUpTo n r A when ord_r(n) > c^2 and A >= numTests n r, from Theorem 4.1 (prime_iff_passes). lt_orderOf_iff decides the order condition by small powers. On the word RAM, OrderSearch finds r with that condition and GcdTotient runs step 3 and computes phi(r), each with a word-size-independent instruction bound. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_015DLg8kwPs8yi7ePPdtoNvG --- Complexity.lean | 3 + .../Algorithms/WordRAM/AKS/GcdTotient.lean | 255 +++++++++++++++ .../Algorithms/WordRAM/AKS/OrderSearch.lean | 294 ++++++++++++++++++ .../NumberTheory/AKS/IntegerVariant.lean | 135 ++++++++ ComplexityTests/WordRAMAKSSearch.lean | 78 +++++ 5 files changed, 765 insertions(+) create mode 100644 Complexity/Algorithms/WordRAM/AKS/GcdTotient.lean create mode 100644 Complexity/Algorithms/WordRAM/AKS/OrderSearch.lean create mode 100644 Complexity/Foundations/NumberTheory/AKS/IntegerVariant.lean create mode 100644 ComplexityTests/WordRAMAKSSearch.lean diff --git a/Complexity.lean b/Complexity.lean index 2af464aa6..ed2aead30 100644 --- a/Complexity.lean +++ b/Complexity.lean @@ -81,6 +81,8 @@ 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.GcdTotient +public import Complexity.Algorithms.WordRAM.AKS.OrderSearch public import Complexity.Algorithms.WordRAM.AKS.PerfectPower public import Complexity.Algorithms.WordRAM.AKS.PolyArray public import Complexity.Algorithms.WordRAM.AKS.PolyMul @@ -2768,6 +2770,7 @@ public import Complexity.Foundations.MeasureTheory.Constructions.BorelSpace.Orde 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.IntegerVariant public import Complexity.Foundations.NumberTheory.AKS.Introspective public import Complexity.Foundations.NumberTheory.AKS.PerfectPower public import Complexity.Foundations.NumberTheory.AKS.SmallR diff --git a/Complexity/Algorithms/WordRAM/AKS/GcdTotient.lean b/Complexity/Algorithms/WordRAM/AKS/GcdTotient.lean new file mode 100644 index 000000000..8e1749171 --- /dev/null +++ b/Complexity/Algorithms/WordRAM/AKS/GcdTotient.lean @@ -0,0 +1,255 @@ +/- +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.NatGcd +public import Complexity.Computability.WordRAM.RegisterFrame +public import Mathlib.Data.Nat.Totient +import Mathlib.Tactic.FinCases + +/-! +# Step 3 of the AKS algorithm and Euler's totient on the word RAM + +Step 3 of the AKS algorithm [Agrawal, Kayal, Saxena, Section 4][AKS2004] outputs COMPOSITE if +`1 < gcd(a, n) < n` for some `a ≤ r`, and step 5 needs `φ(r)`. One loop over `0 ≤ a ≤ r` does +both, with two runs of Euclid's algorithm (`Cslib.Algorithms.WordRAM.NatGcd.program`) per `a`: it +clears a flag when `gcd(n, a)` is neither `1` nor `n`, and counts the `a` with `gcd(r, a) = 1`. +The index `a = 0` is harmless (`gcd(n, 0) = n`), and for `r ≥ 2` neither `0` nor `r` is coprime +to `r`, so the count is `φ(r)`. + +## Main definitions + +* `Cslib.Algorithms.WordRAM.AKS.GcdTotient.program`: the loop. +* `Cslib.Algorithms.WordRAM.AKS.GcdTotient.time`: its instruction bound. + +## Main results + +* `Cslib.Algorithms.WordRAM.AKS.GcdTotient.exec`: the flag records step 3 and the count is `φ(r)`, + within `time r` instructions. + +## References + +* [M. Agrawal, N. Kayal, N. Saxena, *PRIMES is in P*][AKS2004] +-/ + +@[expose] public section + +namespace Cslib.Algorithms.WordRAM.AKS.GcdTotient + +open Cslib.WordRAM + +/-- The number `n`, preserved. -/ +abbrev n : Fin 11 := 0 +/-- The modulus `r`, preserved. -/ +abbrev r : Fin 11 := 1 +/-- The loop index `a`. -/ +abbrev a : Fin 11 := 2 +/-- The loop end `r + 1`. -/ +abbrev lim : Fin 11 := 3 +/-- `1` while every `gcd(n, a)` seen is `1` or `n`. -/ +abbrev ok : Fin 11 := 4 +/-- The number of `a` seen with `gcd(r, a) = 1`. -/ +abbrev phi : Fin 11 := 5 +/-- Constant one. -/ +abbrev one : Fin 11 := 6 +/-- First argument of the gcd, and its result. -/ +abbrev ga : Fin 11 := 7 +/-- Second argument of the gcd. -/ +abbrev gb : Fin 11 := 8 + +/-- Euclid's algorithm runs on `ga`, `gb` and two scratch registers. -/ +def gcdMap : Fin 4 → Fin 11 := ![ga, gb, 9, 10] + +theorem gcdMap_injective : gcdMap.Injective := by decide + +/-- The gcd on `ga` and `gb`. -/ +def gcd : Code .standard 11 := NatGcd.program.rename gcdMap + +/-- Test `gcd(n, a)`, count `gcd(r, a) = 1`, and advance `a`. -/ +def body : Code .standard 11 := + .seq (ram[ ga ← n; gb ← a ]) (.seq gcd + (.seq (ram[ if ga = one then [ ok ← ok ] else [ if ga = n then [ ok ← ok ] else [ ok ← #0 ] ]; + ga ← r; gb ← a ]) (.seq gcd + (ram[ if ga = one then [ phi ← phi + one ] else [ phi ← phi ]; a ← a + one ])))) + +/-- Run the body for `a = 0, 1, …, r`. -/ +def program : Code .standard 11 := + .seq (ram[ one ← #1; lim ← r + one; a ← #0; ok ← #1; phi ← #0 ]) (.loop (.lt a lim) body) + +/-- The instruction bound of `program`. -/ +def time (r : ℕ) : ℕ := (r + 1) * (10 * r + 21) + 7 + +/-- The gcd changes only `ga`, `gb` and its scratch registers. -/ +private theorem gcd_regs {w T : ℕ} {x y : Data 11} (h : ExecWithin w gcd x y T) (j : Fin 11) + (hj : j.val < 7) : y.regs j = x.regs j := by + refine h.regs_eq (Code.preserves_rename_of_forall_ne _ _ _ ?_) + intro b hb + subst hb + fin_cases b <;> simp [gcdMap] at hj + +private theorem gcd_exec {w : ℕ} (x : Data 11) {u v : ℕ} (hu : x.regs ga = u) (hv : x.regs gb = v) + (hfu : Fits w u) (hfv : Fits w v) : + ∃ y, ExecWithin w gcd x y (5 * v + 3) ∧ y.regs ga = Nat.gcd u v ∧ + ∀ j : Fin 11, j.val < 7 → y.regs j = x.regs j := by + obtain ⟨y, hy, hg, -⟩ := NatGcd.exec_rename gcdMap gcdMap_injective x + (by rw [show gcdMap NatGcd.a = ga from rfl, hu]; exact hfu) + (by rw [show gcdMap NatGcd.b = gb from rfl, hv]; exact hfv) + refine ⟨y, ?_, ?_, fun j hj => gcd_regs hy j hj⟩ + · rw [show gcdMap NatGcd.b = gb from rfl, hv] at hy; exact hy + · rw [show ga = gcdMap NatGcd.a from rfl, hg, show gcdMap NatGcd.a = ga from rfl, + show gcdMap NatGcd.b = gb from rfl, hu, hv] + +/-- Counting the residues coprime to `r ≥ 2` up to `r` inclusive gives `φ(r)`. -/ +private theorem card_filter_coprime_range_succ {R : ℕ} (hR : 2 ≤ R) : + ((Finset.range (R + 1)).filter R.Coprime).card = Nat.totient R := by + rw [Finset.range_add_one, Finset.filter_insert] + simp [Nat.coprime_self, show R ≠ 1 by omega, Nat.totient_eq_card_coprime] + +/-- The loop invariant. -/ +private def Inv (N R : ℕ) (y : Data 11) : Prop := + y.regs n = N ∧ y.regs r = R ∧ y.regs one = 1 ∧ y.regs lim = R + 1 ∧ y.regs a ≤ R + 1 ∧ + y.regs ok ≤ 1 ∧ (y.regs ok = 1 ↔ ∀ b < y.regs a, Nat.gcd N b = 1 ∨ Nat.gcd N b = N) ∧ + y.regs phi = ((Finset.range (y.regs a)).filter R.Coprime).card + +private theorem body_step {w N R : ℕ} (hfN : Fits w N) (hfR : Fits w (R + 1)) (y : Data 11) + (hy : Inv N R y) (hc : (Cond.lt a lim).Holds y) : + ∃ z, ExecWithin w body y z (10 * R + 19) ∧ Inv N R z ∧ + R + 1 - z.regs a < R + 1 - y.regs a := by + classical + obtain ⟨hyn, hyr, hy1, hyl, hya, hyok, hyiff, hyphi⟩ := hy + obtain ⟨A, hA⟩ : ∃ A, y.regs a = A := ⟨_, rfl⟩ + have hAR : A ≤ R := by simp only [Cond.holds_lt, hA, hyl] at hc; omega + rw [hA] at hyiff hyphi + have hfA : Fits w A := hfR.of_le (by omega) + have h1f : Fits w 1 := hfR.of_le (by omega) + -- `gcd(n, a)`. + let y₁ := (y.set ga N).set gb A + have hrun₁ : ExecWithin (B := .standard) w (ram[ ga ← n; gb ← a ]) y y₁ 2 := by + ram_eval [y₁, hyn, hA] + have hy₁ (j : Fin 11) (hj : j.val < 7) : y₁.regs j = y.regs j := by + have h9 : j ≠ ga ∧ j ≠ gb := by constructor <;> (intro h; subst h; simp at hj) + simp [y₁, h9.1, h9.2] + obtain ⟨y₂, hrun₂, hg₂, hy₂⟩ := gcd_exec (w := w) y₁ (u := N) (v := A) (by simp +decide [y₁]) + (by simp +decide [y₁]) hfN hfA + have hy₂' (j : Fin 11) (hj : j.val < 7) : y₂.regs j = y.regs j := by rw [hy₂ j hj, hy₁ j hj] + -- Record step 3. + let o := if Nat.gcd N A = 1 ∨ Nat.gcd N A = N then y.regs ok else 0 + let y₃ := ((y₂.set ok o).set ga R).set gb A + have hy₂ok : y₂.regs ok = y.regs ok := hy₂' ok (by simp) + have hy₂n : y₂.regs n = N := by rw [hy₂' n (by simp), hyn] + have hy₂1 : y₂.regs one = 1 := by rw [hy₂' one (by simp), hy1] + have hy₂r : y₂.regs r = R := by rw [hy₂' r (by simp), hyr] + have hy₂a : y₂.regs a = A := by rw [hy₂' a (by simp), hA] + have hrun₃ : ExecWithin (B := .standard) w + (ram[ if ga = one then [ ok ← ok ] else [ if ga = n then [ ok ← ok ] else [ ok ← #0 ] ]; + ga ← r; gb ← a ]) y₂ y₃ 7 := by + by_cases hg1 : Nat.gcd N A = 1 + · refine ExecWithin.mono (T := 4) ?_ (by omega) + ram_eval [y₃, o, hg₂, hg1, hy₂1, hy₂ok, hy₂r, hy₂a] + · by_cases hgN : Nat.gcd N A = N + · have hN1 : ¬ N = 1 := fun h => hg1 (hgN.trans h) + refine ExecWithin.mono (T := 6) ?_ (by omega) + ram_eval [y₃, o, hg₂, hg1, hgN, hN1, hy₂1, hy₂n, hy₂ok, hy₂r, hy₂a] + · ram_eval [y₃, o, hg₂, hg1, hgN, hy₂1, hy₂n, hy₂r, hy₂a] + -- `gcd(r, a)`. + obtain ⟨y₄, hrun₄, hg₄, hy₄⟩ := gcd_exec (w := w) y₃ (u := R) (v := A) (by simp +decide [y₃]) + (by simp +decide [y₃]) (hfR.of_le (by omega)) hfA + have hy₄' (j : Fin 11) (hj : j.val < 7) (hjo : j ≠ ok) : y₄.regs j = y.regs j := by + rw [hy₄ j hj] + have h9 : j ≠ ga ∧ j ≠ gb := by constructor <;> (intro h; subst h; simp at hj) + simp only [y₃, Data.set_regs, Function.update_of_ne h9.1, Function.update_of_ne h9.2, + Function.update_of_ne hjo] + exact hy₂' j hj + have hy₄ok : y₄.regs ok = o := by rw [hy₄ ok (by simp)]; simp +decide [y₃] + have hy₄1 : y₄.regs one = 1 := by rw [hy₄' one (by simp) (by decide), hy1] + have hy₄phi : y₄.regs phi = y.regs phi := hy₄' phi (by simp) (by decide) + have hy₄a : y₄.regs a = A := by rw [hy₄' a (by simp) (by decide), hA] + -- Count and advance. + have hphiA : ((Finset.range A).filter R.Coprime).card ≤ A := + (Finset.card_filter_le _ _).trans (Finset.card_range A).le + let ph := if Nat.gcd R A = 1 then y.regs phi + 1 else y.regs phi + let z := (y₄.set phi ph).set a (A + 1) + have hfph : Fits w (y.regs phi + 1) := hfR.of_le (by omega) + have hfA1 : Fits w (A + 1) := hfR.of_le (by omega) + have hrun₅ : ExecWithin (B := .standard) w + (ram[ if ga = one then [ phi ← phi + one ] else [ phi ← phi ]; a ← a + one ]) y₄ z 4 := by + by_cases hg1 : Nat.gcd R A = 1 + · refine ExecWithin.mono (T := 3) ?_ (by omega) + ram_eval [z, ph, hg₄, hg1, hy₄1, hy₄phi, hy₄a, hfph, hfA1] + · ram_eval [z, ph, hg₄, hg1, hy₄1, hy₄phi, hy₄a, hfA1] + have hza : z.regs a = A + 1 := by simp [z] + refine ⟨z, (hrun₁.seq (hrun₂.seq (hrun₃.seq (hrun₄.seq hrun₅)))).mono (by omega), ?_, ?_⟩ + · refine ⟨?_, ?_, ?_, ?_, ?_, ?_, ?_, ?_⟩ + · simp +decide only [z, Data.set_regs, Function.update_apply, ↓reduceIte] + rw [hy₄' n (by simp) (by decide), hyn] + · simp +decide only [z, Data.set_regs, Function.update_apply, ↓reduceIte] + rw [hy₄' r (by simp) (by decide), hyr] + · simp +decide only [z, Data.set_regs, Function.update_apply, ↓reduceIte]; exact hy₄1 + · simp +decide only [z, Data.set_regs, Function.update_apply, ↓reduceIte] + rw [hy₄' lim (by simp) (by decide), hyl] + · rw [hza]; omega + · simp +decide only [z, Data.set_regs, Function.update_apply, ↓reduceIte] + rw [hy₄ok]; simp only [o]; split_ifs <;> omega + · have hzo : z.regs ok = o := by simp +decide [z, hy₄ok] + rw [hzo, hza] + simp only [o] + constructor + · intro h b hb + split_ifs at h with hc + · rcases Nat.lt_succ_iff_lt_or_eq.mp hb with hb | rfl + · exact hyiff.mp h b hb + · exact hc + · intro h + have hc : Nat.gcd N A = 1 ∨ Nat.gcd N A = N := h A (by omega) + simp only [hc, ↓reduceIte] + exact hyiff.mpr fun b hb => h b (by omega) + · have hzp : z.regs phi = ph := by simp +decide [z] + rw [hzp, hza, Finset.range_add_one, Finset.filter_insert] + simp only [ph, hyphi, Nat.Coprime] + split_ifs with hg + · rw [Finset.card_insert_of_notMem (by simp)] + · rfl + · rw [hza, hA]; omega + +/-- **Step 3 and the totient.** For `r ≥ 2`, with the word size holding `n` and `r + 1`, within +`time r` instructions the flag `ok` is `1` exactly when no `1 ≤ a ≤ r` has `1 < gcd(a, n) < n`, +and `phi` is `φ(r)`. -/ +theorem exec {w : ℕ} (x : Data 11) {N R : ℕ} (hn : x.regs n = N) (hr : x.regs r = R) + (hR : 2 ≤ R) (hfN : Fits w N) (hfR : Fits w (R + 1)) : + ∃ y, ExecWithin w program x y (time R) ∧ y.regs ok ≤ 1 ∧ + (y.regs ok = 1 ↔ ∀ b, 1 ≤ b → b ≤ R → b.gcd N = 1 ∨ b.gcd N = N) ∧ + y.regs phi = Nat.totient R ∧ y.regs n = N ∧ y.regs r = R := by + have h1f : Fits w 1 := hfR.of_le (by omega) + let x₀ := ((((x.set one 1).set lim (R + 1)).set a 0).set ok 1).set phi 0 + have hrun₀ : ExecWithin (B := .standard) w + (ram[ one ← #1; lim ← r + one; a ← #0; ok ← #1; phi ← #0 ]) x x₀ 5 := by + ram_eval [x₀, hr, h1f, hfR] + have hx₀ : Inv N R x₀ := by + refine ⟨?_, ?_, ?_, ?_, ?_, ?_, ?_, ?_⟩ <;> simp +decide [x₀, hn, hr] + obtain ⟨y, hrun, hy, hc⟩ := ExecWithin.loop_invariant (Inv N R) (fun y => R + 1 - y.regs a) + (body_step hfN hfR) hx₀ + obtain ⟨hyn, hyr, -, hyl, hya, hyok, hyiff, hyphi⟩ := hy + have hend : y.regs a = R + 1 := by simp only [Cond.holds_lt, hyl] at hc; omega + rw [hend] at hyiff hyphi + refine ⟨y, (hrun₀.seq hrun).mono ?_, hyok, ?_, ?_, hyn, hyr⟩ + · have : x₀.regs a = 0 := by simp +decide [x₀] + rw [this, Nat.sub_zero, show 10 * R + 19 + 2 = 10 * R + 21 by ring] + unfold time + omega + · rw [hyiff] + constructor + · intro h b hb1 hbR + rw [Nat.gcd_comm] + exact h b (by omega) + · intro h b hb + rcases Nat.eq_zero_or_pos b with rfl | hb0 + · right; exact Nat.gcd_zero_right N + · rw [Nat.gcd_comm]; exact h b hb0 (by omega) + · rw [hyphi, card_filter_coprime_range_succ hR] + +end Cslib.Algorithms.WordRAM.AKS.GcdTotient diff --git a/Complexity/Algorithms/WordRAM/AKS/OrderSearch.lean b/Complexity/Algorithms/WordRAM/AKS/OrderSearch.lean new file mode 100644 index 000000000..df6b205c3 --- /dev/null +++ b/Complexity/Algorithms/WordRAM/AKS/OrderSearch.lean @@ -0,0 +1,294 @@ +/- +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.NatGcd +public import Complexity.Computability.WordRAM.RegisterFrame +import Mathlib.Tactic.Linarith +import Mathlib.Tactic.FinCases + +/-! +# The search for `r` in the AKS algorithm on the word RAM + +Step 2 of the AKS algorithm [Agrawal, Kayal, Saxena, Section 4][AKS2004] finds the least `r` such +that the multiplicative order of `n` modulo `r` exceeds a bound `L`. As in the proof of their +Theorem 5.1, the program tries `r = 2, 3, …` in turn. It accepts `r` when `gcd(n, r) = 1` and +`n^k mod r ≠ 1` for every `1 ≤ k ≤ L`; for `r ≥ 2` this is the order condition +(`Cslib.AKS.lt_orderOf_iff`). The powers are computed by one multiplication modulo `r` each, and +the gcd by Euclid's algorithm (`Cslib.Algorithms.WordRAM.NatGcd.program`). Given any `r₀ ≥ 2` +satisfying the condition, the search stops at some `2 ≤ r ≤ r₀` satisfying it, within a number of +instructions polynomial in `r₀` and `L`. + +## Main definitions + +* `Cslib.Algorithms.WordRAM.AKS.OrderSearch.Good`: the condition tested for one `r`. +* `Cslib.Algorithms.WordRAM.AKS.OrderSearch.program`: the search. +* `Cslib.Algorithms.WordRAM.AKS.OrderSearch.time`: its instruction bound. + +## Main results + +* `Cslib.Algorithms.WordRAM.AKS.OrderSearch.exec`: the search stops at an `r ≤ r₀` satisfying the + condition, within `time L r₀` instructions. + +## References + +* [M. Agrawal, N. Kayal, N. Saxena, *PRIMES is in P*][AKS2004] +-/ + +@[expose] public section + +namespace Cslib.Algorithms.WordRAM.AKS.OrderSearch + +open Cslib.WordRAM + +/-- The condition tested for `r`: `n` is coprime to `r` and no `n^k` with `1 ≤ k ≤ L` is `1` +modulo `r`. -/ +def Good (n L r : ℕ) : Prop := n.Coprime r ∧ ∀ k, 1 ≤ k → k ≤ L → n ^ k % r ≠ 1 + +/-- The number `n`, preserved. -/ +abbrev n : Fin 13 := 0 +/-- The bound `L` on the order, preserved. -/ +abbrev L : Fin 13 := 1 +/-- The candidate `r`, and the output. -/ +abbrev r : Fin 13 := 2 +/-- `1` once a candidate is accepted. -/ +abbrev found : Fin 13 := 3 +/-- Constant one. -/ +abbrev one : Fin 13 := 4 +/-- `n mod r`. -/ +abbrev nr : Fin 13 := 5 +/-- The power `n^k mod r`. -/ +abbrev p : Fin 13 := 6 +/-- The exponent `k`. -/ +abbrev k : Fin 13 := 7 +/-- `1` while no power `n^j` with `1 ≤ j ≤ k` is `1` modulo `r`. -/ +abbrev ok : Fin 13 := 8 +/-- First argument of the gcd, and its result. -/ +abbrev ga : Fin 13 := 9 +/-- Second argument of the gcd. -/ +abbrev gb : Fin 13 := 10 + +/-- Euclid's algorithm runs on `ga`, `gb` and two scratch registers. -/ +def gcdMap : Fin 4 → Fin 13 := ![ga, gb, 11, 12] + +theorem gcdMap_injective : gcdMap.Injective := by decide + +/-- One more power: `p := p · (n mod r) mod r`, and record whether it is `1`. -/ +def powBody : Code .standard 13 := ram[ + p ← p * nr; p ← p % r; k ← k + one; if p = one then [ ok ← #0 ] else [ ok ← ok ] ] + +/-- Test the powers `n^k mod r` for `1 ≤ k ≤ L`. -/ +def powers : Code .standard 13 := + .seq (ram[ nr ← n % r; p ← one; k ← #0; ok ← #1 ]) (.loop (.lt k L) powBody) + +/-- Accept `r` if `gcd(n, r) = 1` and no small power is `1`; otherwise try `r + 1`. -/ +def accept : Code .standard 13 := ram[ + if ga = one then [ if ok = one then [ found ← #1 ] else [ r ← r + one ] ] else [ r ← r + one ] ] + +/-- Test one candidate. -/ +def body : Code .standard 13 := + .seq powers (.seq (ram[ ga ← n; gb ← r ]) (.seq (NatGcd.program.rename gcdMap) accept)) + +/-- Try `r = 2, 3, …` until one is accepted. -/ +def program : Code .standard 13 := + .seq (ram[ one ← #1; r ← #2; found ← #0 ]) (.loop (.lt found one) body) + +/-- The instruction bound of `program`. -/ +def time (L r₀ : ℕ) : ℕ := (2 * r₀ + 1) * (8 * L + 5 * r₀ + 20) + 5 + +/-- The registers that the power test preserves. -/ +private theorem powers_preserves : + ∀ j ∈ [n, L, r, found, one], powers.Preserves j := by decide + +private theorem powers_exec {w : ℕ} (x : Data 13) {N L₀ R : ℕ} (hn : x.regs n = N) + (hL : x.regs L = L₀) (hr : x.regs r = R) (h1 : x.regs one = 1) (hR : 2 ≤ R) + (hfN : Fits w N) (hfR : Fits w (R * R)) (hfL : Fits w (L₀ + 1)) : + ∃ y, ExecWithin w powers x y (8 * L₀ + 6) ∧ y.regs ok ≤ 1 ∧ + (y.regs ok = 1 ↔ ∀ j, 1 ≤ j → j ≤ L₀ → N ^ j % R ≠ 1) ∧ + ∀ j ∈ [n, L, r, found, one], y.regs j = x.regs j := by + have hR0 : 0 < R := by omega + have h1f : Fits w 1 := hfL.of_le (by omega) + let J (y : Data 13) : Prop := y.regs L = L₀ ∧ y.regs r = R ∧ y.regs one = 1 ∧ + y.regs nr = N % R ∧ y.regs k ≤ L₀ ∧ y.regs p = N ^ y.regs k % R ∧ y.regs ok ≤ 1 ∧ + (y.regs ok = 1 ↔ ∀ j, 1 ≤ j → j ≤ y.regs k → N ^ j % R ≠ 1) + have hstep (y : Data 13) (hy : J y) (hc : (Cond.lt k L).Holds y) : + ∃ z, ExecWithin w powBody y z 6 ∧ J z ∧ L₀ - z.regs k < L₀ - y.regs k := by + obtain ⟨hyL, hyr, hy1, hynr, hyk, hyp, hyok, hyiff⟩ := hy + obtain ⟨K, hK⟩ : ∃ K, y.regs k = K := ⟨_, rfl⟩ + have hKL : K < L₀ := by simpa only [Cond.holds_lt, hK, hyL] using hc + rw [hK] at hyp hyiff + have hPR : N ^ K % R < R := Nat.mod_lt _ hR0 + have hNR : N % R < R := Nat.mod_lt _ hR0 + have hmul : N ^ K % R * (N % R) < R * R := Nat.mul_lt_mul'' hPR hNR + have hmod : N ^ K % R * (N % R) % R = N ^ (K + 1) % R := by + rw [← Nat.mul_mod, pow_succ] + have hf₁ : Fits w (N ^ K % R * (N % R)) := hfR.of_le hmul.le + have hf₂ : Fits w (N ^ (K + 1) % R) := hfR.of_le + ((Nat.mod_lt _ hR0).le.trans (Nat.le_mul_self R)) + have hf₃ : Fits w (K + 1) := hfL.of_le (by omega) + have hiff (o : ℕ) : ((∀ j, 1 ≤ j → j ≤ K + 1 → N ^ j % R ≠ 1) ↔ + (∀ j, 1 ≤ j → j ≤ K → N ^ j % R ≠ 1) ∧ N ^ (K + 1) % R ≠ 1) := by + constructor + · intro h; exact ⟨fun j hj hjK => h j hj (by omega), h (K + 1) (by omega) le_rfl⟩ + · rintro ⟨h, hK1⟩ j hj hjK + rcases Nat.lt_or_ge j (K + 1) with hlt | hge + · exact h j hj (by omega) + · rw [show j = K + 1 by omega]; exact hK1 + let o := if N ^ (K + 1) % R = 1 then 0 else y.regs ok + let z := (((y.set p (N ^ K % R * (N % R))).set p (N ^ (K + 1) % R)).set k (K + 1)).set ok o + have hrun : ExecWithin (B := .standard) w powBody y z 6 := by + by_cases hP : N ^ (K + 1) % R = 1 + · refine ExecWithin.mono (T := 5) ?_ (by omega) + ram_eval [powBody, z, o, hyp, hynr, hyr, hK, hy1, hmod, hP, hf₁, hf₂, hf₃, h1f] + · ram_eval [powBody, z, o, hyp, hynr, hyr, hK, hy1, hmod, hP, hf₁, hf₂, hf₃, h1f] + have hzk : z.regs k = K + 1 := by simp +decide [z] + refine ⟨z, hrun, ⟨?_, ?_, ?_, ?_, ?_, ?_, ?_, ?_⟩, ?_⟩ + · simp +decide [z, hyL] + · simp +decide [z, hyr] + · simp +decide [z, hy1] + · simp +decide [z, hynr] + · rw [hzk]; omega + · rw [hzk]; simp +decide [z] + · simp only [z, o, Data.set_regs, Function.update_self] + split_ifs <;> omega + · rw [hzk, hiff 0] + simp only [z, o, Data.set_regs, Function.update_self] + split_ifs with hP + · simp [hP] + · rw [hyiff]; simp [hP] + · rw [hzk, hK]; omega + let x₀ := (((x.set nr (N % R)).set p 1).set k 0).set ok 1 + have hrun₀ : ExecWithin (B := .standard) w (ram[ nr ← n % r; p ← one; k ← #0; ok ← #1 ]) + x x₀ 4 := by + ram_eval [x₀, hn, hr, h1, h1f, hfN.of_le (Nat.mod_le N R)] + have hx₀ : J x₀ := by + have hk : x₀.regs k = 0 := by simp +decide [x₀] + have hp : x₀.regs p = 1 := by simp +decide [x₀] + have ho : x₀.regs ok = 1 := by simp +decide [x₀] + refine ⟨by simp +decide [x₀, hL], by simp +decide [x₀, hr], by simp +decide [x₀, h1], + by simp +decide [x₀], by rw [hk]; omega, ?_, by rw [ho], ?_⟩ + · rw [hk, hp, pow_zero, Nat.mod_eq_of_lt (by omega : 1 < R)] + · rw [ho, hk] + exact ⟨fun _ j hj hj0 => by omega, fun _ => rfl⟩ + obtain ⟨y, hrun, hy, hc⟩ := ExecWithin.loop_invariant J (fun y => L₀ - y.regs k) hstep hx₀ + obtain ⟨hyL, -, -, -, hyk, -, hyok, hyiff⟩ := hy + have hkL : y.regs k = L₀ := by simp only [Cond.holds_lt, hyL] at hc; omega + have hrun' := hrun₀.seq hrun + refine ⟨y, hrun'.mono ?_, hyok, by rw [hyiff, hkL], + fun j hj => hrun'.regs_eq (powers_preserves j hj)⟩ + have : x₀.regs k = 0 := by simp +decide [x₀] + rw [this] + omega + +/-- **The search for `r`.** If some `r₀ ≥ 2` satisfies `Good n L r₀`, and the word size holds +`n`, `r₀²` and `L + 1`, the search stops within `time L r₀` instructions at some `2 ≤ r ≤ r₀` +satisfying `Good n L r`, preserving `n` and `L`. -/ +theorem exec {w : ℕ} (x : Data 13) {N L₀ r₀ : ℕ} (hn : x.regs n = N) (hL : x.regs L = L₀) + (hr₀ : 2 ≤ r₀) (hgood : Good N L₀ r₀) (hfN : Fits w N) (hfR : Fits w (r₀ * r₀)) + (hfL : Fits w (L₀ + 1)) : + ∃ y, ExecWithin w program x y (time L₀ r₀) ∧ 2 ≤ y.regs r ∧ y.regs r ≤ r₀ ∧ + Good N L₀ (y.regs r) ∧ y.regs n = N ∧ y.regs L = L₀ := by + have h1f : Fits w 1 := hfL.of_le (by omega) + have hfr₁ : Fits w (r₀ + 1) := hfR.of_le (by nlinarith) + let I (y : Data 13) : Prop := y.regs n = N ∧ y.regs L = L₀ ∧ y.regs one = 1 ∧ 2 ≤ y.regs r ∧ + y.regs r ≤ r₀ ∧ y.regs found ≤ 1 ∧ (y.regs found = 1 → Good N L₀ (y.regs r)) + have hstep (y : Data 13) (hy : I y) (hc : (Cond.lt found one).Holds y) : + ∃ z, ExecWithin w body y z (8 * L₀ + 5 * r₀ + 16) ∧ I z ∧ + 2 * (r₀ - z.regs r) + (1 - z.regs found) < 2 * (r₀ - y.regs r) + (1 - y.regs found) := by + obtain ⟨hyn, hyL, hy1, hr2, hrr₀, hyf, -⟩ := hy + obtain ⟨R, hR⟩ : ∃ R, y.regs r = R := ⟨_, rfl⟩ + rw [hR] at hr2 hrr₀ + have hf0 : y.regs found = 0 := by simp only [Cond.holds_lt, hy1] at hc; omega + -- The powers. + obtain ⟨y₁, hrun₁, hok1, hokiff, hy₁⟩ := powers_exec (w := w) y hyn hyL hR hy1 hr2 hfN + (hfR.of_le (Nat.mul_le_mul hrr₀ hrr₀)) hfL + have hy₁n : y₁.regs n = N := by rw [hy₁ n (by simp), hyn] + have hy₁r : y₁.regs r = R := by rw [hy₁ r (by simp), hR] + -- The gcd. + let y₂ := (y₁.set ga N).set gb R + have hrun₂ : ExecWithin (B := .standard) w (ram[ ga ← n; gb ← r ]) y₁ y₂ 2 := by + ram_eval [y₂, hy₁n, hy₁r] + obtain ⟨y₃, hrun₃, hg, -⟩ := NatGcd.exec_rename gcdMap gcdMap_injective y₂ + (by simpa +decide [y₂, gcdMap] using hfN) + (by simpa +decide [y₂, gcdMap] using hfR.of_le (hrr₀.trans (Nat.le_mul_self r₀))) + have hgcd : y₃.regs ga = Nat.gcd N R := by + rw [show ga = gcdMap NatGcd.a from rfl, hg]; simp +decide [y₂, gcdMap] + have hy₃ (j : Fin 13) (hj : j.val < 9) : y₃.regs j = y₁.regs j := by + rw [hrun₃.regs_eq (Code.preserves_rename_of_forall_ne _ _ _ ?_)] + · have h9 : j ≠ ga ∧ j ≠ gb := by + constructor <;> (intro h; subst h; simp at hj) + simp [y₂, h9.1, h9.2] + · intro a h; subst h; fin_cases a <;> simp [gcdMap] at hj + have hy₃1 : y₃.regs one = 1 := by rw [hy₃ one (by simp), hy₁ one (by simp), hy1] + have hy₃r : y₃.regs r = R := by rw [hy₃ r (by simp), hy₁r] + have hy₃ok : y₃.regs ok = y₁.regs ok := hy₃ ok (by simp) + -- Accept or advance. + have hfR1 : Fits w (R + 1) := hfr₁.of_le (by omega) + have hgood' : Good N L₀ R ↔ Nat.gcd N R = 1 ∧ y₁.regs ok = 1 := by + rw [hokiff]; rfl + by_cases hG : Good N L₀ R + · obtain ⟨hg1, ho1⟩ := hgood'.mp hG + let z := y₃.set found 1 + have hrun₄ : ExecWithin (B := .standard) w accept y₃ z 5 := by + refine ExecWithin.mono (T := 3) ?_ (by omega) + ram_eval [accept, z, hgcd, hg1, hy₃1, hy₃ok, ho1, h1f] + refine ⟨z, (hrun₁.seq (hrun₂.seq (hrun₃.seq hrun₄))).mono ?_, ?_, ?_⟩ + · rw [show y₂.regs (gcdMap NatGcd.b) = R by simp +decide [y₂, gcdMap]] + omega + · have hz (j : Fin 13) (hj : j ≠ found) : z.regs j = y₃.regs j := by simp [z, hj] + have hzr : z.regs r = R := by rw [hz r (by decide), hy₃r] + refine ⟨?_, ?_, ?_, by rw [hzr]; exact hr2, by rw [hzr]; exact hrr₀, + by simp +decide [z], fun _ => by rw [hzr]; exact hG⟩ + · rw [hz n (by decide), hy₃ n (by simp), hy₁n] + · rw [hz L (by decide), hy₃ L (by simp), hy₁ L (by simp), hyL] + · rw [hz one (by decide), hy₃1] + · simp +decide [z, hy₃r, hR, hf0] + · have hne : R ≠ r₀ := fun h => hG (h ▸ hgood) + let z := y₃.set r (R + 1) + have hrun₄ : ExecWithin (B := .standard) w accept y₃ z 5 := by + by_cases hg1 : Nat.gcd N R = 1 + · have ho1 : ¬ y₁.regs ok = 1 := fun h => hG (hgood'.mpr ⟨hg1, h⟩) + ram_eval [accept, z, hgcd, hg1, hy₃1, hy₃ok, ho1, hy₃r, hfR1] + · refine ExecWithin.mono (T := 3) ?_ (by omega) + ram_eval [accept, z, hgcd, hg1, hy₃1, hy₃r, hfR1] + refine ⟨z, (hrun₁.seq (hrun₂.seq (hrun₃.seq hrun₄))).mono ?_, ?_, ?_⟩ + · rw [show y₂.regs (gcdMap NatGcd.b) = R by simp +decide [y₂, gcdMap]] + omega + · have hf : z.regs found = 0 := by + simp +decide only [z, Data.set_regs, Function.update_apply, ↓reduceIte] + rw [hy₃ found (by simp), hy₁ found (by simp), hf0] + have hz (j : Fin 13) (hj : j ≠ r) : z.regs j = y₃.regs j := by simp [z, hj] + have hzr : z.regs r = R + 1 := by simp [z] + refine ⟨?_, ?_, ?_, by rw [hzr]; omega, by rw [hzr]; omega, by omega, by omega⟩ + · rw [hz n (by decide), hy₃ n (by simp), hy₁n] + · rw [hz L (by decide), hy₃ L (by simp), hy₁ L (by simp), hyL] + · rw [hz one (by decide), hy₃1] + · have hf : z.regs found = 0 := by + simp +decide only [z, Data.set_regs, Function.update_apply, ↓reduceIte] + rw [hy₃ found (by simp), hy₁ found (by simp), hf0] + rw [hf, hf0, hR] + simp [z] + omega + let x₀ := ((x.set one 1).set r 2).set found 0 + have hrun₀ : ExecWithin (B := .standard) w (ram[ one ← #1; r ← #2; found ← #0 ]) x x₀ 3 := by + ram_eval [x₀, h1f, hfr₁.of_le (by omega : 2 ≤ r₀ + 1)] + have hx₀ : I x₀ := by + refine ⟨?_, ?_, ?_, ?_, ?_, ?_, ?_⟩ <;> simp +decide [x₀, hn, hL, hr₀] + obtain ⟨y, hrun, hy, hc⟩ := ExecWithin.loop_invariant I + (fun y => 2 * (r₀ - y.regs r) + (1 - y.regs found)) hstep hx₀ + obtain ⟨hyn, hyL, hy1, hr2, hrr₀, hyf1, hgoody⟩ := hy + have hf1 : y.regs found = 1 := by simp only [Cond.holds_lt, hy1] at hc; omega + refine ⟨y, (hrun₀.seq hrun).mono ?_, hr2, hrr₀, hgoody hf1, hyn, hyL⟩ + have hx₀r : x₀.regs r = 2 := by simp +decide [x₀] + have hx₀f : x₀.regs found = 0 := by simp +decide [x₀] + have hμ : 2 * (r₀ - x₀.regs r) + (1 - x₀.regs found) ≤ 2 * r₀ + 1 := by omega + have := Nat.mul_le_mul hμ (le_refl (8 * L₀ + 5 * r₀ + 16 + 2)) + unfold time + nlinarith + +end Cslib.Algorithms.WordRAM.AKS.OrderSearch diff --git a/Complexity/Foundations/NumberTheory/AKS/IntegerVariant.lean b/Complexity/Foundations/NumberTheory/AKS/IntegerVariant.lean new file mode 100644 index 000000000..515aaed51 --- /dev/null +++ b/Complexity/Foundations/NumberTheory/AKS/IntegerVariant.lean @@ -0,0 +1,135 @@ +/- +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.Correctness +import Mathlib.Data.Nat.Log + +/-! +# The AKS tests with integer thresholds + +The AKS algorithm [Agrawal, Kayal, Saxena, Section 4][AKS2004] compares the multiplicative order +of `n` modulo `r` with `log₂² n` (step 2) and tests the congruence for `a ≤ ⌊√φ(r) log₂ n⌋` +(step 5). A program decides these comparisons more easily with the integer threshold +`c = ⌈log₂ n⌉` in place of `log₂ n`: step 2 asks for `ord_r(n) > c²`, and step 5 tests +`a ≤ ⌊√φ(r) · c⌋ = ⌊√(φ(r) c²)⌋`. This is a standard integer variant. It is still correct, by +Theorem 4.1 of the paper (`Cslib.AKS.prime_iff_passes`): `c² ≥ log₂² n`, so the order condition +is stronger than the paper's, and every prime passes the congruence for every `a`, so testing more +values of `a` loses nothing. + +`Cslib.AKS.PassesUpTo n r A` states the tests with the congruence checked for `1 ≤ a ≤ A`. For +`n ≥ 2`, `ord_r(n) > c²` and `A ≥ ⌊√φ(r) log₂ n⌋`, it holds exactly when `n` is prime +(`Cslib.AKS.prime_iff_passesUpTo`). The order condition is decided by small powers: for `r ≥ 2`, +`ord_r(n) > L` exactly when `n` is coprime to `r` and no `n^k` with `1 ≤ k ≤ L` is `1` modulo `r` +(`Cslib.AKS.lt_orderOf_iff`). + +## Main definitions + +* `Cslib.AKS.numTestsNat`: the integer bound `⌊√(φ(r) ⌈log₂ n⌉²)⌋` of step 5. +* `Cslib.AKS.PassesUpTo`: the tests, with the congruence for `1 ≤ a ≤ A`. + +## Main results + +* `Cslib.AKS.logb_le_clog`: `log₂ n ≤ ⌈log₂ n⌉`. +* `Cslib.AKS.numTests_le_numTestsNat`: the integer bound tests at least the paper's values. +* `Cslib.AKS.lt_orderOf_iff`: the order condition by small powers. +* `Cslib.AKS.prime_iff_passesUpTo`: the tests with integer thresholds characterize the primes. + +## References + +* [M. Agrawal, N. Kayal, N. Saxena, *PRIMES is in P*][AKS2004] +-/ + +@[expose] public section + +namespace Cslib.AKS + +/-- The integer bound `⌊√(φ(r) c²)⌋ = ⌊√φ(r) · c⌋` of step 5, with `c = ⌈log₂ n⌉`. -/ +def numTestsNat (n r : ℕ) : ℕ := Nat.sqrt (Nat.totient r * Nat.clog 2 n ^ 2) + +/-- **The tests with the congruence checked up to `A`**: `n` is not a perfect power (step 1), no +`a ≤ r` has `1 < gcd(a, n) < n` (step 3), and either `n ≤ r` (step 4) or the congruence holds for +every `1 ≤ a ≤ A` (step 5). -/ +def PassesUpTo (n r A : ℕ) : Prop := + ¬ IsPerfectPower n ∧ (∀ a, 1 ≤ a → a ≤ r → a.gcd n = 1 ∨ a.gcd n = n) ∧ + (n ≤ r ∨ ∀ a, 1 ≤ a → a ≤ A → Congruence n r a) + +/-- `log₂ n ≤ ⌈log₂ n⌉`. -/ +theorem logb_le_clog (n : ℕ) : Real.logb 2 n ≤ Nat.clog 2 n := by + rcases Nat.eq_zero_or_pos n with rfl | hn + · simp + have hn0 : (0 : ℝ) < n := by exact_mod_cast hn + rw [Real.logb_le_iff_le_rpow one_lt_two hn0, Real.rpow_natCast] + exact_mod_cast Nat.le_pow_clog one_lt_two n + +/-- The integer bound tests at least the paper's values: `⌊√φ(r) log₂ n⌋ ≤ ⌊√(φ(r) c²)⌋`. -/ +theorem numTests_le_numTestsNat (n r : ℕ) : numTests n r ≤ numTestsNat n r := by + rw [numTests, numTestsNat, ← Real.nat_floor_real_sqrt_eq_nat_sqrt] + refine Nat.floor_le_floor ?_ + rw [Nat.cast_mul, Nat.cast_pow, Real.sqrt_mul (Nat.cast_nonneg _), + Real.sqrt_sq (Nat.cast_nonneg _)] + rcases Nat.eq_zero_or_pos n with rfl | hn + · simp + have hlog : 0 ≤ Real.logb 2 n := + Real.logb_nonneg one_lt_two (by exact_mod_cast hn) + exact mul_le_mul_of_nonneg_left (logb_le_clog n) (Real.sqrt_nonneg _) + +/-- **The order condition by small powers**: for `r ≥ 2`, the multiplicative order of `n` modulo +`r` exceeds `L` exactly when `n` is coprime to `r` and no `n^k` with `1 ≤ k ≤ L` is `1` modulo +`r`. -/ +theorem lt_orderOf_iff {n r L : ℕ} (hr : 2 ≤ r) : + L < orderOf (n : ZMod r) ↔ n.Coprime r ∧ ∀ k, 1 ≤ k → k ≤ L → n ^ k % r ≠ 1 := by + have hpow (k : ℕ) : n ^ k % r = 1 ↔ (n : ZMod r) ^ k = 1 := by + rw [← Nat.cast_pow, ← Nat.cast_one (R := ZMod r), ZMod.natCast_eq_natCast_iff', + Nat.one_mod_eq_one.mpr (by omega)] + constructor + · intro h + refine ⟨coprime_of_orderOf_pos (by omega), fun k hk hkL hk1 => ?_⟩ + have := Nat.le_of_dvd (by omega) (orderOf_dvd_of_pow_eq_one ((hpow k).mp hk1)) + omega + · rintro ⟨hco, hsmall⟩ + have : NeZero r := ⟨by omega⟩ + have hpos : 0 < orderOf (n : ZMod r) := by + rw [← ZMod.coe_unitOfCoprime n hco, orderOf_units] + exact orderOf_pos _ + by_contra hle + exact hsmall _ hpos (by omega) ((hpow _).mpr (pow_orderOf_eq_one _)) + +/-- Every prime passes the tests, for every `r` and `A`. -/ +theorem passesUpTo_of_prime {p : ℕ} (hp : p.Prime) (r A : ℕ) : PassesUpTo p r A := by + obtain ⟨h1, h3, -⟩ := passes_of_prime hp r + exact ⟨h1, h3, Or.inr fun a _ _ => congruence_of_prime hp r a⟩ + +/-- **The tests with integer thresholds characterize the primes**: for `n ≥ 2`, `r` such that the +order of `n` modulo `r` exceeds `⌈log₂ n⌉²`, and `A ≥ ⌊√φ(r) log₂ n⌋`, `n` passes the tests up to +`A` exactly when it is prime. This is [Agrawal, Kayal, Saxena, Theorem 4.1][AKS2004] with integer +thresholds. -/ +theorem prime_iff_passesUpTo {n r A : ℕ} (hn : 2 ≤ n) + (hr : Nat.clog 2 n ^ 2 < orderOf (n : ZMod r)) (hA : numTests n r ≤ A) : + n.Prime ↔ PassesUpTo n r A := by + refine ⟨fun hp => passesUpTo_of_prime hp r A, fun ⟨h1, h3, h5⟩ => ?_⟩ + have hlog : Real.logb 2 n ^ 2 < orderOf (n : ZMod r) := by + have h0 : 0 ≤ Real.logb 2 n := + Real.logb_nonneg one_lt_two (by exact_mod_cast (by omega : 1 ≤ n)) + have h1 : Real.logb 2 n ^ 2 ≤ (Nat.clog 2 n : ℝ) ^ 2 := + pow_le_pow_left₀ h0 (logb_le_clog n) 2 + have h2 : ((Nat.clog 2 n ^ 2 : ℕ) : ℝ) < orderOf (n : ZMod r) := by exact_mod_cast hr + push_cast at h2 + linarith + refine prime_of_passes hn hlog ⟨h1, h3, ?_⟩ + rcases h5 with h | h + · exact Or.inl h + · exact Or.inr fun a ha1 ha => h a ha1 (ha.trans hA) + +/-- The tests with the integer bound `⌊√(φ(r) ⌈log₂ n⌉²)⌋` characterize the primes. -/ +theorem prime_iff_passesUpTo_numTestsNat {n r : ℕ} (hn : 2 ≤ n) + (hr : Nat.clog 2 n ^ 2 < orderOf (n : ZMod r)) : + n.Prime ↔ PassesUpTo n r (numTestsNat n r) := + prime_iff_passesUpTo hn hr (numTests_le_numTestsNat n r) + +end Cslib.AKS diff --git a/ComplexityTests/WordRAMAKSSearch.lean b/ComplexityTests/WordRAMAKSSearch.lean new file mode 100644 index 000000000..cfcac6cae --- /dev/null +++ b/ComplexityTests/WordRAMAKSSearch.lean @@ -0,0 +1,78 @@ +/- +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.GcdTotient +import Complexity.Algorithms.WordRAM.AKS.OrderSearch +import Complexity.Foundations.NumberTheory.AKS.IntegerVariant + +/-! +# The order search, step 3 and the integer thresholds of the AKS test + +Concrete runs of the search for `r` and of the step-3 and totient loop, instances of their +correctness theorems, and checked instances of the integer-threshold conditions. +-/ + +namespace ComplexityTests.WordRAMAKSSearch + +open Cslib.WordRAM Cslib.Algorithms.WordRAM + +/-! ## The search for `r` -/ + +private def searchInput (n L : ℕ) : Data 13 := + (Data.init []).set AKS.OrderSearch.n n |>.set AKS.OrderSearch.L L + +/-- For `n = 5` and `L = 2`: `r = 2` fails (`5 ≡ 1`), `r = 3` fails (`5² ≡ 1`), `r = 4` fails +(`5 ≡ 1`), and `r = 5` fails (`gcd = 5`); `r = 6` fails (`5² ≡ 1`), and `r = 7` is accepted. -/ +example : (Code.eval 16 16 AKS.OrderSearch.program (searchInput 5 2)).map + (·.1.regs AKS.OrderSearch.r) = some 7 := by + decide +kernel + +/-- `Good 5 2 7`: `5` is coprime to `7`, and `5 ≡ 5`, `5² ≡ 4` modulo `7`. -/ +theorem good_5_2_7 : AKS.OrderSearch.Good 5 2 7 := by + refine ⟨by decide, fun k hk1 hk2 => ?_⟩ + interval_cases k <;> decide + +/-- The correctness theorem with `r₀ = 7`: the search stops at a good `2 ≤ r ≤ 7`. -/ +example : ∃ y, ExecWithin 16 AKS.OrderSearch.program (searchInput 5 2) y + (AKS.OrderSearch.time 2 7) ∧ AKS.OrderSearch.Good 5 2 (y.regs AKS.OrderSearch.r) := by + obtain ⟨y, hy, -, -, hgood, -⟩ := AKS.OrderSearch.exec (w := 16) (searchInput 5 2) + (by decide) (by decide) (by decide) good_5_2_7 (by decide) (by decide) (by decide) + exact ⟨y, hy, hgood⟩ + +/-- `Good` is the order condition: `ord₇(5) = 6 > 2`. -/ +example : 2 < orderOf ((5 : ℕ) : ZMod 7) := + (Cslib.AKS.lt_orderOf_iff (by norm_num)).mpr good_5_2_7 + +/-! ## Step 3 and the totient -/ + +private def gcdInput (n r : ℕ) : Data 11 := + (Data.init []).set AKS.GcdTotient.n n |>.set AKS.GcdTotient.r r + +/-- For `n = 15` and `r = 4`, `gcd(3, 15) = 3` fails step 3; for `n = 7` it passes; `φ(4) = 2`. -/ +example : (Code.eval 16 8 AKS.GcdTotient.program (gcdInput 15 4)).map + (fun yn => (yn.1.regs AKS.GcdTotient.ok, yn.1.regs AKS.GcdTotient.phi)) = some (0, 2) ∧ + (Code.eval 16 8 AKS.GcdTotient.program (gcdInput 7 4)).map + (fun yn => (yn.1.regs AKS.GcdTotient.ok, yn.1.regs AKS.GcdTotient.phi)) = some (1, 2) := by + decide +kernel + +/-- The correctness theorem computes `φ(4) = 2`. -/ +example : ∃ y, ExecWithin 16 AKS.GcdTotient.program (gcdInput 7 4) y (AKS.GcdTotient.time 4) ∧ + y.regs AKS.GcdTotient.phi = 2 := by + obtain ⟨y, hy, -, -, hphi, -⟩ := AKS.GcdTotient.exec (w := 16) (gcdInput 7 4) (N := 7) (R := 4) + (by decide) + (by decide) (by decide) (by decide) (by decide) + exact ⟨y, hy, hphi.trans (by decide)⟩ + +/-! ## Integer thresholds -/ + +/-- Step 3 rejects `15` with `r = 3`, since `gcd(3, 15) = 3`. -/ +example (A : ℕ) : ¬ Cslib.AKS.PassesUpTo 15 3 A := by + rintro ⟨-, h3, -⟩ + have := h3 3 (by norm_num) le_rfl + revert this + decide + +end ComplexityTests.WordRAMAKSSearch From 71b9b9c2d0544cef2678c693de8148672c354fe3 Mon Sep 17 00:00:00 2001 From: Samuel Giuseppe Schlesinger Date: Fri, 30 Oct 2026 22:26:22 -0400 Subject: [PATCH 6/8] feat(WordRAM/AKS): the step-5 congruence loop (#893) CongruenceLoop runs the one-a congruence test for a = 1, ..., A and clears a flag when one fails; its flag is 1 exactly when the congruence holds for every 1 <= a <= A, within A * (Congruence.time r n + 6) + 6 instructions. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_015DLg8kwPs8yi7ePPdtoNvG --- Complexity.lean | 1 + .../WordRAM/AKS/CongruenceLoop.lean | 202 ++++++++++++++++++ ComplexityTests/WordRAMAKSSearch.lean | 29 ++- 3 files changed, 230 insertions(+), 2 deletions(-) create mode 100644 Complexity/Algorithms/WordRAM/AKS/CongruenceLoop.lean diff --git a/Complexity.lean b/Complexity.lean index ed2aead30..3ff3c1d6a 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.Congruence +public import Complexity.Algorithms.WordRAM.AKS.CongruenceLoop public import Complexity.Algorithms.WordRAM.AKS.GcdTotient public import Complexity.Algorithms.WordRAM.AKS.OrderSearch public import Complexity.Algorithms.WordRAM.AKS.PerfectPower diff --git a/Complexity/Algorithms/WordRAM/AKS/CongruenceLoop.lean b/Complexity/Algorithms/WordRAM/AKS/CongruenceLoop.lean new file mode 100644 index 000000000..f1b0a8c86 --- /dev/null +++ b/Complexity/Algorithms/WordRAM/AKS/CongruenceLoop.lean @@ -0,0 +1,202 @@ +/- +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.Congruence + +/-! +# Step 5 of the AKS algorithm on the word RAM + +Step 5 of the AKS algorithm [Agrawal, Kayal, Saxena, Section 4][AKS2004] tests the congruence +`(X + a)^n ≡ X^n + a (mod X^r - 1, n)` for `a = 1, …, A`. The program runs the test for one `a` +(`Cslib.Algorithms.WordRAM.AKS.Congruence.program`) for each `a` in turn and clears a flag when +one fails. + +## Main definitions + +* `Cslib.Algorithms.WordRAM.AKS.CongruenceLoop.program`: the loop. +* `Cslib.Algorithms.WordRAM.AKS.CongruenceLoop.time`: its instruction bound. + +## Main results + +* `Cslib.Algorithms.WordRAM.AKS.CongruenceLoop.exec`: the flag is `1` exactly when the congruence + holds for every `1 ≤ a ≤ A`, within `time r n A` instructions. + +## References + +* [M. Agrawal, N. Kayal, N. Saxena, *PRIMES is in P*][AKS2004] +-/ + +@[expose] public section + +namespace Cslib.Algorithms.WordRAM.AKS.CongruenceLoop + +open Cslib.WordRAM Cslib.AKS + +/-- Base address of the result block. -/ +abbrev R : Fin 29 := 0 +/-- Base address of the base block. -/ +abbrev B : Fin 29 := 1 +/-- Base address of the scratch block. -/ +abbrev T : Fin 29 := 2 +/-- The number `r` of coefficients. -/ +abbrev rr : Fin 29 := 3 +/-- The modulus `n`. -/ +abbrev nn : Fin 29 := 4 +/-- The tested value `a`, which is also the loop index. -/ +abbrev av : Fin 29 := 21 +/-- The answer of one test. -/ +abbrev flag : Fin 29 := 24 +/-- The last tested value `A`, preserved. -/ +abbrev A : Fin 29 := 25 +/-- The loop end `A + 1`. -/ +abbrev lim : Fin 29 := 26 +/-- `1` while every test so far has passed. -/ +abbrev ok : Fin 29 := 27 +/-- Constant one. -/ +abbrev c1 : Fin 29 := 28 + +/-- The test for one `a` runs on the first `25` registers. -/ +def congMap : Fin 25 → Fin 29 := Fin.castLE (by norm_num) + +theorem congMap_injective : congMap.Injective := Fin.castLE_injective _ + +/-- The test for one `a`. -/ +def cong : Code .standard 29 := Congruence.program.rename congMap + +/-- Test `a`, record a failure, and advance `a`. -/ +def body : Code .standard 29 := + .seq cong (ram[ if flag = c1 then [ ok ← ok ] else [ ok ← #0 ]; av ← av + c1 ]) + +/-- Test `a = 1, …, A`. -/ +def program : Code .standard 29 := + .seq (ram[ c1 ← #1; lim ← A + c1; av ← c1; ok ← #1 ]) (.loop (.lt av lim) body) + +/-- The instruction bound of `program`. -/ +def time (r n A : ℕ) : ℕ := A * (Congruence.time r n + 6) + 6 + +/-- The registers that the test for one `a` preserves. -/ +def Keep (j : Fin 29) : Prop := + j = R ∨ j = B ∨ j = T ∨ j = rr ∨ j = nn ∨ j = av ∨ j = A ∨ j = lim ∨ j = ok ∨ j = c1 + +instance : DecidablePred Keep := fun _ => by unfold Keep; infer_instance + +private theorem cong_preserves : ∀ j, Keep j → cong.Preserves j := by decide + +private theorem cong_exec {w : ℕ} (x : Data 29) {a b t r n v : ℕ} (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) + (hv : x.regs av = v) (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)) (hfv : Fits w v) : + ∃ y, ExecWithin w cong x y (Congruence.time r n) ∧ y.regs flag ≤ 1 ∧ + (y.regs flag = 1 ↔ Congruence n r v) ∧ ∀ j, Keep j → y.regs j = x.regs j := by + obtain ⟨y, hy, hle, hiff, -⟩ := Congruence.exec (x.comap congMap) ha hb ht hr hn hv hr0 hn0 hab + hta htb hfa hfb hft hf2 hfn hfv + have hy' := hy.rename congMap_injective + have hfl : (y.extend congMap x).regs flag = y.regs Congruence.flag := + congMap_injective.extend_apply _ _ Congruence.flag + exact ⟨_, hy', by rw [hfl]; exact hle, by rw [hfl]; exact hiff, + fun j hj => hy'.regs_eq (cong_preserves j hj)⟩ + +/-- **Step 5.** For `r ≥ 1` and `n ≥ 1`, with the three blocks pairwise disjoint and the word size +holding the block ends, `2 r`, `n² + n` and `A + 1`, the program leaves `1` in `ok` exactly when +`(X + a)^n ≡ X^n + a (mod X^r - 1, n)` for every `1 ≤ a ≤ A`, within `time r n A` instructions. -/ +theorem exec {w : ℕ} (x : Data 29) {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 A = 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₀ + 1)) : + ∃ y, ExecWithin w program x y (time r n A₀) ∧ y.regs ok ≤ 1 ∧ + (y.regs ok = 1 ↔ ∀ v, 1 ≤ v → v ≤ A₀ → Congruence n r v) ∧ + y.regs R = a ∧ y.regs B = b ∧ y.regs T = t ∧ y.regs rr = r ∧ y.regs nn = n := by + classical + have h1 : Fits w 1 := hfA.of_le (by omega) + let I (y : Data 29) : Prop := y.regs R = a ∧ y.regs B = b ∧ y.regs T = t ∧ y.regs rr = r ∧ + y.regs nn = n ∧ y.regs lim = A₀ + 1 ∧ y.regs c1 = 1 ∧ 1 ≤ y.regs av ∧ + y.regs av ≤ A₀ + 1 ∧ y.regs ok ≤ 1 ∧ + (y.regs ok = 1 ↔ ∀ v, 1 ≤ v → v < y.regs av → Congruence n r v) + have hstep (y : Data 29) (hy : I y) (hc : (Cond.lt av lim).Holds y) : + ∃ z, ExecWithin w body y z (Congruence.time r n + 4) ∧ I z ∧ + A₀ + 1 - z.regs av < A₀ + 1 - y.regs av := by + obtain ⟨hyR, hyB, hyT, hyr, hyn, hyl, hy1, hv1, hvA, hyok, hyiff⟩ := hy + obtain ⟨V, hV⟩ : ∃ V, y.regs av = V := ⟨_, rfl⟩ + have hVA : V ≤ A₀ := by simp only [Cond.holds_lt, hV, hyl] at hc; omega + rw [hV] at hv1 hyiff + obtain ⟨y₁, hrun₁, hle, hiff, hkeep⟩ := cong_exec (w := w) y hyR hyB hyT hyr hyn hV hr0 hn0 + hab hta htb hfa hfb hft hf2 hfn (hfA.of_le (by omega)) + have hk (j : Fin 29) (hj : Keep j) : y₁.regs j = y.regs j := hkeep j hj + let o := if Congruence n r V then y.regs ok else 0 + let z := (y₁.set ok o).set av (V + 1) + have hy₁1 : y₁.regs c1 = 1 := by rw [hk c1 (by simp [Keep]), hy1] + have hy₁ok : y₁.regs ok = y.regs ok := hk ok (by simp [Keep]) + have hy₁v : y₁.regs av = V := by rw [hk av (by simp [Keep]), hV] + have hfV : Fits w (V + 1) := hfA.of_le (by omega) + have hrun₂ : ExecWithin (B := .standard) w + (ram[ if flag = c1 then [ ok ← ok ] else [ ok ← #0 ]; av ← av + c1 ]) y₁ z 4 := by + by_cases hC : Congruence n r V + · have hf1 : y₁.regs flag = 1 := hiff.mpr hC + refine ExecWithin.mono (T := 3) ?_ (by omega) + ram_eval [z, o, hC, hf1, hy₁1, hy₁ok, hy₁v, hfV] + · have hf1 : ¬ y₁.regs flag = 1 := fun h => hC (hiff.mp h) + ram_eval [z, o, hC, hf1, hy₁1, hy₁v, hfV] + have hz (j : Fin 29) (hj : Keep j) (hjo : j ≠ ok) (hjv : j ≠ av) : z.regs j = y.regs j := by + simp only [z, Data.set_regs, Function.update_of_ne hjv, Function.update_of_ne hjo] + exact hk j hj + have hzv : z.regs av = V + 1 := by simp [z] + have hzo : z.regs ok = o := by simp +decide [z] + refine ⟨z, hrun₁.seq hrun₂, ⟨?_, ?_, ?_, ?_, ?_, ?_, ?_, ?_, ?_, ?_, ?_⟩, ?_⟩ + · rw [hz R (by simp [Keep]) (by decide) (by decide), hyR] + · rw [hz B (by simp [Keep]) (by decide) (by decide), hyB] + · rw [hz T (by simp [Keep]) (by decide) (by decide), hyT] + · rw [hz rr (by simp [Keep]) (by decide) (by decide), hyr] + · rw [hz nn (by simp [Keep]) (by decide) (by decide), hyn] + · rw [hz lim (by simp [Keep]) (by decide) (by decide), hyl] + · rw [hz c1 (by simp [Keep]) (by decide) (by decide), hy1] + · rw [hzv]; omega + · rw [hzv]; omega + · rw [hzo]; simp only [o]; split_ifs <;> omega + · rw [hzo, hzv] + simp only [o] + constructor + · intro h v hv1' hv + split_ifs at h with hC + rcases Nat.lt_succ_iff_lt_or_eq.mp hv with hv | rfl + · exact hyiff.mp h v hv1' hv + · exact hC + · intro h + have hC : Congruence n r V := h V hv1 (by omega) + simp only [hC, ↓reduceIte] + exact hyiff.mpr fun v hv1' hv => h v hv1' (by omega) + · rw [hzv, hV]; omega + let x₀ := (((x.set c1 1).set lim (A₀ + 1)).set av 1).set ok 1 + have hrun₀ : ExecWithin (B := .standard) w (ram[ c1 ← #1; lim ← A + c1; av ← c1; ok ← #1 ]) + x x₀ 4 := by + ram_eval [x₀, hA, h1, hfA] + have hx₀ : I x₀ := by + refine ⟨by simp +decide [x₀, ha], by simp +decide [x₀, hb], by simp +decide [x₀, ht], + by simp +decide [x₀, hr], by simp +decide [x₀, hn], by simp +decide [x₀], + by simp +decide [x₀], by simp +decide [x₀], by simp +decide [x₀], by simp +decide [x₀], ?_⟩ + have hv : x₀.regs av = 1 := by simp +decide [x₀] + have ho : x₀.regs ok = 1 := by simp +decide [x₀] + rw [hv, ho] + exact ⟨fun _ v hv1 hv => by omega, fun _ => rfl⟩ + obtain ⟨y, hrun, hy, hc⟩ := ExecWithin.loop_invariant I (fun y => A₀ + 1 - y.regs av) hstep hx₀ + obtain ⟨hyR, hyB, hyT, hyr, hyn, hyl, -, -, hvA, hyok, hyiff⟩ := hy + have hend : y.regs av = A₀ + 1 := by simp only [Cond.holds_lt, hyl] at hc; omega + refine ⟨y, (hrun₀.seq hrun).mono ?_, hyok, ?_, hyR, hyB, hyT, hyr, hyn⟩ + · have hv : x₀.regs av = 1 := by simp +decide [x₀] + rw [hv, show A₀ + 1 - 1 = A₀ by omega, show Congruence.time r n + 4 + 2 = + Congruence.time r n + 6 by ring] + unfold time + omega + · rw [hyiff, hend] + exact ⟨fun h v hv1 hv => h v hv1 (by omega), fun h v hv1 hv => h v hv1 (by omega)⟩ + +end Cslib.Algorithms.WordRAM.AKS.CongruenceLoop diff --git a/ComplexityTests/WordRAMAKSSearch.lean b/ComplexityTests/WordRAMAKSSearch.lean index cfcac6cae..2622be1ad 100644 --- a/ComplexityTests/WordRAMAKSSearch.lean +++ b/ComplexityTests/WordRAMAKSSearch.lean @@ -4,6 +4,7 @@ Released under Apache 2.0 license as described in the file LICENSE. Authors: Samuel Schlesinger -/ +import Complexity.Algorithms.WordRAM.AKS.CongruenceLoop import Complexity.Algorithms.WordRAM.AKS.GcdTotient import Complexity.Algorithms.WordRAM.AKS.OrderSearch import Complexity.Foundations.NumberTheory.AKS.IntegerVariant @@ -11,8 +12,9 @@ import Complexity.Foundations.NumberTheory.AKS.IntegerVariant /-! # The order search, step 3 and the integer thresholds of the AKS test -Concrete runs of the search for `r` and of the step-3 and totient loop, instances of their -correctness theorems, and checked instances of the integer-threshold conditions. +Concrete runs of the search for `r`, of the step-3 and totient loop and of the step-5 loop, +instances of their correctness theorems, and checked instances of the integer-threshold +conditions. -/ namespace ComplexityTests.WordRAMAKSSearch @@ -66,6 +68,29 @@ example : ∃ y, ExecWithin 16 AKS.GcdTotient.program (gcdInput 7 4) y (AKS.GcdT (by decide) (by decide) (by decide) (by decide) exact ⟨y, hy, hphi.trans (by decide)⟩ +/-! ## Step 5 -/ + +/-- Blocks at `0`, `r` and `2r`, with `r`, `n` and the last tested value `A` in their registers. -/ +private def loopInput (r n A : ℕ) : Data 29 := + ⟨fun j => if j = AKS.CongruenceLoop.B then r else if j = AKS.CongruenceLoop.T then 2 * r else + if j = AKS.CongruenceLoop.rr then r else if j = AKS.CongruenceLoop.nn then n else + if j = AKS.CongruenceLoop.A then A else 0, + fun _ => 0⟩ + +/-- Modulo `X - 1`, `4` fails at `a = 1` (`2⁴ ≡ 0 ≢ 2`), so the flag is cleared. -/ +example : (Code.eval 16 4 AKS.CongruenceLoop.program (loopInput 1 4 1)).map + (·.1.regs AKS.CongruenceLoop.ok) = some 0 := by + decide +kernel + +/-- The correctness theorem: for the prime `5`, every `a ≤ 3` passes modulo `X² - 1`. -/ +example : ∃ y, ExecWithin 16 AKS.CongruenceLoop.program (loopInput 2 5 3) y + (AKS.CongruenceLoop.time 2 5 3) ∧ y.regs AKS.CongruenceLoop.ok = 1 := by + obtain ⟨y, hy, -, hiff, -⟩ := AKS.CongruenceLoop.exec (w := 16) (loopInput 2 5 3) (a := 0) + (b := 2) (t := 4) (r := 2) (n := 5) (A₀ := 3) 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) + exact ⟨y, hy, hiff.mpr fun v _ _ => Cslib.AKS.congruence_of_prime (by decide) 2 v⟩ + /-! ## Integer thresholds -/ /-- Step 3 rejects `15` with `r = 3`, since `gcd(3, 15) = 3`. -/ From 7bf57c7b71e230c7914cc475754bfe78e3226f41 Mon Sep 17 00:00:00 2001 From: Samuel Giuseppe Schlesinger Date: Fri, 30 Oct 2026 22:56:16 -0400 Subject: [PATCH 7/8] style(WordRAM/AKS): make a non-terminal simp terminal in the order search (#893) Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_015DLg8kwPs8yi7ePPdtoNvG --- Complexity/Algorithms/WordRAM/AKS/OrderSearch.lean | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/Complexity/Algorithms/WordRAM/AKS/OrderSearch.lean b/Complexity/Algorithms/WordRAM/AKS/OrderSearch.lean index df6b205c3..1e6c7f8aa 100644 --- a/Complexity/Algorithms/WordRAM/AKS/OrderSearch.lean +++ b/Complexity/Algorithms/WordRAM/AKS/OrderSearch.lean @@ -271,8 +271,7 @@ theorem exec {w : ℕ} (x : Data 13) {N L₀ r₀ : ℕ} (hn : x.regs n = N) (hL · have hf : z.regs found = 0 := by simp +decide only [z, Data.set_regs, Function.update_apply, ↓reduceIte] rw [hy₃ found (by simp), hy₁ found (by simp), hf0] - rw [hf, hf0, hR] - simp [z] + rw [hf, hf0, hR, show z.regs r = R + 1 by simp [z]] omega let x₀ := ((x.set one 1).set r 2).set found 0 have hrun₀ : ExecWithin (B := .standard) w (ram[ one ← #1; r ← #2; found ← #0 ]) x x₀ 3 := by From dcc17f6d9bb2247c61fa19bc3677f88700a264f7 Mon Sep 17 00:00:00 2001 From: Samuel Giuseppe Schlesinger Date: Sat, 31 Oct 2026 00:28:30 -0400 Subject: [PATCH 8/8] refactor(AKS): define Passes as PassesUpTo at numTests (#893) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Move `PassesUpTo` and `passesUpTo_of_prime` from `IntegerVariant.lean` into `Criterion.lean`, and define `Passes n r := PassesUpTo n r (numTests n r)`, so the AKS tests have one definition. `passes_of_prime` now specializes `passesUpTo_of_prime`. No statement changes; no aliases. Also state the `GcdTotient.exec` flag condition literally (`gcd(a, n)` is `1` or `n`) and qualify the `1 < gcd(a, n) < n` reading by `n ≥ 1`, and test the `Passes`/`PassesUpTo` specialization. Addresses reviewer-codex-edemaine's review of 71b9b9c2 on #1898. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_015DLg8kwPs8yi7ePPdtoNvG --- .../Algorithms/WordRAM/AKS/GcdTotient.lean | 5 ++- .../NumberTheory/AKS/Criterion.lean | 39 ++++++++++++------- .../NumberTheory/AKS/IntegerVariant.lean | 18 ++------- ComplexityTests/AKSCriterion.lean | 9 ++++- 4 files changed, 39 insertions(+), 32 deletions(-) diff --git a/Complexity/Algorithms/WordRAM/AKS/GcdTotient.lean b/Complexity/Algorithms/WordRAM/AKS/GcdTotient.lean index 8e1749171..055f52590 100644 --- a/Complexity/Algorithms/WordRAM/AKS/GcdTotient.lean +++ b/Complexity/Algorithms/WordRAM/AKS/GcdTotient.lean @@ -217,8 +217,9 @@ private theorem body_step {w N R : ℕ} (hfN : Fits w N) (hfR : Fits w (R + 1)) · rw [hza, hA]; omega /-- **Step 3 and the totient.** For `r ≥ 2`, with the word size holding `n` and `r + 1`, within -`time r` instructions the flag `ok` is `1` exactly when no `1 ≤ a ≤ r` has `1 < gcd(a, n) < n`, -and `phi` is `φ(r)`. -/ +`time r` instructions the flag `ok` is `1` exactly when every `1 ≤ a ≤ r` has `gcd(a, n) = 1` or +`gcd(a, n) = n`, and `phi` is `φ(r)`. For `n ≥ 1` the flag condition says that no `1 ≤ a ≤ r` has +`1 < gcd(a, n) < n`; for `n = 0` it is stronger, since `gcd(a, 0) = a`. -/ theorem exec {w : ℕ} (x : Data 11) {N R : ℕ} (hn : x.regs n = N) (hr : x.regs r = R) (hR : 2 ≤ R) (hfN : Fits w N) (hfR : Fits w (R + 1)) : ∃ y, ExecWithin w program x y (time R) ∧ y.regs ok ≤ 1 ∧ diff --git a/Complexity/Foundations/NumberTheory/AKS/Criterion.lean b/Complexity/Foundations/NumberTheory/AKS/Criterion.lean index 845c948c5..1a2382748 100644 --- a/Complexity/Foundations/NumberTheory/AKS/Criterion.lean +++ b/Complexity/Foundations/NumberTheory/AKS/Criterion.lean @@ -29,8 +29,10 @@ Checking the identity directly takes `n` coefficients. The AKS algorithm checks `X^r - 1`, for a small `r`, and for a few values of `a`. Given `r`, the algorithm outputs PRIME exactly when `n` is not a perfect power, no `a ≤ r` has `1 < gcd(a, n) < n`, and either `n ≤ r` or `(X + a)^n ≡ X^n + a (mod X^r - 1, n)` for every `1 ≤ a ≤ ⌊√φ(r) log₂ n⌋` -(`Cslib.AKS.Passes`). Every prime passes, for every `r` (`Cslib.AKS.passes_of_prime`): this is -the easy direction of [Agrawal, Kayal, Saxena, Theorem 4.1][AKS2004]. The converse, for `r` with +(`Cslib.AKS.Passes`). `Cslib.AKS.PassesUpTo n r A` states the same tests with the congruence +checked for `1 ≤ a ≤ A`, and `Passes n r` is its case `A = ⌊√φ(r) log₂ n⌋`. Every prime passes, +for every `r` and `A` (`Cslib.AKS.passesUpTo_of_prime`, `Cslib.AKS.passes_of_prime`): this is the +easy direction of [Agrawal, Kayal, Saxena, Theorem 4.1][AKS2004]. The converse, for `r` with multiplicative order of `n` modulo `r` above `log₂² n`, is the substance of their theorem. ## Main definitions @@ -38,6 +40,7 @@ multiplicative order of `n` modulo `r` above `log₂² n`, is the substance of t * `Cslib.AKS.IsPerfectPower`: `n = a ^ b` for some `b > 1`. * `Cslib.AKS.numTests`: the number `⌊√φ(r) log₂ n⌋` of tested values of `a`. * `Cslib.AKS.Congruence`: `(X + a)^n ≡ X^n + a (mod X^r - 1, n)`. +* `Cslib.AKS.PassesUpTo`: the AKS tests, with the congruence checked for `1 ≤ a ≤ A`. * `Cslib.AKS.Passes`: the conditions under which the AKS algorithm, given `r`, outputs PRIME. ## Main results @@ -45,7 +48,7 @@ multiplicative order of `n` modulo `r` above `log₂² n`, is the substance of t * `Cslib.AKS.not_dvd_choose_of_prime_dvd`: a composite `n` does not divide `C(n, q)` for a prime factor `q < n`. * `Cslib.AKS.prime_iff_add_pow_eq`: the AKS identity characterizes primes. -* `Cslib.AKS.passes_of_prime`: every prime passes the AKS tests. +* `Cslib.AKS.passesUpTo_of_prime`, `Cslib.AKS.passes_of_prime`: every prime passes the AKS tests. ## References @@ -141,14 +144,19 @@ def Congruence (n r a : ℕ) : Prop := AdjoinRoot.mk (X ^ r - 1 : (ZMod n)[X]) ((X + C (a : ZMod n)) ^ n) = AdjoinRoot.mk (X ^ r - 1 : (ZMod n)[X]) (X ^ n + C (a : ZMod n)) -/-- **The conditions under which the AKS algorithm outputs PRIME**, given `r` -[Agrawal, Kayal, Saxena, Section 4][AKS2004]: `n` is not a perfect power (step 1), no `a ≤ r` has -`1 < gcd(a, n) < n` (step 3), and either `n ≤ r` (step 4) or the congruence holds for every -`1 ≤ a ≤ ⌊√φ(r) log₂ n⌋` (step 5). The algorithm chooses `r` as the least number for which the -multiplicative order of `n` modulo `r` exceeds `log₂² n` (step 2). -/ -def Passes (n r : ℕ) : Prop := +/-- **The AKS tests with the congruence checked up to `A`**: `n` is not a perfect power (step 1), +every `1 ≤ a ≤ r` has `gcd(a, n) = 1` or `gcd(a, n) = n` (step 3; for `n ≥ 1`, no `a ≤ r` has +`1 < gcd(a, n) < n`), and either `n ≤ r` (step 4) or the congruence holds for every `1 ≤ a ≤ A` +(step 5). -/ +def PassesUpTo (n r A : ℕ) : Prop := ¬ IsPerfectPower n ∧ (∀ a, 1 ≤ a → a ≤ r → a.gcd n = 1 ∨ a.gcd n = n) ∧ - (n ≤ r ∨ ∀ a, 1 ≤ a → a ≤ numTests n r → Congruence n r a) + (n ≤ r ∨ ∀ a, 1 ≤ a → a ≤ A → Congruence n r a) + +/-- **The conditions under which the AKS algorithm outputs PRIME**, given `r` +[Agrawal, Kayal, Saxena, Section 4][AKS2004]: the tests of `PassesUpTo`, with the congruence +checked for every `1 ≤ a ≤ ⌊√φ(r) log₂ n⌋` in step 5. The algorithm chooses `r` as the least +number for which the multiplicative order of `n` modulo `r` exceeds `log₂² n` (step 2). -/ +def Passes (n r : ℕ) : Prop := PassesUpTo n r (numTests n r) /-- A prime is not a perfect power. -/ theorem not_isPerfectPower_of_prime {p : ℕ} (hp : p.Prime) : ¬ IsPerfectPower p := by @@ -159,13 +167,18 @@ theorem not_isPerfectPower_of_prime {p : ℕ} (hp : p.Prime) : ¬ IsPerfectPower theorem congruence_of_prime {p : ℕ} (hp : p.Prime) (r a : ℕ) : Congruence p r a := by rw [Congruence, add_pow_eq_of_prime hp] -/-- **Every prime passes the AKS tests**, for every `r`: the easy direction of -[Agrawal, Kayal, Saxena, Theorem 4.1][AKS2004]. -/ -theorem passes_of_prime {p : ℕ} (hp : p.Prime) (r : ℕ) : Passes p r := by +/-- **Every prime passes the AKS tests**, for every `r` and every bound `A` on the tested `a`: +the easy direction of [Agrawal, Kayal, Saxena, Theorem 4.1][AKS2004]. -/ +theorem passesUpTo_of_prime {p : ℕ} (hp : p.Prime) (r A : ℕ) : PassesUpTo p r A := by refine ⟨not_isPerfectPower_of_prime hp, fun a _ _ => ?_, Or.inr fun a _ _ => congruence_of_prime hp r a⟩ rcases (Nat.coprime_or_dvd_of_prime hp a) with h | h · exact Or.inl (Nat.Coprime.symm h) · exact Or.inr (Nat.gcd_eq_right h) +/-- **Every prime passes the AKS tests**, for every `r`: the easy direction of +[Agrawal, Kayal, Saxena, Theorem 4.1][AKS2004]. -/ +theorem passes_of_prime {p : ℕ} (hp : p.Prime) (r : ℕ) : Passes p r := + passesUpTo_of_prime hp r (numTests p r) + end Cslib.AKS diff --git a/Complexity/Foundations/NumberTheory/AKS/IntegerVariant.lean b/Complexity/Foundations/NumberTheory/AKS/IntegerVariant.lean index 515aaed51..6281427c9 100644 --- a/Complexity/Foundations/NumberTheory/AKS/IntegerVariant.lean +++ b/Complexity/Foundations/NumberTheory/AKS/IntegerVariant.lean @@ -22,8 +22,9 @@ Theorem 4.1 of the paper (`Cslib.AKS.prime_iff_passes`): `c² ≥ log₂² n`, s is stronger than the paper's, and every prime passes the congruence for every `a`, so testing more values of `a` loses nothing. -`Cslib.AKS.PassesUpTo n r A` states the tests with the congruence checked for `1 ≤ a ≤ A`. For -`n ≥ 2`, `ord_r(n) > c²` and `A ≥ ⌊√φ(r) log₂ n⌋`, it holds exactly when `n` is prime +`Cslib.AKS.PassesUpTo n r A` states the tests with the congruence checked for `1 ≤ a ≤ A`; the +paper's `Cslib.AKS.Passes n r` is its case `A = ⌊√φ(r) log₂ n⌋`. For `n ≥ 2`, `ord_r(n) > c²` and +`A ≥ ⌊√φ(r) log₂ n⌋`, it holds exactly when `n` is prime (`Cslib.AKS.prime_iff_passesUpTo`). The order condition is decided by small powers: for `r ≥ 2`, `ord_r(n) > L` exactly when `n` is coprime to `r` and no `n^k` with `1 ≤ k ≤ L` is `1` modulo `r` (`Cslib.AKS.lt_orderOf_iff`). @@ -31,7 +32,6 @@ values of `a` loses nothing. ## Main definitions * `Cslib.AKS.numTestsNat`: the integer bound `⌊√(φ(r) ⌈log₂ n⌉²)⌋` of step 5. -* `Cslib.AKS.PassesUpTo`: the tests, with the congruence for `1 ≤ a ≤ A`. ## Main results @@ -52,13 +52,6 @@ namespace Cslib.AKS /-- The integer bound `⌊√(φ(r) c²)⌋ = ⌊√φ(r) · c⌋` of step 5, with `c = ⌈log₂ n⌉`. -/ def numTestsNat (n r : ℕ) : ℕ := Nat.sqrt (Nat.totient r * Nat.clog 2 n ^ 2) -/-- **The tests with the congruence checked up to `A`**: `n` is not a perfect power (step 1), no -`a ≤ r` has `1 < gcd(a, n) < n` (step 3), and either `n ≤ r` (step 4) or the congruence holds for -every `1 ≤ a ≤ A` (step 5). -/ -def PassesUpTo (n r A : ℕ) : Prop := - ¬ IsPerfectPower n ∧ (∀ a, 1 ≤ a → a ≤ r → a.gcd n = 1 ∨ a.gcd n = n) ∧ - (n ≤ r ∨ ∀ a, 1 ≤ a → a ≤ A → Congruence n r a) - /-- `log₂ n ≤ ⌈log₂ n⌉`. -/ theorem logb_le_clog (n : ℕ) : Real.logb 2 n ≤ Nat.clog 2 n := by rcases Nat.eq_zero_or_pos n with rfl | hn @@ -100,11 +93,6 @@ theorem lt_orderOf_iff {n r L : ℕ} (hr : 2 ≤ r) : by_contra hle exact hsmall _ hpos (by omega) ((hpow _).mpr (pow_orderOf_eq_one _)) -/-- Every prime passes the tests, for every `r` and `A`. -/ -theorem passesUpTo_of_prime {p : ℕ} (hp : p.Prime) (r A : ℕ) : PassesUpTo p r A := by - obtain ⟨h1, h3, -⟩ := passes_of_prime hp r - exact ⟨h1, h3, Or.inr fun a _ _ => congruence_of_prime hp r a⟩ - /-- **The tests with integer thresholds characterize the primes**: for `n ≥ 2`, `r` such that the order of `n` modulo `r` exceeds `⌈log₂ n⌉²`, and `A ≥ ⌊√φ(r) log₂ n⌋`, `n` passes the tests up to `A` exactly when it is prime. This is [Agrawal, Kayal, Saxena, Theorem 4.1][AKS2004] with integer diff --git a/ComplexityTests/AKSCriterion.lean b/ComplexityTests/AKSCriterion.lean index f9ebb874a..4bb4107a7 100644 --- a/ComplexityTests/AKSCriterion.lean +++ b/ComplexityTests/AKSCriterion.lean @@ -6,8 +6,8 @@ Authors: Samuel Schlesinger import Complexity.Foundations.NumberTheory.AKS.Criterion -/-! The AKS identity on a prime and a composite, the binomial divisibility, and the AKS tests on -a prime. -/ +/-! The AKS identity on a prime and a composite, the binomial divisibility, the AKS tests on a +prime, and `Passes` as `PassesUpTo` at the paper's bound. -/ namespace ComplexityTests.AKSCriterion @@ -27,4 +27,9 @@ example (r : ℕ) : passes_of_prime (by decide) r, fun h => h.1 ⟨2, 3, by norm_num, by norm_num⟩⟩ +-- `Passes` is `PassesUpTo` at the paper's bound `⌊√φ(r) log₂ n⌋`, and the prime `7` passes the +-- tests up to every bound. +example (n r A : ℕ) : (Passes n r ↔ PassesUpTo n r (numTests n r)) ∧ PassesUpTo 7 r A := + ⟨Iff.rfl, passesUpTo_of_prime (by decide) r A⟩ + end ComplexityTests.AKSCriterion