[Analysis Api] Expose klibSourceFile via KtKlibSourceFileProviderMixIn

This `klibSourceFile` information is deserialized from klibs
to retain the information of the original SourceFile location
of a declaration.

^KT-66271 Fixed

analysis/analysis-api-fe10/src/org/jetbrains/kotlin/analysis/api/descriptors/KtFe10AnalysisSession.kt

@@ -65,6 +65,7 @@ class KtFe10AnalysisSession(
     override val resolveExtensionInfoProviderImpl: KtResolveExtensionInfoProvider = KtFe10ResolveExtensionInfoProvider(this)
     override val compilerFacilityImpl: KtCompilerFacility = KtFe10CompilerFacility(this)
     override val dataFlowInfoProviderImpl: KtDataFlowInfoProvider = KtFe10DataFlowInfoProvider(this)
+    override val klibSourceFileProviderImpl: KtKlibSourceFileNameProvider = KtFe10KlibSourceFileNameProvider(this)
 
     override val metadataCalculatorImpl: KtMetadataCalculator
         get() = throw NotSupportedForK1Exception()

analysis/analysis-api-fe10/src/org/jetbrains/kotlin/analysis/api/descriptors/components/KtFe10KlibSourceFileNameProvider.kt

@@ -0,0 +1,20 @@
+/*
+ * Copyright 2010-2024 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.analysis.api.descriptors.components
+
+import org.jetbrains.kotlin.analysis.api.KtAnalysisNonPublicApi
+import org.jetbrains.kotlin.analysis.api.KtAnalysisSession
+import org.jetbrains.kotlin.analysis.api.components.KtKlibSourceFileNameProvider
+import org.jetbrains.kotlin.analysis.api.symbols.KtDeclarationSymbol
+
+@OptIn(KtAnalysisNonPublicApi::class)
+internal class KtFe10KlibSourceFileNameProvider(
+    override val analysisSession: KtAnalysisSession,
+) : KtKlibSourceFileNameProvider() {
+    override fun getKlibSourceFileName(declaration: KtDeclarationSymbol): String? {
+        throw NotImplementedError("Method is not implemented for FE 1.0")
+    }
+}

analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/KtFirAnalysisSession.kt

@@ -128,6 +128,8 @@ private constructor(
 
     override val dataFlowInfoProviderImpl: KtDataFlowInfoProvider = KtFirDataFlowInfoProvider(this)
 
+    override val klibSourceFileProviderImpl: KtKlibSourceFileNameProvider = KtFirKlibSourceFileNameProvider(this)
+
     internal val useSiteSession: FirSession get() = firResolveSession.useSiteFirSession
     internal val firSymbolProvider: FirSymbolProvider get() = useSiteSession.symbolProvider
     internal val targetPlatform: TargetPlatform get() = useSiteSession.moduleData.platform

analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/components/KtFirKlibSourceFileNameProvider.kt

@@ -0,0 +1,24 @@
+/*
+ * Copyright 2010-2024 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.analysis.api.fir.components
+
+import org.jetbrains.kotlin.analysis.api.KtAnalysisNonPublicApi
+import org.jetbrains.kotlin.analysis.api.components.KtKlibSourceFileNameProvider
+import org.jetbrains.kotlin.analysis.api.fir.KtFirAnalysisSession
+import org.jetbrains.kotlin.analysis.api.fir.symbols.KtFirSymbol
+import org.jetbrains.kotlin.analysis.api.symbols.KtDeclarationSymbol
+import org.jetbrains.kotlin.fir.declarations.utils.klibSourceFile
+
+@OptIn(KtAnalysisNonPublicApi::class)
+internal class KtFirKlibSourceFileNameProvider(
+    override val analysisSession: KtFirAnalysisSession,
+) : KtKlibSourceFileNameProvider() {
+    override fun getKlibSourceFileName(declaration: KtDeclarationSymbol): String? {
+        require(declaration is KtFirSymbol<*>)
+        val sourceFile = declaration.firSymbol.klibSourceFile ?: return null
+        return sourceFile.name
+    }
+}

analysis/analysis-api-impl-base/build.gradle.kts

@@ -35,6 +35,7 @@ dependencies {
     testImplementation(project(":analysis:decompiled:decompiler-native"))
     testImplementation(projectTests(":analysis:analysis-test-framework"))
     testImplementation(commonDependency("org.jetbrains.kotlin:kotlin-reflect")) { isTransitive = false }
+    testImplementation(project(":native:analysis-api-klib-reader"))
     testImplementation(toolsJar())
 }
 

analysis/analysis-api-impl-base/tests/org/jetbrains/kotlin/analysis/api/impl/base/test/cases/components/klibSourceFileProvider/AbstractGetKlibSourceFileNameTest.kt

@@ -0,0 +1,55 @@
+/*
+ * Copyright 2010-2024 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.analysis.api.impl.base.test.cases.components.klibSourceFileProvider
+
+import org.jetbrains.kotlin.analysis.api.analyze
+import org.jetbrains.kotlin.analysis.api.klib.reader.getSymbols
+import org.jetbrains.kotlin.analysis.api.klib.reader.readKlibDeclarationAddresses
+import org.jetbrains.kotlin.analysis.api.symbols.KtCallableSymbol
+import org.jetbrains.kotlin.analysis.api.symbols.KtClassOrObjectSymbol
+import org.jetbrains.kotlin.analysis.api.symbols.KtDeclarationSymbol
+import org.jetbrains.kotlin.analysis.project.structure.KtLibraryModule
+import org.jetbrains.kotlin.analysis.test.framework.base.AbstractAnalysisApiBasedTest
+import org.jetbrains.kotlin.analysis.test.framework.project.structure.ktModuleProvider
+import org.jetbrains.kotlin.analysis.test.framework.project.structure.mainModules
+import org.jetbrains.kotlin.test.services.TestModuleStructure
+import org.jetbrains.kotlin.test.services.TestServices
+import org.jetbrains.kotlin.test.services.assertions
+import kotlin.test.fail
+
+/**
+ * Reads through the declarations provided in the .klib and renders their `klibSourceFile`
+ */
+abstract class AbstractGetKlibSourceFileNameTest : AbstractAnalysisApiBasedTest() {
+    override fun doTestByModuleStructure(moduleStructure: TestModuleStructure, testServices: TestServices) {
+        val mainModule = testServices.ktModuleProvider.mainModules
+            .let { modules -> if (modules.size == 1) modules.first() else fail("Expected single main module. Found $modules") }
+
+        val libraryModule = mainModule.ktModule as? KtLibraryModule
+            ?: fail("Expected main module '${mainModule.ktModule}' to be '${KtLibraryModule::class.simpleName}'")
+
+        val actual = StringBuilder()
+        actual.appendLine("klib declarations:")
+
+        analyze(libraryModule) {
+            val klibAddresses = libraryModule.readKlibDeclarationAddresses() ?: fail("Failed reading 'klib addresses' from $libraryModule")
+            klibAddresses.forEach { klibDeclarationAddress ->
+                klibDeclarationAddress.getSymbols().filterIsInstance<KtDeclarationSymbol>().forEach { symbol ->
+                    val sourceFile = symbol.getKlibSourceFileName()
+                    if (symbol is KtCallableSymbol) {
+                        actual.appendLine("Callable: ${symbol.callableIdIfNonLocal}; klibSourceFile: $sourceFile")
+                    }
+
+                    if (symbol is KtClassOrObjectSymbol) {
+                        actual.appendLine("Classifier: ${symbol.classIdIfNonLocal}; klibSourceFile: $sourceFile")
+                    }
+                }
+            }
+        }
+
+        testServices.assertions.assertEqualsToTestDataFileSibling(actual.toString())
+    }
+}

analysis/analysis-api-standalone/build.gradle.kts

@@ -15,10 +15,11 @@ dependencies {
     api(project(":analysis:decompiled:light-classes-for-decompiled"))
     api(project(":analysis:analysis-api-standalone:analysis-api-standalone-base"))
     implementation(project(":analysis:analysis-api-standalone:analysis-api-fir-standalone-base"))
+    testImplementation(project(":native:analysis-api-klib-reader"))
+    testImplementation(projectTests(":analysis:analysis-api-fir"))
+    testImplementation(projectTests(":analysis:analysis-api-impl-base"))
     testImplementation(projectTests(":analysis:analysis-test-framework"))
     testImplementation(projectTests(":analysis:low-level-api-fir"))
-    testImplementation(projectTests(":analysis:analysis-api-impl-base"))
-    testImplementation(projectTests(":analysis:analysis-api-fir"))
 
     testImplementation(kotlinTest("junit"))
     testImplementation(toolsJar())

analysis/analysis-api-standalone/tests-gen/org/jetbrains/kotlin/analysis/api/standalone/fir/test/cases/generated/cases/components/klibSourceFileProvider/FirStandaloneNormalAnalysisSourceModuleGetKlibSourceFileNameTestGenerated.java

@@ -0,0 +1,65 @@
+/*
+ * Copyright 2010-2024 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.analysis.api.standalone.fir.test.cases.generated.cases.components.klibSourceFileProvider;
+
+import com.intellij.testFramework.TestDataPath;
+import org.jetbrains.kotlin.test.util.KtTestUtil;
+import org.jetbrains.annotations.NotNull;
+import org.jetbrains.kotlin.analysis.api.standalone.fir.test.configurators.AnalysisApiFirStandaloneModeTestConfiguratorFactory;
+import org.jetbrains.kotlin.analysis.test.framework.test.configurators.AnalysisApiTestConfiguratorFactoryData;
+import org.jetbrains.kotlin.analysis.test.framework.test.configurators.AnalysisApiTestConfigurator;
+import org.jetbrains.kotlin.analysis.test.framework.test.configurators.TestModuleKind;
+import org.jetbrains.kotlin.analysis.test.framework.test.configurators.FrontendKind;
+import org.jetbrains.kotlin.analysis.test.framework.test.configurators.AnalysisSessionMode;
+import org.jetbrains.kotlin.analysis.test.framework.test.configurators.AnalysisApiMode;
+import org.jetbrains.kotlin.analysis.api.impl.base.test.cases.components.klibSourceFileProvider.AbstractGetKlibSourceFileNameTest;
+import org.jetbrains.kotlin.test.TestMetadata;
+import org.junit.jupiter.api.Test;
+
+import java.io.File;
+import java.util.regex.Pattern;
+
+/** This class is generated by {@link org.jetbrains.kotlin.generators.tests.analysis.api.GenerateAnalysisApiTestsKt}. DO NOT MODIFY MANUALLY */
+@SuppressWarnings("all")
+@TestMetadata("analysis/analysis-api/testData/components/klibSourceFileNameProvider/getKlibSourceFileName")
+@TestDataPath("$PROJECT_ROOT")
+public class FirStandaloneNormalAnalysisSourceModuleGetKlibSourceFileNameTestGenerated extends AbstractGetKlibSourceFileNameTest {
+  @NotNull
+  @Override
+  public AnalysisApiTestConfigurator getConfigurator() {
+    return AnalysisApiFirStandaloneModeTestConfiguratorFactory.INSTANCE.createConfigurator(
+      new AnalysisApiTestConfiguratorFactoryData(
+        FrontendKind.Fir,
+        TestModuleKind.Source,
+        AnalysisSessionMode.Normal,
+        AnalysisApiMode.Standalone
+      )
+    );
+  }
+
+  @Test
+  public void testAllFilesPresentInGetKlibSourceFileName() {
+    KtTestUtil.assertAllTestsPresentByMetadataWithExcluded(this.getClass(), new File("analysis/analysis-api/testData/components/klibSourceFileNameProvider/getKlibSourceFileName"), Pattern.compile("^(.+)\\.kt$"), null, true);
+  }
+
+  @Test
+  @TestMetadata("class.kt")
+  public void testClass() {
+    runTest("analysis/analysis-api/testData/components/klibSourceFileNameProvider/getKlibSourceFileName/class.kt");
+  }
+
+  @Test
+  @TestMetadata("topLevelFunction.kt")
+  public void testTopLevelFunction() {
+    runTest("analysis/analysis-api/testData/components/klibSourceFileNameProvider/getKlibSourceFileName/topLevelFunction.kt");
+  }
+
+  @Test
+  @TestMetadata("topLevelProperty.kt")
+  public void testTopLevelProperty() {
+    runTest("analysis/analysis-api/testData/components/klibSourceFileNameProvider/getKlibSourceFileName/topLevelProperty.kt");
+  }
+}

analysis/analysis-api/src/org/jetbrains/kotlin/analysis/api/KtAnalysisSession.kt

@@ -93,7 +93,8 @@ public abstract class KtAnalysisSession(final override val token: KtLifetimeToke
     KtCompilerFacilityMixIn,
     KtMetadataCalculatorMixIn,
     KtSubstitutorProviderMixIn,
-    KtDataFlowInfoProviderMixin {
+    KtDataFlowInfoProviderMixin,
+    KtKlibSourceFileProviderMixIn {
 
     public abstract val useSiteModule: KtModule
 
@@ -219,6 +220,9 @@ public abstract class KtAnalysisSession(final override val token: KtLifetimeToke
     internal val dataFlowInfoProvider: KtDataFlowInfoProvider get() = dataFlowInfoProviderImpl
     @KtAnalysisNonPublicApi
     protected abstract val dataFlowInfoProviderImpl: KtDataFlowInfoProvider
+
+    internal val klibSourceFileProvider: KtKlibSourceFileNameProvider get() = klibSourceFileProviderImpl
+    protected abstract val klibSourceFileProviderImpl: KtKlibSourceFileNameProvider
 }
 
 public fun KtAnalysisSession.getModule(element: PsiElement): KtModule {

analysis/analysis-api/src/org/jetbrains/kotlin/analysis/api/components/KtKlibSourceFileNameProvider.kt

@@ -0,0 +1,27 @@
+/*
+ * Copyright 2010-2024 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.analysis.api.components
+
+import org.jetbrains.kotlin.analysis.api.KtAnalysisNonPublicApi
+import org.jetbrains.kotlin.analysis.api.lifetime.withValidityAssertion
+import org.jetbrains.kotlin.analysis.api.symbols.KtDeclarationSymbol
+import org.jetbrains.kotlin.descriptors.SourceFile
+import org.jetbrains.kotlin.psi.KtDeclaration
+
+@KtAnalysisNonPublicApi
+public abstract class KtKlibSourceFileNameProvider : KtAnalysisSessionComponent() {
+    public abstract fun getKlibSourceFileName(declaration: KtDeclarationSymbol): String?
+}
+
+@KtAnalysisNonPublicApi
+public interface KtKlibSourceFileProviderMixIn : KtAnalysisSessionMixIn {
+    /**
+     * If [KtDeclaration] is a deserialized, klib based symbol, then information about the original
+     * [SourceFile] might be retained.
+     */
+    public fun KtDeclarationSymbol.getKlibSourceFileName(): String? =
+        withValidityAssertion { analysisSession.klibSourceFileProvider.getKlibSourceFileName(this) }
+}

analysis/analysis-api/testData/components/klibSourceFileNameProvider/getKlibSourceFileName/class.kt

@@ -0,0 +1,10 @@
+// TARGET_PLATFORM: JS
+// MODULE_KIND: LibraryBinary
+
+// FILE: Some.kt
+package some
+class Foo
+
+// FILE: Other.kt
+package other
+class Bar

analysis/analysis-api/testData/components/klibSourceFileNameProvider/getKlibSourceFileName/class.txt

@@ -0,0 +1,3 @@
+klib declarations:
+Classifier: other/Bar; klibSourceFile: Other.kt
+Classifier: some/Foo; klibSourceFile: Some.kt

analysis/analysis-api/testData/components/klibSourceFileNameProvider/getKlibSourceFileName/topLevelFunction.kt

@@ -0,0 +1,10 @@
+// TARGET_PLATFORM: JS
+// MODULE_KIND: LibraryBinary
+
+// FILE: Some.kt
+package some
+fun foo() = 42
+
+// FILE: Other.kt
+package other
+fun bar() = 42

analysis/analysis-api/testData/components/klibSourceFileNameProvider/getKlibSourceFileName/topLevelFunction.txt

@@ -0,0 +1,3 @@
+klib declarations:
+Callable: other/bar; klibSourceFile: Other.kt
+Callable: some/foo; klibSourceFile: Some.kt

analysis/analysis-api/testData/components/klibSourceFileNameProvider/getKlibSourceFileName/topLevelProperty.kt

@@ -0,0 +1,10 @@
+// TARGET_PLATFORM: JS
+// MODULE_KIND: LibraryBinary
+
+// FILE: Some.kt
+package some
+val foo = 42
+
+// FILE: Other.kt
+package other
+val bar = 42

analysis/analysis-api/testData/components/klibSourceFileNameProvider/getKlibSourceFileName/topLevelProperty.txt

@@ -0,0 +1,3 @@
+klib declarations:
+Callable: other/bar; klibSourceFile: Other.kt
+Callable: some/foo; klibSourceFile: Some.kt

generators/analysis-api-generator/tests/org/jetbrains/kotlin/generators/tests/analysis/api/analysisApi.kt

@@ -29,6 +29,7 @@ import org.jetbrains.kotlin.analysis.api.impl.base.test.cases.components.express
 import org.jetbrains.kotlin.analysis.api.impl.base.test.cases.components.expressionTypeProvider.AbstractExpectedExpressionTypeTest
 import org.jetbrains.kotlin.analysis.api.impl.base.test.cases.components.expressionTypeProvider.AbstractHLExpressionTypeTest
 import org.jetbrains.kotlin.analysis.api.impl.base.test.cases.components.importOptimizer.AbstractAnalysisApiImportOptimizerTest
+import org.jetbrains.kotlin.analysis.api.impl.base.test.cases.components.klibSourceFileProvider.AbstractGetKlibSourceFileNameTest
 import org.jetbrains.kotlin.analysis.api.impl.base.test.cases.components.multiplatformInfoProvider.AbstractExpectForActualTest
 import org.jetbrains.kotlin.analysis.api.impl.base.test.cases.components.psiTypeProvider.AbstractAnalysisApiExpressionPsiTypeProviderTest
 import org.jetbrains.kotlin.analysis.api.impl.base.test.cases.components.psiTypeProvider.AbstractAnalysisApiKtTypeByPsiTypeProviderTest
@@ -581,4 +582,10 @@ private fun AnalysisApiTestGroup.generateAnalysisApiComponentsTests() {
             model(it, "exitPointSnapshot")
         }
     }
+
+    component("klibSourceFileNameProvider", filter = frontendIs(FrontendKind.Fir) and analysisApiModeIs(AnalysisApiMode.Standalone)) {
+        test<AbstractGetKlibSourceFileNameTest> {
+            model(it, "getKlibSourceFileName")
+        }
+    }
 }