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 Debug.Trace
  17
  18 import Control.Monad.State
  19 import Data.Either
  20 import Data.List
  21 import Data.Generics
  22 import qualified Data.Map as M
  23
  24 import Language.Core.Core
  25 import Language.Core.CoreUtils
  26 import Language.Core.Env
  27 import Language.Core.Check
  28 import Language.Core.Environments
  29 import Language.Core.Utils
  30
  31 prepModule :: Menv -> Module -> Module
  32 prepModule globalEnv (Module mn tdefs vdefgs) =
  33     Module mn tdefs (snd (evalState
  34       (foldM prepTopVdefg (eempty,[]) vdefgs) initCounter))
  35   where
  36     (tcenv, cenv) = mkTypeEnvsNoChecking tdefs
  37
  38     prepTopVdefg :: (Venv, [Vdefg]) -> Vdefg -> PrepM (Venv, [Vdefg])
  39     prepTopVdefg (venv,vdefgs) vdefg = do
  40          (venv',vdefg') <- prepVdefg (venv,eempty) vdefg
  41          return (venv',vdefgs ++ [vdefg'])
  42
  43     prepVdefg (env@(venv,_)) (Nonrec(Vdef((Nothing,x),t,e))) = do
  44         e' <- prepExp env e
  45         return (eextend venv (x,t), Nonrec(Vdef((Nothing,x),t,e')))
  46     prepVdefg (env@(venv,_))  (Nonrec(Vdef(qx,t,e))) = do
  47         e' <- prepExp env e
  48         return (venv, Nonrec(Vdef(qx,t,e')))
  49     prepVdefg (venv,tvenv) (Rec vdefs) = do
  50         vds' <- mapM (\ (Vdef (qx,t,e)) -> do
  51                          e' <- prepExp (venv',tvenv) e
  52                          return (Vdef (qx,t,e'))) vdefs
  53         return (venv', Rec vds')
  54         where venv' = foldl' eextend venv [(x,t) | Vdef((Nothing,x),t,_) <- vdefs]
  55
  56     prepExp :: (Venv, Tvenv) -> Exp -> PrepM Exp
  57     prepExp _ (Var qv) = return $ Var qv
  58     prepExp _ (Dcon qdc) = return $ Dcon qdc
  59     prepExp _ (Lit l) = return $ Lit l
  60     prepExp env e@(App _ _) = unwindApp env e []
  61     prepExp env e@(Appt _ _) = unwindApp env e []
  62     prepExp (venv,tvenv) (Lam (Vb vb) e) = do
  63        e' <- prepExp (eextend venv vb,tvenv) e
  64        return $ Lam (Vb vb) e'
  65     prepExp (venv,tvenv) (Lam (Tb tb) e) = do
  66        e' <- prepExp (venv,eextend tvenv tb) e
  67        return $ Lam (Tb tb) e'
  68     prepExp env@(venv,tvenv) (Let (Nonrec(Vdef((Nothing,x),t,b))) e)
  69         | (kindOfTy tvenv t `eqKind` Kunlifted && suspends b) = do
  70             -- There are two places where we call the typechecker, one of them
  71             -- here.
  72             -- We need to know the type of the let body in order to construct
  73             -- a case expression.
  74                                 -- need to extend the env with the let-bound var too!
  75             scrut' <- prepExp env b
  76             rhs' <- prepExp (eextend venv (x,t),tvenv) e
  77             return $
  78               let eTy = typeOfExp (eextend venv (x, t)) tvenv e in
  79                 Case scrut' (x,t) eTy [Adefault rhs']
  80     prepExp (venv,tvenv) (Let vdefg e) =  do
  81       (venv',vdefg') <- prepVdefg (venv,tvenv) vdefg
  82       rhs' <- prepExp (venv',tvenv) e
  83       return $ Let vdefg' rhs'
  84     prepExp env@(venv,tvenv) (Case e vb t alts) = do
  85       e' <- prepExp env e
  86       alts' <- mapM (prepAlt (eextend venv vb,tvenv)) alts
  87       return $ Case e' vb t alts'
  88     prepExp env (Cast e t) = do
  89       e' <- prepExp env e
  90       return $ Cast e' t
  91     prepExp env (Note s e) = do
  92       e' <- prepExp env e
  93       return $ Note s e'
  94     prepExp _ (External s t) = return $ External s t
  95
  96     prepAlt :: (Venv,Tvenv) -> Alt -> PrepM Alt
  97     prepAlt (venv,tvenv) (Acon qdc tbs vbs e) = do
  98       rhs' <- prepExp (foldl' eextend venv vbs,foldl' eextend tvenv tbs) e
  99       return $ Acon qdc tbs vbs rhs'
 100     prepAlt env (Alit l e) = (liftM (Alit l)) (prepExp env e)
 101     prepAlt env (Adefault e) = (liftM Adefault) (prepExp env e)
 102
 103     unwindApp :: (Venv, Tvenv) -> Exp -> [Either Exp Ty] -> PrepM Exp
 104     unwindApp env (App e1 e2) as = unwindApp env e1 (Left e2:as)
 105     unwindApp env (Appt e t) as  = unwindApp env e (Right t:as)
 106     unwindApp env (op@(Dcon qdc)) as = do
 107         e' <- rewindApp env op as
 108         -- possibly dubious to assume no type args
 109         etaExpand [] (drop n atys) e'
 110         where (tbs,atys0,_) = splitTy (qlookup cenv_ eempty qdc)
 111               atys = map (substl (map fst tbs) ts) atys0
 112               ts = [t | Right t <- as]
 113               n = length [e | Left e <- as]
 114     unwindApp env (op@(Var qv)) as | isPrimVar qv = do
 115         e' <- rewindApp env op as
 116         etaExpand [] [] e'
 117     unwindApp env (op@(External _ t)) as = do
 118         e' <- rewindApp env op as
 119         etaExpand [] (drop n atys) e'
 120           where (_,atys,_) = splitTy t
 121                 n = length as -- assumes all args are term args
 122     unwindApp env op as = rewindApp env op as
 123
 124
 125     etaExpand :: [Kind] -> [Ty] -> Exp -> PrepM Exp
 126     etaExpand ks ts e = do
 127          -- what a pain
 128          tyvs <- replicateM (length ks) freshVar
 129          termvs <- replicateM (length ts) freshVar
 130          let tyArgs   = zip tyvs ks
 131          let termArgs = zip termvs ts
 132          return $
 133           foldr (\ (t1,k1) e -> Lam (Tb (t1,k1)) e)
 134            (foldr (\ (v,t) e -> Lam (Vb (v,t)) e)
 135               (foldl' (\ e (v,_) -> App e (Var (unqual v)))
 136                  (foldl' (\ e (tv,_) -> Appt e (Tvar tv))
 137                    e tyArgs)
 138               termArgs) termArgs)
 139            tyArgs
 140
 141     rewindApp :: (Venv, Tvenv) -> Exp -> [Either Exp Ty] -> PrepM Exp
 142     rewindApp _ e [] = return e
 143     rewindApp env@(venv,tvenv) e1 (Left e2:as) | kindOfTy tvenv t `eqKind` Kunlifted && suspends e2 = do
 144         v <- freshVar
 145         let venv' = eextend venv (v,t)
 146         rhs <- rewindApp (venv', tvenv) (App e1 (Var (unqual v))) as
 147         newScrut <- prepExp env e2
 148        -- This is the other place where we call the typechecker.
 149         return $ Case newScrut (v,t) (typeOfExp venv' tvenv rhs) [Adefault rhs]
 150         where t = typeOfExp venv tvenv e2
 151     rewindApp env e1 (Left e2:as) = do
 152       e2' <- prepExp env e2
 153       rewindApp env (App e1 e2') as
 154     rewindApp env e (Right t:as) = rewindApp env (Appt e t) as
 155
 156     typeOfExp :: Venv -> Tvenv -> Exp -> Ty
 157     typeOfExp = checkExpr mn globalEnv tcenv cenv
 158
 159     kindOfTy :: Tvenv -> Ty -> Kind
 160     kindOfTy tvenv = checkType mn globalEnv tcenv tvenv
 161
 162     {- Return false for those expressions for which Interp.suspendExp builds a thunk. -}
 163     suspends (Var _) = False
 164     suspends (Lit _) = False
 165     suspends (Lam (Vb _) _) = False
 166     suspends (Lam _ e) = suspends e
 167     suspends (Appt e _) = suspends e
 168     suspends (Cast e _) = suspends e
 169     suspends (Note _ e) = suspends e
 170     suspends (External _ _) = False
 171     suspends _ = True
 172
 173     mlookup :: (Envs -> Env a b) -> Env a b -> Mname -> Env a b
 174     mlookup _ local_env Nothing = local_env
 175     mlookup selector _  (Just m) =
 176       case elookup globalEnv m of
 177         Just env -> selector env
 178         Nothing -> error ("Prep: undefined module name: " ++ show m)
 179
 180     qlookup ::  (Ord a, Show a) => (Envs -> Env a b) -> Env a b -> (Mname,a) -> b
 181     qlookup selector local_env (m,k) =
 182       case elookup (mlookup selector local_env m) k of
 183         Just v -> v
 184         Nothing -> error ("undefined identifier: " ++ show k)
 185
 186 boundVars :: Exp -> [Id]
 187 boundVars (Lam (Vb (v,_)) e) = [v] `union` boundVars e
 188 boundVars (Lam _ e) = boundVars e
 189 boundVars (Let vds e) = (boundVarsVdefs vds) `union` boundVars e
 190 boundVars (Case scrut (v,_) _ alts) =
 191    [v] `union` (boundVars scrut) `union` boundVarsAlts alts
 192 boundVars (Cast e _) = boundVars e
 193 boundVars (Note _ e) = boundVars e
 194 boundVars (App e1 e2) = boundVars e1 `union` boundVars e2
 195 boundVars (Appt e _) = boundVars e
 196 boundVars _ = []
 197
 198 boundVarsVdefs :: Vdefg -> [Id]
 199 boundVarsVdefs (Rec vds) = nub (concatMap boundVarsVdef vds)
 200 boundVarsVdefs (Nonrec vd) = boundVarsVdef vd
 201
 202 boundVarsVdef :: Vdef -> [Id]
 203 boundVarsVdef (Vdef ((_,v),_,e)) = [v] `union` boundVars e
 204
 205 boundVarsAlts :: [Alt] -> [Var]
 206 boundVarsAlts as = nub (concatMap boundVarsAlt as)
 207
 208 boundVarsAlt :: Alt -> [Var]
 209 boundVarsAlt (Acon _ _ vbs e) = (map fst vbs) `union` (boundVars e)
 210 boundVarsAlt (Alit _ e) = boundVars e
 211 boundVarsAlt (Adefault e) = boundVars e
 212
 213 substNewtys :: NtEnv -> Ty -> Ty
 214 substNewtys ntEnv = everywhere'Except (mkT go)
 215                  where go t | Just ((_,tc),args) <- splitTyConApp_maybe t =
 216                          case M.lookup tc ntEnv of
 217                            Just d -> -- trace ("applying newtype: " ++ show t) $
 218                                        (snd (applyNewtype d args))
 219                            Nothing  -> t
 220                        go t = t
 221
 222 newtypeCoercion_maybe :: NtEnv -> Ty -> Maybe CoercionKind
 223 newtypeCoercion_maybe ntEnv t | Just ((_,tc),_) <- splitTyConApp_maybe t =
 224   M.lookup tc ntEnv
 225 newtypeCoercion_maybe _ _ = Nothing
 226
 227 mkTapp :: Ty -> [Ty] -> Ty
 228 mkTapp = foldl Tapp
 229
 230 initCounter :: Int
 231 initCounter = 0
 232
 233 type PrepM = State Int
 234
 235 freshVar :: PrepM String
 236 freshVar = do
 237   i <- get
 238   put (i+1)
 239   return $ ("zd" ++ show i)