From 0add17d9da7149be914ee57fb4e257cabcc18d24 Mon Sep 17 00:00:00 2001
From: Marco Pennekamp <marco.pennekamp@jetbrains.com>
Date: Tue, 10 Oct 2023 19:00:11 +0200
Subject: [PATCH] [AA] Fix KDoc reference resolution of companion object
 members

- The nested scope construction in `getSymbolsFromMemberScope` was
  missing the companion object member scope.

^KTIJ-25995 fixed
---
 .../fir/references/KDocReferenceResolver.kt   | 25 +++++++++++++------
 1 file changed, 18 insertions(+), 7 deletions(-)

diff --git a/analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/references/KDocReferenceResolver.kt b/analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/references/KDocReferenceResolver.kt
index 5c596850bea..0ede2920aad 100644
--- a/analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/references/KDocReferenceResolver.kt
+++ b/analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/references/KDocReferenceResolver.kt
@@ -138,6 +138,13 @@ internal object KDocReferenceResolver {
         }
     }
 
+    /**
+     * Returns the [KtSymbol]s called [fqName] found in the member scope and companion object's member scope of the [KtDeclaration]s that
+     * contain the [contextElement].
+     *
+     * If [fqName] has two or more segments, e.g. `Foo.bar`, the member and companion object scope of the containing [KtDeclaration] will be
+     * queried for a class `Foo` first, and then that class `Foo` will be queried for the member `bar` by short name.
+     */
     context(KtAnalysisSession)
     private fun getSymbolsFromParentMemberScopes(fqName: FqName, contextElement: KtElement): Collection<KtSymbol> {
         for (ktDeclaration in contextElement.parentsOfType<KtDeclaration>(withSelf = true)) {
@@ -147,10 +154,7 @@ internal object KDocReferenceResolver {
             if (ktDeclaration is KtClassOrObject) {
                 val symbol = ktDeclaration.getClassOrObjectSymbol() ?: continue
 
-                val scope = listOfNotNull(
-                    symbol.getMemberScope(),
-                    getCompanionObjectMemberScope(symbol),
-                ).asCompositeScope()
+                val scope = symbol.getCompositeCombinedMemberAndCompanionObjectScope()
 
                 val symbolsFromScope = getSymbolsFromMemberScope(fqName, scope)
                 if (symbolsFromScope.isNotEmpty()) return symbolsFromScope
@@ -160,8 +164,15 @@ internal object KDocReferenceResolver {
     }
 
     context(KtAnalysisSession)
-    private fun getCompanionObjectMemberScope(classSymbol: KtClassOrObjectSymbol): KtScope? {
-        val namedClassSymbol = classSymbol as? KtNamedClassOrObjectSymbol ?: return null
+    private fun KtSymbolWithMembers.getCompositeCombinedMemberAndCompanionObjectScope(): KtScope =
+        listOfNotNull(
+            getCombinedMemberScope(),
+            getCompanionObjectMemberScope(),
+        ).asCompositeScope()
+
+    context(KtAnalysisSession)
+    private fun KtSymbolWithMembers.getCompanionObjectMemberScope(): KtScope? {
+        val namedClassSymbol = this as? KtNamedClassOrObjectSymbol ?: return null
         val companionSymbol = namedClassSymbol.companionObject ?: return null
         return companionSymbol.getMemberScope()
     }
@@ -192,7 +203,7 @@ internal object KDocReferenceResolver {
                 currentScope
                     .getClassifierSymbols(fqNamePart)
                     .filterIsInstance<KtSymbolWithMembers>()
-                    .map { it.getCombinedMemberScope() }
+                    .map { it.getCompositeCombinedMemberAndCompanionObjectScope() }
                     .toList()
                     .asCompositeScope()
             }