FIR: Implement lookup tracking

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

@@ -0,0 +1,41 @@
+/*
+ * Copyright 2010-2020 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
+
+import org.jetbrains.kotlin.diagnostics.DiagnosticUtils.getLineAndColumnInPsiFile
+import org.jetbrains.kotlin.incremental.components.LookupTracker
+import org.jetbrains.kotlin.incremental.components.Position
+import org.jetbrains.kotlin.incremental.components.ScopeKind
+import org.jetbrains.kotlin.name.Name
+import org.jetbrains.kotlin.utils.SmartList
+import java.util.concurrent.ConcurrentHashMap
+
+class IncrementalPassThroughLookupTrackerComponent(
+    private val lookupTracker: LookupTracker,
+    private val sourceToFilePath: (FirSourceElement) -> String
+) : FirLookupTrackerComponent() {
+
+    private val requiresPosition = lookupTracker.requiresPosition
+    private val sourceToFilePathsCache = ConcurrentHashMap<FirSourceElement, String>()
+
+    override fun recordLookup(name: Name, inScopes: List<String>, source: FirSourceElement?, fileSource: FirSourceElement?) {
+        val definedSource = fileSource ?: source ?: throw AssertionError("Cannot record lookup for \"$name\" without a source")
+        val path = sourceToFilePathsCache.getOrPut(definedSource) {
+            sourceToFilePath(definedSource)
+        }
+        val position = if (requiresPosition && source != null && source is FirPsiSourceElement<*>) {
+            getLineAndColumnInPsiFile(source.psi.containingFile, source.psi.textRange).let { Position(it.line, it.column) }
+        } else Position.NO_POSITION
+
+        for (scope in inScopes) {
+            lookupTracker.record(path, position, scope, ScopeKind.PACKAGE, name.asString())
+        }
+    }
+
+    override fun recordLookup(name: Name, inScope: String, source: FirSourceElement?, fileSource: FirSourceElement?) {
+        recordLookup(name, SmartList(inScope), source, fileSource)
+    }
+}

compiler/fir/entrypoint/src/org/jetbrains/kotlin/fir/session/ComponentsContainers.kt

@@ -5,6 +5,7 @@
 
 package org.jetbrains.kotlin.fir.session
 
+import com.intellij.psi.PsiFile
 import org.jetbrains.kotlin.config.LanguageVersionSettings
 import org.jetbrains.kotlin.fir.*
 import org.jetbrains.kotlin.fir.analysis.CheckersComponent
@@ -25,6 +26,7 @@ import org.jetbrains.kotlin.fir.resolve.providers.impl.FirTypeResolverImpl
 import org.jetbrains.kotlin.fir.resolve.transformers.plugin.GeneratedClassIndex
 import org.jetbrains.kotlin.fir.scopes.impl.FirDeclaredMemberScopeProvider
 import org.jetbrains.kotlin.fir.types.FirCorrespondingSupertypesCache
+import org.jetbrains.kotlin.incremental.components.LookupTracker
 
 // -------------------------- Required components --------------------------
 
@@ -55,10 +57,20 @@ fun FirSession.registerThreadUnsafeCaches() {
  * Resolve components which are same on all platforms
  */
 @OptIn(SessionConfiguration::class)
