Add ty_has_free_tyvars predicate.
[matthijs/master-project/cλash.git] / cλash / CLasH / Utils / Core / CoreTools.hs
index e3021adc95e0807109748ad42acc6dacefe03bb6..bf2ca27be3b6717b5262a8f9c1770397924be605 100644 (file)
@@ -23,23 +23,32 @@ import qualified OccName
 import qualified Type
 import qualified Id
 import qualified TyCon
 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
 
 -- Local imports
 import CLasH.Translator.TranslatorTypes
 import CLasH.Utils.GhcTools
 import CLasH.Utils.HsTools
 import CLasH.Utils.Pretty
 
 -- Local imports
 import CLasH.Translator.TranslatorTypes
 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.
@@ -184,6 +193,10 @@ 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
@@ -254,6 +267,22 @@ reduceCoreListToHsList cores app@(CoreSyn.App _ _) = do {
 
 reduceCoreListToHsList _ _ = return []
 
 
 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
 -- | Is the given type a State type?
 isStateType :: Type.Type -> Bool
 -- Resolve any type synonyms remaining
@@ -277,6 +306,29 @@ hasStateType expr = case getType expr of
   Just ty -> isStateType ty
 
 
   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
 -- | 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
@@ -291,3 +343,81 @@ instance TypedThing CoreSyn.CoreBndr where
 
 instance TypedThing Type.Type where
   getType = return . id
 
 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
+    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
+
+-- 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')