External Core lib: lots of cleanup
[ghc.git] / 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 type NtEnv = M.Map Tcon CoercionKind
228
229 mkTapp :: Ty -> [Ty] -> Ty
230 mkTapp = foldl Tapp
231
232 initCounter :: Int
233 initCounter = 0
234
235 type PrepM = State Int
236
237 freshVar :: PrepM String
238 freshVar = do
239 i <- get
240 put (i+1)
241 return $ ("zd" ++ show i)