{-# LANGUAGE CPP #-}

module GHC.StgToCmm.InfoTableProv (emitIpeBufferListNode) where

import Foreign

#if defined(HAVE_LIBZSTD)
import Foreign.C.Types
import qualified Data.ByteString.Internal as BSI
import GHC.IO (unsafePerformIO)
#endif

import GHC.Prelude
import GHC.Platform
import GHC.Types.SrcLoc (pprUserRealSpan, srcSpanFile)
import GHC.Unit.Module
import GHC.Utils.Outputable
import GHC.Data.FastString (fastStringToShortText, unpackFS, LexicalFastString(..))

import GHC.Cmm
import GHC.Cmm.CLabel
import GHC.Cmm.Utils

import GHC.StgToCmm.Config
import GHC.StgToCmm.Monad

import GHC.Data.ShortText (ShortText)
import qualified GHC.Data.ShortText as ST

import Control.Monad.Trans.State.Strict

import qualified Data.ByteString as BS
import qualified Data.ByteString.Builder as BSB
import qualified Data.ByteString.Lazy as BSL
import qualified Data.Map.Strict as M

{-
Note [Compression and Decompression of IPE data]
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

Compiling with `-finfo-table-map` causes build results to include a map from
info tables to source positions called the info table provenance entry (IPE)
map. See Note [Mapping Info Tables to Source Positions]. The IPE information
can grow the size of build results significantly. At the time of writing, a
default build of GHC results in a total of 109M of libHSghc-*.so build results.
A default+ipe build of GHC (see ./hadrian/doc/flavours.md) results in 262M of
libHSghc-*.so build results without compression.

We reduce the impact of IPE data on the size of build results by compressing
the data before it is emitted using the zstd compression library. See
Note [The Info Table Provenance Entry (IPE) Map] for information on the layout
of IPE data on disk and in the RTS. We cannot simply compress all data held in
the IPE entry buffer, as the pointers to info tables must be converted to
memory addresses during linking. Therefore, we can only compress the strings
table and the IPE entries themselves (which essentially only consist of indices
into the strings table).

With compression, a default+ipe build of GHC results in a total of 205M of
libHSghc-*.so build results. This is over a 20% reduction from the uncompressed
case.

Decompression happens lazily, as it only occurs when the IPE map is
constructed (which is also done lazily on first lookup or traversal). During
construction, the 'compressed' field of each IPE buffer list node is examined.
If the field indicates that the data has been compressed, the entry data and
strings table are decompressed before continuing with the normal IPE map
construction.
-}

emitIpeBufferListNode ::
     Module
  -> [InfoProvEnt]
  -> FCode ()
emitIpeBufferListNode :: Module -> [InfoProvEnt] -> FCode ()
emitIpeBufferListNode Module
_ [] = () -> FCode ()
forall a. a -> FCode a
forall (m :: * -> *) a. Monad m => a -> m a
return ()
emitIpeBufferListNode Module
this_mod [InfoProvEnt]
ents = do
    cfg <- FCode StgToCmmConfig
getStgToCmmConfig

    tables_lbl  <- mkStringLitLabel <$> newUnique
    strings_lbl <- mkStringLitLabel <$> newUnique
    entries_lbl <- mkStringLitLabel <$> newUnique

    let ctx      = StgToCmmConfig -> SDocContext
stgToCmmContext StgToCmmConfig
cfg
        platform = StgToCmmConfig -> Platform
stgToCmmPlatform StgToCmmConfig
cfg
        int Int
n    = Platform -> Int -> CmmLit
mkIntCLit Platform
platform Int
n

        ((cg_ipes, unit_id, module_name), strtab) = flip runState emptyStringTable $ do
          unit_id <- lookupStringTable $ ST.pack $ renderWithContext ctx (ppr $ moduleName this_mod)
          module_name <- lookupStringTable $ ST.pack $ renderWithContext ctx (ppr $ moduleUnit this_mod)
          cg_ipes <- mapM (toCgIPE platform ctx) ents
          return (cg_ipes, unit_id, module_name)

        tables :: [CmmStatic]
        tables = (CgInfoProvEnt -> CmmStatic) -> [CgInfoProvEnt] -> [CmmStatic]
forall a b. (a -> b) -> [a] -> [b]
map (CmmLit -> CmmStatic
CmmStaticLit (CmmLit -> CmmStatic)
-> (CgInfoProvEnt -> CmmLit) -> CgInfoProvEnt -> CmmStatic
forall b c a. (b -> c) -> (a -> b) -> a -> c
. CLabel -> CmmLit
CmmLabel (CLabel -> CmmLit)
-> (CgInfoProvEnt -> CLabel) -> CgInfoProvEnt -> CmmLit
forall b c a. (b -> c) -> (a -> b) -> a -> c
. CgInfoProvEnt -> CLabel
ipeInfoTablePtr) [CgInfoProvEnt]
cg_ipes

        uncompressed_strings :: BS.ByteString
        uncompressed_strings = StringTable -> ByteString
getStringTableStrings StringTable
strtab

        strings_bytes :: BS.ByteString
        strings_bytes = Int -> ByteString -> ByteString
compress Int
defaultCompressionLevel ByteString
uncompressed_strings

        strings :: [CmmStatic]
        strings = [ByteString -> CmmStatic
CmmString ByteString
strings_bytes]

        uncompressed_entries :: BS.ByteString
        uncompressed_entries = ByteOrder -> [CgInfoProvEnt] -> ByteString
toIpeBufferEntries (Platform -> ByteOrder
platformByteOrder Platform
platform) [CgInfoProvEnt]
cg_ipes

        entries_bytes :: BS.ByteString
        entries_bytes = Int -> ByteString -> ByteString
compress Int
defaultCompressionLevel ByteString
uncompressed_entries

        entries :: [CmmStatic]
        entries = [ByteString -> CmmStatic
CmmString ByteString
entries_bytes]

        ipe_buffer_lbl :: CLabel
        ipe_buffer_lbl = Module -> CLabel
mkIPELabel Module
this_mod

        ipe_buffer_node :: [CmmStatic]
        ipe_buffer_node = (CmmLit -> CmmStatic) -> [CmmLit] -> [CmmStatic]
forall a b. (a -> b) -> [a] -> [b]
map CmmLit -> CmmStatic
CmmStaticLit
          [ -- 'next' field
            Platform -> CmmLit
zeroCLit Platform
platform

            -- 'compressed' field
          , Int -> CmmLit
int Int
do_compress

            -- 'count' field
          , Int -> CmmLit
int (Int -> CmmLit) -> Int -> CmmLit
forall a b. (a -> b) -> a -> b
$ [CgInfoProvEnt] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [CgInfoProvEnt]
cg_ipes

            -- 'tables' field
          , CLabel -> CmmLit
CmmLabel CLabel
tables_lbl

            -- 'entries' field
          , CLabel -> CmmLit
CmmLabel CLabel
entries_lbl

            -- 'entries_size' field (decompressed size)
          , Int -> CmmLit
int (Int -> CmmLit) -> Int -> CmmLit
forall a b. (a -> b) -> a -> b
$ ByteString -> Int
BS.length ByteString
uncompressed_entries

            -- 'string_table' field
          , CLabel -> CmmLit
CmmLabel CLabel
strings_lbl

            -- 'string_table_size' field (decompressed size)
          , Int -> CmmLit
int (Int -> CmmLit) -> Int -> CmmLit
forall a b. (a -> b) -> a -> b
$ ByteString -> Int
BS.length ByteString
uncompressed_strings

            -- 'module_name' field
          , Integer -> Width -> CmmLit
CmmInt (StrTabOffset -> Integer
forall a b. (Integral a, Num b) => a -> b
fromIntegral StrTabOffset
module_name) Width
W32

            -- 'unit_id' field
          , Integer -> Width -> CmmLit
CmmInt (StrTabOffset -> Integer
forall a b. (Integral a, Num b) => a -> b
fromIntegral StrTabOffset
unit_id) Width
W32
          ]

    -- Emit the list of info table pointers
    emitDecl $ CmmData
      (Section Data tables_lbl)
      (CmmStaticsRaw tables_lbl tables)

    -- Emit the strings table
    emitDecl $ CmmData
      (Section Data strings_lbl)
      (CmmStaticsRaw strings_lbl strings)

    -- Emit the list of IPE buffer entries
    emitDecl $ CmmData
      (Section Data entries_lbl)
      (CmmStaticsRaw entries_lbl entries)

    -- Emit the IPE buffer list node
    emitDecl $ CmmData
      (Section Data ipe_buffer_lbl)
      (CmmStaticsRaw ipe_buffer_lbl ipe_buffer_node)

-- | Emit the fields of an IpeBufferEntry struct for each entry in a given list.
toIpeBufferEntries ::
     ByteOrder       -- ^ Byte order to write the data in
  -> [CgInfoProvEnt] -- ^ List of IPE buffer entries
  -> BS.ByteString
toIpeBufferEntries :: ByteOrder -> [CgInfoProvEnt] -> ByteString
toIpeBufferEntries ByteOrder
byte_order [CgInfoProvEnt]
cg_ipes =
      LazyByteString -> ByteString
BSL.toStrict (LazyByteString -> ByteString)
-> ([Builder] -> LazyByteString) -> [Builder] -> ByteString
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Builder -> LazyByteString
BSB.toLazyByteString (Builder -> LazyByteString)
-> ([Builder] -> Builder) -> [Builder] -> LazyByteString
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Builder] -> Builder
forall a. Monoid a => [a] -> a
mconcat
    ([Builder] -> ByteString) -> [Builder] -> ByteString
