K2 IC: fix classifiers lookups

in particular nested classifiers and companions.

#KTIJ-24833 fixed
#KTIJ-24830 fixed
#KT-55195 fixed

compiler/fir/providers/src/org/jetbrains/kotlin/fir/scopes/impl/FirNestedClassifierScope.kt

@@ -10,6 +10,8 @@ import org.jetbrains.kotlin.fir.declarations.FirClass
 import org.jetbrains.kotlin.fir.declarations.FirRegularClass
 import org.jetbrains.kotlin.fir.declarations.FirTypeAlias
 import org.jetbrains.kotlin.fir.declarations.FirTypeParameterRef
+import org.jetbrains.kotlin.fir.declarations.utils.classId
+import org.jetbrains.kotlin.fir.declarations.utils.isLocal
 import org.jetbrains.kotlin.fir.resolve.substitution.ConeSubstitutor
 import org.jetbrains.kotlin.fir.resolve.substitution.ConeSubstitutorByMap
 import org.jetbrains.kotlin.fir.scopes.FirContainingNamesAwareScope
@@ -21,6 +23,7 @@ import org.jetbrains.kotlin.fir.symbols.impl.FirTypeParameterSymbol
 import org.jetbrains.kotlin.fir.types.ConeTypeParameterType
 import org.jetbrains.kotlin.fir.types.impl.ConeTypeParameterTypeImpl
 import org.jetbrains.kotlin.name.Name
+import org.jetbrains.kotlin.utils.SmartList
 import org.jetbrains.kotlin.utils.addToStdlib.shouldNotBeCalled
 
 abstract class FirNestedClassifierScope(val klass: FirClass, val useSiteSession: FirSession) : FirContainingNamesAwareScope() {
@@ -45,6 +48,10 @@ abstract class FirNestedClassifierScope(val klass: FirClass, val useSiteSession:
     abstract fun isEmpty(): Boolean
 
     override fun getCallableNames(): Set<Name> = emptySet()
+
+    override val scopeOwnerLookupNames: List<String> =
+        if (klass.isLocal) emptyList()
+        else SmartList(klass.classId.asFqNameString())
 }
 
 class FirNestedClassifierScopeImpl(klass: FirClass, useSiteSession: FirSession) : FirNestedClassifierScope(klass, useSiteSession) {

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/QualifiedNameResolution.kt

@@ -43,11 +43,18 @@ fun BodyResolveComponents.resolveRootPartOfQualifier(
             this.nonFatalDiagnostics.addAll(nonFatalDiagnosticsFromExpression.orEmpty())
             annotations += qualifiedAccess.annotations
         }.apply {
-            setTypeOfQualifier(session)
+            setTypeOfQualifier(this@resolveRootPartOfQualifier)
         }
     }
 
-    for (scope in createCurrentScopeList()) {
+    val scopes = createCurrentScopeList()
+    session.lookupTracker?.recordNameLookup(
+        name,
+        scopes.asSequence().flatMap { it.scopeOwnerLookupNames }.asIterable(),
+        qualifiedAccess.source,
+        file.source
+    )
+    for (scope in scopes) {
         scope.getSingleVisibleClassifier(session, this, name)?.let {
             val klass = (it as? FirClassLikeSymbol<*>)?.fullyExpandedClass(session)
                 ?: return@let
@@ -79,7 +86,7 @@ fun BodyResolveComponents.resolveRootPartOfQualifier(
                 )
                 annotations += qualifiedAccess.annotations
             }.apply {
-                setTypeOfQualifier(session)
+                setTypeOfQualifier(this@resolveRootPartOfQualifier)
             }
         }
     }
@@ -104,6 +111,9 @@ fun FirResolvedQualifier.continueQualifier(
         val firClass = outerClassSymbol.fir
         if (firClass !is FirClass) return null
         return firClass.scopeProvider.getNestedClassifierScope(firClass, components.session, components.scopeSession)
+            ?.also {
+                session.lookupTracker?.recordNameLookup(name, it.scopeOwnerLookupNames, qualifiedAccess.source, components.file.source)
+            }
             ?.getSingleVisibleClassifier(session, components, name)
             ?.takeIf { it is FirClassLikeSymbol<*> }
             ?.let { nestedClassSymbol ->
@@ -128,7 +138,7 @@ fun FirResolvedQualifier.continueQualifier(
                         )
                     )
                 }.apply {
-                    setTypeOfQualifier(components.session)
+                    setTypeOfQualifier(components)
                 }
             }
     }
