diff --git a/Biobase/Fasta/Streaming.hs b/Biobase/Fasta/Streaming.hs --- a/Biobase/Fasta/Streaming.hs +++ b/Biobase/Fasta/Streaming.hs @@ -3,6 +3,22 @@ -- -- The functions in here should be streaming in constant memory. -- +-- A typical, slightly complicated is this: +-- @ +-- forEach ∷ forall r . Stream (ByteString m) m r → m (Stream (Of ()) m r) +-- forEach dna = do +-- -- extract the header, but at most 123 characters, dropping the rest +-- hdr SP.:> dta ← extractHeader (Just 123) dna +-- -- create windows @ws@ of a particular type. Include the prefix, the suffix, and make each window 10 characters long +-- let ws = (streamedWindows True True (Just 10) (SequenceIdentifier hdr) PlusStrand dta :: SP.Stream (SP.Of (BioSequenceWindow "DNA" DNA 0)) m r) +-- -- count the number of characters in @dna@, get the return value, print each window +-- count SP.:> r ← SP.mapM_ (liftIO . print) . bswSeqLength $ SP.copy ws +-- liftIO $ print count +-- liftIO $ putStrLn "" +-- -- yield one vacuous @()@ result, return the remainder @r@ from dna. +-- return $ SP.yield () *> return r +-- @ +-- -- TODO Check if this is actually true with some unit tests. module Biobase.Fasta.Streaming @@ -17,7 +33,7 @@ import Data.ByteString.Streaming.Internal (ByteString(..)) import Data.Semigroup as SG import Debug.Trace -import GHC.Generics +import GHC.Generics (Generic) import GHC.TypeLits import Prelude as P import qualified Data.ByteString.Char8 as BS @@ -25,8 +41,11 @@ import Streaming as S import Streaming.Prelude as SP +import Data.ByteString.Streaming.Split + import Biobase.Types.BioSequence import Biobase.Types.Index.Type +import Biobase.Types.Location import Biobase.Types.Strand @@ -73,7 +92,7 @@ -- ^ The size of each window to be processed. → ByteString m r -- ^ A streaming bytestring of Fasta files. - → Stream (Of (BioSequenceWindow w ty k)) m r + → Stream (Of (BioSequenceWindow w ty PartialLocation)) m r -- ^ The outgoing stream of @Current@ windows being processed. {-# Inlinable streamingFasta #-} streamingFasta (HeaderSize hSz) (OverlapSize oSz) (CurrentSize cSz) = go (FindHeader [] 0) where @@ -153,9 +172,149 @@ , _bswPrefix = BioSequence pfx , _bswSequence = BioSequence seq , _bswSuffix = BioSequence BS.empty - , _bswStrand = PlusStrand - , _bswIndex = Index $ entries * cSz + , _bswLocation = PartialLocation PlusStrand (Index $ entries * cSz) (BS.length seq) +-- , _bswStrand = PlusStrand +-- , _bswIndex = Index $ entries * cSz } + +-- | + +streamedFasta ∷ (Monad m) ⇒ ByteString m r → Stream (Stream (ByteString m) m) m r +{-# Inlinable streamedFasta #-} +streamedFasta = S.maps (collapseData) . streamOfStreamedFasta + +-- | Here each individual fasta file will be a stream. +-- +-- TODO Once this works, @streamingFasta@ should be @S.concats . streamOfStreamedFasta@ ... + +streamOfStreamedFasta + ∷ forall m r + . ( Monad m ) + ⇒ ByteString m r + → Stream (Stream (ByteString m) m) m r + -- ^ +{-# Inlinable streamOfStreamedFasta #-} +streamOfStreamedFasta = go . S8.lines where + go = \case + SI.Return r → SI.Return r + SI.Effect m → SI.Effect (fmap go m) + SI.Step fs → SI.Step (SI.Step (fmap (fmap go . splitFasta) fs)) + +-- | Given a 'Stream (ByteString m) m r' which is a 'Stream' of @lines@, split +-- off the first @Fasta@ entry. + +splitFasta ∷ (Monad m) ⇒ Stream (ByteString m) m r → Stream (ByteString m) m (Stream (ByteString m) m r) +{-# Inlinable splitFasta #-} +splitFasta = loop False where + loop hdr = \case + SI.Return r → SI.Return (SI.Return r) + SI.Effect m → SI.Effect (fmap (loop hdr) m) + SI.Step bs → case bs of + Empty r → loop hdr r + Chunk cs xs + | BS.null cs → loop hdr $ SI.Step xs + | h=='>' || h==';' → if hdr then SI.Return (SI.Step bs) else SI.Step $ fmap (loop True) bs + | otherwise → SI.Step $ fmap (loop True) bs + where h = BS.head cs + Go m → SI.Effect $ fmap ((loop hdr) . SI.Step) m + +-- | Given a stream, roughly like @[BS "Header", BS "Data1", BS "Data2", ...]@ +-- create a stream like @[BS "Header", BS "Data"]@. This means that the +-- resulting stream holds exactly two @ByteString@'s. + +collapseData ∷ (Monad m) ⇒ Stream (ByteString m) m r → Stream (ByteString m) m r +{-# Inlinable collapseData #-} +collapseData = loop where + loop = \case + SI.Return r → SI.Return r + SI.Effect m → SI.Effect (fmap loop m) + SI.Step bs → case bs of + Empty r → loop r + Chunk cs xs + | BS.null cs → loop $ SI.Step xs + | h=='>' || h==';' → SI.Step $ fmap (S.yields . S8.concat) bs + | otherwise → SI.Step $ fmap loop bs + where h = BS.head cs + Go m → SI.Effect $ fmap (loop . SI.Step) m + +-- | "Rechunk" a stream of bytestrings. + +reChunkBS ∷ (Monad m) ⇒ Int → Stream (ByteString m) m r → Stream (ByteString m) m r +{-# Inlinable reChunkBS #-} +reChunkBS n = splitsByteStringAt n . S8.concat + +-- | Assuming a "rechunked" stream of bytestrings, create sequence windows. + +chunksToWindows ∷ (Monad m) ⇒ SequenceIdentifier w → Strand → Stream (ByteString m) m r → Stream (Of (BioSequenceWindow w ty PartialLocation)) m r +{-# Inlinable chunksToWindows #-} +chunksToWindows seqId s = SP.map go . SP.drop 1 . SP.scan indexed (BS.empty, 0, 0) (\(bs,i,_) → (bs,i)) . S.mapsM S8.toStrict where + indexed (_,cur,next) bs = (bs,next,next + BS.length bs) + go (bs,i) + = BioSequenceWindow + { _bswIdentifier = seqId + , _bswPrefix = BioSequence "" + , _bswSequence = BioSequence bs + , _bswSuffix = BioSequence "" + , _bswLocation = PartialLocation s (Index i) (BS.length bs) +-- , _bswStrand = s +-- , _bswIndex = Index i + } + +-- | Make it possible to take a fasta stream and produce a stream of +-- 'BioSequenceWindow's. This is a convenience function around +-- 'withSuffix . withPrefix . chunksToWindows . reChunks'. +-- +-- In case of a @Nothing@ window size, a single huge @Fasta@ entry is produced +-- (and materialized!). +-- +-- TODO In case of @Nothing@ window size, we use the 'collapseData' function +-- which has one check too many, and will be slightly slower. However, the +-- check should be once per @ByteString@. + +streamedWindows + ∷ (Monad m) + ⇒ Bool + → Bool + → Maybe Int + -- ^ desired size or a single huge @Fasta@ entry. + → SequenceIdentifier w + → Strand + → (Stream (ByteString m) m) r + → Stream (Of (BioSequenceWindow w ty PartialLocation)) m r +{-# Inlinable streamedWindows #-} +streamedWindows withPrefix withSuffix winSz seqId strnd + = (if withSuffix then attachSuffixes else id) + . (if withPrefix then attachPrefixes else id) + . chunksToWindows seqId strnd + . (case winSz of { Nothing → collapseData; Just sz → reChunkBS sz }) + +-- | Get the full length of a stream of 'BioSequenceWindow's, counted in +-- characters in each 'bswSequence'. +-- +-- To use, start with @bswSeqLength $ SP.copy xs@. Then consume this stream +-- normally. It still provides a 'Stream' of 'BioSequenceWindows's. However, +-- the return type is now not just @r@, but it provides @Int SP.:> r@, where +-- the @Int@ provides the total length of characters within this @Fasta@ entry. +-- +-- This value may then be used to fully update negative strand information. + +bswSeqLength ∷ (Monad m) ⇒ Stream (Of (BioSequenceWindow w ty k)) m r → m (Of Int r) +{-# Inlinable bswSeqLength #-} +bswSeqLength = SP.fold (\x w → x + view (bswSequence._BioSequence.to BS.length) w) 0 id + +-- | As a first function, the header should be extracted from a @Fasta@ stream. Since headers may be malformed / malicious, we make it possible to + +extractHeader + ∷ (Monad m) + ⇒ Maybe Int + → Stream (ByteString m) m r + → m (Of BS.ByteString (Stream (ByteString m) m r)) +{-# Inlinable extractHeader #-} +extractHeader hdrSz = + let go = case hdrSz of { Nothing → id; Just sz → S8.drained . S8.splitAt (fromIntegral sz) } + in S8.toStrict . go . S8.concat . S.splitsAt 1 + +foo = S8.fromStrict ">a\na\na\n>b\nb\nb\n" -- | Control structure for 'streamingFasta'. diff --git a/BiobaseFasta.cabal b/BiobaseFasta.cabal --- a/BiobaseFasta.cabal +++ b/BiobaseFasta.cabal @@ -1,6 +1,6 @@ cabal-version: 2.2 name: BiobaseFasta -version: 0.3.0.0 +version: 0.3.0.1 author: Christian Hoener zu Siederdissen maintainer: choener@bioinf.uni-leipzig.de homepage: https://github.com/choener/BiobaseFasta @@ -44,6 +44,7 @@ , string-conversions >= 0.4 -- , BiobaseTypes == 0.2.0.* + , DPutils == 0.1.0.* default-language: Haskell2010 default-extensions: BangPatterns diff --git a/tests/properties.hs b/tests/properties.hs --- a/tests/properties.hs +++ b/tests/properties.hs @@ -23,6 +23,7 @@ import Test.Tasty.TH import Biobase.Types.BioSequence +import Biobase.Types.Location import Biobase.Types.Strand import Biobase.Fasta.Streaming @@ -76,7 +77,7 @@ , "890" ] -smallTest ∷ Int → Int → Int → Of [BioSequenceWindow Void Void 1] () +smallTest ∷ Int → Int → Int → Of [BioSequenceWindow Void Void PartialLocation] () smallTest h o c = runIdentity . toList -- . SP.map (view windowedFasta) @@ -89,10 +90,10 @@ let res :> r = smallTest 3 3 3 assertEqual "return is null" () r assertEqual "length is 4" 4 (P.length res) - assertEqual "!!0" (BioSequenceWindow "Aaa" "" "123" "" PlusStrand 1) (res!!0) - assertEqual "!!1" (BioSequenceWindow "Bbb" "" "456" "" PlusStrand 1) (res!!1) - assertEqual "!!2" (BioSequenceWindow "Bbb" "456" "7" "" PlusStrand 4) (res!!2) - assertEqual "!!3" (BioSequenceWindow "Ccc" "" "890" "" PlusStrand 1) (res!!3) + assertEqual "!!0" (BioSequenceWindow "Aaa" "" "123" "" (PartialLocation PlusStrand 0 3)) (res!!0) + assertEqual "!!1" (BioSequenceWindow "Bbb" "" "456" "" (PartialLocation PlusStrand 0 3)) (res!!1) + assertEqual "!!2" (BioSequenceWindow "Bbb" "456" "7" "" (PartialLocation PlusStrand 3 1)) (res!!2) + assertEqual "!!3" (BioSequenceWindow "Ccc" "" "890" "" (PartialLocation PlusStrand 0 3)) (res!!3) -- assertEqual "!!0/Fasta" (Fasta "Aaa" "123") (view windowedFasta $ res!!0) assertEqual "!!1/Fasta" (Fasta "Bbb" "456") (view windowedFasta $ res!!1)