[feat] allow flushing diagnostics by source and uri

Author: MangoIV

lsp/src/Language/LSP/Diagnostics.hs

@@ -11,6 +11,7 @@ module Language.LSP.Diagnostics (
   StoreItem (..),
   partitionBySource,
   flushBySource,
+  flushBySourceAndUri,
   updateDiagnostics,
   getDiagnosticParamsFor,
 
@@ -41,8 +42,10 @@ all prior entries for the Uri.
 
 type DiagnosticStore = HM.HashMap J.NormalizedUri StoreItem
 
-data StoreItem
-  = StoreItem (Maybe J.Int32) DiagnosticsBySource
+data StoreItem = StoreItem
+  { documentVersion :: Maybe J.Int32
+  , diagnostics :: DiagnosticsBySource
+  }
   deriving (Show, Eq)
 
 type DiagnosticsBySource = Map.Map (Maybe Text) (SL.SortedList J.Diagnostic)
@@ -60,6 +63,13 @@ flushBySource store (Just source) = HM.map remove store
  where
   remove (StoreItem mv diags) = StoreItem mv (Map.delete (Just source) diags)
 
+flushBySourceAndUri :: DiagnosticStore -> Maybe Text -> J.NormalizedUri -> DiagnosticStore
+flushBySourceAndUri store msource uri = HM.mapWithKey remove store
+ where
+  remove k item
+    | k == uri = item{diagnostics = Map.delete msource $ diagnostics item}
+    | otherwise = item
+
 -- ---------------------------------------------------------------------
 
 updateDiagnostics ::

lsp/src/Language/LSP/Server/Core.hs

@@ -655,6 +655,29 @@ flushDiagnosticsBySource maxDiagnosticCount msource = join $ stateState resDiagn
 
 -- ---------------------------------------------------------------------
 
+{- | Remove all diagnostics from a particular uri and source, and send the updates to
+ the client.
+-}
+flushDiagnosticsBySourceAndUri ::
+  MonadLsp config m =>
+  -- | Max number of diagnostics to send
+  Int ->
+  Maybe Text ->
+  NormalizedUri ->
+  m ()
+flushDiagnosticsBySourceAndUri maxDiagnosticCount msource uri = join $ stateState resDiagnostics $ \oldDiags ->
+  let !newDiags = flushBySourceAndUri oldDiags msource uri
+      -- Send the updated diagnostics to the client
+      act = forM_ (HM.keys newDiags) $ \uri' -> do
+        let mdp = getDiagnosticParamsFor maxDiagnosticCount newDiags uri'
+        case mdp of
+          Nothing -> return ()
+          Just params -> do
+            sendToClient $ L.fromServerNot $ L.TNotificationMessage "2.0" L.SMethod_TextDocumentPublishDiagnostics params
+   in (act, newDiags)
+
+-- ---------------------------------------------------------------------
+
 {- | The changes in a workspace edit should be applied from the end of the file
  toward the start. Sort them into this order.
 -}