diff --git a/biohazard.cabal b/biohazard.cabal --- a/biohazard.cabal +++ b/biohazard.cabal @@ -1,5 +1,5 @@ Name: biohazard -Version: 0.6.10 +Version: 0.6.13 Synopsis: bioinformatics support library Description: This is a collection of modules I separated from various bioinformatics tools. The hope is to make @@ -47,15 +47,25 @@ Bio.Bam.Writer, Bio.Base, Bio.Iteratee, + Bio.Iteratee.Base, Bio.Iteratee.Bgzf, + Bio.Iteratee.Binary, Bio.Iteratee.Builder, + Bio.Iteratee.Char, + Bio.Iteratee.Exception, + Bio.Iteratee.IO, + Bio.Iteratee.IO.Base, + Bio.Iteratee.IO.Fd, + Bio.Iteratee.IO.Handle, + Bio.Iteratee.Iteratee, + Bio.Iteratee.ListLike, + Bio.Iteratee.ReadableChunk, Bio.Iteratee.ZLib, Bio.Prelude, Bio.PriorityQueue, Bio.TwoBit, Bio.Util.Numeric, - Bio.Util.Zlib, - Paths_biohazard + Bio.Util.Zlib Build-depends: aeson >= 0.7 && < 1.1, async >= 2.0 && < 2.2, @@ -70,22 +80,23 @@ exceptions >= 0.6 && < 0.9, filepath >= 1.3 && < 2.0, hashable >= 1.0 && < 1.3, - iteratee >= 0.8.9.6 && < 0.8.10, ListLike >= 3.0 && < 5.0, + monad-control == 1.0.*, primitive >= 0.5 && < 0.7, random >= 1.0 && < 1.2, scientific == 0.3.*, stm == 2.4.*, text >= 1.0 && < 2.0, transformers >= 0.4.1 && < 0.6, + transformers-base >= 0.4 && < 0.6, unix >= 2.5 && < 2.8, unordered-containers >= 0.2.3 && < 0.3, vector == 0.11.*, vector-algorithms >= 0.3 && < 1.0, vector-th-unbox == 0.2.*, - zlib >= 0.5 && < 0.7 + zlib == 0.6.* - Ghc-options: -Wall -fprof-auto + Ghc-options: -Wall Default-Language: Haskell2010 @@ -102,17 +113,21 @@ Other-Extensions: CPP, DeriveGeneric, ExistentialQuantification, + FunctionalDependencies, GeneralizedNewtypeDeriving, PatternGuards, Rank2Types, ScopedTypeVariables, TemplateHaskell, + TupleSections, TypeFamilies, - TypeOperators + TypeOperators, + UndecidableInstances Hs-source-dirs: src Install-Includes: src/cbits/myers_align.h - C-sources: src/cbits/myers_align.c, + C-sources: src/cbits/loops.c, + src/cbits/myers_align.c, src/cbits/trim.c CC-options: -fPIC diff --git a/src/Bio/Adna.hs b/src/Bio/Adna.hs --- a/src/Bio/Adna.hs +++ b/src/Bio/Adna.hs @@ -33,7 +33,7 @@ import Bio.Bam import Bio.Prelude import Bio.TwoBit -import Data.Aeson hiding ( pairs ) +import Data.Aeson import qualified Data.Vector as V import qualified Data.Vector.Generic as G @@ -334,8 +334,8 @@ guard (not $ isUnmapped b) let ref_nm = sq_name $ getRef refs b_rname ref = getFragment tbf ref_nm (b_pos - ctx) (alignedLength b_cigar + 2*ctx) - pairs = aln_from_ref (U.drop ctx ref) b_seq b_cigar - return (b, ft, ref, pairs)) =$ + pps = aln_from_ref (U.drop ctx ref) b_seq b_cigar + return (b, ft, ref, pps)) =$ damagePatternsIter ctx rng it -- | Enumeratee (almost) that computes some statistics from plain BAM @@ -361,9 +361,9 @@ let b@BamRec{..} = unpackBam br guard (not $ isUnmapped b) md <- getMd b - let pairs = aln_from_md b_seq b_cigar md - ref = U.map fromN $ U.filter ((/=) gap . fst) pairs - return (b, ft, ref, pairs)) =$ + let pps = aln_from_md b_seq b_cigar md + ref = U.map fromN $ U.filter ((/=) gap . fst) pps + return (b, ft, ref, pps)) =$ damagePatternsIter 0 rng it where fromN (ns,_) | ns == nucsA = 2 @@ -544,9 +544,9 @@ revcom_both :: ( BamRec, FragType, U.Vector Word8, U.Vector (Nucleotides, Nucleotides) ) -> ( BamRec, FragType, U.Vector Word8, U.Vector (Nucleotides, Nucleotides) ) -revcom_both (b, ft, ref, pairs) - | isReversed b = ( b, ft, revcom_ref ref, revcom_pairs pairs ) - | otherwise = ( b, ft, ref, pairs ) +revcom_both (b, ft, ref, pps) + | isReversed b = ( b, ft, revcom_ref ref, revcom_pairs pps ) + | otherwise = ( b, ft, ref, pps ) where revcom_ref = U.reverse . U.map (\c -> if c > 3 then c else xor c 2) revcom_pairs = U.reverse . U.map (compls *** compls) diff --git a/src/Bio/Bam/Fastq.hs b/src/Bio/Bam/Fastq.hs --- a/src/Bio/Bam/Fastq.hs +++ b/src/Bio/Bam/Fastq.hs @@ -9,7 +9,6 @@ import qualified Data.Attoparsec.ByteString.Char8 as P import qualified Data.ByteString as B import qualified Data.ByteString.Char8 as S -import qualified Data.Iteratee.ListLike as I import qualified Data.Vector.Generic as V -- ^ Parser for @FastA/FastQ@, 'Iteratee' style, based on @@ -100,9 +99,9 @@ | otherwise = Just (i-1) skipJunk :: Monad m => Iteratee Bytes m () -skipJunk = I.peek >>= check +skipJunk = peekStream >>= check where - check (Just c) | bad c = I.dropWhile (c2w '\n' /=) >> I.drop 1 >> skipJunk + check (Just c) | bad c = dropWhileStream (c2w '\n' /=) >> dropStream 1 >> skipJunk check _ = return () bad c = c /= c2w '>' && c /= c2w '@' diff --git a/src/Bio/Bam/Pileup.hs b/src/Bio/Bam/Pileup.hs --- a/src/Bio/Bam/Pileup.hs +++ b/src/Bio/Bam/Pileup.hs @@ -284,7 +284,7 @@ -- | Map quality and a list of encountered bases, with damage -- information and reference base if known. -type BasePile = [( Qual, DamagedBase )] +type BasePile = [DamagedBase] -- | Map quality and a list of encountered indel variants. The deletion -- has the reference sequence, if known, an insertion has the inserted @@ -521,19 +521,19 @@ case pchunks of _ | nwt > 0 -> b' `seq` go k (b':ac) wt bpile ipile bs Seek p' pb' -> go k ac (ins p' pb' wt) (z bpile) ipile bs - Indel nd ni pb' -> go k (pb':ac) wt (z bpile) (y ipile) bs where y = put mq (nd,ni) + Indel nd ni pb' -> go k (pb':ac) wt (z bpile) (y ipile) bs where y = put (,) mq (nd,ni) EndOfRead -> go k ac wt (z bpile) ipile bs where b' = Base (nwt-1) qs mq pchunks - z = put mq qs + z = put (\q x -> x { db_qual = min q (db_qual x) }) mq qs ins q v w = Node q v Empty Empty `unionH` w - put (Q !q) !x (!st,!vs) = ( st { read_depth = read_depth st + 1 - , reads_mapq0 = reads_mapq0 st + (if q == 0 then 1 else 0) - , sum_mapq = sum_mapq st + fromIntegral q - , sum_mapq_squared = sum_mapq_squared st + fromIntegral q * fromIntegral q } - , (Q q, x) : vs ) + put f (Q !q) !x (!st,!vs) = ( st { read_depth = read_depth st + 1 + , reads_mapq0 = reads_mapq0 st + (if q == 0 then 1 else 0) + , sum_mapq = sum_mapq st + fromIntegral q + , sum_mapq_squared = sum_mapq_squared st + fromIntegral q * fromIntegral q } + , f (Q q) x : vs ) -- | We need a simple priority queue. Here's a skew heap (specialized diff --git a/src/Bio/Bam/Reader.hs b/src/Bio/Bam/Reader.hs --- a/src/Bio/Bam/Reader.hs +++ b/src/Bio/Bam/Reader.hs @@ -255,7 +255,7 @@ decodeBam :: Monad m => (BamMeta -> Iteratee [BamRaw] m a) -> Iteratee Block m (Iteratee [BamRaw] m a) decodeBam inner = do meta <- liftBlock get_bam_header refs <- liftBlock get_ref_array - convStream getBamRaw $ inner $! merge meta refs + convStream getBamRaw $ inner $! mmerge meta refs where get_bam_header = do magic <- heads "BAM\SOH" when (magic /= 4) $ do s <- iGetString 10 @@ -275,11 +275,11 @@ -- sequences, so leftovers from the header are discarded. Merging -- is by name. So we merge information from the header into the -- list, then replace the header information. - merge meta refs = + mmerge meta refs = let tbl = M.fromList [ (sq_name sq, sq) | sq <- F.toList (meta_refs meta) ] - in meta { meta_refs = fmap (\s -> maybe s (merge' s) (M.lookup (sq_name s) tbl)) refs } + in meta { meta_refs = fmap (\s -> maybe s (mmerge' s) (M.lookup (sq_name s) tbl)) refs } - merge' l r | sq_length l == sq_length r = l { sq_other_shit = sq_other_shit l ++ sq_other_shit r } + mmerge' l r | sq_length l == sq_length r = l { sq_other_shit = sq_other_shit l ++ sq_other_shit r } | otherwise = l -- contradiction in header, but we'll just ignore it diff --git a/src/Bio/Bam/Rmdup.hs b/src/Bio/Bam/Rmdup.hs --- a/src/Bio/Bam/Rmdup.hs +++ b/src/Bio/Bam/Rmdup.hs @@ -12,7 +12,6 @@ import qualified Data.ByteString as B import qualified Data.ByteString.Char8 as T -import qualified Data.Iteratee as I import qualified Data.Map as M import qualified Data.Vector.Generic as V import qualified Data.Vector.Storable as VS @@ -118,16 +117,16 @@ nice_sort x = sortBy (comparing (V.length . b_seq)) x - mapGroups f o = I.tryHead >>= maybe (return o) (\a -> eneeCheckIfDone (mg1 f a []) o) - mg1 f a acc k = I.tryHead >>= \mb -> case mb of + mapGroups f o = tryHead >>= maybe (return o) (\a -> eneeCheckIfDone (mg1 f a []) o) + mg1 f a acc k = tryHead >>= \mb -> case mb of Nothing -> return . k . Chunk . f $ a:acc Just b | same_pos a b -> mg1 f a (b:acc) k | otherwise -> eneeCheckIfDone (mg1 f b []) . k . Chunk . f $ a:acc check_sort :: Monad m => String -> Enumeratee [BamRec] [BamRec] m a -check_sort msg out = I.tryHead >>= maybe (return out) (\a -> eneeCheckIfDone (step a) out) +check_sort msg out = tryHead >>= maybe (return out) (\a -> eneeCheckIfDone (step a) out) where - step a k = I.tryHead >>= maybe (return . k $ Chunk [a]) (step' a k) + step a k = tryHead >>= maybe (return . k $ Chunk [a]) (step' a k) step' a k b | (b_rname a, b_pos a) > (b_rname b, b_pos b) = fail $ "rmdup: " ++ msg | otherwise = eneeCheckIfDone (step b) . k $ Chunk [a] diff --git a/src/Bio/Bam/Trim.hs b/src/Bio/Bam/Trim.hs --- a/src/Bio/Bam/Trim.hs +++ b/src/Bio/Bam/Trim.hs @@ -224,15 +224,15 @@ merge_seqs v1 v2 v3 v4 = V.zipWith4 zz v1 v2 v3 v4 where - zz !n1 (Q !q1) !n2 (Q !q2) = if n1 == n2 then n1 - else if q1 > q2 then n1 - else n2 + zz !n1 (Q !q1) !n2 (Q !q2) | n1 == compls n2 = n1 + | q1 > q2 = n1 + | otherwise = compls n2 merge_quals qmax v1 v2 v3 v4 = V.zipWith4 zz v1 v2 v3 v4 where - zz !n1 (Q !q1) !n2 (Q !q2) = Q $ if n1 == n2 then min qmax (q1 + q2) - else if q1 > q2 then q1 - q2 - else q2 - q1 + zz !n1 (Q !q1) !n2 (Q !q2) | n1 == compls n2 = Q $ min qmax (q1 + q2) + | q1 > q2 = Q $ q1 - q2 + | otherwise = Q $ q2 - q1 -- | Trimming for a single read: we need one adapter only (the one coming -- /after/ the read), here provided as a list of options, and then we diff --git a/src/Bio/Bam/Writer.hs b/src/Bio/Bam/Writer.hs --- a/src/Bio/Bam/Writer.hs +++ b/src/Bio/Bam/Writer.hs @@ -15,8 +15,9 @@ import Bio.Iteratee.Builder import Bio.Prelude +import Data.ByteString.Builder ( hPutBuilder, Builder, toLazyByteString ) import Data.ByteString.Internal ( ByteString(..) ) -import Data.ByteString.Builder ( hPutBuilder ) +import Data.ByteString.Lazy ( foldrChunks ) import Foreign.Marshal.Alloc ( alloca ) import Foreign.Storable ( pokeByteOff, peek ) import System.IO ( openBinaryFile, IOMode(..) ) @@ -74,7 +75,7 @@ sarr v = conjoin ',' . map shows $ U.toList v class IsBamRec a where - pushBam :: a -> Push + pushBam :: a -> BgzfTokens -> BgzfTokens instance IsBamRec BamRaw where {-# INLINE pushBam #-} @@ -91,29 +92,31 @@ -- | Encodes BAM records straight into a dynamic buffer, the BGZF's it. -- Should be fairly direct and perform well. {-# INLINE encodeBamWith #-} -encodeBamWith :: (MonadIO m, IsBamRec r) => Int -> BamMeta -> Enumeratee [r] B.ByteString m a -encodeBamWith lv meta = joinI . eneeBam . encodeBgzfWith lv +encodeBamWith :: (MonadIO m, IsBamRec r) => Int -> BamMeta -> Enumeratee [r] S.ByteString m () +encodeBamWith lv meta = eneeBam ><> encodeBgzf lv where - eneeBam = eneeCheckIfDone (\k -> mapChunks (foldMap pushBam) . k $ Chunk pushHeader) + eneeBam = eneeCheckIfDone (\k -> mapChunks (foldMap (Endo . pushBam)) . k $ Chunk pushHeader) - pushHeader = pushByteString "BAM\1" - <> setMark -- the length byte - <> pushBuilder (showBamMeta meta) - <> endRecord -- fills the length in - <> pushWord32 (fromIntegral . Z.length $ meta_refs meta) - <> foldMap pushRef (meta_refs meta) + pushHeader :: Endo BgzfTokens + pushHeader = Endo $ TkString "BAM\1" + . TkSetMark -- the length byte + . pushBuilder (showBamMeta meta) + . TkEndRecord -- fills the length in + . TkWord32 (fromIntegral . Z.length $ meta_refs meta) + . appEndo (foldMap (Endo . pushRef) (meta_refs meta)) - pushRef bs = ensureBuffer (fromIntegral $ B.length (sq_name bs) + 9) - <> unsafePushWord32 (fromIntegral $ B.length (sq_name bs) + 1) - <> unsafePushByteString (sq_name bs) - <> unsafePushByte 0 - <> unsafePushWord32 (fromIntegral $ sq_length bs) + pushRef :: BamSQ -> BgzfTokens -> BgzfTokens + pushRef bs = TkWord32 (fromIntegral $ B.length (sq_name bs) + 1) + . TkString (sq_name bs) + . TkWord8 0 + . TkWord32 (fromIntegral $ sq_length bs) + pushBuilder :: Builder -> BgzfTokens -> BgzfTokens + pushBuilder b tk = foldrChunks TkString tk (toLazyByteString b) + {-# INLINE pushBamRaw #-} -pushBamRaw :: BamRaw -> Push -pushBamRaw br = ensureBuffer (B.length (raw_data br) + 4) - <> unsafePushWord32 (fromIntegral $ B.length (raw_data br)) - <> unsafePushByteString (raw_data br) +pushBamRaw :: BamRaw -> BgzfTokens -> BgzfTokens +pushBamRaw = TkLnString . raw_data -- | writes BAM encoded stuff to a file -- XXX This should(!) write indexes on the side---a simple block index @@ -146,75 +149,67 @@ #-} {-# INLINE[1] pushBamRec #-} -pushBamRec :: BamRec -> Push -pushBamRec BamRec{..} = mconcat - [ ensureBuffer minlength - , unsafeSetMark - , unsafePushWord32 $ unRefseq b_rname - , unsafePushWord32 $ fromIntegral b_pos - , unsafePushByte $ fromIntegral $ B.length b_qname + 1 - , unsafePushByte $ unQ b_mapq - , unsafePushWord16 $ fromIntegral bin - , unsafePushWord16 $ fromIntegral $ VS.length b_cigar - , unsafePushWord16 $ fromIntegral b_flag - , unsafePushWord32 $ fromIntegral $ V.length b_seq - , unsafePushWord32 $ unRefseq b_mrnm - , unsafePushWord32 $ fromIntegral b_mpos - , unsafePushWord32 $ fromIntegral b_isize - , unsafePushByteString b_qname - , unsafePushByte 0 - , VS.foldr ((<>) . unsafePushByte) mempty (VS.unsafeCast b_cigar :: VS.Vector Word8) - , pushSeq b_seq - , VS.foldr ((<>) . unsafePushByte . unQ) mempty b_qual - , foldMap pushExt b_exts - , endRecord ] +pushBamRec :: BamRec -> BgzfTokens -> BgzfTokens +pushBamRec BamRec{..} = + TkSetMark + . TkWord32 (unRefseq b_rname) + . TkWord32 (fromIntegral b_pos) + . TkWord8 (fromIntegral $ B.length b_qname + 1) + . TkWord8 (unQ b_mapq) + . TkWord16 (fromIntegral bin) + . TkWord16 (fromIntegral $ VS.length b_cigar) + . TkWord16 (fromIntegral b_flag) + . TkWord32 (fromIntegral $ V.length b_seq) + . TkWord32 (unRefseq b_mrnm) + . TkWord32 (fromIntegral b_mpos) + . TkWord32 (fromIntegral b_isize) + . TkString b_qname + . TkWord8 0 + . VS.foldr ((.) . TkWord8) id (VS.unsafeCast b_cigar :: VS.Vector Word8) + . pushSeq b_seq + . VS.foldr ((.) . TkWord8 . unQ) id b_qual + . foldr ((.) . pushExt) id b_exts + . TkEndRecord where bin = distinctBin b_pos (alignedLength b_cigar) - minlength = 37 + B.length b_qname + 4 * V.length b_cigar + V.length b_qual + (V.length b_seq + 1) `shiftR` 1 - pushSeq :: V.Vector vec Nucleotides => vec Nucleotides -> Push + pushSeq :: V.Vector vec Nucleotides => vec Nucleotides -> BgzfTokens -> BgzfTokens pushSeq v = case v V.!? 0 of - Nothing -> mempty + Nothing -> id Just a -> case v V.!? 1 of - Nothing -> unsafePushByte (unNs a `shiftL` 4) - Just b -> unsafePushByte (unNs a `shiftL` 4 .|. unNs b) - <> pushSeq (V.drop 2 v) + Nothing -> TkWord8 (unNs a `shiftL` 4) + Just b -> TkWord8 (unNs a `shiftL` 4 .|. unNs b) . pushSeq (V.drop 2 v) - pushExt :: (BamKey, Ext) -> Push + pushExt :: (BamKey, Ext) -> BgzfTokens -> BgzfTokens pushExt (BamKey k, e) = case e of - Text t -> common (4 + B.length t) 'Z' $ - unsafePushByteString t <> unsafePushByte 0 - - Bin t -> common (4 + B.length t) 'H' $ - unsafePushByteString t <> unsafePushByte 0 - - Char c -> common 4 'A' $ unsafePushByte c - - Float f -> common 7 'f' $ unsafePushWord32 (fromIntegral $ fromFloat f) + Text t -> common 'Z' . TkString t . TkWord8 0 + Bin t -> common 'H' . TkString t . TkWord8 0 + Char c -> common 'A' . TkWord8 c + Float f -> common 'f' . TkWord32 (fromIntegral $ fromFloat f) Int i -> case put_some_int (U.singleton i) of - (c,op) -> common 7 c (op i) + (c,op) -> common c . op i IntArr ia -> case put_some_int ia of - (c,op) -> common (4 * U.length ia) 'B' $ unsafePushByte (fromIntegral $ ord c) - <> unsafePushWord32 (fromIntegral $ U.length ia-1) - <> U.foldr ((<>) . op) mempty ia + (c,op) -> common 'B' . TkWord8 (fromIntegral $ ord c) + . TkWord32 (fromIntegral $ U.length ia-1) + . U.foldr ((.) . op) id ia - FloatArr fa -> common (4 * U.length fa) 'B' $ unsafePushByte (fromIntegral $ ord 'f') - <> unsafePushWord32 (fromIntegral $ U.length fa-1) - <> U.foldr ((<>) . unsafePushWord32 . fromFloat) mempty fa + FloatArr fa -> common 'B' . TkWord8 (fromIntegral $ ord 'f') + . TkWord32 (fromIntegral $ U.length fa-1) + . U.foldr ((.) . TkWord32 . fromFloat) id fa where - common l z b = ensureBuffer l <> unsafePushWord16 k - <> unsafePushByte (fromIntegral $ ord z) <> b + common :: Char -> BgzfTokens -> BgzfTokens + common z = TkWord16 k . TkWord8 (fromIntegral $ ord z) - put_some_int :: U.Vector Int -> (Char, Int -> Push) + put_some_int :: U.Vector Int -> (Char, Int -> BgzfTokens -> BgzfTokens) put_some_int is - | U.all (between 0 0xff) is = ('C', unsafePushByte . fromIntegral) - | U.all (between (-0x80) 0x7f) is = ('c', unsafePushByte . fromIntegral) - | U.all (between 0 0xffff) is = ('S', unsafePushWord16 . fromIntegral) - | U.all (between (-0x8000) 0x7fff) is = ('s', unsafePushWord16 . fromIntegral) - | U.all (> 0) is = ('I', unsafePushWord32 . fromIntegral) - | otherwise = ('i', unsafePushWord32 . fromIntegral) + | U.all (between 0 0xff) is = ('C', TkWord8 . fromIntegral) + | U.all (between (-0x80) 0x7f) is = ('c', TkWord8 . fromIntegral) + | U.all (between 0 0xffff) is = ('S', TkWord16 . fromIntegral) + | U.all (between (-0x8000) 0x7fff) is = ('s', TkWord16 . fromIntegral) + | U.all (> 0) is = ('I', TkWord32 . fromIntegral) + | otherwise = ('i', TkWord32 . fromIntegral) between :: Int -> Int -> Int -> Bool between l r x = l <= x && x <= r @@ -224,7 +219,12 @@ pokeByteOff buf 0 float >> peek buf packBam :: BamRec -> IO BamRaw -packBam br = do bb' <- case pushBamRec br of Push p -> newBuffer 1000 >>= p - return $ bamRaw 0 (PS (buffer bb') 4 (len bb' - 4)) - +packBam br = do bb <- newBuffer 1000 + (bb', TkEnd) <- store_loop bb (pushBamRec br TkEnd) + return . bamRaw 0 $ PS (buffer bb') 4 (used bb' - 4) + where + store_loop bb tk = do (bb',tk') <- fillBuffer bb tk + case tk' of TkEnd -> return (bb',tk') + _ -> do bb'' <- expandBuffer 0 bb' + store_loop bb'' tk' diff --git a/src/Bio/Iteratee.hs b/src/Bio/Iteratee.hs --- a/src/Bio/Iteratee.hs +++ b/src/Bio/Iteratee.hs @@ -2,55 +2,27 @@ -- with "Prelude" plus a handful of utilities. module Bio.Iteratee ( - groupStreamBy, - groupStreamOn, iGetString, iterGet, iterLoop, iLookAhead, - headStream, - peekStream, - takeStream, - dropStream, - mapChunks, - mapChunksM, - mapStream, - rigidMapStream, - mapStreamM, - mapStreamM_, - filterStream, - filterStreamM, - foldStream, - foldStreamM, - zipStreams, - zipStreams3, + protectTerm, - concatMapStream, - concatMapStreamM, - mapMaybeStream, parMapChunksIO, + parRunIO, progressGen, progressNum, progressPos, - I.takeWhileE, - I.tryHead, - I.isFinished, - I.heads, - I.breakE, - ($==), - mBind, mBind_, ioBind, ioBind_, ListLike, MonadIO, MonadMask, lift, liftIO, - (>=>), (<=<), stdin, stdout, stderr, enumAuxFile, enumInputs, enumDefaultInputs, - defaultBufSize, Ordering'(..), mergeSortStreams, @@ -73,121 +45,38 @@ Fd, withFileFd, - module Data.Iteratee.Binary, - module Data.Iteratee.Char, - module Data.Iteratee.IO, - module Data.Iteratee.Iteratee + + module Bio.Iteratee.Binary, + module Bio.Iteratee.Char, + module Bio.Iteratee.IO, + module Bio.Iteratee.Iteratee, + module Bio.Iteratee.ListLike ) where import Bio.Bam.Header -import Bio.Util.Numeric ( showNum ) +import Bio.Iteratee.Base +import Bio.Iteratee.Binary +import Bio.Iteratee.Char +import Bio.Iteratee.IO +import Bio.Iteratee.Iteratee +import Bio.Iteratee.ListLike import Bio.Prelude +import Bio.Util.Numeric ( showNum ) import Control.Concurrent.Async ( Async, async, wait, cancel ) import Control.Monad.Catch ( MonadMask(..) ) import Control.Monad.IO.Class import Control.Monad.Trans.Class import Data.Binary.Get -import Data.Iteratee.Binary -import Data.Iteratee.Char -import Data.Iteratee.IO hiding ( defaultBufSize ) -import Data.Iteratee.Iteratee hiding ( identity, empty, mapChunks, mapChunksM, (>>>) ) import Data.ListLike ( ListLike ) import System.IO ( hIsTerminalDevice ) import qualified Control.Monad.Catch as CMC import qualified Data.Attoparsec.ByteString as A import qualified Data.ByteString.Char8 as S -import qualified Data.Iteratee as I import qualified Data.ListLike as LL -import qualified Data.NullPoint as N import qualified Data.Vector.Generic as VG import qualified Data.Vector.Generic.Mutable as VM --- | Grouping on 'Iteratee's. @groupStreamOn proj inner outer@ executes --- @inner (proj e)@, where @e@ is the first input element, to obtain an --- 'Iteratee' @i@, then passes elements @e@ to @i@ as long as @proj e@ --- produces the same result. If @proj e@ changes or the input ends, the --- pair of @proj e@ and the result of @run i@ is passed to @outer@. At --- end of input, the resulting @outer@ is returned. -groupStreamOn :: (Monad m, LL.ListLike l e, Eq t1, Nullable l) - => (e -> t1) - -> (t1 -> m (Iteratee l m t2)) - -> Enumeratee l [(t1, t2)] m a -groupStreamOn proj inner = eneeCheckIfDonePass (icont . step) - where - step outer (EOF mx) = idone (liftI outer) $ EOF mx - step outer c@(Chunk as) - | LL.null as = liftI $ step outer - | otherwise = let x = proj (LL.head as) - in lift (inner x) >>= \i -> step' x i outer c - - -- We want to feed a 'Chunk' to the inner 'Iteratee', which might be - -- finished. In that case, we would want to abort, but we cannot, - -- since the outer iteration is still going on. So instead we - -- discard data we would have fed to the inner 'Iteratee'. (Use of - -- 'enumPure1Chunk' is not appropriate, it would accumulate the - -- data, just to have it discarded by the 'run' that eventually - -- happens. - - step' c it outer (Chunk as) - | LL.null as = liftI $ step' c it outer - | (l,r) <- LL.span ((==) c . proj) as, not (LL.null l) = - let od a _str = idoneM a $ EOF Nothing - oc k Nothing = return $ k (Chunk l) - oc k m = icontM k m - in lift (runIter it od oc) >>= \it' -> step' c it' outer (Chunk r) - - step' c it outer str = - lift (run it) >>= \b -> eneeCheckIfDone (`step` str) . outer $ Chunk [(c,b)] - - --- | Grouping on 'Iteratee's. @groupStreamBy cmp inner outer@ executes --- @inner@ to obtain an 'Iteratee' @i@, then passes elements @e@ to @i@ --- as long as @cmp e0 e@, where @e0@ is some preceeding element, is --- true. Else, the result of @run i@ is passed to @outer@ and --- 'groupStreamBy' restarts. At end of input, the resulting @outer@ is --- returned. -groupStreamBy :: (Monad m, LL.ListLike l t, Nullable l) - => (t -> t -> Bool) - -> m (Iteratee l m t2) - -> Enumeratee l [t2] m a -groupStreamBy cmp inner = eneeCheckIfDonePass (icont . step) - where - step outer (EOF mx) = idone (liftI outer) $ EOF mx - step outer c@(Chunk as) - | LL.null as = liftI $ step outer - | otherwise = lift inner >>= \i -> step' (LL.head as) i outer c - - step' c it outer (Chunk as) - | LL.null as = liftI $ step' c it outer - | (l,r) <- LL.span (cmp c) as, not (LL.null l) = - let od a _str = idoneM a $ EOF Nothing - oc k Nothing = return $ k (Chunk l) - oc k m = icontM k m - in lift (runIter it od oc) >>= \it' -> step' (LL.head l) it' outer (Chunk r) - - step' _ it outer str = - lift (run it) >>= \b -> eneeCheckIfDone (`step` str) . outer $ Chunk [b] - - --- | Take a prefix of a stream, the equivalent of 'Data.List.take'. -{-# INLINE takeStream #-} -takeStream :: (Monad m, Nullable s, ListLike s el) => Int -> Enumeratee s s m a -takeStream = I.take - --- | Take first element of a stream or fail. -{-# INLINE headStream #-} -headStream :: ListLike s el => Iteratee s m el -headStream = I.head - -{-# INLINE peekStream #-} -peekStream :: ListLike s el => Iteratee s m (Maybe el) -peekStream = I.peek - -{-# INLINE dropStream #-} -dropStream :: (Nullable s, ListLike s el) => Int -> Iteratee s m () -dropStream = I.drop - -- | Run an Iteratee, collect the input. When it finishes, return the -- result along with *all* input. Effectively allows lookahead. Be -- careful, this will eat memory if the @Iteratee@ doesn't return @@ -215,7 +104,7 @@ -- | Repeatedly apply an 'Iteratee' to a value until end of stream. -- Returns the final value. iterLoop :: (Nullable s, Monad m) => (a -> Iteratee s m a) -> a -> Iteratee s m a -iterLoop it a = do e <- I.isFinished +iterLoop it a = do e <- isFinished if e then return a else it a >>= iterLoop it @@ -234,38 +123,7 @@ Done rest _ a | S.null rest -> idone a (EOF mx) | otherwise -> idone a (Chunk rest) -{-# INLINE mBind #-} --- | Lifts a monadic action and combines it with a continuation. --- @mBind m f@ is the same as @lift m >>= f@, but does not require a --- 'Nullable' constraint on the stream type. -infixl 1 `mBind` -mBind :: Monad m => m a -> (a -> Iteratee s m b) -> Iteratee s m b -mBind m f = Iteratee $ \onDone onCont -> m >>= \a -> runIter (f a) onDone onCont -{-# INLINE mBind_ #-} --- | Lifts a monadic action, ignored the result and combines it with a --- continuation. @mBind_ m f@ is the same as @lift m >>= f@, but does --- not require a 'Nullable' constraint on the stream type. -infixl 1 `mBind_` -mBind_ :: Monad m => m a -> Iteratee s m b -> Iteratee s m b -mBind_ m b = Iteratee $ \onDone onCont -> m >> runIter b onDone onCont - -{-# INLINE ioBind #-} --- | Lifts an IO action and combines it with a continuation. --- @ioBind m f@ is the same as @liftIO m >>= f@, but does not require a --- 'Nullable' constraint on the stream type. -infixl 1 `ioBind` -ioBind :: MonadIO m => IO a -> (a -> Iteratee s m b) -> Iteratee s m b -ioBind m f = Iteratee $ \onDone onCont -> liftIO m >>= \a -> runIter (f a) onDone onCont - -{-# INLINE ioBind_ #-} --- | Lifts an IO action, ignores its result, and combines it with a --- continuation. @ioBind_ m f@ is the same as @liftIO m >> f@, but does --- not require a 'Nullable' constraint on the stream type. -infixl 1 `ioBind_` -ioBind_ :: MonadIO m => IO a -> Iteratee s m b -> Iteratee s m b -ioBind_ m b = Iteratee $ \onDone onCont -> liftIO m >> runIter b onDone onCont - infixl 1 $== {-# INLINE ($==) #-} -- | Compose an 'Enumerator\'' with an 'Enumeratee', giving a new @@ -290,114 +148,6 @@ -> Enumerator' hi s1 m a mergeEnums' e1 e2 etee i = e1 $ \hi -> e2 (\ho -> joinI . etee ho $ ilift lift (i hi)) >>= run --- | Apply a function to the elements of a stream, concatenate the --- results into a stream. No giant intermediate list is produced. -{-# INLINE concatMapStream #-} -concatMapStream :: (Monad m, ListLike s a, NullPoint s) => (a -> t) -> Enumeratee s t m r -concatMapStream f = eneeCheckIfDone (liftI . go) - where - go k (EOF mx) = idone (liftI k) (EOF mx) - go k (Chunk xs) | LL.null xs = liftI (go k) - | otherwise = eneeCheckIfDone (flip go (Chunk (LL.tail xs))) . k . Chunk . f $ LL.head xs - --- | Apply a monadic function to the elements of a stream, concatenate --- the results into a stream. No giant intermediate list is produced. -{-# INLINE concatMapStreamM #-} -concatMapStreamM :: (Monad m, ListLike s a, NullPoint s) => (a -> m t) -> Enumeratee s t m r -concatMapStreamM f = eneeCheckIfDone (liftI . go) - where - go k (EOF mx) = idone (liftI k) (EOF mx) - go k (Chunk xs) | LL.null xs = liftI (go k) - | otherwise = f (LL.head xs) `mBind` - eneeCheckIfDone (flip go (Chunk (LL.tail xs))) . k . Chunk - -{-# INLINE mapMaybeStream #-} -mapMaybeStream :: (ListLike s a, NullPoint s, ListLike t b) => (a -> Maybe b) -> Enumeratee s t m r -mapMaybeStream f = mapChunks mm - where - mm l = if LL.null l then LL.empty else - case f (LL.head l) of Nothing -> mm (LL.tail l) - Just b -> LL.cons b $ mm (LL.tail l) - --- | Apply a filter predicate to an 'Iteratee'. -{-# INLINE filterStream #-} -filterStream :: (ListLike s a, NullPoint s) => (a -> Bool) -> Enumeratee s s m r -filterStream = mapChunks . LL.filter - --- | Apply a monadic filter predicate to an 'Iteratee'. -{-# INLINE filterStreamM #-} -filterStreamM :: (Monad m, ListLike s a, Nullable s) => (a -> m Bool) -> Enumeratee s s m r -filterStreamM k = mapChunksM (go id) - where - go acc s | LL.null s = return $! acc LL.empty - | otherwise = do p <- k (LL.head s) - let acc' = if p then LL.cons (LL.head s) . acc else acc - go acc' (LL.tail s) - -{-# INLINE mapChunks #-} -mapChunks :: NullPoint s => (s -> s') -> Enumeratee s s' m a -mapChunks f = eneeCheckIfDonePass (icont . step) - where - step k (Chunk xs) = eneeCheckIfDonePass (icont . step) . k . Chunk $ f xs - step k str = idone (liftI k) str - -{-# INLINE mapChunksM #-} -mapChunksM :: (Monad m, NullPoint s) => (s -> m s') -> Enumeratee s s' m a -mapChunksM f = eneeCheckIfDonePass (icont . step) - where - step k (Chunk xs) = f xs `mBind` eneeCheckIfDonePass (icont . step) . k . Chunk - step k str = idone (liftI k) str - --- | Map a function over an 'Iteratee'. --- This one is reimplemented and differs from the the one in --- "Data.Iteratee.ListLike" in so far that it doesn't pass on an 'EOF' --- received in the input, which is the expected behavior. -{-# INLINE mapStream #-} -mapStream :: (ListLike (s el) el, ListLike (s el') el', NullPoint (s el)) - => (el -> el') -> Enumeratee (s el) (s el') m a -mapStream = mapChunks . LL.map - --- | Map a function over an 'Iteratee' rigidly. --- This one is reimplemented and differs from the the one in --- "Data.Iteratee.ListLike" in so far that it doesn't pass on an 'EOF' --- received in the input, which is the expected behavior. -{-# INLINE rigidMapStream #-} -rigidMapStream :: (ListLike s el, NullPoint s) => (el -> el) -> Enumeratee s s m a -rigidMapStream = mapChunks . LL.rigidMap - --- | Map a monadic function over an 'Iteratee'. -{-# INLINE mapStreamM #-} -mapStreamM :: (Monad m, ListLike (s el) el, ListLike (s el') el', NullPoint (s el)) - => (el -> m el') -> Enumeratee (s el) (s el') m a -mapStreamM = mapChunksM . LL.mapM - --- | Map a monadic function over an 'Iteratee', discarding the results. -{-# INLINE mapStreamM_ #-} -mapStreamM_ :: (Monad m, Nullable s, ListLike s el) => (el -> m b) -> Iteratee s m () -mapStreamM_ = mapChunksM_ . LL.mapM_ - --- | Fold a monadic function over an 'Iteratee'. -{-# INLINE foldStreamM #-} -foldStreamM :: (Monad m, Nullable s, ListLike s a) => (b -> a -> m b) -> b -> Iteratee s m b -foldStreamM k = foldChunksM go - where - go b s | LL.null s = return b - | otherwise = k b (LL.head s) >>= \b' -> go b' (LL.tail s) - --- | Fold a function over an 'Iteratee'. -foldStream :: (Monad m, Nullable s, ListLike s a) => (b -> a -> b) -> b -> Iteratee s m b -foldStream f = foldChunksM (\b s -> return $! LL.foldl' f b s) - --- | Apply two 'Iteratee's to the same stream. -zipStreams :: (Nullable s, ListLike s el, Monad m) - => Iteratee s m a -> Iteratee s m b -> Iteratee s m (a, b) -zipStreams = I.zip - --- | Apply three 'Iteratee's to the same stream. -zipStreams3 :: (Nullable s, ListLike s el, Monad m) - => Iteratee s m a -> Iteratee s m b -> Iteratee s m c -> Iteratee s m (a, b, c) -zipStreams3 = I.zip3 - type Enumerator' h eo m b = (h -> Iteratee eo m b) -> m (Iteratee eo m b) type Enumeratee' h ei eo m b = (h -> Iteratee eo m b) -> Iteratee ei m (Iteratee eo m b) @@ -414,14 +164,6 @@ go ( f :fs) = enumFile defaultBufSize f >=> go fs go [ ] = return --- | Default buffer size in elements. This is 1024 in "Data.Iteratee", --- which is obviously too small. Since we want to merge many files, a --- read should take more time than a seek. This sets the sensible --- buffer size to more than about one MB. -defaultBufSize :: Int -defaultBufSize = 2*1024*1024 - - data Ordering' a = Less | Equal a | NotLess mergeSortStreams :: (Monad m, ListLike s a, Nullable s) => (a -> a -> Ordering' a) -> Enumeratee s s (Iteratee s m) b @@ -429,12 +171,12 @@ where step out = peekStream >>= \mx -> lift peekStream >>= \my -> case (mx, my) of (Just x, Just y) -> case x `comp` y of - Less -> do I.drop 1 ; eneeCheckIfDone step . out . Chunk $ LL.singleton x - NotLess -> do lift (I.drop 1) ; eneeCheckIfDone step . out . Chunk $ LL.singleton y - Equal z -> do I.drop 1 ; lift (I.drop 1) ; eneeCheckIfDone step . out . Chunk $ LL.singleton z + Less -> do dropStream 1 ; eneeCheckIfDone step . out . Chunk $ LL.singleton x + NotLess -> do lift (dropStream 1) ; eneeCheckIfDone step . out . Chunk $ LL.singleton y + Equal z -> do dropStream 1 ; lift (dropStream 1) ; eneeCheckIfDone step . out . Chunk $ LL.singleton z - (Just x, Nothing) -> do I.drop 1 ; eneeCheckIfDone step . out . Chunk $ LL.singleton x - (Nothing, Just y) -> do lift (I.drop 1) ; eneeCheckIfDone step . out . Chunk $ LL.singleton y + (Just x, Nothing) -> do dropStream 1 ; eneeCheckIfDone step . out . Chunk $ LL.singleton x + (Nothing, Just y) -> do lift (dropStream 1) ; eneeCheckIfDone step . out . Chunk $ LL.singleton y (Nothing, Nothing) -> idone (liftI out) $ EOF Nothing @@ -454,7 +196,7 @@ _ -> liftI $ go' qq k -- we have room for input - go' !qq k (EOF mx) = do a <- liftIO (async (f N.empty)) + go' !qq k (EOF mx) = do a <- liftIO (async (f emptyP)) goE mx (pushQ a qq) k Nothing go' !qq k (Chunk c) = do a <- liftIO (async (f c)) go (pushQ a qq) k Nothing @@ -465,6 +207,26 @@ Nothing -> idone (liftI k) (EOF mx) Just (a,qq') -> liftIO (wait a) >>= eneeCheckIfDonePass (goE mx qq') . k . Chunk +parRunIO :: MonadIO m => Int -> Enumeratee [IO a] a m b +parRunIO np = eneeCheckIfDonePass (go emptyQ) + where + -- check if the queue is full + go !qq k (Just e) = cancelAll qq >> icont (go' emptyQ k) (Just e) + go !qq k Nothing = case popQ qq of + Just (a,qq') | lengthQ qq == np -> liftIO (wait a) >>= eneeCheckIfDonePass (go qq') . k . Chunk + _ -> liftI $ go' qq k + + -- we have room for input + go' !qq k (Chunk (c:cs)) = liftIO (async c) >>= \a -> go' (pushQ a qq) k (Chunk cs) + go' !qq k (Chunk [ ]) = go qq k Nothing + go' !qq k (EOF mx) = goE mx qq k Nothing + + -- input ended, empty the queue + goE _ !qq k (Just e) = cancelAll qq >> icont (go' emptyQ k) (Just e) + goE mx !qq k Nothing = case popQ qq of + Nothing -> idone (liftI k) (EOF mx) + Just (a,qq') -> liftIO (wait a) >>= eneeCheckIfDonePass (goE mx qq') . k . Chunk + -- | Protects the terminal from binary junk. If @i@ is an 'Iteratee' -- that might write binary to 'stdout', then @protectTerm i@ is the same -- 'Iteratee', but it will abort if 'stdout' is a terminal device. @@ -564,7 +326,7 @@ go mv i | i == n = return n | otherwise = - I.tryHead >>= \x -> case x of + tryHead >>= \x -> case x of Nothing -> return i Just a -> liftIO (VM.write mv i a) >> go mv (i+1) @@ -572,7 +334,7 @@ stream2vector :: (MonadIO m, ListLike s a, Nullable s, VG.Vector v a) => Iteratee s m (v a) stream2vector = liftIO (VM.new 1024) >>= go 0 where - go !i !mv = I.tryHead >>= \x -> case x of + go !i !mv = tryHead >>= \x -> case x of Nothing -> liftIO $ VG.unsafeFreeze $ VM.take i mv Just a -> do mv' <- if VM.length mv == i then liftIO (VM.grow mv (VM.length mv)) else return mv when (i `rem` 0x10000 == 0) $ liftIO performGC diff --git a/src/Bio/Iteratee/Base.hs b/src/Bio/Iteratee/Base.hs new file mode 100644 --- /dev/null +++ b/src/Bio/Iteratee/Base.hs @@ -0,0 +1,282 @@ +{-# LANGUAGE TypeFamilies,UndecidableInstances,Rank2Types,ExistentialQuantification #-} + +-- |Monadic Iteratees: +-- incremental input parsers, processors and transformers + +module Bio.Iteratee.Base ( + -- * Types + Stream (..) + ,StreamStatus (..) + -- ** Exception types + ,module Bio.Iteratee.Exception + -- ** Iteratees + ,Iteratee (..) + -- * Functions + -- ** Control functions + ,run + ,tryRun + ,mapIteratee + ,ilift + ,ifold + -- ** Creating Iteratees + ,idone + ,icont + ,liftI + ,idoneM + ,icontM + -- ** Stream Functions + ,setEOF + -- * Classes + ,NullPoint(..) + ,Nullable(..) +) +where + +import Bio.Iteratee.Exception +import Bio.Prelude + +import Control.Monad.Base +import Control.Monad.Catch as CIO +import Control.Monad.IO.Class +import Control.Monad.Trans.Class +import Control.Monad.Trans.Control + +import qualified Control.Exception as E +import qualified Data.ByteString as B +import qualified Data.ByteString.Lazy as L + +-- | NullPoint class. Containers that have a null representation, +-- corresponding to Data.Monoid.mempty. +class NullPoint c where emptyP :: c + +instance NullPoint (Endo a) where emptyP = Endo id +instance NullPoint [a] where emptyP = [] +instance NullPoint B.ByteString where emptyP = B.empty +instance NullPoint L.ByteString where emptyP = L.empty + +-- | Nullable container class +class NullPoint c => Nullable c where nullC :: c -> Bool + +instance Nullable [a] where nullC [] = True ; nullC _ = False +instance Nullable B.ByteString where nullC = B.null +instance Nullable L.ByteString where nullC = L.null + +-- |A stream is a (continuing) sequence of elements bundled in Chunks. +-- The first variant indicates termination of the stream. +-- Chunk a gives the currently available part of the stream. +-- The stream is not terminated yet. +-- The case (null Chunk) signifies a stream with no currently available +-- data but which is still continuing. A stream processor should, +-- informally speaking, ``suspend itself'' and wait for more data +-- to arrive. + +data Stream c = EOF (Maybe SomeException) | Chunk c + deriving (Show, Typeable) + +instance (Eq c) => Eq (Stream c) where + (Chunk c1) == (Chunk c2) = c1 == c2 + (EOF Nothing) == (EOF Nothing) = True + (EOF (Just e1)) == (EOF (Just e2)) = typeOf e1 == typeOf e2 + _ == _ = False + +instance Monoid c => Monoid (Stream c) where + mempty = Chunk mempty + mappend (EOF mErr) _ = EOF mErr + mappend _ (EOF mErr) = EOF mErr + mappend (Chunk s1) (Chunk s2) = Chunk (s1 `mappend` s2) + +-- |Map a function over a stream. +instance Functor Stream where + fmap f (Chunk xs) = Chunk $ f xs + fmap _ (EOF mErr) = EOF mErr + +-- |Describe the status of a stream of data. +data StreamStatus = + DataRemaining + | EofNoError + | EofError SomeException + deriving (Show, Typeable) + +-- ---------------------------------------------- +-- create exception type hierarchy + +-- |Produce the 'EOF' error message. If the stream was terminated because +-- of an error, keep the error message. +setEOF :: Stream c -> SomeException +setEOF (EOF (Just e)) = e +setEOF _ = toException EofException + +-- ---------------------------------------------- +-- | Monadic iteratee +newtype Iteratee s m a = Iteratee{ runIter :: forall r. + (a -> Stream s -> m r) -> + ((Stream s -> Iteratee s m a) -> Maybe SomeException -> m r) -> + m r} + +-- ---------------------------------------------- + +idone :: a -> Stream s -> Iteratee s m a +idone a s = Iteratee $ \onDone _ -> onDone a s + +icont :: (Stream s -> Iteratee s m a) -> Maybe SomeException -> Iteratee s m a +icont k e = Iteratee $ \_ onCont -> onCont k e + +liftI :: (Stream s -> Iteratee s m a) -> Iteratee s m a +liftI k = Iteratee $ \_ onCont -> onCont k Nothing + +-- Monadic versions, frequently used by enumerators +idoneM :: Monad m => a -> Stream s -> m (Iteratee s m a) +idoneM x str = return $ Iteratee $ \onDone _ -> onDone x str + +icontM + :: Monad m => + (Stream s -> Iteratee s m a) + -> Maybe SomeException + -> m (Iteratee s m a) +icontM k e = return $ Iteratee $ \_ onCont -> onCont k e + +instance (Functor m) => Functor (Iteratee s m) where + fmap f m = Iteratee $ \onDone onCont -> + let od = onDone . f + oc = onCont . (fmap f .) + in runIter m od oc + +instance (Functor m, Monad m, Nullable s) => Applicative (Iteratee s m) where + pure x = idone x (Chunk emptyP) + {-# INLINE (<*>) #-} + m <*> a = m >>= flip fmap a + +instance (Monad m, Nullable s) => Monad (Iteratee s m) where + {-# INLINE return #-} + return x = Iteratee $ \onDone _ -> onDone x (Chunk emptyP) + {-# INLINE (>>=) #-} + (>>=) = bindIteratee + +{-# INLINE bindIteratee #-} +bindIteratee :: Nullable s + => Iteratee s m a + -> (a -> Iteratee s m b) + -> Iteratee s m b +bindIteratee = self + where + self m f = Iteratee $ \onDone onCont -> + let m_done a (Chunk s) + | nullC s = runIter (f a) onDone onCont + m_done a stream = runIter (f a) (const . flip onDone stream) f_cont + where f_cont k Nothing = runIter (k stream) onDone onCont + f_cont k e = onCont k e + in runIter m m_done (onCont . (flip self f .)) + +instance NullPoint s => MonadTrans (Iteratee s) where + lift m = Iteratee $ \onDone _ -> m >>= flip onDone (Chunk emptyP) + +instance (MonadBase b m, Nullable s, NullPoint s) => MonadBase b (Iteratee s m) where + liftBase = lift . liftBase + +instance (MonadIO m, Nullable s, NullPoint s) => MonadIO (Iteratee s m) where + liftIO = lift . liftIO + +instance (MonadThrow m, Nullable s, NullPoint s) => + MonadThrow (Iteratee s m) where + throwM e = lift $ CIO.throwM e + +instance (MonadCatch m, Nullable s, NullPoint s) => + MonadCatch (Iteratee s m) where + m `catch` f = Iteratee $ \od oc -> runIter m od oc `CIO.catch` (\e -> runIter (f e) od oc) + +-- prior to exceptions-0.6, these were part of MonadCatch +instance (MonadMask m, Nullable s, NullPoint s) => MonadMask (Iteratee s m) where + mask q = Iteratee $ \od oc -> CIO.mask $ \u -> runIter (q $ ilift u) od oc + uninterruptibleMask q = Iteratee $ \od oc -> CIO.uninterruptibleMask $ \u -> runIter (q $ ilift u) od oc + + +instance forall s. (NullPoint s, Nullable s) => MonadTransControl (Iteratee s) where + type StT (Iteratee s) x = Either (x, Stream s) (Maybe SomeException) + + liftWith f = lift $ f $ \t -> + (runIter t (\x s -> return $ Left (x,s)) + (\_ e -> return $ Right e) ) + restoreT = join . lift . liftM + (either (uncurry idone) + (te . fromMaybe (iterStrExc + "iteratee: error in MonadTransControl instance"))) + where + te :: SomeException -> Iteratee s m a + te e = icont (const (te e)) (Just e) + {-# INLINE liftWith #-} + {-# INLINE restoreT #-} + +instance (MonadBaseControl b m, Nullable s) => MonadBaseControl b (Iteratee s m) where + type StM (Iteratee s m) a = ComposeSt (Iteratee s) m a + liftBaseWith = defaultLiftBaseWith + restoreM = defaultRestoreM + + +-- |Send 'EOF' to the @Iteratee@ and disregard the unconsumed part of the +-- stream. If the iteratee is in an exception state, that exception is +-- thrown with 'Control.Exception.throw'. Iteratees that do not terminate +-- on @EOF@ will throw 'EofException'. +run :: Monad m => Iteratee s m a -> m a +run iter = runIter iter onDone onCont + where + onDone x _ = return x + onCont k Nothing = runIter (k (EOF Nothing)) onDone onCont' + onCont _ (Just e) = E.throw e + onCont' _ Nothing = E.throw EofException + onCont' _ (Just e) = E.throw e + +-- |Run an iteratee, returning either the result or the iteratee exception. +-- Note that only internal iteratee exceptions will be returned; exceptions +-- thrown with @Control.Exception.throw@ or @Control.Monad.CatchIO.throw@ will +-- not be returned. +-- +-- See 'Data.Iteratee.Exception.IFException' for details. +tryRun :: (Exception e, Monad m) => Iteratee s m a -> m (Either e a) +tryRun iter = runIter iter onDone onCont + where + onDone x _ = return $ Right x + onCont k Nothing = runIter (k (EOF Nothing)) onDone onCont' + onCont _ (Just e) = return $ maybeExc e + onCont' _ Nothing = return $ maybeExc (toException EofException) + onCont' _ (Just e) = return $ maybeExc e + maybeExc e = maybe (Left (E.throw e)) Left (fromException e) + +-- |Transform a computation inside an @Iteratee@. +mapIteratee :: (NullPoint s, Monad n, Monad m) => + (m a -> n b) + -> Iteratee s m a + -> Iteratee s n b +mapIteratee f = lift . f . run +{-# DEPRECATED mapIteratee "This function will be removed, compare to 'ilift'" #-} + +-- | Lift a computation in the inner monad of an iteratee. +-- +-- A simple use would be to lift a logger iteratee to a monad stack. +-- +-- > logger :: Iteratee String IO () +-- > logger = mapChunksM_ putStrLn +-- > +-- > loggerG :: MonadIO m => Iteratee String m () +-- > loggerG = ilift liftIO logger +-- +-- A more complex example would involve lifting an iteratee to work with +-- interleaved streams. See the example at 'Data.Iteratee.ListLike.merge'. +ilift :: + (Monad m, Monad n) + => (forall r. m r -> n r) + -> Iteratee s m a + -> Iteratee s n a +ilift f i = Iteratee $ \od oc -> + let onDone a str = return $ Left (a,str) + onCont k mErr = return $ Right (ilift f . k, mErr) + in f (runIter i onDone onCont) >>= either (uncurry od) (uncurry oc) + +-- | Lift a computation in the inner monad of an iteratee, while threading +-- through an accumulator. +ifold :: (Monad m, Monad n) => (forall r. m r -> acc -> n (r, acc)) + -> acc -> Iteratee s m a -> Iteratee s n (a, acc) +ifold f acc i = Iteratee $ \ od oc -> do + (r, acc') <- flip f acc $ + runIter i (curry $ return . Left) (curry $ return . Right) + either (uncurry (od . flip (,) acc')) + (uncurry (oc . (ifold f acc .))) r diff --git a/src/Bio/Iteratee/Bgzf.hsc b/src/Bio/Iteratee/Bgzf.hsc --- a/src/Bio/Iteratee/Bgzf.hsc +++ b/src/Bio/Iteratee/Bgzf.hsc @@ -14,15 +14,14 @@ import Bio.Iteratee import Bio.Prelude import Control.Concurrent.Async ( async, wait ) -import Foreign.Marshal.Alloc ( mallocBytes, free, allocaBytes ) -import Foreign.Storable ( peekByteOff, pokeByteOff ) import Foreign.C.String ( withCAString ) import Foreign.C.Types ( CInt(..), CChar(..), CUInt(..), CULong(..) ) +import Foreign.Marshal.Alloc ( mallocBytes, free, allocaBytes ) import Foreign.Ptr ( nullPtr, castPtr, Ptr, plusPtr, minusPtr ) +import Foreign.Storable ( peekByteOff, pokeByteOff ) import qualified Data.ByteString as S import qualified Data.ByteString.Unsafe as S -import qualified Data.Iteratee.ListLike as I #include @@ -32,8 +31,8 @@ data Block = Block { block_offset :: {-# UNPACK #-} !FileOffset , block_contents :: {-# UNPACK #-} !Bytes } -instance NullPoint Block where empty = mempty -instance Nullable Block where nullC (Block _ s) = S.null s +instance NullPoint Block where emptyP = Block 0 S.empty +instance Nullable Block where nullC = S.null . block_contents instance Monoid Block where mempty = Block 0 S.empty @@ -123,24 +122,24 @@ -- | Decodes a BGZF block header and returns the block size if -- successful. get_bgzf_header :: Monad m => Iteratee Bytes m (Word16, Word16) -get_bgzf_header = do n <- I.heads "\31\139" - _cm <- I.head - flg <- I.head +get_bgzf_header = do n <- heads "\31\139" + _cm <- headStream + flg <- headStream if flg `testBit` 2 then do - I.drop 6 + dropStream 6 xlen <- endianRead2 LSB - it <- I.take (fromIntegral xlen) get_bsize >>= lift . tryRun + it <- takeStream (fromIntegral xlen) get_bsize >>= lift . tryRun case it of Left e -> throwErr e Right s | n == 2 -> return (s,xlen) _ -> throwErr $ iterStrExc "No BGZF" else throwErr $ iterStrExc "No BGZF" where - get_bsize = do i1 <- I.head - i2 <- I.head + get_bsize = do i1 <- headStream + i2 <- headStream len <- endianRead2 LSB if i1 == 66 && i2 == 67 && len == 2 then endianRead2 LSB - else I.drop (fromIntegral len) >> get_bsize + else dropStream (fromIntegral len) >> get_bsize -- | Tests whether a stream is in BGZF format. Does not consume any -- input. @@ -152,9 +151,9 @@ isGzip :: Monad m => Iteratee Bytes m Bool isGzip = liftM (either (const False) id) $ checkErr $ iLookAhead $ test where - test = do n <- I.heads "\31\139" - I.drop 24 - b <- I.isFinished + test = do n <- heads "\31\139" + dropStream 24 + b <- isFinished return $ not b && n == 2 -- ------------------------------------------------------------------------- Output @@ -376,7 +375,7 @@ | LeftoverChunk !Bytes BgzfChunk | NoChunk -instance NullPoint BgzfChunk where empty = NoChunk +instance NullPoint BgzfChunk where emptyP = NoChunk instance Nullable BgzfChunk where nullC NoChunk = True nullC (SpecialChunk s c) = S.null s && nullC c @@ -446,7 +445,7 @@ queue_depth :: Int } deriving Show -compressChunk :: Int -> Ptr CChar -> CUInt -> IO Bytes +compressChunk :: Int -> Ptr Word8 -> CUInt -> IO Bytes compressChunk lv ptr len = allocaBytes (#{const sizeof(z_stream)}) $ \stream -> do buf <- mallocBytes 65536 @@ -473,7 +472,7 @@ z_check "deflateEnd" =<< c_deflateEnd stream crc0 <- c_crc32 0 nullPtr 0 - crc <- c_crc32 crc0 ptr len + crc <- c_crc32 crc0 (castPtr ptr) len compressed_length <- (+) (18+8) `fmap` #{peek z_stream, total_out} stream when (compressed_length > 65536) $ error "produced too big a block" diff --git a/src/Bio/Iteratee/Binary.hs b/src/Bio/Iteratee/Binary.hs new file mode 100644 --- /dev/null +++ b/src/Bio/Iteratee/Binary.hs @@ -0,0 +1,199 @@ +{-# LANGUAGE FlexibleContexts, BangPatterns #-} + +-- |Monadic Iteratees: +-- incremental input parsers, processors, and transformers +-- +-- Iteratees for parsing binary data. + +module Bio.Iteratee.Binary ( + -- * Types + Endian (..) + -- * Endian multi-byte iteratees + ,endianRead2 + ,endianRead3 + ,endianRead3i + ,endianRead4 + ,endianRead8 + -- ** bytestring specializations + -- | In current versions of @iteratee@ there is no difference between the + -- bytestring specializations and polymorphic functions. They exist + -- for compatibility. + ,readWord16be_bs + ,readWord16le_bs + ,readWord32be_bs + ,readWord32le_bs + ,readWord64be_bs + ,readWord64le_bs +) +where + +import Bio.Iteratee.Base +import Data.Bits +import Data.Int +import Data.Word +import Prelude + +import qualified Bio.Iteratee.ListLike as I +import qualified Data.ByteString as B +import qualified Data.ListLike as LL + +-- ------------------------------------------------------------------------ +-- Binary Random IO Iteratees + +-- Iteratees to read unsigned integers written in Big- or Little-endian ways + +-- | Indicate endian-ness. +data Endian = MSB -- ^ Most Significant Byte is first (big-endian) + | LSB -- ^ Least Significan Byte is first (little-endian) + deriving (Eq, Ord, Show, Enum) + +endianRead2 + :: LL.ListLike s Word8 + => Endian + -> Iteratee s m Word16 +endianRead2 e = endianReadN e 2 word16' +{-# INLINE endianRead2 #-} + +endianRead3 + :: LL.ListLike s Word8 + => Endian + -> Iteratee s m Word32 +endianRead3 e = endianReadN e 3 (word32' . (0:)) +{-# INLINE endianRead3 #-} + +-- |Read 3 bytes in an endian manner. If the first bit is set (negative), +-- set the entire first byte so the Int32 will be negative as +-- well. +endianRead3i + :: (Nullable s, LL.ListLike s Word8, Monad m) + => Endian + -> Iteratee s m Int32 +endianRead3i e = do + c1 <- I.headStream + c2 <- I.headStream + c3 <- I.headStream + 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 +{-# INLINE endianRead3i #-} + +endianRead4 + :: LL.ListLike s Word8 + => Endian + -> Iteratee s m Word32 +endianRead4 e = endianReadN e 4 word32' +{-# INLINE endianRead4 #-} + +endianRead8 + :: LL.ListLike s Word8 + => Endian + -> Iteratee s m Word64 +endianRead8 e = endianReadN e 8 word64' +{-# INLINE endianRead8 #-} + +-- This function does all the parsing work, depending upon provided arguments +endianReadN :: + LL.ListLike s Word8 + => Endian + -> Int + -> ([Word8] -> b) + -> Iteratee s m b +endianReadN MSB n0 cnct = liftI (step n0 []) + where + step !n acc (Chunk c) + | LL.null c = liftI (step n acc) + | LL.length c >= n = let (this,next) = LL.splitAt n c + !result = cnct $ acc ++ LL.toList this + in idone result (Chunk next) + | otherwise = liftI (step (n - LL.length c) (acc ++ LL.toList c)) + step !n acc (EOF Nothing) = icont (step n acc) (Just $ toException EofException) + step !n acc (EOF (Just e)) = icont (step n acc) (Just e) +endianReadN LSB n0 cnct = liftI (step n0 []) + where + step !n acc (Chunk c) + | LL.null c = liftI (step n acc) + | LL.length c >= n = let (this,next) = LL.splitAt n c + !result = cnct $ reverse (LL.toList this) ++ acc + in idone result (Chunk next) + | otherwise = liftI (step (n - LL.length c) + (reverse (LL.toList c) ++ acc)) + step !n acc (EOF Nothing) = icont (step n acc) + (Just $ toException EofException) + step !n acc (EOF (Just e)) = icont (step n acc) (Just e) +{-# INLINE endianReadN #-} + +-- As of now, the polymorphic code is as fast as the best specializations +-- I have found, so these just call out. They may be improved in the +-- future, or possibly deprecated. +-- JWL, 2012-01-16 + +readWord16be_bs :: Iteratee B.ByteString m Word16 +readWord16be_bs = endianRead2 MSB +{-# INLINE readWord16be_bs #-} + +readWord16le_bs :: Iteratee B.ByteString m Word16 +readWord16le_bs = endianRead2 LSB +{-# INLINE readWord16le_bs #-} + +readWord32be_bs :: Iteratee B.ByteString m Word32 +readWord32be_bs = endianRead4 MSB +{-# INLINE readWord32be_bs #-} + +readWord32le_bs :: Iteratee B.ByteString m Word32 +readWord32le_bs = endianRead4 LSB +{-# INLINE readWord32le_bs #-} + +readWord64be_bs :: Iteratee B.ByteString m Word64 +readWord64be_bs = endianRead8 MSB +{-# INLINE readWord64be_bs #-} + +readWord64le_bs :: Iteratee B.ByteString m Word64 +readWord64le_bs = endianRead8 LSB +{-# INLINE readWord64le_bs #-} + +word16' :: [Word8] -> Word16 +word16' [c1,c2] = word16 c1 c2 +word16' _ = error "iteratee: internal error in word16'" + +word16 :: Word8 -> Word8 -> Word16 +word16 c1 c2 = (fromIntegral c1 `shiftL` 8) .|. fromIntegral c2 +{-# INLINE word16 #-} + +word32' :: [Word8] -> Word32 +word32' [c1,c2,c3,c4] = word32 c1 c2 c3 c4 +word32' _ = error "iteratee: internal error in word32'" + +word32 :: Word8 -> Word8 -> Word8 -> Word8 -> Word32 +word32 c1 c2 c3 c4 = + (fromIntegral c1 `shiftL` 24) .|. + (fromIntegral c2 `shiftL` 16) .|. + (fromIntegral c3 `shiftL` 8) .|. + fromIntegral c4 +{-# INLINE word32 #-} + +word64' :: [Word8] -> Word64 +word64' [c1,c2,c3,c4,c5,c6,c7,c8] = word64 c1 c2 c3 c4 c5 c6 c7 c8 +word64' _ = error "iteratee: internal error in word64'" +{-# INLINE word64' #-} + +word64 + :: Word8 -> Word8 -> Word8 -> Word8 + -> Word8 -> Word8 -> Word8 -> Word8 + -> Word64 +word64 c1 c2 c3 c4 c5 c6 c7 c8 = + (fromIntegral c1 `shiftL` 56) .|. + (fromIntegral c2 `shiftL` 48) .|. + (fromIntegral c3 `shiftL` 40) .|. + (fromIntegral c4 `shiftL` 32) .|. + (fromIntegral c5 `shiftL` 24) .|. + (fromIntegral c6 `shiftL` 16) .|. + (fromIntegral c7 `shiftL` 8) .|. + fromIntegral c8 +{-# INLINE word64 #-} diff --git a/src/Bio/Iteratee/Builder.hs b/src/Bio/Iteratee/Builder.hs --- a/src/Bio/Iteratee/Builder.hs +++ b/src/Bio/Iteratee/Builder.hs @@ -1,226 +1,281 @@ --- | Buffer builder to assemble Bgzf blocks. (This will probably be --- renamed.) The plan is to serialize stuff (BAM and BCF) into a --- buffer, then Bgzf chunks from the buffer and reuse it. This /should/ --- avoid redundant copying and relieve some pressure from the garbage --- collector. And I hope to plug a mysterious memory leak that doesn't --- show up in the profiler. --- --- Exported functions with @unsafe@ in the name resulting in a type of --- 'Push' omit the bounds checking. To use them safely, an appropriate --- 'ensureBuffer' has to precede them. --- --- XXX This may not be the most clever way to do it. According to the --- reasoning behind the binary-serialise-cbor package, it would be more --- clever to have a representation of the things we can 'Push' that's --- similar to a list, and then a function (an Iteratee?) that consumes --- the list of tokens and fills a buffer. +{-# LANGUAGE ForeignFunctionInterface #-} +-- | Buffer builder to assemble Bgzf blocks. The plan is to serialize +-- stuff (BAM and BCF) into a buffer, then Bgzf chunks from the buffer. +-- We use a large buffer, and we always make sure there is plenty of +-- space in it (to avoid redundant checks). Whenever a block is ready +-- to be compressed, we stick it into a MVar. When we run out of space, +-- we simply use a new buffer. Multiple threads grab pieces from the +-- MVar, compress them, pass them downstream through another MVar. A +-- final thread restores the order and writes the blocks. -module Bio.Iteratee.Builder where +module Bio.Iteratee.Builder ( + BB(..), + newBuffer, + fillBuffer, + expandBuffer, + encodeBgzf, + BgzfTokens(..), + BclArgs(..), + BclSpecialType(..), + int_loop, + loop_bcl_special + ) where -import Bio.Iteratee hiding ( NullPoint ) -import Bio.Iteratee.Bgzf +import Bio.Iteratee +import Bio.Iteratee.Bgzf ( compressChunk, maxBlockSize, bgzfEofMarker ) import Bio.Prelude -import Data.NullPoint ( NullPoint(..) ) -import Foreign.ForeignPtr -import Foreign.Marshal.Alloc -import Foreign.Marshal.Utils -import Foreign.Ptr -import Foreign.Storable +import Foreign.ForeignPtr ( ForeignPtr, withForeignPtr, mallocForeignPtrBytes ) +import Foreign.Marshal.Utils ( copyBytes ) +import Foreign.Ptr ( Ptr, plusPtr ) +import Foreign.Storable ( pokeByteOff ) import qualified Data.ByteString as B import qualified Data.ByteString.Unsafe as B -import qualified Data.ByteString.Builder as B ( Builder, toLazyByteString ) -import qualified Data.ByteString.Lazy as B ( foldrChunks ) +import qualified Data.Vector.Storable as VS --- | The 'MutableByteArray' is garbage collected, so we don't get leaks. --- Once it has grown to a practical size (and the initial 128k should be --- very practical), we don't get fragmentation either. We also avoid --- copies for the most part, since no intermediate 'ByteString's, either --- lazy or strict have to be allocated. +-- | We manage a large buffer (multiple megabytes), of which we fill an +-- initial portion. We remeber the size, the used part, and two marks +-- where we later fill in sizes for the length prefixed BAM or BCF +-- records. We move the buffer down when we yield a piece downstream, +-- and when we run out of space, we simply move to a new buffer. +-- Garbage collection should take care of the rest. Unused 'mark' must +-- be set to (maxBound::Int) so it doesn't interfere with flushing. + data BB = BB { buffer :: {-# UNPACK #-} !(ForeignPtr Word8) - , size :: {-# UNPACK #-} !Int - , len :: {-# UNPACK #-} !Int - , mark :: {-# UNPACK #-} !Int - , mark2 :: {-# UNPACK #-} !Int } + , size :: {-# UNPACK #-} !Int -- total size of buffer + , off :: {-# UNPACK #-} !Int -- offset of active portion + , used :: {-# UNPACK #-} !Int -- used portion (inactive & active) + , mark :: {-# UNPACK #-} !Int -- offset of mark + , mark2 :: {-# UNPACK #-} !Int } -- offset of mark2 --- This still seems to have considerable overhead. Don't know if this --- can be improved by effectively inlining IO and turning the BB into an --- unboxed tuple. XXX -newtype Push = Push (BB -> IO BB) +instance Show BB where + show bb = show (size bb, off bb, used bb, mark bb, mark2 bb) -instance Monoid Push where - {-# INLINE mempty #-} - mempty = Push return - {-# INLINE mappend #-} - Push a `mappend` Push b = Push (a >=> b) +-- | Things we are able to encode. Taking inspiration from +-- binary-serialise-cbor, we define these as a lazy list-like thing and +-- consume it in a interpreter. -instance NullPoint Push where - empty = Push return +data BgzfTokens = TkWord32 {-# UNPACK #-} !Word32 BgzfTokens -- a 4-byte int + | TkWord16 {-# UNPACK #-} !Word16 BgzfTokens -- a 2-byte int + | TkWord8 {-# UNPACK #-} !Word8 BgzfTokens -- a byte + | TkFloat {-# UNPACK #-} !Float BgzfTokens -- a float + | TkDouble {-# UNPACK #-} !Double BgzfTokens -- a double + | TkString {-# UNPACK #-} !B.ByteString BgzfTokens -- a raw string + | TkDecimal {-# UNPACK #-} !Int BgzfTokens -- roughly ':%d' + | TkLnString {-# UNPACK #-} !B.ByteString BgzfTokens -- a length-prefixed string + -- lotsa stuff might be missing here + | TkSetMark BgzfTokens -- sets the first mark + | TkEndRecord BgzfTokens -- completes a BAM record + | TkEndRecordPart1 BgzfTokens -- completes part 1 of a BCF record + | TkEndRecordPart2 BgzfTokens -- completes part 2 of a BCF record + | TkEnd -- nothing more, for now --- | Creates a buffer with a given initial capacity. + -- specialties + | TkBclSpecial !BclArgs BgzfTokens + | TkLowLevel {-# UNPACK #-} !Int (BB -> IO BB) BgzfTokens + +data BclSpecialType = BclNucsBin | BclNucsAsc | BclNucsAscRev | BclQualsBin | BclQualsAsc | BclQualsAscRev + +data BclArgs = BclArgs BclSpecialType + {-# UNPACK #-} !(VS.Vector Word8) -- bcl matrix + {-# UNPACK #-} !Int -- stride + {-# UNPACK #-} !Int -- first cycle + {-# UNPACK #-} !Int -- last cycle + {-# UNPACK #-} !Int -- cluster index + +-- | Creates a buffer. newBuffer :: Int -> IO BB -newBuffer sz = mallocForeignPtrBytes sz >>= \ar -> return $ BB ar sz 0 0 0 +newBuffer sz = mallocForeignPtrBytes sz >>= \ar -> return $ BB ar sz 0 0 maxBound maxBound --- | Ensures a given free space in the buffer by doubling its capacity --- if necessary. -{-# INLINE ensureBuffer #-} -ensureBuffer :: Int -> Push -ensureBuffer n = Push $ \b -> - if len b + n < size b - then return b - else expandBuffer b +-- | Creates a new buffer, copying the content from an old one, with +-- higher capacity. +expandBuffer :: Int -> BB -> IO BB +expandBuffer minsz b = do + let sz' = max (2 * (size b - used b)) minsz + arr1 <- mallocForeignPtrBytes sz' + withForeignPtr arr1 $ \d -> + withForeignPtr (buffer b) $ \s -> + copyBytes d (plusPtr s (off b)) (used b - off b) + return $ BB { buffer = arr1 + , size = sz' + , off = 0 + , used = used b - off b + , mark = if mark b == maxBound then maxBound else mark b - off b + , mark2 = if mark2 b == maxBound then maxBound else mark2 b - off b } -expandBuffer :: BB -> IO BB -expandBuffer b = do arr1 <- mallocForeignPtrBytes (size b + size b) - withForeignPtr arr1 $ \d -> - withForeignPtr (buffer b) $ \s -> - copyBytes d s (len b) - return $ BB { buffer = arr1 - , size = size b + size b - , len = len b - , mark = mark b - , mark2 = mark2 b } +compressChunk' :: Int -> ForeignPtr Word8 -> Int -> Int -> IO B.ByteString +compressChunk' lv fptr off len = + withForeignPtr fptr $ \ptr -> + compressChunk lv (plusPtr ptr off) (fromIntegral len) -{-# INLINE unsafePushByte #-} -unsafePushByte :: Word8 -> Push -unsafePushByte w = Push $ \b -> do - withForeignPtr (buffer b) $ \p -> - pokeByteOff p (len b) w - return $ b { len = len b + 1 } +instance Nullable (Endo BgzfTokens) where + nullC f = case appEndo f TkEnd of TkEnd -> True ; _ -> False -{-# INLINE pushByte #-} -pushByte :: Word8 -> Push -pushByte b = ensureBuffer 1 <> unsafePushByte b +-- | Expand a chain of tokens into a buffer, sending finished pieces +-- downstream as soon as possible. +encodeBgzf :: MonadIO m => Int -> Enumeratee (Endo BgzfTokens) B.ByteString m b +encodeBgzf lv = (\out -> newBuffer (1024*1024) `ioBind` \bb -> eneeCheckIfDone (liftI . go bb) out) + ><> parRunIO (2*numCapabilities) + where + go bb0 k (EOF mx) = final_flush bb0 mx k + go bb0 k (Chunk f) + -- initially, we make sure we have reasonable space. this may not be enough. + | size bb0 - used bb0 < 1024 = expandBuffer (1024*1024) bb0 `ioBind` \bb' -> go' bb' k (appEndo f TkEnd) + | otherwise = go' bb0 k (appEndo f TkEnd) -{-# INLINE unsafePushWord32 #-} -unsafePushWord32 :: Word32 -> Push -unsafePushWord32 w = unsafePushByte (fromIntegral $ w `shiftR` 0) - <> unsafePushByte (fromIntegral $ w `shiftR` 8) - <> unsafePushByte (fromIntegral $ w `shiftR` 16) - <> unsafePushByte (fromIntegral $ w `shiftR` 24) + -- we arrive here because we ran out of buffer space, so we always expand it. + go1 bb0 k tk = expandBuffer (1024*1024) bb0 `ioBind` \bb' -> go' bb' k tk -{-# INLINE unsafePushWord16 #-} -unsafePushWord16 :: Word16 -> Push -unsafePushWord16 w = unsafePushByte (fromIntegral $ w `shiftR` 0) - <> unsafePushByte (fromIntegral $ w `shiftR` 8) + go' bb0 k tk = fillBuffer bb0 tk `ioBind` \(bb',tk') -> flush_blocks tk' bb' k -{-# INLINE pushWord32 #-} -pushWord32 :: Word32 -> Push -pushWord32 w = ensureBuffer 4 <> unsafePushWord32 w -{-# INLINE pushWord16 #-} -pushWord16 :: Word16 -> Push -pushWord16 w = ensureBuffer 2 <> unsafePushWord16 w + -- We can flush anything that is between 'off' and the lower of 'mark' + -- and 'used'. When done, we bump 'off'. + flush_blocks tk bb k + | min (mark bb) (used bb) - off bb < maxBlockSize = + case tk of TkEnd -> liftI $ go bb k + _ -> go1 bb k tk -{-# INLINE unsafePushByteString #-} -unsafePushByteString :: B.ByteString -> Push -unsafePushByteString bs = Push $ \b -> - B.unsafeUseAsCStringLen bs $ \(p,ln) -> - withForeignPtr (buffer b) $ \adr -> - b { len = len b + ln } <$ - copyBytes (adr `plusPtr` len b) p ln + | otherwise = do + eneeCheckIfDone (flush_blocks tk bb { off = off bb + maxBlockSize }) $ + k $ Chunk [compressChunk' lv (buffer bb) (off bb) maxBlockSize] -{-# INLINE pushByteString #-} -pushByteString :: B.ByteString -> Push -pushByteString bs = ensureBuffer (B.length bs) <> unsafePushByteString bs + final_flush bb mx k + | used bb > off bb = + idone (k $ Chunk [ compressChunk' lv (buffer bb) (off bb) (used bb - off bb) + , return bgzfEofMarker ]) (EOF mx) + | otherwise = + idone (k $ Chunk [ return bgzfEofMarker ]) (EOF mx) -{-# INLINE unsafePushFloat #-} -unsafePushFloat :: Float -> Push -unsafePushFloat f = - unsafePushWord32 $ unsafeDupablePerformIO $ - alloca $ \b -> poke (castPtr b) f >> peek b -{-# INLINE pushFloat #-} -pushFloat :: Float -> Push -pushFloat f = ensureBuffer 4 <> unsafePushFloat f +fillBuffer :: BB -> BgzfTokens -> IO (BB, BgzfTokens) +fillBuffer bb0 tk = withForeignPtr (buffer bb0) (\p -> go_slowish p bb0 tk) + where + go_slowish p bb tk1 = go_fast p bb (used bb) tk1 -{-# INLINE pushBuilder #-} -pushBuilder :: B.Builder -> Push -pushBuilder = B.foldrChunks ((<>) . pushByteString) mempty . B.toLazyByteString + go_fast p bb use tk1 = case tk1 of + -- no space? not our job. + _ | size bb - use < 1024 -> return (bb { used = use },tk1) --- | Sets a mark. This can later be filled in with a record length --- (used to create BAM records). -{-# INLINE unsafeSetMark #-} -unsafeSetMark :: Push -unsafeSetMark = Push $ \b -> return $ b { len = len b + 4, mark = len b } + -- the actual end. + TkEnd -> return (bb { used = use },tk1) -{-# INLINE setMark #-} -setMark :: Push -setMark = ensureBuffer 4 <> unsafeSetMark + -- I'm cheating. This stuff works only if the platform allows + -- unaligned accesses, is little-endian and uses IEEE floats. + -- It's true on i386 and ix86_64. + TkWord32 x tk' -> do pokeByteOff p use x + go_fast p bb (use + 4) tk' --- | Ends a record by filling the length into the field that was --- previously marked. Terrible things will happen if this wasn't --- preceded by a corresponding 'setMark'. -{-# INLINE endRecord #-} -endRecord :: Push -endRecord = Push $ \b -> withForeignPtr (buffer b) $ \p -> do - let !l = len b - mark b - 4 - pokeByteOff p (mark b + 0) (fromIntegral $ shiftR l 0 :: Word8) - pokeByteOff p (mark b + 1) (fromIntegral $ shiftR l 8 :: Word8) - pokeByteOff p (mark b + 2) (fromIntegral $ shiftR l 16 :: Word8) - pokeByteOff p (mark b + 3) (fromIntegral $ shiftR l 24 :: Word8) - return b + TkWord16 x tk' -> do pokeByteOff p use x + go_fast p bb (use + 2) tk' --- | Ends the first part of a record. The length is filled in *before* --- the mark, which is specifically done to support the *two* length --- fields in BCF. It also remembers the current position. Horrible --- things happen if this isn't preceeded by *two* succesive invocations --- of 'setMark'. -{-# INLINE endRecordPart1 #-} -endRecordPart1 :: Push -endRecordPart1 = Push $ \b -> withForeignPtr (buffer b) $ \p -> do - let !l = len b - mark b - 4 - pokeByteOff p (mark b - 4) (fromIntegral $ shiftR l 0 :: Word8) - pokeByteOff p (mark b - 3) (fromIntegral $ shiftR l 8 :: Word8) - pokeByteOff p (mark b - 2) (fromIntegral $ shiftR l 16 :: Word8) - pokeByteOff p (mark b - 1) (fromIntegral $ shiftR l 24 :: Word8) - return $ b { mark2 = len b } + TkWord8 x tk' -> do pokeByteOff p use x + go_fast p bb (use + 1) tk' --- | Ends the second part of a record. The length is filled in at the --- mark, but computed from the sencond mark only. This is specifically --- done to support the *two* length fields in BCF. Horrible things --- happen if this isn't preceeded by *two* succesive invocations of --- 'setMark' and one of 'endRecordPart1'. -{-# INLINE endRecordPart2 #-} -endRecordPart2 :: Push -endRecordPart2 = Push $ \b -> withForeignPtr (buffer b) $ \p -> do - let !l = len b - mark2 b - pokeByteOff p (mark b + 0) (fromIntegral $ shiftR l 0 :: Word8) - pokeByteOff p (mark b + 1) (fromIntegral $ shiftR l 8 :: Word8) - pokeByteOff p (mark b + 2) (fromIntegral $ shiftR l 16 :: Word8) - pokeByteOff p (mark b + 3) (fromIntegral $ shiftR l 24 :: Word8) - return b + TkFloat x tk' -> do pokeByteOff p use x + go_fast p bb (use + 4) tk' + TkDouble x tk' -> do pokeByteOff p use x + go_fast p bb (use + 8) tk' -{-# INLINE encodeBgzfWith #-} -encodeBgzfWith :: MonadIO m => Int -> Enumeratee Push B.ByteString m b -encodeBgzfWith lv o = newBuffer 128000 `ioBind` \bb -> eneeCheckIfDone (liftI . step bb) o - where - step bb k (EOF mx) = finalFlush bb k mx - step bb k (Chunk (Push p)) = p bb `ioBind` \bb' -> tryFlush bb' 0 k + TkString s tk' + -- Too big, can't handle. We will get progressively bigger + -- buffers and eventually handle it; for very large strings, + -- it works, but isn't ideal. XXX + | B.length s > size bb - use -> return (bb { used = use },tk') - tryFlush bb off k - | len bb - off < maxBlockSize - = withForeignPtr (buffer bb) - (\p -> moveBytes p (p `plusPtr` off) (len bb - off)) - `ioBind_` liftI (step (bb { len = len bb - off - , mark = mark bb - off `max` 0 }) k) - | otherwise - = withForeignPtr (buffer bb) - (\adr -> compressChunk lv (adr `plusPtr` off) (fromIntegral maxBlockSize)) - `ioBind` eneeCheckIfDone (tryFlush bb (off+maxBlockSize)) . k . Chunk + | otherwise -> do let ln = B.length s + B.unsafeUseAsCString s $ \q -> + copyBytes (p `plusPtr` use) q ln + go_fast p bb (use + ln) tk' - finalFlush bb k mx - | len bb < maxBlockSize - = withForeignPtr (buffer bb) - (\adr -> compressChunk lv (castPtr adr) (fromIntegral $ len bb)) - `ioBind` eneeCheckIfDone (finalFlush2 mx) . k . Chunk + TkDecimal x tk' -> do ln <- int_loop (p `plusPtr` use) x + go_fast p bb (use + ln) tk' - | otherwise - = error "WTF?! This wasn't supposed to happen." + TkLnString s tk' + -- Too big, can't handle. We will get progressively bigger + -- buffers and eventually handle it; for very large strings, + -- it works, but isn't ideal. XXX + | B.length s > size bb - use - 4 -> return (bb { used = use },tk') - finalFlush2 mx k = idone (k $ Chunk bgzfEofMarker) (EOF mx) + | otherwise -> do let ln = B.length s + pokeByteOff p use (fromIntegral ln :: Word32) + B.unsafeUseAsCString s $ \q -> + copyBytes (p `plusPtr` (use + 4)) q ln + go_fast p bb (use + ln + 4) tk' + TkSetMark tk' -> go_slowish p bb { used = use + 4, mark = use } tk' + TkEndRecord tk' -> do let !l = use - mark bb - 4 + pokeByteOff p (mark bb) (fromIntegral l :: Word32) + go_slowish p bb { used = use, mark = maxBound } tk' + + TkEndRecordPart1 tk' -> do let !l = use - mark bb - 4 + pokeByteOff p (mark bb - 4) (fromIntegral l :: Word32) + go_slowish p bb { used = use, mark2 = use } tk' + + TkEndRecordPart2 tk' -> do let !l = use - mark2 bb + pokeByteOff p (mark bb) (fromIntegral l :: Word32) + go_slowish p bb { used = use, mark = maxBound } tk' + + + TkBclSpecial special_args tk' -> do + l <- loop_bcl_special (p `plusPtr` use) special_args + go_fast p bb (use + l) tk' + + TkLowLevel minsize proc tk' + | size bb - use < minsize -> return (bb { used = use },tk1) + | otherwise -> do bb' <- proc bb { used = use } + go_slowish p bb' tk' + + +loop_bcl_special :: Ptr Word8 -> BclArgs -> IO Int +loop_bcl_special p (BclArgs tp vec stride u v i) = + + VS.unsafeWith vec $ \q -> case tp of + BclNucsBin -> do + nuc_loop p stride (plusPtr q i) u v + return $ (v - u + 2) `div` 2 + + BclNucsAsc -> do + nuc_loop_asc p stride (plusPtr q i) u v + return $ v - u + 1 + + BclNucsAscRev -> do + nuc_loop_asc_rev p stride (plusPtr q i) u v + return $ v - u + 1 + + BclQualsBin -> do + qual_loop p stride (plusPtr q i) u v + return $ v - u + 1 + + BclQualsAsc -> do + qual_loop_asc p stride (plusPtr q i) u v + return $ v - u + 1 + + BclQualsAscRev -> do + qual_loop_asc_rev p stride (plusPtr q i) u v + return $ v - u + 1 + +foreign import ccall unsafe "nuc_loop" + nuc_loop :: Ptr Word8 -> Int -> Ptr Word8 -> Int -> Int -> IO () + +foreign import ccall unsafe "nuc_loop_asc" + nuc_loop_asc :: Ptr Word8 -> Int -> Ptr Word8 -> Int -> Int -> IO () + +foreign import ccall unsafe "nuc_loop_asc_rev" + nuc_loop_asc_rev :: Ptr Word8 -> Int -> Ptr Word8 -> Int -> Int -> IO () + +foreign import ccall unsafe "qual_loop" + qual_loop :: Ptr Word8 -> Int -> Ptr Word8 -> Int -> Int -> IO () + +foreign import ccall unsafe "qual_loop_asc" + qual_loop_asc :: Ptr Word8 -> Int -> Ptr Word8 -> Int -> Int -> IO () + +foreign import ccall unsafe "qual_loop_asc_rev" + qual_loop_asc_rev :: Ptr Word8 -> Int -> Ptr Word8 -> Int -> Int -> IO () + +foreign import ccall unsafe "int_loop" + int_loop :: Ptr Word8 -> Int -> IO Int diff --git a/src/Bio/Iteratee/Char.hs b/src/Bio/Iteratee/Char.hs new file mode 100644 --- /dev/null +++ b/src/Bio/Iteratee/Char.hs @@ -0,0 +1,151 @@ +{-# LANGUAGE ScopedTypeVariables, FlexibleContexts #-} + +-- | Utilities for Char-based iteratee processing. + +module Bio.Iteratee.Char ( + -- * Word and Line processors + printLines + ,printLinesUnterminated + ,enumLines + ,enumLinesBS + ,enumWords + ,enumWordsBS +) + +where + +import Bio.Iteratee.Iteratee +import Bio.Iteratee.ListLike +import Data.Char +import Data.Monoid +import qualified Data.ListLike as LL +import Control.Monad (liftM) +import Control.Monad.IO.Class +import qualified Data.ByteString.Char8 as BC +import Prelude + + +-- |Print lines as they are received. This is the first `impure' iteratee +-- with non-trivial actions during chunk processing +-- +-- Only lines ending with a newline are printed, +-- data terminated with EOF is not printed. +printLines :: Iteratee String IO () +printLines = lines' + where + lines' = breakStream (\c -> c == '\r' || c == '\n') >>= \l -> terminators >>= check l + check _ 0 = return () + check "" _ = return () + check l _ = liftIO (putStrLn l) >> lines' + +-- |Print lines as they are received. +-- +-- All lines are printed, including a line with a terminating EOF. +-- If the final line is terminated by EOF without a newline, +-- no newline is printed. +-- this function should be used in preference to printLines when possible, +-- as it is more efficient with long lines. +printLinesUnterminated :: forall s el. + (Eq el, Nullable s, LL.StringLike s, LL.ListLike s el) + => Iteratee s IO () +printLinesUnterminated = lines' + where + lines' = do + joinI $ breakE (`LL.elem` t1) (mapChunksM_ (putStr . LL.toString)) + terminators >>= check + check 0 = return () + check _ = liftIO (putStrLn "") >> lines' + t1 :: s + t1 = LL.fromString "\r\n" + +terminators :: (Eq el, Nullable s, LL.StringLike s, LL.ListLike s el) + => Iteratee s IO Int +terminators = do + l <- heads (LL.fromString "\r\n") + if l == 0 then heads (LL.fromString "\n") else return l + + +-- |Convert the stream of characters to the stream of lines, and +-- apply the given iteratee to enumerate the latter. +-- The stream of lines is normally terminated by the empty line. +-- When the stream of characters is terminated, the stream of lines +-- is also terminated. +-- This is the first proper iteratee-enumerator: it is the iteratee of the +-- character stream and the enumerator of the line stream. + +enumLines + :: (LL.ListLike s el, LL.StringLike s, Nullable s, Monad m) => + Enumeratee s [s] m a +enumLines = convStream getter + where + getter = icont step Nothing + lChar = (== '\n') . last . LL.toString + step (Chunk xs) + | LL.null xs = getter + | lChar xs = idone (LL.lines xs) mempty + | otherwise = icont (step' xs) Nothing + step _str = getter + step' xs (Chunk ys) + | LL.null ys = icont (step' xs) Nothing + | lChar ys = idone (LL.lines . mappend xs $ ys) mempty + | otherwise = let w' = LL.lines $ mappend xs ys + ws = init w' + ck = last w' + in idone ws (Chunk ck) + step' xs str = idone (LL.lines xs) str + +-- |Convert the stream of characters to the stream of words, and +-- apply the given iteratee to enumerate the latter. +-- Words are delimited by white space. +-- This is the analogue of List.words +enumWords :: (LL.ListLike s Char, Nullable s, Monad m) => Enumeratee s [s] m a +enumWords = convStream $ dropWhileStream isSpace >> liftM (:[]) (breakStream isSpace) +{-# INLINE enumWords #-} + +-- Like enumWords, but operates on ByteStrings. +-- This is provided as a higher-performance alternative to enumWords, and +-- is equivalent to treating the stream as a Data.ByteString.Char8.ByteString. +enumWordsBS + :: (Monad m) => Enumeratee BC.ByteString [BC.ByteString] m a +enumWordsBS iter = convStream getter iter + where + getter = liftI step + lChar = isSpace . BC.last + step (Chunk xs) + | BC.null xs = getter + | lChar xs = idone (BC.words xs) (Chunk BC.empty) + | otherwise = icont (step' xs) Nothing + step str = idone mempty str + step' xs (Chunk ys) + | BC.null ys = icont (step' xs) Nothing + | lChar ys = idone (BC.words . BC.append xs $ ys) mempty + | otherwise = let w' = BC.words . BC.append xs $ ys + ws = init w' + ck = last w' + in idone ws (Chunk ck) + step' xs str = idone (BC.words xs) str + +{-# INLINE enumWordsBS #-} + +-- Like enumLines, but operates on ByteStrings. +-- This is provided as a higher-performance alternative to enumLines, and +-- is equivalent to treating the stream as a Data.ByteString.Char8.ByteString. +enumLinesBS :: (Monad m) => Enumeratee BC.ByteString [BC.ByteString] m a +enumLinesBS = convStream getter + where + getter = icont step Nothing + lChar = (== '\n') . BC.last + step (Chunk xs) + | BC.null xs = getter + | lChar xs = idone (BC.lines xs) (Chunk BC.empty) + | otherwise = icont (step' xs) Nothing + step str = idone mempty str + step' xs (Chunk ys) + | BC.null ys = icont (step' xs) Nothing + | lChar ys = idone (BC.lines . BC.append xs $ ys) mempty + | otherwise = let w' = BC.lines $ BC.append xs ys + ws = init w' + ck = last w' + in idone ws (Chunk ck) + step' xs str = idone (BC.lines xs) str + diff --git a/src/Bio/Iteratee/Exception.hs b/src/Bio/Iteratee/Exception.hs new file mode 100644 --- /dev/null +++ b/src/Bio/Iteratee/Exception.hs @@ -0,0 +1,211 @@ +{-# LANGUAGE DeriveDataTypeable, ExistentialQuantification #-} + +-- |Monadic and General Iteratees: +-- Messaging and exception handling. +-- +-- Iteratees use an internal exception handling mechanism that is parallel to +-- that provided by 'Control.Exception'. This allows the iteratee framework +-- to handle its own exceptions outside @IO@. +-- +-- Iteratee exceptions are divided into two categories, 'IterException' and +-- 'EnumException'. @IterExceptions@ are exceptions within an iteratee, and +-- @EnumExceptions@ are exceptions within an enumerator. +-- +-- Enumerators can be constructed to handle an 'IterException' with +-- @Data.Iteratee.Iteratee.enumFromCallbackCatch@. If the enumerator detects +-- an @iteratee exception@, the enumerator calls the provided exception handler. +-- The enumerator is then able to continue feeding data to the iteratee, +-- provided the exception was successfully handled. If the handler could +-- not handle the exception, the 'IterException' is converted to an +-- 'EnumException' and processing aborts. +-- +-- Exceptions can also be cleared by @Data.Iteratee.Iteratee.checkErr@, +-- although in this case the iteratee continuation cannot be recovered. +-- +-- When viewed as Resumable Exceptions, iteratee exceptions provide a means +-- for iteratees to send control messages to enumerators. The @seek@ +-- implementation provides an example. @Data.Iteratee.Iteratee.seek@ stores +-- the current iteratee continuation and throws a 'SeekException', which +-- inherits from 'IterException'. @Data.Iteratee.IO.enumHandleRandom@ is +-- constructed with @enumFromCallbackCatch@ and a handler that performs +-- an @hSeek@. Upon receiving the 'SeekException', @enumHandleRandom@ calls +-- the handler, checks that it executed properly, and then continues with +-- the stored continuation. +-- +-- As the exception hierarchy is open, users can extend it with custom +-- exceptions and exception handlers to implement sophisticated messaging +-- systems based upon resumable exceptions. + + +module Bio.Iteratee.Exception ( + -- * Exception types + IFException (..) + ,Exception (..) -- from Control.Exception + -- ** Enumerator exceptions + ,EnumException (..) + ,DivergentException (..) + ,EnumStringException (..) + ,EnumUnhandledIterException (..) + -- ** Iteratee exceptions + ,IException (..) + ,IterException (..) + ,SeekException (..) + ,EofException (..) + ,IterStringException (..) + -- * Functions + ,enStrExc + ,iterStrExc + ,wrapIterExc + ,iterExceptionToException + ,iterExceptionFromException +) +where + +import Bio.Iteratee.IO.Base +import Control.Exception +import Data.Data +import Prelude + + +-- ---------------------------------------------- +-- create exception type hierarchy + +-- |Root of the Iteratee exception hierarchy. @IFException@ derives from +-- @Control.Exception.SomeException@. 'EnumException', 'IterException', +-- and all inheritants are descendents of 'IFException'. +data IFException = forall e . Exception e => IFException e + deriving Typeable + +instance Show IFException where + show (IFException e) = show e + +instance Exception IFException + +ifExceptionToException :: Exception e => e -> SomeException +ifExceptionToException = toException . IFException + +ifExceptionFromException :: Exception e => SomeException -> Maybe e +ifExceptionFromException x = do + IFException a <- fromException x + cast a + +-- Root of enumerator exceptions. +data EnumException = forall e . Exception e => EnumException e + deriving Typeable + +instance Show EnumException where + show (EnumException e) = show e + +instance Exception EnumException where + toException = ifExceptionToException + fromException = ifExceptionFromException + +enumExceptionToException :: Exception e => e -> SomeException +enumExceptionToException = toException . IterException + +enumExceptionFromException :: Exception e => SomeException -> Maybe e +enumExceptionFromException x = do + IterException a <- fromException x + cast a + +-- |The @iteratee@ diverged upon receiving 'EOF'. +data DivergentException = DivergentException + deriving (Show, Typeable) + +instance Exception DivergentException where + toException = enumExceptionToException + fromException = enumExceptionFromException + +-- |Create an enumerator exception from a @String@. +data EnumStringException = EnumStringException String + deriving (Show, Typeable) + +instance Exception EnumStringException where + toException = enumExceptionToException + fromException = enumExceptionFromException + +-- |Create an 'EnumException' from a string. +enStrExc :: String -> EnumException +enStrExc = EnumException . EnumStringException + +-- |The enumerator received an 'IterException' it could not handle. +data EnumUnhandledIterException = EnumUnhandledIterException IterException + deriving (Show, Typeable) + +instance Exception EnumUnhandledIterException where + toException = enumExceptionToException + fromException = enumExceptionFromException + +-- |Convert an 'IterException' to an 'EnumException'. Meant to be used +-- within an @Enumerator@ to signify that it could not handle the +-- @IterException@. +wrapIterExc :: IterException -> EnumException +wrapIterExc = EnumException . EnumUnhandledIterException + +-- iteratee exceptions + +-- |A class for @iteratee exceptions@. Only inheritants of @IterException@ +-- should be instances of this class. +class Exception e => IException e where + toIterException :: e -> IterException + toIterException = IterException + fromIterException :: IterException -> Maybe e + fromIterException = fromException . toException + +-- |Root of iteratee exceptions. +data IterException = forall e . Exception e => IterException e + deriving Typeable + +instance Show IterException where + show (IterException e) = show e + +instance Exception IterException where + toException = ifExceptionToException + fromException = ifExceptionFromException + +iterExceptionToException :: Exception e => e -> SomeException +iterExceptionToException = toException . IterException + +iterExceptionFromException :: Exception e => SomeException -> Maybe e +iterExceptionFromException x = do + IterException a <- fromException x + cast a + +instance IException IterException where + toIterException = id + fromIterException = Just + +-- |A seek request within an @Iteratee@. +data SeekException = SeekException FileOffset + deriving (Typeable, Show) + +instance Exception SeekException where + toException = iterExceptionToException + fromException = iterExceptionFromException + +instance IException SeekException where + +-- |The @Iteratee@ needs more data but received @EOF@. +data EofException = EofException + deriving (Typeable, Show) + +instance Exception EofException where + toException = iterExceptionToException + fromException = iterExceptionFromException + +instance IException EofException where + +-- |An @Iteratee exception@ specified by a @String@. +data IterStringException = IterStringException String deriving (Typeable, Show) + +instance Exception IterStringException where + toException = iterExceptionToException + fromException = iterExceptionFromException + +instance IException IterStringException where + +-- |Create an @iteratee exception@ from a string. +-- This convenience function wraps 'IterStringException' and 'toException'. +iterStrExc :: String -> SomeException +iterStrExc= toException . IterStringException + diff --git a/src/Bio/Iteratee/IO.hs b/src/Bio/Iteratee/IO.hs new file mode 100644 --- /dev/null +++ b/src/Bio/Iteratee/IO.hs @@ -0,0 +1,97 @@ +{-# LANGUAGE ConstraintKinds #-} + +-- |Random and Binary IO with generic Iteratees. + +module Bio.Iteratee.IO( + -- * Data + defaultBufSize, + -- * File enumerators + -- ** Handle-based enumerators + H.enumHandle, + H.enumHandleRandom, + enumFile, + enumFileRandom, + -- ** FileDescriptor based enumerators + FD.enumFd, + FD.enumFdRandom, + -- * Iteratee drivers + -- These are FileDescriptor-based on POSIX systems, otherwise they are + -- Handle-based. The Handle-based drivers are accessible on POSIX systems + -- at Data.Iteratee.IO.Handle + fileDriver, + fileDriverVBuf, + fileDriverRandom, + fileDriverRandomVBuf, +) + +where + +import Bio.Iteratee.ReadableChunk +import Bio.Iteratee.Iteratee +import Bio.Iteratee.Binary () +import Control.Monad.Catch +import Control.Monad.IO.Class +import Prelude + +import qualified Bio.Iteratee.IO.Handle as H +import qualified Bio.Iteratee.IO.Fd as FD + + +-- | Default buffer size in elements. This was 1024 in "Data.Iteratee", +-- which is obviously too small. Since we often want to merge many +-- files, a read should take more time than a seek. This sets the +-- sensible buffer size to somewhat more than one MB. +defaultBufSize :: Int +defaultBufSize = 2*1024*1024 + + +-- If Posix is available, use the fileDriverRandomFd as fileDriverRandom. Otherwise, use a handle-based variant. +enumFile + :: (MonadIO m, MonadMask m, NullPoint s, ReadableChunk s el) => + Int + -> FilePath + -> Enumerator s m a +enumFile = FD.enumFile + +enumFileRandom + :: (MonadIO m, MonadMask m, NullPoint s, ReadableChunk s el) => + Int + -> FilePath + -> Enumerator s m a +enumFileRandom = FD.enumFileRandom + +-- |Process a file using the given Iteratee. This function wraps +-- enumFd as a convenience. +fileDriver + :: (MonadIO m, MonadMask m, NullPoint s, ReadableChunk s el) => + Iteratee s m a + -> FilePath + -> m a +fileDriver = FD.fileDriverFd defaultBufSize + +-- |A version of fileDriver with a user-specified buffer size (in elements). +fileDriverVBuf + :: (MonadIO m, MonadMask m, NullPoint s, ReadableChunk s el) => + Int + -> Iteratee s m a + -> FilePath + -> m a +fileDriverVBuf = FD.fileDriverFd + +-- |Process a file using the given Iteratee. This function wraps +-- enumFdRandom as a convenience. +fileDriverRandom + :: (MonadIO m, MonadMask m, NullPoint s, ReadableChunk s el) => + Iteratee s m a + -> FilePath + -> m a +fileDriverRandom = FD.fileDriverRandomFd defaultBufSize + +fileDriverRandomVBuf + :: (MonadIO m, MonadMask m, NullPoint s, ReadableChunk s el) => + Int + -> Iteratee s m a + -> FilePath + -> m a +fileDriverRandomVBuf = FD.fileDriverRandomFd + diff --git a/src/Bio/Iteratee/IO/Base.hs b/src/Bio/Iteratee/IO/Base.hs new file mode 100644 --- /dev/null +++ b/src/Bio/Iteratee/IO/Base.hs @@ -0,0 +1,105 @@ +{-# LANGUAGE ForeignFunctionInterface, CPP #-} + +-- Low-level IO operations +-- These operations are either missing from the GHC run-time library, +-- or implemented suboptimally or heavy-handedly + +module Bio.Iteratee.IO.Base ( + FileOffset, + myfdRead, + myfdSeek, + Errno(..), + select'read'pending +) + +where + +import Control.Monad +import Data.Bits -- for select +import Foreign.C +import Foreign.Marshal.Array -- for select +import Foreign.Ptr +import Prelude +import System.IO (SeekMode(..)) +import System.Posix + +-- |Alas, GHC provides no function to read from Fd to an allocated buffer. +-- The library function fdRead is not appropriate as it returns a string +-- already. I'd rather get data from a buffer. +-- Furthermore, fdRead (at least in GHC) allocates a new buffer each +-- time it is called. This is a waste. Yet another problem with fdRead +-- is in raising an exception on any IOError or even EOF. I'd rather +-- avoid exceptions altogether. + +myfdRead :: Fd -> Ptr CChar -> ByteCount -> IO (Either Errno ByteCount) +myfdRead (Fd fd) ptr n = do + n' <- cRead fd ptr n + if n' == -1 then liftM Left getErrno + else return . Right . fromIntegral $ n' + +foreign import ccall unsafe "unistd.h read" cRead + :: CInt -> Ptr CChar -> CSize -> IO CInt + +-- |The following fseek procedure throws no exceptions. +myfdSeek:: Fd -> SeekMode -> FileOffset -> IO (Either Errno FileOffset) +myfdSeek (Fd fd) mode off = do + n' <- cLSeek fd off (mode2Int mode) + if n' == -1 then liftM Left getErrno + else return . Right $ n' + where mode2Int :: SeekMode -> CInt -- From GHC source + mode2Int AbsoluteSeek = 0 + mode2Int RelativeSeek = 1 + mode2Int SeekFromEnd = 2 + +foreign import ccall unsafe "unistd.h lseek" cLSeek + :: CInt -> FileOffset -> CInt -> IO FileOffset + + +-- Darn! GHC doesn't provide the real select over several descriptors! +-- We have to implement it ourselves + +type FDSET = CUInt +type TIMEVAL = CLong -- Two longs +foreign import ccall "unistd.h select" c_select + :: CInt -> Ptr FDSET -> Ptr FDSET -> Ptr FDSET -> Ptr TIMEVAL -> IO CInt + +-- Convert a file descriptor to an FDSet (for use with select) +-- essentially encode a file descriptor in a big-endian notation +fd2fds :: CInt -> [FDSET] +fd2fds fd = replicate nb 0 ++ [setBit 0 off] + where + (nb,off) = quotRem (fromIntegral fd) bitSize_FDSET + +bitSize_FDSET :: Int +#if MIN_VERSION_base(4,7,0) +bitSize_FDSET = finiteBitSize (undefined::FDSET) +#else +bitSize_FDSET = bitSize (undefined::FDSET) +#endif + +fds2mfd :: [FDSET] -> [CInt] +fds2mfd fds = [fromIntegral (j+i*bitSize_FDSET) | + (afds,i) <- zip fds [0..], j <- [0..bitSize_FDSET], + testBit afds j] + +unFd :: Fd -> CInt +unFd (Fd x) = x + +-- |poll if file descriptors have something to read +-- Return the list of read-pending descriptors +select'read'pending :: [Fd] -> IO (Either Errno [Fd]) +select'read'pending mfd = + withArray ([0,1]::[TIMEVAL]) $ \_timeout -> + withArray fds $ \readfs -> do + rc <- c_select (fdmax+1) readfs nullPtr nullPtr nullPtr + if rc == -1 + then liftM Left getErrno + -- because the wait was indefinite, rc must be positive! + else liftM (Right . map Fd . fds2mfd) (peekArray (length fds) readfs) + where + fds :: [FDSET] + fds = foldr ormax [] (map (fd2fds . unFd) mfd) + fdmax = maximum $ map fromIntegral mfd + ormax [] x = x + ormax x [] = x + ormax (a:ar) (b:br) = (a .|. b) : ormax ar br diff --git a/src/Bio/Iteratee/IO/Fd.hs b/src/Bio/Iteratee/IO/Fd.hs new file mode 100644 --- /dev/null +++ b/src/Bio/Iteratee/IO/Fd.hs @@ -0,0 +1,151 @@ +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE ConstraintKinds #-} + +-- |Random and Binary IO with generic Iteratees, using File Descriptors for IO. +-- when available, these are the preferred functions for performing IO as they +-- run in constant space and function properly with sockets, pipes, etc. + +module Bio.Iteratee.IO.Fd( + -- * File enumerators + -- ** FileDescriptor based enumerators for monadic iteratees + enumFd + ,enumFdCatch + ,enumFdRandom + ,enumFile + ,enumFileRandom + -- * Iteratee drivers + ,fileDriverFd + ,fileDriverRandomFd +) + +where + +import Bio.Iteratee.Binary () +import Bio.Iteratee.IO.Base +import Bio.Iteratee.Iteratee +import Bio.Iteratee.ReadableChunk +import Bio.Prelude +import Control.Monad.Catch as CIO +import Control.Monad.IO.Class +import Foreign.Marshal.Alloc +import Foreign.Ptr +import Foreign.Storable +import System.IO (SeekMode(..)) + + +-- ------------------------------------------------------------------------ +-- Binary Random IO enumerators + +makefdCallback :: + (MonadIO m, NullPoint s, ReadableChunk s el) => + Ptr el + -> ByteCount + -> Fd + -> st + -> m (Either SomeException ((Bool, st), s)) +makefdCallback p bufsize fd st = do + n <- liftIO $ myfdRead fd (castPtr p) bufsize + case n of + Left _ -> return $ Left (error "myfdRead failed") + Right 0 -> liftIO yield >> return (Right ((False, st), emptyP)) + Right n' -> liftM (\s -> Right ((True, st), s)) $ + readFromPtr p (fromIntegral n') + +-- |The enumerator of a POSIX File Descriptor. This version enumerates +-- over the entire contents of a file, in order, unless stopped by +-- the iteratee. In particular, seeking is not supported. +enumFd + :: forall s el m a.(NullPoint s, ReadableChunk s el, MonadIO m, MonadMask m) => + Int + -> Fd + -> Enumerator s m a +enumFd bs fd iter = + let bufsize = bs * (sizeOf (undefined :: el)) + in CIO.bracket (liftIO $ mallocBytes bufsize) + (liftIO . free) + (\p -> enumFromCallback (makefdCallback p (fromIntegral bufsize) fd) () iter) + +-- |A variant of enumFd that catches exceptions raised by the @Iteratee@. +enumFdCatch + :: forall e s el m a.(IException e, NullPoint s, ReadableChunk s el, MonadIO m, MonadMask m) + => Int + -> Fd + -> (e -> m (Maybe EnumException)) + -> Enumerator s m a +enumFdCatch bs fd handler iter = + let bufsize = bs * (sizeOf (undefined :: el)) + in CIO.bracket (liftIO $ mallocBytes bufsize) + (liftIO . free) + (\p -> enumFromCallbackCatch (makefdCallback p (fromIntegral bufsize) fd) handler () iter) + + +-- |The enumerator of a POSIX File Descriptor: a variation of @enumFd@ that +-- supports RandomIO (seek requests). +enumFdRandom + :: forall s el m a.(NullPoint s, ReadableChunk s el, MonadIO m, MonadMask m) => + Int + -> Fd + -> Enumerator s m a +enumFdRandom bs fd iter = enumFdCatch bs fd handler iter + where + handler (SeekException off) = + liftM (either + (const . Just $ enStrExc "Error seeking within file descriptor") + (const Nothing)) + . liftIO . myfdSeek fd AbsoluteSeek $ fromIntegral off + +fileDriver + :: (MonadIO m, MonadMask m) => + (Int -> Fd -> Enumerator s m a) + -> Int + -> Iteratee s m a + -> FilePath + -> m a +fileDriver enumf bufsize iter filepath = CIO.bracket + (liftIO $ openFd filepath ReadOnly Nothing defaultFileFlags) + (liftIO . closeFd) + (run <=< flip (enumf bufsize) iter) + +-- |Process a file using the given @Iteratee@. +fileDriverFd + :: (NullPoint s, MonadIO m, MonadMask m, ReadableChunk s el) => + Int -- ^Buffer size (number of elements) + -> Iteratee s m a + -> FilePath + -> m a +fileDriverFd = fileDriver enumFd + +-- |A version of fileDriverFd that supports seeking. +fileDriverRandomFd + :: (NullPoint s, MonadIO m, MonadMask m, ReadableChunk s el) => + Int + -> Iteratee s m a + -> FilePath + -> m a +fileDriverRandomFd = fileDriver enumFdRandom + +enumFile' :: (MonadIO m, MonadMask m) => + (Int -> Fd -> Enumerator s m a) + -> Int -- ^Buffer size + -> FilePath + -> Enumerator s m a +enumFile' enumf bufsize filepath iter = CIO.bracket + (liftIO $ openFd filepath ReadOnly Nothing defaultFileFlags) + (liftIO . closeFd) + (flip (enumf bufsize) iter) + +enumFile :: + (NullPoint s, MonadIO m, MonadMask m, ReadableChunk s el) + => Int -- ^Buffer size + -> FilePath + -> Enumerator s m a +enumFile = enumFile' enumFd + +enumFileRandom :: + (NullPoint s, MonadIO m, MonadMask m, ReadableChunk s el) + => Int -- ^Buffer size + -> FilePath + -> Enumerator s m a +enumFileRandom = enumFile' enumFdRandom + + diff --git a/src/Bio/Iteratee/IO/Handle.hs b/src/Bio/Iteratee/IO/Handle.hs new file mode 100644 --- /dev/null +++ b/src/Bio/Iteratee/IO/Handle.hs @@ -0,0 +1,150 @@ +{-# LANGUAGE ConstraintKinds #-} +{-# LANGUAGE ScopedTypeVariables #-} + +-- |Random and Binary IO with generic Iteratees. These functions use Handles +-- for IO operations, and are provided for compatibility. When available, +-- the File Descriptor based functions are preferred as these wastefully +-- allocate memory rather than running in constant space. + +module Bio.Iteratee.IO.Handle( + -- * File enumerators + enumHandle + ,enumHandleCatch + ,enumHandleRandom + ,enumFile + ,enumFileRandom + -- * Iteratee drivers + ,fileDriverHandle + ,fileDriverRandomHandle +) + +where + +import Bio.Iteratee.Binary () +import Bio.Iteratee.Iteratee +import Bio.Iteratee.ReadableChunk +import Bio.Prelude +import Control.Monad.Catch as CIO +import Control.Monad.IO.Class +import Foreign.Marshal.Alloc +import Foreign.Ptr +import Foreign.Storable +import System.IO + +-- ------------------------------------------------------------------------ +-- Binary Random IO enumerators + +makeHandleCallback :: + (MonadIO m, NullPoint s, ReadableChunk s el) => + Ptr el + -> Int + -> Handle + -> st + -> m (Either SomeException ((Bool, st), s)) +makeHandleCallback p bsize h st = do + n' <- liftIO (CIO.try $ hGetBuf h p bsize :: IO (Either SomeException Int)) + case n' of + Left e -> return $ Left e + Right 0 -> return $ Right ((False, st), emptyP) + Right n -> liftM (\s -> Right ((True, st), s)) $ + readFromPtr p (fromIntegral n) + + +-- |The (monadic) enumerator of a file Handle. This version enumerates +-- over the entire contents of a file, in order, unless stopped by +-- the iteratee. In particular, seeking is not supported. +-- Data is read into a buffer of the specified size. +enumHandle :: + forall s el m a.(NullPoint s, ReadableChunk s el, MonadIO m, MonadMask m) => + Int -- ^Buffer size (number of elements per read) + -> Handle + -> Enumerator s m a +enumHandle bs h i = + let bufsize = bs * sizeOf (undefined :: el) + in CIO.bracket (liftIO $ mallocBytes bufsize) + (liftIO . free) + (\p -> enumFromCallback (makeHandleCallback p bufsize h) () i) + +-- |An enumerator of a file handle that catches exceptions raised by +-- the Iteratee. +enumHandleCatch + :: + forall e s el m a.(IException e, + NullPoint s, + ReadableChunk s el, + MonadIO m, MonadMask m) => + Int -- ^Buffer size (number of elements per read) + -> Handle + -> (e -> m (Maybe EnumException)) + -> Enumerator s m a +enumHandleCatch bs h handler i = + let bufsize = bs * sizeOf (undefined :: el) + in CIO.bracket (liftIO $ mallocBytes bufsize) + (liftIO . free) + (\p -> enumFromCallbackCatch (makeHandleCallback p bufsize h) handler () i) + + +-- |The enumerator of a Handle: a variation of enumHandle that +-- supports RandomIO (seek requests). +-- Data is read into a buffer of the specified size. +enumHandleRandom :: + forall s el m a.(NullPoint s, ReadableChunk s el, MonadIO m, MonadMask m) => + Int -- ^ Buffer size (number of elements per read) + -> Handle + -> Enumerator s m a +enumHandleRandom bs h i = enumHandleCatch bs h handler i + where + handler (SeekException off) = + liftM (either + (Just . EnumException :: IOException -> Maybe EnumException) + (const Nothing)) + . liftIO . CIO.try $ hSeek h AbsoluteSeek $ fromIntegral off + +-- ---------------------------------------------- +-- File Driver wrapper functions. + +enumFile' :: (MonadIO m, MonadMask m) => + (Int -> Handle -> Enumerator s m a) + -> Int -- ^Buffer size + -> FilePath + -> Enumerator s m a +enumFile' enumf bufsize filepath iter = CIO.bracket + (liftIO $ openBinaryFile filepath ReadMode) + (liftIO . hClose) + (flip (enumf bufsize) iter) + +enumFile :: + (NullPoint s, MonadIO m, MonadMask m, ReadableChunk s el) + => Int -- ^Buffer size + -> FilePath + -> Enumerator s m a +enumFile = enumFile' enumHandle + +enumFileRandom :: + (NullPoint s, MonadIO m, MonadMask m, ReadableChunk s el) + => Int -- ^Buffer size + -> FilePath + -> Enumerator s m a +enumFileRandom = enumFile' enumHandleRandom + +-- |Process a file using the given @Iteratee@. This function wraps +-- @enumHandle@ as a convenience. +fileDriverHandle + :: (NullPoint s, MonadIO m, MonadMask m, ReadableChunk s el) => + Int -- ^Buffer size (number of elements) + -> Iteratee s m a + -> FilePath + -> m a +fileDriverHandle bufsize iter filepath = + enumFile bufsize filepath iter >>= run + +-- |A version of @fileDriverHandle@ that supports seeking. +fileDriverRandomHandle + :: (NullPoint s, MonadIO m, MonadMask m, ReadableChunk s el) => + Int -- ^ Buffer size (number of elements) + -> Iteratee s m a + -> FilePath + -> m a +fileDriverRandomHandle bufsize iter filepath = + enumFileRandom bufsize filepath iter >>= run + diff --git a/src/Bio/Iteratee/Iteratee.hs b/src/Bio/Iteratee/Iteratee.hs new file mode 100644 --- /dev/null +++ b/src/Bio/Iteratee/Iteratee.hs @@ -0,0 +1,583 @@ +{-# LANGUAGE KindSignatures + ,RankNTypes + ,FlexibleContexts + ,ScopedTypeVariables + ,BangPatterns + ,DeriveDataTypeable #-} + +-- |Monadic and General Iteratees: +-- incremental input parsers, processors and transformers + +module Bio.Iteratee.Iteratee ( + -- * Types + EnumerateeHandler + -- ** Error handling + ,throwErr + ,throwRecoverableErr + ,checkErr + -- ** Basic Iteratees + ,unitIter + ,skipToEof + ,isStreamFinished + -- ** Iteratee composition + ,mBind + ,mBind_ + ,ioBind + ,ioBind_ + -- ** Chunkwise Iteratees + ,mapChunksM_ + ,foldChunksM + ,getChunk + ,getChunks + -- ** Nested iteratee combinators + ,mapChunks + ,mapChunksM + ,convStream + ,unfoldConvStream + ,unfoldConvStreamCheck + ,joinI + ,joinIM + -- * Enumerators + ,Enumerator + ,Enumeratee + -- ** Basic enumerators + ,enumChunk + ,enumEof + ,enumErr + ,enumPure1Chunk + ,enumList + ,enumCheckIfDone + ,enumFromCallback + ,enumFromCallbackCatch + -- ** Enumerator Combinators + ,eneeCheckIfDone + ,eneeCheckIfDoneHandle + ,eneeCheckIfDoneIgnore + ,eneeCheckIfDonePass + ,mergeEnums + -- ** Enumeratee Combinators + ,($=) + ,(=$) + ,(><>) + ,(<><) + -- * Misc. + ,seek + ,FileOffset + -- * Classes + ,module Bio.Iteratee.Base +) +where + +import Bio.Iteratee.IO.Base +import Bio.Iteratee.Base +import Bio.Prelude hiding (loop) +import Control.Monad.IO.Class +import Control.Monad.Trans.Class + +-- exception helpers +excDivergent :: SomeException +excDivergent = toException DivergentException + +-- ------------------------------------------------------------------------ +-- Primitive iteratees + +-- |Report and propagate an unrecoverable error. +-- Disregard the input first and then propagate the error. This error +-- cannot be handled by 'enumFromCallbackCatch', although it can be cleared +-- by 'checkErr'. +throwErr :: SomeException -> Iteratee s m a +throwErr e = icont (const (throwErr e)) (Just e) + +-- |Report and propagate a recoverable error. This error can be handled by +-- both 'enumFromCallbackCatch' and 'checkErr'. +throwRecoverableErr :: + SomeException + -> (Stream s -> Iteratee s m a) + -> Iteratee s m a +throwRecoverableErr e i = icont i (Just e) + + +-- |Check if an iteratee produces an error. +-- Returns @Right a@ if it completes without errors, otherwise +-- @Left SomeException@. 'checkErr' is useful for iteratees that may not +-- terminate, such as @Data.Iteratee.head@ with an empty stream. +checkErr :: + (NullPoint s) => + Iteratee s m a + -> Iteratee s m (Either SomeException a) +checkErr iter = Iteratee $ \onDone onCont -> + let od = onDone . Right + oc k Nothing = onCont (checkErr . k) Nothing + oc _ (Just e) = onDone (Left e) (Chunk emptyP) + in runIter iter od oc + +-- ------------------------------------------------------------------------ +-- Parser combinators + +-- |The identity iteratee. Doesn't do any processing of input. +unitIter :: (NullPoint s) => Iteratee s m () +unitIter = idone () (Chunk emptyP) + +-- |Get the stream status of an iteratee. +isStreamFinished :: (Nullable s) => Iteratee s m (Maybe SomeException) +isStreamFinished = liftI check + where + check s@(Chunk xs) + | nullC xs = isStreamFinished + | otherwise = idone Nothing s + check s@(EOF e) = idone (Just $ fromMaybe (toException EofException) e) s +{-# INLINE isStreamFinished #-} + + +-- |Skip the rest of the stream +skipToEof :: Iteratee s m () +skipToEof = icont check Nothing + where + check (Chunk _) = skipToEof + check s = idone () s + + +-- |Seek to a position in the stream +seek :: (NullPoint s) => FileOffset -> Iteratee s m () +seek o = throwRecoverableErr (toException $ SeekException o) (const unitIter) + +-- | Map a monadic function over the chunks of the stream and ignore the +-- result. Useful for creating efficient monadic iteratee consumers, e.g. +-- +-- > logger = mapChunksM_ (liftIO . putStrLn) +-- +-- these can be efficiently run in parallel with other iteratees via +-- @Data.Iteratee.ListLike.zip@. +mapChunksM_ :: (Monad m, Nullable s) => (s -> m b) -> Iteratee s m () +mapChunksM_ f = liftI step + where + step (Chunk xs) + | nullC xs = liftI step + | otherwise = lift (f xs) >> liftI step + step s@(EOF _) = idone () s +{-# INLINE mapChunksM_ #-} + +-- | A fold over chunks +foldChunksM :: (Monad m, Nullable s) => (a -> s -> m a) -> a -> Iteratee s m a +foldChunksM f = liftI . go + where + go a (Chunk c) = lift (f a c) >>= liftI . go + go a e = idone a e +{-# INLINE foldChunksM #-} + +-- | Get the current chunk from the stream. +getChunk :: Nullable s => Iteratee s m s +getChunk = liftI step + where + step (Chunk xs) + | nullC xs = liftI step + | otherwise = idone xs $ Chunk emptyP + step (EOF Nothing) = throwErr $ toException EofException + step (EOF (Just e)) = throwErr e +{-# INLINE getChunk #-} + +-- | Get a list of all chunks from the stream. +getChunks :: (Nullable s) => Iteratee s m [s] +getChunks = liftI (step id) + where + step acc (Chunk xs) + | nullC xs = liftI (step acc) + | otherwise = liftI (step $ acc . (xs:)) + step acc stream = idone (acc []) stream +{-# INLINE getChunks #-} + +-- --------------------------------------------------- +-- The converters show a different way of composing two iteratees: +-- `vertical' rather than `horizontal' + +type Enumeratee sFrom sTo (m :: * -> *) a = + Iteratee sTo m a + -> Iteratee sFrom m (Iteratee sTo m a) + +-- The following pattern appears often in Enumeratee code +{-# INLINE eneeCheckIfDone #-} + +-- | Utility function for creating enumeratees. Typical usage is demonstrated +-- by the @breakE@ definition. +-- +-- > breakE +-- > :: (Monad m, LL.ListLike s el, NullPoint s) +-- > => (el -> Bool) +-- > -> Enumeratee s s m a +-- > breakE cpred = eneeCheckIfDone (liftI . step) +-- > where +-- > step k (Chunk s) +-- > | LL.null s = liftI (step k) +-- > | otherwise = case LL.break cpred s of +-- > (str', tail') +-- > | LL.null tail' -> eneeCheckIfDone (liftI . step) . k $ Chunk str' +-- > | otherwise -> idone (k $ Chunk str') (Chunk tail') +-- > step k stream = idone (k stream) stream +-- +eneeCheckIfDone :: + (Monad m, NullPoint elo) => + ((Stream eli -> Iteratee eli m a) -> Iteratee elo m (Iteratee eli m a)) + -> Enumeratee elo eli m a +eneeCheckIfDone f = eneeCheckIfDonePass f' + where + f' k Nothing = f k + f' k (Just e) = throwRecoverableErr e (\s -> joinIM $ enumChunk s $ eneeCheckIfDone f (liftI k)) + +type EnumerateeHandler eli elo m a = + (Stream eli -> Iteratee eli m a) + -> SomeException + -> Iteratee elo m (Iteratee eli m a) + +-- | The same as eneeCheckIfDonePass, with one extra argument: +-- a handler which is used +-- to process any exceptions in a separate method. +eneeCheckIfDoneHandle + :: (NullPoint elo) + => EnumerateeHandler eli elo m a + -> ((Stream eli -> Iteratee eli m a) + -> Maybe SomeException + -> Iteratee elo m (Iteratee eli m a) + ) + -> Enumeratee elo eli m a +eneeCheckIfDoneHandle h f inner = Iteratee $ \od oc -> + let onDone x s = od (idone x s) (Chunk emptyP) + onCont k Nothing = runIter (f k Nothing) od oc + onCont k (Just e) = runIter (h k e) od oc + in runIter inner onDone onCont +{-# INLINABLE eneeCheckIfDoneHandle #-} + +eneeCheckIfDonePass + :: (NullPoint elo) + => ((Stream eli -> Iteratee eli m a) + -> Maybe SomeException + -> Iteratee elo m (Iteratee eli m a) + ) + -> Enumeratee elo eli m a +eneeCheckIfDonePass f = eneeCheckIfDoneHandle (\k e -> f k (Just e)) f +{-# INLINABLE eneeCheckIfDonePass #-} + +eneeCheckIfDoneIgnore + :: (NullPoint elo) + => ((Stream eli -> Iteratee eli m a) + -> Maybe SomeException + -> Iteratee elo m (Iteratee eli m a) + ) + -> Enumeratee elo eli m a +eneeCheckIfDoneIgnore f = eneeCheckIfDoneHandle (\k _ -> f k Nothing) f + +{-# INLINE mBind #-} +-- | Lifts a monadic action and combines it with a continuation. +-- @mBind m f@ is the same as @lift m >>= f@, but does not require a +-- 'Nullable' constraint on the stream type. +infixl 1 `mBind` +mBind :: Monad m => m a -> (a -> Iteratee s m b) -> Iteratee s m b +mBind m f = Iteratee $ \onDone onCont -> m >>= \a -> runIter (f a) onDone onCont + +{-# INLINE mBind_ #-} +-- | Lifts a monadic action, ignored the result and combines it with a +-- continuation. @mBind_ m f@ is the same as @lift m >>= f@, but does +-- not require a 'Nullable' constraint on the stream type. +infixl 1 `mBind_` +mBind_ :: Monad m => m a -> Iteratee s m b -> Iteratee s m b +mBind_ m b = Iteratee $ \onDone onCont -> m >> runIter b onDone onCont + +{-# INLINE ioBind #-} +-- | Lifts an IO action and combines it with a continuation. +-- @ioBind m f@ is the same as @liftIO m >>= f@, but does not require a +-- 'Nullable' constraint on the stream type. +infixl 1 `ioBind` +ioBind :: MonadIO m => IO a -> (a -> Iteratee s m b) -> Iteratee s m b +ioBind m f = Iteratee $ \onDone onCont -> liftIO m >>= \a -> runIter (f a) onDone onCont + +{-# INLINE ioBind_ #-} +-- | Lifts an IO action, ignores its result, and combines it with a +-- continuation. @ioBind_ m f@ is the same as @liftIO m >> f@, but does +-- not require a 'Nullable' constraint on the stream type. +infixl 1 `ioBind_` +ioBind_ :: MonadIO m => IO a -> Iteratee s m b -> Iteratee s m b +ioBind_ m b = Iteratee $ \onDone onCont -> liftIO m >> runIter b onDone onCont + +-- | Convert one stream into another with the supplied mapping function. +-- This function operates on whole chunks at a time, contrasting to +-- @mapStream@ which operates on single elements. +-- +-- > unpacker :: Enumeratee B.ByteString [Word8] m a +-- > unpacker = mapChunks B.unpack +-- +mapChunks :: (NullPoint s) => (s -> s') -> Enumeratee s s' m a +mapChunks f = eneeCheckIfDonePass (icont . step) + where + step k (Chunk xs) = eneeCheckIfDonePass (icont . step) . k . Chunk $ f xs + step k str@(EOF mErr) = idone (k $ EOF mErr) str +{-# INLINE mapChunks #-} + +-- | Convert a stream of @s@ to a stream of @s'@ using the supplied function. +mapChunksM :: (Monad m, NullPoint s) => (s -> m s') -> Enumeratee s s' m a +mapChunksM f = eneeCheckIfDonePass (icont . step) + where + step k (Chunk xs) = f xs `mBind` eneeCheckIfDonePass (icont . step) . k . Chunk + step k str = idone (liftI k) str +{-# INLINE mapChunksM #-} + +-- |Convert one stream into another, not necessarily in lockstep. +-- +-- The transformer mapStream maps one element of the outer stream +-- to one element of the nested stream. The transformer below is more +-- general: it may take several elements of the outer stream to produce +-- one element of the inner stream, or the other way around. +-- The transformation from one stream to the other is specified as +-- Iteratee s m s'. +convStream :: + (Monad m, Nullable s) => + Iteratee s m s' + -> Enumeratee s s' m a +convStream fi = eneeCheckIfDonePass check + where + check k (Just e) = throwRecoverableErr e (const unitIter) >> check k Nothing + check k _ = isStreamFinished >>= maybe (step k) (idone (liftI k) . EOF . Just) + step k = fi >>= eneeCheckIfDonePass check . k . Chunk +{-# INLINABLE convStream #-} + +-- |The most general stream converter. Given a function to produce iteratee +-- transformers and an initial state, convert the stream using iteratees +-- generated by the function while continually updating the internal state. +unfoldConvStream :: + (Monad m, Nullable s) => + (acc -> Iteratee s m (acc, s')) + -> acc + -> Enumeratee s s' m a +unfoldConvStream f acc0 = eneeCheckIfDonePass (check acc0) + where + check acc k (Just e) = throwRecoverableErr e (const unitIter) >> check acc k Nothing + check acc k _ = isStreamFinished >>= + maybe (step acc k) (idone (liftI k) . EOF . Just) + step acc k = f acc >>= \(acc', s') -> + eneeCheckIfDonePass (check acc') . k . Chunk $ s' +{-# INLINABLE unfoldConvStream #-} + +unfoldConvStreamCheck + :: (Monad m, Nullable elo) + => (((Stream eli -> Iteratee eli m a) + -> Maybe SomeException + -> Iteratee elo m (Iteratee eli m a) + ) + -> Enumeratee elo eli m a + ) + -> (acc -> Iteratee elo m (acc, eli)) + -> acc + -> Enumeratee elo eli m a +unfoldConvStreamCheck checkDone f acc0 = checkDone (check acc0) + where + check acc k mX = isStreamFinished >>= + maybe (step acc k mX) (idone (icont k mX) . EOF . Just) + step acc k Nothing = f acc >>= \(acc', s') -> + (checkDone (check acc') . k $ Chunk s') + step acc k (Just ex) = throwRecoverableErr ex $ \str' -> + let i = f acc >>= \(acc', s') -> + (checkDone (check acc') . k $ Chunk s') + in joinIM $ enumChunk str' i +{-# INLINABLE unfoldConvStreamCheck #-} + +-- | Collapse a nested iteratee. The inner iteratee is terminated by @EOF@. +-- Errors are propagated through the result. +-- +-- The stream resumes from the point of the outer iteratee; any remaining +-- input in the inner iteratee will be lost. +-- Differs from 'Control.Monad.join' in that the inner iteratee is terminated, +-- and may have a different stream type than the result. +joinI :: + (Monad m, Nullable s) => + Iteratee s m (Iteratee s' m a) + -> Iteratee s m a +joinI = (>>= + \inner -> Iteratee $ \od oc -> + let onDone x _ = od x (Chunk emptyP) + onCont k Nothing = runIter (k (EOF Nothing)) onDone onCont' + onCont _ (Just e) = runIter (throwErr e) od oc + onCont' _ e = runIter (throwErr (fromMaybe excDivergent e)) od oc + in runIter inner onDone onCont) +{-# INLINE joinI #-} + +-- | Lift an iteratee inside a monad to an iteratee. +joinIM :: (Monad m) => m (Iteratee s m a) -> Iteratee s m a +joinIM mIter = Iteratee $ \od oc -> mIter >>= \iter -> runIter iter od oc + + +-- ------------------------------------------------------------------------ +-- Enumerators +-- | Each enumerator takes an iteratee and returns an iteratee +-- +-- an Enumerator is an iteratee transformer. +-- The enumerator normally stops when the stream is terminated +-- or when the iteratee moves to the done state, whichever comes first. +-- When to stop is of course up to the enumerator... + +type Enumerator s m a = Iteratee s m a -> m (Iteratee s m a) + +-- |Applies the iteratee to the given stream. This wraps 'enumEof', +-- 'enumErr', and 'enumPure1Chunk', calling the appropriate enumerator +-- based upon 'Stream'. +enumChunk :: (Monad m) => Stream s -> Enumerator s m a +enumChunk (Chunk xs) = enumPure1Chunk xs +enumChunk (EOF Nothing) = enumEof +enumChunk (EOF (Just e)) = enumErr e + +-- |The most primitive enumerator: applies the iteratee to the terminated +-- stream. The result is the iteratee in the Done state. It is an error +-- if the iteratee does not terminate on EOF. +enumEof :: (Monad m) => Enumerator s m a +enumEof iter = runIter iter onDone onCont + where + onDone x _str = return $ idone x (EOF Nothing) + onCont k Nothing = runIter (k (EOF Nothing)) onDone onCont' + onCont k e = return $ icont k e + onCont' _ Nothing = return $ throwErr excDivergent + onCont' k e = return $ icont k e + +-- |Another primitive enumerator: tell the Iteratee the stream terminated +-- with an error. +enumErr :: (Exception e, Monad m) => e -> Enumerator s m a +enumErr e iter = runIter iter onDone onCont + where + onDone x _ = return $ idone x (EOF . Just $ toException e) + onCont k Nothing = runIter (k (EOF (Just (toException e)))) onDone onCont' + onCont k e' = return $ icont k e' + onCont' _ Nothing = return $ throwErr excDivergent + onCont' k e' = return $ icont k e' + + +infixr 0 =$ + +-- | Combines an Enumeratee from @s@ to @s'@ and an Iteratee that +-- consumes @s'@ into an Iteratee which consumes @s@ +(=$) + :: (Nullable s, Monad m) + => Enumeratee s s' m a + -> Iteratee s' m a + -> Iteratee s m a +(=$) = (.) joinI + +infixl 1 $= + +-- | Combines Enumerator which produces stream of @s@ and @Enumeratee@ +-- which transforms stream of @s@ to stream +-- of @s'@ to into Enumerator which produces stream of @s'@ +($=) + :: Monad m + => (forall a. Enumerator s m a) + -> Enumeratee s s' m b + -> Enumerator s' m b +($=) enum enee iter = enum (enee iter) >>= run + + +-- | Enumeratee composition +-- Run the second enumeratee within the first. In this example, stream2list +-- is run within the 'takeStream 10', which is itself run within 'takeStream 15', resulting +-- in 15 elements being consumed +-- +-- >>> run =<< enumPure1Chunk [1..1000 :: Int] (joinI $ (I.takeStream 15 ><> I.takeStream 10) I.stream2list) +-- [1,2,3,4,5,6,7,8,9,10] +-- +(><>) :: + (Nullable s1, Monad m) + => (forall x . Enumeratee s1 s2 m x) + -> Enumeratee s2 s3 m a + -> Enumeratee s1 s3 m a +f ><> g = joinI . f . g + +-- | enumeratee composition with the arguments flipped, see '><>' +(<><) :: + (Nullable s1, Monad m) + => Enumeratee s2 s3 m a + -> (forall x. Enumeratee s1 s2 m x) + -> Enumeratee s1 s3 m a +f <>< g = joinI . g . f + +-- | Combine enumeration over two streams. The merging enumeratee would +-- typically be the result of 'Data.Iteratee.ListLike.merge' or +-- 'Data.Iteratee.ListLike.mergeByChunks' (see @merge@ for example). +mergeEnums :: + (Nullable s2, Nullable s1, Monad m) + => Enumerator s1 m a -- ^ inner enumerator + -> Enumerator s2 (Iteratee s1 m) a -- ^ outer enumerator + -> Enumeratee s2 s1 (Iteratee s1 m) a -- ^ merging enumeratee + -> Enumerator s1 m a +mergeEnums e1 e2 etee i = e1 $ e2 (joinI . etee $ ilift lift i) >>= run +{-# INLINE mergeEnums #-} + +-- | The pure 1-chunk enumerator +-- +-- It passes a given list of elements to the iteratee in one chunk +-- This enumerator does no IO and is useful for testing of base parsing +enumPure1Chunk :: (Monad m) => s -> Enumerator s m a +enumPure1Chunk str iter = runIter iter idoneM onCont + where + onCont k Nothing = return $ k $ Chunk str + onCont k e = return $ icont k e + +-- | Enumerate chunks from a list +-- +enumList :: (Monad m) => [s] -> Enumerator s m a +enumList chunks = go chunks + where + go [] i = return i + go xs' i = runIter i idoneM (onCont xs') + where + onCont (x:xs) k Nothing = go xs . k $ Chunk x + onCont _ _ (Just e) = return $ throwErr e + onCont _ k Nothing = return $ icont k Nothing +{-# INLINABLE enumList #-} + +-- | Checks if an iteratee has finished. +-- +-- This enumerator runs the iteratee, performing any monadic actions. +-- If the result is True, the returned iteratee is done. +enumCheckIfDone :: (Monad m) => Iteratee s m a -> m (Bool, Iteratee s m a) +enumCheckIfDone iter = runIter iter onDone onCont + where + onDone x str = return (True, idone x str) + onCont k e = return (False, icont k e) +{-# INLINE enumCheckIfDone #-} + + +-- |Create an enumerator from a callback function +enumFromCallback :: + (Monad m, NullPoint s) => + (st -> m (Either SomeException ((Bool, st), s))) + -> st + -> Enumerator s m a +enumFromCallback c st = + enumFromCallbackCatch c (\NotAnException -> return Nothing) st + +-- Dummy exception to catch in enumFromCallback +-- This never gets thrown, but it lets us +-- share plumbing +data NotAnException = NotAnException + deriving (Show, Typeable) + +instance Exception NotAnException where +instance IException NotAnException where + +-- |Create an enumerator from a callback function with an exception handler. +-- The exception handler is called if an iteratee reports an exception. +enumFromCallbackCatch + :: (IException e, Monad m, NullPoint s) + => (st -> m (Either SomeException ((Bool, st), s))) + -> (e -> m (Maybe EnumException)) + -> st + -> Enumerator s m a +enumFromCallbackCatch c handler = loop + where + loop st iter = runIter iter idoneM (onCont st) + check k (True, st') = loop st' . k . Chunk + check k (False,_st') = return . k . Chunk + onCont st k Nothing = c st >>= + either (return . k . EOF . Just) (uncurry (check k)) + onCont st k j@(Just e) = case fromException e of + Just e' -> handler e' >>= + maybe (loop st . k $ Chunk emptyP) + (return . icont k . Just) . fmap toException + Nothing -> return (icont k j) +{-# INLINE enumFromCallbackCatch #-} + + diff --git a/src/Bio/Iteratee/ListLike.hs b/src/Bio/Iteratee/ListLike.hs new file mode 100644 --- /dev/null +++ b/src/Bio/Iteratee/ListLike.hs @@ -0,0 +1,883 @@ +{-# LANGUAGE TupleSections, ScopedTypeVariables #-} + +-- |Monadic Iteratees: +-- incremental input parsers, processors and transformers +-- +-- This module provides many basic iteratees from which more complicated +-- iteratees can be built. In general these iteratees parallel those in +-- @Data.List@, with some additions. + +module Bio.Iteratee.ListLike ( + -- * Iteratees + -- ** Iteratee Utilities + isFinished + ,stream2list + ,stream2stream + -- ** Basic Iteratees + ,dropWhileStream + ,dropStream + ,headStream + ,tryHead + ,lastStream + ,heads + ,peekStream + ,roll + ,lengthStream + ,chunkLength + ,takeFromChunk + -- ** Nested iteratee combinators + ,breakStream + ,breakE + ,takeStream + ,takeUpTo + ,takeWhileE + ,mapStream + ,concatMapStream + ,concatMapStreamM + ,mapMaybeStream + ,rigidMapStream + ,filterStream + ,filterStreamM + ,groupStreamBy + ,groupStreamOn + ,mergeStreams + ,mergeByChunks + -- ** Folds + ,foldStream + -- * Enumerators + -- ** Basic enumerators + ,enumPureNChunk + -- ** Enumerator Combinators + ,enumWith + ,zipStreams + ,zipStreams3 + ,zipStreams4 + ,zipStreams5 + ,sequenceStreams_ + ,countConsumed + -- ** Monadic functions + ,mapStreamM + ,mapStreamM_ + ,foldStreamM + -- * Re-exported modules + ,module Bio.Iteratee.Iteratee +) +where + +import Bio.Iteratee.Iteratee +import Bio.Prelude +import Control.Monad.Trans.Class + +import qualified Data.ByteString as B +import qualified Data.ListLike as LL +import qualified Data.ListLike.FoldableLL as FLL + + +-- Useful combinators for implementing iteratees and enumerators + +-- | Check if a stream has received 'EOF'. +isFinished :: (Nullable s) => Iteratee s m Bool +isFinished = liftI check + where + check c@(Chunk xs) + | nullC xs = liftI check + | otherwise = idone False c + check s@(EOF _) = idone True s +{-# INLINE isFinished #-} + +-- ------------------------------------------------------------------------ +-- Primitive iteratees + +-- |Read a stream to the end and return all of its elements as a list. +-- This iteratee returns all data from the stream *strictly*. +stream2list :: (Monad m, Nullable s, LL.ListLike s el) => Iteratee s m [el] +stream2list = liftM (concatMap LL.toList) getChunks +{-# INLINE stream2list #-} + +-- |Read a stream to the end and return all of its elements as a stream. +-- This iteratee returns all data from the stream *strictly*. +stream2stream :: (Monad m, Nullable s, Monoid s) => Iteratee s m s +stream2stream = liftM mconcat getChunks +{-# INLINE stream2stream #-} + + +-- ------------------------------------------------------------------------ +-- Parser combinators + +-- |Attempt to read the next element of the stream and return it +-- Raise a (recoverable) error if the stream is terminated. +-- +-- The analogue of @List.head@ +-- +-- Because @head@ can raise an error, it shouldn't be used when constructing +-- iteratees for @convStream@. Use @tryHead@ instead. +headStream :: (LL.ListLike s el) => Iteratee s m el +headStream = liftI step + where + step (Chunk vec) + | LL.null vec = icont step Nothing + | otherwise = idone (LL.head vec) (Chunk $ LL.tail vec) + step stream = icont step (Just (setEOF stream)) +{-# INLINE headStream #-} + +-- | Similar to @headStream@, except it returns @Nothing@ if the stream +-- is terminated. +tryHead :: (LL.ListLike s el) => Iteratee s m (Maybe el) +tryHead = liftI step + where + step (Chunk vec) + | LL.null vec = liftI step + | otherwise = idone (Just $ LL.head vec) (Chunk $ LL.tail vec) + step stream = idone Nothing stream +{-# INLINE tryHead #-} + +-- |Attempt to read the last element of the stream and return it +-- Raise a (recoverable) error if the stream is terminated +-- +-- The analogue of @List.last@ +lastStream :: (LL.ListLike s el, Nullable s) => Iteratee s m el +lastStream = liftI (step Nothing) + where + step l (Chunk xs) + | nullC xs = liftI (step l) + | otherwise = liftI $ step (Just $ LL.last xs) + step l s@(EOF _) = case l of + Nothing -> icont (step l) . Just . setEOF $ s + Just x -> idone x s +{-# INLINE lastStream #-} + + +-- |Given a sequence of characters, attempt to match them against +-- the characters on the stream. Return the count of how many +-- characters matched. The matched characters are removed from the +-- stream. +-- For example, if the stream contains 'abd', then (heads 'abc') +-- will remove the characters 'ab' and return 2. +heads :: (Monad m, Nullable s, LL.ListLike s el, Eq el) => s -> Iteratee s m Int +heads st | nullC st = return 0 +heads st = loopE 0 st + where + loopE cnt xs + | nullC xs = return cnt + | otherwise = liftI (step cnt xs) + step cnt str (Chunk xs) | nullC xs = liftI (step cnt str) + step cnt str stream | nullC str = idone cnt stream + step cnt str s@(Chunk xs) = + if LL.head str == LL.head xs + then step (succ cnt) (LL.tail str) (Chunk $ LL.tail xs) + else idone cnt s + step cnt _ stream = idone cnt stream +{-# INLINE heads #-} + + +-- |Look ahead at the next element of the stream, without removing +-- it from the stream. +-- Return @Just c@ if successful, return @Nothing@ if the stream is +-- terminated by 'EOF'. +peekStream :: (LL.ListLike s el) => Iteratee s m (Maybe el) +peekStream = liftI step + where + step s@(Chunk vec) + | LL.null vec = liftI step + | otherwise = idone (Just $ LL.head vec) s + step stream = idone Nothing stream +{-# INLINE peekStream #-} + +-- | Return a chunk of @t@ elements length while consuming @d@ elements +-- from the stream. Useful for creating a 'rolling average' with +-- 'convStream'. +roll + :: (Monad m, Nullable s, LL.ListLike s el, LL.ListLike s' s) + => Int -- ^ length of chunk (t) + -> Int -- ^ amount to consume (d) + -> Iteratee s m s' +roll t d | t > d = liftI step + where + step (Chunk vec) + | LL.length vec >= t = + idone (LL.singleton $ LL.take t vec) (Chunk $ LL.drop d vec) + | LL.null vec = liftI step + | otherwise = liftI (step' vec) + step stream = idone LL.empty stream + step' v1 (Chunk vec) = step . Chunk $ v1 `mappend` vec + step' v1 stream = idone (LL.singleton v1) stream +roll t d = do r <- joinI (takeStream t stream2stream) + dropStream (d-t) + return $ LL.singleton r + -- d is >= t, so this version works +{-# INLINE roll #-} + + +-- |Drop n elements of the stream, if there are that many. +-- +-- The analogue of @List.drop@ +dropStream :: (Nullable s, LL.ListLike s el) => Int -> Iteratee s m () +dropStream 0 = idone () (Chunk emptyP) +dropStream n' = liftI (step n') + where + step n (Chunk str) + | LL.length str < n = liftI (step (n - LL.length str)) + | otherwise = idone () (Chunk (LL.drop n str)) + step _ stream = idone () stream +{-# INLINE dropStream #-} + +-- |Skip all elements while the predicate is true. +-- +-- The analogue of @List.dropWhile@ +dropWhileStream :: (LL.ListLike s el) => (el -> Bool) -> Iteratee s m () +dropWhileStream p = liftI step + where + step (Chunk str) + | LL.null rest = liftI step + | otherwise = idone () (Chunk rest) + where + rest = LL.dropWhile p str + step stream = idone () stream +{-# INLINE dropWhileStream #-} + + +-- | Return the total length of the remaining part of the stream. +-- +-- This forces evaluation of the entire stream. +-- +-- The analogue of @List.length@ +lengthStream :: (Num a, LL.ListLike s el) => Iteratee s m a +lengthStream = liftI (step 0) + where + step !i (Chunk xs) = liftI (step $ i + fromIntegral (LL.length xs)) + step !i stream = idone i stream +{-# INLINE lengthStream #-} + +-- | Get the length of the current chunk, or @Nothing@ if 'EOF'. +-- +-- This function consumes no input. +chunkLength :: (LL.ListLike s el) => Iteratee s m (Maybe Int) +chunkLength = liftI step + where + step s@(Chunk xs) = idone (Just $ LL.length xs) s + step stream = idone Nothing stream +{-# INLINE chunkLength #-} + +-- | Take @n@ elements from the current chunk, or the whole chunk if +-- @n@ is greater. +takeFromChunk :: + (Nullable s, LL.ListLike s el) + => Int + -> Iteratee s m s +takeFromChunk n | n <= 0 = idone emptyP (Chunk emptyP) +takeFromChunk n = liftI step + where + step (Chunk xs) = let (h,t) = LL.splitAt n xs in idone h $ Chunk t + step stream = idone emptyP stream +{-# INLINE takeFromChunk #-} + +-- |Takes an element predicate and returns the (possibly empty) prefix of +-- the stream. None of the characters in the string satisfy the character +-- predicate. +-- If the stream is not terminated, the first character of the remaining stream +-- satisfies the predicate. +-- +-- N.B. 'breakE' should be used in preference to @breakStream@. +-- @breakStream@ will retain all data until the predicate is met, which may +-- result in a space leak. +-- +-- The analogue of @List.break@ + +breakStream :: (LL.ListLike s el) => (el -> Bool) -> Iteratee s m s +breakStream cpred = icont (step mempty) Nothing + where + step bfr (Chunk str) + | LL.null str = icont (step bfr) Nothing + | otherwise = case LL.break cpred str of + (str', tail') + | LL.null tail' -> icont (step (bfr `mappend` str)) Nothing + | otherwise -> idone (bfr `mappend` str') (Chunk tail') + step bfr stream = idone bfr stream +{-# INLINE breakStream #-} + +-- --------------------------------------------------- +-- The converters show a different way of composing two iteratees: +-- `vertical' rather than `horizontal' + +-- |Takes an element predicate and an iteratee, running the iteratee +-- on all elements of the stream until the predicate is met. +-- +-- the following rule relates @break@ to @breakE@ +-- @break@ pred === @joinI@ (@breakE@ pred stream2stream) +-- +-- @breakE@ should be used in preference to @break@ whenever possible. +breakE + :: (LL.ListLike s el, NullPoint s) + => (el -> Bool) + -> Enumeratee s s m a +breakE cpred = eneeCheckIfDonePass (icont . step) + where + step k (Chunk s) + | LL.null s = liftI (step k) + | otherwise = case LL.break cpred s of + (str', tail') + | LL.null tail' -> eneeCheckIfDonePass (icont . step) . k $ Chunk str' + | otherwise -> idone (k $ Chunk str') (Chunk tail') + step k stream = idone (liftI k) stream +{-# INLINE breakE #-} + +-- |Read n elements from a stream and apply the given iteratee to the +-- stream of the read elements. Unless the stream is terminated early, we +-- read exactly n elements, even if the iteratee has accepted fewer. +-- +-- The analogue of @List.take@ +takeStream :: + (Monad m, Nullable s, LL.ListLike s el) + => Int -- ^ number of elements to consume + -> Enumeratee s s m a +takeStream n' iter + | n' <= 0 = return iter + | otherwise = Iteratee $ \od oc -> runIter iter (on_done od oc) (on_cont od oc) + where + on_done od oc x _ = runIter (dropStream n' >> return (return x)) od oc + on_cont od oc k Nothing = if n' == 0 then od (liftI k) (Chunk mempty) + else runIter (liftI (step n' k)) od oc + on_cont od oc _ (Just e) = runIter (dropStream n' >> throwErr e) od oc + step n k (Chunk str) + | LL.null str = liftI (step n k) + | LL.length str <= n = takeStream (n - LL.length str) $ k (Chunk str) + | otherwise = idone (k (Chunk s1)) (Chunk s2) + where (s1, s2) = LL.splitAt n str + step _n k stream = idone (liftI k) stream +{-# INLINE takeStream #-} + +-- |Read n elements from a stream and apply the given iteratee to the +-- stream of the read elements. If the given iteratee accepted fewer +-- elements, we stop. +-- This is the variation of 'takeStream' with the early termination +-- of processing of the outer stream once the processing of the inner stream +-- finished early. +-- +-- Iteratees composed with 'takeUpTo' will consume only enough elements to +-- reach a done state. Any remaining data will be available in the outer +-- stream. +-- +-- > > let iter = do +-- > h <- joinI $ takeUpTo 5 I.head +-- > t <- stream2list +-- > return (h,t) +-- > +-- > > enumPureNChunk [1..10::Int] 3 iter >>= run >>= print +-- > (1,[2,3,4,5,6,7,8,9,10]) +-- > +-- > > enumPureNChunk [1..10::Int] 7 iter >>= run >>= print +-- > (1,[2,3,4,5,6,7,8,9,10]) +-- +-- in each case, @I.head@ consumes only one element, returning the remaining +-- 4 elements to the outer stream +takeUpTo :: (Monad m, Nullable s, LL.ListLike s el) => Int -> Enumeratee s s m a +takeUpTo i iter + | i <= 0 = idone iter (Chunk emptyP) + | otherwise = Iteratee $ \od oc -> + runIter iter (onDone od oc) (onCont od oc) + where + onDone od oc x str = runIter (idone (return x) str) od oc + onCont od oc k Nothing = if i == 0 then od (liftI k) (Chunk mempty) + else runIter (liftI (step i k)) od oc + onCont od oc _ (Just e) = runIter (throwErr e) od oc + step n k (Chunk str) + | LL.null str = liftI (step n k) + | LL.length str < n = takeUpTo (n - LL.length str) $ k (Chunk str) + | otherwise = + -- check to see if the inner iteratee has completed, and if so, + -- grab any remaining stream to put it in the outer iteratee. + -- the outer iteratee is always complete at this stage, although + -- the inner may not be. + let (s1, s2) = LL.splitAt n str + in Iteratee $ \od' _ -> do + res <- runIter (k (Chunk s1)) (\a s -> return $ Left (a, s)) + (\k' e -> return $ Right (k',e)) + case res of + Left (a,Chunk s1') -> od' (return a) + (Chunk $ s1' `LL.append` s2) + Left (a,s') -> od' (idone a s') (Chunk s2) + Right (k',e) -> od' (icont k' e) (Chunk s2) + step _ k stream = idone (liftI k) stream +{-# INLINE takeUpTo #-} + + +-- |Takes an element predicate and an iteratee, running the iteratee +-- on all elements of the stream while the predicate is met. +-- +-- This is preferred to @takeWhile@. +takeWhileE + :: (LL.ListLike s el, NullPoint s) + => (el -> Bool) + -> Enumeratee s s m a +takeWhileE = breakE . (not .) +{-# INLINEABLE takeWhileE #-} + +-- | Map a function over an 'Iteratee'. +-- This one is reimplemented and differs from the the one in +-- "Data.Iteratee.ListLike" in so far that it doesn't pass on an 'EOF' +-- received in the input, which is the expected behavior. +{-# INLINE mapStream #-} +mapStream :: (LL.ListLike (s el) el, LL.ListLike (s el') el', NullPoint (s el)) + => (el -> el') -> Enumeratee (s el) (s el') m a +mapStream = mapChunks . LL.map + +-- |Map the stream rigidly. +-- +-- Like 'mapStream', but the element type cannot change. +-- This function is necessary for @ByteString@ and similar types +-- that cannot have 'LooseMap' instances, and may be more efficient. +rigidMapStream + :: (LL.ListLike s el, NullPoint s) + => (el -> el) + -> Enumeratee s s m a +rigidMapStream f = mapChunks (LL.rigidMap f) +{-# SPECIALIZE rigidMapStream :: (el -> el) -> Enumeratee [el] [el] m a #-} +{-# SPECIALIZE rigidMapStream :: (Word8 -> Word8) -> Enumeratee B.ByteString B.ByteString m a #-} + + +-- | Apply a function to the elements of a stream, concatenate the +-- results into a stream. No giant intermediate list is produced. +{-# INLINE concatMapStream #-} +concatMapStream :: (Monad m, LL.ListLike s a, NullPoint s) => (a -> t) -> Enumeratee s t m r +concatMapStream f = eneeCheckIfDone (liftI . go) + where + go k (EOF mx) = idone (liftI k) (EOF mx) + go k (Chunk xs) | LL.null xs = liftI (go k) + | otherwise = eneeCheckIfDone (flip go (Chunk (LL.tail xs))) . k . Chunk . f $ LL.head xs + +-- | Apply a monadic function to the elements of a stream, concatenate +-- the results into a stream. No giant intermediate list is produced. +{-# INLINE concatMapStreamM #-} +concatMapStreamM :: (Monad m, LL.ListLike s a, NullPoint s) => (a -> m t) -> Enumeratee s t m r +concatMapStreamM f = eneeCheckIfDone (liftI . go) + where + go k (EOF mx) = idone (liftI k) (EOF mx) + go k (Chunk xs) | LL.null xs = liftI (go k) + | otherwise = f (LL.head xs) `mBind` + eneeCheckIfDone (flip go (Chunk (LL.tail xs))) . k . Chunk + +{-# INLINE mapMaybeStream #-} +mapMaybeStream :: (LL.ListLike s a, NullPoint s, LL.ListLike t b) => (a -> Maybe b) -> Enumeratee s t m r +mapMaybeStream f = mapChunks mm + where + mm l = if LL.null l then LL.empty else + case f (LL.head l) of Nothing -> mm (LL.tail l) + Just b -> LL.cons b $ mm (LL.tail l) + + +-- |Creates an 'enumeratee' with only elements from the stream that +-- satisfy the predicate function. The outer stream is completely consumed. +-- +-- The analogue of @List.filter@ +-- XXX filterStream :: (ListLike s a, NullPoint s) => (a -> Bool) -> Enumeratee s s m r +filterStream + :: (NullPoint s, LL.ListLike s el) + => (el -> Bool) + -> Enumeratee s s m a +filterStream p = mapChunks (LL.filter p) +{-# INLINE filterStream #-} + +-- | Apply a monadic filter predicate to an 'Iteratee'. +{-# INLINE filterStreamM #-} +filterStreamM :: (Monad m, LL.ListLike s a, Nullable s) => (a -> m Bool) -> Enumeratee s s m r +filterStreamM k = mapChunksM (go id) + where + go acc s | LL.null s = return $! acc LL.empty + | otherwise = do p <- k (LL.head s) + let acc' = if p then LL.cons (LL.head s) . acc else acc + go acc' (LL.tail s) + +-- | Grouping on 'Iteratee's. @groupStreamOn proj inner outer@ executes +-- @inner (proj e)@, where @e@ is the first input element, to obtain an +-- 'Iteratee' @i@, then passes elements @e@ to @i@ as long as @proj e@ +-- produces the same result. If @proj e@ changes or the input ends, the +-- pair of @proj e@ and the result of @run i@ is passed to @outer@. At +-- end of input, the resulting @outer@ is returned. +groupStreamOn :: (Monad m, LL.ListLike l e, Eq t1, Nullable l) + => (e -> t1) + -> (t1 -> m (Iteratee l m t2)) + -> Enumeratee l [(t1, t2)] m a +groupStreamOn proj inner = eneeCheckIfDonePass (icont . step) + where + step outer (EOF mx) = idone (liftI outer) $ EOF mx + step outer c@(Chunk as) + | LL.null as = liftI $ step outer + | otherwise = let x = proj (LL.head as) + in lift (inner x) >>= \i -> step' x i outer c + + -- We want to feed a 'Chunk' to the inner 'Iteratee', which might be + -- finished. In that case, we would want to abort, but we cannot, + -- since the outer iteration is still going on. So instead we + -- discard data we would have fed to the inner 'Iteratee'. (Use of + -- 'enumPure1Chunk' is not appropriate, it would accumulate the + -- data, just to have it discarded by the 'run' that eventually + -- happens. + + step' c it outer (Chunk as) + | LL.null as = liftI $ step' c it outer + | (l,r) <- LL.span ((==) c . proj) as, not (LL.null l) = + let od a _str = idoneM a $ EOF Nothing + oc k Nothing = return $ k (Chunk l) + oc k m = icontM k m + in lift (runIter it od oc) >>= \it' -> step' c it' outer (Chunk r) + + step' c it outer str = + lift (run it) >>= \b -> eneeCheckIfDone (`step` str) . outer $ Chunk [(c,b)] + + +-- | Grouping on 'Iteratee's. @groupStreamBy cmp inner outer@ executes +-- @inner@ to obtain an 'Iteratee' @i@, then passes elements @e@ to @i@ +-- as long as @cmp e0 e@, where @e0@ is some preceeding element, is +-- true. Else, the result of @run i@ is passed to @outer@ and +-- 'groupStreamBy' restarts. At end of input, the resulting @outer@ is +-- returned. +groupStreamBy :: (Monad m, LL.ListLike l t, Nullable l) + => (t -> t -> Bool) + -> m (Iteratee l m t2) + -> Enumeratee l [t2] m a +groupStreamBy cmp inner = eneeCheckIfDonePass (icont . step) + where + step outer (EOF mx) = idone (liftI outer) $ EOF mx + step outer c@(Chunk as) + | LL.null as = liftI $ step outer + | otherwise = lift inner >>= \i -> step' (LL.head as) i outer c + + step' c it outer (Chunk as) + | LL.null as = liftI $ step' c it outer + | (l,r) <- LL.span (cmp c) as, not (LL.null l) = + let od a _str = idoneM a $ EOF Nothing + oc k Nothing = return $ k (Chunk l) + oc k m = icontM k m + in lift (runIter it od oc) >>= \it' -> step' (LL.head l) it' outer (Chunk r) + + step' _ it outer str = + lift (run it) >>= \b -> eneeCheckIfDone (`step` str) . outer $ Chunk [b] + + +-- | @mergeStreams@ offers another way to nest iteratees: as a monad stack. +-- This allows for the possibility of interleaving data from multiple +-- streams. +-- +-- > -- print each element from a stream of lines. +-- > logger :: (MonadIO m) => Iteratee [ByteString] m () +-- > logger = mapStreamM_ (liftIO . putStrLn . B.unpack) +-- > +-- > -- combine alternating lines from two sources +-- > -- To see how this was derived, follow the types from +-- > -- 'ileaveLines logger' and work outwards. +-- > run =<< enumFile 10 "file1" (joinI $ enumLinesBS $ +-- > ( enumFile 10 "file2" . joinI . enumLinesBS $ joinI +-- > (ileaveLines logger)) >>= run) +-- > +-- > ileaveLines :: (Functor m, Monad m) +-- > => Enumeratee [ByteString] [ByteString] (Iteratee [ByteString] m) +-- > [ByteString] +-- > ileaveLines = mergeStreams (\l1 l2 -> +-- > [B.pack "f1:\n\t" ,l1 ,B.pack "f2:\n\t" ,l2 ] +-- > +-- > +-- +mergeStreams :: + (LL.ListLike s1 el1 + ,LL.ListLike s2 el2 + ,Nullable s1 + ,Nullable s2 + ,Monad m) + => (el1 -> el2 -> b) + -> Enumeratee s2 b (Iteratee s1 m) a +mergeStreams f = convStream $ liftM2 f (lift headStream) headStream +{-# INLINE mergeStreams #-} + +-- | A version of mergeStreams which operates on chunks instead of +-- elements. +-- +-- mergeByChunks offers more control than 'mergeStreams'. +-- 'mergeStreams' terminates when the first stream terminates, however +-- mergeByChunks will continue until both streams are exhausted. +-- +-- 'mergeByChunks' guarantees that both chunks passed to the merge +-- function will have the same number of elements, although that number +-- may vary between calls. +mergeByChunks :: + (Nullable c2, Nullable c1 + ,LL.ListLike c1 el1, LL.ListLike c2 el2 + , Monad m) + => (c1 -> c2 -> c3) -- ^ merge function + -> (c1 -> c3) + -> (c2 -> c3) + -> Enumeratee c2 c3 (Iteratee c1 m) a +mergeByChunks f f1 f2 = unfoldConvStream iter (0 :: Int) + where + iter 1 = (\x -> (1,f1 x)) `liftM` lift getChunk + iter 2 = (\x -> (2,f2 x)) `liftM` getChunk + iter _ = do + ml1 <- lift chunkLength + ml2 <- chunkLength + case (ml1, ml2) of + (Just l1, Just l2) -> do + let tval = min l1 l2 + c1 <- lift $ takeFromChunk tval + c2 <- takeFromChunk tval + return (0, f c1 c2) + (Just _, Nothing) -> iter 1 + (Nothing, _) -> iter 2 +{-# INLINE mergeByChunks #-} + +-- ------------------------------------------------------------------------ +-- Folds + +-- | Left-associative fold that is strict in the accumulator. +-- This function should be used in preference to 'foldl' whenever possible. +-- +-- The analogue of @List.foldl'@. +foldStream + :: LL.ListLike s el + => (a -> el -> a) + -> a + -> Iteratee s m a +foldStream f i = liftI (step i) + where + step acc (Chunk xs) + | LL.null xs = liftI (step acc) + | otherwise = liftI (step $! FLL.foldl' f acc xs) + step acc stream = idone acc stream +{-# INLINE foldStream #-} + +-- ------------------------------------------------------------------------ +-- Zips + +-- |Enumerate two iteratees over a single stream simultaneously. +-- +-- Compare to @List.zip@. +zipStreams + :: (Monad m, Nullable s, LL.ListLike s el) + => Iteratee s m a + -> Iteratee s m b + -> Iteratee s m (a, b) +zipStreams x0 y0 = do + -- need to check if both iteratees are initially finished. If so, + -- we don't want to push a chunk which will be dropped + (a', x') <- lift $ runIter x0 od oc + (b', y') <- lift $ runIter y0 od oc + case checkDone a' b' of + Just (Right (a,b,s)) -> idone (a,b) s -- 's' may be EOF, needs to stay + Just (Left (Left a)) -> liftM (a,) y' + Just (Left (Right b)) -> liftM (,b) x' + Nothing -> liftI (step x' y') + where + step x y (Chunk xs) | nullC xs = liftI (step x y) + step x y (Chunk xs) = do + (a', x') <- lift $ (\i -> runIter i od oc) =<< enumPure1Chunk xs x + (b', y') <- lift $ (\i -> runIter i od oc) =<< enumPure1Chunk xs y + case checkDone a' b' of + Just (Right (a,b,s)) -> idone (a,b) s + Just (Left (Left a)) -> liftM (a,) y' + Just (Left (Right b)) -> liftM (,b) x' + Nothing -> liftI (step x' y') + step x y (EOF err) = joinIM $ case err of + Nothing -> (liftM2.liftM2) (,) (enumEof x) (enumEof y) + Just e -> (liftM2.liftM2) (,) (enumErr e x) (enumErr e y) + + od a s = return (Just (a, s), idone a s) + oc k e = return (Nothing , icont k e) + + checkDone r1 r2 = case (r1, r2) of + (Just (a, s1), Just (b,s2)) -> Just $ Right (a, b, shorter s1 s2) + (Just (a, _), Nothing) -> Just . Left $ Left a + (Nothing, Just (b, _)) -> Just . Left $ Right b + (Nothing, Nothing) -> Nothing + + shorter c1@(Chunk xs) c2@(Chunk ys) + | LL.length xs < LL.length ys = c1 + | otherwise = c2 + shorter e@(EOF _) _ = e + shorter _ e@(EOF _) = e +{-# INLINE zipStreams #-} + +zipStreams3 + :: (Monad m, Nullable s, LL.ListLike s el) + => Iteratee s m a -> Iteratee s m b + -> Iteratee s m c -> Iteratee s m (a, b, c) +zipStreams3 a b c = zipStreams a (zipStreams b c) >>= + \(r1, (r2, r3)) -> return (r1, r2, r3) +{-# INLINE zipStreams3 #-} + +zipStreams4 + :: (Monad m, Nullable s, LL.ListLike s el) + => Iteratee s m a -> Iteratee s m b + -> Iteratee s m c -> Iteratee s m d + -> Iteratee s m (a, b, c, d) +zipStreams4 a b c d = zipStreams a (zipStreams3 b c d) >>= + \(r1, (r2, r3, r4)) -> return (r1, r2, r3, r4) +{-# INLINE zipStreams4 #-} + +zipStreams5 + :: (Monad m, Nullable s, LL.ListLike s el) + => Iteratee s m a -> Iteratee s m b + -> Iteratee s m c -> Iteratee s m d + -> Iteratee s m e -> Iteratee s m (a, b, c, d, e) +zipStreams5 a b c d e = zipStreams a (zipStreams4 b c d e) >>= + \(r1, (r2, r3, r4, r5)) -> return (r1, r2, r3, r4, r5) +{-# INLINE zipStreams5 #-} + +-- | Enumerate over two iteratees in parallel as long as the first iteratee +-- is still consuming input. The second iteratee will be terminated with EOF +-- when the first iteratee has completed. An example use is to determine +-- how many elements an iteratee has consumed: +-- +-- > snd <$> enumWith (dropWhile (<5)) length +-- +-- Compare to @zipStreams@ +enumWith + :: (Monad m, Nullable s, LL.ListLike s el) + => Iteratee s m a + -> Iteratee s m b + -> Iteratee s m (a, b) +enumWith i1 i2 = do + -- as with zipStreams, first check to see if the initial iteratee is complete, + -- otherwise data would be dropped. + -- running the second iteratee as well to prevent a monadic effect mismatch + -- although I think that would be highly unlikely to happen in common + -- code + (a', x') <- lift $ runIter i1 od oc + (_, y') <- lift $ runIter i2 od oc + case a' of + Just (a, s) -> flip idone s =<< lift (liftM (a,) $ run i2) + Nothing -> go x' y' + where + od a s = return (Just (a, s), idone a s) + oc k e = return (Nothing , icont k e) + + getUsed xs (Chunk ys) = LL.take (LL.length xs - LL.length ys) xs + getUsed xs (EOF _) = xs + + go x y = liftI step + where + step (Chunk xs) | nullC xs = liftI step + step (Chunk xs) = do + (a', x') <- lift $ (\i -> runIter i od oc) =<< enumPure1Chunk xs x + case a' of + Just (a, s) -> do + b <- lift $ run =<< enumPure1Chunk (getUsed xs s) y + idone (a, b) s + Nothing -> lift (enumPure1Chunk xs y) >>= go x' + step (EOF err) = joinIM $ case err of + Nothing -> (liftM2.liftM2) (,) (enumEof x) (enumEof y) + Just e -> (liftM2.liftM2) (,) (enumErr e x) (enumErr e y) +{-# INLINE enumWith #-} + +-- |Enumerate a list of iteratees over a single stream simultaneously +-- and discard the results. This is a different behavior than Prelude's +-- sequence_ which runs iteratees in the list one after the other. +-- +-- Compare to @Prelude.sequence_@. +sequenceStreams_ + :: (Monad m, LL.ListLike s el, Nullable s) + => [Iteratee s m a] + -> Iteratee s m () +sequenceStreams_ = self + where + self is = liftI step + where + step (Chunk xs) | LL.null xs = liftI step + step s@(Chunk _) = do + -- give a chunk to each iteratee + is' <- lift $ mapM (enumChunk s) is + -- filter done iteratees + (done, notDone) <- lift $ partition fst `liftM` mapM enumCheckIfDone is' + if null notDone + then idone () <=< remainingStream $ map snd done + else self $ map snd notDone + step s@(EOF _) = do + s' <- remainingStream <=< lift $ mapM (enumChunk s) is + case s' of + EOF (Just e) -> throwErr e + _ -> idone () s' + + -- returns the unconsumed part of the stream; "sequenceStreams_ is" consumes as + -- much of the stream as the iteratee in is that consumes the most; e.g. + -- sequenceStreams_ [I.head, I.last] consumes whole stream + remainingStream + :: (Monad m, Nullable s, LL.ListLike s el) + => [Iteratee s m a] -> Iteratee s m (Stream s) + remainingStream is = lift $ + return . foldl1 shorter <=< mapM (\i -> runIter i od oc) $ is + where + od _ s = return s + oc _ e = return $ case e of + Nothing -> mempty + _ -> EOF e + + -- return the shorter one of two streams; errors are propagated with the + -- priority given to the "left" + shorter c1@(Chunk xs) c2@(Chunk ys) + | LL.length xs < LL.length ys = c1 + | otherwise = c2 + shorter (EOF e1 ) (EOF e2 ) = EOF (e1 `mplus` e2) + shorter e@(EOF _) _ = e + shorter _ e@(EOF _) = e + +-- |Transform an iteratee into one that keeps track of how much data it +-- consumes. +countConsumed :: forall a s el m n. + (Monad m, LL.ListLike s el, Nullable s, Integral n) => + Iteratee s m a + -> Iteratee s m (a, n) +countConsumed i = go 0 (const i) (Chunk emptyP) + where + go :: n -> (Stream s -> Iteratee s m a) -> Stream s + -> Iteratee s m (a, n) + go !n f str@(EOF _) = (, n) `liftM` f str + go !n f str@(Chunk c) = Iteratee rI + where + newLen = n + fromIntegral (LL.length c) + rI od oc = runIter (f str) onDone onCont + where + onDone a str'@(Chunk c') = + od (a, newLen - fromIntegral (LL.length c')) str' + onDone a str'@(EOF _) = od (a, n) str' + onCont f' mExc = oc (go newLen f') mExc +{-# INLINE countConsumed #-} + +-- ------------------------------------------------------------------------ +-- Enumerators + +-- |The pure n-chunk enumerator +-- It passes a given stream of elements to the iteratee in @n@-sized chunks. +enumPureNChunk :: (Monad m, LL.ListLike s el) => s -> Int -> Enumerator s m a +enumPureNChunk str n iter + | LL.null str = return iter + | n > 0 = enum' str iter + | otherwise = error $ "enumPureNChunk called with n==" ++ show n + where + enum' str' iter' + | LL.null str' = return iter' + | otherwise = let (s1, s2) = LL.splitAt n str' + on_cont k Nothing = enum' s2 . k $ Chunk s1 + on_cont k e = return $ icont k e + in runIter iter' idoneM on_cont +{-# INLINE enumPureNChunk #-} + +-- ------------------------------------------------------------------------ +-- Monadic functions + +-- | Map a monadic function over the elements of the stream and ignore the +-- result. +mapStreamM_ :: (Monad m, Nullable s, LL.ListLike s el) => (el -> m b) -> Iteratee s m () +mapStreamM_ = mapChunksM_ . LL.mapM_ +{-# INLINE mapStreamM_ #-} + +-- | Map a monadic function over an 'Iteratee'. +mapStreamM :: (Monad m, LL.ListLike (s el) el, LL.ListLike (s el') el', NullPoint (s el)) + => (el -> m el') -> Enumeratee (s el) (s el') m a +mapStreamM = mapChunksM . LL.mapM +{-# INLINE mapStreamM #-} + + +-- | Fold a monadic function over an 'Iteratee'. +foldStreamM :: (Monad m, Nullable s, LL.ListLike s a) => (b -> a -> m b) -> b -> Iteratee s m b +foldStreamM k = foldChunksM go + where + go b s | LL.null s = return b + | otherwise = k b (LL.head s) >>= \b' -> go b' (LL.tail s) +{-# INLINE foldStreamM #-} diff --git a/src/Bio/Iteratee/ReadableChunk.hs b/src/Bio/Iteratee/ReadableChunk.hs new file mode 100644 --- /dev/null +++ b/src/Bio/Iteratee/ReadableChunk.hs @@ -0,0 +1,49 @@ +{-# LANGUAGE FunctionalDependencies #-} + +-- | Monadic Iteratees: +-- incremental input parsers, processors and transformers +-- +-- Support for IO enumerators + +module Bio.Iteratee.ReadableChunk ( ReadableChunk(..) ) where + +import Control.Monad.IO.Class +import Data.Word +import Foreign.C +import Foreign.Marshal.Array +import Foreign.Ptr +import Foreign.Storable +import Prelude + +import qualified Data.ByteString as B +import qualified Data.ByteString.Lazy as L + +-- |Class of streams which can be filled from a 'Ptr'. Typically these +-- are streams which can be read from a file, @Handle@, or similar resource. +-- +-- +class (Storable el) => ReadableChunk s el | s -> el where + readFromPtr :: + MonadIO m => + Ptr el + -> Int -- ^ The pointer must not be used after @readFromPtr@ completes. + -> m s -- ^ The Int parameter is the length of the data in *bytes*. + +instance ReadableChunk [Char] Char where + readFromPtr buf l = liftIO $ peekCAStringLen (castPtr buf, l) + +instance ReadableChunk [Word8] Word8 where + readFromPtr buf l = liftIO $ peekArray l buf +instance ReadableChunk [Word16] Word16 where + readFromPtr buf l = liftIO $ peekArray l buf +instance ReadableChunk [Word32] Word32 where + readFromPtr buf l = liftIO $ peekArray l buf +instance ReadableChunk [Word] Word where + readFromPtr buf l = liftIO $ peekArray l buf + +instance ReadableChunk B.ByteString Word8 where + readFromPtr buf l = liftIO $ B.packCStringLen (castPtr buf, l) + +instance ReadableChunk L.ByteString Word8 where + readFromPtr buf l = liftIO $ + return . L.fromChunks . (:[]) =<< readFromPtr buf l diff --git a/src/Bio/Util/Zlib.hs b/src/Bio/Util/Zlib.hs --- a/src/Bio/Util/Zlib.hs +++ b/src/Bio/Util/Zlib.hs @@ -1,3 +1,4 @@ +{-# LANGUAGE CPP #-} module Bio.Util.Zlib ( decompressGzip ) where import Prelude diff --git a/src/cbits/loops.c b/src/cbits/loops.c new file mode 100644 --- /dev/null +++ b/src/cbits/loops.c @@ -0,0 +1,101 @@ +void nuc_loop( char* p, int stride, char* q, int u, int v ) +{ + u *= stride ; + v *= stride ; + + while( u < v ) { + char a = q[ u ] ; + char b = q[ u + stride ] ; + char a1 = a ? 0x10 << (a&3) : 0xf0 ; + char b1 = b ? 0x1 << (b&3) : 0xf ; + *p++ = a1 | b1 ; + u += stride+stride ; + } + if( u == v ) { + char a = q[ u ] ; + char a1 = a ? 0x10 << (a&3) : 0xf0 ; + *p = a1 ; + } +} + +void nuc_loop_asc( char* p, int stride, char* q, int u, int v ) +{ + u *= stride ; + v *= stride ; + + while( u <= v ) { + char a = q[ u ] ; + *p++ = a == 0 ? 'N' : (a&3) == 0 ? 'A' : (a&3) == 1 ? 'C' : (a&3) == 2 ? 'G' : 'T' ; + u += stride ; + } + *p = 0 ; +} + +void nuc_loop_asc_rev( char* p, int stride, char* q, int u, int v ) +{ + u *= stride ; + v *= stride ; + + while( u <= v ) { + char a = q[ v ] ; + *p++ = a == 0 ? 'N' : (a&3) == 0 ? 'T' : (a&3) == 1 ? 'G' : (a&3) == 2 ? 'C' : 'A' ; + v -= stride ; + } + *p = 0 ; +} + +void qual_loop( char* p, int stride, char* q, int u, int v ) +{ + u *= stride ; + v *= stride ; + while( u <= v ) { + *p++ = (q[u] >> 2) & 0x3f ; + u += stride ; + } +} + +void qual_loop_asc( char* p, int stride, char* q, int u, int v ) +{ + u *= stride ; + v *= stride ; + while( u <= v ) { + *p++ = 33 + ((q[u] >> 2) & 0x3f) ; + u += stride ; + } + *p = 0 ; +} + +void qual_loop_asc_rev( char* p, int stride, char* q, int u, int v ) +{ + u *= stride ; + v *= stride ; + while( u <= v ) { + *p++ = 33 + ((q[v] >> 2) & 0x3f) ; + v -= stride ; + } + *p = 0 ; +} + +int int_loop( char* p, int x ) +{ + *p++ = ':' ; + if( x == 0 ) { + *p = '0' ; + return 2 ; + } + char *q = p ; + while( x > 0 ) { + *q++ = '0' + x % 10 ; + x /= 10 ; + } + int r = q-p ; + --q ; + while( p < q ) { + char c = *p ; + *p++ = *q ; + *q-- = c ; + } + return r+1 ; +} + +