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