Managed.hs

{-# language BangPatterns #-}
{-# language DerivingStrategies #-}
{-# language GeneralizedNewtypeDeriving #-}
{-# language MagicHash #-}
{-# language MultiWayIf #-}
{-# language TypeApplications #-}
{-# language ScopedTypeVariables #-}
{-# language UnboxedTuples #-}

module Foreign.C.String.Managed
  ( ManagedCString(..)
  , terminated
  , terminatedU
  , unterminated
  , fromBytes
  , fromLatinString
  , fromShortText
  , pinnedFromBytes
  , pin
  , touch
  , contents
  ) where

import Control.Monad.ST (ST)
import Control.Monad.ST.Run (runByteArrayST)
import Data.Bytes.Types (Bytes(Bytes))
import Data.Char (ord)
import Data.Primitive (ByteArray(..),MutableByteArray)
import Data.Text.Short (ShortText)
import Data.Word (Word8)
import Foreign.C.String (CString)
import Foreign.Ptr (castPtr)
import GHC.Exts (Int(I#),Char(C#),ByteArray#,chr#,touch#)
import GHC.IO (IO(IO))

import qualified Data.Bytes as Bytes
import qualified Data.Bytes.Text.Utf8 as Utf8
import qualified Data.Primitive as PM
import qualified Data.Text.Short as TS
import qualified GHC.Exts as Exts

-- | An unsliced byte sequence with @NUL@ as the final byte.
newtype ManagedCString = ManagedCString ByteArray
  deriving newtype Eq

instance Semigroup ManagedCString where
  ManagedCString a <> ManagedCString b = ManagedCString $ runByteArrayST $ do
    let lenA = PM.sizeofByteArray a
    let lenB = PM.sizeofByteArray b
    dst <- PM.newByteArray (lenA + lenB - 1)
    PM.copyByteArray dst 0 a 0 (lenA - 1)
    PM.copyByteArray dst (lenA - 1) b 0 lenB
    PM.unsafeFreezeByteArray dst

instance Monoid ManagedCString where
  mempty = ManagedCString $ runByteArrayST $ do
    dst <- PM.newByteArray 1
    PM.writeByteArray dst 0 (0 :: Word8)
    PM.unsafeFreezeByteArray dst

instance Exts.IsString ManagedCString where
  fromString = fromLatinString

instance Show ManagedCString where
  showsPrec _ (ManagedCString arr) s0 = PM.foldrByteArray
    ( \(w :: Word8) s ->
      if | w == 0 -> s
         | w < 32 -> '?' : s
         | w > 126 -> '?' : s
         | otherwise -> case fromIntegral @Word8 @Int w of
             I# i -> C# (chr# i) : s
    ) s0 arr

terminatedU :: ManagedCString -> ByteArray
terminatedU (ManagedCString x) = x

terminated :: ManagedCString -> Bytes
terminated (ManagedCString x) = Bytes.fromByteArray x

unterminated :: ManagedCString -> Bytes
unterminated (ManagedCString x) = Bytes x 0 (PM.sizeofByteArray x - 1)

fromShortText :: ShortText -> ManagedCString
fromShortText !ts = fromBytes (Utf8.fromShortText ts)

-- | Copies the slice, appending a @NUL@ byte to the end.
fromBytes :: Bytes -> ManagedCString
fromBytes (Bytes arr off len) = ManagedCString $ runByteArrayST $ do
  dst <- PM.newByteArray (len + 1)
  PM.copyByteArray dst 0 arr off len
  PM.writeByteArray dst len (0 :: Word8)
  PM.unsafeFreezeByteArray dst

-- | Copies the slice into pinned memory, appending a @NUL@ byte to the end.
pinnedFromBytes :: Bytes -> ManagedCString
pinnedFromBytes (Bytes arr off len) = ManagedCString $ runByteArrayST $ do
  dst <- PM.newPinnedByteArray (len + 1)
  PM.copyByteArray dst 0 arr off len
  PM.writeByteArray dst len (0 :: Word8)
  PM.unsafeFreezeByteArray dst

pin :: ManagedCString -> ManagedCString
pin (ManagedCString x) = if PM.isByteArrayPinned x
  then ManagedCString x
  else ManagedCString $ runByteArrayST $ do
    let len = PM.sizeofByteArray x
    dst <- PM.newPinnedByteArray len
    PM.copyByteArray dst 0 x 0 len
    PM.unsafeFreezeByteArray dst

touch :: ManagedCString -> IO ()
touch (ManagedCString (ByteArray x)) = touchByteArray# x

touchByteArray# :: ByteArray# -> IO ()
touchByteArray# x = IO $ \s -> case touch# x s of s' -> (# s', () #)

-- | Convert a 'String' consisting of only characters representable
-- by ISO-8859-1. These are encoded with ISO-8859-1. Any character
-- with a codepoint above @U+00FF@ is replaced by an unspecified byte.
fromLatinString :: String -> ManagedCString
{-# noinline fromLatinString #-}
fromLatinString str = ManagedCString $ runByteArrayST $ do
  let lenPred0 = 63
  dst0 <- PM.newByteArray (lenPred0 + 1)
  go str dst0 0 lenPred0
  where
  go :: forall s. String -> MutableByteArray s -> Int -> Int -> ST s ByteArray
  go [] !dst !ix !_ = do
    PM.writeByteArray dst ix (0 :: Word8)
    PM.resizeMutableByteArray dst (ix + 1) >>= PM.unsafeFreezeByteArray
  go (c:cs) !dst !ix !lenPred = if ix < lenPred
    then do
      PM.writeByteArray dst ix (fromIntegral @Int @Word8 (ord c))
      go cs dst (ix + 1) lenPred
    else do
      let nextLenPred = lenPred * 2
      dst' <- PM.newByteArray (nextLenPred + 1)
      PM.copyMutableByteArray dst' 0 dst 0 ix
      PM.writeByteArray dst' ix (fromIntegral @Int @Word8 (ord c))
      go cs dst' (ix + 1) nextLenPred

-- | Get a pointer to the payload of the managed C string. The behavior is
-- undefined if the argument is not pinned.
contents :: ManagedCString -> CString
contents (ManagedCString x) = castPtr (PM.byteArrayContents x)