Fix diagnostics bug

Given a FOI F with non null typechecking diagnostics D, imagine the following scenario:

1. An edit notification for F is received, creating a new version
2. GetModTime is executed, producing 0 diagnostics.
  2.1 updateFileDiagnostics is called
  2.2 setStageDiagnostics is called
  2.3 LSP.updateDiagnostics is called with a new version, resetting all the
  diagnostics for F
  2.4 newDiags=[] in updateFileDiagnostics, which is different from D (the last
  published diagnostics), which enqueues a new publishDiagnostics [] in the
  Debouncer
3. An edit notification for F is received before typechecking has a chance to
run which undoes the previous edit
4. The debouncer publishes the empty set of diagnostics after waiting 0.1s
5. GetFileContents runs and since the contents of the file haven't changed since
the last time it ran, early cutoff skips everything donwstream

Since TypeCheck is skipped, the empty set of diagnostics stays published until
another edit comes.

The goal of this change is to prevent setStageDiagnostics from losing
diagnostics from other stages. To achieve this, we recover the old diagnostics
for all stages and merge them with the new stage.

src/Development/IDE/Core/Shake.hs

@@ -993,25 +993,19 @@ updateFileDiagnostics :: MonadIO m
 updateFileDiagnostics fp k ShakeExtras{diagnostics, hiddenDiagnostics, publishedDiagnostics, state, debouncer, eventer} current = liftIO $ do
     modTime <- (currentValue =<<) <$> getValues state GetModificationTime fp
     let (currentShown, currentHidden) = partition ((== ShowDiag) . fst) current
+        uri = filePathToUri' fp
+        ver = vfsVersion =<< modTime
+        updateDiagnosticsWithForcing new store = do
+            store' <- evaluate $ setStageDiagnostics uri ver (T.pack $ show k) new store
+            new' <- evaluate $ getUriDiagnostics uri store'
+            return (store', new')
     mask_ $ do
         -- Mask async exceptions to ensure that updated diagnostics are always
         -- published. Otherwise, we might never publish certain diagnostics if
         -- an exception strikes between modifyVar but before
         -- publishDiagnosticsNotification.
-        newDiags <- modifyVar diagnostics $ \old -> do
-            let newDiagsStore = setStageDiagnostics fp (vfsVersion =<< modTime)
-                                  (T.pack $ show k) (map snd currentShown) old
-            let newDiags = getFileDiagnostics fp newDiagsStore
-            _ <- evaluate newDiagsStore
-            _ <- evaluate newDiags
-            pure (newDiagsStore, newDiags)
-        modifyVar_ hiddenDiagnostics $ \old -> do
-            let newDiagsStore = setStageDiagnostics fp (vfsVersion =<< modTime)
-                                  (T.pack $ show k) (map snd currentHidden) old
-            let newDiags = getFileDiagnostics fp newDiagsStore
-            _ <- evaluate newDiagsStore
-            _ <- evaluate newDiags
-            return newDiagsStore
+        newDiags <- modifyVar diagnostics $ updateDiagnosticsWithForcing $ map snd currentShown
+        _ <- modifyVar hiddenDiagnostics $  updateDiagnosticsWithForcing $ map snd currentHidden
         let uri = filePathToUri' fp
         let delay = if null newDiags then 0.1 else 0
         registerEvent debouncer delay uri $ do
@@ -1053,31 +1047,38 @@ getDiagnosticsFromStore (StoreItem _ diags) = concatMap SL.fromSortedList $ Map.
 -- | Sets the diagnostics for a file and compilation step
 --   if you want to clear the diagnostics call this with an empty list
 setStageDiagnostics
-    :: NormalizedFilePath
+    :: NormalizedUri
     -> TextDocumentVersion -- ^ the time that the file these diagnostics originate from was last edited
     -> T.Text
     -> [LSP.Diagnostic]
     -> DiagnosticStore
     -> DiagnosticStore
-setStageDiagnostics fp timeM stage diags ds  =
-    updateDiagnostics ds uri timeM diagsBySource
-    where
-        diagsBySource = Map.singleton (Just stage) (SL.toSortedList diags)
-        uri = filePathToUri' fp
+setStageDiagnostics uri ver stage diags ds = newDiagsStore where
+    -- When 'ver' is a new version, updateDiagnostics throws away diagnostics from all stages
+    -- This interacts bady with early cutoff, so we make sure to preserve diagnostics
+    -- from other stages when calling updateDiagnostics
+    -- But this means that updateDiagnostics cannot be called concurrently
+    -- for different stages anymore
+    updatedDiags = Map.singleton (Just stage) (SL.toSortedList diags) <> oldDiags
+    oldDiags = case HMap.lookup uri ds of
+            Just (StoreItem _ byStage) -> byStage
+            _ -> Map.empty
+    newDiagsStore = updateDiagnostics ds uri ver updatedDiags
+
 
 getAllDiagnostics ::
     DiagnosticStore ->
     [FileDiagnostic]
 getAllDiagnostics =
     concatMap (\(k,v) -> map (fromUri k,ShowDiag,) $ getDiagnosticsFromStore v) . HMap.toList
 
-getFileDiagnostics ::
-    NormalizedFilePath ->
+getUriDiagnostics ::
+    NormalizedUri ->
     DiagnosticStore ->
     [LSP.Diagnostic]
-getFileDiagnostics fp ds =
+getUriDiagnostics uri ds =
     maybe [] getDiagnosticsFromStore $
-    HMap.lookup (filePathToUri' fp) ds
+    HMap.lookup uri ds
 
 filterDiagnostics ::
     (NormalizedFilePath -> Bool) ->