Implement more precise comment placement without ‘ghc-exactprint’

‘ghc-exactprint’ (or perhaps lexer of GHC itself) does a fairly poor job at
associating comments with elements of AST. In many cases the result is not
what you'd expect. We ran into insuperable problems with that to the effect
that correct comment placement were impossible.

The new approach is to exploit the raw position information provided by the
GHC lexer, that is, spans attached to AST elements and comments. This
allowed us to place comments in output in a very precise and satisfactory
fashion.

app/Main.hs

@@ -79,8 +79,7 @@ optsParserInfo = info (helper <*> ver <*> optsParser) . mconcat $
         , $gitBranch
         , $gitHash
         ]
-      , "using ghc-exactprint " ++ VERSION_ghc_exactprint
-      , "using ghc            " ++ VERSION_ghc
+      , "using ghc " ++ VERSION_ghc
       ]
 
 optsParser :: Parser Opts

data/examples/declaration/data/fat-multiline-out.hs

@@ -1,5 +1,4 @@
 -- | Something.
-
 data Foo
   = Foo
       Int

data/examples/declaration/data/gadt/multiline-out.hs

@@ -5,10 +5,10 @@ data Foo a where
   -- | 'Foo' is wonderful.
   Foo
     :: forall a b. ( Show a
-                   , Eq b
-                   ) -- foo
-    -- bar
-    => a
+                   , Eq b -- foo
+                   )
+    => -- bar
+    a
     -> b
     -> Foo 'Int
   -- | But 'Bar' is also not too bad.
@@ -20,3 +20,4 @@ data Foo a where
   Baz
     :: forall a. a
     -> Foo 'String
+  (:~>) :: Foo a -> Foo a -> Foo a

data/examples/declaration/data/gadt/multiline.hs

@@ -13,3 +13,4 @@ data Foo a where
   -- | So is 'Baz'.
   Baz
     :: forall a. a -> Foo 'String
+  (:~>) :: Foo a -> Foo a -> Foo a

data/examples/declaration/data/multiline-arg-parens-out.hs

@@ -0,0 +1,8 @@
+-- | Something.
+data Foo
+  = Foo
+      Bar
+      (Set Baz) -- and here we go
+      -- and that's it
+      Text
+  deriving (Eq)

data/examples/declaration/data/multiline-arg-parens.hs

@@ -0,0 +1,8 @@
+-- | Something.
+
+data Foo
+  = Foo Bar
+        (Set Baz) -- and here we go
+                  -- and that's it
+        Text
+  deriving (Eq)

data/examples/declaration/data/record-out.hs

