Haskellでtailコマンド

conduitを試しに使ってみようという事でtailコマンド書いてみました。
やれることはファイルの後ろから数行を表示するという基本だけ。
引数の処理はテキトーなので行数指定必須です。クソです。

{-# LANGUAGE OverloadedStrings #-}

import Control.Applicative ((<$>))
import Control.Monad.IO.Class (liftIO)
import qualified Data.ByteString.Char8 as BS
import Data.Conduit (($=), ($$))
import qualified Data.Conduit as C
import qualified Data.Conduit.List as CL
import Data.Int
import Data.Word (Word8)
import qualified Data.IORef as I
import System.Environment (getArgs)
import qualified System.IO as SI
import qualified System.IO.MMap as SIM

isEOL :: Char -> Bool
isEOL x = x == '\r' || x == '\n'

mmapSize :: Int64 -> Int
mmapSize fileSize
    | fs > maxSize = maxSize
    | otherwise    = fs
    where
    maxSize = 10 * 1024 * 1024
    fs = fromIntegral fileSize

reverseFile :: C.ResourceIO m => FilePath -> C.Source m BS.ByteString
reverseFile fp = C.sourceIO initialize close f
    where
    initialize = do
        offset <- fromIntegral <$> SI.withBinaryFile fp SI.ReadMode SI.hFileSize
        I.newIORef (mmapSize offset, offset - 1) -- offsetはファイル末端からmmapした領域の末端までの距離
    close = const $ return ()
    f s = do
        (mapSize, offset) <- liftIO (I.readIORef s)
        if offset <= 0
            then return C.Closed
            else do
                let msize = if fromIntegral mapSize > offset then offset else fromIntegral mapSize
                mmap <- liftIO $ SIM.mmapFileByteString fp (Just (offset - msize, fromIntegral msize))
                let (firstl, lastl) = BS.spanEnd (not . isEOL) mmap
                let buf = snd (BS.spanEnd isEOL firstl) `BS.append` lastl
                liftIO $ I.writeIORef s (mapSize, offset - fromIntegral (BS.length buf))
                return $ C.Open $ BS.copy buf


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


main :: IO ()
main = do
    [fp,n] <- getArgs
    l <- BS.dropWhile isEOL <$> C.runResourceT (reverseFile fp $= CL.isolate (read n) $$ stackSink)
    BS.putStrLn l

conduitは書きやすくて再利用性が高いのがいいですね。


suseのtailコマンドと速度を比較してみると30万行のデータを表示するのに111倍遅かったです。