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