perf: sparse case splitting in match compilation (#10823)

This PR lets the match compilation procedure use sparse case analysis
when the patterns only match on some but not all constructors of an
inductive type. This way, less code is produce. Before, code handling
each of the other cases was then optimized and commoned-up by later
compilation pipeline, but that is wasteful to do.

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.

At least temporarily, there is also the option to disable this behaviour
with
```
set_option backwards.match.sparseCases false
```

Author: nomeata

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`.