From 709f22c8e41fd428be588398836dbd1755cfced6 Mon Sep 17 00:00:00 2001
From: Leonardo de Moura <leonardo@microsoft.com>
Date: Mon, 11 Jul 2022 17:49:49 -0700
Subject: [PATCH] feat: add field `dependsOnHigherOrderOutParam` to `ParamInfo`

See issue #1299
---
 src/Lean/Meta/Basic.lean   | 14 ++++++++++++++
 src/Lean/Meta/FunInfo.lean | 27 +++++++++++++++++++++------
 2 files changed, 35 insertions(+), 6 deletions(-)

diff --git a/src/Lean/Meta/Basic.lean b/src/Lean/Meta/Basic.lean
index 65f1957605..eaaf163f1b 100644
--- a/src/Lean/Meta/Basic.lean
+++ b/src/Lean/Meta/Basic.lean
@@ -136,6 +136,20 @@ structure ParamInfo where
     This information affects the generation of congruence theorems.
   -/
   isDecInst      : Bool       := false
+  /--
+    `dependsOnHigherOrderOutParam` is true if the type of this parameter depends on
+    the higher-order output parameter of a previour local instance.
+    Example:
+    ```
+    getElem :
+      {Cont : Type u_1} → {Idx : Type u_2} → {Elem : Type u_3} →
+      {Dom : Cont → Idx → Prop} → [self : GetElem Cont Idx Elem Dom] →
+      (xs : Cont) → (i : Idx) → Dom xs i → Elem
+    ```
+    This flag is true for the parameter with type `Dom xs i` since `Dom` is an output parameter
+    of the instance `[self : GetElem Cont Idx Elem Dom]`
+  -/
+  dependsOnHigherOrderOutParam : Bool := false
   deriving Inhabited
 
 def ParamInfo.isImplicit (p : ParamInfo) : Bool :=
diff --git a/src/Lean/Meta/FunInfo.lean b/src/Lean/Meta/FunInfo.lean
index e1a37d3b54..1e51ff5936 100644
--- a/src/Lean/Meta/FunInfo.lean
+++ b/src/Lean/Meta/FunInfo.lean
@@ -9,9 +9,8 @@ import Lean.Meta.InferType
 namespace Lean.Meta
 
 @[inline] private def checkFunInfoCache (fn : Expr) (maxArgs? : Option Nat) (k : MetaM FunInfo) : MetaM FunInfo := do
-  let s ← get
   let t ← getTransparency
-  match s.cache.funInfo.find? ⟨t, fn, maxArgs?⟩ with
+  match (← get).cache.funInfo.find? ⟨t, fn, maxArgs?⟩ with
   | some finfo => pure finfo
   | none       => do
     let finfo ← k
@@ -58,17 +57,33 @@ private def getFunInfoAux (fn : Expr) (maxArgs? : Option Nat) : MetaM FunInfo :=
     withTransparency TransparencyMode.default do
       forallBoundedTelescope fnType maxArgs? fun fvars type => do
         let mut pinfo := #[]
+        let mut higherOrderOutParams : FVarIdSet := {}
         for i in [:fvars.size] do
           let fvar := fvars[i]!
           let decl ← getFVarLocalDecl fvar
           let backDeps := collectDeps fvars decl.type
+          let dependsOnHigherOrderOutParam :=
+            !higherOrderOutParams.isEmpty
+            && Option.isSome (decl.type.find? fun e => e.isFVar && higherOrderOutParams.contains e.fvarId!)
           pinfo := updateHasFwdDeps pinfo backDeps
           pinfo := pinfo.push {
-            backDeps     := backDeps
-            binderInfo   := decl.binderInfo
-            isProp       := (← isProp decl.type)
-            isDecInst    := (← forallTelescopeReducing decl.type fun _ type => return type.isAppOf ``Decidable)
+            backDeps, dependsOnHigherOrderOutParam
+            binderInfo := decl.binderInfo
+            isProp     := (← isProp decl.type)
+            isDecInst  := (← forallTelescopeReducing decl.type fun _ type => return type.isAppOf ``Decidable)
           }
+          if decl.binderInfo == .instImplicit then
+            /- Collect higher order output parameters of this class -/
+            if let some className ← isClass? decl.type then
+              if let some outParamPositions := getOutParamPositions? (← getEnv) className then
+                unless outParamPositions.isEmpty do
+                  let args := decl.type.getAppArgs
+                  for i in [:args.size] do
+                    if outParamPositions.contains i then
+                      let arg := args[i]!
+                      if fvars.contains arg then
+                        if (← whnf (← inferType arg)).isForall then
+                          higherOrderOutParams := higherOrderOutParams.insert arg.fvarId!
         let resultDeps := collectDeps fvars type
         pinfo := updateHasFwdDeps pinfo resultDeps
         pure { resultDeps := resultDeps, paramInfo := pinfo }