refactor: bv_decide's type analysis to prepare for enum support (#6971)

This PR does some refactoring on bv_decide's type analysis in
preparation for enum support in #6946.

src/Lean/Elab/Tactic/BVDecide/Frontend/Normalize.lean

@@ -13,6 +13,7 @@ import Lean.Elab.Tactic.BVDecide.Frontend.Normalize.EmbeddedConstraint
 import Lean.Elab.Tactic.BVDecide.Frontend.Normalize.AC
 import Lean.Elab.Tactic.BVDecide.Frontend.Normalize.Structures
 import Lean.Elab.Tactic.BVDecide.Frontend.Normalize.IntToBitVec
+import Lean.Elab.Tactic.BVDecide.Frontend.Normalize.TypeAnalysis
 
 /-!
 This module contains the implementation of `bv_normalize`, the preprocessing tactic for `bv_decide`.
@@ -51,6 +52,9 @@ where
     trace[Meta.Tactic.bv] m!"Running preprocessing pipeline on:\n{g}"
     let cfg ← PreProcessM.getConfig
 
+    if cfg.structures || cfg.enums then
+      g := (← typeAnalysisPass.run g).get!
+
     if cfg.structures then
       let some g' ← structuresPass.run g | return none
       g := g'

src/Lean/Elab/Tactic/BVDecide/Frontend/Normalize/Basic.lean

@@ -16,17 +16,39 @@ namespace Frontend.Normalize
 
 open Lean.Meta
 
+/--
+Contains the result of the type analysis to be used in the structures and enums pass.
+-/
+structure TypeAnalysis where
+  /--
+  Structures that are interesting for the structures pass.
+  -/
+  interestingStructures : Std.HashSet Name := {}
+  /--
+  Inductives enums that are interesting for the enums pass.
+  -/
+  interestingEnums : Std.HashSet Name := {}
+  /--
+  Other types that we've seen that are not interesting, currently only used as a cache.
+  -/
+  uninteresting : Std.HashSet Name := {}
+
 structure PreProcessState where
   /--
   Contains `FVarId` that we already know are in `bv_normalize` simp normal form and thus don't
   need to be processed again when we visit the next time.
   -/
   rewriteCache : Std.HashSet FVarId := {}
+  /--
+  Analysis results for the structure and enum pass if required.
+  -/
+  typeAnalysis : TypeAnalysis := {}
 
 abbrev PreProcessM : Type → Type := ReaderT BVDecideConfig <| StateRefT PreProcessState MetaM
 
 namespace PreProcessM
 
+@[inline]
 def getConfig : PreProcessM BVDecideConfig := read
 
 @[inline]
@@ -37,6 +59,38 @@ def checkRewritten (fvar : FVarId) : PreProcessM Bool := do
 def rewriteFinished (fvar : FVarId) : PreProcessM Unit := do
   modify (fun s => { s with rewriteCache := s.rewriteCache.insert fvar })
 
+@[inline]
+def getTypeAnalysis : PreProcessM TypeAnalysis := do
+  return (← get).typeAnalysis
+
+@[inline]
+def lookupInterestingStructure (n : Name) : PreProcessM (Option Bool) := do
+  let s ← getTypeAnalysis
+  if s.uninteresting.contains n then
+    return some false
+  else if s.interestingStructures.contains n then
+    return some true
+  else
+    return none
+
+@[inline]
+def modifyTypeAnalysis (f : TypeAnalysis → TypeAnalysis) :
+    PreProcessM Unit := do
+  modify fun s => { s with typeAnalysis := f s.typeAnalysis }
+
+@[inline]
+def markInterestingStructure (n : Name) : PreProcessM Unit := do
+  modifyTypeAnalysis (fun s => { s with interestingStructures := s.interestingStructures.insert n })
+
+@[inline]
+def markInterestingEnum (n : Name) : PreProcessM Unit := do
+  modifyTypeAnalysis (fun s => { s with interestingEnums := s.interestingEnums.insert n })
+
+@[inline]
+def markUninterestingType (n : Name) : PreProcessM Unit := do
+  modifyTypeAnalysis (fun s => { s with uninteresting := s.uninteresting.insert n })
+
+@[inline]
 def run (cfg : BVDecideConfig) (goal : MVarId) (x : PreProcessM α) : MetaM α := do
   let hyps ← goal.withContext do getPropHyps
   ReaderT.run x cfg |>.run' { rewriteCache := Std.HashSet.empty hyps.size }
