From 1c36ffb062d0b03c2904d2c414b549d9360dd4f9 Mon Sep 17 00:00:00 2001
From: Alexey Utkin <alexey.utkin@jetbrains.com>
Date: Fri, 1 Jul 2016 17:20:39 +0300
Subject: [PATCH] CPP-778 Support for semantic per-variable highlighting
 (visitor fixes)

---
 .../codeInsight/daemon/RainbowVisitor.java    | 29 ++++++++-----------
 1 file changed, 12 insertions(+), 17 deletions(-)

diff --git a/platform/analysis-impl/src/com/intellij/codeInsight/daemon/RainbowVisitor.java b/platform/analysis-impl/src/com/intellij/codeInsight/daemon/RainbowVisitor.java
index e9d359516a85..7f9849cb8b54 100644
--- a/platform/analysis-impl/src/com/intellij/codeInsight/daemon/RainbowVisitor.java
+++ b/platform/analysis-impl/src/com/intellij/codeInsight/daemon/RainbowVisitor.java
@@ -20,25 +20,20 @@ import com.intellij.codeInsight.daemon.impl.HighlightInfo;
 import com.intellij.codeInsight.daemon.impl.HighlightVisitor;
 import com.intellij.codeInsight.daemon.impl.analysis.HighlightInfoHolder;
 import com.intellij.openapi.extensions.Extensions;
-import com.intellij.psi.PsiElement;
 import com.intellij.psi.PsiFile;
-import com.intellij.psi.PsiRecursiveElementWalkingVisitor;
 import org.jetbrains.annotations.NotNull;
 
 public abstract class RainbowVisitor implements HighlightVisitor {
   private HighlightInfoHolder myHolder;
-
-  @NotNull
-  protected abstract PsiRecursiveElementWalkingVisitor getVisitor(@NotNull final RainbowHighlighter highlighter);
+  private RainbowHighlighter myRainbowHighlighter;
 
   @NotNull
   @Override
   public abstract HighlightVisitor clone();
 
-  @Override
-  public void visit(@NotNull PsiElement element) {
-    RainbowHighlighter highlighter = new RainbowHighlighter(myHolder.getColorsScheme());
-    element.accept(getVisitor(highlighter));
+  @NotNull
+  protected RainbowHighlighter getHighlighter() {
+    return myRainbowHighlighter;
   }
 
   @Override
@@ -46,14 +41,14 @@ public abstract class RainbowVisitor implements HighlightVisitor {
                          boolean updateWholeFile,
                          @NotNull HighlightInfoHolder holder,
                          @NotNull Runnable action) {
-    if (RainbowHighlighter.isRainbowEnabled()) {
-      myHolder = holder;
-      try {
-        action.run();
-      }
-      finally {
-        myHolder = null;
-      }
+    myHolder = holder;
+    myRainbowHighlighter = new RainbowHighlighter(myHolder.getColorsScheme());
+    try {
+      action.run();
+    }
+    finally {
+      myHolder = null;
+      myRainbowHighlighter = null;
     }
     return true;
   }