Implement handling of comments and printing of module headers

app/Main.hs

@@ -4,16 +4,17 @@ import Control.Monad
 import Ormolu.Parser
 import Ormolu.Printer
 import System.Environment (getArgs)
+import qualified Data.Text.IO as TIO
 import qualified Outputable as GHC
 
 main :: IO ()
 main = do
   (path:_) <- getArgs
   input <- readFile path
+
   (ws, r) <- parseModule [] path input
   unless (null ws) $
     putStrLn "dynamic option warnings:"
-    -- TODO print ws
   case r of
     Left (srcSpan, err) -> do
       putStrLn (showOutputable srcSpan)
@@ -23,8 +24,7 @@ main = do
       putStrLn (showOutputable anns)
       putStrLn "\nparsed module:\n"
       putStrLn (showOutputable parsedModule)
-
-      printModule anns parsedModule
+      TIO.putStr (printModule anns parsedModule)
 
 showOutputable :: GHC.Outputable o => o -> String
 showOutputable = GHC.showSDocUnsafe . GHC.ppr

data/examples/module-header/double-shebangs-out.hs

@@ -0,0 +1,2 @@
+#!/usr/bin/env stack
+#!/usr/bin/env stack

data/examples/module-header/double-shebangs.hs

@@ -0,0 +1,2 @@
+#!/usr/bin/env stack
+#!/usr/bin/env stack

data/examples/module-header/multiline-out.hs

@@ -0,0 +1,6 @@
+module Foo
+  ( foo
+  , bar
+  , baz
+  )
+where

data/examples/module-header/multiline-with-comments-out.hs

