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 import qualified MkCore
  40
  41 -- Local imports
  42 import CLasH.Translator.TranslatorTypes
  43 import CLasH.Utils.GhcTools
  44 import CLasH.Utils.HsTools
  45 import CLasH.Utils.Pretty
  46
  47 -- | A single binding, used as a shortcut to simplify type signatures.
  48 type Binding = (CoreSyn.CoreBndr, CoreSyn.CoreExpr)
  49
  50 -- | Evaluate a core Type representing type level int from the tfp
  51 -- library to a real int.
  52 eval_tfp_int :: HscTypes.HscEnv -> Type.Type -> Int
  53 eval_tfp_int env ty =
  54   unsafeRunGhc libdir $ do
  55     GHC.setSession env
  56     -- Automatically import modules for any fully qualified identifiers
  57     setDynFlag DynFlags.Opt_ImplicitImportQualified
  58
  59     let from_int_t_name = mkRdrName "Types.Data.Num.Ops" "fromIntegerT"
  60     let from_int_t = SrcLoc.noLoc $ HsExpr.HsVar from_int_t_name
  61     let undef = hsTypedUndef $ coreToHsType ty
  62     let app = SrcLoc.noLoc $ HsExpr.HsApp (from_int_t) (undef)
  63     let int_ty = SrcLoc.noLoc $ HsTypes.HsTyVar TysWiredIn.intTyCon_RDR
  64     let expr = HsExpr.ExprWithTySig app int_ty
  65     core <- toCore expr
  66     execCore core
  67   where
  68     libdir = DynFlags.topDir dynflags
  69     dynflags = HscTypes.hsc_dflags env
  70
  71 normalise_tfp_int :: HscTypes.HscEnv -> Type.Type -> Type.Type
  72 normalise_tfp_int env ty =
  73    unsafePerformIO $ do
  74      nty <- normaliseType env ty
  75      return nty
  76
  77 -- | Get the width of a SizedWord type
  78 -- sized_word_len :: HscTypes.HscEnv -> Type.Type -> Int
  79 -- sized_word_len env ty = eval_tfp_int env (sized_word_len_ty ty)
  80
  81 sized_word_len_ty :: Type.Type -> Type.Type
  82 sized_word_len_ty ty = len
  83   where
  84     args = case Type.splitTyConApp_maybe ty of
  85       Just (tycon, args) -> args
  86       Nothing -> error $ "\nCoreTools.sized_word_len_ty: Not a sized word type: " ++ (pprString ty)
  87     [len]         = args
  88
  89 -- | Get the width of a SizedInt type
  90 -- sized_int_len :: HscTypes.HscEnv -> Type.Type -> Int
  91 -- sized_int_len env ty = eval_tfp_int env (sized_int_len_ty ty)
  92
  93 sized_int_len_ty :: Type.Type -> Type.Type
  94 sized_int_len_ty ty = len
  95   where
  96     args = case Type.splitTyConApp_maybe ty of
  97       Just (tycon, args) -> args
  98       Nothing -> error $ "\nCoreTools.sized_int_len_ty: Not a sized int type: " ++ (pprString ty)
  99     [len]         = args
 100
 101 -- | Get the upperbound of a RangedWord type
 102 -- ranged_word_bound :: HscTypes.HscEnv -> Type.Type -> Int
 103 -- ranged_word_bound env ty = eval_tfp_int env (ranged_word_bound_ty ty)
 104
 105 ranged_word_bound_ty :: Type.Type -> Type.Type
 106 ranged_word_bound_ty ty = len
 107   where
 108     args = case Type.splitTyConApp_maybe ty of
 109       Just (tycon, args) -> args
 110       Nothing -> error $ "\nCoreTools.ranged_word_bound_ty: Not a sized word type: " ++ (pprString ty)
 111     [len]         = args
 112
 113 -- | Evaluate a core Type representing type level int from the TypeLevel
 114 -- library to a real int.
 115 -- eval_type_level_int :: Type.Type -> Int
 116 -- eval_type_level_int ty =
 117 --   unsafeRunGhc $ do
 118 --     -- Automatically import modules for any fully qualified identifiers
 119 --     setDynFlag DynFlags.Opt_ImplicitImportQualified
 120 --
 121 --     let to_int_name = mkRdrName "Data.TypeLevel.Num.Sets" "toInt"
 122 --     let to_int = SrcLoc.noLoc $ HsExpr.HsVar to_int_name
 123 --     let undef = hsTypedUndef $ coreToHsType ty
 124 --     let app = HsExpr.HsApp (to_int) (undef)
 125 --
 126 --     core <- toCore [] app
 127 --     execCore core
 128
 129 -- | Get the length of a FSVec type
 130 -- tfvec_len :: HscTypes.HscEnv -> Type.Type -> Int
 131 -- tfvec_len env ty = eval_tfp_int env (tfvec_len_ty ty)
 132
 133 tfvec_len_ty :: Type.Type -> Type.Type
 134 tfvec_len_ty ty = len
 135   where
 136     args = case Type.splitTyConApp_maybe ty of
 137       Just (tycon, args) -> args
 138       Nothing -> error $ "\nCoreTools.tfvec_len_ty: Not a vector type: " ++ (pprString ty)
 139     [len, el_ty] = args
 140
 141 -- | Get the element type of a TFVec type
 142 tfvec_elem :: Type.Type -> Type.Type
 143 tfvec_elem ty = el_ty
 144   where
 145     args = case Type.splitTyConApp_maybe ty of
 146       Just (tycon, args) -> args
 147       Nothing -> error $ "\nCoreTools.tfvec_len: Not a vector type: " ++ (pprString ty)
 148     [len, el_ty] = args
 149
 150 -- Is the given core expression a lambda abstraction?
 151 is_lam :: CoreSyn.CoreExpr -> Bool
 152 is_lam (CoreSyn.Lam _ _) = True
 153 is_lam _ = False
 154
 155 -- Is the given core expression of a function type?
 156 is_fun :: CoreSyn.CoreExpr -> Bool
 157 -- Treat Type arguments differently, because exprType is not defined for them.
 158 is_fun (CoreSyn.Type _) = False
 159 is_fun expr = (Type.isFunTy . CoreUtils.exprType) expr
 160
 161 -- Is the given core expression polymorphic (i.e., does it accept type
 162 -- arguments?).
 163 is_poly :: CoreSyn.CoreExpr -> Bool
 164 -- Treat Type arguments differently, because exprType is not defined for them.
 165 is_poly (CoreSyn.Type _) = False
 166 is_poly expr = (Maybe.isJust . Type.splitForAllTy_maybe . CoreUtils.exprType) expr
 167
 168 -- Is the given core expression a variable reference?
 169 is_var :: CoreSyn.CoreExpr -> Bool
 170 is_var (CoreSyn.Var _) = True
 171 is_var _ = False
 172
 173 is_lit :: CoreSyn.CoreExpr -> Bool
 174 is_lit (CoreSyn.Lit _) = True
 175 is_lit _ = False
 176
 177 -- Can the given core expression be applied to something? This is true for
 178 -- applying to a value as well as a type.
 179 is_applicable :: CoreSyn.CoreExpr -> Bool
 180 is_applicable expr = is_fun expr || is_poly expr
 181
 182 -- Is the given core expression a variable or an application?
 183 is_simple :: CoreSyn.CoreExpr -> Bool
 184 is_simple (CoreSyn.App _ _) = True
 185 is_simple (CoreSyn.Var _) = True
 186 is_simple (CoreSyn.Cast expr _) = is_simple expr
 187 is_simple _ = False
 188
 189 -- Does the given CoreExpr have any free type vars?
 190 has_free_tyvars :: CoreSyn.CoreExpr -> Bool
 191 has_free_tyvars = not . VarSet.isEmptyVarSet . (CoreFVs.exprSomeFreeVars Var.isTyVar)
 192
 193 -- Does the given CoreExpr have any free local vars?
 194 has_free_vars :: CoreSyn.CoreExpr -> Bool
 195 has_free_vars = not . VarSet.isEmptyVarSet . CoreFVs.exprFreeVars
 196
 197 -- Does the given expression use any of the given binders?
 198 expr_uses_binders :: [CoreSyn.CoreBndr] -> CoreSyn.CoreExpr -> Bool
 199 expr_uses_binders bndrs = not . VarSet.isEmptyVarSet . (CoreFVs.exprSomeFreeVars (`elem` bndrs))
 200
 201 -- Turns a Var CoreExpr into the Id inside it. Will of course only work for
 202 -- simple Var CoreExprs, not complexer ones.
 203 exprToVar :: CoreSyn.CoreExpr -> Var.Id
 204 exprToVar (CoreSyn.Var id) = id
 205 exprToVar expr = error $ "\nCoreTools.exprToVar: Not a var: " ++ show expr
 206
 207 -- Turns a Lit CoreExpr into the Literal inside it.
 208 exprToLit :: CoreSyn.CoreExpr -> Literal.Literal
 209 exprToLit (CoreSyn.Lit lit) = lit
 210 exprToLit expr = error $ "\nCoreTools.exprToLit: Not a lit: " ++ show expr
 211
 212 -- Removes all the type and dictionary arguments from the given argument list,
 213 -- leaving only the normal value arguments. The type given is the type of the
 214 -- expression applied to this argument list.
 215 get_val_args :: Type.Type -> [CoreSyn.CoreExpr] -> [CoreSyn.CoreExpr]
 216 get_val_args ty args = drop n args
 217   where
 218     (tyvars, predtypes, _) = TcType.tcSplitSigmaTy ty
 219     -- The first (length tyvars) arguments should be types, the next
 220     -- (length predtypes) arguments should be dictionaries. We drop this many
 221     -- arguments, to get at the value arguments.
 222     n = length tyvars + length predtypes
 223
 224 getLiterals :: CoreSyn.CoreExpr -> [CoreSyn.CoreExpr]
 225 getLiterals app@(CoreSyn.App _ _) = literals
 226   where
 227     (CoreSyn.Var f, args) = CoreSyn.collectArgs app
 228     literals = filter (is_lit) args
 229
 230 getLiterals lit@(CoreSyn.Lit _) = [lit]
 231
 232 reduceCoreListToHsList ::
 233   [HscTypes.CoreModule] -- ^ The modules where parts of the list are hidden
 234   -> CoreSyn.CoreExpr   -- ^ The refence to atleast one of the nodes
 235   -> TranslatorSession [CoreSyn.CoreExpr]
 236 reduceCoreListToHsList cores app@(CoreSyn.App _ _) = do {
 237   ; let { (fun, args) = CoreSyn.collectArgs app
 238         ; len         = length args
 239         } ;
 240   ; case len of
 241       3 -> do {
 242         ; let topelem = args!!1
 243         ; case (args!!2) of
 244             (varz@(CoreSyn.Var id)) -> do {
 245               ; binds <- mapM (findExpr (isVarName id)) cores
 246               ; otherelems <- reduceCoreListToHsList cores (head (Maybe.catMaybes binds))
 247               ; return (topelem:otherelems)
 248               }
 249             (appz@(CoreSyn.App _ _)) -> do {
 250               ; otherelems <- reduceCoreListToHsList cores appz
 251               ; return (topelem:otherelems)
 252               }
 253             otherwise -> return [topelem]
 254         }
 255       otherwise -> return []
 256   }
 257   where
 258     isVarName :: Monad m => Var.Var -> Var.Var -> m Bool
 259     isVarName lookfor bind = return $ (Var.varName lookfor) == (Var.varName bind)
 260
 261 reduceCoreListToHsList _ _ = return []
 262
 263 -- Is the given var the State data constructor?
 264 isStateCon :: Var.Var -> Bool
 265 isStateCon var = do
 266   -- See if it is a DataConWrapId (not DataConWorkId, since State is a
 267   -- newtype).
 268   case Id.idDetails var of
 269     IdInfo.DataConWrapId dc ->
 270       -- See if the datacon is the State datacon from the State type.
 271       let tycon = DataCon.dataConTyCon dc
 272           tyname = Name.getOccString tycon
 273           dcname = Name.getOccString dc
 274       in case (tyname, dcname) of
 275         ("State", "State") -> True
 276         _ -> False
 277     _ -> False
 278
 279 -- | Is the given type a State type?
 280 isStateType :: Type.Type -> Bool
 281 -- Resolve any type synonyms remaining
 282 isStateType ty | Just ty' <- Type.tcView ty = isStateType ty'
 283 isStateType ty  = Maybe.isJust $ do
 284   -- Split the type. Don't use normal splitAppTy, since that looks through
 285   -- newtypes, and we want to see the State newtype.
 286   (typef, _) <- Type.repSplitAppTy_maybe ty
 287   -- See if the applied type is a type constructor
 288   (tycon, _) <- Type.splitTyConApp_maybe typef
 289   if TyCon.isNewTyCon tycon && Name.getOccString tycon == "State"
 290     then
 291       Just ()
 292     else
 293       Nothing
 294
 295 -- | Does the given TypedThing have a State type?
 296 hasStateType :: (TypedThing t) => t -> Bool
 297 hasStateType expr = case getType expr of
 298   Nothing -> False
 299   Just ty -> isStateType ty
 300
 301
 302 -- | Flattens nested non-recursive lets into a single list of bindings. The
 303 -- expression passed does not have to be a let expression, if it isn't an
 304 -- empty list of bindings is returned.
 305 flattenLets ::
 306   CoreSyn.CoreExpr -- ^ The expression to flatten.
 307   -> ([Binding], CoreSyn.CoreExpr) -- ^ The bindings and resulting expression.
 308 flattenLets (CoreSyn.Let (CoreSyn.NonRec bndr expr) res) =
 309   ((bndr, expr):bindings, res')
 310   where
 311     -- Recursively flatten the contained expression
 312     (bindings, res') = flattenLets res
 313 flattenLets expr = ([], expr)
 314
 315 -- | Create bunch of nested non-recursive let expressions from the given
 316 -- bindings. The first binding is bound at the highest level (and thus
 317 -- available in all other bindings).
 318 mkNonRecLets :: [Binding] -> CoreSyn.CoreExpr -> CoreSyn.CoreExpr
 319 mkNonRecLets bindings expr = MkCore.mkCoreLets binds expr
 320   where
 321     binds = map (uncurry CoreSyn.NonRec) bindings
 322
 323 -- | A class of things that (optionally) have a core Type. The type is
 324 -- optional, since Type expressions don't have a type themselves.
 325 class TypedThing t where
 326   getType :: t -> Maybe Type.Type
 327
 328 instance TypedThing CoreSyn.CoreExpr where
 329   getType (CoreSyn.Type _) = Nothing
 330   getType expr = Just $ CoreUtils.exprType expr
 331
 332 instance TypedThing CoreSyn.CoreBndr where
 333   getType = return . Id.idType
 334
 335 instance TypedThing Type.Type where
 336   getType = return . id