-- | 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
(OrderByStart -> OrderByStart -> Bool)
-> (OrderByStart -> OrderByStart -> Bool) -> Eq OrderByStart
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) =
        Int -> Int -> Ordering
forall a. Ord a => a -> a -> Ordering
compare (SrcSpan -> Int
srcSpanStartLine SrcSpan
l) (SrcSpan -> Int
srcSpanStartLine SrcSpan
r)
        Ordering -> Ordering -> Ordering
forall a. Monoid a => a -> a -> a
`mappend` Int -> Int -> Ordering
forall a. Ord a => a -> a -> Ordering
compare (SrcSpan -> Int
srcSpanStartColumn SrcSpan
l) (SrcSpan -> Int
srcSpanStartColumn SrcSpan
r)
        Ordering -> Ordering -> Ordering
forall a. Monoid a => a -> a -> a
`mappend` Int -> Int -> Ordering
forall a. Ord a => a -> a -> Ordering
compare (SrcSpan -> Int
srcSpanEndLine SrcSpan
r) (SrcSpan -> Int
srcSpanEndLine SrcSpan
l)
        Ordering -> Ordering -> Ordering
forall a. Monoid a => a -> a -> a
`mappend` Int -> Int -> Ordering
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
(OrderByEnd -> OrderByEnd -> Bool)
-> (OrderByEnd -> OrderByEnd -> Bool) -> Eq OrderByEnd
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) =
        Int -> Int -> Ordering
forall a. Ord a => a -> a -> Ordering
compare (SrcSpan -> Int
srcSpanEndLine SrcSpan
l) (SrcSpan -> Int
srcSpanEndLine SrcSpan
r)
        Ordering -> Ordering -> Ordering
forall a. Monoid a => a -> a -> a
`mappend` Int -> Int -> Ordering
forall a. Ord a => a -> a -> Ordering
compare (SrcSpan -> Int
srcSpanEndColumn SrcSpan
l) (SrcSpan -> Int
srcSpanEndColumn SrcSpan
r)
        Ordering -> Ordering -> Ordering
forall a. Monoid a => a -> a -> a
`mappend` Int -> Int -> Ordering
forall a. Ord a => a -> a -> Ordering
compare (SrcSpan -> Int
srcSpanStartLine SrcSpan
r) (SrcSpan -> Int
srcSpanStartLine SrcSpan
l)
        Ordering -> Ordering -> Ordering
forall a. Monoid a => a -> a -> a
`mappend` Int -> Int -> Ordering
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 Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== SrcSpan -> Int
srcSpanStartLine SrcSpan
ss'

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

isAlignedWith :: Comment -> Comment -> Bool
isAlignedWith :: Comment -> Comment -> Bool
isAlignedWith (Comment CommentType
_ SrcSpan
before String
_) (Comment CommentType
_ SrcSpan
after String
_) =
    SrcSpan -> Int
srcSpanEndLine SrcSpan
before Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== SrcSpan -> Int
srcSpanStartLine SrcSpan
after Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1
    Bool -> Bool -> Bool
&& SrcSpan -> Int
srcSpanStartColumn SrcSpan
before Int -> Int -> Bool
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 :: ast SrcSpanInfo -> [Comment] -> ast NodeInfo
annotateWithComments ast SrcSpanInfo
src [Comment]
comments =
    State (Map SrcSpanInfo ([Comment], [Comment])) (ast NodeInfo)
-> Map SrcSpanInfo ([Comment], [Comment]) -> ast NodeInfo
forall s a. State s a -> s -> a
evalState (do
                   (Comment
 -> StateT (Map SrcSpanInfo ([Comment], [Comment])) Identity ())
-> [Comment]
-> StateT (Map SrcSpanInfo ([Comment], [Comment])) Identity ()
forall (t :: * -> *) (f :: * -> *) a b.
(Foldable t, Applicative f) =>
(a -> f b) -> t a -> f ()
traverse_ Comment
-> StateT (Map SrcSpanInfo ([Comment], [Comment])) Identity ()
assignComment [Comment]
comments
                   (SrcSpanInfo
 -> StateT
      (Map SrcSpanInfo ([Comment], [Comment])) Identity NodeInfo)
-> ast SrcSpanInfo
-> State (Map SrcSpanInfo ([Comment], [Comment])) (ast NodeInfo)
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse SrcSpanInfo
-> StateT
     (Map SrcSpanInfo ([Comment], [Comment])) Identity NodeInfo
transferComments ast SrcSpanInfo
src)
              Map SrcSpanInfo ([Comment], [Comment])
nodeinfos
  where
    nodeinfos :: M.Map SrcSpanInfo ([Comment], [Comment])
    nodeinfos :: Map SrcSpanInfo ([Comment], [Comment])
nodeinfos = (SrcSpanInfo
 -> Map SrcSpanInfo ([Comment], [Comment])
 -> Map SrcSpanInfo ([Comment], [Comment]))
