{-# LANGUAGE BangPatterns        #-}
{-# LANGUAGE MagicHash           #-}
{-# LANGUAGE CPP                 #-}
{-# LANGUAGE TypeFamilies        #-}

-- |
-- Module      : Codec.CBOR.ByteArray.Sliced
-- Copyright   : (c) Ben Gamari 2017-2018
-- License     : BSD3-style (see LICENSE.txt)
--
-- Maintainer  : duncan@community.haskell.org
-- Stability   : experimental
-- Portability : non-portable (GHC extensions)
--
-- A ByteArray with more instances than 'Data.Primitive.ByteArray.ByteArray'.
-- Some day when these instances are reliably available from @primitive@ we can
-- likely replace this with 'Data.Primitive.ByteArray.ByteArray'.
--
module Codec.CBOR.ByteArray.Sliced
  ( SlicedByteArray(..)
    -- * Conversions
  , sizeofSlicedByteArray
  , fromShortByteString
  , fromByteString
  , fromByteArray
  , toByteString
  , toBuilder
  ) where

import GHC.Exts
import Data.Char (chr, ord)
import Data.Word
import Foreign.Ptr
import Control.Monad.ST
import System.IO.Unsafe

import qualified Data.Primitive.ByteArray as Prim
#if !MIN_VERSION_primitive(0,7,0)
import           Data.Primitive.Types (Addr(..))
#endif
import qualified Data.ByteString as BS
import qualified Data.ByteString.Unsafe as BS
import qualified Data.ByteString.Short as BSS
import qualified Data.ByteString.Short.Internal as BSS
import qualified Data.ByteString.Builder as BSB
import qualified Data.ByteString.Builder.Internal as BSB

import Codec.CBOR.ByteArray.Internal

data SlicedByteArray = SBA {SlicedByteArray -> ByteArray
unSBA :: !Prim.ByteArray, SlicedByteArray -> Int
offset :: !Int, SlicedByteArray -> Int
length :: !Int}

fromShortByteString :: BSS.ShortByteString -> SlicedByteArray
fromShortByteString :: ShortByteString -> SlicedByteArray
fromShortByteString (BSS.SBS ByteArray#
ba) = ByteArray -> SlicedByteArray
fromByteArray (ByteArray# -> ByteArray
Prim.ByteArray ByteArray#
ba)

fromByteString :: BS.ByteString -> SlicedByteArray
fromByteString :: ByteString -> SlicedByteArray
fromByteString = ShortByteString -> SlicedByteArray
fromShortByteString (ShortByteString -> SlicedByteArray)
-> (ByteString -> ShortByteString) -> ByteString -> SlicedByteArray
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ByteString -> ShortByteString
BSS.toShort

fromByteArray :: Prim.ByteArray -> SlicedByteArray
fromByteArray :: ByteArray -> SlicedByteArray
fromByteArray ByteArray
ba = ByteArray -> Int -> Int -> SlicedByteArray
SBA ByteArray
ba Int
0 (ByteArray -> Int
Prim.sizeofByteArray ByteArray
ba)

sizeofSlicedByteArray :: SlicedByteArray -> Int
sizeofSlicedByteArray :: SlicedByteArray -> Int
sizeofSlicedByteArray (SBA ByteArray
_ Int
_ Int
len) = Int
len

-- | Note that this may require a copy.
toByteString :: SlicedByteArray -> BS.ByteString
toByteString :: SlicedByteArray -> ByteString
toByteString SlicedByteArray
sba =
    IO ByteString -> ByteString
forall a. IO a -> a
unsafePerformIO
    (IO ByteString -> ByteString) -> IO ByteString -> ByteString
forall a b. (a -> b) -> a -> b
$ Ptr Word8 -> Int -> IO () -> IO ByteString
BS.unsafePackCStringFinalizer Ptr Word8
ptr (SlicedByteArray -> Int
sizeofSlicedByteArray SlicedByteArray
sba) (ByteArray -> IO ()
forall a. a -> IO ()
touch ByteArray
pinned)
  where
    pinned :: ByteArray
pinned = SlicedByteArray -> ByteArray
toPinned SlicedByteArray
sba
#if MIN_VERSION_primitive(0,7,0)
    !(Ptr Addr#
addr#) = ByteArray -> Ptr Word8
Prim.byteArrayContents ByteArray
pinned
#else
    !(Addr addr#) = Prim.byteArrayContents pinned
#endif
    ptr :: Ptr Word8
ptr = Addr# -> Ptr Word8
forall a. Addr# -> Ptr a
Ptr Addr#
addr#

toPinned :: SlicedByteArray -> Prim.ByteArray
toPinned :: SlicedByteArray -> ByteArray
toPinned (SBA ByteArray
ba Int
off Int
len)
  | ByteArray -> Bool
isByteArrayPinned ByteArray
ba = ByteArray
ba
  | Bool
otherwise = (forall s. ST s ByteArray) -> ByteArray
forall a. (forall s. ST s a) -> a
runST ((forall s. ST s ByteArray) -> ByteArray)
-> (forall s. ST s ByteArray) -> ByteArray
forall a b. (a -> b) -> a -> b
$ do
        MutableByteArray s
ba' <- Int -> ST s (MutableByteArray (PrimState (ST s)))
forall (m :: * -> *).
PrimMonad m =>
Int -> m (MutableByteArray (PrimState m))
Prim.newPinnedByteArray Int
len
        MutableByteArray (PrimState (ST s))
-> Int -> ByteArray -> Int -> Int -> ST s ()
forall (m :: * -> *).
PrimMonad m =>
MutableByteArray (PrimState m)
-> Int -> ByteArray -> Int -> Int -> m ()
Prim.copyByteArray MutableByteArray s
MutableByteArray (PrimState (ST s))
ba' Int
0 ByteArray
ba Int
off Int
len
        MutableByteArray (PrimState (ST s)) -> ST s ByteArray
forall (m :: * -> *).
PrimMonad m =>
MutableByteArray (PrimState m) -> m ByteArray
Prim.unsafeFreezeByteArray MutableByteArray s
MutableByteArray (PrimState (ST s))
ba'

toBuilder :: SlicedByteArray -> BSB.Builder
toBuilder :: SlicedByteArray -> Builder
toBuilder = \(SBA ByteArray
ba Int
off Int
len) -> (forall r. BuildStep r -> BuildStep r) -> Builder
BSB.builder (ByteArray
-> Int
-> Int
-> (BufferRange -> IO (BuildSignal r))
-> BufferRange
-> IO (BuildSignal r)
forall a.
ByteArray
-> Int
-> Int
-> (BufferRange -> IO (BuildSignal a))
-> BufferRange
-> IO (BuildSignal a)
go ByteArray
ba Int
off Int
len)
  where
    go :: ByteArray
-> Int
-> Int
-> (BufferRange -> IO (BuildSignal a))
-> BufferRange
-> IO (BuildSignal a)
go ByteArray
ba !Int
ip !Int
ipe !BufferRange -> IO (BuildSignal a)
k (BSB.BufferRange Ptr Word8
op Ptr Word8
ope)
      | Int
inpRemaining Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
outRemaining = do
          ByteArray -> Int -> Ptr Word8 -> Int -> IO ()
forall a. ByteArray -> Int -> Ptr a -> Int -> IO ()
copyToAddr ByteArray
ba Int
ip Ptr Word8
op Int
inpRemaining
          let !br' :: BufferRange
br' = Ptr Word8 -> Ptr Word8 -> BufferRange
BSB.BufferRange (Ptr Word8
op Ptr Word8 -> Int -> Ptr Word8
forall a b. Ptr a -> Int -> Ptr b
`plusPtr` Int
inpRemaining) Ptr Word8
ope
          BufferRange -> IO (BuildSignal a)
k BufferRange
br'
      | Bool
otherwise = do
          ByteArray -> Int -> Ptr Word8 -> Int -> IO ()
forall a. ByteArray -> Int -> Ptr a -> Int -> IO ()
copyToAddr ByteArray
ba Int
ip Ptr Word8
op Int
outRemaining
          let !ip' :: Int
ip' = Int
ip Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
outRemaining
          BuildSignal a -> IO (BuildSignal a)
forall (m :: * -> *) a. Monad m => a -> m a
return (BuildSignal a -> IO (BuildSignal a))
-> BuildSignal a -> IO (BuildSignal a)
forall a b. (a -> b) -> a -> b
$ Int
-> Ptr Word8
-> (BufferRange -> IO (BuildSignal a))
-> BuildSignal a
forall a. Int -> Ptr Word8 -> BuildStep a -> BuildSignal a
BSB.bufferFull Int
1 Ptr Word8
ope (ByteArray
-> Int
-> Int
-> (BufferRange -> IO (BuildSignal a))
-> BufferRange
-> IO (BuildSignal a)
go ByteArray
ba Int
ip' Int
ipe BufferRange -> IO (BuildSignal a)
k)
      where
        outRemaining :: Int
outRemaining = Ptr Word8
ope Ptr Word8 -> Ptr Word8 -> Int
forall a b. Ptr a -> Ptr b -> Int
`minusPtr` Ptr Word8
op
        inpRemaining :: Int
inpRemaining = Int
ipe Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
ip

instance IsString SlicedByteArray where
  fromString :: String -> SlicedByteArray
fromString = [Word8] -> SlicedByteArray
forall l. IsList l => [Item l] -> l
fromList ([Word8] -> SlicedByteArray)
-> (String -> [Word8]) -> String -> SlicedByteArray
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Char -> Word8) -> String -> [Word8]
forall a b. (a -> b) -> [a] -> [b]
map Char -> Word8
forall p. Num p => Char -> p
checkedOrd
    where
      checkedOrd :: Char -> p
checkedOrd Char
c
        | Char
c Char -> Char -> Bool
forall a. Ord a => a -> a -> Bool
> Char
'\xff' = String -> p
forall a. HasCallStack => String -> a
error String
"IsString(Codec.CBOR.ByteArray.Sliced): Non-ASCII character"
        | Bool
otherwise  = Int -> p
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int -> p) -> Int -> p
forall a b. (a -> b) -> a -> b
$ Char -> Int
ord Char
c

instance IsList SlicedByteArray where
  type Item SlicedByteArray = Word8
  fromList :: [Item SlicedByteArray] -> SlicedByteArray
fromList [Item SlicedByteArray]
xs = Int -> [Item SlicedByteArray] -> SlicedByteArray
forall l. IsList l => Int -> [Item l] -> l
fromListN ([Word8] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
Prelude.length [Word8]
[Item SlicedByteArray]
xs) [Item SlicedByteArray]
xs
  -- Note that we make no attempt to behave sensibly if @n /= length xs@.
  -- The class definition allows this.
  fromListN :: Int -> [Item SlicedByteArray] -> SlicedByteArray
fromListN Int
n [Item SlicedByteArray]
xs =
      let arr :: ByteArray
arr = Int -> [Word8] -> ByteArray
mkByteArray Int
n [Word8]
[Item SlicedByteArray]
xs
      in ByteArray -> Int -> Int -> SlicedByteArray
SBA ByteArray
arr Int
0 Int
n
  toList :: SlicedByteArray -> [Item SlicedByteArray]
toList (SBA ByteArray
arr Int
off Int
len) =
      (Word8 -> [Word8] -> [Word8])
-> [Word8] -> Int -> Int -> ByteArray -> [Word8]
forall a. (Word8 -> a -> a) -> a -> Int -> Int -> ByteArray -> a
foldrByteArray (:) [] Int
off Int
len ByteArray
arr

instance Show SlicedByteArray where
  showsPrec :: Int -> SlicedByteArray -> ShowS
showsPrec Int
_ = String -> ShowS
forall a. Show a => a -> ShowS
shows (String -> ShowS)
-> (SlicedByteArray -> String) -> SlicedByteArray -> ShowS
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Word8 -> Char) -> [Word8] -> String
forall a b. (a -> b) -> [a] -> [b]
map (Int -> Char
chr (Int -> Char) -> (Word8 -> Int) -> Word8 -> Char
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Word8 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral) ([Word8] -> String)
-> (SlicedByteArray -> [Word8]) -> SlicedByteArray -> String
forall b c a. (b -> c) -> (a -> b) -> a -> c
. SlicedByteArray -> [Word8]
forall l. IsList l => l -> [Item l]
toList

instance Eq SlicedByteArray where
  SBA ByteArray
arr1 Int
off1 Int
len1 == :: SlicedByteArray -> SlicedByteArray -> Bool
== SBA ByteArray
arr2 Int
off2 Int
len2
    | Int
len1 Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
len2
    = Bool
False

    | ByteArray -> ByteArray -> Bool
sameByteArray ByteArray
arr1 ByteArray
arr2
    , Int
off1 Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
off2
    , Int
len1 Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
len2
    = Bool
True

    | Bool
otherwise
    = let (!) :: Prim.ByteArray -> Int -> Word8
          (!) = ByteArray -> Int -> Word8
forall a. Prim a => ByteArray -> Int -> a
Prim.indexByteArray
          go :: Int -> Int -> Bool
go Int
i1 Int
i2
            | Int
i1 Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
len1 Bool -> Bool -> Bool
&& Int
i2 Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
len2   = Bool
True
            | Int
i1 Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
len1 Bool -> Bool -> Bool
|| Int
i2 Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
len2   = Bool
False
            | (ByteArray
arr1 ByteArray -> Int -> Word8
! Int
i1) Word8 -> Word8 -> Bool
forall a. Eq a => a -> a -> Bool
== (ByteArray
arr2 ByteArray -> Int -> Word8
! Int
i2) = Int -> Int -> Bool
go (Int
i1Int -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1) (Int
i2Int -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1)
            | Bool
otherwise                  = Bool
False
      in Int -> Int -> Bool
go Int
off1 Int
off2

instance Ord SlicedByteArray where
  SBA ByteArray
arr1 Int
off1 Int
len1 compare :: SlicedByteArray -> SlicedByteArray -> Ordering
`compare` SBA ByteArray
arr2 Int
off2 Int
len2
    | ByteArray -> ByteArray -> Bool
sameByteArray ByteArray
arr1 ByteArray
arr2
    , Int
off1 Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
off2
    , Int
len1 Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
len2
    = Ordering
EQ

    | Bool
otherwise
    = let (!) :: Prim.ByteArray -> Int -> Word8
          (!) = ByteArray -> Int -> Word8
forall a. Prim a => ByteArray -> Int -> a
Prim.indexByteArray
          go :: Int -> Int -> Ordering
go Int
i1 Int
i2
            | Int
i1 Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
len1 Bool -> Bool -> Bool
&& Int
i2 Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
len2 = Ordering
EQ
            | Int
i1 Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
len1 Bool -> Bool -> Bool
|| Int
i2 Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
len2 = Int
len1 Int -> Int -> Ordering
forall a. Ord a => a -> a -> Ordering
`compare` Int
len2
            | Ordering
EQ <- Ordering
o                  = Int -> Int -> Ordering
go (Int
i1Int -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1) (Int
i2Int -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1)
            | Bool
otherwise                = Ordering
o
            where o :: Ordering
o = (ByteArray
arr1 ByteArray -> Int -> Word8
! Int
i1) Word8 -> Word8 -> Ordering
forall a. Ord a => a -> a -> Ordering
`compare` (ByteArray
arr2 ByteArray -> Int -> Word8
! Int
i2)
      in Int -> Int -> Ordering
go Int
off1 Int
off2