From 593ee17cba14933849e7f4dc03737433f25bf2fc Mon Sep 17 00:00:00 2001
From: Tomasz Pasternak <tomasz.pasternak@jetbrains.com>
Date: Tue, 3 Dec 2024 15:11:55 +0100
Subject: [PATCH] fix(qsync): Hightlighting level is always updated after
 'enable analysis'

---
 base/src/META-INF/blaze-base.xml              |  1 +
 .../QuerySyncHighlightingResetListener.java   | 41 +++++++++++++++++++
 .../QuerySyncHighlightingSettingProvider.java |  4 ++
 3 files changed, 46 insertions(+)
 create mode 100644 base/src/com/google/idea/blaze/base/qsync/QuerySyncHighlightingResetListener.java

diff --git a/base/src/META-INF/blaze-base.xml b/base/src/META-INF/blaze-base.xml
index 48cc519444d..0b411162145 100644
--- a/base/src/META-INF/blaze-base.xml
+++ b/base/src/META-INF/blaze-base.xml
@@ -712,6 +712,7 @@
     <ArtifactFetcher implementation="com.google.idea.blaze.base.qsync.cache.FileApiArtifactFetcher"/>
     <QuerySyncListenerProvider implementation="com.google.idea.blaze.base.qsync.ProjectUpdater$Provider" />
     <QuerySyncListenerProvider implementation="com.google.idea.blaze.base.qsync.ProjectStatsLogger$Provider" />
+    <QuerySyncListenerProvider implementation="com.google.idea.blaze.base.qsync.QuerySyncHighlightingResetListener$Provider"/>
   </extensions>
 
   <projectListeners>
diff --git a/base/src/com/google/idea/blaze/base/qsync/QuerySyncHighlightingResetListener.java b/base/src/com/google/idea/blaze/base/qsync/QuerySyncHighlightingResetListener.java
new file mode 100644
index 00000000000..930d4af12ad
--- /dev/null
+++ b/base/src/com/google/idea/blaze/base/qsync/QuerySyncHighlightingResetListener.java
@@ -0,0 +1,41 @@
+package com.google.idea.blaze.base.qsync;
+
+import com.google.idea.blaze.common.Context;
+import com.google.idea.blaze.qsync.QuerySyncProjectListener;
+import com.google.idea.blaze.qsync.QuerySyncProjectSnapshot;
+import com.intellij.codeInsight.daemon.impl.analysis.FileHighlightingSettingListener;
+import com.intellij.codeInsight.daemon.impl.analysis.HighlightingSettingsPerFile;
+import com.intellij.openapi.fileEditor.FileEditorManager;
+import com.intellij.openapi.project.Project;
+import com.intellij.openapi.vfs.VirtualFile;
+import com.intellij.psi.PsiManager;
+
+public class QuerySyncHighlightingResetListener implements QuerySyncProjectListener {
+
+    private final Project project;
+
+    QuerySyncHighlightingResetListener(Project project) {
+        this.project = project;
+    }
+
+    @Override
+    public void onNewProjectSnapshot(Context<?> context, QuerySyncProjectSnapshot instance) {
+        for (VirtualFile virtualFile : FileEditorManager.getInstance(project).getOpenFiles()) {
+            var psiFile = PsiManager.getInstance(project).findFile(virtualFile);
+            if (psiFile != null) {
+                var newHighlighting = HighlightingSettingsPerFile.getInstance(project).getHighlightingSettingForRoot(psiFile);
+                project.getMessageBus()
+                        .syncPublisher(FileHighlightingSettingListener.SETTING_CHANGE)
+                        .settingChanged(psiFile, newHighlighting);
+            }
+        }
+    }
+
+    public static class Provider implements QuerySyncProjectListenerProvider {
+        @Override
+        public QuerySyncProjectListener createListener(QuerySyncProject querySyncProject) {
+            return new QuerySyncHighlightingResetListener(querySyncProject.getIdeProject());
+        }
+    }
+}
+
diff --git a/base/src/com/google/idea/blaze/base/qsync/QuerySyncHighlightingSettingProvider.java b/base/src/com/google/idea/blaze/base/qsync/QuerySyncHighlightingSettingProvider.java
index 57d76863868..dd5aba6a7e9 100644
--- a/base/src/com/google/idea/blaze/base/qsync/QuerySyncHighlightingSettingProvider.java
+++ b/base/src/com/google/idea/blaze/base/qsync/QuerySyncHighlightingSettingProvider.java
@@ -4,6 +4,7 @@
 import com.google.idea.blaze.base.settings.BlazeImportSettings;
 import com.intellij.codeInsight.daemon.impl.analysis.DefaultHighlightingSettingProvider;
 import com.intellij.codeInsight.daemon.impl.analysis.FileHighlightingSetting;
+import com.intellij.codeInsight.daemon.impl.analysis.HighlightLevelUtil;
 import com.intellij.openapi.project.Project;
 import com.intellij.openapi.vfs.VirtualFile;
 import com.intellij.psi.PsiManager;
@@ -13,10 +14,13 @@
 public class QuerySyncHighlightingSettingProvider extends DefaultHighlightingSettingProvider {
     @Override
     public @Nullable FileHighlightingSetting getDefaultSetting(@NotNull Project project, @NotNull VirtualFile file) {
+
         var psiFile = PsiManager.getInstance(project).findFile(file);
         if (Blaze.getProjectType(psiFile.getProject()) == BlazeImportSettings.ProjectType.QUERY_SYNC) {
             if(!QuerySyncManager.getInstance(psiFile.getProject()).isReadyForAnalysis(psiFile)){
                 return FileHighlightingSetting.ESSENTIAL;
+            } else {
+                return FileHighlightingSetting.FORCE_HIGHLIGHTING;
             }
         }
         return null;