Implement general arrayOf# primop

6 files changed, 69 insertions, 23 deletions

diff --git a/compiler/GHC/Cmm/CLabel.hs b/compiler/GHC/Cmm/CLabel.hs
index b6ad2b3431..b6c6ab4a77 100644
--- a/compiler/GHC/Cmm/CLabel.hs
+++ b/compiler/GHC/Cmm/CLabel.hs
@@ -46,7 +46,7 @@ module GHC.Cmm.CLabel (
         mkAsmTempEndLabel,
         mkAsmTempProcEndLabel,
         mkAsmTempDieLabel,
-
+        mkUnliftedDataLabel,
         mkDirty_MUT_VAR_Label,
         mkNonmovingWriteBarrierEnabledLabel,
         mkUpdInfoLabel,
@@ -273,7 +273,6 @@ data CLabel
   -- | A label before an info table to prevent excessive dead-stripping on darwin
   | DeadStripPreventer CLabel
 
-
   -- | Per-module table of tick locations
   | HpcTicksLabel Module
 
@@ -285,6 +284,9 @@ data CLabel
   | LargeBitmapLabel
         {-# UNPACK #-} !Unique
 
+  -- | Static data from local definitions allocated in the data section,
+  -- arising from a primop, like 'arrayOf#'
+  | UnliftedDataLabel {-# UNPACK #-} !Unique PrimOp
   deriving Eq
 
 instance Show CLabel where
@@ -293,6 +295,9 @@ instance Show CLabel where
 instance Outputable CLabel where
   ppr = text . show
 
+mkUnliftedDataLabel :: Unique -> PrimOp -> CLabel
+mkUnliftedDataLabel = UnliftedDataLabel
+
 isIdLabel :: CLabel -> Bool
 isIdLabel IdLabel{} = True
 isIdLabel _ = False
@@ -369,6 +374,8 @@ instance Ord CLabel where
     nonDetCmpUnique u1 u2
   compare (LargeBitmapLabel u1) (LargeBitmapLabel u2) =
     nonDetCmpUnique u1 u2
+  compare (UnliftedDataLabel u1 _) (UnliftedDataLabel u2 _) =
+    nonDetCmpUnique u1 u2
   compare IdLabel{} _ = LT
   compare _ IdLabel{} = GT
   compare CmmLabel{} _ = LT
@@ -401,6 +408,8 @@ instance Ord CLabel where
   compare _ SRTLabel{} = GT
   compare (IPE_Label {}) _ = LT
   compare  _ (IPE_Label{}) = GT
+  compare UnliftedDataLabel{} _ = LT
+  compare _ UnliftedDataLabel{} = GT
 
 -- | Record where a foreign label is stored.
 data ForeignLabelSource
@@ -731,6 +740,8 @@ isStaticClosureLabel :: CLabel -> Bool
 isStaticClosureLabel (IdLabel _ _ Closure) = True
 -- Closure defined in cmm
 isStaticClosureLabel (CmmLabel _ _ _ CmmClosure) = True
+-- Unlifted data allocated in the data
+isStaticClosureLabel UnliftedDataLabel{} = True
 isStaticClosureLabel _lbl = False
 
 -- | Whether label is a .rodata label
@@ -845,9 +856,10 @@ mkAsmTempDieLabel l = mkAsmTempDerivedLabel l (fsLit "_die")
 
 toClosureLbl :: Platform -> CLabel -> CLabel
 toClosureLbl platform lbl = case lbl of
-   IdLabel n c _        -> IdLabel n c Closure
-   CmmLabel m ext str _ -> CmmLabel m ext str CmmClosure
-   _                    -> pprPanic "toClosureLbl" (pprDebugCLabel platform lbl)
+   IdLabel n c _         -> IdLabel n c Closure
+   CmmLabel m ext str _  -> CmmLabel m ext str CmmClosure
+   l@UnliftedDataLabel{} -> l
+   _                     -> pprPanic "toClosureLbl" (pprDebugCLabel platform lbl)
 
 toSlowEntryLbl :: Platform -> CLabel -> CLabel
 toSlowEntryLbl platform lbl = case lbl of
@@ -914,7 +926,7 @@ hasCAF _                            = False
 -- -----------------------------------------------------------------------------
 -- Does a CLabel need declaring before use or not?
 --
--- See wiki:commentary/compiler/backends/ppr-c#prototypes
+-- See wiki: https://gitlab.haskell.org/ghc/ghc/-/wikis/commentary/compiler/backends/ppr-c#prototypes
 
 needsCDecl :: CLabel -> Bool
   -- False <=> it's pre-declared; don't bother
@@ -946,10 +958,11 @@ needsCDecl (CC_Label _)                 = True
 needsCDecl (CCS_Label _)                = True
 needsCDecl (IPE_Label {})               = True
 needsCDecl (HpcTicksLabel _)            = True
+needsCDecl UnliftedDataLabel{}          = True
+
 needsCDecl (DynamicLinkerLabel {})      = panic "needsCDecl DynamicLinkerLabel"
 needsCDecl PicBaseLabel                 = panic "needsCDecl PicBaseLabel"
 needsCDecl (DeadStripPreventer {})      = panic "needsCDecl DeadStripPreventer"
-
 -- | If a label is a local block label then return just its 'BlockId', otherwise
 -- 'Nothing'.
 maybeLocalBlockLabel :: CLabel -> Maybe BlockId
@@ -1072,6 +1085,7 @@ externallyVisibleCLabel (DynamicLinkerLabel _ _)  = False
 externallyVisibleCLabel (HpcTicksLabel _)       = True
 externallyVisibleCLabel (LargeBitmapLabel _)    = False
 externallyVisibleCLabel (SRTLabel _)            = False
+externallyVisibleCLabel UnliftedDataLabel{}     = False
 externallyVisibleCLabel (PicBaseLabel {}) = panic "externallyVisibleCLabel PicBaseLabel"
 externallyVisibleCLabel (DeadStripPreventer {}) = panic "externallyVisibleCLabel DeadStripPreventer"
 
@@ -1133,6 +1147,7 @@ labelType PicBaseLabel                          = DataLabel
 labelType (DeadStripPreventer _)                = DataLabel
 labelType (HpcTicksLabel _)                     = DataLabel
 labelType (LargeBitmapLabel _)                  = DataLabel
+labelType UnliftedDataLabel{}                   = GcPtrLabel
 
 idInfoLabelType :: IdLabelInfo -> CLabelType
 idInfoLabelType info =
@@ -1402,6 +1417,9 @@ pprCLabel !platform !sty lbl = -- see Note [Bangs in CLabel]
    SRTLabel u
       -> maybe_underscore $ tempLabelPrefixOrUnderscore platform <> pprUniqueAlways u <> pp_cSEP <> text "srt"
 
+   UnliftedDataLabel u op
+      -> maybe_underscore $ tempLabelPrefixOrUnderscore platform <> pprUniqueAlways u <> pp_cSEP <> ppr op
+
    RtsLabel (RtsApFast (NonDetFastString str))
       -> maybe_underscore $ ftext str <> text "_fast"
 
diff --git a/compiler/GHC/Cmm/Info/Build.hs b/compiler/GHC/Cmm/Info/Build.hs
index ab0c32996e..88def15622 100644
--- a/compiler/GHC/Cmm/Info/Build.hs
+++ b/compiler/GHC/Cmm/Info/Build.hs
@@ -56,8 +56,8 @@ import GHC.Types.Name.Set
 
 {- Note [SRTs]
 
-SRTs are the mechanism by which the garbage collector can determine
-the live CAFs in the program.
+Static Reference Tables (SRTs) are the mechanism by which the garbage collector
+can determine the live CAFs in the program.
 
 Representation
 ^^^^^^^^^^^^^^
@@ -1130,7 +1130,6 @@ buildSRTChain profile cafSet =
   where
     mAX_SRT_SIZE = 16
 
-
 buildSRT :: Profile -> [SRTEntry] -> UniqSM (CmmDeclSRTs, SRTEntry)
 buildSRT profile refs = do
   id <- getUniqueM
@@ -1140,6 +1139,7 @@ buildSRT profile refs = do
     srt_n_info = mkSRTInfoLabel (length refs)
     fields =
       mkStaticClosure profile srt_n_info dontCareCCS
+        [] -- no header
         [ CmmLabel lbl | SRTEntry lbl <- refs ]
         [] -- no padding
         [mkIntCLit platform 0] -- link field
diff --git a/compiler/GHC/Cmm/Parser.y b/compiler/GHC/Cmm/Parser.y
index 712a7a5e8a..c6e069279d 100644
--- a/compiler/GHC/Cmm/Parser.y
+++ b/compiler/GHC/Cmm/Parser.y
@@ -437,7 +437,7 @@ static  :: { CmmParse [CmmStatic] }
                         mkStaticClosure profile (mkForeignLabel $3 Nothing ForeignLabelInExternalPackage IsData)
                          -- mkForeignLabel because these are only used
                          -- for CHARLIKE and INTLIKE closures in the RTS.
-                        dontCareCCS (map getLit lits) [] [] [] } }
+                        dontCareCCS (map getLit lits) [] [] [] [] } }
         -- arrays of closures required for the CHARLIKE & INTLIKE arrays
 
 lits    :: { [CmmParse CmmExpr] }