@@ -0,0 +1,21 @@
+-- | Header.
+
+{-# LANGUAGE FlexibleContexts #-}
+{-# LANGUAGE TypeFamilies #-}
+module My.Module
+  ( -- * Something
+    foo
+  , bar
+  , -- * Another thing
+    (<?>)
+  , {- some other thing -} foo2 -- yet another
+  , foo3 -- third one
+  , baz
+  , bar2 -- a multiline comment
+    -- the second line
+  , bar3
+  , module Foo.Bar.Baz
+  )
+where
+
+-- Wow

data/examples/module-header/multiline-with-comments.hs

@@ -0,0 +1,21 @@
+-- | Header.
+
+{-# LANGUAGE FlexibleContexts #-}
+{-# LANGUAGE TypeFamilies     #-}
+
+module My.Module
+  ( -- * Something
+    foo,
+    bar,
+    -- * Another thing
+    (<?>),
+    {- some other thing -} foo2 -- yet another
+    ,foo3 -- third one
+    ,baz,
+    bar2 -- a multiline comment
+         -- the second line
+    ,bar3
+  , module Foo.Bar.Baz )
+where
+
+-- Wow

data/examples/module-header/multiline.hs

@@ -0,0 +1,2 @@
+module Foo (
+  foo, bar, baz) where

data/examples/module-header/multiline2-out.hs

@@ -0,0 +1,6 @@
+module Foo
+  ( foo
+  , bar
+  , baz
+  )
+where

data/examples/module-header/multiline2.hs

@@ -0,0 +1,2 @@
+module Foo
+  (foo, bar, baz) where

data/examples/module-header/simple-out.hs

@@ -0,0 +1 @@
+module Main where

data/examples/module-header/simple-with-comments-out.hs

@@ -0,0 +1,4 @@
+-- | Here we go.
+module Main where
+
+-- Wow.

data/examples/module-header/simple-with-comments.hs

@@ -0,0 +1,5 @@
+-- | Here we go.
+
+module Main where
+
+-- Wow.

data/examples/module-header/simple.hs

@@ -0,0 +1 @@
+module Main where

data/examples/module-header/singleline-out.hs

@@ -0,0 +1 @@
+module Foo (foo, bar, baz) where

data/examples/module-header/singleline.hs

@@ -0,0 +1 @@
+module Foo ( foo, bar, baz ) where

ormolu.cabal

@@ -13,7 +13,8 @@ build-type:           Simple
 description:          A formatter for Haskell source code.
 extra-doc-files:      CHANGELOG.md
                     , README.md
-data-files:           data/printer/*.hs
+data-files:           data/examples/module-header/*.hs
+                    , data/printer/*.hs
 
 source-repository head
   type:               git
@@ -28,17 +29,20 @@ library
   hs-source-dirs:     src
   build-depends:      base           >= 4.8 && < 5.0
                     , containers     >= 0.5 && < 0.7
+                    , exceptions     >= 0.6 && < 0.11
                     , ghc            >= 8.4.3
                     , ghc-boot-th    >= 8.4.3
                     , ghc-exactprint >= 0.5.6
                     , mtl            >= 2.0 && < 3.0
                     , text           >= 0.2 && < 1.3
   exposed-modules:    Ormolu
+                    , Ormolu.Config
+                    , Ormolu.Diff
                     , Ormolu.Parser
                     , Ormolu.Printer
                     , Ormolu.Printer.Combinators
+                    , Ormolu.Printer.Comments
                     , Ormolu.Printer.Internal
-                    , Ormolu.Type
   if flag(dev)
     ghc-options:      -Wall -Werror -Wcompat
                       -Wincomplete-record-updates
@@ -55,12 +59,16 @@ test-suite tests
   type:               exitcode-stdio-1.0
   build-depends:      base           >= 4.8 && < 5.0
                     , containers     >= 0.5 && < 0.7
+                    , filepath       >= 1.2 && < 1.5
                     , hspec          >= 2.0 && < 3.0
                     , ormolu
+                    , path           >= 0.6 && < 0.7
+                    , path-io        >= 1.4.2 && < 2.0
                     , text           >= 0.2 && < 1.3
   build-tools:        hspec-discover >= 2.0 && < 3.0
   other-modules:
                       Ormolu.Printer.CombinatorsSpec
+                    , Ormolu.PrinterSpec
 
   if flag(dev)
     ghc-options:      -Wall -Werror
@@ -74,6 +82,7 @@ executable ormolu
   build-depends:      base           >= 4.8 && < 5.0
                     , ghc            >= 8.4.3
                     , ormolu
+                    , text           >= 0.2 && < 1.3
   if flag(dev)
     ghc-options:      -Wall -Werror -Wcompat
                       -Wincomplete-record-updates

src/Ormolu.hs

@@ -1,3 +1,99 @@
+-- | A formatter for Haskell source code.
+
+{-# LANGUAGE RecordWildCards #-}
+
 module Ormolu
-  (  )
+  ( ormolu
+  , ormoluFile
+  , Config (..)
+  , defaultConfig
+  , DynOption (..)
+  , OrmoluException (..)
+  )
 where
+
+import Control.Exception
+import Control.Monad
+import Control.Monad.Catch (MonadThrow (..))
+import Control.Monad.IO.Class (MonadIO (..))
+import Data.Text (Text)
+import Language.Haskell.GHC.ExactPrint.Types
+import Ormolu.Config
+import Ormolu.Diff
+import Ormolu.Parser
+import Ormolu.Printer
+import qualified Data.Text as T
+import qualified GHC
+
+-- | Format a 'String', return formatted version as 'Text'.
+--
+-- The function
+--
+--     * Takes 'String' because that's what GHC parser accepts.
+--     * Needs 'IO' because some functions from GHC that are necessary to
+--       setup parsing context require 'IO'. There should be no visible
+--       side-effects though.
+--     * Takes file name just to use it in parse error messages.
+--     * Throws 'OrmoluException'.
+
+ormolu
+  :: (MonadIO m, MonadThrow m)
+  => Config                     -- ^ Ormolu configuration
+  -> FilePath                   -- ^ Location of source file
+  -> String                     -- ^ Input to format
+  -> m Text
+ormolu cfg path str = do
+  (anns0, parsedSrc0) <-
+    parseModule' cfg OrmoluParsingFailed path str
+  let txt = printModule anns0 parsedSrc0
+  -- Parse the result of pretty-printing again and make sure that AST is the
+  -- same as AST of original snippet module span positions.
+  (anns1, parsedSrc1) <-
+    parseModule' cfg OrmoluOutputParsingFailed "<rendered>" (T.unpack txt)
+  when (diff (anns0, parsedSrc0) (anns1, parsedSrc1)) $
+    throwM (OrmoluASTDiffers str txt)
+  return txt
+
+-- | Load a file and format it. The file stays intact and the rendered
+-- version is returned as 'Text'.
+--
+-- > ormoluFile cfg path =
+-- >   liftIO (readFile path) >>= ormolu cfg path
+
+ormoluFile
+  :: (MonadIO m, MonadThrow m)
+  => Config                     -- ^ Ormolu configuration
+  -> FilePath                   -- ^ Location of source file
+  -> m Text                     -- ^ Resulting rendition
+ormoluFile cfg path =
+  liftIO (readFile path) >>= ormolu cfg path
+
+-- | A wrapper around 'parseModule'.
+
+parseModule'
+  :: (MonadIO m, MonadThrow m)
+  => Config                     -- ^ Ormolu configuration
+  -> (GHC.SrcSpan -> String -> OrmoluException)
+     -- ^ How to obtain 'OrmoluException' to throw when parsing fails
+  -> FilePath                   -- ^ File name to use in errors
+  -> String                     -- ^ Actual input for the parser
+  -> m (Anns, GHC.ParsedSource) -- ^ Annotations and parsed source
+parseModule' Config {..} mkException path str = do
+  (_, r) <- liftIO (parseModule cfgDynOptions path str)
+  case r of
+    Left (spn, err) -> throwM (mkException spn err)
+    Right x -> return x
+
+-- | Ormolu exception representing all cases when 'ormoluFile' can fail.
+
+data OrmoluException
+  = OrmoluParsingFailed GHC.SrcSpan String
+    -- ^ Parsing of original source code failed
+  | OrmoluOutputParsingFailed GHC.SrcSpan String
+    -- ^ Parsing of formatted source code failed
+  | OrmoluASTDiffers String Text
+    -- ^ Original and resulting ASTs differ, first argument is the original
+    -- source code, second argument is rendered source code
+  deriving (Eq, Show)
+
+instance Exception OrmoluException

src/Ormolu/Config.hs

@@ -1,10 +1,10 @@
--- | Types used in the library.
+-- | Configuration options used by the tool.
 
-module Ormolu.Type
+module Ormolu.Config
   ( Config (..)
   , defaultConfig
   , DynOption (..)
-  , dynOption
+  , dynOptionToLocatedStr
   )
 where
 
@@ -32,5 +32,5 @@ newtype DynOption = DynOption
 
 -- | Convert 'DynOption' to @'GHC.Located' 'String'@.
 
-dynOption :: DynOption -> GHC.Located String
-dynOption (DynOption o) = GHC.L GHC.noSrcSpan o
+dynOptionToLocatedStr :: DynOption -> GHC.Located String
+dynOptionToLocatedStr (DynOption o) = GHC.L GHC.noSrcSpan o

src/Ormolu/Diff.hs

@@ -0,0 +1,18 @@
+-- | Diffing GHC ASTs modulo span positions.
+
+module Ormolu.Diff
+  ( diff
+  )
+where
+
+import GHC
+import Language.Haskell.GHC.ExactPrint.Types
+
+-- | Return 'False' of two annotated ASTs are the same modulo span
+-- positions.
+
+diff
+  :: (Anns, ParsedSource)       -- ^ First annotated AST
+  -> (Anns, ParsedSource)       -- ^ Second annotated AST
+  -> Bool
+diff _ _ = False -- TODO

src/Ormolu/Parser.hs

@@ -6,7 +6,7 @@ where
 
 import Language.Haskell.GHC.ExactPrint.Parsers hiding (parseModule)
 import Language.Haskell.GHC.ExactPrint.Types
-import Ormolu.Type
+import Ormolu.Config
 import qualified CmdLineParser as GHC
 import qualified DynFlags as GHC
 import qualified GHC hiding (parseModule)
@@ -21,5 +21,5 @@ parseModule
 parseModule dynOpts path input = ghcWrapper $ do
   dynFlags0 <- initDynFlagsPure path input
   (dynFlags1, _, ws) <-
-    GHC.parseDynamicFilePragma dynFlags0 (dynOption <$> dynOpts)
+    GHC.parseDynamicFilePragma dynFlags0 (dynOptionToLocatedStr <$> dynOpts)
   return (ws, parseModuleFromStringInternal dynFlags1 path input)

src/Ormolu/Printer.hs

@@ -1,3 +1,4 @@
+{-# LANGUAGE LambdaCase        #-}
 {-# LANGUAGE OverloadedStrings #-}
 {-# LANGUAGE RecordWildCards   #-}
 
@@ -7,29 +8,117 @@ module Ormolu.Printer
   ( printModule )
 where
 
-import GHC
+import Data.Char (isAlphaNum)
+import Data.Text (Text)
+import FieldLabel (FieldLbl (..))
+import GHC hiding (GhcPs, IE)
+import HsImpExp (IE (..))
 import Language.Haskell.GHC.ExactPrint.Types
+import Module (Module (..))
+import OccName (OccName (..))
 import Ormolu.Printer.Combinators
-import qualified Data.Text.IO as TIO
+import Outputable (Outputable (..), showSDocUnsafe)
+import RdrName (RdrName (..), rdrNameOcc)
+import SrcLoc (combineSrcSpans)
 
 -- | Render a module.
 
 printModule
-  :: Anns
-  -> ParsedSource
-  -> IO ()
-printModule anns src = TIO.putStr $
-  runR True (p_HsModule src) anns
+  :: Anns                       -- ^ Annotations
+  -> ParsedSource               -- ^ Parsed source
+  -> Text                       -- ^ Resulting rendition
+printModule anns src =
+  runR False (p_HsModule src) anns
 
 p_HsModule :: ParsedSource -> R ()
-p_HsModule l =
-  located l $ \HsModule {..} ->
-    p_hsmodName hsmodName
+p_HsModule loc@(L moduleSpan hsModule) = do
+  -- NOTE If span of exports in multiline, the whole thing is multiline.
+  -- This is especially important because span of module itself always seems
+  -- to have length zero, so it's not reliable for layout selection.
+  let spn =
+        case hsmodExports hsModule of
+          Nothing -> moduleSpan
+          Just (L exportsSpan _) -> combineSrcSpans moduleSpan exportsSpan
+  locatedVia spn loc $ \HsModule {..} ->
+    case hsmodName of
+      Nothing -> pure ()
+      Just hsmodName' -> line . velt' $
+        [ located hsmodName' p_hsmodName ] ++
+        (case hsmodExports of
+           Nothing -> []
+           Just hsmodExports' ->
+             [ inci (locatedVia spn hsmodExports' p_hsmodExports)
+             ])
+        ++ [ txt "where"
+           ]
 
-p_hsmodName :: Maybe (Located ModuleName) -> R ()
-p_hsmodName Nothing = return ()
-p_hsmodName (Just l) =
-  located l $ \moduleName -> do
-    line $ do
+p_hsmodName :: ModuleName -> R ()
+p_hsmodName mname = do
   txt "module "
-      atom moduleName
+  atom mname
+
+p_hsmodExports :: [LIE GhcPs] -> R ()
+p_hsmodExports xs = do
+  parens . velt $ withSep comma (located' p_lie) xs
+
+p_lie :: IE GhcPs -> R ()
+p_lie = \case
+  IEVar l1 -> located l1 p_ieWrappedName
+  IEThingAbs l1 -> located l1 p_ieWrappedName
+  IEThingAll l1 -> do
+    located l1 p_ieWrappedName
+    txt " (..)"
+  IEThingWith l1 w xs fls -> do
+    located l1 p_ieWrappedName
+    space
+    p_ieWildcard w
+    parens . velt $ withSep comma (located' p_ieWrappedName) xs
+    parens . velt $ withSep comma (located' p_FieldLbl) fls
+  IEModuleContents l1 -> located l1 p_hsmodName
+  -- XXX I have no idea what these things are for.
+  IEGroup _ _ -> return ()
+  IEDoc _ -> return ()
+  IEDocNamed _ -> return ()
+
+p_ieWrappedName :: IEWrappedName RdrName -> R ()
+p_ieWrappedName = \case
+  IEName l2 -> located l2 p_rdrName
+  IEPattern l2 -> located l2 $ \x -> do
+    txt "pattern "
+    p_rdrName x
+  IEType l2 -> located l2 $ \x -> do
+    txt "type "
+    p_rdrName x
+
+p_rdrName :: RdrName -> R ()
+p_rdrName x = opParens (rdrNameOcc x) $ case x of
+  Unqual occName -> atom occName
+  Qual mname occName -> p_qualName mname occName
+  Orig (Module _ mname) occName -> p_qualName mname occName
+  Exact name -> atom name
+
+p_FieldLbl :: FieldLbl RdrName -> R ()
+p_FieldLbl (FieldLabel x _ _) = atom x
+
+p_qualName :: ModuleName -> OccName -> R ()
+p_qualName mname occName = do
+  atom mname
+  txt "."
+  atom occName
+
+p_ieWildcard :: IEWildcard -> R ()
+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.
+
+opParens :: Outputable a => a -> R () -> R ()
+opParens x m =
+  if all (not . isAlphaNum) (showSDocUnsafe (ppr x))
+    then txt "(" >> m >> txt ")"
+    else m

src/Ormolu/Printer/Combinators.hs

@@ -1,3 +1,4 @@
+{-# LANGUAGE LambdaCase        #-}
 {-# LANGUAGE OverloadedStrings #-}
 {-# LANGUAGE RecordWildCards   #-}
 
@@ -14,8 +15,11 @@ module Ormolu.Printer.Combinators
   , newline
   , inci
   , located
+  , locatedVia
+  , located'
   , velt
   , velt'
+  , withSep
     -- ** Wrapping
   , line
   , braces
@@ -26,14 +30,15 @@ module Ormolu.Printer.Combinators
   , ofType
   , sarrow
   , darrow
+  , space
   )
 where
 
 import Data.Data (Data)
 import Data.List (intersperse)
 import Data.Text (Text)
-import Debug.Trace
 import Language.Haskell.GHC.ExactPrint.Types
+import Ormolu.Printer.Comments
 import Ormolu.Printer.Internal
 import Outputable (Outputable (..), showSDocUnsafe)
 import SrcLoc
@@ -57,22 +62,72 @@ atom = txt . T.pack . showSDocUnsafe . ppr
 -- that may be associated with the primitive and sets corresponding layout
 -- for the inner computation.
 
-located :: Data a => Located a -> (a -> R ()) -> R ()
-located loc@(L l a) f = do
-  -- TODO implement handling of comments properly
+located
+  :: Data a
+  => Located a                  -- ^ Thing to enter
+  -> (a -> R ())                -- ^ How to render inner value
+  -> R ()
+located loc@(L l _) = locatedVia l loc
+
+-- | A special version of 'located' that allows to control layout using
+-- externally provided span.
+
+locatedVia
+  :: Data a
+  => SrcSpan                    -- ^ Span that controls layout selection
+  -> Located a                  -- ^ Thing to enter
+  -> (a -> R ())                -- ^ How to renedr inner value
+  -> R ()
+locatedVia l' loc@(L l a) f = do
   mann <- lookupAnn loc
   let m = enterLayout
-        (if isOneLineSpan l
+        (if isOneLineSpan l'
           then SingleLine
           else MultiLine)
         (f a)
   case mann of
     Nothing -> m
     Just Ann {..} ->
-      enterLayout MultiLine $ do
-        traceShow annPriorComments $ mapM_ atom annPriorComments
+      enterLayout MultiLine . sitcc $ 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.
+
+        -- traceShowM (before,after, decoratedElt)
+
+        spitComments before
         m
-        traceShow annFollowingComments $ mapM_ atom annFollowingComments
+        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.
+
+-- | A version of 'located' with arguments flipped.
+
+located'
+  :: Data a
+  => (a -> R ())                -- ^ How to render inner value
+  -> Located a                  -- ^ Thing to enter
+  -> R ()
+located' = flip located
 
 -- | Element of variable layout. This means that the sub-components may be
 -- rendered either on single line or each on its own line depending on
@@ -96,6 +151,19 @@ velt' xs = sitcc $ sequence_ (intersperse sep (sitcc <$> xs))
   where
     sep = vlayout (spit " ") newline
 
+-- | Put separator between renderings of items of a list.
+
+withSep
+  :: R ()                       -- ^ Separator
+  -> (a -> R ())                -- ^ How to render list items
+  -> [a]                        -- ^ List to render
+  -> [R ()]                     -- ^ List of printing actions
+withSep sep f = \case
+  [] -> []
+  (x:xs) ->
+    let g a = sep >> f a
+    in f x : fmap g xs
+
 ----------------------------------------------------------------------------
 -- Wrapping
 
@@ -158,3 +226,8 @@ sarrow = txt "-> "
 
 darrow :: R ()
 darrow = txt "=> "
+
+-- | Print single space.
+
+space :: R ()
+space = txt " "

src/Ormolu/Printer/Comments.hs

@@ -0,0 +1,198 @@
+-- | Helpers for formatting of comments.
+
+{-# LANGUAGE LambdaCase        #-}
+{-# LANGUAGE OverloadedStrings #-}
+
+module Ormolu.Printer.Comments
+  ( -- * Types
+    Decoration (..)
+  , Decorator (..)
+  , Position (..)
+  , CommentMode (..)
+    -- * Functions for working with comments
+  , spitComments
+  , partitionDPs
+  , addDecoration
+  )
+where
+
+import ApiAnnotation (AnnKeywordId (AnnModule))
+import Control.Monad
+import Data.Bifunctor
+import Data.List (isPrefixOf)
+import Data.Maybe (mapMaybe)
+import Language.Haskell.GHC.ExactPrint.Types
+import Ormolu.Printer.Internal
+import SrcLoc
+import qualified Data.Text as T
+
+-- | Placement instructions for comments.
+
+data Decoration
+  = Decoration Decorator Decorator
+  deriving (Eq, Show)
+
+-- | Decorator in this context is a thing to put before\/after a comment.
+
+data Decorator
+  = NoDec                       -- ^ Output nothing
+  | SpaceDec                    -- ^ Output single space
+  | NewlineDec                  -- ^ Output single newline
+  deriving (Eq, Show)
+
+-- | Position: before vs after.
+
+data Position
+  = Before                      -- ^ Before
+  | After                       -- ^ After
+  deriving (Eq, Show)
+
+-- | For which type of AST leaf we're preparing the comments.
+
+data CommentMode
+  = Module                      -- ^ Module
+  | Other                       -- ^ Other element
+  deriving (Eq, Show)
+
+-- | Output a bunch of 'Comment's. 'DeltaPos'es are used to insert extra
+-- space between the comments when necessary.
+
+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
+
+-- | Output a 'Comment'.
+
+spitComment :: Comment -> R ()
+spitComment (Comment str _ _) =
+  if isMultiline str
+    then if isPragma str
+           then handleOne (normalizePragma str)
+           else forM_ (normalizeIndent str) handleOne
+    else handleOne str
+  where
+    handleOne x = do
+      ensureIndent
+      spit (T.pack x)
+    isMultiline x = not ("--" `isPrefixOf` x)
+    isPragma x = "{-#" `isPrefixOf` x
+    normalizeIndent = fmap (dropWhile (== ' ')) . lines
+    normalizePragma = unwords . words
+
+-- | 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
+  = bimap (takeComments Before) (takeComments After)
+  . break ((== G AnnModule) . fst)
+  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
+
+-- | If 'KeywordId' is a comment, extract it.
+
+annComment :: KeywordId -> Maybe Comment
+annComment (AnnComment x) = Just x
+annComment _ = Nothing
+
+-- | 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) ->
+              srcSpanStartLine spn0 == srcSpanStartLine spn1
+            (_, _) -> False

src/Ormolu/Printer/Internal.hs

@@ -200,7 +200,7 @@ lookupAnn :: Data a => Located a -> R (Maybe Annotation)
 lookupAnn (L s d) = do
   anns <- R (asks rcAnns)
   let k = AnnKey s (annGetConstr d)
-  traceShow k $ return (M.lookup k anns)
+  return (M.lookup k anns)
 
 ----------------------------------------------------------------------------
 -- Debug helpers

tests/Ormolu/Printer/CombinatorsSpec.hs

@@ -35,11 +35,9 @@ stdTest
 stdTest name m path = describe name $ do
   let spath = "data/printer/" ++ path ++ "-single.hs"
       mpath = "data/printer/" ++ path ++ "-multi.hs"
-  context "single-line" $
-    it "works as expected" $
+  it "single-line works" $
       singleLine m `shouldRender` spath
-  context "multi-line" $
-    it "works as expected" $
+  it "multi-line works" $
       multiLine m `shouldRender` mpath
 
 ----------------------------------------------------------------------------

tests/Ormolu/PrinterSpec.hs

@@ -0,0 +1,76 @@
+{-# LANGUAGE TemplateHaskell #-}
+
+module Ormolu.PrinterSpec (spec) where
+
+import Control.Monad
+import Control.Monad.IO.Class
+import Data.List (isSuffixOf)
+import Data.Text (Text)
+import Ormolu
+import Path
+import Path.IO
+import System.FilePath (addExtension, dropExtensions, splitExtensions)
+import Test.Hspec
+import qualified Data.Text as T
+import qualified Data.Text.IO as T
+
+spec :: Spec
+spec = do
+  es <- runIO locateExamples
+  forM_ es checkExample
+
+-- | Check a single given example.
+
+checkExample :: Path Rel File -> Spec
+checkExample srcPath' = it (fromRelFile srcPath' ++ " works") $ do
+    let srcPath = examplesDir </> srcPath'
+    expectedOutputPath <- deriveOutput srcPath
+    -- 1. Given input snippet of source code parse it and pretty print it.
+    -- 2. Parse the result of pretty-printing again and make sure that AST
+    -- is the same as AST of the original snippet. (This happens in
+    -- 'ormoluFile' automatically.)
+    formatted0 <- ormoluFile defaultConfig (fromRelFile srcPath)
+    -- 3. Check the output against expected output. Thus all tests should
+    -- include two files: input and expected output.
+    expected <- (liftIO . T.readFile . fromRelFile) expectedOutputPath
+    formatted0 `shouldMatch` expected
+    -- 4. Check that running the formatter on the output produces the same
+    -- output again (the transformation is idempotent).
+    formatted1 <- ormolu defaultConfig "<formatted>" (T.unpack formatted0)
+    formatted1 `shouldMatch` formatted0
+
+-- | Build list of examples for testing.
+
+locateExamples :: IO [Path Rel File]
+locateExamples =
+  filter isInput . snd <$> listDirRecurRel examplesDir
+
+-- | Does given path look like input path (as opposed to expected output
+-- path)?
+
+isInput :: Path Rel File -> Bool
+isInput path =
+  let s = fromRelFile path
+      (s', exts) = splitExtensions s
+  in exts == ".hs" && not ("-out" `isSuffixOf` s')
+
+-- | For given path of input file return expected name of output.
+
+deriveOutput :: Path Rel File -> IO (Path Rel File)
+deriveOutput path = parseRelFile $
+  addExtension (dropExtensions (fromRelFile path) ++ "-out") "hs"
+
+-- | A version of 'shouldBe' that is specialized to comparing 'Text' values.
+-- It also prints multi-line snippets in a more readable form.
+
+shouldMatch :: Text -> Text -> Expectation
+shouldMatch actual expected  =
+  when (actual /= expected) . expectationFailure $ unlines
+    [ "expected:"
+    , T.unpack expected
+    , "bot got:"
+    , T.unpack actual
+    ]
+
+examplesDir :: Path Rel Dir
+examplesDir = $(mkRelDir "data/examples")