eval: more robust way to extract comments from ParsedModule

plugins/hls-eval-plugin/src/Ide/Plugin/Eval/Rules.hs

@@ -1,13 +1,13 @@
 {-# LANGUAGE CPP             #-}
 {-# LANGUAGE LambdaCase      #-}
 {-# LANGUAGE PatternSynonyms #-}
-{-# LANGUAGE RecordWildCards #-}
 
--- To avoid warning "Pattern match has inaccessible right hand side"
-{-# OPTIONS_GHC -Wno-overlapping-patterns #-}
 module Ide.Plugin.Eval.Rules (GetEvalComments(..), rules,queueForEvaluation, unqueueForEvaluation, Log) where
 
+import           Control.Lens                         (toListOf)
 import           Control.Monad.IO.Class               (MonadIO (liftIO))
+import qualified Data.ByteString                      as BS
+import           Data.Data.Lens                       (biplate)
 import           Data.HashSet                         (HashSet)
 import qualified Data.HashSet                         as Set
 import           Data.IORef
@@ -24,8 +24,7 @@ import           Development.IDE                      (GetModSummaryWithoutTimes
                                                        fromNormalizedFilePath,
                                                        msrModSummary,
                                                        realSrcSpanToRange,
-                                                       useWithStale_,
-                                                       use_)
+                                                       useWithStale_, use_)
 import           Development.IDE.Core.PositionMapping (toCurrentRange)
 import           Development.IDE.Core.Rules           (computeLinkableTypeForDynFlags,
                                                        needsCompilationRule)
@@ -39,14 +38,12 @@ import           Development.IDE.GHC.Compat
 import qualified Development.IDE.GHC.Compat           as SrcLoc
 import qualified Development.IDE.GHC.Compat.Util      as FastString
 import           Development.IDE.Graph                (alwaysRerun)
-import           Ide.Logger         (Pretty (pretty),
+import           GHC.Parser.Annotation
+import           Ide.Logger                           (Pretty (pretty),
                                                        Recorder, WithPriority,
                                                        cmapWithPrio)
-import           GHC.Parser.Annotation
 import           Ide.Plugin.Eval.Types
 
-import qualified Data.ByteString                      as BS
-
 newtype Log = LogShake Shake.Log deriving Show
 
 instance Pretty Log where
@@ -74,28 +71,17 @@ unqueueForEvaluation ide nfp = do
     -- remove the module from the Evaluating state, so that next time it won't evaluate to True
     atomicModifyIORef' var $ \fs -> (Set.delete nfp fs, ())
 
-#if MIN_VERSION_ghc(9,5,0)
-getAnnotations :: Development.IDE.GHC.Compat.Located (HsModule GhcPs) -> [LEpaComment]
-getAnnotations (L _ m@(HsModule { hsmodExt = XModulePs {hsmodAnn = anns'}})) =
-#else
-getAnnotations :: Development.IDE.GHC.Compat.Located HsModule -> [LEpaComment]
-getAnnotations (L _ m@(HsModule { hsmodAnn = anns'})) =
-#endif
-    priorComments annComments <> getFollowingComments annComments
-     <> concatMap getCommentsForDecl (hsmodImports m)
-     <> concatMap getCommentsForDecl (hsmodDecls m)
-       where
-         annComments = epAnnComments anns'
-
-getCommentsForDecl :: GenLocated (SrcSpanAnn' (EpAnn ann)) e
-                            -> [LEpaComment]
-getCommentsForDecl (L (SrcSpanAnn (EpAnn _ _ cs) _) _) = priorComments cs <> getFollowingComments cs
-getCommentsForDecl (L (SrcSpanAnn (EpAnnNotUsed) _) _) = []
-
 apiAnnComments' :: ParsedModule -> [SrcLoc.RealLocated EpaCommentTok]
 apiAnnComments' pm = do
-  L span (EpaComment c _) <- getAnnotations $ pm_parsed_source pm
+  L span (EpaComment c _) <- getEpaComments $ pm_parsed_source pm
   pure (L (anchor span) c)
+  where
+#if MIN_VERSION_ghc(9,5,0)
+    getEpaComments :: Development.IDE.GHC.Compat.Located (HsModule GhcPs) -> [LEpaComment]
+#else
+    getEpaComments :: Development.IDE.GHC.Compat.Located HsModule -> [LEpaComment]
+#endif
+    getEpaComments = toListOf biplate
 
 pattern RealSrcSpanAlready :: SrcLoc.RealSrcSpan -> SrcLoc.RealSrcSpan
 pattern RealSrcSpanAlready x = x

plugins/hls-eval-plugin/test/Main.hs

@@ -112,6 +112,7 @@ tests =
   , goldenWithEval ":kind treats a multilined result properly" "T25" "hs"
   , goldenWithEvalAndFs "local imports" (FS.directProjectMulti ["T26.hs", "Util.hs"]) "T26" "hs"
   , goldenWithEval "Preserves one empty comment line after prompt" "T27" "hs"
+  , goldenWithEval "Evaluate comment after multiline function definition" "T28" "hs"
   , goldenWithEval "Multi line comments" "TMulti" "hs"
   , goldenWithEval "Multi line comments, with the last test line ends without newline" "TEndingMulti" "hs"
   , goldenWithEval "Evaluate expressions in Plain comments in both single line and multi line format" "TPlainComment" "hs"

plugins/hls-eval-plugin/test/testdata/T28.expected.hs

@@ -0,0 +1,7 @@
+module T28 where
+
+f True = True
+f False = False
+
+-- >>> 1+1
+-- 2

plugins/hls-eval-plugin/test/testdata/T28.hs

@@ -0,0 +1,6 @@
+module T28 where
+
+f True = True
+f False = False
+
+-- >>> 1+1