F#でのメモ化 - クロージャからスレッドセーフ、複数の引数対応まで

F#での再帰処理と一緒に言及される、定番の話題がメモ化だ。よくあるのが、フィボナッチ数を求める処理をメモ化するのではなく、その関数自体をメモ化する話題だ。昨日の話題*1は、前者に言及している。今回の話題は、後者についてだ。

特別な工夫なく、ロジック通りにフィボナッチ数を求めるコードを書けば、その処理は無用な重複処理を孕むことになる。一度計算した順位のフィボナッチ数は、再計算するのではなく、その計算結果を再利用して出力することで、処理を効率化できる。つまり処理のメモ化だ。

同様に、複数回呼び出しされる関数について、その関数の出力結果と、その時の引数の組み合わせをメモして再利用すれば、やはり処理を効率化することができる。これが今回の話題となる、関数のメモ化だ。

関数のメモ化、メモ化関数
コード
参照

関数のメモ化、メモ化関数

次の関数はmemoizationは、引数fnとして与えられた関数をメモ化し、返り値としてメモ化関数を返す。

open System.Collections.Generic

let memoization fn =
  let cache = new Dictionary<_,_>()
  printfn "counter: %A" cache.Count
  
  fun arg ->
    printfn "%A is received" arg
    match cache.TryGetValue arg with
    | true, value -> value
    | false, _ ->
      let value = fn arg
      cache.Add(arg, value)
      value

関数memoizationeは、引数fnとして関数を受け取る。
ジェネリック型のkey, valueを持つDictionaryを定義する。
- ジェネリック型なので、Dictionaryは汎用的に機能する。

無名関数を定義する。→この関数がメモ化関数
- この関数は引数argを受け取る。→メモ化関数の引数となる
1. 引数argをキーとして、Dictionaryから値を取得する。
2. キーが存在したら、値を返す。
3. キーが存在しなかったら、関数fnに引数argを与えて実行する。
4. 引数argをキー、実行結果を値として、Dictionaryに追加する。
5. 実行結果を返す。

クロージャ

その実行結果は、一見、不思議な振る舞いに見える。まずcacheが生成され、「counter: 0」が出力される。これは初回のみ実行され、2度目以降は、続く無名関数から実行される。メモ化が実現しているのは、このcacheが永続しているからだ。そのため「counter: 」は出力されず、受け取った引数を出力し続ける。これがクロージャーと呼ばれる仕組みだ。

クロージャは関数とその関数が作られた環境という 2 つのものの組み合わせです。この環境は、クロージャが作られた時点でスコープ内部にあったあらゆる変数によって構成されています。この場合、myFunc は makeFunc が実行された時に作られた displayName 関数のインスタンスへの参照です。displayName のインスタンスはレキシカル環境への参照を保持し、そこに name 変数が存在します。このため、makeFunc が実行された時に、name 変数が残っていて "Mozilla" が console.log に渡されます。

クロージャ - JavaScript | MDN

実際にフィボナッチ数を求める関数をメモ化し、40番目のフィボナッチ数を求める処理を3回繰り返してみる。2回目、3回目の処理時間が、初回に比べて遥かに短いことが分かる。

🔎test.fsx

#load "fibonacci.fsx"
#load "memoization.fsx"

open fibonacci
open memoization
open System.Diagnostics

let memo_fn = memoization fib

let sw = new Stopwatch()

for i in 1 .. 3 do
  sw.Restart()
  memo_fn 40 |> printfn "%A"
  sw.Stop()
  sw.ToString() |> printfn "%s"

つまり、次に示すコードで”memoization expensive_fn”が返すのは、関数memoizationの”fun arg ->”で定義される無名関数であり、memo_fnは、その関数オブジェクト、あるいはそのインスタンス的な存在となる。

let memo_fn = memoization expensive_fn

