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 Control.Monad.State
  17 import Data.Either
  18 import Data.List
  19 import Data.Generics
  20 import qualified Data.Map as M
  21
  22 import Language.Core.Core
  23 import Language.Core.Env
  24 import Language.Core.Check
  25 import Language.Core.Environments
  26 import Language.Core.Utils
  27
  28 prepModule :: Menv -> Module -> Module
  29 prepModule globalEnv (Module mn tdefs vdefgs) =
  30     Module mn tdefs (snd (evalState
  31       (foldM prepTopVdefg (eempty,[]) vdefgs) initCounter))
  32   where
  33     (tcenv, cenv) = mkTypeEnvsNoChecking tdefs
  34
  35     prepTopVdefg :: (Venv, [Vdefg]) -> Vdefg -> PrepM (Venv, [Vdefg])
  36     prepTopVdefg (venv,vdefgs) vdefg = do
  37          (venv',vdefg') <- prepVdefg (venv,eempty) vdefg
  38          return (venv',vdefgs ++ [vdefg'])
  39
  40     prepVdefg (env@(venv,_)) (Nonrec(Vdef((Nothing,x),t,e))) = do
  41         e' <- prepExp env e
  42         return (eextend venv (x,t), Nonrec(Vdef((Nothing,x),t,e')))
  43     prepVdefg (env@(venv,_))  (Nonrec(Vdef(qx,t,e))) = do
  44         e' <- prepExp env e
  45         return (venv, Nonrec(Vdef(qx,t,e')))
  46     prepVdefg (venv,tvenv) (Rec vdefs) = do
  47         vds' <- mapM (\ (Vdef (qx,t,e)) -> do
  48                          e' <- prepExp (venv',tvenv) e
  49                          return (Vdef (qx,t,e'))) vdefs
  50         return (venv', Rec vds')
  51         where venv' = foldl' eextend venv [(x,t) | Vdef((Nothing,x),t,_) <- vdefs]
  52
  53     prepExp :: (Venv, Tvenv) -> Exp -> PrepM Exp
  54     prepExp _ (Var qv) = return $ Var qv
  55     prepExp _ (Dcon qdc) = return $ Dcon qdc
  56     prepExp _ (Lit l) = return $ Lit l
  57     prepExp env e@(App _ _) = unwindApp env e []
  58     prepExp env e@(Appt _ _) = unwindApp env e []
  59     prepExp (venv,tvenv) (Lam (Vb vb) e) = do
  60        e' <- prepExp (eextend venv vb,tvenv) e
  61        return $ Lam (Vb vb) e'
  62     prepExp (venv,tvenv) (Lam (Tb tb) e) = do
  63        e' <- prepExp (venv,eextend tvenv tb) e
  64        return $ Lam (Tb tb) e'
  65     prepExp env@(venv,tvenv) (Let (Nonrec(Vdef((Nothing,x),t,b))) e)
  66         | (kindOfTy tvenv t `eqKind` Kunlifted && suspends b) = do
  67             -- There are two places where we call the typechecker, one of them
  68             -- here.
  69             -- We need to know the type of the let body in order to construct
  70             -- a case expression.
  71                                 -- need to extend the env with the let-bound var too!
  72             scrut' <- prepExp env b
  73             rhs' <- prepExp (eextend venv (x,t),tvenv) e
  74             return $
  75               let eTy = typeOfExp (eextend venv (x, t)) tvenv e in
  76                 Case scrut' (x,t) eTy [Adefault rhs']
  77     prepExp (venv,tvenv) (Let vdefg e) =  do
  78       (venv',vdefg') <- prepVdefg (venv,tvenv) vdefg
  79       rhs' <- prepExp (venv',tvenv) e
  80       return $ Let vdefg' rhs'
  81     prepExp env@(venv,tvenv) (Case e vb t alts) = do
  82       e' <- prepExp env e
  83       alts' <- mapM (prepAlt (eextend venv vb,tvenv)) alts
  84       return $ Case e' vb t alts'
  85     prepExp env (Cast e t) = do
  86       e' <- prepExp env e
  87       return $ Cast e' t
  88     prepExp env (Note s e) = do
  89       e' <- prepExp env e
  90       return $ Note s e'
  91     prepExp _ (External s t) = return $ External s t
  92
  93     prepAlt :: (Venv,Tvenv) -> Alt -> PrepM Alt
  94     prepAlt (venv,tvenv) (Acon qdc tbs vbs e) = do
  95       rhs' <- prepExp (foldl' eextend venv vbs,foldl' eextend tvenv tbs) e
  96       return $ Acon qdc tbs vbs rhs'
  97     prepAlt env (Alit l e) = (liftM (Alit l)) (prepExp env e)
  98     prepAlt env (Adefault e) = (liftM Adefault) (prepExp env e)
  99
 100     ntEnv = mkNtEnv globalEnv
 101
 102     unwindApp :: (Venv, Tvenv) -> Exp -> [Either Exp Ty] -> PrepM Exp
 103     unwindApp env (App e1 e2) as = unwindApp env e1 (Left e2:as)
 104     unwindApp env (Appt e t) as  = unwindApp env e (Right t:as)
 105     unwindApp env (op@(Dcon qdc)) as = do
 106         e' <- rewindApp env op as
 107         -- possibly dubious to assume no type args
 108         etaExpand [] (drop n atys) e'
 109         where (tbs,atys0,_) = splitTy (qlookup cenv_ eempty qdc)
 110               atys = map (substl (map fst tbs) ts) atys0
 111               ts = [t | Right t <- as]
 112               n = length [e | Left e <- as]
 113     unwindApp env (op@(Var(qv@(_,p)))) as | isPrimVar qv = do
 114         e' <- rewindApp env op as
 115         (liftM k) $ etaExpand (snd (unzip extraTbs)) (drop n atys) (k1 e')
 116         where -- TODO: avoid copying code. these two cases are the same
 117
 118               -- etaExpand needs to add the type arguments too! Bah!
 119               primEnv = case elookup globalEnv primMname of
 120                               Just es -> venv_ es
 121                               _       -> error "eek"
 122               (_, _, resTy') = (maybe (error "unwindApp") splitTy (elookup primEnv p))
 123               (tbs, atys0, _resTy) = (maybe (error "unwindApp") (splitTy . (substNewtys ntEnv)) (elookup primEnv p))
 124               -- The magic here is so we know to eta-expand applications of
 125               -- primops whose return types are newtypes.
 126               -- There are no actual GHC primops that have this property, but
 127               -- a back-end tool writer (for example: me) might want to add
 128               -- such a primop.
 129               -- If this code wasn't here, and we had a primop
 130               -- foo# :: Int -> IO (),
 131               -- we would see (foo# 5) and think it was fully applied, when
 132               -- actually we need to rewrite it as:
 133               -- (\ (s::State# RealWorld#) -> foo# 5 s)
 134               -- (This code may be a very good case against introducing such
 135               -- primops.)
 136                   -- tim 10/29/2008: I think this is no longer necessary.
 137                   -- hPutChar now has a (#wub,blub#) return type.
 138               (k,k1) = case newtypeCoercion_maybe ntEnv resTy' of
 139                          Just co -> case splitTyConApp_maybe resTy' of
 140                                       Just (_, args) -> ((\ e -> Cast e (SymCoercion (mkTapp co args))), (\ e1 -> Cast e1 (mkTapp co args)))
 141                                       _ -> ((\ e -> Cast e (SymCoercion co)), (\ e1 -> Cast e1 co))
 142                          _       -> (id,id)
 143               n_args = length ts
 144               (appliedTbs, extraTbs) = (take n_args tbs, drop n_args tbs)
 145               atys = map (substl (map fst appliedTbs) ts) atys0
 146               ts = [t | Right t <- as]
 147               n = length [e | Left e <- as]
 148     unwindApp env (op@(External _ t)) as = do
 149         e' <- rewindApp env op as
 150         etaExpand [] (drop n atys) e'
 151           where (_,atys,_) = splitTy t
 152                 n = length as -- assumes all args are term args
 153     unwindApp env op as = rewindApp env op as
 154
 155
 156     etaExpand :: [Kind] -> [Ty] -> Exp -> PrepM Exp
 157     etaExpand ks ts e = do
 158          -- what a pain
 159          tyvs <- replicateM (length ks) freshVar
 160          termvs <- replicateM (length ts) freshVar
 161          let tyArgs   = zip tyvs ks
 162          let termArgs = zip termvs ts
 163          return $
 164           foldr (\ (t1,k1) e -> Lam (Tb (t1,k1)) e)
 165            (foldr (\ (v,t) e -> Lam (Vb (v,t)) e)
 166               (foldl' (\ e (v,_) -> App e (Var (unqual v)))
 167                  (foldl' (\ e (tv,_) -> Appt e (Tvar tv))
 168                    e tyArgs)
 169               termArgs) termArgs)
 170            tyArgs
 171
 172     rewindApp :: (Venv, Tvenv) -> Exp -> [Either Exp Ty] -> PrepM Exp
 173     rewindApp _ e [] = return e
 174     rewindApp env@(venv,tvenv) e1 (Left e2:as) | kindOfTy tvenv t `eqKind` Kunlifted && suspends e2 = do
 175         v <- freshVar
 176         let venv' = eextend venv (v,t)
 177         rhs <- rewindApp (venv', tvenv) (App e1 (Var (unqual v))) as
 178         newScrut <- prepExp env e2
 179        -- This is the other place where we call the typechecker.
 180         return $ Case newScrut (v,t) (typeOfExp venv' tvenv rhs) [Adefault rhs]
 181         where t = typeOfExp venv tvenv e2
 182     rewindApp env e1 (Left e2:as) = do
 183       e2' <- prepExp env e2
 184       rewindApp env (App e1 e2') as
 185     rewindApp env e (Right t:as) = rewindApp env (Appt e t) as
 186
 187     typeOfExp :: Venv -> Tvenv -> Exp -> Ty
 188     typeOfExp = checkExpr mn globalEnv tcenv cenv
 189
 190     kindOfTy :: Tvenv -> Ty -> Kind
 191     kindOfTy tvenv = checkType mn globalEnv tcenv tvenv
 192
 193     {- Return false for those expressions for which Interp.suspendExp builds a thunk. -}
 194     suspends (Var _) = False
 195     suspends (Lit _) = False
 196     suspends (Lam (Vb _) _) = False
 197     suspends (Lam _ e) = suspends e
 198     suspends (Appt e _) = suspends e
 199     suspends (Cast e _) = suspends e
 200     suspends (Note _ e) = suspends e
 201     suspends (External _ _) = False
 202     suspends _ = True
 203
 204     mlookup :: (Envs -> Env a b) -> Env a b -> Mname -> Env a b
 205     mlookup _ local_env Nothing = local_env
 206     mlookup selector _  (Just m) =
 207       case elookup globalEnv m of
 208         Just env -> selector env
 209         Nothing -> error ("Prep: undefined module name: " ++ show m)
 210
 211     qlookup ::  (Ord a, Show a) => (Envs -> Env a b) -> Env a b -> (Mname,a) -> b
 212     qlookup selector local_env (m,k) =
 213       case elookup (mlookup selector local_env m) k of
 214         Just v -> v
 215         Nothing -> error ("undefined identifier: " ++ show k)
 216
 217 boundVars :: Exp -> [Id]
 218 boundVars (Lam (Vb (v,_)) e) = [v] `union` boundVars e
 219 boundVars (Lam _ e) = boundVars e
 220 boundVars (Let vds e) = (boundVarsVdefs vds) `union` boundVars e
 221 boundVars (Case scrut (v,_) _ alts) =
 222    [v] `union` (boundVars scrut) `union` boundVarsAlts alts
 223 boundVars (Cast e _) = boundVars e
 224 boundVars (Note _ e) = boundVars e
 225 boundVars (App e1 e2) = boundVars e1 `union` boundVars e2
 226 boundVars (Appt e _) = boundVars e
 227 boundVars _ = []
 228
 229 boundVarsVdefs :: Vdefg -> [Id]
 230 boundVarsVdefs (Rec vds) = nub (concatMap boundVarsVdef vds)
 231 boundVarsVdefs (Nonrec vd) = boundVarsVdef vd
 232
 233 boundVarsVdef :: Vdef -> [Id]
 234 boundVarsVdef (Vdef ((_,v),_,e)) = [v] `union` boundVars e
 235
 236 boundVarsAlts :: [Alt] -> [Var]
 237 boundVarsAlts as = nub (concatMap boundVarsAlt as)
 238
 239 boundVarsAlt :: Alt -> [Var]
 240 boundVarsAlt (Acon _ _ vbs e) = (map fst vbs) `union` (boundVars e)
 241 boundVarsAlt (Alit _ e) = boundVars e
 242 boundVarsAlt (Adefault e) = boundVars e
 243
 244 mkNtEnv :: Menv -> NtEnv
 245 mkNtEnv menv =
 246   foldl M.union M.empty $
 247         map (\ (mn,e) ->
 248                  foldr (\ (key,thing) rest ->
 249                             case thing of
 250                               Kind _ -> rest
 251                               Coercion (DefinedCoercion _ (lhs,rhs)) ->
 252                                   case splitTyConApp_maybe lhs of
 253                                     Just ((_,tc1),_) -> M.insert tc1 (rhs,Tcon (Just mn, key)) rest
 254                                     _ -> rest) M.empty (etolist (tcenv_ e))) (etolist menv)
 255
 256 substNewtys :: NtEnv -> Ty -> Ty
 257 substNewtys ntEnv = everywhere'Except (mkT go)
 258                  where go t | Just ((_,tc),_) <- splitTyConApp_maybe t =
 259                          case M.lookup tc ntEnv of
 260                            Just (rhs,_) -> rhs
 261                            Nothing  -> t
 262                        go t = t
 263
 264 newtypeCoercion_maybe :: NtEnv -> Ty -> Maybe Ty
 265 newtypeCoercion_maybe ntEnv t | Just ((_,tc),_) <- splitTyConApp_maybe t =
 266   case M.lookup tc ntEnv of
 267     Just (_, coercion) -> Just coercion
 268     Nothing               -> Nothing
 269 newtypeCoercion_maybe _ _ = Nothing
 270
 271 -- first element: rep type
 272 -- second element: coercion tcon
 273 type NtEnv  = M.Map Tcon (Ty, Ty)
 274
 275 mkTapp :: Ty -> [Ty] -> Ty
 276 mkTapp = foldl Tapp
 277
 278 initCounter :: Int
 279 initCounter = 0
 280
 281 type PrepM = State Int
 282
 283 freshVar :: PrepM String
 284 freshVar = do
 285   i <- get
 286   put (i+1)
 287   return $ ("zd" ++ show i)