-- | Comment handling.
module Floskell.Comments ( annotateWithComments ) where

import           Control.Arrow                ( first, second )
import           Control.Monad.State.Strict

import           Data.Foldable                ( traverse_ )
import           Data.List                    ( isPrefixOf )
import qualified Data.Map.Strict              as M

import           Floskell.Types

import           Language.Haskell.Exts.SrcLoc ( SrcSpanInfo(..) )

-- Order by start of span, larger spans before smaller spans.
newtype OrderByStart = OrderByStart SrcSpan
    deriving ( OrderByStart -> OrderByStart -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: OrderByStart -> OrderByStart -> Bool
$c/= :: OrderByStart -> OrderByStart -> Bool
== :: OrderByStart -> OrderByStart -> Bool
$c== :: OrderByStart -> OrderByStart -> Bool
Eq )

instance Ord OrderByStart where
    compare :: OrderByStart -> OrderByStart -> Ordering
compare (OrderByStart SrcSpan
l) (OrderByStart SrcSpan
r) =
        forall a. Ord a => a -> a -> Ordering
compare (SrcSpan -> Int
srcSpanStartLine SrcSpan
l) (SrcSpan -> Int
srcSpanStartLine SrcSpan
r)
        forall a. Monoid a => a -> a -> a
`mappend` forall a. Ord a => a -> a -> Ordering
compare (SrcSpan -> Int
srcSpanStartColumn SrcSpan
l) (SrcSpan -> Int
srcSpanStartColumn SrcSpan
r)
        forall a. Monoid a => a -> a -> a
`mappend` forall a. Ord a => a -> a -> Ordering
compare (SrcSpan -> Int
srcSpanEndLine SrcSpan
r) (SrcSpan -> Int
srcSpanEndLine SrcSpan
l)
        forall a. Monoid a => a -> a -> a
`mappend` forall a. Ord a => a -> a -> Ordering
compare (SrcSpan -> Int
srcSpanEndColumn SrcSpan
r) (SrcSpan -> Int
srcSpanEndColumn SrcSpan
l)

-- Order by end of span, smaller spans before larger spans.
newtype OrderByEnd = OrderByEnd SrcSpan
    deriving ( OrderByEnd -> OrderByEnd -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: OrderByEnd -> OrderByEnd -> Bool
$c/= :: OrderByEnd -> OrderByEnd -> Bool
== :: OrderByEnd -> OrderByEnd -> Bool
$c== :: OrderByEnd -> OrderByEnd -> Bool
Eq )

instance Ord OrderByEnd where
    compare :: OrderByEnd -> OrderByEnd -> Ordering
compare (OrderByEnd SrcSpan
l) (OrderByEnd SrcSpan
r) =
        forall a. Ord a => a -> a -> Ordering
compare (SrcSpan -> Int
srcSpanEndLine SrcSpan
l) (SrcSpan -> Int
srcSpanEndLine SrcSpan
r)
        forall a. Monoid a => a -> a -> a
`mappend` forall a. Ord a => a -> a -> Ordering
compare (SrcSpan -> Int
srcSpanEndColumn SrcSpan
l) (SrcSpan -> Int
srcSpanEndColumn SrcSpan
r)
        forall a. Monoid a => a -> a -> a
`mappend` forall a. Ord a => a -> a -> Ordering
compare (SrcSpan -> Int
srcSpanStartLine SrcSpan
r) (SrcSpan -> Int
srcSpanStartLine SrcSpan
l)
        forall a. Monoid a => a -> a -> a
`mappend` forall a. Ord a => a -> a -> Ordering
compare (SrcSpan -> Int
srcSpanStartColumn SrcSpan
r) (SrcSpan -> Int
srcSpanStartColumn SrcSpan
l)

onSameLine :: SrcSpan -> SrcSpan -> Bool
onSameLine :: SrcSpan -> SrcSpan -> Bool
onSameLine SrcSpan
ss SrcSpan
ss' = SrcSpan -> Int
srcSpanEndLine SrcSpan
ss forall a. Eq a => a -> a -> Bool
== SrcSpan -> Int
srcSpanStartLine SrcSpan
ss'

isAfterComment :: Comment -> Bool
isAfterComment :: Comment -> Bool
isAfterComment (Comment CommentType
PreprocessorDirective SrcSpan
_ [Char]
str) = [Char]
"#endif" forall a. Eq a => [a] -> [a] -> Bool
`isPrefixOf` [Char]
str
isAfterComment (Comment CommentType
_ SrcSpan
_ [Char]
str) =
    forall a. Int -> [a] -> [a]
