content: {-# LANGUAGE MultiWayIf #-}
  {-# LANGUAGE OverloadedStrings #-}
  {-# LANGUAGE RecordWildCards #-}
  
  -- | Comment handling.
  module Floskell.Comments ( filterCommentLike, annotateWithComments ) where
  
  import           Control.Arrow                ( first, second )
  import           Control.Monad.State.Strict
  
  import           Data.Foldable                ( traverse_ )
  import           Data.List                    ( foldl', isPrefixOf )
  import qualified Data.Map.Strict              as M
  import           Data.Text.Lazy               ( Text )
  import qualified Data.Text.Lazy               as TL
  
  import           Floskell.Types
  
  import           Language.Haskell.Exts.SrcLoc ( SrcSpanInfo(..) )
  
  data FilterMode = Normal | CppContinuation | Unformatted
      deriving ( Eq, Ord, Enum, Show )
  
  data FilterState = FilterState { stMode     :: !FilterMode
                                 , stLines    :: [Text]
                                 , stComments :: [Comment]
                                 }
      deriving ( Show )
  
  -- Order by start of span, larger spans before smaller spans.
  newtype OrderByStart = OrderByStart SrcSpan
      deriving ( Eq )
  
  instance Ord OrderByStart where
      compare (OrderByStart l) (OrderByStart r) =
          compare (srcSpanStartLine l) (srcSpanStartLine r)
          `mappend` compare (srcSpanStartColumn l) (srcSpanStartColumn r)
          `mappend` compare (srcSpanEndLine r) (srcSpanEndLine l)
          `mappend` compare (srcSpanEndColumn r) (srcSpanEndColumn l)
  
  -- Order by end of span, smaller spans before larger spans.
  newtype OrderByEnd = OrderByEnd SrcSpan
      deriving ( Eq )
  
  instance Ord OrderByEnd where
      compare (OrderByEnd l) (OrderByEnd r) =
          compare (srcSpanEndLine l) (srcSpanEndLine r)
          `mappend` compare (srcSpanEndColumn l) (srcSpanEndColumn r)
          `mappend` compare (srcSpanStartLine r) (srcSpanStartLine l)
          `mappend` compare (srcSpanStartColumn r) (srcSpanStartColumn l)
  
  -- | Remove comment-like blocks from input source, replacing them with
  -- blank likes to keep SrcSpan information intact.
  filterCommentLike :: [Text] -> ([Text], [Comment])
  filterCommentLike = finish . foldl' go start . zip [ 1 .. ]
    where
      start = FilterState Normal [] []
  
      go s@FilterState{..} (n, l) = case stMode of
          Normal -> if
              | isShebangLine l -> addComment s Normal IgnoredLine n l
              | isCppLine l ->
                  let newMode =
                          if isCppContinuation l then CppContinuation else Normal
                  in
                      addComment s newMode PreprocessorDirective n l
              | isBeginIgnore l -> addComment s Unformatted IgnoredLine n l
              | otherwise -> addLine s l
  
          CppContinuation ->
              let newMode =
                      if isCppContinuation l then CppContinuation else Normal
              in
                  addComment s newMode PreprocessorDirective n l
  
          Unformatted ->
              let newMode = if isEndIgnore l then Normal else Unformatted
              in
                  addComment s newMode IgnoredLine n l
  
      finish s = (reverse $ stLines s, reverse $ stComments s)
  
      addLine s@FilterState{..} l = s { stLines = l : stLines }
  
      addComment s@FilterState{..} mode t n l =
          s { stMode     = mode
            , stLines    = "" : stLines
            , stComments = makeComment t n l : stComments
            }
  
      makeComment t n l =
          Comment t
                  (SrcSpan "" n 1 n (fromIntegral $ TL.length l + 1))
                  (TL.unpack l)
  
      isShebangLine = TL.isPrefixOf "#!"
  
      isCppLine src =
          any (`TL.isPrefixOf` src)
              [ "#define"
              , "#elif"
              , "#else"
              , "#end"
              , "#enum"
              , "#error"
              , "#if"
              , "#include"
              , "#undef"
              , "#warning"
              ]
  
      isCppContinuation = TL.isSuffixOf "\\"
  
      isBeginIgnore src = any (`TL.isPrefixOf` src)
                              [ "-- floskell-begin-disable-region"
                              , "-- floskell-disable"
                              , "{- floskell-disable"
                              ]
  
      isEndIgnore src = any (`TL.isPrefixOf` src)
                            [ "-- floskell-end-disable-region"
                            , "-- floskell-enable"
                            , "{- floskell-enable"
                            ]
  
  onSameLine :: SrcSpan -> SrcSpan -> Bool
  onSameLine ss ss' = srcSpanEndLine ss == srcSpanStartLine ss'
  
  isAfterComment :: Comment -> Bool
  isAfterComment (Comment PreprocessorDirective _ str) = "#endif" `isPrefixOf` str
  isAfterComment (Comment _ _ str) =
      take 1 (dropWhile (== ' ') $ dropWhile (== '-') str) == "^"
  
  isAlignedWith :: Comment -> Comment -> Bool
  isAlignedWith (Comment _ before _) (Comment _ after _) =
      srcSpanEndLine before == srcSpanStartLine after - 1
      && srcSpanStartColumn before == srcSpanStartColumn after
  
  -- | Annotate the AST with comments.
  annotateWithComments
      :: Traversable ast => ast SrcSpanInfo -> [Comment] -> ast NodeInfo
  annotateWithComments src comments =
      evalState (do
                     traverse_ assignComment comments
                     traverse transferComments src)
                nodeinfos
    where
      nodeinfos :: M.Map SrcSpanInfo ([Comment], [Comment])
      nodeinfos = foldr (\ssi -> M.insert ssi ([], [])) M.empty src
  
      -- Assign a single comment to the right AST node
      assignComment
          :: Comment -> State (M.Map SrcSpanInfo ([Comment], [Comment])) ()
      assignComment comment@(Comment _ cspan _) = case surrounding comment of
          (Nothing, Nothing) -> error "No target nodes for comment"
          (Just before, Nothing) -> insertComment After before
          (Nothing, Just after) -> insertComment Before after
          (Just before, Just after) ->
              if srcInfoSpan before `onSameLine` cspan || isAfterComment comment
              then insertComment After before
              else do
                  cmts <- gets (M.! before)
                  case cmts of
                      -- We've already collected comments for this
                      -- node and this comment is a continuation.
                      (_, c' : _)
                          | c' `isAlignedWith` comment ->
                              insertComment After before
  
                      -- The comment does not belong to this node.
                      -- If there is a node following this comment,
                      -- assign it to that node, else keep it here,
                      -- anyway.
                      _ -> insertComment Before after
        where
          insertComment :: Location
                        -> SrcSpanInfo
                        -> State (M.Map SrcSpanInfo ([Comment], [Comment])) ()
          insertComment Before ssi = modify $ M.adjust (first (comment :)) ssi
          insertComment After ssi = modify $ M.adjust (second (comment :)) ssi
  
      -- Transfer collected comments into the AST.
      transferComments
          :: SrcSpanInfo
          -> State (M.Map SrcSpanInfo ([Comment], [Comment])) NodeInfo
      transferComments ssi = do
          (c, c') <- gets (M.! ssi)
          -- Sometimes, there are multiple AST nodes with the same
          -- SrcSpan.  Make sure we assign comments to only one of
          -- them.
          modify $ M.insert ssi ([], [])
          return $ NodeInfo (srcInfoSpan ssi) (reverse c) (reverse c')
  
      surrounding (Comment _ ss _) = (nodeBefore ss, nodeAfter ss)
  
      nodeBefore ss = fmap snd $ OrderByEnd ss `M.lookupLT` spansByEnd
  
      nodeAfter ss = fmap snd $ OrderByStart ss `M.lookupGT` spansByStart
  
      spansByStart = foldr (\ssi -> M.insert (OrderByStart $ srcInfoSpan ssi) ssi)
                           M.empty
                           src
  
      spansByEnd =
          foldr (\ssi -> M.insert (OrderByEnd $ srcInfoSpan ssi) ssi) M.empty src
  