diff --git a/LICENSE b/LICENSE new file mode 100644 --- /dev/null +++ b/LICENSE @@ -0,0 +1,30 @@ +Copyright John W. Lato 2010 + +All rights reserved. + +Redistribution and use in source and binary forms, with or without +modification, are permitted provided that the following conditions are met: + + * Redistributions of source code must retain the above copyright + notice, this list of conditions and the following disclaimer. + + * Redistributions in binary form must reproduce the above + copyright notice, this list of conditions and the following + disclaimer in the documentation and/or other materials provided + with the distribution. + + * Neither the name of John W. Lato nor the names of other + contributors may be used to endorse or promote products derived + from this software without specific prior written permission. + +THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS +"AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT +LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR +A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT +OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, +SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT +LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, +DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY +THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT +(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. diff --git a/Setup.hs b/Setup.hs new file mode 100644 --- /dev/null +++ b/Setup.hs @@ -0,0 +1,3 @@ +#!/usr/bin/env runhaskell +import Distribution.Simple +main = defaultMain diff --git a/mutable-iter.cabal b/mutable-iter.cabal new file mode 100644 --- /dev/null +++ b/mutable-iter.cabal @@ -0,0 +1,34 @@ +Name: mutable-iter +Version: 0.2 +Synopsis: iteratees based upon mutable buffers +Description: Provides iteratees backed by mutable buffers. This enables iteratees to run without any extra memory allocations. +Homepage: http://tanimoto.us/~jwlato/haskell/mutable-iter +License: BSD3 +License-file: LICENSE +Author: John W. Lato +Maintainer: jwlato@gmail.com +Copyright: John W. Lato, 2010 +Stability: Experimental +Tested-with: GHC == 6.12.1 +Category: Data +Build-type: Simple + +-- Extra files to be distributed with the package, such as examples or +-- a README. +-- Extra-source-files: + +-- Constraint on the version of Cabal needed to build this package. +Cabal-version: >=1.2 + + +Library + hs-source-dirs: src + Exposed-modules: Data.MutableIter + ,Data.MutableIter.Binary + ,Data.MutableIter.IOBuffer + Build-depends: + base >= 4 && < 6, + iteratee >= 0.5 && < 0.6, + MonadCatchIO-transformers >= 0.2 && < 0.3, + transformers >= 0.2 && < 0.3, + vector >= 0.4 && < 0.5 diff --git a/src/Data/MutableIter.hs b/src/Data/MutableIter.hs new file mode 100644 --- /dev/null +++ b/src/Data/MutableIter.hs @@ -0,0 +1,349 @@ +{-# LANGUAGE + FlexibleInstances + ,DeriveFunctor + ,RankNTypes + ,ScopedTypeVariables #-} + +module Data.MutableIter ( + MIteratee (..) + ,IOBuffer (..) + ,IB.createIOBuffer + ,MEnumerator + ,MEnumeratee + ,joinIob + ,joinIM + ,wrapEnum + ,liftI + ,idone + ,icont + ,guardNull + ,head + ,heads + ,peek + ,drop + ,dropWhile + ,foldl' + ,mapStream + ,mapAccum + ,convStream + ,takeUpTo + ,enumHandleRandom + ,fileDriverRandom + ,newFp + ) + +where + +import Prelude hiding (head, null, drop, dropWhile, catch) +import qualified Prelude as P + +import qualified Data.MutableIter.IOBuffer as IB +import Data.MutableIter.IOBuffer (IOBuffer, null, empty, pop) + +import qualified Data.Iteratee as I +import Data.Maybe + +import Control.Exception (SomeException, IOException) +import Control.Monad +import Control.Monad.CatchIO as CIO +import Control.Monad.IO.Class +import Control.Monad.Trans.Class +import Foreign.ForeignPtr +import Foreign.Storable (Storable, sizeOf, poke) +import Foreign.Marshal.Array + +import System.IO + +isNullChunk :: (MonadCatchIO s) => I.Stream (IOBuffer r el) -> s Bool +isNullChunk (I.EOF _) = return False +isNullChunk (I.Chunk s) = liftIO $ null s + +newtype MIteratee s m a = MIteratee {unwrap :: I.Iteratee s m a} deriving (Functor) + +instance (MonadCatchIO m, Storable el) => Monad (MIteratee (IOBuffer r el) m) where + {-# INLINE return #-} + return x = MIteratee (I.Iteratee $ \onDone _ -> onDone x (I.Chunk empty)) + -- {-# INLINE (>>=) #-} -- this inline makes things a bit slower with 6.12. Seems fixed in ghc-7 + m >>= f = MIteratee (I.Iteratee $ \onDone onCont -> + let mDone a str = do + isNull <- isNullChunk str + if isNull + then I.runIter (unwrap $ f a) onDone onCont + else I.runIter (unwrap $ f a) (const . flip onDone str) (fCont str) + fCont str k Nothing = I.runIter (k str) onDone onCont + fCont _ k e = onCont k e + in I.runIter (unwrap m) mDone (\k -> onCont (unwrap . (>>= f) . MIteratee . k))) + +instance (Storable el) => MonadTrans (MIteratee (IOBuffer s el)) where + lift m = MIteratee $ I.Iteratee $ \od _ -> m >>= flip od (I.Chunk empty) + +instance (MonadCatchIO m, Storable el) + => MonadIO (MIteratee (IOBuffer r el) m) where + liftIO = lift . liftIO + +instance (MonadCatchIO m, Storable el) => + MonadCatchIO (MIteratee (IOBuffer r el) m) where + catch m f = MIteratee $ I.Iteratee $ \od oc -> + I.runIter (unwrap m) od oc `catch` (\e -> I.runIter (unwrap $ f e) od oc) + block = mapIteratee block + unblock = mapIteratee unblock + +mapIteratee :: (Monad m, Storable el) => + (m a -> m b) + -> MIteratee (IOBuffer r el) m a + -> MIteratee (IOBuffer r el) m b +mapIteratee f = lift . f . I.run . unwrap + +joinIob :: (MonadCatchIO m, Storable el) => + MIteratee (IOBuffer r el) m (MIteratee (IOBuffer r el') m a) + -> MIteratee (IOBuffer r el) m a +joinIob outer = outer >>= \inner -> MIteratee $ I.Iteratee $ \od oc -> + let onDone x _ = od x (I.Chunk empty) + onCont k Nothing = I.runIter (k (I.EOF Nothing)) onDone onCont' + onCont _ (Just e) = I.runIter (I.throwErr e) od oc + onCont' _ e = I.runIter (I.throwErr (fromMaybe excDiv e)) od oc + in I.runIter (unwrap inner) onDone onCont + +joinIM :: (Monad m) => + m (MIteratee (IOBuffer r el) m a) + -> MIteratee (IOBuffer r el) m a +joinIM mIter = MIteratee (I.Iteratee (\od oc -> mIter >>= (\iter -> + I.runIter (unwrap iter) od oc))) + +excDiv :: SomeException +excDiv = toException I.DivergentException + +type MEnumerator s m a = MIteratee s m a -> m (MIteratee s m a) + +wrapEnum :: (Monad m) => I.Enumerator s m a -> MEnumerator s m a +wrapEnum enum = liftM MIteratee . enum . unwrap + +liftI :: (Monad m) => (I.Stream s -> MIteratee s m a) -> MIteratee s m a +liftI = MIteratee . I.liftI . (unwrap .) + +idone x str = MIteratee $ I.idone x str + +icont k mErr = MIteratee $ I.icont (unwrap . k) mErr + +guardNull :: (MonadCatchIO m, Storable el) => + IOBuffer r el + -> MIteratee (IOBuffer r el) m a + -> MIteratee (IOBuffer r el) m a + -> MIteratee (IOBuffer r el) m a +guardNull buf onEmpty onFull = do + isNull <- liftIO $ null buf + if isNull then onEmpty else onFull + +head :: (MonadCatchIO m, Storable el) => MIteratee (IOBuffer r el) m el +head = liftI step + where + step (I.Chunk buf) = guardNull buf (icont step Nothing) $ do + x <- liftIO $ pop buf + idone x (I.Chunk buf) + step str = MIteratee . I.throwErr . toException $ I.EofException + +peek :: (MonadCatchIO m, Storable el) => MIteratee (IOBuffer r el) m (Maybe el) +peek = liftI step + where + step c@(I.Chunk buf) = guardNull buf (icont step Nothing) $ do + x <- liftIO $ IB.lookAtHead buf + idone x c + step str = idone Nothing str + +heads :: (MonadCatchIO m, Storable el, Eq el) => + [el] + -> MIteratee (IOBuffer r el) m Int +heads st = loop 0 st + where + loop cnt [] = return cnt + loop cnt str = liftI (step cnt str) + step cnt s@(x:xs) (I.Chunk buf) = do + mx <- liftIO $ IB.lookAtHead buf + maybe (icont (step cnt s) Nothing) (\h -> if h == x + then liftIO (pop buf) >> step (succ cnt) xs (I.Chunk buf) + else idone cnt (I.Chunk buf)) mx + step cnt s str = idone cnt str + +drop :: (MonadCatchIO m, Storable el) => Int -> MIteratee (IOBuffer r el) m () +drop n = liftI (step n) + where + step 0 str = idone () str + step n' (I.Chunk buf) = guardNull buf (icont (step n') Nothing) $ do + l <- liftIO $ IB.length buf + if n' < l + then liftIO (IB.drop n' buf) >> idone () (I.Chunk buf) + else liftIO (IB.drop l buf) >> liftI (step (n' - l)) + step _ str = idone () str + +dropWhile :: (MonadCatchIO m, Storable el) => + (el -> Bool) + -> MIteratee (IOBuffer r el) m () +dropWhile pred = liftI step + where + step (I.Chunk buf) = guardNull buf (icont step Nothing) $ do + liftIO $ IB.dropWhile pred buf + l <- liftIO $ IB.length buf + if l == 0 then icont step Nothing else idone () (I.Chunk buf) + step str = idone () str + +foldl' :: (MonadCatchIO m, Storable el, Show a) => + (a -> el -> a) + -> a + -> MIteratee (IOBuffer r el) m a +foldl' f acc' = liftI (step acc') + where + step acc (I.Chunk buf) = guardNull buf (liftI (step acc)) $ do + newacc <- liftIO $ IB.foldl' f acc buf + icont (step newacc) Nothing + step acc str = idone acc str + + +-- ----------------------------------------------------------- +-- Enumeratees + +type MEnumeratee sFrom sTo m a = MIteratee sTo m a -> MIteratee sFrom m (MIteratee sTo m a) + +eneeCheckIfDone :: + (Monad m, Storable elo) => + ((I.Stream eli -> MIteratee eli m a) + -> MIteratee (IOBuffer r elo) m (MIteratee eli m a)) + -> MEnumeratee (IOBuffer r elo) eli m a +eneeCheckIfDone f inner = MIteratee $ I.Iteratee $ \od oc -> + let onDone x s = od (idone x s) (I.Chunk empty) + onCont k Nothing = I.runIter (unwrap $ f (MIteratee . k)) od oc + onCont _ (Just e) = I.runIter (I.throwErr e) od oc + in I.runIter (unwrap inner) onDone (onCont) + +mapStream :: (MonadCatchIO pr, Storable elo, Storable eli) => + Int + -> (eli -> elo) + -> MEnumeratee (IOBuffer r eli) + (IOBuffer r elo) + pr + a +mapStream maxlen f i = do + offp <- liftIO $ newFp 0 + bufp <- liftIO $ mallocForeignPtrArray maxlen + goMap offp bufp i + where + goMap offp bufp = eneeCheckIfDone (liftI . step offp bufp) + step offp bufp k (I.Chunk buf) = + guardNull buf (liftI (step offp bufp k)) $ do + newIOBuf <- liftIO $ IB.mapBuffer f offp bufp buf + goMap offp bufp $ k (I.Chunk newIOBuf) + step _ _ k s = idone (liftI k) s + +mapAccum :: (MonadCatchIO pr, Storable eli, Storable elo) => + Int + -> (b -> eli -> (b,elo)) + -> b + -> MEnumeratee (IOBuffer r eli) (IOBuffer r elo) pr a +mapAccum maxlen f acc i = do + offp <- liftIO $ newFp 0 + bufp <- liftIO $ mallocForeignPtrArray maxlen + goMap offp bufp acc i + where + goMap offp bufp acc = eneeCheckIfDone (liftI . step offp bufp acc) + step offp bufp acc k (I.Chunk buf) = + guardNull buf (liftI (step offp bufp acc k)) $ do + (newAcc, newBuf) <- liftIO $ IB.mapAccumBuffer f offp bufp acc buf + goMap offp bufp newAcc $ k (I.Chunk newBuf) + step _ _ _ k s = idone (liftI k) s + +convStream :: (MonadCatchIO pr, Storable elo, Storable eli) => + MIteratee (IOBuffer r eli) pr (IOBuffer r elo) + -> MEnumeratee (IOBuffer r eli) + (IOBuffer r elo) + pr + a +convStream fi = eneeCheckIfDone check + where + check k = MIteratee (I.isStreamFinished) >>= + maybe (step k) (idone (liftI k) . I.EOF . Just) + step k = fi >>= convStream fi . k . I.Chunk + +takeUpTo :: (MonadCatchIO pr, Storable el) => + Int + -> MEnumeratee (IOBuffer r el) (IOBuffer r el) pr a +takeUpTo 0 iter = return iter +takeUpTo i iter = MIteratee $ I.Iteratee $ \od oc -> + I.runIter (unwrap iter) (onDone od oc) (onCont od oc) + where + onDone od oc x _ = od (return x) (I.Chunk empty) + onCont od oc k Nothing = + if i == 0 then od (liftI (MIteratee . k)) (I.Chunk empty) + else I.runIter (unwrap . liftI $ step i (MIteratee . k)) od oc + onCont od oc _ (Just e) = I.runIter (I.throwErr e) od oc + step n k (I.Chunk buf) = guardNull buf (liftI (step n (k))) $ do + blen <- liftIO $ IB.length buf + if blen <= n then takeUpTo (n - blen) $ k (I.Chunk buf) + else do + (s1, s2) <- liftIO $ IB.splitAt buf n + idone (k (I.Chunk s1)) (I.Chunk s2) + step _ k str = idone (k str) str + +-- --------------------------------------------- +-- drivers and enumerators + +makeHandleCallback :: (MonadCatchIO m, Storable el) => + ForeignPtr Int + -> ForeignPtr el + -> Int + -> Handle + -> () + -> m (Either SomeException ((Bool, ()), (IOBuffer r el))) +makeHandleCallback offp fp bsize h () = liftIO $ do + n' <- withForeignPtr fp $ \p -> (CIO.try $ (hGetBuf h p bsize) :: + IO (Either SomeException Int)) + case n' of + Left e -> return $ Left e + Right 0 -> return $ Right ((False, ()), empty) + Right n -> do + withForeignPtr offp $ flip poke 0 + return . (\s -> Right ((True, ()), s)) $ + IB.createIOBuffer (fromIntegral n) offp fp + +enumHandleCatch + :: + forall e r el m a.(I.IException e, MonadCatchIO m, Storable el) => + Int + -> Handle + -> (e -> m (Maybe I.EnumException)) + -> MIteratee (IOBuffer r el) m a + -> m (MIteratee (IOBuffer r el) m a) +enumHandleCatch bs h handler i = do + let numbytes = bs * sizeOf (undefined :: el) + bufp <- liftIO $ mallocForeignPtrArray bs + offp <- liftIO $ newFp 0 + liftM MIteratee $ I.enumFromCallbackCatch + (makeHandleCallback offp bufp numbytes h) + handler + () + (unwrap i) + +enumHandleRandom :: forall r el m a.(MonadCatchIO m, Storable el) => + Int -- ^ Buffer size (number of elements per read) + -> Handle + -> MIteratee (IOBuffer r el) m a + -> m (MIteratee (IOBuffer r el) m a) +enumHandleRandom bs h i = enumHandleCatch bs h handler i + where + handler (I.SeekException off) = + liftM (either + (Just . I.EnumException :: IOException -> Maybe I.EnumException) + (const Nothing)) + . liftIO . CIO.try $ hSeek h AbsoluteSeek $ fromIntegral off + +fileDriverRandom :: (MonadCatchIO m, Storable el) => + Int + -> (forall r. MIteratee (IOBuffer r el) m a) + -> FilePath + -> m a +fileDriverRandom bufsize iter filepath = CIO.bracket + (liftIO $ openBinaryFile filepath ReadMode) + (liftIO . hClose) + ((I.run . unwrap) <=< flip (enumHandleRandom bufsize) iter) + +newFp :: Storable a => a -> IO (ForeignPtr a) +newFp a = mallocForeignPtr >>= \fp -> + withForeignPtr fp (flip poke a) >> return fp diff --git a/src/Data/MutableIter/Binary.hs b/src/Data/MutableIter/Binary.hs new file mode 100644 --- /dev/null +++ b/src/Data/MutableIter/Binary.hs @@ -0,0 +1,80 @@ +-- |Iteratees for parsing binary data. +module Data.MutableIter.Binary ( + -- * Types + Endian (..), + -- * Endian multi-byte iteratees + endianRead2, + endianRead3, + endianRead4 +) +where + +import Data.MutableIter as I +import qualified Data.MutableIter.IOBuffer as IB +import Data.MutableIter.IOBuffer (IOBuffer) +import Data.Iteratee.Binary (Endian (..)) + +import Data.Word +import Data.Bits +import Data.Int +import Control.Monad.CatchIO + + +-- ------------------------------------------------------------------------ +-- Binary Random IO Iteratees + +-- Iteratees to read unsigned integers written in Big- or Little-endian ways + +endianRead2 + :: (MonadCatchIO m) => + Endian + -> MIteratee (IOBuffer r Word8) m Word16 +endianRead2 e = do + c1 <- I.head + c2 <- I.head + case e of + MSB -> return $ (fromIntegral c1 `shiftL` 8) .|. fromIntegral c2 + LSB -> return $ (fromIntegral c2 `shiftL` 8) .|. fromIntegral c1 + +-- |read 3 bytes in an endian manner. If the first bit is set (negative), +-- set the entire first byte so the Word32 can be properly set negative as +-- well. +endianRead3 + :: (MonadCatchIO m) => + Endian + -> MIteratee (IOBuffer r Word8) m Word32 +endianRead3 e = do + c1 <- I.head + c2 <- I.head + c3 <- I.head + case e of + MSB -> return $ (((fromIntegral c1 + `shiftL` 8) .|. fromIntegral c2) + `shiftL` 8) .|. fromIntegral c3 + LSB -> + let m :: Int32 + m = shiftR (shiftL (fromIntegral c3) 24) 8 in + return $ (((fromIntegral c3 + `shiftL` 8) .|. fromIntegral c2) + `shiftL` 8) .|. fromIntegral m + +endianRead4 + :: (MonadCatchIO m) => + Endian + -> MIteratee (IOBuffer r Word8) m Word32 +endianRead4 e = do + c1 <- I.head + c2 <- I.head + c3 <- I.head + c4 <- I.head + case e of + MSB -> return $ + (((((fromIntegral c1 + `shiftL` 8) .|. fromIntegral c2) + `shiftL` 8) .|. fromIntegral c3) + `shiftL` 8) .|. fromIntegral c4 + LSB -> return $ + (((((fromIntegral c4 + `shiftL` 8) .|. fromIntegral c3) + `shiftL` 8) .|. fromIntegral c2) + `shiftL` 8) .|. fromIntegral c1 diff --git a/src/Data/MutableIter/IOBuffer.hs b/src/Data/MutableIter/IOBuffer.hs new file mode 100644 --- /dev/null +++ b/src/Data/MutableIter/IOBuffer.hs @@ -0,0 +1,286 @@ +{-# LANGUAGE BangPatterns, ScopedTypeVariables #-} + +module Data.MutableIter.IOBuffer ( + IOBuffer + ,createIOBuffer + ,null + ,empty + ,copyBuffer + ,append + ,length + ,pop + ,lookAtHead + ,drop + ,dropWhile + ,take + ,splitAt + ,mapBuffer + ,mapAccumBuffer + ,foldl' + ,castBuffer + ,freeze + ,hPut + ,unsafeToForeignPtr +) + +where + +import Prelude hiding (length, take, drop, null, splitAt, dropWhile) + +import qualified Data.Vector.Storable as V +import qualified Data.Iteratee as I +import Control.Monad +import Control.Monad.CatchIO + +import Foreign +import System.IO + +import System.IO.Unsafe (unsafePerformIO) + + +-- |A mutable buffer to hold storable elements. This data type supports +-- memory recycling. +data IOBuffer r el = IOBuffer {-# UNPACK #-} !Int + {-# UNPACK #-} !(ForeignPtr Int) + {-# UNPACK #-} !(ForeignPtr el) + +instance Storable el => I.NullPoint (IOBuffer r el) where + empty = empty + +fpPeek :: Storable el => ForeignPtr el -> IO el +fpPeek = flip withForeignPtr peek + +fpPoke :: Storable el => ForeignPtr el -> el -> IO () +fpPoke fp el = withForeignPtr fp (flip poke el) + +withBuf :: IOBuffer r el -> (Int -> Ptr Int -> Ptr el -> IO a) -> IO a +withBuf (IOBuffer l ofp ofb) f = withForeignPtr ofp (\op -> + withForeignPtr ofb (\ob -> f l op ob)) +{-# INLINE withBuf #-} + +newFp :: Storable a => a -> IO (ForeignPtr a) +newFp a = mallocForeignPtr >>= \fp -> fpPoke fp a >> return fp + +-- |Create a buffer from a length and data array. +createIOBuffer :: (Storable el) => + Int + -> ForeignPtr Int + -> ForeignPtr el + -> IOBuffer r el +createIOBuffer len op buf = IOBuffer len op buf + +-- |Empty buffer. +empty :: Storable el => IOBuffer r el +empty = IOBuffer 0 nullForeignPtr nullForeignPtr + +nullForeignPtr = unsafePerformIO (newForeignPtr_ nullPtr) + +-- |Check if the buffer is empty. +null :: IOBuffer r el -> IO Bool +null (IOBuffer 0 _ _) = return True +null buf = withBuf buf $ \l po _ -> liftM (>= l) $ peek po +{-# INLINE null #-} + +-- |IOBuffer length. +length :: IOBuffer r el -> IO Int +length buf = withBuf buf $ \l po _ -> liftM (l -) $ peek po +{-# INLINE length #-} + +-- |Retrieve the front element from the buffer and advance the internal pointer. +-- It is an error to call this on an empty buffer. +pop :: (Storable el) => IOBuffer r el -> IO el +pop (IOBuffer 0 _ _ ) = error "Can't pop head off of empty buffer" +pop buf = withBuf buf $ \l po pb -> do + off <- peek po + if off >= l then error "Can't pop head from empty buffer" + else poke po (off+1) >> peek (pb `advancePtr` off) +{-# INLINE pop #-} + +-- |Retrieve the first element, if it exists. +-- This function does not advance the buffer pointer. +lookAtHead :: (Storable el) => + IOBuffer r el -> IO (Maybe el) +lookAtHead (IOBuffer 0 _ _ ) = return Nothing +lookAtHead buf = withBuf buf $ \l po pb -> do + off <- peek po + if off >= l then return Nothing + else liftM Just $ peek (pb `advancePtr` off) +{-# INLINE lookAtHead #-} + +-- |Drop n elements from the front of the buffer. +-- if the buffer has fewer elements, all are dropped. +drop :: Int -> IOBuffer r el -> IO () +drop n_drop buf = withBuf buf $ \l po pb -> do + off <- peek po + poke po (min l (off+n_drop)) +{-# INLINE drop #-} + +dropWhile :: (Storable el) => + (el -> Bool) + -> IOBuffer r el + -> IO () +dropWhile pred buf = withBuf buf $ \l po pb -> do + off <- peek po + let len = l-off + let go cnt p | cnt < len = do + this <- peek p + if pred this then go (succ cnt) (p `advancePtr` 1) else return cnt + let go cnt _ = return cnt --off the end of the buffer + n <- go 0 (pb `advancePtr` off) + poke po n +{-# INLINE dropWhile #-} + +-- |Create a new buffer from the first n elements, sharing data. +-- This function advances the pointer of the original buffer. +take :: (Storable el) => + IOBuffer r el + -> Int + -> IO (IOBuffer r el) +take (IOBuffer 0 _ _ ) _ = return empty +take buf@(IOBuffer _ fpo fpb) n_take = withBuf buf $ \l po pb -> do + off <- peek po + po' <- newFp off + poke po $ min l (off+n_take) + return $ IOBuffer l po' fpb + +-- |Split one buffer to two, sharing storage. +splitAt :: (Storable el) => + IOBuffer r el + -> Int + -> IO (IOBuffer r el, IOBuffer r el) +splitAt (IOBuffer 0 _ _) _ = return (empty, empty) +splitAt buf 0 = return (empty, buf) +splitAt buf@(IOBuffer l fpo fpb) n + | n>0 && n <= l = withForeignPtr fpo $ \po -> withForeignPtr fpb $ \pb -> do + let ib1 = IOBuffer n fpo fpb + off <- peek po + po2 <- newFp (off+n) + return (ib1, IOBuffer l po2 fpb) + | True = return (buf, empty) + +-- |Copy data from one buffer to another. +copyBuffer :: (Storable el) => + IOBuffer r el -> IO (IOBuffer r el) +copyBuffer (IOBuffer 0 _ _) = return empty +copyBuffer buf = withBuf buf $ \l po pb -> do + off <- peek po + let len' = l-off + if len' > 0 then do + po' <- newFp 0 + pb' <- mallocForeignPtrArray len' + withForeignPtr pb' $ \p -> copyArray p (pb `advancePtr` off) len' + return $ IOBuffer len' po' pb' + else return empty + +-- |Append two buffers. Copies data from both into a new buffer. +append :: (Storable el) => + IOBuffer r el + -> IOBuffer r el + -> IO (IOBuffer r el) +append ib1 (IOBuffer 0 _ _) = copyBuffer ib1 +append (IOBuffer 0 _ _) ib2 = copyBuffer ib2 +append ib1 ib2 = withBuf ib1 (\l1 po1 pb1 -> withBuf ib2 (\l2 po2 pb2 -> do + len1 <- length ib1 + len2 <- length ib2 + let len = len1 + len2 + off1 <- peek po1 + off2 <- peek po2 + po <- newFp 0 + pb <- mallocForeignPtrArray len + withForeignPtr pb $ \p -> copyArray p (pb1 `advancePtr` off1) len1 + withForeignPtr pb $ \p -> copyArray (p `advancePtr` len1) + (pb2 `advancePtr` off2) len2 + return $ IOBuffer len po pb)) + +-- |Safely convert an IOBuffer to a Vector. +-- TODO: finish this implementation, look at Vector docs. +freeze :: (Storable el) => IOBuffer r el -> IO (V.Vector el) +freeze (IOBuffer 0 _ _) = return V.empty +freeze buf = withBuf buf $ \l po pb -> do + off <- peek po + return undefined + +-- |Write out the contents of the IOBuffer to a handle. This operation +-- drains the buffer. +hPut :: forall r el. (Storable el) => Handle -> IOBuffer r el -> IO () +hPut h (IOBuffer 0 _ _) = return () +hPut h buf = withBuf buf $ \l po pb -> do + off <- peek po + (hPutBuf h) (pb `advancePtr` off) (bytemult * (l-off)) + poke po l + where + bytemult = sizeOf (undefined :: el) + +-- |copy data from one buffer to another with the specified map function. +-- this operation drains the original buffer. +mapBuffer :: (Storable el, Storable el') => + (el -> el') + -> ForeignPtr Int + -> ForeignPtr el' + -> IOBuffer r el + -> IO (IOBuffer r el') +mapBuffer f dp datap (IOBuffer 0 _ _) = return empty +mapBuffer f fdp fdatap buf = withForeignPtr fdp $ \dp -> + withForeignPtr fdatap $ \datap -> withBuf buf $ \l po pb -> + let go !cnt ip' op' = when (cnt>0) $ do + poke op' . f =<< peek ip' + go (cnt-1) (ip' `advancePtr` 1) (op' `advancePtr` 1) + in do + off <- peek po + go (l-off) (pb `advancePtr` off) datap + poke dp 0 + poke po l + return $ createIOBuffer (l-off) fdp fdatap +{-# INLINE mapBuffer #-} + +mapAccumBuffer :: (Storable el, Storable el') => + (acc -> el -> (acc,el')) + -> ForeignPtr Int + -> ForeignPtr el' + -> acc + -> IOBuffer r el + -> IO (acc, IOBuffer r el') +mapAccumBuffer f fdp fdatap acc (IOBuffer 0 _ _) = return (acc, empty) +mapAccumBuffer f fdp fdatap acc buf = withForeignPtr fdp $ \dp -> + withForeignPtr fdatap $ \datap -> withBuf buf $ \l po pb -> + let go !acc' cnt !ip' !op' + | cnt == 0 = return acc' + | True = do + (acc2, el) <- liftM (f acc') $ peek ip' + poke op' el + go acc2 (cnt-1) (ip' `advancePtr` 1) (op' `advancePtr` 1) + in do + off <- peek po + acc' <- go acc (l-off) (pb `advancePtr` off) datap + poke dp 0 + poke po l + return $ (acc', createIOBuffer (l-off) fdp fdatap) +{-# INLINE mapAccumBuffer #-} + + +-- |Cast a buffer to a different type. Any extra data is truncated. +-- This is not safe unless the buffer offset is 0. +castBuffer :: forall r m el el'. (Storable el, Storable el') => + IOBuffer r el -> IO (IOBuffer r el') +castBuffer (IOBuffer 0 _ _) = return empty +castBuffer (IOBuffer l po pb) = do + off <- fpPeek po + when (off /= 0) (error "castBuffer called with non-zero offset") + return $ IOBuffer l' po (castForeignPtr pb) + where + l' = (l * sizeOf (undefined :: el)) `div` sizeOf (undefined :: el') + +foldl' :: (Storable b) => (a -> b -> a) -> a -> IOBuffer r b -> IO a +foldl' f acc (IOBuffer 0 _ _) = return acc +foldl' f i0 buf = withBuf buf $ \l po pb -> + let go !n !acc p | n>0 = do + el <- peek p + go (n-1) (f acc el) (p `advancePtr` 1) + go _ acc _ = return acc + in do + off <- peek po + go (l-off) i0 (pb `advancePtr` off) +{-# INLINE foldl' #-} + +unsafeToForeignPtr :: IOBuffer r el -> (Int, ForeignPtr Int, ForeignPtr el) +unsafeToForeignPtr (IOBuffer l po pb) = (l,po,pb)