module Data.Progress (progress, progressWithFile, progressWithCalls, progress', progressWithCalls') where
import System.IO.Unsafe
import System.IO
import System.Console.ANSI
import Data.Data
import Control.Monad.Identity
import Control.Monad
import Control.Concurrent
import Control.Exception
import Data.List
newtype Size t = Size { unSize :: Integer }
size :: (Data t) => t -> Integer
size x = unSize (gfoldl (\(Size n) y -> Size (n + size y)) (const (Size 1)) x)
active :: MVar Integer
active = unsafePerformIO (newMVar $ 1)
putBar :: Integer -> IO ()
putBar percent = modifyMVar_ active $ \n ->
if n < percent then do
setCursorColumn 0
let s = show $ min 100 percent
putStr $ replicate (3 length s) ' ' ++ s
return percent
else
return n
withActive m = do
b <- modifyMVar active (\n -> return $ if n == 1 then (0, False) else (n, True))
if b then
m
else
finally m (modifyMVar_ active (const $ return $ 1))
fork' m = do
caps <- getNumCapabilities
if caps >= 2 then
void $ forkIO m
else
void m
progress f dat = withActive $ do
putStr "\n 0%"
sz <- newMVar $ 1
count <- newMVar 0
fork' $ do
let s = size dat
evaluate s
modifyMVar_ sz $ const $ return s
let
rec :: (Data t) => t -> t
rec dat = runIdentity $ gfoldl
(\(Identity f) x -> unsafePerformIO $ do
s <- readMVar sz
modifyMVar_ count $ \n ->
if n /= 1 then do
let n' = n + 1
unless (s == 1) $ putBar (n' * 100 `quot` s)
return n'
else
return n
return $ Identity $ f $ rec x)
Identity
dat
finally
(do
res <- f $ rec dat
evaluate res
return res)
(do
modifyMVar_ count $ const $ return $ 1
putBar 100
putStrLn "")
try' :: IO t -> IO (Either SomeException t)
try' = try
progressWithFile f hdl = withActive $ do
putStr "\n 0%"
thd <- try' $ do
sz <- liftM fromInteger $ hFileSize hdl
forkIO $ foldM_ (\_ _ -> do
n <- liftM fromInteger $ hTell hdl
putBar (n * 100 `quot` sz)
threadDelay 500000
return ())
()
(repeat ())
finally
(f hdl)
(do
try' $ either (\_ -> return ()) killThread thd
putBar 100
putStrLn "")
progressWithCalls f x = withActive $ do
putStr "\n 0%"
parms <- newMVar (1, 20000, 1, 0)
let rec depth count x = do
count' <- newMVar (1, 1)
res <- f (rec (depth + 1) count') x
evaluate res
(x, y) <- readMVar count'
modifyMVar_ count $ \(_, z) -> return (y, z + y)
modifyMVar_ parms $ \tup@(dep, rPrev, yPrev, total) -> do
let tup'@(dep', ratio', y', total') = if total < 0 then
tup
else if dep == 1 || depth <= dep then
(depth, (4 * rPrev + y * 65536 `quot` x) `quot` 5, y, total + 1)
else
(dep, rPrev, yPrev, total + 1)
when (y' >= 100) $ putBar (total' * 100 * 65536 ^ dep' `quot` (y' * ratio' ^ dep'))
return tup'
return res
count <- newMVar (1, 1)
finally
(do
res <- rec 0 count x
return res)
(do
modifyMVar_ parms $ const $ return (0, 0, 0, 1)
putBar 100
putStrLn "")
progress' f = progress (return . f)
progressWithCalls' f = progressWithCalls (\g -> return . f (unsafePerformIO . g))