perf: sparse case splitting in match compilation

src/Init/Data/Int/Lemmas.lean

@@ -552,7 +552,7 @@ protected theorem mul_eq_zero {a b : Int} : a * b = 0 ↔ a = 0 ∨ b = 0 := by
   exact match a, b, h with
   | .ofNat 0, _, _ => by simp
   | _, .ofNat 0, _ => by simp
-  | .ofNat (a+1), .negSucc b, h => by cases h
+  | .ofNat (_+1), .negSucc _, h => by cases h
 
 protected theorem mul_ne_zero {a b : Int} (a0 : a ≠ 0) (b0 : b ≠ 0) : a * b ≠ 0 :=
   Or.rec a0 b0 ∘ Int.mul_eq_zero.mp

src/Lean/Compiler/LCNF/Util.lean

@@ -9,6 +9,7 @@ prelude
 public import Init.Data.FloatArray.Basic
 public import Lean.CoreM
 public import Lean.Util.Recognizers
+import Lean.Meta.Basic
 
 public section
 

src/Lean/Elab/Tactic/RCases.lean

@@ -323,7 +323,7 @@ partial def rcasesCore (g : MVarId) (fs : FVarSubst) (clears : Array FVarId) (e
           let (v, g) ← g.intro x
           let (varsOut, g) ← g.introNP vars.size
           let fs' := (vars.zip varsOut).foldl (init := fs) fun fs (v, w) => fs.insert v (mkFVar w)
-          pure ([(n, ps)], #[⟨⟨g, #[mkFVar v], fs'⟩, n⟩])
+          pure ([(n, ps)], #[{mvarId := g, fields := #[mkFVar v], subst := fs', ctorName := n }])
       | ConstantInfo.inductInfo info, _ => do
         let (altVarNames, r) ← processConstructors pat.ref info.numParams #[] info.ctors pat.asAlts
         (r, ·) <$> g.cases e.fvarId! altVarNames (useNatCasesAuxOn := true)

src/Lean/Meta/Match/Match.lean

@@ -17,15 +17,22 @@ public section
 
 namespace Lean.Meta.Match
 
+register_builtin_option backwards.match.sparseCases : Bool := {
+  defValue := true
+  descr := "if true (the default), generate and use sparse case constructs when splitting inductive
+    types. In some cases this will prevent Lean from noticing that a match statement is complete
+    because it performs less case-splitting for the unreachable case. In this case, give explicit
+    patterns to perform the deeper split with `by contradiction` as the right-hand side.
+     ,"
+}
+
 register_builtin_option backwards.match.rowMajor : Bool := {
   defValue := true
-  group := "bootstrap"
   descr := "If true (the default), match compilation will split the discrimnants based \
     on position of the first constructor pattern in the first alternative. If false, \
     it splits them from left to right, which can lead to unnecessary code bloat."
 }
 
-
 private def mkIncorrectNumberOfPatternsMsg [ToMessageData α]
     (discrepancyKind : String) (expected actual : Nat) (pats : List α) :=
   let patternsMsg := MessageData.joinSep (pats.map toMessageData) ", "
@@ -162,6 +169,12 @@ private def hasArrayLitPattern (p : Problem) : Bool :=
     | .arrayLit .. :: _ => true
     | _                 => false
 
+private def hasVarOrInaccessiblePattern (p : Problem) : Bool :=
+  p.alts.any fun alt => match alt.patterns with
+    | .inaccessible _ :: _ => true
+    | .var _ :: _          => true
+    | _                    => false
+
 private def isVariableTransition (p : Problem) : Bool :=
   p.alts.all fun alt => match alt.patterns with
     | .inaccessible _ :: _ => true
@@ -341,6 +354,35 @@ where
           msg := msg ++ m!"\n  {lhs} ≋ {rhs}"
         throwErrorAt alt.ref msg
 
+private abbrev isCtorIdxIneq? (e : Expr) : Option FVarId := do
+  if let some (_, lhs, _rhs) := e.ne? then
+    if
+      lhs.isApp &&
+      lhs.getAppFn.isConst &&
+      (`ctorIdx).isSuffixOf lhs.getAppFn.constName! && -- This should be an env extension maybe
+      lhs.appArg!.isFVar
+    then
+      return lhs.appArg!.fvarId!
+  none
+
+private partial def contradiction (mvarId : MVarId) : MetaM Bool := do
+  mvarId.withContext do
+    withTraceNode `Meta.Match.match (msg := (return m!"{exceptBoolEmoji ·} Match.contradiction")) do
+    trace[Meta.Match.match] m!"Match.contradiction:\n{mvarId}"
+    if (← mvarId.contradictionCore {}) then
+      trace[Meta.Match.match] "Contradiction found!"
+      return true
+    else
+      -- Try harder by splitting `ctorIdx x ≠ 23` assumptions
+      for localDecl in (← getLCtx) do
+        if let some fvarId := isCtorIdxIneq? localDecl.type then
+          trace[Meta.Match.match] "splitting ctorIdx assumption {localDecl.type}"
+          let subgoals ← mvarId.cases fvarId
+          return ← subgoals.allM (contradiction ·.mvarId)
+
+      mvarId.admit
+      return false
+
 /--
 Try to solve the problem by using the first alternative whose pending constraints can be resolved.
 -/
@@ -352,10 +394,10 @@ where
   go (alts : List Alt) : StateRefT State MetaM Unit := do
     match alts with
     | [] =>
+      let mvarId ← p.mvarId.exfalso
       /- TODO: allow users to configure which tactic is used to close leaves. -/
-      unless (← p.mvarId.contradictionCore {}) do
-        trace[Meta.Match.match] "missing alternative"
-        p.mvarId.admit
+      unless (← contradiction mvarId) do
+        trace[Meta.Match.match] "contradiction failed, missing alternative"
         modify fun s => { s with counterExamples := p.examples :: s.counterExamples }
     | alt :: _ =>
       solveCnstrs p.mvarId alt
@@ -516,13 +558,29 @@ private def throwCasesException (p : Problem) (ex : Exception) : MetaM α := do
               "- <constructor> = <constructor>, examples: List.cons x xs = List.cons y ys, and List.cons x xs = List.nil"
   | _ => throw ex
 
+private def collectCtors (p : Problem) : Array Name :=
+  p.alts.foldl (init := #[]) fun ctors alt =>
+    match alt.patterns with
+    | .ctor n _ _ _ :: _ => if ctors.contains n then ctors else ctors.push n
+    | _                  => ctors
+
 private def processConstructor (p : Problem) : MetaM (Array Problem) := do
   trace[Meta.Match.match] "constructor step"
   let x :: xs := p.vars | unreachable!
+  let interestingCtors? ←
+    -- We use a sparse case analysis only if there is at least one non-constructor pattern,
+    -- but not just because there are constructors missing (in that case we benefit from
+    -- the eager split in ruling out constructors by type or by a more explicit error message)
+    if backwards.match.sparseCases.get (← getOptions) && hasVarOrInaccessiblePattern p then
+      let ctors := collectCtors p
+      trace[Meta.Match.match] "using sparse cases: {ctors}"
+      pure (some ctors)
+    else
+      pure none
   let subgoals? ← commitWhenSome? do
      let subgoals ←
        try
-         p.mvarId.cases x.fvarId!
+         p.mvarId.cases x.fvarId! (interestingCtors? := interestingCtors?)
        catch ex =>
          if p.alts.isEmpty then
            /- If we have no alternatives and dependent pattern matching fails, then a "missing cases" error is better than a "stuck" error message. -/
@@ -545,28 +603,42 @@ private def processConstructor (p : Problem) : MetaM (Array Problem) := do
          return some subgoals
   let some subgoals := subgoals? | return #[{ p with vars := xs }]
   subgoals.mapM fun subgoal => subgoal.mvarId.withContext do
-    let subst    := subgoal.subst
-    let fields   := subgoal.fields.toList
-    let newVars  := fields ++ xs
-    let newVars  := newVars.map fun x => x.applyFVarSubst subst
-    let subex    := Example.ctor subgoal.ctorName <| fields.map fun field => match field with
-      | .fvar fvarId => Example.var fvarId
-      | _            => Example.underscore -- This case can happen due to dependent elimination
-    let examples := p.examples.map <| Example.replaceFVarId x.fvarId! subex
-    let examples := examples.map <| Example.applyFVarSubst subst
-    let newAlts  := p.alts.filter fun alt => match alt.patterns with
-      | .ctor n .. :: _       => n == subgoal.ctorName
-      | .var _ :: _           => true
-      | .inaccessible _ :: _  => true
-      | _                     => false
-    let newAlts  := newAlts.map fun alt => alt.applyFVarSubst subst
-    let newAlts ← newAlts.mapM fun alt => do
-      match alt.patterns with
-      | .ctor _ _ _ fields :: ps  => return { alt with patterns := fields ++ ps }
-      | .var _ :: _               => expandVarIntoCtor alt subgoal.ctorName
-      | .inaccessible _ :: _      => processInaccessibleAsCtor alt subgoal.ctorName
-      | _                         => unreachable!
-    return { mvarId := subgoal.mvarId, vars := newVars, alts := newAlts, examples := examples }
+    -- withTraceNode `Meta.Match.match (msg := (return m!"{exceptEmoji ·} case {subgoal.ctorName}")) do
+    if let some ctorName := subgoal.ctorName then
+      -- A normal constructor case
+      let subst    := subgoal.subst
+      let fields   := subgoal.fields.toList
+      let newVars  := fields ++ xs
+      let newVars  := newVars.map fun x => x.applyFVarSubst subst
+      let subex    := Example.ctor ctorName <| fields.map fun field => match field with
+        | .fvar fvarId => Example.var fvarId
+        | _            => Example.underscore -- This case can happen due to dependent elimination
+      let examples := p.examples.map <| Example.replaceFVarId x.fvarId! subex
+      let examples := examples.map <| Example.applyFVarSubst subst
+      let newAlts  := p.alts.filter fun alt => match alt.patterns with
+        | .ctor n .. :: _       => n == ctorName
+        | .var _ :: _           => true
+        | .inaccessible _ :: _  => true
+        | _                     => false
+      let newAlts  := newAlts.map fun alt => alt.applyFVarSubst subst
+      let newAlts ← newAlts.mapM fun alt => do
+        match alt.patterns with
+        | .ctor _ _ _ fields :: ps  => return { alt with patterns := fields ++ ps }
+        | .var _ :: _               => expandVarIntoCtor alt ctorName
+        | .inaccessible _ :: _      => processInaccessibleAsCtor alt ctorName
+        | _                         => unreachable!
+      return { mvarId := subgoal.mvarId, vars := newVars, alts := newAlts, examples := examples }
+    else
+      -- A catch-all case
+      let subst := subgoal.subst
+      trace[Meta.Match.match] "constructor catch-all case"
+      let examples := p.examples.map <| Example.applyFVarSubst subst
+      let newVars := p.vars.map fun x => x.applyFVarSubst subst
+      let newAlts := p.alts.filter fun alt => match alt.patterns with
+        | .ctor .. :: _ => false
+        | _             => true
+      let newAlts := newAlts.map fun alt => alt.applyFVarSubst subst
+      return { mvarId := subgoal.mvarId, alts := newAlts, vars := newVars, examples := examples }
 
 private def processNonVariable (p : Problem) : MetaM Problem := withGoalOf p do
   let x :: xs := p.vars | unreachable!
@@ -808,6 +880,15 @@ def processExFalso (p : Problem) : MetaM Problem := do
   let mvarId' ← p.mvarId.exfalso
   return { p with mvarId := mvarId' }
 
+private def tracedForM (xs : Array α) (process : α → StateRefT State MetaM Unit) : StateRefT State MetaM Unit :=
+  if xs.size > 1 then
+    for x in xs, i in [:xs.size] do
+      withTraceNode `Meta.Match.match (msg := (return m!"{exceptEmoji ·} subgoal {i+1}/{xs.size}")) do
+        process x
+  else
+    for x in xs do
+      process x
+
 private partial def process (p : Problem) : StateRefT State MetaM Unit := do
   traceState p
   if isDone p then
@@ -870,7 +951,7 @@ private partial def process (p : Problem) : StateRefT State MetaM Unit := do
 
   if (← isConstructorTransition p) then
     let ps ← processConstructor p
-    ps.forM process
+    tracedForM ps process
     return
 
   if isVariableTransition p then
@@ -881,12 +962,12 @@ private partial def process (p : Problem) : StateRefT State MetaM Unit := do
 
   if isValueTransition p then
     let ps ← processValue p
-    ps.forM process
+    tracedForM ps process
     return
 
   if isArrayLitTransition p then
     let ps ← processArrayLit p
-    ps.forM process
+    tracedForM ps process
     return
 
   if (← hasNatValPattern p) then

src/Lean/Meta/Tactic/Cases.lean

@@ -9,6 +9,8 @@ prelude
 public import Lean.Meta.Tactic.Induction
 public import Lean.Meta.Tactic.Acyclic
 public import Lean.Meta.Tactic.UnifyEq
+import Lean.Meta.Constructions.SparseCasesOn
+import Lean.Meta.Constructions.CtorIdx
 
 public section
 
@@ -149,7 +151,8 @@ def generalizeIndices (mvarId : MVarId) (fvarId : FVarId) : MetaM GeneralizeIndi
     generalizeIndices' mvarId fvarDecl.toExpr fvarDecl.userName
 
 structure CasesSubgoal extends InductionSubgoal where
-  ctorName : Name
+  /-- The constructor of this subgoal. Is `none` in the catch-all of a sparse case match -/
+  ctorName : Option Name
 
 namespace Cases
 
@@ -216,11 +219,13 @@ private def elimAuxIndices (s₁ : GeneralizeIndicesSubgoal) (s₂ : Array Cases
 private def toCasesSubgoals (s : Array InductionSubgoal) (ctorNames : Array Name) (majorFVarId : FVarId) (us : List Level) (params : Array Expr)
     : Array CasesSubgoal :=
   s.mapIdx fun i s =>
-    let ctorName := ctorNames[i]!
-    let ctorApp  := mkAppN (mkAppN (mkConst ctorName us) params) s.fields
-    let s        := { s with subst := s.subst.insert majorFVarId ctorApp }
-    { ctorName           := ctorName,
-      toInductionSubgoal := s }
+    if _ : i < ctorNames.size then
+      let ctorName := ctorNames[i]
+      let ctorApp  := mkAppN (mkAppN (mkConst ctorName us) params) s.fields
+      let subst := s.subst.erase majorFVarId |>.insert majorFVarId ctorApp
+      { s with ctorName := ctorName, subst}
+    else
+      { s with ctorName := none }
 
 partial def unifyEqs? (numEqs : Nat) (mvarId : MVarId) (subst : FVarSubst) (caseName? : Option Name := none): MetaM (Option (MVarId × FVarSubst)) := withIncRecDepth do
   if numEqs == 0 then
@@ -244,17 +249,33 @@ private def unifyCasesEqs (numEqs : Nat) (subgoals : Array CasesSubgoal) : MetaM
       }
 
 private def inductionCasesOn (mvarId : MVarId) (majorFVarId : FVarId) (givenNames : Array AltVarNames) (ctx : Context)
-    (useNatCasesAuxOn : Bool := false) : MetaM (Array CasesSubgoal) := mvarId.withContext do
+    (useNatCasesAuxOn : Bool := false) (interestingCtors? : Option (Array Name) := none) :
+    MetaM (Array CasesSubgoal) := mvarId.withContext do
   let majorType ← inferType (mkFVar majorFVarId)
   let (us, params) ← getInductiveUniverseAndParams majorType
-  let mut casesOn := mkCasesOnName ctx.inductiveVal.name
-  if useNatCasesAuxOn && ctx.inductiveVal.name == ``Nat && (← getEnv).contains ``Nat.casesAuxOn then
-    casesOn := ``Nat.casesAuxOn
+
+  if let some interestingCtors := interestingCtors? then
+    -- Avoid Init.Prelude complications
+    let hasNE := (← getEnv).contains ``Ne
+    -- We can only create a sparse casesOn if we have `ctorIdx` (in particular, if it is a type)
+    let hasCtorIdx := (← getEnv).contains (mkCtorIdxName ctx.inductiveVal.name)
+    if hasNE && hasCtorIdx && !interestingCtors.isEmpty &&
+      interestingCtors.size < ctx.inductiveVal.ctors.length then
+      let casesOn ← Lean.Meta.mkSparseCasesOn ctx.inductiveVal.name interestingCtors
+      let s ← mvarId.induction majorFVarId casesOn givenNames
+      return toCasesSubgoals s interestingCtors majorFVarId us params
+
+  let casesOn :=
+    if useNatCasesAuxOn && ctx.inductiveVal.name == ``Nat && (← getEnv).contains ``Nat.casesAuxOn then
+      ``Nat.casesAuxOn
+    else
+      mkCasesOnName ctx.inductiveVal.name
   let ctors   := ctx.inductiveVal.ctors.toArray
   let s ← mvarId.induction majorFVarId casesOn givenNames
   return toCasesSubgoals s ctors majorFVarId us params
 
-def cases (mvarId : MVarId) (majorFVarId : FVarId) (givenNames : Array AltVarNames := #[]) (useNatCasesAuxOn : Bool := false) : MetaM (Array CasesSubgoal) := do
+def cases (mvarId : MVarId) (majorFVarId : FVarId) (givenNames : Array AltVarNames := #[])
+    (useNatCasesAuxOn : Bool := false) (interestingCtors? : Option (Array Name) := none) : MetaM (Array CasesSubgoal) := do
   try
     mvarId.withContext do
       mvarId.checkNotAssigned `cases
@@ -268,10 +289,11 @@ def cases (mvarId : MVarId) (majorFVarId : FVarId) (givenNames : Array AltVarNam
         if (← hasIndepIndices ctx) then
           -- Simple case
           inductionCasesOn mvarId majorFVarId givenNames ctx (useNatCasesAuxOn := useNatCasesAuxOn)
+            (interestingCtors? := interestingCtors?)
         else
           let s₁ ← generalizeIndices mvarId majorFVarId
           trace[Meta.Tactic.cases] "after generalizeIndices\n{MessageData.ofGoal s₁.mvarId}"
-          let s₂ ← inductionCasesOn s₁.mvarId s₁.fvarId givenNames ctx
+          let s₂ ← inductionCasesOn s₁.mvarId s₁.fvarId givenNames ctx (interestingCtors? := interestingCtors?)
           let s₂ ← elimAuxIndices s₁ s₂
           unifyCasesEqs s₁.numEqs s₂
   catch ex =>
@@ -288,8 +310,9 @@ Apply `casesOn` using the free variable `majorFVarId` as the major premise (aka
   It enables using `Nat.casesAuxOn` instead of `Nat.casesOn`,
   which causes case splits on `n : Nat` to be represented as `0` and `n' + 1` rather than as `Nat.zero` and `Nat.succ n'`.
 -/
-def _root_.Lean.MVarId.cases (mvarId : MVarId) (majorFVarId : FVarId) (givenNames : Array AltVarNames := #[]) (useNatCasesAuxOn : Bool := false) : MetaM (Array CasesSubgoal) :=
-  Cases.cases mvarId majorFVarId givenNames (useNatCasesAuxOn := useNatCasesAuxOn)
+def _root_.Lean.MVarId.cases (mvarId : MVarId) (majorFVarId : FVarId) (givenNames : Array AltVarNames := #[]) (useNatCasesAuxOn : Bool := false)
+  (interestingCtors? : Option (Array Name) := none) : MetaM (Array CasesSubgoal) :=
+  Cases.cases mvarId majorFVarId givenNames (useNatCasesAuxOn := useNatCasesAuxOn) (interestingCtors? := interestingCtors?)
 
 /--
 Keep applying `cases` on any hypothesis that satisfies `p`.