cλash/CLasH/Utils/Core/CoreTools.hs

   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
   7
   8 --Standard modules
   9 import qualified Maybe
  10 import qualified System.IO.Unsafe
  11
  12 -- GHC API
  13 import qualified GHC
  14 import qualified Type
  15 import qualified TcType
  16 import qualified HsExpr
  17 import qualified HsTypes
  18 import qualified HscTypes
  19 import qualified Name
  20 import qualified Id
  21 import qualified TyCon
  22 import qualified DataCon
  23 import qualified TysWiredIn
  24 import qualified DynFlags
  25 import qualified SrcLoc
  26 import qualified CoreSyn
  27 import qualified Var
  28 import qualified IdInfo
  29 import qualified VarSet
  30 import qualified CoreUtils
  31 import qualified CoreFVs
  32 import qualified Literal
  33 import qualified MkCore
  34 import qualified VarEnv
  35
  36 -- Local imports
  37 import CLasH.Translator.TranslatorTypes
  38 import CLasH.Utils.GhcTools
  39 import CLasH.Utils.HsTools
  40 import CLasH.Utils.Pretty
  41 import CLasH.Utils
  42 import qualified CLasH.Utils.Core.BinderTools as BinderTools
  43
  44 -- | A single binding, used as a shortcut to simplify type signatures.
  45 type Binding = (CoreSyn.CoreBndr, CoreSyn.CoreExpr)
  46
  47 -- | Evaluate a core Type representing type level int from the tfp
  48 -- library to a real int.
  49 eval_tfp_int :: HscTypes.HscEnv -> Type.Type -> Int
  50 eval_tfp_int env ty =
  51   unsafeRunGhc libdir $ do
  52     GHC.setSession env
  53     -- Automatically import modules for any fully qualified identifiers
  54     setDynFlag DynFlags.Opt_ImplicitImportQualified
  55
  56     let from_int_t_name = mkRdrName "Types.Data.Num.Ops" "fromIntegerT"
  57     let from_int_t = SrcLoc.noLoc $ HsExpr.HsVar from_int_t_name
  58     let undef = hsTypedUndef $ coreToHsType ty
  59     let app = SrcLoc.noLoc $ HsExpr.HsApp (from_int_t) (undef)
  60     let int_ty = SrcLoc.noLoc $ HsTypes.HsTyVar TysWiredIn.intTyCon_RDR
  61     let expr = HsExpr.ExprWithTySig app int_ty
  62     core <- toCore expr
  63     execCore core
  64   where
  65     libdir = DynFlags.topDir dynflags
  66     dynflags = HscTypes.hsc_dflags env
  67
  68 normalise_tfp_int :: HscTypes.HscEnv -> Type.Type -> Type.Type
  69 normalise_tfp_int env ty =
  70    System.IO.Unsafe.unsafePerformIO $
  71      normaliseType env ty
  72
  73 -- | Get the width of a SizedWord type
  74 -- sized_word_len :: HscTypes.HscEnv -> Type.Type -> Int
  75 -- sized_word_len env ty = eval_tfp_int env (sized_word_len_ty ty)
  76
  77 sized_word_len_ty :: Type.Type -> Type.Type
  78 sized_word_len_ty ty = len
  79   where
  80     args = case Type.splitTyConApp_maybe ty of
  81       Just (tycon, args) -> args
  82       Nothing -> error $ "\nCoreTools.sized_word_len_ty: Not a sized word type: " ++ (pprString ty)
  83     [len]         = args
  84
  85 -- | Get the width of a SizedInt type
  86 -- sized_int_len :: HscTypes.HscEnv -> Type.Type -> Int
  87 -- sized_int_len env ty = eval_tfp_int env (sized_int_len_ty ty)
  88
  89 sized_int_len_ty :: Type.Type -> Type.Type
  90 sized_int_len_ty ty = len
  91   where
  92     args = case Type.splitTyConApp_maybe ty of
  93       Just (tycon, args) -> args
  94       Nothing -> error $ "\nCoreTools.sized_int_len_ty: Not a sized int type: " ++ (pprString ty)
  95     [len]         = args
  96
  97 -- | Get the upperbound of a RangedWord type
  98 -- ranged_word_bound :: HscTypes.HscEnv -> Type.Type -> Int
  99 -- ranged_word_bound env ty = eval_tfp_int env (ranged_word_bound_ty ty)
 100
 101 ranged_word_bound_ty :: Type.Type -> Type.Type
 102 ranged_word_bound_ty ty = len
 103   where
 104     args = case Type.splitTyConApp_maybe ty of
 105       Just (tycon, args) -> args
 106       Nothing -> error $ "\nCoreTools.ranged_word_bound_ty: Not a sized word type: " ++ (pprString ty)
 107     [len]         = args
 108
 109 -- | Evaluate a core Type representing type level int from the TypeLevel
 110 -- library to a real int.
 111 -- eval_type_level_int :: Type.Type -> Int
 112 -- eval_type_level_int ty =
 113 --   unsafeRunGhc $ do
 114 --     -- Automatically import modules for any fully qualified identifiers
 115 --     setDynFlag DynFlags.Opt_ImplicitImportQualified
 116 --
 117 --     let to_int_name = mkRdrName "Data.TypeLevel.Num.Sets" "toInt"
 118 --     let to_int = SrcLoc.noLoc $ HsExpr.HsVar to_int_name
 119 --     let undef = hsTypedUndef $ coreToHsType ty
 120 --     let app = HsExpr.HsApp (to_int) (undef)
 121 --
 122 --     core <- toCore [] app
 123 --     execCore core
 124
 125 -- | Get the length of a FSVec type
 126 -- tfvec_len :: HscTypes.HscEnv -> Type.Type -> Int
 127 -- tfvec_len env ty = eval_tfp_int env (tfvec_len_ty ty)
 128
 129 tfvec_len_ty :: Type.Type -> Type.Type
 130 tfvec_len_ty ty = len
 131   where
 132     args = case Type.splitTyConApp_maybe ty of
 133       Just (tycon, args) -> args
 134       Nothing -> error $ "\nCoreTools.tfvec_len_ty: Not a vector type: " ++ (pprString ty)
 135     [len, el_ty] = args
 136
 137 -- | Get the element type of a TFVec type
 138 tfvec_elem :: Type.Type -> Type.Type
 139 tfvec_elem ty = el_ty
 140   where
 141     args = case Type.splitTyConApp_maybe ty of
 142       Just (tycon, args) -> args
 143       Nothing -> error $ "\nCoreTools.tfvec_len: Not a vector type: " ++ (pprString ty)
 144     [len, el_ty] = args
 145
 146 -- Is the given core expression a lambda abstraction?
 147 is_lam :: CoreSyn.CoreExpr -> Bool
 148 is_lam (CoreSyn.Lam _ _) = True
 149 is_lam _ = False
 150
 151 -- Is the given core expression a let expression?
 152 is_let :: CoreSyn.CoreExpr -> Bool
 153 is_let (CoreSyn.Let _ _) = True
 154 is_let _ = False
 155
 156 -- Is the given core expression of a function type?
 157 is_fun :: CoreSyn.CoreExpr -> Bool
 158 -- Treat Type arguments differently, because exprType is not defined for them.
 159 is_fun (CoreSyn.Type _) = False
 160 is_fun expr = (Type.isFunTy . CoreUtils.exprType) expr
 161
 162 -- Is the given core expression polymorphic (i.e., does it accept type
 163 -- arguments?).
 164 is_poly :: CoreSyn.CoreExpr -> Bool
 165 -- Treat Type arguments differently, because exprType is not defined for them.
 166 is_poly (CoreSyn.Type _) = False
 167 is_poly expr = (Maybe.isJust . Type.splitForAllTy_maybe . CoreUtils.exprType) expr
 168
 169 -- Is the given core expression a variable reference?
 170 is_var :: CoreSyn.CoreExpr -> Bool
 171 is_var (CoreSyn.Var _) = True
 172 is_var _ = False
 173
 174 is_lit :: CoreSyn.CoreExpr -> Bool
 175 is_lit (CoreSyn.Lit _) = True
 176 is_lit _ = False
 177
 178 -- Can the given core expression be applied to something? This is true for
 179 -- applying to a value as well as a type.
 180 is_applicable :: CoreSyn.CoreExpr -> Bool
 181 is_applicable expr = is_fun expr || is_poly expr
 182
 183 -- Is the given core expression a variable or an application?
 184 is_simple :: CoreSyn.CoreExpr -> Bool
 185 is_simple (CoreSyn.App _ _) = True
 186 is_simple (CoreSyn.Var _) = True
 187 is_simple (CoreSyn.Cast expr _) = is_simple expr
 188 is_simple _ = False
 189
 190 -- Does the given CoreExpr have any free type vars?
 191 has_free_tyvars :: CoreSyn.CoreExpr -> Bool
 192 has_free_tyvars = not . VarSet.isEmptyVarSet . (CoreFVs.exprSomeFreeVars Var.isTyVar)
 193
 194 -- Does the given type have any free type vars?
 195 ty_has_free_tyvars :: Type.Type -> Bool
 196 ty_has_free_tyvars = not . VarSet.isEmptyVarSet . Type.tyVarsOfType
 197
 198 -- Does the given CoreExpr have any free local vars?
 199 has_free_vars :: CoreSyn.CoreExpr -> Bool
 200 has_free_vars = not . VarSet.isEmptyVarSet . CoreFVs.exprFreeVars
 201
 202 -- Does the given expression use any of the given binders?
 203 expr_uses_binders :: [CoreSyn.CoreBndr] -> CoreSyn.CoreExpr -> Bool
 204 expr_uses_binders bndrs = not . VarSet.isEmptyVarSet . (CoreFVs.exprSomeFreeVars (`elem` bndrs))
 205
 206 -- Turns a Var CoreExpr into the Id inside it. Will of course only work for
 207 -- simple Var CoreExprs, not complexer ones.
 208 exprToVar :: CoreSyn.CoreExpr -> Var.Id
 209 exprToVar (CoreSyn.Var id) = id
 210 exprToVar expr = error $ "\nCoreTools.exprToVar: Not a var: " ++ show expr
 211
 212 -- Turns a Lit CoreExpr into the Literal inside it.
 213 exprToLit :: CoreSyn.CoreExpr -> Literal.Literal
 214 exprToLit (CoreSyn.Lit lit) = lit
 215 exprToLit expr = error $ "\nCoreTools.exprToLit: Not a lit: " ++ show expr
 216
 217 -- Removes all the type and dictionary arguments from the given argument list,
 218 -- leaving only the normal value arguments. The type given is the type of the
 219 -- expression applied to this argument list.
 220 get_val_args :: Type.Type -> [CoreSyn.CoreExpr] -> [CoreSyn.CoreExpr]
 221 get_val_args ty args = drop n args
 222   where
 223     (tyvars, predtypes, _) = TcType.tcSplitSigmaTy ty
 224     -- The first (length tyvars) arguments should be types, the next
 225     -- (length predtypes) arguments should be dictionaries. We drop this many
 226     -- arguments, to get at the value arguments.
 227     n = length tyvars + length predtypes
 228
 229 getLiterals :: HscTypes.HscEnv -> CoreSyn.CoreExpr -> [CoreSyn.CoreExpr]
 230 getLiterals _ app@(CoreSyn.App _ _) = literals
 231   where
 232     (CoreSyn.Var f, args) = CoreSyn.collectArgs app
 233     literals = filter (is_lit) args
 234
 235 getLiterals _ lit@(CoreSyn.Lit _) = [lit]
 236
 237 getLiterals hscenv letrec@(CoreSyn.Let (CoreSyn.NonRec letBind (letExpr)) letRes) = [lit]
 238   where
 239     ty     = Var.varType letBind
 240     litInt = eval_tfp_int hscenv ty
 241     lit    = CoreSyn.Lit (Literal.mkMachInt (toInteger litInt))
 242
 243 getLiterals _ expr = error $ "\nCoreTools.getLiterals: Not a known Lit: " ++ pprString expr
 244
 245 reduceCoreListToHsList ::
 246   [HscTypes.CoreModule] -- ^ The modules where parts of the list are hidden
 247   -> CoreSyn.CoreExpr   -- ^ The refence to atleast one of the nodes
 248   -> TranslatorSession [CoreSyn.CoreExpr]
 249 reduceCoreListToHsList cores app@(CoreSyn.App _ _) = do {
 250   ; let { (fun, args) = CoreSyn.collectArgs app
 251         ; len         = length args
 252         } ;
 253   ; case len of
 254       3 -> do {
 255         ; let topelem = args!!1
 256         ; case (args!!2) of
 257             (varz@(CoreSyn.Var id)) -> do {
 258               ; binds <- mapM (findExpr (isVarName id)) cores
 259               ; otherelems <- reduceCoreListToHsList cores (head (Maybe.catMaybes binds))
 260               ; return (topelem:otherelems)
 261               }
 262             (appz@(CoreSyn.App _ _)) -> do {
 263               ; otherelems <- reduceCoreListToHsList cores appz
 264               ; return (topelem:otherelems)
 265               }
 266             otherwise -> return [topelem]
 267         }
 268       otherwise -> return []
 269   }
 270   where
 271     isVarName :: Monad m => Var.Var -> Var.Var -> m Bool
 272     isVarName lookfor bind = return $ (Var.varName lookfor) == (Var.varName bind)
 273
 274 reduceCoreListToHsList _ _ = return []
 275
 276 -- Is the given var the State data constructor?
 277 isStateCon :: Var.Var -> Bool
 278 isStateCon var =
 279   -- See if it is a DataConWrapId (not DataConWorkId, since State is a
 280   -- newtype).
 281   case Id.idDetails var of
 282     IdInfo.DataConWrapId dc ->
 283       -- See if the datacon is the State datacon from the State type.
 284       let tycon = DataCon.dataConTyCon dc
 285           tyname = Name.getOccString tycon
 286           dcname = Name.getOccString dc
 287       in case (tyname, dcname) of
 288         ("State", "State") -> True
 289         _ -> False
 290     _ -> False
 291
 292 -- | Is the given type a State type?
 293 isStateType :: Type.Type -> Bool
 294 -- Resolve any type synonyms remaining
 295 isStateType ty | Just ty' <- Type.tcView ty = isStateType ty'
 296 isStateType ty  = Maybe.isJust $ do
 297   -- Split the type. Don't use normal splitAppTy, since that looks through
 298   -- newtypes, and we want to see the State newtype.
 299   (typef, _) <- Type.repSplitAppTy_maybe ty
 300   -- See if the applied type is a type constructor
 301   (tycon, _) <- Type.splitTyConApp_maybe typef
 302   if TyCon.isNewTyCon tycon && Name.getOccString tycon == "State"
 303     then
 304       Just ()
 305     else
 306       Nothing
 307
 308 -- | Does the given TypedThing have a State type?
 309 hasStateType :: (TypedThing t) => t -> Bool
 310 hasStateType expr = case getType expr of
 311   Nothing -> False
 312   Just ty -> isStateType ty
 313
 314
 315 -- | Flattens nested lets into a single list of bindings. The expression
 316 --   passed does not have to be a let expression, if it isn't an empty list of
 317 --   bindings is returned.
 318 flattenLets ::
 319   CoreSyn.CoreExpr -- ^ The expression to flatten.
 320   -> ([Binding], CoreSyn.CoreExpr) -- ^ The bindings and resulting expression.
 321 flattenLets (CoreSyn.Let binds expr) =
 322   (bindings ++ bindings', expr')
 323   where
 324     -- Recursively flatten the contained expression
 325     (bindings', expr') =flattenLets expr
 326     -- Flatten our own bindings to remove the Rec / NonRec constructors
 327     bindings = CoreSyn.flattenBinds [binds]
 328 flattenLets expr = ([], expr)
 329
 330 -- | Create bunch of nested non-recursive let expressions from the given
 331 -- bindings. The first binding is bound at the highest level (and thus
 332 -- available in all other bindings).
 333 mkNonRecLets :: [Binding] -> CoreSyn.CoreExpr -> CoreSyn.CoreExpr
 334 mkNonRecLets bindings expr = MkCore.mkCoreLets binds expr
 335   where
 336     binds = map (uncurry CoreSyn.NonRec) bindings
 337
 338 -- | A class of things that (optionally) have a core Type. The type is
 339 -- optional, since Type expressions don't have a type themselves.
 340 class TypedThing t where
 341   getType :: t -> Maybe Type.Type
 342
 343 instance TypedThing CoreSyn.CoreExpr where
 344   getType (CoreSyn.Type _) = Nothing
 345   getType expr = Just $ CoreUtils.exprType expr
 346
 347 instance TypedThing CoreSyn.CoreBndr where
 348   getType = return . Id.idType
 349
 350 instance TypedThing Type.Type where
 351   getType = return . id
 352
 353 -- | Generate new uniques for all binders in the given expression.
 354 -- Does not support making type variables unique, though this could be
 355 -- supported if required (by passing a CoreSubst.Subst instead of VarEnv to
 356 -- genUniques' below).
 357 genUniques :: CoreSyn.CoreExpr -> TranslatorSession CoreSyn.CoreExpr
 358 genUniques = genUniques' VarEnv.emptyVarEnv
 359
 360 -- | A helper function to generate uniques, that takes a VarEnv containing the
 361 --   substitutions already performed.
 362 genUniques' :: VarEnv.VarEnv CoreSyn.CoreBndr -> CoreSyn.CoreExpr -> TranslatorSession CoreSyn.CoreExpr
 363 genUniques' subst (CoreSyn.Var f) = do
 364   -- Replace the binder with its new value, if applicable.
 365   let f' = VarEnv.lookupWithDefaultVarEnv subst f f
 366   return (CoreSyn.Var f')
 367 -- Leave literals untouched
 368 genUniques' subst (CoreSyn.Lit l) = return $ CoreSyn.Lit l
 369 genUniques' subst (CoreSyn.App f arg) = do
 370   -- Only work on subexpressions
 371   f' <- genUniques' subst f
 372   arg' <- genUniques' subst arg
 373   return (CoreSyn.App f' arg')
 374 -- Don't change type abstractions
 375 genUniques' subst expr@(CoreSyn.Lam bndr res) | CoreSyn.isTyVar bndr = return expr
 376 genUniques' subst (CoreSyn.Lam bndr res) = do
 377   -- Generate a new unique for the bound variable
 378   (subst', bndr') <- genUnique subst bndr
 379   res' <- genUniques' subst' res
 380   return (CoreSyn.Lam bndr' res')
 381 genUniques' subst (CoreSyn.Let (CoreSyn.NonRec bndr bound) res) = do
 382   -- Make the binders unique
 383   (subst', bndr') <- genUnique subst bndr
 384   bound' <- genUniques' subst' bound
 385   res' <- genUniques' subst' res
 386   return $ CoreSyn.Let (CoreSyn.NonRec bndr' bound') res'
 387 genUniques' subst (CoreSyn.Let (CoreSyn.Rec binds) res) = do
 388   -- Make each of the binders unique
 389   (subst', bndrs') <- mapAccumLM genUnique subst (map fst binds)
 390   bounds' <- mapM (genUniques' subst' . snd) binds
 391   res' <- genUniques' subst' res
 392   let binds' = zip bndrs' bounds'
 393   return $ CoreSyn.Let (CoreSyn.Rec binds') res'
 394 genUniques' subst (CoreSyn.Case scrut bndr ty alts) = do
 395   -- Process the scrutinee with the original substitution, since non of the
 396   -- binders bound in the Case statement is in scope in the scrutinee.
 397   scrut' <- genUniques' subst scrut
 398   -- Generate a new binder for the scrutinee
 399   (subst', bndr') <- genUnique subst bndr
 400   -- Process each of the alts
 401   alts' <- mapM (doalt subst') alts
 402   return $ CoreSyn.Case scrut' bndr' ty alts'
 403   where
 404     doalt subst (con, bndrs, expr) = do
 405       (subst', bndrs') <- mapAccumLM genUnique subst bndrs
 406       expr' <- genUniques' subst' expr
 407       -- Note that we don't return subst', since bndrs are only in scope in
 408       -- expr.
 409       return (con, bndrs', expr')
 410 genUniques' subst (CoreSyn.Cast expr coercion) = do
 411   expr' <- genUniques' subst expr
 412   -- Just process the casted expression
 413   return $ CoreSyn.Cast expr' coercion
 414 genUniques' subst (CoreSyn.Note note expr) = do
 415   expr' <- genUniques' subst expr
 416   -- Just process the annotated expression
 417   return $ CoreSyn.Note note expr'
 418 -- Leave types untouched
 419 genUniques' subst expr@(CoreSyn.Type _) = return expr
 420
 421 -- Generate a new unique for the given binder, and extend the given
 422 -- substitution to reflect this.
 423 genUnique :: VarEnv.VarEnv CoreSyn.CoreBndr -> CoreSyn.CoreBndr -> TranslatorSession (VarEnv.VarEnv CoreSyn.CoreBndr, CoreSyn.CoreBndr)
 424 genUnique subst bndr = do
 425   bndr' <- BinderTools.cloneVar bndr
 426   -- Replace all occurences of the old binder with a reference to the new
 427   -- binder.
 428   let subst' = VarEnv.extendVarEnv subst bndr bndr'
 429   return (subst', bndr')