Linter for generalizing type class hypotheses

Stream: mathlib4

[Jovan Gerbscheid]

I made a simple linter that tries to find type class assumptions that can be generalized.

Code

import Mathlib

section
open Lean Meta

unsafe def extractAux (root : Bool) (excluded : FVarIdSet) (e : Expr) : StateRefT (Std.HashMap FVarId (Option Expr)) (StateT (PtrSet Expr) IO) Unit := do
  if (← getThe (PtrSet Expr)).contains e then
    return
  else
  -- if !root then
  modifyThe (PtrSet Expr) (·.insert e)
  for arg in e.getAppArgs do
    if let .fvar fvarId := arg then
      if !excluded.contains fvarId then
        if e.hasLooseBVars || root then
          modify (·.insert fvarId none)
        else
          modify (·.alter fvarId fun | none => e | some none => some none | some (some e') => if e == e' then e else some none)
  match e with
  | .app ..          => extractAux false excluded e.getAppFn; e.getAppArgs.forM (extractAux false excluded)
  | .forallE _ d b _
  | .lam _ d b _     => extractAux false excluded d; extractAux false excluded b
  | .letE _ t v b _  => extractAux false excluded t; extractAux false excluded v; extractAux false excluded b
  | .mdata _ expr    => extractAux false excluded expr
  | .proj _ _ struct =>
    if !root then
      if let .fvar fvarId := struct then
        modify (·.alter fvarId fun | none => e | some e' => if e == e' then e else some none)
    extractAux false excluded struct
  | _ => return


unsafe def extract (binderTypes : Array Expr) (type value : Expr) : StateRefT (Std.HashMap FVarId (Option Expr)) (StateT (PtrSet Expr) IO) Unit := do
  let excluded := binderTypes.foldl (init := {}) fun s binderType => collectFVars s binderType
  let excluded := excluded.fvarSet
  extractAux false excluded type
  extractAux true excluded value


def getTypeClass (e : Expr) : MetaM (Option Name) := do
  let .const n _ := e.getAppFn | return none
  if isClass (← getEnv) n then
    return n
  else
    return none


def excludedPairs : Std.HashSet (Name × Name) := .ofList [
  (``Monad, ``Alternative), (``Monad, ``Applicative),
  (``EquivBEq, ``PartialEquivBEq), (``EquivBEq, ``ReflBEq),
  (``IsEmpty, ``Unique), (``Subsingleton, ``Unique), (``Subsingleton, ``IsEmpty), (``Finite, ``Nonempty), (``Finite, ``Fintype), (``Finite, ``IsEmpty),
  (``Nontrivial, ``NeZero),
  (``Algebra.IsAlgebraic, ``Algebra.IsIntegral),
  (``SecondCountableTopology, ``SecondCountableTopologyEither),
  (``NoZeroDivisors, ``NoZeroSMulDivisors),

  (``Add, ``HAdd), (``Sub, ``HSub), (``Mul, ``HMul), (``Div, ``HDiv), (``Pow, ``HPow), (``SMul, ``HSMul), (``VAdd, ``HVAdd),
  (``Zero, ``OfNat), (``One, ``OfNat),

  (``NoMinOrder, ``NoMaxOrder), (``NoMaxOrder, ``NoMinOrder),
  (``IsPredArchimedean, ``IsSuccArchimedean), (``IsSuccArchimedean, ``IsPredArchimedean),
  (``IsStronglyCoatomic, ``IsStronglyAtomic), (``IsStronglyAtomic, ``IsStronglyCoatomic),
  (``IsCoatomic, ``IsAtomic), (``IsAtomic, ``IsCoatomic),
  (``IsUpperModularLattice, ``IsLowerModularLattice), (``IsLowerModularLattice, ``IsUpperModularLattice),
  (``OrderTop, ``OrderBot), (``OrderBot, ``OrderTop),
  (``BoundedGENhdsClass, ``BoundedLENhdsClass), (``BoundedLENhdsClass, ``BoundedGENhdsClass),
  (``ClosedIicTopology, ``ClosedIciTopology), (``ClosedIciTopology, ``ClosedIicTopology),
  (``ContinuousInf, ``ContinuousSup), (``ContinuousSup, ``ContinuousInf),
  (``Topology.IsLower, ``Topology.IsUpper), (``Topology.IsUpper, ``Topology.IsLower),
  (``Topology.IsLowerSet, ``Topology.IsUpperSet), (``Topology.IsUpperSet, ``Topology.IsLowerSet),
  (``InfConvergenceClass, ``SupConvergenceClass), (``SupConvergenceClass, ``InfConvergenceClass),
  (``NoTopOrder, ``NoBotOrder), (``NoBotOrder, ``NoTopOrder),

  (``IsLeftCancelMul, ``IsRightCancelMul), (``IsRightCancelMul, ``IsLeftCancelMul),

  (``CategoryTheory.RepresentablyCoflat, ``CategoryTheory.RepresentablyFlat), (``CategoryTheory.RepresentablyFlat, ``CategoryTheory.RepresentablyCoflat),
  (``CategoryTheory.Epi, ``CategoryTheory.Mono), (``CategoryTheory.Mono, ``CategoryTheory.Epi),
  (``CategoryTheory.IsCofiltered, ``CategoryTheory.IsFiltered), (``CategoryTheory.IsFiltered, ``CategoryTheory.IsCofiltered),
  (``CategoryTheory.IsCofilteredOrEmpty, ``CategoryTheory.IsFilteredOrEmpty), (``CategoryTheory.IsFilteredOrEmpty, ``CategoryTheory.IsCofilteredOrEmpty),
  (``CategoryTheory.Functor.Initial, ``CategoryTheory.Functor.Final), (``CategoryTheory.Functor.Final, ``CategoryTheory.Functor.Initial),
  (``CategoryTheory.MorphismProperty.HasRightCalculusOfFractions, ``CategoryTheory.MorphismProperty.HasLeftCalculusOfFractions), (``CategoryTheory.MorphismProperty.HasLeftCalculusOfFractions, ``CategoryTheory.MorphismProperty.HasRightCalculusOfFractions),
  (``ComplexShape.Embedding.IsTruncLE, ``ComplexShape.Embedding.IsTruncGE), (``ComplexShape.Embedding.IsTruncGE, ``ComplexShape.Embedding.IsTruncLE),
  (``CategoryTheory.Limits.HasFiniteColimits, ``CategoryTheory.Limits.HasFiniteLimits), (``CategoryTheory.Limits.HasFiniteLimits, ``CategoryTheory.Limits.HasFiniteColimits),
  (``CategoryTheory.Limits.HasFiniteCoproducts, ``CategoryTheory.Limits.HasFiniteProducts), (``CategoryTheory.Limits.HasFiniteProducts, ``CategoryTheory.Limits.HasFiniteCoproducts),
  (``CategoryTheory.Limits.HasColimitsOfSize, ``CategoryTheory.Limits.HasLimitsOfSize), (``CategoryTheory.Limits.HasLimitsOfSize, ``CategoryTheory.Limits.HasColimitsOfSize),
  (``CategoryTheory.Limits.HasFilteredColimitsOfSize, ``CategoryTheory.Limits.HasCofilteredLimitsOfSize), (``CategoryTheory.Limits.HasCofilteredLimitsOfSize, ``CategoryTheory.Limits.HasFilteredColimitsOfSize),
  (``CategoryTheory.EnoughProjectives, ``CategoryTheory.EnoughInjectives), (``CategoryTheory.EnoughInjectives, ``CategoryTheory.EnoughProjectives),
]
def excludedOld : List Name := [
  ``Nonempty, ``Countable
]
def excludedNew : List Name := [
  ``CovariantClass, ``ContravariantClass
]


unsafe def replace (e k r : Expr) : (StateT (PtrMap Expr Expr) IO) Expr :=
  let rec visit (e : Expr) : (StateT (PtrMap Expr Expr) IO) Expr := do
    if e == k then
      return r
    else
      if let some e' := (← get).find? e then
        return e'
      else
        match e with
        | .app f a         => return e.updateApp! (← visit f) (← visit a)
        | .mdata _ b       => return e.updateMData! (← visit b)
        | .proj _ _ b      => return e.updateProj! (← visit b)
        | .letE _ t v b _  => return e.updateLet! (← visit t) (← visit v) (← visit b)
        | .lam _ d b _     => return e.updateLambdaE! (← visit d) (← visit b)
        | .forallE _ d b _ => return e.updateForallE! (← visit d) (← visit b)
        | e                => return e
  visit e

def checkAbstraction (type value k kType : Expr) : MetaM Bool := do
  withLocalDeclD `_a kType fun fvar => do
    let type ← unsafe ((replace type k fvar).run' mkPtrMap)
    unless ← isTypeCorrect type do
      return false
    let value ← unsafe ((replace value k fvar).run' mkPtrMap)
    unless ← isTypeCorrect value do
      return false
    isDefEq type (← inferType value)

def generalizations (type value : Expr) : MetaM (Option MessageData) := do
  forallTelescope type fun fvars type => do
    let types ← fvars.mapM (inferType ·)
    let value := value.beta fvars
    let (_, map) ← unsafe ((extract types type value).run {}).run' mkPtrSet
    let mut msgs := #[]
    for (fvarId, term) in map do
      if let some term := term then
        if !term.isProj then
          let some instName := term.getAppFn.constName? | continue
          if !(← isInstance instName) then
            continue
        let newType ← inferType term
        let decl ← fvarId.getDecl
        let oldType := decl.type
        let some oldCls ← getTypeClass oldType | continue
        let some newCls ← getTypeClass newType | continue
        if oldCls != newCls && !excludedPairs.contains (oldCls, newCls) && !excludedOld.contains oldCls && !excludedNew.contains newCls then
          if ← checkAbstraction type value term newType then
            msgs := msgs.push (← AddMessageContext.addMessageContext
              m! "[{oldType}] can be generalized to [{newType}]")
    if msgs.isEmpty then return none else
      return MessageData.ofArray msgs


namespace Batteries.Tactic.Lint

/-- A linter for commutativity lemmas that are marked simp. -/
@[env_linter]
def myLinter : Linter where
  noErrorsFound := "No error found"
  errorsFound := "Error found"
  test := fun declName => do
    -- throwError declName
    unless ← isAutoDecl declName do
      let .thmInfo c ← getConstInfo declName | return none
      return ← generalizations c.type c.value
    return none

end Batteries.Tactic.Lint
end

#lint only myLinter in Mathlib

It works by asking whether a hypothesis inst always appears inside the same application or projection. If so, then replace this expression by a new free variable, and check if the result is still type correct. So far it only looks at hypotheses that don't appear in other hypotheses. It also won't necessarily give the most general generalization.

This is similar to #mathlib4 > Elab to generalize type classes for theorems , except this looks at the underlying expressions, instead of modifying the syntax.

Running it on Lean and Batteries yields a few theorems that can be generalized (lean#7611 and batteries#1171)

Running it on Mathlib gives an enormous ~2500 results, after filtering out most unwanted results. Here they are in a file:

Generalizable_hypotheses.md

The filtering is necessary because not everying that the linter finds is a reasonable generalization, such as replacing Nonempty _ by Inhabited _ or by NeZero 1, or replacing NoMaxOrder α by NoMinOrder αᵒᵈ

[Vlad Tsyrklevich]

Cool! I think an approach working on Exprs has a lot of advantages, speed being a very important one, and I had begun to regret my original Syntax approach for how slow it was. This seems to be strictly more general than my implementation, since it also covers type classes in variable blocks, and finds about ~3x as many as mine does. One thing I had found I needed was to add some filtering for what generalizations are performed since some generalizations are not meaningful (this was specifically to cover Algebra, Module, and IsTopological* classes, though there may be more.)

[Vlad Tsyrklevich]

It also won't necessarily give the most general generalization.

Could you give an example when this is the case?

[Vlad Tsyrklevich]

Based on just skimming through the list, some of the following look to me like generalizations that should be excluded:

#check @measurableSet_lineDifferentiableAt_uncurry /- [[SecondCountableTopology
   E] can be generalized to [SecondCountableTopologyEither E E]] -/
#check @Set.instNoZeroDivisors /- [[NoZeroDivisors α] can be generalized to [NoZeroSMulDivisors α α]] -/

Applying this didn't work:

-- Mathlib.Analysis.CStarAlgebra.ContinuousFunctionalCalculus.Isometric
#check @isometry_cfcHom /- [[IsometricContinuousFunctionalCalculus R A
   p] can be generalized to [ContinuousFunctionalCalculus R A p]] -/

It may be worth trying to apply all of these and see if something breaks to find bugs. In my experience, I found that you may want to exclude Mathlib/Tactic, Mathlib/Algebra/Field/IsField.lean, and Mathlib/Data/Int/Cast/Field.lean

[Vlad Tsyrklevich]

This also works on instances which is cool, though the volume of hard-to-tell-if-its-useful generalizations with the added difficulty of generalizing type class assumptions in variable blocks led me to not get very far when I tried to see if I could split any out from your list.

[Jovan Gerbscheid]

One example of not giving the most general generalization is where Monad m can be replaced by Functor m, but it suggests to replace it with Applicative m.

[Vlad Tsyrklevich]

It's not immediately clear to me why that would be the case, e.g. what is the reason why a less general class may be picked?

[Jovan Gerbscheid]

I could change the code to also look for the more general one, but currently it just looks for type classes that are "one step away"

[Jovan Gerbscheid]

The code doesn't actually use anything about type classes. It is more of a hypothesis generalizer. And at the end it checks that it generalized a type class to a type class.

[Jovan Gerbscheid]

Here's an updated version of the document, now with 2092 theorems.

Generalizable_hypotheses.md