wc そのよん

id:ha-tan:20061114:1163507847の続き。ざっとリファクタリングしてみました。出力桁数は、前もってファイルサイズから計算しておくようにしました。GNU wcも似たような感じになってました。

module Main (main) where

import Control.Monad (when)
import Data.Char (isSpace)
import System.Environment (getArgs)
import System.IO (stdin, hGetContents, openBinaryFile, openFile, hClose, 
                  IOMode(ReadMode), hFileSize, hSetBinaryMode)
import Text.Printf (printf)

readBinaryFile :: FilePath -> IO String
readBinaryFile path =
  hGetContents =<< openBinaryFile path ReadMode

getBinaryContents :: IO String
getBinaryContents = do
  hSetBinaryMode stdin True
  getContents

computeWidth :: [FilePath] -> IO Int
computeWidth files = return . length . show . sum =<< mapM size files
  where
    size file = do
      f <- openFile file ReadMode
      b <- hFileSize f
      hClose f
      return b

wc :: String -> (Int, Int, Int)
wc = wc' (0, 0, 0) False
  where
    wc' (l, w, b) sp []
      | sp        = (l, w,     b)
      | otherwise = (l, w + 1, b)
    wc' (l, w, b) sp (x : xs)
      | x == '\n'       = seq ln $ seq bn $ wc' (ln, w,  bn) True  xs
      | sp && isSpace x =          seq bn $ wc' (l,  w,  bn) True  xs
      | isSpace x       = seq wn $ seq bn $ wc' (l,  wn, bn) True  xs
      | otherwise       =          seq bn $ wc' (l,  w,  bn) False xs
      where
        (ln, wn, bn) = (l + 1, w + 1, b + 1)

wcStdin :: IO ()
wcStdin = putStr . wcFormat 7 "" . wc =<< getBinaryContents

wcFile :: Int -> FilePath -> IO (Int, Int, Int)
wcFile width path = do
  s <- readBinaryFile path
  let w = wc s
  putStr $ wcFormat width path w
  return w

wcTotal :: Int -> [(Int, Int, Int)] -> IO ()
wcTotal width ws = do
  let l = sum $ map (\ (l, _, _) -> l) ws
      w = sum $ map (\ (_, w, _) -> w) ws
      b = sum $ map (\ (_, _, b) -> b) ws
  putStr $ wcFormat width "total" (l, w, b)

wcFormat :: Int -> String -> (Int, Int, Int) -> String
wcFormat len path (l, w, b) =
  let s   = show len
      fmt = "%" ++ s ++ "d %" ++ s ++ "d %" ++ s ++ "d %s\n"
  in printf fmt l w b path

main :: IO ()
main = do
  args <- getArgs
  if (null args)
    then wcStdin
    else do
      width <- computeWidth args
      ws <- mapM (wcFile width) args
      when (length ws > 1) $ wcTotal width ws

BUGS:

• 普通のファイル以外(ディレクトリとか)を指定したときにエラーで終了します。うーん、そのうち何とかしたいなぁ。