module Language.Haskell.Tools.Transform.PlaceComments where
import Control.Monad.State
import Control.Monad.Writer
import Control.Reference hiding (element)
import Data.Char (isSpace, isAlphaNum)
import qualified Data.Map as Map
import Data.Maybe
import qualified Data.Set as Set (lookupLE, lookupGE, fromList)
import ApiAnnotation (AnnotationComment(..))
import SrcLoc
import Language.Haskell.Tools.AST
getNormalComments :: Map.Map SrcSpan [Located AnnotationComment] -> Map.Map SrcSpan [Located AnnotationComment]
getNormalComments = Map.map (filter (not . isPragma . unLoc))
getPragmaComments :: Map.Map SrcSpan [Located AnnotationComment] -> Map.Map String [Located String]
getPragmaComments comms = Map.fromListWith (++) $ map (\(L l (AnnBlockComment str)) -> (getPragmaCommand str, [L l str]))
$ filter (isPragma . unLoc) $ concatMap snd $ Map.toList comms
where getPragmaCommand = takeWhile (\c -> isAlphaNum c || c == '_') . dropWhile isSpace . drop 3
isPragma :: AnnotationComment -> Bool
isPragma (AnnBlockComment str) = take 3 str == "{-#" && take 3 (reverse str) == "}-#"
isPragma _ = False
placeComments :: RangeInfo stage => Map.Map SrcSpan [Located AnnotationComment]
-> Ann UModule dom stage
-> Ann UModule dom stage
placeComments comms mod
= resizeAnnots (concatMap (map nextSrcLoc . snd) (Map.toList comms)) mod
where spans = allElemSpans mod
sortedElemStarts = Set.fromList $ map srcSpanStart spans
sortedElemEnds = Set.fromList $ map srcSpanEnd spans
nextSrcLoc comm@(L sp _)
= let after = fromMaybe noSrcLoc (Set.lookupLE (srcSpanStart sp) sortedElemEnds)
before = fromMaybe noSrcLoc (Set.lookupGE (srcSpanEnd sp) sortedElemStarts)
in ((after,before),comm)
allElemSpans :: (SourceInfoTraversal node, RangeInfo stage) => Ann node dom stage -> [SrcSpan]
allElemSpans = execWriter . sourceInfoTraverse (SourceInfoTrf (\ni -> tell [ni ^. nodeSpan] >> pure ni) pure pure)
resizeAnnots :: RangeInfo stage => [((SrcLoc, SrcLoc), Located AnnotationComment)]
-> Ann UModule dom stage
-> Ann UModule dom stage
resizeAnnots comments elem
= flip evalState comments $
modImports&annList !~ expandAnnot
>=> modDecl&annList !~ expandTopLevelDecl
>=> expandAnnot
$ elem
type ExpandType elem dom stage = Ann elem dom stage -> State [((SrcLoc, SrcLoc), Located AnnotationComment)] (Ann elem dom stage)
expandTopLevelDecl :: RangeInfo stage => ExpandType UDecl dom stage
expandTopLevelDecl
= declBody & annJust & cbElements & annList !~ expandClsElement
>=> declCons & annList !~ expandConDecl
>=> declGadt & annList !~ expandGadtConDecl
>=> declTypeSig !~ expandTypeSig
>=> expandAnnot
expandTypeSig :: RangeInfo stage => ExpandType UTypeSignature dom stage
expandTypeSig
= tsType & typeParams !~ expandAnnot >=> expandAnnot
expandClsElement :: RangeInfo stage => ExpandType UClassElement dom stage
expandClsElement
= ceTypeSig !~ expandTypeSig
>=> ceBind !~ expandValueBind
>=> expandAnnot
expandValueBind :: RangeInfo stage => ExpandType UValueBind dom stage
expandValueBind
= valBindLocals & annJust & localBinds & annList !~ expandLocalBind
>=> funBindMatches & annList & matchBinds & annJust & localBinds & annList !~ expandLocalBind
>=> expandAnnot
expandLocalBind :: RangeInfo stage => ExpandType ULocalBind dom stage
expandLocalBind
= localVal !~ expandValueBind
>=> localSig !~ expandTypeSig
>=> expandAnnot
expandConDecl :: RangeInfo stage => ExpandType UConDecl dom stage
expandConDecl
= conDeclFields & annList !~ expandAnnot >=> expandAnnot
expandGadtConDecl :: RangeInfo stage => ExpandType UGadtConDecl dom stage
expandGadtConDecl
= gadtConType & gadtConRecordFields & annList !~ expandAnnot >=> expandAnnot
expandAnnot :: forall elem dom stage . RangeInfo stage => ExpandType elem dom stage
expandAnnot elem
= do let Just sp = elem ^? annotation&sourceInfo&nodeSpan
applicable <- gets (applicableComments (srcSpanStart sp) (srcSpanEnd sp))
if not (null applicable) then do
let newSp@(RealSrcSpan newSpan)
= foldl combineSrcSpans (fromJust $ elem ^? nodeSp) (map (getLoc . snd) applicable)
modify (filter (not . (\case RealSrcSpan s -> newSpan `containsSpan` s; _ -> True) . getLoc . snd))
return $ nodeSp .= newSp $ elem
else return elem
where nodeSp :: Simple Partial (Ann elem dom stage) SrcSpan
nodeSp = annotation&sourceInfo&nodeSpan
applicableComments :: SrcLoc -> SrcLoc
-> [((SrcLoc, SrcLoc), Located AnnotationComment)]
-> [((SrcLoc, SrcLoc), Located AnnotationComment)]
applicableComments start end = filter applicableComment
where
applicableComment ((_, before), L _ comm)
| isCommentOnNext comm = before == start
applicableComment ((after, _), L _ comm)
| isCommentOnPrev comm = after == end
applicableComment ((after, _), L (RealSrcSpan loc) _)
| after == end && srcLocLine (realSrcSpanStart loc) == getLineLocDefault end = True
applicableComment ((_, before), L (RealSrcSpan loc) _)
| before == start && srcLocLine (realSrcSpanEnd loc) + 1 == getLineLocDefault start
&& srcLocCol (realSrcSpanStart loc) == getLineColDefault start
= True
applicableComment _ = False
getLineLocDefault (RealSrcLoc l) = srcLocLine l
getLineLocDefault _ = 1
getLineColDefault (RealSrcLoc l) = srcLocCol l
getLineColDefault _ = 1
isCommentOnNext :: AnnotationComment -> Bool
isCommentOnNext (AnnDocCommentNext _) = True
isCommentOnNext (AnnLineComment s) = firstNonspaceCharIs '|' s
isCommentOnNext (AnnBlockComment s) = firstNonspaceCharIs '|' s
isCommentOnNext _ = False
isCommentOnPrev :: AnnotationComment -> Bool
isCommentOnPrev (AnnDocCommentPrev _) = True
isCommentOnPrev (AnnLineComment s) = firstNonspaceCharIs '^' s
isCommentOnPrev (AnnBlockComment s) = firstNonspaceCharIs '^' s
isCommentOnPrev _ = False
firstNonspaceCharIs :: Char -> String -> Bool
firstNonspaceCharIs c s = Just c == listToMaybe (dropWhile isSpace (drop 2 s))