take Int
1 (forall a. (a -> Bool) -> [a] -> [a]
dropWhile (forall a. Eq a => a -> a -> Bool
== Char
' ') forall a b. (a -> b) -> a -> b
$ forall a. (a -> Bool) -> [a] -> [a]
dropWhile (forall a. Eq a => a -> a -> Bool
== Char
'-') [Char]
str) forall a. Eq a => a -> a -> Bool
== [Char]
"^"

isAlignedWith :: Comment -> Comment -> Bool
isAlignedWith :: Comment -> Comment -> Bool
isAlignedWith (Comment CommentType
_ SrcSpan
before [Char]
_) (Comment CommentType
_ SrcSpan
after [Char]
_) =
    SrcSpan -> Int
srcSpanEndLine SrcSpan
before forall a. Eq a => a -> a -> Bool
== SrcSpan -> Int
srcSpanStartLine SrcSpan
after forall a. Num a => a -> a -> a
- Int
1
    Bool -> Bool -> Bool
&& SrcSpan -> Int
srcSpanStartColumn SrcSpan
before forall a. Eq a => a -> a -> Bool
== SrcSpan -> Int
srcSpanStartColumn SrcSpan
after

-- | Annotate the AST with comments.
annotateWithComments
    :: Traversable ast => ast SrcSpanInfo -> [Comment] -> ast NodeInfo
annotateWithComments :: forall (ast :: * -> *).
Traversable ast =>
ast SrcSpanInfo -> [Comment] -> ast NodeInfo
annotateWithComments ast SrcSpanInfo
src [Comment]
comments =
    forall s a. State s a -> s -> a
evalState (do
                   forall (t :: * -> *) (f :: * -> *) a b.
(Foldable t, Applicative f) =>
(a -> f b) -> t a -> f ()
traverse_ Comment -> State (Map SrcSpanInfo ([Comment], [Comment])) ()
assignComment [Comment]
comments
                   forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse SrcSpanInfo
-> State (Map SrcSpanInfo ([Comment], [Comment])) NodeInfo
transferComments ast SrcSpanInfo
src)
              Map SrcSpanInfo ([Comment], [Comment])
nodeinfos
  where
    nodeinfos :: M.Map SrcSpanInfo ([Comment], [Comment])
    nodeinfos :: Map SrcSpanInfo ([Comment], [Comment])
nodeinfos = forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr (\SrcSpanInfo
ssi -> forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert SrcSpanInfo
ssi ([], [])) forall k a. Map k a
M.empty ast SrcSpanInfo
src

    -- Assign a single comment to the right AST node
    assignComment
        :: Comment -> State (M.Map SrcSpanInfo ([Comment], [Comment])) ()
    assignComment :: Comment -> State (Map SrcSpanInfo ([Comment], [Comment])) ()
assignComment comment :: Comment
comment@(Comment CommentType
_ SrcSpan
cspan [Char]
_) = case Comment -> (Maybe SrcSpanInfo, Maybe SrcSpanInfo)
surrounding Comment
comment of
        (Maybe SrcSpanInfo
Nothing, Maybe SrcSpanInfo
Nothing) -> forall a. HasCallStack => [Char] -> a
error [Char]
"No target nodes for comment"
        (Just SrcSpanInfo
before, Maybe SrcSpanInfo
Nothing) -> Location
-> SrcSpanInfo -> State (Map SrcSpanInfo ([Comment], [Comment])) ()
insertComment Location
After SrcSpanInfo
before
        (Maybe SrcSpanInfo
Nothing, Just SrcSpanInfo
after) -> Location
-> SrcSpanInfo -> State (Map SrcSpanInfo ([Comment], [Comment])) ()
insertComment Location
Before SrcSpanInfo
after
        (Just SrcSpanInfo
before, Just SrcSpanInfo
after) ->
            if SrcSpanInfo -> SrcSpan
srcInfoSpan SrcSpanInfo
before SrcSpan -> SrcSpan -> Bool
`onSameLine` SrcSpan
cspan Bool -> Bool -> Bool
|| Comment -> Bool
isAfterComment Comment
comment
            then Location
-> SrcSpanInfo -> State (Map SrcSpanInfo ([Comment], [Comment])) ()
insertComment Location
After SrcSpanInfo
before
            else do
                ([Comment], [Comment])
cmts <- forall s (m :: * -> *) a. MonadState s m => (s -> a) -> m a
gets (forall k a. Ord k => Map k a -> k -> a
M.! SrcSpanInfo
before)
                case ([Comment], [Comment])
cmts of
                    -- We've already collected comments for this
                    -- node and this comment is a continuation.
                    ([Comment]
_, Comment
c' : [Comment]
_)
                        | Comment
c' Comment -> Comment -> Bool
`isAlignedWith` Comment
comment ->
                            Location
-> SrcSpanInfo -> State (Map SrcSpanInfo ([Comment], [Comment])) ()
insertComment Location
After SrcSpanInfo
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.
                    ([Comment], [Comment])
