feat: tc: generalize instances when applying

gebner · gebner · commit 182e3436269e · 2023-03-28T11:50:16.000-07:00
diff --git a/src/Lean/Meta/SynthInstance.lean b/src/Lean/Meta/SynthInstance.lean
@@ -26,7 +26,7 @@ register_builtin_option synthInstance.maxSize : Nat := {
 }
 
 register_builtin_option synthInstance.etaExperiment : Bool := {
-  defValue := false
+  defValue := true
   descr := "[DO NOT USE EXCEPT FOR TESTING] enable structure eta for type-classes during type-class search"
 }
 
@@ -48,11 +48,15 @@ structure GeneratorNode where
   currInstanceIdx : Nat
   deriving Inhabited
 
+inductive Subgoal where
+  | tc (cls : Expr)
+  | defEq (a b : Expr)
+
 structure ConsumerNode where
   mvar     : Expr
   key      : Expr
   mctx     : MetavarContext
-  subgoals : List Expr
+  subgoals : List Subgoal
   size     : Nat -- instance size so far
   deriving Inhabited
 
@@ -328,10 +332,40 @@ def preprocessOutParam (type : Expr) : MetaM Expr :=
    Given `getSubgoalsAux args j subgoals instVal type`,
    we have that `type.instantiateRevRange j args.size args` does not have loose bound variables. -/
 structure SubgoalsResult where
-  subgoals     : List Expr
+  subgoals     : Array Subgoal
   instVal      : Expr
   instTypeBody : Expr
 
+partial def generalizeTypeClassInstances (e : Expr) (outerLCtx : LocalContext) (outerLocalInsts : LocalInstances) :
+    MetaM (Expr × Array Subgoal) := do
+  let subgoals ← IO.mkRef #[]
+  let rec
+    go (e : Expr) : MonadCacheT ExprStructEq Expr MetaM Expr :=
+      checkCache (ExprStructEq.mk e) fun _ => withIncRecDepth do
+        let e := e.headBeta -- (fun x => t) 1
+        if let some e := e.etaExpandedStrict? then go e else
+        if (← isClass? (← inferType e)).isSome then
+          if !e.getAppFn.isMVar && e.hasExprMVar then
+            return ← mkMVarForExpr e
+        match e with
+        | .lam .. | .letE .. => lambdaLetTelescope e fun xs e' => do mkLambdaFVars xs (← go e')
+        | .forallE .. => forallTelescope e fun xs e' => do mkForallFVars xs (← go e')
+        | .mdata d e => return .mdata d (← go e)
+        | _ => e.withApp fun f as => do
+          return mkAppN f (← as.mapM go),
+    mkMVarForExpr (e : Expr) := do
+      let (_, usedFVars) ← e.collectFVars.run {}
+      let usedFVars ← usedFVars.addDependencies
+      let mut args := #[]
+      for ldecl in ← getLCtx do
+        if !outerLCtx.contains ldecl.fvarId && usedFVars.fvarSet.contains ldecl.fvarId then
+          args := args.push (.fvar ldecl.fvarId)
+      let mvarType ← mkForallFVars args (← go (← inferType e))
+      let mvar ← mkFreshExprMVarAt outerLCtx outerLocalInsts mvarType
+      subgoals.modify (·.push (.defEq (← mkLambdaFVars args e) mvar))
+      return mkAppN mvar args
+  return (← (go e).run, ← subgoals.get)
+
 /--
   `getSubgoals lctx localInsts xs inst` creates the subgoals for the instance `inst`.
   The subgoals are in the context of the free variables `xs`, and