@@ -53,6 +107,7 @@ structure Pass where
 
 namespace Pass
 
+@[inline]
 def run (pass : Pass) (goal : MVarId) : PreProcessM (Option MVarId) := do
   withTraceNode `bv (fun _ => return m!"Running pass: {pass.name} on\n{goal}") do
     pass.run' goal

src/Lean/Elab/Tactic/BVDecide/Frontend/Normalize/Structures.lean

@@ -27,43 +27,11 @@ namespace Frontend.Normalize
 
 open Lean.Meta
 
-/--
-Contains a cache for interesting and uninteresting types such that we don't duplicate work in the
-structures pass.
--/
-structure InterestingStructures where
-  interesting : Std.HashSet Name := {}
-  uninteresting : Std.HashSet Name := {}
-
-private abbrev M := StateRefT InterestingStructures MetaM
-
-namespace M
-
-@[inline]
-def lookup (n : Name) : M (Option Bool) := do
-  let s ← get
-  if s.uninteresting.contains n then
-    return some false
-  else if s.interesting.contains n then
-    return some true
-  else
-    return none
-
-@[inline]
-def markInteresting (n : Name) : M Unit := do
-  modify (fun s => {s with interesting := s.interesting.insert n })
-
-@[inline]
-def markUninteresting (n : Name) : M Unit := do
-  modify (fun s => {s with uninteresting := s.uninteresting.insert n })
-
-end M
-
 partial def structuresPass : Pass where
   name := `structures
   run' goal := do
-    let (_, { interesting, .. }) ← checkContext goal |>.run {}
-
+    let interesting := (← PreProcessM.getTypeAnalysis).interestingStructures
+    if interesting.isEmpty then return goal
     let goals ← goal.casesRec fun decl => do
       if decl.isLet || decl.isImplementationDetail then
         return false
@@ -77,6 +45,7 @@ where
   postprocess (goal : MVarId) (interesting : Std.HashSet Name) : PreProcessM (Option MVarId) := do
     goal.withContext do
       let mut relevantLemmas : SimpTheoremsArray := #[]
+      relevantLemmas ← relevantLemmas.addTheorem (.decl ``ne_eq) (← mkConstWithLevelParams ``ne_eq)
       for const in interesting do
         let constInfo ← getConstInfoInduct const
         let ctorName := (← getConstInfoCtor constInfo.ctors.head!).name
@@ -94,50 +63,5 @@ where
       let some (_, newGoal) := result? | return none
       return newGoal
 
-  checkContext (goal : MVarId) : M Unit := do
-    goal.withContext do
-      for decl in ← getLCtx do
-        if !decl.isLet && !decl.isImplementationDetail then
-          discard <| typeInteresting decl.type
-
-  constInterestingCached (n : Name) : M Bool := do
-    if let some cached ← M.lookup n then
-      return cached
-
-    let interesting ← constInteresting n
-    if interesting then
-      M.markInteresting n
-      return true
-    else
-      M.markUninteresting n
-      return false
-
-  constInteresting (n : Name) : M Bool := do
-    let env ← getEnv
-    if !isStructure env n then
-      return false
-    let constInfo ← getConstInfoInduct n
-    if constInfo.isRec then
-      return false
-
-    let ctorTyp := (← getConstInfoCtor constInfo.ctors.head!).type
-    let analyzer state arg := do
-      return state || (← typeInteresting (← arg.fvarId!.getType))
-    forallTelescope ctorTyp fun args _ => args.foldlM (init := false) analyzer
-
-  typeInteresting (expr : Expr) : M Bool := do
-    match_expr expr with
-    | BitVec n => return (← getNatValue? n).isSome
-    | UInt8 => return true
-    | UInt16 => return true
-    | UInt32 => return true
-    | UInt64 => return true
-    | USize => return true
-    | Bool => return true
-    | _ =>
-      let some const := expr.getAppFn.constName? | return false
-      constInterestingCached const
-
-
 end Frontend.Normalize
 end Lean.Elab.Tactic.BVDecide

