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