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 CoreExpr have any free local vars?
197 has_free_vars :: CoreSyn.CoreExpr -> Bool
198 has_free_vars = not . VarSet.isEmptyVarSet . CoreFVs.exprFreeVars
200 -- Does the given expression use any of the given binders?
201 expr_uses_binders :: [CoreSyn.CoreBndr] -> CoreSyn.CoreExpr -> Bool
202 expr_uses_binders bndrs = not . VarSet.isEmptyVarSet . (CoreFVs.exprSomeFreeVars (`elem` bndrs))
204 -- Turns a Var CoreExpr into the Id inside it. Will of course only work for
205 -- simple Var CoreExprs, not complexer ones.
206 exprToVar :: CoreSyn.CoreExpr -> Var.Id
207 exprToVar (CoreSyn.Var id) = id
208 exprToVar expr = error $ "\nCoreTools.exprToVar: Not a var: " ++ show expr
210 -- Turns a Lit CoreExpr into the Literal inside it.
211 exprToLit :: CoreSyn.CoreExpr -> Literal.Literal
212 exprToLit (CoreSyn.Lit lit) = lit
213 exprToLit expr = error $ "\nCoreTools.exprToLit: Not a lit: " ++ show expr
215 -- Removes all the type and dictionary arguments from the given argument list,
216 -- leaving only the normal value arguments. The type given is the type of the
217 -- expression applied to this argument list.
218 get_val_args :: Type.Type -> [CoreSyn.CoreExpr] -> [CoreSyn.CoreExpr]
219 get_val_args ty args = drop n args
221 (tyvars, predtypes, _) = TcType.tcSplitSigmaTy ty
222 -- The first (length tyvars) arguments should be types, the next
223 -- (length predtypes) arguments should be dictionaries. We drop this many
224 -- arguments, to get at the value arguments.
225 n = length tyvars + length predtypes
227 getLiterals :: CoreSyn.CoreExpr -> [CoreSyn.CoreExpr]
228 getLiterals app@(CoreSyn.App _ _) = literals
230 (CoreSyn.Var f, args) = CoreSyn.collectArgs app
231 literals = filter (is_lit) args
233 getLiterals lit@(CoreSyn.Lit _) = [lit]
235 reduceCoreListToHsList ::
236 [HscTypes.CoreModule] -- ^ The modules where parts of the list are hidden
237 -> CoreSyn.CoreExpr -- ^ The refence to atleast one of the nodes
238 -> TranslatorSession [CoreSyn.CoreExpr]
239 reduceCoreListToHsList cores app@(CoreSyn.App _ _) = do {
240 ; let { (fun, args) = CoreSyn.collectArgs app
245 ; let topelem = args!!1
247 (varz@(CoreSyn.Var id)) -> do {
248 ; binds <- mapM (findExpr (isVarName id)) cores
249 ; otherelems <- reduceCoreListToHsList cores (head (Maybe.catMaybes binds))
250 ; return (topelem:otherelems)
252 (appz@(CoreSyn.App _ _)) -> do {
253 ; otherelems <- reduceCoreListToHsList cores appz
254 ; return (topelem:otherelems)
256 otherwise -> return [topelem]
258 otherwise -> return []
261 isVarName :: Monad m => Var.Var -> Var.Var -> m Bool
262 isVarName lookfor bind = return $ (Var.varName lookfor) == (Var.varName bind)
264 reduceCoreListToHsList _ _ = return []
266 -- Is the given var the State data constructor?
267 isStateCon :: Var.Var -> Bool
269 -- See if it is a DataConWrapId (not DataConWorkId, since State is a
271 case Id.idDetails var of
272 IdInfo.DataConWrapId dc ->
273 -- See if the datacon is the State datacon from the State type.
274 let tycon = DataCon.dataConTyCon dc
275 tyname = Name.getOccString tycon
276 dcname = Name.getOccString dc
277 in case (tyname, dcname) of
278 ("State", "State") -> True
282 -- | Is the given type a State type?
283 isStateType :: Type.Type -> Bool
284 -- Resolve any type synonyms remaining
285 isStateType ty | Just ty' <- Type.tcView ty = isStateType ty'
286 isStateType ty = Maybe.isJust $ do
287 -- Split the type. Don't use normal splitAppTy, since that looks through
288 -- newtypes, and we want to see the State newtype.
289 (typef, _) <- Type.repSplitAppTy_maybe ty
290 -- See if the applied type is a type constructor
291 (tycon, _) <- Type.splitTyConApp_maybe typef
292 if TyCon.isNewTyCon tycon && Name.getOccString tycon == "State"
298 -- | Does the given TypedThing have a State type?
299 hasStateType :: (TypedThing t) => t -> Bool
300 hasStateType expr = case getType expr of
302 Just ty -> isStateType ty
305 -- | Flattens nested lets into a single list of bindings. The expression
306 -- passed does not have to be a let expression, if it isn't an empty list of
307 -- bindings is returned.
309 CoreSyn.CoreExpr -- ^ The expression to flatten.
310 -> ([Binding], CoreSyn.CoreExpr) -- ^ The bindings and resulting expression.
311 flattenLets (CoreSyn.Let binds expr) =
312 (bindings ++ bindings', expr')
314 -- Recursively flatten the contained expression
315 (bindings', expr') =flattenLets expr
316 -- Flatten our own bindings to remove the Rec / NonRec constructors
317 bindings = CoreSyn.flattenBinds [binds]
318 flattenLets expr = ([], expr)
320 -- | Create bunch of nested non-recursive let expressions from the given
321 -- bindings. The first binding is bound at the highest level (and thus
322 -- available in all other bindings).
323 mkNonRecLets :: [Binding] -> CoreSyn.CoreExpr -> CoreSyn.CoreExpr
324 mkNonRecLets bindings expr = MkCore.mkCoreLets binds expr
326 binds = map (uncurry CoreSyn.NonRec) bindings
328 -- | A class of things that (optionally) have a core Type. The type is
329 -- optional, since Type expressions don't have a type themselves.
330 class TypedThing t where
331 getType :: t -> Maybe Type.Type
333 instance TypedThing CoreSyn.CoreExpr where
334 getType (CoreSyn.Type _) = Nothing
335 getType expr = Just $ CoreUtils.exprType expr
337 instance TypedThing CoreSyn.CoreBndr where
338 getType = return . Id.idType
340 instance TypedThing Type.Type where
341 getType = return . id
343 -- | Generate new uniques for all binders in the given expression.
344 -- Does not support making type variables unique, though this could be
345 -- supported if required (by passing a CoreSubst.Subst instead of VarEnv to
346 -- genUniques' below).
347 genUniques :: CoreSyn.CoreExpr -> TranslatorSession CoreSyn.CoreExpr
348 genUniques = genUniques' VarEnv.emptyVarEnv
350 -- | A helper function to generate uniques, that takes a VarEnv containing the
351 -- substitutions already performed.
352 genUniques' :: VarEnv.VarEnv CoreSyn.CoreBndr -> CoreSyn.CoreExpr -> TranslatorSession CoreSyn.CoreExpr
353 genUniques' subst (CoreSyn.Var f) = do
354 -- Replace the binder with its new value, if applicable.
355 let f' = VarEnv.lookupWithDefaultVarEnv subst f f
356 return (CoreSyn.Var f')
357 -- Leave literals untouched
358 genUniques' subst (CoreSyn.Lit l) = return $ CoreSyn.Lit l
359 genUniques' subst (CoreSyn.App f arg) = do
360 -- Only work on subexpressions
361 f' <- genUniques' subst f
362 arg' <- genUniques' subst arg
363 return (CoreSyn.App f' arg')
364 -- Don't change type abstractions
365 genUniques' subst expr@(CoreSyn.Lam bndr res) | CoreSyn.isTyVar bndr = return expr
366 genUniques' subst (CoreSyn.Lam bndr res) = do
367 -- Generate a new unique for the bound variable
368 (subst', bndr') <- genUnique subst bndr
369 res' <- genUniques' subst' res
370 return (CoreSyn.Lam bndr' res')
371 genUniques' subst (CoreSyn.Let (CoreSyn.NonRec bndr bound) res) = do
372 -- Make the binders unique
373 (subst', bndr') <- genUnique subst bndr
374 bound' <- genUniques' subst' bound
375 res' <- genUniques' subst' res
376 return $ CoreSyn.Let (CoreSyn.NonRec bndr' bound') res'
377 genUniques' subst (CoreSyn.Let (CoreSyn.Rec binds) res) = do
378 -- Make each of the binders unique
379 (subst', bndrs') <- mapAccumLM genUnique subst (map fst binds)
380 bounds' <- mapM (genUniques' subst') (map snd binds)
381 res' <- genUniques' subst' res
382 let binds' = zip bndrs' bounds'
383 return $ CoreSyn.Let (CoreSyn.Rec binds') res'
384 genUniques' subst (CoreSyn.Case scrut bndr ty alts) = do
385 -- Process the scrutinee with the original substitution, since non of the
386 -- binders bound in the Case statement is in scope in the scrutinee.
387 scrut' <- genUniques' subst scrut
388 -- Generate a new binder for the scrutinee
389 (subst', bndr') <- genUnique subst bndr
390 -- Process each of the alts
391 alts' <- mapM (doalt subst') alts
392 return $ CoreSyn.Case scrut' bndr' ty alts'
394 doalt subst (con, bndrs, expr) = do
395 (subst', bndrs') <- mapAccumLM genUnique subst bndrs
396 expr' <- genUniques' subst' expr
397 -- Note that we don't return subst', since bndrs are only in scope in
399 return (con, bndrs', expr')
400 genUniques' subst (CoreSyn.Cast expr coercion) = do
401 expr' <- genUniques' subst expr
402 -- Just process the casted expression
403 return $ CoreSyn.Cast expr' coercion
404 genUniques' subst (CoreSyn.Note note expr) = do
405 expr' <- genUniques' subst expr
406 -- Just process the annotated expression
407 return $ CoreSyn.Note note expr'
408 -- Leave types untouched
409 genUniques' subst expr@(CoreSyn.Type _) = return expr
411 -- Generate a new unique for the given binder, and extend the given
412 -- substitution to reflect this.
413 genUnique :: VarEnv.VarEnv CoreSyn.CoreBndr -> CoreSyn.CoreBndr -> TranslatorSession (VarEnv.VarEnv CoreSyn.CoreBndr, CoreSyn.CoreBndr)
414 genUnique subst bndr = do
415 bndr' <- BinderTools.cloneVar bndr
416 -- Replace all occurences of the old binder with a reference to the new
418 let subst' = VarEnv.extendVarEnv subst bndr bndr'
419 return (subst', bndr')