Split off selector case creation code into CoreTools.
[matthijs/master-project/cλash.git] / cλash / CLasH / Utils / Core / CoreTools.hs
index bf2ca27be3b6717b5262a8f9c1770397924be605..9b276d659eabac9731b8f1fcb6c6205c00ebb1cc 100644 (file)
@@ -7,7 +7,7 @@ module CLasH.Utils.Core.CoreTools where
 
 --Standard modules
 import qualified Maybe
 
 --Standard modules
 import qualified Maybe
-import System.IO.Unsafe
+import qualified System.IO.Unsafe
 
 -- GHC API
 import qualified GHC
 
 -- GHC API
 import qualified GHC
@@ -15,24 +15,18 @@ import qualified Type
 import qualified TcType
 import qualified HsExpr
 import qualified HsTypes
 import qualified TcType
 import qualified HsExpr
 import qualified HsTypes
-import qualified HsBinds
 import qualified HscTypes
 import qualified HscTypes
-import qualified RdrName
 import qualified Name
 import qualified Name
-import qualified OccName
-import qualified Type
 import qualified Id
 import qualified TyCon
 import qualified DataCon
 import qualified TysWiredIn
 import qualified Id
 import qualified TyCon
 import qualified DataCon
 import qualified TysWiredIn
-import qualified Bag
 import qualified DynFlags
 import qualified SrcLoc
 import qualified CoreSyn
 import qualified Var
 import qualified IdInfo
 import qualified VarSet
 import qualified DynFlags
 import qualified SrcLoc
 import qualified CoreSyn
 import qualified Var
 import qualified IdInfo
 import qualified VarSet
-import qualified Unique
 import qualified CoreUtils
 import qualified CoreFVs
 import qualified Literal
 import qualified CoreUtils
 import qualified CoreFVs
 import qualified Literal
@@ -42,6 +36,7 @@ import qualified VarEnv
 -- Local imports
 import CLasH.Translator.TranslatorTypes
 import CLasH.Utils.GhcTools
 -- Local imports
 import CLasH.Translator.TranslatorTypes
 import CLasH.Utils.GhcTools
+import CLasH.Utils.Core.BinderTools
 import CLasH.Utils.HsTools
 import CLasH.Utils.Pretty
 import CLasH.Utils
 import CLasH.Utils.HsTools
 import CLasH.Utils.Pretty
 import CLasH.Utils
@@ -73,9 +68,8 @@ eval_tfp_int env ty =
 
 normalise_tfp_int :: HscTypes.HscEnv -> Type.Type -> Type.Type
 normalise_tfp_int env ty =
 
 normalise_tfp_int :: HscTypes.HscEnv -> Type.Type -> Type.Type
 normalise_tfp_int env ty =
-   unsafePerformIO $ do
-     nty <- normaliseType env ty
-     return nty
+   System.IO.Unsafe.unsafePerformIO $
+     normaliseType env ty
 
 -- | Get the width of a SizedWord type
 -- sized_word_len :: HscTypes.HscEnv -> Type.Type -> Int
 
 -- | Get the width of a SizedWord type
 -- sized_word_len :: HscTypes.HscEnv -> Type.Type -> Int
@@ -155,6 +149,11 @@ is_lam :: CoreSyn.CoreExpr -> Bool
 is_lam (CoreSyn.Lam _ _) = True
 is_lam _ = False
 
 is_lam (CoreSyn.Lam _ _) = True
 is_lam _ = False
 
+-- Is the given core expression a let expression?
+is_let :: CoreSyn.CoreExpr -> Bool
+is_let (CoreSyn.Let _ _) = True
+is_let _ = False
+
 -- Is the given core expression of a function type?
 is_fun :: CoreSyn.CoreExpr -> Bool
 -- Treat Type arguments differently, because exprType is not defined for them.
 -- Is the given core expression of a function type?
 is_fun :: CoreSyn.CoreExpr -> Bool
 -- Treat Type arguments differently, because exprType is not defined for them.
@@ -228,13 +227,21 @@ get_val_args ty args = drop n args
     -- arguments, to get at the value arguments.
     n = length tyvars + length predtypes
 
     -- arguments, to get at the value arguments.
     n = length tyvars + length predtypes
 
-getLiterals :: CoreSyn.CoreExpr -> [CoreSyn.CoreExpr]
-getLiterals app@(CoreSyn.App _ _) = literals
+getLiterals :: HscTypes.HscEnv -> CoreSyn.CoreExpr -> [CoreSyn.CoreExpr]
+getLiterals app@(CoreSyn.App _ _) = literals
   where
     (CoreSyn.Var f, args) = CoreSyn.collectArgs app
     literals = filter (is_lit) args
 
   where
     (CoreSyn.Var f, args) = CoreSyn.collectArgs app
     literals = filter (is_lit) args
 
