module Language.Haskell.Tools.AST.FromGHC.SourceMap where
import ApiAnnotation
import Data.Maybe
import Data.Map as Map
import Data.List as List
import Safe
import SrcLoc as GHC
import FastString as GHC
type SourceMap = (Map AnnKeywordId (Map SrcLoc SrcLoc), Map SrcLoc (SrcSpan, AnnKeywordId))
getKeywordAnywhere :: AnnKeywordId -> SourceMap -> Maybe SrcSpan
getKeywordAnywhere keyw srcmap = return . uncurry mkSrcSpan =<< headMay . assocs =<< (Map.lookup keyw (fst srcmap))
getKeywordInside :: AnnKeywordId -> SrcSpan -> SourceMap -> Maybe SrcSpan
getKeywordInside keyw sr srcmap = getSourceElementInside True sr =<< Map.lookup keyw (fst srcmap)
getKeywordsInside :: AnnKeywordId -> SrcSpan -> SourceMap -> [SrcSpan]
getKeywordsInside keyw sr srcmap
= let tokensOfType = Map.lookup keyw (fst srcmap)
(_, startsAtBegin, startAfterBegin) = Map.splitLookup (srcSpanStart sr) $ fromMaybe empty tokensOfType
(startsBeforeEnd, _) = Map.split (srcSpanEnd sr) $ maybe id (Map.insert (srcSpanStart sr)) startsAtBegin startAfterBegin
in List.map (uncurry mkSrcSpan) $ List.filter (\(_, end) -> end <= srcSpanEnd sr) $ assocs startsBeforeEnd
getKeywordInsideBack :: AnnKeywordId -> SrcSpan -> SourceMap -> Maybe SrcSpan
getKeywordInsideBack keyw sr srcmap = getSourceElementInside False sr =<< Map.lookup keyw (fst srcmap)
getSourceElementInside :: Bool -> SrcSpan -> Map SrcLoc SrcLoc -> Maybe SrcSpan
getSourceElementInside b sr srcmap =
case (if b then lookupGE (srcSpanStart sr) else lookupLT (srcSpanEnd sr)) srcmap of
Just (k, v) -> let sp = mkSrcSpan k v in if sp `isSubspanOf` sr then Just sp else Nothing
Nothing -> Nothing
getNextToken :: SrcLoc -> SourceMap -> Maybe (SrcSpan, AnnKeywordId)
getNextToken loc srcmap = fmap snd $ Map.lookupGE loc $ snd srcmap
getTokensAfter :: SrcLoc -> SourceMap -> [(SrcSpan, AnnKeywordId)]
getTokensAfter loc srcmap = case Map.splitLookup loc $ snd srcmap of
(_, Just elem, after) -> elem : elems after
(_, Nothing, after) -> elems after
annotationsToSrcMap :: Map ApiAnnKey [SrcSpan] -> SourceMap
annotationsToSrcMap anns = (Map.map (List.foldr addToSrcRanges Map.empty) $ mapKeysWith (++) snd anns, tokenMap)
where
addToSrcRanges :: SrcSpan -> Map SrcLoc SrcLoc -> Map SrcLoc SrcLoc
addToSrcRanges span srcmap = Map.insert (srcSpanStart span) (srcSpanEnd span) srcmap
tokenMap = Map.fromList $ List.map (\(k,v) -> (srcSpanStart k, (k, v))) $ concatMap (\(key,vals) -> List.map ((, snd key)) vals) $ Map.assocs anns