Multilinear folding completeness

ArkLib.lean

@@ -77,6 +77,7 @@ import ArkLib.Data.CodingTheory.ProximityGap.DG25.Basic
 import ArkLib.Data.CodingTheory.ProximityGap.DG25.MainResults
 import ArkLib.Data.CodingTheory.ProximityGap.DG25.ReedSolomon
 import ArkLib.Data.CodingTheory.ProximityGap.Folding
+import ArkLib.Data.CodingTheory.ProximityGap.Folding.Multilinear
 import ArkLib.Data.CodingTheory.ProximityGap.MCAGenerator
 import ArkLib.Data.CodingTheory.ProximityGap.ProximityGenerators
 import ArkLib.Data.CodingTheory.ReedSolomon
@@ -142,6 +143,7 @@ import ArkLib.Data.Matrix.Sparse
 import ArkLib.Data.Matrix.Vandermonde
 import ArkLib.Data.Misc.Basic
 import ArkLib.Data.MvPolynomial.Degrees
+import ArkLib.Data.MvPolynomial.EvenAndOdd
 import ArkLib.Data.MvPolynomial.Interpolation
 import ArkLib.Data.MvPolynomial.LinearMvExtension
 import ArkLib.Data.MvPolynomial.Multilinear

ArkLib/Data/CodingTheory/ProximityGap/Folding.lean

@@ -262,6 +262,34 @@ theorem foldWord_k_1 [NeZero n] {i : Fin (2 ^ (n - 1))} {α : F} :
     ((f i + f i') / 2) + α * ((f i - f i') / (2 * x)) := by
   simp [foldWord, foldValue_k_1]
 
+/-- An explicit formula for `foldWord` when `k = 1` that
+  does not use Lagrange interpolation. Function-level version. -/
+theorem foldWord_k_1' [NeZero n] {α : F} :
+  foldWord domain f 1 α = fun i ↦
+    let x : domain := CosetFftDomain.twoNthRoot (i := 1)
+        ⟨domain.subdomain 1 i, by simp⟩
+    let i := domain.log x
+    let i' := domain.log ⟨-x.1, by obtain ⟨x, hx⟩ := x; simpa using hx⟩
+    ((f i + f i') / 2) + α * ((f i - f i') / (2 * x)) := by aesop (add simp [foldWord_k_1])
+
+/-- The version of a folding where
+  k steps are achieved via iterated application
+  of k=1 folding. -/
+noncomputable def iteratedFoldWord (domain : SmoothCosetFftDomain n F)
+  (f : Word F (Fin (2 ^ n))) (k : ℕ) (α : Fin k → F) :
+  Word F (Fin (2 ^ (n - k))) :=
+  match k with
+  | 0 => f
+  | Nat.succ k =>
+    let prev := iteratedFoldWord domain f k (fun i ↦ α ⟨i.val, by omega⟩)
+    let foldedPrev :=
+      foldWord (domain.subdomain k) prev 1 (α ⟨k, by omega⟩)
+    fun i ↦ foldedPrev ⟨i.val, by aesop (add safe cases Fin)⟩
+
+@[simp]
+lemma iteratedFoldWord_zero {α : Fin 0 → F} :
+  iteratedFoldWord domain f 0 α = f := rfl
+
 omit [DecidableEq F] in
 /-- TODO: this will go once this https://github.com/Verified-zkEVM/CompPoly/pull/203
   is merged. -/
@@ -417,6 +445,49 @@ theorem foldWord_mem_code_of_mem_code {d : ℕ}
         rw [ReedSolomon.toPolynomial_eval_at_domain]
         simp [evalOnPoints]
 
