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 = normalize_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 = normalize_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 normalize_tfp_int :: HscTypes.HscEnv -> Type.Type -> Type.Type
105 normalize_tfp_int env ty =
106 System.IO.Unsafe.unsafePerformIO $
109 sized_word_len_ty :: Type.Type -> Type.Type
110 sized_word_len_ty ty = len
112 args = case Type.splitTyConApp_maybe ty of
113 Just (tycon, args) -> args
114 Nothing -> error $ "\nCoreTools.sized_word_len_ty: Not a sized word type: " ++ (pprString ty)
117 sized_int_len_ty :: Type.Type -> Type.Type
118 sized_int_len_ty ty = len
120 args = case Type.splitTyConApp_maybe ty of
121 Just (tycon, args) -> args
122 Nothing -> error $ "\nCoreTools.sized_int_len_ty: Not a sized int type: " ++ (pprString ty)
125 ranged_word_bound_ty :: Type.Type -> Type.Type
126 ranged_word_bound_ty ty = len
128 args = case Type.splitTyConApp_maybe ty of
129 Just (tycon, args) -> args
130 Nothing -> error $ "\nCoreTools.ranged_word_bound_ty: Not a sized word type: " ++ (pprString ty)
133 tfvec_len_ty :: Type.Type -> Type.Type
134 tfvec_len_ty ty = len
136 args = case Type.splitTyConApp_maybe ty of
137 Just (tycon, args) -> args
138 Nothing -> error $ "\nCoreTools.tfvec_len_ty: Not a vector type: " ++ (pprString ty)
141 -- | Get the element type of a TFVec type
142 tfvec_elem :: Type.Type -> Type.Type
143 tfvec_elem ty = el_ty
145 args = case Type.splitTyConApp_maybe ty of
146 Just (tycon, args) -> args
147 Nothing -> error $ "\nCoreTools.tfvec_len: Not a vector type: " ++ (pprString ty)
150 -- Is the given core expression a lambda abstraction?
151 is_lam :: CoreSyn.CoreExpr -> Bool
152 is_lam (CoreSyn.Lam _ _) = True
155 -- Is the given core expression a let expression?
156 is_let :: CoreSyn.CoreExpr -> Bool
157 is_let (CoreSyn.Let _ _) = True
160 -- Is the given core expression of a function type?
161 is_fun :: CoreSyn.CoreExpr -> Bool
162 -- Treat Type arguments differently, because exprType is not defined for them.
163 is_fun (CoreSyn.Type _) = False
164 is_fun expr = (Type.isFunTy . CoreUtils.exprType) expr
166 -- Is the given core expression polymorphic (i.e., does it accept type
168 is_poly :: CoreSyn.CoreExpr -> Bool
169 -- Treat Type arguments differently, because exprType is not defined for them.
170 is_poly (CoreSyn.Type _) = False
171 is_poly expr = (Maybe.isJust . Type.splitForAllTy_maybe . CoreUtils.exprType) expr
173 -- Is the given core expression a variable reference?
174 is_var :: CoreSyn.CoreExpr -> Bool
175 is_var (CoreSyn.Var _) = True
178 is_lit :: CoreSyn.CoreExpr -> Bool
179 is_lit (CoreSyn.Lit _) = True
182 -- Can the given core expression be applied to something? This is true for
183 -- applying to a value as well as a type.
184 is_applicable :: CoreSyn.CoreExpr -> Bool
185 is_applicable expr = is_fun expr || is_poly expr
187 -- Is the given core expression a variable or an application?
188 is_simple :: CoreSyn.CoreExpr -> Bool
189 is_simple (CoreSyn.App _ _) = True
190 is_simple (CoreSyn.Var _) = True
191 is_simple (CoreSyn.Cast expr _) = is_simple expr
194 -- Does the given CoreExpr have any free type vars?
195 has_free_tyvars :: CoreSyn.CoreExpr -> Bool
196 has_free_tyvars = not . VarSet.isEmptyVarSet . (CoreFVs.exprSomeFreeVars Var.isTyVar)
198 -- Does the given type have any free type vars?
199 ty_has_free_tyvars :: Type.Type -> Bool
200 ty_has_free_tyvars = not . VarSet.isEmptyVarSet . Type.tyVarsOfType
202 -- Does the given CoreExpr have any free local vars?
203 has_free_vars :: CoreSyn.CoreExpr -> Bool
204 has_free_vars = not . VarSet.isEmptyVarSet . CoreFVs.exprFreeVars
206 -- Does the given expression use any of the given binders?
207 expr_uses_binders :: [CoreSyn.CoreBndr] -> CoreSyn.CoreExpr -> Bool
208 expr_uses_binders bndrs = not . VarSet.isEmptyVarSet . (CoreFVs.exprSomeFreeVars (`elem` bndrs))
210 -- Turns a Var CoreExpr into the Id inside it. Will of course only work for
211 -- simple Var CoreExprs, not complexer ones.
212 exprToVar :: CoreSyn.CoreExpr -> Var.Id
213 exprToVar (CoreSyn.Var id) = id
214 exprToVar expr = error $ "\nCoreTools.exprToVar: Not a var: " ++ show expr
216 -- Turns a Lit CoreExpr into the Literal inside it.
217 exprToLit :: CoreSyn.CoreExpr -> Literal.Literal
218 exprToLit (CoreSyn.Lit lit) = lit
219 exprToLit expr = error $ "\nCoreTools.exprToLit: Not a lit: " ++ show expr
221 -- Removes all the type and dictionary arguments from the given argument list,
222 -- leaving only the normal value arguments. The type given is the type of the
223 -- expression applied to this argument list.
224 get_val_args :: Type.Type -> [CoreSyn.CoreExpr] -> [CoreSyn.CoreExpr]
225 get_val_args ty args = drop n args
227 (tyvars, predtypes, _) = TcType.tcSplitSigmaTy ty
228 -- The first (length tyvars) arguments should be types, the next
229 -- (length predtypes) arguments should be dictionaries. We drop this many
230 -- arguments, to get at the value arguments.
231 n = length tyvars + length predtypes
233 -- Finds out what literal Integer this expression represents.
234 getIntegerLiteral :: CoreSyn.CoreExpr -> TranslatorSession Integer
235 getIntegerLiteral expr =
236 case CoreSyn.collectArgs expr of
237 (CoreSyn.Var f, [CoreSyn.Lit (Literal.MachInt integer)])
238 | getFullString f == "GHC.Integer.smallInteger" -> return integer
239 (CoreSyn.Var f, [CoreSyn.Lit (Literal.MachInt64 integer)])
240 | getFullString f == "GHC.Integer.int64ToInteger" -> return integer
241 (CoreSyn.Var f, [CoreSyn.Lit (Literal.MachWord integer)])
242 | getFullString f == "GHC.Integer.wordToInteger" -> return integer
243 (CoreSyn.Var f, [CoreSyn.Lit (Literal.MachWord64 integer)])
244 | getFullString f == "GHC.Integer.word64ToInteger" -> return integer
245 -- fromIntegerT returns the integer corresponding to the type of its
246 -- (third) argument. Since it is polymorphic, the type of that
247 -- argument is passed as the first argument, so we can just use that
249 (CoreSyn.Var f, [CoreSyn.Type dec_ty, dec_dict, CoreSyn.Type num_ty, num_dict, arg])
250 | getFullString f == "Types.Data.Num.Ops.fromIntegerT" -> do
251 int <- MonadState.lift tsType $ tfp_to_int dec_ty
252 return $ toInteger int
253 _ -> error $ "CoreTools.getIntegerLiteral: Unsupported Integer literal: " ++ pprString expr
255 reduceCoreListToHsList ::
256 [HscTypes.CoreModule] -- ^ The modules where parts of the list are hidden
257 -> CoreSyn.CoreExpr -- ^ The refence to atleast one of the nodes
258 -> TranslatorSession [CoreSyn.CoreExpr]
259 reduceCoreListToHsList cores app@(CoreSyn.App _ _) = do {
260 ; let { (fun, args) = CoreSyn.collectArgs app
265 ; let topelem = args!!1
267 (varz@(CoreSyn.Var id)) -> do {
268 ; binds <- mapM (findExpr (isVarName id)) cores
269 ; otherelems <- reduceCoreListToHsList cores (head (Maybe.catMaybes binds))
270 ; return (topelem:otherelems)
272 (appz@(CoreSyn.App _ _)) -> do {
273 ; otherelems <- reduceCoreListToHsList cores appz
274 ; return (topelem:otherelems)
276 otherwise -> return [topelem]
278 otherwise -> return []
281 isVarName :: Monad m => Var.Var -> Var.Var -> m Bool
282 isVarName lookfor bind = return $ (Var.varName lookfor) == (Var.varName bind)
284 reduceCoreListToHsList _ _ = return []
286 -- Is the given var the State data constructor?
287 isStateCon :: Var.Var -> Bool
289 -- See if it is a DataConWrapId (not DataConWorkId, since State is a
291 case Id.idDetails var of
292 IdInfo.DataConWrapId dc ->
293 -- See if the datacon is the State datacon from the State type.
294 let tycon = DataCon.dataConTyCon dc
295 tyname = Name.getOccString tycon
296 dcname = Name.getOccString dc
297 in case (tyname, dcname) of
298 ("State", "State") -> True
302 -- | Is the given type a State type?
303 isStateType :: Type.Type -> Bool
304 -- Resolve any type synonyms remaining
305 isStateType ty | Just ty' <- Type.tcView ty = isStateType ty'
306 isStateType ty = Maybe.isJust $ do
307 -- Split the type. Don't use normal splitAppTy, since that looks through
308 -- newtypes, and we want to see the State newtype.
309 (typef, _) <- Type.repSplitAppTy_maybe ty
310 -- See if the applied type is a type constructor
311 (tycon, _) <- Type.splitTyConApp_maybe typef
312 if TyCon.isNewTyCon tycon && Name.getOccString tycon == "State"
318 -- | Does the given TypedThing have a State type?
319 hasStateType :: (TypedThing t) => t -> Bool
320 hasStateType expr = case getType expr of
322 Just ty -> isStateType ty
325 -- | Flattens nested lets into a single list of bindings. The expression
326 -- passed does not have to be a let expression, if it isn't an empty list of
327 -- bindings is returned.
329 CoreSyn.CoreExpr -- ^ The expression to flatten.
330 -> ([Binding], CoreSyn.CoreExpr) -- ^ The bindings and resulting expression.
331 flattenLets (CoreSyn.Let binds expr) =
332 (bindings ++ bindings', expr')
334 -- Recursively flatten the contained expression
335 (bindings', expr') =flattenLets expr
336 -- Flatten our own bindings to remove the Rec / NonRec constructors
337 bindings = CoreSyn.flattenBinds [binds]
338 flattenLets expr = ([], expr)
340 -- | Create bunch of nested non-recursive let expressions from the given
341 -- bindings. The first binding is bound at the highest level (and thus
342 -- available in all other bindings).
343 mkNonRecLets :: [Binding] -> CoreSyn.CoreExpr -> CoreSyn.CoreExpr
344 mkNonRecLets bindings expr = MkCore.mkCoreLets binds expr
346 binds = map (uncurry CoreSyn.NonRec) bindings
348 -- | A class of things that (optionally) have a core Type. The type is
349 -- optional, since Type expressions don't have a type themselves.
350 class TypedThing t where
351 getType :: t -> Maybe Type.Type
353 instance TypedThing CoreSyn.CoreExpr where
354 getType (CoreSyn.Type _) = Nothing
355 getType expr = Just $ CoreUtils.exprType expr
357 instance TypedThing CoreSyn.CoreBndr where
358 getType = return . Id.idType
360 instance TypedThing Type.Type where
361 getType = return . id
363 -- | Generate new uniques for all binders in the given expression.
364 -- Does not support making type variables unique, though this could be
365 -- supported if required (by passing a CoreSubst.Subst instead of VarEnv to
366 -- genUniques' below).
367 genUniques :: CoreSyn.CoreExpr -> TranslatorSession CoreSyn.CoreExpr
368 genUniques = genUniques' VarEnv.emptyVarEnv
370 -- | A helper function to generate uniques, that takes a VarEnv containing the
371 -- substitutions already performed.
372 genUniques' :: VarEnv.VarEnv CoreSyn.CoreBndr -> CoreSyn.CoreExpr -> TranslatorSession CoreSyn.CoreExpr
373 genUniques' subst (CoreSyn.Var f) = do
374 -- Replace the binder with its new value, if applicable.
375 let f' = VarEnv.lookupWithDefaultVarEnv subst f f
376 return (CoreSyn.Var f')
377 -- Leave literals untouched
378 genUniques' subst (CoreSyn.Lit l) = return $ CoreSyn.Lit l
379 genUniques' subst (CoreSyn.App f arg) = do
380 -- Only work on subexpressions
381 f' <- genUniques' subst f
382 arg' <- genUniques' subst arg
383 return (CoreSyn.App f' arg')
384 -- Don't change type abstractions
385 genUniques' subst expr@(CoreSyn.Lam bndr res) | CoreSyn.isTyVar bndr = return expr
386 genUniques' subst (CoreSyn.Lam bndr res) = do
387 -- Generate a new unique for the bound variable
388 (subst', bndr') <- genUnique subst bndr
389 res' <- genUniques' subst' res
390 return (CoreSyn.Lam bndr' res')
391 genUniques' subst (CoreSyn.Let (CoreSyn.NonRec bndr bound) res) = do
392 -- Make the binders unique
393 (subst', bndr') <- genUnique subst bndr
394 bound' <- genUniques' subst' bound
395 res' <- genUniques' subst' res
396 return $ CoreSyn.Let (CoreSyn.NonRec bndr' bound') res'
397 genUniques' subst (CoreSyn.Let (CoreSyn.Rec binds) res) = do
398 -- Make each of the binders unique
399 (subst', bndrs') <- mapAccumLM genUnique subst (map fst binds)
400 bounds' <- mapM (genUniques' subst' . snd) binds
401 res' <- genUniques' subst' res
402 let binds' = zip bndrs' bounds'
403 return $ CoreSyn.Let (CoreSyn.Rec binds') res'
404 genUniques' subst (CoreSyn.Case scrut bndr ty alts) = do
405 -- Process the scrutinee with the original substitution, since non of the
406 -- binders bound in the Case statement is in scope in the scrutinee.
407 scrut' <- genUniques' subst scrut
408 -- Generate a new binder for the scrutinee
409 (subst', bndr') <- genUnique subst bndr
410 -- Process each of the alts
411 alts' <- mapM (doalt subst') alts
412 return $ CoreSyn.Case scrut' bndr' ty alts'
414 doalt subst (con, bndrs, expr) = do
415 (subst', bndrs') <- mapAccumLM genUnique subst bndrs
416 expr' <- genUniques' subst' expr
417 -- Note that we don't return subst', since bndrs are only in scope in
419 return (con, bndrs', expr')
420 genUniques' subst (CoreSyn.Cast expr coercion) = do
421 expr' <- genUniques' subst expr
422 -- Just process the casted expression
423 return $ CoreSyn.Cast expr' coercion
424 genUniques' subst (CoreSyn.Note note expr) = do
425 expr' <- genUniques' subst expr
426 -- Just process the annotated expression
427 return $ CoreSyn.Note note expr'
428 -- Leave types untouched
429 genUniques' subst expr@(CoreSyn.Type _) = return expr
431 -- Generate a new unique for the given binder, and extend the given
432 -- substitution to reflect this.
433 genUnique :: VarEnv.VarEnv CoreSyn.CoreBndr -> CoreSyn.CoreBndr -> TranslatorSession (VarEnv.VarEnv CoreSyn.CoreBndr, CoreSyn.CoreBndr)
434 genUnique subst bndr = do
435 bndr' <- BinderTools.cloneVar bndr
436 -- Replace all occurences of the old binder with a reference to the new
438 let subst' = VarEnv.extendVarEnv subst bndr bndr'
439 return (subst', bndr')
441 -- Create a "selector" case that selects the ith field from a datacon
442 mkSelCase :: CoreSyn.CoreExpr -> Int -> TranslatorSession CoreSyn.CoreExpr
443 mkSelCase scrut i = do
444 let scrut_ty = CoreUtils.exprType scrut
445 case Type.splitTyConApp_maybe scrut_ty of
446 -- The scrutinee should have a type constructor. We keep the type
447 -- arguments around so we can instantiate the field types below
448 Just (tycon, tyargs) -> case TyCon.tyConDataCons tycon of
449 -- The scrutinee type should have a single dataconstructor,
450 -- otherwise we can't construct a valid selector case.
452 let field_tys = DataCon.dataConInstOrigArgTys datacon tyargs
453 -- Create a list of wild binders for the fields we don't want
454 let wildbndrs = map MkCore.mkWildBinder field_tys
455 -- Create a single binder for the field we want
456 sel_bndr <- mkInternalVar "sel" (field_tys!!i)
457 -- Create a wild binder for the scrutinee
458 let scrut_bndr = MkCore.mkWildBinder scrut_ty
459 -- Create the case expression
460 let binders = take i wildbndrs ++ [sel_bndr] ++ drop (i+1) wildbndrs
461 return $ CoreSyn.Case scrut scrut_bndr scrut_ty [(CoreSyn.DataAlt datacon, binders, CoreSyn.Var sel_bndr)]
462 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)
463 Nothing -> error $ "CoreTools.mkSelCase: Creating extractor case, but scrutinee has no tycon? Extracting element " ++ (show i) ++ " from '" ++ pprString scrut ++ "'" ++ " Type: " ++ (pprString scrut_ty)