KT-4017 Syntax highlighting drops constantly for ~0.5s with every little change in the source

#KT-4017 Fixed

idea/src/META-INF/plugin.xml

@@ -214,6 +214,8 @@
     <liveTemplateMacro implementation="org.jetbrains.jet.plugin.liveTemplates.macro.JetSuggestVariableNameMacro"/>
 
     <annotator language="jet" implementationClass="org.jetbrains.jet.plugin.highlighter.JetPsiChecker"/>
+    <highlightRangeExtension implementation="org.jetbrains.jet.plugin.highlighter.JetPsiChecker"/>
+
     <annotator language="jet" implementationClass="org.jetbrains.jet.plugin.highlighter.DebugInfoAnnotator"/>
     <problemHighlightFilter implementation="org.jetbrains.jet.plugin.highlighter.KotlinProblemHighlightFilter"/>
 

idea/src/org/jetbrains/jet/plugin/highlighter/AfterAnalysisHighlightingVisitor.java

@@ -17,7 +17,6 @@
 package org.jetbrains.jet.plugin.highlighter;
 
 import com.intellij.lang.annotation.AnnotationHolder;
-import org.jetbrains.jet.lang.psi.JetElement;
 import org.jetbrains.jet.lang.resolve.BindingContext;
 
 abstract class AfterAnalysisHighlightingVisitor extends HighlightingVisitor {
@@ -27,9 +26,4 @@ abstract class AfterAnalysisHighlightingVisitor extends HighlightingVisitor {
         super(holder);
         this.bindingContext = bindingContext;
     }
-
-    @Override
-    public void visitJetElement(JetElement element) {
-        element.acceptChildren(this);
-    }
 }

idea/src/org/jetbrains/jet/plugin/highlighter/JetPsiChecker.java

@@ -17,6 +17,7 @@
 package org.jetbrains.jet.plugin.highlighter;
 
 import com.google.common.collect.Sets;
+import com.intellij.codeInsight.daemon.impl.HighlightRangeExtension;
 import com.intellij.codeInsight.intention.EmptyIntentionAction;
 import com.intellij.codeInsight.intention.IntentionAction;
 import com.intellij.codeInspection.ProblemHighlightType;
@@ -26,15 +27,14 @@ import com.intellij.lang.annotation.Annotator;
 import com.intellij.openapi.application.ApplicationManager;
 import com.intellij.openapi.diagnostic.Logger;
 import com.intellij.openapi.editor.colors.TextAttributesKey;
-import com.intellij.openapi.progress.ProcessCanceledException;
 import com.intellij.openapi.util.TextRange;
-import com.intellij.psi.MultiRangeReference;
+import com.intellij.psi.*;
-import com.intellij.psi.PsiElement;
+import com.intellij.util.containers.MultiMap;
-import com.intellij.psi.PsiReference;
 import com.intellij.xml.util.XmlStringUtil;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 import org.jetbrains.annotations.TestOnly;
+import org.jetbrains.jet.analyzer.AnalyzeExhaust;
 import org.jetbrains.jet.lang.diagnostics.Diagnostic;
 import org.jetbrains.jet.lang.diagnostics.Errors;
 import org.jetbrains.jet.lang.diagnostics.Severity;
@@ -51,11 +51,12 @@ import java.util.Collection;
 import java.util.List;
 import java.util.Set;
 