forall a b. (a -> b) -> a -> b
$ (CgInfoProvEnt -> Builder) -> [CgInfoProvEnt] -> [Builder]
forall a b. (a -> b) -> [a] -> [b]
map ([Builder] -> Builder
forall a. Monoid a => [a] -> a
mconcat ([Builder] -> Builder)
-> (CgInfoProvEnt -> [Builder]) -> CgInfoProvEnt -> Builder
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (StrTabOffset -> Builder) -> [StrTabOffset] -> [Builder]
forall a b. (a -> b) -> [a] -> [b]
map StrTabOffset -> Builder
word32Builder ([StrTabOffset] -> [Builder])
-> (CgInfoProvEnt -> [StrTabOffset]) -> CgInfoProvEnt -> [Builder]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. CgInfoProvEnt -> [StrTabOffset]
to_ipe_buf_ent) [CgInfoProvEnt]
cg_ipes
  where
    to_ipe_buf_ent :: CgInfoProvEnt -> [Word32]
    to_ipe_buf_ent :: CgInfoProvEnt -> [StrTabOffset]
to_ipe_buf_ent CgInfoProvEnt
cg_ipe =
      [ CgInfoProvEnt -> StrTabOffset
ipeTableName CgInfoProvEnt
cg_ipe
      , CgInfoProvEnt -> StrTabOffset
ipeClosureDesc CgInfoProvEnt
cg_ipe
      , CgInfoProvEnt -> StrTabOffset
ipeTypeDesc CgInfoProvEnt
cg_ipe
      , CgInfoProvEnt -> StrTabOffset
ipeLabel CgInfoProvEnt
cg_ipe
      , CgInfoProvEnt -> StrTabOffset
ipeSrcFile CgInfoProvEnt
cg_ipe
      , CgInfoProvEnt -> StrTabOffset
ipeSrcSpan CgInfoProvEnt
cg_ipe
      ]

    word32Builder :: Word32 -> BSB.Builder
    word32Builder :: StrTabOffset -> Builder
