utils/ext-core/Language/Core/CoreUtils.hs

   1 module Language.Core.CoreUtils where
   2
   3 import Language.Core.Core
   4 import Language.Core.Utils
   5 import Language.Core.Printer()
   6
   7 --import Debug.Trace
   8
   9 import Data.Generics
  10 import Data.List
  11 import Data.Maybe
  12
  13 splitDataConApp_maybe :: Exp -> Maybe (Qual Dcon, [Ty], [Exp])
  14 splitDataConApp_maybe (Dcon d) = Just (d, [], [])
  15 splitDataConApp_maybe (Appt rator t) =
  16    case splitDataConApp_maybe rator of
  17      Just (r, ts, rs) -> Just (r, ts ++ [t], rs)
  18      Nothing          -> Nothing
  19 splitDataConApp_maybe (App rator rand) =
  20   case splitDataConApp_maybe rator of
  21     Just (r, ts, rs) -> Just (r, ts, rs++[rand])
  22     Nothing -> Nothing
  23 splitDataConApp_maybe _ = Nothing
  24
  25 splitApp :: Exp -> (Exp, [Exp])
  26 splitApp (Appt rator _) = splitApp rator
  27 splitApp (App rator rand) =
  28   case splitApp rator of
  29     (r, rs) -> (r, rs++[rand])
  30 splitApp e = (e, [])
  31
  32 splitAppIgnoreCasts :: Exp -> (Exp, [Exp])
  33 splitAppIgnoreCasts (Appt rator _) = splitApp rator
  34 splitAppIgnoreCasts (App (Cast rator _) rand) = splitApp (App rator rand)
  35 splitAppIgnoreCasts (App rator rand) =
  36   case splitApp rator of
  37     (r, rs) -> (r, rs++[rand])
  38 splitAppIgnoreCasts e = (e, [])
  39
  40 splitFunTy_maybe :: Ty -> Maybe ([Ty], Ty)
  41 splitFunTy_maybe (Tforall _ t) = splitFunTy_maybe t
  42 splitFunTy_maybe t =
  43   case splitFunTy2_maybe t of
  44     Just (rator, rand) -> case splitFunTy_maybe rand of
  45                             Just (r,s) -> Just (rator:r, s)
  46                             Nothing -> Just ([rator], rand)
  47     Nothing -> Nothing
  48
  49 splitFunTy2_maybe :: Ty -> Maybe (Ty,Ty)
  50 splitFunTy2_maybe (Tapp (Tapp (Tcon c) t) u) | c == tcArrow = Just (t, u)
  51 splitFunTy2_maybe _ = Nothing
  52
  53 vdefNamesQ :: [Vdef] -> [Qual Var]
  54 vdefNamesQ = map (\ (Vdef (v,_,_)) -> v)
  55
  56 vdefNames :: [Vdef] -> [Var]
  57 vdefNames = snd . unzip . vdefNamesQ
  58
  59 vdefTys :: [Vdef] -> [Ty]
  60 vdefTys = map (\ (Vdef (_,t,_)) -> t)
  61
  62 vdefgNames :: Vdefg -> [Var]
  63 vdefgNames = snd . unzip . vdefgNamesQ
  64
  65 vdefgNamesQ :: Vdefg -> [Qual Var]
  66 vdefgNamesQ (Rec vds) = map (\ (Vdef (v,_,_)) -> v) vds
  67 vdefgNamesQ (Nonrec (Vdef (v,_,_))) = [v]
  68
  69 vdefgTys :: Vdefg -> [Ty]
  70 vdefgTys (Rec vds) = map (\ (Vdef (_,t,_)) -> t) vds
  71 vdefgTys (Nonrec (Vdef (_,t,_))) = [t]
  72 vdefgBodies :: Vdefg -> [Exp]
  73 vdefgBodies (Rec vds) = map (\ (Vdef (_,_,e)) -> e) vds
  74 vdefgBodies (Nonrec (Vdef (_,_,e))) = [e]
  75
  76 vbNames :: [Vbind] -> [Var]
  77 vbNames = fst . unzip
  78
  79 -- assumes v is not bound in e
  80 substIn :: Data a => Var -> Var -> a -> a
  81 substIn v newV = everywhereExcept (mkT frob)
  82   where frob (Var (Nothing,v1)) | v == v1   = Var (Nothing,newV)
  83         frob e                              = e
  84
  85 substVars :: Data a => [Var] -> [Var] -> a -> a
  86 substVars oldVars newVars e = foldl' (\ e1 (old,new) -> substIn old new e1)
  87   e (zip oldVars newVars)
  88
  89
  90 tdefNames :: [Tdef] -> [Qual Var]
  91 tdefNames = concatMap doOne
  92   where doOne (Data qtc _ cds) = qtc:(concatMap doCdef cds)
  93         doOne (Newtype qtc qtc1 _ _) = [qtc, qtc1]
  94         doCdef (Constr qdc _ _) = [qdc]
  95
  96 tdefDcons :: [Tdef] -> [Qual Var]
  97 tdefDcons = concatMap doOne
  98   where doOne (Data _ _ cds) = concatMap doCdef cds
  99         doOne _ = []
 100         doCdef (Constr qdc _ _) = [qdc]
 101
 102 tdefTcons :: [Tdef] -> [Qual Var]
 103 tdefTcons = concatMap doOne
 104   where doOne (Data qtc _ _) = [qtc]
 105         doOne (Newtype qtc qtc1 _ _) = [qtc, qtc1]
 106
 107 filterVdefgs :: (Vdef -> Bool) -> [Vdefg] -> [Vdefg]
 108 filterVdefgs ok = catMaybes . (map dropNames)
 109   where dropNames (Nonrec v) | not (ok v) = Nothing
 110         dropNames v@(Nonrec _) = Just v
 111         dropNames (Rec bs) = case filter ok bs of
 112            [] -> Nothing
 113            newBs -> Just (Rec newBs)
 114
 115 applyNewtype :: CoercionKind -> [Ty] -> (Ty,Ty)
 116 applyNewtype _d@(DefinedCoercion tbs (from,to)) tys =
 117   let (tvs,_) = unzip tbs in
 118     let res = (substl tvs tys from,substl tvs tys to) in
 119       -- trace ("co = " ++ show d ++ " args  = " ++ show tys ++ " res = " ++ show res) $
 120         res
 121
 122 {- Simultaneous substitution on types for type variables,
 123    renaming as neceessary to avoid capture.
 124    No checks for correct kindedness. -}
 125 substl :: [Tvar] -> [Ty] -> Ty -> Ty
 126 substl tvs ts t = f (zip tvs ts) t
 127   where
 128     f env t0 =
 129      case t0 of
 130        Tcon _ -> t0
 131        Tvar v -> case lookup v env of
 132                    Just t1 -> t1
 133                    Nothing -> t0
 134        Tapp t1 t2 -> Tapp (f env t1) (f env t2)
 135        Tforall (tv,k) t1 ->
 136          if tv `elem` free then
 137            Tforall (t',k) (f ((tv,Tvar t'):env) t1)
 138          else
 139            Tforall (tv,k) (f (filter ((/=tv).fst) env) t1)
 140        TransCoercion t1 t2 -> TransCoercion (f env t1) (f env t2)
 141        SymCoercion t1 -> SymCoercion (f env t1)
 142        UnsafeCoercion t1 t2 -> UnsafeCoercion (f env t1) (f env t2)
 143        LeftCoercion t1 -> LeftCoercion (f env t1)
 144        RightCoercion t1 -> RightCoercion (f env t1)
 145        InstCoercion t1 t2 -> InstCoercion (f env t1) (f env t2)
 146      where free = foldr union [] (map (freeTvars.snd) env)
 147            t' = freshTvar free
 148
 149
 150 {- Return free tvars in a type -}
 151 freeTvars :: Ty -> [Tvar]
 152 freeTvars (Tcon _) = []
 153 freeTvars (Tvar v) = [v]
 154 freeTvars (Tapp t1 t2) = freeTvars t1 `union` freeTvars t2
 155 freeTvars (Tforall (t,_) t1) = delete t (freeTvars t1)
 156 freeTvars (TransCoercion t1 t2) = freeTvars t1 `union` freeTvars t2
 157 freeTvars (SymCoercion t) = freeTvars t
 158 freeTvars (UnsafeCoercion t1 t2) = freeTvars t1 `union` freeTvars t2
 159 freeTvars (LeftCoercion t) = freeTvars t
 160 freeTvars (RightCoercion t) = freeTvars t
 161 freeTvars (InstCoercion t1 t2) = freeTvars t1 `union` freeTvars t2
 162
 163 {- Return any tvar *not* in the argument list. -}
 164 freshTvar :: [Tvar] -> Tvar
 165 freshTvar tvs = maximum ("":tvs) ++ "x" -- one simple way!
 166
 167 splitLambda :: Exp -> ([Bind],Exp)
 168 splitLambda (Lam vb e) = case splitLambda e of
 169   (vbs,rhs) -> (vb:vbs,rhs)
 170 splitLambda (Note _ e) = splitLambda e
 171 splitLambda e          = ([],e)
 172
 173 vbinds :: [Bind] -> [(Var,Ty)]
 174 vbinds = foldl' stuff []
 175   where stuff :: [(Var,Ty)] -> Bind -> [(Var,Ty)]
 176         stuff rest (Tb _) = rest
 177         stuff rest (Vb p) = p:rest
 178
 179 splitBinds :: [Bind] -> ([(Tvar,Kind)],[(Var,Ty)])
 180 splitBinds = foldr stuff ([],[])
 181   where stuff (Tb t) (tbs,vbs) = (t:tbs,vbs)
 182         stuff (Vb v) (tbs,vbs) = (tbs,v:vbs)
 183
 184 freeVars :: Exp -> [Qual Var]
 185 freeVars (Var v)                    = [v]
 186 freeVars (Dcon _)                   = []
 187 freeVars (Lit _)                    = []
 188 freeVars (App f g)                  = freeVars f `union` freeVars g
 189 freeVars (Appt e _)                 = freeVars e
 190 freeVars (Lam (Tb _) e)             = freeVars e
 191 freeVars (Lam (Vb (v,_)) e)         = delete (unqual v) (freeVars e)
 192 freeVars (Let (Nonrec (Vdef (v,_,rhs))) e) = freeVars rhs `union` (delete v (freeVars e))
 193 freeVars (Let r@(Rec _) e)         = (freeVars e \\ boundVars) `union` (freeVarss rhss \\ boundVars)
 194   where boundVars = map unqual $ vdefgNames r
 195         rhss      = vdefgBodies r
 196 freeVars (Case e (v,_) _ alts)      = freeVars e `union` (delete v1 (boundVarsAlts alts))
 197   where v1 = unqual v
 198         boundVarsAlts as = freeVarss rhss \\ (v1:caseVars)
 199           where rhss = map (\ a -> case a of
 200                              Acon _ _ _ r -> r
 201                              Alit _ r     -> r
 202                              Adefault r   -> r) as
 203                 caseVars = foldl' union [] (map (\ a -> case a of
 204                                                Acon _ _ vbs _ ->
 205                                                  (map unqual (fst (unzip vbs)))
 206                                                _              -> []) as)
 207 freeVars (Cast e _)                 = freeVars e
 208 freeVars (Note _ e)                 = freeVars e
 209 freeVars (External {})              = []
 210
 211 freeVarss :: [Exp] -> [Qual Var]
 212 freeVarss = foldl' union [] . map freeVars