フューチャーAdvent Calendar 2022 の24日目です。

はじめに

マルチコア対応のOCaml 5.0がリリースされましたね！¹
この記事ではOCaml 5.0でサポートされる目玉機能？の一つであるEffect Handlerについて調べてみました。

私はOCamlをほぼ何も知りません。
したがってOCamlそのものに関する記述の相当怪しいと思います。
指摘は大歓迎です。

Effect Handlerとは？

Effect Handlerとは一般化された例外です。何のことか分からないですよね。。

皆さんのよく知る例外ではハンドリングする側は例外クラスのインスタンスのような値を取れると思いまが、Effect Handlerでは値に加えて継続も受け取ることができます。継続という用語は聞き馴染みがない方も多いかも知れません。継続とは一言でいうと残りの計算のことです。昔にしょぼい記事を書いたのでよければご覧ください。

OCamlのEffect Handlerは継続といっても、One shot delimited continuationという継続を扱います。One shotというのはその継続は最大一回しか実行できないということです。delimitedというのは継続の範囲が区切られているということで、OCamlでは「perform(javaでいうthrowに相当)したところから、performを呼んだ関数の終わりまで」です。delimitedでない継続というのは「performされてからプログラムが終了するまで」を指します。よく分からないですよね。この後に紹介する例を見れば多少イメージがつくのではないかと思います。

Effect Handlerの嬉しさ

さてEffect Handlerが何かはふんわり分かったとします。疑問に思うのがこれはなにが嬉しいのかということです。私は正直よく分かっていませんが、OCaml 5の開発を推進してきたSivaramakrishnanさんの動画によるとOne shot continautionsを使うと多くの並行処理のメカニズムを表現できるらしいです。相当表現力豊かなプリミティブということなのでしょう。

なのでEffect Handlerをベースに並行処理のライブラリを作ることができます。Ocaml 5ではそのようなライブラリが実際にいろいろ用意されているようです。SivaramakrishnanさんはOcaml 5の並行処理ライブラリについて以下のようなメリットがあると主張していました。

GoやHaskellと違い、ランタイムに並行処理の仕組みを組み込まないことで、ランタイムの開発の負担になることを避けることができ、またケースバイケースで適したライブラリを使える。
Effect Hanlderを使うと非同期IOをネイティブサポートできるらしく、それはRustやJavaScriptのasyncのように同期、非同期の色がつかない。

Effect Handlerを使ったサンプル

環境

WSLのUbuntu 20.04.5 LTS (GNU/Linux 5.10.102.1-microsoft-standard-WSL2 x86_64)でプログラムを実行しました。

環境構築はチュートリアルのInstallationに習っています。チュートリアルではocaml-base-compiler.5.0.0~alpha1を使っていますが、私はocaml-base-compiler.5.0.0~beta1で動作させました。

例外風に使う

open Effect
open Effect.Deep

type _ Effect.t += Foo: int -> int Effect.t

let f () =  1  + (perform (Foo 1))

let () = print_int (try_with f () {
    effc = fun (type a) (eff: a t) ->
        match eff with
        | Foo n-> Some (fun (k: (a, _) continuation) -> continue k (n+1))
        | _ -> None
})

こちらはSivaramakrishnanさんのスライドに載っていた擬似コードを実際に動くように修正したものです。
上記は例外を投げている風のコードです。
fという関数を定義して、それをtry_with関数に渡しています。
このコードを実行されると何が表示されるでしょうか？

前提知識を説明しながら、コードの動作を説明します。

type _ Effect.t += Foo: int -> int Effect.t

Effectを定義しているところがtype _ Effect.t += Foo: int -> int Effect.tです。これはExtensible variant typesと呼ばれる構文のようです。私の理解ではこれは拡張可能なEnumです。拡張可能なので、未知のケースがあるかもしれず、そのためマッチするときは常に_ ->のケースが必要になるようです。

let f () =  1  + (perform (Foo 1))

次に注目するのはf関数内で呼ばれているperformです。performはEffectを発生させる関数です。例外で言えばraiseやthrowに相当するものです。なのでperformが呼ばれると制御がハンドリング側に移ります。このコードで言えば try_withです。
performの型は'a t -> 'aです。Fooがint -> int Effect.tなので、Foo 1はint Effect.t、ゆえにperform (Foo 1)はintを返します。この値はその継続を実行するときに渡す値です。

let () = print_int (try_with f () {
    effc = fun (type a) (eff: a t) ->
        match eff with
        | Foo n-> Some (fun (k: (a, _) continuation) -> continue k (n+1))
        | _ -> None
})

try_withはtry-catchに相当するものようです。Ocaml lsp Serverを導入しているVSCodeでカーソルを当てるとtry_with f v h runs the computation f v under the handler hという説明が出てきます。try_withに渡す関数内でperform (Foo 1)とすることでeffはFoo 1になります。

match eff with部分で注目すべき点はkとcontinueです。kは(a, _) continuation)という型で、performが呼ばれた時点から関数の終わりまでの継続です。具体的にいうと1とpeform (Foo 1)の戻り値を足すという処理です。この継続は型aを引数にとって型_を返します。今回の場合はaはintで_もintです。

(type a)が何なのかや、(a, _) continuation)のアンダーバーがどういう意味なのかは分かっていません。
だれか教えてください。。

continueは('a, 'b) continuation -> 'a -> 'bという型を持っています。継続に引数を渡して実行を再開し、戻り値を得ます。
continueにn+1(=2)を渡しているので、perform (Foo 1)の戻り値は2となり、継続の戻り値は1+2で3になります。
よってtry_withの戻り値は3となり、print_int 3が実行されることになります。

