{-# LANGUAGE CPP #-}
{-# LANGUAGE Rank2Types #-}
{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE PatternGuards #-}
{-# LANGUAGE FlexibleContexts #-}
#ifdef TRUSTWORTHY
{-# LANGUAGE Trustworthy #-}
#endif
#ifndef MIN_VERSION_base
#define MIN_VERSION_base(x,y,z) 1
#endif
module Control.Lens.Internal.ByteString
( unpackStrict, traversedStrict, traversedStrictTree
, unpackStrict8, traversedStrict8, traversedStrictTree8
, unpackLazy, traversedLazy
, unpackLazy8, traversedLazy8
) where
import Control.Applicative
import Control.Lens
import qualified Data.ByteString as B
import qualified Data.ByteString.Lazy as BL
import qualified Data.ByteString.Lazy.Internal as BLI
import qualified Data.ByteString.Lazy.Char8 as BL8
import qualified Data.ByteString.Internal as BI
import Data.Bits
import Data.Char
import Data.Int (Int64)
import Data.Word (Word8)
import Foreign.Ptr
import Foreign.Storable
#if MIN_VERSION_base(4,4,0)
import Foreign.ForeignPtr.Safe
import Foreign.ForeignPtr.Unsafe
#else
import Foreign.ForeignPtr
#endif
import GHC.Base (unsafeChr)
import GHC.ForeignPtr (mallocPlainForeignPtrBytes)
import GHC.IO (unsafeDupablePerformIO)
traversedStrict :: Int -> IndexedTraversal' Int B.ByteString Word8
traversedStrict i0 pafb (BI.PS fp off len) =
let p = unsafeForeignPtrToPtr fp
in fmap (rebuild len) (go i0 (p `plusPtr` off) (p `plusPtr` (off+len)))
where
rebuild n = \xs -> unsafeCreate n $ \p -> go2 p xs
go2 !p (x:xs) = poke p x >> go2 (p `plusPtr` 1) xs
go2 _ [] = return ()
go !i !p !q
| p == q = pure []
| otherwise = let !x = BI.inlinePerformIO $ do
x' <- peek p
touchForeignPtr fp
return x'
in (:) <$> indexed pafb (i :: Int) x <*> go (i + 1) (p `plusPtr` 1) q
{-# INLINE traversedStrict #-}
traversedStrict8 :: Int -> IndexedTraversal' Int B.ByteString Char
traversedStrict8 i0 pafb (BI.PS fp off len) =
let p = unsafeForeignPtrToPtr fp
in fmap (rebuild len) (go i0 (p `plusPtr` off) (p `plusPtr` (off+len)))
where
rebuild n = \xs -> unsafeCreate n $ \p -> go2 p xs
go2 !p (x:xs) = poke p (c2w x) >> go2 (p `plusPtr` 1) xs
go2 _ [] = return ()
go !i !p !q
| p == q = pure []
| otherwise = let !x = BI.inlinePerformIO $ do
x' <- peek p
touchForeignPtr fp
return x'
in (:) <$> indexed pafb (i :: Int) (w2c x) <*> go (i + 1) (p `plusPtr` 1) q
{-# INLINE traversedStrict8 #-}
grain :: Int
grain = 32
{-# INLINE grain #-}
traversedStrictTree :: Int -> IndexedTraversal' Int B.ByteString Word8
traversedStrictTree i0 pafb (BI.PS fp off len) = rebuild len <$> go (unsafeForeignPtrToPtr fp `plusPtr` (off - i0)) i0 (i0 + len)
where
rebuild n f = unsafeCreate n $ \q -> f $! (q `plusPtr` (off - i0))
go !p !i !j
| i + grain < j, k <- i + shiftR (j - i) 1 = (\l r q -> l q >> r q) <$> go p i k <*> go p k j
| otherwise = run p i j
run !p !i !j
| i == j = pure (\_ -> return ())
| otherwise = let !x = BI.inlinePerformIO $ do
x' <- peekByteOff p i
touchForeignPtr fp
return x'
in (\y ys !q -> pokeByteOff q i y >> ys q) <$> indexed pafb (i :: Int) x <*> run p (i + 1) j
{-# INLINE traversedStrictTree #-}
traversedStrictTree8 :: Int -> IndexedTraversal' Int B.ByteString Char
traversedStrictTree8 i0 pafb (BI.PS fp off len) = rebuild len <$> go i0 (i0 + len)
where
p = unsafeForeignPtrToPtr fp `plusPtr` (off - i0)
rebuild n f = unsafeCreate n $ \q -> f (q `plusPtr` (off - i0))
go !i !j
| i + grain < j, k <- i + shiftR (j - i) 1 = (\l r q -> l q >> r q) <$> go i k <*> go k j
| otherwise = run i j
run !i !j
| i == j = pure (\_ -> return ())
| otherwise = let !x = BI.inlinePerformIO $ do
x' <- peekByteOff p i
touchForeignPtr fp
return x'
in (\y ys q -> poke (q `plusPtr` i) (c2w y) >> ys q) <$> indexed pafb (i :: Int) (w2c x) <*> run (i + 1) j
{-# INLINE traversedStrictTree8 #-}
unpackLazy :: BL.ByteString -> [Word8]
unpackLazy = BL.unpack
{-# INLINE unpackLazy #-}
traversedLazy :: IndexedTraversal' Int64 BL.ByteString Word8
traversedLazy pafb = go 0 where
go _ BLI.Empty = pure BLI.Empty
go i (BLI.Chunk b bs) = BLI.Chunk <$> reindexed (fromIntegral :: Int -> Int64) (traversedStrictTree (fromIntegral i)) pafb b <*> go i' bs
where !i' = i + B.length b
{-# INLINE traversedLazy #-}
unpackLazy8 :: BL.ByteString -> String
unpackLazy8 = BL8.unpack
{-# INLINE unpackLazy8 #-}
traversedLazy8 :: IndexedTraversal' Int64 BL.ByteString Char
traversedLazy8 pafb = go 0 where
go _ BLI.Empty = pure BLI.Empty
go i (BLI.Chunk b bs) = BLI.Chunk <$> reindexed (fromIntegral :: Int -> Int64) (traversedStrictTree8 (fromIntegral i)) pafb b <*> go i' bs
where !i' = i + B.length b
{-# INLINE traversedLazy8 #-}
w2c :: Word8 -> Char
w2c = unsafeChr . fromIntegral
{-# INLINE w2c #-}
c2w :: Char -> Word8
c2w = fromIntegral . ord
{-# INLINE c2w #-}
unpackStrict :: B.ByteString -> [Word8]
unpackStrict (BI.PS fp off len) =
let p = unsafeForeignPtrToPtr fp
in go (p `plusPtr` off) (p `plusPtr` (off+len))
where
go !p !q | p == q = []
| otherwise = let !x = BI.inlinePerformIO $ do
x' <- peek p
touchForeignPtr fp
return x'
in x : go (p `plusPtr` 1) q
{-# INLINE unpackStrict #-}
unpackStrict8 :: B.ByteString -> String
unpackStrict8 (BI.PS fp off len) =
let p = unsafeForeignPtrToPtr fp
in go (p `plusPtr` off) (p `plusPtr` (off+len))
where
go !p !q | p == q = []
| otherwise = let !x = BI.inlinePerformIO $ do
x' <- peek p
touchForeignPtr fp
return x'
in w2c x : go (p `plusPtr` 1) q
{-# INLINE unpackStrict8 #-}
unsafeCreate :: Int -> (Ptr Word8 -> IO ()) -> B.ByteString
unsafeCreate l f = unsafeDupablePerformIO (create l f)
{-# INLINE unsafeCreate #-}
create :: Int -> (Ptr Word8 -> IO ()) -> IO B.ByteString
create l f = do
fp <- mallocPlainForeignPtrBytes l
withForeignPtr fp $ \p -> f p
return $! BI.PS fp 0 l
{-# INLINE create #-}