-> Map SrcSpanInfo ([Comment], [Comment])
-> ast SrcSpanInfo
-> Map SrcSpanInfo ([Comment], [Comment])
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr (\SrcSpanInfo
ssi -> SrcSpanInfo
-> ([Comment], [Comment])
-> Map SrcSpanInfo ([Comment], [Comment])
-> Map SrcSpanInfo ([Comment], [Comment])
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert SrcSpanInfo
ssi ([], [])) Map SrcSpanInfo ([Comment], [Comment])
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
-> StateT (Map SrcSpanInfo ([Comment], [Comment])) Identity ()
assignComment comment :: Comment
comment@(Comment CommentType
_ SrcSpan
cspan String
_) = case Comment -> (Maybe SrcSpanInfo, Maybe SrcSpanInfo)
surrounding Comment
comment of
        (Maybe SrcSpanInfo
Nothing, Maybe SrcSpanInfo
Nothing) -> String
-> StateT (Map SrcSpanInfo ([Comment], [Comment])) Identity ()
forall a. HasCallStack => String -> a
error String
"No target nodes for comment"
        (Just SrcSpanInfo
before, Maybe SrcSpanInfo
Nothing) -> Location
-> SrcSpanInfo
-> StateT (Map SrcSpanInfo ([Comment], [Comment])) Identity ()
insertComment Location
After SrcSpanInfo
before
        (Maybe SrcSpanInfo
Nothing, Just SrcSpanInfo
after) -> Location
-> SrcSpanInfo
-> StateT (Map SrcSpanInfo ([Comment], [Comment])) Identity ()
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
-> StateT (Map SrcSpanInfo ([Comment], [Comment])) Identity ()
insertComment Location
After SrcSpanInfo
before
            else do
                ([Comment], [Comment])
cmts <- (Map SrcSpanInfo ([Comment], [Comment]) -> ([Comment], [Comment]))
-> StateT
     (Map SrcSpanInfo ([Comment], [Comment]))
     Identity
     ([Comment], [Comment])
