{-# LANGUAGE ScopedTypeVariables #-}
-- |This module defines core functions for tracking the consumption of a
-- ByteString, as well as several helper functions for making tracking
-- ByteStrings easier.
module Data.ByteString.Lazy.Progress(
         trackProgress
       , trackProgressWithChunkSize
       --
       , trackProgressString 
       , trackProgressStringWithChunkSize 
       --
       , bytesToUnittedStr
       )
 where

import           Control.Applicative ((<$>))
import qualified Data.ByteString      as BSS
import           Data.ByteString.Lazy(ByteString)
import qualified Data.ByteString.Lazy as BS
import           Data.Maybe          (isJust)
import           Data.Time.Clock     (getCurrentTime,diffUTCTime,UTCTime)
import           Data.Word           (Word64)
import           System.IO.Unsafe    (unsafeInterleaveIO)

-- |Given a function, return a bytestring that will call that function when it
-- is partially consumed. The Words provided to the function will be the number
-- of bytes that were just consumed and the total bytes consumed thus far.
trackProgress :: (Word64 -> Word64 -> IO ()) ->
                 ByteString ->
                 IO ByteString
trackProgress :: (Word64 -> Word64 -> IO ()) -> ByteString -> IO ByteString
trackProgress Word64 -> Word64 -> IO ()
tracker ByteString
inputBS =
  [ByteString] -> ByteString
BS.fromChunks ([ByteString] -> ByteString) -> IO [ByteString] -> IO ByteString
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Word64 -> [ByteString] -> IO [ByteString]
runTrack Word64
0 (ByteString -> [ByteString]
BS.toChunks ByteString
inputBS)
 where
  runTrack :: Word64 -> [ByteString] -> IO [ByteString]
runTrack Word64
_ [] = [ByteString] -> IO [ByteString]
forall (m :: * -> *) a. Monad m => a -> m a
return []
  runTrack Word64
x (ByteString
fst:[ByteString]
rest) = IO [ByteString] -> IO [ByteString]
forall a. IO a -> IO a
unsafeInterleaveIO (IO [ByteString] -> IO [ByteString])
-> IO [ByteString] -> IO [ByteString]
forall a b. (a -> b) -> a -> b
$ do
    let amtRead :: Word64
amtRead = Int -> Word64
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int -> Word64) -> Int -> Word64
forall a b. (a -> b) -> a -> b
$ ByteString -> Int
BSS.length ByteString
fst
    Word64 -> Word64 -> IO ()
tracker Word64
amtRead (Word64
x Word64 -> Word64 -> Word64
forall a. Num a => a -> a -> a
+ Word64
amtRead)
    (ByteString
fst ByteString -> [ByteString] -> [ByteString]
forall a. a -> [a] -> [a]
:) ([ByteString] -> [ByteString])
-> IO [ByteString] -> IO [ByteString]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Word64 -> [ByteString] -> IO [ByteString]
runTrack (Word64
x Word64 -> Word64 -> Word64
forall a. Num a => a -> a -> a
+ Word64
amtRead) [ByteString]
rest

-- |Works like 'trackProgress', except uses fixed-size chunks of the given
-- size.  Thus, for this function, the first number passed to your function
-- will always be the given size *except* for the last call to the function,
-- which will be less then or equal to the final size.
trackProgressWithChunkSize :: Word64 -> (Word64 -> Word64 -> IO ()) ->
                              ByteString ->
                              IO ByteString
trackProgressWithChunkSize :: Word64
-> (Word64 -> Word64 -> IO ()) -> ByteString -> IO ByteString
trackProgressWithChunkSize Word64
chunkSize Word64 -> Word64 -> IO ()
tracker ByteString
inputBS = Word64 -> ByteString -> IO ByteString
runLoop Word64
0  ByteString
inputBS
 where
  runLoop :: Word64 -> ByteString -> IO ByteString
