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

There was a leftover big lambda in the CorePrep'd code, which confused
the bytecode generator.  Actually big lambdas are harmless.  This patch
refactors ByteCodeGen so that it systemantically used 'bcView' to eliminate
junk.  I did a little clean up in CorePrep too.

See comments in Trac #2898.

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)

diff --git a/compiler/ghci/ByteCodeGen.lhs b/compiler/ghci/ByteCodeGen.lhs
index 5dd63d3..a9e3c07 100644 (file)
--- a/compiler/ghci/ByteCodeGen.lhs
+++ b/compiler/ghci/ByteCodeGen.lhs
@@ -256,13 +256,14 @@ schemeR fvs (nm, rhs)
    = undefined
    | otherwise
 -}
-   = schemeR_wrk fvs nm rhs (collect [] rhs)
+   = schemeR_wrk fvs nm rhs (collect rhs)
 
-collect :: [Var] -> AnnExpr Id VarSet -> ([Var], AnnExpr' Id VarSet)
-collect xs (_, AnnNote _ e) = collect xs e
-collect xs (_, AnnCast e _) = collect xs e
-collect xs (_, AnnLam x e)  = collect (if isTyVar x then xs else (x:xs)) e
-collect xs (_, not_lambda)  = (reverse xs, not_lambda)
+collect :: AnnExpr Id VarSet -> ([Var], AnnExpr' Id VarSet)
+collect (_, e) = go [] e
+  where
+    go xs e | Just e' <- bcView e = go xs e'
+    go xs (AnnLam x (_,e))        = go (x:xs) e
+    go xs not_lambda              = (reverse xs, not_lambda)
 
 schemeR_wrk :: [Id] -> Id -> AnnExpr Id VarSet -> ([Var], AnnExpr' Var VarSet) -> BcM (ProtoBCO Name) 
 schemeR_wrk fvs nm original_body (args, body)
@@ -346,6 +347,10 @@ instance Outputable TickInfo where
 -- on the stack, returning a HNF.
 schemeE :: Int -> Sequel -> BCEnv -> AnnExpr' Id VarSet -> BcM BCInstrList
 
+schemeE d s p e
+   | Just e' <- bcView e
+   = schemeE d s p e'
+
 -- Delegate tail-calls to schemeT.
 schemeE d s p e@(AnnApp _ _) 
    = schemeT d s p e
@@ -397,7 +402,7 @@ schemeE d s p (AnnLet binds (_,body))
          sizes = map (\rhs_fvs -> sum (map idSizeW rhs_fvs)) fvss
 
         -- the arity of each rhs
-        arities = map (length . fst . collect []) rhss
+        arities = map (length . fst . collect) rhss
 
          -- This p', d' defn is safe because all the items being pushed
          -- are ptrs, so all have size 1.  d' and p' reflect the stack
@@ -491,12 +496,6 @@ schemeE d s p (AnnCase scrut _ _ [(DataAlt dc, [bind1], rhs)])
 schemeE d s p (AnnCase scrut bndr _ alts)
    = doCase d s p scrut bndr alts False{-not an unboxed tuple-} 
 
-schemeE d s p (AnnNote _ (_, body))
-   = schemeE d s p body
-
-schemeE d s p (AnnCast (_, body) _)
-   = schemeE d s p body
-
 schemeE _ _ _ expr
    = pprPanic "ByteCodeGen.schemeE: unhandled case" 
                (pprCoreExpr (deAnnotate' expr))
@@ -1169,18 +1168,11 @@ implement_tagToId names
 
 pushAtom :: Int -> BCEnv -> AnnExpr' Id VarSet -> BcM (BCInstrList, Int)
 
-pushAtom d p (AnnApp f (_, AnnType _))
-   = pushAtom d p (snd f)
-
-pushAtom d p (AnnNote _ e)
-   = pushAtom d p (snd e)
-
-pushAtom d p (AnnLam x e) 
-   | isTyVar x 
-   = pushAtom d p (snd e)
+pushAtom d p e 
+   | Just e' <- bcView e 
+   = pushAtom d p e'
 
 pushAtom d p (AnnVar v)
-
    | idCgRep v == VoidArg
    = return (nilOL, 0)
 
@@ -1411,34 +1403,37 @@ unboxedTupleException
 mkSLIDE :: Int -> Int -> OrdList BCInstr
 mkSLIDE n d = if d == 0 then nilOL else unitOL (SLIDE n d)
 
-splitApp :: AnnExpr' id ann -> (AnnExpr' id ann, [AnnExpr' id ann])
+splitApp :: AnnExpr' Var ann -> (AnnExpr' Var ann, [AnnExpr' Var ann])
        -- The arguments are returned in *right-to-left* order
-splitApp (AnnApp (_,f) (_,a))
-              | isTypeAtom a = splitApp f
-              | otherwise    = case splitApp f of 
-                                    (f', as) -> (f', a:as)
-splitApp (AnnNote _ (_,e))    = splitApp e
-splitApp (AnnCast (_,e) _)    = splitApp e
-splitApp e                   = (e, [])
-
-
-isTypeAtom :: AnnExpr' id ann -> Bool
-isTypeAtom (AnnType _) = True
-isTypeAtom _           = False
-
-isVoidArgAtom :: AnnExpr' id ann -> Bool
-isVoidArgAtom (AnnVar v)        = typePrimRep (idType v) == VoidRep
-isVoidArgAtom (AnnNote _ (_,e)) = isVoidArgAtom e
-isVoidArgAtom (AnnCast (_,e) _) = isVoidArgAtom e
-isVoidArgAtom _                = False
+splitApp e | Just e' <- bcView e = splitApp e'
+splitApp (AnnApp (_,f) (_,a))   = case splitApp f of 
+                                     (f', as) -> (f', a:as)
+splitApp e                      = (e, [])
+
+
+bcView :: AnnExpr' Var ann -> Maybe (AnnExpr' Var ann)
+-- The "bytecode view" of a term discards
+--  a) type abstractions
+--  b) type applications
+--  c) casts
+--  d) notes
+-- Type lambdas *can* occur in random expressions, 
+-- whereas value lambdas cannot; that is why they are nuked here
+bcView (AnnNote _ (_,e))            = Just e
+bcView (AnnCast (_,e) _)            = Just e
+bcView (AnnLam v (_,e)) | isTyVar v  = Just e
+bcView (AnnApp (_,e) (_, AnnType _)) = Just e
+bcView _                             = Nothing
+
+isVoidArgAtom :: AnnExpr' Var ann -> Bool
+isVoidArgAtom e | Just e' <- bcView e = isVoidArgAtom e'
+isVoidArgAtom (AnnVar v)              = typePrimRep (idType v) == VoidRep
+isVoidArgAtom _                      = False
 
 atomPrimRep :: AnnExpr' Id ann -> PrimRep
-atomPrimRep (AnnVar v)    = typePrimRep (idType v)
-atomPrimRep (AnnLit l)    = typePrimRep (literalType l)
-atomPrimRep (AnnNote _ b) = atomPrimRep (snd b)
-atomPrimRep (AnnApp f (_, AnnType _)) = atomPrimRep (snd f)
-atomPrimRep (AnnLam x e) | isTyVar x = atomPrimRep (snd e)
-atomPrimRep (AnnCast b _) = atomPrimRep (snd b)
+atomPrimRep e | Just e' <- bcView e = atomPrimRep e'
+atomPrimRep (AnnVar v)             = typePrimRep (idType v)
+atomPrimRep (AnnLit l)             = typePrimRep (literalType l)
 atomPrimRep other = pprPanic "atomPrimRep" (ppr (deAnnotate (undefined,other)))
 
 atomRep :: AnnExpr' Id ann -> CgRep