Fixed Trial and Error Removal sometimes not reacting to finished comp…

…ilation Fixes #41
Wumpf · Aug 22, 2017 · 3944cca · 3944cca
1 parent b8ea6ab
commit 3944cca
Showing 1 changed file with 21 additions and 5 deletions.
diff --git a/IncludeToolbox/TrialAndErrorRemoval.cs b/IncludeToolbox/TrialAndErrorRemoval.cs
@@ -25,6 +25,11 @@ internal sealed class TrialAndErrorRemoval
         private AutoResetEvent outputWaitEvent = new AutoResetEvent(false);
         private const int timeoutMS = 30000; // 30 seconds
 
+        /// <summary>
+        /// Need to keep build events object around as long as it is used, otherwise the events may not be fired!
+        /// </summary>
+        private BuildEvents buildEvents;
+
         public void PerformTrialAndErrorIncludeRemoval(EnvDTE.Document document, TrialAndErrorRemovalOptionsPage settings)
         {
             if (document == null)
@@ -64,8 +69,7 @@ public void PerformTrialAndErrorIncludeRemoval(EnvDTE.Document document, TrialAn
             }
 
             // Hook into build events.
-            document.DTE.Events.BuildEvents.OnBuildProjConfigDone += OnBuildConfigFinished;
-            document.DTE.Events.BuildEvents.OnBuildDone += OnBuildFinished;
+            SubscribeBuildEvents();
 
             // The rest runs in a separate thread since the compile function is non blocking and we want to use BuildEvents
             // We are not using Task, since we want to make use of WaitHandles - using this together with Task is a bit more complicated to get right.
@@ -115,7 +119,6 @@ private void ExtractSelectionAndIncludes(EnvDTE.Document document, TrialAndError
             string documentText = documentTextView.TextView.TextSnapshot.GetText();
             IEnumerable<Formatter.IncludeLineInfo> includeLines = Formatter.IncludeLineInfo.ParseIncludes(documentText, Formatter.ParseOptions.IgnoreIncludesInPreprocessorConditionals | Formatter.ParseOptions.KeepOnlyValidIncludes);
 
-
             // Optionally skip top most include.
             if (settings.IgnoreFirstInclude)
                 includeLines = includeLines.Skip(1);
@@ -249,8 +252,7 @@ private void OnTrialAndErrorRemovalDone(IVsThreadedWaitDialog2 progressDialog, E
             progressDialog.EndWaitDialog();
 
             // Remove build hook again.
-            document.DTE.Events.BuildEvents.OnBuildDone -= OnBuildFinished;
-            document.DTE.Events.BuildEvents.OnBuildProjConfigDone -= OnBuildConfigFinished;
+            UnsubscribeBuildEvents();
 
             // Message.
             Output.Instance.WriteLine("Removed {0} #include directives from '{1}'", numRemovedIncludes, document.Name);
@@ -261,6 +263,20 @@ private void OnTrialAndErrorRemovalDone(IVsThreadedWaitDialog2 progressDialog, E
             OnFileFinished?.Invoke(numRemovedIncludes, canceled);
         }
 
+        private void SubscribeBuildEvents()
+        {
+            buildEvents = VSUtils.GetDTE().Events.BuildEvents;
+            buildEvents.OnBuildDone += OnBuildFinished;
+            buildEvents.OnBuildProjConfigDone += OnBuildConfigFinished;
+        }
+
+        private void UnsubscribeBuildEvents()
+        {
+            buildEvents.OnBuildDone -= OnBuildFinished;
+            buildEvents.OnBuildProjConfigDone -= OnBuildConfigFinished;
+            buildEvents = null;
+        }
+
         private void OnBuildFinished(vsBuildScope scope, vsBuildAction action)
         {
             //Output.Instance.WriteLine("OnBuildFinished. scope: {0}, action: {1}", scope, action);