fix: do not look for split-tactic candidates in proof terms

leodemoura · leodemoura · commit b06cc0d11a1f · 2024-07-01T12:09:14.000-07:00
We also skip candidates in implicit arguments and binders of lambda expressions

See new test.
diff --git a/src/Lean/Meta/Basic.lean b/src/Lean/Meta/Basic.lean
@@ -129,7 +129,7 @@ structure ParamInfo where
   hasFwdDeps     : Bool       := false
   /-- `backDeps` contains the backwards dependencies. That is, the (0-indexed) position of previous parameters that this one depends on. -/
   backDeps       : Array Nat  := #[]
-  /-- `isProp` is true if the parameter is always a proposition. -/
+  /-- `isProp` is true if the parameter type is always a proposition. -/
   isProp         : Bool       := false
   /--
     `isDecInst` is true if the parameter's type is of the form `Decidable ...`.
diff --git a/src/Lean/Meta/Tactic/Split.lean b/src/Lean/Meta/Tactic/Split.lean
@@ -17,7 +17,7 @@ def getSimpMatchContext : MetaM Simp.Context :=
    return {
       simpTheorems   := {}
       congrTheorems := (← getSimpCongrTheorems)
-      config        := { Simp.neutralConfig with dsimp := false, implicitDefEqProofs := true }
+      config        := { Simp.neutralConfig with dsimp := false }
    }
 
 def simpMatch (e : Expr) : MetaM Simp.Result := do
@@ -270,7 +270,7 @@ def mkDiscrGenErrorMsg (e : Expr) : MessageData :=
 def throwDiscrGenError (e : Expr) : MetaM α :=
   throwError (mkDiscrGenErrorMsg e)
 
-def splitMatch (mvarId : MVarId) (e : Expr) : MetaM (List MVarId) := do
+def splitMatch (mvarId : MVarId) (e : Expr) : MetaM (List MVarId) := mvarId.withContext do
   let some app ← matchMatcherApp? e | throwError "internal error in `split` tactic: match application expected{indentExpr e}\nthis error typically occurs when the `split` tactic internal functions have been used in a new meta-program"
   let matchEqns ← Match.getEquationsFor app.matcherName
   let mvarIds ← applyMatchSplitter mvarId app.matcherName app.matcherLevels app.params app.discrs
@@ -279,43 +279,14 @@ def splitMatch (mvarId : MVarId) (e : Expr) : MetaM (List MVarId) := do
     return (i+1, mvarId::mvarIds)
   return mvarIds.reverse
 
-/-- Return an `if-then-else` or `match-expr` to split. -/
-partial def findSplit? (env : Environment) (e : Expr) (splitIte := true) (exceptionSet : ExprSet := {}) : Option Expr :=
-  go e
-where
-  go (e : Expr) : Option Expr :=
-    if let some target := e.find? isCandidate then
-      if e.isIte || e.isDIte then
-        let cond := target.getArg! 1 5
-        -- Try to find a nested `if` in `cond`
-        go cond |>.getD target
-      else
-        some target
-    else
-      none
-
-  isCandidate (e : Expr) : Bool := Id.run do
-    if exceptionSet.contains e then
-      false
-    else if splitIte && (e.isIte || e.isDIte) then
-      !(e.getArg! 1 5).hasLooseBVars
-    else if let some info := isMatcherAppCore? env e then
-      let args := e.getAppArgs
-      for i in [info.getFirstDiscrPos : info.getFirstDiscrPos + info.numDiscrs] do
-        if args[i]!.hasLooseBVars then
-          return false
-      return true
-    else
-      false
-
 end Split
 
 open Split
 
-partial def splitTarget? (mvarId : MVarId) (splitIte := true) : MetaM (Option (List MVarId)) := commitWhenSome? do
+partial def splitTarget? (mvarId : MVarId) (splitIte := true) : MetaM (Option (List MVarId)) := commitWhenSome? do mvarId.withContext do
   let target ← instantiateMVars (← mvarId.getType)
   let rec go (badCases : ExprSet) : MetaM (Option (List MVarId)) := do
-    if let some e := findSplit? (← getEnv) target splitIte badCases then
+    if let some e ← findSplit? target (if splitIte then .both else .match) badCases then
       if e.isIte || e.isDIte then
         return (← splitIfTarget? mvarId).map fun (s₁, s₂) => [s₁.mvarId, s₂.mvarId]
       else
