getEnclosingElementForLocalDeclaration() made stub-friendly

compiler/frontend/src/org/jetbrains/kotlin/psi/JetPsiUtil.java

@@ -834,15 +834,24 @@ public class JetPsiUtil {
             }
         }
 
-        //noinspection unchecked
+        // No appropriate stub-tolerant method in PsiTreeUtil, nor JetStubbedPsiUtil, writing manually
-        JetElement container = PsiTreeUtil.getParentOfType(
+        PsiElement current = PsiTreeUtil.getStubOrPsiParent(declaration);
-                declaration,
+        while (current != null) {
-                JetBlockExpression.class, JetClassInitializer.class, JetProperty.class, JetFunction.class, JetParameter.class
+            PsiElement parent = PsiTreeUtil.getStubOrPsiParent(current);
-        );
+            if (parent instanceof JetScript) return null;
-        if (container == null) return null;
+            if (current instanceof JetClassInitializer) {
-        if (container.getParent() instanceof JetScript) return null;
+                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) {