Preserve comments on language pragmas and similar

CHANGELOG.md

@@ -3,6 +3,9 @@
 * Grouping of statements in `do`-blocks is now respected. [Issue
   74](https://github.com/tweag/ormolu/issues/74).
 
+* Comments on pragmas are now preserved. [Issue
+  216](https://github.com/tweag/ormolu/issues/216).
+
 ## Ormolu 0.0.4.0
 
 * When given several files to format, Ormolu does not stop on the first

data/examples/other/pragma-comments-out.hs

@@ -0,0 +1,11 @@
+{-# LANGUAGE OverloadedStrings #-}
+-- TODO This extension is probably too dangerous, remove it.
+{-# LANGUAGE RecordWildCards #-}
+-- Avoid warning produced by TH.
+{-# OPTIONS_GHC -fno-warn-overlapping-patterns #-}
+
+-- | Header comment.
+module Foo
+  (
+  )
+where

data/examples/other/pragma-comments.hs

@@ -0,0 +1,11 @@
+-- | Header comment.
+
+{-# LANGUAGE OverloadedStrings #-}
+
+-- Avoid warning produced by TH.
+{-# OPTIONS_GHC -fno-warn-overlapping-patterns #-}
+
+-- TODO This extension is probably too dangerous, remove it.
+{-# LANGUAGE RecordWildCards #-}
+
+module Foo () where

data/examples/other/pragma-sorting-out.hs

@@ -1,10 +1,9 @@
 {-# LANGUAGE NondecreasingIndentation #-}
 {-# LANGUAGE ScopedTypeVariables #-}
 {-# LANGUAGE TypeFamilies #-}
-{-# LANGUAGE NoMonoLocalBinds #-}
-
 -- This gap is necessary for stylish Haskell not to re-arrange
 -- NoMonoLocalBinds before TypeFamilies
+{-# LANGUAGE NoMonoLocalBinds #-}
 
 module Foo
   ( bar,

expected-failures/Agda.txt

@@ -11,7 +11,7 @@ Formatting is not idempotent:
 Please, consider reporting the bug.
 
 Parsing of formatted code failed:
-  src/full/Agda/Syntax/Translation/ConcreteToAbstract.hs<rendered>:525:7-13
+  src/full/Agda/Syntax/Translation/ConcreteToAbstract.hs<rendered>:524:7-13
   parse error on input `C.QName'
 Please, consider reporting the bug.
 

expected-failures/idris.txt

@@ -83,7 +83,7 @@ Formatting is not idempotent:
 Please, consider reporting the bug.
 
 Formatting is not idempotent:
-  src/Idris/REPL.hs<rendered>:1271:33
+  src/Idris/REPL.hs<rendered>:1270:33
   before: "ht c) _) = -- consta"
   after:  "ht c) _) =\n  -- cons"
 Please, consider reporting the bug.

ormolu.cabal

@@ -82,6 +82,7 @@ library
                     , Ormolu.Parser.CommentStream
                     , Ormolu.Parser.Pragma
                     , Ormolu.Parser.Result
+                    , Ormolu.Parser.Shebang
                     , Ormolu.Printer
                     , Ormolu.Printer.Combinators
                     , Ormolu.Printer.Comments

src/Ormolu/Parser.hs

@@ -32,6 +32,7 @@ import Ormolu.Exception
 import Ormolu.Parser.Anns
 import Ormolu.Parser.CommentStream
 import Ormolu.Parser.Result
+import Ormolu.Parser.Shebang
 import qualified Panic as GHC
 import qualified Parser as GHC
 import qualified StringBuffer as GHC
@@ -93,14 +94,16 @@ parseModule Config {..} path input' = liftIO $ do
             -- later stages; but we fail in those cases.
             Just err -> Left err
             Nothing ->
-              let (comments, exts, shebangs) = mkCommentStream extraComments pstate
+              let (stackHeader, shebangs, pragmas, comments) =
+                    mkCommentStream extraComments pstate
                in Right
                     ParseResult
                       { prParsedSource = pmod,
                         prAnns = mkAnns pstate,
-                        prCommentStream = comments,
-                        prExtensions = exts,
+                        prStackHeader = stackHeader,
                         prShebangs = shebangs,
+                        prPragmas = pragmas,
+                        prCommentStream = comments,
                         prUseRecordDot = useRecordDot,
                         prImportQualifiedPost =
                           GHC.xopt ImportQualifiedPost dynFlags
@@ -154,14 +157,14 @@ runParser parser flags filename input = GHC.unP parser parseState
     buffer = GHC.stringToStringBuffer input
     parseState = GHC.mkPState flags buffer location
 
--- | Transform given lines possibly returning comments extracted from them.
+-- | Transform given input possibly returning comments extracted from it.
 -- This handles LINE pragmas and shebangs.
 extractCommentsFromLines ::
   -- | File name, just to use in the spans
   FilePath ->
-  -- | List of lines from that file
+  -- | Contents of that file
   String ->
-  -- | Adjusted lines together with comments extracted from them
+  -- | Adjusted input with comments extracted from it
   (String, [Located String])
 extractCommentsFromLines path =
   unlines' . unzip . zipWith (extractCommentFromLine path) [1 ..] . lines

src/Ormolu/Parser/CommentStream.hs

@@ -7,7 +7,6 @@ module Ormolu.Parser.CommentStream
   ( CommentStream (..),
     Comment (..),
     mkCommentStream,
-    isShebang,
     isPrevHaddock,
     isMultilineComment,
     showCommentStream,
@@ -16,14 +15,14 @@ where
 
 import Data.Char (isSpace)
 import Data.Data (Data)
-import Data.Either (partitionEithers)
-import Data.List (dropWhileEnd, isPrefixOf, sortOn)
+import qualified Data.List as L
 import Data.List.NonEmpty (NonEmpty (..))
 import qualified Data.List.NonEmpty as NE
 import Data.Maybe (mapMaybe)
 import qualified GHC
 import qualified Lexer as GHC
 import Ormolu.Parser.Pragma
+import Ormolu.Parser.Shebang
 import Ormolu.Utils (showOutputable)
 import SrcLoc
 
@@ -45,37 +44,38 @@ mkCommentStream ::
   [Located String] ->
   -- | Parser state to use for comment extraction
   GHC.PState ->
-  -- | Comment stream, a set of extracted pragmas, and extracted shebangs
-  (CommentStream, [Pragma], [Located String])
+  -- | Stack header, shebangs, pragmas, and comment stream
+  ( Maybe (RealLocated Comment),
+    [Shebang],
+    [([RealLocated Comment], Pragma)],
+    CommentStream
+  )
 mkCommentStream extraComments pstate =
-  ( CommentStream $
-      mkComment <$> sortOn (realSrcSpanStart . getRealSrcSpan) comments,
+  ( mstackHeader,
+    shebangs,
     pragmas,
-    shebangs
+    CommentStream comments
   )
   where
-    (comments, pragmas) = partitionEithers (partitionComments <$> rawComments)
-    rawComments =
-      mapMaybe toRealSpan $
+    (comments, pragmas) = extractPragmas rawComments1
+    (rawComments1, mstackHeader) = extractStackHeader rawComments0
+    rawComments0 =
+      L.sortOn (realSrcSpanStart . getRealSrcSpan) . mapMaybe toRealSpan $
         otherExtraComments
           ++ mapMaybe (liftMaybe . fmap unAnnotationComment) (GHC.comment_q pstate)
           ++ concatMap
             (mapMaybe (liftMaybe . fmap unAnnotationComment) . snd)
             (GHC.annotations_comments pstate)
-    (shebangs, otherExtraComments) = span (isShebang . unLoc) extraComments
-
--- | Return 'True' if given 'String' is a shebang.
-isShebang :: String -> Bool
-isShebang str = "#!" `isPrefixOf` str
+    (shebangs, otherExtraComments) = extractShebangs extraComments
 
 -- | Test whether a 'Comment' looks like a Haddock following a definition,
 -- i.e. something starting with @-- ^@.
 isPrevHaddock :: Comment -> Bool
-isPrevHaddock (Comment (x :| _)) = "-- ^" `isPrefixOf` x
+isPrevHaddock (Comment (x :| _)) = "-- ^" `L.isPrefixOf` x
 
 -- | Is this comment multiline-style?
 isMultilineComment :: Comment -> Bool
-isMultilineComment (Comment (x :| _)) = "{-" `isPrefixOf` x
+isMultilineComment (Comment (x :| _)) = "{-" `L.isPrefixOf` x
 
 -- | Pretty-print a 'CommentStream'.
 showCommentStream :: CommentStream -> String
@@ -94,7 +94,7 @@ showCommentStream (CommentStream xs) =
 mkComment :: RealLocated String -> RealLocated Comment
 mkComment (L l s) =
   L l . Comment . fmap dropTrailing $
-    if "{-" `isPrefixOf` s
+    if "{-" `L.isPrefixOf` s
       then case NE.nonEmpty (lines s) of
         Nothing -> s :| []
         Just (x :| xs) ->
@@ -106,7 +106,7 @@ mkComment (L l s) =
            in x :| (drop n <$> xs)
       else s :| []
   where
-    dropTrailing = dropWhileEnd isSpace
+    dropTrailing = L.dropWhileEnd isSpace
     startIndent = srcSpanStartCol l - 1
 
 -- | Get a 'String' from 'GHC.AnnotationComment'.
@@ -129,12 +129,32 @@ toRealSpan :: Located a -> Maybe (RealLocated a)
 toRealSpan (L (RealSrcSpan l) a) = Just (L l a)
 toRealSpan _ = Nothing
 
--- | If a given comment is a pragma, return it in parsed form in 'Right'.
--- Otherwise return the original comment unchanged.
-partitionComments ::
-  RealLocated String ->
-  Either (RealLocated String) Pragma
-partitionComments input =
-  case parsePragma (unRealSrcSpan input) of
-    Nothing -> Left input
-    Just pragma -> Right pragma
+-- | Detect and extract stack header if it is present.
+extractStackHeader ::
+  [RealLocated String] ->
+  ([RealLocated String], Maybe (RealLocated Comment))
+extractStackHeader = \case
+  [] -> ([], Nothing)
+  (x : xs) ->
+    let comment = mkComment x
+     in if isStackHeader (unRealSrcSpan comment)
+          then (xs, Just comment)
+          else (x : xs, Nothing)
+  where
+    isStackHeader (Comment (x :| _)) =
+      "stack" `L.isPrefixOf` dropWhile isSpace (drop 2 x)
+
+-- | Extract pragmas and their associated comments.
+extractPragmas ::
+  [RealLocated String] ->
+  ([RealLocated Comment], [([RealLocated Comment], Pragma)])
+extractPragmas = go id id
+  where
+    go csSoFar pragmasSoFar = \case
+      [] -> (csSoFar [], pragmasSoFar [])
+      (x : xs) ->
+        case parsePragma (unRealSrcSpan x) of
+          Nothing -> go (csSoFar . (mkComment x :)) pragmasSoFar xs
+          Just pragma ->
+            let combined = (csSoFar [], pragma)
+             in go id (pragmasSoFar . (combined :)) xs

src/Ormolu/Parser/Result.hs

@@ -11,19 +11,22 @@ import GHC
 import Ormolu.Parser.Anns
 import Ormolu.Parser.CommentStream
 import Ormolu.Parser.Pragma (Pragma)
+import Ormolu.Parser.Shebang (Shebang)
 
 -- | A collection of data that represents a parsed module in Ormolu.
 data ParseResult = ParseResult
   { -- | 'ParsedSource' from GHC
     prParsedSource :: ParsedSource,
     -- | Ormolu-specfic representation of annotations
     prAnns :: Anns,
+    -- | Stack header
+    prStackHeader :: Maybe (RealLocated Comment),
+    -- | Shebangs found in the input
+    prShebangs :: [Shebang],
+    -- | Pragmas and the associated comments
+    prPragmas :: [([RealLocated Comment], Pragma)],
     -- | Comment stream
     prCommentStream :: CommentStream,
-    -- | Extensions enabled in that module
-    prExtensions :: [Pragma],
-    -- | Shebangs found in the input
-    prShebangs :: [Located String],
     -- | Whether or not record dot syntax is enabled
     prUseRecordDot :: Bool,
     -- | Whether or not ImportQualifiedPost is enabled

src/Ormolu/Parser/Shebang.hs

@@ -0,0 +1,27 @@
+{-# LANGUAGE DeriveDataTypeable #-}
+
+-- | A module for dealing with shebangs.
+module Ormolu.Parser.Shebang
+  ( Shebang (..),
+    extractShebangs,
+    isShebang,
+  )
+where
+
+import Data.Data (Data)
+import qualified Data.List as L
+import SrcLoc
+
+-- | A wrapper for a shebang.
+newtype Shebang = Shebang (Located String)
+  deriving (Eq, Data)
+
+-- | Extract shebangs from the beginning of a comment stream.
+extractShebangs :: [Located String] -> ([Shebang], [Located String])
+extractShebangs comments = (Shebang <$> shebangs, rest)
+  where
+    (shebangs, rest) = span (isShebang . unLoc) comments
+
+-- | Return 'True' if given 'String' is a shebang.
+isShebang :: String -> Bool
+isShebang str = "#!" `L.isPrefixOf` str

src/Ormolu/Printer.hs

@@ -21,8 +21,9 @@ printModule ::
 printModule ParseResult {..} =
   runR
     ( p_hsModule
+        prStackHeader
         prShebangs
-        prExtensions
+        prPragmas
         prImportQualifiedPost
         prParsedSource
     )

src/Ormolu/Printer/Comments.hs

@@ -7,17 +7,15 @@ module Ormolu.Printer.Comments
   ( spitPrecedingComments,
     spitFollowingComments,
     spitRemainingComments,
-    spitStackHeader,
+    spitCommentNow,
+    spitCommentPending,
   )
 where
 
 import Control.Monad
-import Data.Char (isSpace)
 import Data.Coerce (coerce)
 import Data.Data (Data)
-import Data.List (isPrefixOf)
 import qualified Data.List.NonEmpty as NE
-import Data.List.NonEmpty (NonEmpty (..))
 import qualified Data.Text as T
 import Ormolu.Parser.CommentStream
 import Ormolu.Printer.Internal
@@ -55,16 +53,6 @@ spitFollowingComments ref = do
 spitRemainingComments :: R ()
 spitRemainingComments = void $ handleCommentSeries spitRemainingComment
 
--- | If there is a stack header in the comment stream, print it.
-spitStackHeader :: R ()
-spitStackHeader = do
-  let isStackHeader (Comment (x :| _)) =
-        "stack" `isPrefixOf` dropWhile isSpace (drop 2 x)
-  mstackHeader <- popComment (isStackHeader . unRealSrcSpan)
-  forM_ mstackHeader $ \(L spn x) -> do
-    spitCommentNow spn x
-    newline
-
 ----------------------------------------------------------------------------
 -- Single-comment functions