+private lemma div_two_pow_div_two (d k : ℕ) :
+    d / 2 ^ k / 2 ^ 1 = d / 2 ^ (k + 1) := by
+  rw [pow_one, Nat.div_div_eq_div_mul, ←pow_succ]
+
+/-- Perfect completeness of iterated folding: if a word belongs to an RS-code
+  then its `iteratedFoldWord` belongs to a folded RS-code.
+-/
+theorem iteratedFoldWord_mem_code_of_mem_code {d : ℕ}
+  {α : Fin k → F}
+  (hk : k ≤ n)
+  (hk_d_dvd : 2 ^ k ∣ d)
+  {f : Word F (Fin (2 ^ n))}
+  (hf : f ∈ ReedSolomon.code (domain : Fin (2 ^ n) ↪ F) d) :
+  iteratedFoldWord domain f k α ∈
+    ReedSolomon.code (domain.subdomain k : Fin (2 ^ (n - k)) ↪ F) (d / (2 ^ k)) := by
+  induction k with
+  | zero => simp [hf]
+  | succ k ih =>
+    have hdvd_k : 2 ^ k ∣ d := dvd_trans (pow_dvd_pow 2 (Nat.le_succ k)) hk_d_dvd
+    have hprev := ih (α := fun i ↦ α ⟨i.val, by omega⟩) (by omega) hdvd_k
+    have hk1 : (1 : ℕ) ≤ n - k := by omega
+    have hdvd1 : (2 : ℕ) ^ 1 ∣ d / 2 ^ k := by
+      obtain ⟨c, rfl⟩ := hk_d_dvd
+      refine ⟨c, ?_⟩
+      rw [pow_succ, mul_assoc, Nat.mul_div_cancel_left _ (by positivity), pow_one]
+    have hfold := foldWord_mem_code_of_mem_code (domain := domain.subdomain k) (k := 1)
+        (α := α ⟨k, by omega⟩) hk1 hdvd1 hprev
+    rw [ReedSolomon.mem_code_iff_exists_polynomial] at hfold
+    obtain ⟨p, hpdeg, hpeval⟩ := hfold
+    rw [div_two_pow_div_two] at hpdeg
+    have hunfold : iteratedFoldWord domain f (k + 1) α =
+        fun i : Fin (2 ^ (n - (k + 1))) ↦
+            foldWord (domain.subdomain k)
+              (iteratedFoldWord domain f k (fun j ↦ α ⟨j.val, by omega⟩)) 1
+              (α ⟨k, by omega⟩) ⟨i.val, by rw [Nat.sub_sub]; exact i.isLt⟩ := rfl
+    rw [hunfold, ReedSolomon.mem_code_iff_exists_polynomial]
+    refine ⟨p, hpdeg, ?_⟩
+    funext i
+    have hb := subdomain_one_comp (ω := domain) hk ⟨i.val, by rw [Nat.sub_sub]; exact i.isLt⟩ i rfl
+    simp only [hpeval, ReedSolomon.evalOnPoints, Function.Embedding.coeFn_mk, LinearMap.coe_mk,
+      AddHom.coe_mk]
+    rw [hb]
+
 private noncomputable def foldWordAuxCoeff (domain : SmoothCosetFftDomain n F)
   (f : Word F (Fin (2 ^ n))) (k : ℕ) (i : Fin k) (x : F) : F :=
   (foldWordAux domain f k x).coeff i

ArkLib/Data/CodingTheory/ProximityGap/Folding/Multilinear.lean

