From e4992176c19759c7d6c381c1449f83bc4adf51de Mon Sep 17 00:00:00 2001
From: Ilya Kirillov <ilya.kirillov@jetbrains.com>
Date: Thu, 29 Jul 2021 14:17:14 +0200
Subject: [PATCH] FIR LC: fix inheritance checking for light classes

Before it was not able to check inheritance with type parameters substitutions like

class A<T>
class B: A<Int>
---
 .../KtSymbolDeclarationOverridesProvider.kt   |  9 ++++
 .../idea/asJava/classes/firLightClassUtils.kt | 12 +++---
 ...KtFirSymbolDeclarationOverridesProvider.kt | 41 ++++++++++++++++++-
 3 files changed, 54 insertions(+), 8 deletions(-)

diff --git a/idea/idea-frontend-api/src/org/jetbrains/kotlin/idea/frontend/api/components/KtSymbolDeclarationOverridesProvider.kt b/idea/idea-frontend-api/src/org/jetbrains/kotlin/idea/frontend/api/components/KtSymbolDeclarationOverridesProvider.kt
index 9effaa02510..4600fb6a6a3 100644
--- a/idea/idea-frontend-api/src/org/jetbrains/kotlin/idea/frontend/api/components/KtSymbolDeclarationOverridesProvider.kt
+++ b/idea/idea-frontend-api/src/org/jetbrains/kotlin/idea/frontend/api/components/KtSymbolDeclarationOverridesProvider.kt
@@ -13,6 +13,9 @@ public abstract class KtSymbolDeclarationOverridesProvider : KtAnalysisSessionCo
     public abstract fun <T : KtSymbol> getAllOverriddenSymbols(callableSymbol: T): List<KtCallableSymbol>
     public abstract fun <T : KtSymbol> getDirectlyOverriddenSymbols(callableSymbol: T): List<KtCallableSymbol>
 
+    public abstract fun isSubClassOf(subClass: KtClassOrObjectSymbol, superClass: KtClassOrObjectSymbol): Boolean
+    public abstract fun isDirectSubClassOf(subClass: KtClassOrObjectSymbol, superClass: KtClassOrObjectSymbol): Boolean
+
     public abstract fun getIntersectionOverriddenSymbols(symbol: KtCallableSymbol): Collection<KtCallableSymbol>
 }
 
@@ -41,6 +44,12 @@ public interface KtSymbolDeclarationOverridesProviderMixIn : KtAnalysisSessionMi
     public fun KtCallableSymbol.getDirectlyOverriddenSymbols(): List<KtCallableSymbol> =
         analysisSession.symbolDeclarationOverridesProvider.getDirectlyOverriddenSymbols(this)
 
+    public fun KtClassOrObjectSymbol.isSubClassOf(superClass: KtClassOrObjectSymbol): Boolean =
+        analysisSession.symbolDeclarationOverridesProvider.isSubClassOf(this, superClass)
+
+    public fun KtClassOrObjectSymbol.isDirectSubClassOf(superClass: KtClassOrObjectSymbol): Boolean =
+        analysisSession.symbolDeclarationOverridesProvider.isDirectSubClassOf(this, superClass)
+
     public fun KtCallableSymbol.getIntersectionOverriddenSymbols(): Collection<KtCallableSymbol> =
         analysisSession.symbolDeclarationOverridesProvider.getIntersectionOverriddenSymbols(this)
 }
\ No newline at end of file
diff --git a/idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/asJava/classes/firLightClassUtils.kt b/idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/asJava/classes/firLightClassUtils.kt
index bee48e12919..9aa45d9d571 100644
--- a/idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/asJava/classes/firLightClassUtils.kt
+++ b/idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/asJava/classes/firLightClassUtils.kt
@@ -332,17 +332,15 @@ internal fun KtSymbolWithMembers.createInnerClasses(manager: PsiManager): List<F
 
 internal fun KtClassOrObject.checkIsInheritor(baseClassOrigin: KtClassOrObject, checkDeep: Boolean): Boolean {
     return analyseForLightClasses(this) {
-        val thisSymbol = this@checkIsInheritor.getNamedClassOrObjectSymbol() ?: return false
-        val baseSymbol = baseClassOrigin.getNamedClassOrObjectSymbol() ?: return false
+        val subClassSymbol = this@checkIsInheritor.getNamedClassOrObjectSymbol() ?: return false
+        val superClassSymbol = baseClassOrigin.getNamedClassOrObjectSymbol() ?: return false
 
-        if (thisSymbol == baseSymbol) return@analyseForLightClasses false
-
-        val baseType = baseSymbol.buildSelfClassType()
+        if (subClassSymbol == superClassSymbol) return@analyseForLightClasses false
 
         if (checkDeep) {
-            thisSymbol.buildSelfClassType().isSubTypeOf(baseType)
+            subClassSymbol.isSubClassOf(superClassSymbol)
         } else {
-            thisSymbol.superTypes.any { baseType.isEqualTo(it.type) }
+            subClassSymbol.isDirectSubClassOf(superClassSymbol)
         }
     }
 }
