[Analysis API] Allow obtaining original for fake declaration/file

It is required to remove dependency on low-level API from completion

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

@@ -53,6 +53,7 @@ class KtFe10AnalysisSession(
     override val visibilityCheckerImpl: KtVisibilityChecker = KtFe10VisibilityChecker(this)
     override val overrideInfoProviderImpl: KtOverrideInfoProvider = KtFe10OverrideInfoProvider(this)
     override val multiplatformInfoProviderImpl: KtMultiplatformInfoProvider = KtFe10MultiplatformInfoProvider(this)
+    override val originalPsiProviderImpl: KtOriginalPsiProvider = KtFe10OriginalPsiProvider(this)
     override val inheritorsProviderImpl: KtInheritorsProvider = KtFe10InheritorsProvider(this)
     override val typesCreatorImpl: KtTypeCreator = KtFe10TypeCreator(this)
     override val samResolverImpl: KtSamResolver = KtFe10SamResolver(this)

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

@@ -0,0 +1,24 @@
+/*
+ * Copyright 2010-2023 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.components.KtOriginalPsiProvider
+import org.jetbrains.kotlin.analysis.api.descriptors.KtFe10AnalysisSession
+import org.jetbrains.kotlin.analysis.api.descriptors.components.base.Fe10KtAnalysisSessionComponent
+import org.jetbrains.kotlin.psi.KtDeclaration
+import org.jetbrains.kotlin.psi.KtFile
+
+internal class KtFe10OriginalPsiProvider(
+    override val analysisSession: KtFe10AnalysisSession
+) : KtOriginalPsiProvider(), Fe10KtAnalysisSessionComponent {
+    override fun getOriginalDeclaration(declaration: KtDeclaration): KtDeclaration? = null
+
+    override fun getOriginalKtFile(file: KtFile): KtFile? = null
+
+    override fun recordOriginalDeclaration(fakeDeclaration: KtDeclaration, originalDeclaration: KtDeclaration) {}
+
+    override fun recordOriginalKtFile(fakeFile: KtFile, originalFile: KtFile) {}
+}

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

@@ -108,6 +108,8 @@ private constructor(
 
     override val multiplatformInfoProviderImpl: KtMultiplatformInfoProvider = KtFirMultiplatformInfoProvider(this, token)
 
+    override val originalPsiProviderImpl: KtOriginalPsiProvider = KtFirOriginalPsiProvider(this, token)
+
     override val symbolInfoProviderImpl: KtSymbolInfoProvider = KtFirSymbolInfoProvider(this, token)
 
     override val typesCreatorImpl: KtTypeCreator = KtFirTypeCreator(this, token)

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

@@ -0,0 +1,31 @@
+/*
+ * Copyright 2010-2023 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.components.KtOriginalPsiProvider
+import org.jetbrains.kotlin.analysis.api.fir.KtFirAnalysisSession
+import org.jetbrains.kotlin.analysis.api.lifetime.KtLifetimeToken
+import org.jetbrains.kotlin.analysis.low.level.api.fir.util.originalDeclaration
+import org.jetbrains.kotlin.analysis.low.level.api.fir.util.originalKtFile
+import org.jetbrains.kotlin.psi.KtDeclaration
+import org.jetbrains.kotlin.psi.KtFile
+
+internal class KtFirOriginalPsiProvider(
+    override val analysisSession: KtFirAnalysisSession,
+    override val token: KtLifetimeToken,
+) : KtOriginalPsiProvider(), KtFirAnalysisSessionComponent {
+    override fun getOriginalDeclaration(declaration: KtDeclaration): KtDeclaration? = declaration.originalDeclaration
+
+    override fun getOriginalKtFile(file: KtFile): KtFile? = file.originalKtFile
+
+    override fun recordOriginalDeclaration(fakeDeclaration: KtDeclaration, originalDeclaration: KtDeclaration) {
+        fakeDeclaration.originalDeclaration = originalDeclaration
+    }
+
+    override fun recordOriginalKtFile(fakeFile: KtFile, originalFile: KtFile) {
+        fakeFile.originalKtFile = originalFile
+    }
+}

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

@@ -55,6 +55,7 @@ public abstract class KtAnalysisSession(final override val token: KtLifetimeToke
     KtVisibilityCheckerMixIn,
     KtMemberSymbolProviderMixin,
     KtMultiplatformInfoProviderMixin,
+    KtOriginalPsiProviderMixIn,
     KtInheritorsProviderMixIn,
     KtTypeCreatorMixIn,
     KtAnalysisScopeProviderMixIn,
@@ -141,6 +142,9 @@ public abstract class KtAnalysisSession(final override val token: KtLifetimeToke
     internal val multiplatformInfoProvider: KtMultiplatformInfoProvider get() = multiplatformInfoProviderImpl
     protected abstract val multiplatformInfoProviderImpl: KtMultiplatformInfoProvider
 
+    internal val originalPsiProvider: KtOriginalPsiProvider get() = originalPsiProviderImpl
+    protected abstract val originalPsiProviderImpl: KtOriginalPsiProvider
+
     internal val symbolInfoProvider: KtSymbolInfoProvider get() = symbolInfoProviderImpl
     protected abstract val symbolInfoProviderImpl: KtSymbolInfoProvider
 

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

@@ -0,0 +1,47 @@
+/*
+ * Copyright 2010-2023 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.lifetime.withValidityAssertion
+import org.jetbrains.kotlin.psi.KtDeclaration
+import org.jetbrains.kotlin.psi.KtFile
+
+public abstract class KtOriginalPsiProvider : KtAnalysisSessionComponent() {
+    public abstract fun getOriginalDeclaration(declaration: KtDeclaration): KtDeclaration?
+    public abstract fun getOriginalKtFile(file: KtFile): KtFile?
+
+    public abstract fun recordOriginalDeclaration(fakeDeclaration: KtDeclaration, originalDeclaration: KtDeclaration)
+    public abstract fun recordOriginalKtFile(fakeFile: KtFile, originalFile: KtFile)
+}
+
+public interface KtOriginalPsiProviderMixIn : KtAnalysisSessionMixIn {
+    /**
+     * If [KtDeclaration] is a non-local declaration in a fake file analyzed in dependent session, returns the original declaration
+     * for [this]. Otherwise, returns `null`.
+     */
+    public fun KtDeclaration.getOriginalDeclaration(): KtDeclaration? =
+        withValidityAssertion { analysisSession.originalPsiProvider.getOriginalDeclaration(this) }
+
+    /**
+     * If [this] is a fake file analyzed in dependent session, returns the original file for [this]. Otherwise, returns `null`.
+     */
+    public fun KtFile.getOriginalKtFile(): KtFile? =
+        withValidityAssertion { analysisSession.originalPsiProvider.getOriginalKtFile(this) }
+
+    /**
+     * Records [declaration] as an original declaration for [this].
+     */
+    public fun KtDeclaration.recordOriginalDeclaration(declaration: KtDeclaration) {
+        withValidityAssertion { analysisSession.originalPsiProvider.recordOriginalDeclaration(this, declaration) }
+    }
+
+    /**
+     * Records [file] as an original file for [this].
+     */
+    public fun KtFile.recordOriginalKtFile(file: KtFile) {
+        withValidityAssertion { analysisSession.originalPsiProvider.recordOriginalKtFile(this, file) }
+    }
+}