content: {- Copyright 2013-2019 NGLess Authors
   - License: MIT -}
  {-# LANGUAGE ScopedTypeVariables, FlexibleContexts, CPP #-}
  
  module Utils.Conduit
      ( ByteLine(..)
      , byteLineSinkHandle
      , byteLineVSinkHandle
      , linesC
      , linesVC
      , zipSink2
      , zipSource2
      ) where
  
  import qualified Data.ByteString as B
  
  import qualified Data.Conduit.List as CL
  import qualified Data.Conduit as C
  import qualified Data.Vector         as V
  import qualified Data.Vector.Mutable as VM
  
  import           Control.Monad (when)
  import           Control.Monad.IO.Class (MonadIO, liftIO)
  import           Control.Monad.Error.Class (MonadError(..))
  import           System.IO
  
  import NGLess.NGError
  
  -- | This just signals that a "line" is expected.
  newtype ByteLine = ByteLine { unwrapByteLine :: B.ByteString }
                  deriving (Show)
  
  -- A limit was introduced to avoid running out of memory on corrupt files, but
  -- it needs to be large enough to accommodate nanopore reads, see:
  -- https://groups.google.com/forum/#!topic/ngless/-ovfYW8hfAs
  maxLineSize :: Int
  maxLineSize = 1024 * 1024 * 1024
  
  concatrevline :: B.ByteString -> [B.ByteString] -> ByteLine
  concatrevline line [] = ByteLine $ lineWindowsTerminated line
  concatrevline line toks = ByteLine . lineWindowsTerminated $ B.concat (reverse (line:toks))
  {-# INLINE concatrevline #-}
  
  linesC:: (MonadError NGError m) => C.ConduitT B.ByteString ByteLine m ()
  linesC = continue 0 []
      where
          continue n toks
              | n > maxLineSize = throwDataError ("Line too long (length is " ++ show n ++ " characters).")
              | otherwise = C.await >>= maybe
                                          (when (n > 0) $ C.yield (concatrevline B.empty toks))
                                          (emit n toks)
          emit n toks tok = case B.elemIndex 10 tok of
                  Nothing -> continue (n + B.length tok) (tok:toks)
                  Just ix -> let (start,rest) = B.splitAt ix tok in do
                                  C.yield (concatrevline start toks)
                                  emit 0 [] (B.tail rest)
  
  {-# INLINE linesC #-}
  
  -- | Remove trailing \r present when the original line terminator was \r\n (windows)
  lineWindowsTerminated :: B.ByteString -> B.ByteString
  lineWindowsTerminated line = if not (B.null line) && B.index line (B.length line - 1) == carriage_return
                                  then B.take (B.length line - 1) line
                                  else line
                                      where carriage_return = 13
  {-# INLINE lineWindowsTerminated #-}
  
  -- | Equivalent to 'linesC .| CC.conduitVector nlines'
  linesVC :: (MonadIO m, MonadError NGError m) => Int -> C.ConduitT B.ByteString (V.Vector ByteLine) m ()
  linesVC nlines = do
              vec <- liftIO $ VM.new nlines
              continue vec 0 0 []
          where
              continue vec vix n toks
                  | n > maxLineSize = throwDataError ("Line too long (length is " ++ show n ++ " characters).")
                  | otherwise = C.await >>= maybe (finish vec vix toks) (emit vec vix n toks)
  
              finish   _   0 [] = return ()
              finish vec vix [] = C.yield =<< liftIO (V.unsafeSlice 0 vix <$> V.unsafeFreeze vec)
              finish vec vix toks = do
                      liftIO $ VM.write vec vix (ByteLine . lineWindowsTerminated . B.concat $ reverse toks)
                      finish vec (vix + 1) []
  
              emit vec vix n toks tok = do
                  (done, vec', vix', n', toks') <- liftIO $ splitWrite [] vec vix n toks tok
                  CL.sourceList (reverse done)
                  continue vec' vix' n' toks'
  
              -- splitWrite is in IO. This is a micro-optimization, but this code
              -- can be in the inner loop, so it's worthwhile to micro-optimize.
              --
              -- Basically, moving up and down the transformer stack (with lift &
              -- friends) can be expensive. Doing the inner loop in IO is
              -- measurably faster.
              splitWrite done vec vix !n toks tok
                  | vix >= nlines = do
                      f <- V.unsafeFreeze vec
                      vec' <- VM.new nlines
                      splitWrite (f:done) vec' 0 n toks tok
                  | otherwise = case B.elemIndex 10 tok of
                      Nothing -> return (done, vec, vix, n + B.length tok, (if not (B.null tok) then (tok:toks) else toks))
                      Just ix -> do
                          let (start, rest) = B.splitAt ix tok
                          VM.write vec vix (concatrevline start toks)
                          splitWrite done vec (vix + 1) 0 [] (B.tail rest)
  {-# INLINE linesVC #-}
  
  
  byteLineSinkHandle :: (MonadIO m) => Handle -> C.ConduitT ByteLine C.Void m ()
  byteLineSinkHandle h = CL.mapM_ (\(ByteLine val) -> liftIO (B.hPut h val >> B.hPut h nl))
      where
          nl = B.singleton 10
  {-# INLINE byteLineSinkHandle #-}
  
  
  byteLineVSinkHandle :: (MonadIO m) => Handle -> C.ConduitT (V.Vector ByteLine) C.Void m ()
  byteLineVSinkHandle h = CL.mapM_ $ liftIO . V.mapM_ (\(ByteLine val) -> B.hPut h val >> B.hPut h nl)
      where
          nl = B.singleton 10
  {-# INLINE byteLineVSinkHandle #-}
  
  
  zipSource2 :: Monad m => C.ConduitT () a m () -> C.ConduitT () b m () -> C.ConduitT () (a,b) m ()
  zipSource2 a b = C.getZipSource ((,) <$> C.ZipSource a <*> C.ZipSource b)
  {-# INLINE zipSource2 #-}
  
  
  zipSink2 :: (Monad m) => C.ConduitT i C.Void m a -> C.ConduitT i C.Void m b -> C.ConduitT i C.Void m (a,b)
  zipSink2 a b = C.getZipSink((,) <$> C.ZipSink a <*> C.ZipSink b)
  {-# INLINE zipSink2 #-}
  
  