word32Builder = case ByteOrder
byte_order of
      ByteOrder
BigEndian    -> StrTabOffset -> Builder
BSB.word32BE
      ByteOrder
LittleEndian -> StrTabOffset -> Builder
BSB.word32LE

toCgIPE :: Platform -> SDocContext -> InfoProvEnt -> State StringTable CgInfoProvEnt
toCgIPE :: Platform
-> SDocContext
-> InfoProvEnt
-> StateT StringTable Identity CgInfoProvEnt
toCgIPE Platform
platform SDocContext
ctx InfoProvEnt
ipe = do
    table_name <- ShortText -> State StringTable StrTabOffset
lookupStringTable (ShortText -> State StringTable StrTabOffset)
-> ShortText -> State StringTable StrTabOffset
forall a b. (a -> b) -> a -> b
$ String -> ShortText
ST.pack (String -> ShortText) -> String -> ShortText
forall a b. (a -> b) -> a -> b
$ SDocContext -> SDoc -> String
renderWithContext SDocContext
ctx (Platform -> CLabel -> SDoc
forall doc. IsLine doc => Platform -> CLabel -> doc
pprCLabel Platform
platform (InfoProvEnt -> CLabel
infoTablePtr InfoProvEnt
ipe))
    closure_desc <- lookupStringTable $ ST.pack $ show (infoProvEntClosureType ipe)
    type_desc <- lookupStringTable $ ST.pack $ infoTableType ipe
    let label_str = String
