Perf: use SmallArray for primops' Ids cache (#20857)

SmallArray doesn't perform bounds check (faster).

Make primop tags start at 0 to avoid index arithmetic.

4 files changed, 112 insertions, 11 deletions

diff --git a/compiler/GHC/Builtin/Utils.hs b/compiler/GHC/Builtin/Utils.hs
index 9d91b1246d..4428716681 100644
--- a/compiler/GHC/Builtin/Utils.hs
+++ b/compiler/GHC/Builtin/Utils.hs
@@ -78,10 +78,10 @@ import GHC.Hs.Doc
 import GHC.Unit.Module.ModIface (IfaceExport)
 
 import GHC.Data.List.SetOps
+import GHC.Data.SmallArray
 
 import Control.Applicative ((<|>))
 import Data.List        ( intercalate , find )
-import Data.Array
 import Data.Maybe
 import qualified Data.Map as Map
 
@@ -133,7 +133,7 @@ knownKeyNames
              , concatMap wired_tycon_kk_names wiredInTyCons
              , concatMap wired_tycon_kk_names typeNatTyCons
              , map idName wiredInIds
-             , map (idName . primOpId) allThePrimOps
+             , map idName allThePrimOpIds
              , map (idName . primOpWrapperId) allThePrimOps
              , basicKnownKeyNames
              , templateHaskellNames
@@ -238,13 +238,21 @@ sense of them in interface pragmas. It's cool, though they all have
 ************************************************************************
 -}
 