_ -> Location
-> SrcSpanInfo -> State (Map SrcSpanInfo ([Comment], [Comment])) ()
insertComment Location
Before SrcSpanInfo
after
      where
        insertComment :: Location
                      -> SrcSpanInfo
                      -> State (M.Map SrcSpanInfo ([Comment], [Comment])) ()
        insertComment :: Location
-> SrcSpanInfo -> State (Map SrcSpanInfo ([Comment], [Comment])) ()
insertComment Location
Before SrcSpanInfo
ssi = forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify forall a b. (a -> b) -> a -> b
$ forall k a. Ord k => (a -> a) -> k -> Map k a -> Map k a
M.adjust (forall (a :: * -> * -> *) b c d.
Arrow a =>
a b c -> a (b, d) (c, d)
first (Comment
comment forall a. a -> [a] -> [a]
:)) SrcSpanInfo
ssi
        insertComment Location
After SrcSpanInfo
ssi = forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify forall a b. (a -> b) -> a -> b
$ forall k a. Ord k => (a -> a) -> k -> Map k a -> Map k a
M.adjust (forall (a :: * -> * -> *) b c d.
Arrow a =>
a b c -> a (d, b) (d, c)
second (Comment
comment forall a. a -> [a] -> [a]
:)) SrcSpanInfo
ssi

    -- Transfer collected comments into the AST.
    transferComments
        :: SrcSpanInfo
        -> State (M.Map SrcSpanInfo ([Comment], [Comment])) NodeInfo
    transferComments :: SrcSpanInfo
-> State (Map SrcSpanInfo ([Comment], [Comment])) NodeInfo
transferComments SrcSpanInfo
ssi = do
        ([Comment]
c, [Comment]
c') <- forall s (m :: * -> *) a. MonadState s m => (s -> a) -> m a
gets (forall k a. Ord k => Map k a -> k -> a
M.! SrcSpanInfo
ssi)
        -- Sometimes, there are multiple AST nodes with the same
        -- SrcSpan.  Make sure we assign comments to only one of
        -- them.
        forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify forall a b. (a -> b) -> a -> b
$ forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert SrcSpanInfo
ssi ([], [])
        forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ SrcSpan -> [Comment] -> [Comment] -> NodeInfo
NodeInfo (SrcSpanInfo -> SrcSpan
srcInfoSpan SrcSpanInfo
ssi) (forall a. [a] -> [a]
reverse [Comment]
c) (forall a. [a] -> [a]
reverse [Comment]
c')

    surrounding :: Comment -> (Maybe SrcSpanInfo, Maybe SrcSpanInfo)
surrounding (Comment CommentType
_ SrcSpan
ss [Char]
_) = (SrcSpan -> Maybe SrcSpanInfo
nodeBefore SrcSpan
ss, SrcSpan -> Maybe SrcSpanInfo
nodeAfter SrcSpan
ss)

    nodeBefore :: SrcSpan -> Maybe SrcSpanInfo
nodeBefore SrcSpan
ss = forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall a b. (a, b) -> b
snd forall a b. (a -> b) -> a -> b
$ SrcSpan -> OrderByEnd
OrderByEnd SrcSpan
ss forall k v. Ord k => k -> Map k v -> Maybe (k, v)
`M.lookupLT` Map OrderByEnd SrcSpanInfo
spansByEnd

    nodeAfter :: SrcSpan -> Maybe SrcSpanInfo
nodeAfter SrcSpan
ss = forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall a b. (a, b) -> b
snd forall a b. (a -> b) -> a -> b
$ SrcSpan -> OrderByStart
OrderByStart SrcSpan
ss forall k v. Ord k => k -> Map k v -> Maybe (k, v)
`M.lookupGT` Map OrderByStart SrcSpanInfo
spansByStart

    spansByStart :: Map OrderByStart SrcSpanInfo
spansByStart = forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr (\SrcSpanInfo
ssi -> forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert (SrcSpan -> OrderByStart
OrderByStart forall a b. (a -> b) -> a -> b
$ SrcSpanInfo -> SrcSpan
srcInfoSpan SrcSpanInfo
ssi) SrcSpanInfo
ssi)
                         forall k a. Map k a
M.empty
                         ast SrcSpanInfo
src

    spansByEnd :: Map OrderByEnd SrcSpanInfo
spansByEnd =
        forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr (\SrcSpanInfo
ssi -> forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert (SrcSpan -> OrderByEnd
OrderByEnd forall a b. (a -> b) -> a -> b
$ SrcSpanInfo -> SrcSpan
srcInfoSpan SrcSpanInfo
ssi) SrcSpanInfo
ssi) forall k a. Map k a
M.empty ast SrcSpanInfo
src