content: {-# LANGUAGE DataKinds #-}
  {-# LANGUAGE LambdaCase #-}
  {-# LANGUAGE OverloadedStrings #-}
  
  -- | In most cases import "Ormolu.Printer.Combinators" instead, these
  -- functions are the low-level building blocks and should not be used on
  -- their own. The 'R' monad is re-exported from "Ormolu.Printer.Combinators"
  -- as well.
  module Ormolu.Printer.Internal
    ( -- * The 'R' monad
      R,
      runR,
  
      -- * Internal functions
      txt,
      interferingTxt,
      atom,
      space,
      newline,
      askSourceType,
      askModuleFixityMap,
      askDebug,
      inci,
      sitcc,
      Layout (..),
      enterLayout,
      vlayout,
      getLayout,
  
      -- * Helpers for braces
      useBraces,
      dontUseBraces,
      canUseBraces,
  
      -- * Special helpers for comment placement
      CommentPosition (..),
      registerPendingCommentLine,
      trimSpanStream,
      nextEltSpan,
      popComment,
      getEnclosingComments,
      getEnclosingSpan,
      getEnclosingSpanWhere,
      withEnclosingSpan,
      thisLineSpans,
  
      -- * Stateful markers
      SpanMark (..),
      spanMarkSpan,
      HaddockStyle (..),
      setSpanMark,
      getSpanMark,
  
      -- * Extensions
      isExtensionEnabled,
    )
  where
  
  import Control.Monad
  import Control.Monad.Reader
  import Control.Monad.State.Strict
  import Data.Bool (bool)
  import Data.Choice (Choice)
  import Data.Choice qualified as Choice
  import Data.Coerce
  import Data.Functor ((<&>))
  import Data.List (find)
  import Data.Maybe (listToMaybe)
  import Data.Text (Text)
  import Data.Text qualified as T
  import Data.Text.Lazy qualified as TL
  import Data.Text.Lazy.Builder
  import GHC.Data.EnumSet (EnumSet)
  import GHC.Data.EnumSet qualified as EnumSet
  import GHC.LanguageExtensions.Type
  import GHC.Types.SrcLoc
  import GHC.Utils.Outputable (Outputable)
  import Ormolu.Config (SourceType (..))
  import Ormolu.Fixity (ModuleFixityMap)
  import Ormolu.Parser.CommentStream
  import Ormolu.Printer.SpanStream
  import Ormolu.Utils (showOutputable)
  
  ----------------------------------------------------------------------------
  -- The 'R' monad
  
  -- | The 'R' monad hosts combinators that allow us to describe how to render
  -- AST.
  newtype R a = R (ReaderT RC (State SC) a)
    deriving (Functor, Applicative, Monad)
  
  -- | Reader context of 'R'. This should be used when we control rendering by
  -- enclosing certain expressions with wrappers.
  data RC = RC
    { -- | Indentation level, as the column index we need to start from after
      -- a newline if we break lines
      rcIndent :: !Int,
      -- | Current layout
      rcLayout :: Layout,
      -- | Spans of enclosing elements of AST
      rcEnclosingSpans :: [RealSrcSpan],
      -- | Whether the last expression in the layout can use braces
      rcCanUseBraces :: Bool,
      -- | Enabled extensions
      rcExtensions :: EnumSet Extension,
      -- | Whether the source is a signature or a regular module
      rcSourceType :: SourceType,
      -- | Module fixity map
      rcModuleFixityMap :: ModuleFixityMap,
      -- | Whether to print out debug information during printing
      rcDebug :: !Bool
    }
  
  -- | State context of 'R'.
  data SC = SC
    { -- | Index of the next column to render
      scColumn :: !Int,
      -- | Indentation level that was used for the current line
      scIndent :: !Int,
      -- | Rendered source code so far
      scBuilder :: Builder,
      -- | Span stream
      scSpanStream :: SpanStream,
      -- | Spans of atoms that have been printed on the current line so far
      scThisLineSpans :: [RealSrcSpan],
      -- | Comment stream
      scCommentStream :: CommentStream,
      -- | Pending comment lines (in reverse order) to be inserted before next
      -- newline, 'Int' is the indentation level
      scPendingComments :: ![(CommentPosition, Text)],
      -- | Whether to output a space before the next output
      scRequestedDelimiter :: !RequestedDelimiter,
      -- | An auxiliary marker for keeping track of last output element
      scSpanMark :: !(Maybe SpanMark)
    }
  
  -- | Make sure next output is delimited by one of the following.
  data RequestedDelimiter
    = -- | A space
      RequestedSpace
    | -- | A newline
      RequestedNewline
    | -- | Nothing
      RequestedNothing
    | -- | We just output a newline
      AfterNewline
    | -- | We haven't printed anything yet
      VeryBeginning
    deriving (Eq, Show)
  
  -- | 'Layout' options.
  data Layout
    = -- | Put everything on single line
      SingleLine
    | -- | Use multiple lines
      MultiLine
    deriving (Eq, Show)
  
  -- | Modes for rendering of pending comments.
  data CommentPosition
    = -- | Put the comment on the same line
      OnTheSameLine
    | -- | Put the comment on next line
      OnNextLine
    deriving (Eq, Show)
  
  -- | Run 'R' monad.
  runR ::
    -- | Monad to run
    R () ->
    -- | Span stream
    SpanStream ->
    -- | Comment stream
    CommentStream ->
    -- | Whether the source is a signature or a regular module
    SourceType ->
    -- | Enabled extensions
    EnumSet Extension ->
    -- | Module fixity map
    ModuleFixityMap ->
    -- | Resulting rendition
    Bool ->
    Text
  runR (R m) sstream cstream sourceType extensions moduleFixityMap debug =
    TL.toStrict . toLazyText . scBuilder $ execState (runReaderT m rc) sc
    where
      rc =
        RC
          { rcIndent = 0,
            rcLayout = MultiLine,
            rcEnclosingSpans = [],
            rcCanUseBraces = False,
            rcExtensions = extensions,
            rcSourceType = sourceType,
            rcModuleFixityMap = moduleFixityMap,
            rcDebug = debug
          }
      sc =
        SC
          { scColumn = 0,
            scIndent = 0,
            scBuilder = mempty,
            scSpanStream = sstream,
            scThisLineSpans = [],
            scCommentStream = cstream,
            scPendingComments = [],
            scRequestedDelimiter = VeryBeginning,
            scSpanMark = Nothing
          }
  
  ----------------------------------------------------------------------------
  -- Internal functions
  
  -- | Type of the thing to output. Influences the primary low-level rendering
  -- function 'spit'.
  data SpitType
    = -- | Simple opaque text that breaks comment series.
      SimpleText
    | -- | Like 'SimpleText', but assume that when this text is inserted it
      -- will separate an 'Atom' and its pending comments, so insert an extra
      -- 'newline' in that case to force the pending comments and continue on
      -- a fresh line.
      InterferingText
    | -- | An atom that typically have span information in the AST and can
      -- have comments attached to it.
      Atom
    | -- | Used for rendering comment lines.
      CommentPart
    deriving (Show, Eq)
  
  -- | Output a fixed 'Text' fragment. The argument may not contain any line
  -- breaks. 'txt' is used to output all sorts of “fixed” bits of syntax like
  -- keywords and pipes @|@ in functional dependencies.
  --
  -- To separate various bits of syntax with white space use 'space' instead
  -- of @'txt' " "@. To output 'Outputable' Haskell entities like numbers use
  -- 'atom'.
  txt ::
    -- | 'Text' to output
    Text ->
    R ()
  txt = spit SimpleText
  
  -- | Similar to 'txt' but the text inserted this way is assumed to break the
  -- “link” between the preceding atom and its pending comments.
  interferingTxt ::
    -- | 'Text' to output
    Text ->
    R ()
  interferingTxt = spit InterferingText
  
  -- | Output 'Outputable' fragment of AST. This can be used to output numeric
  -- literals and similar. Everything that doesn't have inner structure but
  -- does have an 'Outputable' instance.
  atom ::
    (Outputable a) =>
    a ->
    R ()
  atom = spit Atom . T.pack . showOutputable
  
  -- | Low-level non-public helper to define 'txt' and 'atom'.
  spit ::
    -- | Type of the thing to spit
    SpitType ->
    -- | 'Text' to output
    Text ->
    R ()
  spit _ "" = return ()
  spit stype text = do
    requestedDel <- R (gets scRequestedDelimiter)
    pendingComments <- R (gets scPendingComments)
    when (stype == InterferingText && not (null pendingComments)) newline
    case requestedDel of
      RequestedNewline -> do
        R . modify $ \sc ->
          sc
            { scRequestedDelimiter = RequestedNothing
            }
        case stype of
          CommentPart -> newlineRaw
          _ -> newline
      _ -> return ()
    R $ do
      i <- asks rcIndent
      c <- gets scColumn
      closestEnclosing <- asks (listToMaybe . rcEnclosingSpans)
      let indentedTxt = spaces <> text
          spaces = T.replicate spacesN " "
          spacesN =
            if c == 0
              then i
              else bool 0 1 (requestedDel == RequestedSpace)
      modify $ \sc ->
        sc
          { scBuilder = scBuilder sc <> fromText indentedTxt,
            scColumn = scColumn sc + T.length indentedTxt,
            scIndent =
              if c == 0
                then i
                else scIndent sc,
            scThisLineSpans =
              let xs = scThisLineSpans sc
               in case stype of
                    Atom -> case closestEnclosing of
                      Nothing -> xs
                      Just x -> x : xs
                    _ -> xs,
            scRequestedDelimiter = RequestedNothing,
            scSpanMark =
              -- If there are pending comments, do not reset last comment
              -- location.
              if (stype == CommentPart) || (not . null . scPendingComments) sc
                then scSpanMark sc
                else Nothing
          }
  
  -- | This primitive /does not/ necessarily output a space. It just ensures
  -- that the next thing that will be printed on the same line will be
  -- separated by a single space from the previous output. Using this
  -- combinator twice results in at most one space.
  --
  -- In practice this design prevents trailing white space and makes it hard
  -- to output more than one delimiting space in a row, which is what we
  -- usually want.
  space :: R ()
  space = R . modify $ \sc ->
    sc
      { scRequestedDelimiter = case scRequestedDelimiter sc of
          RequestedNothing -> RequestedSpace
          other -> other
      }
  
  -- | Output a newline. First time 'newline' is used after some non-'newline'
  -- output it gets inserted immediately. Second use of 'newline' does not
  -- output anything but makes sure that the next non-white space output will
  -- be prefixed by a newline. Using 'newline' more than twice in a row has no
  -- effect. Also, using 'newline' at the very beginning has no effect, this
  -- is to avoid leading whitespace.
  --
  -- Similarly to 'space', this design prevents trailing newlines and makes it
  -- hard to output more than one blank newline in a row.
  newline :: R ()
  newline = do
    indent <- R (gets scIndent)
    cs <- reverse <$> R (gets scPendingComments)
    case cs of
      [] -> newlineRaw
      ((position, _) : _) -> do
        case position of
          OnTheSameLine -> space
          OnNextLine -> newlineRaw
        R . forM_ cs $ \(_, text) ->
          let modRC rc =
                rc
                  { rcIndent = indent
                  }
              R m = do
                unless (T.null text) $
                  spit CommentPart text
                newlineRaw
           in local modRC m
        R . modify $ \sc ->
          sc
            { scPendingComments = []
            }
  
  -- | Low-level newline primitive. This one always just inserts a newline, no
  -- hooks can be attached.
  newlineRaw :: R ()
  newlineRaw = R . modify $ \sc ->
    let requestedDel = scRequestedDelimiter sc
        builderSoFar = scBuilder sc
     in sc
          { scBuilder = case requestedDel of
              AfterNewline -> builderSoFar
              RequestedNewline -> builderSoFar
              VeryBeginning -> builderSoFar
              _ -> builderSoFar <> "\n",
            scColumn = 0,
            scIndent = 0,
            scThisLineSpans = [],
            scRequestedDelimiter = case scRequestedDelimiter sc of
              AfterNewline -> RequestedNewline
              RequestedNewline -> RequestedNewline
              VeryBeginning -> VeryBeginning
              _ -> AfterNewline
          }
  
  -- | Return the source type.
  askSourceType :: R SourceType
  askSourceType = R (asks rcSourceType)
  
  -- | Retrieve the module fixity map.
  askModuleFixityMap :: R ModuleFixityMap
  askModuleFixityMap = R (asks rcModuleFixityMap)
  
  -- | Retrieve whether we should print out certain debug information while
  -- printing.
  askDebug :: R (Choice "debug")
  askDebug = R (asks (Choice.fromBool . rcDebug))
  
  inciBy :: Int -> R () -> R ()
  inciBy step (R m) = R (local modRC m)
    where
      modRC rc =
        rc
          { rcIndent = rcIndent rc + step
          }
  
  -- | Increase indentation level by one indentation step for the inner
  -- computation. 'inci' should be used when a part of code must be more
  -- indented relative to the parts outside of 'inci' in order for the output
  -- to be valid Haskell. When layout is single-line there is no obvious
  -- effect, but with multi-line layout correct indentation levels matter.
  inci :: R () -> R ()
  inci = inciBy indentStep
  
  -- | Set indentation level for the inner computation equal to current
  -- column. This makes sure that the entire inner block is uniformly
  -- \"shifted\" to the right.
  sitcc :: R () -> R ()
  sitcc (R m) = do
    requestedDel <- R (gets scRequestedDelimiter)
    i <- R (asks rcIndent)
    c <- R (gets scColumn)
    let modRC rc =
          rc
            { rcIndent = max i (c + bool 0 1 (requestedDel == RequestedSpace))
            }
    R (local modRC m)
  
  -- | Set 'Layout' for internal computation.
  enterLayout :: Layout -> R () -> R ()
  enterLayout l (R m) = R (local modRC m)
    where
      modRC rc =
        rc
          { rcLayout = l
          }
  
  -- | Do one or another thing depending on current 'Layout'.
  vlayout ::
    -- | Single line
    R a ->
    -- | Multi line
    R a ->
    R a
  vlayout sline mline = do
    l <- getLayout
    case l of
      SingleLine -> sline
      MultiLine -> mline
  
  -- | Get current 'Layout'.
  getLayout :: R Layout
  getLayout = R (asks rcLayout)
  
  ----------------------------------------------------------------------------
  -- Special helpers for comment placement
  
  -- | Register a comment line for outputting. It will be inserted right
  -- before next newline. When the comment goes after something else on the
  -- same line, a space will be inserted between preceding text and the
  -- comment when necessary.
  registerPendingCommentLine ::
    -- | Comment position
    CommentPosition ->
    -- | 'Text' to output
    Text ->
    R ()
  registerPendingCommentLine position text = R $ do
    modify $ \sc ->
      sc
        { scPendingComments = (position, text) : scPendingComments sc
        }
  
  -- | Drop elements that begin before or at the same place as given
  -- 'SrcSpan'.
  trimSpanStream ::
    -- | Reference span
    RealSrcSpan ->
    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 ::
    (LComment -> Bool) ->
    R (Maybe LComment)
  popComment f = R $ do
    CommentStream cstream <- gets scCommentStream
    case cstream of
      (x : xs) | f x -> do
        modify $ \sc -> sc {scCommentStream = CommentStream xs}
        return $ Just x
      _ -> return Nothing
  
  -- | Get the comments contained in the enclosing span.
  getEnclosingComments :: R [LComment]
  getEnclosingComments = do
    isEnclosed <-
      getEnclosingSpan <&> \case
        Just enclSpan -> containsSpan enclSpan
        Nothing -> const False
    CommentStream cstream <- R $ gets scCommentStream
    pure $ takeWhile (isEnclosed . getLoc) cstream
  
  -- | Get the immediately enclosing 'RealSrcSpan'.
  getEnclosingSpan :: R (Maybe RealSrcSpan)
  getEnclosingSpan = getEnclosingSpanWhere (const True)
  
  -- | Get the first enclosing 'RealSrcSpan' that satisfies given predicate.
  getEnclosingSpanWhere ::
    -- | Predicate to use
    (RealSrcSpan -> Bool) ->
    R (Maybe RealSrcSpan)
  getEnclosingSpanWhere f =
    find f <$> R (asks rcEnclosingSpans)
  
  -- | Set 'RealSrcSpan' of enclosing span for the given computation.
  withEnclosingSpan :: RealSrcSpan -> R () -> R ()
  withEnclosingSpan spn (R m) = R (local modRC m)
    where
      modRC rc =
        rc
          { rcEnclosingSpans = spn : rcEnclosingSpans rc
          }
  
  -- | Get spans on this line so far.
  thisLineSpans :: R [RealSrcSpan]
  thisLineSpans = R (gets scThisLineSpans)
  
  ----------------------------------------------------------------------------
  -- Stateful markers
  
  -- | An auxiliary marker for keeping track of last output element.
  data SpanMark
    = -- | Haddock comment
      HaddockSpan HaddockStyle RealSrcSpan
    | -- | Non-haddock comment
      CommentSpan RealSrcSpan
    | -- | A statement in a do-block and such span
      StatementSpan RealSrcSpan
  
  -- | Project 'RealSrcSpan' from 'SpanMark'.
  spanMarkSpan :: SpanMark -> RealSrcSpan
  spanMarkSpan = \case
    HaddockSpan _ s -> s
    CommentSpan s -> s
    StatementSpan s -> s
  
  -- | Haddock string style.
  data HaddockStyle
    = -- | @-- |@
      Pipe
    | -- | @-- ^@
      Caret
    | -- | @-- *@
      Asterisk Int
    | -- | @-- $@
      Named String
  
  -- | Set span of last output comment.
  setSpanMark ::
    -- | Span mark to set
    SpanMark ->
    R ()
  setSpanMark spnMark = R . modify $ \sc ->
    sc
      { scSpanMark = Just spnMark
      }
  
  -- | Get span of last output comment.
  getSpanMark :: R (Maybe SpanMark)
  getSpanMark = R (gets scSpanMark)
  
  ----------------------------------------------------------------------------
  -- Helpers for braces
  
  -- | Make the inner computation use braces around single-line layouts.
  useBraces :: R () -> R ()
  useBraces (R r) = R (local (\i -> i {rcCanUseBraces = True}) r)
  
  -- | Make the inner computation omit braces around single-line layouts.
  dontUseBraces :: R () -> R ()
  dontUseBraces (R r) = R (local (\i -> i {rcCanUseBraces = False}) r)
  
  -- | Return 'True' if we can use braces in this context.
  canUseBraces :: R Bool
  canUseBraces = R (asks rcCanUseBraces)
  
  ----------------------------------------------------------------------------
  -- Constants
  
  -- | Indentation step.
  indentStep :: Int
  indentStep = 2
  
  ----------------------------------------------------------------------------
  -- Extensions
  
  isExtensionEnabled :: Extension -> R Bool
  isExtensionEnabled ext = R . asks $ EnumSet.member ext . rcExtensions
  