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