Refactor inline array allocation

- Move array representation knowledge into SMRep

- Separate out low-level heap-object allocation so that we can reuse
  it from doNewArrayOp

- remove card-table initialisation, we can safely ignore the card
  table for newly allocated arrays.

4 files changed, 67 insertions, 96 deletions

diff --git a/compiler/codeGen/StgCmmHeap.hs b/compiler/codeGen/StgCmmHeap.hs
index 75ad8b40f4..2a0eaf9da7 100644
--- a/compiler/codeGen/StgCmmHeap.hs
+++ b/compiler/codeGen/StgCmmHeap.hs
@@ -16,7 +16,7 @@ module StgCmmHeap (
 
         mkStaticClosureFields, mkStaticClosure,
 
-        allocDynClosure, allocDynClosureCmm,
+        allocDynClosure, allocDynClosureCmm, allocHeapClosure,
         emitSetDynHdr
     ) where
 
@@ -88,61 +88,69 @@ allocDynClosureCmm
 -- significant - see test T4801.
 
 
-allocDynClosure mb_id info_tbl lf_info use_cc _blame_cc args_w_offsets
-  = do  { let (args, offsets) = unzip args_w_offsets
-        ; cmm_args <- mapM getArgAmode args     -- No void args
-        ; allocDynClosureCmm mb_id info_tbl lf_info
-                             use_cc _blame_cc (zip cmm_args offsets)
-        }
+allocDynClosure mb_id info_tbl lf_info use_cc _blame_cc args_w_offsets = do
+  let (args, offsets) = unzip args_w_offsets
+  cmm_args <- mapM getArgAmode args     -- No void args
+  allocDynClosureCmm mb_id info_tbl lf_info
+                     use_cc _blame_cc (zip cmm_args offsets)
 
-allocDynClosureCmm mb_id info_tbl lf_info use_cc _blame_cc amodes_w_offsets
-  = do  { virt_hp <- getVirtHp
 
-        -- SAY WHAT WE ARE ABOUT TO DO
-        ; let rep = cit_rep info_tbl
-        ; tickyDynAlloc mb_id rep lf_info
-        ; profDynAlloc rep use_cc
+allocDynClosureCmm mb_id info_tbl lf_info use_cc _blame_cc amodes_w_offsets = do
+  -- SAY WHAT WE ARE ABOUT TO DO
+  let rep = cit_rep info_tbl
+  tickyDynAlloc mb_id rep lf_info
+  profDynAlloc rep use_cc
+  let info_ptr = CmmLit (CmmLabel (cit_lbl info_tbl))
+  allocHeapClosure rep info_ptr use_cc amodes_w_offsets
 
-        -- FIND THE OFFSET OF THE INFO-PTR WORD
-        ; let   info_offset = virt_hp + 1
-                -- info_offset is the VirtualHpOffset of the first
-                -- word of the new object
-                -- Remember, virtHp points to last allocated word,
-                -- ie 1 *before* the info-ptr word of new object.
 
-                info_ptr = CmmLit (CmmLabel (cit_lbl info_tbl))
+-- | Low-level heap object allocation.
+allocHeapClosure
+  :: SMRep                            -- ^ representation of the object
+  -> CmmExpr                          -- ^ info pointer
+  -> CmmExpr                          -- ^ cost centre
+  -> [(CmmExpr,ByteOff)]              -- ^ payload
+  -> FCode CmmExpr                    -- ^ returns the address of the object
+allocHeapClosure rep info_ptr use_cc payload = do
+  virt_hp <- getVirtHp
 
-        -- ALLOCATE THE OBJECT
-        ; base <- getHpRelOffset info_offset
-        ; emitComment $ mkFastString "allocDynClosure"
-        ; emitSetDynHdr base info_ptr  use_cc
-        ; let (cmm_args, offsets) = unzip amodes_w_offsets
-        ; hpStore base cmm_args offsets
+  -- Find the offset of the info-ptr word
+  let info_offset = virt_hp + 1
+            -- info_offset is the VirtualHpOffset of the first
+            -- word of the new object
+            -- Remember, virtHp points to last allocated word,
+            -- ie 1 *before* the info-ptr word of new object.
 
-        -- BUMP THE VIRTUAL HEAP POINTER
-        ; dflags <- getDynFlags
-        ; setVirtHp (virt_hp + heapClosureSize dflags rep)
+  base <- getHpRelOffset info_offset
+  emitComment $ mkFastString "allocDynClosure"
+  emitSetDynHdr base info_ptr use_cc
+
+  -- Fill in the fields
+  hpStore base payload
+
+  -- Bump the virtual heap pointer
+  dflags <- getDynFlags
+  setVirtHp (virt_hp + heapClosureSizeW dflags rep)
+
+  return base
 
-        ; getHpRelOffset info_offset
-        }
 
 emitSetDynHdr :: CmmExpr -> CmmExpr -> CmmExpr -> FCode ()
 emitSetDynHdr base info_ptr ccs
   = do dflags <- getDynFlags