@@ -0,0 +1,182 @@
+/-
+Copyright (c) 2024-2026 ArkLib Contributors. All rights reserved.
+Released under Apache 2.0 license as described in the file LICENSE.
+Authors: Ilia Vlasov, Aristotle (Harmonic)
+-/
+
+import Mathlib.Algebra.Polynomial.Roots
+import Mathlib.LinearAlgebra.Lagrange
+
+import ArkLib.Data.CodingTheory.ProximityGap.Basic
+import ArkLib.Data.CodingTheory.ProximityGap.BCIKS20.Curves
+import ArkLib.Data.CodingTheory.ProximityGap.Folding
+import ArkLib.Data.Domain.CosetFftDomain.Subdomain
+import ArkLib.Data.Domain.CosetFftDomain.Log
+import ArkLib.Data.MvPolynomial.EvenAndOdd
+import CompPoly.Data.MvPolynomial.Notation
+
+/-! This module provides an equivalent statement
+  of folding completeness of RS-codes in terms of multilinear polynomials
+  as can be found in [ACFY24].
+
+## References
+
+  * [Arnon, G., Chiesa, A., Fenzi, G., and Yogev, E., *WHIR: Reed–Solomon Proximity Testing
+      with Super-Fast Verification*][ACFY24]
+-/
+
+namespace ProximityGap
+
+open NNReal Finset Function
+open scoped ProbabilityTheory
+open scoped BigOperators LinearCode
+open Code Affine ReedSolomon
+open Domain
+open CosetFftDomain CosetFftDomainClass
+open MvPolynomial LinearMvExtension
+
+variable {F : Type} [Field F] [DecidableEq F]
+variable {n : ℕ}
+variable {domain : SmoothCosetFftDomain n F} {f : Word F (Fin (2 ^ n))}
+variable {k : ℕ} {x : F}
+
+/-- One step of lemma 4.15 from [ACFY24]. -/
+lemma foldWord_eq_evalOnPoints_powAlgHom [NeZero n] {α : F}
+  {g : F⦃≤ 1⦄[X (Fin n)]}
+  (hf : f = evalOnPoints domain (powAlgHom g.1)) :
+  foldWord domain f 1 α =
+    evalOnPoints
+      (domain.subdomain 1)
+      (powAlgHom (g.1.aeval (fun i ↦
+          if h : i = 0 then C α else MvPolynomial.X (⟨i.val - 1, by omega⟩ : Fin (n - 1))))) := by
+  have hchar := CosetFftDomainClass.domain_implies_char_ne_2 domain
+  have h2ne0 : (2 : F) ≠ 0 := fun contra ↦ hchar <|
+    ringChar.of_eq (CharP.ringChar_of_prime_eq_zero Nat.prime_two contra)
+  subst hf
+  conv_lhs =>
+    rw [powAlgHom_eq_even_add_odd_powAlgHom hchar]
+  rw [even_and_odd_eval hchar, foldWord_k_1']
+  ext u
+  extract_lets x j j'
+  have : x.val ≠ 0 := fun contra ↦ by
+    have := x.2
+    simp_all
+  aesop
+    (add safe (by field_simp))
+    (add simp
+      [evalOnPoints,
+       subdomain_sqFoldMapGen_eq_pow_domain,
+       evalOnPoints_sq_eq_evalOnPoints_subdomain])
+    (add unsafe
+      [(by ring_nf),
+       (by rw [add_comm, mul_comm]),
+       sqFoldMapGen_eq_sqFoldMapGen_of_pow_apply_eq_pow_apply])
+
+private noncomputable def substFun (m : ℕ) (β : Fin m → F) (i : Fin n) :
+    MvPolynomial (Fin (n - m)) F :=
+  if h : i.val < m then MvPolynomial.C (β ⟨i.val, h⟩)
+  else MvPolynomial.X ⟨i.val - m, by omega⟩
+
+omit [DecidableEq F] in
+private lemma aeval_substFun_comp {k : ℕ} [NeZero (n - k)] (γ : Fin (k + 1) → F)
+    (g0 : MvPolynomial (Fin n) F) :
+    (MvPolynomial.aeval (substFun k (fun j ↦ γ ⟨j.val, by omega⟩)) g0).aeval
+        (fun i : Fin (n - k) ↦ if h : i = 0 then MvPolynomial.C (γ ⟨k, by omega⟩)
+          else MvPolynomial.X (⟨i.val - 1, by omega⟩ : Fin (n - k - 1)))
+      = MvPolynomial.aeval (substFun (k + 1) γ) g0 := by
+  rw [MvPolynomial.comp_aeval_apply]
+  refine congrArg (fun φ ↦ MvPolynomial.aeval φ g0) ?_
+  funext i
+  unfold substFun
+  by_cases h1 : i.val < k
+  · rw [dif_pos h1, dif_pos (show i.val < k + 1 by omega)]
+    simp
+  · rw [dif_neg h1]
+    by_cases h2 : i.val = k
+      <;> aesop (add safe (by grind))
+
+private lemma aeval_split_mem {n : ℕ} [NeZero n] {R : Type} [Field R]
+  (hchar : ¬CharP R 2)
+  (p : R⦃≤ 1⦄[X (Fin n)]) (α : R) :
+  p.1.aeval
+    (fun i ↦ if h : i = 0 then C α else (MvPolynomial.X ⟨i.val - 1, by omega⟩ : R[X (Fin (n - 1))]))
+      ∈ restrictDegree (Fin (n - 1)) R 1 := by
+  rw [even_and_odd_eval hchar]
+  exact Submodule.add_mem _ (even_pred p).2
+    (by rw [MvPolynomial.C_mul']; exact Submodule.smul_mem _ _ (odd_pred p).2)
+
+omit [DecidableEq F] in
+private lemma aeval_substFun_mem [NeZero n] {gg : F⦃≤ 1⦄[X (Fin n)]}
+  {domain : SmoothCosetFftDomain n F} :
+  ∀ (m : ℕ), m ≤ n → ∀ (β : Fin m → F),
+    MvPolynomial.aeval (substFun m β) gg.1 ∈ MvPolynomial.restrictDegree (Fin (n - m)) F 1 := by
+  intro m
+  induction m with
+  | zero =>
+    intro hm β
+    have : (substFun (n := n) 0 β) = MvPolynomial.X := by aesop
+    aesop
+  | succ m ih =>
+    intro hm β
+    haveI : NeZero (n - m) := ⟨by omega⟩
+    have hchar : ¬CharP F 2 := CosetFftDomainClass.domain_implies_char_ne_2 domain
+    have hq : MvPolynomial.aeval (substFun m (fun j ↦ β ⟨j.val, by omega⟩)) gg.1 ∈
+      MvPolynomial.restrictDegree (Fin (n - m)) F 1 := ih (by omega) _
+    have hmem :
+        (MvPolynomial.aeval (substFun m (fun j ↦ β ⟨j.val, by omega⟩)) gg.1).aeval
+          (fun i : Fin (n - m) ↦ if h : i = 0 then MvPolynomial.C (β ⟨m, by omega⟩)
+            else MvPolynomial.X (⟨i.val - 1, by omega⟩ : Fin (n - m - 1)))
+          ∈ MvPolynomial.restrictDegree (Fin (n - m - 1)) F 1 :=
+      aeval_split_mem hchar ⟨_, hq⟩ (β ⟨m, by omega⟩)
+    rw [←aeval_substFun_comp (k := m) β gg.1]
+    exact hmem
+
+/-- Lemma 4.15 from [ACFY24]. Provides a way to
+  compute the corresponding multilinear extension
+  for the interated folding of codewords. -/
+theorem iteratedFoldWord_eq_evalOnPoints_powAlgHom [NeZero n] {α : Fin k → F}
+  {g : F⦃≤ 1⦄[X (Fin n)]}
+  (hk : k ≤ n)
+  (hf : f = evalOnPoints domain (powAlgHom g.1)) :
+  iteratedFoldWord domain f k α =
+      evalOnPoints
+        (domain.subdomain k)
+        (powAlgHom (g.1.aeval (fun i ↦
+          if h : i.val < k then C (α ⟨i.val, h⟩) else MvPolynomial.X
+            (⟨i.val - k, by omega⟩ : Fin (n - k))))) := by
+  suffices H : ∀ (k : ℕ), k ≤ n → ∀ (α : Fin k → F),
+      iteratedFoldWord domain f k α
+        = evalOnPoints (domain.subdomain k) (powAlgHom (g.1.aeval (substFun k α))) by
+    exact H k hk α
+  intro k
+  induction k with
+  | zero =>
+    intro _ α
+    have : (substFun (n := n) 0 α) = MvPolynomial.X := by aesop
+    aesop
+  | succ k ih =>
+    intro hk α
+    haveI : NeZero (n - k) := ⟨by omega⟩
+    have hprev : iteratedFoldWord domain f k (fun j ↦ α ⟨j.val, by omega⟩)
+        = evalOnPoints (domain.subdomain k)
+            (powAlgHom (g.1.aeval (substFun k (fun j ↦ α ⟨j.val, by omega⟩)))) :=
+      ih (by omega) _
+    have hmem := aeval_substFun_mem (domain := domain) (gg := g) k (by omega)
+      (fun j ↦ α ⟨j.val, by omega⟩)
+    have hfold := foldWord_eq_evalOnPoints_powAlgHom
+        (domain := domain.subdomain k)
+        (f := iteratedFoldWord domain f k (fun j ↦ α ⟨j.val, by omega⟩))
+        (α := α ⟨k, by omega⟩)
+        (g := ⟨g.1.aeval (substFun k (fun j ↦ α ⟨j.val, by omega⟩)), hmem⟩)
+        hprev
+    funext i
+    have hi2 : i.val < 2 ^ (n - k - 1) := by grind
+    change foldWord (domain.subdomain k)
+          (iteratedFoldWord domain f k (fun j ↦ α ⟨j.val, by omega⟩)) 1 (α ⟨k, by omega⟩)
+          ⟨i.val, hi2⟩ = _
+    rw [hfold]
+    simp only [evalOnPoints, Function.Embedding.coeFn_mk, LinearMap.coe_mk, AddHom.coe_mk]
+    rw [aeval_substFun_comp (k := k) (γ := α) g.1,
+        subdomain_one_comp (ω := domain) (by omega) ⟨i.val, hi2⟩ i rfl]
+
+end ProximityGap