@@ -1245,7 +1245,7 @@ profilingInfo profile desc_str ty_str
 staticClosure :: UnitId -> FastString -> FastString -> [CmmLit] -> CmmParse ()
 staticClosure pkg cl_label info payload
   = do profile <- getProfile
-       let lits = mkStaticClosure profile (mkCmmInfoLabel pkg info) dontCareCCS payload [] [] []
+       let lits = mkStaticClosure profile (mkCmmInfoLabel pkg info) dontCareCCS payload [] [] [] []
        code $ emitDataLits (mkCmmDataLabel pkg (NeedExternDecl True) cl_label) lits
 
 foreignCall
diff --git a/compiler/GHC/StgToCmm/Closure.hs b/compiler/GHC/StgToCmm/Closure.hs
index 3ff745a719..bba8d91643 100644
--- a/compiler/GHC/StgToCmm/Closure.hs
+++ b/compiler/GHC/StgToCmm/Closure.hs
@@ -66,6 +66,7 @@ module GHC.StgToCmm.Closure (
         cafBlackHoleInfoTable,
         indStaticInfoTable,
         staticClosureNeedsLink,
+        smallArrayStaticInfoTable,
     ) where
 
 import GHC.Prelude
@@ -948,6 +949,14 @@ indStaticInfoTable
                  , cit_srt  = Nothing
                  , cit_clo  = Nothing }
 
