[IR] Add an API for writing lookups from IrGenerationExtension

^KT-79504 Fixed

Author: demiurg906

compiler/fir/fir2ir/src/org/jetbrains/kotlin/fir/backend/Fir2IrPluginContext.kt

@@ -5,6 +5,7 @@
 
 package org.jetbrains.kotlin.fir.backend
 
+import org.jetbrains.kotlin.KtSourceElement
 import org.jetbrains.kotlin.backend.common.extensions.FirIncompatiblePluginAPI
 import org.jetbrains.kotlin.backend.common.extensions.IrPluginContext
 import org.jetbrains.kotlin.backend.common.ir.BuiltinSymbolsBase
@@ -25,10 +26,21 @@ import org.jetbrains.kotlin.ir.IrDiagnosticReporter
 import org.jetbrains.kotlin.ir.KtDiagnosticReporterWithImplicitIrBasedContext
 import org.jetbrains.kotlin.ir.ObsoleteDescriptorBasedAPI
 import org.jetbrains.kotlin.backend.common.linkage.IrDeserializer
+import org.jetbrains.kotlin.fir.lookupTracker
+import org.jetbrains.kotlin.fir.recordFqNameLookup
+import org.jetbrains.kotlin.ir.declarations.IrClass
+import org.jetbrains.kotlin.ir.declarations.IrConstructor
+import org.jetbrains.kotlin.ir.declarations.IrDeclarationWithName
+import org.jetbrains.kotlin.ir.declarations.IrFile
+import org.jetbrains.kotlin.ir.declarations.IrProperty
+import org.jetbrains.kotlin.ir.declarations.IrSimpleFunction
+import org.jetbrains.kotlin.ir.declarations.IrTypeAlias
 import org.jetbrains.kotlin.ir.symbols.*
 import org.jetbrains.kotlin.ir.util.IdSignature
 import org.jetbrains.kotlin.ir.util.ReferenceSymbolTable
 import org.jetbrains.kotlin.ir.util.TypeTranslator
+import org.jetbrains.kotlin.ir.util.fqNameWhenAvailable
+import org.jetbrains.kotlin.ir.util.sourceElement
 import org.jetbrains.kotlin.name.CallableId
 import org.jetbrains.kotlin.name.ClassId
 import org.jetbrains.kotlin.name.FqName
@@ -135,6 +147,14 @@ class Fir2IrPluginContext(
         return callables.mapNotNull { c.declarationStorage.irExtractor(it) }.filterIsInstance<R>()
     }
 
+    override fun recordLookup(declaration: IrDeclarationWithName, fromFile: IrFile) {
+        val lookupTracker = c.session.lookupTracker ?: return
+        val fqName = declaration.fqNameWhenAvailable ?: return
+        val firFile = (fromFile.metadata as? FirMetadataSource.File)?.fir ?: return
+        val fileSource = firFile.source ?: return
+        lookupTracker.recordFqNameLookup(fqName, source = null, fileSource = fileSource)
+    }
+
     @Deprecated("Use messageCollector or diagnosticReporter properties instead", level = DeprecationLevel.ERROR)
     override fun createDiagnosticReporter(pluginId: String): MessageCollector {
         error(ERROR_MESSAGE)

compiler/ir/backend.common/src/org/jetbrains/kotlin/backend/common/extensions/IrPluginContext.kt

@@ -13,6 +13,9 @@ import org.jetbrains.kotlin.descriptors.ModuleDescriptor
 import org.jetbrains.kotlin.ir.IrDiagnosticReporter
 import org.jetbrains.kotlin.ir.ObsoleteDescriptorBasedAPI
 import org.jetbrains.kotlin.ir.builders.IrGeneratorContext
+import org.jetbrains.kotlin.ir.declarations.IrDeclaration
+import org.jetbrains.kotlin.ir.declarations.IrDeclarationWithName
+import org.jetbrains.kotlin.ir.declarations.IrFile
 import org.jetbrains.kotlin.ir.symbols.*
 import org.jetbrains.kotlin.ir.util.IdSignature
 import org.jetbrains.kotlin.ir.util.ReferenceSymbolTable
@@ -65,6 +68,15 @@ interface IrPluginContext : IrGeneratorContext {
     fun referenceFunctions(callableId: CallableId): Collection<IrSimpleFunctionSymbol>
     fun referenceProperties(callableId: CallableId): Collection<IrPropertySymbol>
 
+    // ------------------------------------ IC API ------------------------------------
+
+    /**
+     * Records information that [declaration] was referenced during modification of file [fromFile].
+     * This information later will be used by the Incremental compilation to correctly invalidate
+     * compiled files on source changes.
+     */
+    fun recordLookup(declaration: IrDeclarationWithName, fromFile: IrFile)
+
     // ------------------------------------ Deprecated API ------------------------------------
 
     @Deprecated("Use diagnosticReporter instead", level = DeprecationLevel.ERROR)

compiler/ir/backend.common/src/org/jetbrains/kotlin/backend/common/extensions/IrPluginContextImpl.kt

@@ -25,6 +25,8 @@ import org.jetbrains.kotlin.ir.declarations.IrDeclaration
 import org.jetbrains.kotlin.ir.declarations.IrSimpleFunction
 import org.jetbrains.kotlin.ir.expressions.IrConstructorCall
 import org.jetbrains.kotlin.backend.common.linkage.IrDeserializer
+import org.jetbrains.kotlin.ir.declarations.IrDeclarationWithName
+import org.jetbrains.kotlin.ir.declarations.IrFile
 import org.jetbrains.kotlin.ir.symbols.*
 import org.jetbrains.kotlin.ir.util.IdSignature
 import org.jetbrains.kotlin.ir.util.ReferenceSymbolTable
@@ -195,6 +197,8 @@ open class IrPluginContextImpl(
         return referenceProperties(callableId.asSingleFqName())
     }
 
+    override fun recordLookup(declaration: IrDeclarationWithName, fromFile: IrFile) {}
+
     @Deprecated("This API is deprecated. It will be removed after the 2.3 release", level = DeprecationLevel.WARNING)
     @FirIncompatiblePluginAPI
     override fun referenceTopLevel(

plugins/plugin-sandbox/plugin-sandbox-ic-test/testData/pureKotlin/referenceFunctionFromIrPlugin/build.log

@@ -2,6 +2,7 @@
 
 Compiling files:
   src/a.kt
+  src/b.kt
 End of files
 Exit code: OK
 

plugins/plugin-sandbox/plugin-sandbox-ic-test/testData/pureKotlin/referenceFunctionFromIrPlugin/fail.txt

@@ -15,6 +15,7 @@ import org.jetbrains.kotlin.ir.declarations.IrSimpleFunction
 import org.jetbrains.kotlin.ir.expressions.IrBlockBody
 import org.jetbrains.kotlin.ir.expressions.IrConst
 import org.jetbrains.kotlin.ir.expressions.impl.IrCallImpl
+import org.jetbrains.kotlin.ir.util.file
 import org.jetbrains.kotlin.ir.util.getAnnotation
 import org.jetbrains.kotlin.ir.visitors.IrVisitorVoid
 import org.jetbrains.kotlin.ir.visitors.acceptChildrenVoid
@@ -60,5 +61,6 @@ class IrTransformerForICTesting(private val context: IrPluginContext) : IrVisito
             superQualifierSymbol = null
         )
         body.statements.add(0, functionCall)
+        context.recordLookup(functionToCall, fromFile = declaration.file)
     }
 }