diff --git a/compiler/fir/entrypoint/src/org/jetbrains/kotlin/fir/IncrementalPassThroughLookupTrackerComponent.kt b/compiler/fir/entrypoint/src/org/jetbrains/kotlin/fir/IncrementalPassThroughLookupTrackerComponent.kt
index 1b0e6c108b1..e00d3554ecf 100644
--- a/compiler/fir/entrypoint/src/org/jetbrains/kotlin/fir/IncrementalPassThroughLookupTrackerComponent.kt
+++ b/compiler/fir/entrypoint/src/org/jetbrains/kotlin/fir/IncrementalPassThroughLookupTrackerComponent.kt
@@ -17,7 +17,7 @@ import java.util.concurrent.ConcurrentHashMap
 
 class IncrementalPassThroughLookupTrackerComponent(
     private val lookupTracker: LookupTracker,
-    private val sourceToFilePath: (KtSourceElement) -> String
+    private val sourceToFilePath: (KtSourceElement) -> String?
 ) : FirLookupTrackerComponent() {
 
     private val requiresPosition = lookupTracker.requiresPosition
@@ -26,7 +26,8 @@ class IncrementalPassThroughLookupTrackerComponent(
     override fun recordLookup(name: Name, inScopes: List<String>, source: KtSourceElement?, fileSource: KtSourceElement?) {
         val definedSource = fileSource ?: source ?: throw AssertionError("Cannot record lookup for \"$name\" without a source")
         val path = sourceToFilePathsCache.getOrPut(definedSource) {
-            sourceToFilePath(definedSource)
+            sourceToFilePath(definedSource) ?:
+            return
         }
         val position = if (requiresPosition && source != null && source is KtPsiSourceElement) {
             getLineAndColumnInPsiFile(source.psi.containingFile, source.psi.textRange).let { Position(it.line, it.column) }
diff --git a/compiler/fir/entrypoint/src/org/jetbrains/kotlin/fir/pipeline/buildFir.kt b/compiler/fir/entrypoint/src/org/jetbrains/kotlin/fir/pipeline/buildFir.kt
index d8b20ba8dad..bb5b32eb709 100644
--- a/compiler/fir/entrypoint/src/org/jetbrains/kotlin/fir/pipeline/buildFir.kt
+++ b/compiler/fir/entrypoint/src/org/jetbrains/kotlin/fir/pipeline/buildFir.kt
@@ -12,15 +12,18 @@ import org.jetbrains.kotlin.fir.declarations.FirFile
 import org.jetbrains.kotlin.fir.lightTree.LightTree2Fir
 import org.jetbrains.kotlin.fir.resolve.providers.firProvider
 import org.jetbrains.kotlin.fir.resolve.providers.impl.FirProviderImpl
+import org.jetbrains.kotlin.fir.session.sourcesToPathsMapper
 import org.jetbrains.kotlin.psi.KtFile
 import java.io.File
 
 fun FirSession.buildFirViaLightTree(files: Collection<File>): List<FirFile> {
     val firProvider = (firProvider as FirProviderImpl)
+    val sourcesToPathsMapper = sourcesToPathsMapper
     val builder = LightTree2Fir(this, firProvider.kotlinScopeProvider)
     return files.map {
         builder.buildFirFile(it).also { firFile ->
             firProvider.recordFile(firFile)
+            sourcesToPathsMapper.registerFileSource(firFile.source!!, it.path)
         }
     }
 }
diff --git a/compiler/fir/entrypoint/src/org/jetbrains/kotlin/fir/session/ComponentsContainers.kt b/compiler/fir/entrypoint/src/org/jetbrains/kotlin/fir/session/ComponentsContainers.kt
index 0c91ca6be60..3af9bc04bc9 100644
--- a/compiler/fir/entrypoint/src/org/jetbrains/kotlin/fir/session/ComponentsContainers.kt
+++ b/compiler/fir/entrypoint/src/org/jetbrains/kotlin/fir/session/ComponentsContainers.kt
@@ -6,6 +6,7 @@
 package org.jetbrains.kotlin.fir.session
 
 import com.intellij.psi.PsiFile
+import org.jetbrains.kotlin.KtNodeTypes
 import org.jetbrains.kotlin.KtPsiSourceElement
 import org.jetbrains.kotlin.KtSourceElement
 import org.jetbrains.kotlin.config.JvmAnalysisFlags
@@ -29,6 +30,7 @@ import org.jetbrains.kotlin.fir.resolve.calls.ConeCallConflictResolverFactory
 import org.jetbrains.kotlin.fir.resolve.calls.FirSyntheticNamesProvider
 import org.jetbrains.kotlin.fir.resolve.calls.jvm.JvmCallConflictResolverFactory
 import org.jetbrains.kotlin.fir.resolve.inference.InferenceComponents
+import org.jetbrains.kotlin.fir.resolve.providers.firProvider
 import org.jetbrains.kotlin.fir.resolve.providers.impl.FirQualifierResolverImpl
 import org.jetbrains.kotlin.fir.resolve.providers.impl.FirTypeResolverImpl
 import org.jetbrains.kotlin.fir.resolve.transformers.FirPhaseCheckingPhaseManager
@@ -96,10 +98,10 @@ fun FirSession.registerResolveComponents(lookupTracker: LookupTracker? = null) {
     register(CheckersComponent::class, CheckersComponent())
     register(FirNameConflictsTrackerComponent::class, FirNameConflictsTracker())
     register(FirModuleVisibilityChecker::class, FirModuleVisibilityChecker.Standard(this))
+    register(SourcesToPathsMapper::class, SourcesToPathsMapper())
     if (lookupTracker != null) {
-        val firFileToPath: (KtSourceElement) -> String = {
-            val psiSource = (it as? KtPsiSourceElement) ?: TODO("Not implemented for non-FirPsiSourceElement")
-            ((psiSource.psi as? PsiFile) ?: psiSource.psi.containingFile).virtualFile.path
+        val firFileToPath: (KtSourceElement) -> String? = {
+            sourcesToPathsMapper.getSourceFilePath(it)
         }
         register(
             FirLookupTrackerComponent::class,
diff --git a/compiler/fir/entrypoint/src/org/jetbrains/kotlin/fir/session/SourcesToPathsMapper.kt b/compiler/fir/entrypoint/src/org/jetbrains/kotlin/fir/session/SourcesToPathsMapper.kt
new file mode 100644
index 00000000000..e0e7cb9801c
--- /dev/null
+++ b/compiler/fir/entrypoint/src/org/jetbrains/kotlin/fir/session/SourcesToPathsMapper.kt
@@ -0,0 +1,31 @@
+/*
+ * Copyright 2010-2021 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
+ */
+
+package org.jetbrains.kotlin.fir.session
+
+import com.intellij.lang.LighterASTNode
+import org.jetbrains.kotlin.KtPsiSourceElement
+import org.jetbrains.kotlin.KtSourceElement
+import org.jetbrains.kotlin.fir.FirSession
+import org.jetbrains.kotlin.fir.FirSessionComponent
+import org.jetbrains.kotlin.psi
+
+class SourcesToPathsMapper : FirSessionComponent {
+
+    private val sourcesToPath = mutableMapOf<LighterASTNode, String>()
+
+    fun registerFileSource(sourceElement: KtSourceElement, path: String) {
+        if (sourceElement !is KtPsiSourceElement) {
+            sourcesToPath[sourceElement.treeStructure.root] = path
+        }
+    }
+
+    fun getSourceFilePath(sourceElement: KtSourceElement): String? {
+        val psi = sourceElement.psi
+        return if (psi != null) psi.containingFile?.virtualFile?.path else sourcesToPath[sourceElement.treeStructure.root]
+    }
+}
+
+val FirSession.sourcesToPathsMapper: SourcesToPathsMapper by FirSession.sessionComponentAccessor()