Rewrite fromInteger and literal generation.

[matthijs/master-project/cλash.git] / cλash / CLasH / Utils / Core / CoreTools.hs

diff --git a/cλash/CLasH/Utils/Core/CoreTools.hs b/cλash/CLasH/Utils/Core/CoreTools.hs
index acc2fa630c416a32f6a99a15c0df18281b6526da..09595702570cfe8745e49ec8190c4afd6d3fe44b 100644 (file)
--- a/cλash/CLasH/Utils/Core/CoreTools.hs
+++ b/cλash/CLasH/Utils/Core/CoreTools.hs
@@ -7,7 +7,9 @@ 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
+import qualified Data.Map as Map
+import qualified Data.Accessor.Monad.Trans.State as MonadState

 
 -- GHC API
 import qualified GHC
 
 -- GHC API
 import qualified GHC


@@ -15,34 +17,28 @@ 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 MkCore
 import qualified VarEnv
 import qualified CoreUtils
 import qualified CoreFVs
 import qualified Literal
 import qualified MkCore
 import qualified VarEnv

-import qualified Literal

 
 -- 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


@@ -52,7 +48,40 @@ import qualified CLasH.Utils.Core.BinderTools as BinderTools
 type Binding = (CoreSyn.CoreBndr, CoreSyn.CoreExpr)
 
 -- | Evaluate a core Type representing type level int from the tfp
 type Binding = (CoreSyn.CoreBndr, CoreSyn.CoreExpr)
 
 -- | Evaluate a core Type representing type level int from the tfp

--- library to a real int.
+-- library to a real int. Checks if the type really is a Dec type and
+-- caches the results.
+tfp_to_int :: Type.Type -> TypeSession Int
+tfp_to_int ty = do
+  hscenv <- MonadState.get tsHscEnv
+  let norm_ty = normalise_tfp_int hscenv ty
+  case Type.splitTyConApp_maybe norm_ty of
+    Just (tycon, args) -> do
+      let name = Name.getOccString (TyCon.tyConName tycon)
+      case name of
+        "Dec" ->
+          tfp_to_int' ty
+        otherwise -> do
+          return $ error ("Callin tfp_to_int on non-dec:" ++ (show ty))
+    Nothing -> return $ error ("Callin tfp_to_int on non-dec:" ++ (show ty))
+
+-- | Evaluate a core Type representing type level int from the tfp
+-- library to a real int. Caches the results. Do not use directly, use
+-- tfp_to_int instead.
+tfp_to_int' :: Type.Type -> TypeSession Int
+tfp_to_int' ty = do
+  lens <- MonadState.get tsTfpInts
+  hscenv <- MonadState.get tsHscEnv
+  let norm_ty = normalise_tfp_int hscenv ty
+  let existing_len = Map.lookup (OrdType norm_ty) lens
+  case existing_len of
+    Just len -> return len
+    Nothing -> do
+      let new_len = eval_tfp_int hscenv ty
+      MonadState.modify tsTfpInts (Map.insert (OrdType norm_ty) (new_len))
+      return new_len
+      
+-- | Evaluate a core Type representing type level int from the tfp
+-- library to a real int. Do not use directly, use tfp_to_int instead.

 eval_tfp_int :: HscTypes.HscEnv -> Type.Type -> Int
 eval_tfp_int env ty =
   unsafeRunGhc libdir $ do
 eval_tfp_int :: HscTypes.HscEnv -> Type.Type -> Int
 eval_tfp_int env ty =
   unsafeRunGhc libdir $ do


@@ -74,9 +103,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


@@ -156,6 +184,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.


@@ -229,21 +262,27 @@ 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 :: HscTypes.HscEnv -> CoreSyn.CoreExpr -> [CoreSyn.CoreExpr]
-getLiterals _ app@(CoreSyn.App _ _) = literals
-  where
-    (CoreSyn.Var f, args) = CoreSyn.collectArgs app
-    literals = filter (is_lit) args
-
-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
+-- Finds out what literal Integer this expression represents.
+getIntegerLiteral :: CoreSyn.CoreExpr -> TranslatorSession Integer
+getIntegerLiteral expr =
+  case CoreSyn.collectArgs expr of
+    (CoreSyn.Var f, [CoreSyn.Lit (Literal.MachInt integer)]) 
+      | getFullString f == "GHC.Integer.smallInteger" -> return integer
+    (CoreSyn.Var f, [CoreSyn.Lit (Literal.MachInt64 integer)]) 
+      | getFullString f == "GHC.Integer.int64ToInteger" -> return integer
+    (CoreSyn.Var f, [CoreSyn.Lit (Literal.MachWord integer)]) 
+      | getFullString f == "GHC.Integer.wordToInteger" -> return integer
+    (CoreSyn.Var f, [CoreSyn.Lit (Literal.MachWord64 integer)]) 
+      | getFullString f == "GHC.Integer.word64ToInteger" -> return integer
+    -- fromIntegerT returns the integer corresponding to the type of its
+    -- (third) argument. Since it is polymorphic, the type of that
+    -- argument is passed as the first argument, so we can just use that
+    -- one.
+    (CoreSyn.Var f, [CoreSyn.Type dec_ty, dec_dict, CoreSyn.Type num_ty, num_dict, arg]) 
+      | getFullString f == "Types.Data.Num.Ops.fromIntegerT" -> do
+          int <- MonadState.lift tsType $ tfp_to_int dec_ty
+          return $ toInteger int
+    _ -> error $ "CoreTools.getIntegerLiteral: Unsupported Integer literal: " ++ 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


@@ -278,7 +317,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


@@ -390,7 +429,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'


@@ -430,3 +469,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 ++ "'"