efficient Bytestring mapM_ for IO/ST?

[wren ng thornton]

On 3/21/11 4:40 PM, Brandon Moore wrote:
> Is there an efficient way to iterate over the bytes of a ByteString?

The code I've been using (rather similar to your unsafe map) is:

     import qualified Data.ByteString.Internal as BSI
     import qualified Foreign.ForeignPtr       as FFI

     foldIO :: (a -> Word8 -> IO a) -> a -> ByteString -> IO a
     foldIO f z0 (BSI.PS fp off len) =
         FFI.withForeignPtr fp $ \p0 -> do
             let q = p0 `plusPtr` (off+len)
             let go z p
                     | z `seq` p `seq` False = undefined
                     | p == q    = return z
                     | otherwise = do
                         w  <- peek p
                         z' <- f z w
                         go z' (p `plusPtr` 1)
             go z0 (p0 `plusPtr` off)
     {-# INLINE foldIO #-}

Some things to note:

* It's a left fold rather than a right fold, just like foldM, except 
that we can't generalize it to work for all monads. (We could do a right 
fold just as easily by starting with p0`plusPtr`(off+len) and counting 
down to p0`plusPtr`off if desired.)

* Because we're just keeping the head pointer, we can increment it as we 
go instead of using peekElemOff. This improves the performance by only 
performing one addition per loop (the p++) instead of two (ix++ and 
*(p+ix)), and by requiring one less register (for ix).

* The inline pragma helps performance in a major way.

* I haven't actually looked at Core nor tried much to optimize it. This 
just seems like the easiest way to allow the accumulator to perform IO 
instead of being pure. (For pure code there's foldl'.)

-- 
Live well,
~wren