すでにディクショナリが生成された状態でインスタンス化されるので、続くmemo_fnの複数回の呼び出しでは、次のコードは実行されないのだ。結果として、出力も初回分のみ「counter: 0」しか表示されない。

let cache = new Dictionary<_,_>()
printfn "counter: %A" cache.Count

次のサイトで紹介されているサンプル関数memoizeをよく観察すると、匿名関数が()で括られているのが分かる。コメントなどで明示されていることではないのだが、この無名関数が関数memoizeの戻り値であることを、暗に指し示す意図なのかもしれない。

  (fun x ->
    match cache.TryGetValue x with
    | true, v -> v
    | false, _ -> let v = fn (x)
                  cache.Add(x,v)
                  v)

memoize | F# Snippets

スレッドセーフ

immutableであるという特性は、スレッドごとの整合性を担保できるため、特に並列処理において有効に作用する。メモ化関数をスレッドセーフとするには、キャッシュとして利用するDictionaryに代えて、ConcurrentDictionaryを利用すればよい。こうすることで複数のスレッドから参照可能なキャッシュが生成される。

同時に、メモ化関数からの関数呼び出しでは、キーワードlazyを用いる。こうすることで同時に呼び出しされたときに、関数が2度実行されることを防ぐ。

open System.Collections.Concurrent

let thread_safe_memoization fn =
    let cache = new ConcurrentDictionary<_,_>()
    printfn "counter: %A" cache.Count

    fun arg ->
        printfn "%A is received" arg
        match cache.TryGetValue arg with
        | true, value -> value
        | false, _ ->
            let value = lazy fn arg
            cache.TryAdd(arg, value.Force()) |> ignore
            value.Force()

そして、この処理は次のように短く書き直すこともできる。

open System.Collections.Concurrent

let shorter_memoization fn = 
  let cache = new ConcurrentDictionary<_,_>()
  fun arg -> cache.GetOrAdd(arg, lazy fn arg).Value

複数の引数への対応

let memoization_2args fn = shorter_memoization (fun (arg1, arg2) -> fn arg1 arg2) |> fun fn' arg1 arg2 -> fn' (arg1, arg2)

let memoization_3args fn = shorter_memoization (fun (arg1, arg2, arg3) -> fn arg1 arg2 arg3) |> fun fn' arg1 arg2 arg3 -> fn' (arg1, arg2, arg3)

そして複数の引数を取る関数についての対応だ。引数が何個あろうが、基本的な考えは変わらない。

1つ目の無名関数
- 単一のtupleを引数とする
  - このtupleの各要素が、メモ化関数の引数となる
- メモ化関数が、2つ目の無名関数の引数となる

fun (arg1, arg2) -> fn arg1 arg2

2つ目の無名関数
- メモ化関数を引数とする
- メモ化関数の引数を、単一のtupleへまとめる
- 単一のtupleを引数とするメモ化関数を返す

fun fn' arg1 arg2 -> fn' (arg1, arg2)

3つの引数を取る関数のメモ化をテストしてみる。その関数のコンセプトは、

第1引数：開始メッセージ
第2引数：終了メッセージ
第3引数：待機時間（ミリ秒）

開始メッセージを表示し、待機時間停止後に終了メッセージを表示する。最後に文字列”return”を返す。この関数をメモ化し、ループで3回実行する。

🔎test_3args.fsx

#load "memoization.fsx"
#load "fibonacci.fsx"

open memoization
open fibonacci
open System.Diagnostics
open System.Threading

let expensiveFunction start_msg end_msg (sleep_time : int) =
  printfn "%s" start_msg
  Thread.Sleep sleep_time
  printfn "%s" end_msg
  "return"

let test_3args = memoization_3args expensiveFunction
let sw = new Stopwatch()

for i in 1 .. 3 do
  sw.Restart()
  test_3args "start" "end" 3000 |> printf "%s\t\t"
  sw.Stop()
  sw.ToString() |> printfn "%s"