Remote GHCi: Optimize the serialization/deserialization of byte code

Summary: This cuts allocations by about a quarter.

Test Plan:
* validate
* `ghci -fexternal-interpreter` in `nofib/real/anna`

Reviewers: niteria, bgamari, ezyang, austin, hvr, erikd

Subscribers: thomie

Differential Revision: https://phabricator.haskell.org/D1875

GHC Trac Issues: #11100

2 files changed, 67 insertions, 11 deletions

diff --git a/compiler/ghci/ByteCodeLink.hs b/compiler/ghci/ByteCodeLink.hs
index 74f490b8fd..c108d1665c 100644
--- a/compiler/ghci/ByteCodeLink.hs
+++ b/compiler/ghci/ByteCodeLink.hs
@@ -28,6 +28,7 @@ import SizedSeq
 import GHCi
 import ByteCodeTypes
 import HscTypes
+import DynFlags
 import Name
 import NameEnv
 import PrimOp
@@ -39,6 +40,8 @@ import Util
 
 -- Standard libraries
 import Data.Array.Unboxed
+import Data.Array.Base
+import Data.Word
 import Foreign.Ptr
 import GHC.IO           ( IO(..) )
 import GHC.Exts
@@ -68,10 +71,19 @@ linkBCO hsc_env ie ce bco_ix breakarray
            (UnlinkedBCO _ arity insns bitmap lits0 ptrs0) = do
   lits <- mapM (lookupLiteral hsc_env ie) (ssElts lits0)
   ptrs <- mapM (resolvePtr hsc_env ie ce bco_ix breakarray) (ssElts ptrs0)
-  return (ResolvedBCO arity insns bitmap
+  let dflags = hsc_dflags hsc_env
+  return (ResolvedBCO arity (toWordArray dflags insns) bitmap
               (listArray (0, fromIntegral (sizeSS lits0)-1) lits)
               (addListToSS emptySS ptrs))
 
+-- Turn the insns array from a Word16 array into a Word array.  The
+-- latter is much faster to serialize/deserialize. Assumes the input
+-- array is zero-indexed.
+toWordArray :: DynFlags -> UArray Int Word16 -> UArray Int Word
+toWordArray dflags (UArray _ _ n arr) = UArray 0 (n'-1) n' arr
+  where n' = (n + w16s_per_word - 1) `quot` w16s_per_word
+        w16s_per_word = wORD_SIZE dflags `quot` 2
+
 lookupLiteral :: HscEnv -> ItblEnv -> BCONPtr -> IO Word
 lookupLiteral _ _ (BCONPtrWord lit) = return lit
 lookupLiteral hsc_env _ (BCONPtrLbl  sym) = do
diff --git a/libraries/ghci/GHCi/ResolvedBCO.hs b/libraries/ghci/GHCi/ResolvedBCO.hs
index a349dedaba..aa63d36c9c 100644
--- a/libraries/ghci/GHCi/ResolvedBCO.hs
+++ b/libraries/ghci/GHCi/ResolvedBCO.hs
@@ -1,4 +1,5 @@
-{-# LANGUAGE RecordWildCards, DeriveGeneric, GeneralizedNewtypeDeriving #-}
+{-# LANGUAGE RecordWildCards, DeriveGeneric, GeneralizedNewtypeDeriving,
+    BangPatterns #-}
 module GHCi.ResolvedBCO
   ( ResolvedBCO(..)
   , ResolvedBCOPtr(..)
@@ -8,38 +9,81 @@ import SizedSeq
 import GHCi.RemoteTypes
 import GHCi.BreakArray
 
+import Control.Monad.ST
 import Data.Array.Unboxed
+import Data.Array.Base
 import Data.Binary
 import GHC.Generics
 
 -- -----------------------------------------------------------------------------
 -- ResolvedBCO
 
--- A ResolvedBCO is one in which all the Name references have been
+-- A A ResolvedBCO is one in which all the Name references have been
 -- resolved to actual addresses or RemoteHValues.
-
+--
+-- Note, all arrays are zero-indexed (we assume this when
+-- serializing/deserializing)
 data ResolvedBCO
    = ResolvedBCO {
-        resolvedBCOArity  :: Int,
-        resolvedBCOInstrs :: UArray Int Word16,         -- insns
+        resolvedBCOArity  :: {-# UNPACK #-} !Int,
+        resolvedBCOInstrs :: UArray Int Word,           -- insns
         resolvedBCOBitmap :: UArray Int Word,           -- bitmap
         resolvedBCOLits   :: UArray Int Word,           -- non-ptrs
         resolvedBCOPtrs   :: (SizedSeq ResolvedBCOPtr)  -- ptrs
    }
    deriving (Generic, Show)
 
-instance Binary ResolvedBCO
+instance Binary ResolvedBCO where
+  put ResolvedBCO{..} = do
+    put resolvedBCOArity
+    putArray resolvedBCOInstrs
+    putArray resolvedBCOBitmap
+    putArray resolvedBCOLits
+    put resolvedBCOPtrs
+  get = ResolvedBCO <$> get <*> getArray <*> getArray <*> getArray <*> get
+
+-- Specialized versions of the binary get/put for UArray Int Word.
+-- This saves a bit of time and allocation over using the default
+-- get/put, because we get specialisd code and also avoid serializing
+-- the bounds.
+putArray :: UArray Int Word -> Put
+putArray a@(UArray _ _ n _) = do
+  put n
+  mapM_ put (elems a)
+
+getArray :: Get (UArray Int Word)
+getArray = do
+  n  <- get
+  xs <- gets n []
+  return $! mkArray n xs
+ where
+  gets 0 xs = return xs
+  gets n xs = do
+    x <- get
+    gets (n-1) (x:xs)
+
+  mkArray :: Int -> [Word] -> UArray Int Word
+  mkArray n0 xs0 = runST $ do
+    !marr <- newArray (0,n0-1) 0
+    let go 0 _ = return ()
+        go _ [] = error "mkArray"
+        go n (x:xs) = do
+          let n' = n-1
+          unsafeWrite marr n' x
+          go n' xs
+    go n0 xs0
+    unsafeFreezeSTUArray marr
 
 data ResolvedBCOPtr
-  = ResolvedBCORef Int
+  = ResolvedBCORef {-# UNPACK #-} !Int
       -- ^ reference to the Nth BCO in the current set
-  | ResolvedBCOPtr (RemoteRef HValue)
+  | ResolvedBCOPtr {-# UNPACK #-} !(RemoteRef HValue)
       -- ^ reference to a previously created BCO
-  | ResolvedBCOStaticPtr (RemotePtr ())
+  | ResolvedBCOStaticPtr {-# UNPACK #-} !(RemotePtr ())
       -- ^ reference to a static ptr
   | ResolvedBCOPtrBCO ResolvedBCO
       -- ^ a nested BCO
-  | ResolvedBCOPtrBreakArray (RemoteRef BreakArray)
+  | ResolvedBCOPtrBreakArray {-# UNPACK #-} !(RemoteRef BreakArray)
       -- ^ Resolves to the MutableArray# inside the BreakArray
   deriving (Generic, Show)