perf: in match splitters, thunk alts if needed (#11239)

This PR adds a `Unit` assumption to alternatives of the splitter that
would otherwise not have arguments. This fixes #11211.

In practice these argument-less alternatives did not cause wrong
behavior, as the motive when used with `split` is always a function
type. But it is better to be safe here (maybe someone uses splitters in
other ways), it may increase the effectiveness of #10184 and simplifies
#11220.

The perf impact is insignificant in the grand scheme of things on
stdlib, but the change is effective:
```
~/lean4 $ build/release/stage1/bin/lean tests/lean/run/matchSplitStats.lean 
969 splitters found
455 splitters are const defs
~/lean4 $ build/release/stage2/bin/lean tests/lean/run/matchSplitStats.lean 
969 splitters found
829 splitters are const defs
```

Author: nomeata

src/Lean/Meta/Match/MatchEqs.lean

@@ -279,8 +279,7 @@ private abbrev ConvertM := ReaderT (FVarIdMap (Expr × Nat × Array Bool)) $ Sta
   - `altNews` are the new free variables which contains additional hypotheses that ensure they are only used
      when the previous overlapping alternatives are not applicable. -/
 private partial def mkSplitterProof (matchDeclName : Name) (template : Expr) (alts altsNew : Array Expr)
-    (altsNewNumParams : Array Nat)
-    (altArgMasks : Array (Array Bool)) : MetaM Expr := do
+    (altsNewNumParams : Array Nat) (altArgMasks : Array (Array Bool)) (numDiscrEqs : Nat) : MetaM Expr := do
   trace[Meta.Match.matchEqs] "proof template: {template}"
   let map := mkMap
   let (proof, mvarIds) ← convertTemplate template |>.run map |>.run #[]
@@ -454,7 +453,10 @@ where
       else
         let Expr.fvar fvarId .. := e.getAppFn | return .continue
         let some (altNew, numParams, argMask) := (← read).get? fvarId | return .continue
-        trace[Meta.Match.matchEqs] ">> argMask: {argMask}, e: {e}, {altNew}"
+        trace[Meta.Match.matchEqs] ">> argMask: {argMask}, numParams: {numParams}, e: {e}, alsNew: {altNew}, "
+        if numParams + numDiscrEqs = 0 then
+          let eNew := mkApp altNew (mkConst ``Unit.unit)
+          return TransformStep.done eNew
         let mut newArgs := #[]
         let argMask := trimFalseTrail argMask
         unless e.getAppNumArgs ≥ argMask.size do
@@ -463,7 +465,7 @@ where
           if includeArg then
             newArgs := newArgs.push arg
         let eNew := mkAppN altNew newArgs
-        /- Recall that `numParams` does not include the equalities associated with discriminants of the form `h : discr`. -/
+        /- Recall that `numParams` does not include the `numDiscrEqs` equalities associated with discriminants of the form `h : discr`. -/
         let (mvars, _, _) ← forallMetaBoundedTelescope (← inferType eNew) (numParams - newArgs.size) (kind := MetavarKind.syntheticOpaque)
         modify fun s => s ++ (mvars.map (·.mvarId!))
         let eNew := mkAppN eNew mvars
@@ -543,7 +545,13 @@ where go baseName splitterName := withConfig (fun c => { c with etaStruct := .no
           if let some h ← simpH? h patterns.size then
             hs := hs.push h
         trace[Meta.Match.matchEqs] "hs: {hs}"
-        let splitterAltType ← mkForallFVars ys (← hs.foldrM (init := (← mkForallFVars eqs altResultType)) (mkArrow · ·))
+        let splitterAltType ← mkForallFVars eqs altResultType
+        let splitterAltType ← mkArrowN hs splitterAltType
+        let splitterAltType ← mkForallFVars ys splitterAltType
+        let splitterAltType ← if splitterAltType == altResultType then
+          mkArrow (mkConst ``Unit) splitterAltType
+        else
+          pure splitterAltType
         let splitterAltType ← unfoldNamedPattern splitterAltType
         let splitterAltNumParam := hs.size + ys.size
         -- Create a proposition for representing terms that do not match `patterns`
@@ -580,14 +588,14 @@ where go baseName splitterName := withConfig (fun c => { c with etaStruct := .no
       let splitterParams := params.toArray ++ #[motive] ++ discrs.toArray ++ altsNew
       let splitterType ← mkForallFVars splitterParams matchResultType
       trace[Meta.Match.matchEqs] "splitterType: {splitterType}"
-      let template := mkAppN (mkConst constInfo.name us) (params ++ #[motive] ++ discrs ++ alts)
-      let template ← deltaExpand template (· == constInfo.name)
-      let template := template.headBeta
       let splitterVal ←
         if (← isDefEq splitterType constInfo.type) then
           pure <| mkConst constInfo.name us
         else
-          mkLambdaFVars splitterParams (← mkSplitterProof matchDeclName template alts altsNew splitterAltNumParams altArgMasks)
+          let template := mkAppN (mkConst constInfo.name us) (params ++ #[motive] ++ discrs ++ alts)
+          let template ← deltaExpand template (· == constInfo.name)
+          let template := template.headBeta
+          mkLambdaFVars splitterParams (← mkSplitterProof matchDeclName template alts altsNew splitterAltNumParams altArgMasks numDiscrEqs)
       addAndCompile <| Declaration.defnDecl {
         name        := splitterName
         levelParams := constInfo.levelParams

src/Lean/Meta/Match/MatcherApp/Transform.lean

@@ -311,16 +311,27 @@ def transform
         altType in altTypes do
       let alt' ← forallAltTelescope' origAltType (numParams - numDiscrEqs) 0 fun ys args => do
         let altType ← instantiateForall altType ys
+        -- Look past the thunking unit parameter, if present
+        let altType ← if splitterNumParams + numDiscrEqs = 0 then
+            instantiateForall altType #[mkConst ``Unit.unit]
+          else
+            pure altType
         -- The splitter inserts its extra parameters after the first ys.size parameters, before
         -- the parameters for the numDiscrEqs
-        forallBoundedTelescope altType (splitterNumParams - ys.size) fun ys2 altType => do
+        let alt' ← forallBoundedTelescope altType (splitterNumParams - ys.size) fun ys2 altType => do
           forallBoundedTelescope altType numDiscrEqs fun ys3 altType => do
             forallBoundedTelescope altType extraEqualities fun ys4 altType => do
               let altParams := args ++ ys3
               let alt ← try instantiateLambda alt altParams
                         catch _ => throwError "unexpected matcher application, insufficient number of parameters in alternative"
               let alt' ← onAlt altIdx altType altParams alt
               mkLambdaFVars (ys ++ ys2 ++ ys3 ++ ys4) alt'
+        let alt' ← if splitterNumParams + numDiscrEqs = 0 then
+          -- The splitter expects a thunked alternative, but we don't want the `x : Unit` to be in
+          -- the context (e.g. in functional induction), so use Function.const rather than a lambda
+          mkAppM ``Function.const #[mkConst ``Unit, alt']
+        else
+          pure alt'
       alts' := alts'.push alt'
 
     remaining' := remaining' ++ (← onRemaining matcherApp.remaining)

src/Lean/Meta/Tactic/Split.lean

@@ -285,9 +285,16 @@ def applyMatchSplitter (mvarId : MVarId) (matcherDeclName : Name) (us : Array Le
     let mvarIds ← mvarId.applyN splitter matchEqns.size
     let (_, mvarIds) ← mvarIds.foldlM (init := (0, [])) fun (i, mvarIds) mvarId => do
       let numParams := matchEqns.splitterAltNumParams[i]!
-      let (_, mvarId) ← mvarId.introN numParams
+      let mvarId ←
+        if numParams + info.getNumDiscrEqs = 0 then
+          trace[split.debug] "introducing unit param for alt {(i : Nat)}"
+          let (unitFvarId, mvarId) ← mvarId.intro1
+          mvarId.tryClear unitFvarId
+        else
+          let (_, mvarId) ← mvarId.introN numParams
+          pure mvarId
       trace[split.debug] "before unifyEqs\n{mvarId}"
-      match (← Cases.unifyEqs? (numEqs + info.getNumDiscrEqs) mvarId {}) with
+      match (← Cases.unifyEqs? (info.getNumDiscrEqs + numEqs) mvarId {}) with
       | none   => return (i+1, mvarIds) -- case was solved
       | some (mvarId, fvarSubst) =>
         trace[split.debug] "after unifyEqs\n{mvarId}"

tests/lean/run/casesOnSameCtor.lean

@@ -38,16 +38,18 @@ info: Vec.match_on_same_ctor.{u_1, u} {α : Type u}
 /--
 info: Vec.match_on_same_ctor.splitter.{u_1, u} {α : Type u}
   {motive : {a : Nat} → (t t_1 : Vec α a) → t.ctorIdx = t_1.ctorIdx → Sort u_1} {a✝ : Nat} (t t✝ : Vec α a✝)
-  (h : t.ctorIdx = t✝.ctorIdx) (h_1 : motive nil nil ⋯)
+  (h : t.ctorIdx = t✝.ctorIdx) (h_1 : Unit → motive nil nil ⋯)
   (h_2 : (a : α) → (n : Nat) → (a_1 : Vec α n) → (a' : α) → (a'_1 : Vec α n) → motive (cons a a_1) (cons a' a'_1) ⋯) :
   motive t t✝ h
 -/
 #guard_msgs in
 #check Vec.match_on_same_ctor.splitter
 
--- Since there is no overlap, the splitter is equal to the matcher
--- (I wonder if we should use this in general in MatchEq)
-example : @Vec.match_on_same_ctor = @Vec.match_on_same_ctor.splitter := by rfl
+-- After #11211 this is no longer true. Should we thunk the same-ctor-construction?
+
+-- -- Since there is no overlap, the splitter is equal to the matcher
+-- -- (I wonder if we should use this in general in MatchEq)
+-- example : @Vec.match_on_same_ctor = @Vec.match_on_same_ctor.splitter := by rfl
 
 /--
 info: Vec.match_on_same_ctor.eq_2.{u_1, u} {α : Type u}

tests/lean/run/issue11211.lean

@@ -0,0 +1,99 @@
+/-!
+Checks that splitters have `Unit →` thunks and that nothing is confused because of that.
+-/
+
+set_option linter.unusedVariables false
+
+
+-- set_option trace.Meta.Match.matchEqs true
+
+def f (xs : List Nat) : Nat :=
+  match xs with
+  | [] => 1
+  | _ => 2
+
+/--
+info: def f.match_1.{u_1} : (motive : List Nat → Sort u_1) →
+  (xs : List Nat) → (Unit → motive []) → ((x : List Nat) → motive x) → motive xs
+-/
+#guard_msgs in
+#print sig f.match_1
+
+/--
+info: private def f.match_1.splitter.{u_1} : (motive : List Nat → Sort u_1) →
+  (xs : List Nat) → (Unit → motive []) → ((x : List Nat) → (x = [] → False) → motive x) → motive xs
+-/
+#guard_msgs(pass trace, all) in
+#print sig f.match_1.splitter
+
+
+/--
+info: private theorem f.match_1.congr_eq_1.{u_1} : ∀ (motive : List Nat → Sort u_1) (xs : List Nat) (h_1 : Unit → motive [])
+  (h_2 : (x : List Nat) → motive x),
+  xs = [] →
+    (match xs with
+      | [] => h_1 ()
+      | x => h_2 x) ≍
+      h_1 ()
+-/
+#guard_msgs(pass trace, all) in
+#print sig f.match_1.congr_eq_1
+
+--  set_option trace.split.debug true
+
+theorem test1: f n ≤ 2 := by
+  unfold f
+  split <;> grind
+
+
+theorem test2 : f n ≤ 2 := by
+  unfold f
+  grind
+
+/--
+info: theorem f.fun_cases : ∀ (motive : List Nat → Prop),
+  motive [] → (∀ (xs : List Nat), (xs = [] → False) → motive xs) → ∀ (xs : List Nat), motive xs
+-/
+#guard_msgs(pass trace, all) in
+#print sig f.fun_cases
+
+def Option_map (f : α → β) : Option α → Option β
+  | some x => some (f x)
+  | none   => none
+
+/--
+info: def Option_map.match_1.{u_1, u_2} : {α : Type u_1} →
+  (motive : Option α → Sort u_2) → (x : Option α) → ((x : α) → motive (some x)) → (Unit → motive none) → motive x
+-/
+#guard_msgs in
+#print sig Option_map.match_1
+
+/--
+info: private def Option_map.match_1.splitter.{u_1, u_2} : {α : Type u_1} →
+  (motive : Option α → Sort u_2) → (x : Option α) → ((x : α) → motive (some x)) → (Unit → motive none) → motive x :=
+@Option_map.match_1
+-/
+#guard_msgs in
+#print Option_map.match_1.splitter
+
+/--
+info: theorem Option_map.fun_cases.{u_1} : ∀ {α : Type u_1} (motive : Option α → Prop),
+  (∀ (x : α), motive (some x)) → motive none → ∀ (x : Option α), motive x
+-/
+#guard_msgs(pass trace, all) in
+#print sig Option_map.fun_cases
+
+def List_map (f : α → β) (l : List α) : List β := match _ : l with
+  | x::xs => f x :: List_map f xs
+  | []   => []
+termination_by l
+
+def foo₁ (a : Nat) (ha : a = 37) :=
+    (match h : a with | 42 => 23 | n => n) = 37
+
+/--
+info: private def foo₁.match_1.splitter.{u_1} : (motive : Nat → Sort u_1) →
+  (a : Nat) → (a = 42 → motive 42) → ((n : Nat) → (n = 42 → False) → a = n → motive n) → motive a
+-/
+#guard_msgs in
+#print sig foo₁.match_1.splitter

tests/lean/run/matchSparse.lean

@@ -37,8 +37,9 @@ info: expensive.match_1.{u_1} (motive : Expr → Expr → Sort u_1) (x✝ x✝¹
 /--
 info: expensive.match_1.splitter.{u_1} (motive : Expr → Expr → Sort u_1) (x✝ x✝¹ : Expr)
   (h_1 :
-    motive (((sort zero.succ).app (sort zero.succ)).app (sort zero.succ))
-      (((sort zero.succ).app (sort zero.succ)).app (sort zero.succ)))
+    Unit →
+      motive (((sort zero.succ).app (sort zero.succ)).app (sort zero.succ))
+        (((sort zero.succ).app (sort zero.succ)).app (sort zero.succ)))
   (h_2 :
     (x x_1 : Expr) →
       (x = ((sort zero.succ).app (sort zero.succ)).app (sort zero.succ) →