src/Lean/Elab/Tactic/BVDecide/Frontend/Normalize/TypeAnalysis.lean

@@ -0,0 +1,73 @@
+/-
+Copyright (c) 2025 Lean FRO, LLC. All rights reserved.
+Released under Apache 2.0 license as described in the file LICENSE.
+Authors: Henrik Böving
+-/
+prelude
+import Lean.Elab.Tactic.BVDecide.Frontend.Normalize.Basic
+
+/-!
+This file implements the type analysis pass for the structures and enum inductives pass. It figures
+out which types that occur either directly or transitively (e.g. through being contained in a
+structure) qualify for further treatment by the structures or enum pass.
+-/
+
+namespace Lean.Elab.Tactic.BVDecide
+namespace Frontend.Normalize
+
+open Lean.Meta
+
+partial def typeAnalysisPass : Pass where
+  name := `typeAnalysis
+  run' goal := do
+    checkContext goal
+    return goal
+where
+  checkContext (goal : MVarId) : PreProcessM Unit := do
+    goal.withContext do
+      for decl in ← getLCtx do
+        if !decl.isLet && !decl.isImplementationDetail then
+          discard <| typeInteresting decl.type
+
+  constInteresting (n : Name) : PreProcessM Bool := do
+    let analysis ← PreProcessM.getTypeAnalysis
+    if analysis.interestingStructures.contains n || analysis.interestingEnums.contains n then
+      return true
+    else if analysis.uninteresting.contains n then
+      return false
+
+    let env ← getEnv
+    if isStructure env n then
+      let constInfo ← getConstInfoInduct n
+      if constInfo.isRec then
+        PreProcessM.markUninterestingType n
+        return false
+
+      let ctorTyp := (← getConstInfoCtor constInfo.ctors.head!).type
+      let analyzer state arg := do return state || (← typeInteresting (← arg.fvarId!.getType))
+      let interesting ← forallTelescope ctorTyp fun args _ => args.foldlM (init := false) analyzer
+      if interesting then
+        PreProcessM.markInterestingStructure n
+      return interesting
+    else if ← isEnumType n then
+      PreProcessM.markInterestingEnum n
+      return true
+    else
+      PreProcessM.markUninterestingType n
+      return false
+
+  typeInteresting (expr : Expr) : PreProcessM Bool := do
+    match_expr expr with
+    | BitVec n => return (← getNatValue? n).isSome
+    | UInt8 => return true
+    | UInt16 => return true
+    | UInt32 => return true
+    | UInt64 => return true
+    | USize => return true
+    | Bool => return true
+    | _ =>
+      let some const := expr.getAppFn.constName? | return false
+      constInteresting const
+
+end Frontend.Normalize
+end Lean.Elab.Tactic.BVDecide

src/Std/Tactic/BVDecide/Syntax.lean

@@ -49,6 +49,10 @@ structure BVDecideConfig where
   -/
   fixedInt : Bool := true
   /--
+  Handle equality on enum inductives by turning them into `BitVec`.
+  -/
+  enums : Bool := true
+  /--
   Output the AIG of bv_decide as graphviz into a file called aig.gv in the working directory of the
   Lean process.
   -/

tests/lean/run/bv_structures.lean

@@ -109,3 +109,31 @@ example (a b : Pair) (h : a = b) : a.x = b.x ∧ a.z = b.z := by
   bv_decide
 
 end Ex8
+
+
+namespace Ex9
+
+inductive Enum where
+  | a
+  | b
+
+structure Pair where
+  x : BitVec 16
+  e : Enum
+
+-- We don't accidentally split up enums even though they are considered interesting
+example (a b c : Pair) (h1 : a.x < b.x) (h2 : b.x ≤ c.x) : a.x < c.x := by
+  bv_decide
+
+end Ex9
+
+namespace Ex10
+
+structure Pair where
+  x : BitVec 16
+  y : BitVec 16
+
+example (a b c : Pair) (h1 : a ≠ b) (h2 : b = c) : a ≠ c := by
+  bv_decide
+
+end Ex10