refactor: IndGroupInfo and IndGroupInst (#4738)

This adds the types
* `IndGroupInfo`, a variant of `InductiveVal` with information that
   applies to a whole group of mutual inductives and
* `IndGroupInst` which extends `IndGroupInfo` with levels and parameters
   to indicate a instantiation of the group.

One purpose of this abstraction is to make it clear when a fuction
operates on a group as a whole, rather than a specific inductive within
the group.

This is extracted from #4718 and #4733 to reduce PR size and improve
bisectability.

src/Lean/Elab/PreDefinition/Structural/BRecOn.lean

@@ -10,6 +10,7 @@ import Lean.Elab.RecAppSyntax
 import Lean.Elab.PreDefinition.Basic
 import Lean.Elab.PreDefinition.Structural.Basic
 import Lean.Elab.PreDefinition.Structural.FunPacker
+import Lean.Elab.PreDefinition.Structural.RecArgInfo
 
 namespace Lean.Elab.Structural
 open Meta
@@ -145,26 +146,16 @@ private partial def replaceRecApps (recArgInfos : Array RecArgInfo) (positions :
         e.withApp fun f args => do
           if let .some fnIdx := recArgInfos.findIdx? (f.isConstOf ·.fnName) then
             let recArgInfo := recArgInfos[fnIdx]!
-            let numFixed  := recArgInfo.numFixed
-            let recArgPos := recArgInfo.recArgPos
-            if recArgPos >= args.size then
-              throwError "insufficient number of parameters at recursive application {indentExpr e}"
-            let recArg := args[recArgPos]!
+            let some recArg := args[recArgInfo.recArgPos]?
+              | throwError "insufficient number of parameters at recursive application {indentExpr e}"
             -- For reflexive type, we may have nested recursive applications in recArg
             let recArg ← loop below recArg
             let f ←
-              try toBelow below recArgInfo.indParams.size positions fnIdx recArg
+              try toBelow below recArgInfo.indGroupInst.params.size positions fnIdx recArg
               catch _ => throwError "failed to eliminate recursive application{indentExpr e}"
-            -- Recall that the fixed parameters are not in the scope of the `brecOn`. So, we skip them.
-            let argsNonFixed := args.extract numFixed args.size
-            -- The function `f` does not explicitly take `recArg` and its indices as arguments. So, we skip them too.
-            let mut fArgs := #[]
-            for i in [:argsNonFixed.size] do
-              let j := i + numFixed
-              if recArgInfo.recArgPos != j && !recArgInfo.indicesPos.contains j then
-                let arg := argsNonFixed[i]!
-                let arg ← replaceRecApps recArgInfos positions below arg
-                fArgs := fArgs.push arg
+            -- We don't pass the fixed parameters, the indices and the major arg to `f`, only the rest
+            let (_, fArgs) := recArgInfo.pickIndicesMajor args[recArgInfo.numFixed:]
+            let fArgs ← fArgs.mapM (replaceRecApps recArgInfos positions below ·)
             return mkAppN f fArgs
           else
             return mkAppN (← loop below f) (← args.mapM (loop below))
@@ -237,35 +228,39 @@ def mkBRecOnF (recArgInfos : Array RecArgInfo) (positions : Positions)
       let valueNew   ← replaceRecApps recArgInfos positions below value
       mkLambdaFVars (indexMajorArgs ++ #[below] ++ otherArgs) valueNew
 
-
 /--
 Given the `motives`, figures out whether to use `.brecOn` or `.binductionOn`, pass
 the right universe levels, the parameters, and the motives.
 It was already checked earlier in `checkCodomainsLevel` that the functions live in the same universe.
 -/
 def mkBRecOnConst (recArgInfos : Array RecArgInfo) (positions : Positions)
-   (motives : Array Expr) : MetaM (Name → Expr) := do
-  -- For now, just look at the first
-  let recArgInfo := recArgInfos[0]!
+   (motives : Array Expr) : MetaM (Nat → Expr) := do
+  let indGroup := recArgInfos[0]!.indGroupInst
   let motive := motives[0]!
   let brecOnUniv ← lambdaTelescope motive fun _ type => getLevel type
-  let indInfo ← getConstInfoInduct recArgInfo.indName
+  let indInfo ← getConstInfoInduct indGroup.all[0]!
   let useBInductionOn := indInfo.isReflexive && brecOnUniv == levelZero
   let brecOnUniv ←
     if indInfo.isReflexive && brecOnUniv != levelZero then
       decLevel brecOnUniv
     else
       pure brecOnUniv
-  let brecOnCons := fun n =>
+  let brecOnCons := fun idx  =>
     let brecOn :=
-      if useBInductionOn then .const (mkBInductionOnName n) recArgInfo.indLevels
-      else                    .const (mkBRecOnName n) (brecOnUniv :: recArgInfo.indLevels)
-    mkAppN brecOn recArgInfo.indParams
+      if let .some n := indGroup.all[idx]? then
+        if useBInductionOn then .const (mkBInductionOnName n) indGroup.levels
+        else                    .const (mkBRecOnName n) (brecOnUniv :: indGroup.levels)
+      else
+        let n := indGroup.all[0]!
+        let j := idx - indGroup.all.size + 1
+        if useBInductionOn then .const (mkBInductionOnName n |>.appendIndexAfter j) indGroup.levels
+        else                    .const (mkBRecOnName n |>.appendIndexAfter j) (brecOnUniv :: indGroup.levels)
+    mkAppN brecOn indGroup.params
 
   -- Pick one as a prototype
-  let brecOnAux := brecOnCons recArgInfo.indName
+  let brecOnAux := brecOnCons 0
   -- Infer the type of the packed motive arguments
-  let packedMotiveTypes ← inferArgumentTypesN recArgInfo.indAll.size brecOnAux
+  let packedMotiveTypes ← inferArgumentTypesN indGroup.numMotives brecOnAux
   let packedMotives ← positions.mapMwith packMotives packedMotiveTypes motives
 
   return fun n => mkAppN (brecOnCons n) packedMotives
@@ -277,10 +272,10 @@ combinators. This assumes that all `.brecOn` functions of a mutual inductive hav
 It also undoes the permutation and packing done by `packMotives`
 -/
 def inferBRecOnFTypes (recArgInfos : Array RecArgInfo) (positions : Positions)
-    (brecOnConst : Name → Expr) : MetaM (Array Expr) := do
+    (brecOnConst : Nat → Expr) : MetaM (Array Expr) := do
   let numTypeFormers := positions.size
   let recArgInfo := recArgInfos[0]! -- pick an arbitrary one
-  let brecOn := brecOnConst recArgInfo.indName
+  let brecOn := brecOnConst 0
   check brecOn
   let brecOnType ← inferType brecOn
   -- Skip the indices and major argument
@@ -298,11 +293,11 @@ def inferBRecOnFTypes (recArgInfos : Array RecArgInfo) (positions : Positions)
 Completes the `.brecOn` for the given function.
 The `value` is the function with (only) the fixed parameters moved into the context.
 -/
-def mkBrecOnApp (positions : Positions) (fnIdx : Nat) (brecOnConst : Name → Expr)
+def mkBrecOnApp (positions : Positions) (fnIdx : Nat) (brecOnConst : Nat → Expr)
     (FArgs : Array Expr) (recArgInfo : RecArgInfo) (value : Expr) : MetaM Expr := do
   lambdaTelescope value fun ys _value => do
     let (indexMajorArgs, otherArgs) := recArgInfo.pickIndicesMajor ys
-    let brecOn := brecOnConst recArgInfo.indName
+    let brecOn := brecOnConst recArgInfo.indIdx
     let brecOn := mkAppN brecOn indexMajorArgs
     let packedFTypes ← inferArgumentTypesN positions.size brecOn
     let packedFArgs ← positions.mapMwith packFArgs packedFTypes FArgs

src/Lean/Elab/PreDefinition/Structural/Basic.lean

@@ -9,46 +9,6 @@ import Lean.Meta.ForEachExpr
 
 namespace Lean.Elab.Structural
 
-/--
-Information about the argument of interest of a structurally recursive function.
-
-The `Expr`s in this data structure expect the `fixedParams` to be in scope, but not the other
-parameters of the function. This ensures that this data structure makes sense in the other functions
-of a mutually recursive group.
--/
-structure RecArgInfo where
-  /-- the name of the recursive function -/
-  fnName      : Name
-  /-- the fixed prefix of arguments of the function we are trying to justify termination using structural recursion. -/
-  numFixed    : Nat
-  /-- position of the argument (counted including fixed prefix) we are recursing on -/
-  recArgPos   : Nat
-  /-- position of the indices (counted including fixed prefix) of the inductive datatype indices we are recursing on -/
-  indicesPos  : Array Nat
-  /-- inductive datatype name of the argument we are recursing on -/
-  indName     : Name
-  /-- inductive datatype universe levels of the argument we are recursing on -/
-  indLevels   : List Level
-  /-- inductive datatype parameters of the argument we are recursing on -/
-  indParams   : Array Expr
-  /-- The types mutually inductive with indName -/
-  indAll      : Array Name
-deriving Inhabited
-/--
-If `xs` are the parameters of the functions (excluding fixed prefix), partitions them
-into indices and major arguments, and other parameters.
--/
-def RecArgInfo.pickIndicesMajor (info : RecArgInfo) (xs : Array Expr) : (Array Expr × Array Expr) := Id.run do
-  let mut indexMajorArgs := #[]
-  let mut otherArgs := #[]
-  for h : i in [:xs.size] do
-    let j := i + info.numFixed
-    if j = info.recArgPos || info.indicesPos.contains j then
-      indexMajorArgs := indexMajorArgs.push xs[i]
-    else
-      otherArgs := otherArgs.push xs[i]
-  return (indexMajorArgs, otherArgs)
-
 structure State where
   /-- As part of the inductive predicates case, we keep adding more and more discriminants from the
      local context and build up a bigger matcher application until we reach a fixed point.
@@ -128,3 +88,6 @@ def Positions.mapMwith {α β m} [Monad m] [Inhabited β] (f : α → Array β 
   (Array.zip ys positions).mapM fun ⟨y, poss⟩ => f y (poss.map (xs[·]!))
 
 end Lean.Elab.Structural
+
+builtin_initialize
+  Lean.registerTraceClass `Elab.definition.structural

src/Lean/Elab/PreDefinition/Structural/FindRecArg.lean

@@ -5,6 +5,7 @@ Authors: Leonardo de Moura, Joachim Breitner
 -/
 prelude
 import Lean.Elab.PreDefinition.Structural.Basic
+import Lean.Elab.PreDefinition.Structural.RecArgInfo
 
 namespace Lean.Elab.Structural
 open Meta
@@ -74,15 +75,19 @@ def getRecArgInfo (fnName : Name) (numFixed : Nat) (xs : Array Expr) (i : Nat) :
           | some (indParam, y) =>
             throwError "its type is an inductive datatype{indentExpr xType}\nand the datatype parameter{indentExpr indParam}\ndepends on the function parameter{indentExpr y}\nwhich does not come before the varying parameters and before the indices of the recursion parameter."
           | none =>
+            let indAll := indInfo.all.toArray
+            let .some indIdx := indAll.indexOf? indInfo.name | panic! "{indInfo.name} not in {indInfo.all}"
             let indicesPos := indIndices.map fun index => match xs.indexOf? index with | some i => i.val | none => unreachable!
-            return { fnName      := fnName
-                     numFixed    := numFixed
-                     recArgPos   := i
-                     indicesPos  := indicesPos
-                     indName     := indInfo.name
-                     indLevels   := us
-                     indParams   := indParams
-                     indAll      := indInfo.all.toArray }
+            let indGroupInst := {
+              IndGroupInfo.ofInductiveVal indInfo with
+              levels := us
+              params := indParams }
+            return { fnName       := fnName
+                     numFixed     := numFixed
+                     recArgPos    := i
+                     indicesPos   := indicesPos
+                     indGroupInst := indGroupInst
+                     indIdx       := indIdx }
     else
       throwError "the index #{i+1} exceeds {xs.size}, the number of parameters"
 

src/Lean/Elab/PreDefinition/Structural/IndGroupInfo.lean

@@ -0,0 +1,65 @@
+/-
+Copyright (c) 2024 Lean FRO. All rights reserved.
+Released under Apache 2.0 license as described in the file LICENSE.
+Authors: Joachim Breitner
+-/
+prelude
+import Lean.Meta.InferType
+
+/-!
+This module contains the types
+* `IndGroupInfo`, a variant of `InductiveVal` with information that
+   applies to a whole group of mutual inductives and
+* `IndGroupInst` which extends `IndGroupInfo` with levels and parameters
+   to indicate a instantiation of the group.
+
+One purpose of this abstraction is to make it clear when a fuction operates on a group as
+a whole, rather than a specific inductive within the group.
+-/
+
+namespace Lean.Elab.Structural
+open Lean Meta
+
+/--
+A mutually inductive group, identified by the `all` array of the `InductiveVal` of its
+constituents.
+-/
+structure IndGroupInfo where
+  all       : Array Name
+  numNested : Nat
+deriving BEq, Inhabited
+
+def IndGroupInfo.ofInductiveVal (indInfo : InductiveVal) : IndGroupInfo where
+  all       := indInfo.all.toArray
+  numNested := indInfo.numNested
+
+def IndGroupInfo.numMotives (group : IndGroupInfo) : Nat :=
+  group.all.size + group.numNested
+
+/--
+An instance of an mutually inductive group of inductives, identified by the `all` array
+and the level and expressions parameters.
+
+For example this distinguishes between `List α` and `List β` so that we will not even attempt
+mutual structural recursion on such incompatible types.
+-/
+structure IndGroupInst extends IndGroupInfo where
+  levels    : List Level
+  params    : Array Expr
+deriving Inhabited
+
+def IndGroupInst.toMessageData (igi : IndGroupInst) : MessageData :=
+  mkAppN (.const igi.all[0]! igi.levels) igi.params
+
+instance : ToMessageData IndGroupInst where
+  toMessageData := IndGroupInst.toMessageData
+
+def IndGroupInst.isDefEq (igi1 igi2 : IndGroupInst) : MetaM Bool := do
+  unless igi1.toIndGroupInfo == igi2.toIndGroupInfo do return false
+  unless igi1.levels.length = igi2.levels.length do return false
+  unless (igi1.levels.zip igi2.levels).all (fun (l₁, l₂) => Level.isEquiv l₁ l₂) do return false
+  unless igi1.params.size = igi2.params.size do return false
+  unless (← (igi1.params.zip igi2.params).allM (fun (e₁, e₂) => Meta.isDefEqGuarded e₁ e₂)) do return false
+  return true
+
+end Lean.Elab.Structural

src/Lean/Elab/PreDefinition/Structural/IndPred.lean

@@ -7,6 +7,7 @@ prelude
 import Lean.Meta.IndPredBelow
 import Lean.Elab.PreDefinition.Basic
 import Lean.Elab.PreDefinition.Structural.Basic
+import Lean.Elab.PreDefinition.Structural.RecArgInfo
 
 namespace Lean.Elab.Structural
 open Meta
@@ -81,8 +82,8 @@ def mkIndPredBRecOn (recArgInfo : RecArgInfo) (value : Expr) : M Expr := do
     let motive ← mkForallFVars otherArgs type
     let motive ← mkLambdaFVars indexMajorArgs motive
     trace[Elab.definition.structural] "brecOn motive: {motive}"
-    let brecOn := Lean.mkConst (mkBRecOnName recArgInfo.indName) recArgInfo.indLevels
-    let brecOn := mkAppN brecOn recArgInfo.indParams
+    let brecOn := Lean.mkConst (mkBRecOnName recArgInfo.indName!) recArgInfo.indGroupInst.levels
+    let brecOn := mkAppN brecOn recArgInfo.indGroupInst.params
     let brecOn := mkApp brecOn motive
     let brecOn := mkAppN brecOn indexMajorArgs
     check brecOn

src/Lean/Elab/PreDefinition/Structural/Main.lean

@@ -92,11 +92,11 @@ def getMutualFixedPrefix (preDefs : Array PreDefinition) : M Nat :=
 /-- Checks that all parameter types are mutually inductive -/
 private def checkAllFromSameClique (recArgInfos : Array RecArgInfo) : MetaM Unit := do
   for recArgInfo in recArgInfos do
-    unless recArgInfos[0]!.indAll.contains recArgInfo.indName do
+    unless recArgInfos[0]!.indGroupInst.all.contains recArgInfo.indName! do
       throwError m!"Cannot use structural mutual recursion: The recursive argument of " ++
-        m!"{recArgInfos[0]!.fnName} is of type {recArgInfos[0]!.indName}, " ++
+        m!"{recArgInfos[0]!.fnName} is of type {recArgInfos[0]!.indName!}, " ++
         m!"the recursive argument of {recArgInfo.fnName} is of type " ++
-            m!"{recArgInfo.indName}, and these are not mutually recursive."
+            m!"{recArgInfo.indName!}, and these are not mutually recursive."
 
 private def elimMutualRecursion (preDefs : Array PreDefinition) (recArgPoss : Array Nat) : M (Array PreDefinition) := do
   withoutModifyingEnv do
@@ -120,17 +120,17 @@ private def elimMutualRecursion (preDefs : Array PreDefinition) (recArgPoss : Ar
         lambdaTelescope value fun ys _value => do
           getRecArgInfo preDef.declName maxNumFixed (xs ++ ys) recArgPos
 
+      checkAllFromSameClique recArgInfos
+
       -- Check that the inductive parameters agree. This also ensures that they only depend on the
       -- trimmed prefix (minimum of all `.numFixed`).
       for recArgInfo in recArgInfos[1:] do
-        let ok ← Array.allM
-          (fun (e₁, e₂) => isDefEqGuarded e₁ e₂)
-          (Array.zip recArgInfo.indParams recArgInfos[0]!.indParams)
-        unless recArgInfos[0]!.indParams.size = recArgInfo.indParams.size && ok do
+        let ok ← IndGroupInst.isDefEq recArgInfo.indGroupInst recArgInfos[0]!.indGroupInst
+        unless ok do
           throwError m!"The inductive type of the recursive parameter of {recArgInfos[0]!.fnName} " ++
             m!"and {recArgInfo.fnName} have different parameters:" ++
-            indentD m!"{recArgInfos[0]!.indParams}" ++
-            indentD m!"{recArgInfo.indParams}"
+            indentD m!"{recArgInfos[0]!.indGroupInst.params}" ++
+            indentD m!"{recArgInfo.indGroupInst.params}"
 
       return (recArgInfos.map (·.numFixed)).foldl Nat.min maxNumFixed
 
@@ -151,7 +151,7 @@ private def elimMutualRecursion (preDefs : Array PreDefinition) (recArgPoss : Ar
       -- Two passes should suffice
       assert! recArgInfos.all (·.numFixed = numFixed)
 
-      let indInfo ← getConstInfoInduct recArgInfos[0]!.indName
+      let indInfo ← getConstInfoInduct recArgInfos[0]!.indName!
       if ← isInductivePredicate indInfo.name then
         -- Here we branch off to the IndPred construction, but only for non-mutual functions
         unless preDefs.size = 1 do
@@ -166,9 +166,8 @@ private def elimMutualRecursion (preDefs : Array PreDefinition) (recArgPoss : Ar
         check valueNew
         return #[{ preDef with value := valueNew }]
 
-      checkAllFromSameClique recArgInfos
       -- Sort the (indices of the) definitions by their position in indInfo.all
-      let positions : Positions := .groupAndSort (·.indName) recArgInfos indInfo.all.toArray
+      let positions : Positions := .groupAndSort (·.indName!) recArgInfos indInfo.all.toArray
 
       -- Construct the common `.brecOn` arguments
       let motives ← (Array.zip recArgInfos values).mapM fun (r, v) => mkBRecOnMotive r v
@@ -236,8 +235,6 @@ def structuralRecursion (preDefs : Array PreDefinition) (termArgs? : Option Term
     markAsRecursive preDef.declName
   applyAttributesOf preDefsNonRec AttributeApplicationTime.afterCompilation
 
-builtin_initialize
-  registerTraceClass `Elab.definition.structural
 
 end Structural