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 DataCon
  27 import qualified TysWiredIn
  28 import qualified Bag
  29 import qualified DynFlags
  30 import qualified SrcLoc
  31 import qualified CoreSyn
  32 import qualified Var
  33 import qualified IdInfo
  34 import qualified VarSet
  35 import qualified Unique
  36 import qualified CoreUtils
  37 import qualified CoreFVs
  38 import qualified Literal
  39
  40 -- Local imports
  41 import CLasH.Translator.TranslatorTypes
  42 import CLasH.Utils.GhcTools
  43 import CLasH.Utils.HsTools
  44 import CLasH.Utils.Pretty
  45
  46 -- | Evaluate a core Type representing type level int from the tfp
  47 -- library to a real int.
  48 eval_tfp_int :: HscTypes.HscEnv -> Type.Type -> Int
  49 eval_tfp_int env ty =
  50   unsafeRunGhc libdir $ do
  51     GHC.setSession env
  52     -- Automatically import modules for any fully qualified identifiers
  53     setDynFlag DynFlags.Opt_ImplicitImportQualified
  54
  55     let from_int_t_name = mkRdrName "Types.Data.Num.Ops" "fromIntegerT"
  56     let from_int_t = SrcLoc.noLoc $ HsExpr.HsVar from_int_t_name
  57     let undef = hsTypedUndef $ coreToHsType ty
  58     let app = SrcLoc.noLoc $ HsExpr.HsApp (from_int_t) (undef)
  59     let int_ty = SrcLoc.noLoc $ HsTypes.HsTyVar TysWiredIn.intTyCon_RDR
  60     let expr = HsExpr.ExprWithTySig app int_ty
  61     core <- toCore expr
  62     execCore core
  63   where
  64     libdir = DynFlags.topDir dynflags
  65     dynflags = HscTypes.hsc_dflags env
  66
  67 normalise_tfp_int :: HscTypes.HscEnv -> Type.Type -> Type.Type
  68 normalise_tfp_int env ty =
  69    unsafePerformIO $ do
  70      nty <- normaliseType env ty
  71      return nty
  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 of a function type?
 152 is_fun :: CoreSyn.CoreExpr -> Bool
 153 -- Treat Type arguments differently, because exprType is not defined for them.
 154 is_fun (CoreSyn.Type _) = False
 155 is_fun expr = (Type.isFunTy . CoreUtils.exprType) expr
 156
 157 -- Is the given core expression polymorphic (i.e., does it accept type
 158 -- arguments?).
 159 is_poly :: CoreSyn.CoreExpr -> Bool
 160 -- Treat Type arguments differently, because exprType is not defined for them.
 161 is_poly (CoreSyn.Type _) = False
 162 is_poly expr = (Maybe.isJust . Type.splitForAllTy_maybe . CoreUtils.exprType) expr
 163
 164 -- Is the given core expression a variable reference?
 165 is_var :: CoreSyn.CoreExpr -> Bool
 166 is_var (CoreSyn.Var _) = True
 167 is_var _ = False
 168
 169 is_lit :: CoreSyn.CoreExpr -> Bool
 170 is_lit (CoreSyn.Lit _) = True
 171 is_lit _ = False
 172
 173 -- Can the given core expression be applied to something? This is true for
 174 -- applying to a value as well as a type.
 175 is_applicable :: CoreSyn.CoreExpr -> Bool
 176 is_applicable expr = is_fun expr || is_poly expr
 177
 178 -- Is the given core expression a variable or an application?
 179 is_simple :: CoreSyn.CoreExpr -> Bool
 180 is_simple (CoreSyn.App _ _) = True
 181 is_simple (CoreSyn.Var _) = True
 182 is_simple (CoreSyn.Cast expr _) = is_simple expr
 183 is_simple _ = False
 184
 185 -- Does the given CoreExpr have any free type vars?
 186 has_free_tyvars :: CoreSyn.CoreExpr -> Bool
 187 has_free_tyvars = not . VarSet.isEmptyVarSet . (CoreFVs.exprSomeFreeVars Var.isTyVar)
 188
 189 -- Does the given CoreExpr have any free local vars?
 190 has_free_vars :: CoreSyn.CoreExpr -> Bool
 191 has_free_vars = not . VarSet.isEmptyVarSet . CoreFVs.exprFreeVars
 192
 193 -- Does the given expression use any of the given binders?
 194 expr_uses_binders :: [CoreSyn.CoreBndr] -> CoreSyn.CoreExpr -> Bool
 195 expr_uses_binders bndrs = not . VarSet.isEmptyVarSet . (CoreFVs.exprSomeFreeVars (`elem` bndrs))
 196
 197 -- Turns a Var CoreExpr into the Id inside it. Will of course only work for
 198 -- simple Var CoreExprs, not complexer ones.
 199 exprToVar :: CoreSyn.CoreExpr -> Var.Id
 200 exprToVar (CoreSyn.Var id) = id
 201 exprToVar expr = error $ "\nCoreTools.exprToVar: Not a var: " ++ show expr
 202
 203 -- Turns a Lit CoreExpr into the Literal inside it.
 204 exprToLit :: CoreSyn.CoreExpr -> Literal.Literal
 205 exprToLit (CoreSyn.Lit lit) = lit
 206 exprToLit expr = error $ "\nCoreTools.exprToLit: Not a lit: " ++ show expr
 207
 208 -- Removes all the type and dictionary arguments from the given argument list,
 209 -- leaving only the normal value arguments. The type given is the type of the
 210 -- expression applied to this argument list.
 211 get_val_args :: Type.Type -> [CoreSyn.CoreExpr] -> [CoreSyn.CoreExpr]
 212 get_val_args ty args = drop n args
 213   where
 214     (tyvars, predtypes, _) = TcType.tcSplitSigmaTy ty
 215     -- The first (length tyvars) arguments should be types, the next
 216     -- (length predtypes) arguments should be dictionaries. We drop this many
 217     -- arguments, to get at the value arguments.
 218     n = length tyvars + length predtypes
 219
 220 getLiterals :: CoreSyn.CoreExpr -> [CoreSyn.CoreExpr]
 221 getLiterals app@(CoreSyn.App _ _) = literals
 222   where
 223     (CoreSyn.Var f, args) = CoreSyn.collectArgs app
 224     literals = filter (is_lit) args
 225
 226 getLiterals lit@(CoreSyn.Lit _) = [lit]
 227
 228 reduceCoreListToHsList ::
 229   [HscTypes.CoreModule] -- ^ The modules where parts of the list are hidden
 230   -> CoreSyn.CoreExpr   -- ^ The refence to atleast one of the nodes
 231   -> TranslatorSession [CoreSyn.CoreExpr]
 232 reduceCoreListToHsList cores app@(CoreSyn.App _ _) = do {
 233   ; let { (fun, args) = CoreSyn.collectArgs app
 234         ; len         = length args
 235         } ;
 236   ; case len of
 237       3 -> do {
 238         ; let topelem = args!!1
 239         ; case (args!!2) of
 240             (varz@(CoreSyn.Var id)) -> do {
 241               ; binds <- mapM (findExpr (isVarName id)) cores
 242               ; otherelems <- reduceCoreListToHsList cores (head (Maybe.catMaybes binds))
 243               ; return (topelem:otherelems)
 244               }
 245             (appz@(CoreSyn.App _ _)) -> do {
 246               ; otherelems <- reduceCoreListToHsList cores appz
 247               ; return (topelem:otherelems)
 248               }
 249             otherwise -> return [topelem]
 250         }
 251       otherwise -> return []
 252   }
 253   where
 254     isVarName :: Monad m => Var.Var -> Var.Var -> m Bool
 255     isVarName lookfor bind = return $ (Var.varName lookfor) == (Var.varName bind)
 256
 257 reduceCoreListToHsList _ _ = return []
 258
 259 -- Is the given var the State data constructor?
 260 isStateCon :: Var.Var -> Bool
 261 isStateCon var = do
 262   -- See if it is a DataConWrapId (not DataConWorkId, since State is a
 263   -- newtype).
 264   case Id.idDetails var of
 265     IdInfo.DataConWrapId dc ->
 266       -- See if the datacon is the State datacon from the State type.
 267       let tycon = DataCon.dataConTyCon dc
 268           tyname = Name.getOccString tycon
 269           dcname = Name.getOccString dc
 270       in case (tyname, dcname) of
 271         ("State", "State") -> True
 272         _ -> False
 273     _ -> False
 274
 275 -- | Is the given type a State type?
 276 isStateType :: Type.Type -> Bool
 277 -- Resolve any type synonyms remaining
 278 isStateType ty | Just ty' <- Type.tcView ty = isStateType ty'
 279 isStateType ty  = Maybe.isJust $ do
 280   -- Split the type. Don't use normal splitAppTy, since that looks through
 281   -- newtypes, and we want to see the State newtype.
 282   (typef, _) <- Type.repSplitAppTy_maybe ty
 283   -- See if the applied type is a type constructor
 284   (tycon, _) <- Type.splitTyConApp_maybe typef
 285   if TyCon.isNewTyCon tycon && Name.getOccString tycon == "State"
 286     then
 287       Just ()
 288     else
 289       Nothing
 290
 291 -- | Does the given TypedThing have a State type?
 292 hasStateType :: (TypedThing t) => t -> Bool
 293 hasStateType expr = case getType expr of
 294   Nothing -> False
 295   Just ty -> isStateType ty
 296
 297
 298 -- | A class of things that (optionally) have a core Type. The type is
 299 -- optional, since Type expressions don't have a type themselves.
 300 class TypedThing t where
 301   getType :: t -> Maybe Type.Type
 302
 303 instance TypedThing CoreSyn.CoreExpr where
 304   getType (CoreSyn.Type _) = Nothing
 305   getType expr = Just $ CoreUtils.exprType expr
 306
 307 instance TypedThing CoreSyn.CoreBndr where
 308   getType = return . Id.idType
 309
 310 instance TypedThing Type.Type where
 311   getType = return . id