再帰関数のスタックオーバーフローを倒す話 その2.5

F#

連載目次

はじめに

再帰関数のスタックオーバーフローを倒す話 その1ではCPS変換について、 再帰関数のスタックオーバーフローを倒す話 その1.5では末尾呼び出しについて、 再帰関数のスタックオーバーフローを倒す話 その2では2種類の末尾呼び出しの最適化について話しました。

今回は、継続渡しスタイルがつらい人*1をモナドで救う話をします。 最初にネタバレすると、結局F#では救われないことになりますが。 まぁ、これはあくまでその1のおまけとして読んでください。

前提知識は、上記記事(特にその1)を理解していることです。

継続渡しスタイルつらい話

階乗を求める関数factを考えましょう。

let rec fact n =
  if n = 0I then 1I
  else
    n * (fact (n - 1I))

これを使って、nの階乗とmの階乗の和を求める関数は、このように書けます。

let f n m =
  let factN = fact n
  let factM = fact m
  factN + factM

このfactを継続渡しスタイルにしたfactCpsを考えてみます。

let rec factCps n cont =
  if n = 0I then 1I |> cont
  else
    factCps (n - 1I) (fun pre ->
    n * pre |> cont)

これを使って、同様にnの階乗とmの階乗の和を求める関数を書くと、こうなります。

let fCps n m cont =
  factCps n (fun factN ->
  factCps m (fun factM ->
  factN + factM |> cont))

元の普通のF#コードに比べると、このコードをすらっと読み下せる人は格段に減ります。 このコードをすらっと書ける人となるとさらに減ります。

このあたりが継続渡しスタイルのつらいところです。

コンピュテーション式が世界を救う

階乗を求める関数を使った2つのコードをちょっと並べてみましょう。

(* 普通のやつ *)
let f n m =
  let factN = fact n
  let factM = fact m
  factN + factM

(* CPS版 *)
let fCps n m cont =
  factCps n (fun factN ->
  factCps m (fun factM ->
  factN + factM |> cont))

その1でもみましたけど、letのネスト*2がfunのネストになっています。 コンピュテーション式の変換規則を知っている人であれば、これはそのままlet!の変換規則に見えるでしょう。

そうです、継続渡しスタイルはコンピュテーション式の後ろに隠すことができるのです！

継続もにゃど

type ContinuationBuilder() =
  member __.Return(x) = (fun cont -> x |> cont)
  member __.Bind(cpsFun, rest) = (fun cont -> cpsFun (fun x -> rest x cont))

let cont = ContinuationBuilder()

さぁ準備は整いました！ contコンピュテーション式を使ってfCpsを書き直してみます。

(* contコンピュテーション式版 *)
let fCps n m = cont {
  let! factN = factCps n
  let! factM = factCps m
  return factN + factM
}

これなら、継続渡しスタイルの関数でも使える気がしませんか？

注目すべき点はいくつかあります。 まずは、fCpsから引数が一つ減ったように見えるようになりました。 実際は減っていないんですが、contコンピュテーション式を使う限り、隠された引数を意識する必要はありません。

次に、factCpsは最後の引数として継続を受け取るはずなのに、あたかもそんな引数ないかのように呼び出しているように見えます。 これがcontコンピュテーション式の力で、継続を引き回す処理はContinuationBuilderのBindの中に隠蔽されています。

let! factN = factCps n
...

は

cont.Bind(factCps n, fun factN -> ...)

のように展開されます。

member __.Bind(cpsFun, rest) = (fun cont -> cpsFun (fun x -> rest x cont))

Bindの定義はこのようになっており、factCps nの部分はcpsFunが受け取ります。 このcpsFunの部分にfactCps nを埋め込んでみると・・・

factCps n (fun factN -> ...)

と、どこかで見たことのある形が現れますね。 このように、継続渡しスタイルの関数に継続を渡す部分は、Bindが裏でやってくれるのです。

これで継続渡しスタイルもこわくない！

factCpsをcontコンピュテーション式で書き直す

ここまで来たらfactCpsも書き直してしまいましょう。

let rec factCps n cont =
  if n = 0I then 1I |> cont
  else
    factCps (n - 1I) (fun pre ->
    n * pre |> cont)

これがこうです！

let rec factCps n = cont {
  if n = 0I then
    return 1I
  else
    let! pre = factCps (n - 1I)
    return n * pre
}

だいぶ普通のコードっぽく読めるようになったのではないでしょうか？ しかし、これには罠があります。

contコンピュテーション式の罠

このfactCpsに50000Iとid関数を渡して起動すると、スタックオーバーフローが起きます。 なぜこのようなことが起こるのでしょうか？ factCpsのコンピュテーション式を展開してみましょう。

let rec factCps n =
  let b = cont
  if n = 0I then
    b.Return(1I)
  else
    b.Bind(factCps (n - 1I), (fun pre -> b.Return(n * pre)))

あれれ、factCpsの再帰呼び出し部分が末尾じゃなくなってます！ そらスタックオーバーフロー起きるわー。 Bindにinlineを付けてみましょう。

Bindのinline化

ContinuationBuilderを修正します。

type ContinuationBuilder() =
  member __.Return(x) = (fun cont -> x |> cont)
  member inline __.Bind(cpsFun, rest) = (fun cont -> cpsFun (fun x -> rest x cont))

let cont = ContinuationBuilder()

これで、factCpsは以下のように展開されるようになります。

let rec factCps n =
  let b = cont
  if n = 0I then
    b.Return(1I)
  else
    let f = factCps (n - 1I)
    (fun cont -> f (fun x -> b.Return(n * x) cont))

factCpsが！ラムダ式の！外にいる！！！ これはもうどうしようもありませんね。

まとめ

ということで、継続渡しスタイルのつらみをコンピュテーション式で軽減する話でした。 そして、継続渡しスタイルの末尾再帰関数をコンピュテーション式を使って書き直すと、末尾再帰にならずに死ぬ、という話もしました。 みなさん、注意しましょう。

*1:読めない、書けない

*2:letがネストしているように見えないかもしれませんが、let式のネストになっています