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