nptools-0.2.1: cp-rescue.hs

{-# LANGUAGE ScopedTypeVariables #-}
-- This script is originally from Neil Mitchell's blog

import Data.ByteString (hGet, hPut)
import Data.Char (isDigit)
import System.IO
import System.Environment
import Control.Monad
import Control.Exception

type Offset = Integer
type Length = Integer
type ChunkLength = Length

defaultChunkLength :: ChunkLength
defaultChunkLength = 10000000

main :: IO ()
main = do
  args <- getArgs

  (clen,src,dest) <-
    case args of
      [src,dest]
         -> return (defaultChunkLength,src,dest)
      ["--chunk-length",clen,src,dest]
         | all isDigit clen
         -> return (read clen,src,dest)
      _  -> fail "Usage: cp-rescue [--chunk-length <num>] <src> <dest>"

  withBinaryFile src ReadMode $ \hSrc ->
    withBinaryFile dest WriteMode $ \hDest -> do
      nSrc   <- hFileSize hSrc
      nDest  <- hFileSize hDest
      when (nSrc /= nDest) $ hSetFileSize hDest nSrc
      copy hSrc hDest $ split clen 0 nSrc

copyChunk :: Handle -> Handle -> Offset -> Length -> IO ()
copyChunk hSrc hDest from len = do
  hSeek hSrc AbsoluteSeek from
  hSeek hDest AbsoluteSeek from
  bs <- hGet hSrc $ fromIntegral len
  hPut hDest bs

copy :: Handle -> Handle -> [(Offset,Length)] -> IO ()
copy _     _      []      = return ()
copy hSrc  hDest  chunks  = do
  putStrLn $ "Copying " ++ show (length chunks) ++ " of at most " ++ show (snd $ head chunks)
  chunks' <- forM chunks $ \(from,len) -> do
    res <- Control.Exception.try $ copyChunk hSrc hDest from len
    case res of
      Left (_ :: IOException) -> putChar '#' >> return (split (len `div` 5) from len)
      Right _ -> putChar '.' >> return []
  putChar '\n'
  copy hSrc hDest $ concat chunks'

split :: ChunkLength -> Offset -> Length -> [(Offset,Length)]
split clen off len
  | clen <= 0    = []
  | clen >= len  = [(off,len)]
  | otherwise    = (off,clen) : split clen (off+clen) (len-clen)