From cf830887ec13b4b5e0d9f7b6ce46b6c6990ca1e0 Mon Sep 17 00:00:00 2001
From: Mikhail Glukhikh <mikhail.glukhikh@jetbrains.com>
Date: Sat, 20 Feb 2021 12:55:29 +0300
Subject: [PATCH] FIR2IR: optimize override binding for
 Fir2IrLazySimpleFunction

---
 .../generators/FakeOverrideGenerator.kt       | 21 ++++++++++++++++++-
 .../fir/lazy/Fir2IrLazySimpleFunction.kt      |  5 ++++-
 2 files changed, 24 insertions(+), 2 deletions(-)

diff --git a/compiler/fir/fir2ir/src/org/jetbrains/kotlin/fir/backend/generators/FakeOverrideGenerator.kt b/compiler/fir/fir2ir/src/org/jetbrains/kotlin/fir/backend/generators/FakeOverrideGenerator.kt
index 6d25e0adb44..3b6fe8e370d 100644
--- a/compiler/fir/fir2ir/src/org/jetbrains/kotlin/fir/backend/generators/FakeOverrideGenerator.kt
+++ b/compiler/fir/fir2ir/src/org/jetbrains/kotlin/fir/backend/generators/FakeOverrideGenerator.kt
@@ -124,6 +124,25 @@ class FakeOverrideGenerator(
         return result
     }
 
+    internal fun calcBaseSymbolsForFakeOverrideFunction(
+        klass: FirClass<*>,
+        fakeOverride: IrSimpleFunction,
+        originalSymbol: FirNamedFunctionSymbol,
+    ) {
+        val scope = klass.unsubstitutedScope(session, scopeSession, withForcedTypeCalculator = true)
+        val classLookupTag = klass.symbol.toLookupTag()
+        val baseFirSymbolsForFakeOverride =
+            if (originalSymbol.shouldHaveComputedBaseSymbolsForClass(classLookupTag)) {
+                computeBaseSymbols(originalSymbol, FirTypeScope::getDirectOverriddenFunctions, scope, classLookupTag)
+            } else {
+                listOf(originalSymbol)
+            }
+        baseFunctionSymbols[fakeOverride] = baseFirSymbolsForFakeOverride
+    }
+
+    private fun FirCallableSymbol<*>.shouldHaveComputedBaseSymbolsForClass(classLookupTag: ConeClassLikeLookupTag): Boolean =
+        fir.origin.fromSupertypes && dispatchReceiverClassOrNull() == classLookupTag
+
     private inline fun <reified D : FirCallableMemberDeclaration<D>, reified S : FirCallableSymbol<D>, reified I : IrDeclaration> createFakeOverriddenIfNeeded(
         klass: FirClass<*>,
         irClass: IrClass,
@@ -149,7 +168,7 @@ class FakeOverrideGenerator(
         val baseSymbol = originalSymbol.unwrapSubstitutionAndIntersectionOverrides() as S
 
         val (fakeOverrideFirDeclaration, baseFirSymbolsForFakeOverride) = when {
-            originalSymbol.fir.origin.fromSupertypes && originalSymbol.dispatchReceiverClassOrNull() == classLookupTag -> {
+            originalSymbol.shouldHaveComputedBaseSymbolsForClass(classLookupTag) -> {
                 // Substitution or intersection case
                 // We have already a FIR declaration for such fake override
                 originalDeclaration to computeBaseSymbols(originalSymbol, computeDirectOverridden, scope, classLookupTag)
diff --git a/compiler/fir/fir2ir/src/org/jetbrains/kotlin/fir/lazy/Fir2IrLazySimpleFunction.kt b/compiler/fir/fir2ir/src/org/jetbrains/kotlin/fir/lazy/Fir2IrLazySimpleFunction.kt
index f464e8caf81..2b347aaee2b 100644
--- a/compiler/fir/fir2ir/src/org/jetbrains/kotlin/fir/lazy/Fir2IrLazySimpleFunction.kt
+++ b/compiler/fir/fir2ir/src/org/jetbrains/kotlin/fir/lazy/Fir2IrLazySimpleFunction.kt
@@ -8,6 +8,7 @@ package org.jetbrains.kotlin.fir.lazy
 import org.jetbrains.kotlin.fir.backend.*
 import org.jetbrains.kotlin.fir.declarations.*
 import org.jetbrains.kotlin.fir.initialSignatureAttr
+import org.jetbrains.kotlin.fir.originalIfFakeOverride
 import org.jetbrains.kotlin.fir.symbols.Fir2IrSimpleFunctionSymbol
 import org.jetbrains.kotlin.ir.declarations.*
 import org.jetbrains.kotlin.ir.declarations.lazy.lazyVar
@@ -70,7 +71,9 @@ class Fir2IrLazySimpleFunction(
     override var overriddenSymbols: List<IrSimpleFunctionSymbol> by lazyVar {
         val parent = parent
         if (isFakeOverride && parent is Fir2IrLazyClass) {
-            parent.declarations
+            fakeOverrideGenerator.calcBaseSymbolsForFakeOverrideFunction(
+                firParent, this, fir.symbol.originalIfFakeOverride()!!
+            )
             fakeOverrideGenerator.getOverriddenSymbolsForFakeOverride(this)?.let { return@lazyVar it }
         }
         fir.generateOverriddenFunctionSymbols(firParent, session, scopeSession, declarationStorage, fakeOverrideGenerator)