@@ -156,7 +166,7 @@ private fun FqName.continueQualifierInPackage(
             this.nonFatalDiagnostics.addAll(nonFatalDiagnosticsFromExpression.orEmpty())
             annotations += qualifiedAccess.annotations
         }.apply {
-            setTypeOfQualifier(components.session)
+            setTypeOfQualifier(components)
         }
     }
 
@@ -181,7 +191,7 @@ private fun FqName.continueQualifierInPackage(
         isFullyQualified = true
         annotations += qualifiedAccess.annotations
     }.apply {
-        setTypeOfQualifier(components.session)
+        setTypeOfQualifier(components)
     }
 }
 

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/ResolveUtils.kt

@@ -213,22 +213,22 @@ fun BodyResolveComponents.buildResolvedQualifierForClass(
         this.annotations.addAll(annotations)
     }.build().apply {
         if (classId.isLocal) {
-            resultType = typeForQualifierByDeclaration(regularClass.fir, session)
+            resultType = typeForQualifierByDeclaration(regularClass.fir, session, element = this@apply, file)
                 ?.also { replaceCanBeValue(true) }
                 ?: session.builtinTypes.unitType.type
         } else {
-            setTypeOfQualifier(session)
+            setTypeOfQualifier(this@buildResolvedQualifierForClass)
         }
     }
 }
 
