Fix error message genNegation for non-SizedWord types
[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 import qualified VarEnv
41 import qualified Literal
42
43 -- Local imports
44 import CLasH.Translator.TranslatorTypes
45 import CLasH.Utils.GhcTools
46 import CLasH.Utils.HsTools
47 import CLasH.Utils.Pretty
48 import CLasH.Utils
49 import qualified CLasH.Utils.Core.BinderTools as BinderTools
50
51 -- | A single binding, used as a shortcut to simplify type signatures.
52 type Binding = (CoreSyn.CoreBndr, CoreSyn.CoreExpr)
53
54 -- | Evaluate a core Type representing type level int from the tfp
55 -- library to a real int.
56 eval_tfp_int :: HscTypes.HscEnv -> Type.Type -> Int
57 eval_tfp_int env ty =
58   unsafeRunGhc libdir $ do
59     GHC.setSession env
60     -- Automatically import modules for any fully qualified identifiers
61     setDynFlag DynFlags.Opt_ImplicitImportQualified
62
63     let from_int_t_name = mkRdrName "Types.Data.Num.Ops" "fromIntegerT"
64     let from_int_t = SrcLoc.noLoc $ HsExpr.HsVar from_int_t_name
65     let undef = hsTypedUndef $ coreToHsType ty
66     let app = SrcLoc.noLoc $ HsExpr.HsApp (from_int_t) (undef)
67     let int_ty = SrcLoc.noLoc $ HsTypes.HsTyVar TysWiredIn.intTyCon_RDR
68     let expr = HsExpr.ExprWithTySig app int_ty
69     core <- toCore expr
70     execCore core
71   where
72     libdir = DynFlags.topDir dynflags
73     dynflags = HscTypes.hsc_dflags env
74
75 normalise_tfp_int :: HscTypes.HscEnv -> Type.Type -> Type.Type
76 normalise_tfp_int env ty =
77    unsafePerformIO $ do
78      nty <- normaliseType env ty
79      return nty
80
81 -- | Get the width of a SizedWord type
82 -- sized_word_len :: HscTypes.HscEnv -> Type.Type -> Int
83 -- sized_word_len env ty = eval_tfp_int env (sized_word_len_ty ty)
84     
85 sized_word_len_ty :: Type.Type -> Type.Type
86 sized_word_len_ty ty = len
87   where
88     args = case Type.splitTyConApp_maybe ty of
89       Just (tycon, args) -> args
90       Nothing -> error $ "\nCoreTools.sized_word_len_ty: Not a sized word type: " ++ (pprString ty)
91     [len]         = args
92
93 -- | Get the width of a SizedInt type
94 -- sized_int_len :: HscTypes.HscEnv -> Type.Type -> Int
95 -- sized_int_len env ty = eval_tfp_int env (sized_int_len_ty ty)
96
97 sized_int_len_ty :: Type.Type -> Type.Type
98 sized_int_len_ty ty = len
99   where
100     args = case Type.splitTyConApp_maybe ty of
101       Just (tycon, args) -> args
102       Nothing -> error $ "\nCoreTools.sized_int_len_ty: Not a sized int type: " ++ (pprString ty)
103     [len]         = args
104     
105 -- | Get the upperbound of a RangedWord type
106 -- ranged_word_bound :: HscTypes.HscEnv -> Type.Type -> Int
107 -- ranged_word_bound env ty = eval_tfp_int env (ranged_word_bound_ty ty)
108     
109 ranged_word_bound_ty :: Type.Type -> Type.Type
110 ranged_word_bound_ty ty = len
111   where
112     args = case Type.splitTyConApp_maybe ty of
113       Just (tycon, args) -> args
114       Nothing -> error $ "\nCoreTools.ranged_word_bound_ty: Not a sized word type: " ++ (pprString ty)
115     [len]         = args
116
117 -- | Evaluate a core Type representing type level int from the TypeLevel
118 -- library to a real int.
119 -- eval_type_level_int :: Type.Type -> Int
120 -- eval_type_level_int ty =
121 --   unsafeRunGhc $ do
122 --     -- Automatically import modules for any fully qualified identifiers
123 --     setDynFlag DynFlags.Opt_ImplicitImportQualified
124 -- 
125 --     let to_int_name = mkRdrName "Data.TypeLevel.Num.Sets" "toInt"
126 --     let to_int = SrcLoc.noLoc $ HsExpr.HsVar to_int_name
127 --     let undef = hsTypedUndef $ coreToHsType ty
128 --     let app = HsExpr.HsApp (to_int) (undef)
129 -- 
130 --     core <- toCore [] app
131 --     execCore core 
132
133 -- | Get the length of a FSVec type
134 -- tfvec_len :: HscTypes.HscEnv -> Type.Type -> Int
135 -- tfvec_len env ty = eval_tfp_int env (tfvec_len_ty ty)
136
137 tfvec_len_ty :: Type.Type -> Type.Type
138 tfvec_len_ty ty = len
139   where  
140     args = case Type.splitTyConApp_maybe ty of
141       Just (tycon, args) -> args
142       Nothing -> error $ "\nCoreTools.tfvec_len_ty: Not a vector type: " ++ (pprString ty)
143     [len, el_ty] = args
144     
145 -- | Get the element type of a TFVec type
146 tfvec_elem :: Type.Type -> Type.Type
147 tfvec_elem ty = el_ty
148   where
149     args = case Type.splitTyConApp_maybe ty of
150       Just (tycon, args) -> args
151       Nothing -> error $ "\nCoreTools.tfvec_len: Not a vector type: " ++ (pprString ty)
152     [len, el_ty] = args
153
154 -- Is the given core expression a lambda abstraction?
155 is_lam :: CoreSyn.CoreExpr -> Bool
156 is_lam (CoreSyn.Lam _ _) = True
157 is_lam _ = False
158
159 -- Is the given core expression of a function type?
160 is_fun :: CoreSyn.CoreExpr -> Bool
161 -- Treat Type arguments differently, because exprType is not defined for them.
162 is_fun (CoreSyn.Type _) = False
163 is_fun expr = (Type.isFunTy . CoreUtils.exprType) expr
164
165 -- Is the given core expression polymorphic (i.e., does it accept type
166 -- arguments?).
167 is_poly :: CoreSyn.CoreExpr -> Bool
168 -- Treat Type arguments differently, because exprType is not defined for them.
169 is_poly (CoreSyn.Type _) = False
170 is_poly expr = (Maybe.isJust . Type.splitForAllTy_maybe . CoreUtils.exprType) expr
171
172 -- Is the given core expression a variable reference?
173 is_var :: CoreSyn.CoreExpr -> Bool
174 is_var (CoreSyn.Var _) = True
175 is_var _ = False
176
177 is_lit :: CoreSyn.CoreExpr -> Bool
178 is_lit (CoreSyn.Lit _) = True
179 is_lit _ = False
180
181 -- Can the given core expression be applied to something? This is true for
182 -- applying to a value as well as a type.
183 is_applicable :: CoreSyn.CoreExpr -> Bool
184 is_applicable expr = is_fun expr || is_poly expr
185
186 -- Is the given core expression a variable or an application?
187 is_simple :: CoreSyn.CoreExpr -> Bool
188 is_simple (CoreSyn.App _ _) = True
189 is_simple (CoreSyn.Var _) = True
190 is_simple (CoreSyn.Cast expr _) = is_simple expr
191 is_simple _ = False
192
193 -- Does the given CoreExpr have any free type vars?
194 has_free_tyvars :: CoreSyn.CoreExpr -> Bool
195 has_free_tyvars = not . VarSet.isEmptyVarSet . (CoreFVs.exprSomeFreeVars Var.isTyVar)
196
197 -- Does the given type have any free type vars?
198 ty_has_free_tyvars :: Type.Type -> Bool
199 ty_has_free_tyvars = not . VarSet.isEmptyVarSet . Type.tyVarsOfType
200
201 -- Does the given CoreExpr have any free local vars?
202 has_free_vars :: CoreSyn.CoreExpr -> Bool
203 has_free_vars = not . VarSet.isEmptyVarSet . CoreFVs.exprFreeVars
204
205 -- Does the given expression use any of the given binders?
206 expr_uses_binders :: [CoreSyn.CoreBndr] -> CoreSyn.CoreExpr -> Bool
207 expr_uses_binders bndrs = not . VarSet.isEmptyVarSet . (CoreFVs.exprSomeFreeVars (`elem` bndrs))
208
209 -- Turns a Var CoreExpr into the Id inside it. Will of course only work for
210 -- simple Var CoreExprs, not complexer ones.
211 exprToVar :: CoreSyn.CoreExpr -> Var.Id
212 exprToVar (CoreSyn.Var id) = id
213 exprToVar expr = error $ "\nCoreTools.exprToVar: Not a var: " ++ show expr
214
215 -- Turns a Lit CoreExpr into the Literal inside it.
216 exprToLit :: CoreSyn.CoreExpr -> Literal.Literal
217 exprToLit (CoreSyn.Lit lit) = lit
218 exprToLit expr = error $ "\nCoreTools.exprToLit: Not a lit: " ++ show expr
219
220 -- Removes all the type and dictionary arguments from the given argument list,
221 -- leaving only the normal value arguments. The type given is the type of the
222 -- expression applied to this argument list.
223 get_val_args :: Type.Type -> [CoreSyn.CoreExpr] -> [CoreSyn.CoreExpr]
224 get_val_args ty args = drop n args
225   where
226     (tyvars, predtypes, _) = TcType.tcSplitSigmaTy ty
227     -- The first (length tyvars) arguments should be types, the next 
228     -- (length predtypes) arguments should be dictionaries. We drop this many
229     -- arguments, to get at the value arguments.
230     n = length tyvars + length predtypes
231
232 getLiterals :: HscTypes.HscEnv -> CoreSyn.CoreExpr -> [CoreSyn.CoreExpr]
233 getLiterals _ app@(CoreSyn.App _ _) = literals
234   where
235     (CoreSyn.Var f, args) = CoreSyn.collectArgs app
236     literals = filter (is_lit) args
237
238 getLiterals _ lit@(CoreSyn.Lit _) = [lit]
239
240 getLiterals hscenv letrec@(CoreSyn.Let (CoreSyn.NonRec letBind (letExpr)) letRes) = [lit]
241   where
242     ty     = Var.varType letBind
243     litInt = eval_tfp_int hscenv ty
244     lit    = CoreSyn.Lit (Literal.mkMachInt (toInteger litInt))
245
246 getLiterals _ expr = error $ "\nCoreTools.getLiterals: Not a known Lit: " ++ pprString expr
247
248 reduceCoreListToHsList :: 
249   [HscTypes.CoreModule] -- ^ The modules where parts of the list are hidden
250   -> CoreSyn.CoreExpr   -- ^ The refence to atleast one of the nodes
251   -> TranslatorSession [CoreSyn.CoreExpr]
252 reduceCoreListToHsList cores app@(CoreSyn.App _ _) = do {
253   ; let { (fun, args) = CoreSyn.collectArgs app
254         ; len         = length args 
255         } ;
256   ; case len of
257       3 -> do {
258         ; let topelem = args!!1
259         ; case (args!!2) of
260             (varz@(CoreSyn.Var id)) -> do {
261               ; binds <- mapM (findExpr (isVarName id)) cores
262               ; otherelems <- reduceCoreListToHsList cores (head (Maybe.catMaybes binds))
263               ; return (topelem:otherelems)
264               }
265             (appz@(CoreSyn.App _ _)) -> do {
266               ; otherelems <- reduceCoreListToHsList cores appz
267               ; return (topelem:otherelems)
268               }
269             otherwise -> return [topelem]
270         }
271       otherwise -> return []
272   }
273   where
274     isVarName :: Monad m => Var.Var -> Var.Var -> m Bool
275     isVarName lookfor bind = return $ (Var.varName lookfor) == (Var.varName bind)
276
277 reduceCoreListToHsList _ _ = return []
278
279 -- Is the given var the State data constructor?
280 isStateCon :: Var.Var -> Bool
281 isStateCon var = do
282   -- See if it is a DataConWrapId (not DataConWorkId, since State is a
283   -- newtype).
284   case Id.idDetails var of
285     IdInfo.DataConWrapId dc -> 
286       -- See if the datacon is the State datacon from the State type.
287       let tycon = DataCon.dataConTyCon dc
288           tyname = Name.getOccString tycon
289           dcname = Name.getOccString dc
290       in case (tyname, dcname) of
291         ("State", "State") -> True
292         _ -> False
293     _ -> False
294
295 -- | Is the given type a State type?
296 isStateType :: Type.Type -> Bool
297 -- Resolve any type synonyms remaining
298 isStateType ty | Just ty' <- Type.tcView ty = isStateType ty'
299 isStateType ty  = Maybe.isJust $ do
300   -- Split the type. Don't use normal splitAppTy, since that looks through
301   -- newtypes, and we want to see the State newtype.
302   (typef, _) <- Type.repSplitAppTy_maybe ty
303   -- See if the applied type is a type constructor
304   (tycon, _) <- Type.splitTyConApp_maybe typef
305   if TyCon.isNewTyCon tycon && Name.getOccString tycon == "State"
306     then
307       Just ()
308     else
309       Nothing
310
311 -- | Does the given TypedThing have a State type?
312 hasStateType :: (TypedThing t) => t -> Bool
313 hasStateType expr = case getType expr of
314   Nothing -> False
315   Just ty -> isStateType ty
316
317
318 -- | Flattens nested lets into a single list of bindings. The expression
319 --   passed does not have to be a let expression, if it isn't an empty list of
320 --   bindings is returned.
321 flattenLets ::
322   CoreSyn.CoreExpr -- ^ The expression to flatten.
323   -> ([Binding], CoreSyn.CoreExpr) -- ^ The bindings and resulting expression.
324 flattenLets (CoreSyn.Let binds expr) = 
325   (bindings ++ bindings', expr')
326   where
327     -- Recursively flatten the contained expression
328     (bindings', expr') =flattenLets expr
329     -- Flatten our own bindings to remove the Rec / NonRec constructors
330     bindings = CoreSyn.flattenBinds [binds]
331 flattenLets expr = ([], expr)
332
333 -- | Create bunch of nested non-recursive let expressions from the given
334 -- bindings. The first binding is bound at the highest level (and thus
335 -- available in all other bindings).
336 mkNonRecLets :: [Binding] -> CoreSyn.CoreExpr -> CoreSyn.CoreExpr
337 mkNonRecLets bindings expr = MkCore.mkCoreLets binds expr
338   where
339     binds = map (uncurry CoreSyn.NonRec) bindings
340
341 -- | A class of things that (optionally) have a core Type. The type is
342 -- optional, since Type expressions don't have a type themselves.
343 class TypedThing t where
344   getType :: t -> Maybe Type.Type
345
346 instance TypedThing CoreSyn.CoreExpr where
347   getType (CoreSyn.Type _) = Nothing
348   getType expr = Just $ CoreUtils.exprType expr
349
350 instance TypedThing CoreSyn.CoreBndr where
351   getType = return . Id.idType
352
353 instance TypedThing Type.Type where
354   getType = return . id
355
356 -- | Generate new uniques for all binders in the given expression.
357 -- Does not support making type variables unique, though this could be
358 -- supported if required (by passing a CoreSubst.Subst instead of VarEnv to
359 -- genUniques' below).
360 genUniques :: CoreSyn.CoreExpr -> TranslatorSession CoreSyn.CoreExpr
361 genUniques = genUniques' VarEnv.emptyVarEnv
362
363 -- | A helper function to generate uniques, that takes a VarEnv containing the
364 --   substitutions already performed.
365 genUniques' :: VarEnv.VarEnv CoreSyn.CoreBndr -> CoreSyn.CoreExpr -> TranslatorSession CoreSyn.CoreExpr
366 genUniques' subst (CoreSyn.Var f) = do
367   -- Replace the binder with its new value, if applicable.
368   let f' = VarEnv.lookupWithDefaultVarEnv subst f f
369   return (CoreSyn.Var f')
370 -- Leave literals untouched
371 genUniques' subst (CoreSyn.Lit l) = return $ CoreSyn.Lit l
372 genUniques' subst (CoreSyn.App f arg) = do
373   -- Only work on subexpressions
374   f' <- genUniques' subst f
375   arg' <- genUniques' subst arg
376   return (CoreSyn.App f' arg')
377 -- Don't change type abstractions
378 genUniques' subst expr@(CoreSyn.Lam bndr res) | CoreSyn.isTyVar bndr = return expr
379 genUniques' subst (CoreSyn.Lam bndr res) = do
380   -- Generate a new unique for the bound variable
381   (subst', bndr') <- genUnique subst bndr
382   res' <- genUniques' subst' res
383   return (CoreSyn.Lam bndr' res')
384 genUniques' subst (CoreSyn.Let (CoreSyn.NonRec bndr bound) res) = do
385   -- Make the binders unique
386   (subst', bndr') <- genUnique subst bndr
387   bound' <- genUniques' subst' bound
388   res' <- genUniques' subst' res
389   return $ CoreSyn.Let (CoreSyn.NonRec bndr' bound') res'
390 genUniques' subst (CoreSyn.Let (CoreSyn.Rec binds) res) = do
391   -- Make each of the binders unique
392   (subst', bndrs') <- mapAccumLM genUnique subst (map fst binds)
393   bounds' <- mapM (genUniques' subst') (map snd binds)
394   res' <- genUniques' subst' res
395   let binds' = zip bndrs' bounds'
396   return $ CoreSyn.Let (CoreSyn.Rec binds') res'
397 genUniques' subst (CoreSyn.Case scrut bndr ty alts) = do
398   -- Process the scrutinee with the original substitution, since non of the
399   -- binders bound in the Case statement is in scope in the scrutinee.
400   scrut' <- genUniques' subst scrut
401   -- Generate a new binder for the scrutinee
402   (subst', bndr') <- genUnique subst bndr
403   -- Process each of the alts
404   alts' <- mapM (doalt subst') alts
405   return $ CoreSyn.Case scrut' bndr' ty alts'
406   where
407     doalt subst (con, bndrs, expr) = do
408       (subst', bndrs') <- mapAccumLM genUnique subst bndrs
409       expr' <- genUniques' subst' expr
410       -- Note that we don't return subst', since bndrs are only in scope in
411       -- expr.
412       return (con, bndrs', expr')
413 genUniques' subst (CoreSyn.Cast expr coercion) = do
414   expr' <- genUniques' subst expr
415   -- Just process the casted expression
416   return $ CoreSyn.Cast expr' coercion
417 genUniques' subst (CoreSyn.Note note expr) = do
418   expr' <- genUniques' subst expr
419   -- Just process the annotated expression
420   return $ CoreSyn.Note note expr'
421 -- Leave types untouched
422 genUniques' subst expr@(CoreSyn.Type _) = return expr
423
424 -- Generate a new unique for the given binder, and extend the given
425 -- substitution to reflect this.
426 genUnique :: VarEnv.VarEnv CoreSyn.CoreBndr -> CoreSyn.CoreBndr -> TranslatorSession (VarEnv.VarEnv CoreSyn.CoreBndr, CoreSyn.CoreBndr)
427 genUnique subst bndr = do
428   bndr' <- BinderTools.cloneVar bndr
429   -- Replace all occurences of the old binder with a reference to the new
430   -- binder.
431   let subst' = VarEnv.extendVarEnv subst bndr bndr'
432   return (subst', bndr')