+smallArrayStaticInfoTable :: WordOff -> CmmInfoTable
+smallArrayStaticInfoTable n
+  = CmmInfoTable { cit_lbl  = mkSMAP_FROZEN_DIRTY_infoLabel
+                 , cit_rep  = smallArrPtrsRep (fromIntegral n)
+                 , cit_prof = NoProfilingInfo
+                 , cit_srt  = Nothing
+                 , cit_clo  = Nothing }
+
 staticClosureNeedsLink :: Bool -> CmmInfoTable -> Bool
 -- A static closure needs a link field to aid the GC when traversing
 -- the static closure graph.  But it only needs such a field if either
diff --git a/compiler/GHC/StgToCmm/Heap.hs b/compiler/GHC/StgToCmm/Heap.hs
index 16161cb028..7660890951 100644
--- a/compiler/GHC/StgToCmm/Heap.hs
+++ b/compiler/GHC/StgToCmm/Heap.hs
@@ -173,11 +173,14 @@ mkStaticClosureFields
         -> [CmmLit]             -- Payload
         -> [CmmLit]             -- The full closure
 mkStaticClosureFields profile info_tbl ccs caf_refs payload
-  = mkStaticClosure profile info_lbl ccs payload padding
+  = mkStaticClosure profile info_lbl ccs header payload padding
         static_link_field saved_info_field
   where
     platform = profilePlatform profile
     info_lbl = cit_lbl info_tbl
