Implement Resolver.asMemberOf for Properties and Functions (#110)

* Initial implementation and test setup for AsMemberOf

This is just a copy of the implementation in Room.
Still need more tests, support for other declarations
and possible optimizations

* use compiler utils for substitution

* AsMemberOf for function, initial setup

* Use new type substitutor

Old type substitutor seems to have problems with * projection
where as new one seems to work (failing test case: receiveArgs method).

I've replaced the code to use the new substitutor (even with fields) and
also implemented a fallback implementation for KSFunctionType that derives
from a KSFunctionDeclaration.

* add more documentation and support for java

* code cleanup pre-review

* update java to ksp reference docs

* Throw errors for illegal arguments, rename KSFunctionType to KSFunction

* remove unused imports from as member of test

* cache as member of calls

* implement equals in KSFunction

api/src/main/kotlin/com/google/devtools/ksp/processing/Resolver.kt

@@ -83,4 +83,62 @@ interface Resolver {
      * Calling [overrides] is expensive and should be avoided if possible.
      */
     fun overrides(overrider: KSDeclaration, overridee: KSDeclaration): Boolean
-}
+
+    /**
+     * Returns the type of the [property] when it is viewed as member of the [containing] type.
+     *
+     * For instance, for the following input:
+     * ```
+     * class Base<T>(val x:T)
+     * val foo: Base<Int>
+     * val bar: Base<String>
+     * ```
+     * When `x` is viewed as member of `foo`, this method will return the [KSType] for `Int`
+     * whereas when `x` is viewed as member of `bar`, this method will return the [KSType]
+     * representing `String`.
+     *
+     * If the substitution fails (e.g. if [containing] is an error type, a [KSType] with [KSType.isError] `true` is
+     * returned.
+     *
+     * @param property The property whose type will be returned
+     * @param containing The type that contains [property]
+     * @throws IllegalArgumentException Throws [IllegalArgumentException] when [containing] does not contain
+     * [property] or if the [property] is not declared in a class, object or interface.
+     */
+    fun asMemberOf(
+        property: KSPropertyDeclaration,
+        containing: KSType
+    ): KSType
+
+    /**
+     * Returns the type of the [function] when it is viewed as member of the [containing] type.
+     *
+     * For instance, for the following input:
+     * ```
+     * interface Base<T> {
+     *   fun f(t:T?):T
+     * }
+     * val foo: Base<Int>
+     * val bar: Base<String>
+     * ```
+     * When `f()` is viewed as member of `foo`, this method will return a [KSFunction] where
+     * the [KSFunction.returnType] is `Int` and the parameter `t` is of type `Int?`.
+     * When `f()` is viewed as member of `bar`, this method will return a [KSFunction]
+     * where the [KSFunction.returnType] is `String` and the parameter `t` is of type `String?`.
+     *
+     * If the function has type parameters, they'll not be resolved and can be read from
+     * [KSFunction.typeParameters].
+     *
+     * If the substitution fails (e.g. if [containing] is an error type, a [KSFunction] with [KSFunction.isError] `true`
+     * is returned.
+     *
+     * @param function The function whose types will be resolved
+     * @param containing The type that contains [function].
+     * @throws IllegalArgumentException Throws [IllegalArgumentException] when [containing] does not contain [function]
+     * or if the [function] is not declared in a class, object or interface.
+     */
+    fun asMemberOf(
+        function: KSFunctionDeclaration,
+        containing: KSType
+    ): KSFunction
+}

api/src/main/kotlin/com/google/devtools/ksp/symbol/KSFunction.kt

@@ -0,0 +1,60 @@
+/*
+ * Copyright 2020 Google LLC
+ * Copyright 2010-2020 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.google.devtools.ksp.symbol
+import com.google.devtools.ksp.processing.Resolver
+/**
+ * Holds the information for a [KSFunctionDeclaration] where type arguments are resolved as member
+ * of a specific [KSType].
+ *
+ * @see Resolver.asMemberOf
+ */
+interface KSFunction {
+    /**
+     * The return type of the function. Note that this might be `null` if an error happened when
+     * the type is resolved.
+     *
+     * @see KSFunctionDeclaration.returnType
+     */
+    val returnType: KSType?
+
+    /**
+     * The types of the value parameters of the function. Note that this list might have `null`
+     * values in it if the type of a parameter could not be resolved.
+     *
+     * @see KSFunctionDeclaration.parameters
+     */
+    val parameterTypes: List<KSType?>
+
+    /**
+     * The type parameters of the function.
+     *
+     * @see KSFunctionDeclaration.typeParameters
+     */
+    val typeParameters: List<KSTypeParameter>
+
+    /**
+     * The receiver type of the function.
+     *
+     * @see KSFunctionDeclaration.extensionReceiver
+     */
+    val extensionReceiverType: KSType?
+
+    /**
+     * True if the compiler couldn't resolve the function.
+     */
+    val isError: Boolean
+}

api/src/main/kotlin/com/google/devtools/ksp/symbol/KSFunctionDeclaration.kt