@@ -334,7 +305,7 @@ partial def splitTarget? (mvarId : MVarId) (splitIte := true) : MetaM (Option (L
 
 def splitLocalDecl? (mvarId : MVarId) (fvarId : FVarId) : MetaM (Option (List MVarId)) := commitWhenSome? do
   mvarId.withContext do
-    if let some e := findSplit? (← getEnv) (← instantiateMVars (← inferType (mkFVar fvarId))) then
+    if let some e ← findSplit? (← instantiateMVars (← inferType (mkFVar fvarId))) then
       if e.isIte || e.isDIte then
         return (← splitIfLocalDecl? mvarId fvarId).map fun (mvarId₁, mvarId₂) => [mvarId₁, mvarId₂]
       else
diff --git a/src/Lean/Meta/Tactic/SplitIf.lean b/src/Lean/Meta/Tactic/SplitIf.lean
@@ -9,6 +9,110 @@ import Lean.Meta.Tactic.Cases
 import Lean.Meta.Tactic.Simp.Main
 
 namespace Lean.Meta
+
+inductive SplitKind where
+  | ite | match | both
+
+def SplitKind.considerIte : SplitKind → Bool
+  | .ite | .both => true
+  | _ => false
+
+def SplitKind.considerMatch : SplitKind → Bool
+  | .match | .both => true
+  | _ => false
+
+namespace FindSplitImpl
+
+structure Context where
+  exceptionSet : ExprSet := {}
+  kind : SplitKind := .both
+
+unsafe abbrev FindM := ReaderT Context $ StateT (PtrSet Expr) MetaM
+
+private def isCandidate (env : Environment) (ctx : Context) (e : Expr) : Bool := Id.run do
+  if ctx.exceptionSet.contains e then
+    return false
+  if ctx.kind.considerIte && (e.isIte || e.isDIte) then
+    return !(e.getArg! 1 5).hasLooseBVars
+  if ctx.kind.considerMatch then
+    if let some info := isMatcherAppCore? env e then
+      let args := e.getAppArgs
+      for i in [info.getFirstDiscrPos : info.getFirstDiscrPos + info.numDiscrs] do
+        if args[i]!.hasLooseBVars then
+          return false
+      return true
+  return false
+
+@[inline] unsafe def checkVisited (e : Expr) : OptionT FindM Unit := do
+  if (← get).contains e then
+    failure
+  modify fun s => s.insert e
+
+unsafe def visit (e : Expr) : OptionT FindM Expr := do
+  checkVisited e
+  if isCandidate (← getEnv) (← read) e then
+    return e
+  else
+    -- We do not look for split candidates in proofs.
+    unless e.hasLooseBVars do
+      if (← isProof e) then
+        failure
+    match e with
+    | .lam _ _ b _ | .proj _ _ b -- We do not look for split candidates in the binder of lambdas.
+    | .mdata _ b       => visit b
+    | .forallE _ d b _ => visit d <|> visit b -- We want to look for candidates at `A → B`
+    | .letE _ _ v b _  => visit v <|> visit b
+    | .app ..          => visitApp? e
+    | _                => failure
+where
+  visitApp? (e : Expr) : FindM (Option Expr) :=
+    e.withApp fun f args => do
+    let info ← getFunInfo f
+    for u : i in [0:args.size] do
+      let arg := args[i]
+      if h : i < info.paramInfo.size then
+        let info := info.paramInfo[i]
+        unless info.isProp do
+          if info.isExplicit then
+            let some found ← visit arg | pure ()
+            return found
+      else
+        let some found ← visit arg | pure ()
+        return found
+    visit f
+
+end FindSplitImpl
+
+/-- Return an `if-then-else` or `match-expr` to split. -/
+partial def findSplit? (e : Expr) (kind : SplitKind := .both) (exceptionSet : ExprSet := {}) : MetaM (Option Expr) := do
+  go (← instantiateMVars e)
+where
+  go (e : Expr) : MetaM (Option Expr) := do
+    if let some target ← find? e then
+      if e.isIte || e.isDIte then
+        let cond := target.getArg! 1 5
+        -- Try to find a nested `if` in `cond`
+        return (← go cond).getD target
+      else
+        return some target
+    else
+      return none
+
+  find? (e : Expr) : MetaM (Option Expr) := do
+    let some candidate ← unsafe FindSplitImpl.visit e { kind, exceptionSet } |>.run' mkPtrSet
+      | return none
+    trace[split.debug] "candidate:{indentExpr candidate}"
+    return some candidate
+
+/-- Return the condition and decidable instance of an `if` expression to case split. -/
+private partial def findIfToSplit? (e : Expr) : MetaM (Option (Expr × Expr)) := do
+  if let some iteApp ← findSplit? e .ite then
+    let cond := iteApp.getArg! 1 5
+    let dec := iteApp.getArg! 2 5
+    return (cond, dec)
+  else
+    return none
+
 namespace SplitIf
 
 builtin_initialize ext : LazyInitExtension MetaM Simp.Context ←
@@ -21,7 +125,7 @@ builtin_initialize ext : LazyInitExtension MetaM Simp.Context ←
     return {
       simpTheorems  := #[s]
       congrTheorems := (← getSimpCongrTheorems)
-      config        := { Simp.neutralConfig with dsimp := false, implicitDefEqProofs := true }
+      config        := { Simp.neutralConfig with dsimp := false }
     }
 
 /--
@@ -68,19 +172,9 @@ private def discharge? (numIndices : Nat) (useDecide : Bool) : Simp.Discharge :=
 def mkDischarge? (useDecide := false) : MetaM Simp.Discharge :=
   return discharge? (← getLCtx).numIndices useDecide
 
-/-- Return the condition and decidable instance of an `if` expression to case split. -/
-private partial def findIfToSplit? (e : Expr) : Option (Expr × Expr) :=
-  if let some iteApp := e.find? fun e => (e.isIte || e.isDIte) && !(e.getArg! 1 5).hasLooseBVars then
-    let cond := iteApp.getArg! 1 5
-    let dec := iteApp.getArg! 2 5
-    -- Try to find a nested `if` in `cond`
-    findIfToSplit? cond |>.getD (cond, dec)
-  else
-    none
-
-def splitIfAt? (mvarId : MVarId) (e : Expr) (hName? : Option Name) : MetaM (Option (ByCasesSubgoal × ByCasesSubgoal)) := do
+def splitIfAt? (mvarId : MVarId) (e : Expr) (hName? : Option Name) : MetaM (Option (ByCasesSubgoal × ByCasesSubgoal)) := mvarId.withContext do
   let e ← instantiateMVars e
-  if let some (cond, decInst) := findIfToSplit? e then
+  if let some (cond, decInst) ← findIfToSplit? e then
     let hName ← match hName? with
       | none       => mkFreshUserName `h
       | some hName => pure hName
diff --git a/tests/lean/run/splitIssue2.lean b/tests/lean/run/splitIssue2.lean
@@ -0,0 +1,71 @@
+namespace Batteries
+
+/-- Union-find node type -/
+structure UFNode where
+  /-- Parent of node -/
+  parent : Nat
+
+namespace UnionFind
+
+/-- Parent of a union-find node, defaults to self when the node is a root -/
+def parentD (arr : Array UFNode) (i : Nat) : Nat :=
+  if h : i < arr.size then (arr.get ⟨i, h⟩).parent else i
+
+/-- Rank of a union-find node, defaults to 0 when the node is a root -/
+def rankD (arr : Array UFNode) (i : Nat) : Nat := 0
+
+theorem parentD_of_not_lt : ¬i < arr.size → parentD arr i = i := sorry
+
+theorem parentD_set {arr : Array UFNode} {x v i} :
+    parentD (arr.set x v) i = if x.1 = i then v.parent else parentD arr i := by
+  rw [parentD]
+  sorry
+
+end UnionFind
+
+open UnionFind
+
+structure UnionFind where
+  arr : Array UFNode
+
+namespace UnionFind
+
+/-- Size of union-find structure. -/
+@[inline] abbrev size (self : UnionFind) := self.arr.size
+
+/-- Parent of union-find node -/
+abbrev parent (self : UnionFind) (i : Nat) : Nat := parentD self.arr i
+
+theorem parent_lt (self : UnionFind) (i : Nat) : self.parent i < self.size ↔ i < self.size :=
+  sorry
+
+/-- Rank of union-find node -/
+abbrev rank (self : UnionFind) (i : Nat) : Nat := rankD self.arr i
+
+/-- Maximum rank of nodes in a union-find structure -/
+noncomputable def rankMax (self : UnionFind) := 0
+
+/-- Root of a union-find node. -/
+def root (self : UnionFind) (x : Fin self.size) : Fin self.size :=
+  let y := (self.arr.get x).parent
+  if h : y = x then
+    x
+  else
+    have : self.rankMax - self.rank (self.arr.get x).parent < self.rankMax - self.rank x :=
+      sorry
+    self.root ⟨y, sorry⟩
+termination_by self.rankMax - self.rank x
+
+/-- Root of a union-find node. Returns input if index is out of bounds. -/
+def rootD (self : UnionFind) (x : Nat) : Nat :=
+  if h : x < self.size then self.root ⟨x, h⟩ else x
+
+theorem rootD_parent (self : UnionFind) (x : Nat) : self.rootD (self.parent x) = self.rootD x := by
+  simp only [rootD, Array.data_length, parent_lt]
+  split
+  · simp only [parentD, ↓reduceDIte, *]
+    conv => rhs; rw [root]
+    split
+    · rw [root, dif_pos] <;> simp_all
+    · simp
+  · simp only [not_false_eq_true, parentD_of_not_lt, *]
