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
11 import qualified System.IO.Unsafe
12 import qualified Data.Map as Map
13 import qualified Data.Accessor.Monad.Trans.State as MonadState
18 import qualified TcType
19 import qualified HsExpr
20 import qualified HsTypes
21 import qualified HscTypes
24 import qualified TyCon
25 import qualified DataCon
26 import qualified TysWiredIn
27 import qualified DynFlags
28 import qualified SrcLoc
29 import qualified CoreSyn
31 import qualified IdInfo
32 import qualified VarSet
33 import qualified CoreUtils
34 import qualified CoreFVs
35 import qualified Literal
36 import qualified MkCore
37 import qualified VarEnv
38 import qualified Outputable
41 import CLasH.Translator.TranslatorTypes
42 import CLasH.Utils.GhcTools
43 import CLasH.Utils.Core.BinderTools
44 import CLasH.Utils.HsTools
45 import CLasH.Utils.Pretty
47 import qualified CLasH.Utils.Core.BinderTools as BinderTools
49 -- | A single binding, used as a shortcut to simplify type signatures.
50 type Binding = (CoreSyn.CoreBndr, CoreSyn.CoreExpr)
52 -- | Evaluate a core Type representing type level int from the tfp
53 -- library to a real int. Checks if the type really is a Dec type and
54 -- caches the results.
55 tfp_to_int :: Type.Type -> TypeSession Int
57 hscenv <- MonadState.get tsHscEnv
58 let norm_ty = normalize_tfp_int hscenv ty
59 case Type.splitTyConApp_maybe norm_ty of
60 Just (tycon, args) -> do
61 let name = Name.getOccString (TyCon.tyConName tycon)
66 return $ error ("Callin tfp_to_int on non-dec:" ++ (show ty))
67 Nothing -> return $ error ("Callin tfp_to_int on non-dec:" ++ (show ty))
69 -- | Evaluate a core Type representing type level int from the tfp
70 -- library to a real int. Caches the results. Do not use directly, use
71 -- tfp_to_int instead.
72 tfp_to_int' :: Type.Type -> TypeSession Int
74 lens <- MonadState.get tsTfpInts
75 hscenv <- MonadState.get tsHscEnv
76 let norm_ty = normalize_tfp_int hscenv ty
77 let existing_len = Map.lookup (OrdType norm_ty) lens
79 Just len -> return len
81 let new_len = eval_tfp_int hscenv ty
82 MonadState.modify tsTfpInts (Map.insert (OrdType norm_ty) (new_len))
85 -- | Evaluate a core Type representing type level int from the tfp
86 -- library to a real int. Do not use directly, use tfp_to_int instead.
87 eval_tfp_int :: HscTypes.HscEnv -> Type.Type -> Int
89 unsafeRunGhc libdir $ do
91 -- Automatically import modules for any fully qualified identifiers
92 setDynFlag DynFlags.Opt_ImplicitImportQualified
94 let from_int_t_name = mkRdrName "Types.Data.Num.Ops" "fromIntegerT"
95 let from_int_t = SrcLoc.noLoc $ HsExpr.HsVar from_int_t_name
96 let undef = hsTypedUndef $ coreToHsType ty
97 let app = SrcLoc.noLoc $ HsExpr.HsApp (from_int_t) (undef)
98 let int_ty = SrcLoc.noLoc $ HsTypes.HsTyVar TysWiredIn.intTyCon_RDR
99 let expr = HsExpr.ExprWithTySig app int_ty
103 libdir = DynFlags.topDir dynflags
104 dynflags = HscTypes.hsc_dflags env
106 normalize_tfp_int :: HscTypes.HscEnv -> Type.Type -> Type.Type
107 normalize_tfp_int env ty =
108 System.IO.Unsafe.unsafePerformIO $
111 sized_word_len_ty :: Type.Type -> Type.Type
112 sized_word_len_ty ty = len
114 args = case Type.splitTyConApp_maybe ty of
115 Just (tycon, args) -> args
116 Nothing -> error $ "\nCoreTools.sized_word_len_ty: Not a sized word type: " ++ (pprString ty)
119 sized_int_len_ty :: Type.Type -> Type.Type
120 sized_int_len_ty ty = len
122 args = case Type.splitTyConApp_maybe ty of
123 Just (tycon, args) -> args
124 Nothing -> error $ "\nCoreTools.sized_int_len_ty: Not a sized int type: " ++ (pprString ty)
127 ranged_word_bound_ty :: Type.Type -> Type.Type
128 ranged_word_bound_ty ty = len
130 args = case Type.splitTyConApp_maybe ty of
131 Just (tycon, args) -> args
132 Nothing -> error $ "\nCoreTools.ranged_word_bound_ty: Not a sized word type: " ++ (pprString ty)
135 tfvec_len_ty :: Type.Type -> Type.Type
136 tfvec_len_ty ty = len
138 args = case Type.splitTyConApp_maybe ty of
139 Just (tycon, args) -> args
140 Nothing -> error $ "\nCoreTools.tfvec_len_ty: Not a vector type: " ++ (pprString ty)
143 -- | Get the element type of a TFVec type
144 tfvec_elem :: Type.Type -> Type.Type
145 tfvec_elem ty = el_ty
147 args = case Type.splitTyConApp_maybe ty of
148 Just (tycon, args) -> args
149 Nothing -> error $ "\nCoreTools.tfvec_len: Not a vector type: " ++ (pprString ty)
152 -- | Gets the index of the given datacon in the given typed thing.
153 -- Errors out if it does not occur or if the type is not an ADT.
154 datacon_index :: TypedThing t => t -> DataCon.DataCon -> Int
155 datacon_index tt dc =
157 Nothing -> error $ "Getting datacon index of untyped thing? " ++ pprString tt
158 Just ty -> case Type.splitTyConApp_maybe ty of
159 Nothing -> error $ "Trying to find datacon in a type without a tycon?" ++ pprString ty
160 Just (tycon, _) -> case TyCon.tyConDataCons_maybe tycon of
161 Nothing -> error $ "Trying to find datacon in a type without datacons?" ++ pprString ty
162 Just dcs -> case List.elemIndex dc dcs of
163 Nothing -> error $ "Datacon " ++ pprString dc ++ " does not occur in type: " ++ pprString ty
166 -- Is the given core expression a lambda abstraction?
167 is_lam :: CoreSyn.CoreExpr -> Bool
168 is_lam (CoreSyn.Lam _ _) = True
171 -- Is the given core expression a let expression?
172 is_let :: CoreSyn.CoreExpr -> Bool
173 is_let (CoreSyn.Let _ _) = True
176 -- Is the given core expression of a function type?
177 is_fun :: CoreSyn.CoreExpr -> Bool
178 -- Treat Type arguments differently, because exprType is not defined for them.
179 is_fun (CoreSyn.Type _) = False
180 is_fun expr = (Type.isFunTy . CoreUtils.exprType) expr
182 -- Is the given core expression polymorphic (i.e., does it accept type
184 is_poly :: CoreSyn.CoreExpr -> Bool
185 -- Treat Type arguments differently, because exprType is not defined for them.
186 is_poly (CoreSyn.Type _) = False
187 is_poly expr = (Maybe.isJust . Type.splitForAllTy_maybe . CoreUtils.exprType) expr
189 -- Is the given core expression a variable reference?
190 is_var :: CoreSyn.CoreExpr -> Bool
191 is_var (CoreSyn.Var _) = True
194 is_lit :: CoreSyn.CoreExpr -> Bool
195 is_lit (CoreSyn.Lit _) = True
198 -- Can the given core expression be applied to something? This is true for
199 -- applying to a value as well as a type.
200 is_applicable :: CoreSyn.CoreExpr -> Bool
201 is_applicable expr = is_fun expr || is_poly expr
203 -- Is the given core expression a variable or an application?
204 is_simple :: CoreSyn.CoreExpr -> Bool
205 is_simple (CoreSyn.App _ _) = True
206 is_simple (CoreSyn.Var _) = True
207 is_simple (CoreSyn.Cast expr _) = is_simple expr
210 -- Does the given CoreExpr have any free type vars?
211 has_free_tyvars :: CoreSyn.CoreExpr -> Bool
212 has_free_tyvars = not . VarSet.isEmptyVarSet . (CoreFVs.exprSomeFreeVars Var.isTyVar)
214 -- Does the given type have any free type vars?
215 ty_has_free_tyvars :: Type.Type -> Bool
216 ty_has_free_tyvars = not . VarSet.isEmptyVarSet . Type.tyVarsOfType
218 -- Does the given CoreExpr have any free local vars?
219 has_free_vars :: CoreSyn.CoreExpr -> Bool
220 has_free_vars = not . VarSet.isEmptyVarSet . CoreFVs.exprFreeVars
222 -- Does the given expression use any of the given binders?
223 expr_uses_binders :: [CoreSyn.CoreBndr] -> CoreSyn.CoreExpr -> Bool
224 expr_uses_binders bndrs = not . VarSet.isEmptyVarSet . (CoreFVs.exprSomeFreeVars (`elem` bndrs))
226 -- Turns a Var CoreExpr into the Id inside it. Will of course only work for
227 -- simple Var CoreExprs, not complexer ones.
228 exprToVar :: CoreSyn.CoreExpr -> Var.Id
229 exprToVar (CoreSyn.Var id) = id
230 exprToVar expr = error $ "\nCoreTools.exprToVar: Not a var: " ++ show expr
232 -- Turns a Lit CoreExpr into the Literal inside it.
233 exprToLit :: CoreSyn.CoreExpr -> Literal.Literal
234 exprToLit (CoreSyn.Lit lit) = lit
235 exprToLit expr = error $ "\nCoreTools.exprToLit: Not a lit: " ++ show expr
237 -- Removes all the type and dictionary arguments from the given argument list,
238 -- leaving only the normal value arguments. The type given is the type of the
239 -- expression applied to this argument list.
240 get_val_args :: Type.Type -> [CoreSyn.CoreExpr] -> [CoreSyn.CoreExpr]
241 get_val_args ty args = drop n args
243 (tyvars, predtypes, _) = TcType.tcSplitSigmaTy ty
244 -- The first (length tyvars) arguments should be types, the next
245 -- (length predtypes) arguments should be dictionaries. We drop this many
246 -- arguments, to get at the value arguments.
247 n = length tyvars + length predtypes
249 -- Finds out what literal Integer this expression represents.
250 getIntegerLiteral :: CoreSyn.CoreExpr -> TranslatorSession Integer
251 getIntegerLiteral expr =
252 case CoreSyn.collectArgs expr of
253 (CoreSyn.Var f, [CoreSyn.Lit (Literal.MachInt integer)])
254 | getFullString f == "GHC.Integer.smallInteger" -> return integer
255 (CoreSyn.Var f, [CoreSyn.Lit (Literal.MachInt64 integer)])
256 | getFullString f == "GHC.Integer.int64ToInteger" -> return integer
257 (CoreSyn.Var f, [CoreSyn.Lit (Literal.MachWord integer)])
258 | getFullString f == "GHC.Integer.wordToInteger" -> return integer
259 (CoreSyn.Var f, [CoreSyn.Lit (Literal.MachWord64 integer)])
260 | getFullString f == "GHC.Integer.word64ToInteger" -> return integer
261 -- fromIntegerT returns the integer corresponding to the type of its
262 -- (third) argument. Since it is polymorphic, the type of that
263 -- argument is passed as the first argument, so we can just use that
265 (CoreSyn.Var f, [CoreSyn.Type dec_ty, dec_dict, CoreSyn.Type num_ty, num_dict, arg])
266 | getFullString f == "Types.Data.Num.Ops.fromIntegerT" -> do
267 int <- MonadState.lift tsType $ tfp_to_int dec_ty
268 return $ toInteger int
269 _ -> error $ "CoreTools.getIntegerLiteral: Unsupported Integer literal: " ++ pprString expr
271 reduceCoreListToHsList ::
272 [HscTypes.CoreModule] -- ^ The modules where parts of the list are hidden
273 -> CoreSyn.CoreExpr -- ^ The refence to atleast one of the nodes
274 -> TranslatorSession [CoreSyn.CoreExpr]
275 reduceCoreListToHsList cores app@(CoreSyn.App _ _) = do {
276 ; let { (fun, args) = CoreSyn.collectArgs app
281 ; let topelem = args!!1
283 (varz@(CoreSyn.Var id)) -> do {
284 ; binds <- mapM (findExpr (isVarName id)) cores
285 ; otherelems <- reduceCoreListToHsList cores (head (Maybe.catMaybes binds))
286 ; return (topelem:otherelems)
288 (appz@(CoreSyn.App _ _)) -> do {
289 ; otherelems <- reduceCoreListToHsList cores appz
290 ; return (topelem:otherelems)
292 otherwise -> return [topelem]
294 otherwise -> return []
297 isVarName :: Monad m => Var.Var -> Var.Var -> m Bool
298 isVarName lookfor bind = return $ (Var.varName lookfor) == (Var.varName bind)
300 reduceCoreListToHsList _ _ = return []
302 -- Is the given var the State data constructor?
303 isStateCon :: Var.Var -> Bool
305 -- See if it is a DataConWrapId (not DataConWorkId, since State is a
307 case Id.idDetails var of
308 IdInfo.DataConWrapId dc ->
309 -- See if the datacon is the State datacon from the State type.
310 let tycon = DataCon.dataConTyCon dc
311 tyname = Name.getOccString tycon
312 dcname = Name.getOccString dc
313 in case (tyname, dcname) of
314 ("State", "State") -> True
318 -- | Is the given type a State type?
319 isStateType :: Type.Type -> Bool
320 -- Resolve any type synonyms remaining
321 isStateType ty | Just ty' <- Type.tcView ty = isStateType ty'
322 isStateType ty = Maybe.isJust $ do
323 -- Split the type. Don't use normal splitAppTy, since that looks through
324 -- newtypes, and we want to see the State newtype.
325 (typef, _) <- Type.repSplitAppTy_maybe ty
326 -- See if the applied type is a type constructor
327 (tycon, _) <- Type.splitTyConApp_maybe typef
328 if TyCon.isNewTyCon tycon && Name.getOccString tycon == "State"
334 -- | Does the given TypedThing have a State type?
335 hasStateType :: (TypedThing t) => t -> Bool
336 hasStateType expr = case getType expr of
338 Just ty -> isStateType ty
341 -- | Flattens nested lets into a single list of bindings. The expression
342 -- passed does not have to be a let expression, if it isn't an empty list of
343 -- bindings is returned.
345 CoreSyn.CoreExpr -- ^ The expression to flatten.
346 -> ([Binding], CoreSyn.CoreExpr) -- ^ The bindings and resulting expression.
347 flattenLets (CoreSyn.Let binds expr) =
348 (bindings ++ bindings', expr')
350 -- Recursively flatten the contained expression
351 (bindings', expr') =flattenLets expr
352 -- Flatten our own bindings to remove the Rec / NonRec constructors
353 bindings = CoreSyn.flattenBinds [binds]
354 flattenLets expr = ([], expr)
356 -- | Create bunch of nested non-recursive let expressions from the given
357 -- bindings. The first binding is bound at the highest level (and thus
358 -- available in all other bindings).
359 mkNonRecLets :: [Binding] -> CoreSyn.CoreExpr -> CoreSyn.CoreExpr
360 mkNonRecLets bindings expr = MkCore.mkCoreLets binds expr
362 binds = map (uncurry CoreSyn.NonRec) bindings
364 -- | A class of things that (optionally) have a core Type. The type is
365 -- optional, since Type expressions don't have a type themselves.
366 class Outputable.Outputable t => TypedThing t where
367 getType :: t -> Maybe Type.Type
369 instance TypedThing CoreSyn.CoreExpr where
370 getType (CoreSyn.Type _) = Nothing
371 getType expr = Just $ CoreUtils.exprType expr
373 instance TypedThing CoreSyn.CoreBndr where
374 getType = return . Id.idType
376 instance TypedThing Type.Type where
377 getType = return . id
379 -- | Generate new uniques for all binders in the given expression.
380 -- Does not support making type variables unique, though this could be
381 -- supported if required (by passing a CoreSubst.Subst instead of VarEnv to
382 -- genUniques' below).
383 genUniques :: CoreSyn.CoreExpr -> TranslatorSession CoreSyn.CoreExpr
384 genUniques = genUniques' VarEnv.emptyVarEnv
386 -- | A helper function to generate uniques, that takes a VarEnv containing the
387 -- substitutions already performed.
388 genUniques' :: VarEnv.VarEnv CoreSyn.CoreBndr -> CoreSyn.CoreExpr -> TranslatorSession CoreSyn.CoreExpr
389 genUniques' subst (CoreSyn.Var f) = do
390 -- Replace the binder with its new value, if applicable.
391 let f' = VarEnv.lookupWithDefaultVarEnv subst f f
392 return (CoreSyn.Var f')
393 -- Leave literals untouched
394 genUniques' subst (CoreSyn.Lit l) = return $ CoreSyn.Lit l
395 genUniques' subst (CoreSyn.App f arg) = do
396 -- Only work on subexpressions
397 f' <- genUniques' subst f
398 arg' <- genUniques' subst arg
399 return (CoreSyn.App f' arg')
400 -- Don't change type abstractions
401 genUniques' subst expr@(CoreSyn.Lam bndr res) | CoreSyn.isTyVar bndr = return expr
402 genUniques' subst (CoreSyn.Lam bndr res) = do
403 -- Generate a new unique for the bound variable
404 (subst', bndr') <- genUnique subst bndr
405 res' <- genUniques' subst' res
406 return (CoreSyn.Lam bndr' res')
407 genUniques' subst (CoreSyn.Let (CoreSyn.NonRec bndr bound) res) = do
408 -- Make the binders unique
409 (subst', bndr') <- genUnique subst bndr
410 bound' <- genUniques' subst' bound
411 res' <- genUniques' subst' res
412 return $ CoreSyn.Let (CoreSyn.NonRec bndr' bound') res'
413 genUniques' subst (CoreSyn.Let (CoreSyn.Rec binds) res) = do
414 -- Make each of the binders unique
415 (subst', bndrs') <- mapAccumLM genUnique subst (map fst binds)
416 bounds' <- mapM (genUniques' subst' . snd) binds
417 res' <- genUniques' subst' res
418 let binds' = zip bndrs' bounds'
419 return $ CoreSyn.Let (CoreSyn.Rec binds') res'
420 genUniques' subst (CoreSyn.Case scrut bndr ty alts) = do
421 -- Process the scrutinee with the original substitution, since non of the
422 -- binders bound in the Case statement is in scope in the scrutinee.
423 scrut' <- genUniques' subst scrut
424 -- Generate a new binder for the scrutinee
425 (subst', bndr') <- genUnique subst bndr
426 -- Process each of the alts
427 alts' <- mapM (doalt subst') alts
428 return $ CoreSyn.Case scrut' bndr' ty alts'
430 doalt subst (con, bndrs, expr) = do
431 (subst', bndrs') <- mapAccumLM genUnique subst bndrs
432 expr' <- genUniques' subst' expr
433 -- Note that we don't return subst', since bndrs are only in scope in
435 return (con, bndrs', expr')
436 genUniques' subst (CoreSyn.Cast expr coercion) = do
437 expr' <- genUniques' subst expr
438 -- Just process the casted expression
439 return $ CoreSyn.Cast expr' coercion
440 genUniques' subst (CoreSyn.Note note expr) = do
441 expr' <- genUniques' subst expr
442 -- Just process the annotated expression
443 return $ CoreSyn.Note note expr'
444 -- Leave types untouched
445 genUniques' subst expr@(CoreSyn.Type _) = return expr
447 -- Generate a new unique for the given binder, and extend the given
448 -- substitution to reflect this.
449 genUnique :: VarEnv.VarEnv CoreSyn.CoreBndr -> CoreSyn.CoreBndr -> TranslatorSession (VarEnv.VarEnv CoreSyn.CoreBndr, CoreSyn.CoreBndr)
450 genUnique subst bndr = do
451 bndr' <- BinderTools.cloneVar bndr
452 -- Replace all occurences of the old binder with a reference to the new
454 let subst' = VarEnv.extendVarEnv subst bndr bndr'
455 return (subst', bndr')
457 -- Create a "selector" case that selects the ith field from dc_ith
459 mkSelCase :: CoreSyn.CoreExpr -> Int -> Int -> TranslatorSession CoreSyn.CoreExpr
460 mkSelCase scrut dc_i i = do
461 case Type.splitTyConApp_maybe scrut_ty of
462 -- The scrutinee should have a type constructor. We keep the type
463 -- arguments around so we can instantiate the field types below
464 Just (tycon, tyargs) -> case TyCon.tyConDataCons_maybe tycon of
465 -- The scrutinee type should have a single dataconstructor,
466 -- otherwise we can't construct a valid selector case.
467 Just dcs | dc_i < 0 || dc_i >= length dcs -> error $ "\nCoreTools.mkSelCase: Creating extractor case, but datacon index is invalid." ++ error_msg
469 let datacon = (dcs!!dc_i)
470 let field_tys = DataCon.dataConInstOrigArgTys datacon tyargs
471 if i < 0 || i >= length field_tys
472 then error $ "\nCoreTools.mkSelCase: Creating extractor case, but field index is invalid." ++ error_msg
474 -- Create a list of wild binders for the fields we don't want
475 let wildbndrs = map MkCore.mkWildBinder field_tys
476 -- Create a single binder for the field we want
477 sel_bndr <- mkInternalVar "sel" (field_tys!!i)
478 -- Create a wild binder for the scrutinee
479 let scrut_bndr = MkCore.mkWildBinder scrut_ty
480 -- Create the case expression
481 let binders = take i wildbndrs ++ [sel_bndr] ++ drop (i+1) wildbndrs
482 return $ CoreSyn.Case scrut scrut_bndr scrut_ty [(CoreSyn.DataAlt datacon, binders, CoreSyn.Var sel_bndr)]
483 Nothing -> error $ "CoreTools.mkSelCase: Creating extractor case, but scrutinee has no datacons?" ++ error_msg
484 Nothing -> error $ "CoreTools.mkSelCase: Creating extractor case, but scrutinee has no tycon?" ++ error_msg
486 scrut_ty = CoreUtils.exprType scrut
487 error_msg = " Extracting element " ++ (show i) ++ " from datacon " ++ (show dc_i) ++ " from '" ++ pprString scrut ++ "'" ++ " Type: " ++ (pprString scrut_ty)