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