runLoop Word64
x ByteString
bstr | ByteString -> Bool
BS.null ByteString
bstr = ByteString -> IO ByteString
forall (m :: * -> *) a. Monad m => a -> m a
return ByteString
BS.empty
                 | Bool
otherwise    = IO ByteString -> IO ByteString
forall a. IO a -> IO a
unsafeInterleaveIO (IO ByteString -> IO ByteString) -> IO ByteString -> IO ByteString
forall a b. (a -> b) -> a -> b
$ do
    let (ByteString
first,ByteString
rest) = Int64 -> ByteString -> (ByteString, ByteString)
BS.splitAt (Word64 -> Int64
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word64
chunkSize) ByteString
bstr
        amtRead :: Word64
amtRead      = Int64 -> Word64
forall a b. (Integral a, Num b) => a -> b
fromIntegral (ByteString -> Int64
BS.length ByteString
first)
    Word64 -> Word64 -> IO ()
tracker Word64
amtRead (Word64
x Word64 -> Word64 -> Word64
forall a. Num a => a -> a -> a
+ Word64
amtRead)
    (ByteString
first ByteString -> ByteString -> ByteString
`BS.append`) (ByteString -> ByteString) -> IO ByteString -> IO ByteString
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Word64 -> ByteString -> IO ByteString
runLoop (Word64
x Word64 -> Word64 -> Word64
forall a. Num a => a -> a -> a
+ Word64
amtRead) ByteString
rest

-- |Given a format string (described below), track the progress of a function.
-- The argument to the callback will be the string expanded with the given
-- progress information.
--
-- Format string items:
--
--   * %b is the number of bytes read
--
--   * %B is the number of bytes read, formatted into a human-readable string
--
--   * %c is the size of the last chunk read
--
--   * %C is the size of the last chunk read, formatted human-readably
--
--   * %r is the rate in bytes per second
--
--   * %R is the rate, formatted human-readably
--
--   * %% is the character '%'
--
-- If you provide a total size (the maybe argument, in bytes), then you may
-- also use the following items:
--
--   * %t is the estimated time to completion in seconds
--
--   * %T is the estimated time to completion, formatted as HH:MM:SS
--
--   * %p is the percentage complete
--
trackProgressString :: String -> Maybe Word64 -> (String -> IO ()) ->
                       IO (ByteString -> IO ByteString)
trackProgressString :: String
-> Maybe Word64
-> (String -> IO ())
-> IO (ByteString -> IO ByteString)
trackProgressString String
formatStr Maybe Word64
mTotal String -> IO ()
tracker = do
  UTCTime
startTime <- IO UTCTime
getCurrentTime
  (ByteString -> IO ByteString) -> IO (ByteString -> IO ByteString)
forall (m :: * -> *) a. Monad m => a -> m a
return ((Word64 -> Word64 -> IO ()) -> ByteString -> IO ByteString
trackProgress (UTCTime -> Word64 -> Word64 -> IO ()
handler UTCTime
startTime))
 where
  handler :: UTCTime -> Word64 -> Word64 -> IO ()
handler UTCTime
startTime Word64
chunkSize Word64
total = do
    UTCTime
now <- IO UTCTime
getCurrentTime
    String -> IO ()
tracker (String
-> UTCTime -> UTCTime -> Maybe Word64 -> Word64 -> Word64 -> String
buildString String
formatStr UTCTime
startTime UTCTime
now Maybe Word64
mTotal Word64
chunkSize Word64
total)

-- |Exactly as 'trackProgressString', but use the given chunkSize instead
-- of the default chunk size.
trackProgressStringWithChunkSize :: String -- ^the format string
                                    -> Word64 -- ^the chunk size
                                    -> Maybe Word64 -- ^total size (opt.)
                                    -> (String -> IO ()) -- ^the action
                                    -> IO (ByteString -> IO ByteString)
