45721a891a7bf6d1662465322795daa2995a67ff
[matthijs/master-project/cλash.git] / cλash / CLasH / Utils / Core / CoreTools.hs
1 -- | This module provides a number of functions to find out things about Core
2 -- programs. This module does not provide the actual plumbing to work with
3 -- Core and Haskell (it uses HsTools for this), but only the functions that
4 -- know about various libraries and know which functions to call.
5 module CLasH.Utils.Core.CoreTools where
6
7 --Standard modules
8 import qualified Maybe
9 import System.IO.Unsafe
10
11 -- GHC API
12 import qualified GHC
13 import qualified Type
14 import qualified TcType
15 import qualified HsExpr
16 import qualified HsTypes
17 import qualified HsBinds
18 import qualified HscTypes
19 import qualified RdrName
20 import qualified Name
21 import qualified OccName
22 import qualified TysWiredIn
23 import qualified Bag
24 import qualified DynFlags
25 import qualified SrcLoc
26 import qualified CoreSyn
27 import qualified Var
28 import qualified VarSet
29 import qualified Unique
30 import qualified CoreUtils
31 import qualified CoreFVs
32 import qualified Literal
33
34 -- Local imports
35 import CLasH.Utils.GhcTools
36 import CLasH.Utils.HsTools
37 import CLasH.Utils.Pretty
38
39 -- | Evaluate a core Type representing type level int from the tfp
40 -- library to a real int.
41 eval_tfp_int :: HscTypes.HscEnv -> Type.Type -> Int
42 eval_tfp_int env ty =
43   unsafeRunGhc $ do
44     GHC.setSession env
45     -- Automatically import modules for any fully qualified identifiers
46     setDynFlag DynFlags.Opt_ImplicitImportQualified
47
48     let from_int_t_name = mkRdrName "Types.Data.Num" "fromIntegerT"
49     let from_int_t = SrcLoc.noLoc $ HsExpr.HsVar from_int_t_name
50     let undef = hsTypedUndef $ coreToHsType ty
51     let app = SrcLoc.noLoc $ HsExpr.HsApp (from_int_t) (undef)
52     let int_ty = SrcLoc.noLoc $ HsTypes.HsTyVar TysWiredIn.intTyCon_RDR
53     let expr = HsExpr.ExprWithTySig app int_ty
54     core <- toCore expr
55     execCore core 
56
57 normalise_tfp_int :: HscTypes.HscEnv -> Type.Type -> Type.Type
58 normalise_tfp_int env ty =
59    unsafePerformIO $ do
60      nty <- normaliseType env ty
61      return nty
62
63 -- | Get the width of a SizedWord type
64 -- sized_word_len :: HscTypes.HscEnv -> Type.Type -> Int
65 -- sized_word_len env ty = eval_tfp_int env (sized_word_len_ty ty)
66     
67 sized_word_len_ty :: Type.Type -> Type.Type
68 sized_word_len_ty ty = len
69   where
70     args = case Type.splitTyConApp_maybe ty of
71       Just (tycon, args) -> args
72       Nothing -> error $ "\nCoreTools.sized_word_len_ty: Not a sized word type: " ++ (pprString ty)
73     [len]         = args
74
75 -- | Get the width of a SizedInt type
76 -- sized_int_len :: HscTypes.HscEnv -> Type.Type -> Int
77 -- sized_int_len env ty = eval_tfp_int env (sized_int_len_ty ty)
78
79 sized_int_len_ty :: Type.Type -> Type.Type
80 sized_int_len_ty ty = len
81   where
82     args = case Type.splitTyConApp_maybe ty of
83       Just (tycon, args) -> args
84       Nothing -> error $ "\nCoreTools.sized_int_len_ty: Not a sized int type: " ++ (pprString ty)
85     [len]         = args
86     
87 -- | Get the upperbound of a RangedWord type
88 -- ranged_word_bound :: HscTypes.HscEnv -> Type.Type -> Int
89 -- ranged_word_bound env ty = eval_tfp_int env (ranged_word_bound_ty ty)
90     
91 ranged_word_bound_ty :: Type.Type -> Type.Type
92 ranged_word_bound_ty ty = len
93   where
94     args = case Type.splitTyConApp_maybe ty of
95       Just (tycon, args) -> args
96       Nothing -> error $ "\nCoreTools.ranged_word_bound_ty: Not a sized word type: " ++ (pprString ty)
97     [len]         = args
98
99 -- | Evaluate a core Type representing type level int from the TypeLevel
100 -- library to a real int.
101 -- eval_type_level_int :: Type.Type -> Int
102 -- eval_type_level_int ty =
103 --   unsafeRunGhc $ do
104 --     -- Automatically import modules for any fully qualified identifiers
105 --     setDynFlag DynFlags.Opt_ImplicitImportQualified
106 -- 
107 --     let to_int_name = mkRdrName "Data.TypeLevel.Num.Sets" "toInt"
108 --     let to_int = SrcLoc.noLoc $ HsExpr.HsVar to_int_name
109 --     let undef = hsTypedUndef $ coreToHsType ty
110 --     let app = HsExpr.HsApp (to_int) (undef)
111 -- 
112 --     core <- toCore [] app
113 --     execCore core 
114
115 -- | Get the length of a FSVec type
116 -- tfvec_len :: HscTypes.HscEnv -> Type.Type -> Int
117 -- tfvec_len env ty = eval_tfp_int env (tfvec_len_ty ty)
118
119 tfvec_len_ty :: Type.Type -> Type.Type
120 tfvec_len_ty ty = len
121   where  
122     args = case Type.splitTyConApp_maybe ty of
123       Just (tycon, args) -> args
124       Nothing -> error $ "\nCoreTools.tfvec_len_ty: Not a vector type: " ++ (pprString ty)
125     [len, el_ty] = args
126     
127 -- | Get the element type of a TFVec type
128 tfvec_elem :: Type.Type -> Type.Type
129 tfvec_elem ty = el_ty
130   where
131     args = case Type.splitTyConApp_maybe ty of
132       Just (tycon, args) -> args
133       Nothing -> error $ "\nCoreTools.tfvec_len: Not a vector type: " ++ (pprString ty)
134     [len, el_ty] = args
135
136 -- Is the given core expression a lambda abstraction?
137 is_lam :: CoreSyn.CoreExpr -> Bool
138 is_lam (CoreSyn.Lam _ _) = True
139 is_lam _ = False
140
141 -- Is the given core expression of a function type?
142 is_fun :: CoreSyn.CoreExpr -> Bool
143 -- Treat Type arguments differently, because exprType is not defined for them.
144 is_fun (CoreSyn.Type _) = False
145 is_fun expr = (Type.isFunTy . CoreUtils.exprType) expr
146
147 -- Is the given core expression polymorphic (i.e., does it accept type
148 -- arguments?).
149 is_poly :: CoreSyn.CoreExpr -> Bool
150 -- Treat Type arguments differently, because exprType is not defined for them.
151 is_poly (CoreSyn.Type _) = False
152 is_poly expr = (Maybe.isJust . Type.splitForAllTy_maybe . CoreUtils.exprType) expr
153
154 -- Is the given core expression a variable reference?
155 is_var :: CoreSyn.CoreExpr -> Bool
156 is_var (CoreSyn.Var _) = True
157 is_var _ = False
158
159 is_lit :: CoreSyn.CoreExpr -> Bool
160 is_lit (CoreSyn.Lit _) = True
161 is_lit _ = False
162
163 -- Can the given core expression be applied to something? This is true for
164 -- applying to a value as well as a type.
165 is_applicable :: CoreSyn.CoreExpr -> Bool
166 is_applicable expr = is_fun expr || is_poly expr
167
168 -- Is the given core expression a variable or an application?
169 is_simple :: CoreSyn.CoreExpr -> Bool
170 is_simple (CoreSyn.App _ _) = True
171 is_simple (CoreSyn.Var _) = True
172 is_simple (CoreSyn.Cast expr _) = is_simple expr
173 is_simple _ = False
174
175 -- Does the given CoreExpr have any free type vars?
176 has_free_tyvars :: CoreSyn.CoreExpr -> Bool
177 has_free_tyvars = not . VarSet.isEmptyVarSet . (CoreFVs.exprSomeFreeVars Var.isTyVar)
178
179 -- Does the given CoreExpr have any free local vars?
180 has_free_vars :: CoreSyn.CoreExpr -> Bool
181 has_free_vars = not . VarSet.isEmptyVarSet . CoreFVs.exprFreeVars
182
183 -- Turns a Var CoreExpr into the Id inside it. Will of course only work for
184 -- simple Var CoreExprs, not complexer ones.
185 exprToVar :: CoreSyn.CoreExpr -> Var.Id
186 exprToVar (CoreSyn.Var id) = id
187 exprToVar expr = error $ "\nCoreTools.exprToVar: Not a var: " ++ show expr
188
189 -- Turns a Lit CoreExpr into the Literal inside it.
190 exprToLit :: CoreSyn.CoreExpr -> Literal.Literal
191 exprToLit (CoreSyn.Lit lit) = lit
192 exprToLit expr = error $ "\nCoreTools.exprToLit: Not a lit: " ++ show expr
193
194 -- Removes all the type and dictionary arguments from the given argument list,
195 -- leaving only the normal value arguments. The type given is the type of the
196 -- expression applied to this argument list.
197 get_val_args :: Type.Type -> [CoreSyn.CoreExpr] -> [CoreSyn.CoreExpr]
198 get_val_args ty args = drop n args
199   where
200     (tyvars, predtypes, _) = TcType.tcSplitSigmaTy ty
201     -- The first (length tyvars) arguments should be types, the next 
202     -- (length predtypes) arguments should be dictionaries. We drop this many
203     -- arguments, to get at the value arguments.
204     n = length tyvars + length predtypes
205
206 getLiterals :: CoreSyn.CoreExpr -> [CoreSyn.CoreExpr]
207 getLiterals app@(CoreSyn.App _ _) = literals
208   where
209     (CoreSyn.Var f, args) = CoreSyn.collectArgs app
210     literals = filter (is_lit) args