{-# LANGUAGE MagicHash #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE CPP #-}
{-# LANGUAGE UnliftedFFITypes #-}
{-# LANGUAGE DeriveDataTypeable #-}
{-# LANGUAGE BangPatterns #-}
module Z.Data.CBytes
( CBytes
, create
, pack
, unpack
, null , length
, empty, append, concat, intercalate, intercalateElem
, toBytes, fromBytes, toText, toTextMaybe, fromText
, fromCString, fromCString', fromCStringN
, withCBytes, allocCBytes
, V.w2c, V.c2w
, NullPointerException(..)
, CString
) where
import Control.DeepSeq
import Control.Exception (Exception, throwIO)
import Control.Monad
import Control.Monad.Primitive
import Control.Monad.ST
import Data.Bits
import Data.Foldable (foldlM)
import Data.Hashable (Hashable(..))
import qualified Data.List as List
import Data.String (IsString (..))
import Data.Typeable
import Data.Primitive.PrimArray
import Data.Word
import Foreign.C
import Foreign.Storable (peekElemOff)
import GHC.CString
import GHC.Ptr
import GHC.Stack
import Prelude hiding (all, any, appendFile, break,
concat, concatMap, drop, dropWhile,
elem, filter, foldl, foldl1, foldr,
foldr1, getContents, getLine, head,
init, interact, last, length, lines,
map, maximum, minimum, notElem, null,
putStr, putStrLn, readFile, replicate,
reverse, scanl, scanl1, scanr, scanr1,
span, splitAt, tail, take, takeWhile,
unlines, unzip, writeFile, zip,
zipWith)
import Z.Data.Array
import qualified Z.Data.Text as T
import Z.Data.Text.UTF8Codec (encodeCharModifiedUTF8)
import qualified Z.Data.Vector.Base as V
import System.IO.Unsafe (unsafeDupablePerformIO)
data CBytes
= CBytesOnHeap {-# UNPACK #-} !(PrimArray Word8)
| CBytesLiteral {-# UNPACK #-} !CString
create :: HasCallStack
=> Int
-> (CString -> IO Int)
-> IO CBytes
{-# INLINE create #-}
create n fill | n <= 0 = fill nullPtr >> return empty
| otherwise = do
mba <- newPinnedPrimArray n :: IO (MutablePrimArray RealWorld Word8)
l <- withMutablePrimArrayContents mba (fill . castPtr)
when (l+1>n) (throwIO (NULLTerminatorNotFound callStack))
writePrimArray mba l 0
shrinkMutablePrimArray mba (l+1)
CBytesOnHeap <$> unsafeFreezePrimArray mba
data CBytesException = NULLTerminatorNotFound CallStack
deriving (Show, Typeable)
instance Exception CBytesException
instance Show CBytes where
show = unpack
instance Read CBytes where
readsPrec p s = [(pack x, r) | (x, r) <- readsPrec p s]
instance NFData CBytes where
{-# INLINE rnf #-}
rnf (CBytesOnHeap _) = ()
rnf (CBytesLiteral _) = ()
instance Eq CBytes where
{-# INLINE (==) #-}
cbyteA == cbyteB = unsafeDupablePerformIO $
withCBytes cbyteA $ \ pA ->
withCBytes cbyteB $ \ pB ->
if pA == pB
then return True
else do
r <- c_strcmp pA pB
return (r == 0)
instance Ord CBytes where
{-# INLINE compare #-}
cbyteA `compare` cbyteB = unsafeDupablePerformIO $
withCBytes cbyteA $ \ pA ->
withCBytes cbyteB $ \ pB ->
if pA == pB
then return EQ
else do
r <- c_strcmp pA pB
return (r `compare` 0)
instance Semigroup CBytes where
(<>) = append
instance Monoid CBytes where
{-# INLINE mempty #-}
mempty = empty
{-# INLINE mappend #-}
mappend = append
{-# INLINE mconcat #-}
mconcat = concat
instance Hashable CBytes where
hashWithSalt salt (CBytesOnHeap pa@(PrimArray ba#)) = unsafeDupablePerformIO $ do
V.c_fnv_hash_ba ba# 0 (sizeofPrimArray pa - 1) salt
hashWithSalt salt (CBytesLiteral p@(Ptr addr#)) = unsafeDupablePerformIO $ do
len <- c_strlen p
V.c_fnv_hash_addr addr# (fromIntegral len) salt
append :: CBytes -> CBytes -> CBytes
{-# INLINABLE append #-}
append strA strB
| lenA == 0 = strB
| lenB == 0 = strA
| otherwise = unsafeDupablePerformIO $ do
mpa <- newPinnedPrimArray (lenA+lenB+1)
withCBytes strA $ \ pa ->
withCBytes strB $ \ pb -> do
copyPtrToMutablePrimArray mpa 0 (castPtr pa) lenA
copyPtrToMutablePrimArray mpa lenA (castPtr pb) lenB
writePrimArray mpa (lenA + lenB) 0
pa' <- unsafeFreezePrimArray mpa
return (CBytesOnHeap pa')
where
lenA = length strA
lenB = length strB
empty :: CBytes
{-# NOINLINE empty #-}
empty = CBytesLiteral (Ptr "\0"#)
concat :: [CBytes] -> CBytes
{-# INLINABLE concat #-}
concat bss = case pre 0 0 bss of
(0, _) -> empty
(1, _) -> let Just b = List.find (not . null) bss in b
(_, l) -> runST $ do
buf <- newPinnedPrimArray (l+1)
copy bss 0 buf
writePrimArray buf l 0
CBytesOnHeap <$> unsafeFreezePrimArray buf
where
pre :: Int -> Int -> [CBytes] -> (Int, Int)
pre !nacc !lacc [] = (nacc, lacc)
pre !nacc !lacc (b:bs)
| l <= 0 = pre nacc lacc bs
| otherwise = pre (nacc+1) (l + lacc) bs
where !l = length b
copy :: [CBytes] -> Int -> MutablePrimArray s Word8 -> ST s ()
copy [] !_ !_ = return ()
copy (b:bs) !i !mba = do
let l = length b
when (l /= 0) (case b of
CBytesOnHeap ba ->
copyPrimArray mba i ba 0 l
CBytesLiteral p ->
copyPtrToMutablePrimArray mba i (castPtr p) l)
copy bs (i+l) mba
intercalate :: CBytes -> [CBytes] -> CBytes
{-# INLINE intercalate #-}
intercalate s = concat . List.intersperse s
intercalateElem :: Word8 -> [CBytes] -> CBytes
{-# INLINABLE intercalateElem #-}
intercalateElem w8 bss = case len bss 0 of
0 -> empty
l -> runST $ do
buf <- newPinnedPrimArray (l+1)
copy bss 0 buf
writePrimArray buf l 0
CBytesOnHeap <$> unsafeFreezePrimArray buf
where
len [] !acc = acc
len [b] !acc = length b + acc
len (b:bs) !acc = len bs (acc + length b + 1)
copy :: [CBytes] -> Int -> MutablePrimArray s Word8 -> ST s ()
copy (b:bs) !i !mba = do
let l = length b
when (l /= 0) (case b of
CBytesOnHeap ba ->
copyPrimArray mba i ba 0 l
CBytesLiteral p ->
copyPtrToMutablePrimArray mba i (castPtr p) l)
case bs of
[] -> return ()
_ -> do
let i' = i + l
writePrimArray mba i' w8
copy bs (i'+1) mba
instance IsString CBytes where
{-# INLINE fromString #-}
fromString = pack
{-# RULES
"CBytes pack/unpackCString#" forall addr# .
pack (unpackCString# addr#) = CBytesLiteral (Ptr addr#)
#-}
{-# RULES
"CBytes pack/unpackCStringUtf8#" forall addr# .
pack (unpackCStringUtf8# addr#) = CBytesLiteral (Ptr addr#)
#-}
pack :: String -> CBytes
{-# INLINE CONLIKE [1] pack #-}
pack s = runST $ do
mba <- newPinnedPrimArray V.defaultInitSize
(SP2 i mba') <- foldlM go (SP2 0 mba) s
writePrimArray mba' i 0
shrinkMutablePrimArray mba' (i+1)
ba <- unsafeFreezePrimArray mba'
return (CBytesOnHeap ba)
where
go :: SP2 s -> Char -> ST s (SP2 s)
go (SP2 i mba) !c = do
siz <- getSizeofMutablePrimArray mba
if i < siz - 4
then do
i' <- encodeCharModifiedUTF8 mba i c
return (SP2 i' mba)
else do
let !siz' = siz `shiftL` 1
!mba' <- resizeMutablePrimArray mba siz'
i' <- encodeCharModifiedUTF8 mba' i c
return (SP2 i' mba')
data SP2 s = SP2 {-# UNPACK #-}!Int {-# UNPACK #-}!(MutablePrimArray s Word8)
unpack :: CBytes -> String
{-# INLINABLE unpack #-}
unpack cbytes = unsafeDupablePerformIO . withCBytes cbytes $ \ (Ptr addr#) ->
return (unpackCStringUtf8# addr#)
null :: CBytes -> Bool
{-# INLINE null #-}
null (CBytesOnHeap pa) = indexPrimArray pa 0 == 0
null (CBytesLiteral p) = unsafeDupablePerformIO (peekElemOff p 0) == 0
length :: CBytes -> Int
{-# INLINE length #-}
length (CBytesOnHeap pa) = sizeofPrimArray pa - 1
length (CBytesLiteral p) = fromIntegral $ unsafeDupablePerformIO (c_strlen p)
toBytes :: CBytes -> V.Bytes
{-# INLINABLE toBytes #-}
toBytes cbytes@(CBytesOnHeap pa) = V.PrimVector pa 0 l
where l = length cbytes
toBytes cbytes@(CBytesLiteral p) = V.create (l+1) (\ mpa -> do
copyPtrToMutablePrimArray mpa 0 (castPtr p) l
writePrimArray mpa l 0)
where l = length cbytes
fromBytes :: V.Bytes -> CBytes
{-# INLINABLE fromBytes #-}
fromBytes (V.Vec arr s l) = runST (do
mpa <- newPinnedPrimArray (l+1)
copyPrimArray mpa 0 arr s l
writePrimArray mpa l 0
pa <- unsafeFreezePrimArray mpa
return (CBytesOnHeap pa))
toText :: CBytes -> T.Text
{-# INLINABLE toText #-}
toText = T.validate . toBytes
toTextMaybe :: CBytes -> Maybe T.Text
{-# INLINABLE toTextMaybe #-}
toTextMaybe = T.validateMaybe . toBytes
fromText :: T.Text -> CBytes
{-# INLINABLE fromText #-}
fromText = fromBytes . T.getUTF8Bytes
fromCString :: CString -> IO CBytes
{-# INLINABLE fromCString #-}
fromCString cstring = do
if cstring == nullPtr
then return empty
else do
len <- fromIntegral <$> c_strlen cstring
mpa <- newPinnedPrimArray (len+1)
copyPtrToMutablePrimArray mpa 0 (castPtr cstring) len
writePrimArray mpa len 0
pa <- unsafeFreezePrimArray mpa
return (CBytesOnHeap pa)
fromCString' :: HasCallStack => CString -> IO (Maybe CBytes)
{-# INLINABLE fromCString' #-}
fromCString' cstring =
if cstring == nullPtr
then throwIO (NullPointerException callStack)
else do
len <- fromIntegral <$> c_strlen cstring
mpa <- newPinnedPrimArray (len+1)
copyPtrToMutablePrimArray mpa 0 (castPtr cstring) len
writePrimArray mpa len 0
pa <- unsafeFreezePrimArray mpa
return (Just (CBytesOnHeap pa))
fromCStringN :: CString -> Int -> IO CBytes
{-# INLINABLE fromCStringN #-}
fromCStringN cstring len = do
if cstring == nullPtr || len == 0
then return empty
else do
mpa <- newPinnedPrimArray (len+1)
copyPtrToMutablePrimArray mpa 0 (castPtr cstring) len
writePrimArray mpa len 0
pa <- unsafeFreezePrimArray mpa
return (CBytesOnHeap pa)
data NullPointerException = NullPointerException CallStack deriving (Show, Typeable)
instance Exception NullPointerException
withCBytes :: CBytes -> (CString -> IO a) -> IO a
{-# INLINABLE withCBytes #-}
withCBytes (CBytesOnHeap pa) f = withPrimArrayContents pa (f . castPtr)
withCBytes (CBytesLiteral ptr) f = f ptr
allocCBytes :: HasCallStack
=> Int
-> (CString -> IO a)
-> IO (CBytes, a)
{-# INLINABLE allocCBytes #-}
allocCBytes n fill | n <= 0 = fill nullPtr >>= \ a -> return (empty, a)
| otherwise = do
mba <- newPinnedPrimArray n :: IO (MutablePrimArray RealWorld Word8)
a <- withMutablePrimArrayContents mba (fill . castPtr)
l <- fromIntegral <$> withMutablePrimArrayContents mba (c_strlen . castPtr)
when (l+1>n) (throwIO (NULLTerminatorNotFound callStack))
shrinkMutablePrimArray mba (l+1)
bs <- unsafeFreezePrimArray mba
return (CBytesOnHeap bs, a)
c_strcmp :: CString -> CString -> IO CInt
{-# INLINE c_strcmp #-}
c_strcmp (Ptr a#) (Ptr b#) = V.c_strcmp a# b#
c_strlen :: CString -> IO CSize
{-# INLINE c_strlen #-}
c_strlen (Ptr a#) = V.c_strlen a#