+    header = case cit_rep info_tbl of
+      SmallArrayPtrsRep size -> [mkIntCLit (targetPlatform dflags) size]
+      _ -> []
 
     -- CAFs must have consistent layout, regardless of whether they
     -- are actually updatable or not.  The layout of a CAF is:
@@ -218,11 +221,12 @@ mkStaticClosureFields profile info_tbl ccs caf_refs payload
                                       -- See Note [STATIC_LINK fields]
                                       -- in rts/sm/Storage.h
 
-mkStaticClosure :: Profile -> CLabel -> CostCentreStack -> [CmmLit]
+mkStaticClosure :: Profile -> CLabel -> CostCentreStack -> [CmmLit] -> [CmmLit]
   -> [CmmLit] -> [CmmLit] -> [CmmLit] -> [CmmLit]
-mkStaticClosure profile info_lbl ccs payload padding static_link_field saved_info_field
+mkStaticClosure dflags info_lbl ccs header payload padding static_link_field saved_info_field
   =  [CmmLabel info_lbl]
   ++ staticProfHdr profile ccs
+  ++ header
   ++ payload
   ++ padding
   ++ static_link_field
diff --git a/compiler/GHC/StgToCmm/Prim.hs b/compiler/GHC/StgToCmm/Prim.hs
index 99a6960a68..49cf10e33d 100644
--- a/compiler/GHC/StgToCmm/Prim.hs
+++ b/compiler/GHC/StgToCmm/Prim.hs
@@ -53,6 +53,9 @@ import Data.Maybe
 
 import Control.Monad (liftM, when, unless)
 
+import GHC.Types.CostCentre (dontCareCCS)
+import GHC.StgToCmm.Closure
+
 ------------------------------------------------------------------------
 --      Primitive operations and foreign calls
 ------------------------------------------------------------------------
@@ -223,15 +226,27 @@ emitPrimOp dflags primop = case primop of
         (replicate (fromIntegral n) init)
     _ -> PrimopCmmEmit_External
 
-  SmallArrayOfOp -> \elems -> opAllDone $ \[res] ->
+  op@SmallArrayOfOp -> \elems -> opAllDone $ \[res] -> do
     let n = length elems
-      in doNewArrayOp
-          res
-          (smallArrPtrsRep (fromIntegral n))
-          mkSMAP_FROZEN_DIRTY_infoLabel
-          [ ( mkIntExpr platform n
-            , fixedHdrSize dflags + oFFSET_StgSmallMutArrPtrs_ptrs dflags ) ]
-          elems
+    case allStatic elems of
+      Just known -> do
+        u <- newUnique
+        let lbl = mkUnliftedDataLabel u op
+        emitDataCon lbl (smallArrayStaticInfoTable n) dontCareCCS known
+        emit $ mkAssign (CmmLocal res) (CmmLit $ CmmLabel lbl)
+      Nothing -> doNewArrayOp
+        res
+        (smallArrPtrsRep (fromIntegral n))
+        mkSMAP_FROZEN_DIRTY_infoLabel
+        [ ( mkIntExpr platform n
+          , fixedHdrSize dflags + oFFSET_StgSmallMutArrPtrs_ptrs dflags ) ]
+        elems
+      where
+        -- todo: comment
+        allStatic = foldr step (Just [])
+
+        step (CmmLit l) (Just acc) = Just (l : acc) -- c.f. XXX getLit
+        step _ _ = Nothing
 
   CopySmallArrayOp -> \case
     [src, src_off, dst, dst_off, (CmmLit (CmmInt n _))] ->