K2 IC: fix classifiers lookups

in particular nested classifiers and companions.

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

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()
             }
         }