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 = case List.elemIndex dc dcs of
156 Nothing -> error $ "Datacon " ++ pprString dc ++ " does not occur in typed thing: " ++ pprString tt
159 dcs = datacons_for tt
161 -- | Gets all datacons for the given typed thing. Errors out if the
162 -- typed thing is not ADT typed.
163 datacons_for :: TypedThing t => t -> [DataCon.DataCon]
166 Nothing -> error $ "Getting datacon index of untyped thing? " ++ pprString tt
167 Just ty -> case Type.splitTyConApp_maybe ty of
168 Nothing -> error $ "Trying to find datacon in a type without a tycon?" ++ pprString ty
169 Just (tycon, _) -> case TyCon.tyConDataCons_maybe tycon of
170 Nothing -> error $ "Trying to find datacon in a type without datacons?" ++ pprString ty
173 -- Is the given core expression a lambda abstraction?
174 is_lam :: CoreSyn.CoreExpr -> Bool
175 is_lam (CoreSyn.Lam _ _) = True
178 -- Is the given core expression a let expression?
179 is_let :: CoreSyn.CoreExpr -> Bool
180 is_let (CoreSyn.Let _ _) = True
183 -- Is the given core expression of a function type?
184 is_fun :: CoreSyn.CoreExpr -> Bool
185 -- Treat Type arguments differently, because exprType is not defined for them.
186 is_fun (CoreSyn.Type _) = False
187 is_fun expr = (Type.isFunTy . CoreUtils.exprType) expr
189 -- Is the given core expression polymorphic (i.e., does it accept type
191 is_poly :: CoreSyn.CoreExpr -> Bool
192 -- Treat Type arguments differently, because exprType is not defined for them.
193 is_poly (CoreSyn.Type _) = False
194 is_poly expr = (Maybe.isJust . Type.splitForAllTy_maybe . CoreUtils.exprType) expr
196 -- Is the given core expression a variable reference?
197 is_var :: CoreSyn.CoreExpr -> Bool
198 is_var (CoreSyn.Var _) = True
201 is_lit :: CoreSyn.CoreExpr -> Bool
202 is_lit (CoreSyn.Lit _) = True
205 -- Can the given core expression be applied to something? This is true for
206 -- applying to a value as well as a type.
207 is_applicable :: CoreSyn.CoreExpr -> Bool
208 is_applicable expr = is_fun expr || is_poly expr
210 -- Is the given core expression a variable or an application?
211 is_simple :: CoreSyn.CoreExpr -> Bool
212 is_simple (CoreSyn.App _ _) = True
213 is_simple (CoreSyn.Var _) = True
214 is_simple (CoreSyn.Cast expr _) = is_simple expr
217 -- Does the given CoreExpr have any free type vars?
218 has_free_tyvars :: CoreSyn.CoreExpr -> Bool
219 has_free_tyvars = not . VarSet.isEmptyVarSet . (CoreFVs.exprSomeFreeVars Var.isTyVar)
221 -- Does the given type have any free type vars?
222 ty_has_free_tyvars :: Type.Type -> Bool
223 ty_has_free_tyvars = not . VarSet.isEmptyVarSet . Type.tyVarsOfType
225 -- Does the given CoreExpr have any free local vars?
226 has_free_vars :: CoreSyn.CoreExpr -> Bool
227 has_free_vars = not . VarSet.isEmptyVarSet . CoreFVs.exprFreeVars
229 -- Does the given expression use any of the given binders?
230 expr_uses_binders :: [CoreSyn.CoreBndr] -> CoreSyn.CoreExpr -> Bool
231 expr_uses_binders bndrs = not . VarSet.isEmptyVarSet . (CoreFVs.exprSomeFreeVars (`elem` bndrs))
233 -- Turns a Var CoreExpr into the Id inside it. Will of course only work for
234 -- simple Var CoreExprs, not complexer ones.
235 exprToVar :: CoreSyn.CoreExpr -> Var.Id
236 exprToVar (CoreSyn.Var id) = id
237 exprToVar expr = error $ "\nCoreTools.exprToVar: Not a var: " ++ show expr
239 -- Turns a Lit CoreExpr into the Literal inside it.
240 exprToLit :: CoreSyn.CoreExpr -> Literal.Literal
241 exprToLit (CoreSyn.Lit lit) = lit
242 exprToLit expr = error $ "\nCoreTools.exprToLit: Not a lit: " ++ show expr
244 -- Removes all the type and dictionary arguments from the given argument list,
245 -- leaving only the normal value arguments. The type given is the type of the
246 -- expression applied to this argument list.
247 get_val_args :: Type.Type -> [CoreSyn.CoreExpr] -> [CoreSyn.CoreExpr]
248 get_val_args ty args = drop n args
250 (tyvars, predtypes, _) = TcType.tcSplitSigmaTy ty
251 -- The first (length tyvars) arguments should be types, the next
252 -- (length predtypes) arguments should be dictionaries. We drop this many
253 -- arguments, to get at the value arguments.
254 n = length tyvars + length predtypes
256 -- Finds out what literal Integer this expression represents.
257 getIntegerLiteral :: CoreSyn.CoreExpr -> TranslatorSession Integer
258 getIntegerLiteral expr =
259 case CoreSyn.collectArgs expr of
260 (CoreSyn.Var f, [CoreSyn.Lit (Literal.MachInt integer)])
261 | getFullString f == "GHC.Integer.smallInteger" -> return integer
262 (CoreSyn.Var f, [CoreSyn.Lit (Literal.MachInt64 integer)])
263 | getFullString f == "GHC.Integer.int64ToInteger" -> return integer
264 (CoreSyn.Var f, [CoreSyn.Lit (Literal.MachWord integer)])
265 | getFullString f == "GHC.Integer.wordToInteger" -> return integer
266 (CoreSyn.Var f, [CoreSyn.Lit (Literal.MachWord64 integer)])
267 | getFullString f == "GHC.Integer.word64ToInteger" -> return integer
268 -- fromIntegerT returns the integer corresponding to the type of its
269 -- (third) argument. Since it is polymorphic, the type of that
270 -- argument is passed as the first argument, so we can just use that
272 (CoreSyn.Var f, [CoreSyn.Type dec_ty, dec_dict, CoreSyn.Type num_ty, num_dict, arg])
273 | getFullString f == "Types.Data.Num.Ops.fromIntegerT" -> do
274 int <- MonadState.lift tsType $ tfp_to_int dec_ty
275 return $ toInteger int
276 _ -> error $ "CoreTools.getIntegerLiteral: Unsupported Integer literal: " ++ pprString expr
278 reduceCoreListToHsList ::
279 [HscTypes.CoreModule] -- ^ The modules where parts of the list are hidden
280 -> CoreSyn.CoreExpr -- ^ The refence to atleast one of the nodes
281 -> TranslatorSession [CoreSyn.CoreExpr]
282 reduceCoreListToHsList cores app@(CoreSyn.App _ _) = do {
283 ; let { (fun, args) = CoreSyn.collectArgs app
288 ; let topelem = args!!1
290 (varz@(CoreSyn.Var id)) -> do {
291 ; binds <- mapM (findExpr (isVarName id)) cores
292 ; otherelems <- reduceCoreListToHsList cores (head (Maybe.catMaybes binds))
293 ; return (topelem:otherelems)
295 (appz@(CoreSyn.App _ _)) -> do {
296 ; otherelems <- reduceCoreListToHsList cores appz
297 ; return (topelem:otherelems)
299 otherwise -> return [topelem]
301 otherwise -> return []
304 isVarName :: Monad m => Var.Var -> Var.Var -> m Bool
305 isVarName lookfor bind = return $ (Var.varName lookfor) == (Var.varName bind)
307 reduceCoreListToHsList _ _ = return []
309 -- Is the given var the State data constructor?
310 isStateCon :: Var.Var -> Bool
312 -- See if it is a DataConWrapId (not DataConWorkId, since State is a
314 case Id.idDetails var of
315 IdInfo.DataConWrapId dc ->
316 -- See if the datacon is the State datacon from the State type.
317 let tycon = DataCon.dataConTyCon dc
318 tyname = Name.getOccString tycon
319 dcname = Name.getOccString dc
320 in case (tyname, dcname) of
321 ("State", "State") -> True
325 -- | Is the given type a State type?
326 isStateType :: Type.Type -> Bool
327 -- Resolve any type synonyms remaining
328 isStateType ty | Just ty' <- Type.tcView ty = isStateType ty'
329 isStateType ty = Maybe.isJust $ do
330 -- Split the type. Don't use normal splitAppTy, since that looks through
331 -- newtypes, and we want to see the State newtype.
332 (typef, _) <- Type.repSplitAppTy_maybe ty
333 -- See if the applied type is a type constructor
334 (tycon, _) <- Type.splitTyConApp_maybe typef
335 if TyCon.isNewTyCon tycon && Name.getOccString tycon == "State"
341 -- | Does the given TypedThing have a State type?
342 hasStateType :: (TypedThing t) => t -> Bool
343 hasStateType expr = case getType expr of
345 Just ty -> isStateType ty
348 -- | Flattens nested lets into a single list of bindings. The expression
349 -- passed does not have to be a let expression, if it isn't an empty list of
350 -- bindings is returned.
352 CoreSyn.CoreExpr -- ^ The expression to flatten.
353 -> ([Binding], CoreSyn.CoreExpr) -- ^ The bindings and resulting expression.
354 flattenLets (CoreSyn.Let binds expr) =
355 (bindings ++ bindings', expr')
357 -- Recursively flatten the contained expression
358 (bindings', expr') =flattenLets expr
359 -- Flatten our own bindings to remove the Rec / NonRec constructors
360 bindings = CoreSyn.flattenBinds [binds]
361 flattenLets expr = ([], expr)
363 -- | Create bunch of nested non-recursive let expressions from the given
364 -- bindings. The first binding is bound at the highest level (and thus
365 -- available in all other bindings).
366 mkNonRecLets :: [Binding] -> CoreSyn.CoreExpr -> CoreSyn.CoreExpr
367 mkNonRecLets bindings expr = MkCore.mkCoreLets binds expr
369 binds = map (uncurry CoreSyn.NonRec) bindings
371 -- | A class of things that (optionally) have a core Type. The type is
372 -- optional, since Type expressions don't have a type themselves.
373 class Outputable.Outputable t => TypedThing t where
374 getType :: t -> Maybe Type.Type
376 instance TypedThing CoreSyn.CoreExpr where
377 getType (CoreSyn.Type _) = Nothing
378 getType expr = Just $ CoreUtils.exprType expr
380 instance TypedThing CoreSyn.CoreBndr where
381 getType = return . Id.idType
383 instance TypedThing Type.Type where
384 getType = return . id
386 -- | Generate new uniques for all binders in the given expression.
387 -- Does not support making type variables unique, though this could be
388 -- supported if required (by passing a CoreSubst.Subst instead of VarEnv to
389 -- genUniques' below).
390 genUniques :: CoreSyn.CoreExpr -> TranslatorSession CoreSyn.CoreExpr
391 genUniques = genUniques' VarEnv.emptyVarEnv
393 -- | A helper function to generate uniques, that takes a VarEnv containing the
394 -- substitutions already performed.
395 genUniques' :: VarEnv.VarEnv CoreSyn.CoreBndr -> CoreSyn.CoreExpr -> TranslatorSession CoreSyn.CoreExpr
396 genUniques' subst (CoreSyn.Var f) = do
397 -- Replace the binder with its new value, if applicable.
398 let f' = VarEnv.lookupWithDefaultVarEnv subst f f
399 return (CoreSyn.Var f')
400 -- Leave literals untouched
401 genUniques' subst (CoreSyn.Lit l) = return $ CoreSyn.Lit l
402 genUniques' subst (CoreSyn.App f arg) = do
403 -- Only work on subexpressions
404 f' <- genUniques' subst f
405 arg' <- genUniques' subst arg
406 return (CoreSyn.App f' arg')
407 -- Don't change type abstractions
408 genUniques' subst expr@(CoreSyn.Lam bndr res) | CoreSyn.isTyVar bndr = return expr
409 genUniques' subst (CoreSyn.Lam bndr res) = do
410 -- Generate a new unique for the bound variable
411 (subst', bndr') <- genUnique subst bndr
412 res' <- genUniques' subst' res
413 return (CoreSyn.Lam bndr' res')
414 genUniques' subst (CoreSyn.Let (CoreSyn.NonRec bndr bound) res) = do
415 -- Make the binders unique
416 (subst', bndr') <- genUnique subst bndr
417 bound' <- genUniques' subst' bound
418 res' <- genUniques' subst' res
419 return $ CoreSyn.Let (CoreSyn.NonRec bndr' bound') res'
420 genUniques' subst (CoreSyn.Let (CoreSyn.Rec binds) res) = do
421 -- Make each of the binders unique
422 (subst', bndrs') <- mapAccumLM genUnique subst (map fst binds)
423 bounds' <- mapM (genUniques' subst' . snd) binds
424 res' <- genUniques' subst' res
425 let binds' = zip bndrs' bounds'
426 return $ CoreSyn.Let (CoreSyn.Rec binds') res'
427 genUniques' subst (CoreSyn.Case scrut bndr ty alts) = do
428 -- Process the scrutinee with the original substitution, since non of the
429 -- binders bound in the Case statement is in scope in the scrutinee.
430 scrut' <- genUniques' subst scrut
431 -- Generate a new binder for the scrutinee
432 (subst', bndr') <- genUnique subst bndr
433 -- Process each of the alts
434 alts' <- mapM (doalt subst') alts
435 return $ CoreSyn.Case scrut' bndr' ty alts'
437 doalt subst (con, bndrs, expr) = do
438 (subst', bndrs') <- mapAccumLM genUnique subst bndrs
439 expr' <- genUniques' subst' expr
440 -- Note that we don't return subst', since bndrs are only in scope in
442 return (con, bndrs', expr')
443 genUniques' subst (CoreSyn.Cast expr coercion) = do
444 expr' <- genUniques' subst expr
445 -- Just process the casted expression
446 return $ CoreSyn.Cast expr' coercion
447 genUniques' subst (CoreSyn.Note note expr) = do
448 expr' <- genUniques' subst expr
449 -- Just process the annotated expression
450 return $ CoreSyn.Note note expr'
451 -- Leave types untouched
452 genUniques' subst expr@(CoreSyn.Type _) = return expr
454 -- Generate a new unique for the given binder, and extend the given
455 -- substitution to reflect this.
456 genUnique :: VarEnv.VarEnv CoreSyn.CoreBndr -> CoreSyn.CoreBndr -> TranslatorSession (VarEnv.VarEnv CoreSyn.CoreBndr, CoreSyn.CoreBndr)
457 genUnique subst bndr = do
458 bndr' <- BinderTools.cloneVar bndr
459 -- Replace all occurences of the old binder with a reference to the new
461 let subst' = VarEnv.extendVarEnv subst bndr bndr'
462 return (subst', bndr')
464 -- Create a "selector" case that selects the ith field from dc_ith
466 mkSelCase :: CoreSyn.CoreExpr -> Int -> Int -> TranslatorSession CoreSyn.CoreExpr
467 mkSelCase scrut dc_i i = do
468 case Type.splitTyConApp_maybe scrut_ty of
469 -- The scrutinee should have a type constructor. We keep the type
470 -- arguments around so we can instantiate the field types below
471 Just (tycon, tyargs) -> case TyCon.tyConDataCons_maybe tycon of
472 -- The scrutinee type should have a single dataconstructor,
473 -- otherwise we can't construct a valid selector case.
474 Just dcs | dc_i < 0 || dc_i >= length dcs -> error $ "\nCoreTools.mkSelCase: Creating extractor case, but datacon index is invalid." ++ error_msg
476 let datacon = (dcs!!dc_i)
477 let field_tys = DataCon.dataConInstOrigArgTys datacon tyargs
478 if i < 0 || i >= length field_tys
479 then error $ "\nCoreTools.mkSelCase: Creating extractor case, but field index is invalid." ++ error_msg
481 -- Create a list of wild binders for the fields we don't want
482 let wildbndrs = map MkCore.mkWildBinder field_tys
483 -- Create a single binder for the field we want
484 sel_bndr <- mkInternalVar "sel" (field_tys!!i)
485 -- Create a wild binder for the scrutinee
486 let scrut_bndr = MkCore.mkWildBinder scrut_ty
487 -- Create the case expression
488 let binders = take i wildbndrs ++ [sel_bndr] ++ drop (i+1) wildbndrs
489 return $ CoreSyn.Case scrut scrut_bndr scrut_ty [(CoreSyn.DataAlt datacon, binders, CoreSyn.Var sel_bndr)]
490 Nothing -> error $ "CoreTools.mkSelCase: Creating extractor case, but scrutinee has no datacons?" ++ error_msg
491 Nothing -> error $ "CoreTools.mkSelCase: Creating extractor case, but scrutinee has no tycon?" ++ error_msg
493 scrut_ty = CoreUtils.exprType scrut
494 error_msg = " Extracting element " ++ (show i) ++ " from datacon " ++ (show dc_i) ++ " from '" ++ pprString scrut ++ "'" ++ " Type: " ++ (pprString scrut_ty)