Analysis API: introduce KtSignature

analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/KtSymbolByFirBuilder.kt

@@ -9,7 +9,10 @@ import com.intellij.openapi.project.Project
 import com.intellij.psi.JavaPsiFacade
 import com.intellij.psi.search.GlobalSearchScope
 import com.intellij.util.containers.ContainerUtil
-import org.jetbrains.kotlin.analysis.api.*
+import org.jetbrains.kotlin.analysis.api.KtStarProjectionTypeArgument
+import org.jetbrains.kotlin.analysis.api.KtTypeArgument
+import org.jetbrains.kotlin.analysis.api.KtTypeArgumentWithVariance
+import org.jetbrains.kotlin.analysis.api.ValidityTokenOwner
 import org.jetbrains.kotlin.analysis.api.fir.symbols.*
 import org.jetbrains.kotlin.analysis.api.fir.types.*
 import org.jetbrains.kotlin.analysis.api.fir.utils.weakRef
@@ -34,6 +37,7 @@ import org.jetbrains.kotlin.fir.resolve.substitution.ConeSubstitutor
 import org.jetbrains.kotlin.fir.resolve.substitution.ConeSubstitutorByMap
 import org.jetbrains.kotlin.fir.symbols.ConeClassLikeLookupTag
 import org.jetbrains.kotlin.fir.symbols.ConeTypeParameterLookupTag
+import org.jetbrains.kotlin.fir.symbols.ensureResolved
 import org.jetbrains.kotlin.fir.symbols.impl.FirClassifierSymbol
 import org.jetbrains.kotlin.fir.symbols.impl.FirTypeParameterSymbol
 import org.jetbrains.kotlin.fir.types.*
@@ -207,6 +211,12 @@ internal class KtSymbolByFirBuilder private constructor(
             }
         }
 
