[FIR] Introduce FirTypeScope.processOverriddenFunctionWithDepth

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

@@ -73,8 +73,8 @@ class JavaClassMembersEnhancementScope(
         }
     }
 
-    override fun processOverriddenFunctions(
+    override fun processOverriddenFunctionsWithDepth(
         functionSymbol: FirFunctionSymbol<*>,
-        processor: (FirFunctionSymbol<*>) -> ProcessorAction
+        processor: (FirFunctionSymbol<*>, Int) -> ProcessorAction
     ): ProcessorAction = doProcessOverriddenFunctions(functionSymbol, processor, overriddenFunctions, useSiteMemberScope)
 }

compiler/fir/jvm/src/org/jetbrains/kotlin/fir/scopes/jvm/JvmMappedScope.kt

@@ -42,9 +42,9 @@ class JvmMappedScope(
     // JvmMappedScope is basically used as declaration scope but it's being wrapped into substitution scopes where
     // a use-site (FirOverrideAwareScope) is expected
     // So, we put here a stub implementation
-    override fun processOverriddenFunctions(
+    override fun processOverriddenFunctionsWithDepth(
         functionSymbol: FirFunctionSymbol<*>,
-        processor: (FirFunctionSymbol<*>) -> ProcessorAction
+        processor: (FirFunctionSymbol<*>, Int) -> ProcessorAction
     ) = ProcessorAction.NONE
 
     override fun processDeclaredConstructors(processor: (FirConstructorSymbol) -> Unit) {

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/scopes/KotlinScopeProvider.kt

@@ -93,7 +93,7 @@ data class ConeSubstitutionScopeKey(
     val classId: ClassId?, val isFromExpectClass: Boolean, val substitutor: ConeSubstitutor
 ) : ScopeSessionKey<FirClass<*>, FirClassSubstitutionScope>()
 
-fun FirClass<*>.unsubstitutedScope(useSiteSession: FirSession, scopeSession: ScopeSession): FirScope {
+fun FirClass<*>.unsubstitutedScope(useSiteSession: FirSession, scopeSession: ScopeSession): FirTypeScope {
     return scopeProvider.getUseSiteMemberScope(this, useSiteSession, scopeSession)
 }
 

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

@@ -137,9 +137,9 @@ abstract class AbstractFirUseSiteMemberScope(
             isVararg = parameter.isVararg
         }
 
-    override fun processOverriddenFunctions(
+    override fun processOverriddenFunctionsWithDepth(
         functionSymbol: FirFunctionSymbol<*>,
-        processor: (FirFunctionSymbol<*>) -> ProcessorAction
+        processor: (FirFunctionSymbol<*>, Int) -> ProcessorAction
     ): ProcessorAction = doProcessOverriddenFunctions(functionSymbol, processor, directOverridden, superTypesScope)
 
     override fun processClassifiersByNameWithSubstitution(name: Name, processor: (FirClassifierSymbol<*>, ConeSubstitutor) -> Unit) {

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

@@ -64,12 +64,18 @@ class FirClassSubstitutionScope(
         return super.processFunctionsByName(name, processor)
     }
 
-    override fun processOverriddenFunctions(
+    override fun processOverriddenFunctionsWithDepth(
         functionSymbol: FirFunctionSymbol<*>,
-        processor: (FirFunctionSymbol<*>) -> ProcessorAction
+        processor: (FirFunctionSymbol<*>, Int) -> ProcessorAction
     ): ProcessorAction {
         val unwrapped = functionSymbol.overriddenSymbol as FirFunctionSymbol<*>? ?: functionSymbol
-        return useSiteMemberScope.processOverriddenFunctions(unwrapped, processor)
+        val overriddenDepth = if (unwrapped != functionSymbol) 1 else 0
+        if (!processor(unwrapped, overriddenDepth)) {
+            return ProcessorAction.STOP
+        }
+        return useSiteMemberScope.processOverriddenFunctionsWithDepth(unwrapped) { symbol, depth ->
+            processor(symbol, depth + overriddenDepth)
+        }
     }
 
     override fun processPropertiesByName(name: Name, processor: (FirVariableSymbol<*>) -> Unit) {

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

@@ -96,9 +96,9 @@ class FirIntegerLiteralTypeScope(private val session: FirSession, val isUnsigned
     override fun processPropertiesByName(name: Name, processor: (FirVariableSymbol<*>) -> Unit) {
     }
 
-    override fun processOverriddenFunctions(
+    override fun processOverriddenFunctionsWithDepth(
         functionSymbol: FirFunctionSymbol<*>,
-        processor: (FirFunctionSymbol<*>) -> ProcessorAction
+        processor: (FirFunctionSymbol<*>, Int) -> ProcessorAction
     ): ProcessorAction = ProcessorAction.NEXT
 }
 

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

@@ -232,9 +232,9 @@ class FirTypeIntersectionScope private constructor(
         super.processClassifiersByNameWithSubstitution(name, processor)
     }
 
-    override fun processOverriddenFunctions(
+    override fun processOverriddenFunctionsWithDepth(
         functionSymbol: FirFunctionSymbol<*>,
-        processor: (FirFunctionSymbol<*>) -> ProcessorAction
+        processor: (FirFunctionSymbol<*>, Int) -> ProcessorAction
     ): ProcessorAction {
         @Suppress("UNCHECKED_CAST")
         val directOverriddenSymbols =
@@ -244,8 +244,11 @@ class FirTypeIntersectionScope private constructor(
         for (directOverridden in directOverriddenSymbols) {
             // TODO: Preserve the scope where directOverridden came from
             for (scope in scopes) {
-                if (!processor(directOverridden)) return ProcessorAction.STOP
+                if (!processor(directOverridden, 0)) return ProcessorAction.STOP
-                if (!scope.processOverriddenFunctions(directOverridden, processor)) return ProcessorAction.STOP
+                if (!scope.processOverriddenFunctionsWithDepth(directOverridden) { symbol, depth ->
+                        processor(symbol, depth)
+                    }
+                ) return ProcessorAction.STOP
             }
         }
 

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

@@ -6,6 +6,7 @@
 package org.jetbrains.kotlin.fir.scopes
 
 import org.jetbrains.kotlin.fir.symbols.impl.FirFunctionSymbol
+import org.jetbrains.kotlin.fir.symbols.impl.FirNamedFunctionSymbol
 
 abstract class FirTypeScope : FirScope() {
     // Currently, this function has very weak guarantees
@@ -14,33 +15,55 @@ abstract class FirTypeScope : FirScope() {
     // - 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
-    abstract fun processOverriddenFunctions(
+    abstract fun processOverriddenFunctionsWithDepth(
         functionSymbol: FirFunctionSymbol<*>,
-        processor: (FirFunctionSymbol<*>) -> ProcessorAction
+        processor: (FirFunctionSymbol<*>, Int) -> ProcessorAction
     ): ProcessorAction
 
+    inline fun processOverriddenFunctions(
+        functionSymbol: FirFunctionSymbol<*>,
+        crossinline processor: (FirFunctionSymbol<*>) -> ProcessorAction
+    ): ProcessorAction = processOverriddenFunctionsWithDepth(functionSymbol) { symbol, _ ->
+        processor(symbol)
+    }
+
+    inline fun processDirectlyOverriddenFunctions(
+        functionSymbol: FirFunctionSymbol<*>,
+        crossinline processor: (FirFunctionSymbol<*>) -> ProcessorAction
+    ): ProcessorAction = processOverriddenFunctionsWithDepth(functionSymbol) { symbol, depth ->
+        if (depth == 1) {
+            processor(symbol)
+        } else {
+            ProcessorAction.NEXT
+        }
+    }
+
     // This is just a helper for a common implementation
     protected fun doProcessOverriddenFunctions(
         functionSymbol: FirFunctionSymbol<*>,
-        processor: (FirFunctionSymbol<*>) -> ProcessorAction,
+        processor: (FirFunctionSymbol<*>, Int) -> ProcessorAction,
         directOverriddenMap: Map<FirFunctionSymbol<*>, Collection<FirFunctionSymbol<*>>>,
         baseScope: FirTypeScope
     ): ProcessorAction {
         val directOverridden =
-            directOverriddenMap[functionSymbol] ?: return baseScope.processOverriddenFunctions(functionSymbol, processor)
+            directOverriddenMap[functionSymbol] ?: return baseScope.processOverriddenFunctionsWithDepth(functionSymbol, processor)
 
         for (overridden in directOverridden) {
-            if (!processor(overridden)) return ProcessorAction.STOP
+            val overriddenDepth = if (overridden is FirNamedFunctionSymbol && overridden.isFakeOverride) 0 else 1
-            if (!baseScope.processOverriddenFunctions(overridden, processor)) return ProcessorAction.STOP
+            if (!processor(overridden, overriddenDepth)) return ProcessorAction.STOP
+            if (!baseScope.processOverriddenFunctionsWithDepth(overridden) { symbol, depth ->
+                    processor(symbol, depth + overriddenDepth)
+                }
+            ) return ProcessorAction.STOP
         }
 
-        return baseScope.processOverriddenFunctions(functionSymbol, processor)
+        return baseScope.processOverriddenFunctionsWithDepth(functionSymbol, processor)
     }
 
     object Empty : FirTypeScope() {
-        override fun processOverriddenFunctions(
+        override fun processOverriddenFunctionsWithDepth(
             functionSymbol: FirFunctionSymbol<*>,
-            processor: (FirFunctionSymbol<*>) -> ProcessorAction
+            processor: (FirFunctionSymbol<*>, Int) -> ProcessorAction
         ): ProcessorAction = ProcessorAction.NEXT
     }
 }