small reworking of the loop-breaker-choosing algorithm
[ghc-hetmet.git] / compiler / simplCore / LiberateCase.lhs
index c29a5b9..8b3d91b 100644 (file)
@@ -4,22 +4,33 @@
 \section[LiberateCase]{Unroll recursion to allow evals to be lifted from a loop}
 
 \begin{code}
 \section[LiberateCase]{Unroll recursion to allow evals to be lifted from a loop}
 
 \begin{code}
+{-# OPTIONS -w #-}
+-- The above warning supression flag is a temporary kludge.
+-- While working on this module you are encouraged to remove it and fix
+-- any warnings in the module. See
+--     http://hackage.haskell.org/trac/ghc/wiki/Commentary/CodingStyle#Warnings
+-- for details
+
 module LiberateCase ( liberateCase ) where
 
 #include "HsVersions.h"
 
 module LiberateCase ( liberateCase ) where
 
 #include "HsVersions.h"
 
-import DynFlags                ( DynFlags, DynFlag(..) )
-import StaticFlags     ( opt_LiberateCaseThreshold )
+import DynFlags
+import HscTypes
 import CoreLint                ( showPass, endPass )
 import CoreSyn
 import CoreUnfold      ( couldBeSmallEnoughToInline )
 import CoreLint                ( showPass, endPass )
 import CoreSyn
 import CoreUnfold      ( couldBeSmallEnoughToInline )
-import Id              ( Id, setIdName, idName, setIdNotExported )
+import Rules           ( RuleBase )
+import UniqSupply      ( UniqSupply )
+import SimplMonad      ( SimplCount, zeroSimplCount )
+import Id
 import VarEnv
 import Name            ( localiseName )
 import VarEnv
 import Name            ( localiseName )
-import Outputable
 import Util             ( notNull )
 \end{code}
 
 import Util             ( notNull )
 \end{code}
 
+The liberate-case transformation
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 This module walks over @Core@, and looks for @case@ on free variables.
 The criterion is:
        if there is case on a free on the route to the recursive call,
 This module walks over @Core@, and looks for @case@ on free variables.
 The criterion is:
        if there is case on a free on the route to the recursive call,
@@ -27,30 +38,24 @@ The criterion is:
 
 Example
 
 
 Example
 
-\begin{verbatim}
-f = \ t -> case v of
-              V a b -> a : f t
-\end{verbatim}
+   f = \ t -> case v of
+                V a b -> a : f t
 
 => the inner f is replaced.
 
 
 => the inner f is replaced.
 
-\begin{verbatim}
-f = \ t -> case v of
-              V a b -> a : (letrec
+   f = \ t -> case v of
+                V a b -> a : (letrec
                                f =  \ t -> case v of
                                               V a b -> a : f t
                                f =  \ t -> case v of
                                               V a b -> a : f t
-                            in f) t
-\end{verbatim}
+                              in f) t
 (note the NEED for shadowing)
 
 => Simplify
 
 (note the NEED for shadowing)
 
 => Simplify
 
-\begin{verbatim}
-f = \ t -> case v of
-              V a b -> a : (letrec
+  f = \ t -> case v of
+                V a b -> a : (letrec
                                f = \ t -> a : f t
                                f = \ t -> a : f t
-                            in f t)
-\begin{verbatim}
+                              in f t)
 
 Better code, because 'a' is  free inside the inner letrec, rather
 than needing projection from v.
 
 Better code, because 'a' is  free inside the inner letrec, rather
 than needing projection from v.
@@ -72,7 +77,6 @@ We'd like to avoid the redundant pattern match, transforming to
 
        (is this necessarily an improvement)
 
 
        (is this necessarily an improvement)
 
-
 Similarly drop:
 
        drop n [] = []
 Similarly drop:
 
        drop n [] = []
@@ -81,10 +85,36 @@ Similarly drop:
 
 Would like to pass n along unboxed.
        
 
 Would like to pass n along unboxed.
        
+Note [Scrutinee with cast]
+~~~~~~~~~~~~~~~~~~~~~~~~~~
+Consider this:
+    f = \ t -> case (v `cast` co) of
+                V a b -> a : f t
+
+Exactly the same optimisation (unrolling one call to f) will work here, 
+despite the cast.  See mk_alt_env in the Case branch of libCase.
+
+
+Note [Only functions!]
+~~~~~~~~~~~~~~~~~~~~~~
+Consider the following code
+
+       f = g (case v of V a b -> a : t f)
+
+where g is expensive. If we aren't careful, liberate case will turn this into
+
+       f = g (case v of
+               V a b -> a : t (letrec f = g (case v of V a b -> a : f t)
+                                in f)
+             )
+
+Yikes! We evaluate g twice. This leads to a O(2^n) explosion
+if g calls back to the same code recursively.
+
+Solution: make sure that we only do the liberate-case thing on *functions*
 
 To think about (Apr 94)
 ~~~~~~~~~~~~~~
 
 To think about (Apr 94)
 ~~~~~~~~~~~~~~
-
 Main worry: duplicating code excessively.  At the moment we duplicate
 the entire binding group once at each recursive call.  But there may
 be a group of recursive calls which share a common set of evaluated
 Main worry: duplicating code excessively.  At the moment we duplicate
 the entire binding group once at each recursive call.  But there may
 be a group of recursive calls which share a common set of evaluated
@@ -96,7 +126,6 @@ big.
 
 Data types
 ~~~~~~~~~~
 
 Data types
 ~~~~~~~~~~
-
 The ``level'' of a binder tells how many
 recursive defns lexically enclose the binding
 A recursive defn "encloses" its RHS, not its
 The ``level'' of a binder tells how many
 recursive defns lexically enclose the binding
 A recursive defn "encloses" its RHS, not its
@@ -110,57 +139,25 @@ scope.  For example:
 Here, the level of @f@ is zero, the level of @g@ is one,
 and the level of @h@ is zero (NB not one).
 
 Here, the level of @f@ is zero, the level of @g@ is one,
 and the level of @h@ is zero (NB not one).
 
-\begin{code}
-type LibCaseLevel = Int
-
-topLevel :: LibCaseLevel
-topLevel = 0
-\end{code}
-
-\begin{code}
-data LibCaseEnv
-  = LibCaseEnv
-       Int                     -- Bomb-out size for deciding if
-                               -- potential liberatees are too big.
-                               -- (passed in from cmd-line args)
-
-       LibCaseLevel            -- Current level
-
-       (IdEnv LibCaseLevel)    -- Binds all non-top-level in-scope Ids
-                               -- (top-level and imported things have
-                               -- a level of zero)
-
-       (IdEnv CoreBind)        -- Binds *only* recursively defined
-                               -- Ids, to their own binding group,
-                               -- and *only* in their own RHSs
-
-       [(Id,LibCaseLevel)]     -- Each of these Ids was scrutinised by an
-                               -- enclosing case expression, with the
-                               -- specified number of enclosing
-                               -- recursive bindings; furthermore,
-                               -- the Id is bound at a lower level
-                               -- than the case expression.  The
-                               -- order is insignificant; it's a bag
-                               -- really
 
 
-initEnv :: Int -> LibCaseEnv
-initEnv bomb_size = LibCaseEnv bomb_size 0 emptyVarEnv emptyVarEnv []
+%************************************************************************
+%*                                                                     *
+        Top-level code
+%*                                                                     *
+%************************************************************************
 
 
-bombOutSize (LibCaseEnv bomb_size _ _ _ _) = bomb_size
-\end{code}
-
-
-Programs
-~~~~~~~~
 \begin{code}
 \begin{code}
-liberateCase :: DynFlags -> [CoreBind] -> IO [CoreBind]
-liberateCase dflags binds
-  = do {
-       showPass dflags "Liberate case" ;
-       let { binds' = do_prog (initEnv opt_LiberateCaseThreshold) binds } ;
-       endPass dflags "Liberate case" Opt_D_verbose_core2core binds'
-                               {- no specific flag for dumping -} 
-    }
+liberateCase :: HscEnv -> UniqSupply -> RuleBase -> ModGuts
+            -> IO (SimplCount, ModGuts)
+liberateCase hsc_env _ _ guts
+  = do { let dflags = hsc_dflags hsc_env
+
+       ; showPass dflags "Liberate case"
+       ; let { env = initEnv dflags
+             ; binds' = do_prog env (mg_binds guts) }
+       ; endPass dflags "Liberate case" Opt_D_verbose_core2core binds'
+                       {- no specific flag for dumping -} 
+       ; return (zeroSimplCount dflags, guts { mg_binds = binds' }) }
   where
     do_prog env [] = []
     do_prog env (bind:binds) = bind' : do_prog env' binds
   where
     do_prog env [] = []
     do_prog env (bind:binds) = bind' : do_prog env' binds
@@ -168,9 +165,15 @@ liberateCase dflags binds
                               (env', bind') = libCaseBind env bind
 \end{code}
 
                               (env', bind') = libCaseBind env bind
 \end{code}
 
+
+%************************************************************************
+%*                                                                     *
+        Main payload
+%*                                                                     *
+%************************************************************************
+
 Bindings
 ~~~~~~~~
 Bindings
 ~~~~~~~~
-
 \begin{code}
 libCaseBind :: LibCaseEnv -> CoreBind -> (LibCaseEnv, CoreBind)
 
 \begin{code}
 libCaseBind :: LibCaseEnv -> CoreBind -> (LibCaseEnv, CoreBind)
 
@@ -186,7 +189,7 @@ libCaseBind env (Rec pairs)
 
     pairs' = [(binder, libCase env_rhs rhs) | (binder,rhs) <- pairs]
 
 
     pairs' = [(binder, libCase env_rhs rhs) | (binder,rhs) <- pairs]
 
-    env_rhs = if all rhs_small_enough rhss then extended_env else env
+    env_rhs = if all rhs_small_enough pairs then extended_env else env
 
        -- We extend the rec-env by binding each Id to its rhs, first
        -- processing the rhs with an *un-extended* environment, so
 
        -- We extend the rec-env by binding each Id to its rhs, first
        -- processing the rhs with an *un-extended* environment, so
@@ -207,8 +210,9 @@ libCaseBind env (Rec pairs)
        --      clash at code generation time.
     adjust bndr = setIdNotExported (setIdName bndr (localiseName (idName bndr)))
 
        --      clash at code generation time.
     adjust bndr = setIdNotExported (setIdName bndr (localiseName (idName bndr)))
 
-    rhs_small_enough rhs = couldBeSmallEnoughToInline lIBERATE_BOMB_SIZE rhs
-    lIBERATE_BOMB_SIZE   = bombOutSize env
+    rhs_small_enough (id,rhs)
+       =  idArity id > 0       -- Note [Only functions!]
+       && couldBeSmallEnoughToInline (bombOutSize env) rhs
 \end{code}
 
 
 \end{code}
 
 
@@ -225,6 +229,7 @@ libCase env (Lit lit)               = Lit lit
 libCase env (Type ty)          = Type ty
 libCase env (App fun arg)       = App (libCase env fun) (libCase env arg)
 libCase env (Note note body)    = Note note (libCase env body)
 libCase env (Type ty)          = Type ty
 libCase env (App fun arg)       = App (libCase env fun) (libCase env arg)
 libCase env (Note note body)    = Note note (libCase env body)
+libCase env (Cast e co)         = Cast (libCase env e) co
 
 libCase env (Lam binder body)
   = Lam binder (libCase (addBinders env [binder]) body)
 
 libCase env (Lam binder body)
   = Lam binder (libCase (addBinders env [binder]) body)
@@ -237,14 +242,15 @@ libCase env (Let bind body)
 libCase env (Case scrut bndr ty alts)
   = Case (libCase env scrut) bndr ty (map (libCaseAlt env_alts) alts)
   where
 libCase env (Case scrut bndr ty alts)
   = Case (libCase env scrut) bndr ty (map (libCaseAlt env_alts) alts)
   where
-    env_alts = addBinders env_with_scrut [bndr]
-    env_with_scrut = case scrut of
-                       Var scrut_var -> addScrutedVar env scrut_var
-                       other         -> env
+    env_alts = addBinders (mk_alt_env scrut) [bndr]
+    mk_alt_env (Var scrut_var) = addScrutedVar env scrut_var
+    mk_alt_env (Cast scrut _)  = mk_alt_env scrut      -- Note [Scrutinee with cast]
+    mk_alt_env otehr          = env
 
 libCaseAlt env (con,args,rhs) = (con, args, libCase (addBinders env args) rhs)
 \end{code}
 
 
 libCaseAlt env (con,args,rhs) = (con, args, libCase (addBinders env args) rhs)
 \end{code}
 
+
 Ids
 ~~~
 \begin{code}
 Ids
 ~~~
 \begin{code}
@@ -263,19 +269,23 @@ libCaseId env v
 \end{code}
 
 
 \end{code}
 
 
+%************************************************************************
+%*                                                                     *
+       Utility functions
+%*                                                                     *
+%************************************************************************
 
 
-Utility functions
-~~~~~~~~~~~~~~~~~
 \begin{code}
 addBinders :: LibCaseEnv -> [CoreBndr] -> LibCaseEnv
 \begin{code}
 addBinders :: LibCaseEnv -> [CoreBndr] -> LibCaseEnv
-addBinders (LibCaseEnv bomb lvl lvl_env rec_env scruts) binders
-  = LibCaseEnv bomb lvl lvl_env' rec_env scruts
+addBinders env@(LibCaseEnv { lc_lvl = lvl, lc_lvl_env = lvl_env }) binders
+  = env { lc_lvl_env = lvl_env' }
   where
     lvl_env' = extendVarEnvList lvl_env (binders `zip` repeat lvl)
 
 addRecBinds :: LibCaseEnv -> [(Id,CoreExpr)] -> LibCaseEnv
   where
     lvl_env' = extendVarEnvList lvl_env (binders `zip` repeat lvl)
 
 addRecBinds :: LibCaseEnv -> [(Id,CoreExpr)] -> LibCaseEnv
-addRecBinds (LibCaseEnv bomb lvl lvl_env rec_env scruts) pairs
-  = LibCaseEnv bomb lvl' lvl_env' rec_env' scruts
+addRecBinds env@(LibCaseEnv {lc_lvl = lvl, lc_lvl_env = lvl_env, 
+                            lc_rec_env = rec_env}) pairs
+  = env { lc_lvl = lvl', lc_lvl_env = lvl_env', lc_rec_env = rec_env' }
   where
     lvl'     = lvl + 1
     lvl_env' = extendVarEnvList lvl_env [(binder,lvl) | (binder,_) <- pairs]
   where
     lvl'     = lvl + 1
     lvl_env' = extendVarEnvList lvl_env [(binder,lvl) | (binder,_) <- pairs]
@@ -285,9 +295,10 @@ addScrutedVar :: LibCaseEnv
              -> Id             -- This Id is being scrutinised by a case expression
              -> LibCaseEnv
 
              -> Id             -- This Id is being scrutinised by a case expression
              -> LibCaseEnv
 
-addScrutedVar env@(LibCaseEnv bomb lvl lvl_env rec_env scruts) scrut_var
+addScrutedVar env@(LibCaseEnv { lc_lvl = lvl, lc_lvl_env = lvl_env, 
+                               lc_scruts = scruts }) scrut_var
   | bind_lvl < lvl
   | bind_lvl < lvl
-  = LibCaseEnv bomb lvl lvl_env rec_env scruts'
+  = env { lc_scruts = scruts' }
        -- Add to scruts iff the scrut_var is being scrutinised at
        -- a deeper level than its defn
 
        -- Add to scruts iff the scrut_var is being scrutinised at
        -- a deeper level than its defn
 
@@ -299,12 +310,11 @@ addScrutedVar env@(LibCaseEnv bomb lvl lvl_env rec_env scruts) scrut_var
                 Nothing  -> topLevel
 
 lookupRecId :: LibCaseEnv -> Id -> Maybe CoreBind
                 Nothing  -> topLevel
 
 lookupRecId :: LibCaseEnv -> Id -> Maybe CoreBind
-lookupRecId (LibCaseEnv bomb lvl lvl_env rec_env scruts) id
-  = lookupVarEnv rec_env id
+lookupRecId env id = lookupVarEnv (lc_rec_env env) id
 
 lookupLevel :: LibCaseEnv -> Id -> LibCaseLevel
 
 lookupLevel :: LibCaseEnv -> Id -> LibCaseLevel
-lookupLevel (LibCaseEnv bomb lvl lvl_env rec_env scruts) id
-  = case lookupVarEnv lvl_env id of
+lookupLevel env id
+  = case lookupVarEnv (lc_lvl_env env) id of
       Just lvl -> lvl
       Nothing  -> topLevel
 
       Just lvl -> lvl
       Nothing  -> topLevel
 
@@ -312,6 +322,61 @@ freeScruts :: LibCaseEnv
           -> LibCaseLevel      -- Level of the recursive Id
           -> [Id]              -- Ids that are scrutinised between the binding
                                -- of the recursive Id and here
           -> LibCaseLevel      -- Level of the recursive Id
           -> [Id]              -- Ids that are scrutinised between the binding
                                -- of the recursive Id and here
-freeScruts (LibCaseEnv bomb lvl lvl_env rec_env scruts) rec_bind_lvl
-  = [v | (v,scrut_lvl) <- scruts, scrut_lvl > rec_bind_lvl]
+freeScruts env rec_bind_lvl
+  = [v | (v,scrut_lvl) <- lc_scruts env, scrut_lvl > rec_bind_lvl]
 \end{code}
 \end{code}
+
+%************************************************************************
+%*                                                                     *
+        The environment
+%*                                                                     *
+%************************************************************************
+
+\begin{code}
+type LibCaseLevel = Int
+
+topLevel :: LibCaseLevel
+topLevel = 0
+\end{code}
+
+\begin{code}
+data LibCaseEnv
+  = LibCaseEnv {
+       lc_size :: Int,         -- Bomb-out size for deciding if
+                               -- potential liberatees are too big.
+                               -- (passed in from cmd-line args)
+
+       lc_lvl :: LibCaseLevel, -- Current level
+
+       lc_lvl_env :: IdEnv LibCaseLevel,  
+                       -- Binds all non-top-level in-scope Ids
+                       -- (top-level and imported things have
+                       -- a level of zero)
+
+       lc_rec_env :: IdEnv CoreBind, 
+                       -- Binds *only* recursively defined ids, 
+                       -- to their own binding group,
+                       -- and *only* in their own RHSs
+
+       lc_scruts :: [(Id,LibCaseLevel)]
+                       -- Each of these Ids was scrutinised by an
+                       -- enclosing case expression, with the
+                       -- specified number of enclosing
+                       -- recursive bindings; furthermore,
+                       -- the Id is bound at a lower level
+                       -- than the case expression.  The order is
+                       -- insignificant; it's a bag really
+       }
+
+initEnv :: DynFlags -> LibCaseEnv
+initEnv dflags 
+  = LibCaseEnv { lc_size = specThreshold dflags,
+                lc_lvl = 0,
+                lc_lvl_env = emptyVarEnv, 
+                lc_rec_env = emptyVarEnv,
+                lc_scruts = [] }
+
+bombOutSize = lc_size
+\end{code}
+
+