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

   1 -- | This module provides a number of functions to find out things about Core
   2 -- programs. This module does not provide the actual plumbing to work with
   3 -- Core and Haskell (it uses HsTools for this), but only the functions that
   4 -- know about various libraries and know which functions to call.
   5 module CLasH.Utils.Core.CoreTools where
   6
   7 --Standard modules
   8 import qualified Maybe
   9 import System.IO.Unsafe
  10
  11 -- GHC API
  12 import qualified GHC
  13 import qualified Type
  14 import qualified TcType
  15 import qualified HsExpr
  16 import qualified HsTypes
  17 import qualified HsBinds
  18 import qualified HscTypes
  19 import qualified RdrName
  20 import qualified Name
  21 import qualified OccName
  22 import qualified TysWiredIn
  23 import qualified Bag
  24 import qualified DynFlags
  25 import qualified SrcLoc
  26 import qualified CoreSyn
  27 import qualified Var
  28 import qualified VarSet
  29 import qualified Unique
  30 import qualified CoreUtils
  31 import qualified CoreFVs
  32 import qualified Literal
  33
  34 -- Local imports
  35 import CLasH.Utils.GhcTools
  36 import CLasH.Utils.HsTools
  37 import CLasH.Utils.Pretty
  38
  39 -- | Evaluate a core Type representing type level int from the tfp
  40 -- library to a real int.
  41 eval_tfp_int :: HscTypes.HscEnv -> Type.Type -> Int
  42 eval_tfp_int env ty =
  43   unsafeRunGhc libdir $ do
  44     GHC.setSession env
  45     -- Automatically import modules for any fully qualified identifiers
  46     setDynFlag DynFlags.Opt_ImplicitImportQualified
  47
  48     let from_int_t_name = mkRdrName "Types.Data.Num" "fromIntegerT"
  49     let from_int_t = SrcLoc.noLoc $ HsExpr.HsVar from_int_t_name
  50     let undef = hsTypedUndef $ coreToHsType ty
  51     let app = SrcLoc.noLoc $ HsExpr.HsApp (from_int_t) (undef)
  52     let int_ty = SrcLoc.noLoc $ HsTypes.HsTyVar TysWiredIn.intTyCon_RDR
  53     let expr = HsExpr.ExprWithTySig app int_ty
  54     core <- toCore expr
  55     execCore core
  56   where
  57     libdir = DynFlags.topDir dynflags
  58     dynflags = HscTypes.hsc_dflags env
  59
  60 normalise_tfp_int :: HscTypes.HscEnv -> Type.Type -> Type.Type
  61 normalise_tfp_int env ty =
  62    unsafePerformIO $ do
  63      nty <- normaliseType env ty
  64      return nty
  65
  66 -- | Get the width of a SizedWord type
  67 -- sized_word_len :: HscTypes.HscEnv -> Type.Type -> Int
  68 -- sized_word_len env ty = eval_tfp_int env (sized_word_len_ty ty)
  69
  70 sized_word_len_ty :: Type.Type -> Type.Type
  71 sized_word_len_ty ty = len
  72   where
  73     args = case Type.splitTyConApp_maybe ty of
  74       Just (tycon, args) -> args
  75       Nothing -> error $ "\nCoreTools.sized_word_len_ty: Not a sized word type: " ++ (pprString ty)
  76     [len]         = args
  77
  78 -- | Get the width of a SizedInt type
  79 -- sized_int_len :: HscTypes.HscEnv -> Type.Type -> Int
  80 -- sized_int_len env ty = eval_tfp_int env (sized_int_len_ty ty)
  81
  82 sized_int_len_ty :: Type.Type -> Type.Type
  83 sized_int_len_ty ty = len
  84   where
  85     args = case Type.splitTyConApp_maybe ty of
  86       Just (tycon, args) -> args
  87       Nothing -> error $ "\nCoreTools.sized_int_len_ty: Not a sized int type: " ++ (pprString ty)
  88     [len]         = args
  89
  90 -- | Get the upperbound of a RangedWord type
  91 -- ranged_word_bound :: HscTypes.HscEnv -> Type.Type -> Int
  92 -- ranged_word_bound env ty = eval_tfp_int env (ranged_word_bound_ty ty)
  93
  94 ranged_word_bound_ty :: Type.Type -> Type.Type
  95 ranged_word_bound_ty ty = len
  96   where
  97     args = case Type.splitTyConApp_maybe ty of
  98       Just (tycon, args) -> args
  99       Nothing -> error $ "\nCoreTools.ranged_word_bound_ty: Not a sized word type: " ++ (pprString ty)
 100     [len]         = args
 101
 102 -- | Evaluate a core Type representing type level int from the TypeLevel
 103 -- library to a real int.
 104 -- eval_type_level_int :: Type.Type -> Int
 105 -- eval_type_level_int ty =
 106 --   unsafeRunGhc $ do
 107 --     -- Automatically import modules for any fully qualified identifiers
 108 --     setDynFlag DynFlags.Opt_ImplicitImportQualified
 109 --
 110 --     let to_int_name = mkRdrName "Data.TypeLevel.Num.Sets" "toInt"
 111 --     let to_int = SrcLoc.noLoc $ HsExpr.HsVar to_int_name
 112 --     let undef = hsTypedUndef $ coreToHsType ty
 113 --     let app = HsExpr.HsApp (to_int) (undef)
 114 --
 115 --     core <- toCore [] app
 116 --     execCore core
 117
 118 -- | Get the length of a FSVec type
 119 -- tfvec_len :: HscTypes.HscEnv -> Type.Type -> Int
 120 -- tfvec_len env ty = eval_tfp_int env (tfvec_len_ty ty)
 121
 122 tfvec_len_ty :: Type.Type -> Type.Type
 123 tfvec_len_ty ty = len
 124   where
 125     args = case Type.splitTyConApp_maybe ty of
 126       Just (tycon, args) -> args
 127       Nothing -> error $ "\nCoreTools.tfvec_len_ty: Not a vector type: " ++ (pprString ty)
 128     [len, el_ty] = args
 129
 130 -- | Get the element type of a TFVec type
 131 tfvec_elem :: Type.Type -> Type.Type
 132 tfvec_elem ty = el_ty
 133   where
 134     args = case Type.splitTyConApp_maybe ty of
 135       Just (tycon, args) -> args
 136       Nothing -> error $ "\nCoreTools.tfvec_len: Not a vector type: " ++ (pprString ty)
 137     [len, el_ty] = args
 138
 139 -- Is the given core expression a lambda abstraction?
 140 is_lam :: CoreSyn.CoreExpr -> Bool
 141 is_lam (CoreSyn.Lam _ _) = True
 142 is_lam _ = False
 143
 144 -- Is the given core expression of a function type?
 145 is_fun :: CoreSyn.CoreExpr -> Bool
 146 -- Treat Type arguments differently, because exprType is not defined for them.
 147 is_fun (CoreSyn.Type _) = False
 148 is_fun expr = (Type.isFunTy . CoreUtils.exprType) expr
 149
 150 -- Is the given core expression polymorphic (i.e., does it accept type
 151 -- arguments?).
 152 is_poly :: CoreSyn.CoreExpr -> Bool
 153 -- Treat Type arguments differently, because exprType is not defined for them.
 154 is_poly (CoreSyn.Type _) = False
 155 is_poly expr = (Maybe.isJust . Type.splitForAllTy_maybe . CoreUtils.exprType) expr
 156
 157 -- Is the given core expression a variable reference?
 158 is_var :: CoreSyn.CoreExpr -> Bool
 159 is_var (CoreSyn.Var _) = True
 160 is_var _ = False
 161
 162 is_lit :: CoreSyn.CoreExpr -> Bool
 163 is_lit (CoreSyn.Lit _) = True
 164 is_lit _ = False
 165
 166 -- Can the given core expression be applied to something? This is true for
 167 -- applying to a value as well as a type.
 168 is_applicable :: CoreSyn.CoreExpr -> Bool
 169 is_applicable expr = is_fun expr || is_poly expr
 170
 171 -- Is the given core expression a variable or an application?
 172 is_simple :: CoreSyn.CoreExpr -> Bool
 173 is_simple (CoreSyn.App _ _) = True
 174 is_simple (CoreSyn.Var _) = True
 175 is_simple (CoreSyn.Cast expr _) = is_simple expr
 176 is_simple _ = False
 177
 178 -- Does the given CoreExpr have any free type vars?
 179 has_free_tyvars :: CoreSyn.CoreExpr -> Bool
 180 has_free_tyvars = not . VarSet.isEmptyVarSet . (CoreFVs.exprSomeFreeVars Var.isTyVar)
 181
 182 -- Does the given CoreExpr have any free local vars?
 183 has_free_vars :: CoreSyn.CoreExpr -> Bool
 184 has_free_vars = not . VarSet.isEmptyVarSet . CoreFVs.exprFreeVars
 185
 186 -- Turns a Var CoreExpr into the Id inside it. Will of course only work for
 187 -- simple Var CoreExprs, not complexer ones.
 188 exprToVar :: CoreSyn.CoreExpr -> Var.Id
 189 exprToVar (CoreSyn.Var id) = id
 190 exprToVar expr = error $ "\nCoreTools.exprToVar: Not a var: " ++ show expr
 191
 192 -- Turns a Lit CoreExpr into the Literal inside it.
 193 exprToLit :: CoreSyn.CoreExpr -> Literal.Literal
 194 exprToLit (CoreSyn.Lit lit) = lit
 195 exprToLit expr = error $ "\nCoreTools.exprToLit: Not a lit: " ++ show expr
 196
 197 -- Removes all the type and dictionary arguments from the given argument list,
 198 -- leaving only the normal value arguments. The type given is the type of the
 199 -- expression applied to this argument list.
 200 get_val_args :: Type.Type -> [CoreSyn.CoreExpr] -> [CoreSyn.CoreExpr]
 201 get_val_args ty args = drop n args
 202   where
 203     (tyvars, predtypes, _) = TcType.tcSplitSigmaTy ty
 204     -- The first (length tyvars) arguments should be types, the next
 205     -- (length predtypes) arguments should be dictionaries. We drop this many
 206     -- arguments, to get at the value arguments.
 207     n = length tyvars + length predtypes
 208
 209 getLiterals :: CoreSyn.CoreExpr -> [CoreSyn.CoreExpr]
 210 getLiterals app@(CoreSyn.App _ _) = literals
 211   where
 212     (CoreSyn.Var f, args) = CoreSyn.collectArgs app
 213     literals = filter (is_lit) args