Polynomial computations by reflection in Lean 4, using typeclass inference to build computable representations for expressions involving +,-,* etc.

PolynomialReflection.lean

import Mathlib.Data.Polynomial.Eval

open scoped Polynomial
open Polynomial (C X)

variable {R} [Semiring R] (p q : R[X])

class Polynomial.ComputableRepr where
  coeff : ℕ → R
  degreeBound : ℕ
  coeff_eq ⦃n : ℕ⦄ : p.coeff n = coeff n
  support_subset : p.support ⊆ Finset.range degreeBound

open Polynomial.ComputableRepr

instance : (0 : R[X]).ComputableRepr where
  coeff _ := 0
  degreeBound := 0
  coeff_eq := by simp
  support_subset := by simp

-- Can't make this work:
/- noncomputable instance (m : ℕ) [m.AtLeastTwo] : (OfNat.ofNat m : R[X]).ComputableRepr where
  coeff n := if n = 0 then m else 0
  degreeBound := 1
  coeff_eq := by sorry
  support_subset := sorry
example : (2 : ℤ[X]).ComputableRepr := inferInstance -- fails -/

instance (r : R) : (C r).ComputableRepr where
  coeff n := if n = 0 then r else 0
  degreeBound := 1
  coeff_eq := by simp [Polynomial.coeff_C]
  support_subset := by rw [← Polynomial.monomial_zero_left]; apply Polynomial.support_monomial'