-primOpIds :: Array Int Id
--- A cache of the PrimOp Ids, indexed by PrimOp tag
-primOpIds = array (1,maxPrimOpTag) [ (primOpTag op, mkPrimOpId op)
-                                   | op <- allThePrimOps ]
+-- | A cache of the PrimOp Ids, indexed by PrimOp tag (0 indexed)
+primOpIds :: SmallArray Id
+{-# NOINLINE primOpIds #-}
+primOpIds = listToArray (maxPrimOpTag+1) primOpTag mkPrimOpId allThePrimOps
 
+-- | Get primop id.
+--
+-- Retrieve it from `primOpIds` cache without performing bounds checking.
 primOpId :: PrimOp -> Id
-primOpId op = primOpIds ! primOpTag op
+primOpId op = indexSmallArray primOpIds (primOpTag op)
+
+-- | All the primop ids, as a list
+allThePrimOpIds :: [Id]
+{-# INLINE allThePrimOpIds #-}
+allThePrimOpIds = map (indexSmallArray primOpIds) [0..maxPrimOpTag]
 
 {-
 ************************************************************************
@@ -257,7 +265,7 @@ primOpId op = primOpIds ! primOpTag op
 ghcPrimExports :: [IfaceExport]
 ghcPrimExports
  = map (avail . idName) ghcPrimIds ++
-   map (avail . idName . primOpId) allThePrimOps ++
+   map (avail . idName) allThePrimOpIds ++
    [ availTC n [n] []
    | tc <- exposedPrimTyCons, let n = tyConName tc  ]
 
@@ -265,7 +273,7 @@ ghcPrimDeclDocs :: DeclDocMap
 ghcPrimDeclDocs = DeclDocMap $ Map.fromList $ mapMaybe findName primOpDocs
   where
     names = map idName ghcPrimIds ++
-            map (idName . primOpId) allThePrimOps ++
+            map idName allThePrimOpIds ++
             map tyConName exposedPrimTyCons
     findName (nameStr, doc)
       | Just name <- find ((nameStr ==) . getOccString) names
diff --git a/compiler/GHC/Data/SmallArray.hs b/compiler/GHC/Data/SmallArray.hs
new file mode 100644
index 0000000000..2697c8380b
--- /dev/null
+++ b/compiler/GHC/Data/SmallArray.hs
@@ -0,0 +1,92 @@
+{-# LANGUAGE MagicHash #-}
+{-# LANGUAGE UnboxedTuples #-}
+{-# LANGUAGE BlockArguments #-}
+
+-- | Small-array
+module GHC.Data.SmallArray
+  ( SmallMutableArray (..)
+  , SmallArray (..)
+  , newSmallArray
+  , writeSmallArray
+  , freezeSmallArray
+  , unsafeFreezeSmallArray
+  , indexSmallArray
+  , listToArray
+  )
+where
+
+import GHC.Exts
+import GHC.Prelude
+import GHC.ST
+
+data SmallArray a = SmallArray (SmallArray# a)
+
+data SmallMutableArray s a = SmallMutableArray (SmallMutableArray# s a)
+
+newSmallArray
+  :: Int  -- ^ size
+  -> a    -- ^ initial contents
+  -> State# s
+  -> (# State# s, SmallMutableArray s a #)
+{-# INLINE newSmallArray #-}
+newSmallArray (I# sz) x s = case newSmallArray# sz x s of
+  (# s', a #) -> (# s', SmallMutableArray a #)
+
+writeSmallArray
+  :: SmallMutableArray s a -- ^ array
+  -> Int                   -- ^ index
+  -> a                     -- ^ new element
+  -> State# s
+  -> State# s
+{-# INLINE writeSmallArray #-}
+writeSmallArray (SmallMutableArray a) (I# i) x = writeSmallArray# a i x
+
+
+-- | Copy and freeze a slice of a mutable array.
+freezeSmallArray
+  :: SmallMutableArray s a -- ^ source
+  -> Int                   -- ^ offset
+  -> Int                   -- ^ length
+  -> State# s
+  -> (# State# s, SmallArray a #)
+{-# INLINE freezeSmallArray #-}
+freezeSmallArray (SmallMutableArray ma) (I# offset) (I# len) s =
+  case freezeSmallArray# ma offset len s of
+    (# s', a #) -> (# s', SmallArray a #)
+
+-- | Freeze a mutable array (no copy!)
+unsafeFreezeSmallArray
+  :: SmallMutableArray s a
+  -> State# s
+  -> (# State# s, SmallArray a #)
+{-# INLINE unsafeFreezeSmallArray #-}
+unsafeFreezeSmallArray (SmallMutableArray ma) s =
+  case unsafeFreezeSmallArray# ma s of
+    (# s', a #) -> (# s', SmallArray a #)
+
+
+-- | Index a small-array (no bounds checking!)
+indexSmallArray
+  :: SmallArray a -- ^ array
+  -> Int          -- ^ index
+  -> a
+{-# INLINE indexSmallArray #-}
+indexSmallArray (SmallArray sa#) (I# i) = case indexSmallArray# sa# i of
+  (# v #) -> v
+
+
+-- | Convert a list into an array.
+listToArray :: Int -> (e -> Int) -> (e -> a) -> [e] -> SmallArray a
+{-# INLINE listToArray #-}
+listToArray (I# size) index_of value_of xs = runST $ ST \s ->
+  let
+    index_of' e = case index_of e of I# i -> i
+    write_elems ma es s = case es of
+      []    -> s
+      e:es' -> case writeSmallArray# ma (index_of' e) (value_of e) s of
+                 s' -> write_elems ma es' s'
+  in
+  case newSmallArray# size undefined s of
+    (# s', ma #) -> case write_elems ma xs s' of
+      s'' -> case unsafeFreezeSmallArray# ma s'' of
+        (# s''', a #) -> (# s''', SmallArray a #)
diff --git a/compiler/ghc.cabal.in b/compiler/ghc.cabal.in
index 487fd7971c..89c6bfb51f 100644
--- a/compiler/ghc.cabal.in
+++ b/compiler/ghc.cabal.in
@@ -376,6 +376,7 @@ Library
         GHC.Data.Maybe
         GHC.Data.OrdList
         GHC.Data.Pair
+        GHC.Data.SmallArray
         GHC.Data.Stream
         GHC.Data.Strict
         GHC.Data.StringBuffer
diff --git a/utils/genprimopcode/Main.hs b/utils/genprimopcode/Main.hs
index c13447e527..2e0886e59b 100644
--- a/utils/genprimopcode/Main.hs
+++ b/utils/genprimopcode/Main.hs
@@ -741,13 +741,13 @@ gen_primop_vector_tycons (Info _ entries)
 gen_primop_tag :: Info -> String
 gen_primop_tag (Info _ entries)
    = unlines (max_def_type : max_def :
-              tagOf_type : zipWith f primop_entries [1 :: Int ..])
+              tagOf_type : zipWith f primop_entries [0 :: Int ..])
      where
         primop_entries = concatMap desugarVectorSpec $ filter is_primop entries
         tagOf_type = "primOpTag :: PrimOp -> Int"
         f i n = "primOpTag " ++ cons i ++ " = " ++ show n
         max_def_type = "maxPrimOpTag :: Int"
-        max_def      = "maxPrimOpTag = " ++ show (length primop_entries)
+        max_def      = "maxPrimOpTag = " ++ show (length primop_entries - 1)
 
 gen_data_decl :: Info -> String
 gen_data_decl (Info _ entries) =