FIR IC: provide file sources to lookup where possible

because detecting file source by element source is possible for LT
only in some limited cases.

compiler/fir/entrypoint/src/org/jetbrains/kotlin/fir/IncrementalPassThroughLookupTrackerComponent.kt

@@ -24,10 +24,13 @@ class IncrementalPassThroughLookupTrackerComponent(
     private val sourceToFilePathsCache = ConcurrentHashMap<KtSourceElement, String>()
 
     override fun recordLookup(name: Name, inScopes: List<String>, source: KtSourceElement?, fileSource: KtSourceElement?) {
+        assert(fileSource != null || source is KtPsiSourceElement) // finding file for a source only possible for PSI, here it means
+        // that we allow null for file source only for PSI-only "sources", currently - java ones
         val definedSource = fileSource ?: source ?: throw AssertionError("Cannot record lookup for \"$name\" without a source")
         val path = sourceToFilePathsCache.getOrPut(definedSource) {
             sourceToFilePath(definedSource) ?:
-            return
+                // TODO: the lookup by non-file source mostly doesn't work for the LT, so we cannot afford null file sources here
+                return
         }
         val position = if (requiresPosition && source != null && source is KtPsiSourceElement) {
             getLineAndColumnInPsiFile(source.psi.containingFile, source.psi.textRange).let { Position(it.line, it.column) }

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

@@ -374,7 +374,11 @@ private inline fun <T : FirExpression> BodyResolveComponents.transformExpression
     }
 }
 
-fun FirCheckedSafeCallSubject.propagateTypeFromOriginalReceiver(nullableReceiverExpression: FirExpression, session: FirSession) {
+fun FirCheckedSafeCallSubject.propagateTypeFromOriginalReceiver(
+    nullableReceiverExpression: FirExpression,
+    session: FirSession,
+    file: FirFile
+) {
     // If the receiver expression is smartcast to `null`, it would have `Nothing?` as its type, which may not have members called by user
     // code. Hence, we fallback to the type before intersecting with `Nothing?`.
     val receiverType = ((nullableReceiverExpression as? FirExpressionWithSmartcastToNull)
@@ -388,12 +392,13 @@ fun FirCheckedSafeCallSubject.propagateTypeFromOriginalReceiver(nullableReceiver
     val resolvedTypeRef =
         typeRef.resolvedTypeFromPrototype(expandedReceiverType.makeConeTypeDefinitelyNotNullOrNotNull(session.typeContext))
     replaceTypeRef(resolvedTypeRef)
-    session.lookupTracker?.recordTypeResolveAsLookup(resolvedTypeRef, source, null)
+    session.lookupTracker?.recordTypeResolveAsLookup(resolvedTypeRef, source, file.source)
 }
 
 fun FirSafeCallExpression.propagateTypeFromQualifiedAccessAfterNullCheck(
     nullableReceiverExpression: FirExpression,
     session: FirSession,
+    file: FirFile,
 ) {
     val receiverType = nullableReceiverExpression.typeRef.coneTypeSafe<ConeKotlinType>()
     val typeAfterNullCheck = selector.expressionTypeOrUnitForAssignment() ?: return
@@ -410,7 +415,7 @@ fun FirSafeCallExpression.propagateTypeFromQualifiedAccessAfterNullCheck(
 
     val resolvedTypeRef = typeRef.resolvedTypeFromPrototype(resultingType)
     replaceTypeRef(resolvedTypeRef)
-    session.lookupTracker?.recordTypeResolveAsLookup(resolvedTypeRef, source, null)
+    session.lookupTracker?.recordTypeResolveAsLookup(resolvedTypeRef, source, file.source)
 }
 
 private fun FirStatement.expressionTypeOrUnitForAssignment(): ConeKotlinType? {

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/inference/FirCallCompleter.kt

@@ -97,7 +97,7 @@ class FirCallCompleter(
         if (call is FirExpression) {
             val resolvedTypeRef = typeRef.resolvedTypeFromPrototype(initialType)
             call.resultType = resolvedTypeRef
-            session.lookupTracker?.recordTypeResolveAsLookup(resolvedTypeRef, call.source, null)
+            session.lookupTracker?.recordTypeResolveAsLookup(resolvedTypeRef, call.source, components.context.file.source)
         }
 
         addConstraintFromExpectedType(
@@ -127,7 +127,8 @@ class FirCallCompleter(
                             components.returnTypeCalculator,
                             session.typeApproximator,
                             components.dataFlowAnalyzer,
-                            components.integerLiteralAndOperatorApproximationTransformer
+                            components.integerLiteralAndOperatorApproximationTransformer,
+                            components.context
                         ),
                         null
                     )
@@ -239,6 +240,7 @@ class FirCallCompleter(
             session.typeApproximator,
             components.dataFlowAnalyzer,
             components.integerLiteralAndOperatorApproximationTransformer,
+            components.context,
             mode
         )
     }
@@ -306,19 +308,20 @@ class FirCallCompleter(
             )
 
             val lookupTracker = session.lookupTracker
+            val fileSource = components.file.source
             lambdaArgument.valueParameters.forEachIndexed { index, parameter ->
                 val newReturnType = parameters[index].approximateLambdaInputType()
                 val newReturnTypeRef = if (parameter.returnTypeRef is FirImplicitTypeRef) {
                     newReturnType.toFirResolvedTypeRef(parameter.source)
                 } else parameter.returnTypeRef.resolvedTypeFromPrototype(newReturnType)
                 parameter.replaceReturnTypeRef(newReturnTypeRef)
-                lookupTracker?.recordTypeResolveAsLookup(newReturnTypeRef, parameter.source, null)
+                lookupTracker?.recordTypeResolveAsLookup(newReturnTypeRef, parameter.source, fileSource)
             }
 
             lambdaArgument.replaceValueParameters(lambdaArgument.valueParameters + listOfNotNull(itParam))
             lambdaArgument.replaceReturnTypeRef(
                 expectedReturnTypeRef?.also {
-                    lookupTracker?.recordTypeResolveAsLookup(it, lambdaArgument.source, null)
+                    lookupTracker?.recordTypeResolveAsLookup(it, lambdaArgument.source, fileSource)
                 } ?: components.noExpectedType
             )
 

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/inference/PostponedArgumentsAnalyzer.kt

@@ -102,7 +102,9 @@ class PostponedArgumentsAnalyzer(
                 }
             }
             replaceTypeRef(resolvedTypeRef)
-            resolutionContext.session.lookupTracker?.recordTypeResolveAsLookup(resolvedTypeRef, source, null)
+            resolutionContext.session.lookupTracker?.recordTypeResolveAsLookup(
+                resolvedTypeRef, source, resolutionContext.bodyResolveComponents.file.source
+            )
         }
     }
 

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

@@ -34,6 +34,7 @@ import org.jetbrains.kotlin.fir.resolve.diagnostics.ConeTypeParameterInQualified
 import org.jetbrains.kotlin.fir.resolve.inference.ResolvedLambdaAtom
 import org.jetbrains.kotlin.fir.resolve.providers.symbolProvider
 import org.jetbrains.kotlin.fir.resolve.substitution.ConeSubstitutor
+import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.*
 import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.FirArrayOfCallTransformer
 import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.remapArgumentsWithVararg
 import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.resultType
@@ -68,6 +69,7 @@ class FirCallCompletionResultsWriterTransformer(
     private val typeApproximator: ConeTypeApproximator,
     private val dataFlowAnalyzer: FirDataFlowAnalyzer<*>,
     private val integerOperatorApproximator: IntegerLiteralAndOperatorApproximationTransformer,
+    private val context: BodyResolveContext,
     private val mode: Mode = Mode.Normal
 ) : FirAbstractTreeTransformer<ExpectedArgumentType?>(phase = FirResolvePhase.IMPLICIT_TYPES_BODY_RESOLVE) {
 
@@ -158,7 +160,7 @@ class FirCallCompletionResultsWriterTransformer(
         if (declaration !is FirErrorFunction) {
             result.replaceTypeArguments(typeArguments)
         }
-        session.lookupTracker?.recordTypeResolveAsLookup(typeRef, qualifiedAccessExpression.source, null)
+        session.lookupTracker?.recordTypeResolveAsLookup(typeRef, qualifiedAccessExpression.source, context.file.source)
         return result
     }
 
@@ -181,7 +183,7 @@ class FirCallCompletionResultsWriterTransformer(
         val resultType = typeRef.substituteTypeRef(subCandidate)
         resultType.ensureResolvedTypeDeclaration(session)
         result.replaceTypeRef(resultType)
-        session.lookupTracker?.recordTypeResolveAsLookup(resultType, qualifiedAccessExpression.source, null)
+        session.lookupTracker?.recordTypeResolveAsLookup(resultType, qualifiedAccessExpression.source, context.file.source)
 
         if (mode == Mode.DelegatedPropertyCompletion) {
             subCandidate.symbol.fir.transformSingle(
@@ -237,7 +239,7 @@ class FirCallCompletionResultsWriterTransformer(
         val expectedArgumentsTypeMapping = runIf(!calleeReference.isError) { subCandidate.createArgumentsMapping() }
         result.argumentList.transformArguments(this, expectedArgumentsTypeMapping)
         result.replaceTypeRef(resultType)
-        session.lookupTracker?.recordTypeResolveAsLookup(resultType, functionCall.source, null)
+        session.lookupTracker?.recordTypeResolveAsLookup(resultType, functionCall.source, context.file.source)
 
         if (mode == Mode.DelegatedPropertyCompletion) {
             subCandidate.symbol.fir.transformSingle(
@@ -316,7 +318,7 @@ class FirCallCompletionResultsWriterTransformer(
     ): D {
         val resultTypeRef = typeRef.substituteTypeRef(calleeReference.candidate)
         replaceTypeRef(resultTypeRef)
-        session.lookupTracker?.recordTypeResolveAsLookup(resultTypeRef, source, null)
+        session.lookupTracker?.recordTypeResolveAsLookup(resultTypeRef, source, context.file.source)
         return this
     }
 
@@ -358,7 +360,7 @@ class FirCallCompletionResultsWriterTransformer(
             )?.toExpectedType()
         )
 
-        safeCallExpression.propagateTypeFromQualifiedAccessAfterNullCheck(safeCallExpression.receiver, session)
+        safeCallExpression.propagateTypeFromQualifiedAccessAfterNullCheck(safeCallExpression.receiver, session, context.file)
 
         return safeCallExpression
     }
@@ -380,7 +382,7 @@ class FirCallCompletionResultsWriterTransformer(
         val resultType = typeRef.withReplacedConeType(finalType)
         callableReferenceAccess.replaceTypeRef(resultType)
         callableReferenceAccess.replaceTypeArguments(typeArguments)
-        session.lookupTracker?.recordTypeResolveAsLookup(resultType, typeRef.source ?: callableReferenceAccess.source, null)
+        session.lookupTracker?.recordTypeResolveAsLookup(resultType, typeRef.source ?: callableReferenceAccess.source, context.file.source)
 
         val resolvedReference = when (calleeReference) {
             is FirErrorReferenceWithCandidate -> buildErrorNamedReference {
@@ -417,7 +419,7 @@ class FirCallCompletionResultsWriterTransformer(
         val resultLValueType = lValueTypeRef.substituteTypeRef(calleeReference.candidate)
         resultLValueType.ensureResolvedTypeDeclaration(session)
         variableAssignment.replaceLValueTypeRef(resultLValueType)
-        session.lookupTracker?.recordTypeResolveAsLookup(resultLValueType, variableAssignment.lValue.source, null)
+        session.lookupTracker?.recordTypeResolveAsLookup(resultLValueType, variableAssignment.lValue.source, context.file.source)
 
         return variableAssignment.transformCalleeReference(
             StoreCalleeReference,
@@ -438,7 +440,7 @@ class FirCallCompletionResultsWriterTransformer(
             val substitutedReceiverType = finalSubstitutor.substituteOrNull(originalType) ?: return qualifiedAccessExpression
             val resolvedTypeRef = qualifiedAccessExpression.typeRef.resolvedTypeFromPrototype(substitutedReceiverType)
             qualifiedAccessExpression.replaceTypeRef(resolvedTypeRef)
-            session.lookupTracker?.recordTypeResolveAsLookup(resolvedTypeRef, qualifiedAccessExpression.source, null)
+            session.lookupTracker?.recordTypeResolveAsLookup(resolvedTypeRef, qualifiedAccessExpression.source, context.file.source)
             return qualifiedAccessExpression
         }
 
@@ -639,7 +641,7 @@ class FirCallCompletionResultsWriterTransformer(
                             (expectedType == null && anonymousFunction.isSuspendFunctionType())
                 )
             anonymousFunction.replaceTypeRef(resolvedTypeRef)
-            session.lookupTracker?.recordTypeResolveAsLookup(resolvedTypeRef, anonymousFunction.source, null)
+            session.lookupTracker?.recordTypeResolveAsLookup(resolvedTypeRef, anonymousFunction.source, context.file.source)
         }
 
         val result = transformElement(anonymousFunction, null)
@@ -662,8 +664,8 @@ class FirCallCompletionResultsWriterTransformer(
                 resultFunction.constructFunctionalTypeRef(isSuspend = expectedType?.isSuspendFunctionType(session) == true)
             resultFunction.replaceTypeRef(resolvedTypeRef)
             session.lookupTracker?.let {
-                it.recordTypeResolveAsLookup(newReturnTypeRef, anonymousFunction.source, null)
+                it.recordTypeResolveAsLookup(newReturnTypeRef, anonymousFunction.source, context.file.source)
-                it.recordTypeResolveAsLookup(resolvedTypeRef, anonymousFunction.source, null)
+                it.recordTypeResolveAsLookup(resolvedTypeRef, anonymousFunction.source, context.file.source)
             }
         }
 
@@ -722,7 +724,7 @@ class FirCallCompletionResultsWriterTransformer(
                 resultType = resultType.resolvedTypeFromPrototype(it.getApproximatedType(data?.getExpectedType(block)))
             }
             block.replaceTypeRef(resultType)
-            session.lookupTracker?.recordTypeResolveAsLookup(resultType, block.source, null)
+            session.lookupTracker?.recordTypeResolveAsLookup(resultType, block.source, context.file.source)
         }
         transformElement(block, data)
         if (block.resultType is FirErrorTypeRef) {

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

@@ -249,7 +249,7 @@ open class FirDeclarationsResolveTransformer(transformer: FirBodyResolveTransfor
                         isNullable = false
                     )
                 }.also {
-                    session.lookupTracker?.recordTypeResolveAsLookup(it, propertyReferenceAccess.source ?: source, null)
+                    session.lookupTracker?.recordTypeResolveAsLookup(it, propertyReferenceAccess.source ?: source, components.file.source)
                 }
             )
         }
@@ -829,7 +829,8 @@ open class FirDeclarationsResolveTransformer(transformer: FirBodyResolveTransfor
             components.returnTypeCalculator,
             session.typeApproximator,
             dataFlowAnalyzer,
-            components.integerLiteralAndOperatorApproximationTransformer
+            components.integerLiteralAndOperatorApproximationTransformer,
+            components.context
         )
         lambda.transformSingle(writer, expectedTypeRef.coneTypeSafe<ConeKotlinType>()?.toExpectedType())
 
@@ -861,14 +862,14 @@ open class FirDeclarationsResolveTransformer(transformer: FirBodyResolveTransfor
         }
         lambda.replaceReturnTypeRef(
             lambda.returnTypeRef.resolvedTypeFromPrototype(returnType).also {
-                session.lookupTracker?.recordTypeResolveAsLookup(it, lambda.source, null)
+                session.lookupTracker?.recordTypeResolveAsLookup(it, lambda.source, components.file.source)
             }
         )
         lambda.replaceTypeRef(
             lambda.constructFunctionalTypeRef(
                 isSuspend = expectedTypeRef.coneTypeSafe<ConeKotlinType>()?.isSuspendFunctionType(session) == true
             ).also {
-                session.lookupTracker?.recordTypeResolveAsLookup(it, lambda.source, null)
+                session.lookupTracker?.recordTypeResolveAsLookup(it, lambda.source, components.file.source)
             }
         )
         return lambda.addReturn()

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

@@ -317,9 +317,9 @@ open class FirExpressionsResolveTransformer(transformer: FirBodyResolveTransform
             dataFlowAnalyzer.enterSafeCallAfterNullCheck(safeCallExpression)
 
             safeCallExpression.apply {
-                checkedSubjectRef.value.propagateTypeFromOriginalReceiver(receiver, components.session)
+                checkedSubjectRef.value.propagateTypeFromOriginalReceiver(receiver, components.session, components.file)
                 transformSelector(this@FirExpressionsResolveTransformer, data)
-                propagateTypeFromQualifiedAccessAfterNullCheck(receiver, session)
+                propagateTypeFromQualifiedAccessAfterNullCheck(receiver, session, context.file)
             }
 
             dataFlowAnalyzer.exitSafeCall(safeCallExpression)
@@ -927,7 +927,7 @@ open class FirExpressionsResolveTransformer(transformer: FirBodyResolveTransform
                 if (typeRef != null) {
                     lhs.replaceTypeRef(
                         buildResolvedTypeRef { type = typeRef }.also {
-                            session.lookupTracker?.recordTypeResolveAsLookup(it, getClassCall.source, null)
+                            session.lookupTracker?.recordTypeResolveAsLookup(it, getClassCall.source, components.file.source)
                         }
                     )
                     typeRef

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

@@ -245,7 +245,9 @@ private class ReturnTypeCalculatorWithJump(
                 val baseDeclaration = baseSymbol.fir as FirTypedDeclaration
                 val baseReturnTypeRef = tryCalculateReturnType(baseDeclaration)
                 val baseReturnType = baseReturnTypeRef.type
-                session.lookupTracker?.recordTypeResolveAsLookup(baseReturnTypeRef, declaration.source, null)
+                session.lookupTracker?.recordTypeResolveAsLookup(
+                    baseReturnTypeRef, declaration.source, session.firProvider.getFirCallableContainerFile(declaration.symbol)?.source
+                )
                 val coneType = substitutor.substituteOrSelf(baseReturnType)
                 val returnType = declaration.returnTypeRef.resolvedTypeFromPrototype(coneType)
                 declaration.replaceReturnTypeRef(returnType)