GhcApiAstTraversals: Utils.hs

{-# LANGUAGE RankNTypes #-}
module Utils where

import Data.Generics

-- import qualified GHC.Paths
import PprTyThing
import DynFlags
import GHC
import Outputable
import SrcLoc
import qualified OccName(occNameString)
import Bag(Bag,bagToList)
import Var(Var)
import FastString(FastString)
import NameSet(NameSet,nameSetToList)

import Instances
import Data.List

-- for tagging data with its source
data Stage = Parser | Renamer | TypeChecker deriving (Eq,Ord,Show)

-- generic Data-based show, with special cases for GHC Ast types,
-- showing abstract types abstractly and avoiding known potholes
showData :: Data a => Stage -> Int -> a -> String
showData stage n = 
  generic `ext1Q` list `extQ` string `extQ` fastString `extQ` srcSpan 
          `extQ` name `extQ` occName `extQ` moduleName `extQ` var `extQ` dataCon
          `extQ` bagName `extQ` bagRdrName `extQ` bagVar `extQ` nameSet
          `extQ` postTcType `extQ` fixity
  where generic :: Data a => a -> String
        generic t = indent n ++ "(" ++ showConstr (toConstr t)
                 ++ space (concat (intersperse " " (gmapQ (showData stage (n+1)) t))) ++ ")"
        space "" = ""
        space s  = ' ':s
        indent n = "\n" ++ replicate n ' ' 
        string     = show :: String -> String
        fastString = ("{FastString: "++) . (++"}") . show :: FastString -> String
        list l     = indent n ++ "[" 
                              ++ concat (intersperse "," (map (showData stage (n+1)) l)) ++ "]"

        name       = ("{Name: "++) . (++"}") . showSDoc . ppr :: Name -> String
        occName    = ("{OccName: "++) . (++"}") .  OccName.occNameString 
        moduleName = ("{ModuleName: "++) . (++"}") . showSDoc . ppr :: ModuleName -> String
        srcSpan    = ("{"++) . (++"}") . showSDoc . ppr :: SrcSpan -> String
        var        = ("{Var: "++) . (++"}") . showSDoc . ppr :: Var -> String
        dataCon    = ("{DataCon: "++) . (++"}") . showSDoc . ppr :: DataCon -> String

        bagRdrName:: Bag (Located (HsBind RdrName)) -> String
        bagRdrName = ("{Bag(Located (HsBind RdrName)): "++) . (++"}") . list . bagToList 
        bagName   :: Bag (Located (HsBind Name)) -> String
        bagName    = ("{Bag(Located (HsBind Name)): "++) . (++"}") . list . bagToList 
        bagVar    :: Bag (Located (HsBind Var)) -> String
        bagVar     = ("{Bag(Located (HsBind Var)): "++) . (++"}") . list . bagToList 

        nameSet | stage `elem` [Parser,TypeChecker] 
                = const ("{!NameSet placeholder here!}") :: NameSet -> String
                | otherwise     
                = ("{NameSet: "++) . (++"}") . list . nameSetToList 

        postTcType | stage<TypeChecker = const "{!type placeholder here?!}" :: PostTcType -> String
                   | otherwise     = showSDoc . ppr :: Type -> String

        fixity | stage<Renamer = const "{!fixity placeholder here?!}" :: GHC.Fixity -> String
               | otherwise     = ("{Fixity: "++) . (++"}") . showSDoc . ppr :: GHC.Fixity -> String

-- like everything, but avoid known potholes
everythingStaged :: Stage -> (r -> r -> r) -> r -> GenericQ r -> GenericQ r
everythingStaged stage k z f x 
  | (const False `extQ` postTcType `extQ` fixity `extQ` nameSet) x = z
  | otherwise = foldl k (f x) (gmapQ (everythingStaged stage k z f) x)
  where nameSet    = const (stage `elem` [Parser,TypeChecker]) :: NameSet -> Bool
        postTcType = const (stage<TypeChecker)                 :: PostTcType -> Bool
        fixity     = const (stage<Renamer)                     :: GHC.Fixity -> Bool

everythingBut :: GenericQ Bool -> (r -> r -> r) -> r -> GenericQ r -> GenericQ r
everythingBut q k z f x 
  | q x       = z
  | otherwise = foldl k (f x) (gmapQ (everythingBut q k z f) x)