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.
2021-12-24 15:37:24 +01:00
parent 11d6dccf51
commit 2aba465494
8 changed files with 47 additions and 29 deletions
@@ -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) }
@@ -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? {
@@ -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
            )

@@ -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
+            )
        }
    }

@@ -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(resolvedTypeRef, anonymousFunction.source, null)
+                it.recordTypeResolveAsLookup(newReturnTypeRef, anonymousFunction.source, context.file.source)
+                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) {
@@ -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()
@@ -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
@@ -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)