-getLiterals lit@(CoreSyn.Lit _) = [lit]
+getLiterals _ lit@(CoreSyn.Lit _) = [lit]
+
+getLiterals hscenv letrec@(CoreSyn.Let (CoreSyn.NonRec letBind (letExpr)) letRes) = [lit]
+  where
+    ty     = Var.varType letBind
+    litInt = eval_tfp_int hscenv ty
+    lit    = CoreSyn.Lit (Literal.mkMachInt (toInteger litInt))
+
+getLiterals _ expr = error $ "\nCoreTools.getLiterals: Not a known Lit: " ++ pprString expr
 
 reduceCoreListToHsList :: 
   [HscTypes.CoreModule] -- ^ The modules where parts of the list are hidden
 
 reduceCoreListToHsList :: 
   [HscTypes.CoreModule] -- ^ The modules where parts of the list are hidden
@@ -269,7 +276,7 @@ reduceCoreListToHsList _ _ = return []
 
 -- Is the given var the State data constructor?
 isStateCon :: Var.Var -> Bool
 
 -- Is the given var the State data constructor?
 isStateCon :: Var.Var -> Bool
-isStateCon var = do
+isStateCon var =
   -- See if it is a DataConWrapId (not DataConWorkId, since State is a
   -- newtype).
   case Id.idDetails var of
   -- See if it is a DataConWrapId (not DataConWorkId, since State is a
   -- newtype).
   case Id.idDetails var of
@@ -381,7 +388,7 @@ genUniques' subst (CoreSyn.Let (CoreSyn.NonRec bndr bound) res) = do
 genUniques' subst (CoreSyn.Let (CoreSyn.Rec binds) res) = do
   -- Make each of the binders unique
   (subst', bndrs') <- mapAccumLM genUnique subst (map fst binds)
 genUniques' subst (CoreSyn.Let (CoreSyn.Rec binds) res) = do
   -- Make each of the binders unique
   (subst', bndrs') <- mapAccumLM genUnique subst (map fst binds)
-  bounds' <- mapM (genUniques' subst') (map snd binds)
+  bounds' <- mapM (genUniques' subst' . snd) binds
   res' <- genUniques' subst' res
   let binds' = zip bndrs' bounds'
   return $ CoreSyn.Let (CoreSyn.Rec binds') res'
   res' <- genUniques' subst' res
   let binds' = zip bndrs' bounds'
   return $ CoreSyn.Let (CoreSyn.Rec binds') res'
@@ -421,3 +428,27 @@ genUnique subst bndr = do
   -- binder.
   let subst' = VarEnv.extendVarEnv subst bndr bndr'
   return (subst', bndr')
   -- binder.
   let subst' = VarEnv.extendVarEnv subst bndr bndr'
   return (subst', bndr')
+
+-- Create a "selector" case that selects the ith field from a datacon
+mkSelCase :: CoreSyn.CoreExpr -> Int -> TranslatorSession CoreSyn.CoreExpr
+mkSelCase scrut i = do
+  let scrut_ty = CoreUtils.exprType scrut
+  case Type.splitTyConApp_maybe scrut_ty of
+    -- The scrutinee should have a type constructor. We keep the type
+    -- arguments around so we can instantiate the field types below
+    Just (tycon, tyargs) -> case TyCon.tyConDataCons tycon of
+      -- The scrutinee type should have a single dataconstructor,
+      -- otherwise we can't construct a valid selector case.
+      [datacon] -> do
+        let field_tys = DataCon.dataConInstOrigArgTys datacon tyargs
+        -- Create a list of wild binders for the fields we don't want
+        let wildbndrs = map MkCore.mkWildBinder field_tys
+        -- Create a single binder for the field we want
+        sel_bndr <- mkInternalVar "sel" (field_tys!!i)
+        -- Create a wild binder for the scrutinee
+        let scrut_bndr = MkCore.mkWildBinder scrut_ty
+        -- Create the case expression
+        let binders = take i wildbndrs ++ [sel_bndr] ++ drop (i+1) wildbndrs
+        return $ CoreSyn.Case scrut scrut_bndr scrut_ty [(CoreSyn.DataAlt datacon, binders, CoreSyn.Var sel_bndr)]
+      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)
+    Nothing -> error $ "CoreTools.mkSelCase: Creating extractor case, but scrutinee has no tycon? Extracting element " ++ (show i) ++ " from '" ++ pprString scrut ++ "'"