Fix casesimpl and caseremove wrt scrutinee binders.
[matthijs/master-project/cλash.git] / cλash / CLasH / Utils / Core / CoreTools.hs
1 {-# LANGUAGE PatternGuards, TypeSynonymInstances #-}
2 -- | This module provides a number of functions to find out things about Core
3 -- programs. This module does not provide the actual plumbing to work with
4 -- Core and Haskell (it uses HsTools for this), but only the functions that
5 -- know about various libraries and know which functions to call.
6 module CLasH.Utils.Core.CoreTools where
7
8 --Standard modules
9 import qualified Maybe
10 import qualified System.IO.Unsafe
11
12 -- GHC API
13 import qualified GHC
14 import qualified Type
15 import qualified TcType
16 import qualified HsExpr
17 import qualified HsTypes
18 import qualified HscTypes
19 import qualified Name
20 import qualified Id
21 import qualified TyCon
22 import qualified DataCon
23 import qualified TysWiredIn
24 import qualified DynFlags
25 import qualified SrcLoc
26 import qualified CoreSyn
27 import qualified Var
28 import qualified IdInfo
29 import qualified VarSet
30 import qualified CoreUtils
31 import qualified CoreFVs
32 import qualified Literal
33 import qualified MkCore
34 import qualified VarEnv
35
36 -- Local imports
37 import CLasH.Translator.TranslatorTypes
38 import CLasH.Utils.GhcTools
39 import CLasH.Utils.HsTools
40 import CLasH.Utils.Pretty
41 import CLasH.Utils
42 import qualified CLasH.Utils.Core.BinderTools as BinderTools
43
44 -- | A single binding, used as a shortcut to simplify type signatures.
45 type Binding = (CoreSyn.CoreBndr, CoreSyn.CoreExpr)
46
47 -- | Evaluate a core Type representing type level int from the tfp
48 -- library to a real int.
49 eval_tfp_int :: HscTypes.HscEnv -> Type.Type -> Int
50 eval_tfp_int env ty =
51   unsafeRunGhc libdir $ do
52     GHC.setSession env
53     -- Automatically import modules for any fully qualified identifiers
54     setDynFlag DynFlags.Opt_ImplicitImportQualified
55
56     let from_int_t_name = mkRdrName "Types.Data.Num.Ops" "fromIntegerT"
57     let from_int_t = SrcLoc.noLoc $ HsExpr.HsVar from_int_t_name
58     let undef = hsTypedUndef $ coreToHsType ty
59     let app = SrcLoc.noLoc $ HsExpr.HsApp (from_int_t) (undef)
60     let int_ty = SrcLoc.noLoc $ HsTypes.HsTyVar TysWiredIn.intTyCon_RDR
61     let expr = HsExpr.ExprWithTySig app int_ty
62     core <- toCore expr
63     execCore core
64   where
65     libdir = DynFlags.topDir dynflags
66     dynflags = HscTypes.hsc_dflags env
67
68 normalise_tfp_int :: HscTypes.HscEnv -> Type.Type -> Type.Type
69 normalise_tfp_int env ty =
70    System.IO.Unsafe.unsafePerformIO $
71      normaliseType env ty
72
73 -- | Get the width of a SizedWord type
74 -- sized_word_len :: HscTypes.HscEnv -> Type.Type -> Int
75 -- sized_word_len env ty = eval_tfp_int env (sized_word_len_ty ty)
76     
77 sized_word_len_ty :: Type.Type -> Type.Type
78 sized_word_len_ty ty = len
79   where
80     args = case Type.splitTyConApp_maybe ty of
81       Just (tycon, args) -> args
82       Nothing -> error $ "\nCoreTools.sized_word_len_ty: Not a sized word type: " ++ (pprString ty)
83     [len]         = args
84
85 -- | Get the width of a SizedInt type
86 -- sized_int_len :: HscTypes.HscEnv -> Type.Type -> Int
87 -- sized_int_len env ty = eval_tfp_int env (sized_int_len_ty ty)
88
89 sized_int_len_ty :: Type.Type -> Type.Type
90 sized_int_len_ty ty = len
91   where
92     args = case Type.splitTyConApp_maybe ty of
93       Just (tycon, args) -> args
94       Nothing -> error $ "\nCoreTools.sized_int_len_ty: Not a sized int type: " ++ (pprString ty)
95     [len]         = args
96     
97 -- | Get the upperbound of a RangedWord type
98 -- ranged_word_bound :: HscTypes.HscEnv -> Type.Type -> Int
99 -- ranged_word_bound env ty = eval_tfp_int env (ranged_word_bound_ty ty)
100     
101 ranged_word_bound_ty :: Type.Type -> Type.Type
102 ranged_word_bound_ty ty = len
103   where
104     args = case Type.splitTyConApp_maybe ty of
105       Just (tycon, args) -> args
106       Nothing -> error $ "\nCoreTools.ranged_word_bound_ty: Not a sized word type: " ++ (pprString ty)
107     [len]         = args
108
109 -- | Evaluate a core Type representing type level int from the TypeLevel
110 -- library to a real int.
111 -- eval_type_level_int :: Type.Type -> Int
112 -- eval_type_level_int ty =
113 --   unsafeRunGhc $ do
114 --     -- Automatically import modules for any fully qualified identifiers
115 --     setDynFlag DynFlags.Opt_ImplicitImportQualified
116 -- 
117 --     let to_int_name = mkRdrName "Data.TypeLevel.Num.Sets" "toInt"
118 --     let to_int = SrcLoc.noLoc $ HsExpr.HsVar to_int_name
119 --     let undef = hsTypedUndef $ coreToHsType ty
120 --     let app = HsExpr.HsApp (to_int) (undef)
121 -- 
122 --     core <- toCore [] app
123 --     execCore core 
124
125 -- | Get the length of a FSVec type
126 -- tfvec_len :: HscTypes.HscEnv -> Type.Type -> Int
127 -- tfvec_len env ty = eval_tfp_int env (tfvec_len_ty ty)
128
129 tfvec_len_ty :: Type.Type -> Type.Type
130 tfvec_len_ty ty = len
131   where  
132     args = case Type.splitTyConApp_maybe ty of
133       Just (tycon, args) -> args
134       Nothing -> error $ "\nCoreTools.tfvec_len_ty: Not a vector type: " ++ (pprString ty)
135     [len, el_ty] = args
136     
137 -- | Get the element type of a TFVec type
138 tfvec_elem :: Type.Type -> Type.Type
139 tfvec_elem ty = el_ty
140   where
141     args = case Type.splitTyConApp_maybe ty of
142       Just (tycon, args) -> args
143       Nothing -> error $ "\nCoreTools.tfvec_len: Not a vector type: " ++ (pprString ty)
144     [len, el_ty] = args
145
146 -- Is the given core expression a lambda abstraction?
147 is_lam :: CoreSyn.CoreExpr -> Bool
148 is_lam (CoreSyn.Lam _ _) = True
149 is_lam _ = False
150
151 -- Is the given core expression of a function type?
152 is_fun :: CoreSyn.CoreExpr -> Bool
153 -- Treat Type arguments differently, because exprType is not defined for them.
154 is_fun (CoreSyn.Type _) = False
155 is_fun expr = (Type.isFunTy . CoreUtils.exprType) expr
156
157 -- Is the given core expression polymorphic (i.e., does it accept type
158 -- arguments?).
159 is_poly :: CoreSyn.CoreExpr -> Bool
160 -- Treat Type arguments differently, because exprType is not defined for them.
161 is_poly (CoreSyn.Type _) = False
162 is_poly expr = (Maybe.isJust . Type.splitForAllTy_maybe . CoreUtils.exprType) expr
163
164 -- Is the given core expression a variable reference?
165 is_var :: CoreSyn.CoreExpr -> Bool
166 is_var (CoreSyn.Var _) = True
167 is_var _ = False
168
169 is_lit :: CoreSyn.CoreExpr -> Bool
170 is_lit (CoreSyn.Lit _) = True
171 is_lit _ = False
172
173 -- Can the given core expression be applied to something? This is true for
174 -- applying to a value as well as a type.
175 is_applicable :: CoreSyn.CoreExpr -> Bool
176 is_applicable expr = is_fun expr || is_poly expr
177
178 -- Is the given core expression a variable or an application?
179 is_simple :: CoreSyn.CoreExpr -> Bool
180 is_simple (CoreSyn.App _ _) = True
181 is_simple (CoreSyn.Var _) = True
182 is_simple (CoreSyn.Cast expr _) = is_simple expr
183 is_simple _ = False
184
185 -- Does the given CoreExpr have any free type vars?
186 has_free_tyvars :: CoreSyn.CoreExpr -> Bool
187 has_free_tyvars = not . VarSet.isEmptyVarSet . (CoreFVs.exprSomeFreeVars Var.isTyVar)
188
189 -- Does the given type have any free type vars?
190 ty_has_free_tyvars :: Type.Type -> Bool
191 ty_has_free_tyvars = not . VarSet.isEmptyVarSet . Type.tyVarsOfType
192
193 -- Does the given CoreExpr have any free local vars?
194 has_free_vars :: CoreSyn.CoreExpr -> Bool
195 has_free_vars = not . VarSet.isEmptyVarSet . CoreFVs.exprFreeVars
196
197 -- Does the given expression use any of the given binders?
198 expr_uses_binders :: [CoreSyn.CoreBndr] -> CoreSyn.CoreExpr -> Bool
199 expr_uses_binders bndrs = not . VarSet.isEmptyVarSet . (CoreFVs.exprSomeFreeVars (`elem` bndrs))
200
201 -- Turns a Var CoreExpr into the Id inside it. Will of course only work for
202 -- simple Var CoreExprs, not complexer ones.
203 exprToVar :: CoreSyn.CoreExpr -> Var.Id
204 exprToVar (CoreSyn.Var id) = id
205 exprToVar expr = error $ "\nCoreTools.exprToVar: Not a var: " ++ show expr
206
207 -- Turns a Lit CoreExpr into the Literal inside it.
208 exprToLit :: CoreSyn.CoreExpr -> Literal.Literal
209 exprToLit (CoreSyn.Lit lit) = lit
210 exprToLit expr = error $ "\nCoreTools.exprToLit: Not a lit: " ++ show expr
211
212 -- Removes all the type and dictionary arguments from the given argument list,
213 -- leaving only the normal value arguments. The type given is the type of the
214 -- expression applied to this argument list.
215 get_val_args :: Type.Type -> [CoreSyn.CoreExpr] -> [CoreSyn.CoreExpr]
216 get_val_args ty args = drop n args
217   where
218     (tyvars, predtypes, _) = TcType.tcSplitSigmaTy ty
219     -- The first (length tyvars) arguments should be types, the next 
220     -- (length predtypes) arguments should be dictionaries. We drop this many
221     -- arguments, to get at the value arguments.
222     n = length tyvars + length predtypes
223
224 getLiterals :: HscTypes.HscEnv -> CoreSyn.CoreExpr -> [CoreSyn.CoreExpr]
225 getLiterals _ app@(CoreSyn.App _ _) = literals
226   where
227     (CoreSyn.Var f, args) = CoreSyn.collectArgs app
228     literals = filter (is_lit) args
229
230 getLiterals _ lit@(CoreSyn.Lit _) = [lit]
231
232 getLiterals hscenv letrec@(CoreSyn.Let (CoreSyn.NonRec letBind (letExpr)) letRes) = [lit]
233   where
234     ty     = Var.varType letBind
235     litInt = eval_tfp_int hscenv ty
236     lit    = CoreSyn.Lit (Literal.mkMachInt (toInteger litInt))
237
238 getLiterals _ expr = error $ "\nCoreTools.getLiterals: Not a known Lit: " ++ pprString expr
239
240 reduceCoreListToHsList :: 
241   [HscTypes.CoreModule] -- ^ The modules where parts of the list are hidden
242   -> CoreSyn.CoreExpr   -- ^ The refence to atleast one of the nodes
243   -> TranslatorSession [CoreSyn.CoreExpr]
244 reduceCoreListToHsList cores app@(CoreSyn.App _ _) = do {
245   ; let { (fun, args) = CoreSyn.collectArgs app
246         ; len         = length args 
247         } ;
248   ; case len of
249       3 -> do {
250         ; let topelem = args!!1
251         ; case (args!!2) of
252             (varz@(CoreSyn.Var id)) -> do {
253               ; binds <- mapM (findExpr (isVarName id)) cores
254               ; otherelems <- reduceCoreListToHsList cores (head (Maybe.catMaybes binds))
255               ; return (topelem:otherelems)
256               }
257             (appz@(CoreSyn.App _ _)) -> do {
258               ; otherelems <- reduceCoreListToHsList cores appz
259               ; return (topelem:otherelems)
260               }
261             otherwise -> return [topelem]
262         }
263       otherwise -> return []
264   }
265   where
266     isVarName :: Monad m => Var.Var -> Var.Var -> m Bool
267     isVarName lookfor bind = return $ (Var.varName lookfor) == (Var.varName bind)
268
269 reduceCoreListToHsList _ _ = return []
270
271 -- Is the given var the State data constructor?
272 isStateCon :: Var.Var -> Bool
273 isStateCon var =
274   -- See if it is a DataConWrapId (not DataConWorkId, since State is a
275   -- newtype).
276   case Id.idDetails var of
277     IdInfo.DataConWrapId dc -> 
278       -- See if the datacon is the State datacon from the State type.
279       let tycon = DataCon.dataConTyCon dc
280           tyname = Name.getOccString tycon
281           dcname = Name.getOccString dc
282       in case (tyname, dcname) of
283         ("State", "State") -> True
284         _ -> False
285     _ -> False
286
287 -- | Is the given type a State type?
288 isStateType :: Type.Type -> Bool
289 -- Resolve any type synonyms remaining
290 isStateType ty | Just ty' <- Type.tcView ty = isStateType ty'
291 isStateType ty  = Maybe.isJust $ do
292   -- Split the type. Don't use normal splitAppTy, since that looks through
293   -- newtypes, and we want to see the State newtype.
294   (typef, _) <- Type.repSplitAppTy_maybe ty
295   -- See if the applied type is a type constructor
296   (tycon, _) <- Type.splitTyConApp_maybe typef
297   if TyCon.isNewTyCon tycon && Name.getOccString tycon == "State"
298     then
299       Just ()
300     else
301       Nothing
302
303 -- | Does the given TypedThing have a State type?
304 hasStateType :: (TypedThing t) => t -> Bool
305 hasStateType expr = case getType expr of
306   Nothing -> False
307   Just ty -> isStateType ty
308
309
310 -- | Flattens nested lets into a single list of bindings. The expression
311 --   passed does not have to be a let expression, if it isn't an empty list of
312 --   bindings is returned.
313 flattenLets ::
314   CoreSyn.CoreExpr -- ^ The expression to flatten.
315   -> ([Binding], CoreSyn.CoreExpr) -- ^ The bindings and resulting expression.
316 flattenLets (CoreSyn.Let binds expr) = 
317   (bindings ++ bindings', expr')
318   where
319     -- Recursively flatten the contained expression
320     (bindings', expr') =flattenLets expr
321     -- Flatten our own bindings to remove the Rec / NonRec constructors
322     bindings = CoreSyn.flattenBinds [binds]
323 flattenLets expr = ([], expr)
324
325 -- | Create bunch of nested non-recursive let expressions from the given
326 -- bindings. The first binding is bound at the highest level (and thus
327 -- available in all other bindings).
328 mkNonRecLets :: [Binding] -> CoreSyn.CoreExpr -> CoreSyn.CoreExpr
329 mkNonRecLets bindings expr = MkCore.mkCoreLets binds expr
330   where
331     binds = map (uncurry CoreSyn.NonRec) bindings
332
333 -- | A class of things that (optionally) have a core Type. The type is
334 -- optional, since Type expressions don't have a type themselves.
335 class TypedThing t where
336   getType :: t -> Maybe Type.Type
337
338 instance TypedThing CoreSyn.CoreExpr where
339   getType (CoreSyn.Type _) = Nothing
340   getType expr = Just $ CoreUtils.exprType expr
341
342 instance TypedThing CoreSyn.CoreBndr where
343   getType = return . Id.idType
344
345 instance TypedThing Type.Type where
346   getType = return . id
347
348 -- | Generate new uniques for all binders in the given expression.
349 -- Does not support making type variables unique, though this could be
350 -- supported if required (by passing a CoreSubst.Subst instead of VarEnv to
351 -- genUniques' below).
352 genUniques :: CoreSyn.CoreExpr -> TranslatorSession CoreSyn.CoreExpr
353 genUniques = genUniques' VarEnv.emptyVarEnv
354
355 -- | A helper function to generate uniques, that takes a VarEnv containing the
356 --   substitutions already performed.
357 genUniques' :: VarEnv.VarEnv CoreSyn.CoreBndr -> CoreSyn.CoreExpr -> TranslatorSession CoreSyn.CoreExpr
358 genUniques' subst (CoreSyn.Var f) = do
359   -- Replace the binder with its new value, if applicable.
360   let f' = VarEnv.lookupWithDefaultVarEnv subst f f
361   return (CoreSyn.Var f')
362 -- Leave literals untouched
363 genUniques' subst (CoreSyn.Lit l) = return $ CoreSyn.Lit l
364 genUniques' subst (CoreSyn.App f arg) = do
365   -- Only work on subexpressions
366   f' <- genUniques' subst f
367   arg' <- genUniques' subst arg
368   return (CoreSyn.App f' arg')
369 -- Don't change type abstractions
370 genUniques' subst expr@(CoreSyn.Lam bndr res) | CoreSyn.isTyVar bndr = return expr
371 genUniques' subst (CoreSyn.Lam bndr res) = do
372   -- Generate a new unique for the bound variable
373   (subst', bndr') <- genUnique subst bndr
374   res' <- genUniques' subst' res
375   return (CoreSyn.Lam bndr' res')
376 genUniques' subst (CoreSyn.Let (CoreSyn.NonRec bndr bound) res) = do
377   -- Make the binders unique
378   (subst', bndr') <- genUnique subst bndr
379   bound' <- genUniques' subst' bound
380   res' <- genUniques' subst' res
381   return $ CoreSyn.Let (CoreSyn.NonRec bndr' bound') res'
382 genUniques' subst (CoreSyn.Let (CoreSyn.Rec binds) res) = do
383   -- Make each of the binders unique
384   (subst', bndrs') <- mapAccumLM genUnique subst (map fst binds)
385   bounds' <- mapM (genUniques' subst' . snd) binds
386   res' <- genUniques' subst' res
387   let binds' = zip bndrs' bounds'
388   return $ CoreSyn.Let (CoreSyn.Rec binds') res'
389 genUniques' subst (CoreSyn.Case scrut bndr ty alts) = do
390   -- Process the scrutinee with the original substitution, since non of the
391   -- binders bound in the Case statement is in scope in the scrutinee.
392   scrut' <- genUniques' subst scrut
393   -- Generate a new binder for the scrutinee
394   (subst', bndr') <- genUnique subst bndr
395   -- Process each of the alts
396   alts' <- mapM (doalt subst') alts
397   return $ CoreSyn.Case scrut' bndr' ty alts'
398   where
399     doalt subst (con, bndrs, expr) = do
400       (subst', bndrs') <- mapAccumLM genUnique subst bndrs
401       expr' <- genUniques' subst' expr
402       -- Note that we don't return subst', since bndrs are only in scope in
403       -- expr.
404       return (con, bndrs', expr')
405 genUniques' subst (CoreSyn.Cast expr coercion) = do
406   expr' <- genUniques' subst expr
407   -- Just process the casted expression
408   return $ CoreSyn.Cast expr' coercion
409 genUniques' subst (CoreSyn.Note note expr) = do
410   expr' <- genUniques' subst expr
411   -- Just process the annotated expression
412   return $ CoreSyn.Note note expr'
413 -- Leave types untouched
414 genUniques' subst expr@(CoreSyn.Type _) = return expr
415
416 -- Generate a new unique for the given binder, and extend the given
417 -- substitution to reflect this.
418 genUnique :: VarEnv.VarEnv CoreSyn.CoreBndr -> CoreSyn.CoreBndr -> TranslatorSession (VarEnv.VarEnv CoreSyn.CoreBndr, CoreSyn.CoreBndr)
419 genUnique subst bndr = do
420   bndr' <- BinderTools.cloneVar bndr
421   -- Replace all occurences of the old binder with a reference to the new
422   -- binder.
423   let subst' = VarEnv.extendVarEnv subst bndr bndr'
424   return (subst', bndr')