@@ -17,11 +17,16 @@
 
 
 package com.google.devtools.ksp.symbol
-
+import com.google.devtools.ksp.processing.Resolver
 /**
  * A function definition
  *
  * Dispatch receiver can be obtained through [parentDeclaration].
+ *
+ * To obtain the function signature where type arguments are resolved as member of a given [KSType],
+ * use [Resolver.asMemberOf].
+ *
+ * @see KSFunctionType
  */
 interface KSFunctionDeclaration : KSDeclaration, KSDeclarationContainer {
     /**
@@ -57,4 +62,4 @@ interface KSFunctionDeclaration : KSDeclaration, KSDeclarationContainer {
      * Calling [findOverridee] is expensive and should be avoided if possible.
      */
     fun findOverridee(): KSFunctionDeclaration?
-}
+}

compiler-plugin/src/main/kotlin/com/google/devtools/ksp/processing/impl/ResolverImpl.kt

@@ -31,6 +31,7 @@ import org.jetbrains.kotlin.container.get
 import org.jetbrains.kotlin.descriptors.*
 import org.jetbrains.kotlin.descriptors.annotations.AnnotationDescriptor
 import org.jetbrains.kotlin.incremental.components.NoLookupLocation
+import com.google.devtools.ksp.closestClassDeclaration
 import com.google.devtools.ksp.processing.KSBuiltIns
 import com.google.devtools.ksp.processing.Resolver
 import com.google.devtools.ksp.symbol.*
@@ -50,13 +51,17 @@ import org.jetbrains.kotlin.load.java.lazy.descriptors.LazyJavaTypeParameterDesc
 import org.jetbrains.kotlin.load.java.lazy.types.JavaTypeResolver
 import org.jetbrains.kotlin.load.java.lazy.types.toAttributes
 import org.jetbrains.kotlin.load.java.structure.impl.JavaClassImpl
+import org.jetbrains.kotlin.load.java.structure.impl.JavaFieldImpl
 import org.jetbrains.kotlin.load.java.structure.impl.JavaMethodImpl
 import org.jetbrains.kotlin.load.java.structure.impl.JavaTypeImpl
 import org.jetbrains.kotlin.load.java.structure.impl.JavaTypeParameterImpl
 import org.jetbrains.kotlin.metadata.jvm.deserialization.JvmProtoBufUtil
 import org.jetbrains.kotlin.name.FqName
 import org.jetbrains.kotlin.psi.*
 import org.jetbrains.kotlin.resolve.*
+import org.jetbrains.kotlin.resolve.calls.inference.components.NewTypeSubstitutor
+import org.jetbrains.kotlin.resolve.calls.inference.components.composeWith
+import org.jetbrains.kotlin.resolve.calls.inference.substitute
 import org.jetbrains.kotlin.resolve.calls.smartcasts.DataFlowInfo
 import org.jetbrains.kotlin.resolve.constants.ConstantValue
 import org.jetbrains.kotlin.resolve.constants.evaluate.ConstantExpressionEvaluator
