hls-notes-plugin: Allow to see where a note is referenced from (#4624)

Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>

Author: jvanbruegge

plugins/hls-notes-plugin/src/Ide/Plugin/Notes.hs

@@ -1,17 +1,21 @@
 module Ide.Plugin.Notes (descriptor, Log) where
 
 import           Control.Lens                     ((^.))
-import           Control.Monad.Except             (throwError)
+import           Control.Monad.Except             (ExceptT, MonadError,
+                                                   throwError)
 import           Control.Monad.IO.Class           (liftIO)
 import qualified Data.Array                       as A
+import           Data.Foldable                    (foldl')
 import           Data.HashMap.Strict              (HashMap)
 import qualified Data.HashMap.Strict              as HM
 import qualified Data.HashSet                     as HS
+import           Data.List                        (uncons)
 import           Data.Maybe                       (catMaybes, listToMaybe,
                                                    mapMaybe)
 import           Data.Text                        (Text, intercalate)
 import qualified Data.Text                        as T
 import qualified Data.Text.Utf16.Rope.Mixed       as Rope
+import           Data.Traversable                 (for)
 import           Development.IDE                  hiding (line)
 import           Development.IDE.Core.PluginUtils (runActionE, useE)
 import           Development.IDE.Core.Shake       (toKnownFiles)
@@ -21,8 +25,8 @@ import           GHC.Generics                     (Generic)
 import           Ide.Plugin.Error                 (PluginError (..))
 import           Ide.Types
 import qualified Language.LSP.Protocol.Lens       as L
-import           Language.LSP.Protocol.Message    (Method (Method_TextDocumentDefinition),
-                                                   SMethod (SMethod_TextDocumentDefinition))
+import           Language.LSP.Protocol.Message    (Method (Method_TextDocumentDefinition, Method_TextDocumentReferences),
+                                                   SMethod (SMethod_TextDocumentDefinition, SMethod_TextDocumentReferences))
 import           Language.LSP.Protocol.Types
 import           Text.Regex.TDFA                  (Regex, caseSensitive,
                                                    defaultCompOpt,
@@ -31,25 +35,39 @@ import           Text.Regex.TDFA                  (Regex, caseSensitive,
 
 data Log
     = LogShake Shake.Log
-    | LogNotesFound NormalizedFilePath [(Text, Position)]
+    | LogNotesFound NormalizedFilePath [(Text, [Position])]
+    | LogNoteReferencesFound NormalizedFilePath [(Text, [Position])]
     deriving Show
 
 data GetNotesInFile = MkGetNotesInFile
     deriving (Show, Generic, Eq, Ord)
     deriving anyclass (Hashable, NFData)
-type instance RuleResult GetNotesInFile = HM.HashMap Text Position
+-- The GetNotesInFile action scans the source file and extracts a map of note
+-- definitions (note name -> position) and a map of note references
+-- (note name -> [position]).
+type instance RuleResult GetNotesInFile = (HM.HashMap Text Position, HM.HashMap Text [Position])
 
 data GetNotes = MkGetNotes
     deriving (Show, Generic, Eq, Ord)
     deriving anyclass (Hashable, NFData)
+-- GetNotes collects all note definition across all files in the
+-- project. It returns a map from note name to pair of (filepath, position).
 type instance RuleResult GetNotes = HashMap Text (NormalizedFilePath, Position)
 
+data GetNoteReferences = MkGetNoteReferences
+    deriving (Show, Generic, Eq, Ord)
+    deriving anyclass (Hashable, NFData)
+-- GetNoteReferences collects all note references across all files in the
+-- project. It returns a map from note name to list of (filepath, position).
+type instance RuleResult GetNoteReferences = HashMap Text [(NormalizedFilePath, Position)]
+
 instance Pretty Log where
     pretty = \case
-        LogShake l -> pretty l
-        LogNotesFound file notes ->
-            "Found notes in " <> pretty (show file) <> ": ["
-            <> pretty (intercalate ", " (fmap (\(s, p) -> "\"" <> s <> "\" at " <> T.pack (show p)) notes)) <> "]"
+            LogShake l -> pretty l
+            LogNoteReferencesFound file refs -> "Found note references in " <> prettyNotes file refs
+            LogNotesFound file notes -> "Found notes in " <> prettyNotes file notes
+        where prettyNotes file hm = pretty (show file) <> ": ["
+                <> pretty (intercalate ", " (fmap (\(s, p) -> "\"" <> s <> "\" at " <> intercalate ", " (map (T.pack . show) p)) hm)) <> "]"
 
 {-
 The first time the user requests a jump-to-definition on a note reference, the
@@ -59,7 +77,9 @@ title is then saved in the HLS database to be retrieved for all future requests.
 descriptor :: Recorder (WithPriority Log) -> PluginId -> PluginDescriptor IdeState
 descriptor recorder plId = (defaultPluginDescriptor plId "Provides goto definition support for GHC-style notes")
     { Ide.Types.pluginRules = findNotesRules recorder
-    , Ide.Types.pluginHandlers = mkPluginHandler SMethod_TextDocumentDefinition jumpToNote
+    , Ide.Types.pluginHandlers =
+        mkPluginHandler SMethod_TextDocumentDefinition jumpToNote
+        <> mkPluginHandler SMethod_TextDocumentReferences listReferences
     }
 
 findNotesRules :: Recorder (WithPriority Log) -> Rules ()
@@ -69,20 +89,59 @@ findNotesRules recorder = do
 
     defineNoDiagnostics (cmapWithPrio LogShake recorder) $ \MkGetNotes _ -> do
         targets <- toKnownFiles <$> useNoFile_ GetKnownTargets
-        definedNotes <- catMaybes <$> mapM (\nfp -> fmap (HM.map (nfp,)) <$> use MkGetNotesInFile nfp) (HS.toList targets)
+        definedNotes <- catMaybes <$> mapM (\nfp -> fmap (HM.map (nfp,) . fst) <$> use MkGetNotesInFile nfp) (HS.toList targets)
         pure $ Just $ HM.unions definedNotes
 
+    defineNoDiagnostics (cmapWithPrio LogShake recorder) $ \MkGetNoteReferences _ -> do
+        targets <- toKnownFiles <$> useNoFile_ GetKnownTargets
+        definedReferences <- catMaybes <$> for (HS.toList targets) (\nfp -> do
+                references <- fmap snd <$> use MkGetNotesInFile nfp
+                pure $ fmap (HM.map (fmap (nfp,))) references
+            )
+        pure $ Just $ foldl' (HM.unionWith (<>)) HM.empty definedReferences
+
+err :: MonadError PluginError m => Text -> Maybe a -> m a
+err s = maybe (throwError $ PluginInternalError s) pure
+
+getNote :: NormalizedFilePath -> IdeState -> Position -> ExceptT PluginError (HandlerM c) (Maybe Text)
+getNote nfp state (Position l c) = do
+    contents <-
+        err "Error getting file contents"
+        =<< liftIO (runAction "notes.getfileContents" state (getFileContents nfp))
+    line <- err "Line not found in file" (listToMaybe $ Rope.lines $ fst
+        (Rope.splitAtLine 1 $ snd $ Rope.splitAtLine (fromIntegral l) contents))
+    pure $ listToMaybe $ mapMaybe (atPos $ fromIntegral c) $ matchAllText noteRefRegex line
+  where
+    atPos c arr = case arr A.! 0 of
+        -- We check if the line we are currently at contains a note
+        -- reference. However, we need to know if the cursor is within the
+        -- match or somewhere else. The second entry of the array contains
+        -- the title of the note as extracted by the regex.
+        (_, (c', len)) -> if c' <= c && c <= c' + len
+            then Just (fst (arr A.! 1)) else Nothing
+
+listReferences :: PluginMethodHandler IdeState Method_TextDocumentReferences
+listReferences state _ param
+    | Just nfp <- uriToNormalizedFilePath uriOrig
+    = do
+        let pos@(Position l _) = param ^. L.position
+        noteOpt <- getNote nfp state pos
+        case noteOpt of
+            Nothing -> pure (InR Null)
+            Just note -> do
+                notes <- runActionE "notes.definedNoteReferencess" state $ useE MkGetNoteReferences nfp
+                poss <- err ("Note reference (a comment of the form `{- Note [" <> note <> "] -}`) not found") (HM.lookup note notes)
+                pure $ InL (mapMaybe (\(noteFp, pos@(Position l' _)) -> if l' == l then Nothing else Just (
+                        Location (fromNormalizedUri $ normalizedFilePathToUri noteFp) (Range pos pos))) poss)
+    where
+        uriOrig = toNormalizedUri $ param ^. (L.textDocument . L.uri)
+listReferences _ _ _ = throwError $ PluginInternalError "conversion to normalized file path failed"
+
 jumpToNote :: PluginMethodHandler IdeState Method_TextDocumentDefinition
 jumpToNote state _ param
     | Just nfp <- uriToNormalizedFilePath uriOrig
     = do
-        let Position l c = param ^. L.position
-        contents <-
-            err "Error getting file contents"
-            =<< liftIO (runAction "notes.getfileContents" state (getFileContents nfp))
-        line <- err "Line not found in file" (listToMaybe $ Rope.lines $ fst
-            (Rope.splitAtLine 1 $ snd $ Rope.splitAtLine (fromIntegral l) contents))
-        let noteOpt = listToMaybe $ mapMaybe (atPos $ fromIntegral c) $ matchAllText noteRefRegex line
+        noteOpt <- getNote nfp state (param ^. L.position)
         case noteOpt of
             Nothing -> pure (InR (InR Null))
             Just note -> do
@@ -93,28 +152,23 @@ jumpToNote state _ param
                     ))
     where
         uriOrig = toNormalizedUri $ param ^. (L.textDocument . L.uri)
-        err s = maybe (throwError $ PluginInternalError s) pure
-        atPos c arr = case arr A.! 0 of
-            -- We check if the line we are currently at contains a note
-            -- reference. However, we need to know if the cursor is within the
-            -- match or somewhere else. The second entry of the array contains
-            -- the title of the note as extracted by the regex.
-            (_, (c', len)) -> if c' <= c && c <= c' + len
-                then Just (fst (arr A.! 1)) else Nothing
 jumpToNote _ _ _ = throwError $ PluginInternalError "conversion to normalized file path failed"
 
-findNotesInFile :: NormalizedFilePath -> Recorder (WithPriority Log) -> Action (Maybe (HM.HashMap Text Position))
+findNotesInFile :: NormalizedFilePath -> Recorder (WithPriority Log) -> Action (Maybe (HM.HashMap Text Position, HM.HashMap Text [Position]))
 findNotesInFile file recorder = do
     -- GetFileContents only returns a value if the file is open in the editor of
     -- the user. If not, we need to read it from disk.
     contentOpt <- (snd =<<) <$> use GetFileContents file
     content <- case contentOpt of
         Just x  -> pure $ Rope.toText x
         Nothing -> liftIO $ readFileUtf8 $ fromNormalizedFilePath file
-    let matches = (A.! 1) <$> matchAllText noteRegex content
-        m = toPositions matches content
-    logWith recorder Debug $ LogNotesFound file (HM.toList m)
-    pure $ Just m
+    let noteMatches = (A.! 1) <$> matchAllText noteRegex content
+        notes = toPositions noteMatches content
+    logWith recorder Debug $ LogNotesFound file (HM.toList notes)
+    let refMatches = (A.! 1) <$> matchAllText noteRefRegex content
+        refs = toPositions refMatches content
+    logWith recorder Debug $ LogNoteReferencesFound file (HM.toList refs)
+    pure $ Just (HM.mapMaybe (fmap fst . uncons) notes, refs)
     where
         uint = fromIntegral . toInteger
         -- the regex library returns the character index of the match. However
@@ -129,7 +183,7 @@ findNotesInFile file recorder = do
                 let !c' = c + 1
                     (!n', !nc') = if char' == '\n' then (n + 1, c') else (n, nc)
                     p@(!_, !_) = if char == c then
-                            (xs, HM.insert name (Position (uint n') (uint (char - nc'))) m)
+                            (xs, HM.insertWith (<>) name [Position (uint n') (uint (char - nc'))] m)
                         else (x:xs, m)
                 in (p, (n', nc', c'))
             ) ((matches, HM.empty), (0, 0, 0))

plugins/hls-notes-plugin/test/NotesTest.hs

@@ -11,6 +11,7 @@ main :: IO ()
 main = defaultTestRunner $
   testGroup "Notes"
     [ gotoNoteTests
+    , noteReferenceTests
     ]
 
 runSessionWithServer' :: FilePath -> (FilePath -> Session a) -> IO a
@@ -21,6 +22,21 @@ runSessionWithServer' fp act =
         , testDirLocation = Left fp
         } act
 
+noteReferenceTests :: TestTree
+noteReferenceTests = testGroup "Note References"
+   [
+      testCase "multi_file" $ runSessionWithServer' testDataDir $ \dir -> do
+        doc <- openDoc "NoteDef.hs" "haskell"
+        waitForKickDone
+        refs <- getReferences doc (Position 21 15) False
+        let fp = dir </> "NoteDef.hs"
+        liftIO $ refs @?= [
+            Location (filePathToUri (dir </> "Other.hs")) (Range (Position 6 13) (Position 6 13)),
+            Location (filePathToUri fp) (Range (Position 9 9) (Position 9 9)),
+            Location (filePathToUri fp) (Range (Position 5 67) (Position 5 67))
+          ]
+    ]
+
 gotoNoteTests :: TestTree
 gotoNoteTests = testGroup "Goto Note Definition"
     [
@@ -29,13 +45,13 @@ gotoNoteTests = testGroup "Goto Note Definition"
         waitForKickDone
         defs <- getDefinitions doc (Position 3 41)
         let fp = dir </> "NoteDef.hs"
-        liftIO $ defs @?= InL (Definition (InR [Location (filePathToUri fp) (Range (Position 8 9) (Position 8 9))]))
+        liftIO $ defs @?= InL (Definition (InR [Location (filePathToUri fp) (Range (Position 11 9) (Position 11 9))]))
     , testCase "liberal_format" $ runSessionWithServer' testDataDir $ \dir -> do
         doc <- openDoc "NoteDef.hs" "haskell"
         waitForKickDone
         defs <- getDefinitions doc (Position 5 64)
         let fp = dir </> "NoteDef.hs"
-        liftIO $ defs @?= InL (Definition (InR [Location (filePathToUri fp) (Range (Position 18 11) (Position 18 11))]))
+        liftIO $ defs @?= InL (Definition (InR [Location (filePathToUri fp) (Range (Position 21 11) (Position 21 11))]))
 
     , testCase "invalid_note" $ runSessionWithServer' testDataDir $ const $ do
         doc <- openDoc "NoteDef.hs" "haskell"
@@ -54,7 +70,7 @@ gotoNoteTests = testGroup "Goto Note Definition"
         waitForKickDone
         defs <- getDefinitions doc (Position 5 20)
         let fp = dir </> "NoteDef.hs"
-        liftIO $ defs @?= InL (Definition (InR [Location (filePathToUri fp) (Range (Position 12 6) (Position 12 6))]))
+        liftIO $ defs @?= InL (Definition (InR [Location (filePathToUri fp) (Range (Position 15 6) (Position 15 6))]))
     ]
 
 testDataDir :: FilePath

plugins/hls-notes-plugin/test/testdata/NoteDef.hs

@@ -6,6 +6,9 @@ foo _ = 0 -- We always return zero, see Note [Returning zero from foo]
 -- The plugin is more liberal with the note definitions, see Note [Single line comments]
 -- It does not work on wrong note definitions, see Note [Not a valid Note]
 
+-- We can also have multiple references to the same note, see
+-- Note [Single line comments]
+
 {- Note [Returning zero from foo]
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 This is a big long form note, with very important info

plugins/hls-notes-plugin/test/testdata/Other.hs

@@ -4,3 +4,4 @@ import NoteDef
 
 bar :: Int
 bar = 4 -- See @Note [Multiple notes in comment]@ in NoteDef
+-- See Note [Single line comments]