-public class JetPsiChecker implements Annotator {
+public class JetPsiChecker implements Annotator, HighlightRangeExtension {
     private static boolean namesHighlightingTest;
-
     private static final Logger LOG = Logger.getInstance(JetPsiChecker.class);
 
+    private HighlightingPassCache passCache = null;
+
     @TestOnly
     public static void setNamesHighlightingTest(boolean namesHighlightingTest) {
         JetPsiChecker.namesHighlightingTest = namesHighlightingTest;
@@ -92,8 +93,7 @@ public class JetPsiChecker implements Annotator {
 
     @Override
     public void annotate(@NotNull PsiElement element, @NotNull AnnotationHolder holder) {
-        if (!JetPluginUtil.isInSource(element) ||
+        if (!JetPluginUtil.isInSource(element) || JetPluginUtil.isKtFileInGradleProjectInWrongFolder(element)) {
-                JetPluginUtil.isKtFileInGradleProjectInWrongFolder(element)) {
             return;
         }
 
@@ -101,35 +101,31 @@ public class JetPsiChecker implements Annotator {
             element.accept(visitor);
         }
 
-        if (element instanceof JetFile) {
+        JetFile file = (JetFile) element.getContainingFile();
-            JetFile file = (JetFile) element;
 
-            try {
+        if (passCache == null || passCache.isOutdated(file)) {
-                BindingContext bindingContext = AnalyzerFacadeWithCache.analyzeFileWithCache(file).getBindingContext();
+            passCache = new HighlightingPassCache(AnalyzerFacadeWithCache.analyzeFileWithCache(file), file);
+        }
 
-                if (JetPluginUtil.isInSource(element, /* includeLibrarySources = */ false)) {
+        if (!passCache.analyzeExhaust.isError()) {
-                    Collection<Diagnostic> diagnostics = Sets.newLinkedHashSet(bindingContext.getDiagnostics());
+            BindingContext bindingContext = passCache.analyzeExhaust.getBindingContext();
-                    Set<PsiElement> redeclarations = Sets.newHashSet();
+            for (HighlightingVisitor visitor : getAfterAnalysisVisitor(holder, bindingContext)) {
-                    for (Diagnostic diagnostic : diagnostics) {
+                element.accept(visitor);
-                        // This is needed because we have the same context for all files
+            }
-                        if (diagnostic.getPsiFile() != file) continue;
 
-                        registerDiagnosticAnnotations(diagnostic, redeclarations, holder);
+            if (JetPluginUtil.isInSource(element, /* includeLibrarySources = */ false)) {
-                    }
+                for (Diagnostic diagnostic : passCache.elementToDiagnostic.get(element)) {
-                }
+                    registerDiagnosticAnnotations(diagnostic, passCache.redeclarations, holder);
-
-                for (HighlightingVisitor visitor : getAfterAnalysisVisitor(holder, bindingContext)) {
-                    file.acceptChildren(visitor);
                 }
             }
-            catch (ProcessCanceledException e) {
+        }
-                throw e;
+        else if (element instanceof JetFile) {
-            }
+            Throwable error = passCache.analyzeExhaust.getError();
-            catch (Throwable e) {
+            if (JetPluginUtil.isInSource(element, /* includeLibrarySources = */ false)) {
-                // For failing tests and to notify about idea internal error in -ea mode
+                holder.createErrorAnnotation(file, error.getClass().getCanonicalName() + ": " + error.getMessage());
-                holder.createErrorAnnotation(element, e.getClass().getCanonicalName() + ": " + e.getMessage());
-                LOG.error(e);
             }
+
+            LOG.error(error);
         }
     }
 
@@ -277,4 +273,42 @@ public class JetPsiChecker implements Annotator {
         annotation.setTooltip(getMessage(redeclarationDiagnostic));
         return annotation;
     }
+
+    @Override
+    public boolean isForceHighlightParents(@NotNull PsiFile file) {
+        return file instanceof JetFile;
+    }
+
+    private static class HighlightingPassCache {
+        private final AnalyzeExhaust analyzeExhaust;
+        private final MultiMap<PsiElement, Diagnostic> elementToDiagnostic;
+
+        private final Set<PsiElement> redeclarations = Sets.newHashSet();
+        private final JetFile jetFile;
+        private final long modificationCount;
+
+        public HighlightingPassCache(AnalyzeExhaust analyzeExhaust, JetFile jetFile) {
+            this.analyzeExhaust = analyzeExhaust;
+            this.jetFile = jetFile;
+            this.elementToDiagnostic = buildElementToDiagnosticCache(analyzeExhaust, jetFile);
+            this.modificationCount = PsiManager.getInstance(jetFile.getProject()).getModificationTracker().getModificationCount();
+        }
+
+        public boolean isOutdated(JetFile jetFile) {
+            return this.jetFile != jetFile || PsiManager.getInstance(jetFile.getProject()).getModificationTracker().getModificationCount() != modificationCount;
+        }
+
+        private static MultiMap<PsiElement, Diagnostic> buildElementToDiagnosticCache(AnalyzeExhaust analyzeExhaust, JetFile jetFile) {
+            MultiMap<PsiElement, Diagnostic> elementToDiagnostic = MultiMap.create();
+            Collection<Diagnostic> diagnostics = Sets.newLinkedHashSet(analyzeExhaust.getBindingContext().getDiagnostics());
+
+            for (Diagnostic diagnostic : diagnostics) {
+                if (diagnostic.getPsiFile() == jetFile) {
+                    elementToDiagnostic.putValue(diagnostic.getPsiElement(), diagnostic);
+                }
+            }
+
+            return elementToDiagnostic;
+        }
+    }
 }