Enumeratorは終了条件の検査からの解放だ

長年の疑問の答えが Enumerator だった。今日はそんなお話です。

findを実装する

Real World Haskell の第9章に UNIX の find コマンドを作る例があります。最初は安直に実装してみましょう。まず、ディレクトリの内容を得る補助関数と、ディレクトリが辿れるか調べる補助関数を用意します。

getValidContents :: FilePath -> IO [String]
getValidContents path = 
    filter (`notElem` [".", "..", ".git", ".svn"]) <$> getDirectoryContents path

isSearchableDir :: FilePath -> IO Bool
isSearchableDir dir =
    (&&) <$> doesDirectoryExist dir
         <*> (searchable <$> getPermissions dir)

次に、ディレクトリを再帰的に走査して、ファイル名のリストを返す関数を作ります。

getRecursiveContents :: FilePath -> IO [FilePath]
getRecursiveContents dir = do
  cnts <- map (dir </>) <$> getValidContents dir
  concat <$> forM cnts $. \path -> do
    isDirectory <- isSearchableDir path
    if isDirectory
      then getRecursiveContents path
      else return [path]

ある文字列を名前の一部として持つファイルのみを選別する関数 grep と、これらを組み合わせる関数 find を以下のように実装します。

grep :: String -> [FilePath] -> [FilePath]
grep pattern = filter (pattern `isInfixOf`)

find :: FilePath -> String -> IO ()
find dir pattern = grep pattern <$> getRecursiveContents dir
               >>= mapM_ putStrLn

これで、find "ディレクトリ" "パターン" を評価すれば、一応目的は達成できます。しかし、getRecursiveContents は、与えられたディレクトリの隅々を走査し終わった後でないと、ファイル名のリストを返しません。

ですから、この find はしばらくユーザーを待たせてから、一気にファイル名を表示します。

Enumerator で実装した find

RWH では、ファイルを列挙する者と消費する者の間に簡単な API を定義して、この問題を解決しています。これは、よく考えると Enumerator そのものです。そこで、その API を Enumerator の API に置き換えて実装してみることにします。

消費者の方は簡単です。

grepE :: String -> Enumeratee String String IO b
grepE pattern = E.filter (pattern `isInfixOf`)

printI :: Iteratee String IO ()
printI = EL.head >>= maybe (return ()) $. \file -> do
    liftIO . putStrLn $ file
    printI

生産者の方は、Enumerator が (<==<) で合成できることを知っていれば、以下のように実装できるでしょう。

enumDir :: FilePath -> Enumerator String IO b
enumDir dir = list
  where
    list (Continue k) = do
        (files,dirs) <- liftIO getFilesDirs
        if null dirs
           then k (Chunks files)
           else k (Chunks files) >>== walk dirs
    list step = returnI step
    walk dirs = foldr1 (<==<) $ map enumDir dirs
    getFilesDirs = do
        cnts <- map (dir </>) <$> getValidContents dir
        (,) <$> filterM doesFileExist cnts
            <*> filterM isSearchableDir cnts

これで Enumerator 版 find の完成です。

findE :: FilePath -> String -> IO ()
findE dir pattern = run_ $ enumDir dir
                        $$ grepE pattern
                        =$ printI

findE は、望み通りファイルを見つたタイミングで、ファイルを表示します。

終了条件の検査からの解放

上記の findE を N 個表示して終わるように、改造しましょう。EL.isolate を間に差し込むだけです。

findE' :: FilePath -> String -> Integer -> IO ()
findE' dir pattern n = run_ $ enumDir dir
                           $$ grepE pattern
                           =$ EL.isolate n
                           =$ printI

この例から、Enumerator は終了条件の検査から解放されていることがはっきり分かりますね。IO が絡んだときに、終了条件の検査から解放されるコードをどう書くのか。長年の疑問の答えは、Enumerator でした。

おまけ

必要はモジュールは以下の通り。

import Control.Applicative ((<$>),(<*>))
import Control.Monad (filterM, forM)
import Control.Monad.IO.Class (liftIO)
import Data.Enumerator (Iteratee, Enumeratee, Enumerator, Stream(..), Step(..), (>>==), (<==<), returnI, yield, run_, joinI, ($$), (=$))
import qualified Data.Enumerator as E (filter)
import Data.Enumerator.List as EL
import Data.List (isInfixOf)
import System.Directory (getDirectoryContents, doesFileExist, doesDirectoryExist, getPermissions, Permissions(..))
import System.FilePath ((</>))

infixr 5 $.
($.) :: (a -> b) -> a -> b
($.) = ($)