-fun FirSession.registerResolveComponents() {
+fun FirSession.registerResolveComponents(lookupTracker: LookupTracker? = null) {
     register(FirQualifierResolver::class, FirQualifierResolverImpl(this))
     register(FirTypeResolver::class, FirTypeResolverImpl(this))
     register(CheckersComponent::class, CheckersComponent())
+    if (lookupTracker != null) {
+        val firFileToPath: (FirSourceElement) -> String = {
+            val psiSource = (it as? FirPsiSourceElement<*>) ?: TODO("Not implemented for non-FirPsiSourceElement")
+            ((psiSource.psi as? PsiFile) ?: psiSource.psi.containingFile).virtualFile.path
+        }
+        register(
+            FirLookupTrackerComponent::class,
+            IncrementalPassThroughLookupTrackerComponent(lookupTracker, firFileToPath)
+        )
+    }
 }
 
 /*

compiler/fir/entrypoint/src/org/jetbrains/kotlin/fir/session/FirSessionFactory.kt

@@ -33,6 +33,7 @@ import org.jetbrains.kotlin.fir.resolve.providers.FirSymbolProvider
 import org.jetbrains.kotlin.fir.resolve.providers.impl.*
 import org.jetbrains.kotlin.fir.resolve.scopes.wrapScopeWithJvmMapped
 import org.jetbrains.kotlin.fir.scopes.KotlinScopeProvider
+import org.jetbrains.kotlin.incremental.components.LookupTracker
 import org.jetbrains.kotlin.load.java.JavaClassFinderImpl
 import org.jetbrains.kotlin.load.kotlin.PackagePartProvider
 import org.jetbrains.kotlin.load.kotlin.VirtualFileFinderFactory
@@ -68,12 +69,13 @@ object FirSessionFactory {
         project: Project,
         dependenciesProvider: FirSymbolProvider? = null,
         languageVersionSettings: LanguageVersionSettings = LanguageVersionSettingsImpl.DEFAULT,
+        lookupTracker: LookupTracker? = null,
         init: FirSessionConfigurator.() -> Unit = {}
     ): FirJavaModuleBasedSession {
         return FirJavaModuleBasedSession(moduleInfo, sessionProvider).apply {
             registerThreadUnsafeCaches()
             registerCommonComponents(languageVersionSettings)
-            registerResolveComponents()
+            registerResolveComponents(lookupTracker)
             registerJavaSpecificResolveComponents()
 
             val kotlinScopeProvider = KotlinScopeProvider(::wrapScopeWithJvmMapped)

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

@@ -12,6 +12,7 @@ import org.jetbrains.kotlin.config.LanguageVersionSettings
 import org.jetbrains.kotlin.fir.java.FirProjectSessionProvider
 import org.jetbrains.kotlin.fir.session.FirJvmModuleInfo
 import org.jetbrains.kotlin.fir.session.FirSessionFactory
+import org.jetbrains.kotlin.incremental.components.LookupTracker
 import org.jetbrains.kotlin.load.kotlin.PackagePartProvider
 import org.jetbrains.kotlin.modules.Module
 import org.jetbrains.kotlin.name.Name
@@ -24,6 +25,7 @@ fun createSessionWithDependencies(
     languageVersionSettings: LanguageVersionSettings,
     sourceScope: GlobalSearchScope,
     librariesScope: GlobalSearchScope,
+    lookupTracker: LookupTracker?,
     packagePartProvider: (GlobalSearchScope) -> PackagePartProvider,
     sessionConfigurator: FirSessionFactory.FirSessionConfigurator.() -> Unit = {}
 ): FirSession {
@@ -33,6 +35,7 @@ fun createSessionWithDependencies(
         languageVersionSettings,
         sourceScope,
         librariesScope,
+        lookupTracker,
         packagePartProvider,
         sessionConfigurator
     ) {
@@ -46,6 +49,7 @@ fun createSessionWithDependencies(
     languageVersionSettings: LanguageVersionSettings,
     sourceScope: GlobalSearchScope,
     librariesScope: GlobalSearchScope,
+    lookupTracker: LookupTracker?,
     packagePartProvider: (GlobalSearchScope) -> PackagePartProvider,
     sessionConfigurator: FirSessionFactory.FirSessionConfigurator.() -> Unit = {}
 ): FirSession {
@@ -55,6 +59,7 @@ fun createSessionWithDependencies(
         languageVersionSettings,
         sourceScope,
         librariesScope,
+        lookupTracker,
         packagePartProvider,
         sessionConfigurator
     ) {
@@ -68,6 +73,7 @@ private inline fun createSessionWithDependencies(
     languageVersionSettings: LanguageVersionSettings,
     sourceScope: GlobalSearchScope,
     librariesScope: GlobalSearchScope,
+    lookupTracker: LookupTracker?,
     packagePartProvider: (GlobalSearchScope) -> PackagePartProvider,
     noinline sessionConfigurator: FirSessionFactory.FirSessionConfigurator.() -> Unit,
     moduleInfoProvider: (dependencies: List<ModuleInfo>) -> ModuleInfo,
@@ -84,6 +90,7 @@ private inline fun createSessionWithDependencies(
         sourceScope,
         project,
         languageVersionSettings = languageVersionSettings,
+        lookupTracker = lookupTracker,
         init = sessionConfigurator
     )
 }