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