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