asMemberOf() should always substitute type parameters for methods using the relevant substitutor

(cherry picked from commit dfadd17)

(cherry picked from commit 49b07a7)

plugins/annotation-processing/src/org/jetbrains/kotlin/annotation/processing/impl/KotlinTypes.kt

@@ -19,6 +19,7 @@ package org.jetbrains.kotlin.annotation.processing.impl
 import com.intellij.psi.*
 import com.intellij.psi.search.GlobalSearchScope
 import com.intellij.psi.util.*
+import org.jetbrains.kotlin.java.model.JeElement
 import org.jetbrains.kotlin.java.model.elements.JeClassInitializerExecutableElement
 import org.jetbrains.kotlin.java.model.elements.JeMethodExecutableElement
 import org.jetbrains.kotlin.java.model.elements.JeTypeElement
@@ -219,27 +220,33 @@ class KotlinTypes(val javaPsiFacade: JavaPsiFacade, val psiManager: PsiManager,
         return JeDeclaredType(psiType, psiClass, containing)
     }
 
-    override fun asMemberOf(containing: DeclaredType, element: Element): TypeMirror {
-        val substitutor = when (containing) {
-            is JeDeclaredType -> {
-                val result = containing.psiType.resolveGenerics()
-                if (result.isValidResult) result.substitutor else PsiSubstitutor.EMPTY
-            }
-            else -> throw IllegalArgumentException("Invalid containing type: $containing")
+    private fun Array<out PsiType>.findSuperType(superTypeClass: PsiClass): PsiClassType? {
+        for (supertype in this) {
+            if (supertype is PsiClassType && supertype.resolve() == superTypeClass) return supertype
+            supertype.superTypes.findSuperType(superTypeClass)?.let { return it }
         }
-        
+        return null
+    }
+
+    override fun asMemberOf(containing: DeclaredType, element: Element): TypeMirror {
+        if (containing !is JeDeclaredType || element is JeClassInitializerExecutableElement) return element.asType()
+        val containingType = containing.psiType
+
+        val member = (element as JeElement).psi as? PsiMember ?: return element.asType()
+        val methodContainingClass = member.containingClass ?: return element.asType()
+
+        val relevantSuperType = containingType.superTypes.findSuperType(methodContainingClass) ?: return element.asType()
+        val resolveResult = relevantSuperType.resolveGenerics()
+        if (!resolveResult.isValidResult) return element.asType()
+        val substitutor = resolveResult.substitutor
+
         return when (element) {
             is JeMethodExecutableElement -> {
                 val method = element.psi
-                if (method.hasModifierProperty(PsiModifier.STATIC) || !method.hasTypeParameters()) {
-                    JeMethodExecutableTypeMirror(method)
-                } else {
-                    val signature = method.getSignature(substitutor)
-                    val returnType = substitutor.substitute(element.psi.returnType)
-                    JeMethodExecutableTypeMirror(method, signature, returnType)
-                }
+                val signature = method.getSignature(substitutor)
+                val returnType = substitutor.substitute(element.psi.returnType)
+                JeMethodExecutableTypeMirror(method, signature, returnType)
             }
-            is JeClassInitializerExecutableElement -> element.asType()
             is JeVariableElement -> substitutor.substitute(element.psi.type).toJeType(psiManager)
             else -> throw IllegalArgumentException("Invalid element type: $element")
         }

plugins/annotation-processing/testData/processors/AsMemberOf.kt

@@ -0,0 +1,15 @@
+open class Base<T> {
+    @JvmField
+    val f: T = null!!
+
+    fun m(t: T): T = null!!
+}
+
+class Impl<T> : Base<T>()
+
+annotation class Anno
+
+@Anno
+class Test {
+    val f = Impl<String>()
+}

plugins/plugins-tests/tests/org/jetbrains/kotlin/annotation/processing/test/processor/AbstractProcessorTest.kt

@@ -25,8 +25,7 @@ import org.jetbrains.kotlin.incremental.SourceRetentionAnnotationHandlerImpl
 import org.jetbrains.kotlin.java.model.elements.JeAnnotationMirror
 import org.jetbrains.kotlin.java.model.elements.JeMethodExecutableElement
 import org.jetbrains.kotlin.java.model.elements.JeTypeElement
-import org.jetbrains.kotlin.load.kotlin.incremental.components.IncrementalCompilationComponents
-import org.jetbrains.kotlin.modules.TargetId
+import org.jetbrains.kotlin.java.model.elements.JeVariableElement
 import org.jetbrains.kotlin.resolve.jvm.extensions.AnalysisCompletedHandlerExtension
 import org.jetbrains.kotlin.test.ConfigurationKind
 import org.jetbrains.kotlin.test.KotlinTestUtils
@@ -152,8 +151,12 @@ abstract class AbstractProcessorTest : AbstractBytecodeTextTest() {
             parameterTypes.zip(it.parameters).all { it.first == it.second.asType().toString() }
         } as JeMethodExecutableElement
     }
-    
-    protected fun ProcessingEnvironment.findClass(fqName: String) = elementUtils.getTypeElement(fqName) as JeTypeElement 
+
+    protected fun TypeElement.findField(name: String): JeVariableElement {
+        return enclosedElements.first { it is JeVariableElement && it.simpleName.toString() == name } as JeVariableElement
+    }
+
+    protected fun ProcessingEnvironment.findClass(fqName: String) = elementUtils.getTypeElement(fqName) as JeTypeElement
 
     protected fun assertEquals(expected: String, actual: Name) = assertEquals(expected, actual.toString())
 }

plugins/plugins-tests/tests/org/jetbrains/kotlin/annotation/processing/test/processor/ProcessorTests.kt

@@ -24,6 +24,7 @@ import org.jetbrains.kotlin.java.model.types.JeMethodExecutableTypeMirror
 import org.jetbrains.kotlin.resolve.jvm.extensions.AnalysisCompletedHandlerExtension
 import org.jetbrains.kotlin.utils.addToStdlib.firstIsInstance
 import javax.lang.model.element.AnnotationMirror
+import javax.lang.model.element.Element
 import javax.lang.model.type.DeclaredType
 import javax.lang.model.type.TypeMirror
 import javax.lang.model.type.TypeVariable
@@ -249,4 +250,23 @@ class ProcessorTests : AbstractProcessorTest() {
         check(true, "TestTrue")
         check(false, "TestFalse")
     }
+
+    fun testAsMemberOf() = test("AsMemberOf", "*") { set, roundEnv, env ->
+        val f = env.findClass("Test").findField("f")
+        val fType = f.asType() as JeDeclaredType
+
+        val base = env.findClass("Base")
+        val baseF = base.findField("f")
+        val baseM = base.findMethod("m", "T")
+
+        fun check(element: Element, expectedTypeSignature: String) {
+            assertEquals(expectedTypeSignature, env.typeUtils.asMemberOf(fType, element).toString())
+        }
+
+        assertEquals("(T)T", baseM.asType().toString())
+        check(baseM, "(java.lang.String)java.lang.String")
+
+        assertEquals("T", baseF.asType().toString())
+        check(baseF, "java.lang.String")
+    }
 }