FIR: Introduce FirScope.processOverriddenFunctions

^KT-35495 In Progress

compiler/fir/java/src/org/jetbrains/kotlin/fir/java/scopes/JavaClassEnhancementScope.kt

@@ -56,6 +56,7 @@ class JavaClassEnhancementScope(
         FirJavaEnhancementContext(session) { null }.copyWithNewDefaultTypeQualifiers(typeQualifierResolver, jsr305State, owner.annotations)
 
     private val enhancements = mutableMapOf<FirCallableSymbol<*>, FirCallableSymbol<*>>()
+    private val overriddenFunctions = mutableMapOf<FirFunctionSymbol<*>, Collection<FirFunctionSymbol<*>>>()
 
     override fun processPropertiesByName(name: Name, processor: (FirVariableSymbol<*>) -> Unit) {
         useSiteMemberScope.processPropertiesByName(name) process@{ original ->
@@ -253,7 +254,9 @@ class JavaClassEnhancementScope(
             else -> throw AssertionError("Unknown Java method to enhance: ${firMethod.render()}")
         }.apply {
             annotations += firMethod.annotations
-        }.build()
+        }.build().also {
+            overriddenFunctions[it.symbol] = overriddenMembers.mapNotNull { it.symbol as? FirFunctionSymbol<*> }
+        }
         return function.symbol
     }
 
@@ -412,4 +415,8 @@ class JavaClassEnhancementScope(
         }
     }
 
+    override fun processOverriddenFunctions(
+        functionSymbol: FirFunctionSymbol<*>,
+        processor: (FirFunctionSymbol<*>) -> ProcessorAction
+    ): ProcessorAction = doProcessOverriddenFunctions(functionSymbol, processor, overriddenFunctions, useSiteMemberScope)
 }

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/scopes/impl/AbstractFirUseSiteMemberScope.kt

