refactor: treat casesOn and matcher applications uniformly

src/Lean/Compiler/LCNF.lean

@@ -78,20 +78,6 @@ where
     else
       k
 
-  /--
-  Visit a `matcher`/`casesOn` alternative.
-  -/
-  visitAlt (e : Expr) (numParams : Nat) : M Expr := do
-    withNewRootScope do
-      let mut (as, e) ← Compiler.visitBoundedLambda e numParams
-      if as.size < numParams then
-        e ← etaExpandN e (numParams - as.size)
-        let (as', e') ← Compiler.visitLambda e
-        as := as ++ as'
-        e := e'
-      e ← mkLetUsingScope (← visit e)
-      mkLambda as e
-
   /--
   If `args.size == arity`, then just return `app`.
   Otherwise return
@@ -125,25 +111,29 @@ where
     else
       mkOverApplication (mkAppN f args[:arity]) args arity
 
-  visitMatcher (matcherInfo : Meta.MatcherInfo) (e : Expr) : M Expr :=
-    etaIfUnderApplied e matcherInfo.arity do
-      let mut args := e.getAppArgs
-      for i in matcherInfo.getDiscrRange do
-        args ← args.modifyM i visitChild
-      for i in matcherInfo.getAltRange, altIdx in [: matcherInfo.numAlts] do
-        let alt ← visitAlt args[i]! matcherInfo.altNumParams[altIdx]!
-        args := args.set! i alt
-      mkAppWithArity e.getAppFn args matcherInfo.arity
+  /--
+  Visit a `matcher`/`casesOn` alternative.
+  -/
+  visitAlt (e : Expr) (numParams : Nat) : M Expr := do
+    withNewRootScope do
+      let mut (as, e) ← Compiler.visitBoundedLambda e numParams
+      if as.size < numParams then
+        e ← etaExpandN e (numParams - as.size)
+        let (as', e') ← Compiler.visitLambda e
+        as := as ++ as'
+        e := e'
+      e ← mkLetUsingScope (← visit e)
+      mkLambda as e
 
-  visitCasesOn (casesOnInfo : CasesOnInfo) (e : Expr) : M Expr :=
-    etaIfUnderApplied e casesOnInfo.arity do
+  visitCases (casesInfo : CasesInfo) (e : Expr) : M Expr :=
+    etaIfUnderApplied e casesInfo.arity do
       let mut args := e.getAppArgs
-      args ← args.modifyM casesOnInfo.majorPos visitChild
-      for i in casesOnInfo.minorsRange, ctor in casesOnInfo.ctors do
-        let .ctorInfo ctorVal ← getConstInfo ctor | unreachable!
-        let alt ← visitAlt args[i]! ctorVal.numFields
+      for i in casesInfo.discrsRange do
+        args ← args.modifyM i visitChild
+      for i in casesInfo.altsRange, numParams in casesInfo.altNumParams do
+        let alt ← visitAlt args[i]! numParams
         args := args.set! i alt
-      mkAppWithArity e.getAppFn args casesOnInfo.arity
+      mkAppWithArity e.getAppFn args casesInfo.arity
 
   visitCtor (arity : Nat) (e : Expr) : M Expr :=
     etaIfUnderApplied e arity do
@@ -239,9 +229,7 @@ where
 
   visitApp (e : Expr) : M Expr := do
     if let .const declName _ := e.getAppFn then
-      if let some matcherInfo ← Meta.getMatcherInfo? declName then
-        visitMatcher matcherInfo e
-      else if declName == ``Quot.lift then
+      if declName == ``Quot.lift then
         visitQuotLift e
       else if declName == ``Quot.mk then
         visitCtor 3 e
@@ -251,8 +239,8 @@ where
         visitAndRec e
       else if declName == ``False.rec || declName == ``Empty.rec || declName == ``False.casesOn || declName == ``Empty.casesOn then
         visitFalseRec e
-      else if let some casesOnInfo ← getCasesOnInfo? declName then
-        visitCasesOn casesOnInfo e
+      else if let some casesInfo ← getCasesInfo? declName then
+        visitCases casesInfo e
       else if let some arity ← getCtorArity? declName then
         visitCtor arity e
       else if isNoConfusion (← getEnv) declName then

src/Lean/Compiler/Util.lean

@@ -4,6 +4,7 @@ Released under Apache 2.0 license as described in the file LICENSE.
 Authors: Leonardo de Moura
 -/
 import Lean.Compiler.CompilerM
+import Lean.Meta.Match.MatcherInfo
 
 namespace Lean.Compiler
 
@@ -41,23 +42,43 @@ def mkLcUnreachable (type : Expr) : CompilerM Expr := do
     let u ← Meta.getLevel type
     return .app (.const ``lcUnreachable [u]) type
 
-structure CasesOnInfo where
-  arity       : Nat
-  majorPos    : Nat
-  minorsRange : Std.Range
-  ctors       : List Name
+/--
+Store information about `matcher` and `casesOn` declarations.
 
-def getCasesOnInductiveVal? (declName : Name) : CoreM (Option InductiveVal) := do
+We treat them uniformly in the code generator.
+-/
+structure CasesInfo where
+  arity        : Nat
+  discrsRange  : Std.Range
+  altsRange    : Std.Range
+  altNumParams : Array Nat
+
+private def getCasesOnInductiveVal? (declName : Name) : CoreM (Option InductiveVal) := do
   unless isCasesOnRecursor (← getEnv) declName do return none
   let .inductInfo val ← getConstInfo declName.getPrefix | return none
   return some val
 
-def getCasesOnInfo? (declName : Name) : CoreM (Option CasesOnInfo) := do
+private def getCasesOnInfo? (declName : Name) : CoreM (Option CasesInfo) := do
   let some val ← getCasesOnInductiveVal? declName | return none
-  let arity       := val.numIndices + val.numParams + 1 /- motive -/ + 1 /- major -/ + val.numCtors
-  let majorPos    := val.numIndices + val.numParams + 1 /- motive -/
-  let minorsRange := { start := majorPos + 1, stop := arity }
-  return some { arity, majorPos, minorsRange, ctors := val.ctors }
+  let arity        := val.numIndices + val.numParams + 1 /- motive -/ + 1 /- major -/ + val.numCtors
+  let majorPos     := val.numIndices + val.numParams + 1 /- motive -/
+  let discrsRange  := { start := majorPos, stop := majorPos + 1 }
+  let altsRange    := { start := majorPos + 1, stop := arity }
+  let altNumParams ← val.ctors.toArray.mapM fun ctor => do
+    let .ctorInfo ctorVal ← getConstInfo ctor | unreachable!
+    return ctorVal.numFields
+  return some { arity, discrsRange, altsRange, altNumParams }
+
+def getCasesInfo? (declName : Name) : CoreM (Option CasesInfo) := do
+  if let some matcherInfo ← Meta.getMatcherInfo? declName then
+    return some {
+      arity        := matcherInfo.arity
+      discrsRange  := matcherInfo.getDiscrRange
+      altsRange    := matcherInfo.getAltRange
+      altNumParams := matcherInfo.altNumParams
+    }
+  else
+    getCasesOnInfo? declName
 
 def getCtorArity? (declName : Name) : CoreM (Option Nat) := do
   let .ctorInfo val ← getConstInfo declName | return none