From eebe013de2b3955b8d075cece08d40403f9e1b45 Mon Sep 17 00:00:00 2001
From: Andrey Breslav <andrey.breslav@jetbrains.com>
Date: Tue, 27 Jan 2015 16:03:21 +0300
Subject: [PATCH] getEnclosingElementForLocalDeclaration() made stub-friendly

---
 .../org/jetbrains/kotlin/psi/JetPsiUtil.java  | 25 +++++++++++++------
 1 file changed, 17 insertions(+), 8 deletions(-)

diff --git a/compiler/frontend/src/org/jetbrains/kotlin/psi/JetPsiUtil.java b/compiler/frontend/src/org/jetbrains/kotlin/psi/JetPsiUtil.java
index b7a8a967099..052152421db 100644
--- a/compiler/frontend/src/org/jetbrains/kotlin/psi/JetPsiUtil.java
+++ b/compiler/frontend/src/org/jetbrains/kotlin/psi/JetPsiUtil.java
@@ -834,15 +834,24 @@ public class JetPsiUtil {
             }
         }
 
-        //noinspection unchecked
-        JetElement container = PsiTreeUtil.getParentOfType(
-                declaration,
-                JetBlockExpression.class, JetClassInitializer.class, JetProperty.class, JetFunction.class, JetParameter.class
-        );
-        if (container == null) return null;
-        if (container.getParent() instanceof JetScript) return null;
+        // No appropriate stub-tolerant method in PsiTreeUtil, nor JetStubbedPsiUtil, writing manually
+        PsiElement current = PsiTreeUtil.getStubOrPsiParent(declaration);
+        while (current != null) {
+            PsiElement parent = PsiTreeUtil.getStubOrPsiParent(current);
+            if (parent instanceof JetScript) return null;
+            if (current instanceof JetClassInitializer) {
+                return ((JetClassInitializer) current).getBody();
+            }
+            if (current instanceof JetBlockExpression ||
+                current instanceof JetProperty ||
+                current instanceof JetParameter ||
+                current instanceof JetFunction) {
+                return (JetElement) current;
+            }
 
-        return (container instanceof JetClassInitializer) ? ((JetClassInitializer) container).getBody() : container;
+            current = parent;
+        }
+        return null;
     }
 
     public static boolean isLocal(@NotNull JetDeclaration declaration) {