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
12 import qualified Data.Accessor.Monad.Trans.State as MonadState
17 import qualified TcType
18 import qualified HsExpr
19 import qualified HsTypes
20 import qualified HscTypes
23 import qualified TyCon
24 import qualified DataCon
25 import qualified TysWiredIn
26 import qualified DynFlags
27 import qualified SrcLoc
28 import qualified CoreSyn
30 import qualified IdInfo
31 import qualified VarSet
32 import qualified CoreUtils
33 import qualified CoreFVs
34 import qualified Literal
35 import qualified MkCore
36 import qualified VarEnv
39 import CLasH.Translator.TranslatorTypes
40 import CLasH.Utils.GhcTools
41 import CLasH.Utils.Core.BinderTools
42 import CLasH.Utils.HsTools
43 import CLasH.Utils.Pretty
45 import qualified CLasH.Utils.Core.BinderTools as BinderTools
47 -- | A single binding, used as a shortcut to simplify type signatures.
48 type Binding = (CoreSyn.CoreBndr, CoreSyn.CoreExpr)
50 -- | Evaluate a core Type representing type level int from the tfp
51 -- library to a real int. Checks if the type really is a Dec type and
52 -- caches the results.
53 tfp_to_int :: Type.Type -> TypeSession Int
55 hscenv <- MonadState.get tsHscEnv
56 let norm_ty = normalise_tfp_int hscenv ty
57 case Type.splitTyConApp_maybe norm_ty of
58 Just (tycon, args) -> do
59 let name = Name.getOccString (TyCon.tyConName tycon)
64 return $ error ("Callin tfp_to_int on non-dec:" ++ (show ty))
65 Nothing -> return $ error ("Callin tfp_to_int on non-dec:" ++ (show ty))
67 -- | Evaluate a core Type representing type level int from the tfp
68 -- library to a real int. Caches the results. Do not use directly, use
69 -- tfp_to_int instead.
70 tfp_to_int' :: Type.Type -> TypeSession Int
72 lens <- MonadState.get tsTfpInts
73 hscenv <- MonadState.get tsHscEnv
74 let norm_ty = normalise_tfp_int hscenv ty
75 let existing_len = Map.lookup (OrdType norm_ty) lens
77 Just len -> return len
79 let new_len = eval_tfp_int hscenv ty
80 MonadState.modify tsTfpInts (Map.insert (OrdType norm_ty) (new_len))
83 -- | Evaluate a core Type representing type level int from the tfp
84 -- library to a real int. Do not use directly, use tfp_to_int instead.
85 eval_tfp_int :: HscTypes.HscEnv -> Type.Type -> Int
87 unsafeRunGhc libdir $ do
89 -- Automatically import modules for any fully qualified identifiers
90 setDynFlag DynFlags.Opt_ImplicitImportQualified
92 let from_int_t_name = mkRdrName "Types.Data.Num.Ops" "fromIntegerT"
93 let from_int_t = SrcLoc.noLoc $ HsExpr.HsVar from_int_t_name
94 let undef = hsTypedUndef $ coreToHsType ty
95 let app = SrcLoc.noLoc $ HsExpr.HsApp (from_int_t) (undef)
96 let int_ty = SrcLoc.noLoc $ HsTypes.HsTyVar TysWiredIn.intTyCon_RDR
97 let expr = HsExpr.ExprWithTySig app int_ty
101 libdir = DynFlags.topDir dynflags
102 dynflags = HscTypes.hsc_dflags env
104 normalise_tfp_int :: HscTypes.HscEnv -> Type.Type -> Type.Type
105 normalise_tfp_int env ty =
106 System.IO.Unsafe.unsafePerformIO $
109 -- | Get the width of a SizedWord type
110 -- sized_word_len :: HscTypes.HscEnv -> Type.Type -> Int
111 -- sized_word_len env ty = eval_tfp_int env (sized_word_len_ty ty)
113 sized_word_len_ty :: Type.Type -> Type.Type
114 sized_word_len_ty ty = len
116 args = case Type.splitTyConApp_maybe ty of
117 Just (tycon, args) -> args
118 Nothing -> error $ "\nCoreTools.sized_word_len_ty: Not a sized word type: " ++ (pprString ty)
121 -- | Get the width of a SizedInt type
122 -- sized_int_len :: HscTypes.HscEnv -> Type.Type -> Int
123 -- sized_int_len env ty = eval_tfp_int env (sized_int_len_ty ty)
125 sized_int_len_ty :: Type.Type -> Type.Type
126 sized_int_len_ty ty = len
128 args = case Type.splitTyConApp_maybe ty of
129 Just (tycon, args) -> args
130 Nothing -> error $ "\nCoreTools.sized_int_len_ty: Not a sized int type: " ++ (pprString ty)
133 -- | Get the upperbound of a RangedWord type
134 -- ranged_word_bound :: HscTypes.HscEnv -> Type.Type -> Int
135 -- ranged_word_bound env ty = eval_tfp_int env (ranged_word_bound_ty ty)
137 ranged_word_bound_ty :: Type.Type -> Type.Type
138 ranged_word_bound_ty ty = len
140 args = case Type.splitTyConApp_maybe ty of
141 Just (tycon, args) -> args
142 Nothing -> error $ "\nCoreTools.ranged_word_bound_ty: Not a sized word type: " ++ (pprString ty)
145 -- | Evaluate a core Type representing type level int from the TypeLevel
146 -- library to a real int.
147 -- eval_type_level_int :: Type.Type -> Int
148 -- eval_type_level_int ty =
150 -- -- Automatically import modules for any fully qualified identifiers
151 -- setDynFlag DynFlags.Opt_ImplicitImportQualified
153 -- let to_int_name = mkRdrName "Data.TypeLevel.Num.Sets" "toInt"
154 -- let to_int = SrcLoc.noLoc $ HsExpr.HsVar to_int_name
155 -- let undef = hsTypedUndef $ coreToHsType ty
156 -- let app = HsExpr.HsApp (to_int) (undef)
158 -- core <- toCore [] app
161 -- | Get the length of a FSVec type
162 -- tfvec_len :: HscTypes.HscEnv -> Type.Type -> Int
163 -- tfvec_len env ty = eval_tfp_int env (tfvec_len_ty ty)
165 tfvec_len_ty :: Type.Type -> Type.Type
166 tfvec_len_ty ty = len
168 args = case Type.splitTyConApp_maybe ty of
169 Just (tycon, args) -> args
170 Nothing -> error $ "\nCoreTools.tfvec_len_ty: Not a vector type: " ++ (pprString ty)
173 -- | Get the element type of a TFVec type
174 tfvec_elem :: Type.Type -> Type.Type
175 tfvec_elem ty = el_ty
177 args = case Type.splitTyConApp_maybe ty of
178 Just (tycon, args) -> args
179 Nothing -> error $ "\nCoreTools.tfvec_len: Not a vector type: " ++ (pprString ty)
182 -- Is the given core expression a lambda abstraction?
183 is_lam :: CoreSyn.CoreExpr -> Bool
184 is_lam (CoreSyn.Lam _ _) = True
187 -- Is the given core expression a let expression?
188 is_let :: CoreSyn.CoreExpr -> Bool
189 is_let (CoreSyn.Let _ _) = True
192 -- Is the given core expression of a function type?
193 is_fun :: CoreSyn.CoreExpr -> Bool
194 -- Treat Type arguments differently, because exprType is not defined for them.
195 is_fun (CoreSyn.Type _) = False
196 is_fun expr = (Type.isFunTy . CoreUtils.exprType) expr
198 -- Is the given core expression polymorphic (i.e., does it accept type
200 is_poly :: CoreSyn.CoreExpr -> Bool
201 -- Treat Type arguments differently, because exprType is not defined for them.
202 is_poly (CoreSyn.Type _) = False
203 is_poly expr = (Maybe.isJust . Type.splitForAllTy_maybe . CoreUtils.exprType) expr
205 -- Is the given core expression a variable reference?
206 is_var :: CoreSyn.CoreExpr -> Bool
207 is_var (CoreSyn.Var _) = True
210 is_lit :: CoreSyn.CoreExpr -> Bool
211 is_lit (CoreSyn.Lit _) = True
214 -- Can the given core expression be applied to something? This is true for
215 -- applying to a value as well as a type.
216 is_applicable :: CoreSyn.CoreExpr -> Bool
217 is_applicable expr = is_fun expr || is_poly expr
219 -- Is the given core expression a variable or an application?
220 is_simple :: CoreSyn.CoreExpr -> Bool
221 is_simple (CoreSyn.App _ _) = True
222 is_simple (CoreSyn.Var _) = True
223 is_simple (CoreSyn.Cast expr _) = is_simple expr
226 -- Does the given CoreExpr have any free type vars?
227 has_free_tyvars :: CoreSyn.CoreExpr -> Bool
228 has_free_tyvars = not . VarSet.isEmptyVarSet . (CoreFVs.exprSomeFreeVars Var.isTyVar)
230 -- Does the given type have any free type vars?
231 ty_has_free_tyvars :: Type.Type -> Bool
232 ty_has_free_tyvars = not . VarSet.isEmptyVarSet . Type.tyVarsOfType
234 -- Does the given CoreExpr have any free local vars?
235 has_free_vars :: CoreSyn.CoreExpr -> Bool
236 has_free_vars = not . VarSet.isEmptyVarSet . CoreFVs.exprFreeVars
238 -- Does the given expression use any of the given binders?
239 expr_uses_binders :: [CoreSyn.CoreBndr] -> CoreSyn.CoreExpr -> Bool
240 expr_uses_binders bndrs = not . VarSet.isEmptyVarSet . (CoreFVs.exprSomeFreeVars (`elem` bndrs))
242 -- Turns a Var CoreExpr into the Id inside it. Will of course only work for
243 -- simple Var CoreExprs, not complexer ones.
244 exprToVar :: CoreSyn.CoreExpr -> Var.Id
245 exprToVar (CoreSyn.Var id) = id
246 exprToVar expr = error $ "\nCoreTools.exprToVar: Not a var: " ++ show expr
248 -- Turns a Lit CoreExpr into the Literal inside it.
249 exprToLit :: CoreSyn.CoreExpr -> Literal.Literal
250 exprToLit (CoreSyn.Lit lit) = lit
251 exprToLit expr = error $ "\nCoreTools.exprToLit: Not a lit: " ++ show expr
253 -- Removes all the type and dictionary arguments from the given argument list,
254 -- leaving only the normal value arguments. The type given is the type of the
255 -- expression applied to this argument list.
256 get_val_args :: Type.Type -> [CoreSyn.CoreExpr] -> [CoreSyn.CoreExpr]
257 get_val_args ty args = drop n args
259 (tyvars, predtypes, _) = TcType.tcSplitSigmaTy ty
260 -- The first (length tyvars) arguments should be types, the next
261 -- (length predtypes) arguments should be dictionaries. We drop this many
262 -- arguments, to get at the value arguments.
263 n = length tyvars + length predtypes
265 -- Finds out what literal Integer this expression represents.
266 getIntegerLiteral :: CoreSyn.CoreExpr -> TranslatorSession Integer
267 getIntegerLiteral expr =
268 case CoreSyn.collectArgs expr of
269 (CoreSyn.Var f, [CoreSyn.Lit (Literal.MachInt integer)])
270 | getFullString f == "GHC.Integer.smallInteger" -> return integer
271 (CoreSyn.Var f, [CoreSyn.Lit (Literal.MachInt64 integer)])
272 | getFullString f == "GHC.Integer.int64ToInteger" -> return integer
273 (CoreSyn.Var f, [CoreSyn.Lit (Literal.MachWord integer)])
274 | getFullString f == "GHC.Integer.wordToInteger" -> return integer
275 (CoreSyn.Var f, [CoreSyn.Lit (Literal.MachWord64 integer)])
276 | getFullString f == "GHC.Integer.word64ToInteger" -> return integer
277 -- fromIntegerT returns the integer corresponding to the type of its
278 -- (third) argument. Since it is polymorphic, the type of that
279 -- argument is passed as the first argument, so we can just use that
281 (CoreSyn.Var f, [CoreSyn.Type dec_ty, dec_dict, CoreSyn.Type num_ty, num_dict, arg])
282 | getFullString f == "Types.Data.Num.Ops.fromIntegerT" -> do
283 int <- MonadState.lift tsType $ tfp_to_int dec_ty
284 return $ toInteger int
285 _ -> error $ "CoreTools.getIntegerLiteral: Unsupported Integer literal: " ++ pprString expr
287 reduceCoreListToHsList ::
288 [HscTypes.CoreModule] -- ^ The modules where parts of the list are hidden
289 -> CoreSyn.CoreExpr -- ^ The refence to atleast one of the nodes
290 -> TranslatorSession [CoreSyn.CoreExpr]
291 reduceCoreListToHsList cores app@(CoreSyn.App _ _) = do {
292 ; let { (fun, args) = CoreSyn.collectArgs app
297 ; let topelem = args!!1
299 (varz@(CoreSyn.Var id)) -> do {
300 ; binds <- mapM (findExpr (isVarName id)) cores
301 ; otherelems <- reduceCoreListToHsList cores (head (Maybe.catMaybes binds))
302 ; return (topelem:otherelems)
304 (appz@(CoreSyn.App _ _)) -> do {
305 ; otherelems <- reduceCoreListToHsList cores appz
306 ; return (topelem:otherelems)
308 otherwise -> return [topelem]
310 otherwise -> return []
313 isVarName :: Monad m => Var.Var -> Var.Var -> m Bool
314 isVarName lookfor bind = return $ (Var.varName lookfor) == (Var.varName bind)
316 reduceCoreListToHsList _ _ = return []
318 -- Is the given var the State data constructor?
319 isStateCon :: Var.Var -> Bool
321 -- See if it is a DataConWrapId (not DataConWorkId, since State is a
323 case Id.idDetails var of
324 IdInfo.DataConWrapId dc ->
325 -- See if the datacon is the State datacon from the State type.
326 let tycon = DataCon.dataConTyCon dc
327 tyname = Name.getOccString tycon
328 dcname = Name.getOccString dc
329 in case (tyname, dcname) of
330 ("State", "State") -> True
334 -- | Is the given type a State type?
335 isStateType :: Type.Type -> Bool
336 -- Resolve any type synonyms remaining
337 isStateType ty | Just ty' <- Type.tcView ty = isStateType ty'
338 isStateType ty = Maybe.isJust $ do
339 -- Split the type. Don't use normal splitAppTy, since that looks through
340 -- newtypes, and we want to see the State newtype.
341 (typef, _) <- Type.repSplitAppTy_maybe ty
342 -- See if the applied type is a type constructor
343 (tycon, _) <- Type.splitTyConApp_maybe typef
344 if TyCon.isNewTyCon tycon && Name.getOccString tycon == "State"
350 -- | Does the given TypedThing have a State type?
351 hasStateType :: (TypedThing t) => t -> Bool
352 hasStateType expr = case getType expr of
354 Just ty -> isStateType ty
357 -- | Flattens nested lets into a single list of bindings. The expression
358 -- passed does not have to be a let expression, if it isn't an empty list of
359 -- bindings is returned.
361 CoreSyn.CoreExpr -- ^ The expression to flatten.
362 -> ([Binding], CoreSyn.CoreExpr) -- ^ The bindings and resulting expression.
363 flattenLets (CoreSyn.Let binds expr) =
364 (bindings ++ bindings', expr')
366 -- Recursively flatten the contained expression
367 (bindings', expr') =flattenLets expr
368 -- Flatten our own bindings to remove the Rec / NonRec constructors
369 bindings = CoreSyn.flattenBinds [binds]
370 flattenLets expr = ([], expr)
372 -- | Create bunch of nested non-recursive let expressions from the given
373 -- bindings. The first binding is bound at the highest level (and thus
374 -- available in all other bindings).
375 mkNonRecLets :: [Binding] -> CoreSyn.CoreExpr -> CoreSyn.CoreExpr
376 mkNonRecLets bindings expr = MkCore.mkCoreLets binds expr
378 binds = map (uncurry CoreSyn.NonRec) bindings
380 -- | A class of things that (optionally) have a core Type. The type is
381 -- optional, since Type expressions don't have a type themselves.
382 class TypedThing t where
383 getType :: t -> Maybe Type.Type
385 instance TypedThing CoreSyn.CoreExpr where
386 getType (CoreSyn.Type _) = Nothing
387 getType expr = Just $ CoreUtils.exprType expr
389 instance TypedThing CoreSyn.CoreBndr where
390 getType = return . Id.idType
392 instance TypedThing Type.Type where
393 getType = return . id
395 -- | Generate new uniques for all binders in the given expression.
396 -- Does not support making type variables unique, though this could be
397 -- supported if required (by passing a CoreSubst.Subst instead of VarEnv to
398 -- genUniques' below).
399 genUniques :: CoreSyn.CoreExpr -> TranslatorSession CoreSyn.CoreExpr
400 genUniques = genUniques' VarEnv.emptyVarEnv
402 -- | A helper function to generate uniques, that takes a VarEnv containing the
403 -- substitutions already performed.
404 genUniques' :: VarEnv.VarEnv CoreSyn.CoreBndr -> CoreSyn.CoreExpr -> TranslatorSession CoreSyn.CoreExpr
405 genUniques' subst (CoreSyn.Var f) = do
406 -- Replace the binder with its new value, if applicable.
407 let f' = VarEnv.lookupWithDefaultVarEnv subst f f
408 return (CoreSyn.Var f')
409 -- Leave literals untouched
410 genUniques' subst (CoreSyn.Lit l) = return $ CoreSyn.Lit l
411 genUniques' subst (CoreSyn.App f arg) = do
412 -- Only work on subexpressions
413 f' <- genUniques' subst f
414 arg' <- genUniques' subst arg
415 return (CoreSyn.App f' arg')
416 -- Don't change type abstractions
417 genUniques' subst expr@(CoreSyn.Lam bndr res) | CoreSyn.isTyVar bndr = return expr
418 genUniques' subst (CoreSyn.Lam bndr res) = do
419 -- Generate a new unique for the bound variable
420 (subst', bndr') <- genUnique subst bndr
421 res' <- genUniques' subst' res
422 return (CoreSyn.Lam bndr' res')
423 genUniques' subst (CoreSyn.Let (CoreSyn.NonRec bndr bound) res) = do
424 -- Make the binders unique
425 (subst', bndr') <- genUnique subst bndr
426 bound' <- genUniques' subst' bound
427 res' <- genUniques' subst' res
428 return $ CoreSyn.Let (CoreSyn.NonRec bndr' bound') res'
429 genUniques' subst (CoreSyn.Let (CoreSyn.Rec binds) res) = do
430 -- Make each of the binders unique
431 (subst', bndrs') <- mapAccumLM genUnique subst (map fst binds)
432 bounds' <- mapM (genUniques' subst' . snd) binds
433 res' <- genUniques' subst' res
434 let binds' = zip bndrs' bounds'
435 return $ CoreSyn.Let (CoreSyn.Rec binds') res'
436 genUniques' subst (CoreSyn.Case scrut bndr ty alts) = do
437 -- Process the scrutinee with the original substitution, since non of the
438 -- binders bound in the Case statement is in scope in the scrutinee.
439 scrut' <- genUniques' subst scrut
440 -- Generate a new binder for the scrutinee
441 (subst', bndr') <- genUnique subst bndr
442 -- Process each of the alts
443 alts' <- mapM (doalt subst') alts
444 return $ CoreSyn.Case scrut' bndr' ty alts'
446 doalt subst (con, bndrs, expr) = do
447 (subst', bndrs') <- mapAccumLM genUnique subst bndrs
448 expr' <- genUniques' subst' expr
449 -- Note that we don't return subst', since bndrs are only in scope in
451 return (con, bndrs', expr')
452 genUniques' subst (CoreSyn.Cast expr coercion) = do
453 expr' <- genUniques' subst expr
454 -- Just process the casted expression
455 return $ CoreSyn.Cast expr' coercion
456 genUniques' subst (CoreSyn.Note note expr) = do
457 expr' <- genUniques' subst expr
458 -- Just process the annotated expression
459 return $ CoreSyn.Note note expr'
460 -- Leave types untouched
461 genUniques' subst expr@(CoreSyn.Type _) = return expr
463 -- Generate a new unique for the given binder, and extend the given
464 -- substitution to reflect this.
465 genUnique :: VarEnv.VarEnv CoreSyn.CoreBndr -> CoreSyn.CoreBndr -> TranslatorSession (VarEnv.VarEnv CoreSyn.CoreBndr, CoreSyn.CoreBndr)
466 genUnique subst bndr = do
467 bndr' <- BinderTools.cloneVar bndr
468 -- Replace all occurences of the old binder with a reference to the new
470 let subst' = VarEnv.extendVarEnv subst bndr bndr'
471 return (subst', bndr')
473 -- Create a "selector" case that selects the ith field from a datacon
474 mkSelCase :: CoreSyn.CoreExpr -> Int -> TranslatorSession CoreSyn.CoreExpr
475 mkSelCase scrut i = do
476 let scrut_ty = CoreUtils.exprType scrut
477 case Type.splitTyConApp_maybe scrut_ty of
478 -- The scrutinee should have a type constructor. We keep the type
479 -- arguments around so we can instantiate the field types below
480 Just (tycon, tyargs) -> case TyCon.tyConDataCons tycon of
481 -- The scrutinee type should have a single dataconstructor,
482 -- otherwise we can't construct a valid selector case.
484 let field_tys = DataCon.dataConInstOrigArgTys datacon tyargs
485 -- Create a list of wild binders for the fields we don't want
486 let wildbndrs = map MkCore.mkWildBinder field_tys
487 -- Create a single binder for the field we want
488 sel_bndr <- mkInternalVar "sel" (field_tys!!i)
489 -- Create a wild binder for the scrutinee
490 let scrut_bndr = MkCore.mkWildBinder scrut_ty
491 -- Create the case expression
492 let binders = take i wildbndrs ++ [sel_bndr] ++ drop (i+1) wildbndrs
493 return $ CoreSyn.Case scrut scrut_bndr scrut_ty [(CoreSyn.DataAlt datacon, binders, CoreSyn.Var sel_bndr)]
494 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)
495 Nothing -> error $ "CoreTools.mkSelCase: Creating extractor case, but scrutinee has no tycon? Extracting element " ++ (show i) ++ " from '" ++ pprString scrut ++ "'"