まとめると、以下のような動作となります。

perform (Foo 1)がハンドリングされて1足されて2が返ってくる
1+2で3になり、それが (continue k (n+1))の戻り値となる
よって3がプリントされる

継続がなんなのかふんわり分かったでしょうか？

String Effect

open Effect
open Effect.Deep

type _ Effect.t += E : string Effect.t

let print () =
  print_string "0";
  print_string (perform E);
  print_string "3"

let main () =
  match_with print () {
    retc = (fun res -> res);
    exnc = raise;
    effc = fun (type a) (eff: a t) ->
      match eff with
      | E ->
          Some (fun (k: (a, _) continuation) ->
           print_string "1";
           continue k "2";
           print_endline "4")
      | _ -> None
  }

let () =
  main ()

こちらはicfp22というカンファレンスで行われていたらしいチュートリアルの資料のコードの引用です。

match_withとtry_withは似たようなものに見えますが、使い分けはよくわかりません。

さてコードはなにを出力するでしょうか？
答えは01234です。

処理の動作を説明します。

まずprint関数のprint_string "0";が走るので0が出力されます。
その後、print_string (perform E)が走るのでEffectが発生します。
そしてEffectがeffc = fun (type a) (eff: a t) -> ~~でハンドリングされます。
1が表示されてcontinue k "2";で継続に2を渡して実行します。
すると(perform E)が2を返すのでprint_string (perform E);は2をプリントします。
その後3を表示し、継続の実行を終了します。継続はunitを返して、continue k "2";の戻り値はunitになります。
最後に4を改行とともに出力して処理が終了します。

Round-robin execution scheduler

最後にやや実用的な例を紹介して終わりにします。並行処理を実行するプログラムです。

open Effect
open Effect.Deep

type _ Effect.t += Fork : (unit -> unit) -> unit Effect.t
type _ Effect.t += Yield : unit Effect.t

let fork f = perform (Fork f)
let yield () = perform Yield

let run main =
  let run_q = Queue.create () in
    let enqueue k = Queue.push k run_q in
      let deque () =
        if Queue.is_empty run_q then ()
        else continue (Queue.pop run_q) ()
      in
      let rec spawn f =
        match_with f () {
          retc = (fun () -> deque ());
          exnc = (fun e ->
            print_string (Printexc.to_string e);
            deque ());
          effc = fun (type a) (e: a Effect.t) ->
            match e with
            | Yield -> Some (fun (k: (a, unit) continuation) -> enqueue k; deque ()) 
            | Fork f -> Some (fun (k: (a, unit) continuation) -> enqueue k; spawn f)
            | _ -> None
        }
          in
          spawn main

let log = Printf.printf

let rec f id depth =
  log "Starting number %i\n!" id;
  if depth > 0 then begin
    log "Forking number %i\n!" (id * 2 + 1);
    fork (fun () -> f (id * 2 + 1) (depth - 1));
    log "Forking number %i\n!" (id * 2 + 2);
    fork (fun () -> f (id * 2 + 2) (depth - 1));
  end else begin
    log "Yielding number %i\n!" id;
    yield ();
    log "Resumed number %i\n!" id;
  end;
  log "Finishing number %i\n!" id


let () = run (fun () -> f 0 2)

こちらはEffect Hanlderの応用コードを紹介する記事から引用しています。

run関数の概要を説明します。

Queueを持っていて、このQueueには継続を詰めていきます。
enqueueとdequeueという関数が定義されています。enqueueは継続をpushします。dequeueは継続を取り出して実行します。
メインの関数はspawnという関数です。このプログラムにわたす関数はForkとYieldをperformすることが期待されていて、Yieldの場合は継続をqueueにいれてdequeueを呼び、Fork fの場合は継続をqueueに入れて、spawn fを再帰的に呼び出します。つまり処理を一旦やめて、Queueに詰まっている継続を実行したいときはYield、処理を一旦やめて新しい処理を呼びたい場合はFork fをperformします。

コード末尾のlet () = run (fun () -> f 0 2)ではスケジュールプログラムを実行しています。
fという関数は上記のような木構造をforkしながら巡って、木の子供に達したらyieldするプログラムです。データ構造はなく、fに渡す数値で木のある位置にいると考えています。
以下のような木を走査しています。

実行結果は以下のようになります。並行して処理が実行できているのが分かると思います。私は紙にかきながら処理を辿っていましたが「Starting number 2」あたりで頭がこんがらがってきて辿るのをやめました。

実行結果

Starting number 0
!Forking number 1
!Starting number 1
!Forking number 3
!Starting number 3
!Yielding number 3
!Forking number 2
!Starting number 2
!Forking number 5
!Starting number 5
!Yielding number 5
!Forking number 4
!Starting number 4
!Yielding number 4
!Resumed number 3
!Finishing number 3
!Finishing number 0
!Forking number 6
!Starting number 6
!Yielding number 6
!Resumed number 5
!Finishing number 5
!Finishing number 1
!Resumed number 4
!Finishing number 4
!Finishing number 2
!Resumed number 6
!Finishing number 6
!⏎

ちなみに以下のようにQueueをStackにすると以下のように深さ優先探索でたどるようになります。

+ let run main =
+  let run_q = Stack.create () in
+    let enqueue k = Stack.push k run_q in
+      let deque () =
+        if Stack.is_empty run_q then ()
-  let run_q = Queue.create () in
-    let enqueue k = Queue.push k run_q in
-      let deque () =
-        if Queue.is_empty run_q then ()
-        else continue (Stack.pop run_q) ()