find relevant enums in type analysis

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

@@ -17,11 +17,20 @@ 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.
+Contains the result of the type analysis to be used in the structures and enums pass.
 -/
-structure InterestingStructures where
-  interesting : Std.HashSet Name := {}
+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
@@ -31,9 +40,9 @@ structure PreProcessState where
   -/
   rewriteCache : Std.HashSet FVarId := {}
   /--
-  Analysis results for the structure pass if required.
+  Analysis results for the structure and enum pass if required.
   -/
-  structuresAnalysis : InterestingStructures := {}
+  typeAnalysis : TypeAnalysis := {}
 
 abbrev PreProcessM : Type → Type := ReaderT BVDecideConfig <| StateRefT PreProcessState MetaM
 
@@ -51,33 +60,37 @@ def rewriteFinished (fvar : FVarId) : PreProcessM Unit := do
   modify (fun s => { s with rewriteCache := s.rewriteCache.insert fvar })
 
 @[inline]
-def getStructureAnalysis : PreProcessM InterestingStructures := do
-  return (← get).structuresAnalysis
+def getTypeAnalysis : PreProcessM TypeAnalysis := do
+  return (← get).typeAnalysis
 
 @[inline]
 def lookupInterestingStructure (n : Name) : PreProcessM (Option Bool) := do
-  let s ← getStructureAnalysis
+  let s ← getTypeAnalysis
   if s.uninteresting.contains n then
     return some false
-  else if s.interesting.contains n then
+  else if s.interestingStructures.contains n then
     return some true
   else
     return none
 
 @[inline]
-def modifyStructureAnalysis (f : InterestingStructures → InterestingStructures) :
+def modifyTypeAnalysis (f : TypeAnalysis → TypeAnalysis) :
     PreProcessM Unit := do
-  modify fun s => { s with structuresAnalysis := f s.structuresAnalysis }
+  modify fun s => { s with typeAnalysis := f s.typeAnalysis }
 
 @[inline]
 def markInterestingStructure (n : Name) : PreProcessM Unit := do
-  modifyStructureAnalysis (fun s => { s with interesting := s.interesting.insert n })
+  modifyTypeAnalysis (fun s => { s with interestingStructures := s.interestingStructures.insert n })
 
 @[inline]
-def markUninterestingStructure (n : Name) : PreProcessM Unit := do
-  modifyStructureAnalysis (fun s => { s with uninteresting := s.uninteresting.insert n })
+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 }
@@ -94,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

@@ -30,7 +30,7 @@ open Lean.Meta
 partial def structuresPass : Pass where
   name := `structures
   run' goal := do
-    let interesting := (← PreProcessM.getStructureAnalysis).interesting
+    let interesting := (← PreProcessM.getTypeAnalysis).interestingStructures
     let goals ← goal.casesRec fun decl => do
       if decl.isLet || decl.isImplementationDetail then
         return false

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

@@ -29,30 +29,32 @@ where
         if !decl.isLet && !decl.isImplementationDetail then
           discard <| typeInteresting decl.type
 
-  constInterestingCached (n : Name) : PreProcessM Bool := do
-    if let some cached ← PreProcessM.lookupInterestingStructure n then
-      return cached
-
-    let interesting ← constInteresting n
-    if interesting then
-      PreProcessM.markInterestingStructure n
+  constInteresting (n : Name) : PreProcessM Bool := do
+    let analysis ← PreProcessM.getTypeAnalysis
+    if analysis.interestingStructures.contains n || analysis.interestingEnums.contains n then
       return true
-    else
-      PreProcessM.markUninterestingStructure n
+    else if analysis.uninteresting.contains n then
       return false
 
-  constInteresting (n : Name) : PreProcessM Bool := do
     let env ← getEnv
-    if !isStructure env n then
-      return false
-    let constInfo ← getConstInfoInduct n
-    if constInfo.isRec then
-      return false
+    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))
-    forallTelescope ctorTyp fun args _ => args.foldlM (init := false) analyzer
+      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
@@ -65,7 +67,7 @@ where
     | Bool => return true
     | _ =>
       let some const := expr.getAppFn.constName? | return false
-      constInterestingCached const
+      constInteresting const
 
 end Frontend.Normalize
 end Lean.Elab.Tactic.BVDecide

tests/lean/run/bv_enums.lean

@@ -0,0 +1,18 @@
+import Std.Tactic.BVDecide
+
+
+namespace Ex1
+
+inductive State where
+  | s1
+  | s2
+
+structure Pair where
+  x : BitVec 16
+  s : State
+
+example (a b c : Pair) (h1 : a = b) (h2 : b.x < c.x) (h3 : b.s = c.s) : a.s = c.s ∧ a.x < c.x := by
+  bv_normalize
+  sorry
+
+end Ex1

tests/lean/run/bv_structures.lean

@@ -109,3 +109,20 @@ 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