Working with Qq

Stream: metaprogramming / tactics

[Heather Macbeth]

I wrote a norm_num extension in Lean 4. It's for teaching purposes so its scope is very limited, and paradoxically this makes it messier to write: I want it to disprove $a \equiv b \mod n$ only in the "obvious" case when $0 \le a < n$ and $0 \le b < n$ and $a\ne b$.

My code is basically copy-pasta and I feel like I have made rather heavy weather of it, with clunky handling of Qq. Does anyone with some experience have time to look at it and teach me some tricks?

import Mathlib.Data.Int.ModEq
import Mathlib.Tactic.NormNum
import Mathlib.Tactic.Linarith

open Lean hiding Rat mkRat
open Lean.Meta Qq Lean.Elab Term
open Lean.Parser.Tactic Mathlib.Meta.NormNum

namespace Mathlib.Meta.NormNum

theorem isInt_ModEq_true : {a b a' b' n : ℤ} → IsInt a a' → IsInt b b' → decide (a' = b') = true →
    Int.ModEq n a b
  | _, _, a', b', n, ⟨rfl⟩, ⟨rfl⟩, hab =>
    by
      dsimp
      replace hab := of_decide_eq_true hab
      rw [hab]
      apply Int.ModEq.refl

theorem isInt_ModEq_false : {a b a' b' n : ℤ} → IsInt a a' → IsInt b b' → decide (0 < n) = true →
    decide (a' < n) = true → decide (b' < n) = true → decide (0 ≤ a') = true →
    decide (0 ≤ b') = true → decide (a' ≠ b') = true → ¬ Int.ModEq n a b
  | _, _, a', b', n, ⟨rfl⟩, ⟨rfl⟩, hn, han, hbn, ha, hb, hab =>
    by
      change a' % n ≠ b' % n
      replace hn := of_decide_eq_true hn
      replace han := of_decide_eq_true han
      replace hbn := of_decide_eq_true hbn
      replace ha := of_decide_eq_true ha
      replace hb := of_decide_eq_true hb
      replace hab := of_decide_eq_true hab
      intro H
      apply hab
      obtain ⟨h₁, -, -⟩ := (Int.ediv_emod_unique hn).mp ⟨Int.ediv_eq_zero_of_lt ha han, H⟩
      have h₂ := Int.emod_add_ediv b' n
      have h₃ := Int.ediv_eq_zero_of_lt hb hbn
      rw [h₃] at h₂
      linarith

end Mathlib.Meta.NormNum

/-- The `norm_num` extension which identifies expressions of the form `a ≡ b [ZMOD n]`,
such that `norm_num` successfully recognises both `a` and `b` and they are small compared to `n`. -/
@[norm_num Int.ModEq _ _ _] def evalModEq : NormNumExt where eval (e : Q(Prop)) := do
  let .app (.app (.app f (n : Q(ℤ))) (a : Q(ℤ))) (b : Q(ℤ)) ← whnfR e | failure
  guard <|← withNewMCtxDepth <| isDefEq f q(Int.ModEq)
  let ra : Result a ← derive a
  let rb : Result b ← derive b
  let rn : Result q($n) ← derive n
  let i : Q(Ring ℤ) := q(Int.instRingInt)
  let ⟨za, na, pa⟩ ← ra.toInt
  let ⟨zb, nb, pb⟩ ← rb.toInt
  let ⟨zn, _, _⟩ ← rn.toInt i
  if za = zb then
    -- reduce `a ≡ b [ZMOD n]` to `true` if `a` and `b` reduce to the same integer
    let pab : Q(decide ($na = $nb) = true) := (q(Eq.refl true) : Expr)
    let r : Q(Int.ModEq $n $a $b) := q(isInt_ModEq_true $pa $pb $pab)
    return (.isTrue r : Result q(Int.ModEq $n $a $b))
  else
    -- reduce `a ≡ b [ZMOD n]` to `false` if `0 < n`, `a` reduces to `a'` with `0 ≤ a' < n`,
    -- and `b` reduces to `b'` with `0 ≤ b' < n`
    let pab : Q(decide ($na ≠ $nb) = true) := (q(Eq.refl true) : Expr)
    if zn = 0 then failure
    let pn : Q(decide (0 < $n) = true) := (q(Eq.refl true) : Expr)
    if zn ≤ za then failure
    let pan : Q(decide ($na < $n) = true) := (q(Eq.refl true) : Expr)
    if zn ≤ zb then failure
    let pbn : Q(decide ($nb < $n) = true) := (q(Eq.refl true) : Expr)
    if za < 0 then failure
    let pa0 : Q(decide (0 ≤ $na) = true) := (q(Eq.refl true) : Expr)
    if zb < 0 then failure
    let pb0 : Q(decide (0 ≤ $nb) = true) := (q(Eq.refl true) : Expr)
    let r : Q(¬Int.ModEq $n $a $b) := q(isInt_ModEq_false $pa $pb $pn $pan $pbn $pa0 $pb0 $pab)
    return (.isFalse r : Result q(¬Int.ModEq $n $a $b))