diohsc-0.1.3: Slurp.hs

-- This file is part of Diohsc
-- Copyright (C) 2020 Martin Bays <mbays@sdf.org>
--
-- This program is free software: you can redistribute it and/or modify
-- it under the terms of version 3 of the GNU General Public License as
-- published by the Free Software Foundation, or any later version.
--
-- You should have received a copy of the GNU General Public License
-- along with this program.  If not, see http://www.gnu.org/licenses/.

module Slurp where

import           Control.Monad        (when)
import           Data.Hourglass       (timeDiffP)
import           System.IO
import           System.IO.Unsafe     (unsafeInterleaveIO)
import           Time.System          (timeCurrentP)
import           Time.Types           (ElapsedP, NanoSeconds (..), Seconds (..))

import qualified Data.ByteString      as BS
import qualified Data.ByteString.Lazy as BL

-- incorporate lazy IO into a bytestring such that it will print progress as
-- it is forced, if this takes more than a second.
interleaveProgress :: ElapsedP -> BL.ByteString -> IO BL.ByteString
interleaveProgress t0 bs = do
    t1 <- timeCurrentP
    let slurp _ _ n' [] = do
            t' <- timeCurrentP
            when (t' `timeDiffP` t1 > (Seconds 1, 0)) $ do
                putStrLn $ "\r\ESC[KReceived: " ++ humanBytes n'
                    ++ "  " ++ humanRate n' (t' `timeDiffP` t0)
                hFlush stdout
            return []
        slurp n t n' (c:cs) = unsafeInterleaveIO $ do
            let n'' = n' + fromIntegral (BS.length c)
            t' <- timeCurrentP
            if t' `timeDiffP` t > (Seconds 1, 0)
            then do
                when (t' `timeDiffP` t1 > (Seconds 1, 0)) $ do
                    putStr $   "\r\ESC[KProgress: " ++ humanBytes n'
                        ++ "  " ++ humanRate (n' - n) (t' `timeDiffP` t)
                    hFlush stdout
                (c:) <$> slurp n'' t' n'' cs
            else (c:) <$> slurp n t n'' cs
        humanBytes n | n < 1024 = show n ++ "B"
        humanBytes n | n < 1024*1024 =
            let (n',p) = (`divMod` 10) $ (n*10) `div` 1024
            in show n' ++ "." ++ show p ++ "KB"
        humanBytes n =
            let (n',p) = (`divMod` 10) $ (n*10) `div` (1024*1024)
            in show n' ++ "." ++ show p ++ "MB"
        humanRate _ (0,0) = ""
        humanRate n (Seconds s, NanoSeconds ns) = humanBytes r ++ "/s" where
            r = (billion*n)`div`((billion*s)+ns)
            billion = 1000000000
    BL.fromChunks <$> slurp 0 t0 0 (BL.toChunks bs)

-- |force bs, printing progress
slurpNoisily :: ElapsedP -> BL.ByteString -> IO ()
slurpNoisily t0 bs = do
    bs' <- BL.toStrict <$> interleaveProgress t0 bs
    seq bs' $ return ()