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