-> ((RealSrcSpan, LexicalFastString) -> String)
-> Maybe (RealSrcSpan, LexicalFastString)
-> String
forall b a. b -> (a -> b) -> Maybe a -> b
maybe String
"" ((\(LexicalFastString FastString
s) -> FastString -> String
unpackFS FastString
s) (LexicalFastString -> String)
-> ((RealSrcSpan, LexicalFastString) -> LexicalFastString)
-> (RealSrcSpan, LexicalFastString)
-> String
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (RealSrcSpan, LexicalFastString) -> LexicalFastString
forall a b. (a, b) -> b
snd) (InfoProvEnt -> Maybe (RealSrcSpan, LexicalFastString)
infoTableProv InfoProvEnt
ipe)
    let (src_loc_file, src_loc_span) =
            case infoTableProv ipe of
              Maybe (RealSrcSpan, LexicalFastString)
Nothing -> (ShortText
forall a. Monoid a => a
mempty, String
"")
              Just (RealSrcSpan
span, LexicalFastString
_) ->
                  let file :: ShortText
file = FastString -> ShortText
fastStringToShortText (FastString -> ShortText) -> FastString -> ShortText
forall a b. (a -> b) -> a -> b
$ RealSrcSpan -> FastString
srcSpanFile RealSrcSpan
span
                      coords :: String
coords = SDocContext -> SDoc -> String
renderWithContext SDocContext
ctx (Bool -> RealSrcSpan -> SDoc
pprUserRealSpan Bool
False RealSrcSpan
span)
                  in (ShortText
file, String
coords)
    label    <- lookupStringTable $ ST.pack label_str
    src_file <- lookupStringTable src_loc_file
    src_span <- lookupStringTable $ ST.pack src_loc_span
    return $ CgInfoProvEnt { ipeInfoTablePtr = infoTablePtr ipe
                           , ipeTableName = table_name
                           , ipeClosureDesc = closure_desc
                           , ipeTypeDesc = type_desc
                           , ipeLabel = label
                           , ipeSrcFile = src_file
                           , ipeSrcSpan = src_span
                           }

data CgInfoProvEnt = CgInfoProvEnt
                               { CgInfoProvEnt -> CLabel
ipeInfoTablePtr :: !CLabel
                               , CgInfoProvEnt -> StrTabOffset
ipeTableName :: !StrTabOffset
                               , CgInfoProvEnt -> StrTabOffset
ipeClosureDesc :: !StrTabOffset
                               , CgInfoProvEnt -> StrTabOffset
ipeTypeDesc :: !StrTabOffset
                               , CgInfoProvEnt -> StrTabOffset
ipeLabel :: !StrTabOffset
                               , CgInfoProvEnt -> StrTabOffset
ipeSrcFile :: !StrTabOffset
                               , CgInfoProvEnt -> StrTabOffset
ipeSrcSpan :: !StrTabOffset
                               }

data StringTable = StringTable { StringTable -> DList ShortText
stStrings :: DList ShortText
                               , StringTable -> Int
stLength :: !Int
                               , StringTable -> Map ShortText StrTabOffset
stLookup :: !(M.Map ShortText StrTabOffset)
                               }

type StrTabOffset = Word32

emptyStringTable :: StringTable
emptyStringTable :: StringTable
emptyStringTable =
    StringTable { stStrings :: DList ShortText
stStrings = DList ShortText
forall a. DList a
emptyDList
                , stLength :: Int
stLength = Int
0
                , stLookup :: Map ShortText StrTabOffset
stLookup = Map ShortText StrTabOffset
forall k a. Map k a
M.empty
                }

getStringTableStrings :: StringTable -> BS.ByteString
getStringTableStrings :: StringTable -> ByteString
getStringTableStrings StringTable
st =
    LazyByteString -> ByteString
BSL.toStrict (LazyByteString -> ByteString) -> LazyByteString -> ByteString
forall a b. (a -> b) -> a -> b
$ Builder -> LazyByteString
BSB.toLazyByteString
    (Builder -> LazyByteString) -> Builder -> LazyByteString
forall a b. (a -> b) -> a -> b
$ (ShortText -> Builder) -> [ShortText] -> Builder
forall m a. Monoid m => (a -> m) -> [a] -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap ShortText -> Builder
f ([ShortText] -> Builder) -> [ShortText] -> Builder
forall a b. (a -> b) -> a -> b
$ DList ShortText -> [ShortText]
forall a. DList a -> [a]
dlistToList (StringTable -> DList ShortText
stStrings StringTable
st)
  where
    f :: ShortText -> Builder
