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

   1 {-# OPTIONS -fno-warn-name-shadowing #-}
   2 {-
   3 Preprocess a module to normalize it in the following ways:
   4         (1) Saturate all constructor and primop applications.
   5               (as well as external calls; this is probably already
   6                guaranteed, but paranoia is good)
   7         (2) Arrange that any non-trivial expression of unlifted kind ('#')
   8              is turned into the scrutinee of a Case.
   9 After these preprocessing steps, Core can be interpreted (or given an operational semantics)
  10       ignoring type information almost completely.
  11 -}
  12
  13
  14 module Language.Core.Prep where
  15
  16 import Data.Either
  17 import Data.List
  18 import Data.Generics
  19 import qualified Data.Map as M
  20
  21 import Language.Core.Core
  22 import Language.Core.Env
  23 import Language.Core.Check
  24 import Language.Core.Environments
  25 import Language.Core.Encoding
  26 import Language.Core.Utils
  27
  28 prepModule :: Menv -> Module -> Module
  29 prepModule globalEnv (Module mn tdefs vdefgs) =
  30     Module mn tdefs vdefgs'
  31   where
  32
  33     (tcenv, cenv) = mkTypeEnvsNoChecking tdefs
  34     (_,vdefgs') = foldl' prepTopVdefg (eempty,[]) vdefgs
  35
  36     prepTopVdefg (venv,vdefgs) vdefg = (venv',vdefgs ++ [vdefg'])
  37        where (venv',vdefg') = prepVdefg (venv,eempty) vdefg
  38
  39     prepVdefg (env@(venv,_)) (Nonrec(Vdef((Nothing,x),t,e))) =
  40         (eextend venv (x,t), Nonrec(Vdef((Nothing,x),t,prepExp env e)))
  41     prepVdefg (env@(venv,_))  (Nonrec(Vdef(qx,t,e))) =
  42         (venv, Nonrec(Vdef(qx,t,prepExp env e)))
  43     prepVdefg (venv,tvenv) (Rec vdefs) =
  44         (venv',Rec [ Vdef(qx,t,prepExp (venv',tvenv) e) | Vdef(qx,t,e) <- vdefs])
  45         where venv' = foldl' eextend venv [(x,t) | Vdef((Nothing,x),t,_) <- vdefs]
  46
  47     prepExp _ (Var qv) = Var qv
  48     prepExp _ (Dcon qdc) = Dcon qdc
  49     prepExp _ (Lit l) = Lit l
  50     prepExp env e@(App _ _) = unwindApp env e []
  51     prepExp env e@(Appt _ _) = unwindApp env e []
  52     prepExp (venv,tvenv) (Lam (Vb vb) e) = Lam (Vb vb) (prepExp (eextend venv vb,tvenv) e)
  53     prepExp (venv,tvenv) (Lam (Tb tb) e) = Lam (Tb tb) (prepExp (venv,eextend tvenv tb) e)
  54     prepExp env@(venv,tvenv) (Let (Nonrec(Vdef((Nothing,x),t,b))) e)
  55         | (kindOfTy tvenv t `eqKind` Kunlifted && suspends b) =
  56             -- There are two places where we call the typechecker, one of them
  57             -- here.
  58             -- We need to know the type of the let body in order to construct
  59             -- a case expression.
  60                                 -- need to extend the env with the let-bound var too!
  61             let eTy = typeOfExp (eextend venv (x, t)) tvenv e in
  62                Case (prepExp env b) (x,t)
  63                   eTy
  64                   [Adefault (prepExp (eextend venv (x,t),tvenv) e)]
  65     prepExp (venv,tvenv) (Let vdefg e) =  Let vdefg' (prepExp (venv',tvenv) e)
  66                 where (venv',vdefg') = prepVdefg (venv,tvenv) vdefg
  67     prepExp env@(venv,tvenv) (Case e vb t alts) = Case (prepExp env e) vb t (map (prepAlt (eextend venv vb,tvenv)) alts)
  68     prepExp env (Cast e t) = Cast (prepExp env e) t
  69     prepExp env (Note s e) = Note s (prepExp env e)
  70     prepExp _ (External s t) = External s t
  71
  72     prepAlt (venv,tvenv) (Acon qdc tbs vbs e) = Acon qdc tbs vbs (prepExp (foldl' eextend venv vbs,foldl' eextend tvenv tbs) e)
  73     prepAlt env (Alit l e) = Alit l (prepExp env e)
  74     prepAlt env (Adefault e) = Adefault (prepExp env e)
  75
  76     ntEnv = mkNtEnv globalEnv
  77
  78     unwindApp env (App e1 e2) as = unwindApp env e1 (Left e2:as)
  79     unwindApp env (Appt e t) as  = unwindApp env e (Right t:as)
  80     unwindApp env (op@(Dcon qdc)) as =
  81         -- possibly dubious to assume no type args
  82         etaExpand [] (drop n atys) (rewindApp env op as)
  83         where (tbs,atys0,_) = splitTy (qlookup cenv_ eempty qdc)
  84               atys = map (substl (map fst tbs) ts) atys0
  85               ts = [t | Right t <- as]
  86               n = length [e | Left e <- as]
  87     unwindApp env (op@(Var(qv@(_,p)))) as | isPrimVar qv =
  88         k $ etaExpand (snd (unzip extraTbs)) (drop n atys) (k1 (rewindApp env op as))
  89         where -- TODO: avoid copying code. these two cases are the same
  90
  91               -- etaExpand needs to add the type arguments too! Bah!
  92               primEnv = case elookup globalEnv primMname of
  93                               Just es -> venv_ es
  94                               _       -> error "eek"
  95               (_, _, resTy') = (maybe (error "unwindApp") splitTy (elookup primEnv p))
  96               (tbs, atys0, _resTy) = (maybe (error "unwindApp") (splitTy . (substNewtys ntEnv)) (elookup primEnv p))
  97               -- The magic here is so we know to eta-expand applications of
  98               -- primops whose return types are newtypes.
  99               -- There are no actual GHC primops that have this property, but
 100               -- a back-end tool writer (for example: me) might want to add
 101               -- such a primop.
 102               -- If this code wasn't here, and we had a primop
 103               -- foo# :: Int -> IO (),
 104               -- we would see (foo# 5) and think it was fully applied, when
 105               -- actually we need to rewrite it as:
 106               -- (\ (s::State# RealWorld#) -> foo# 5 s)
 107               -- (This code may be a very good case against introducing such
 108               -- primops.)
 109               (k,k1) = case newtypeCoercion_maybe ntEnv resTy' of
 110                          Just co -> case splitTyConApp_maybe resTy' of
 111                                       Just (_, args) -> ((\ e -> Cast e (SymCoercion (mkTapp co args))), (\ e1 -> Cast e1 (mkTapp co args)))
 112                                       _ -> ((\ e -> Cast e (SymCoercion co)), (\ e1 -> Cast e1 co))
 113                          _       -> (id,id)
 114               n_args = length ts
 115               (appliedTbs, extraTbs) = (take n_args tbs, drop n_args tbs)
 116               atys = map (substl (map fst appliedTbs) ts) atys0
 117               ts = [t | Right t <- as]
 118               n = length [e | Left e <- as]
 119     unwindApp env (op@(External _ t)) as =
 120         etaExpand [] (drop n atys) (rewindApp env op as)
 121           where (_,atys,_) = splitTy t
 122                 n = length as -- assumes all args are term args
 123     unwindApp env op as = rewindApp env op as
 124
 125
 126     etaExpand :: [Kind] -> [Ty] -> Exp -> Exp
 127     etaExpand ks ts e =
 128          -- what a pain
 129          let tyArgs = [(zEncodeString $ "$t_"++(show i),k) | (i, k) <- zip [(1::Integer)..] ks]
 130              termArgs = [ (zEncodeString $ '$':(show i),t) | (i,t) <- zip [(1::Integer)..] ts] in
 131           foldr (\ (t1,k1) e -> Lam (Tb (t1,k1)) e)
 132            (foldr (\ (v,t) e -> Lam (Vb (v,t)) e)
 133               (foldl' (\ e (v,_) -> App e (Var (unqual v)))
 134                  (foldl' (\ e (tv,_) -> Appt e (Tvar tv))
 135                    e tyArgs)
 136               termArgs) termArgs)
 137            tyArgs
 138
 139     rewindApp _ e [] = e
 140     rewindApp env@(venv,tvenv) e1 (Left e2:as) | kindOfTy tvenv t `eqKind` Kunlifted && suspends e2 =
 141        -- This is the other place where we call the typechecker.
 142         Case newScrut (v,t) (typeOfExp venv' tvenv rhs) [Adefault rhs]
 143         where newScrut = prepExp env e2
 144               rhs = (rewindApp (venv', tvenv) (App e1 (Var (unqual v))) as)
 145                  -- note:
 146                  -- e1 gets moved inside rhs. so if we pick a case
 147                  -- var name (outside e1) equal to a name bound *inside*
 148                  -- e1, the binding *inside* e1 will shadow "v"
 149                  -- Which would be name capture!
 150                  -- So, we pass the bound vars of e1 to freshVar along with
 151                  -- the domain of the current env.
 152               v = freshVar (edomain venv `union` (boundVars e1))
 153               t = typeOfExp venv tvenv e2
 154               venv' = eextend venv (v,t)
 155     rewindApp env e1 (Left e2:as) = rewindApp env (App e1 (prepExp env e2)) as
 156     rewindApp env e (Right t:as) = rewindApp env (Appt e t) as
 157
 158     freshVar vs = maximum ("":vs) ++ "x" -- one simple way!
 159
 160     typeOfExp :: Venv -> Tvenv -> Exp -> Ty
 161     typeOfExp = checkExpr mn globalEnv tcenv cenv
 162
 163     kindOfTy :: Tvenv -> Ty -> Kind
 164     kindOfTy tvenv = checkType mn globalEnv tcenv tvenv
 165
 166     {- Return false for those expressions for which Interp.suspendExp builds a thunk. -}
 167     suspends (Var _) = False
 168     suspends (Lit _) = False
 169     suspends (Lam (Vb _) _) = False
 170     suspends (Lam _ e) = suspends e
 171     suspends (Appt e _) = suspends e
 172     suspends (Cast e _) = suspends e
 173     suspends (Note _ e) = suspends e
 174     suspends (External _ _) = False
 175     suspends _ = True
 176
 177     mlookup :: (Envs -> Env a b) -> Env a b -> Mname -> Env a b
 178     mlookup _ local_env Nothing = local_env
 179     mlookup selector _  (Just m) =
 180       case elookup globalEnv m of
 181         Just env -> selector env
 182         Nothing -> error ("Prep: undefined module name: " ++ show m)
 183
 184     qlookup ::  (Ord a, Show a) => (Envs -> Env a b) -> Env a b -> (Mname,a) -> b
 185     qlookup selector local_env (m,k) =
 186       case elookup (mlookup selector local_env m) k of
 187         Just v -> v
 188         Nothing -> error ("undefined identifier: " ++ show k)
 189
 190 boundVars :: Exp -> [Id]
 191 boundVars (Lam (Vb (v,_)) e) = [v] `union` boundVars e
 192 boundVars (Lam _ e) = boundVars e
 193 boundVars (Let vds e) = (boundVarsVdefs vds) `union` boundVars e
 194 boundVars (Case scrut (v,_) _ alts) =
 195    [v] `union` (boundVars scrut) `union` boundVarsAlts alts
 196 boundVars (Cast e _) = boundVars e
 197 boundVars (Note _ e) = boundVars e
 198 boundVars (App e1 e2) = boundVars e1 `union` boundVars e2
 199 boundVars (Appt e _) = boundVars e
 200 boundVars _ = []
 201
 202 boundVarsVdefs :: Vdefg -> [Id]
 203 boundVarsVdefs (Rec vds) = nub (concatMap boundVarsVdef vds)
 204 boundVarsVdefs (Nonrec vd) = boundVarsVdef vd
 205
 206 boundVarsVdef :: Vdef -> [Id]
 207 boundVarsVdef (Vdef ((_,v),_,e)) = [v] `union` boundVars e
 208
 209 boundVarsAlts :: [Alt] -> [Var]
 210 boundVarsAlts as = nub (concatMap boundVarsAlt as)
 211
 212 boundVarsAlt :: Alt -> [Var]
 213 boundVarsAlt (Acon _ _ vbs e) = (map fst vbs) `union` (boundVars e)
 214 boundVarsAlt (Alit _ e) = boundVars e
 215 boundVarsAlt (Adefault e) = boundVars e
 216
 217 mkNtEnv :: Menv -> NtEnv
 218 mkNtEnv menv =
 219   foldl M.union M.empty $
 220         map (\ (mn,e) ->
 221                  foldr (\ (key,thing) rest ->
 222                             case thing of
 223                               Kind _ -> rest
 224                               Coercion (DefinedCoercion _ (lhs,rhs)) ->
 225                                   case splitTyConApp_maybe lhs of
 226                                     Just ((_,tc1),_) -> M.insert tc1 (rhs,Tcon (Just mn, key)) rest
 227                                     _ -> rest) M.empty (etolist (tcenv_ e))) (etolist menv)
 228
 229 substNewtys :: NtEnv -> Ty -> Ty
 230 substNewtys ntEnv = everywhere'Except (mkT go)
 231                  where go t | Just ((_,tc),_) <- splitTyConApp_maybe t =
 232                          case M.lookup tc ntEnv of
 233                            Just (rhs,_) -> rhs
 234                            Nothing  -> t
 235                        go t = t
 236
 237 newtypeCoercion_maybe :: NtEnv -> Ty -> Maybe Ty
 238 newtypeCoercion_maybe ntEnv t | Just ((_,tc),_) <- splitTyConApp_maybe t =
 239   case M.lookup tc ntEnv of
 240     Just (_, coercion) -> Just coercion
 241     Nothing               -> Nothing
 242 newtypeCoercion_maybe _ _ = Nothing
 243
 244 -- first element: rep type
 245 -- second element: coercion tcon
 246 type NtEnv  = M.Map Tcon (Ty, Ty)
 247
 248 mkTapp :: Ty -> [Ty] -> Ty
 249 mkTapp = foldl Tapp