@@ -3,7 +3,14 @@ data Foo
   = Foo
       { fooX :: Int -- ^ X
       , fooY :: Int -- ^ Y
-      , fooBar, fooBaz :: Bool -- ^ BarBaz
+      , fooBar, fooBaz :: NonEmpty (Identity Bool) -- ^ BarBaz
+      , fooGag
+        , fooGog
+          :: NonEmpty
+               ( Indentity
+                   Bool
+               )
+        -- ^ GagGog
       , fooFoo
         , barBar
           :: Int -- ^ Huh!

data/examples/declaration/data/record.hs

@@ -3,7 +3,10 @@
 data Foo = Foo
   { fooX :: Int -- ^ X
   , fooY :: Int -- ^ Y
-  , fooBar, fooBaz :: Bool -- ^ BarBaz
+  , fooBar, fooBaz :: NonEmpty (Identity Bool) -- ^ BarBaz
+  , fooGag, fooGog :: NonEmpty (Indentity
+                                  Bool)
+    -- ^ GagGog
   , fooFoo
       , barBar :: Int -- ^ Huh!
   } deriving (Eq, Show)

data/examples/declaration/data/simple-broken-out.hs

@@ -0,0 +1,8 @@
+-- | Here we go.
+data Foo
+  = Foo {unFoo :: Int}
+  deriving (Eq)
+
+-- | And once again.
+data Bar = Bar {unBar :: Int}
+  deriving (Eq)

data/examples/declaration/data/simple-broken.hs

@@ -0,0 +1,10 @@
+-- | Here we go.
+
+data Foo
+  = Foo { unFoo :: Int }
+  deriving (Eq)
+
+-- | And once again.
+
+data Bar = Bar { unBar :: Int }
+  deriving (Eq)

data/examples/import/comments-per-import-out.hs

@@ -1,4 +1,4 @@
-import Bar
-import Baz
+-- (1)
+import Bar -- (2)
+import Baz -- (3)
 import Foo
--- (3)

data/examples/import/explicit-imports-with-comments-out.hs

@@ -1,5 +1,6 @@
 import qualified MegaModule as M
-  ( (<<<)
+  ( -- (1)
+    (<<<) -- (2)
   , (>>>)
-  , Either
+  , Either -- (3)
   )

data/examples/other/multiline-comments-reindent-out.hs

@@ -0,0 +1,6 @@
+{-
+And so here we have a
+multiline comment.
+
+Indeed.
+-}

data/examples/other/multiline-comments-reindent.hs

@@ -0,0 +1,6 @@
+              {-
+                 And so here we have a
+                   multiline comment.
+
+                 Indeed.
+              -}

ormolu.cabal

@@ -29,16 +29,17 @@ library
   hs-source-dirs:     src
   build-depends:      base           >= 4.8 && < 5.0
                     , containers     >= 0.5 && < 0.7
+                    , dlist          >= 0.8 && < 0.9
                     , exceptions     >= 0.6 && < 0.11
                     , ghc            >= 8.4.3
                     , ghc-boot-th    >= 8.4.3
-                    , ghc-exactprint >= 0.5.6
+                    , ghc-paths      >= 0.1 && < 0.2
                     , mtl            >= 2.0 && < 3.0
                     , syb            >= 0.7 && < 0.8
                     , text           >= 0.2 && < 1.3
                     , yaml           >= 0.8 && < 0.12
   exposed-modules:    Ormolu
-                    , Ormolu.Comments
+                    , Ormolu.CommentStream
                     , Ormolu.Config
                     , Ormolu.Diff
                     , Ormolu.Exception
@@ -56,6 +57,8 @@ library
                     , Ormolu.Printer.Meat.ImportExport
                     , Ormolu.Printer.Meat.Module
                     , Ormolu.Printer.Meat.Type
+                    , Ormolu.SpanStream
+                    , Ormolu.Utils
   if flag(dev)
     ghc-options:      -Wall -Werror -Wcompat
                       -Wincomplete-record-updates
@@ -94,7 +97,6 @@ executable ormolu
   hs-source-dirs:     app
   build-depends:      base           >= 4.8 && < 5.0
                     , ghc            >= 8.4.3
-                    , ghc-exactprint >= 0.5.6
                     , gitrev         >= 1.3 && < 1.4
                     , optparse-applicative >= 0.14 && < 0.15
                     , ormolu

src/Ormolu.hs

@@ -18,7 +18,7 @@ import Control.Monad
 import Control.Monad.IO.Class (MonadIO (..))
 import Data.Text (Text)
 import Debug.Trace
-import Language.Haskell.GHC.ExactPrint.Types
+import Ormolu.CommentStream
 import Ormolu.Config
 import Ormolu.Diff
 import Ormolu.Exception
@@ -47,20 +47,20 @@ ormolu
   -> String                     -- ^ Input to format
   -> m Text
 ormolu cfg path str = do
-  (ws, (anns0, parsedSrc0)) <-
+  (ws, (cstream0, parsedSrc0)) <-
     parseModule' cfg OrmoluParsingFailed path str
   when (cfgDebug cfg) $ do
     traceM "warnings:\n"
     traceM (concatMap showWarn ws)
-    traceM "anns:\n"
-    traceM (showOutputable anns0)
-  let txt = printModule (cfgDebug cfg) anns0 parsedSrc0
+    traceM "comment stream:\n"
+    traceM (showCommentStream cstream0)
+  let txt = printModule (cfgDebug cfg) cstream0 parsedSrc0
   -- Parse the result of pretty-printing again and make sure that AST is the
   -- same as AST of original snippet module span positions.
   unless (cfgUnsafe cfg) $ do
-    (_, (anns1, parsedSrc1)) <-
+    (_, (cstream1, parsedSrc1)) <-
       parseModule' cfg OrmoluOutputParsingFailed "<rendered>" (T.unpack txt)
-    when (diff (anns0, parsedSrc0) (anns1, parsedSrc1)) $
+    when (diff (cstream0, parsedSrc0) (cstream1, parsedSrc1)) $
       liftIO $ throwIO (OrmoluASTDiffers str txt)
   return txt
 
@@ -90,7 +90,8 @@ parseModule'
      -- ^ How to obtain 'OrmoluException' to throw when parsing fails
   -> FilePath                   -- ^ File name to use in errors
   -> String                     -- ^ Actual input for the parser
-  -> m ([GHC.Warn], (Anns, GHC.ParsedSource)) -- ^ Annotations and parsed source
+  -> m ([GHC.Warn], (CommentStream, GHC.ParsedSource))
+     -- ^ Comment stream and parsed source
 parseModule' Config {..} mkException path str = do
   (ws, r) <- parseModule cfgDynOptions path str
   case r of
@@ -103,6 +104,14 @@ showWarn :: GHC.Warn -> String
 showWarn (GHC.Warn reason l) =
   showOutputable reason ++ "\n" ++ showOutputable l ++ "\n"
 
+-- | Pretty-print a 'CommentStream'.
+
+showCommentStream :: CommentStream -> String
+showCommentStream (CommentStream xs) = unlines $
+  showComment <$> xs
+  where
+    showComment (GHC.L l str) = showOutputable l ++ " " ++ show str
+
 -- | Pretty-print an 'GHC.Outputable' thing.
 
 showOutputable :: GHC.Outputable o => o -> String

src/Ormolu/CommentStream.hs

@@ -0,0 +1,93 @@
+{-# LANGUAGE DeriveDataTypeable         #-}
+{-# LANGUAGE GeneralizedNewtypeDeriving #-}
+{-# LANGUAGE LambdaCase                 #-}
+
+-- | Functions for working with comment stream.
+
+module Ormolu.CommentStream
+  ( CommentStream (..)
+  , Comment (..)
+  , mkCommentStream
+  , isPrevHaddock
+  )
+where
+
+import Data.Data (Data)
+import Data.List (isPrefixOf, sortOn)
+import Data.List.NonEmpty (NonEmpty (..))
+import Data.Maybe (mapMaybe)
+import SrcLoc
+import qualified Data.List.NonEmpty as NE
+import qualified GHC
+import qualified Lexer as GHC
+
+-- | A stream of 'RealLocated' comment strings in ascending order with
+-- respect to beginning of corresponding spans.
+
+newtype CommentStream = CommentStream [RealLocated Comment]
+  deriving (Eq, Data, Semigroup, Monoid)
+
+-- | A wrapper for a single comment.
+
+newtype Comment = Comment (NonEmpty String)
+  deriving (Eq, Show, Data)
+
+-- | Create 'CommentStream' from 'GHC.PState'.
+
+mkCommentStream
+  :: [Located String]           -- ^ Extra comments to include
+  -> GHC.PState                 -- ^ Parser state to use for comment extraction
+  -> CommentStream
+mkCommentStream extraComments pstate
+  = CommentStream
+  -- NOTE It's easier to normalize pragmas right when we construct comment
+  -- streams. Because this way we need to do it only once and when we
+  -- perform checking later they'll automatically match.
+  . fmap (fmap (Comment . normalizeComment . normalizePragma))
+  . sortOn startOfSpan
+  . mapMaybe toRealSpan $
+      extraComments ++
+      (fmap unAnnotationComment <$> GHC.comment_q pstate) ++
+      concatMap (fmap (fmap unAnnotationComment) . snd) (GHC.annotations_comments pstate)
+  where
+    startOfSpan (L l _) = realSrcSpanStart l
+    toRealSpan (L (RealSrcSpan l) a) = Just (L l a)
+    toRealSpan _ = Nothing
+
+-- | Test whether a 'Comment' looks like a Haddock following a definition.
+
+isPrevHaddock :: Comment -> Bool
+isPrevHaddock (Comment (x :| _)) = "-- ^" `isPrefixOf` x
+
+----------------------------------------------------------------------------
+-- Helpers
+
+-- | Normalize pragmas by deleting extra white space.
+
+normalizePragma :: String -> String
+normalizePragma x =
+  if "{-#" `isPrefixOf` x
+    then unwords (words x)
+    else x
+
+-- | Normalize comment string. Sometimes one multi-line comment is turned
+-- into several lines for subsequent outputting with correct indentation for
+-- each line.
+
+normalizeComment :: String -> NonEmpty String
+normalizeComment s =
+  case NE.nonEmpty (lines s) of
+    Nothing -> s :| []
+    Just xs -> dropWhile (== ' ') <$> xs
+
+-- | Get a 'String' from 'GHC.AnnotationComment'.
+
+unAnnotationComment :: GHC.AnnotationComment -> String
+unAnnotationComment = \case
+  GHC.AnnDocCommentNext s -> s
+  GHC.AnnDocCommentPrev s -> s
+  GHC.AnnDocCommentNamed s -> s
+  GHC.AnnDocSection _ s -> s
+  GHC.AnnDocOptions s -> s
+  GHC.AnnLineComment s -> s
+  GHC.AnnBlockComment s -> s

src/Ormolu/Comments.hs

@@ -1,33 +0,0 @@
--- | Various functions for manipulation of 'Comment's.
-
-module Ormolu.Comments
-  ( annComment
-  , normalizeComment
-  )
-where
-
-import Data.List (isPrefixOf)
-import Language.Haskell.GHC.ExactPrint.Types
-
--- | If 'KeywordId' is a comment, extract it.
-
-annComment :: KeywordId -> Maybe Comment
-annComment (AnnComment x) = Just x
-annComment _ = Nothing
-
--- | Normalize comment string. Sometimes one multi-line comment is turned
--- into several lines for subsequent outputting with correct indentation for
--- each line.
-
-normalizeComment :: String -> [String]
-normalizeComment s =
-  if isMultiline s
-    then if isPragma s
-           then [normalizePragma s]
-           else normalizeIndent s
-    else [s]
-  where
-    isMultiline x = not ("--" `isPrefixOf` x)
-    isPragma x = "{-#" `isPrefixOf` x
-    normalizeIndent = fmap (dropWhile (== ' ')) . lines
-    normalizePragma = unwords . words

src/Ormolu/Config.hs

@@ -1,9 +1,9 @@
--- | Configuration options used by the tool.
-
 {-# LANGUAGE GeneralizedNewtypeDeriving #-}
 {-# LANGUAGE OverloadedStrings          #-}
 {-# LANGUAGE RecordWildCards            #-}
 
+-- | Configuration options used by the tool.
+
 module Ormolu.Config
   ( Config (..)
   , defaultConfig

src/Ormolu/Diff.hs

@@ -10,26 +10,23 @@ module Ormolu.Diff
 where
 
 import Data.Generics
-import Data.Maybe (mapMaybe)
-import GHC hiding (GhcPs)
-import Language.Haskell.GHC.ExactPrint.Types
-import Ormolu.Comments
+import GHC
+import Ormolu.CommentStream
 import Ormolu.Imports (sortImports)
-import qualified Data.Map.Strict as M
 
 -- | Return 'False' if two annotated ASTs are the same modulo span
 -- positions.
 
 diff
-  :: (Anns, ParsedSource)       -- ^ First annotated AST
-  -> (Anns, ParsedSource)       -- ^ Second annotated AST
+  :: (CommentStream, ParsedSource) -- ^ First annotated AST
+  -> (CommentStream, ParsedSource) -- ^ Second annotated AST
   -> Bool
-diff (anns0, ps0) (anns1, ps1) =
-  not (matchIgnoringSrcSpans (simplifyAnns anns0) (simplifyAnns anns1))
+diff (cstream0, ps0) (cstream1, ps1) =
+  not (matchIgnoringSrcSpans cstream0 cstream1)
   || not (matchIgnoringSrcSpans ps0 ps1)
 
--- | Compare two 'ParsedSource' values disregarding differences in
--- 'SrcSpan's and the ordering of import lists.
+-- | Compare two values for equality disregarding differences in 'SrcSpan's
+-- and the ordering of import lists.
 
 matchIgnoringSrcSpans :: Data a => a -> a -> Bool
 matchIgnoringSrcSpans = genericQuery
@@ -49,24 +46,3 @@ matchIgnoringSrcSpans = genericQuery
           matchIgnoringSrcSpans
             hs0 { hsmodImports = sortImports (hsmodImports hs0) }
             hs1 { hsmodImports = sortImports (hsmodImports hs1) }
-
--- | Simplified collection of 'Comment's.
-
-data Comments = Comments [Comment] [Comment] [Comment]
-  deriving (Show, Typeable, Data)
-
--- | Simplify a collection of annotations.
-
-simplifyAnns :: Anns -> [Comments]
-simplifyAnns = fmap simplifyAnn . M.elems
-
--- | Simplify single annotation.
-
-simplifyAnn :: Annotation -> Comments
-simplifyAnn Ann {..} = Comments
-  (f . fst <$> annPriorComments)
-  (f . fst <$> annFollowingComments)
-  (mapMaybe (fmap f . annComment) (fst <$> annsDP))
-  where
-    f (Comment str i o) =
-      Comment (unlines $ normalizeComment str) i o

src/Ormolu/Exception.hs

@@ -1,7 +1,7 @@
--- | 'OrmoluException' type and surrounding definitions.
-
 {-# LANGUAGE LambdaCase #-}
 
+-- | 'OrmoluException' type and surrounding definitions.
+
 module Ormolu.Exception
   ( OrmoluException (..)
   , withPrettyOrmoluExceptions

src/Ormolu/Imports.hs

@@ -12,8 +12,9 @@ import Data.Bifunctor
 import Data.Function (on)
 import Data.List (sortBy)
 import GHC hiding (GhcPs, IE)
+import HsExtension
 import HsImpExp (IE (..))
-import Language.Haskell.GHC.ExactPrint.Types
+import Ormolu.Utils (unL)
 
 -- | Sort imports by module name. This also sorts explicit import lists for
 -- each declaration.
@@ -83,8 +84,3 @@ compareIewn (IEPattern _) (IEType _) = LT
 compareIewn (IEType _) (IEName _) = GT
 compareIewn (IEType _) (IEPattern _) = GT
 compareIewn (IEType x) (IEType y) = unL x `compare` unL y
-
--- | Exact inner value from 'Located'.
-
-unL :: Located e -> e
-unL (L _ e) = e

src/Ormolu/Parser.hs

@@ -10,16 +10,23 @@ where
 import Control.Exception
 import Control.Monad
 import Control.Monad.IO.Class
-import Data.Generics
-import GHC hiding (GhcPs, IE, parseModule)
+import Data.List (isPrefixOf)
+import Data.Maybe (catMaybes)
+import GHC hiding (IE, parseModule, parser)
 import GHC.LanguageExtensions.Type (Extension (Cpp))
-import Language.Haskell.GHC.ExactPrint.Parsers hiding (parseModule)
-import Language.Haskell.GHC.ExactPrint.Types
+import GHC.Paths (libdir)
+import Ormolu.CommentStream
 import Ormolu.Config
 import Ormolu.Exception
 import qualified CmdLineParser as GHC
-import qualified Data.Map.Strict as M
 import qualified DynFlags as GHC
+import qualified FastString as GHC
+import qualified HeaderInfo as GHC
+import qualified Lexer as GHC
+import qualified Outputable as GHC
+import qualified Parser as GHC
+import qualified SrcLoc as GHC
+import qualified StringBuffer as GHC
 
 -- | Parse a complete module from string.
 
@@ -28,8 +35,9 @@ parseModule
   => [DynOption]        -- ^ Dynamic options that affect parsing
   -> FilePath           -- ^ File name (only for source location annotations)
   -> String             -- ^ Input for parser
-  -> m ([GHC.Warn], Either (SrcSpan, String) (Anns, ParsedSource))
-parseModule dynOpts path input = liftIO $ do
+  -> m ([GHC.Warn], Either (SrcSpan, String) (CommentStream, ParsedSource))
+parseModule dynOpts path input' = liftIO $ do
+  let (input, extraComments) = stripLinePragmas input'
   (ws, dynFlags) <- ghcWrapper $ do
     dynFlags0 <- initDynFlagsPure path input
     (dynFlags1, _, ws) <-
@@ -40,64 +48,97 @@ parseModule dynOpts path input = liftIO $ do
   -- want.
   when (GHC.xopt Cpp dynFlags) $
    throwIO OrmoluCppEnabled
-  let r = case parseModuleFromStringInternal dynFlags path input of
-            Left e -> Left e
-            Right (anns, psrc) ->
-              Right (dropImportComments anns psrc, psrc)
+  let r = case runParser GHC.parseModule dynFlags path input of
+        GHC.PFailed _ ss m -> Left (ss, GHC.showSDoc dynFlags m)
+        GHC.POk x pmod -> Right (mkCommentStream extraComments x, pmod)
   return (ws, r)
 
--- | Drop comments associated with elements of import lists. The reason we
--- do this is that those comments are not associated as user would expect.
---
--- For example:
---
--- > import Foo -- (1)
--- > import Bar -- (2)
--- > import Baz -- (3)
---
--- Here, @(1)@ is considered preceding comment of @Baz@ import, in fact, it
--- has nothing to do with @Foo@. Similarly, @(2)@ is associted with @Baz@,
--- not @Bar@. Finally @(3)@ is actually preceeding comment for the thing
--- that follows the import list, or if nothing follows, it's a trailing
--- comment of the whole module.
---
--- Since we do sorting of imports when we print them, we would need to do
--- sorting of corresponding annotations as well (because we do checking of
--- annotations and AST in "Ormolu.Diff" as part of the self-check), and this
--- is hard. Even if we manage to do that the result will be totally
--- confusing, e.g. for the example above:
---
--- > -- (1)
--- > import Bar
--- > -- (2)
--- > import Baz
--- > import Foo
--- > -- (3)
---
--- The solution is to drop the tricky comments right after parsing, so we
--- don't need to deal with them at all, including the need to update
--- annotations for the correctness checking we do in "Ormolu.Diff".
---
--- The solution is not perfect, but practical.
+----------------------------------------------------------------------------
+-- Helpers (taken from ghc-exactprint)
 
-dropImportComments
-  :: Anns
-  -> ParsedSource
-  -> Anns
-dropImportComments anns (L _ HsModule {..}) =
-  foldr M.delete anns (hsmodImports >>= keyAnnsFor)
+-- | Requires GhcMonad constraint because there is no pure variant of
+-- 'parseDynamicFilePragma'. Yet, in constrast to 'initDynFlags', it does
+-- not (try to) read the file at filepath, but solely depends on the module
+-- source in the input string.
+--
+-- Passes "-hide-all-packages" to the GHC API to prevent parsing of package
+-- environment files. However this only works if there is no invocation of
+-- 'setSessionDynFlags' before calling 'initDynFlagsPure'. See GHC tickets
+-- #15513, #15541.
 
--- | Extract all 'AnnKey's from given 'Data'.
+initDynFlagsPure
+  :: GHC.GhcMonad m
+  => FilePath                   -- ^ Module path
+  -> String                     -- ^ Module contents
+  -> m GHC.DynFlags             -- ^ Dynamic flags for that module
+initDynFlagsPure fp input = do
+  -- I was told we could get away with using the 'unsafeGlobalDynFlags'. as
+  -- long as 'parseDynamicFilePragma' is impure there seems to be no reason
+  -- to use it.
+  dflags0 <- GHC.getSessionDynFlags
+  let pragmaInfo = GHC.getOptions dflags0 (GHC.stringToStringBuffer input) fp
+  (dflags1, _, _) <- GHC.parseDynamicFilePragma dflags0 pragmaInfo
+  -- Turn this on last to avoid T10942
+  let dflags2 = dflags1 `GHC.gopt_set` GHC.Opt_KeepRawTokenStream
+  -- Prevent parsing of .ghc.environment.* "package environment files"
+  (dflags3, _, _) <- GHC.parseDynamicFlagsCmdLine
+    dflags2
+    [GHC.noLoc "-hide-all-packages"]
+  _ <- GHC.setSessionDynFlags dflags3
+  return dflags3
 
-keyAnnsFor :: GenericQ [AnnKey]
-keyAnnsFor a = everything mappend (const id `ext2Q` queryLocated) a []
+-- | Default runner of 'GHC.Ghc' action in 'IO'.
+
+ghcWrapper :: GHC.Ghc a -> IO a
+ghcWrapper
+  = GHC.defaultErrorHandler GHC.defaultFatalMessager GHC.defaultFlushOut
+  . GHC.runGhc (Just libdir)
+
+-- | Run a 'GHC.P' computation.
+
+runParser
+  :: GHC.P a                    -- ^ Computation to run
+  -> GHC.DynFlags               -- ^ Dynamic flags
+  -> FilePath                   -- ^ Module path
+  -> String                     -- ^ Module contents
+  -> GHC.ParseResult a          -- ^ Parse result
+runParser parser flags filename input = GHC.unP parser parseState
   where
-    queryLocated
-      :: (Data e0, Data e1)
-      => GenLocated e0 e1
-      -> [AnnKey]
-      -> [AnnKey]
-    queryLocated (L mspn x) =
-      case cast mspn :: Maybe SrcSpan of
-        Nothing -> id
-        Just spn -> (mkAnnKey (L spn x) :)
+    location = GHC.mkRealSrcLoc (GHC.mkFastString filename) 1 1
+    buffer = GHC.stringToStringBuffer input
+    parseState = GHC.mkPState flags buffer location
+
+-- | Remove GHC style line pragams (@{-# LINE .. #-}@) and convert them into
+-- comments.
+
+stripLinePragmas :: String -> (String, [Located String])
+stripLinePragmas = unlines' . unzip . findLines . lines
+  where
+    unlines' (a, b) = (unlines a, catMaybes b)
+
+findLines :: [String] -> [(String, Maybe (Located String))]
+findLines = zipWith checkLine [1..]
+
+checkLine :: Int -> String -> (String, Maybe (Located String))
+checkLine line s
+  | "{-# LINE" `isPrefixOf` s =
+      let (pragma, res) = getPragma s
+          size = length pragma
+          mSrcLoc = mkSrcLoc (GHC.mkFastString "LINE")
+          ss = mkSrcSpan (mSrcLoc line 1) (mSrcLoc line (size+1))
+      in (res, Just $ L ss pragma)
+  -- Deal with shebang/cpp directives too
+  -- x |  "#" `isPrefixOf` s = ("",Just $ Comment ((line, 1), (line, length s)) s)
+  | "#!" `isPrefixOf` s =
+      let mSrcLoc = mkSrcLoc (GHC.mkFastString "SHEBANG")
+          ss = mkSrcSpan (mSrcLoc line 1) (mSrcLoc line (length s))
+      in ("",Just $ L ss s)
+  | otherwise = (s, Nothing)
+
+getPragma :: String -> (String, String)
+getPragma [] = error "Ormolu.Parser.getPragma: input must not be empty"
+getPragma s@(x:xs)
+  | "#-}" `isPrefixOf` s = ("#-}", "   " ++ drop 3 s)
+  | otherwise =
+      let (prag, remline) = getPragma xs
+      in (x:prag, ' ':remline)

src/Ormolu/Printer.hs

@@ -10,17 +10,18 @@ module Ormolu.Printer
 where
 
 import Data.Text (Text)
-import GHC (ParsedSource)
-import Language.Haskell.GHC.ExactPrint.Types
+import GHC
+import Ormolu.CommentStream
 import Ormolu.Printer.Combinators
 import Ormolu.Printer.Meat.Module
+import Ormolu.SpanStream
 
 -- | Render a module.
 
 printModule
-  :: Bool                       -- ^ Trace debugging information
-  -> Anns                       -- ^ Annotations
+  :: Bool                       -- ^ Whether to trace debugging information
+  -> CommentStream              -- ^ Comment stream
   -> ParsedSource               -- ^ Parsed source
   -> Text                       -- ^ Resulting rendition
-printModule debugOn anns src =
-  runR debugOn (p_hsModule src) anns
+printModule debugOn cstream src =
+  runR debugOn (p_hsModule src) (mkSpanStream src) cstream

src/Ormolu/Printer/Combinators.hs

@@ -15,6 +15,7 @@ module Ormolu.Printer.Combinators
   , newline
   , inci
   , relaxComments
+  , hasMoreComments
   , located
   , locatedVia
   , located'
@@ -44,9 +45,9 @@ import Data.Bool (bool)
 import Data.Data (Data)
 import Data.List (intersperse)
 import Data.Text (Text)
-import Language.Haskell.GHC.ExactPrint.Types
 import Ormolu.Printer.Comments
 import Ormolu.Printer.Internal
+import Ormolu.Utils (unL, getSpan, isModule)
 import Outputable (Outputable (..), showSDocUnsafe)
 import SrcLoc
 import qualified Data.Text as T
@@ -85,43 +86,25 @@ locatedVia
   -> Located a                  -- ^ Thing to enter
   -> (a -> R ())                -- ^ How to render inner value
   -> R ()
-locatedVia ml loc@(L l a) f = do
-  mann <- lookupAnn loc
+locatedVia ml loc f = do
   relaxed <- relaxedComments
-  let m = case ml of
-       Nothing -> f a
-       Just l' -> switchLayout l' (f a)
-  case mann of
-    Nothing -> m
-    Just Ann {..} -> bool sitcc id relaxed $ do
-      -- There are three things in 'Ann' which contain comments:
-
-      let cmode =
-            if annGetConstr a == CN "HsModule"
-              then Module
-              else Other
-          (before, after) = partitionDPs cmode l annsDP
-
-      -- 'annPriorComments' contains comments that were directly placed
-      -- before entities such as comments (in both styles) before function
-      -- definitions and inline comments before smaller things like types
-      -- and literals.
-
-      spitComments (addDecoration cmode Before l <$> annPriorComments)
-
-      -- Comments inside 'annsDP' marked with 'AnnComment' are trickier,
-      -- they seem to contain everything that goes after the thing they
-      -- are attached to and in some cases (e.g. for modules) they contain
-      -- comments that go before things. Exact location can only be
-      -- deduced by analyzing the associated span.
-
-      spitComments before
-      m
-      spitComments after
-
-      -- I wasn't able to find any case when 'annFollowingComments' is
-      -- populated, so we'll ignore that one for now and fix it when we
-      -- have an example of source code where it matters.
+  bool sitcc id relaxed $ do
+    let withRealLocated (L l a) g =
+          case l of
+            UnhelpfulSpan _ -> return ()
+            RealSrcSpan l' -> g (L l' a)
+    withRealLocated loc spitPrecedingComments
+    let setEnclosingSpan =
+          case getSpan loc of
+            UnhelpfulSpan _ -> id
+            RealSrcSpan orf ->
+              if isModule (unL loc)
+                then id
+                else withEnclosingSpan orf
+    setEnclosingSpan $ case ml of
+       Nothing -> f (unL loc)
+       Just l' -> switchLayout l' (f (unL loc))
+    withRealLocated loc spitFollowingComments
 
 -- | A version of 'located' with arguments flipped.
 

src/Ormolu/Printer/Comments.hs

@@ -1,186 +1,205 @@
--- | Helpers for formatting of comments.
-
 {-# LANGUAGE LambdaCase        #-}
 {-# LANGUAGE OverloadedStrings #-}
 
+-- | Helpers for formatting of comments.
+
 module Ormolu.Printer.Comments
-  ( -- * Types
-    Decoration (..)
-  , Decorator (..)
-  , Position (..)
-  , CommentMode (..)
-    -- * Functions for working with comments
-  , spitComments
-  , partitionDPs
-  , addDecoration
+  ( spitPrecedingComments
+  , spitFollowingComments
+  , spitRemainingComments
   )
 where
 
-import ApiAnnotation (AnnKeywordId (AnnModule))
 import Control.Monad
-import Data.Bifunctor
-import Data.Maybe (mapMaybe)
-import Language.Haskell.GHC.ExactPrint.Types
-import Ormolu.Comments
+import Data.Coerce (coerce)
+import Data.Data (Data)
+import Data.Maybe (isJust)
+import Ormolu.CommentStream
 import Ormolu.Printer.Internal
+import Ormolu.Utils (isModule, getSpan)
 import SrcLoc
+import qualified Data.List.NonEmpty as NE
 import qualified Data.Text as T
 
--- | Placement instructions for comments.
+----------------------------------------------------------------------------
+-- Top-level
 
-data Decoration
-  = Decoration Decorator Decorator
-  deriving (Eq, Show)
+-- | Output all preceding comments for an element at given location.
 
--- | Decorator in this context is a thing to put before\/after a comment.
+spitPrecedingComments
+  :: Data a
+  => RealLocated a              -- ^ AST element to attach comments to
+  -> R ()
+spitPrecedingComments = handleCommentSeries . spitPrecedingComment
 
-data Decorator
-  = NoDec                       -- ^ Output nothing
-  | SpaceDec                    -- ^ Output single space
-  | NewlineDec                  -- ^ Output single newline
-  deriving (Eq, Show)
+-- | Output all comments following an element at given location.
 
--- | Position: before vs after.
+spitFollowingComments
+  :: Data a
+  => RealLocated a              -- ^ AST element of attach comments to
+  -> R ()
+spitFollowingComments ref = do
+  trimSpanStream (getSpan ref)
+  handleCommentSeries (spitFollowingComment ref)
 
-data Position
-  = Before                      -- ^ Before
-  | After                       -- ^ After
-  deriving (Eq, Show)
+-- | Output all remaining comments in the comment stream.
 
--- | For which type of AST leaf we're preparing the comments.
+spitRemainingComments :: R ()
+spitRemainingComments = handleCommentSeries spitRemainingComment
 
-data CommentMode
-  = Module                      -- ^ Module
-  | Other                       -- ^ Other element
-  deriving (Eq, Show)
+----------------------------------------------------------------------------
+-- Single-comment functions
 
--- | Output a bunch of 'Comment's. 'DeltaPos'es are used to insert extra
--- space between the comments when necessary.
+-- | Output a single preceding comment for an element at given location.
 
-spitComments :: [(Comment, Decoration)] -> R ()
-spitComments = mapM_ $ \(comment, (Decoration d0 d1)) -> do
-  let spitDecorator = \case
-        NoDec -> return ()
-        SpaceDec -> spit " "
-        NewlineDec -> newline
-  spitDecorator d0
-  spitComment comment
-  spitDecorator d1
+spitPrecedingComment
+  :: Data a
+  => RealLocated a              -- ^ AST element to attach comments to
+  -> Maybe RealSrcSpan          -- ^ Location of last comment in the series
+  -> R (Maybe RealSrcSpan)      -- ^ Location of this comment
+spitPrecedingComment (L ref a) mlastSpn = do
+  let p (L l _) = realSrcSpanEnd l <= realSrcSpanStart ref
+  withPoppedComment p $ \l comment -> do
+    when (needsNewlineBefore l mlastSpn) newline
+    spitComment comment
+    if theSameLine l ref && not (isModule a)
+      then spit " "
+      else newline
 
--- | Output a 'Comment'.
+-- | Output a comment that follows element at given location immediately on
+-- the same line, if there is any.
+
+spitFollowingComment
+  :: Data a
+  => RealLocated a              -- ^ AST element to attach comments to
+  -> Maybe RealSrcSpan          -- ^ Location of last comment in the series
+  -> R (Maybe RealSrcSpan)      -- ^ Location of this comment
+spitFollowingComment (L ref a) mlastSpn = do
+  mnSpn <- nextEltSpan
+  meSpn <- getEnclosingSpan
+  i <- getIndent
+  withPoppedComment (commentFollowsElt ref mnSpn meSpn mlastSpn) $ \l comment ->
+    if theSameLine l ref && not (isModule a)
+      then modNextline $ \m -> setIndent i $ do
+        spit " "
+        spitComment comment
+        m
+      else modNextline $ \m -> setIndent i $ do
+        m
+        when (needsNewlineBefore l mlastSpn) newline
+        spitComment comment
+        newline
+
+-- | Output a single remaining comment from the comment stream.
+
+spitRemainingComment
+  :: Maybe RealSrcSpan          -- ^ Location of last comment in the series
+  -> R (Maybe RealSrcSpan)      -- ^ Location of this comment
+spitRemainingComment mlastSpn =
+  withPoppedComment (const True) $ \l comment -> do
+    when (needsNewlineBefore l mlastSpn) newline
+    spitComment comment
+    newline
+
+----------------------------------------------------------------------------
+-- Helpers
+
+-- | Output series of comments.
+
+handleCommentSeries
+  :: (Maybe RealSrcSpan -> R (Maybe RealSrcSpan))
+     -- ^ Given location of previous comment, output the next comment
+     -- returning its location, or 'Nothing' if we are done
+  -> R ()
+handleCommentSeries f = go Nothing
+  where
+    go mlastSpn = do
+      r <- f mlastSpn
+      case r of
+        Nothing -> return ()
+        Just spn -> go (Just spn)
+
+-- | Try to pop a comment using given predicate and if there is a comment
+-- matching the predicate, print it out.
+
+withPoppedComment
+  :: (RealLocated Comment -> Bool) -- ^ Comment predicate
+  -> (RealSrcSpan -> Comment -> R ()) -- ^ Priting function
+  -> R (Maybe RealSrcSpan)
+withPoppedComment p f = do
+  r <- popComment p
+  case r of
+    Nothing -> return Nothing
+    Just (L l comment) -> Just l <$ f l comment
+
+-- | Determine if we need to insert a newline between current comment and
+-- last printed comment.
+
+needsNewlineBefore
+  :: RealSrcSpan                -- ^ Current comment span
+  -> Maybe RealSrcSpan          -- ^ Last printed comment span
+  -> Bool
+needsNewlineBefore l mlastSpn =
+  case mlastSpn of
+    Nothing -> False
+    Just lastSpn ->
+      srcSpanStartLine l > srcSpanEndLine lastSpn + 1
+
+-- | Is the comment and AST element are on the same line?
+
+theSameLine
+  :: RealSrcSpan                -- ^ Current comment span
+  -> RealSrcSpan                -- ^ AST element location
+  -> Bool
+theSameLine l ref =
+  srcSpanEndLine l == srcSpanStartLine ref
+
+-- | Determine if given comment follows AST element.
+
+commentFollowsElt
+  :: RealSrcSpan                -- ^ Location of AST element
+  -> Maybe RealSrcSpan          -- ^ Location of next AST element
+  -> Maybe RealSrcSpan          -- ^ Location of enclosing AST element
+  -> Maybe RealSrcSpan          -- ^ Location of last comment in the series
+  -> RealLocated Comment        -- ^ Comment to test
+  -> Bool
+commentFollowsElt ref mnSpn meSpn mlastSpn (L l comment) =
+  -- A comment follows a AST element if all 4 conditions are satisfied:
+  goesAfter && logicallyFollows && noEltBetween && supersedesParentElt
+  where
+    -- 1) The comment starts after end of the AST element:
+    goesAfter =
+      realSrcSpanStart l >= realSrcSpanEnd ref
+    -- 2) The comment logically belongs to the element, three cases:
+    logicallyFollows
+      = theSameLine l ref -- a) it's on the same line
+      || isPrevHaddock comment -- b) it's a Haddock string starting with -- ^
+      || isJust mlastSpn -- c) it's a continuation of a comment block
+    -- 3) There is no other AST element between this element and the comment:
+    noEltBetween =
+      case mnSpn of
+        Nothing -> True
+        Just nspn ->
+          realSrcSpanStart nspn >= realSrcSpanEnd l
+    -- Less obvious: if column of comment is closer to the start of
+    -- enclosing element, it probably related to that parent element, not to
+    -- the current child element. This rule is important because otherwise
+    -- all comments would end up assigned to closest inner elements, and
+    -- parent elements won't have a chance to get any comments assigned to
+    -- them. This is not OK because comments will get indented according to
+    -- the AST elements they are attached to.
+    supersedesParentElt =
+      case meSpn of
+        Nothing -> True
+        Just espn ->
+          let startColumn = srcLocCol . realSrcSpanStart
+          in abs (startColumn espn - startColumn l)
+               > abs (startColumn ref - startColumn l)
+
+-- | Output a 'Comment'. This is a low-level printing function.
 
 spitComment :: Comment -> R ()
-spitComment (Comment str _ _) =
-  forM_ (normalizeComment str) $ \x -> do
-    ensureIndent
-    spit (T.pack x)
-
--- | Partition annotations to get a collection of 'Comment's preceding a
--- definition and following it. Every 'Comment' has corresponding
--- 'Decoration' which is used to understand how to decorate it.
-
-partitionDPs
-  :: CommentMode           -- ^ For which type of element we prepare comments
-  -> SrcSpan               -- ^ Span of element the comments are attached to
-  -> [(KeywordId, DeltaPos)]    -- ^ Annotations
-  -> ([(Comment, Decoration)], [(Comment, Decoration)])
-partitionDPs cmode refSpan anns =
-  case cmode of
-    Module -> partitionDPsModule refSpan anns
-    Other -> partitionDPsOther refSpan anns
-
--- | Try to partition comments as if for a module.
-
-partitionDPsModule
-  :: SrcSpan               -- ^ Span of element the comments are attached to
-  -> [(KeywordId, DeltaPos)]    -- ^ Annotations
-  -> ([(Comment, Decoration)], [(Comment, Decoration)])
-partitionDPsModule refSpan xs
-  = bimap (takeComments Before) (takeComments After) $
-  -- NOTE If there is no annotation corresponding to module keyword, then
-  -- the module doesn't have a header and all comments shoud go after.
-  if G AnnModule `elem` fmap fst xs
-    then break ((== G AnnModule) . fst) xs
-    else ([], xs)
+spitComment =
+  sequence_ . NE.intersperse newline . fmap f . coerce
   where
-    takeComments pos = mapMaybe $ \(keywordId, dpos) -> do
-      c <- annComment keywordId
-      return (c, getDecoration Module pos refSpan (c, dpos))
-
--- | Partition comments according to their spans (works for everything but
--- modules).
-
-partitionDPsOther
-  :: SrcSpan               -- ^ Span of element the comments are attached to
-  -> [(KeywordId, DeltaPos)]    -- ^ Annotations
-  -> ([(Comment, Decoration)], [(Comment, Decoration)])
-partitionDPsOther refSpan
-  = bimap (fmap (addDecoration Other Before refSpan))
-          (fixupLastDec . fmap (addDecoration Other After refSpan))
-  . break (followedBySpan refSpan  . commentIdentifier . fst)
-  . mapMaybe annComment'
-  where
-    annComment' :: (KeywordId, DeltaPos) -> Maybe (Comment, DeltaPos)
-    annComment' (keywordId, dpos) = do
-      c <- annComment keywordId
-      return (c, dpos)
-    followedBySpan :: SrcSpan -> SrcSpan -> Bool
-    followedBySpan spn0 spn1 =
-      if srcSpanEnd spn0 < srcSpanStart spn1
-        then True
-        else False
-
--- | Last following comment cannot be standalone because in that case we get
--- redundant newlines.
-
-fixupLastDec :: [(Comment, Decoration)] -> [(Comment, Decoration)]
-fixupLastDec [] = []
-fixupLastDec [(c, Decoration d0 _)] = [(c, Decoration d0 NoDec)]
-fixupLastDec (c:cs) = c : fixupLastDec cs
-
--- | Replace 'DeltaPos' with 'Decoration'.
-
-addDecoration
-  :: CommentMode           -- ^ For which type of element we prepare comments
-  -> Position              -- ^ Is this for comments before or after?
-  -> SrcSpan               -- ^ Span of element the comments are attached to
-  -> (Comment, DeltaPos)   -- ^ Thing to decorate
-  -> (Comment, Decoration)
-addDecoration cmode pos refSpan (comment, dpos) =
-  ( comment
-  , getDecoration cmode pos refSpan (comment, dpos)
-  )
-
--- | Calculate decoration for a comment.
-
-getDecoration
-  :: CommentMode           -- ^ For which type of element we prepare comments
-  -> Position              -- ^ Is this for comment before or after?
-  -> SrcSpan               -- ^ Span of element the comments are attached to
-  -> (Comment, DeltaPos)   -- ^ Thing to decorate
-  -> Decoration
-getDecoration cmode pos refSpan (c, (DP (r, _))) =
-  Decoration preceedingDec followingDec
-  where
-    preceedingDec =
-      if sameLine
-        then case pos of
-               Before -> NoDec
-               After -> SpaceDec
-        else if r > 1
-               then NewlineDec
-               else NoDec
-    followingDec =
-      if sameLine
-        then case pos of
-               Before -> SpaceDec
-               After -> NewlineDec
-        else NewlineDec
-    sameLine =
-      case cmode of
-        Module -> False
-        Other ->
-          case (refSpan, commentIdentifier c) of
-            (RealSrcSpan spn0, RealSrcSpan spn1) ->
-              srcSpanEndLine spn0 == srcSpanStartLine spn1
-            (_, _) -> False
+    f x = ensureIndent >> spit (T.pack x)

src/Ormolu/Printer/Internal.hs

@@ -1,5 +1,6 @@
 {-# LANGUAGE GeneralizedNewtypeDeriving #-}
 {-# LANGUAGE OverloadedStrings          #-}
+{-# LANGUAGE RankNTypes                 #-}
 
 -- | In most cases import "Ormolu.Printer.Combinators" instead, these
 -- functions are the low-level building blocks and should not be used on
@@ -13,6 +14,7 @@ module Ormolu.Printer.Internal
     -- * Internal functions
   , spit
   , newline
+  , modNextline
   , ensureIndent
   , inci
   , sitcc
@@ -21,21 +23,31 @@ module Ormolu.Printer.Internal
   , enterLayout
   , vlayout
   , relaxedComments
-  , lookupAnn
+    -- * Special helpers for comment placement
+  , trimSpanStream
+  , nextEltSpan
+  , popComment
+  , hasMoreComments
+  , getIndent
+  , setIndent
+  , getEnclosingSpan
+  , withEnclosingSpan
   )
 where
 
 import Control.Monad.Reader
 import Control.Monad.State.Strict
-import Data.Data (Data)
+import Data.Coerce
+import Data.Maybe (listToMaybe)
 import Data.Text (Text)
 import Data.Text.Lazy.Builder
 import Debug.Trace
-import Language.Haskell.GHC.ExactPrint.Types
+import GHC
+import Ormolu.CommentStream
+import Ormolu.SpanStream
 import SrcLoc
-import qualified Data.Text.Lazy as TL
-import qualified Data.Map  as M
 import qualified Data.Text as T
+import qualified Data.Text.Lazy as TL
 
 ----------------------------------------------------------------------------
 -- The 'R' monad
@@ -56,10 +68,10 @@ data RC = RC
     -- ^ Current layout
   , rcRelaxedComments :: Bool
     -- ^ Whether to relax aligning rules for comments
-  , rcAnns :: Anns
-    -- ^ The collection of annotations obtained after parsing
   , rcDebug :: Bool
     -- ^ Whether to print debugging info as we go
+  , rcEnclosingSpan :: Maybe RealSrcSpan
+    -- ^ Span of enclosing element of AST
   }
 
 -- | State context of 'R'.
@@ -69,6 +81,12 @@ data SC = SC
     -- ^ Index of the next column to render
   , scBuilder :: Builder
     -- ^ Rendered source code so far
+  , scSpanStream :: SpanStream
+    -- ^ Span stream
+  , scCommentStream :: CommentStream
+    -- ^ Comment stream
+  , scNewline :: R ()
+    -- ^ What to render as newline
   }
 
 -- | 'Layout' options.
@@ -83,21 +101,25 @@ data Layout
 runR
   :: Bool                       -- ^ Whether to print debugging info
   -> R ()                       -- ^ Monad to run
-  -> Anns                       -- ^ Annotations to use
+  -> SpanStream                 -- ^ Span stream
+  -> CommentStream              -- ^ Comment stream
   -> Text                       -- ^ Resulting rendition
-runR debug (R m) anns =
+runR debug (R m) sstream cstream =
   TL.toStrict . toLazyText . scBuilder $ execState (runReaderT m rc) sc
   where
     rc = RC
       { rcIndent = 0
       , rcLayout = MultiLine
       , rcRelaxedComments = False
-      , rcAnns = anns
       , rcDebug = debug
+      , rcEnclosingSpan = Nothing
       }
     sc = SC
       { scColumn = 0
       , scBuilder = mempty
+      , scSpanStream = sstream
+      , scCommentStream = cstream
+      , scNewline = newlineRaw
       }
 
 ----------------------------------------------------------------------------
@@ -109,26 +131,49 @@ runR debug (R m) anns =
 spit :: Text -> R ()
 spit x = do
   traceR "spit_before" (Just x)
-  R (modify modSC)
+  R . modify $ \sc -> sc
+    { scBuilder = scBuilder sc <> fromText x
+    , scColumn = scColumn sc + T.length x
+    }
   traceR "spit_after" Nothing
-  where
-    modSC sc = sc
-      { scBuilder = scBuilder sc <> fromText x
-      , scColumn = scColumn sc + T.length x
-      }
 
--- | Output a newline.
+-- | Output a newline. The 'modNewline' function can be used to alter what
+-- will be inserted. This is used to output comments following an element of
+-- AST because we cannot output comments immediately, e.g. because we need
+-- to close parentheses first, etc.
+--
+-- 'newline' auto-resets its modifications so the changes introduced with
+-- 'modNewline' only have effect once.
 
 newline :: R ()
 newline = do
+  n <- R (gets scNewline)
+  R . modify $ \sc -> sc
+    { scNewline = newlineRaw
+    }
+  n
+
+-- | Low-level newline primitive. This one always just inserts a newline, no
+-- hooks can be attached.
+
+newlineRaw :: R ()
+newlineRaw = do
   traceR "newline_before" (Just "\n")
-  R (modify modSC)
+  R . modify $ \sc -> sc
+    { scBuilder = scBuilder sc <> "\n"
+    , scColumn = 0
+    }
   traceR "newline_after" Nothing
-  where
-    modSC sc = sc
-      { scBuilder = scBuilder sc <> "\n"
-      , scColumn = 0
-      }
+
+-- | Modify how next newline will be output. The argument of call-back is
+-- the version of 'newline' built so far.
+
+modNextline :: (R () -> R ()) -> R ()
+modNextline f = R $ do
+  old <- gets scNewline
+  modify $ \sc -> sc
+    { scNewline = f old
+    }
 
 -- | Ensure that indentation level is satisfied. Insert correct number of
 -- spaces if it isn't.
@@ -148,12 +193,11 @@ inci :: R () -> R ()
 inci m' = do
   traceR "inci_before" Nothing
   let R m = traceR "inci_inside" Nothing >> m'
+      modRC rc = rc
+        { rcIndent = rcIndent rc + indentStep
+        }
   R (local modRC m)
   traceR "inci_ended" Nothing
-  where
-    modRC x = x
-      { rcIndent = rcIndent x + indentStep
-      }
 
 -- | Set indentation level for the inner computation equal to current
 -- column. This makes sure that the entire inner block is uniformly
@@ -165,10 +209,10 @@ sitcc m' = do
   traceR "sitcc_before" Nothing
   i <- R (asks rcIndent)
   c <- R (gets scColumn)
-  let modRC x = x
+  let R m = traceR "sitcc_inside" Nothing >> m'
+      modRC rc = rc
         { rcIndent = max i c
         }
-      R m = traceR "sitcc_inside" Nothing >> m'
   vlayout m' (R (local modRC m))
   traceR "sitcc_ended" Nothing
 
@@ -179,7 +223,7 @@ sitcc m' = do
 relaxComments :: R () -> R ()
 relaxComments (R m) = do
   traceR "relax_start" Nothing
-  let modRC x = x
+  let modRC rc = rc
         { rcRelaxedComments = True
         }
   R (local modRC m)
@@ -194,7 +238,7 @@ enterLayout l (R m) = do
           SingleLine -> "single_line"
           MultiLine -> "multi_line"
   traceR ("lstart_" ++ label) Nothing
-  let modRC x = x
+  let modRC rc = rc
         { rcLayout = l
         }
   R (local modRC m)
@@ -207,25 +251,91 @@ vlayout
   -> R ()                       -- ^ Multi line
   -> R ()
 vlayout sline mline = do
-  l <- currentLayout
+  l <- R (asks rcLayout)
   case l of
     SingleLine -> sline
     MultiLine -> mline
 
--- | Return current layout.
-
-currentLayout :: R Layout
-currentLayout = R (asks rcLayout)
-
 -- | Check whether we're in a region with relaxed comments placement.
 
 relaxedComments :: R Bool
 relaxedComments = R (asks rcRelaxedComments)
 
--- | Lookup an annotation.
+----------------------------------------------------------------------------
+-- Special helpers for comment placement
 
-lookupAnn :: Data a => Located a -> R (Maybe Annotation)
-lookupAnn l = M.lookup (mkAnnKey l) <$> R (asks rcAnns)
+-- | Drop elements that begin before or at the same place as given
+-- 'SrcSpan'.
+
+trimSpanStream
+  :: RealSrcSpan                -- ^ Reference span
+  -> R ()
+trimSpanStream ref = do
+  let leRef :: RealSrcSpan -> Bool
+      leRef x = realSrcSpanStart x <= realSrcSpanStart ref
+  R . modify $ \sc -> sc
+    { scSpanStream = coerce (dropWhile leRef) (scSpanStream sc)
+    }
+
+-- | Get location of next element in AST.
+
+nextEltSpan :: R (Maybe RealSrcSpan)
+nextEltSpan = listToMaybe . coerce <$> R (gets scSpanStream)
+
+-- | Pop a 'Comment' from the 'CommentStream' if given predicate is
+-- satisfied and there are comments in the stream.
+
+popComment
+  :: (RealLocated Comment -> Bool)
+  -> R (Maybe (RealLocated Comment))
+popComment f = R $ do
+  CommentStream cstream <- gets scCommentStream
+  case cstream of
+    [] -> return Nothing
+    (x:xs) ->
+      if f x
+        then Just x <$ modify (\sc -> sc
+               { scCommentStream = CommentStream xs
+               })
+        else return Nothing
+
+-- | Return 'True' if there are more comments in the 'CommentStream'.
+
+hasMoreComments :: R Bool
+hasMoreComments = R $ do
+  CommentStream cstream <- gets scCommentStream
+  (return . not . null) cstream
+
+-- | Current indentation level.
+
+getIndent :: R Int
+getIndent = R (asks rcIndent)
+
+-- | Set indentation level.
+
+setIndent :: Int -> R () -> R ()
+setIndent i m' = do
+  traceR "set_indent_before" Nothing
+  let R m = traceR "set_indent_inside" Nothing >> m'
+      modRC rc = rc
+        { rcIndent = i
+        }
+  R (local modRC m)
+  traceR "set_indent_after" Nothing
+
+-- | Get 'RealSrcSpan' of enclosing span, if any.
+
+getEnclosingSpan :: R (Maybe RealSrcSpan)
+getEnclosingSpan = R (asks rcEnclosingSpan)
+
+-- | Set 'RealSrcSpan' of enclosing span for the given computation.
+
+withEnclosingSpan :: RealSrcSpan -> R () -> R ()
+withEnclosingSpan spn (R m) = do
+  let modRC rc = rc
+        { rcEnclosingSpan = Just spn
+        }
+  R (local modRC m)
 
 ----------------------------------------------------------------------------
 -- Debug helpers

src/Ormolu/Printer/Meat/Common.hs

@@ -10,23 +10,17 @@ module Ormolu.Printer.Meat.Common
   , p_rdrName'
   , p_qualName
   , p_ieWildcard
-    -- * Helpers
   , opParens
-  , combineSrcSpans'
-  , getSpan
-  , unL
   )
 where
 
 import Data.Char (isAlphaNum)
-import Data.List.NonEmpty (NonEmpty (..))
 import GHC hiding (GhcPs, IE)
 import Module (Module (..))
 import OccName (OccName (..))
 import Ormolu.Printer.Combinators
 import Outputable (Outputable (..), showSDocUnsafe)
 import RdrName (RdrName (..), rdrNameOcc)
-import SrcLoc (combineSrcSpans)
 
 p_hsmodName :: ModuleName -> R ()
 p_hsmodName mname = do
@@ -64,9 +58,6 @@ p_ieWildcard = \case
   NoIEWildcard -> return ()
   IEWildcard n -> parens (atom n)
 
-----------------------------------------------------------------------------
--- Helpers
-
 -- | Put parentheses around the second argument if the 'Outputable' thing
 -- consists only of punctuation characters.
 
@@ -75,18 +66,3 @@ opParens x m =
   if all (not . isAlphaNum) (showSDocUnsafe (ppr x))
     then txt "(" >> m >> txt ")"
     else m
-
--- | Combine all source spans from the given list.
-
-combineSrcSpans' :: NonEmpty SrcSpan -> SrcSpan
-combineSrcSpans' (x:|xs) = foldr combineSrcSpans x xs
-
--- | Get source span from a 'Located' thing.
-
-getSpan :: Located e -> SrcSpan
-getSpan (L spn _) = spn
-
--- | Exact inner value from 'Located'.
-
-unL :: Located e -> e
-unL (L _ e) = e

src/Ormolu/Printer/Meat/Declaration.hs

@@ -8,8 +8,7 @@ module Ormolu.Printer.Meat.Declaration
   )
 where
 
-import HsDecls
-import Language.Haskell.GHC.ExactPrint.Types
+import GHC
 import Ormolu.Printer.Combinators
 import Ormolu.Printer.Meat.Declaration.Data
 import Ormolu.Printer.Meat.Declaration.Type

src/Ormolu/Printer/Meat/Declaration/Data.hs

@@ -2,6 +2,8 @@
 {-# LANGUAGE OverloadedStrings #-}
 {-# LANGUAGE RecordWildCards   #-}
 
+-- | Renedring of data type declarations.
+
 module Ormolu.Printer.Meat.Declaration.Data
   ( p_dataDecl
   )
@@ -9,13 +11,13 @@ where
 
 import BasicTypes (DerivStrategy (..))
 import Control.Monad
-import Data.List.NonEmpty (nonEmpty)
+import Data.List.NonEmpty (NonEmpty (..))
 import Data.Maybe (isJust)
-import GHC hiding (GhcPs, IE)
-import Language.Haskell.GHC.ExactPrint.Types
+import GHC
 import Ormolu.Printer.Combinators
 import Ormolu.Printer.Meat.Common
 import Ormolu.Printer.Meat.Type
+import Ormolu.Utils (unL, getSpan, combineSrcSpans')
 import RdrName (RdrName (..))
 import SrcLoc (Located)
 
@@ -39,20 +41,17 @@ p_dataDecl name tvars HsDataDefn {..} = do
       txt ":: "
       relaxComments (located k p_hsType)
   let gadt = isJust dd_kindSig || any (isGadt . unL) dd_cons
-  case nonEmpty dd_cons of
-    Nothing -> pure ()
-    Just dd_cons_ne ->
-      if gadt
-        then do
-          txt " where"
-          newline
-          inci $ newlineSep (located' p_conDecl) dd_cons
-        else switchLayout (combineSrcSpans' (getSpan <$> dd_cons_ne)) $ do
-          breakpoint
-          inci $ do
-            txt "= "
-            let sep = vlayout (txt " | ") (txt "| ")
-            velt $ withSep sep (located' p_conDecl) dd_cons
+  if gadt
+    then do
+      txt " where"
+      newline
+      inci $ newlineSep (located' p_conDecl) dd_cons
+    else switchLayout (combineSrcSpans' (getSpan name :| (getSpan <$> dd_cons))) $ do
+      breakpoint
+      inci $ do
+        txt "= "
+        let sep = vlayout (txt " | ") (txt "| ")
+        velt $ withSep sep (located' p_conDecl) dd_cons
   newline
   inci . located dd_derivs $ \xs ->
     forM_ xs (line . located' p_hsDerivingClause)
@@ -60,7 +59,7 @@ p_dataDecl name tvars HsDataDefn {..} = do
 p_conDecl :: ConDecl GhcPs -> R ()
 p_conDecl = \case
   ConDeclGADT {..} -> velt'
-    [ spaceSep (located' p_rdrName') con_names
+    [ spaceSep (located' p_rdrName) con_names
     , inci $ do
         txt ":: "
         relaxComments (locatedVia Nothing (hsib_body con_type) p_hsType)
@@ -83,17 +82,17 @@ p_conDecl = \case
           txt "=> "
     case con_details of
       PrefixCon xs -> do
-        located con_name p_rdrName'
+        located con_name p_rdrName
         unless (null xs) breakpoint
         inci $ velt' (located' p_hsType <$> xs)
       RecCon l -> do
-        located con_name p_rdrName'
+        located con_name p_rdrName
         breakpoint
         inci $ located l p_conDeclFields
       InfixCon x y -> velt'
         [ located x p_hsType
         , inci $ velt'
-          [ backticks (located con_name p_rdrName')
+          [ backticks (located con_name p_rdrName)
           , inci $ located y p_hsType
           ]
         ]

src/Ormolu/Printer/Meat/Declaration/Type.hs

@@ -9,8 +9,7 @@ module Ormolu.Printer.Meat.Declaration.Type
 where
 
 import Control.Monad
-import HsTypes
-import Language.Haskell.GHC.ExactPrint.Types
+import GHC
 import Ormolu.Printer.Combinators
 import Ormolu.Printer.Meat.Common
 import Ormolu.Printer.Meat.Type

src/Ormolu/Printer/Meat/Declaration/TypeFamily.hs

@@ -11,9 +11,7 @@ where
 
 import Control.Monad
 import Data.Maybe (maybeToList, isJust)
-import HsDecls
-import HsTypes
-import Language.Haskell.GHC.ExactPrint.Types
+import GHC
 import Ormolu.Printer.Combinators
 import Ormolu.Printer.Meat.Common
 import Ormolu.Printer.Meat.Type

src/Ormolu/Printer/Meat/ImportExport.hs

@@ -11,9 +11,8 @@ module Ormolu.Printer.Meat.ImportExport
 where
 
 import Control.Monad
-import GHC hiding (GhcPs, IE)
+import GHC
 import HsImpExp (IE (..))
-import Language.Haskell.GHC.ExactPrint.Types
 import Ormolu.Printer.Combinators
 import Ormolu.Printer.Meat.Common
 

src/Ormolu/Printer/Meat/Module.hs

@@ -9,9 +9,11 @@ module Ormolu.Printer.Meat.Module
 where
 
 import Control.Monad
+import Data.Maybe (isJust)
 import GHC hiding (GhcPs, IE)
 import Ormolu.Imports
 import Ormolu.Printer.Combinators
+import Ormolu.Printer.Comments
 import Ormolu.Printer.Meat.Common
 import Ormolu.Printer.Meat.Declaration
 import Ormolu.Printer.Meat.ImportExport
@@ -38,7 +40,10 @@ p_hsModule loc@(L moduleSpan hsModule) = do
             inci (locatedVia Nothing hsmodExports' p_hsmodExports)
         breakpoint
         txt "where"
-        unless (null hsmodImports) newline
+        when (not (null hsmodImports) || not (null hsmodDecls)) newline
     forM_ (sortImports hsmodImports) (located' p_hsmodImport)
     when (not (null hsmodImports) && not (null hsmodDecls)) newline
     newlineSep (located' p_hsDecl) hsmodDecls
+    trailingComments <- hasMoreComments
+    when (trailingComments && isJust hsmodName) newline
+    spitRemainingComments

src/Ormolu/Printer/Meat/Type.hs

@@ -12,8 +12,7 @@ module Ormolu.Printer.Meat.Type
   )
 where
 
-import GHC hiding (GhcPs, IE)
-import Language.Haskell.GHC.ExactPrint.Types
+import GHC
 import Ormolu.Printer.Combinators
 import Ormolu.Printer.Meat.Common
 

src/Ormolu/SpanStream.hs

@@ -0,0 +1,47 @@
+{-# LANGUAGE DeriveDataTypeable         #-}
+{-# LANGUAGE GeneralizedNewtypeDeriving #-}
+
+-- | Build span stream of AST elements.
+
+module Ormolu.SpanStream
+  ( SpanStream (..)
+  , mkSpanStream
+  )
+where
+
+import Data.DList (DList)
+import Data.Data (Data)
+import Data.Generics (everything, ext2Q)
+import Data.List (sortOn)
+import Data.Typeable (cast)
+import SrcLoc
+import qualified Data.DList as D
+
+-- | A stream of 'RealSrcSpan's in ascending order. This allows us to tell
+-- e.g. whether there is another \"located\" element of AST between current
+-- element and comment we're considering for printing.
+
+newtype SpanStream = SpanStream [RealSrcSpan]
+  deriving (Eq, Show, Data, Semigroup, Monoid)
+
+-- | Create 'SpanStream' from a data structure containing 'RealSrcSpan's.
+
+mkSpanStream
+  :: Data a
+  => a                          -- ^ Data structure to inspect (AST)
+  -> SpanStream
+mkSpanStream a
+  = SpanStream
+  . sortOn realSrcSpanStart
+  . D.toList
+  $ everything mappend (const mempty `ext2Q` queryLocated) a
+  where
+    queryLocated
+      :: (Data e0, Data e1)
+      => GenLocated e0 e1
+      -> DList RealSrcSpan
+    queryLocated (L mspn _) =
+      case cast mspn :: Maybe SrcSpan of
+        Nothing -> mempty
+        Just (UnhelpfulSpan _) -> mempty
+        Just (RealSrcSpan spn) -> D.singleton spn

src/Ormolu/Utils.hs

@@ -0,0 +1,31 @@
+module Ormolu.Utils
+  ( combineSrcSpans'
+  , isModule
+  , unL
+  , getSpan
+  )
+where
+
+import Data.Data (Data, showConstr, toConstr)
+import Data.List.NonEmpty (NonEmpty (..))
+import SrcLoc
+
+-- | Combine all source spans from the given list.
+
+combineSrcSpans' :: NonEmpty SrcSpan -> SrcSpan
+combineSrcSpans' (x:|xs) = foldr combineSrcSpans x xs
+
+-- | Return 'True' if given element of AST is module.
+
+isModule :: Data a => a -> Bool
+isModule x = showConstr (toConstr x) == "HsModule"
+
+-- | Exact inner value from 'Located'.
+
+unL :: Located e -> e
+unL (L _ e) = e
+
+-- | Get source span from a 'Located' thing.
+
+getSpan :: GenLocated l e -> l
+getSpan (L spn _) = spn

tests/Ormolu/Printer/CombinatorsSpec.hs

@@ -5,7 +5,6 @@ module Ormolu.Printer.CombinatorsSpec (spec) where
 import Ormolu.Printer.Combinators
 import Ormolu.Printer.Internal
 import Test.Hspec
-import qualified Data.Map.Strict as M
 import qualified Data.Text.IO as T
 
 -- NOTE Testing combinators in separation is easy, but it's not very useful
@@ -97,7 +96,7 @@ rModuleHeader = do
 shouldRender :: R () -> FilePath -> Expectation
 shouldRender m path = do
   expected <- T.readFile path
-  runR False m M.empty `shouldBe` expected
+  runR False m mempty mempty `shouldBe` expected
 
 -- | Render using single-line layout.