trackProgressStringWithChunkSize :: String
-> Word64
-> Maybe Word64
-> (String -> IO ())
-> IO (ByteString -> IO ByteString)
trackProgressStringWithChunkSize String
formatStr Word64
chunk Maybe Word64
mTotal String -> IO ()
tracker = do
  UTCTime
startTime <- IO UTCTime
getCurrentTime
  (ByteString -> IO ByteString) -> IO (ByteString -> IO ByteString)
forall (m :: * -> *) a. Monad m => a -> m a
return (Word64
-> (Word64 -> Word64 -> IO ()) -> ByteString -> IO ByteString
trackProgressWithChunkSize Word64
chunk (UTCTime -> Word64 -> Word64 -> IO ()
handler UTCTime
startTime))
 where
  handler :: UTCTime -> Word64 -> Word64 -> IO ()
handler UTCTime
startTime Word64
chunkSize Word64
total = do
    UTCTime
now <- IO UTCTime
getCurrentTime
    String -> IO ()
tracker (String
-> UTCTime -> UTCTime -> Maybe Word64 -> Word64 -> Word64 -> String
buildString String
formatStr UTCTime
startTime UTCTime
now Maybe Word64
mTotal Word64
chunkSize Word64
total)

-- build a progress string for trackProgressString et al
buildString :: String ->
               UTCTime -> UTCTime -> Maybe Word64 -> Word64 -> Word64 ->
               String
buildString :: String
-> UTCTime -> UTCTime -> Maybe Word64 -> Word64 -> Word64 -> String
buildString String
form UTCTime
startTime UTCTime
curTime Maybe Word64
mTotal Word64
chunkSize Word64
amtRead = String -> String
subPercents String
form
 where
  per_b :: String
per_b = Word64 -> String
forall a. Show a => a -> String
show Word64
amtRead
  per_B :: String
per_B = Word64 -> String
bytesToUnittedStr Word64
amtRead
  per_c :: String
per_c = Word64 -> String
forall a. Show a => a -> String
show Word64
chunkSize
  per_C :: String
per_C = Word64 -> String
bytesToUnittedStr Word64
chunkSize
  diff :: Word64
diff  = Word64 -> Word64 -> Word64
forall a. Ord a => a -> a -> a
max Word64
1 (Rational -> Word64
forall a b. (RealFrac a, Integral b) => a -> b
round (Rational -> Word64) -> Rational -> Word64
forall a b. (a -> b) -> a -> b
$ NominalDiffTime -> Rational
forall a. Real a => a -> Rational
toRational (NominalDiffTime -> Rational) -> NominalDiffTime -> Rational
forall a b. (a -> b) -> a -> b
$ UTCTime -> UTCTime -> NominalDiffTime
diffUTCTime UTCTime
curTime UTCTime
startTime)
  rate :: Word64
rate  = Word64
amtRead Word64 -> Word64 -> Word64
forall a. Integral a => a -> a -> a
`div` Word64
diff
  per_r :: String
per_r = Word64 -> String
forall a. Show a => a -> String
show Word64
rate
  per_R :: String
per_R = Word64 -> String
bytesToUnittedStr Word64
rate String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
"ps"
  total :: Word64
total = case Maybe Word64
mTotal of
            Just Word64
t  -> Word64
t
            Maybe Word64
Nothing -> String -> Word64
forall a. HasCallStack => String -> a
error String
"INTERNAL ERROR (needed total w/ Nothing)"
  tleft :: Word64
tleft = (Word64
total Word64 -> Word64 -> Word64
forall a. Num a => a -> a -> a
- Word64
amtRead) Word64 -> Word64 -> Word64
forall a. Integral a => a -> a -> a
`div` Word64
rate
  per_t :: String
per_t = Word64 -> String
forall a. Show a => a -> String
show Word64
tleft
  hLeft :: Word64
hLeft = Word64
tleft Word64 -> Word64 -> Word64
forall a. Integral a => a -> a -> a
`div` (Word64
60 Word64 -> Word64 -> Word64
forall a. Num a => a -> a -> a
* Word64
60)
  mLeft :: Word64
