112 lines
3.3 KiB
Haskell
112 lines
3.3 KiB
Haskell
{-# language BangPatterns #-}
|
|
{-# language DerivingStrategies #-}
|
|
{-# language DuplicateRecordFields #-}
|
|
{-# language TypeFamilies #-}
|
|
{-# language MagicHash #-}
|
|
{-# language UnboxedTuples #-}
|
|
{-# language NamedFieldPuns #-}
|
|
|
|
module Data.Bytes.Chunks
|
|
( Chunks(..)
|
|
, concat
|
|
, reverse
|
|
, reverseOnto
|
|
, length
|
|
) where
|
|
|
|
import Prelude hiding (length,concat,reverse)
|
|
|
|
import GHC.ST (ST(..))
|
|
import Data.Bytes.Types (Bytes(Bytes))
|
|
import Data.Primitive (ByteArray(..),MutableByteArray(..))
|
|
import GHC.Exts (ByteArray#,MutableByteArray#)
|
|
import GHC.Exts (Int#,State#,Int(I#),(+#))
|
|
import Control.Monad.ST.Run (runByteArrayST)
|
|
|
|
import qualified GHC.Exts as Exts
|
|
import qualified Data.Primitive as PM
|
|
import qualified Data.Bytes.Types as B
|
|
|
|
data Chunks
|
|
= ChunksCons {-# UNPACK #-} !Bytes !Chunks
|
|
| ChunksNil
|
|
deriving stock (Show)
|
|
|
|
instance Semigroup Chunks where
|
|
ChunksNil <> a = a
|
|
cs@(ChunksCons _ _) <> ChunksNil = cs
|
|
as@(ChunksCons _ _) <> bs@(ChunksCons _ _) =
|
|
reverseOnto bs (reverse as)
|
|
|
|
instance Monoid Chunks where
|
|
mempty = ChunksNil
|
|
|
|
instance Eq Chunks where
|
|
-- TODO: There is a more efficient way to do this, but
|
|
-- it is tedious.
|
|
a == b = concat a == concat b
|
|
|
|
concat :: Chunks -> ByteArray
|
|
concat x = ByteArray (concat# x)
|
|
|
|
concat# :: Chunks -> ByteArray#
|
|
{-# noinline concat# #-}
|
|
concat# ChunksNil = case mempty of {ByteArray x -> x}
|
|
concat# (ChunksCons (Bytes{array=c,offset=coff,length=szc}) cs) = case cs of
|
|
ChunksNil -> case c of {ByteArray x -> x}
|
|
ChunksCons (Bytes{array=d,offset=doff,length=szd}) ds ->
|
|
unBa $ runByteArrayST $ do
|
|
let szboth = szc + szd
|
|
len = chunksLengthGo szboth ds
|
|
dst <- PM.newByteArray len
|
|
PM.copyByteArray dst 0 c coff szc
|
|
PM.copyByteArray dst szc d doff szd
|
|
_ <- copy dst szboth ds
|
|
PM.unsafeFreezeByteArray dst
|
|
|
|
length :: Chunks -> Int
|
|
length = chunksLengthGo 0
|
|
|
|
chunksLengthGo :: Int -> Chunks -> Int
|
|
chunksLengthGo !n ChunksNil = n
|
|
chunksLengthGo !n (ChunksCons (Bytes{B.length=len}) cs) =
|
|
chunksLengthGo (n + len) cs
|
|
|
|
-- | Copy the contents of the chunks into a mutable array.
|
|
-- Precondition: The destination must have enough space to
|
|
-- house the contents. This is not checked.
|
|
copy ::
|
|
MutableByteArray s -- ^ Destination
|
|
-> Int -- ^ Destination offset
|
|
-> Chunks -- ^ Source
|
|
-> ST s Int -- ^ Returns the next index into the destination after the payload
|
|
{-# inline copy #-}
|
|
copy (MutableByteArray dst) (I# off) cs = ST
|
|
(\s0 -> case copy# dst off cs s0 of
|
|
(# s1, nextOff #) -> (# s1, I# nextOff #)
|
|
)
|
|
|
|
copy# :: MutableByteArray# s -> Int# -> Chunks -> State# s -> (# State# s, Int# #)
|
|
copy# _ off ChunksNil s0 = (# s0, off #)
|
|
copy# marr off (ChunksCons (Bytes{B.array,B.offset,B.length=len}) cs) s0 =
|
|
case Exts.copyByteArray# (unBa array) (unI offset) marr off (unI len) s0 of
|
|
s1 -> copy# marr (off +# unI len) cs s1
|
|
|
|
|
|
-- | Reverse chunks but not the bytes within each chunk.
|
|
reverse :: Chunks -> Chunks
|
|
reverse = reverseOnto ChunksNil
|
|
|
|
-- | Variant of 'reverse' that allows the caller to provide
|
|
-- an initial list of chunks that the reversed chunks will
|
|
-- be pushed onto.
|
|
reverseOnto :: Chunks -> Chunks -> Chunks
|
|
reverseOnto !x ChunksNil = x
|
|
reverseOnto !x (ChunksCons y ys) =
|
|
reverseOnto (ChunksCons y x) ys
|
|
|
|
unI :: Int -> Int#
|
|
unI (I# i) = i
|
|
|
|
unBa :: ByteArray -> ByteArray#
|
|
unBa (ByteArray x) = x
|