Removed cache for labeled elements

Resolve labels on the fly

compiler/frontend/src/org/jetbrains/jet/lang/psi/JetPsiUtil.java

@@ -136,6 +136,14 @@ public class JetPsiUtil {
         return JetTokens.LABELS.contains(expression.getOperationReference().getReferencedNameElementType());
     }
 
+    @NotNull
+    public static String getLabelName(@NotNull JetPrefixExpression expression) {
+        assert isLabeledExpression(expression);
+        String labelName = expression.getOperationReference().getReferencedName();
+        assert labelName.startsWith("@") : "Incorrect label name " + expression.getText();
+        return labelName.substring(1);
+    }
+
     @NotNull
     public static Name safeName(@Nullable String name) {
         return name == null ? SpecialNames.NO_NAME_PROVIDED : Name.identifier(name);

compiler/frontend/src/org/jetbrains/jet/lang/types/expressions/BasicExpressionTypingVisitor.java

@@ -748,12 +748,7 @@ public class BasicExpressionTypingVisitor extends ExpressionTypingVisitor {
         JetSimpleNameExpression operationSign = expression.getOperationReference();
         assert JetTokens.LABELS.contains(operationSign.getReferencedNameElementType());
 
-        String referencedName = operationSign.getReferencedName();
+        return facade.getTypeInfo(baseExpression, context, isStatement);
-        context.labelResolver.enterLabeledElement(Name.identifierForLabel(referencedName.substring(1)), baseExpression);
-        // TODO : Some processing for the label?
-        JetTypeInfo typeInfo = facade.getTypeInfo(baseExpression, context, isStatement);
-        context.labelResolver.exitLabeledElement(baseExpression);
-        return typeInfo;
     }
 
     private static boolean isKnownToBeNotNull(JetExpression expression, ExpressionTypingContext context) {

compiler/frontend/src/org/jetbrains/jet/lang/types/expressions/ClosureExpressionsTypingVisitor.java

@@ -103,11 +103,6 @@ public class ClosureExpressionsTypingVisitor extends ExpressionTypingVisitor {
         JetBlockExpression bodyExpression = expression.getFunctionLiteral().getBodyExpression();
         if (bodyExpression == null) return null;
 
-        Name callerName = getCallerName(expression);
-        if (callerName != null) {
-            context.labelResolver.enterLabeledElement(callerName, expression);
-        }
-
         JetType expectedType = context.expectedType;
         boolean functionTypeExpected = !noExpectedType(expectedType) && KotlinBuiltIns.getInstance().isFunctionOrExtensionFunctionType(
                 expectedType);
@@ -125,48 +120,9 @@ public class ClosureExpressionsTypingVisitor extends ExpressionTypingVisitor {
             return JetTypeInfo.create(resultType, context.dataFlowInfo);
         }
 
-        if (callerName != null) {
-            context.labelResolver.exitLabeledElement(expression);
-        }
-
         return DataFlowUtils.checkType(resultType, expression, context, context.dataFlowInfo);
     }
 
-    @Nullable
-    private static Name getCallerName(@NotNull JetFunctionLiteralExpression expression) {
-        JetCallExpression callExpression = getContainingCallExpression(expression);
-        if (callExpression == null) return null;
-
-        JetExpression calleeExpression = callExpression.getCalleeExpression();
-        if (calleeExpression instanceof JetSimpleNameExpression) {
-            JetSimpleNameExpression nameExpression = (JetSimpleNameExpression) calleeExpression;
-            return nameExpression.getReferencedNameAsName();
-        }
-
-        return null;
-    }
-
-    @Nullable
-    private static JetCallExpression getContainingCallExpression(JetFunctionLiteralExpression expression) {
-        PsiElement parent = expression.getParent();
-        if (parent instanceof JetCallExpression) {
-            // f {}
-            return (JetCallExpression) parent;
-        }
-
-        if (parent instanceof JetValueArgument) {
-            // f ({}) or f(p = {})
-            JetValueArgument argument = (JetValueArgument) parent;
-            PsiElement argList = argument.getParent();
-            if (argList == null) return null;
-            PsiElement call = argList.getParent();
-            if (call instanceof JetCallExpression) {
-                return (JetCallExpression) call;
-            }
-        }
-        return null;
-    }
-
     @NotNull
     private AnonymousFunctionDescriptor createFunctionDescriptor(
             @NotNull JetFunctionLiteralExpression expression,

compiler/frontend/src/org/jetbrains/jet/lang/types/expressions/LabelResolver.java

@@ -16,8 +16,8 @@
 
 package org.jetbrains.jet.lang.types.expressions;
 
+import com.google.common.collect.Lists;
 import com.intellij.psi.PsiElement;
-import com.intellij.util.containers.Stack;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 import org.jetbrains.jet.lang.descriptors.*;
@@ -28,9 +28,7 @@ import org.jetbrains.jet.lang.resolve.DescriptorResolver;
 import org.jetbrains.jet.lang.resolve.name.Name;
 
 import java.util.Collection;
-import java.util.HashMap;
+import java.util.List;
-import java.util.Iterator;
-import java.util.Map;
 
 import static org.jetbrains.jet.lang.diagnostics.Errors.LABEL_NAME_CLASH;
 import static org.jetbrains.jet.lang.diagnostics.Errors.UNRESOLVED_REFERENCE;
@@ -44,50 +42,80 @@ public class LabelResolver {
         return new LabelResolver();
     }
 
-    private final Map<Name, Stack<JetElement>> labeledElements = new HashMap<Name, Stack<JetElement>>();
-
     private LabelResolver() {}
 
-    public void enterLabeledElement(@NotNull Name labelName, @NotNull JetExpression labeledExpression) {
+    @NotNull
-        JetExpression cacheExpression = getCachingExpression(labeledExpression);
+    private List<JetElement> getElementsByLabelName(@NotNull Name labelName, @NotNull JetSimpleNameExpression labelExpression) {
-        if (cacheExpression != null) {
+        List<JetElement> elements = Lists.newArrayList();
-            Stack<JetElement> stack = labeledElements.get(labelName);
+        PsiElement parent = labelExpression.getParent();
-            if (stack == null) {
+        while (parent != null) {
-                stack = new Stack<JetElement>();
+            Name name = getLabelNameIfAny(parent);
-                labeledElements.put(labelName, stack);
+            if (name != null && name.equals(labelName)) {
+                elements.add(getExpressionUnderLabel((JetExpression) parent));
             }
-            stack.push(cacheExpression);
+            parent = parent.getParent();
         }
+        return elements;
     }
 
-    public void exitLabeledElement(@NotNull JetExpression expression) {
+    @Nullable
-        JetExpression cacheExpression = getCachingExpression(expression);
+    private Name getLabelNameIfAny(@NotNull PsiElement element) {
-
+        if (element instanceof JetPrefixExpression) {
-        // TODO : really suboptimal
+            JetPrefixExpression prefixExpression = (JetPrefixExpression) element;
-        for (Iterator<Map.Entry<Name,Stack<JetElement>>> mapIter = labeledElements.entrySet().iterator(); mapIter.hasNext(); ) {
+            if (JetPsiUtil.isLabeledExpression(prefixExpression)) {
-            Map.Entry<Name, Stack<JetElement>> entry = mapIter.next();
+                return Name.identifierForLabel(JetPsiUtil.getLabelName(prefixExpression));
-            Stack<JetElement> stack = entry.getValue();
-            for (Iterator<JetElement> stackIter = stack.iterator(); stackIter.hasNext(); ) {
-                JetElement recorded = stackIter.next();
-                if (recorded == cacheExpression) {
-                    stackIter.remove();
-                }
-            }
-            if (stack.isEmpty()) {
-                mapIter.remove();
             }
         }
+        else if (element instanceof JetFunctionLiteralExpression) {
+            return getCallerName((JetFunctionLiteralExpression) element);
+        }
+        return null;
     }
 
     @NotNull
-    private JetExpression getCachingExpression(@NotNull JetExpression labeledExpression) {
+    private JetExpression getExpressionUnderLabel(@NotNull JetExpression labeledExpression) {
-        JetExpression expression = JetPsiUtil.deparenthesize(labeledExpression);
+        JetExpression expression = JetPsiUtil.safeDeparenthesize(labeledExpression, true);
         if (expression instanceof JetFunctionLiteralExpression) {
-            expression = ((JetFunctionLiteralExpression) expression).getFunctionLiteral();
+            return ((JetFunctionLiteralExpression) expression).getFunctionLiteral();
         }
         return expression;
     }
 
+    @Nullable
+    private Name getCallerName(@NotNull JetFunctionLiteralExpression expression) {
+        JetCallExpression callExpression = getContainingCallExpression(expression);
+        if (callExpression == null) return null;
+
+        JetExpression calleeExpression = callExpression.getCalleeExpression();
+        if (calleeExpression instanceof JetSimpleNameExpression) {
+            JetSimpleNameExpression nameExpression = (JetSimpleNameExpression) calleeExpression;
+            return nameExpression.getReferencedNameAsName();
+        }
+
+        return null;
+    }
+
+    @Nullable
+    private JetCallExpression getContainingCallExpression(JetFunctionLiteralExpression expression) {
+        PsiElement parent = expression.getParent();
+        if (parent instanceof JetCallExpression) {
+            // f {}
+            return (JetCallExpression) parent;
+        }
+
+        if (parent instanceof JetValueArgument) {
+            // f ({}) or f(p = {})
+            JetValueArgument argument = (JetValueArgument) parent;
+            PsiElement argList = argument.getParent();
+            if (argList == null) return null;
+            PsiElement call = argList.getParent();
+            if (call instanceof JetCallExpression) {
+                return (JetCallExpression) call;
+            }
+        }
+        return null;
+    }
+
     @Nullable
     private JetElement resolveControlLabel(@NotNull Name labelName, @NotNull JetSimpleNameExpression labelExpression, boolean reportUnresolved, ExpressionTypingContext context) {
         Collection<DeclarationDescriptor> declarationsByLabel = context.scope.getDeclarationsByLabel(labelName);
@@ -123,18 +151,18 @@ public class LabelResolver {
     }
 
     private JetElement resolveNamedLabel(@NotNull Name labelName, @NotNull JetSimpleNameExpression labelExpression, boolean reportUnresolved, ExpressionTypingContext context) {
-        Stack<JetElement> stack = labeledElements.get(labelName);
+        List<JetElement> list = getElementsByLabelName(labelName, labelExpression);
-        if (stack == null || stack.isEmpty()) {
+        if (list.isEmpty()) {
             if (reportUnresolved) {
                 context.trace.report(UNRESOLVED_REFERENCE.on(labelExpression, labelExpression));
             }
             return null;
         }
-        else if (stack.size() > 1) {
+        else if (list.size() > 1) {
             context.trace.report(LABEL_NAME_CLASH.on(labelExpression));
         }
 
-        JetElement result = stack.peek();
+        JetElement result = list.get(0);
         context.trace.record(LABEL_TARGET, labelExpression, result);
         return result;
     }