mLeft = (Word64
tleft Word64 -> Word64 -> Word64
forall a. Integral a => a -> a -> a
`div` Word64
60) Word64 -> Word64 -> Word64
forall a. Integral a => a -> a -> a
`mod` Word64
60
  sLeft :: Word64
sLeft =  Word64
tleft           Word64 -> Word64 -> Word64
forall a. Integral a => a -> a -> a
`mod` Word64
60
  per_T :: String
per_T = Word64 -> String
forall a. Show a => a -> String
showPadded Word64
hLeft String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
":" String -> String -> String
forall a. [a] -> [a] -> [a]
++ Word64 -> String
forall a. Show a => a -> String
showPadded Word64
mLeft String -> String -> String
forall a. [a] -> [a] -> [a]
++
          String
":" String -> String -> String
forall a. [a] -> [a] -> [a]
++  Word64 -> String
forall a. Show a => a -> String
showPadded Word64
sLeft
  perc :: Double
perc  = Double
100 Double -> Double -> Double
forall a. Num a => a -> a -> a
* (Word64 -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word64
amtRead Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Word64 -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word64
total) :: Double
  per_p :: String
per_p = Integer -> String
forall a. Show a => a -> String
show (Double -> Integer
forall a b. (RealFrac a, Integral b) => a -> b
round Double
perc) String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
"%"
  oktot :: Bool
oktot = Maybe Word64 -> Bool
forall a. Maybe a -> Bool
isJust Maybe Word64
mTotal
  --
  subPercents :: String -> String
subPercents []         = []
  subPercents (Char
'%':String
rest) = String -> String
subPercents' String
rest
  subPercents (Char
x:String
rest)   = Char
x Char -> String -> String
forall a. a -> [a] -> [a]
: String -> String
subPercents String
rest
  --
  subPercents' :: String -> String
subPercents' []                 = []
  subPercents' (Char
'b':String
rest)         = String
per_b String -> String -> String
forall a. [a] -> [a] -> [a]
++ String -> String
subPercents String
rest
  subPercents' (Char
'B':String
rest)         = String
per_B String -> String -> String
forall a. [a] -> [a] -> [a]
++ String -> String
subPercents String
rest
  subPercents' (Char
'c':String
rest)         = String
per_c String -> String -> String
forall a. [a] -> [a] -> [a]
++ String -> String
subPercents String
rest
  subPercents' (Char
'C':String
rest)         = String
per_C String -> String -> String
forall a. [a] -> [a] -> [a]
++ String -> String
subPercents String
rest
  subPercents' (Char
'r':String
rest)         = String
per_r String -> String -> String
forall a. [a] -> [a] -> [a]
++ String -> String
subPercents String
rest
  subPercents' (Char
'R':String
rest)         = String
per_R String -> String -> String
forall a. [a] -> [a] -> [a]
++ String -> String
subPercents String
rest
  subPercents' (Char
't':String
rest) | Bool
oktot = String
per_t String -> String -> String
forall a. [a] -> [a] -> [a]
++ String -> String
subPercents String
rest
  subPercents' (Char
'T':String
rest) | Bool
oktot = String
per_T String -> String -> String
forall a. [a] -> [a] -> [a]
++ String -> String
subPercents String
rest
  subPercents' (Char
'p':String
rest) | Bool
oktot = String
per_p String -> String -> String
forall a. [a] -> [a] -> [a]
++ String -> String
subPercents String
rest
  subPercents' (Char
'%':String
rest)         = String
"%"   String -> String -> String
forall a. [a] -> [a] -> [a]
++ String -> String
subPercents String
rest
  subPercents' (Char
x:String
rest)           = Char
'%' Char -> String -> String
forall a. a -> [a] -> [a]
: (Char
'x' Char -> String -> String
forall a. a -> [a] -> [a]
: String -> String
subPercents String
rest)

