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 -- Turns a Var CoreExpr into the Id inside it. Will of course only work for
 190 -- simple Var CoreExprs, not complexer ones.
 191 exprToVar :: CoreSyn.CoreExpr -> Var.Id
 192 exprToVar (CoreSyn.Var id) = id
 193 exprToVar expr = error $ "\nCoreTools.exprToVar: Not a var: " ++ show expr
 194
 195 -- Turns a Lit CoreExpr into the Literal inside it.
 196 exprToLit :: CoreSyn.CoreExpr -> Literal.Literal
 197 exprToLit (CoreSyn.Lit lit) = lit
 198 exprToLit expr = error $ "\nCoreTools.exprToLit: Not a lit: " ++ show expr
 199
 200 -- Removes all the type and dictionary arguments from the given argument list,
 201 -- leaving only the normal value arguments. The type given is the type of the
 202 -- expression applied to this argument list.
 203 get_val_args :: Type.Type -> [CoreSyn.CoreExpr] -> [CoreSyn.CoreExpr]
 204 get_val_args ty args = drop n args
 205   where
 206     (tyvars, predtypes, _) = TcType.tcSplitSigmaTy ty
 207     -- The first (length tyvars) arguments should be types, the next
 208     -- (length predtypes) arguments should be dictionaries. We drop this many
 209     -- arguments, to get at the value arguments.
 210     n = length tyvars + length predtypes
 211
 212 getLiterals :: CoreSyn.CoreExpr -> [CoreSyn.CoreExpr]
 213 getLiterals app@(CoreSyn.App _ _) = literals
 214   where
 215     (CoreSyn.Var f, args) = CoreSyn.collectArgs app
 216     literals = filter (is_lit) args
 217
 218 reduceCoreListToHsList :: CoreSyn.CoreExpr -> [CoreSyn.CoreExpr]
 219 reduceCoreListToHsList app@(CoreSyn.App _ _) = out
 220   where
 221     (fun, args) = CoreSyn.collectArgs app
 222     len = length args
 223     out = case len of
 224           3 -> ((args!!1) : (reduceCoreListToHsList (args!!2)))
 225           otherwise -> []
 226
 227 reduceCoreListToHsList _ = []
 228
 229 -- | Is the given type a State type?
 230 isStateType :: Type.Type -> Bool
 231 -- Resolve any type synonyms remaining
 232 isStateType ty | Just ty' <- Type.tcView ty = isStateType ty'
 233 isStateType ty  = Maybe.isJust $ do
 234   -- Split the type. Don't use normal splitAppTy, since that looks through
 235   -- newtypes, and we want to see the State newtype.
 236   (typef, _) <- Type.repSplitAppTy_maybe ty
 237   -- See if the applied type is a type constructor
 238   (tycon, _) <- Type.splitTyConApp_maybe typef
 239   if TyCon.isNewTyCon tycon && Name.getOccString tycon == "State"
 240     then
 241       Just ()
 242     else
 243       Nothing
 244
 245 -- | Does the given expression have a State type?
 246 hasStateType :: CoreSyn.CoreExpr -> Bool
 247 hasStateType (CoreSyn.Type _) = False
 248 hasStateType expr = (isStateType . CoreUtils.exprType) expr