7d5423ccc16b39ccf2c96e8cd2ee8e4fc0eab242
[matthijs/master-project/cλash.git] / 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 = case List.elemIndex dc dcs of
156     Nothing -> error $ "Datacon " ++ pprString dc ++ " does not occur in typed thing: " ++ pprString tt
157     Just i -> i
158   where
159     dcs = datacons_for tt
160
161 -- | Gets all datacons for the given typed thing. Errors out if the
162 -- typed thing is not ADT typed.
163 datacons_for :: TypedThing t => t -> [DataCon.DataCon]
164 datacons_for tt =
165   case getType tt of
166     Nothing -> error $ "Getting datacon index of untyped thing? " ++ pprString tt
167     Just ty -> case Type.splitTyConApp_maybe ty of
168       Nothing -> error $ "Trying to find datacon in a type without a tycon?" ++ pprString ty
169       Just (tycon, _) -> case TyCon.tyConDataCons_maybe tycon of
170         Nothing -> error $ "Trying to find datacon in a type without datacons?" ++ pprString ty
171         Just dcs -> dcs
172
173 -- Is the given core expression a lambda abstraction?
174 is_lam :: CoreSyn.CoreExpr -> Bool
175 is_lam (CoreSyn.Lam _ _) = True
176 is_lam _ = False
177
178 -- Is the given core expression a let expression?
179 is_let :: CoreSyn.CoreExpr -> Bool
180 is_let (CoreSyn.Let _ _) = True
181 is_let _ = False
182
183 -- Is the given core expression of a function type?
184 is_fun :: CoreSyn.CoreExpr -> Bool
185 -- Treat Type arguments differently, because exprType is not defined for them.
186 is_fun (CoreSyn.Type _) = False
187 is_fun expr = (Type.isFunTy . CoreUtils.exprType) expr
188
189 -- Is the given core expression polymorphic (i.e., does it accept type
190 -- arguments?).
191 is_poly :: CoreSyn.CoreExpr -> Bool
192 -- Treat Type arguments differently, because exprType is not defined for them.
193 is_poly (CoreSyn.Type _) = False
194 is_poly expr = (Maybe.isJust . Type.splitForAllTy_maybe . CoreUtils.exprType) expr
195
196 -- Is the given core expression a variable reference?
197 is_var :: CoreSyn.CoreExpr -> Bool
198 is_var (CoreSyn.Var _) = True
199 is_var _ = False
200
201 is_lit :: CoreSyn.CoreExpr -> Bool
202 is_lit (CoreSyn.Lit _) = True
203 is_lit _ = False
204
205 -- Can the given core expression be applied to something? This is true for
206 -- applying to a value as well as a type.
207 is_applicable :: CoreSyn.CoreExpr -> Bool
208 is_applicable expr = is_fun expr || is_poly expr
209
210 -- Is the given core expression a variable or an application?
211 is_simple :: CoreSyn.CoreExpr -> Bool
212 is_simple (CoreSyn.App _ _) = True
213 is_simple (CoreSyn.Var _) = True
214 is_simple (CoreSyn.Cast expr _) = is_simple expr
215 is_simple _ = False
216
217 -- Does the given CoreExpr have any free type vars?
218 has_free_tyvars :: CoreSyn.CoreExpr -> Bool
219 has_free_tyvars = not . VarSet.isEmptyVarSet . (CoreFVs.exprSomeFreeVars Var.isTyVar)
220
221 -- Does the given type have any free type vars?
222 ty_has_free_tyvars :: Type.Type -> Bool
223 ty_has_free_tyvars = not . VarSet.isEmptyVarSet . Type.tyVarsOfType
224
225 -- Does the given CoreExpr have any free local vars?
226 has_free_vars :: CoreSyn.CoreExpr -> Bool
227 has_free_vars = not . VarSet.isEmptyVarSet . CoreFVs.exprFreeVars
228
229 -- Does the given expression use any of the given binders?
230 expr_uses_binders :: [CoreSyn.CoreBndr] -> CoreSyn.CoreExpr -> Bool
231 expr_uses_binders bndrs = not . VarSet.isEmptyVarSet . (CoreFVs.exprSomeFreeVars (`elem` bndrs))
232
233 -- Turns a Var CoreExpr into the Id inside it. Will of course only work for
234 -- simple Var CoreExprs, not complexer ones.
235 exprToVar :: CoreSyn.CoreExpr -> Var.Id
236 exprToVar (CoreSyn.Var id) = id
237 exprToVar expr = error $ "\nCoreTools.exprToVar: Not a var: " ++ show expr
238
239 -- Turns a Lit CoreExpr into the Literal inside it.
240 exprToLit :: CoreSyn.CoreExpr -> Literal.Literal
241 exprToLit (CoreSyn.Lit lit) = lit
242 exprToLit expr = error $ "\nCoreTools.exprToLit: Not a lit: " ++ show expr
243
244 -- Removes all the type and dictionary arguments from the given argument list,
245 -- leaving only the normal value arguments. The type given is the type of the
246 -- expression applied to this argument list.
247 get_val_args :: Type.Type -> [CoreSyn.CoreExpr] -> [CoreSyn.CoreExpr]
248 get_val_args ty args = drop n args
249   where
250     (tyvars, predtypes, _) = TcType.tcSplitSigmaTy ty
251     -- The first (length tyvars) arguments should be types, the next 
252     -- (length predtypes) arguments should be dictionaries. We drop this many
253     -- arguments, to get at the value arguments.
254     n = length tyvars + length predtypes
255
256 -- Finds out what literal Integer this expression represents.
257 getIntegerLiteral :: CoreSyn.CoreExpr -> TranslatorSession Integer
258 getIntegerLiteral expr =
259   case CoreSyn.collectArgs expr of
260     (CoreSyn.Var f, [CoreSyn.Lit (Literal.MachInt integer)]) 
261       | getFullString f == "GHC.Integer.smallInteger" -> return integer
262     (CoreSyn.Var f, [CoreSyn.Lit (Literal.MachInt64 integer)]) 
263       | getFullString f == "GHC.Integer.int64ToInteger" -> return integer
264     (CoreSyn.Var f, [CoreSyn.Lit (Literal.MachWord integer)]) 
265       | getFullString f == "GHC.Integer.wordToInteger" -> return integer
266     (CoreSyn.Var f, [CoreSyn.Lit (Literal.MachWord64 integer)]) 
267       | getFullString f == "GHC.Integer.word64ToInteger" -> return integer
268     -- fromIntegerT returns the integer corresponding to the type of its
269     -- (third) argument. Since it is polymorphic, the type of that
270     -- argument is passed as the first argument, so we can just use that
271     -- one.
272     (CoreSyn.Var f, [CoreSyn.Type dec_ty, dec_dict, CoreSyn.Type num_ty, num_dict, arg]) 
273       | getFullString f == "Types.Data.Num.Ops.fromIntegerT" -> do
274           int <- MonadState.lift tsType $ tfp_to_int dec_ty
275           return $ toInteger int
276     _ -> error $ "CoreTools.getIntegerLiteral: Unsupported Integer literal: " ++ pprString expr
277
278 reduceCoreListToHsList :: 
279   [HscTypes.CoreModule] -- ^ The modules where parts of the list are hidden
280   -> CoreSyn.CoreExpr   -- ^ The refence to atleast one of the nodes
281   -> TranslatorSession [CoreSyn.CoreExpr]
282 reduceCoreListToHsList cores app@(CoreSyn.App _ _) = do {
283   ; let { (fun, args) = CoreSyn.collectArgs app
284         ; len         = length args 
285         } ;
286   ; case len of
287       3 -> do {
288         ; let topelem = args!!1
289         ; case (args!!2) of
290             (varz@(CoreSyn.Var id)) -> do {
291               ; binds <- mapM (findExpr (isVarName id)) cores
292               ; otherelems <- reduceCoreListToHsList cores (head (Maybe.catMaybes binds))
293               ; return (topelem:otherelems)
294               }
295             (appz@(CoreSyn.App _ _)) -> do {
296               ; otherelems <- reduceCoreListToHsList cores appz
297               ; return (topelem:otherelems)
298               }
299             otherwise -> return [topelem]
300         }
301       otherwise -> return []
302   }
303   where
304     isVarName :: Monad m => Var.Var -> Var.Var -> m Bool
305     isVarName lookfor bind = return $ (Var.varName lookfor) == (Var.varName bind)
306
307 reduceCoreListToHsList _ _ = return []
308
309 -- Is the given var the State data constructor?
310 isStateCon :: Var.Var -> Bool
311 isStateCon var =
312   -- See if it is a DataConWrapId (not DataConWorkId, since State is a
313   -- newtype).
314   case Id.idDetails var of
315     IdInfo.DataConWrapId dc -> 
316       -- See if the datacon is the State datacon from the State type.
317       let tycon = DataCon.dataConTyCon dc
318           tyname = Name.getOccString tycon
319           dcname = Name.getOccString dc
320       in case (tyname, dcname) of
321         ("State", "State") -> True
322         _ -> False
323     _ -> False
324
325 -- | Is the given type a State type?
326 isStateType :: Type.Type -> Bool
327 -- Resolve any type synonyms remaining
328 isStateType ty | Just ty' <- Type.tcView ty = isStateType ty'
329 isStateType ty  = Maybe.isJust $ do
330   -- Split the type. Don't use normal splitAppTy, since that looks through
331   -- newtypes, and we want to see the State newtype.
332   (typef, _) <- Type.repSplitAppTy_maybe ty
333   -- See if the applied type is a type constructor
334   (tycon, _) <- Type.splitTyConApp_maybe typef
335   if TyCon.isNewTyCon tycon && Name.getOccString tycon == "State"
336     then
337       Just ()
338     else
339       Nothing
340
341 -- | Does the given TypedThing have a State type?
342 hasStateType :: (TypedThing t) => t -> Bool
343 hasStateType expr = case getType expr of
344   Nothing -> False
345   Just ty -> isStateType ty
346
347
348 -- | Flattens nested lets into a single list of bindings. The expression
349 --   passed does not have to be a let expression, if it isn't an empty list of
350 --   bindings is returned.
351 flattenLets ::
352   CoreSyn.CoreExpr -- ^ The expression to flatten.
353   -> ([Binding], CoreSyn.CoreExpr) -- ^ The bindings and resulting expression.
354 flattenLets (CoreSyn.Let binds expr) = 
355   (bindings ++ bindings', expr')
356   where
357     -- Recursively flatten the contained expression
358     (bindings', expr') =flattenLets expr
359     -- Flatten our own bindings to remove the Rec / NonRec constructors
360     bindings = CoreSyn.flattenBinds [binds]
361 flattenLets expr = ([], expr)
362
363 -- | Create bunch of nested non-recursive let expressions from the given
364 -- bindings. The first binding is bound at the highest level (and thus
365 -- available in all other bindings).
366 mkNonRecLets :: [Binding] -> CoreSyn.CoreExpr -> CoreSyn.CoreExpr
367 mkNonRecLets bindings expr = MkCore.mkCoreLets binds expr
368   where
369     binds = map (uncurry CoreSyn.NonRec) bindings
370
371 -- | A class of things that (optionally) have a core Type. The type is
372 -- optional, since Type expressions don't have a type themselves.
373 class Outputable.Outputable t => TypedThing t where
374   getType :: t -> Maybe Type.Type
375
376 instance TypedThing CoreSyn.CoreExpr where
377   getType (CoreSyn.Type _) = Nothing
378   getType expr = Just $ CoreUtils.exprType expr
379
380 instance TypedThing CoreSyn.CoreBndr where
381   getType = return . Id.idType
382
383 instance TypedThing Type.Type where
384   getType = return . id
385
386 -- | Generate new uniques for all binders in the given expression.
387 -- Does not support making type variables unique, though this could be
388 -- supported if required (by passing a CoreSubst.Subst instead of VarEnv to
389 -- genUniques' below).
390 genUniques :: CoreSyn.CoreExpr -> TranslatorSession CoreSyn.CoreExpr
391 genUniques = genUniques' VarEnv.emptyVarEnv
392
393 -- | A helper function to generate uniques, that takes a VarEnv containing the
394 --   substitutions already performed.
395 genUniques' :: VarEnv.VarEnv CoreSyn.CoreBndr -> CoreSyn.CoreExpr -> TranslatorSession CoreSyn.CoreExpr
396 genUniques' subst (CoreSyn.Var f) = do
397   -- Replace the binder with its new value, if applicable.
398   let f' = VarEnv.lookupWithDefaultVarEnv subst f f
399   return (CoreSyn.Var f')
400 -- Leave literals untouched
401 genUniques' subst (CoreSyn.Lit l) = return $ CoreSyn.Lit l
402 genUniques' subst (CoreSyn.App f arg) = do
403   -- Only work on subexpressions
404   f' <- genUniques' subst f
405   arg' <- genUniques' subst arg
406   return (CoreSyn.App f' arg')
407 -- Don't change type abstractions
408 genUniques' subst expr@(CoreSyn.Lam bndr res) | CoreSyn.isTyVar bndr = return expr
409 genUniques' subst (CoreSyn.Lam bndr res) = do
410   -- Generate a new unique for the bound variable
411   (subst', bndr') <- genUnique subst bndr
412   res' <- genUniques' subst' res
413   return (CoreSyn.Lam bndr' res')
414 genUniques' subst (CoreSyn.Let (CoreSyn.NonRec bndr bound) res) = do
415   -- Make the binders unique
416   (subst', bndr') <- genUnique subst bndr
417   bound' <- genUniques' subst' bound
418   res' <- genUniques' subst' res
419   return $ CoreSyn.Let (CoreSyn.NonRec bndr' bound') res'
420 genUniques' subst (CoreSyn.Let (CoreSyn.Rec binds) res) = do
421   -- Make each of the binders unique
422   (subst', bndrs') <- mapAccumLM genUnique subst (map fst binds)
423   bounds' <- mapM (genUniques' subst' . snd) binds
424   res' <- genUniques' subst' res
425   let binds' = zip bndrs' bounds'
426   return $ CoreSyn.Let (CoreSyn.Rec binds') res'
427 genUniques' subst (CoreSyn.Case scrut bndr ty alts) = do
428   -- Process the scrutinee with the original substitution, since non of the
429   -- binders bound in the Case statement is in scope in the scrutinee.
430   scrut' <- genUniques' subst scrut
431   -- Generate a new binder for the scrutinee
432   (subst', bndr') <- genUnique subst bndr
433   -- Process each of the alts
434   alts' <- mapM (doalt subst') alts
435   return $ CoreSyn.Case scrut' bndr' ty alts'
436   where
437     doalt subst (con, bndrs, expr) = do
438       (subst', bndrs') <- mapAccumLM genUnique subst bndrs
439       expr' <- genUniques' subst' expr
440       -- Note that we don't return subst', since bndrs are only in scope in
441       -- expr.
442       return (con, bndrs', expr')
443 genUniques' subst (CoreSyn.Cast expr coercion) = do
444   expr' <- genUniques' subst expr
445   -- Just process the casted expression
446   return $ CoreSyn.Cast expr' coercion
447 genUniques' subst (CoreSyn.Note note expr) = do
448   expr' <- genUniques' subst expr
449   -- Just process the annotated expression
450   return $ CoreSyn.Note note expr'
451 -- Leave types untouched
452 genUniques' subst expr@(CoreSyn.Type _) = return expr
453
454 -- Generate a new unique for the given binder, and extend the given
455 -- substitution to reflect this.
456 genUnique :: VarEnv.VarEnv CoreSyn.CoreBndr -> CoreSyn.CoreBndr -> TranslatorSession (VarEnv.VarEnv CoreSyn.CoreBndr, CoreSyn.CoreBndr)
457 genUnique subst bndr = do
458   bndr' <- BinderTools.cloneVar bndr
459   -- Replace all occurences of the old binder with a reference to the new
460   -- binder.
461   let subst' = VarEnv.extendVarEnv subst bndr bndr'
462   return (subst', bndr')
463
464 -- Create a "selector" case that selects the ith field from dc_ith
465 -- datacon
466 mkSelCase :: CoreSyn.CoreExpr -> Int -> Int -> TranslatorSession CoreSyn.CoreExpr
467 mkSelCase scrut dc_i i = do
468   case Type.splitTyConApp_maybe scrut_ty of
469     -- The scrutinee should have a type constructor. We keep the type
470     -- arguments around so we can instantiate the field types below
471     Just (tycon, tyargs) -> case TyCon.tyConDataCons_maybe tycon of
472       -- The scrutinee type should have a single dataconstructor,
473       -- otherwise we can't construct a valid selector case.
474       Just dcs | dc_i < 0 || dc_i >= length dcs -> error $ "\nCoreTools.mkSelCase: Creating extractor case, but datacon index is invalid." ++ error_msg
475                | otherwise -> do
476         let datacon = (dcs!!dc_i)
477         let field_tys = DataCon.dataConInstOrigArgTys datacon  tyargs
478         if i < 0 || i >= length field_tys
479           then error $ "\nCoreTools.mkSelCase: Creating extractor case, but field index is invalid." ++ error_msg
480           else do
481             -- Create a list of wild binders for the fields we don't want
482             let wildbndrs = map MkCore.mkWildBinder field_tys
483             -- Create a single binder for the field we want
484             sel_bndr <- mkInternalVar "sel" (field_tys!!i)
485             -- Create a wild binder for the scrutinee
486             let scrut_bndr = MkCore.mkWildBinder scrut_ty
487             -- Create the case expression
488             let binders = take i wildbndrs ++ [sel_bndr] ++ drop (i+1) wildbndrs
489             return $ CoreSyn.Case scrut scrut_bndr scrut_ty [(CoreSyn.DataAlt datacon, binders, CoreSyn.Var sel_bndr)]
490       Nothing -> error $ "CoreTools.mkSelCase: Creating extractor case, but scrutinee has no datacons?" ++ error_msg
491     Nothing -> error $ "CoreTools.mkSelCase: Creating extractor case, but scrutinee has no tycon?" ++ error_msg
492   where
493     scrut_ty = CoreUtils.exprType scrut
494     error_msg = " Extracting element " ++ (show i) ++ " from datacon " ++ (show dc_i) ++ " from '" ++ pprString scrut ++ "'" ++ " Type: " ++ (pprString scrut_ty)