@@ -365,29 +399,29 @@ def getSubgoals (lctx : LocalContext) (localInsts : LocalInstances) (xs : Array
       instVal := instVal.instantiateRev subst
       subst := #[]
       unless instType.isForall do break
+  let (instType', defEqSubgoals) ← generalizeTypeClassInstances (instType.instantiateRev subst) lctx localInsts
   return {
     instVal := instVal.instantiateRev subst
-    instTypeBody := instType.instantiateRev subst
-    subgoals := inst.synthOrder.map (mvars[·]!) |>.toList
+    instTypeBody := instType'
+    subgoals := inst.synthOrder.map (Subgoal.tc mvars[·]!) ++ defEqSubgoals
   }
 
 /--
   Try to synthesize metavariable `mvar` using the instance `inst`.
   Remark: `mctx` is set using `withMCtx`.
   If it succeeds, the result is a new updated metavariable context and a new list of subgoals.
   A subgoal is created for each instance implicit parameter of `inst`. -/
-def tryResolve (mvar : Expr) (inst : Instance) : MetaM (Option (MetavarContext × List Expr)) := do
+def tryResolve (mvar : Expr) (inst : Instance) : MetaM (Option (MetavarContext × List Subgoal)) := do
   let mvarType   ← inferType mvar
   let lctx       ← getLCtx
   let localInsts ← getLocalInstances
   forallTelescopeReducing mvarType fun xs mvarTypeBody => do
-    let ⟨subgoals, instVal, instTypeBody⟩ ← getSubgoals lctx localInsts xs inst
+    let {subgoals, instTypeBody, instVal} ← getSubgoals lctx localInsts xs inst
     withTraceNode `Meta.synthInstance.tryResolve (withMCtx (← getMCtx) do
         return m!"{exceptOptionEmoji ·} {← instantiateMVars mvarTypeBody} ≟ {← instantiateMVars instTypeBody}") do
     if (← isDefEq mvarTypeBody instTypeBody) then
       let instVal ← mkLambdaFVars xs instVal
-      if (← isDefEq mvar instVal) then
-        return some ((← getMCtx), subgoals)
+      return some ((← getMCtx), (subgoals.push (.defEq mvar instVal)).toList)
     return none
 
 /--
@@ -463,11 +497,11 @@ def inprocessOutParams (mvar : Expr) : MetaM Expr := do
   let type ← instantiateMVars (← inferType mvar)
   let outerLCtx ← getLCtx
   let outerLocalInsts ← getLocalInstances
-  forallTelescopeReducing type fun xs typeBody => do
-  let typeBody ← whnf typeBody
+  forallTelescopeReducing type fun xs typeBody0 => do
+  let typeBody ← whnf typeBody0
   let c@(.const className _) := typeBody.getAppFn | return mvar
-  let some outParamsPos := getOutParamPositions? (← getEnv) className | return mvar
-  if outParamsPos.isEmpty then return mvar
+  let outParamsPos := (getOutParamPositions? (← getEnv) className).getD #[]
+  if outParamsPos.isEmpty && typeBody == typeBody0 then return mvar
   let args := typeBody.getAppArgs
   let cType ← inferType c
   let args ← preprocessArgs cType 0 args outParamsPos
@@ -488,7 +522,15 @@ def removeUnusedArguments (mvar : Expr) : MetaM Expr := do
   return mvar'
 
 /-- Process the next subgoal in the given consumer node. -/
-def consume (cNode : ConsumerNode) : SynthM Unit := do
+partial def consume (cNode : ConsumerNode) : SynthM Unit := do
+  match cNode.subgoals with
+  | []      => addAnswer cNode
+  | .defEq a b :: subgoals =>
+    if let (true, mctx) ← withMCtx cNode.mctx do return (← withDefault (isDefEq a b), ← getMCtx) then
+      consume { cNode with subgoals, mctx }
+    else
+      pure ()
+  | .tc mvar :: mvars =>
   /- Filter out subgoals that have already been assigned when solving typing constraints.
     This may happen when a local instance type depends on other local instances.
     For example, in Mathlib, we have
@@ -500,18 +542,12 @@ def consume (cNode : ConsumerNode) : SynthM Unit := do
       SetLike (@Submodule R M _inst_1 _inst_2 _inst_3) M
     ```
   -/
-  let cNode := { cNode with
-    subgoals := ← withMCtx cNode.mctx do
-      cNode.subgoals.filterM (not <$> ·.mvarId!.isAssigned)
-  }
-  match cNode.subgoals with
-  | []      => addAnswer cNode
-  | mvar::mvars =>
+    if ← mvar.mvarId!.isAssigned then consume { cNode with subgoals := mvars } else
     let (mvar, mvar', mctx) ← withMCtx cNode.mctx do
       let mvar ← removeUnusedArguments mvar
       let mvar' ← inprocessOutParams mvar
       return (mvar, mvar', ← getMCtx)
-    let cNode := { cNode with subgoals := mvar::mvars, mctx }
+    let cNode := { cNode with subgoals := .tc mvar :: mvars, mctx }
     let waiter := Waiter.consumerNode cNode
     let key ← mkTableKeyFor cNode.mctx mvar'
     match (← findEntry? key) with
@@ -559,7 +595,8 @@ def resume : SynthM Unit := do
   let (cNode, answer) ← getNextToResume
   match cNode.subgoals with
   | []         => panic! "resume found no remaining subgoals"
-  | mvar::rest =>
+  | .defEq .. :: _ => panic! "resume found defeq constraint"
+  | .tc mvar :: rest =>
     match (← tryAnswer cNode.mctx mvar answer) with
     | none      => return ()
     | some mctx =>
