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