@@ -87,6 +92,11 @@ class ResolverImpl(
     val javaActualAnnotationArgumentExtractor = JavaActualAnnotationArgumentExtractor()
     private val nameToKSMap: MutableMap<KSName, KSClassDeclaration>
 
+    /**
+     * Checking as member of is an expensive operation, hence we cache result values in this map.
+     */
+    private val functionAsMemberOfCache: MutableMap<Pair<KSFunctionDeclaration, KSType>, KSFunction>
+
     private val typeMapper = KotlinTypeMapper(
         BindingContext.EMPTY, ClassBuilderMode.LIGHT_CLASSES,
         JvmProtoBufUtil.DEFAULT_MODULE_NAME,
@@ -119,10 +129,10 @@ class ResolverImpl(
         lazyJavaResolverContext = LazyJavaResolverContext(javaResolverComponents, TypeParameterResolver.EMPTY) { null }
         javaTypeResolver = lazyJavaResolverContext.typeResolver
         moduleClassResolver = lazyJavaResolverContext.components.moduleClassResolver
-
         instance = this
 
         nameToKSMap = mutableMapOf()
+        functionAsMemberOfCache = mutableMapOf()
 
         val visitor = object : KSVisitorVoid() {
             override fun visitFile(file: KSFile, data: Unit) {
@@ -286,7 +296,7 @@ class ResolverImpl(
         }
     }
 
-    // TODO: Resolve Java fields/variables is not supported by this function. Not needed currently.
+    // TODO: Resolve Java variables is not supported by this function. Not needed currently.
     fun resolveJavaDeclaration(psi: PsiElement): DeclarationDescriptor? {
         return when (psi) {
             is PsiClass -> moduleClassResolver.resolveClass(JavaClassImpl(psi))
@@ -302,6 +312,8 @@ class ResolverImpl(
                             ?.unsubstitutedMemberScope!!.getDescriptorsFiltered().single { it.findPsi() == psi }
                 }
             }
+            is PsiField -> moduleClassResolver.resolveClass(JavaFieldImpl(psi).containingClass)
+                ?.unsubstitutedMemberScope!!.getDescriptorsFiltered().single { it.findPsi() == psi } as CallableMemberDescriptor
             else -> throw IllegalStateException("unhandled psi element kind: ${psi.javaClass}")
         }
     }
@@ -330,7 +342,7 @@ class ResolverImpl(
             is KSPropertyDeclarationImpl -> resolveDeclaration(property.ktProperty)
             is KSPropertyDeclarationParameterImpl -> resolveDeclaration(property.ktParameter)
             is KSPropertyDeclarationDescriptorImpl -> property.descriptor
-            is KSPropertyDeclarationJavaImpl -> throw IllegalStateException("should not invoke resolve on Java fields")
+            is KSPropertyDeclarationJavaImpl -> resolveJavaDeclaration(property.psi)
             else -> throw IllegalStateException("unexpected class: ${property.javaClass}")
         } as PropertyDescriptor?
     }
@@ -467,6 +479,76 @@ class ResolverImpl(
         return KSTypeArgumentLiteImpl.getCached(typeRef, variance)
     }
 
+    override fun asMemberOf(
+        property: KSPropertyDeclaration,
+        containing: KSType
+    ): KSType {
+        val propertyDeclaredIn = property.closestClassDeclaration()
+            ?: throw IllegalArgumentException("Cannot call asMemberOf with a property that is " +
+                "not declared in a class or an interface")
+        val declaration = resolvePropertyDeclaration(property)
+        if (declaration != null && containing is KSTypeImpl && !containing.isError) {
+            if (!containing.kotlinType.isSubtypeOf(propertyDeclaredIn)) {
+                throw IllegalArgumentException(
+                    "$containing is not a sub type of the class/interface that contains `$property` ($propertyDeclaredIn)"
+                )
+            }
+            val typeSubstitutor = containing.kotlinType.createTypeSubstitutor()
+            val substituted = declaration.substitute(typeSubstitutor) as? ValueDescriptor
+            substituted?.let {
+                return KSTypeImpl.getCached(substituted.type)
+            }
+        }
+        // if substitution fails, fallback to the type from the property
+        return KSErrorType
+    }
+
+    override fun asMemberOf(
+        function: KSFunctionDeclaration,
+        containing: KSType
+    ): KSFunction {
+        val key = function to containing
+        return functionAsMemberOfCache.getOrPut(key) {
+            computeAsMemberOf(function, containing)
+        }
+    }
+
+    private fun computeAsMemberOf(
+        function: KSFunctionDeclaration,
+        containing: KSType
+    ) : KSFunction {
+        val functionDeclaredIn = function.closestClassDeclaration()
+            ?: throw IllegalArgumentException("Cannot call asMemberOf with a function that is " +
+                "not declared in a class or an interface")
+        val declaration = resolveFunctionDeclaration(function)
+        if (declaration != null && containing is KSTypeImpl && !containing.isError) {
+            if (!containing.kotlinType.isSubtypeOf(functionDeclaredIn)) {
+                throw IllegalArgumentException(
+                    "$containing is not a sub type of the class/interface that contains " +
+                        "`$function` ($functionDeclaredIn)"
+                )
+            }
+            val typeSubstitutor = containing.kotlinType.createTypeSubstitutor()
+            val substituted = declaration.substitute(typeSubstitutor)
+            return KSFunctionImpl(substituted)
+        }
+        // if substitution fails, return an error function that resembles the original declaration
+        return KSFunctionErrorImpl(function)
+    }
+
+    private fun KotlinType.isSubtypeOf(declaration: KSClassDeclaration): Boolean {
+        val classDeclaration = resolveClassDeclaration(declaration)
+        if (classDeclaration == null) {
+            throw IllegalArgumentException(
+                "Cannot find the declaration for class $classDeclaration"
+            )
+        }
+        return constructor
+            .declarationDescriptor
+            ?.getAllSuperClassifiers()
+            ?.any { it == classDeclaration } == true
+    }
+
     override val builtIns: KSBuiltIns by lazy {
         val builtIns = module.builtIns
         object : KSBuiltIns {
@@ -535,3 +617,15 @@ fun MemberDescriptor.toKSDeclaration(): KSDeclaration =
             else -> throw IllegalStateException("Unknown expect/actual implementation")
         }
     }
+
+/**
+ * [NewTypeSubstitutor] handles variance better than [TypeSubstitutor] so we use it when subtituting
+ * types in [ResolverImpl.asMemberOf] implementations.
+ */
+private fun TypeSubstitutor.toNewSubstitutor() = composeWith(
+    org.jetbrains.kotlin.resolve.calls.inference.components.EmptySubstitutor
+)
+
+private fun KotlinType.createTypeSubstitutor(): NewTypeSubstitutor {
+    return SubstitutionUtils.buildDeepSubstitutor(this).toNewSubstitutor()
+}