Allow the argument to 'reserve' to be a compile-time expression

By using the constant-folder to reduce it to an integer.

3 files changed, 29 insertions, 13 deletions

diff --git a/compiler/cmm/CmmOpt.hs b/compiler/cmm/CmmOpt.hs
index acaed28acf..54dbbebae6 100644
--- a/compiler/cmm/CmmOpt.hs
+++ b/compiler/cmm/CmmOpt.hs
@@ -7,6 +7,8 @@
 -----------------------------------------------------------------------------
 
 module CmmOpt (
+        constantFoldNode,
+        constantFoldExpr,
         cmmMachOpFold,
         cmmMachOpFoldM
  ) where
@@ -24,6 +26,16 @@ import Platform
 import Data.Bits
 import Data.Maybe
 
+
+constantFoldNode :: DynFlags -> CmmNode e x -> CmmNode e x
+constantFoldNode dflags = mapExp (constantFoldExpr dflags)
+
+constantFoldExpr :: DynFlags -> CmmExpr -> CmmExpr
+constantFoldExpr dflags = wrapRecExp f
+  where f (CmmMachOp op args) = cmmMachOpFold dflags op args
+        f (CmmRegOff r 0) = CmmReg r
+        f e = e
+
 -- -----------------------------------------------------------------------------
 -- MachOp constant folder
 
diff --git a/compiler/cmm/CmmParse.y b/compiler/cmm/CmmParse.y
index 8438198f7d..5f2c4d86be 100644
--- a/compiler/cmm/CmmParse.y
+++ b/compiler/cmm/CmmParse.y
@@ -221,6 +221,7 @@ import StgCmmLayout     hiding (ArgRep(..))
 import StgCmmTicky
 import StgCmmBind       ( emitBlackHoleCode, emitUpdateFrame )
 
+import CmmOpt
 import MkGraph
 import Cmm
 import CmmUtils
@@ -628,8 +629,8 @@ stmt    :: { CmmParse () }
                 { cmmIfThenElse $2 $4 $6 }
         | 'push' '(' exprs0 ')' maybe_body
                 { pushStackFrame $3 $5 }
-        | 'reserve' INT '=' lreg maybe_body
-                { reserveStackFrame (fromIntegral $2) $4 $5 }
+        | 'reserve' expr '=' lreg maybe_body
+                { reserveStackFrame $2 $4 $5 }
 
 foreignLabel     :: { CmmParse CmmExpr }
         : NAME                          { return (CmmLit (CmmLabel (mkForeignLabel $1 Nothing ForeignLabelInThisPackage IsFunction))) }
@@ -1076,12 +1077,21 @@ pushStackFrame fields body = do
   emit g
   withUpdFrameOff new_updfr_off body
 
-reserveStackFrame :: Int -> CmmParse CmmReg -> CmmParse () -> CmmParse ()
-reserveStackFrame size preg body = do
+reserveStackFrame
+  :: CmmParse CmmExpr
+  -> CmmParse CmmReg
+  -> CmmParse ()
+  -> CmmParse ()
+reserveStackFrame psize preg body = do
   dflags <- getDynFlags
   old_updfr_off <- getUpdFrameOff
   reg <- preg
-  let frame = old_updfr_off + wORD_SIZE dflags * size
+  esize <- psize
+  let size = case constantFoldExpr dflags esize of
+               CmmLit (CmmInt n _) -> n
+               _other -> pprPanic "CmmParse: not a compile-time integer: "
+                            (ppr esize)
+  let frame = old_updfr_off + wORD_SIZE dflags * fromIntegral size
   emitAssign reg (CmmStackSlot Old frame)
   withUpdFrameOff frame body
 
diff --git a/compiler/cmm/CmmSink.hs b/compiler/cmm/CmmSink.hs
index 6a3bcb7840..c404a2e932 100644
--- a/compiler/cmm/CmmSink.hs
+++ b/compiler/cmm/CmmSink.hs
@@ -171,7 +171,7 @@ cmmSink dflags graph = ofBlockList (g_entry graph) $ sink mapEmpty $ blocks
 
       -- Now sink and inline in this block
       (middle', assigs) = walk dflags ann_middles (mapFindWithDefault [] lbl sunk)
-      fold_last = constantFold dflags last
+      fold_last = constantFoldNode dflags last
       (final_last, assigs') = tryToInline dflags live fold_last assigs
 
       -- We cannot sink into join points (successors with more than
@@ -311,7 +311,7 @@ walk dflags nodes assigs = go nodes emptyBlock assigs
     | Just a <- shouldSink dflags node2 = go ns block (a : as1)
     | otherwise                         = go ns block' as'
     where
-      node1 = constantFold dflags node
+      node1 = constantFoldNode dflags node
 
       (node2, as1) = tryToInline dflags live node1 as
 
@@ -321,12 +321,6 @@ walk dflags nodes assigs = go nodes emptyBlock assigs
       block' = foldl blockSnoc block dropped `blockSnoc` node2
 
 
-constantFold :: DynFlags -> CmmNode e x -> CmmNode e x
-constantFold dflags node = mapExpDeep f node
-  where f (CmmMachOp op args) = cmmMachOpFold dflags op args
-        f (CmmRegOff r 0) = CmmReg r
-        f e = e
-
 --
 -- Heuristic to decide whether to pick up and sink an assignment
 -- Currently we pick up all assignments to local registers.  It might