Diagnostics with psi elements added

compiler/frontend/src/org/jetbrains/jet/lang/diagnostics/DiagnosticWithPsiElement.java

@@ -0,0 +1,22 @@
+package org.jetbrains.jet.lang.diagnostics;
+
+import com.intellij.openapi.util.TextRange;
+import com.intellij.psi.PsiElement;
+import org.jetbrains.annotations.NotNull;
+
+/**
+ * @author svtk
+ */
+public class DiagnosticWithPsiElement<T extends PsiElement> extends GenericDiagnostic {
+    private final T psiElement;
+
+    public DiagnosticWithPsiElement(DiagnosticFactory factory, Severity severity, String message, T psiElement) {
+        super(factory, severity, message, psiElement.getTextRange());
+        this.psiElement = psiElement;
+    }
+
+    @NotNull
+    public T getPsiElement() {
+        return psiElement;
+    }
+}

compiler/frontend/src/org/jetbrains/jet/lang/diagnostics/Errors.java

@@ -63,7 +63,7 @@ public interface Errors {
 
         @NotNull
         public Diagnostic on(@NotNull PsiElement element) {
-            return on(element.getTextRange());
+            return new DiagnosticWithPsiElement<PsiElement>(this, severity, message, element);
         }
 
         @NotNull
@@ -102,7 +102,7 @@ public interface Errors {
 
         @NotNull
         public Diagnostic on(@NotNull PsiElement element, @NotNull T argument) {
-            return on(element.getTextRange(), argument);
+            return new DiagnosticWithPsiElement<PsiElement>(this, severity, makeMessage(argument), element);
         }
     }
 
@@ -139,7 +139,7 @@ public interface Errors {
 
         @NotNull
         public Diagnostic on(@NotNull PsiElement element, @NotNull A a, @NotNull B b) {
-            return on(element.getTextRange(), a, b);
+            return new DiagnosticWithPsiElement<PsiElement>(this, severity, makeMessage(a, b), element);
         }
     }
     
@@ -180,7 +180,7 @@ public interface Errors {
 
         @NotNull
         public Diagnostic on(@NotNull PsiElement element, @NotNull A a, @NotNull B b, @NotNull C c) {
-            return on(element.getTextRange(), a, b, c);
+            return new DiagnosticWithPsiElement<PsiElement>(this, severity, makeMessage(a, b, c), element);
         }
     }
 
@@ -537,18 +537,11 @@ public interface Errors {
 
     RedeclarationDiagnosticFactory REDECLARATION = RedeclarationDiagnosticFactory.INSTANCE;
 
-    public class UnresolvedReferenceDiagnostic extends GenericDiagnostic {
-
-        private final JetReferenceExpression reference;
+    public class UnresolvedReferenceDiagnostic extends DiagnosticWithPsiElement<JetReferenceExpression> {
 
         public UnresolvedReferenceDiagnostic(JetReferenceExpression referenceExpression) {
-            super(UNRESOLVED_REFERENCE, ERROR, "Unresolved reference", referenceExpression.getTextRange());
-            this.reference = referenceExpression;
+            super(UNRESOLVED_REFERENCE, ERROR, "Unresolved reference", referenceExpression);
         }
 
-        @NotNull
-        public JetReferenceExpression getReference() {
-            return reference;
-        }
     }
 }

idea/src/org/jetbrains/jet/plugin/annotations/DebugInfoAnnotator.java

@@ -57,7 +57,7 @@ public class DebugInfoAnnotator implements Annotator {
                 final Set<JetReferenceExpression> unresolvedReferences = Sets.newHashSet();
                 for (Diagnostic diagnostic : bindingContext.getDiagnostics()) {
                     if (diagnostic instanceof Errors.UnresolvedReferenceDiagnostic) {
-                        unresolvedReferences.add(((Errors.UnresolvedReferenceDiagnostic) diagnostic).getReference());
+                        unresolvedReferences.add(((Errors.UnresolvedReferenceDiagnostic) diagnostic).getPsiElement());
                     }
                 }
 

idea/tests/org/jetbrains/jet/JetTestUtils.java

@@ -51,7 +51,7 @@ public class JetTestUtils {
         public void report(@NotNull Diagnostic diagnostic) {
             if (diagnostic instanceof Errors.UnresolvedReferenceDiagnostic) {
                 Errors.UnresolvedReferenceDiagnostic unresolvedReferenceDiagnostic = (Errors.UnresolvedReferenceDiagnostic) diagnostic;
-                throw new IllegalStateException("Unresolved: " + unresolvedReferenceDiagnostic.getReference().getText());
+                throw new IllegalStateException("Unresolved: " + unresolvedReferenceDiagnostic.getPsiElement().getText());
             }
         }
     };

idea/tests/org/jetbrains/jet/resolve/ExpectedResolveData.java

@@ -100,7 +100,7 @@ public class ExpectedResolveData {
         for (Diagnostic diagnostic : bindingContext.getDiagnostics()) {
             if (diagnostic instanceof Errors.UnresolvedReferenceDiagnostic) {
                 Errors.UnresolvedReferenceDiagnostic unresolvedReferenceDiagnostic = (Errors.UnresolvedReferenceDiagnostic) diagnostic;
-                unresolvedReferences.add(unresolvedReferenceDiagnostic.getReference());
+                unresolvedReferences.add(unresolvedReferenceDiagnostic.getPsiElement());
             }
         }