diff --git a/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/element/builder/FirTowerDataContextCollector.kt b/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/element/builder/FirTowerDataContextCollector.kt
index dc0f7a73bc7..9ce4911a613 100644
--- a/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/element/builder/FirTowerDataContextCollector.kt
+++ b/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/element/builder/FirTowerDataContextCollector.kt
@@ -50,11 +50,7 @@ internal class FirTowerDataContextAllElementsCollector : FirTowerDataContextColl
         var current: PsiElement? = ktElement
         while (current != null) {
             if (current is KtElement) {
-                elementsToContext[current]?.let { return it }
-            }
-            if (current is KtDeclaration) {
-                val originalDeclaration = current.originalDeclaration
-                originalDeclaration?.let { elementsToContext[it] }?.let { return it }
+                getContextForElement(current)?.let { return it }
             }
             if (current is KtFile) {
                 break
@@ -63,6 +59,17 @@ internal class FirTowerDataContextAllElementsCollector : FirTowerDataContextColl
         }
         return null
     }
+
+    private fun getContextForElement(element: KtElement): FirTowerDataContext? {
+        elementsToContext[element]?.let { return it }
+
+        if (element is KtDeclaration) {
+            val originalDeclaration = element.originalDeclaration
+            originalDeclaration?.let { elementsToContext[it] }?.let { return it }
+        }
+
+        return null
+    }
 }
 
 private tailrec fun PsiElement.closestBlockLevelOrInitializerExpression(): KtExpression? =