-fun FirResolvedQualifier.setTypeOfQualifier(session: FirSession) {
+fun FirResolvedQualifier.setTypeOfQualifier(components: BodyResolveComponents) {
     val classSymbol = symbol
     if (classSymbol != null) {
         classSymbol.lazyResolveToPhase(FirResolvePhase.TYPES)
         val declaration = classSymbol.fir
         if (declaration !is FirTypeAlias || typeArguments.isEmpty()) {
-            val typeByDeclaration = typeForQualifierByDeclaration(declaration, session)
+            val typeByDeclaration = typeForQualifierByDeclaration(declaration, components.session, element = this, components.file)
             if (typeByDeclaration != null) {
                 this.resultType = typeByDeclaration
                 replaceCanBeValue(true)
@@ -236,7 +236,7 @@ fun FirResolvedQualifier.setTypeOfQualifier(session: FirSession) {
             }
         }
     }
-    this.resultType = session.builtinTypes.unitType.type
+    this.resultType = components.session.builtinTypes.unitType.type
 }
 
 internal fun typeForReifiedParameterReference(parameterReferenceBuilder: FirResolvedReifiedParameterReferenceBuilder): ConeLookupTagBasedType {
@@ -244,10 +244,12 @@ internal fun typeForReifiedParameterReference(parameterReferenceBuilder: FirReso
     return typeParameterSymbol.constructType(emptyArray(), false)
 }
 
-internal fun typeForQualifierByDeclaration(declaration: FirDeclaration, session: FirSession): ConeKotlinType? {
+internal fun typeForQualifierByDeclaration(
+    declaration: FirDeclaration, session: FirSession, element: FirElement, file: FirFile
+): ConeKotlinType? {
     if (declaration is FirTypeAlias) {
         val expandedDeclaration = declaration.expandedConeType?.lookupTag?.toSymbol(session)?.fir ?: return null
-        return typeForQualifierByDeclaration(expandedDeclaration, session)
+        return typeForQualifierByDeclaration(expandedDeclaration, session, element, file)
     }
     if (declaration is FirRegularClass) {
         if (declaration.classKind == ClassKind.OBJECT) {
@@ -255,6 +257,7 @@ internal fun typeForQualifierByDeclaration(declaration: FirDeclaration, session:
         } else {
             val companionObjectSymbol = declaration.companionObjectSymbol
             if (companionObjectSymbol != null) {
+                session.lookupTracker?.recordCompanionLookup(companionObjectSymbol.classId, element.source, file.source)
                 return companionObjectSymbol.constructType(emptyArray(), false)
             }
         }

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/calls/tower/FirTowerResolveTask.kt

@@ -119,7 +119,7 @@ internal abstract class FirBaseTowerResolveTask(
                 this.symbol = it
                 this.source = source?.fakeElement(KtFakeSourceElementKind.ImplicitReceiver)
             }.apply {
-                setTypeOfQualifier(components.session)
+                setTypeOfQualifier(components)
             }
             ExpressionReceiverValue(resolvedQualifier)
         }

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/calls/tower/TowerLevels.kt

@@ -363,7 +363,7 @@ internal class ScopeTowerLevel(
             this.symbol = this@toResolvedQualifierExpressionReceiver
             this.source = source?.fakeElement(KtFakeSourceElementKind.ImplicitReceiver)
         }.apply {
-            setTypeOfQualifier(bodyResolveComponents.session)
+            setTypeOfQualifier(bodyResolveComponents)
         }
         return ExpressionReceiverValue(resolvedQualifier)
     }

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/providers/impl/FirQualifierResolverImpl.kt

@@ -18,7 +18,10 @@ import org.jetbrains.kotlin.name.FqName
 @NoMutableState
 class FirQualifierResolverImpl(val session: FirSession) : FirQualifierResolver() {
 
-    override fun resolveSymbolWithPrefix(parts: List<FirQualifierPart>, prefix: ClassId): FirClassifierSymbol<*>? {
+    override fun resolveSymbolWithPrefix(
+        parts: List<FirQualifierPart>, prefix: ClassId
+    ): FirClassifierSymbol<*>? {
+
         val symbolProvider = session.symbolProvider
 
         val fqName = ClassId(

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/providers/impl/FirTypeResolverImpl.kt

@@ -118,6 +118,10 @@ class FirTypeResolverImpl(private val session: FirSession) : FirTypeResolver() {
             }
         }
 
+        session.lookupTracker?.recordUserTypeRefLookup(
+            typeRef, scopes.asSequence().flatMap { it.scopeOwnerLookupNames }.asIterable(), useSiteFile?.source
+        )
+
         for (scope in scopes) {
             if (applicability == CandidateApplicability.RESOLVED) break
             scope.processClassifiersByNameWithSubstitution(qualifier.first().name) { symbol, substitutorFromScope ->
@@ -205,7 +209,7 @@ class FirTypeResolverImpl(private val session: FirSession) : FirTypeResolver() {
     @OptIn(SymbolInternals::class)
     private fun FirQualifierResolver.resolveEnumEntrySymbol(
         qualifier: List<FirQualifierPart>,
-        classId: ClassId
+        classId: ClassId,
     ): FirVariableSymbol<FirEnumEntry>? {
         // Assuming the current qualifier refers to an enum entry, we drop the last part so we get a reference to the enum class.
         val enumClassSymbol = resolveSymbolWithPrefix(qualifier.dropLast(1), classId) ?: return null
@@ -453,6 +457,8 @@ class FirTypeResolverImpl(private val session: FirSession) : FirTypeResolver() {
                 }
             }
             else -> error(typeRef.render())
+        }.also {
+            session.lookupTracker?.recordTypeResolveAsLookup(it.type, typeRef.source, useSiteFile?.source)
         }
     }
 

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/FirSpecificTypeResolverTransformer.kt

@@ -15,8 +15,6 @@ import org.jetbrains.kotlin.fir.diagnostics.ConeDiagnostic
 import org.jetbrains.kotlin.fir.diagnostics.ConeSimpleDiagnostic
 import org.jetbrains.kotlin.fir.diagnostics.ConeUnexpectedTypeArgumentsError
 import org.jetbrains.kotlin.fir.expressions.FirStatement
-import org.jetbrains.kotlin.fir.lookupTracker
-import org.jetbrains.kotlin.fir.recordTypeLookup
 import org.jetbrains.kotlin.fir.render
 import org.jetbrains.kotlin.fir.resolve.FirTypeResolutionResult
 import org.jetbrains.kotlin.fir.resolve.SupertypeSupplier
@@ -87,11 +85,11 @@ class FirSpecificTypeResolverTransformer(
 
     @OptIn(PrivateForInline::class)
     override fun transformTypeRef(typeRef: FirTypeRef, data: ScopeClassDeclaration): FirResolvedTypeRef {
-        session.lookupTracker?.recordTypeLookup(typeRef, data.scopes.flatMap { it.scopeOwnerLookupNames }, currentFile?.source)
         withBareTypes(allowed = false) {
             typeRef.transformChildren(this, data)
         }
         val (resolvedType, diagnostic) = resolveType(typeRef, data)
+
         return transformType(typeRef, resolvedType, diagnostic, data)
     }
 
@@ -101,7 +99,6 @@ class FirSpecificTypeResolverTransformer(
         data: ScopeClassDeclaration
     ): FirResolvedTypeRef {
         functionTypeRef.transformChildren(this, data)
-        session.lookupTracker?.recordTypeLookup(functionTypeRef, data.scopes.flatMap { it.scopeOwnerLookupNames }, currentFile?.source)
         val resolvedTypeWithDiagnostic = resolveType(functionTypeRef, data)
         val resolvedType = resolvedTypeWithDiagnostic.type.takeIfAcceptable()
         val diagnostic = resolvedTypeWithDiagnostic.diagnostic
@@ -254,6 +251,7 @@ class FirSpecificTypeResolverTransformer(
             val typeRefToTry = buildUserTypeRef {
                 qualifier += qualifiersToTry
                 isMarkedNullable = false
+                source = typeRef.source
             }
             val (resolvedType, diagnostic) = resolveType(typeRefToTry, data)
             if (resolvedType is ConeErrorType || diagnostic != null) continue

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/FirSupertypesResolution.kt

@@ -8,7 +8,8 @@ package org.jetbrains.kotlin.fir.resolve.transformers
 import kotlinx.collections.immutable.PersistentList
 import kotlinx.collections.immutable.persistentListOf
 import kotlinx.collections.immutable.toPersistentList
-import org.jetbrains.kotlin.fir.*
+import org.jetbrains.kotlin.fir.FirElement
+import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.declarations.*
 import org.jetbrains.kotlin.fir.declarations.utils.expandedConeType
 import org.jetbrains.kotlin.fir.declarations.utils.isCompanion
@@ -20,6 +21,7 @@ import org.jetbrains.kotlin.fir.expressions.FirStatement
 import org.jetbrains.kotlin.fir.extensions.FirSupertypeGenerationExtension
 import org.jetbrains.kotlin.fir.extensions.extensionService
 import org.jetbrains.kotlin.fir.extensions.supertypeGenerators
+import org.jetbrains.kotlin.fir.render
 import org.jetbrains.kotlin.fir.resolve.*
 import org.jetbrains.kotlin.fir.resolve.dfa.cfg.isLocalClassOrAnonymousObject
 import org.jetbrains.kotlin.fir.resolve.diagnostics.ConeTypeParameterSupertype
@@ -41,7 +43,7 @@ import org.jetbrains.kotlin.fir.types.builder.buildErrorTypeRef
 import org.jetbrains.kotlin.fir.types.impl.FirImplicitBuiltinTypeRef
 import org.jetbrains.kotlin.fir.visitors.FirDefaultTransformer
 import org.jetbrains.kotlin.fir.visitors.FirDefaultVisitor
-import org.jetbrains.kotlin.fir.visitors.FirTransformer
+import org.jetbrains.kotlin.fir.withFileAnalysisExceptionWrapping
 import org.jetbrains.kotlin.name.ClassId
 import org.jetbrains.kotlin.types.model.TypeArgumentMarker
 import org.jetbrains.kotlin.util.PrivateForInline
@@ -450,16 +452,6 @@ open class FirSupertypeResolverVisitor(
         supertypeRefs: List<FirTypeRef>,
     ): List<FirResolvedTypeRef> {
         return resolveSpecificClassLikeSupertypes(classLikeDeclaration) { transformer, scopeDeclaration ->
-            if (!classLikeDeclaration.isLocalClassOrAnonymousObject()) {
-                session.lookupTracker?.let {
-                    val fileSource = getFirClassifierContainerFileIfAny(classLikeDeclaration.symbol)?.source
-                    val scopeOwnerLookupNames = scopeDeclaration.scopes.flatMap { scope -> scope.scopeOwnerLookupNames }
-                    for (supertypeRef in supertypeRefs) {
-                        it.recordTypeLookup(supertypeRef, scopeOwnerLookupNames, fileSource)
-                    }
-                }
-            }
-
             supertypeRefs.mapTo(mutableListOf()) {
                 val superTypeRef = it.transform<FirTypeRef, ScopeClassDeclaration>(transformer, scopeDeclaration)
                 val typeParameterType = superTypeRef.coneTypeSafe<ConeTypeParameterType>()

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/body/resolve/FirExpressionsResolveTransformer.kt

@@ -158,16 +158,25 @@ open class FirExpressionsResolveTransformer(transformer: FirAbstractBodyResolveT
                     },
                     data,
                 )
+
+                fun FirStatement.alsoRecordLookup() = also {
+                    if (transformedCallee is FirExpression && transformedCallee.isResolved) {
+                        session.lookupTracker?.recordTypeResolveAsLookup(
+                            transformedCallee.resolvedType, callee.source, components.file.source
+                        )
+                    }
+                }
+
                 // NB: here we can get raw expression because of dropped qualifiers (see transform callee),
                 // so candidate existence must be checked before calling completion
                 if (transformedCallee is FirQualifiedAccessExpression && transformedCallee.candidate() != null) {
                     if (!transformedCallee.isAcceptableResolvedQualifiedAccess()) {
-                        return qualifiedAccessExpression
+                        return qualifiedAccessExpression.alsoRecordLookup()
                     }
                     callCompleter.completeCall(transformedCallee, data)
                 } else {
                     transformedCallee
-                }
+                }.alsoRecordLookup()
             }
         }
 

compiler/fir/semantics/src/org/jetbrains/kotlin/fir/FirLookupTrackerComponent.kt

@@ -40,20 +40,22 @@ fun FirLookupTrackerComponent.recordCallLookup(callInfo: AbstractCallInfo, inSco
     recordNameLookup(callInfo.name, inScopes, callInfo.callSite.source, callInfo.containingFile.source)
 }
 
-fun FirLookupTrackerComponent.recordTypeLookup(typeRef: FirTypeRef, inScopes: Iterable<String>, fileSource: KtSourceElement?) {
-    if (typeRef is FirUserTypeRef) recordNameLookup(typeRef.qualifier.first().name, inScopes, typeRef.source, fileSource)
-}
-
 fun FirLookupTrackerComponent.recordClassLikeLookup(classId: ClassId, source: KtSourceElement?, fileSource: KtSourceElement?) {
     if (!classId.isLocal && classId !in StandardClassIds.allBuiltinTypes) {
-        recordLookup(classId.relativeClassName.asString(), classId.packageFqName.asString(), source, fileSource)
-        if (classId.shortClassName == DEFAULT_NAME_FOR_COMPANION_OBJECT) {
-            recordLookup(
-                classId.outerClassId!!.relativeClassName.asString(),
-                classId.outerClassId!!.packageFqName.asString(),
-                source, fileSource
-            )
-        }
+        val classFqName = classId.asSingleFqName()
+        recordLookup(classFqName.shortName().asString(), classFqName.parent().asString(), source, fileSource)
+    }
+}
+
+fun FirLookupTrackerComponent.recordCompanionLookup(classId: ClassId, source: KtSourceElement?, fileSource: KtSourceElement?) {
+    if (!classId.isLocal && classId !in StandardClassIds.allBuiltinTypes) {
+        val classFqName = classId.asSingleFqName()
+        recordLookup(classFqName.shortName().asString(), classFqName.parent().asString(), source, fileSource)
+        recordLookup(
+            classFqName.parent().shortName().asString(),
+            classFqName.parent().parent().asString(),
+            source, fileSource
+        )
     }
 }
 
@@ -82,6 +84,15 @@ fun FirLookupTrackerComponent.recordTypeResolveAsLookup(typeRef: FirTypeRef, sou
     recordTypeResolveAsLookup(typeRef.type, source, fileSource)
 }
 
+fun FirLookupTrackerComponent.recordUserTypeRefLookup(typeRef: FirUserTypeRef, inScopes: Iterable<String>, fileSource: KtSourceElement?) {
+    inScopes.forEach { scope ->
+        typeRef.qualifier.fold(FqName(scope)) { result, suffix ->
+            recordLookup(suffix.name.asString(), result.asString(), typeRef.source, fileSource)
+            result.child(suffix.name)
+        }
+    }
+}
+
 // TODO: review all places that record resolved type as lookup and consider minimize the number of them; see #KT-66366
 fun FirLookupTrackerComponent.recordTypeResolveAsLookup(type: ConeKotlinType?, source: KtSourceElement?, fileSource: KtSourceElement?) {
     if (type == null) return