-- show a number padded to force at least two digits.
showPadded :: Show a => a -> String
showPadded :: forall a. Show a => a -> String
showPadded a
x = String
prefix String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
base
 where
  base :: String
base   = a -> String
forall a. Show a => a -> String
show a
x
  prefix :: String
prefix = case String
base of
             []  -> String
"00"
             [Char
x] ->  String
"0"
             String
_   ->   String
""

-- |Convert a number of bytes to a string represenation that uses a reasonable
-- unit to make the number human-readable.
bytesToUnittedStr :: Word64 -> String
bytesToUnittedStr :: Word64 -> String
bytesToUnittedStr Word64
x
  | Word64
x Word64 -> Word64 -> Bool
forall a. Ord a => a -> a -> Bool
< Word64
bk_brk = Word64 -> String
forall a. Show a => a -> String
show Word64
x String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
"b"
  | Word64
x Word64 -> Word64 -> Bool
forall a. Ord a => a -> a -> Bool
< Word64
km_brk = Word64 -> Word64 -> String
forall {a} {p}. (Integral p, Integral a) => p -> a -> String
showHundredthsDiv Word64
x Word64
k String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
"k"
  | Word64
x Word64 -> Word64 -> Bool
forall a. Ord a => a -> a -> Bool
< Word64
mg_brk = Word64 -> Word64 -> String
forall {a} {p}. (Integral p, Integral a) => p -> a -> String
showHundredthsDiv Word64
x Word64
m String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
"m"
  | Bool
otherwise  = Word64 -> Word64 -> String
forall {a} {p}. (Integral p, Integral a) => p -> a -> String
showHundredthsDiv Word64
x Word64
g String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
"g"
 where
  bk_brk :: Word64
bk_brk = Word64
4096
  km_brk :: Word64
km_brk = Word64
768 Word64 -> Word64 -> Word64
forall a. Num a => a -> a -> a
* Word64
k
  mg_brk :: Word64
mg_brk = Word64
768 Word64 -> Word64 -> Word64
forall a. Num a => a -> a -> a
* Word64
m
  --
  k :: Word64
k      = Word64
1024
  m :: Word64
m      = Word64
1024 Word64 -> Word64 -> Word64
forall a. Num a => a -> a -> a
* Word64
k
  g :: Word64
g      = Word64
1024 Word64 -> Word64 -> Word64
forall a. Num a => a -> a -> a
* Word64
m

-- Divide the first number by the second, and convert to a string showing two
-- decimal places.  
showHundredthsDiv :: p -> a -> String
showHundredthsDiv   p
_    a
0 = String -> String
forall a. HasCallStack => String -> a
error String
"Should never happen!"
showHundredthsDiv p
amt a
size = Integer -> String
forall a. Show a => a -> String
show Integer
ones String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
"." String -> String -> String
forall a. [a] -> [a] -> [a]
++ Integer -> String
forall a. Show a => a -> String
show Integer
tenths String -> String -> String
forall a. [a] -> [a] -> [a]
++ Integer -> String
forall a. Show a => a -> String
show Integer
hundreths
 where
  Double
divRes :: Double = p -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral p
amt Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ a -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral a
size
  divRes100 :: Integer
divRes100        = Double -> Integer
forall a b. (RealFrac a, Integral b) => a -> b
round (Double
divRes Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
100)
  ones :: Integer
ones             =  Integer
divRes100 Integer -> Integer -> Integer
forall a. Integral a => a -> a -> a
`div` Integer
100
  tenths :: Integer
tenths           = (Integer
divRes100 Integer -> Integer -> Integer
forall a. Integral a => a -> a -> a
`div` Integer
10) Integer -> Integer -> Integer
forall a. Integral a => a -> a -> a
`mod` Integer
10
  hundreths :: Integer
hundreths        =  Integer
divRes100           Integer -> Integer -> Integer
forall a. Integral a => a -> a -> a
`mod` Integer
10