FIR: Introduce FirScope.processOverriddenFunctions

^KT-35495 In Progress

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<*> {
-        val firSimpleFunction = symbol.fir as? FirSimpleFunction ?: return symbol
-        if (firSimpleFunction.valueParameters.isEmpty() || firSimpleFunction.valueParameters.any { it.defaultValue != null }) return symbol
-
-        var foundFir: FirFunction<*>? = null
+    private fun computeDirectOverridden(symbol: FirFunctionSymbol<*>): Collection<FirFunctionSymbol<*>> {
+        val result = mutableListOf<FirFunctionSymbol<*>>()
+        val firSimpleFunction = symbol.fir as? FirSimpleFunction ?: return emptyList()
         superTypesScope.processFunctionsByName(symbol.callableId.callableName) { superSymbol ->
             val superFunctionFir = superSymbol.fir
-            if (foundFir == null &&
-                superFunctionFir is FirSimpleFunction &&
-                overrideChecker.isOverriddenFunction(firSimpleFunction, superFunctionFir) &&
-                superFunctionFir.valueParameters.any { parameter -> parameter.defaultValue != null }
+            if (superFunctionFir is FirSimpleFunction &&
+                overrideChecker.isOverriddenFunction(firSimpleFunction, superFunctionFir)
             ) {
-                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,