From 84baee1ef1e3b422c785b484e1a38af0cb6a651c Mon Sep 17 00:00:00 2001
From: Yan Zhulanow <yan.zhulanow@jetbrains.com>
Date: Fri, 20 Jan 2023 18:01:46 +0900
Subject: [PATCH] [LL API] Fallback for designation path calculation for
 special cases

Currently, designation path requires all intermediate classes to be
regular named classes. This is not certainly possible, especially for
invalid cases such as:

---
class Foo {
    object {
        class Bar {
            fun foo() {} // What's the designation path for 'foo()'?
        }
    }
}
---

If there was a way to get a parent for a given FIR element, we could
just follow parents until the root FIR node (which would be a FirFile).
---
 .../low/level/api/fir/api/FirDesignation.kt   | 55 ++++++++++++++++++-
 1 file changed, 53 insertions(+), 2 deletions(-)

diff --git a/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/api/FirDesignation.kt b/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/api/FirDesignation.kt
index 4a021d13b30..f4dca647954 100644
--- a/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/api/FirDesignation.kt
+++ b/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/api/FirDesignation.kt
@@ -5,7 +5,11 @@
 
 package org.jetbrains.kotlin.analysis.low.level.api.fir.api
 
+import org.jetbrains.kotlin.KtRealPsiSourceElement
+import org.jetbrains.kotlin.analysis.low.level.api.fir.project.structure.llFirModuleData
 import org.jetbrains.kotlin.analysis.low.level.api.fir.providers.nullableJavaSymbolProvider
+import org.jetbrains.kotlin.analysis.low.level.api.fir.sessions.llFirResolvableSession
+import org.jetbrains.kotlin.analysis.low.level.api.fir.sessions.llFirSession
 import org.jetbrains.kotlin.analysis.low.level.api.fir.util.FirElementFinder
 import org.jetbrains.kotlin.analysis.low.level.api.fir.util.getContainingFile
 import org.jetbrains.kotlin.analysis.low.level.api.fir.util.withFirEntry
@@ -21,7 +25,10 @@ import org.jetbrains.kotlin.fir.declarations.utils.isLocal
 import org.jetbrains.kotlin.fir.diagnostics.ConeDestructuringDeclarationsOnTopLevel
 import org.jetbrains.kotlin.fir.resolve.providers.firProvider
 import org.jetbrains.kotlin.fir.resolve.toSymbol
+import org.jetbrains.kotlin.fir.visitors.FirVisitorVoid
 import org.jetbrains.kotlin.name.ClassId
+import org.jetbrains.kotlin.psi.KtClassOrObject
+import org.jetbrains.kotlin.psi.psiUtil.parents
 
 class FirDesignationWithFile(
     path: List<FirRegularClass>,
@@ -124,8 +131,52 @@ private fun collectDesignationPathWithContainingClass(target: FirDeclaration, co
         return declaration
     }
 
-    val chain = generateSequence(containingClassId) { it.outerClassId }.map { resolveChunk(it) }
-    return chain.toMutableList().also { it.reverse() }
+    val chunks = generateSequence(containingClassId) { it.outerClassId }.toList()
+
+    if (chunks.any { it.shortClassName.isSpecial }) {
+        val fallbackResult = collectDesignationPathWithTreeTraversal(target)
+        if (fallbackResult != null) {
+            return fallbackResult
+        }
+    }
+
+    return chunks
+        .dropWhile { it.shortClassName.isSpecial }
+        .map { resolveChunk(it) }
+        .asReversed()
+}
+
+/*
+    This implementation is certainly inefficient, however there seem to be no better way to implement designation collection for
+    anonymous outer classes unless FIR tree gets a way to get an element parent.
+ */
+private fun collectDesignationPathWithTreeTraversal(target: FirDeclaration): List<FirRegularClass>? {
+    val containingFile = target.getContainingFile() ?: return null
+
+    val path = ArrayDeque<FirElement>()
+    path.addLast(containingFile)
+
+    var result: List<FirRegularClass>? = null
+
+    val visitor = object : FirVisitorVoid() {
+        override fun visitElement(element: FirElement) {
+            if (result != null) {
+                return
+            } else if (element === target) {
+                result = path.filterIsInstance<FirRegularClass>()
+            } else {
+                try {
+                    path.addLast(element)
+                    element.acceptChildren(this)
+                } finally {
+                    path.removeLast()
+                }
+            }
+        }
+    }
+
+    containingFile.accept(visitor)
+    return result
 }
 
 private fun getTargetSession(target: FirDeclaration): FirSession {