+        fun buildFunctionLikeSignature(fir: FirFunction): KtFunctionLikeSignature<KtFunctionLikeSymbol> {
+            if (fir is FirSimpleFunction && fir.origin != FirDeclarationOrigin.SamConstructor)
+                return buildFunctionSignature(fir)
+            return buildFunctionLikeSymbol(fir).toSignature()
+        }
+
         fun buildFunctionSymbol(fir: FirSimpleFunction): KtFirFunctionSymbol {
             fir.unwrapSubstitutionOverrideIfNeeded()?.let {
                 return buildFunctionSymbol(it)
@@ -216,6 +226,23 @@ internal class KtSymbolByFirBuilder private constructor(
             return symbolsCache.cache(fir) { KtFirFunctionSymbol(fir, resolveState, token, this@KtSymbolByFirBuilder) }
         }
 
+        fun buildFunctionSignature(fir: FirSimpleFunction): KtFunctionLikeSignature<KtFirFunctionSymbol> {
+            fir.symbol.ensureResolved(FirResolvePhase.IMPLICIT_TYPES_BODY_RESOLVE)
+            val functionSymbol = buildFunctionSymbol(fir)
+            return KtFunctionLikeSignature(
+                functionSymbol,
+                typeBuilder.buildKtType(fir.returnTypeRef),
+                fir.receiverTypeRef?.let { typeBuilder.buildKtType(it) },
+                functionSymbol.valueParameters.zip(fir.valueParameters).map { (ktSymbol, fir) ->
+                    var type = fir.returnTypeRef.coneType
+                    if (fir.isVararg) {
+                        type = type.arrayElementType() ?: type
+                    }
+                    KtVariableLikeSignature(ktSymbol, typeBuilder.buildKtType(type), null)
+                }
+            )
+        }
+
         fun buildAnonymousFunctionSymbol(fir: FirAnonymousFunction): KtFirAnonymousFunctionSymbol {
             return symbolsCache.cache(fir) { KtFirAnonymousFunctionSymbol(fir, resolveState, token, this@KtSymbolByFirBuilder) }
         }
@@ -256,6 +283,11 @@ internal class KtSymbolByFirBuilder private constructor(
             }
         }
 
+        fun buildVariableLikeSignature(fir: FirVariable): KtVariableLikeSignature<KtVariableLikeSymbol> {
+            if (fir is FirProperty && !fir.isLocal && fir !is FirSyntheticProperty) return buildPropertySignature(fir)
+            return buildVariableLikeSymbol(fir).toSignature()
+        }
+
         fun buildVariableSymbol(fir: FirProperty): KtVariableSymbol {
             return when {
                 fir.isLocal -> buildLocalVariableSymbol(fir)
@@ -277,6 +309,15 @@ internal class KtSymbolByFirBuilder private constructor(
             }
         }
 
+        fun buildPropertySignature(fir: FirProperty): KtVariableLikeSignature<KtVariableSymbol> {
+            fir.symbol.ensureResolved(FirResolvePhase.IMPLICIT_TYPES_BODY_RESOLVE)
+            return KtVariableLikeSignature(
+                buildPropertySymbol(fir),
+                typeBuilder.buildKtType(fir.returnTypeRef),
+                fir.receiverTypeRef?.let { typeBuilder.buildKtType(it) }
+            )
+        }
+
         fun buildLocalVariableSymbol(fir: FirProperty): KtFirLocalVariableSymbol {
             checkRequirementForBuildingSymbol<KtFirLocalVariableSymbol>(fir, fir.isLocal)
             return symbolsCache.cache(fir) {
@@ -331,6 +372,15 @@ internal class KtSymbolByFirBuilder private constructor(
             }
         }
 
+        fun buildCallableSignature(fir: FirCallableDeclaration): KtSignature<KtCallableSymbol> {
+            return when (fir) {
+                is FirPropertyAccessor -> buildPropertyAccessorSymbol(fir).toSignature()
+                is FirFunction -> functionLikeBuilder.buildFunctionLikeSignature(fir)
+                is FirVariable -> variableLikeBuilder.buildVariableLikeSignature(fir)
+                else -> throwUnexpectedElementError(fir)
+            }
+        }
+
 
         fun buildPropertyAccessorSymbol(fir: FirPropertyAccessor): KtPropertyAccessorSymbol {
             return when {

analysis/analysis-api/src/org/jetbrains/kotlin/analysis/api/symbols/KtSignature.kt

@@ -0,0 +1,130 @@
+/*
+ * Copyright 2010-2021 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
+ */
+
+package org.jetbrains.kotlin.analysis.api.symbols
+
+import org.jetbrains.kotlin.analysis.api.ValidityTokenOwner
+import org.jetbrains.kotlin.analysis.api.tokens.ValidityToken
+import org.jetbrains.kotlin.analysis.api.types.KtSubstitutor
+import org.jetbrains.kotlin.analysis.api.types.KtType
+import org.jetbrains.kotlin.analysis.api.withValidityAssertion
+
+/**
+ * A signature for a callable symbol. Comparing to a `KtCallableSymbol`, a signature can carry use-site type information. For example
+ * ```
+ * fun test(l: List<String>) {
+ *   l.get(1) // The symbol `get` has type `(Int) -> T` where is the type parameter declared in `List`. On the other hand, a KtSignature
+ *            // carries instantiated type information `(Int) -> String`.
+ * }
+ * ```
+ *
+ * Equality of [KtSignature] is derived from its content.
+ */
+public sealed class KtSignature<out S : KtCallableSymbol> : ValidityTokenOwner {
+    /**
+     * The original symbol for this signature.
+     */
+    public abstract val symbol: S
+
+    /**
+     * The use-site-substituted return type.
+     */
+    public abstract val returnType: KtType
+
+    /**
+     * The use-site-substituted extension receiver type.
+     */
+    public abstract val receiverType: KtType?
+
+    abstract override fun equals(other: Any?): Boolean
+    abstract override fun hashCode(): Int
+}
+
+/**
+ * A signature of a function-like symbol.
+ */
+public data class KtFunctionLikeSignature<out S : KtFunctionLikeSymbol>(
+    private val _symbol: S,
+    private val _returnType: KtType,
+    private val _receiverType: KtType?,
+    private val _valueParameters: List<KtVariableLikeSignature<KtValueParameterSymbol>>,
+) : KtSignature<S>() {
+    override val token: ValidityToken
+        get() = _symbol.token
+    override val symbol: S
+        get() = withValidityAssertion { _symbol }
+    override val returnType: KtType
+        get() = withValidityAssertion { _returnType }
+    override val receiverType: KtType?
+        get() = withValidityAssertion { _receiverType }
+
+    /**
+     * The use-site-substituted value parameters.
+     */
+    public val valueParameters: List<KtVariableLikeSignature<KtValueParameterSymbol>>
+        get() = withValidityAssertion { _valueParameters }
+}
+
+/**
+ * A signature of a variable-like symbol.
+ */
+public data class KtVariableLikeSignature<out S : KtVariableLikeSymbol>(
+    private val _symbol: S,
+    private val _returnType: KtType,
+    private val _receiverType: KtType?,
+) : KtSignature<S>() {
+    override val token: ValidityToken
+        get() = _symbol.token
+    override val symbol: S
+        get() = withValidityAssertion { _symbol }
+    override val returnType: KtType
+        get() = withValidityAssertion { _returnType }
+    override val receiverType: KtType?
+        get() = withValidityAssertion { _receiverType }
+}
+
+public fun <S : KtCallableSymbol> S.toSignature(substitutor: KtSubstitutor = KtSubstitutor.Empty(token)): KtSignature<S> {
+    return when (this) {
+        is KtVariableLikeSymbol -> toSignature(substitutor)
+        is KtFunctionLikeSymbol -> toSignature(substitutor)
+        else -> error("unexpected callable symbol $this")
+    }
+}
+
+/**
+ * Creates a signature with the same type as the given symbol.
+ */
+public fun <S : KtVariableLikeSymbol> S.toSignature(substitutor: KtSubstitutor = KtSubstitutor.Empty(token)): KtVariableLikeSignature<S> {
+    return KtVariableLikeSignature(this, returnType, receiverType).substitute(substitutor)
+}
+
+/**
+ * Creates a signature with the same type as the given symbol.
+ */
+public fun <S : KtFunctionLikeSymbol> S.toSignature(substitutor: KtSubstitutor = KtSubstitutor.Empty(token)): KtFunctionLikeSignature<S> {
+    return KtFunctionLikeSignature(this, returnType, receiverType, valueParameters.map { it.toSignature() }).substitute(substitutor)
+}
+
+/**
+ * Applies a substitutor to the given signature and return a new signature with substituted types.
+ */
+@Suppress("UNCHECKED_CAST")
+public fun <C : KtSignature<S>, S : KtCallableSymbol> C.substitute(substitutor: KtSubstitutor): C {
+    if (substitutor is KtSubstitutor.Empty) return this
+    return when (this) {
+        is KtFunctionLikeSignature<*> -> KtFunctionLikeSignature(
+            symbol,
+            substitutor.substituteOrSelf(returnType),
+            receiverType?.let { substitutor.substituteOrSelf(it) },
+            valueParameters.map { it.substitute(substitutor) }
+        ) as C
+        is KtVariableLikeSignature<*> -> KtVariableLikeSignature(
+            symbol,
+            substitutor.substituteOrSelf(returnType),
+            receiverType?.let { substitutor.substituteOrSelf(it) },
+        ) as C
+        else -> error("impossible since KtSignature is sealed")
+    }
+}