@@ -27,6 +27,7 @@ abstract class AbstractFirUseSiteMemberScope(
 ) : AbstractFirOverrideScope(session, overrideChecker) {
 
     private val functions = hashMapOf<Name, Collection<FirFunctionSymbol<*>>>()
+    private val directOverridden = hashMapOf<FirFunctionSymbol<*>, Collection<FirFunctionSymbol<*>>>()
 
     override fun processFunctionsByName(name: Name, processor: (FirFunctionSymbol<*>) -> Unit) {
         functions.getOrPut(name) {
@@ -41,7 +42,9 @@ abstract class AbstractFirUseSiteMemberScope(
     ): Collection<FirFunctionSymbol<*>> = mutableListOf<FirFunctionSymbol<*>>().apply {
         val overrideCandidates = mutableSetOf<FirFunctionSymbol<*>>()
         declaredMemberScope.processFunctionsByName(name) {
-            val symbol = processInheritedDefaultParameters(it)
+            val directOverridden = computeDirectOverridden(it)
+            this@AbstractFirUseSiteMemberScope.directOverridden[it] = directOverridden
+            val symbol = processInheritedDefaultParameters(it, directOverridden)
             overrideCandidates += symbol
             add(symbol)
         }
@@ -56,30 +59,39 @@ abstract class AbstractFirUseSiteMemberScope(
         }
     }
 
-    private fun processInheritedDefaultParameters(symbol: FirFunctionSymbol<*>): FirFunctionSymbol<*> {
+    private fun computeDirectOverridden(symbol: FirFunctionSymbol<*>): Collection<FirFunctionSymbol<*>> {
-        val firSimpleFunction = symbol.fir as? FirSimpleFunction ?: return symbol
+        val result = mutableListOf<FirFunctionSymbol<*>>()
-        if (firSimpleFunction.valueParameters.isEmpty() || firSimpleFunction.valueParameters.any { it.defaultValue != null }) return symbol
+        val firSimpleFunction = symbol.fir as? FirSimpleFunction ?: return emptyList()
-
-        var foundFir: FirFunction<*>? = null
         superTypesScope.processFunctionsByName(symbol.callableId.callableName) { superSymbol ->
             val superFunctionFir = superSymbol.fir
-            if (foundFir == null &&
+            if (superFunctionFir is FirSimpleFunction &&
-                superFunctionFir is FirSimpleFunction &&
+                overrideChecker.isOverriddenFunction(firSimpleFunction, superFunctionFir)
-                overrideChecker.isOverriddenFunction(firSimpleFunction, superFunctionFir) &&
-                superFunctionFir.valueParameters.any { parameter -> parameter.defaultValue != null }
             ) {
-                foundFir = superFunctionFir
+                result.add(superSymbol)
             }
         }
 
-        if (foundFir == null) return symbol
+        return result
+    }
+
+    private fun processInheritedDefaultParameters(
+        symbol: FirFunctionSymbol<*>,
+        directOverridden: Collection<FirFunctionSymbol<*>>
+    ): FirFunctionSymbol<*> {
+        val firSimpleFunction = symbol.fir as? FirSimpleFunction ?: return symbol
+        if (firSimpleFunction.valueParameters.isEmpty() || firSimpleFunction.valueParameters.any { it.defaultValue != null }) return symbol
+
+        val overriddenWithDefault: FirFunction<*> =
+            directOverridden.singleOrNull {
+                it.fir.valueParameters.any { parameter -> parameter.defaultValue != null }
+            }?.fir ?: return symbol
 
         val newSymbol = FirNamedFunctionSymbol(symbol.callableId, false, null)
 
         createFunctionCopy(firSimpleFunction, newSymbol).apply {
             resolvePhase = firSimpleFunction.resolvePhase
             typeParameters += firSimpleFunction.typeParameters
-            valueParameters += firSimpleFunction.valueParameters.zip(foundFir!!.valueParameters)
+            valueParameters += firSimpleFunction.valueParameters.zip(overriddenWithDefault.valueParameters)
                 .map { (overrideParameter, overriddenParameter) ->
                     if (overriddenParameter.defaultValue != null)
                         createValueParameterCopy(overrideParameter, overriddenParameter.defaultValue).apply {
@@ -93,7 +105,10 @@ abstract class AbstractFirUseSiteMemberScope(
         return newSymbol
     }
 
-    protected open fun createFunctionCopy(firSimpleFunction: FirSimpleFunction, newSymbol: FirNamedFunctionSymbol): FirSimpleFunctionBuilder =
+    protected open fun createFunctionCopy(
+        firSimpleFunction: FirSimpleFunction,
+        newSymbol: FirNamedFunctionSymbol
+    ): FirSimpleFunctionBuilder =
         FirSimpleFunctionBuilder().apply {
             source = firSimpleFunction.source
             session = firSimpleFunction.session
@@ -119,6 +134,11 @@ abstract class AbstractFirUseSiteMemberScope(
             isVararg = parameter.isVararg
         }
 
+    override fun processOverriddenFunctions(
+        functionSymbol: FirFunctionSymbol<*>,
+        processor: (FirFunctionSymbol<*>) -> ProcessorAction
+    ): ProcessorAction = doProcessOverriddenFunctions(functionSymbol, processor, directOverridden, superTypesScope)
+
     override fun processClassifiersByNameWithSubstitution(name: Name, processor: (FirClassifierSymbol<*>, ConeSubstitutor) -> Unit) {
         declaredMemberScope.processClassifiersByNameWithSubstitution(name, processor)
         superTypesScope.processClassifiersByNameWithSubstitution(name, processor)

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/scopes/impl/FirClassSubstitutionScope.kt

@@ -64,6 +64,14 @@ class FirClassSubstitutionScope(
         return super.processFunctionsByName(name, processor)
     }
 
+    override fun processOverriddenFunctions(
+        functionSymbol: FirFunctionSymbol<*>,
+        processor: (FirFunctionSymbol<*>) -> ProcessorAction
+    ): ProcessorAction {
+        val unwrapped = functionSymbol.unwrapOverriddenOnce()
+        return useSiteMemberScope.processOverriddenFunctions(unwrapped, processor)
+    }
+
     override fun processPropertiesByName(name: Name, processor: (FirVariableSymbol<*>) -> Unit) {
         return useSiteMemberScope.processPropertiesByName(name) process@{ original ->
             when (original) {

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/scopes/impl/FirSuperTypeScope.kt

@@ -239,6 +239,17 @@ class FirSuperTypeScope private constructor(
         super.processClassifiersByNameWithSubstitution(name, processor)
     }
 
+    override fun processOverriddenFunctions(
+        functionSymbol: FirFunctionSymbol<*>,
+        processor: (FirFunctionSymbol<*>) -> ProcessorAction
+    ): ProcessorAction {
+        for (scope in scopes) {
+            if (!scope.processOverriddenFunctions(functionSymbol, processor)) return ProcessorAction.STOP
+        }
+
+        return ProcessorAction.NEXT
+    }
+
     companion object {
         fun prepareSupertypeScope(
             session: FirSession,

compiler/fir/tree/src/org/jetbrains/kotlin/fir/scopes/FirIterableScope.kt

@@ -45,4 +45,12 @@ abstract class FirIterableScope : FirScope() {
     override fun processPropertiesByName(name: Name, processor: (FirVariableSymbol<*>) -> Unit) {
         return processComposite(FirScope::processPropertiesByName, name, processor)
     }
+
+    override fun processOverriddenFunctions(functionSymbol: FirFunctionSymbol<*>, processor: (FirFunctionSymbol<*>) -> ProcessorAction): ProcessorAction {
+        for (scope in scopes) {
+            if (!scope.processOverriddenFunctions(functionSymbol, processor)) return ProcessorAction.STOP
+        }
+
+        return ProcessorAction.NEXT
+    }
 }

compiler/fir/tree/src/org/jetbrains/kotlin/fir/scopes/FirScope.kt

@@ -33,6 +33,46 @@ abstract class FirScope {
     ) {}
 
     open fun mayContainName(name: Name) = true
+
+    // Currently, this function has very weak guarantees
+    // - It may silently do nothing on symbols originated from different scope instance
+    // - It may return the same overridden symbols more then once in case of substitution
+    // - It doesn't guarantee any specific order in which overridden tree will be traversed
+    // But if the scope instance is the same as the one from which the symbol was originated, this function will enumarate all members
+    // of the overridden tree
+    // TODO: Consider extracting this function to a separate abstract class/interface, so only limited types of scopes would be supporing it
+    // on interface level
+    open fun processOverriddenFunctions(
+        functionSymbol: FirFunctionSymbol<*>,
+        processor: (FirFunctionSymbol<*>) -> ProcessorAction
+    ): ProcessorAction = ProcessorAction.NONE
+
+    // This is just a helper for a common implementation
+    protected fun doProcessOverriddenFunctions(
+        functionSymbol: FirFunctionSymbol<*>,
+        processor: (FirFunctionSymbol<*>) -> ProcessorAction,
+        directOverriddenMap: Map<FirFunctionSymbol<*>, Collection<FirFunctionSymbol<*>>>,
+        baseScope: FirScope
+    ): ProcessorAction {
+        val directOverridden =
+            directOverriddenMap[functionSymbol] ?: return baseScope.processOverriddenFunctions(functionSymbol, processor)
+
+        for (overridden in directOverridden) {
+            if (!processor(overridden)) return ProcessorAction.STOP
+            if (!baseScope.processOverriddenFunctions(overridden, processor)) return ProcessorAction.STOP
+        }
+
+        return ProcessorAction.NEXT
+    }
+}
+
+fun FirScope.processOverriddenFunctionsAndSelf(
+    functionSymbol: FirFunctionSymbol<*>,
+    processor: (FirFunctionSymbol<*>) -> ProcessorAction
+): ProcessorAction {
+    if (!processor(functionSymbol)) return ProcessorAction.STOP
+
+    return processOverriddenFunctions(functionSymbol, processor)
 }
 
 enum class ProcessorAction {

compiler/fir/tree/src/org/jetbrains/kotlin/fir/symbols/impl/FirCallableSymbol.kt

@@ -15,3 +15,9 @@ abstract class FirCallableSymbol<D : FirCallableDeclaration<D>> : AbstractFirBas
     open val overriddenSymbol: FirCallableSymbol<D>?
         get() = null
 }
+
+inline fun <reified E : FirCallableSymbol<*>> E.unwrapOverriddenOnce(): E {
+    overriddenSymbol?.let { return it as E }
+
+    return this
+}