1 files changed, 458 insertions, 0 deletions
diff --git a/lib/Language/Haskell/Stylish/Printer.hs b/lib/Language/Haskell/Stylish/Printer.hs
new file mode 100644
index 0000000..a7ddf5e
--- /dev/null
+++ b/lib/Language/Haskell/Stylish/Printer.hs
@@ -0,0 +1,458 @@
+{-# LANGUAGE BangPatterns #-}
+{-# LANGUAGE BlockArguments #-}
+{-# LANGUAGE DeriveGeneric #-}
+{-# LANGUAGE DerivingStrategies #-}
+{-# LANGUAGE DoAndIfThenElse #-}
+{-# LANGUAGE GeneralizedNewtypeDeriving #-}
+{-# LANGUAGE LambdaCase #-}
+module Language.Haskell.Stylish.Printer
+  ( Printer(..)
+  , PrinterConfig(..)
+  , PrinterState(..)
+
+    -- * Alias
+  , P
+
+    -- * Functions to use the printer
+  , runPrinter
+  , runPrinter_
+
+    -- ** Combinators
+  , comma
+  , dot
+  , getAnnot
+  , getCurrentLine
+  , getCurrentLineLength
+  , getDocstrPrev
+  , newline
+  , parenthesize
+  , peekNextCommentPos
+  , prefix
+  , putComment
+  , putEolComment
+  , putOutputable
+  , putAllSpanComments
+  , putCond
+  , putType
+  , putRdrName
+  , putText
+  , removeCommentTo
+  , removeCommentToEnd
+  , removeLineComment
+  , sep
+  , groupAttachedComments
+  , space
+  , spaces
+  , suffix
+  , pad
+
+    -- ** Advanced combinators
+  , withColumns
+  , modifyCurrentLine
+  , wrapping
+  ) where
+
+--------------------------------------------------------------------------------
+import           Prelude                         hiding (lines)
+
+--------------------------------------------------------------------------------
+import           ApiAnnotation                   (AnnKeywordId(..), AnnotationComment(..))
+import           GHC.Hs.Extension                (GhcPs, NoExtField(..))
+import           GHC.Hs.Types                    (HsType(..))
+import           Module                          (ModuleName, moduleNameString)
+import           RdrName                         (RdrName(..))
+import           SrcLoc                          (GenLocated(..), RealLocated)
+import           SrcLoc                          (Located, SrcSpan(..))
+import           SrcLoc                          (srcSpanStartLine, srcSpanEndLine)
+import           Outputable                      (Outputable)
+
+--------------------------------------------------------------------------------
+import           Control.Monad                   (forM_, replicateM_)
+import           Control.Monad.Reader            (MonadReader, ReaderT(..), asks, local)
+import           Control.Monad.State             (MonadState, State)
+import           Control.Monad.State             (runState)
+import           Control.Monad.State             (get, gets, modify, put)
+import           Data.Foldable                   (find)
+import           Data.Functor                    ((<&>))
+import           Data.List                       (delete, isPrefixOf)
+import           Data.List.NonEmpty              (NonEmpty(..))
+
+--------------------------------------------------------------------------------
+import           Language.Haskell.Stylish.Module (Module, Lines, lookupAnnotation)
+import           Language.Haskell.Stylish.GHC    (showOutputable, unLocated)
+
+-- | Shorthand for 'Printer' monad
+type P = Printer
+
+-- | Printer that keeps state of file
+newtype Printer a = Printer (ReaderT PrinterConfig (State PrinterState) a)
+  deriving (Applicative, Functor, Monad, MonadReader PrinterConfig, MonadState PrinterState)
+
+-- | Configuration for printer, currently empty
+data PrinterConfig = PrinterConfig
+    { columns :: !(Maybe Int)
+    }
+
+-- | State of printer
+data PrinterState = PrinterState
+  { lines :: !Lines
+  , linePos :: !Int
+  , currentLine :: !String
+  , pendingComments :: ![RealLocated AnnotationComment]
+  , parsedModule :: !Module
+  }
+
+-- | Run printer to get printed lines out of module as well as return value of monad
+runPrinter :: PrinterConfig -> [RealLocated AnnotationComment] -> Module -> Printer a -> (a, Lines)
+runPrinter cfg comments m (Printer printer) =
+  let
+    (a, PrinterState parsedLines _ startedLine _ _) = runReaderT printer cfg `runState` PrinterState [] 0 "" comments m
+  in
+    (a, parsedLines <> if startedLine == [] then [] else [startedLine])
+
+-- | Run printer to get printed lines only
+runPrinter_ :: PrinterConfig -> [RealLocated AnnotationComment] -> Module -> Printer a -> Lines
+runPrinter_ cfg comments m printer = snd (runPrinter cfg comments m printer)
+
+-- | Print text
+putText :: String -> P ()
+putText txt = do
+  l <- gets currentLine
+  modify \s -> s { currentLine = l <> txt }
+
+-- | Check condition post action, and use fallback if false
+putCond :: (PrinterState -> Bool) -> P b -> P b -> P b
+putCond p action fallback = do
+  prevState <- get
+  res <- action
+  currState <- get
+  if p currState then pure res
+  else put prevState >> fallback
+
+-- | Print an 'Outputable'
+putOutputable :: Outputable a => a -> P ()
+putOutputable = putText . showOutputable
+
+-- | Put all comments that has positions within 'SrcSpan' and separate by
+--   passed @P ()@
+putAllSpanComments :: P () -> SrcSpan -> P ()
+putAllSpanComments suff = \case
+  UnhelpfulSpan _ -> pure ()
+  RealSrcSpan rspan -> do
+    cmts <- removeComments \(L rloc _) ->
+      srcSpanStartLine rloc >= srcSpanStartLine rspan &&
+      srcSpanEndLine rloc <= srcSpanEndLine rspan
+
+    forM_ cmts (\c -> putComment c >> suff)
+
+-- | Print any comment
+putComment :: AnnotationComment -> P ()
+putComment = \case
+  AnnLineComment s -> putText s
+  AnnDocCommentNext s -> putText s
+  AnnDocCommentPrev s -> putText s
+  AnnDocCommentNamed s -> putText s
+  AnnDocSection _ s -> putText s
+  AnnDocOptions s -> putText s
+  AnnBlockComment s -> putText s
+
+-- | Given the current start line of 'SrcSpan', remove and put EOL comment for same line
+putEolComment :: SrcSpan -> P ()
+putEolComment = \case
+  RealSrcSpan rspan -> do
+    cmt <- removeComment \case
+      L rloc (AnnLineComment s) ->
+        and
+          [ srcSpanStartLine rspan == srcSpanStartLine rloc
+          , not ("-- ^" `isPrefixOf` s)
+          , not ("-- |" `isPrefixOf` s)
+          ]
+      _ -> False
+    forM_ cmt (\c -> space >> putComment c)
+  UnhelpfulSpan _ -> pure ()
+
+-- | Print a 'RdrName'
+putRdrName :: Located RdrName -> P ()
+putRdrName (L pos n) = case n of
+  Unqual name -> do
+    annots <- getAnnot pos
+    if AnnOpenP `elem` annots then do
+      putText "("
+      putText (showOutputable name)
+      putText ")"
+    else if AnnBackquote `elem` annots then do
+      putText "`"
+      putText (showOutputable name)
+      putText "`"
+    else if AnnSimpleQuote `elem` annots then do
+      putText "'"
+      putText (showOutputable name)
+    else
+      putText (showOutputable name)
+  Qual modulePrefix name ->
+    putModuleName modulePrefix >> dot >> putText (showOutputable name)
+  Orig _ name ->
+    putText (showOutputable name)
+  Exact name ->
+    putText (showOutputable name)
+
+-- | Print module name
+putModuleName :: ModuleName -> P ()
+putModuleName = putText . moduleNameString
+
+-- | Print type
+putType :: Located (HsType GhcPs) -> P ()
+putType ltp = case unLocated ltp of
+  HsFunTy NoExtField argTp funTp -> do
+    putOutputable argTp
+    space
+    putText "->"
+    space
+    putType funTp
+  HsAppTy NoExtField t1 t2 ->
+    putType t1 >> space >> putType t2
+  HsExplicitListTy NoExtField _ xs -> do
+    putText "'["
+    sep
+      (comma >> space)
+      (fmap putType xs)
+    putText "]"
+  HsExplicitTupleTy NoExtField xs -> do
+    putText "'("
+    sep
+      (comma >> space)
+      (fmap putType xs)
+    putText ")"
+  HsOpTy NoExtField lhs op rhs -> do
+    putType lhs
+    space
+    putRdrName op
+    space
+    putType rhs
+  HsTyVar NoExtField _ rdrName ->
+    putRdrName rdrName
+  HsTyLit _ tp ->
+    putOutputable tp
+  HsParTy _ tp -> do
+    putText "("
+    putType tp
+    putText ")"
+  HsTupleTy NoExtField _ xs -> do
+    putText "("
+    sep
+      (comma >> space)
+      (fmap putType xs)
+    putText ")"
+  HsForAllTy NoExtField _ _ _ ->
+    putOutputable ltp
+  HsQualTy NoExtField _ _ ->
+    putOutputable ltp
+  HsAppKindTy _ _ _ ->
+    putOutputable ltp
+  HsListTy _ _ ->
+    putOutputable ltp
+  HsSumTy _ _ ->
+    putOutputable ltp
+  HsIParamTy _ _ _ ->
+    putOutputable ltp
+  HsKindSig _ _ _ ->
+    putOutputable ltp
+  HsStarTy _ _ ->
+    putOutputable ltp
+  HsSpliceTy _ _ ->
+    putOutputable ltp
+  HsDocTy _ _ _ ->
+    putOutputable ltp
+  HsBangTy _ _ _ ->
+    putOutputable ltp
+  HsRecTy _ _ ->
+    putOutputable ltp
+  HsWildCardTy _ ->
+    putOutputable ltp
+  XHsType _ ->
+    putOutputable ltp
+
+-- | Get a docstring on the start line of 'SrcSpan' that is a @-- ^@ comment
+getDocstrPrev :: SrcSpan -> P (Maybe AnnotationComment)
+getDocstrPrev = \case
+  UnhelpfulSpan _ -> pure Nothing
+  RealSrcSpan rspan -> do
+    removeComment \case
+      L rloc (AnnLineComment s) ->
+        and
+          [ srcSpanStartLine rspan == srcSpanStartLine rloc
+          , "-- ^" `isPrefixOf` s
+          ]
+      _ -> False
+
+-- | Print a newline
+newline :: P ()
+newline = do
+  l <- gets currentLine
+  modify \s -> s { currentLine = "", linePos = 0, lines = lines s <> [l] }
+
+-- | Print a space
+space :: P ()
+space = putText " "
+
+-- | Print a number of spaces
+spaces :: Int -> P ()
+spaces i = replicateM_ i space
+
+-- | Print a dot
+dot :: P ()
+dot = putText "."
+
+-- | Print a comma
+comma :: P ()
+comma = putText ","
+
+-- | Add parens around a printed action
+parenthesize :: P a -> P a
+parenthesize action = putText "(" *> action <* putText ")"
+
+-- | Add separator between each element of the given printers
+sep :: P a -> [P a] -> P ()
+sep _ [] = pure ()
+sep s (first : rest) = first >> forM_ rest ((>>) s)
+
+-- | Prefix a printer with another one
+prefix :: P a -> P b -> P b
+prefix pa pb = pa >> pb
+
+-- | Suffix a printer with another one
+suffix :: P a -> P b -> P a
+suffix pa pb = pb >> pa
+
+-- | Indent to a given number of spaces.  If the current line already exceeds
+-- that number in length, nothing happens.
+pad :: Int -> P ()
+pad n = do
+    len <- length <$> getCurrentLine
+    spaces $ n - len
+
+-- | Gets comment on supplied 'line' and removes it from the state
+removeLineComment :: Int -> P (Maybe AnnotationComment)
+removeLineComment line =
+  removeComment (\(L rloc _) -> srcSpanStartLine rloc == line)
+
+-- | Removes comments from the state up to start line of 'SrcSpan' and returns
+--   the ones that were removed
+removeCommentTo :: SrcSpan -> P [AnnotationComment]
+removeCommentTo = \case
+  UnhelpfulSpan _ -> pure []
+  RealSrcSpan rspan -> removeCommentTo' (srcSpanStartLine rspan)
+
+-- | Removes comments from the state up to end line of 'SrcSpan' and returns
+--   the ones that were removed
+removeCommentToEnd :: SrcSpan -> P [AnnotationComment]
+removeCommentToEnd = \case
+  UnhelpfulSpan _ -> pure []
+  RealSrcSpan rspan -> removeCommentTo' (srcSpanEndLine rspan)
+
+-- | Removes comments to the line number given and returns the ones removed
+removeCommentTo' :: Int -> P [AnnotationComment]
+removeCommentTo' line =
+  removeComment (\(L rloc _) -> srcSpanStartLine rloc < line) >>= \case
+    Nothing -> pure []
+    Just c -> do
+      rest <- removeCommentTo' line
+      pure (c : rest)
+
+-- | Removes comments from the state while given predicate 'p' is true
+removeComments :: (RealLocated AnnotationComment -> Bool) -> P [AnnotationComment]
+removeComments p =
+  removeComment p >>= \case
+    Just c -> do
+      rest <- removeComments p
+      pure (c : rest)
+    Nothing -> pure []
+
+-- | Remove a comment from the state given predicate 'p'
+removeComment :: (RealLocated AnnotationComment -> Bool) -> P (Maybe AnnotationComment)
+removeComment p = do
+  comments <- gets pendingComments
+
+  let
+    foundComment =
+      find p comments
+
+    newPendingComments =
+      maybe comments (`delete` comments) foundComment
+
+  modify \s -> s { pendingComments = newPendingComments }
+  pure $ fmap (\(L _ c) -> c) foundComment
+
+-- | Get all annotations for 'SrcSpan'
+getAnnot :: SrcSpan -> P [AnnKeywordId]
+getAnnot spn = gets (lookupAnnotation spn . parsedModule)
+
+-- | Get current line
+getCurrentLine :: P String
+getCurrentLine = gets currentLine
+
+-- | Get current line length
+getCurrentLineLength :: P Int
+getCurrentLineLength = fmap length getCurrentLine
+
+-- | Peek at the next comment in the state
+peekNextCommentPos :: P (Maybe SrcSpan)
+peekNextCommentPos = do
+  gets pendingComments <&> \case
+    (L next _ : _) -> Just (RealSrcSpan next)
+    [] -> Nothing
+
+-- | Get attached comments belonging to '[Located a]' given
+groupAttachedComments :: [Located a] -> P [([AnnotationComment], NonEmpty (Located a))]
+groupAttachedComments = go
+  where
+    go :: [Located a] -> P [([AnnotationComment], NonEmpty (Located a))]
+    go (L rspan x : xs) = do
+      comments <- removeCommentTo rspan
+      nextGroupStartM <- peekNextCommentPos
+
+      let
+        sameGroupOf = maybe xs \nextGroupStart ->
+          takeWhile (\(L p _)-> p < nextGroupStart) xs
+
+        restOf = maybe [] \nextGroupStart ->
+          dropWhile (\(L p _) -> p <= nextGroupStart) xs
+
+      restGroups <- go (restOf nextGroupStartM)
+      pure $ (comments, L rspan x :| sameGroupOf nextGroupStartM) : restGroups
+
+    go _ = pure []
+
+modifyCurrentLine :: (String -> String) -> P ()
+modifyCurrentLine f = do
+    s0 <- get
+    put s0 {currentLine = f $ currentLine s0}
+
+wrapping
+    :: P a  -- ^ First printer to run
+    -> P a  -- ^ Printer to run if first printer violates max columns
+    -> P a  -- ^ Result of either the first or the second printer
+wrapping p1 p2 = do
+    maxCols <- asks columns
+    case maxCols of
+        -- No wrapping
+        Nothing -> p1
+        Just c  -> do
+            s0 <- get
+            x <- p1
+            s1 <- get
+            if length (currentLine s1) <= c
+                -- No need to wrap
+                then pure x
+                else do
+                    put s0
+                    y <- p2
+                    s2 <- get
+                    if length (currentLine s1) == length (currentLine s2)
+                        -- Wrapping didn't help!
+                        then put s1 >> pure x
+                        -- Wrapped
+                        else pure y
+
+withColumns :: Maybe Int -> P a -> P a
+withColumns c = local $ \pc -> pc {columns = c}