From 954926e0deebb698790a293ee4afa4afe33c5083 Mon Sep 17 00:00:00 2001
From: Tianyu Geng <tgeng@google.com>
Date: Tue, 7 Sep 2021 17:16:37 -0700
Subject: [PATCH] FIR IDE: get containing symbol of KtTypeParameterSymbol

As part of this change, we now also allow HL API to get containing
symbol for any KtSymbols.
---
 ...tFirSymbolContainingDeclarationProvider.kt | 33 ++++++++++++-------
 .../KtSymbolContainingDeclarationProvider.kt  |  5 +--
 2 files changed, 25 insertions(+), 13 deletions(-)

diff --git a/analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/components/KtFirSymbolContainingDeclarationProvider.kt b/analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/components/KtFirSymbolContainingDeclarationProvider.kt
index a3e20463e48..831764dcc71 100644
--- a/analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/components/KtFirSymbolContainingDeclarationProvider.kt
+++ b/analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/components/KtFirSymbolContainingDeclarationProvider.kt
@@ -15,6 +15,7 @@ import org.jetbrains.kotlin.analysis.api.tokens.ValidityToken
 import org.jetbrains.kotlin.analysis.api.components.KtSymbolContainingDeclarationProvider
 import org.jetbrains.kotlin.analysis.api.fir.KtFirAnalysisSession
 import org.jetbrains.kotlin.analysis.api.fir.symbols.KtFirSymbol
+import org.jetbrains.kotlin.analysis.api.fir.symbols.KtFirTypeParameterSymbol
 import org.jetbrains.kotlin.analysis.api.symbols.*
 import org.jetbrains.kotlin.analysis.api.symbols.markers.KtSymbolKind
 import org.jetbrains.kotlin.analysis.api.symbols.markers.KtSymbolWithKind
@@ -24,9 +25,9 @@ internal class KtFirSymbolContainingDeclarationProvider(
     override val analysisSession: KtFirAnalysisSession,
     override val token: ValidityToken,
 ) : KtSymbolContainingDeclarationProvider(), KtFirAnalysisSessionComponent {
-    override fun getContainingDeclaration(symbol: KtSymbolWithKind): KtSymbolWithKind? {
+    override fun getContainingDeclaration(symbol: KtSymbol): KtSymbolWithKind? {
         if (symbol is KtPackageSymbol) return null
-        if (symbol.symbolKind == KtSymbolKind.TOP_LEVEL) return null
+        if (symbol is KtSymbolWithKind && symbol.symbolKind == KtSymbolKind.TOP_LEVEL) return null
         if (symbol is KtCallableSymbol) {
             val classId = symbol.callableIdIfNonLocal?.classId
             if (classId != null) {
@@ -35,15 +36,25 @@ internal class KtFirSymbolContainingDeclarationProvider(
                 }
             }
         }
-        return when (symbol.origin) {
-            KtSymbolOrigin.SOURCE, KtSymbolOrigin.SOURCE_MEMBER_GENERATED ->
-                getContainingDeclarationForKotlinInSourceSymbol(symbol)
-            KtSymbolOrigin.LIBRARY, KtSymbolOrigin.JAVA, KtSymbolOrigin.JAVA_SYNTHETIC_PROPERTY ->
-                getContainingDeclarationForLibrarySymbol(symbol)
-            KtSymbolOrigin.PROPERTY_BACKING_FIELD -> getContainingDeclarationForBackingFieldSymbol(symbol)
-            KtSymbolOrigin.INTERSECTION_OVERRIDE -> TODO()
-            KtSymbolOrigin.SAM_CONSTRUCTOR -> null
-            KtSymbolOrigin.DELEGATED -> TODO()
+        return when (symbol) {
+            is KtFirTypeParameterSymbol -> {
+                symbol.firRef.withFir { fir ->
+                    fir.containingDeclarationSymbol?.fir?.let { containingDeclaration ->
+                        firSymbolBuilder.buildSymbol(containingDeclaration) as KtSymbolWithKind
+                    }
+                }
+            }
+            is KtSymbolWithKind -> when (symbol.origin) {
+                KtSymbolOrigin.SOURCE, KtSymbolOrigin.SOURCE_MEMBER_GENERATED ->
+                    getContainingDeclarationForKotlinInSourceSymbol(symbol)
+                KtSymbolOrigin.LIBRARY, KtSymbolOrigin.JAVA, KtSymbolOrigin.JAVA_SYNTHETIC_PROPERTY ->
+                    getContainingDeclarationForLibrarySymbol(symbol)
+                KtSymbolOrigin.PROPERTY_BACKING_FIELD -> getContainingDeclarationForBackingFieldSymbol(symbol)
+                KtSymbolOrigin.INTERSECTION_OVERRIDE -> TODO()
+                KtSymbolOrigin.SAM_CONSTRUCTOR -> null
+                KtSymbolOrigin.DELEGATED -> TODO()
+            }
+            else -> null
         }
     }
 
diff --git a/analysis/analysis-api/src/org/jetbrains/kotlin/analysis/api/components/KtSymbolContainingDeclarationProvider.kt b/analysis/analysis-api/src/org/jetbrains/kotlin/analysis/api/components/KtSymbolContainingDeclarationProvider.kt
index c3b8781b3a5..fdac6210f59 100644
--- a/analysis/analysis-api/src/org/jetbrains/kotlin/analysis/api/components/KtSymbolContainingDeclarationProvider.kt
+++ b/analysis/analysis-api/src/org/jetbrains/kotlin/analysis/api/components/KtSymbolContainingDeclarationProvider.kt
@@ -5,10 +5,11 @@
 
 package org.jetbrains.kotlin.analysis.api.components
 
+import org.jetbrains.kotlin.analysis.api.symbols.KtSymbol
 import org.jetbrains.kotlin.analysis.api.symbols.markers.KtSymbolWithKind
 
 public abstract class KtSymbolContainingDeclarationProvider : KtAnalysisSessionComponent() {
-    public abstract fun getContainingDeclaration(symbol: KtSymbolWithKind): KtSymbolWithKind?
+    public abstract fun getContainingDeclaration(symbol: KtSymbol): KtSymbolWithKind?
 }
 
 public interface KtSymbolContainingDeclarationProviderMixIn : KtAnalysisSessionMixIn {
@@ -18,6 +19,6 @@ public interface KtSymbolContainingDeclarationProviderMixIn : KtAnalysisSessionM
      *   for class members returns containing class
      *   for local declaration returns declaration it was declared it
      */
-    public fun KtSymbolWithKind.getContainingSymbol(): KtSymbolWithKind? =
+    public fun KtSymbol.getContainingSymbol(): KtSymbolWithKind? =
         analysisSession.containingDeclarationProvider.getContainingDeclaration(this)
 }
\ No newline at end of file