+-- Extract the filename, stringbuffer content and dynflags associed to a module
+--
+-- XXX: Explain pre-conditions
+getModuleSourceAndFlags :: GhcMonad m => Module -> m (String, StringBuffer, DynFlags)
+getModuleSourceAndFlags mod = do
+ m <- getModSummary (moduleName mod)
+ case ml_hs_file $ ms_location m of
+ Nothing -> throw $ mkApiErr (text "No source available for module " <+> ppr mod)
+ Just sourceFile -> do
+ source <- liftIO $ hGetStringBuffer sourceFile
+ return (sourceFile, source, ms_hspp_opts m)
+
+
+-- | Return module source as token stream, including comments.
+--
+-- The module must be in the module graph and its source must be available.
+-- Throws a 'HscTypes.SourceError' on parse error.
+getTokenStream :: GhcMonad m => Module -> m [Located Token]
+getTokenStream mod = do
+ (sourceFile, source, flags) <- getModuleSourceAndFlags mod
+ let startLoc = mkSrcLoc (mkFastString sourceFile) 1 1
+ case lexTokenStream source startLoc flags of
+ POk _ ts -> return ts
+ PFailed span err -> throw $ mkSrcErr (unitBag $ mkPlainErrMsg span err)
+
+-- | Give even more information on the source than 'getTokenStream'
+-- This function allows reconstructing the source completely with
+-- 'showRichTokenStream'.
+getRichTokenStream :: GhcMonad m => Module -> m [(Located Token, String)]
+getRichTokenStream mod = do
+ (sourceFile, source, flags) <- getModuleSourceAndFlags mod
+ let startLoc = mkSrcLoc (mkFastString sourceFile) 1 1
+ case lexTokenStream source startLoc flags of
+ POk _ ts -> return $ addSourceToTokens startLoc source ts
+ PFailed span err -> throw $ mkSrcErr (unitBag $ mkPlainErrMsg span err)
+
+-- | Given a source location and a StringBuffer corresponding to this
+-- location, return a rich token stream with the source associated to the
+-- tokens.
+addSourceToTokens :: SrcLoc -> StringBuffer -> [Located Token]
+ -> [(Located Token, String)]
+addSourceToTokens _ _ [] = []
+addSourceToTokens loc buf (t@(L span _) : ts)
+ | not (isGoodSrcSpan span) = (t,"") : addSourceToTokens loc buf ts
+ | otherwise = (t,str) : addSourceToTokens newLoc newBuf ts
+ where
+ (newLoc, newBuf, str) = go "" loc buf
+ start = srcSpanStart span
+ end = srcSpanEnd span
+ go acc loc buf | loc < start = go acc nLoc nBuf
+ | start <= loc && loc < end = go (ch:acc) nLoc nBuf
+ | otherwise = (loc, buf, reverse acc)
+ where (ch, nBuf) = nextChar buf
+ nLoc = advanceSrcLoc loc ch
+
+
+-- | Take a rich token stream such as produced from 'getRichTokenStream' and
+-- return source code almost identical to the original code (except for
+-- insignificant whitespace.)
+showRichTokenStream :: [(Located Token, String)] -> String
+showRichTokenStream ts = go startLoc ts ""
+ where sourceFile = srcSpanFile (getLoc . fst . head $ ts)
+ startLoc = mkSrcLoc sourceFile 1 1
+ go _ [] = id
+ go loc ((L span _, str):ts)
+ | not (isGoodSrcSpan span) = go loc ts
+ | locLine == tokLine = ((replicate (tokCol - locCol) ' ') ++)
+ . (str ++)
+ . go tokEnd ts
+ | otherwise = ((replicate (tokLine - locLine) '\n') ++)
+ . ((replicate tokCol ' ') ++)
+ . (str ++)
+ . go tokEnd ts
+ where (locLine, locCol) = (srcLocLine loc, srcLocCol loc)
+ (tokLine, tokCol) = (srcSpanStartLine span, srcSpanStartCol span)
+ tokEnd = srcSpanEnd span
+