Haskellでtailコマンド - その2

前回作ったtailがあまりにも遅過ぎて酷かったので、全面的に書き直しました。
mmapは使ってません。conduitは使ってます。

{-
 - 方針
 -  Sourceでは、seekして後ろから4096バイトずつ(又はファイルサイズの1/10バイトずつ)取得する。
 -  Conduitで、後ろからn個目の改行がある所までで入力を切る。
 -  Sinkは後ろから順に積んでいく。
 -}

import Control.Applicative ((<$>))
import Control.Monad (when)
import Control.Monad.IO.Class (liftIO)
import qualified Data.ByteString as BS
import qualified Data.ByteString.Char8 as BC
import Data.Conduit ( ($$), ($=) )
import qualified Data.Conduit as C
import System.Environment (getArgs)
import qualified System.IO as SI

type FileSize = Integer


readLen :: FileSize -> Integer
readLen s
    | s < std  = std
    | otherwise = if x > limit then limit else x
    where
    x = s `div` 10
    std = 4096
    limit = 10 * 1024 * 1024


reverseFile :: C.ResourceIO m => FilePath -> C.Source m BS.ByteString
reverseFile fp = C.sourceIO initialize close pull
    where
    initialize = do
        h <- SI.openBinaryFile fp SI.ReadMode
        s <- readLen <$> SI.hFileSize h
        b <- SI.hIsSeekable h
        when b $ SI.hSeek h SI.SeekFromEnd (-1)
        return $ if b then Just (h,s) else Nothing
    close st
        | Nothing    <- st = return ()
        | Just (h,_) <- st = SI.hClose h
    pull st
        | Nothing     <- st = return C.Closed
        | Just (h,l)  <- st = do
            pos <- liftIO $ SI.hTell h
            if pos == 0
                then return C.Closed
                else do
                    let len = if pos < l then pos else l
                    liftIO $ SI.hSeek h SI.RelativeSeek (-len)
                    x <- liftIO $ BS.hGetSome h $ fromIntegral len
                    liftIO $ SI.hSeek h SI.RelativeSeek (-len)
                    return $ C.Open x


stackSink :: C.Resource m => C.Sink BS.ByteString m BS.ByteString
stackSink = C.sinkState BS.empty push close
    where
    close = return
    push st i = return (i `BS.append` st, C.Processing)
        

limitLinesReverse :: C.Resource m => Int -> C.Conduit BS.ByteString m BS.ByteString
limitLinesReverse count = C.conduitState count push close
    where
    close _ = return []
    push n i
        | n == 0    = return (n, C.Finished (Just i) [])
        | otherwise =  do
            let (n',bs) = loop n i
            if n' == 0
                then return (n', C.Finished Nothing [bs])
                else return (n', C.Producing [bs])
    loop n x
        | c <= n = (n - c, x)
        | c > n  = (0, BC.drop (pos - 1) x)
        where
        c = BC.count '\n' x
        pos = BC.elemIndices '\n' x !! (c - n)


main :: IO ()
main = do
    [fp, n] <- getArgs
    x <- C.runResourceT $ reverseFile fp $= limitLinesReverse (read n) $$ stackSink
    BC.putStrLn $ BC.drop 1 x

これをghc -O2 -o mytail でコンパイルしてOpenSuseのtailと比較してみました。
結果：

対象データ
    18MiB / 381091行 のテキストファイル(.emacsをずらずら繋げて作成)

500行読み込み
    自作tail    : ./mytail ~/test.txt 500  0.00s user 0.00s system 60% cpu 0.012 total
    Suseのtail  : tail -n 500 ~/test.txt  0.00s user 0.00s system 15% cpu 0.007 total

1000行読み込み
    自作tail    : ./mytail ~/test.txt 1000  0.00s user 0.01s system 41% cpu 0.029 total
    Suseのtail  : tail -n 1000 ~/test.txt  0.00s user 0.00s system 19% cpu 0.016 total

5000行読み込み
    自作tail    : ./mytail ~/test.txt 5000  0.00s user 0.01s system 18% cpu 0.082 total
    Suseのtail  : tail -n 5000 ~/test.txt  0.00s user 0.01s system 8% cpu 0.068 total

100000行読み込み
    自作tail    : ./mytail ~/test.txt 100000  0.01s user 0.21s system 14% cpu 1.438 total
    Suseのtail  : tail -n 100000 ~/test.txt  0.01s user 0.18s system 12% cpu 1.433 total

381091行(全部)読み込み
    自作tail    : ./mytail ~/test.txt 381091  0.04s user 0.69s system 13% cpu 5.546 total
    Suseのtail  : tail -n 381091 ~/test.txt  0.01s user 0.82s system 15% cpu 5.446 total

今日は気持ちよく寝れそうです。