forall s (m :: * -> *) a. MonadState s m => (s -> a) -> m a
gets (Map SrcSpanInfo ([Comment], [Comment])
-> SrcSpanInfo -> ([Comment], [Comment])
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
-> StateT (Map SrcSpanInfo ([Comment], [Comment])) Identity ()
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
-> StateT (Map SrcSpanInfo ([Comment], [Comment])) Identity ()
insertComment Location
Before SrcSpanInfo
after
      where
        insertComment :: Location
                      -> SrcSpanInfo
                      -> State (M.Map SrcSpanInfo ([Comment], [Comment])) ()
        insertComment :: Location
-> SrcSpanInfo
-> StateT (Map SrcSpanInfo ([Comment], [Comment])) Identity ()
insertComment Location
Before SrcSpanInfo
ssi = (Map SrcSpanInfo ([Comment], [Comment])
 -> Map SrcSpanInfo ([Comment], [Comment]))
-> StateT (Map SrcSpanInfo ([Comment], [Comment])) Identity ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify ((Map SrcSpanInfo ([Comment], [Comment])
  -> Map SrcSpanInfo ([Comment], [Comment]))
 -> StateT (Map SrcSpanInfo ([Comment], [Comment])) Identity ())
-> (Map SrcSpanInfo ([Comment], [Comment])
    -> Map SrcSpanInfo ([Comment], [Comment]))
-> StateT (Map SrcSpanInfo ([Comment], [Comment])) Identity ()
forall a b. (a -> b) -> a -> b
$ (([Comment], [Comment]) -> ([Comment], [Comment]))
-> SrcSpanInfo
-> Map SrcSpanInfo ([Comment], [Comment])
-> Map SrcSpanInfo ([Comment], [Comment])
forall k a. Ord k => (a -> a) -> k -> Map k a -> Map k a
M.adjust (([Comment] -> [Comment])
-> ([Comment], [Comment]) -> ([Comment], [Comment])
forall (a :: * -> * -> *) b c d.
Arrow a =>
a b c -> a (b, d) (c, d)
first (Comment
comment Comment -> [Comment] -> [Comment]
forall a. a -> [a] -> [a]
:)) SrcSpanInfo
ssi
        insertComment Location
After SrcSpanInfo
ssi = (Map SrcSpanInfo ([Comment], [Comment])
 -> Map SrcSpanInfo ([Comment], [Comment]))
-> StateT (Map SrcSpanInfo ([Comment], [Comment])) Identity ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify ((Map SrcSpanInfo ([Comment], [Comment])
  -> Map SrcSpanInfo ([Comment], [Comment]))
 -> StateT (Map SrcSpanInfo ([Comment], [Comment])) Identity ())
-> (Map SrcSpanInfo ([Comment], [Comment])
    -> Map SrcSpanInfo ([Comment], [Comment]))
-> StateT (Map SrcSpanInfo ([Comment], [Comment])) Identity ()
forall a b. (a -> b) -> a -> b
$ (([Comment], [Comment]) -> ([Comment], [Comment]))
-> SrcSpanInfo
-> Map SrcSpanInfo ([Comment], [Comment])
-> Map SrcSpanInfo ([Comment], [Comment])
forall k a. Ord k => (a -> a) -> k -> Map k a -> Map k a
M.adjust (([Comment] -> [Comment])
-> ([Comment], [Comment]) -> ([Comment], [Comment])
forall (a :: * -> * -> *) b c d.
Arrow a =>
a b c -> a (d, b) (d, c)
second (Comment
comment Comment -> [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
-> StateT
     (Map SrcSpanInfo ([Comment], [Comment])) Identity NodeInfo
transferComments SrcSpanInfo
ssi = do
        ([Comment]
c, [Comment]
c') <- (Map SrcSpanInfo ([Comment], [Comment]) -> ([Comment], [Comment]))
-> StateT
     (Map SrcSpanInfo ([Comment], [Comment]))
     Identity
     ([Comment], [Comment])
forall s (m :: * -> *) a. MonadState s m => (s -> a) -> m a
gets (Map SrcSpanInfo ([Comment], [Comment])
-> SrcSpanInfo -> ([Comment], [Comment])
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.
        (Map SrcSpanInfo ([Comment], [Comment])
 -> Map SrcSpanInfo ([Comment], [Comment]))
-> StateT (Map SrcSpanInfo ([Comment], [Comment])) Identity ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify ((Map SrcSpanInfo ([Comment], [Comment])
  -> Map SrcSpanInfo ([Comment], [Comment]))
 -> StateT (Map SrcSpanInfo ([Comment], [Comment])) Identity ())
-> (Map SrcSpanInfo ([Comment], [Comment])
    -> Map SrcSpanInfo ([Comment], [Comment]))
-> StateT (Map SrcSpanInfo ([Comment], [Comment])) Identity ()
forall a b. (a -> b) -> a -> b
$ SrcSpanInfo
-> ([Comment], [Comment])
-> Map SrcSpanInfo ([Comment], [Comment])
-> Map SrcSpanInfo ([Comment], [Comment])
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert SrcSpanInfo
ssi ([], [])
        NodeInfo
-> StateT
     (Map SrcSpanInfo ([Comment], [Comment])) Identity NodeInfo
forall (m :: * -> *) a. Monad m => a -> m a
return (NodeInfo
 -> StateT
      (Map SrcSpanInfo ([Comment], [Comment])) Identity NodeInfo)
-> NodeInfo
-> StateT
     (Map SrcSpanInfo ([Comment], [Comment])) Identity NodeInfo
forall a b. (a -> b) -> a -> b
$ SrcSpan -> [Comment] -> [Comment] -> NodeInfo
NodeInfo (SrcSpanInfo -> SrcSpan
srcInfoSpan SrcSpanInfo
ssi) ([Comment] -> [Comment]
forall a. [a] -> [a]
reverse [Comment]
c) ([Comment] -> [Comment]
forall a. [a] -> [a]
reverse [Comment]
c')

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

    nodeBefore :: SrcSpan -> Maybe SrcSpanInfo
nodeBefore SrcSpan
ss = ((OrderByEnd, SrcSpanInfo) -> SrcSpanInfo)
-> Maybe (OrderByEnd, SrcSpanInfo) -> Maybe SrcSpanInfo
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (OrderByEnd, SrcSpanInfo) -> SrcSpanInfo
forall a b. (a, b) -> b
snd (Maybe (OrderByEnd, SrcSpanInfo) -> Maybe SrcSpanInfo)
-> Maybe (OrderByEnd, SrcSpanInfo) -> Maybe SrcSpanInfo
forall a b. (a -> b) -> a -> b
$ SrcSpan -> OrderByEnd
OrderByEnd SrcSpan
ss OrderByEnd
-> Map OrderByEnd SrcSpanInfo -> Maybe (OrderByEnd, SrcSpanInfo)
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 = ((OrderByStart, SrcSpanInfo) -> SrcSpanInfo)
-> Maybe (OrderByStart, SrcSpanInfo) -> Maybe SrcSpanInfo
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (OrderByStart, SrcSpanInfo) -> SrcSpanInfo
forall a b. (a, b) -> b
snd (Maybe (OrderByStart, SrcSpanInfo) -> Maybe SrcSpanInfo)
-> Maybe (OrderByStart, SrcSpanInfo) -> Maybe SrcSpanInfo
forall a b. (a -> b) -> a -> b
$ SrcSpan -> OrderByStart
OrderByStart SrcSpan
ss OrderByStart
-> Map OrderByStart SrcSpanInfo
-> Maybe (OrderByStart, SrcSpanInfo)
forall k v. Ord k => k -> Map k v -> Maybe (k, v)
`M.lookupGT` Map OrderByStart SrcSpanInfo
spansByStart

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

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