Add support for fromIntegerT
[matthijs/master-project/cλash.git] / cλash / CLasH / Utils / Core / CoreTools.hs
index 42373a4eace28d337b193d9eb1376f7666061d7f..acc2fa630c416a32f6a99a15c0df18281b6526da 100644 (file)
@@ -1,3 +1,4 @@
+{-# LANGUAGE PatternGuards, TypeSynonymInstances #-}
 -- | This module provides a number of functions to find out things about Core
 -- programs. This module does not provide the actual plumbing to work with
 -- Core and Haskell (it uses HsTools for this), but only the functions that
 -- | This module provides a number of functions to find out things about Core
 -- programs. This module does not provide the actual plumbing to work with
 -- Core and Haskell (it uses HsTools for this), but only the functions that
@@ -19,22 +20,36 @@ import qualified HscTypes
 import qualified RdrName
 import qualified Name
 import qualified OccName
 import qualified RdrName
 import qualified Name
 import qualified OccName
+import qualified Type
+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 TysWiredIn
 import qualified Bag
 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 VarSet
 import qualified Unique
 import qualified CoreUtils
 import qualified CoreFVs
 import qualified Literal
+import qualified MkCore
+import qualified VarEnv
+import qualified Literal
 
 -- Local imports
 
 -- Local imports
+import CLasH.Translator.TranslatorTypes
 import CLasH.Utils.GhcTools
 import CLasH.Utils.HsTools
 import CLasH.Utils.Pretty
 import CLasH.Utils.GhcTools
 import CLasH.Utils.HsTools
 import CLasH.Utils.Pretty
+import CLasH.Utils
+import qualified CLasH.Utils.Core.BinderTools as BinderTools
+
+-- | A single binding, used as a shortcut to simplify type signatures.
+type Binding = (CoreSyn.CoreBndr, CoreSyn.CoreExpr)
 
 -- | Evaluate a core Type representing type level int from the tfp
 -- library to a real int.
 
 -- | Evaluate a core Type representing type level int from the tfp
 -- library to a real int.
@@ -45,7 +60,7 @@ eval_tfp_int env ty =
     -- Automatically import modules for any fully qualified identifiers
     setDynFlag DynFlags.Opt_ImplicitImportQualified
 
     -- Automatically import modules for any fully qualified identifiers
     setDynFlag DynFlags.Opt_ImplicitImportQualified
 
-    let from_int_t_name = mkRdrName "Types.Data.Num" "fromIntegerT"
+    let from_int_t_name = mkRdrName "Types.Data.Num.Ops" "fromIntegerT"
     let from_int_t = SrcLoc.noLoc $ HsExpr.HsVar from_int_t_name
     let undef = hsTypedUndef $ coreToHsType ty
     let app = SrcLoc.noLoc $ HsExpr.HsApp (from_int_t) (undef)
     let from_int_t = SrcLoc.noLoc $ HsExpr.HsVar from_int_t_name
     let undef = hsTypedUndef $ coreToHsType ty
     let app = SrcLoc.noLoc $ HsExpr.HsApp (from_int_t) (undef)
@@ -179,10 +194,18 @@ is_simple _ = False
 has_free_tyvars :: CoreSyn.CoreExpr -> Bool
 has_free_tyvars = not . VarSet.isEmptyVarSet . (CoreFVs.exprSomeFreeVars Var.isTyVar)
 
 has_free_tyvars :: CoreSyn.CoreExpr -> Bool
 has_free_tyvars = not . VarSet.isEmptyVarSet . (CoreFVs.exprSomeFreeVars Var.isTyVar)
 
+-- Does the given type have any free type vars?
+ty_has_free_tyvars :: Type.Type -> Bool
+ty_has_free_tyvars = not . VarSet.isEmptyVarSet . Type.tyVarsOfType
+
 -- Does the given CoreExpr have any free local vars?
 has_free_vars :: CoreSyn.CoreExpr -> Bool
 has_free_vars = not . VarSet.isEmptyVarSet . CoreFVs.exprFreeVars
 
 -- Does the given CoreExpr have any free local vars?
 has_free_vars :: CoreSyn.CoreExpr -> Bool
 has_free_vars = not . VarSet.isEmptyVarSet . CoreFVs.exprFreeVars
 
+-- Does the given expression use any of the given binders?
+expr_uses_binders :: [CoreSyn.CoreBndr] -> CoreSyn.CoreExpr -> Bool
+expr_uses_binders bndrs = not . VarSet.isEmptyVarSet . (CoreFVs.exprSomeFreeVars (`elem` bndrs))
+
 -- Turns a Var CoreExpr into the Id inside it. Will of course only work for
 -- simple Var CoreExprs, not complexer ones.
 exprToVar :: CoreSyn.CoreExpr -> Var.Id
 -- Turns a Var CoreExpr into the Id inside it. Will of course only work for
 -- simple Var CoreExprs, not complexer ones.
 exprToVar :: CoreSyn.CoreExpr -> Var.Id
@@ -206,19 +229,204 @@ 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
 
--- reduceCoreListToHsList :: CoreExpr -> [a]
-reduceCoreListToHsList app@(CoreSyn.App _ _) = out
+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
+  -> CoreSyn.CoreExpr   -- ^ The refence to atleast one of the nodes
+  -> TranslatorSession [CoreSyn.CoreExpr]
+reduceCoreListToHsList cores app@(CoreSyn.App _ _) = do {
+  ; let { (fun, args) = CoreSyn.collectArgs app
+        ; len         = length args 
+        } ;
+  ; case len of
+      3 -> do {
+        ; let topelem = args!!1
+        ; case (args!!2) of
+            (varz@(CoreSyn.Var id)) -> do {
+              ; binds <- mapM (findExpr (isVarName id)) cores
+              ; otherelems <- reduceCoreListToHsList cores (head (Maybe.catMaybes binds))
+              ; return (topelem:otherelems)
+              }
+            (appz@(CoreSyn.App _ _)) -> do {
+              ; otherelems <- reduceCoreListToHsList cores appz
+              ; return (topelem:otherelems)
+              }
+            otherwise -> return [topelem]
+        }
+      otherwise -> return []
+  }
+  where
+    isVarName :: Monad m => Var.Var -> Var.Var -> m Bool
+    isVarName lookfor bind = return $ (Var.varName lookfor) == (Var.varName bind)
+
+reduceCoreListToHsList _ _ = return []
+
+-- Is the given var the State data constructor?
+isStateCon :: Var.Var -> Bool
+isStateCon var = do
+  -- See if it is a DataConWrapId (not DataConWorkId, since State is a
+  -- newtype).
+  case Id.idDetails var of
+    IdInfo.DataConWrapId dc -> 
+      -- See if the datacon is the State datacon from the State type.
+      let tycon = DataCon.dataConTyCon dc
+          tyname = Name.getOccString tycon
+          dcname = Name.getOccString dc
+      in case (tyname, dcname) of
+        ("State", "State") -> True
+        _ -> False
+    _ -> False
+
+-- | Is the given type a State type?
+isStateType :: Type.Type -> Bool
+-- Resolve any type synonyms remaining
+isStateType ty | Just ty' <- Type.tcView ty = isStateType ty'
+isStateType ty  = Maybe.isJust $ do
+  -- Split the type. Don't use normal splitAppTy, since that looks through
+  -- newtypes, and we want to see the State newtype.
+  (typef, _) <- Type.repSplitAppTy_maybe ty
+  -- See if the applied type is a type constructor
+  (tycon, _) <- Type.splitTyConApp_maybe typef
+  if TyCon.isNewTyCon tycon && Name.getOccString tycon == "State"
+    then
+      Just ()
+    else
+      Nothing
+
+-- | Does the given TypedThing have a State type?
+hasStateType :: (TypedThing t) => t -> Bool
+hasStateType expr = case getType expr of
+  Nothing -> False
+  Just ty -> isStateType ty
+
+
+-- | Flattens nested lets into a single list of bindings. The expression
+--   passed does not have to be a let expression, if it isn't an empty list of
+--   bindings is returned.
+flattenLets ::
+  CoreSyn.CoreExpr -- ^ The expression to flatten.
+  -> ([Binding], CoreSyn.CoreExpr) -- ^ The bindings and resulting expression.
+flattenLets (CoreSyn.Let binds expr) = 
+  (bindings ++ bindings', expr')
+  where
+    -- Recursively flatten the contained expression
+    (bindings', expr') =flattenLets expr
+    -- Flatten our own bindings to remove the Rec / NonRec constructors
+    bindings = CoreSyn.flattenBinds [binds]
+flattenLets expr = ([], expr)
+
+-- | Create bunch of nested non-recursive let expressions from the given
+-- bindings. The first binding is bound at the highest level (and thus
+-- available in all other bindings).
+mkNonRecLets :: [Binding] -> CoreSyn.CoreExpr -> CoreSyn.CoreExpr
+mkNonRecLets bindings expr = MkCore.mkCoreLets binds expr
+  where
+    binds = map (uncurry CoreSyn.NonRec) bindings
+
+-- | A class of things that (optionally) have a core Type. The type is
+-- optional, since Type expressions don't have a type themselves.
+class TypedThing t where
+  getType :: t -> Maybe Type.Type
+
+instance TypedThing CoreSyn.CoreExpr where
+  getType (CoreSyn.Type _) = Nothing
+  getType expr = Just $ CoreUtils.exprType expr
+
+instance TypedThing CoreSyn.CoreBndr where
+  getType = return . Id.idType
+
+instance TypedThing Type.Type where
+  getType = return . id
+
+-- | Generate new uniques for all binders in the given expression.
+-- Does not support making type variables unique, though this could be
+-- supported if required (by passing a CoreSubst.Subst instead of VarEnv to
+-- genUniques' below).
+genUniques :: CoreSyn.CoreExpr -> TranslatorSession CoreSyn.CoreExpr
+genUniques = genUniques' VarEnv.emptyVarEnv
+
+-- | A helper function to generate uniques, that takes a VarEnv containing the
+--   substitutions already performed.
+genUniques' :: VarEnv.VarEnv CoreSyn.CoreBndr -> CoreSyn.CoreExpr -> TranslatorSession CoreSyn.CoreExpr
+genUniques' subst (CoreSyn.Var f) = do
+  -- Replace the binder with its new value, if applicable.
+  let f' = VarEnv.lookupWithDefaultVarEnv subst f f
+  return (CoreSyn.Var f')
+-- Leave literals untouched
+genUniques' subst (CoreSyn.Lit l) = return $ CoreSyn.Lit l
+genUniques' subst (CoreSyn.App f arg) = do
+  -- Only work on subexpressions
+  f' <- genUniques' subst f
+  arg' <- genUniques' subst arg
+  return (CoreSyn.App f' arg')
+-- Don't change type abstractions
+genUniques' subst expr@(CoreSyn.Lam bndr res) | CoreSyn.isTyVar bndr = return expr
+genUniques' subst (CoreSyn.Lam bndr res) = do
+  -- Generate a new unique for the bound variable
+  (subst', bndr') <- genUnique subst bndr
+  res' <- genUniques' subst' res
+  return (CoreSyn.Lam bndr' res')
+genUniques' subst (CoreSyn.Let (CoreSyn.NonRec bndr bound) res) = do
+  -- Make the binders unique
+  (subst', bndr') <- genUnique subst bndr
+  bound' <- genUniques' subst' bound
+  res' <- genUniques' subst' res
+  return $ CoreSyn.Let (CoreSyn.NonRec bndr' bound') res'
+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)
+  res' <- genUniques' subst' res
+  let binds' = zip bndrs' bounds'
+  return $ CoreSyn.Let (CoreSyn.Rec binds') res'
+genUniques' subst (CoreSyn.Case scrut bndr ty alts) = do
+  -- Process the scrutinee with the original substitution, since non of the
+  -- binders bound in the Case statement is in scope in the scrutinee.
+  scrut' <- genUniques' subst scrut
+  -- Generate a new binder for the scrutinee
+  (subst', bndr') <- genUnique subst bndr
+  -- Process each of the alts
+  alts' <- mapM (doalt subst') alts
+  return $ CoreSyn.Case scrut' bndr' ty alts'
   where
   where
-    (fun, args) = CoreSyn.collectArgs app
-    len = length args
-    out = case len of
-          3 -> ((args!!1) : (reduceCoreListToHsList (args!!2)))
-          otherwise -> []
+    doalt subst (con, bndrs, expr) = do
+      (subst', bndrs') <- mapAccumLM genUnique subst bndrs
+      expr' <- genUniques' subst' expr
+      -- Note that we don't return subst', since bndrs are only in scope in
+      -- expr.
+      return (con, bndrs', expr')
+genUniques' subst (CoreSyn.Cast expr coercion) = do
+  expr' <- genUniques' subst expr
+  -- Just process the casted expression
+  return $ CoreSyn.Cast expr' coercion
+genUniques' subst (CoreSyn.Note note expr) = do
+  expr' <- genUniques' subst expr
+  -- Just process the annotated expression
+  return $ CoreSyn.Note note expr'
+-- Leave types untouched
+genUniques' subst expr@(CoreSyn.Type _) = return expr
 
 
-reduceCoreListToHsList _ = []
+-- Generate a new unique for the given binder, and extend the given
+-- substitution to reflect this.
+genUnique :: VarEnv.VarEnv CoreSyn.CoreBndr -> CoreSyn.CoreBndr -> TranslatorSession (VarEnv.VarEnv CoreSyn.CoreBndr, CoreSyn.CoreBndr)
+genUnique subst bndr = do
+  bndr' <- BinderTools.cloneVar bndr
+  -- Replace all occurences of the old binder with a reference to the new
+  -- binder.
+  let subst' = VarEnv.extendVarEnv subst bndr bndr'
+  return (subst', bndr')