fb-util-0.1.0.1: Util/ByteString.hs

{-
  Copyright (c) Meta Platforms, Inc. and affiliates.
  All rights reserved.

  This source code is licensed under the BSD-style license found in the
  LICENSE file in the root directory of this source tree.
-}

{-# OPTIONS_GHC -fno-warn-unused-imports #-}

{-# LANGUAGE CPP #-}

module Util.ByteString
  ( newCStringFromLazyByteString
  , useLazyByteStringAsCString
  , intToByteString
  , useAsCString
  , useAsCStringLen
  , useByteStringsAsCStrings
  , bsListAsCStrLenArr
  , unsafeBsListAsCStrLenArr
  ) where

import Control.Arrow (left, second)
import Control.Monad
import qualified Data.ByteString.Unsafe as BS
import qualified Data.ByteString as BS
import Foreign
import Foreign.C

import Data.ByteString.Builder
import Data.ByteString.Builder.Extra
import qualified Data.ByteString as Strict
import qualified Data.ByteString.Lazy as Lazy
import qualified Data.Attoparsec.ByteString.Char8 as Atto
import qualified Data.Text as Text

import qualified Util.ASan as ASan

-- | Uses a lazy 'ByteString' as a NUL-terminated 'CString'.
useLazyByteStringAsCString
  :: Lazy.ByteString -> (CString -> IO a) -> IO a
useLazyByteStringAsCString = allocateAndCopy allocaBytes

-- | Copies a lazy 'ByteString' and adds a terminating NUL. The
-- resulting 'CString' must be 'free'd.
newCStringFromLazyByteString :: Lazy.ByteString -> IO CString
newCStringFromLazyByteString = flip (allocateAndCopy mallocBytes') return

useAsCStringLen :: BS.ByteString -> (CStringLen -> IO a) -> IO a
-- TODO(T24195918): We need to do this for ABI compatibility between
-- sigma.service and sigma.service.asan. Codemod to __SANITIZE_ADDRESS__
-- to enable ASAN checks for Haskell allocations.
#ifdef __HASKELL_SANITIZE_ADDRESS__
useAsCStringLen = ASan.byteStringWithCStringLen
#else
useAsCStringLen = BS.useAsCStringLen
#endif

useAsCString :: BS.ByteString -> (CString -> IO a) -> IO a
#ifdef __HASKELL_SANITIZE_ADDRESS__
useAsCString = ASan.byteStringWithCString
#else
useAsCString = BS.useAsCString
#endif

-- | Allocates space for a lazy 'ByteString' and NUL terminator, then
-- copies all its chunks into the buffer.
allocateAndCopy
  :: (Int -> (CString -> IO b) -> IO b)
  -> Lazy.ByteString
  -> (CString -> IO b)
  -> IO b
allocateAndCopy allocate lbs action = allocate (len + 1) $ \dst -> do
  copyChunksTo dst lbs
  pokeByteOff dst len (0::CChar)
  action dst

  where
  len :: Int
  len = fromIntegral (Lazy.length lbs)

  copyChunksTo :: Ptr a -> Lazy.ByteString -> IO ()
  copyChunksTo dst = foldM_ copyChunkTo 0 . Lazy.toChunks
    where
    copyChunkTo :: Int -> Strict.ByteString -> IO Int
    copyChunkTo off bs = BS.unsafeUseAsCStringLen bs $ \(src, chunkLen) -> do
      copyBytes (dst `plusPtr` off) src chunkLen
      return (off + chunkLen)

mallocBytes' :: Int -> (Ptr a -> IO b) -> IO b
mallocBytes' size action = mallocBytes size >>= action

lengthMinInt :: Int
lengthMinInt = length . show $ (minBound :: Int)

-- | Cheap (fast, low memory) conversion of Ints into an
-- ASCII/UTF8-encoded lazy ByteString.
intToByteString :: Int -> Lazy.ByteString
intToByteString =
  toLazyByteStringWith (untrimmedStrategy lengthMinInt 0) Lazy.empty . intDec

useByteStringsAsCStrings :: [BS.ByteString] -> (Ptr CString -> IO a) -> IO a
useByteStringsAsCStrings bs f = go bs []
  where go (t:ts) acc = useAsCString t $ go ts . (:acc)
        go []     acc = withArray0 nullPtr (reverse acc) f

bsListAsCStrLenArr :: [BS.ByteString]
                   -> (Ptr CString -> Ptr CSize -> CSize -> IO a)
                   -> IO a
bsListAsCStrLenArr bs f = go bs []
  where go (t:ts) acc = useAsCStringLen t $ go ts . (:acc)
        go []     acc = withArrayLen strs $ \len strPtr ->
                        withArray lens $ \lenPtr ->
                            f strPtr lenPtr (fromIntegral len)
            where (strs, lens) = second (map fromIntegral) $ unzip $ reverse acc

unsafeBsListAsCStrLenArr :: [BS.ByteString]
                   -> (Ptr CString -> Ptr CSize -> CSize -> IO a)
                   -> IO a
unsafeBsListAsCStrLenArr bs f = go bs []
  where go (t:ts) acc = BS.unsafeUseAsCStringLen t $ go ts . (:acc)
        go []     acc = withArrayLen strs $ \len strPtr ->
                        withArray lens $ \lenPtr ->
                            f strPtr lenPtr (fromIntegral len)
            where (strs, lens) = second (map fromIntegral) $ unzip $ reverse acc