module SST (SrcSpanType, knownLoc, sstLHsBind) where
import GHC
import DynFlags
import Outputable
import TypeRep
import PprTyThing
import SrcLoc
import HsBinds
import Var
import Bag
import TcHsSyn

data SrcSpanType = SST SrcSpan Type

instance Eq SrcSpanType where
  SST s _ == SST s' _ = s == s'

instance Ord SrcSpanType where
  SST s _ `compare` SST s' _ = s `compare` s'
  
instance Outputable SrcSpanType where
  ppr (SST s ty) = ppr s $$ nest 2 (ppr ty)

knownLoc (SST s _) = isGoodSrcSpan s
  
sstLHsBind :: Located (HsBind Id) -> [SrcSpanType]
sstLHsBind (L s (AbsBinds { abs_binds = binds })) = 
  concatMap sstLHsBind . bagToList $ binds
sstLHsBind (L s (FunBind { fun_id = lid, fun_matches = (MatchGroup lms ty) })) =
  SST s ty:(concatMap sstLMatch lms)
sstLHsBind (L s (VarBind { var_id = id, var_rhs = expr })) = SST s (idType id):(sstLHsExpr expr)
sstLHsBind (L s (PatBind { pat_lhs = (L _ pat), pat_rhs = GRHSs lgrhs lbinds})) = SST s (hsPatType pat):(concatMap sstLGRHS lgrhs ++ sstHsLocalBinds lbinds)

sstLMatch :: LMatch Id -> [SrcSpanType]
sstLMatch (L s (Match lpats _ (GRHSs lgrhs localBinds))) = 
  map (\ (L s pat) -> SST s (hsPatType pat)) lpats ++
  concatMap sstLGRHS lgrhs ++ sstHsLocalBinds localBinds

sstLGRHS :: LGRHS Id -> [SrcSpanType]
sstLGRHS (L s (GRHS guards expr)) = sstLHsExpr expr

sstHsLocalBinds :: HsLocalBinds Id -> [SrcSpanType]
sstHsLocalBinds localBinds = case localBinds of
  EmptyLocalBinds                   -> []
  HsValBinds (ValBindsIn binds _)   -> concatMap sstLHsBind (bagToList binds)
  HsValBinds (ValBindsOut rbinds _) -> concatMap sstLHsBind . concatMap (bagToList . snd) $ rbinds
  HsIPBinds _                       -> []


-- HsExpr

sstLHsExpr :: Located (HsExpr Id) -> [SrcSpanType]
sstLHsExpr (L s hsExpr)  = case hsExpr of
  HsVar id                  -> [SST s (idType id)]
  HsIPVar _                 -> []
  HsOverLit _               -> []
  HsLit _                   -> []
  HsLam (MatchGroup lms ty) -> concatMap sstLMatch lms
  HsApp e1 e2               -> sstLHsExpr e1 ++ sstLHsExpr e2
  OpApp e1 e2 _ e3          -> sstLHsExpr e1 ++ sstLHsExpr e2 ++ sstLHsExpr e3
  NegApp e _                -> sstLHsExpr e
  HsPar e                   -> sstLHsExpr e
  SectionL e1 e2            -> sstLHsExpr e1 ++ sstLHsExpr e2
  SectionR e1 e2            -> sstLHsExpr e1 ++ sstLHsExpr e2
  HsCase e (MatchGroup lms _) -> sstLHsExpr e ++ concatMap sstLMatch lms
  HsIf e1 e2 e3 -> sstLHsExpr e1 ++ sstLHsExpr e2 ++ sstLHsExpr e3
  HsLet lbinds e -> sstHsLocalBinds lbinds ++ sstLHsExpr e
  HsDo _ stmts e _ -> foldr sstLStmt (sstLHsExpr e) stmts  
  ExplicitList _ es -> concatMap sstLHsExpr es
  ExplicitPArr _ es -> concatMap sstLHsExpr es
  ExplicitTuple es _ -> concatMap sstLHsExpr es
  RecordCon _ _ recbinds -> [] -- FIXME
  RecordUpd e recbinds _ _ _ -> sstLHsExpr e -- FIXME
  ExprWithTySig e _ -> sstLHsExpr e
  ExprWithTySigOut e _ -> sstLHsExpr e
  ArithSeq _ _ -> [] -- FIXME
  PArrSeq _ _ -> [] -- FIXME
  HsSCC _ e -> sstLHsExpr e
  HsCoreAnn _ e -> sstLHsExpr e
  HsBracket _ -> [] -- FIXME
  HsBracketOut _ _ -> [] -- FIXME
  HsSpliceE _ -> [] -- FIXME
  HsProc _ _ -> [] -- FIXME
  HsArrApp e1 e2 _ _ _ -> sstLHsExpr e1 ++ sstLHsExpr e2
  HsArrForm e _ _ -> sstLHsExpr e
  HsTick _ _ e -> sstLHsExpr e
  HsBinTick _ _ e -> sstLHsExpr e
  HsTickPragma _ e -> sstLHsExpr e
  HsWrap _ e -> sstLHsExpr (L s e) -- FIXME
  
-- LStmt

sstLStmt :: LStmt Id -> [SrcSpanType] -> [SrcSpanType]
sstLStmt (L s stmt) i = case stmt of
  BindStmt (L s' pat) e _ _ -> sstLHsExpr e ++ [SST s' (hsPatType pat)] ++ i
  ExprStmt e _ _ -> sstLHsExpr e ++ i
  LetStmt lbinds -> sstHsLocalBinds lbinds ++ i
  _ -> i -- FIXME