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 qualified System.IO.Unsafe
11 import qualified Data.Map as Map
16 import qualified TcType
17 import qualified HsExpr
18 import qualified HsTypes
19 import qualified HscTypes
22 import qualified TyCon
23 import qualified DataCon
24 import qualified TysWiredIn
25 import qualified DynFlags
26 import qualified SrcLoc
27 import qualified CoreSyn
29 import qualified IdInfo
30 import qualified VarSet
31 import qualified CoreUtils
32 import qualified CoreFVs
33 import qualified Literal
34 import qualified MkCore
35 import qualified VarEnv
38 import CLasH.Translator.TranslatorTypes
39 import CLasH.Utils.GhcTools
40 import CLasH.Utils.Core.BinderTools
41 import CLasH.Utils.HsTools
42 import CLasH.Utils.Pretty
44 import qualified CLasH.Utils.Core.BinderTools as BinderTools
46 -- | A single binding, used as a shortcut to simplify type signatures.
47 type Binding = (CoreSyn.CoreBndr, CoreSyn.CoreExpr)
49 -- | Evaluate a core Type representing type level int from the tfp
50 -- library to a real int. Checks if the type really is a Dec type and
51 -- caches the results.
52 tfp_to_int :: Type.Type -> TypeSession Int
54 hscenv <- MonadState.get tsHscEnv
55 let norm_ty = normalise_tfp_int hscenv ty
56 case Type.splitTyConApp_maybe norm_ty of
57 Just (tycon, args) -> do
58 let name = Name.getOccString (TyCon.tyConName tycon)
63 return $ error ("Callin tfp_to_int on non-dec:" ++ (show ty))
64 Nothing -> return $ error ("Callin tfp_to_int on non-dec:" ++ (show ty))
66 -- | Evaluate a core Type representing type level int from the tfp
67 -- library to a real int. Caches the results. Do not use directly, use
68 -- tfp_to_int instead.
69 tfp_to_int' :: Type.Type -> TypeSession Int
71 lens <- MonadState.get tsTfpInts
72 hscenv <- MonadState.get tsHscEnv
73 let norm_ty = normalise_tfp_int hscenv ty
74 let existing_len = Map.lookup (OrdType norm_ty) lens
76 Just len -> return len
78 let new_len = eval_tfp_int hscenv ty
79 MonadState.modify tsTfpInts (Map.insert (OrdType norm_ty) (new_len))
82 -- | Evaluate a core Type representing type level int from the tfp
83 -- library to a real int. Do not use directly, use tfp_to_int instead.
84 eval_tfp_int :: HscTypes.HscEnv -> Type.Type -> Int
86 unsafeRunGhc libdir $ do
88 -- Automatically import modules for any fully qualified identifiers
89 setDynFlag DynFlags.Opt_ImplicitImportQualified
91 let from_int_t_name = mkRdrName "Types.Data.Num.Ops" "fromIntegerT"
92 let from_int_t = SrcLoc.noLoc $ HsExpr.HsVar from_int_t_name
93 let undef = hsTypedUndef $ coreToHsType ty
94 let app = SrcLoc.noLoc $ HsExpr.HsApp (from_int_t) (undef)
95 let int_ty = SrcLoc.noLoc $ HsTypes.HsTyVar TysWiredIn.intTyCon_RDR
96 let expr = HsExpr.ExprWithTySig app int_ty
100 libdir = DynFlags.topDir dynflags
101 dynflags = HscTypes.hsc_dflags env
103 normalise_tfp_int :: HscTypes.HscEnv -> Type.Type -> Type.Type
104 normalise_tfp_int env ty =
105 System.IO.Unsafe.unsafePerformIO $
108 -- | Get the width of a SizedWord type
109 -- sized_word_len :: HscTypes.HscEnv -> Type.Type -> Int
110 -- sized_word_len env ty = eval_tfp_int env (sized_word_len_ty ty)
112 sized_word_len_ty :: Type.Type -> Type.Type
113 sized_word_len_ty ty = len
115 args = case Type.splitTyConApp_maybe ty of
116 Just (tycon, args) -> args
117 Nothing -> error $ "\nCoreTools.sized_word_len_ty: Not a sized word type: " ++ (pprString ty)
120 -- | Get the width of a SizedInt type
121 -- sized_int_len :: HscTypes.HscEnv -> Type.Type -> Int
122 -- sized_int_len env ty = eval_tfp_int env (sized_int_len_ty ty)
124 sized_int_len_ty :: Type.Type -> Type.Type
125 sized_int_len_ty ty = len
127 args = case Type.splitTyConApp_maybe ty of
128 Just (tycon, args) -> args
129 Nothing -> error $ "\nCoreTools.sized_int_len_ty: Not a sized int type: " ++ (pprString ty)
132 -- | Get the upperbound of a RangedWord type
133 -- ranged_word_bound :: HscTypes.HscEnv -> Type.Type -> Int
134 -- ranged_word_bound env ty = eval_tfp_int env (ranged_word_bound_ty ty)
136 ranged_word_bound_ty :: Type.Type -> Type.Type
137 ranged_word_bound_ty ty = len
139 args = case Type.splitTyConApp_maybe ty of
140 Just (tycon, args) -> args
141 Nothing -> error $ "\nCoreTools.ranged_word_bound_ty: Not a sized word type: " ++ (pprString ty)
144 -- | Evaluate a core Type representing type level int from the TypeLevel
145 -- library to a real int.
146 -- eval_type_level_int :: Type.Type -> Int
147 -- eval_type_level_int ty =
149 -- -- Automatically import modules for any fully qualified identifiers
150 -- setDynFlag DynFlags.Opt_ImplicitImportQualified
152 -- let to_int_name = mkRdrName "Data.TypeLevel.Num.Sets" "toInt"
153 -- let to_int = SrcLoc.noLoc $ HsExpr.HsVar to_int_name
154 -- let undef = hsTypedUndef $ coreToHsType ty
155 -- let app = HsExpr.HsApp (to_int) (undef)
157 -- core <- toCore [] app
160 -- | Get the length of a FSVec type
161 -- tfvec_len :: HscTypes.HscEnv -> Type.Type -> Int
162 -- tfvec_len env ty = eval_tfp_int env (tfvec_len_ty ty)
164 tfvec_len_ty :: Type.Type -> Type.Type
165 tfvec_len_ty ty = len
167 args = case Type.splitTyConApp_maybe ty of
168 Just (tycon, args) -> args
169 Nothing -> error $ "\nCoreTools.tfvec_len_ty: Not a vector type: " ++ (pprString ty)
172 -- | Get the element type of a TFVec type
173 tfvec_elem :: Type.Type -> Type.Type
174 tfvec_elem ty = el_ty
176 args = case Type.splitTyConApp_maybe ty of
177 Just (tycon, args) -> args
178 Nothing -> error $ "\nCoreTools.tfvec_len: Not a vector type: " ++ (pprString ty)
181 -- Is the given core expression a lambda abstraction?
182 is_lam :: CoreSyn.CoreExpr -> Bool
183 is_lam (CoreSyn.Lam _ _) = True
186 -- Is the given core expression a let expression?
187 is_let :: CoreSyn.CoreExpr -> Bool
188 is_let (CoreSyn.Let _ _) = True
191 -- Is the given core expression of a function type?
192 is_fun :: CoreSyn.CoreExpr -> Bool
193 -- Treat Type arguments differently, because exprType is not defined for them.
194 is_fun (CoreSyn.Type _) = False
195 is_fun expr = (Type.isFunTy . CoreUtils.exprType) expr
197 -- Is the given core expression polymorphic (i.e., does it accept type
199 is_poly :: CoreSyn.CoreExpr -> Bool
200 -- Treat Type arguments differently, because exprType is not defined for them.
201 is_poly (CoreSyn.Type _) = False
202 is_poly expr = (Maybe.isJust . Type.splitForAllTy_maybe . CoreUtils.exprType) expr
204 -- Is the given core expression a variable reference?
205 is_var :: CoreSyn.CoreExpr -> Bool
206 is_var (CoreSyn.Var _) = True
209 is_lit :: CoreSyn.CoreExpr -> Bool
210 is_lit (CoreSyn.Lit _) = True
213 -- Can the given core expression be applied to something? This is true for
214 -- applying to a value as well as a type.
215 is_applicable :: CoreSyn.CoreExpr -> Bool
216 is_applicable expr = is_fun expr || is_poly expr
218 -- Is the given core expression a variable or an application?
219 is_simple :: CoreSyn.CoreExpr -> Bool
220 is_simple (CoreSyn.App _ _) = True
221 is_simple (CoreSyn.Var _) = True
222 is_simple (CoreSyn.Cast expr _) = is_simple expr
225 -- Does the given CoreExpr have any free type vars?
226 has_free_tyvars :: CoreSyn.CoreExpr -> Bool
227 has_free_tyvars = not . VarSet.isEmptyVarSet . (CoreFVs.exprSomeFreeVars Var.isTyVar)
229 -- Does the given type have any free type vars?
230 ty_has_free_tyvars :: Type.Type -> Bool
231 ty_has_free_tyvars = not . VarSet.isEmptyVarSet . Type.tyVarsOfType
233 -- Does the given CoreExpr have any free local vars?
234 has_free_vars :: CoreSyn.CoreExpr -> Bool
235 has_free_vars = not . VarSet.isEmptyVarSet . CoreFVs.exprFreeVars
237 -- Does the given expression use any of the given binders?
238 expr_uses_binders :: [CoreSyn.CoreBndr] -> CoreSyn.CoreExpr -> Bool
239 expr_uses_binders bndrs = not . VarSet.isEmptyVarSet . (CoreFVs.exprSomeFreeVars (`elem` bndrs))
241 -- Turns a Var CoreExpr into the Id inside it. Will of course only work for
242 -- simple Var CoreExprs, not complexer ones.
243 exprToVar :: CoreSyn.CoreExpr -> Var.Id
244 exprToVar (CoreSyn.Var id) = id
245 exprToVar expr = error $ "\nCoreTools.exprToVar: Not a var: " ++ show expr
247 -- Turns a Lit CoreExpr into the Literal inside it.
248 exprToLit :: CoreSyn.CoreExpr -> Literal.Literal
249 exprToLit (CoreSyn.Lit lit) = lit
250 exprToLit expr = error $ "\nCoreTools.exprToLit: Not a lit: " ++ show expr
252 -- Removes all the type and dictionary arguments from the given argument list,
253 -- leaving only the normal value arguments. The type given is the type of the
254 -- expression applied to this argument list.
255 get_val_args :: Type.Type -> [CoreSyn.CoreExpr] -> [CoreSyn.CoreExpr]
256 get_val_args ty args = drop n args
258 (tyvars, predtypes, _) = TcType.tcSplitSigmaTy ty
259 -- The first (length tyvars) arguments should be types, the next
260 -- (length predtypes) arguments should be dictionaries. We drop this many
261 -- arguments, to get at the value arguments.
262 n = length tyvars + length predtypes
264 getLiterals :: HscTypes.HscEnv -> CoreSyn.CoreExpr -> [CoreSyn.CoreExpr]
265 getLiterals _ app@(CoreSyn.App _ _) = literals
267 (CoreSyn.Var f, args) = CoreSyn.collectArgs app
268 literals = filter (is_lit) args
270 getLiterals _ lit@(CoreSyn.Lit _) = [lit]
272 getLiterals hscenv letrec@(CoreSyn.Let (CoreSyn.NonRec letBind (letExpr)) letRes) = [lit]
274 ty = Var.varType letBind
275 litInt = eval_tfp_int hscenv ty
276 lit = CoreSyn.Lit (Literal.mkMachInt (toInteger litInt))
278 getLiterals _ expr = error $ "\nCoreTools.getLiterals: Not a known Lit: " ++ pprString expr
280 reduceCoreListToHsList ::
281 [HscTypes.CoreModule] -- ^ The modules where parts of the list are hidden
282 -> CoreSyn.CoreExpr -- ^ The refence to atleast one of the nodes
283 -> TranslatorSession [CoreSyn.CoreExpr]
284 reduceCoreListToHsList cores app@(CoreSyn.App _ _) = do {
285 ; let { (fun, args) = CoreSyn.collectArgs app
290 ; let topelem = args!!1
292 (varz@(CoreSyn.Var id)) -> do {
293 ; binds <- mapM (findExpr (isVarName id)) cores
294 ; otherelems <- reduceCoreListToHsList cores (head (Maybe.catMaybes binds))
295 ; return (topelem:otherelems)
297 (appz@(CoreSyn.App _ _)) -> do {
298 ; otherelems <- reduceCoreListToHsList cores appz
299 ; return (topelem:otherelems)
301 otherwise -> return [topelem]
303 otherwise -> return []
306 isVarName :: Monad m => Var.Var -> Var.Var -> m Bool
307 isVarName lookfor bind = return $ (Var.varName lookfor) == (Var.varName bind)
309 reduceCoreListToHsList _ _ = return []
311 -- Is the given var the State data constructor?
312 isStateCon :: Var.Var -> Bool
314 -- See if it is a DataConWrapId (not DataConWorkId, since State is a
316 case Id.idDetails var of
317 IdInfo.DataConWrapId dc ->
318 -- See if the datacon is the State datacon from the State type.
319 let tycon = DataCon.dataConTyCon dc
320 tyname = Name.getOccString tycon
321 dcname = Name.getOccString dc
322 in case (tyname, dcname) of
323 ("State", "State") -> True
327 -- | Is the given type a State type?
328 isStateType :: Type.Type -> Bool
329 -- Resolve any type synonyms remaining
330 isStateType ty | Just ty' <- Type.tcView ty = isStateType ty'
331 isStateType ty = Maybe.isJust $ do
332 -- Split the type. Don't use normal splitAppTy, since that looks through
333 -- newtypes, and we want to see the State newtype.
334 (typef, _) <- Type.repSplitAppTy_maybe ty
335 -- See if the applied type is a type constructor
336 (tycon, _) <- Type.splitTyConApp_maybe typef
337 if TyCon.isNewTyCon tycon && Name.getOccString tycon == "State"
343 -- | Does the given TypedThing have a State type?
344 hasStateType :: (TypedThing t) => t -> Bool
345 hasStateType expr = case getType expr of
347 Just ty -> isStateType ty
350 -- | Flattens nested lets into a single list of bindings. The expression
351 -- passed does not have to be a let expression, if it isn't an empty list of
352 -- bindings is returned.
354 CoreSyn.CoreExpr -- ^ The expression to flatten.
355 -> ([Binding], CoreSyn.CoreExpr) -- ^ The bindings and resulting expression.
356 flattenLets (CoreSyn.Let binds expr) =
357 (bindings ++ bindings', expr')
359 -- Recursively flatten the contained expression
360 (bindings', expr') =flattenLets expr
361 -- Flatten our own bindings to remove the Rec / NonRec constructors
362 bindings = CoreSyn.flattenBinds [binds]
363 flattenLets expr = ([], expr)
365 -- | Create bunch of nested non-recursive let expressions from the given
366 -- bindings. The first binding is bound at the highest level (and thus
367 -- available in all other bindings).
368 mkNonRecLets :: [Binding] -> CoreSyn.CoreExpr -> CoreSyn.CoreExpr
369 mkNonRecLets bindings expr = MkCore.mkCoreLets binds expr
371 binds = map (uncurry CoreSyn.NonRec) bindings
373 -- | A class of things that (optionally) have a core Type. The type is
374 -- optional, since Type expressions don't have a type themselves.
375 class TypedThing t where
376 getType :: t -> Maybe Type.Type
378 instance TypedThing CoreSyn.CoreExpr where
379 getType (CoreSyn.Type _) = Nothing
380 getType expr = Just $ CoreUtils.exprType expr
382 instance TypedThing CoreSyn.CoreBndr where
383 getType = return . Id.idType
385 instance TypedThing Type.Type where
386 getType = return . id
388 -- | Generate new uniques for all binders in the given expression.
389 -- Does not support making type variables unique, though this could be
390 -- supported if required (by passing a CoreSubst.Subst instead of VarEnv to
391 -- genUniques' below).
392 genUniques :: CoreSyn.CoreExpr -> TranslatorSession CoreSyn.CoreExpr
393 genUniques = genUniques' VarEnv.emptyVarEnv
395 -- | A helper function to generate uniques, that takes a VarEnv containing the
396 -- substitutions already performed.
397 genUniques' :: VarEnv.VarEnv CoreSyn.CoreBndr -> CoreSyn.CoreExpr -> TranslatorSession CoreSyn.CoreExpr
398 genUniques' subst (CoreSyn.Var f) = do
399 -- Replace the binder with its new value, if applicable.
400 let f' = VarEnv.lookupWithDefaultVarEnv subst f f
401 return (CoreSyn.Var f')
402 -- Leave literals untouched
403 genUniques' subst (CoreSyn.Lit l) = return $ CoreSyn.Lit l
404 genUniques' subst (CoreSyn.App f arg) = do
405 -- Only work on subexpressions
406 f' <- genUniques' subst f
407 arg' <- genUniques' subst arg
408 return (CoreSyn.App f' arg')
409 -- Don't change type abstractions
410 genUniques' subst expr@(CoreSyn.Lam bndr res) | CoreSyn.isTyVar bndr = return expr
411 genUniques' subst (CoreSyn.Lam bndr res) = do
412 -- Generate a new unique for the bound variable
413 (subst', bndr') <- genUnique subst bndr
414 res' <- genUniques' subst' res
415 return (CoreSyn.Lam bndr' res')
416 genUniques' subst (CoreSyn.Let (CoreSyn.NonRec bndr bound) res) = do
417 -- Make the binders unique
418 (subst', bndr') <- genUnique subst bndr
419 bound' <- genUniques' subst' bound
420 res' <- genUniques' subst' res
421 return $ CoreSyn.Let (CoreSyn.NonRec bndr' bound') res'
422 genUniques' subst (CoreSyn.Let (CoreSyn.Rec binds) res) = do
423 -- Make each of the binders unique
424 (subst', bndrs') <- mapAccumLM genUnique subst (map fst binds)
425 bounds' <- mapM (genUniques' subst' . snd) binds
426 res' <- genUniques' subst' res
427 let binds' = zip bndrs' bounds'
428 return $ CoreSyn.Let (CoreSyn.Rec binds') res'
429 genUniques' subst (CoreSyn.Case scrut bndr ty alts) = do
430 -- Process the scrutinee with the original substitution, since non of the
431 -- binders bound in the Case statement is in scope in the scrutinee.
432 scrut' <- genUniques' subst scrut
433 -- Generate a new binder for the scrutinee
434 (subst', bndr') <- genUnique subst bndr
435 -- Process each of the alts
436 alts' <- mapM (doalt subst') alts
437 return $ CoreSyn.Case scrut' bndr' ty alts'
439 doalt subst (con, bndrs, expr) = do
440 (subst', bndrs') <- mapAccumLM genUnique subst bndrs
441 expr' <- genUniques' subst' expr
442 -- Note that we don't return subst', since bndrs are only in scope in
444 return (con, bndrs', expr')
445 genUniques' subst (CoreSyn.Cast expr coercion) = do
446 expr' <- genUniques' subst expr
447 -- Just process the casted expression
448 return $ CoreSyn.Cast expr' coercion
449 genUniques' subst (CoreSyn.Note note expr) = do
450 expr' <- genUniques' subst expr
451 -- Just process the annotated expression
452 return $ CoreSyn.Note note expr'
453 -- Leave types untouched
454 genUniques' subst expr@(CoreSyn.Type _) = return expr
456 -- Generate a new unique for the given binder, and extend the given
457 -- substitution to reflect this.
458 genUnique :: VarEnv.VarEnv CoreSyn.CoreBndr -> CoreSyn.CoreBndr -> TranslatorSession (VarEnv.VarEnv CoreSyn.CoreBndr, CoreSyn.CoreBndr)
459 genUnique subst bndr = do
460 bndr' <- BinderTools.cloneVar bndr
461 -- Replace all occurences of the old binder with a reference to the new
463 let subst' = VarEnv.extendVarEnv subst bndr bndr'
464 return (subst', bndr')
466 -- Create a "selector" case that selects the ith field from a datacon
467 mkSelCase :: CoreSyn.CoreExpr -> Int -> TranslatorSession CoreSyn.CoreExpr
468 mkSelCase scrut i = do
469 let scrut_ty = CoreUtils.exprType scrut
470 case Type.splitTyConApp_maybe scrut_ty of
471 -- The scrutinee should have a type constructor. We keep the type
472 -- arguments around so we can instantiate the field types below
473 Just (tycon, tyargs) -> case TyCon.tyConDataCons tycon of
474 -- The scrutinee type should have a single dataconstructor,
475 -- otherwise we can't construct a valid selector case.
477 let field_tys = DataCon.dataConInstOrigArgTys datacon tyargs
478 -- Create a list of wild binders for the fields we don't want
479 let wildbndrs = map MkCore.mkWildBinder field_tys
480 -- Create a single binder for the field we want
481 sel_bndr <- mkInternalVar "sel" (field_tys!!i)
482 -- Create a wild binder for the scrutinee
483 let scrut_bndr = MkCore.mkWildBinder scrut_ty
484 -- Create the case expression
485 let binders = take i wildbndrs ++ [sel_bndr] ++ drop (i+1) wildbndrs
486 return $ CoreSyn.Case scrut scrut_bndr scrut_ty [(CoreSyn.DataAlt datacon, binders, CoreSyn.Var sel_bndr)]
487 dcs -> error $ "CoreTools.mkSelCase: Scrutinee type must have exactly one datacon. Extracting element " ++ (show i) ++ " from '" ++ pprString scrut ++ "' Datacons: " ++ (show dcs) ++ " Type: " ++ (pprString scrut_ty)
488 Nothing -> error $ "CoreTools.mkSelCase: Creating extractor case, but scrutinee has no tycon? Extracting element " ++ (show i) ++ " from '" ++ pprString scrut ++ "'"