Get diagnostics for particular element

compiler/frontend/src/org/jetbrains/jet/lang/resolve/Diagnostics.java

@@ -1,6 +1,24 @@
+/*
+ * Copyright 2010-2013 JetBrains s.r.o.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
 package org.jetbrains.jet.lang.resolve;
 
+import com.intellij.psi.PsiElement;
 import org.jetbrains.annotations.NotNull;
+import org.jetbrains.annotations.Nullable;
 import org.jetbrains.jet.lang.diagnostics.Diagnostic;
 
 import java.util.Collection;
@@ -11,6 +29,9 @@ public interface Diagnostics extends Iterable<Diagnostic> {
     @NotNull
     Collection<Diagnostic> all();
 
+    @NotNull
+    Collection<Diagnostic> forElement(@NotNull PsiElement psiElement);
+
     boolean isEmpty();
 
     @NotNull
@@ -23,6 +44,12 @@ public interface Diagnostics extends Iterable<Diagnostic> {
             return Collections.emptyList();
         }
 
+        @NotNull
+        @Override
+        public Collection<Diagnostic> forElement(@NotNull PsiElement psiElement) {
+            return Collections.emptyList();
+        }
+
         @Override
         public boolean isEmpty() {
             return true;

compiler/frontend/src/org/jetbrains/jet/lang/resolve/DiagnosticsElementsCache.java

@@ -0,0 +1,57 @@
+/*
+ * Copyright 2010-2013 JetBrains s.r.o.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.jetbrains.jet.lang.resolve;
+
+import com.intellij.openapi.util.AtomicNotNullLazyValue;
+import com.intellij.psi.PsiElement;
+import com.intellij.util.containers.ConcurrentMultiMap;
+import com.intellij.util.containers.MultiMap;
+import org.jetbrains.annotations.NotNull;
+import org.jetbrains.annotations.Nullable;
+import org.jetbrains.jet.lang.diagnostics.Diagnostic;
+
+import java.util.Collection;
+
+public class DiagnosticsElementsCache {
+    private final Diagnostics diagnostics;
+
+    private final AtomicNotNullLazyValue<MultiMap<PsiElement, Diagnostic>> elementToDiagnostic = new AtomicNotNullLazyValue<MultiMap<PsiElement, Diagnostic>>() {
+        @NotNull
+        @Override
+        protected MultiMap<PsiElement, Diagnostic> compute() {
+            return buildElementToDiagnosticCache(diagnostics);
+        }
+    };
+
+    public DiagnosticsElementsCache(Diagnostics diagnostics) {
+        this.diagnostics = diagnostics;
+    }
+
+    @NotNull
+    public Collection<Diagnostic> getDiagnostics(@NotNull PsiElement psiElement) {
+        return elementToDiagnostic.getValue().get(psiElement);
+    }
+
+    private static MultiMap<PsiElement, Diagnostic> buildElementToDiagnosticCache(Diagnostics diagnostics) {
+        MultiMap<PsiElement, Diagnostic> elementToDiagnostic = new ConcurrentMultiMap<PsiElement, Diagnostic>();
+        for (Diagnostic diagnostic : diagnostics) {
+            elementToDiagnostic.putValue(diagnostic.getPsiElement(), diagnostic);
+        }
+
+        return elementToDiagnostic;
+    }
+}

compiler/frontend/src/org/jetbrains/jet/lang/resolve/DiagnosticsWithSuppression.java

@@ -1,3 +1,19 @@
+/*
+ * Copyright 2010-2013 JetBrains s.r.o.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
 package org.jetbrains.jet.lang.resolve;
 
 import com.google.common.collect.ImmutableSet;
@@ -43,6 +59,7 @@ public class DiagnosticsWithSuppression implements Diagnostics {
             return !isSuppressed(diagnostic);
         }
     };
+    private final DiagnosticsElementsCache elementsCache = new DiagnosticsElementsCache(this);
 
     public DiagnosticsWithSuppression(@NotNull BindingContext context, @NotNull Collection<Diagnostic> diagnostics) {
         this.context = context;
@@ -72,6 +89,12 @@ public class DiagnosticsWithSuppression implements Diagnostics {
         return ContainerUtil.filter(diagnostics, filter);
     }
 
+    @NotNull
+    @Override
+    public Collection<Diagnostic> forElement(@NotNull PsiElement psiElement) {
+        return elementsCache.getDiagnostics(psiElement);
+    }
+
     @Override
     public boolean isEmpty() {
         return all().isEmpty();

compiler/frontend/src/org/jetbrains/jet/lang/resolve/SimpleDiagnostics.java

@@ -1,5 +1,22 @@
+/*
+ * Copyright 2010-2013 JetBrains s.r.o.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
 package org.jetbrains.jet.lang.resolve;
 
+import com.intellij.psi.PsiElement;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.jet.lang.diagnostics.Diagnostic;
 
@@ -8,6 +25,7 @@ import java.util.Iterator;
 
 public class SimpleDiagnostics implements Diagnostics {
     private final Collection<Diagnostic> diagnostics;
+    private final DiagnosticsElementsCache elementsCache = new DiagnosticsElementsCache(this);
 
     public SimpleDiagnostics(@NotNull Collection<Diagnostic> diagnostics) {
         this.diagnostics = diagnostics;
@@ -19,6 +37,12 @@ public class SimpleDiagnostics implements Diagnostics {
         return diagnostics;
     }
 
+    @NotNull
+    @Override
+    public Collection<Diagnostic> forElement(@NotNull PsiElement psiElement) {
+        return elementsCache.getDiagnostics(psiElement);
+    }
+
     @NotNull
     @Override
     public Diagnostics noSuppression() {