diff --git a/idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/frontend/api/fir/components/KtFirSymbolDeclarationOverridesProvider.kt b/idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/frontend/api/fir/components/KtFirSymbolDeclarationOverridesProvider.kt
index 03d60a08830..4d4a6f82923 100644
--- a/idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/frontend/api/fir/components/KtFirSymbolDeclarationOverridesProvider.kt
+++ b/idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/frontend/api/fir/components/KtFirSymbolDeclarationOverridesProvider.kt
@@ -5,18 +5,23 @@
 
 package org.jetbrains.kotlin.idea.frontend.api.fir.components
 
+import org.jetbrains.kotlin.fir.analysis.checkers.toRegularClassSymbol
 import org.jetbrains.kotlin.fir.declarations.*
+import org.jetbrains.kotlin.fir.declarations.utils.superConeTypes
 import org.jetbrains.kotlin.fir.resolve.ScopeSession
 import org.jetbrains.kotlin.fir.scopes.*
 import org.jetbrains.kotlin.fir.symbols.FirBasedSymbol
+import org.jetbrains.kotlin.fir.symbols.ensureResolved
 import org.jetbrains.kotlin.fir.symbols.impl.FirCallableSymbol
 import org.jetbrains.kotlin.fir.symbols.impl.FirIntersectionOverrideFunctionSymbol
 import org.jetbrains.kotlin.fir.symbols.impl.FirIntersectionOverridePropertySymbol
+import org.jetbrains.kotlin.idea.fir.low.level.api.lazy.resolve.ResolveType
 import org.jetbrains.kotlin.idea.frontend.api.components.KtSymbolDeclarationOverridesProvider
 import org.jetbrains.kotlin.idea.frontend.api.fir.KtFirAnalysisSession
 import org.jetbrains.kotlin.idea.frontend.api.fir.symbols.KtFirNamedClassOrObjectSymbol
 import org.jetbrains.kotlin.idea.frontend.api.fir.symbols.KtFirSymbol
 import org.jetbrains.kotlin.idea.frontend.api.symbols.KtCallableSymbol
+import org.jetbrains.kotlin.idea.frontend.api.symbols.KtClassOrObjectSymbol
 import org.jetbrains.kotlin.idea.frontend.api.symbols.KtSymbol
 import org.jetbrains.kotlin.idea.frontend.api.symbols.KtSymbolOrigin
 import org.jetbrains.kotlin.idea.frontend.api.tokens.ValidityToken
@@ -129,12 +134,46 @@ internal class KtFirSymbolDeclarationOverridesProvider(
         }
     }
 
+    override fun isSubClassOf(subClass: KtClassOrObjectSymbol, superClass: KtClassOrObjectSymbol): Boolean {
+        return isSubClassOf(subClass, superClass, checkDeep = true)
+    }
+
+    override fun isDirectSubClassOf(subClass: KtClassOrObjectSymbol, superClass: KtClassOrObjectSymbol): Boolean {
+        return isSubClassOf(subClass, superClass, checkDeep = false)
+    }
+
+    private fun isSubClassOf(subClass: KtClassOrObjectSymbol, superClass: KtClassOrObjectSymbol, checkDeep: Boolean): Boolean {
+        require(subClass is KtFirSymbol<*>)
+        require(superClass is KtFirSymbol<*>)
+
+        if (subClass == superClass) return false
+        return subClass.firRef.withFirByType(ResolveType.ClassSuperTypes) { subClassFir ->
+            check(subClassFir is FirRegularClass)
+            superClass.firRef.withFir { superClassFir ->
+                check(superClassFir is FirRegularClass)
+                isSubClassOf(subClassFir, superClassFir, checkDeep)
+            }
+        }
+    }
+
+    private fun isSubClassOf(subClass: FirRegularClass, superClass: FirRegularClass, checkDeep: Boolean): Boolean {
+        if (subClass.superConeTypes.any { it.toRegularClassSymbol(rootModuleSession) == superClass.symbol }) return true
+        if (!checkDeep) return false
+        subClass.superConeTypes.forEach { superType ->
+            val superOfSub = superType.toRegularClassSymbol(rootModuleSession) ?: return@forEach
+            superOfSub.ensureResolved(FirResolvePhase.SUPER_TYPES)
+            if (isSubClassOf(superOfSub.fir, superClass, checkDeep = true)) return true
+        }
+        return false
+    }
+
     override fun getIntersectionOverriddenSymbols(symbol: KtCallableSymbol): Collection<KtCallableSymbol> {
         require(symbol is KtFirSymbol<*>)
         if (symbol.origin != KtSymbolOrigin.INTERSECTION_OVERRIDE) return emptyList()
         return symbol.firRef.withFir { fir ->
             val firSymbol = fir.symbol
-            firSymbol.getIntersectionOverriddenSymbols().map { analysisSession.firSymbolBuilder.callableBuilder.buildCallableSymbol(it.fir) }
+            firSymbol.getIntersectionOverriddenSymbols()
+                .map { analysisSession.firSymbolBuilder.callableBuilder.buildCallableSymbol(it.fir) }
         }
     }