From 9355fa6cccbc89460fcb0a6664ce7248ba88dedf Mon Sep 17 00:00:00 2001
From: Daniil Ovchinnikov <daniil.ovchinnikov@jetbrains.com>
Date: Wed, 7 Nov 2018 16:43:08 +0300
Subject: [PATCH] [groovy] inference: store nested sessions by resolve result
 instead of reference expression

---
 .../lang/resolve/MethodResolveResult.kt       |  2 +-
 .../inference/GroovyInferenceSession.kt       | 36 ++++++++++++-------
 .../ReferenceExpressionConstraint.kt          |  2 +-
 3 files changed, 25 insertions(+), 15 deletions(-)

diff --git a/plugins/groovy/groovy-psi/src/org/jetbrains/plugins/groovy/lang/resolve/MethodResolveResult.kt b/plugins/groovy/groovy-psi/src/org/jetbrains/plugins/groovy/lang/resolve/MethodResolveResult.kt
index a36478ac0592..0cf1d9093f2d 100644
--- a/plugins/groovy/groovy-psi/src/org/jetbrains/plugins/groovy/lang/resolve/MethodResolveResult.kt
+++ b/plugins/groovy/groovy-psi/src/org/jetbrains/plugins/groovy/lang/resolve/MethodResolveResult.kt
@@ -53,7 +53,7 @@ class MethodResolveResult(
         .addReturnConstraint()
         .resolveMode(false)
         .startFromTop(true)
-        .build().inferSubst(ref)
+        .build().inferSubst(this)
     }
   }
 
diff --git a/plugins/groovy/groovy-psi/src/org/jetbrains/plugins/groovy/lang/resolve/processors/inference/GroovyInferenceSession.kt b/plugins/groovy/groovy-psi/src/org/jetbrains/plugins/groovy/lang/resolve/processors/inference/GroovyInferenceSession.kt
index 2579ad416abf..89421392e4c2 100644
--- a/plugins/groovy/groovy-psi/src/org/jetbrains/plugins/groovy/lang/resolve/processors/inference/GroovyInferenceSession.kt
+++ b/plugins/groovy/groovy-psi/src/org/jetbrains/plugins/groovy/lang/resolve/processors/inference/GroovyInferenceSession.kt
@@ -5,18 +5,20 @@ import com.intellij.psi.PsiElement
 import com.intellij.psi.PsiSubstitutor
 import com.intellij.psi.PsiTypeParameter
 import com.intellij.psi.impl.source.resolve.graphInference.InferenceSession
+import org.jetbrains.plugins.groovy.lang.psi.api.GroovyResolveResult
 import org.jetbrains.plugins.groovy.lang.psi.api.statements.expressions.GrMethodCall
-import org.jetbrains.plugins.groovy.lang.psi.api.statements.expressions.GrReferenceExpression
 
-class GroovyInferenceSession(typeParams: Array<PsiTypeParameter>,
-                             val siteSubstitutor: PsiSubstitutor,
-                             context: PsiElement,
-                             val closureSkipList: List<GrMethodCall> = emptyList(),
-                             val skipClosureBlock: Boolean = true) : InferenceSession(typeParams, siteSubstitutor, context.manager, context) {
+class GroovyInferenceSession(
+  typeParams: Array<PsiTypeParameter>,
+  val siteSubstitutor: PsiSubstitutor,
+  context: PsiElement,
+  val closureSkipList: List<GrMethodCall> = emptyList(),
+  val skipClosureBlock: Boolean = true
+) : InferenceSession(typeParams, siteSubstitutor, context.manager, context) {
 
-  val myNestedSessions = mutableMapOf<GrReferenceExpression, GroovyInferenceSession>()
+  val nestedSessions = mutableMapOf<GroovyResolveResult, GroovyInferenceSession>()
 
-  fun result(): PsiSubstitutor {
+  private fun result(): PsiSubstitutor {
     resolveBounds(myInferenceVariables, siteSubstitutor)
     return prepareSubstitution()
   }
@@ -26,15 +28,23 @@ class GroovyInferenceSession(typeParams: Array<PsiTypeParameter>,
     return result()
   }
 
-  fun inferSubst(ref: GrReferenceExpression): PsiSubstitutor {
+  fun inferSubst(result: GroovyResolveResult): PsiSubstitutor {
     repeatInferencePhases()
-    findSession(ref)?.let { return it.result() }
+    findSession(result)?.let {
+      return it.result()
+    }
     return PsiSubstitutor.EMPTY
   }
 
-  private fun findSession(ref: GrReferenceExpression): GroovyInferenceSession? {
-    myNestedSessions[ref]?.let { return it }
-    myNestedSessions.values.forEach { nested -> nested.findSession(ref)?.let { return it }}
+  private fun findSession(result: GroovyResolveResult): GroovyInferenceSession? {
+    nestedSessions[result]?.let {
+      return it
+    }
+    for (nested in nestedSessions.values) {
+      nested.findSession(result)?.let {
+        return it
+      }
+    }
     return null
   }
 }
diff --git a/plugins/groovy/groovy-psi/src/org/jetbrains/plugins/groovy/lang/resolve/processors/inference/ReferenceExpressionConstraint.kt b/plugins/groovy/groovy-psi/src/org/jetbrains/plugins/groovy/lang/resolve/processors/inference/ReferenceExpressionConstraint.kt
index f0bac2094d5d..e3ab2b47bde4 100644
--- a/plugins/groovy/groovy-psi/src/org/jetbrains/plugins/groovy/lang/resolve/processors/inference/ReferenceExpressionConstraint.kt
+++ b/plugins/groovy/groovy-psi/src/org/jetbrains/plugins/groovy/lang/resolve/processors/inference/ReferenceExpressionConstraint.kt
@@ -19,7 +19,7 @@ class ReferenceExpressionConstraint(private val callRef: GrReferenceExpression,
     resolved.candidate?.let {
       val typeParameters = it.method.typeParameters
       val nestedSession = GroovyInferenceSession(typeParameters, it.siteSubstitutor, callRef, emptyList(), session.skipClosureBlock)
-      session.myNestedSessions[callRef] = nestedSession
+      session.nestedSessions[resolved] = nestedSession
       nestedSession.propagateVariables(session)
 
       nestedSession.addConstraint(MethodCallConstraint(callRef, it))