f ShortText
x = ShortByteString -> Builder
BSB.shortByteString (ShortText -> ShortByteString
ST.contents ShortText
x) Builder -> Builder -> Builder
forall a. Monoid a => a -> a -> a
`mappend` Word8 -> Builder
BSB.word8 Word8
0

lookupStringTable :: ShortText -> State StringTable StrTabOffset
lookupStringTable :: ShortText -> State StringTable StrTabOffset
lookupStringTable ShortText
str = (StringTable -> (StrTabOffset, StringTable))
-> State StringTable StrTabOffset
forall (m :: * -> *) s a. Monad m => (s -> (a, s)) -> StateT s m a
state ((StringTable -> (StrTabOffset, StringTable))
 -> State StringTable StrTabOffset)
-> (StringTable -> (StrTabOffset, StringTable))
-> State StringTable StrTabOffset
forall a b. (a -> b) -> a -> b
$ \StringTable
st ->
    case ShortText -> Map ShortText StrTabOffset -> Maybe StrTabOffset
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup ShortText
str (StringTable -> Map ShortText StrTabOffset
stLookup StringTable
st) of
      Just StrTabOffset
off -> (StrTabOffset
off, StringTable
st)
      Maybe StrTabOffset
Nothing ->
          let !st' :: StringTable
st' = StringTable
st { stStrings = stStrings st `snoc` str
                        , stLength  = stLength st + ST.byteLength str + 1
                        , stLookup  = M.insert str res (stLookup st)
                        }
              res :: StrTabOffset
res = Int -> StrTabOffset
forall a b. (Integral a, Num b) => a -> b
fromIntegral (StringTable -> Int
stLength StringTable
st)
          in (StrTabOffset
res, StringTable
st')

do_compress :: Int
compress    :: Int -> BS.ByteString -> BS.ByteString
#if !defined(HAVE_LIBZSTD)
do_compress :: Int
do_compress   = Int
0
compress :: Int -> ByteString -> ByteString
compress Int
_ ByteString
bs = ByteString
bs
#else
do_compress = 1

compress clvl (BSI.PS srcForeignPtr off len) = unsafePerformIO $
    withForeignPtr srcForeignPtr $ \srcPtr -> do
      maxCompressedSize <- zstd_compress_bound $ fromIntegral len
      dstForeignPtr <- BSI.mallocByteString (fromIntegral maxCompressedSize)
      withForeignPtr dstForeignPtr $ \dstPtr -> do
        compressedSize <- fromIntegral <$>
          zstd_compress
            dstPtr
            maxCompressedSize
            (srcPtr `plusPtr` off)
            (fromIntegral len)
            (fromIntegral clvl)
        BSI.create compressedSize $ \p -> copyBytes p dstPtr compressedSize

foreign import ccall unsafe "ZSTD_compress"
    zstd_compress ::
         Ptr dst -- ^ Destination buffer
      -> CSize   -- ^ Capacity of destination buffer
      -> Ptr src -- ^ Source buffer
      -> CSize   -- ^ Size of source buffer
      -> CInt    -- ^ Compression level
      -> IO CSize

-- | Compute the maximum compressed size for a given source buffer size
foreign import ccall unsafe "ZSTD_compressBound"
    zstd_compress_bound ::
         CSize -- ^ Size of source buffer
      -> IO CSize
#endif

defaultCompressionLevel :: Int
defaultCompressionLevel :: Int
defaultCompressionLevel = Int
3

newtype DList a = DList ([a] -> [a])

emptyDList :: DList a
emptyDList :: forall a. DList a
emptyDList = ([a] -> [a]) -> DList a
forall a. ([a] -> [a]) -> DList a
DList [a] -> [a]
forall a. a -> a
id

snoc :: DList a -> a -> DList a
snoc :: forall a. DList a -> a -> DList a
snoc (DList [a] -> [a]
f) a
x = ([a] -> [a]) -> DList a
forall a. ([a] -> [a]) -> DList a
DList ([a] -> [a]
f ([a] -> [a]) -> ([a] -> [a]) -> [a] -> [a]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a
xa -> [a] -> [a]
forall a. a -> [a] -> [a]
:))

dlistToList :: DList a -> [a]
dlistToList :: forall a. DList a -> [a]
dlistToList (DList [a] -> [a]
f) = [a] -> [a]
f []