foldrが効率的に動作する条件

haskell

eagerな言語ではfoldrはリストの長さに比例するスタックを消費する。GHCでも(+)のような正格な関数で畳み込もうとすれば同様にO(n)の振る舞いになる。

{-# OPTIONS_GHC -O0 #-}
import Data.List(foldl')
main = do
  n <- readLn
  print $ foldr (+) 0 [1..n] -- O(n)スタック、O(1)ヒープ
  print $ foldl (+) 0 [1..n] -- O(n)スタック、O(n)ヒープ
  print $ foldl' (+) 0 [1..n] -- O(1)スタック、O(1)ヒープ

一方、foldrを使って効率的に実装できる関数もHaskellには存在する。例えば以下はconcatの立派な定義であり、O(1)空間で動作する*1。

concat :: [[a]] -> [a]
concat = foldr (++) []

このようなことが可能な条件は何か。それは、畳み込みに使用する関数(この例では(++))が、第二引数をすぐに評価しようとしないことである。具体的には、それを捨てるか、そのまま返すか、非正格なデータ構造に入れて返すかのどれかを選ぶことになる。

foldrに渡される関数の立場から見てみる。

myfunc = foldr f z [x0, x1, x2...]
  where
    f x y = ...

fがyとして受け取るのは通常の値ではなく、f x1 (f x2 (f x3...))のような未評価の式で、これは「foldrの残りの部分」を表す。つまりyは(無引数の)継続だと思うことができる。この観点から見ると、foldrは継続渡し形式を使ったループを一般化した関数だ。

このように考えると、「第二引数をすぐに評価しようとしない」という制約の意味がはっきりする。

  • yを捨てるのは、すなわち継続を捨てることで、ループからの途中脱出に対応する。
  • yをそのまま返すのは、継続を末尾呼び出しすることで、ループを続けることに対応する。
  • yをテータ構造に入れて返すのは、継続を保存することで、ループの一時中断に対応する。
  • yを加工して返すのは、継続の非末尾呼び出しに対応し、スタックを消費する。

なお、「データ構造に入れて返す」だけでなく「そのまま返す」「捨てる」が許されているので、例えばfindもO(1)空間で動く。

find :: (a -> Bool) -> [a] -> Maybe a
find cond = foldr f Nothing
  where
    f x cont = if cond x then Just x else cont

値を取る継続

fの第二引数を継続だと思った場合、すぐ思い付く一般化は、それが値を取るようにすることだろう。

loop :: (acc -> a -> (acc -> r) -> r) -> acc -> r -> [a] -> r
loop _ _ z [] = z
loop f acc z (x:xs) = f acc x $ \acc' -> loop f acc' z xs

これで、foldlのように値を蓄積しながらループすることが可能になった。「値の蓄積」「途中脱出」の両方が可能という点で、典型的な手続き型言語のループ構文に相当する能力を持っている。

加えて「一時中断」が可能なので、これを利用すると例えばtakeを実装できる。

take :: Int -> [a] -> [a]
take len = loop f len []
  where
    f n x cont
      | n <= 0 = []
      | otherwise = x : cont (n - 1)

難点は複雑なことだが、一ヶ月くらい使っていたら慣れるような気もする。

モナド上のfoldr

Control.Monadには、foldlのモナド版であるfoldMという関数が定義されている。

foldM :: (Monad m) => (a -> b -> m a) -> a -> [b] -> m a
foldM _ acc [] = return acc
foldM f acc (x:xs) = f acc x >>= \acc' -> foldM f acc' xs

これに倣って、foldrのモナド版を定義できる。Data.Foldableをインポートすれば実際に使用可能。

foldrM :: (Monad m) => (a -> b -> m b) -> b -> [a] -> m b
foldrM _ z [] = return z
foldrM f z (x:xs) = foldrM f z xs >>= f x

しかしこの定義は正格な言語のものに近く、IOなどの正格なモナドではfにかかわらずO(n)スタックを消費する。そこで、ここまでの流れに沿ってfの第二引数を継続だと思ってみる。するとこの引数の型はbではなく未実行の計算m bであるべきだろう。

foldrM' :: (Monad m) => (a -> m b -> m b) -> m b -> [a] -> m b
foldrM' _ z [] = z
foldrM' f z (x:xs) = f x (foldrM' f z xs)

foldrの定義と同じになってしまった。つまり、foldrはそのままでモナド上の(途中脱出可能な)ループを表現できる。例として、探索パスのリストから指定された名前の実行可能ファイルを探すIO関数は次のように書ける。

import Control.Applicative
import Control.Exception (tryJust)
import Control.Monad (guard)
import System.Directory hiding (findExecutable)
import System.FilePath
import System.IO.Error (isDoesNotExistError)

findExecutable :: [FilePath] -> String -> IO (Maybe FilePath)
findExecutable dirs name = foldr f (return Nothing) dirs
  where
    f dir cont = do
      found <- isExec file
      if found then return (Just file) else cont
      where file = dir </> name

    isExec file = either (const False) executable <$>
      tryJust (guard . isDoesNotExistError) (getPermissions file)

同様に上で定義したloopもモナド上で使える。たぶんloopが便利なのはモナドが絡んでいる場合だけだろう。

悪ノリ

リストいらなくね？

namedlet :: acc -> (acc -> (acc -> r) -> r) -> r
namedlet acc f = f acc $ \acc' -> namedlet acc' f

main = namedlet 0 $ \n cont ->
  if n > 100
    then return ()
    else do
      ln <- getLine
      putStrLn ln
      cont $ n + length ln

(この記事は2010-08-24を参考にして書きました)

*1:結果のリストを先頭から読み捨てていくという前提で