Proof of algorithm 6 from https://kannwischer.eu/thesis/phd-thesis-2023-01-03.pdf

Montgomery.lean

import Lean
import Mathlib.Data.ZMod.Algebra
import Mathlib.Data.Polynomial.Basic
import Mathlib.Data.Int.ModEq
import Mathlib.Algebra.Order.WithZero
import Std.Data.Int.DivMod
-- import Mathlib.Tactic.Linarith

/- Q1: This theorem is in DivMod, but currently commented out. Is it not ported yet?  -/
theorem div_lt_of_lt_mul {a b c : Int} (H : 0 < c) (H' : a < b * c) : a / c < b :=
  sorry

/- This seems a lot of work for something that ought to be somewhere in the libraries already.
   Q2: Is there a better way? If not, does this deserve inclusion into mathlib? -/
theorem mod_inv_cancels_out {a b : Nat} (H: NeZero b) (H_c: Nat.Coprime a b)
  : ↑(ZMod.inv b ↑a) * ↑a ≡ 1 [ZMOD ↑b] := by
  have eq a n b: (a % n = b % n) = (a ≡ b [ZMOD n]) := by rfl
  have h: ↑(ZMod.inv b ↑a) * ↑a ≡ ↑(ZMod.inv b ↑a) % b * (↑a % b) [ZMOD b] := by
    rw [ ← eq ]
    rw [ Int.mul_emod ]
  apply (Int.ModEq.trans h)
  rw [ ← @ZMod.val_int_cast (↑b) (↑a) ]
  rw [ ← ZMod.val_int_cast (↑ (ZMod.inv b ↑a)) ]
  rw [ ← eq ]
  have h (x y z: Nat): (↑ x: Int) * (↑ y: Int) % (↑ z: Int) = ↑ (x * y % z) := by
    apply Int.ofNat_mod_ofNat
  rw [ h, ← ZMod.val_mul ]
  simp
  have h_r: ↑a * (↑a)⁻¹ = ZMod.inv b (↑ a) * ↑ a := by
    ring_nf
    trivial
  rw [ ← h_r, ZMod.coe_mul_inv_eq_one (↑ a) ]
  rw [ ← ZMod.val_int_cast ]
  simp
  assumption

/- Montgomery Reduction, Algorithm 6 from https://kannwischer.eu/thesis/phd-thesis-2023-01-03.pdf -/
def mont_reduction
  (q: Nat)
  (R: Nat)
  (minus_q_minus_1: Int)
  (a: Nat)
  (h_R: R > q ∧ exists n, R = 2 ^ n)
  (h_q_minus_1: minus_q_minus_1 = - ZMod.inv R q)
  (h_q: 0 < q)
  (h_a: a < q * R)
  (h_q_R: Nat.Coprime q R):
  { t: Int // t % q = (a * ZMod.inv q R) % q ∧ 0 ≤ t ∧ t < 2 * q } 
:=
  -- the algorithm itself, carefully crafted to minimize field operations and
  -- rather operate on regular ints, since this is mostly an algorithm (and not
  -- a math object)
  let f := (a * minus_q_minus_1) % R
  let t := (a + f * q) / R
  -- Modulo facts are spread across ModEq.lean (nice notation, usually in
  -- Int.ModEq) and DivMod.lean (normal notation, usually in Int).
  --
  -- Q3: I ended up doing a lot of switching back and forth, and Int.ModEq.eq
  -- goes only one way. Is there a canonical way of moving in and out of the ≡
  -- notation? 
  have eq a n b: (a % n = b % n) = (a ≡ b [ZMOD n]) := by rfl
  -- main goal
  have h_t: t % q = a * ZMod.inv q R % q := by
    -- step 1: multiply by R both sides
    apply Int.ModEq.eq
    have h_gcd: (↑ q: Int) / Int.gcd q R = q := by
      rw [ ‹ Int.gcd q R = 1 › ] 
      simp
    rw [ ← h_gcd ]
    apply (Int.ModEq.cancel_right_div_gcd (by simp; assumption))

    -- step 2: multiplication cancels out, now
    have h: ↑(ZMod.inv q ↑R) * ↑R ≡ 1 [ZMOD ↑q] := by
      apply mod_inv_cancels_out
      apply NeZero.of_pos h_q
      apply Nat.Coprime.symm
      assumption
    have h' := Int.ModEq.mul_left a h
    ring_nf at h'
    symm
    apply (Int.ModEq.trans h')
    simp

    have h_q: ↑(ZMod.inv R ↑q) * ↑q ≡ 1 [ZMOD ↑R] := by
      apply mod_inv_cancels_out
      have h_R_pos : 0 < R := by linarith
      apply NeZero.of_pos h_R_pos
      assumption

    -- key divisibility hypothesis
    have h_div: a + f * q ≡ 0 [ZMOD R] := by
      simp
      rw [ ← eq ]

      have h_simp: minus_q_minus_1 * ↑q % R = -1 % R := by
        rw [h_q_minus_1]
        -- Lots of legwork to put this into the form expected by Int.ModEq.mul_right
        -- Q4: Any more efficient way of doing this?
        have h_tmp : -1 = 1 * -1 := by simp 
        have h_tmp2 :  -↑(ZMod.inv R ↑q) * ↑q = (↑(ZMod.inv R ↑q) * ↑q) * (-1:Int) := by
          ring_nf
        rw [eq]
        rw [h_tmp, h_tmp2]
        apply (Int.ModEq.mul_right)
        apply h_q

      rw [Int.add_emod]
      have h: ((a * minus_q_minus_1 % R) * ↑q) % ↑R = - a % ↑R := by
        rw [Int.mul_emod]
        simp
        rw [← Int.mul_emod]
        rw [Int.mul_assoc]
        rw [Int.mul_emod]
        rw [h_simp]
        rw [←Int.mul_emod]
        simp
      rw [h]
      rw [← Int.add_emod]
      simp

    rw [ Int.modEq_iff_add_fac ] at h_div
    simp at h_div
    obtain ⟨ x ⟩ := h_div
    have h: a + (a * minus_q_minus_1 % R) * ↑q = - ↑R * x := by linarith
    rw [ h ]

    -- apply divisibility to prove the goal
    have h_mod: (-↑R * x) % ↑R = 0 := by
      have h: -↑R*x = ↑R * (-x) := by ring
      rw [ h ]
      apply Int.mul_emod_right
    rw [ Int.ediv_mul_cancel_of_emod_eq_zero ]
    rw [ ←h ] 
    symm
    have h: a + (a * minus_q_minus_1 % R) * ↑q = a + ↑q * (a * minus_q_minus_1 % R) := by ring
    rw [ h ]
    simp [ Int.modEq_add_fac, Int.ModEq.refl ]
    assumption

  -- secondary goal
  have h_t': 0 ≤ t ∧ t < 2 * q := by
    apply And.intro
    case left =>
      simp
      apply Int.ediv_nonneg
      case Ha =>
        have h' : 0 ≤ ↑a * minus_q_minus_1 % ↑R * ↑q := by
          apply Int.mul_nonneg
          case ha =>
            apply Int.emod_nonneg
            linarith
          case hb => linarith
        apply Int.le_trans h'
        linarith
      case Hb => linarith
    case right =>
      simp
      have h': (↑a + ↑a * minus_q_minus_1 % ↑R * ↑q) < 2 * R * q := by
        apply Int.lt_of_lt_of_le 
        apply @Int.add_lt_add a (q * R) (↑a * minus_q_minus_1 % ↑R * ↑q) (q * R)
        linarith
        apply @Int.lt_of_lt_of_le _ (R * q)
        apply mul_lt_mul_of_pos_right
        apply Int.emod_lt_of_pos
        linarith
        linarith
        linarith
        linarith
      apply div_lt_of_lt_mul
      linarith
      linarith
  ⟨ t, ⟨ h_t, h_t' ⟩ ⟩