instance : (X : R[X]).ComputableRepr where
  coeff n := if n = 1 then 1 else 0
  degreeBound := 2
  coeff_eq _ := by simp_rw [Polynomial.coeff_X, eq_comm]
  support_subset := by
    rw [← Polynomial.monomial_one_one_eq_X]
    exact (Polynomial.support_monomial' _ _).trans (by simp)

def Nat.withBotAdd : ℕ → ℕ → ℕ
  | 0, _ => 0
  | _, 0 => 0
  | m + 1, n + 1 => m + n + 1

@[simps] def WithBot.natSuccIso : WithBot ℕ ≃o ℕ where
  toFun m := m.recBotCoe 0 (· + 1)
  invFun m := m.rec ⊥ (↑)
  left_inv m := by cases m <;> rfl
  right_inv m := by cases m <;> rfl
  map_rel_iff' {m n} := by
    cases m; · simp [recBotCoe]
    cases n; · simpa [recBotCoe] using WithBot.bot_lt_coe _
    simp [recBotCoe]

def WithBot.natSucc (n : WithBot ℕ) : ℕ := WithBot.natSuccIso n

lemma Polynomial.natSucc_degree (hp : p ≠ 0) : p.degree.natSucc = p.natDegree + 1 := by
  rw [natDegree, ← WithBot.coe_unbot _ (degree_eq_bot.not.mpr hp)]; rfl

lemma WithBot.natSucc_withBotAdd (m n : WithBot ℕ) :
    m.natSucc.withBotAdd n.natSucc = (m + n).natSucc := by cases m <;> cases n <;> rfl

lemma WithBot.natSuccIso_symm_withBotAdd (m n : ℕ) :
    WithBot.natSuccIso.symm (m.withBotAdd n) =
    WithBot.natSuccIso.symm m + WithBot.natSuccIso.symm n := by cases m <;> cases n <;> rfl

variable [p.ComputableRepr]

lemma natSucc_degree_le_degreeBound : p.degree.natSucc ≤ degreeBound p := by
  obtain rfl | hp := eq_or_ne p 0; · exact bot_le
  have := Polynomial.leadingCoeff_ne_zero.mpr hp
  rw [Polynomial.leadingCoeff, ← Polynomial.mem_support_iff] at this
  rw [Polynomial.natSucc_degree p hp]
  exact Finset.mem_range.mp (support_subset this)

lemma Polynomial.support_subset_range : p.support ⊆ Finset.range p.degree.natSucc := fun n h ↦ by
  by_cases hp : degree p = ⊥
  · simp_all
  · rw [natSucc_degree p (degree_eq_bot.not.mp hp), Finset.mem_range_succ_iff]
    contrapose! h
    exact not_mem_support_iff.mpr (coeff_eq_zero_of_natDegree_lt h)

-- TODO: replace assumptions of `support_smul` and here by `[SMulZeroClass S R]`
instance {S} [Monoid S] [DistribMulAction S R] (s : S) : (s • p).ComputableRepr where
  coeff n := s • coeff p n
  degreeBound := degreeBound p
  coeff_eq _ := by rw [Polynomial.coeff_smul, coeff_eq]
  support_subset := (Polynomial.support_smul _ _).trans support_subset

namespace Polynomial.ComputableRepr

variable {p} in
theorem coeff_eq_zero_of_le {n : ℕ} (hn : degreeBound p ≤ n) : p.coeff n = 0 :=
  not_mem_support_iff.mp fun h ↦ hn.not_lt (Finset.mem_range.mp <| support_subset h)

@[reducible] def ofEq (h : p = q) : q.ComputableRepr where
  coeff := coeff p
  degreeBound := degreeBound p
  coeff_eq _ := by rw [← coeff_eq, h]
  support_subset := h ▸ support_subset

section degreeAux

def degreeAux (P : ℕ → Prop) [DecidablePred P] : ℕ → WithBot ℕ
  | 0 => ⊥
  | n + 1 => if P n then n else degreeAux P n

variable (P : ℕ → Prop) [DecidablePred P] (n : ℕ)

-- not used
lemma degreeAux_lt : degreeAux P n < n := by
  induction' n with n ih; · apply WithBot.bot_lt_coe
  rw [degreeAux]; split_ifs
  on_goal 2 => apply ih.trans
  all_goals exact WithBot.coe_lt_coe.mpr n.lt_succ_self

lemma degreeAux_of_forall_not (h : ∀ m < n, ¬ P m) : degreeAux P n = ⊥ := by
  induction' n with n ih; · rfl
  rw [degreeAux, if_neg (h n n.lt_succ_self)]; exact ih fun m hm ↦ h m (hm.trans n.lt_succ_self)

lemma degreeAux_spec_of_exists (h : ∃ m < n, P m) :
    ∃ m < n, P m ∧ (∀ k < n, m < k → ¬ P k) ∧ degreeAux P n = m := by
  induction' n with n ih
  · obtain ⟨m, hmn, -⟩ := h; cases hmn
  by_cases hn : P n
  · refine ⟨n, n.lt_succ_self, hn, fun k hkn hnk ↦ (hnk.not_le <| Nat.lt_succ.mp hkn).elim, ?_⟩
    rw [degreeAux, if_pos hn]
  obtain ⟨m, hmn, hm⟩ := h
  obtain rfl | hmn := (Nat.lt_succ.mp hmn).eq_or_lt
  · exact (hn hm).elim
  obtain ⟨m, hmn, hm, ih, _⟩ := ih ⟨m, hmn, hm⟩
  refine ⟨m, hmn.trans n.lt_succ_self, hm, fun k hkn hmk ↦ ?_, ?_⟩
  · obtain rfl | hkn := (Nat.lt_succ.mp hkn).eq_or_lt; exacts [hn, ih k hkn hmk]
  · rwa [degreeAux, if_neg hn]

end degreeAux

variable [DecidableEq R]

def degree : WithBot ℕ := degreeAux (coeff p · ≠ 0) (degreeBound p)
def natDegree : ℕ := (degree p).unbot' 0
def leadingCoeff : R := coeff p (natDegree p)
abbrev Monic : Prop := leadingCoeff p = 1

lemma degree_eq : p.degree = degree p := by
  obtain rfl | hp := eq_or_ne p 0
  · rw [degree, degreeAux_of_forall_not, degree_zero]
    rintro m -; rw [← coeff_eq]; exact not_not_intro (coeff_zero m)
  have := leadingCoeff_ne_zero.mpr hp
  obtain ⟨m, -, hpm, hmax, eq⟩ := degreeAux_spec_of_exists (coeff p · ≠ 0) (degreeBound p)
    ⟨p.natDegree, Finset.mem_range.mp (support_subset <| mem_support_iff.mpr this),
      by simp_rw [← coeff_eq]; exact this⟩
  rw [degree, eq]
  refine degree_eq_of_le_of_coeff_ne_zero ?_ (by rwa [coeff_eq])
  rw [degree_le_iff_coeff_zero]
  intro n hmn
  obtain hnd | hdn := lt_or_le n (degreeBound p)
  · rw [coeff_eq]; exact of_not_not (hmax n hnd <| WithBot.coe_lt_coe.mp hmn)
  · exact coeff_eq_zero_of_le hdn

lemma natDegree_eq : p.natDegree = natDegree p := by rw [natDegree, ← degree_eq]; rfl
lemma leadingCoeff_eq : p.leadingCoeff = leadingCoeff p := by
  rw [leadingCoeff, ← coeff_eq, ← natDegree_eq]; rfl
lemma monic_iff : p.Monic ↔ Monic p := by rw [Monic, ← leadingCoeff_eq]; rfl

instance : Decidable p.Monic := decidable_of_iff _ (monic_iff p).symm

end Polynomial.ComputableRepr

noncomputable instance : (1 : R[X]).ComputableRepr := .ofEq (C 1) _ (map_one C)
noncomputable instance (m : ℕ) : (m : R[X]).ComputableRepr := .ofEq _ _ (Polynomial.C_eq_nat_cast m)

variable [q.ComputableRepr]

instance : (p + q).ComputableRepr where
  coeff n := coeff p n + coeff q n
  degreeBound := max (degreeBound p) (degreeBound q)
  coeff_eq _ := by simp_rw [Polynomial.coeff_add, coeff_eq]
  support_subset := Polynomial.support_add.trans <|
    (Finset.union_subset_union support_subset support_subset).trans fun _ _ ↦ by simp_all

open scoped BigOperators

instance : (p * q).ComputableRepr where
  coeff n := ∑ i in Finset.range (min (degreeBound p) (n + 1)), coeff p i * coeff q (n - i)
  degreeBound := (degreeBound p).withBotAdd (degreeBound q)
  coeff_eq n := by
    simp_rw [← coeff_eq, Polynomial.coeff_mul, Finset.Nat.sum_antidiagonal_eq_sum_range_succ_mk]
    rw [← Finset.sum_range_add_sum_Ico _ (min_le_right _ _)]
    convert add_zero _
    refine Finset.sum_eq_zero fun n h ↦ ?_
    rw [Finset.mem_Ico, min_le_iff, ← not_lt, or_iff_left h.2.not_le, ← Finset.mem_range] at h
    rw [Polynomial.not_mem_support_iff.mp fun hn ↦ h.1 (support_subset hn), zero_mul]
  support_subset := (p * q).support_subset_range.trans <| Finset.range_mono <|
    (WithBot.natSuccIso.monotone <| p.degree_mul_le q).trans <| by
      rw [← WithBot.natSuccIso.symm.map_rel_iff, OrderIso.symm_apply_apply,
        WithBot.natSuccIso_symm_withBotAdd]
      apply add_le_add <;> rw [OrderIso.le_symm_apply] <;> apply natSucc_degree_le_degreeBound

-- TODO: exponentiation by repeated squaring
noncomputable instance (priority := mid) : (n : ℕ) → (p ^ n).ComputableRepr :=
  Nat.rec (ofEq 1 _ (pow_zero _).symm) fun n _ ↦ ofEq (p * p ^ n) _ (pow_succ _ _).symm

instance (m : ℕ) : (X ^ m : R[X]).ComputableRepr where
  coeff n := if n = m then 1 else 0
  degreeBound := m + 1
  coeff_eq _ := by simp_rw [Polynomial.coeff_X_pow]
  support_subset := by
    rw [← Polynomial.monomial_one_right_eq_X_pow]
    exact (Polynomial.support_monomial' _ _).trans (by simp)

noncomputable instance {S} [Semiring S] (f : R →+* S) : (p.map f).ComputableRepr where
  coeff n := f (coeff p n)
  degreeBound := degreeBound p
  coeff_eq _ := by rw [Polynomial.coeff_map, coeff_eq]
  support_subset := (Polynomial.support_map_subset _ _).trans support_subset

section

variable {S} [Ring S] (p q : S[X]) [p.ComputableRepr] [q.ComputableRepr]

instance : (-p).ComputableRepr where
  coeff n := - coeff p n
  degreeBound := degreeBound p
  coeff_eq _ := by rw [Polynomial.coeff_neg, coeff_eq]
  support_subset := by rw [Polynomial.support_neg]; exact support_subset

-- `coeff p n - coeff q n` more efficient?
noncomputable instance : (p - q).ComputableRepr := .ofEq (p + -q) _ (sub_eq_add_neg _ _)

end

theorem Polynomial.ComputableRepr.eq_iff : p = q ↔
    ∀ i ∈ Finset.range (max (degreeBound p) (degreeBound q)), coeff p i = coeff q i :=
  ⟨fun h _ _ ↦ by simp_rw [← coeff_eq, h], fun h ↦ ext fun n ↦ by
    simp only [Finset.mem_range, lt_max_iff] at h
    by_cases hn : ?_
    · simp_rw [coeff_eq]; exact h n hn
    · simp_rw [not_or, not_lt] at hn; rw [coeff_eq_zero_of_le hn.1, coeff_eq_zero_of_le hn.2]⟩

instance [DecidableEq R] : Decidable (p = q) := decidable_of_iff _ (eq_iff p q).symm

example : (X + X : ℤ[X]) = (2 * X : ℤ[X]) := by
  rw [← Nat.cast_eq_ofNat] -- have to do this
  decide

noncomputable def H : ℚ[X] := X ^ 3 - (C (6/8) * X) - C (1/8)
noncomputable def H' : ℤ[X] := C 8 * X ^ 3 - (C 6 * X) - C 1

example : DecidableEq ℚ := inferInstance -- okay
-- example : (6 / 8 : ℚ) * 8 = 6 := by decide -- doesn't work because operations on ℚ are irreducible
-- see https://leanprover.zulipchat.com/#narrow/stream/113489-new-members/topic/how.20to.20explicitly.20rfl.20for.20.E2.84.9A.3F/near/406961128

lemma H_H' : (C 8) * H = H'.map (algebraMap ℤ ℚ) := by rw [H, H']; apply of_decide_eq_true; rfl
-- `of_decide_eq_true <| by rw [H, H']; rfl` doesn't work. Maybe the rewritten Decidable instance is not good?

lemma H'_degree : H'.natDegree = 3 := by rw [H', natDegree_eq]; rfl
lemma H_degree : H.natDegree = 3 := by rw [H, natDegree_eq]; rfl
lemma H_monic : H.Monic := by rw [H, monic_iff]; rfl -- decide doesn't work (because it's in ℚ)
lemma H_ne_zero : H ≠ 0 := by rw [H]; apply of_decide_eq_true; rfl
lemma H'_coeff_zero : H'.coeff 0 = -1 := by rw [H', coeff_eq]; rfl
lemma H'_leading : H'.leadingCoeff = 8 := by rw [H', leadingCoeff_eq]; rfl

example :
    (X ^ 1000000000000 + 1 : ℚ[X]).degree = 1000000000000 := by rw [degree_eq]; rfl

set_option profiler true
example :
    (1 + X + X ^ 2 - X ^ 5 - X ^ 6 - 2 * X ^ 7 - X ^ 8 - X ^ 9 + X ^ 12 + X ^ 13 + X ^ 14 +
        X ^ 15 + X ^ 16 + X ^ 17 - X ^ 20 - X ^ 22 - X ^ 24 - X ^ 26 - X ^ 28 + X ^ 31 + X ^ 32 +
        X ^ 33 + X ^ 34 + X ^ 35 + X ^ 36 - X ^ 39 - X ^ 40 - 2 * X ^ 41 - X ^ 42 - X ^ 43 +
        X ^ 46 + X ^ 47 + X ^ 48 : ℤ[X]).degree = 48 := by
  rw [← Nat.cast_eq_ofNat (R := ℤ[X]) (n := 2), degree_eq]; rfl
/- typeclass inference of Polynomial.ComputableRepr took 159ms
   elaboration took 314ms -/

example :
    (1 + X + X ^ 2 - X ^ 5 - X ^ 6 - 2 * X ^ 7 - X ^ 8 - X ^ 9 + X ^ 12 + X ^ 13 + X ^ 14 +
        X ^ 15 + X ^ 16 + X ^ 17 - X ^ 20 - X ^ 22 - X ^ 24 - X ^ 26 - X ^ 28 + X ^ 31 + X ^ 32 +
        X ^ 33 + X ^ 34 + X ^ 35 + X ^ 36 - X ^ 39 - X ^ 40 - 2 * X ^ 41 - X ^ 42 - X ^ 43 +
        X ^ 46 + X ^ 47 + X ^ 48 : ℚ[X]).degree = 48 := by
  rw [← Nat.cast_eq_ofNat (R := ℚ[X]) (n := 2), degree_eq]; rfl
/- typeclass inference of Polynomial.ComputableRepr took 172ms
   elaboration took 352ms -/