Haskell の ArrowLoop を、正格評価である F# に移植する

まえがき

Arrow を F# に移植しようとして ArrowLoop で相当悩んだので、自分用のメモとして残します。

移植する関数

初見で理解できなさすぎる、ArrowLoop を構成する関数 loop がこちら。

class Arrow a => ArrowLoop a where
    loop :: a (b, d) (c, d) -> a b c

instance ArrowLoop (->) where
    loop :: ((b, d) -> (c, d)) -> b -> c
    loop f b = let (c, d) = f (b, d) in c

一見何の変哲もなさそうですが、Haskell が非正格評価だからこそ実現できている↓この部分を、正格評価であるF#に移植しようとして悩みました。

loop f b = let (c, d) = f (b, d) in c

Phase1: 再帰関数として移植する

Step1: 手元の言語で書いてみる

上記関数をとりあえずF#で書きつつ、役割に応じて変数名をつけます。

let loop (f: 'T * 'Feedback -> 'U * 'Feedback) (arg: 'T) : 'U =
    let result, feedback = f (arg, feedback)
    result

コンパイルエラー : 値またはコンストラクター'feedback'が存在しません。

見た目はほぼ同じですが、コンパイルエラーが出ました。
それもそのはず。正格評価の言語では、変数を定義するより前に使うことはできません。

Step2: 構造を理解し分解する

コンパイルできなかったのは、変数feedbackが fの戻り値から引数に文字通りフィードバックされている ことが原因です。

fの呼び出し後にしか得られないfeedbackをfの呼び出し前に用意しなければならない。
はい、無理です。
そもそもこの関数fは非正格評価であるのをいいことに、「フィードバックの生成」と「結果の生成」を同時にやっているので、役割ごとに関数を分割して問題をシンプルにします。

let loop (f: 'T * 'Feedback -> 'Feedback) (g: 'T * 'Feedback -> 'U) (arg: 'T) : 'U =
    let feedback = f (arg, feedback)
    g (arg, feedback)

こうしても元の動作は維持されているはず。
分解したことで関数gは何の不安もない普通の関数になりましたので、今後は無視します。
(元の定義のままでも実装できますが、ループ部分だけを取り出したほうがHaskellの動作に近くなります¹)

これ以降で扱うloop

let loop (f: 'T * 'Feedback -> 'Feedback) (arg: 'T) : 'Feedback =
    let feedback = f (arg, feedback)
    feedback

Step3: 不動点コンビネータに変換する

変数feedbackを再帰的に参照する代わりに、関数loopを再帰関数にすることでコンパイルできるようになります。

let rec loop (f: 'T * 'Feedback -> 'Feedback) (arg: 'T) : 'Feedback =
    f (arg, loop f arg) 

ただし、このままだとloopの実行序盤でloop自身が呼び出されてしまい、無限ループが起きてしまいます。
この解決策として遅延評価を使います²。

let rec loop (f: 'T * Lazy<'Feedback> -> 'Feedback) (arg: 'T) : 'Feedback =
    f (arg, lazy (loop f arg))

または、 'Feedback が関数型 'A -> 'R であると決め打ちしてしまえばこのようにも書けます³。無限ループするコードと見た目がほぼ同じですが、 loop f arg1 が部分適用になったことで遅延評価されるようになり、無限ループを回避しています。

let rec loop (f: 'T * ('A -> 'R) -> 'A -> 'R) (arg1: 'T) (arg2: 'A) : 'R =
    f (arg1, loop f arg1) arg2

この関数 loop を使って階乗を求める関数を作ってみました。

let factorial x =
    loop (fun (_, next) x -> if x <= 0 then 1 else x * next (x - 1)) () x

Phase2: 末尾最適化がかかる形で移植する

再帰関数として移植できれば基本的には満足なのですが、F#においてはPhase1の書き方だと末尾再帰最適化が効かなくなってしまいます。Phase2ではこれを末尾最適化が効く形に直します。
まずは再帰関数における末端の呼び出しを区別すべく、ループを続けるか抜けるかを表す型を用意します⁴。

type LoopState<'Result, 'State> =
    | Break of result: 'Result
    | Continue of state: 'State

あとは流れで実装できます。
末尾最適化がかかる形に変形するための規則やらなんやらにより、初期値がない Feedback ではなく初期値を与える State を使います(追記: アキュムレータと書いたほうが分かりやすかったかもしれません)。

[<TailCall>]
let rec loop (f: 'T * 'State -> LoopState<'Result, 'State>) (stat: 'State) (arg: 'T) : 'Result =
    match f (arg, state) with
    | Break(result) -> result
    | Continue(state) -> loop f state arg

これで末尾最適化がかかる関数に変形できました。正格評価の言語で使うなら、現実的にはこちらの定義を採用するほうがよいでしょう。
このバージョンを使った場合の階乗を求める関数は次のようになります。

let factorial x =
    loop
        (fun (x, (i, acm)) -> if x < i then Break(acm) else Continue(i + 1, acm * i))
        (1, 1)
        x

この関数を定義に採用した場合のHaskellによるArrowLoopは次の通りです⁵。Stateの初期値はふつう定数になるのと、元のArrowLoopでもフィードバック型は戻り値のArrowから消えていたことを考え、stateの初期値はArrow内で受け取るのではなくloopの引数にしました。

class Arrow a => ArrowLoop a where
    loop :: a (b, d) (Either c d) -> d -> a b c

1. feedbackとresultの"スコープ"が異なるため。feedbackが参照されるたびにresultも生成されてしまうのは、HaskellのArrowLoopとしても想定している挙動ではないでしょう。 ↩

2. 遅延評価を使うならHaskellでいい気がします。F#は遅延評価するために型が変わって不便ですね。 ↩

3. 'Feedback の有用な使い方のひとつが再帰関数なので、これで十分なことが多いです。どうせ正格評価の言語で無限リストなんて使いません。 ↩

4. Result 型を流用することもできますが、Errorでループを抜けるのは気持ち悪いので別途用意しました。 ↩

5. 本当はF#で書きたかったけど高階型引数に対応していないのでね… ↩