PprCore: Add size annotations for top-level bindings

1 files changed, 44 insertions, 21 deletions

diff --git a/compiler/coreSyn/PprCore.hs b/compiler/coreSyn/PprCore.hs
index c0af96886d..e33c1157cf 100644
--- a/compiler/coreSyn/PprCore.hs
+++ b/compiler/coreSyn/PprCore.hs
@@ -10,10 +10,12 @@ Printing of Core syntax
 module PprCore (
         pprCoreExpr, pprParendExpr,
         pprCoreBinding, pprCoreBindings, pprCoreAlt,
+        pprCoreBindingWithSize, pprCoreBindingsWithSize,
         pprRules
     ) where
 
 import CoreSyn
+import CoreStats (exprStats)
 import Literal( pprLiteral )
 import Name( pprInfixName, pprPrefixName )
 import Var
@@ -46,11 +48,17 @@ pprCoreBinding  :: OutputableBndr b => Bind b  -> SDoc
 pprCoreExpr     :: OutputableBndr b => Expr b  -> SDoc
 pprParendExpr   :: OutputableBndr b => Expr b  -> SDoc
 
-pprCoreBindings = pprTopBinds
-pprCoreBinding  = pprTopBind
+pprCoreBindings = pprTopBinds noAnn
+pprCoreBinding  = pprTopBind noAnn
+
+pprCoreBindingsWithSize :: [CoreBind] -> SDoc
+pprCoreBindingWithSize  :: CoreBind  -> SDoc
+
+pprCoreBindingsWithSize = pprTopBinds sizeAnn
+pprCoreBindingWithSize = pprTopBind sizeAnn
 
 instance OutputableBndr b => Outputable (Bind b) where
-    ppr bind = ppr_bind bind
+    ppr bind = ppr_bind noAnn bind
 
 instance OutputableBndr b => Outputable (Expr b) where
     ppr expr = pprCoreExpr expr
@@ -63,32 +71,47 @@ instance OutputableBndr b => Outputable (Expr b) where
 ************************************************************************
 -}
 
-pprTopBinds :: OutputableBndr a => [Bind a] -> SDoc
-pprTopBinds binds = vcat (map pprTopBind binds)
+-- | A function to produce an annotation for a given right-hand-side
+type Annotation b = Expr b -> SDoc
+
+-- | Annotate with the size of the right-hand-side
+sizeAnn :: CoreExpr -> SDoc
+sizeAnn e = ptext (sLit "-- RHS size:") <+> ppr (exprStats e)
+
+-- | No annotation
+noAnn :: Expr b -> SDoc
+noAnn _ = empty
+
+pprTopBinds :: OutputableBndr a
+            => Annotation a -- ^ generate an annotation to place before the
+                            -- binding
+            -> [Bind a]     -- ^ bindings to show
+            -> SDoc         -- ^ the pretty result
+pprTopBinds ann binds = vcat (map (pprTopBind ann) binds)
 
-pprTopBind :: OutputableBndr a => Bind a -> SDoc
-pprTopBind (NonRec binder expr)
- = ppr_binding (binder,expr) $$ blankLine
+pprTopBind :: OutputableBndr a => Annotation a -> Bind a -> SDoc
+pprTopBind ann (NonRec binder expr)
+ = ppr_binding ann (binder,expr) $$ blankLine
 
-pprTopBind (Rec [])
+pprTopBind _ (Rec [])
   = ptext (sLit "Rec { }")
-pprTopBind (Rec (b:bs))
+pprTopBind ann (Rec (b:bs))
   = vcat [ptext (sLit "Rec {"),
-          ppr_binding b,
-          vcat [blankLine $$ ppr_binding b | b <- bs],
+          ppr_binding ann b,
+          vcat [blankLine $$ ppr_binding ann b | b <- bs],
           ptext (sLit "end Rec }"),
           blankLine]
 
-ppr_bind :: OutputableBndr b => Bind b -> SDoc
+ppr_bind :: OutputableBndr b => Annotation b -> Bind b -> SDoc
 
-ppr_bind (NonRec val_bdr expr) = ppr_binding (val_bdr, expr)
-ppr_bind (Rec binds)           = vcat (map pp binds)
-                               where
-                                 pp bind = ppr_binding bind <> semi
+ppr_bind ann (NonRec val_bdr expr) = ppr_binding ann (val_bdr, expr)
+ppr_bind ann (Rec binds)           = vcat (map pp binds)
+                                    where
+                                      pp bind = ppr_binding ann bind <> semi
 
-ppr_binding :: OutputableBndr b => (b, Expr b) -> SDoc
-ppr_binding (val_bdr, expr)
-  = pprBndr LetBind val_bdr $$
+ppr_binding :: OutputableBndr b => Annotation b -> (b, Expr b) -> SDoc
+ppr_binding ann (val_bdr, expr)
+  = ann expr $$ pprBndr LetBind val_bdr $$
     hang (ppr val_bdr <+> equals) 2 (pprCoreExpr expr)
 
 pprParendExpr expr = ppr_expr parens expr
@@ -210,7 +233,7 @@ ppr_expr add_par (Let bind@(NonRec val_bdr rhs) expr@(Let _ _))
 -- General case (recursive case, too)
 ppr_expr add_par (Let bind expr)
   = add_par $
-    sep [hang (ptext keyword) 2 (ppr_bind bind <+> ptext (sLit "} in")),
+    sep [hang (ptext keyword) 2 (ppr_bind noAnn bind <+> ptext (sLit "} in")),
          pprCoreExpr expr]
   where
     keyword = case bind of