refactor: use Match.AltParamInfo also for splitters (#11261)

This PR continues the homogenization between matchers and splitters,
following up on #11256. In particular it removes the ambiguity whether
`numParams` includes the `discrEqns` or not.

Author: nomeata

src/Lean/Meta/Match/MatchEqs.lean

@@ -85,19 +85,23 @@ def unfoldNamedPattern (e : Expr) : MetaM Expr := do
 
   This can be used to use the alternative of a match expression in its splitter.
 -/
-partial def forallAltVarsTelescope (altType : Expr) (altNumParams numDiscrEqs : Nat)
+partial def forallAltVarsTelescope (altType : Expr) (altInfo : AltParamInfo)
   (k : (patVars : Array Expr) → (args : Array Expr) → (mask : Array Bool) → (type : Expr) → MetaM α) : MetaM α := do
-  go #[] #[] #[] 0 altType
+  assert! altInfo.numOverlaps = 0
+  if altInfo.hasUnitThunk then
+    let type ← whnfForall altType
+    let type ← Match.unfoldNamedPattern type
+    let type ← instantiateForall type #[mkConst ``Unit.unit]
+    k #[] #[mkConst ``Unit.unit] #[false] type
+  else
+    go #[] #[] #[] 0 altType
 where
   go (ys : Array Expr) (args : Array Expr) (mask : Array Bool) (i : Nat) (type : Expr) : MetaM α := do
     let type ← whnfForall type
-    if i < altNumParams - numDiscrEqs then
-      let Expr.forallE n d b .. := type
-        | throwError "expecting {altNumParams} parameters, excluding {numDiscrEqs} equalities, but found type{indentExpr altType}"
 
-      -- Handle the special case of `Unit` parameters.
-      if i = 0 && altNumParams - numDiscrEqs = 1 && d.isConstOf ``Unit && !b.hasLooseBVars then
-        return ← k #[] #[mkConst ``Unit.unit] #[false] b
+    if i < altInfo.numFields then
+      let Expr.forallE n d b .. := type
+        | throwError "expecting {altInfo.numFields} parameters, but found type{indentExpr altType}"
 
       let d ← Match.unfoldNamedPattern d
       withLocalDeclD n d fun y => do
@@ -145,17 +149,17 @@ where
 
   This can be used to use the alternative of a match expression in its splitter.
 -/
-partial def forallAltTelescope (altType : Expr) (altNumParams numDiscrEqs : Nat)
+partial def forallAltTelescope (altType : Expr) (altInfo : AltParamInfo) (numDiscrEqs : Nat)
     (k : (ys : Array Expr) → (eqs : Array Expr) → (args : Array Expr) → (mask : Array Bool) → (type : Expr) → MetaM α)
     : MetaM α := do
-  forallAltVarsTelescope altType altNumParams numDiscrEqs fun ys args mask altType => do
+  forallAltVarsTelescope altType altInfo fun ys args mask altType => do
     go ys #[] args mask 0 altType
 where
   go (ys : Array Expr) (eqs : Array Expr) (args : Array Expr) (mask : Array Bool) (i : Nat) (type : Expr) : MetaM α := do
     let type ← whnfForall type
     if i < numDiscrEqs then
       let Expr.forallE n d b .. := type
-        | throwError "expecting {altNumParams} parameters, including {numDiscrEqs} equalities, but found type{indentExpr altType}"
+        | throwError "expecting {numDiscrEqs} equalities, but found type{indentExpr altType}"
       let arg ← if let some (_, _, rhs) ← matchEq? d then
         mkEqRefl rhs
       else if let some (_, _, _, rhs) ← matchHEq? d then
@@ -270,16 +274,17 @@ private partial def withSplitterAlts (altTypes : Array Expr) (f : Array Expr →
       f xs
   go 0 #[]
 
-private abbrev ConvertM := ReaderT (FVarIdMap (Expr × Nat × Array Bool)) $ StateRefT (Array MVarId) MetaM
+private abbrev ConvertM := ReaderT (FVarIdMap (Expr × AltParamInfo × Array Bool)) $ StateRefT (Array MVarId) MetaM
 
 /--
   Construct a proof for the splitter generated by `mkEquationsFor`.
   The proof uses the definition of the `match`-declaration as a template (argument `template`).
   - `alts` are free variables corresponding to alternatives of the `match` auxiliary declaration being processed.
   - `altNews` are the new free variables which contains additional hypotheses that ensure they are only used
-     when the previous overlapping alternatives are not applicable. -/
+     when the previous overlapping alternatives are not applicable.
+  - `altInfos` refers to the splitter -/
 private partial def mkSplitterProof (matchDeclName : Name) (template : Expr) (alts altsNew : Array Expr)
-    (altsNewNumParams : Array Nat) (altArgMasks : Array (Array Bool)) (numDiscrEqs : Nat) : MetaM Expr := do
+    (altInfos : Array AltParamInfo) (altArgMasks : Array (Array Bool)) : MetaM Expr := do
   trace[Meta.Match.matchEqs] "proof template: {template}"
   let map := mkMap
   let (proof, mvarIds) ← convertTemplate template |>.run map |>.run #[]
@@ -289,10 +294,10 @@ private partial def mkSplitterProof (matchDeclName : Name) (template : Expr) (al
     solveOverlap mvarId
   instantiateMVars proof
 where
-  mkMap : FVarIdMap (Expr × Nat × Array Bool) := Id.run do
+  mkMap : FVarIdMap (Expr × AltParamInfo × Array Bool) := Id.run do
     let mut m := {}
-    for alt in alts, altNew in altsNew, numParams in altsNewNumParams, argMask in altArgMasks do
-      m := m.insert alt.fvarId! (altNew, numParams, argMask)
+    for alt in alts, altNew in altsNew, altInfo in altInfos, argMask in altArgMasks do
+      m := m.insert alt.fvarId! (altNew, altInfo, argMask)
     return m
 
   trimFalseTrail (argMask : Array Bool) : Array Bool :=
@@ -452,9 +457,9 @@ where
         return .done (← convertCastEqRec e)
       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}, numParams: {numParams}, e: {e}, alsNew: {altNew}, "
-        if numParams + numDiscrEqs = 0 then
+        let some (altNew, altParamInfo, argMask) := (← read).get? fvarId | return .continue
+        trace[Meta.Match.matchEqs] ">> argMask: {argMask}, altParamInfo: {repr altParamInfo}, e: {e}, alsNew: {altNew}, "
+        if altParamInfo.hasUnitThunk then
           let eNew := mkApp altNew (mkConst ``Unit.unit)
           return TransformStep.done eNew
         let mut newArgs := #[]
@@ -465,8 +470,7 @@ where
           if includeArg then
             newArgs := newArgs.push arg
         let eNew := mkAppN altNew newArgs
-        /- 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)
+        let (mvars, _, _) ← forallMetaBoundedTelescope (← inferType eNew) altParamInfo.numOverlaps (kind := MetavarKind.syntheticOpaque)
         modify fun s => s ++ (mvars.map (·.mvarId!))
         let eNew := mkAppN eNew mvars
         return TransformStep.done eNew
@@ -528,14 +532,14 @@ where go baseName splitterName := withConfig (fun c => { c with etaStruct := .no
     let mut notAlts := #[]
     let mut idx := 1
     let mut splitterAltTypes := #[]
-    let mut splitterAltNumParams := #[]
+    let mut splitterAltInfos := #[]
     let mut altArgMasks := #[] -- masks produced by `forallAltTelescope`
     for i in *...alts.size do
-      let altNumParams := matchInfo.altNumParams[i]!
+      let altInfo := matchInfo.altInfos[i]!
       let thmName := Name.str baseName eqnThmSuffixBase |>.appendIndexAfter idx
       eqnNames := eqnNames.push thmName
-      let (notAlt, splitterAltType, splitterAltNumParam, argMask) ←
-          forallAltTelescope (← inferType alts[i]!) altNumParams numDiscrEqs
+      let (notAlt, splitterAltType, splitterAltInfo, argMask) ←
+          forallAltTelescope (← inferType alts[i]!) altInfo numDiscrEqs
           fun ys eqs rhsArgs argMask altResultType => do
         let patterns := altResultType.getAppArgs
         let mut hs := #[]
@@ -548,12 +552,13 @@ where go baseName splitterName := withConfig (fun c => { c with etaStruct := .no
         let splitterAltType ← mkForallFVars eqs altResultType
         let splitterAltType ← mkArrowN hs splitterAltType
         let splitterAltType ← mkForallFVars ys splitterAltType
-        let splitterAltType ← if splitterAltType == altResultType then
+        let hasUnitThunk := splitterAltType == altResultType
+        let splitterAltType ← if hasUnitThunk then
           mkArrow (mkConst ``Unit) splitterAltType
         else
           pure splitterAltType
         let splitterAltType ← unfoldNamedPattern splitterAltType
-        let splitterAltNumParam := hs.size + ys.size
+        let splitterAltInfo := { numFields := ys.size, numOverlaps := hs.size, hasUnitThunk }
         -- Create a proposition for representing terms that do not match `patterns`
         let mut notAlt := mkConst ``False
         for discr in discrs.toArray.reverse, pattern in patterns.reverse do
@@ -576,10 +581,10 @@ where go baseName splitterName := withConfig (fun c => { c with etaStruct := .no
             type        := thmType
             value       := thmVal
           }
-          return (notAlt, splitterAltType, splitterAltNumParam, argMask)
+          return (notAlt, splitterAltType, splitterAltInfo, argMask)
       notAlts := notAlts.push notAlt
       splitterAltTypes := splitterAltTypes.push splitterAltType
-      splitterAltNumParams := splitterAltNumParams.push splitterAltNumParam
+      splitterAltInfos := splitterAltInfos.push splitterAltInfo
       altArgMasks := altArgMasks.push argMask
       trace[Meta.Match.matchEqs] "splitterAltType: {splitterAltType}"
       idx := idx + 1
@@ -595,7 +600,7 @@ where go baseName splitterName := withConfig (fun c => { c with etaStruct := .no
           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)
+          mkLambdaFVars splitterParams (← mkSplitterProof matchDeclName template alts altsNew splitterAltInfos altArgMasks)
       addAndCompile <| Declaration.defnDecl {
         name        := splitterName
         levelParams := constInfo.levelParams
@@ -605,7 +610,8 @@ where go baseName splitterName := withConfig (fun c => { c with etaStruct := .no
         safety      := .safe
       }
       setInlineAttribute splitterName
-      let result := { eqnNames, splitterName, splitterAltNumParams }
+      let splitterMatchInfo := { matchInfo with altInfos := splitterAltInfos }
+      let result := { eqnNames, splitterName, splitterMatchInfo }
       registerMatchEqns matchDeclName result
 
 /- We generate the equations and splitter on demand, and do not save them on .olean files. -/
@@ -651,12 +657,12 @@ where go baseName := withConfig (fun c => { c with etaStruct := .none }) do
     let mut notAlts := #[]
     let mut idx := 1
     for i in *...alts.size do
-      let altNumParams := matchInfo.altNumParams[i]!
+      let altInfo := matchInfo.altInfos[i]!
       let thmName := (Name.str baseName congrEqnThmSuffixBase).appendIndexAfter idx
       eqnNames := eqnNames.push thmName
       let notAlt ← do
         let alt := alts[i]!
-        Match.forallAltVarsTelescope (← inferType alt) altNumParams numDiscrEqs fun altVars args _mask altResultType => do
+        Match.forallAltVarsTelescope (← inferType alt) altInfo fun altVars args _mask altResultType => do
         let patterns ← forallTelescope altResultType fun _ t => pure t.getAppArgs
         let mut heqsTypes := #[]
         assert! patterns.size == discrs.size

src/Lean/Meta/Match/MatchEqsExt.lean

@@ -7,6 +7,7 @@ module
 
 prelude
 public import Lean.Meta.Basic
+public import Lean.Meta.Match.Basic
 import Lean.Meta.Eqns
 
 public section
@@ -16,8 +17,8 @@ namespace Lean.Meta.Match
 structure MatchEqns where
   eqnNames             : Array Name
   splitterName         : Name
-  splitterAltNumParams : Array Nat
-  deriving Inhabited, Repr
+  splitterMatchInfo    : MatcherInfo
+deriving Inhabited, Repr
 
 def MatchEqns.size (e : MatchEqns) : Nat :=
   e.eqnNames.size

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

@@ -190,10 +190,10 @@ def withUserNames {n} [MonadControlT MetaM n] [Monad n]
 -/
 private def forallAltTelescope'
     {n} [Monad n] [MonadControlT MetaM n]
-    {α} (origAltType : Expr) (numParams numDiscrEqs : Nat)
+    {α} (origAltType : Expr) (altInfo : Match.AltParamInfo)
     (k : Array Expr → Array Expr → n α) : n α := do
   map2MetaM (fun k =>
-    Match.forallAltVarsTelescope origAltType numParams numDiscrEqs
+    Match.forallAltVarsTelescope origAltType altInfo
       fun ys args _mask _bodyType => k ys args
   ) k
 
@@ -308,28 +308,30 @@ def transform
     let mut alts' := #[]
     for altIdx in *...matcherApp.alts.size,
         alt in matcherApp.alts,
-        numParams in matcherApp.altNumParams,
-        splitterNumParams in matchEqns.splitterAltNumParams,
+        altInfo in matcherApp.altInfos,
+        splitterAltInfo in matchEqns.splitterMatchInfo.altInfos,
         origAltType in origAltTypes,
         altType in altTypes do
-      let alt' ← forallAltTelescope' origAltType (numParams - numDiscrEqs) 0 fun ys args => do
+      assert! altInfo.numOverlaps = 0
+      let alt' ← forallAltTelescope' origAltType altInfo fun ys args => do
+        assert! ys.size == splitterAltInfo.numFields
         let altType ← instantiateForall altType ys
         -- Look past the thunking unit parameter, if present
-        let altType ← if splitterNumParams + numDiscrEqs = 0 then
+        let altType ← if altInfo.hasUnitThunk 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
-        let alt' ← forallBoundedTelescope altType (splitterNumParams - ys.size) fun ys2 altType => do
+        let alt' ← forallBoundedTelescope altType splitterAltInfo.numOverlaps 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
+        let alt' ← if splitterAltInfo.hasUnitThunk 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']

src/Lean/Meta/Match/MatcherInfo.lean

@@ -16,11 +16,12 @@ namespace Match
 structure DiscrInfo where
   /-- `some h` if the discriminant is annotated with `h:` -/
   hName? : Option Name := none
-  deriving Inhabited
+deriving Inhabited, Repr
 
 
 structure Overlaps where
   map : Std.HashMap Nat (Std.TreeSet Nat) := {}
+deriving Inhabited, Repr
 
 def Overlaps.insert (o : Overlaps) (overlapping overlapped : Nat) : Overlaps where
   map := o.map.alter overlapped fun s? => some ((s?.getD {}).insert overlapping)
@@ -40,7 +41,7 @@ structure AltParamInfo where
   numOverlaps : Nat
   /-- Whether this alternatie has an artifcial `Unit` parameter -/
   hasUnitThunk : Bool
-deriving Inhabited
+deriving Inhabited, Repr
 
 /--
 A "matcher" auxiliary declaration has the following structure:
@@ -63,6 +64,7 @@ structure MatcherInfo where
   -/
   discrInfos   : Array DiscrInfo
   overlaps     : Overlaps := {}
+deriving Inhabited, Repr
 
 @[expose] def MatcherInfo.numAlts (info : MatcherInfo) : Nat :=
   info.altInfos.size

src/Lean/Meta/Tactic/Split.lean

@@ -284,14 +284,14 @@ def applyMatchSplitter (mvarId : MVarId) (matcherDeclName : Name) (us : Array Le
     trace[split.debug] "after check splitter"
     let mvarIds ← mvarId.applyN splitter matchEqns.size
     let (_, mvarIds) ← mvarIds.foldlM (init := (0, [])) fun (i, mvarIds) mvarId => do
-      let numParams := matchEqns.splitterAltNumParams[i]!
+      let altInfo := matchEqns.splitterMatchInfo.altInfos[i]!
       let mvarId ←
-        if numParams + info.getNumDiscrEqs = 0 then
+        if altInfo.hasUnitThunk 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
+          let (_, mvarId) ← mvarId.introN (altInfo.numFields + altInfo.numOverlaps)
           pure mvarId
       trace[split.debug] "before unifyEqs\n{mvarId}"
       match (← Cases.unifyEqs? (info.getNumDiscrEqs + numEqs) mvarId {}) with