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