df664a5fb8682e4d4dcde8fb967d4280448e32a8
[ghc-hetmet.git] / utils / ext-core / Prep.hs
1 {-# OPTIONS -Wall -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         (2) Arrange that any non-trivial expression of unlifted kind ('#')
6              is turned into the scrutinee of a Case.
7 After these preprocessing steps, Core can be interpreted (or given an operational semantics)
8       ignoring type information almost completely.
9 -}
10
11
12 module Prep where
13
14 import Data.Either
15
16 import Prims
17 import Core
18 import Env
19 import Check
20
21 import Data.List
22
23 prepModule :: Menv -> Module -> Module
24 prepModule globalEnv (Module mn tdefs vdefgs) = 
25     Module mn tdefs vdefgs' 
26   where
27
28     (_,vdefgs') = foldl prepTopVdefg (eempty,[]) vdefgs
29
30     prepTopVdefg (venv,vdefgs) vdefg = (venv',vdefgs ++ [vdefg'])
31        where (venv',vdefg') = prepVdefg (venv,eempty) vdefg
32  
33     prepVdefg (env@(venv,_)) (Nonrec(Vdef((Nothing,x),t,e))) = 
34         (eextend venv (x,t), Nonrec(Vdef((Nothing,x),t,prepExp env e)))
35     prepVdefg (env@(venv,_))  (Nonrec(Vdef(qx,t,e))) =
36         (venv, Nonrec(Vdef(qx,t,prepExp env e)))
37     prepVdefg (venv,tvenv) (Rec vdefs) = 
38         (venv',Rec [ Vdef(qx,t,prepExp (venv',tvenv) e) | Vdef(qx,t,e) <- vdefs])
39         where venv' = foldl eextend venv [(x,t) | Vdef((Nothing,x),t,_) <- vdefs]
40
41     prepExp _ (Var qv) = Var qv
42     prepExp _ (Dcon qdc) = Dcon qdc
43     prepExp _ (Lit l) = Lit l
44     prepExp env e@(App _ _) = unwindApp env e []
45     prepExp env e@(Appt _ _) = unwindApp env e []
46     prepExp (venv,tvenv) (Lam (Vb vb) e) = Lam (Vb vb) (prepExp (eextend venv vb,tvenv) e)
47     prepExp (venv,tvenv) (Lam (Tb tb) e) = Lam (Tb tb) (prepExp (venv,eextend tvenv tb) e)
48     prepExp env@(venv,tvenv) (Let (Nonrec(Vdef((Nothing,x),t,b))) e) 
49         | (kindOfTy tvenv t `eqKind` Kunlifted && suspends b) = 
50             -- There are two places where we call the typechecker, one of them
51             -- here.
52             -- We need to know the type of the let body in order to construct
53             -- a case expression. 
54                                 -- need to extend the env with the let-bound var too!
55             let eTy = typeOfExp (eextend venv (x, t), tvenv) e in
56                Case (prepExp env b) (x,t) 
57                   eTy
58                   [Adefault (prepExp (eextend venv (x,t),tvenv) e)] 
59     prepExp (venv,tvenv) (Let vdefg e) =  Let vdefg' (prepExp (venv',tvenv) e)
60                 where (venv',vdefg') = prepVdefg (venv,tvenv) vdefg
61     prepExp env@(venv,tvenv) (Case e vb t alts) = Case (prepExp env e) vb t (map (prepAlt (eextend venv vb,tvenv)) alts)
62     prepExp env (Cast e t) = Cast (prepExp env e) t
63     prepExp env (Note s e) = Note s (prepExp env e)
64     prepExp _ (External s t) = External s t
65
66     prepAlt (venv,tvenv) (Acon qdc tbs vbs e) = Acon qdc tbs vbs (prepExp (foldl eextend venv vbs,foldl eextend tvenv tbs) e)
67     prepAlt env (Alit l e) = Alit l (prepExp env e)
68     prepAlt env (Adefault e) = Adefault (prepExp env e)
69
70
71     unwindApp env (App e1 e2) as = unwindApp env e1 (Left e2:as)
72     unwindApp env (Appt e t) as  = unwindApp env e (Right t:as)
73     unwindApp env (op@(Dcon qdc)) as = 
74         -- possibly dubious to assume no type args
75         etaExpand [] (drop n atys) (rewindApp env op as)
76         where (tbs,atys0,_) = splitTy (qlookup cenv_ eempty qdc)
77               atys = map (substl (map fst tbs) ts) atys0
78               ts = [t | Right t <- as]
79               n = length [e | Left e <- as]
80     unwindApp env (op@(Var(qv@(_,p)))) as | isPrimVar qv =
81         etaExpand (snd (unzip extraTbs)) (drop n atys) (rewindApp env op as)
82         where -- TODO: avoid copying code. these two cases are the same
83
84               -- etaExpand needs to add the type arguments too! Bah!
85               (tbs, atys0, _) = (maybe (error "unwindApp") splitTy (elookup (venv_ primEnv) p))
86               n_args = length ts
87               (appliedTbs, extraTbs) = (take n_args tbs, drop n_args tbs)
88               atys = map (substl (map fst appliedTbs) ts) atys0
89               ts = [t | Right t <- as]
90               n = length [e | Left e <- as]
91     unwindApp env op as = rewindApp env op as
92
93
94     etaExpand :: [Kind] -> [Ty] -> Exp -> Exp
95     etaExpand ks ts e = 
96          -- what a pain
97          let tyArgs = [("$t_"++(show i),k) | (i, k) <- zip [(1::Integer)..] ks]   
98              termArgs = [ ('$':(show i),t) | (i,t) <- zip [(1::Integer)..] ts] in
99           foldr (\ (t1,k1) e -> Lam (Tb (t1,k1)) e)
100            (foldr (\ (v,t) e -> Lam (Vb (v,t)) e)
101               (foldl (\ e (v,_) -> App e (Var (unqual v)))
102                  (foldl (\ e (tv,_) -> Appt e (Tvar tv))
103                    e tyArgs)
104               termArgs) termArgs)
105            tyArgs
106
107     rewindApp _ e [] = e
108     rewindApp env@(venv,tvenv) e1 (Left e2:as) | kindOfTy tvenv t `eqKind` Kunlifted && suspends e2 =
109        -- This is the other place where we call the typechecker.
110         Case newScrut (v,t) (typeOfExp env' rhs) [Adefault rhs]
111         where newScrut = prepExp env e2
112               rhs = (rewindApp env' (App e1 (Var (unqual v))) as)
113                  -- note:
114                  -- e1 gets moved inside rhs. so if we pick a case
115                  -- var name (outside e1) equal to a name bound *inside*
116                  -- e1, the binding *inside* e1 will shadow "v"
117                  -- Which would be name capture!
118                  -- So, we pass the bound vars of e1 to freshVar along with
119                  -- the domain of the current env.
120               v = freshVar (edomain venv `union` (boundVars e1))
121               t = typeOfExp env e2
122               env' = (eextend venv (v,t),tvenv)
123     rewindApp env e1 (Left e2:as) = rewindApp env (App e1 (prepExp env e2)) as
124     rewindApp env e (Right t:as) = rewindApp env (Appt e t) as
125
126     freshVar vs = maximum ("":vs) ++ "x" -- one simple way!
127
128     typeOfExp :: (Venv, Tvenv) -> Exp -> Ty
129     typeOfExp = uncurry (checkExpr mn globalEnv tdefs)
130
131     kindOfTy :: Tvenv -> Ty -> Kind
132     kindOfTy tvenv = checkType mn globalEnv tdefs tvenv
133
134     {- Return false for those expressions for which Interp.suspendExp builds a thunk. -}
135     suspends (Var _) = False
136     suspends (Lit _) = False
137     suspends (Lam (Vb _) _) = False
138     suspends (Lam _ e) = suspends e
139     suspends (Appt e _) = suspends e
140     suspends (Cast e _) = suspends e
141     suspends (Note _ e) = suspends e
142     suspends (External _ _) = False
143     suspends _ = True
144
145     mlookup :: (Envs -> Env a b) -> Env a b -> Mname -> Env a b
146     mlookup _ local_env Nothing = local_env
147     mlookup selector _  (Just m) =   
148       case elookup globalEnv m of
149         Just env -> selector env
150         Nothing -> error ("Prep: undefined module name: " ++ show m)
151
152     qlookup ::  (Ord a, Show a) => (Envs -> Env a b) -> Env a b -> (Mname,a) -> b
153     qlookup selector local_env (m,k) =   
154       case elookup (mlookup selector local_env m) k of
155         Just v -> v
156         Nothing -> error ("undefined identifier: " ++ show k)
157
158 boundVars :: Exp -> [Id]
159 boundVars (Lam (Vb (v,_)) e) = [v] `union` boundVars e
160 boundVars (Lam _ e) = boundVars e
161 boundVars (Let vds e) = (boundVarsVdefs vds) `union` boundVars e
162 boundVars (Case scrut (v,_) _ alts) = 
163    [v] `union` (boundVars scrut) `union` boundVarsAlts alts
164 boundVars (Cast e _) = boundVars e
165 boundVars (Note _ e) = boundVars e
166 boundVars (App e1 e2) = boundVars e1 `union` boundVars e2
167 boundVars (Appt e _) = boundVars e
168 boundVars _ = []
169
170 boundVarsVdefs :: Vdefg -> [Id]
171 boundVarsVdefs (Rec vds) = nub (concatMap boundVarsVdef vds)
172 boundVarsVdefs (Nonrec vd) = boundVarsVdef vd
173
174 boundVarsVdef :: Vdef -> [Id]
175 boundVarsVdef (Vdef ((_,v),_,e)) = [v] `union` boundVars e
176
177 boundVarsAlts :: [Alt] -> [Var]
178 boundVarsAlts as = nub (concatMap boundVarsAlt as)
179
180 boundVarsAlt :: Alt -> [Var]
181 boundVarsAlt (Acon _ _ vbs e) = (map fst vbs) `union` (boundVars e)
182 boundVarsAlt (Alit _ e) = boundVars e
183 boundVarsAlt (Adefault e) = boundVars e