-       hpStore base (header dflags) [0, wORD_SIZE dflags ..]
+       hpStore base (zip (header dflags) [0, wORD_SIZE dflags ..])
   where
     header :: DynFlags -> [CmmExpr]
     header dflags = [info_ptr] ++ dynProfHdr dflags ccs
         -- ToDof: Parallel stuff
         -- No ticky header
 
-hpStore :: CmmExpr -> [CmmExpr] -> [ByteOff] -> FCode ()
 -- Store the item (expr,off) in base[off]
-hpStore base vals offs
-  = do dflags <- getDynFlags
-       let mk_store val off = mkStore (cmmOffsetB dflags base off) val
-       emit (catAGraphs (zipWith mk_store vals offs))
-
+hpStore :: CmmExpr -> [(CmmExpr, ByteOff)] -> FCode ()
+hpStore base vals = do
+  dflags <- getDynFlags
+  sequence_ $
+    [ emitStore (cmmOffsetB dflags base off) val | (val,off) <- vals ]
 
 -----------------------------------------------------------
 --              Layout of static closures
diff --git a/compiler/codeGen/StgCmmPrim.hs b/compiler/codeGen/StgCmmPrim.hs
index 504510c359..a4327c4064 100644
--- a/compiler/codeGen/StgCmmPrim.hs
+++ b/compiler/codeGen/StgCmmPrim.hs
@@ -90,10 +90,11 @@ cgOpApp (StgPrimOp primop) args res_ty = do
     dflags <- getDynFlags
     cmm_args <- getNonVoidArgAmodes args
     case shouldInlinePrimOp dflags primop cmm_args of
-        Nothing -> do let fun = CmmLit (CmmLabel (mkRtsPrimOpLabel primop))
-                      emitCall (NativeNodeCall, NativeReturn) fun cmm_args
+        Nothing -> do  -- out-of-line
+          let fun = CmmLit (CmmLabel (mkRtsPrimOpLabel primop))
+          emitCall (NativeNodeCall, NativeReturn) fun cmm_args
 
-        Just f
+        Just f  -- inline
           | ReturnsPrim VoidRep <- result_info
           -> do f []
                 emitReturn []
@@ -1533,36 +1534,24 @@ doNewArrayOp :: CmmFormal -> Integer -> CmmExpr -> FCode ()
 doNewArrayOp res_r n init = do
     dflags <- getDynFlags
 
-    let card_bytes = cardRoundUp dflags (fromInteger n)
-        size = fromInteger n + bytesToWordsRoundUp dflags card_bytes
-        words = arrPtrsHdrSizeWords dflags + size
-
-    -- If the allocation is of small, statically-known size, we reuse
-    -- the existing heap check to allocate inline.
-    virt_hp <- getVirtHp
-
-    -- FIND THE OFFSET OF THE INFO-PTR WORD
-    let   info_offset = virt_hp + 1
-          -- info_offset is the VirtualHpOffset of the first
-          -- word of the new object
-          -- Remember, virtHp points to last allocated word,
-          -- ie 1 *before* the info-ptr word of new object.
-    base <- getHpRelOffset info_offset
-    setVirtHp (virt_hp + fromIntegral words)  -- check n < big
-    arr <- CmmLocal `fmap` newTemp (bWord dflags)
-    emit $ mkAssign arr base
+    let info_ptr = mkLblExpr mkMAP_DIRTY_infoLabel
+
+    -- ToDo: this probably isn't right (card size?)
     tickyAllocPrim (mkIntExpr dflags (arrPtrsHdrSize dflags))
-        (cmmMulWord dflags (mkIntExpr dflags (fromInteger n)) (wordSize dflags))
+        (mkIntExpr dflags (fromInteger n * wORD_SIZE dflags))
         (zeroExpr dflags)
 
-    emitSetDynHdr base (mkLblExpr mkMAP_DIRTY_infoLabel) curCCS
-    emit $ mkStore (cmmOffsetB dflags base
-                    (fixedHdrSize dflags * wORD_SIZE dflags +
-                     oFFSET_StgMutArrPtrs_ptrs dflags))
-                   (mkIntExpr dflags (fromInteger n))
-    emit $ mkStore (cmmOffsetB dflags base
-                    (fixedHdrSize dflags * wORD_SIZE dflags +
-                     oFFSET_StgMutArrPtrs_size dflags)) (mkIntExpr dflags size)
+    let rep = arrPtrsRep dflags (fromIntegral n)
+        hdr_size = fixedHdrSize dflags * wORD_SIZE dflags
+    base <- allocHeapClosure rep info_ptr curCCS
+                     [ (mkIntExpr dflags (fromInteger n),
+                        hdr_size + oFFSET_StgMutArrPtrs_ptrs dflags)
+                     , (mkIntExpr dflags (nonHdrSizeW rep),
+                        hdr_size + oFFSET_StgMutArrPtrs_size dflags)
+                     ]
+
+    arr <- CmmLocal `fmap` newTemp (bWord dflags)
+    emit $ mkAssign arr base
 
     -- Initialise all elements of the the array
     p <- assignTemp $ cmmOffsetB dflags (CmmReg arr) (arrPtrsHdrSize dflags)
