d8f4289d28b25253544126ba0424f90f6fc842de
[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 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 HsBinds
19 import qualified HscTypes
20 import qualified RdrName
21 import qualified Name
22 import qualified OccName
23 import qualified Type
24 import qualified Id
25 import qualified TyCon
26 import qualified DataCon
27 import qualified TysWiredIn
28 import qualified Bag
29 import qualified DynFlags
30 import qualified SrcLoc
31 import qualified CoreSyn
32 import qualified Var
33 import qualified IdInfo
34 import qualified VarSet
35 import qualified Unique
36 import qualified CoreUtils
37 import qualified CoreFVs
38 import qualified Literal
39
40 -- Local imports
41 import CLasH.Translator.TranslatorTypes
42 import CLasH.Utils.GhcTools
43 import CLasH.Utils.HsTools
44 import CLasH.Utils.Pretty
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.
51 eval_tfp_int :: HscTypes.HscEnv -> Type.Type -> Int
52 eval_tfp_int env ty =
53   unsafeRunGhc libdir $ do
54     GHC.setSession env
55     -- Automatically import modules for any fully qualified identifiers
56     setDynFlag DynFlags.Opt_ImplicitImportQualified
57
58     let from_int_t_name = mkRdrName "Types.Data.Num.Ops" "fromIntegerT"
59     let from_int_t = SrcLoc.noLoc $ HsExpr.HsVar from_int_t_name
60     let undef = hsTypedUndef $ coreToHsType ty
61     let app = SrcLoc.noLoc $ HsExpr.HsApp (from_int_t) (undef)
62     let int_ty = SrcLoc.noLoc $ HsTypes.HsTyVar TysWiredIn.intTyCon_RDR
63     let expr = HsExpr.ExprWithTySig app int_ty
64     core <- toCore expr
65     execCore core
66   where
67     libdir = DynFlags.topDir dynflags
68     dynflags = HscTypes.hsc_dflags env
69
70 normalise_tfp_int :: HscTypes.HscEnv -> Type.Type -> Type.Type
71 normalise_tfp_int env ty =
72    unsafePerformIO $ do
73      nty <- normaliseType env ty
74      return nty
75
76 -- | Get the width of a SizedWord type
77 -- sized_word_len :: HscTypes.HscEnv -> Type.Type -> Int
78 -- sized_word_len env ty = eval_tfp_int env (sized_word_len_ty ty)
79     
80 sized_word_len_ty :: Type.Type -> Type.Type
81 sized_word_len_ty ty = len
82   where
83     args = case Type.splitTyConApp_maybe ty of
84       Just (tycon, args) -> args
85       Nothing -> error $ "\nCoreTools.sized_word_len_ty: Not a sized word type: " ++ (pprString ty)
86     [len]         = args
87
88 -- | Get the width of a SizedInt type
89 -- sized_int_len :: HscTypes.HscEnv -> Type.Type -> Int
90 -- sized_int_len env ty = eval_tfp_int env (sized_int_len_ty ty)
91
92 sized_int_len_ty :: Type.Type -> Type.Type
93 sized_int_len_ty ty = len
94   where
95     args = case Type.splitTyConApp_maybe ty of
96       Just (tycon, args) -> args
97       Nothing -> error $ "\nCoreTools.sized_int_len_ty: Not a sized int type: " ++ (pprString ty)
98     [len]         = args
99     
100 -- | Get the upperbound of a RangedWord type
101 -- ranged_word_bound :: HscTypes.HscEnv -> Type.Type -> Int
102 -- ranged_word_bound env ty = eval_tfp_int env (ranged_word_bound_ty ty)
103     
104 ranged_word_bound_ty :: Type.Type -> Type.Type
105 ranged_word_bound_ty ty = len
106   where
107     args = case Type.splitTyConApp_maybe ty of
108       Just (tycon, args) -> args
109       Nothing -> error $ "\nCoreTools.ranged_word_bound_ty: Not a sized word type: " ++ (pprString ty)
110     [len]         = args
111
112 -- | Evaluate a core Type representing type level int from the TypeLevel
113 -- library to a real int.
114 -- eval_type_level_int :: Type.Type -> Int
115 -- eval_type_level_int ty =
116 --   unsafeRunGhc $ do
117 --     -- Automatically import modules for any fully qualified identifiers
118 --     setDynFlag DynFlags.Opt_ImplicitImportQualified
119 -- 
120 --     let to_int_name = mkRdrName "Data.TypeLevel.Num.Sets" "toInt"
121 --     let to_int = SrcLoc.noLoc $ HsExpr.HsVar to_int_name
122 --     let undef = hsTypedUndef $ coreToHsType ty
123 --     let app = HsExpr.HsApp (to_int) (undef)
124 -- 
125 --     core <- toCore [] app
126 --     execCore core 
127
128 -- | Get the length of a FSVec type
129 -- tfvec_len :: HscTypes.HscEnv -> Type.Type -> Int
130 -- tfvec_len env ty = eval_tfp_int env (tfvec_len_ty ty)
131
132 tfvec_len_ty :: Type.Type -> Type.Type
133 tfvec_len_ty ty = len
134   where  
135     args = case Type.splitTyConApp_maybe ty of
136       Just (tycon, args) -> args
137       Nothing -> error $ "\nCoreTools.tfvec_len_ty: Not a vector type: " ++ (pprString ty)
138     [len, el_ty] = args
139     
140 -- | Get the element type of a TFVec type
141 tfvec_elem :: Type.Type -> Type.Type
142 tfvec_elem ty = el_ty
143   where
144     args = case Type.splitTyConApp_maybe ty of
145       Just (tycon, args) -> args
146       Nothing -> error $ "\nCoreTools.tfvec_len: Not a vector type: " ++ (pprString ty)
147     [len, el_ty] = args
148
149 -- Is the given core expression a lambda abstraction?
150 is_lam :: CoreSyn.CoreExpr -> Bool
151 is_lam (CoreSyn.Lam _ _) = True
152 is_lam _ = False
153
154 -- Is the given core expression of a function type?
155 is_fun :: CoreSyn.CoreExpr -> Bool
156 -- Treat Type arguments differently, because exprType is not defined for them.
157 is_fun (CoreSyn.Type _) = False
158 is_fun expr = (Type.isFunTy . CoreUtils.exprType) expr
159
160 -- Is the given core expression polymorphic (i.e., does it accept type
161 -- arguments?).
162 is_poly :: CoreSyn.CoreExpr -> Bool
163 -- Treat Type arguments differently, because exprType is not defined for them.
164 is_poly (CoreSyn.Type _) = False
165 is_poly expr = (Maybe.isJust . Type.splitForAllTy_maybe . CoreUtils.exprType) expr
166
167 -- Is the given core expression a variable reference?
168 is_var :: CoreSyn.CoreExpr -> Bool
169 is_var (CoreSyn.Var _) = True
170 is_var _ = False
171
172 is_lit :: CoreSyn.CoreExpr -> Bool
173 is_lit (CoreSyn.Lit _) = True
174 is_lit _ = False
175
176 -- Can the given core expression be applied to something? This is true for
177 -- applying to a value as well as a type.
178 is_applicable :: CoreSyn.CoreExpr -> Bool
179 is_applicable expr = is_fun expr || is_poly expr
180
181 -- Is the given core expression a variable or an application?
182 is_simple :: CoreSyn.CoreExpr -> Bool
183 is_simple (CoreSyn.App _ _) = True
184 is_simple (CoreSyn.Var _) = True
185 is_simple (CoreSyn.Cast expr _) = is_simple expr
186 is_simple _ = False
187
188 -- Does the given CoreExpr have any free type vars?
189 has_free_tyvars :: CoreSyn.CoreExpr -> Bool
190 has_free_tyvars = not . VarSet.isEmptyVarSet . (CoreFVs.exprSomeFreeVars Var.isTyVar)
191
192 -- Does the given CoreExpr have any free local vars?
193 has_free_vars :: CoreSyn.CoreExpr -> Bool
194 has_free_vars = not . VarSet.isEmptyVarSet . CoreFVs.exprFreeVars
195
196 -- Does the given expression use any of the given binders?
197 expr_uses_binders :: [CoreSyn.CoreBndr] -> CoreSyn.CoreExpr -> Bool
198 expr_uses_binders bndrs = not . VarSet.isEmptyVarSet . (CoreFVs.exprSomeFreeVars (`elem` bndrs))
199
200 -- Turns a Var CoreExpr into the Id inside it. Will of course only work for
201 -- simple Var CoreExprs, not complexer ones.
202 exprToVar :: CoreSyn.CoreExpr -> Var.Id
203 exprToVar (CoreSyn.Var id) = id
204 exprToVar expr = error $ "\nCoreTools.exprToVar: Not a var: " ++ show expr
205
206 -- Turns a Lit CoreExpr into the Literal inside it.
207 exprToLit :: CoreSyn.CoreExpr -> Literal.Literal
208 exprToLit (CoreSyn.Lit lit) = lit
209 exprToLit expr = error $ "\nCoreTools.exprToLit: Not a lit: " ++ show expr
210
211 -- Removes all the type and dictionary arguments from the given argument list,
212 -- leaving only the normal value arguments. The type given is the type of the
213 -- expression applied to this argument list.
214 get_val_args :: Type.Type -> [CoreSyn.CoreExpr] -> [CoreSyn.CoreExpr]
215 get_val_args ty args = drop n args
216   where
217     (tyvars, predtypes, _) = TcType.tcSplitSigmaTy ty
218     -- The first (length tyvars) arguments should be types, the next 
219     -- (length predtypes) arguments should be dictionaries. We drop this many
220     -- arguments, to get at the value arguments.
221     n = length tyvars + length predtypes
222
223 getLiterals :: CoreSyn.CoreExpr -> [CoreSyn.CoreExpr]
224 getLiterals app@(CoreSyn.App _ _) = literals
225   where
226     (CoreSyn.Var f, args) = CoreSyn.collectArgs app
227     literals = filter (is_lit) args
228
229 getLiterals lit@(CoreSyn.Lit _) = [lit]
230
231 reduceCoreListToHsList :: 
232   [HscTypes.CoreModule] -- ^ The modules where parts of the list are hidden
233   -> CoreSyn.CoreExpr   -- ^ The refence to atleast one of the nodes
234   -> TranslatorSession [CoreSyn.CoreExpr]
235 reduceCoreListToHsList cores app@(CoreSyn.App _ _) = do {
236   ; let { (fun, args) = CoreSyn.collectArgs app
237         ; len         = length args 
238         } ;
239   ; case len of
240       3 -> do {
241         ; let topelem = args!!1
242         ; case (args!!2) of
243             (varz@(CoreSyn.Var id)) -> do {
244               ; binds <- mapM (findExpr (isVarName id)) cores
245               ; otherelems <- reduceCoreListToHsList cores (head (Maybe.catMaybes binds))
246               ; return (topelem:otherelems)
247               }
248             (appz@(CoreSyn.App _ _)) -> do {
249               ; otherelems <- reduceCoreListToHsList cores appz
250               ; return (topelem:otherelems)
251               }
252             otherwise -> return [topelem]
253         }
254       otherwise -> return []
255   }
256   where
257     isVarName :: Monad m => Var.Var -> Var.Var -> m Bool
258     isVarName lookfor bind = return $ (Var.varName lookfor) == (Var.varName bind)
259
260 reduceCoreListToHsList _ _ = return []
261
262 -- Is the given var the State data constructor?
263 isStateCon :: Var.Var -> Bool
264 isStateCon var = do
265   -- See if it is a DataConWrapId (not DataConWorkId, since State is a
266   -- newtype).
267   case Id.idDetails var of
268     IdInfo.DataConWrapId dc -> 
269       -- See if the datacon is the State datacon from the State type.
270       let tycon = DataCon.dataConTyCon dc
271           tyname = Name.getOccString tycon
272           dcname = Name.getOccString dc
273       in case (tyname, dcname) of
274         ("State", "State") -> True
275         _ -> False
276     _ -> False
277
278 -- | Is the given type a State type?
279 isStateType :: Type.Type -> Bool
280 -- Resolve any type synonyms remaining
281 isStateType ty | Just ty' <- Type.tcView ty = isStateType ty'
282 isStateType ty  = Maybe.isJust $ do
283   -- Split the type. Don't use normal splitAppTy, since that looks through
284   -- newtypes, and we want to see the State newtype.
285   (typef, _) <- Type.repSplitAppTy_maybe ty
286   -- See if the applied type is a type constructor
287   (tycon, _) <- Type.splitTyConApp_maybe typef
288   if TyCon.isNewTyCon tycon && Name.getOccString tycon == "State"
289     then
290       Just ()
291     else
292       Nothing
293
294 -- | Does the given TypedThing have a State type?
295 hasStateType :: (TypedThing t) => t -> Bool
296 hasStateType expr = case getType expr of
297   Nothing -> False
298   Just ty -> isStateType ty
299
300
301 -- | A class of things that (optionally) have a core Type. The type is
302 -- optional, since Type expressions don't have a type themselves.
303 class TypedThing t where
304   getType :: t -> Maybe Type.Type
305
306 instance TypedThing CoreSyn.CoreExpr where
307   getType (CoreSyn.Type _) = Nothing
308   getType expr = Just $ CoreUtils.exprType expr
309
310 instance TypedThing CoreSyn.CoreBndr where
311   getType = return . Id.idType
312
313 instance TypedThing Type.Type where
314   getType = return . id