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 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 HsBinds
  19 import qualified HscTypes
  20 import qualified RdrName
  21 import qualified Name
  22 import qualified OccName
  23 import qualified Type
  24 import qualified Id
  25 import qualified TyCon
  26 import qualified TysWiredIn
  27 import qualified Bag
  28 import qualified DynFlags
  29 import qualified SrcLoc
  30 import qualified CoreSyn
  31 import qualified Var
  32 import qualified VarSet
  33 import qualified Unique
  34 import qualified CoreUtils
  35 import qualified CoreFVs
  36 import qualified Literal
  37
  38 -- Local imports
  39 import CLasH.Utils.GhcTools
  40 import CLasH.Utils.HsTools
  41 import CLasH.Utils.Pretty
  42
  43 -- | Evaluate a core Type representing type level int from the tfp
  44 -- library to a real int.
  45 eval_tfp_int :: HscTypes.HscEnv -> Type.Type -> Int
  46 eval_tfp_int env ty =
  47   unsafeRunGhc libdir $ do
  48     GHC.setSession env
  49     -- Automatically import modules for any fully qualified identifiers
  50     setDynFlag DynFlags.Opt_ImplicitImportQualified
  51
  52     let from_int_t_name = mkRdrName "Types.Data.Num.Ops" "fromIntegerT"
  53     let from_int_t = SrcLoc.noLoc $ HsExpr.HsVar from_int_t_name
  54     let undef = hsTypedUndef $ coreToHsType ty
  55     let app = SrcLoc.noLoc $ HsExpr.HsApp (from_int_t) (undef)
  56     let int_ty = SrcLoc.noLoc $ HsTypes.HsTyVar TysWiredIn.intTyCon_RDR
  57     let expr = HsExpr.ExprWithTySig app int_ty
  58     core <- toCore expr
  59     execCore core
  60   where
  61     libdir = DynFlags.topDir dynflags
  62     dynflags = HscTypes.hsc_dflags env
  63
  64 normalise_tfp_int :: HscTypes.HscEnv -> Type.Type -> Type.Type
  65 normalise_tfp_int env ty =
  66    unsafePerformIO $ do
  67      nty <- normaliseType env ty
  68      return nty
  69
  70 -- | Get the width of a SizedWord type
  71 -- sized_word_len :: HscTypes.HscEnv -> Type.Type -> Int
  72 -- sized_word_len env ty = eval_tfp_int env (sized_word_len_ty ty)
  73
  74 sized_word_len_ty :: Type.Type -> Type.Type
  75 sized_word_len_ty ty = len
  76   where
  77     args = case Type.splitTyConApp_maybe ty of
  78       Just (tycon, args) -> args
  79       Nothing -> error $ "\nCoreTools.sized_word_len_ty: Not a sized word type: " ++ (pprString ty)
  80     [len]         = args
  81
  82 -- | Get the width of a SizedInt type
  83 -- sized_int_len :: HscTypes.HscEnv -> Type.Type -> Int
  84 -- sized_int_len env ty = eval_tfp_int env (sized_int_len_ty ty)
  85
  86 sized_int_len_ty :: Type.Type -> Type.Type
  87 sized_int_len_ty ty = len
  88   where
  89     args = case Type.splitTyConApp_maybe ty of
  90       Just (tycon, args) -> args
  91       Nothing -> error $ "\nCoreTools.sized_int_len_ty: Not a sized int type: " ++ (pprString ty)
  92     [len]         = args
  93
  94 -- | Get the upperbound of a RangedWord type
  95 -- ranged_word_bound :: HscTypes.HscEnv -> Type.Type -> Int
  96 -- ranged_word_bound env ty = eval_tfp_int env (ranged_word_bound_ty ty)
  97
  98 ranged_word_bound_ty :: Type.Type -> Type.Type
  99 ranged_word_bound_ty ty = len
 100   where
 101     args = case Type.splitTyConApp_maybe ty of
 102       Just (tycon, args) -> args
 103       Nothing -> error $ "\nCoreTools.ranged_word_bound_ty: Not a sized word type: " ++ (pprString ty)
 104     [len]         = args
 105
 106 -- | Evaluate a core Type representing type level int from the TypeLevel
 107 -- library to a real int.
 108 -- eval_type_level_int :: Type.Type -> Int
 109 -- eval_type_level_int ty =
 110 --   unsafeRunGhc $ do
 111 --     -- Automatically import modules for any fully qualified identifiers
 112 --     setDynFlag DynFlags.Opt_ImplicitImportQualified
 113 --
 114 --     let to_int_name = mkRdrName "Data.TypeLevel.Num.Sets" "toInt"
 115 --     let to_int = SrcLoc.noLoc $ HsExpr.HsVar to_int_name
 116 --     let undef = hsTypedUndef $ coreToHsType ty
 117 --     let app = HsExpr.HsApp (to_int) (undef)
 118 --
 119 --     core <- toCore [] app
 120 --     execCore core
 121
 122 -- | Get the length of a FSVec type
 123 -- tfvec_len :: HscTypes.HscEnv -> Type.Type -> Int
 124 -- tfvec_len env ty = eval_tfp_int env (tfvec_len_ty ty)
 125
 126 tfvec_len_ty :: Type.Type -> Type.Type
 127 tfvec_len_ty ty = len
 128   where
 129     args = case Type.splitTyConApp_maybe ty of
 130       Just (tycon, args) -> args
 131       Nothing -> error $ "\nCoreTools.tfvec_len_ty: Not a vector type: " ++ (pprString ty)
 132     [len, el_ty] = args
 133
 134 -- | Get the element type of a TFVec type
 135 tfvec_elem :: Type.Type -> Type.Type
 136 tfvec_elem ty = el_ty
 137   where
 138     args = case Type.splitTyConApp_maybe ty of
 139       Just (tycon, args) -> args
 140       Nothing -> error $ "\nCoreTools.tfvec_len: Not a vector type: " ++ (pprString ty)
 141     [len, el_ty] = args
 142
 143 -- Is the given core expression a lambda abstraction?
 144 is_lam :: CoreSyn.CoreExpr -> Bool
 145 is_lam (CoreSyn.Lam _ _) = True
 146 is_lam _ = False
 147
 148 -- Is the given core expression of a function type?
 149 is_fun :: CoreSyn.CoreExpr -> Bool
 150 -- Treat Type arguments differently, because exprType is not defined for them.
 151 is_fun (CoreSyn.Type _) = False
 152 is_fun expr = (Type.isFunTy . CoreUtils.exprType) expr
 153
 154 -- Is the given core expression polymorphic (i.e., does it accept type
 155 -- arguments?).
 156 is_poly :: CoreSyn.CoreExpr -> Bool
 157 -- Treat Type arguments differently, because exprType is not defined for them.
 158 is_poly (CoreSyn.Type _) = False
 159 is_poly expr = (Maybe.isJust . Type.splitForAllTy_maybe . CoreUtils.exprType) expr
 160
 161 -- Is the given core expression a variable reference?
 162 is_var :: CoreSyn.CoreExpr -> Bool
 163 is_var (CoreSyn.Var _) = True
 164 is_var _ = False
 165
 166 is_lit :: CoreSyn.CoreExpr -> Bool
 167 is_lit (CoreSyn.Lit _) = True
 168 is_lit _ = False
 169
 170 -- Can the given core expression be applied to something? This is true for
 171 -- applying to a value as well as a type.
 172 is_applicable :: CoreSyn.CoreExpr -> Bool
 173 is_applicable expr = is_fun expr || is_poly expr
 174
 175 -- Is the given core expression a variable or an application?
 176 is_simple :: CoreSyn.CoreExpr -> Bool
 177 is_simple (CoreSyn.App _ _) = True
 178 is_simple (CoreSyn.Var _) = True
 179 is_simple (CoreSyn.Cast expr _) = is_simple expr
 180 is_simple _ = False
 181
 182 -- Does the given CoreExpr have any free type vars?
 183 has_free_tyvars :: CoreSyn.CoreExpr -> Bool
 184 has_free_tyvars = not . VarSet.isEmptyVarSet . (CoreFVs.exprSomeFreeVars Var.isTyVar)
 185
 186 -- Does the given CoreExpr have any free local vars?
 187 has_free_vars :: CoreSyn.CoreExpr -> Bool
 188 has_free_vars = not . VarSet.isEmptyVarSet . CoreFVs.exprFreeVars
 189
 190 -- Does the given expression use any of the given binders?
 191 expr_uses_binders :: [CoreSyn.CoreBndr] -> CoreSyn.CoreExpr -> Bool
 192 expr_uses_binders bndrs = not . VarSet.isEmptyVarSet . (CoreFVs.exprSomeFreeVars (`elem` bndrs))
 193
 194 -- Turns a Var CoreExpr into the Id inside it. Will of course only work for
 195 -- simple Var CoreExprs, not complexer ones.
 196 exprToVar :: CoreSyn.CoreExpr -> Var.Id
 197 exprToVar (CoreSyn.Var id) = id
 198 exprToVar expr = error $ "\nCoreTools.exprToVar: Not a var: " ++ show expr
 199
 200 -- Turns a Lit CoreExpr into the Literal inside it.
 201 exprToLit :: CoreSyn.CoreExpr -> Literal.Literal
 202 exprToLit (CoreSyn.Lit lit) = lit
 203 exprToLit expr = error $ "\nCoreTools.exprToLit: Not a lit: " ++ show expr
 204
 205 -- Removes all the type and dictionary arguments from the given argument list,
 206 -- leaving only the normal value arguments. The type given is the type of the
 207 -- expression applied to this argument list.
 208 get_val_args :: Type.Type -> [CoreSyn.CoreExpr] -> [CoreSyn.CoreExpr]
 209 get_val_args ty args = drop n args
 210   where
 211     (tyvars, predtypes, _) = TcType.tcSplitSigmaTy ty
 212     -- The first (length tyvars) arguments should be types, the next
 213     -- (length predtypes) arguments should be dictionaries. We drop this many
 214     -- arguments, to get at the value arguments.
 215     n = length tyvars + length predtypes
 216
 217 getLiterals :: CoreSyn.CoreExpr -> [CoreSyn.CoreExpr]
 218 getLiterals app@(CoreSyn.App _ _) = literals
 219   where
 220     (CoreSyn.Var f, args) = CoreSyn.collectArgs app
 221     literals = filter (is_lit) args
 222
 223 getLiterals lit@(CoreSyn.Lit _) = [lit]
 224
 225 reduceCoreListToHsList :: CoreSyn.CoreExpr -> [CoreSyn.CoreExpr]
 226 reduceCoreListToHsList app@(CoreSyn.App _ _) = out
 227   where
 228     (fun, args) = CoreSyn.collectArgs app
 229     len = length args
 230     out = case len of
 231           3 -> ((args!!1) : (reduceCoreListToHsList (args!!2)))
 232           otherwise -> []
 233
 234 reduceCoreListToHsList _ = []
 235
 236 -- | Is the given type a State type?
 237 isStateType :: Type.Type -> Bool
 238 -- Resolve any type synonyms remaining
 239 isStateType ty | Just ty' <- Type.tcView ty = isStateType ty'
 240 isStateType ty  = Maybe.isJust $ do
 241   -- Split the type. Don't use normal splitAppTy, since that looks through
 242   -- newtypes, and we want to see the State newtype.
 243   (typef, _) <- Type.repSplitAppTy_maybe ty
 244   -- See if the applied type is a type constructor
 245   (tycon, _) <- Type.splitTyConApp_maybe typef
 246   if TyCon.isNewTyCon tycon && Name.getOccString tycon == "State"
 247     then
 248       Just ()
 249     else
 250       Nothing
 251
 252 -- | Does the given TypedThing have a State type?
 253 hasStateType :: (TypedThing t) => t -> Bool
 254 hasStateType expr = case getType expr of
 255   Nothing -> False
 256   Just ty -> isStateType ty
 257
 258
 259 -- | A class of things that (optionally) have a core Type. The type is
 260 -- optional, since Type expressions don't have a type themselves.
 261 class TypedThing t where
 262   getType :: t -> Maybe Type.Type
 263
 264 instance TypedThing CoreSyn.CoreExpr where
 265   getType (CoreSyn.Type _) = Nothing
 266   getType expr = Just $ CoreUtils.exprType expr
 267
 268 instance TypedThing CoreSyn.CoreBndr where
 269   getType = return . Id.idType
 270
 271 instance TypedThing Type.Type where
 272   getType = return . id