@@ -1577,26 +1566,12 @@ doNewArrayOp res_r n init = do
          (cmmOffsetW dflags (CmmReg arr) (fromInteger n)))
         (catAGraphs loopBody)
 
-    -- Initialise the mark bits with 0. This will be unrolled in the
-    -- backend to e.g. a single assignment since the arguments are
-    -- statically known.
-    emitMemsetCall
-        (cmmOffsetExprW dflags (CmmReg (CmmLocal p))
-         (mkIntExpr dflags (fromInteger n)))
-        (mkIntExpr dflags 0)
-        (mkIntExpr dflags card_bytes)
-        (mkIntExpr dflags (wORD_SIZE dflags))
     emit $ mkAssign (CmmLocal res_r) (CmmReg arr)
 
 -- | The inline allocation limit is 128 bytes, expressed in words.
 maxInlineAllocThreshold :: DynFlags -> Integer
 maxInlineAllocThreshold dflags = toInteger (128 `quot` wORD_SIZE dflags)
 
-arrPtrsHdrSizeWords :: DynFlags -> WordOff
-arrPtrsHdrSizeWords dflags =
-    fixedHdrSize dflags +
-    (sIZEOF_StgMutArrPtrs_NoHdr dflags `div` wORD_SIZE dflags)
-
 -- ----------------------------------------------------------------------------
 -- Copying pointer arrays
 
@@ -1724,18 +1699,6 @@ emitCloneArray info_p res_r src0 src_off0 n0 = do
         (mkIntExpr dflags (wORD_SIZE dflags))
     emit $ mkAssign (CmmLocal res_r) arr
 
-card :: DynFlags -> Int -> Int
-card dflags i = i `shiftR` mUT_ARR_PTRS_CARD_BITS dflags
-
--- Convert a number of elements to a number of cards, rounding up
-cardRoundUp :: DynFlags -> Int -> Int
-cardRoundUp dflags i =
-    card dflags (i + ((1 `shiftL` mUT_ARR_PTRS_CARD_BITS dflags) - 1))
-
-bytesToWordsRoundUp :: DynFlags -> Int -> Int
-bytesToWordsRoundUp dflags e =
-    (e + wORD_SIZE dflags - 1) `quot` (wORD_SIZE dflags)
-
 -- | Takes and offset in the destination array, the base address of
 -- the card table, and the number of elements affected (*not* the
 -- number of cards). The number of elements may not be zero.
diff --git a/compiler/codeGen/StgCmmProf.hs b/compiler/codeGen/StgCmmProf.hs
index e8a2a10fdd..f858c5a0b6 100644
--- a/compiler/codeGen/StgCmmProf.hs
+++ b/compiler/codeGen/StgCmmProf.hs
@@ -149,7 +149,7 @@ profDynAlloc :: SMRep -> CmmExpr -> FCode ()
 profDynAlloc rep ccs
   = ifProfiling $
     do dflags <- getDynFlags
-       profAlloc (mkIntExpr dflags (heapClosureSize dflags rep)) ccs
+       profAlloc (mkIntExpr dflags (heapClosureSizeW dflags rep)) ccs
 
 -- | Record the allocation of a closure (size is given by a CmmExpr)
 -- The size must be in words, because the allocation counter in a CCS counts
diff --git a/compiler/codeGen/StgCmmTicky.hs b/compiler/codeGen/StgCmmTicky.hs
index 3f3c3c5a19..50112f1ef8 100644
--- a/compiler/codeGen/StgCmmTicky.hs
+++ b/compiler/codeGen/StgCmmTicky.hs
@@ -415,7 +415,7 @@ tickyDynAlloc :: Maybe Id -> SMRep -> LambdaFormInfo -> FCode ()
 --
 -- TODO what else to count while we're here?
 tickyDynAlloc mb_id rep lf = ifTicky $ getDynFlags >>= \dflags ->
-  let bytes = wORD_SIZE dflags * heapClosureSize dflags rep
+  let bytes = wORD_SIZE dflags * heapClosureSizeW dflags rep
 
       countGlobal tot ctr = do
         bumpTickyCounterBy tot bytes