Program that walks the HsSyn AST to create a list of all names and their source locations. Usage: 1. Place Main.hs and Test.hs in same directory. 2. Compile and run: ghc Main.hs ./Main

Main.hs

{-# LANGUAGE BangPatterns, PatternGuards #-}
module Main (main) where

import Control.Monad (forM_, unless)
import Prelude hiding (id, mod)
import System.Environment (getArgs)
import System.Exit (ExitCode(ExitFailure), exitWith)

import Bag
import Digraph (flattenSCCs)
import DynFlags (defaultFatalMessager, defaultFlushOut)
import GHC hiding (flags)
import GHC.Paths (libdir)
import HscTypes (isBootSummary, msHsFilePath)
import Id (idName)
import MonadUtils (liftIO)
import Name (getOccString, nameModule_maybe)

main :: IO ()
main = do
    args <- getArgs
    processTargets args ["Test.hs"]

processTargets :: [String] -> [FilePath] -> IO ()
processTargets ghcArgs filenames =
    defaultErrorHandler defaultFatalMessager defaultFlushOut $
        runGhc (Just libdir) $ do
            dflags <- getSessionDynFlags
            (pflags, unrec, warns) <- parseDynamicFlags dflags
                                      (map noLoc ghcArgs)
            unless (null unrec) $
                liftIO $ putStrLn $ "Unrecognised options:\n" ++
                show (map unLoc unrec)
            liftIO $ mapM_ (putStrLn . unLoc) warns
            let dflags2 = pflags { hscTarget = HscNothing }
            _ <- setSessionDynFlags dflags2
            defaultCleanupHandler dflags2 $ do
                targets <- mapM (\f -> guessTarget f Nothing) filenames
                setTargets targets
                modgraph <- depanal [] False
                let mods = flattenSCCs $ topSortModuleGraph False modgraph
                           Nothing
                indexSymbols mods

indexSymbols :: ModuleGraph -> Ghc ()
indexSymbols graph = forM_ graph $ \ ms -> do
    let filename = msHsFilePath ms
    handleSourceError printErrorAndExit $ do
        liftIO $ putStrLn ("Loading " ++ filename ++ " ...")
        mod <- loadModule =<< typecheckModule =<< parseModule ms
        case mod of
            _ | isBootSummary ms -> return ()
            _ | Just (group, _, _, _) <- renamedSource mod -> do
                let syms = bagToList $ symbols group $ typecheckedSource mod
                liftIO $ print $ map (\ (Symbol name _) -> nameToString name) syms
            _  -> liftIO $ exitWith (ExitFailure 1)
  where
    -- Convert Name to String.
    nameToString name = (maybe "" (moduleNameString . moduleName) . nameModule_maybe $ name)
                        ++ "." ++ getOccString name
    printErrorAndExit e = do
        printException e
        liftIO $ exitWith (ExitFailure 1)

------------------------------------------------------------------------
-- Extract all interesting symbols from the AST

-- | Data type used to store names with their location. We only want
-- real names with real locations, not compiler derived ones.
data Symbol = Symbol !Name !RealSrcSpan
            deriving (Eq)

-- TODO: This only traverses a small part of the AST. Eventually we
-- want to traverse it all.

symbols :: HsGroup Name -> TypecheckedSource -> Bag Symbol
symbols _group src = concatMapBag lHsBindSymbols src

lHsBindSymbols :: LHsBind Id -> Bag Symbol
lHsBindSymbols lbinding = case unLoc lbinding of
    b@FunBind {} -> unitBag (symbol (unLoc $ fun_id b) lbinding) `unionBags`
                    matchGroupSymbols (fun_matches b)
    PatBind { pat_lhs = _lhs } -> emptyBag  -- patThings lhs []
    VarBind { var_id = id } -> unitBag (symbol id lbinding)
    AbsBinds { abs_binds = lbindings } -> concatMapBag lHsBindSymbols lbindings
  where
    matchGroupSymbols (MatchGroup lmatches _) =
        unionManyBags $ map lMatchSymbols lmatches

lMatchSymbols :: LMatch Id -> Bag Symbol
lMatchSymbols lmatch = case unLoc lmatch of
    Match _ _ (GRHSs grhss _) -> unionManyBags $ map lGRHSSymbols grhss
  where
    lGRHSSymbols lgrhs = case unLoc lgrhs of
        GRHS _ lHsExpr -> lHsExprSymbols lHsExpr

-- TODO: Only use HsVars with non-local variables
lHsExprSymbols :: LHsExpr Id -> Bag Symbol
lHsExprSymbols lexpr = case unLoc lexpr of
    HsVar id
        | isExternalName (idName id) -> unitBag $ symbol id lexpr
        | otherwise -> emptyBag
    HsApp e1 e2 -> lHsExprSymbols e1 `unionBags` lHsExprSymbols e2
    HsWrap _ e -> lHsExprSymbols (L (getLoc lexpr) e)
    _ -> emptyBag

------------------------------------------------------------------------
-- Utilities

symbol :: Id -> Located a -> Symbol
symbol id located = Symbol (idName id) (realLocation located)

realLocation :: Located a -> RealSrcSpan
realLocation lHs = case getLoc lHs of
    RealSrcSpan l -> l
    UnhelpfulSpan _ -> error "realLocation: UnhelpfulSpan"

concatMapBag :: (a -> Bag b) -> Bag a -> Bag b
concatMapBag f = concatBag . mapBag f

Test.hs

module Test
    ( mysum
    ) where

import Data.List (foldl')

mysum :: [Int] -> Int
mysum xs = foldl' (+) 0 xs