Tidy up treatment of big lambda (fixes Trac #2898)

[ghc-hetmet.git] / compiler / coreSyn / CorePrep.lhs

diff --git a/compiler/coreSyn/CorePrep.lhs b/compiler/coreSyn/CorePrep.lhs
index 5fa5002..db8bebc 100644 (file)
--- a/compiler/coreSyn/CorePrep.lhs
+++ b/compiler/coreSyn/CorePrep.lhs
@@ -61,8 +61,9 @@ The goal of this pass is to prepare for code generation.
     [I'm experimenting with leaving 'ok-for-speculation' 
      rhss in let-form right up to this point.]
 
-4.  Ensure that lambdas only occur as the RHS of a binding
+4.  Ensure that *value* lambdas only occur as the RHS of a binding
     (The code generator can't deal with anything else.)
+    Type lambdas are ok, however, because the code gen discards them.
 
 5.  [Not any more; nuked Jun 2002] Do the seq/par munging.
 
@@ -159,6 +160,7 @@ mkDataConWorkers data_tycons
 
 data FloatingBind = FloatLet CoreBind
                  | FloatCase Id CoreExpr Bool
+                       -- Invariant: the expression is not a lambda
                        -- The bool indicates "ok-for-speculation"
 
 data Floats = Floats OkToSpec (OrdList FloatingBind)
@@ -400,12 +402,6 @@ corePrepExprFloat env (Note n@(SCC _) expr) = do
     (floats, expr2) <- deLamFloat expr1
     return (floats, Note n expr2)
 
-corePrepExprFloat env (Case (Var id) bndr ty [(DEFAULT,[],expr)])
-  | Just (TickBox {}) <- isTickBoxOp_maybe id = do
-    expr1 <- corePrepAnExpr env expr
-    (floats, expr2) <- deLamFloat expr1
-    return (floats, Case (Var id) bndr ty [(DEFAULT,[],expr2)])
-
 corePrepExprFloat env (Note other_note expr) = do
     (floats, expr') <- corePrepExprFloat env expr
     return (floats, Note other_note expr')
@@ -421,6 +417,12 @@ corePrepExprFloat env expr@(Lam _ _) = do
   where
     (bndrs,body) = collectBinders expr
 
+corePrepExprFloat env (Case (Var id) bndr ty [(DEFAULT,[],expr)])
+  | Just (TickBox {}) <- isTickBoxOp_maybe id = do
+    expr1 <- corePrepAnExpr env expr
+    (floats, expr2) <- deLamFloat expr1
+    return (floats, Case (Var id) bndr ty [(DEFAULT,[],expr2)])
+
 corePrepExprFloat env (Case scrut bndr ty alts) = do
     (floats1, scrut1) <- corePrepExprFloat env scrut
     (floats2, scrut2) <- deLamFloat scrut1
@@ -639,15 +641,20 @@ mkLocalNonRec bndr dem floats rhs
 
 
 mkBinds :: Floats -> CoreExpr -> UniqSM CoreExpr
+-- Lambdas are not allowed as the body of a 'let'
 mkBinds (Floats _ binds) body 
   | isNilOL binds = return body
-  | otherwise    = do body' <- deLam body
-                        -- Lambdas are not allowed as the body of a 'let'
-                       return (foldrOL mk_bind body' binds)
+  | otherwise    = do { body' <- deLam body
+                       ; return (wrapBinds binds body') }
+
+wrapBinds :: OrdList FloatingBind -> CoreExpr -> CoreExpr
+wrapBinds binds body
+  = foldrOL mk_bind body binds
   where
     mk_bind (FloatCase bndr rhs _) body = Case rhs bndr (exprType body) [(DEFAULT, [], body)]
     mk_bind (FloatLet bind)        body = Let bind body
 
+---------------------
 etaExpandRhs :: CoreBndr -> CoreExpr -> UniqSM CoreExpr
 etaExpandRhs bndr rhs = do
        -- Eta expand to match the arity claimed by the binder
@@ -703,8 +710,8 @@ deLam :: CoreExpr -> UniqSM CoreExpr
 -- and returns one that definitely isn't:
 --     (\x.e) ==>  let f = \x.e in f
 deLam expr = do
-    (floats, expr) <- deLamFloat expr
-    mkBinds floats expr
+    (Floats _ binds, expr) <- deLamFloat expr
+    return (wrapBinds binds expr)
 
 
 deLamFloat :: CoreExpr -> UniqSM (Floats, CoreExpr)