Apply a Word32 function to a ByteString

Question

I have a ByteString b with a length of the form 4*n (with integer n) and want to use map with a function f::Word32->Word32 on b (so that f gets applied to "b[0..3]", "b[4..7]", etc). How can this be done in an efficient (and elegant) way?

Answer 1

You can chunk up the ByteString very efficiently, given that B.take and B.drop are both O(1) operations:

import Data.ByteString (ByteString)
import qualified Data.ByteString as B

chunk :: Int -> ByteString -> [ByteString]
chunk k = takeWhile (not . B.null) . map (B.take k) . iterate (B.drop k) 

then:

\> :set -XOverloadedStrings
\> chunk 4 "abcdefghijkl"
["abcd","efgh","ijkl"]

the rest would be to map over the list converting to and from desired type, and a single call to B.concat at the end.

A possible fromByteString could be implemented using bit shifts and left fold:

import Data.Bits (Bits, shiftL, (.|.))

fromByteString :: (Num a, Bits a) => ByteString -> a
fromByteString = B.foldl go 0
    where go acc i = (acc  `shiftL` 8) .|. (fromIntegral i)

then:

\> map fromByteString $ chunk 4 "abcdefghijkl" :: [Word32]
[1633837924,1701209960,1768581996]

Answer 2

A hackish but efficient way is to convert to a storable vector (which does not require a copy!), map over that, and convert back:

import Data.Vector.Storable.ByteString
import qualified Data.Vector.Storable as VS
import qualified Data.ByteString as BS

mapBSChunks :: (VS.Storable a, VS.Storable b)
              => (a->b) -> BS.ByteString -> BS.ByteString
mapBSChunks f = vectorToByteString . VS.map f . byteStringToVector

As per Michael's comment, you can easily define those hackish conversion functions locally:

bytestringToVector bs = runST
   ( V.unsafeThaw (toByteVector bs) >>= V.unsafeFreeze . M.unsafeCast )
vectorToByteString v = runST
   ( V.unsafeThaw v >>= fmap fromByteVector . V.unsafeFreeze . M.unsafeCast )

though I'd rather depend on a library to provide this, in particular because the unsafe casting is a bit fishy.