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

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