Inside-Python - lru_cache

functools --- 高階関数と呼び出し可能オブジェクトの操作 — Python 3.7.3 ドキュメント

pythonのlru_cacheはアノテーションとしてユーザの関数をデコレーションすることができ、 その呼び出しを監視し、結果をキャッシュするものである。

この実装はなかなか面白い。

単純なキャッシュアルゴリズムであれば、ハッシュテーブルでよいのだが、lru_cacheはデータのセット時に使われていないデータの追い出し作業が発生する。

つまり、get時には取得されたデータは「使われた」として順位の修正の必要があり、またセット時にキャッシュの限界に達したら最下位のデータを消す必要がある。

https://github.com/python/cpython/blob/master/Lib/functools.py#L496

実装自体はここにあるが、アノテーションのための実装もそこそこあり、またデータが無制限に入る場合、および1件もキャッシュできない場合には面白くないので 一番コアな部分を抽出する。

また、同期処理に関するところ, statsに関するところはコアな部分ではないのでいったん削った。

        def wrapper(*args, **kwds):
            key = make_key(args, kwds, typed)
            link = cache_get(key)
            if link is not None:
                # Move the link to the front of the circular queue
                link_prev, link_next, _key, result = link
                link_prev[NEXT] = link_next
                link_next[PREV] = link_prev
                last = root[PREV]
                last[NEXT] = root[PREV] = link
                link[PREV] = last
                link[NEXT] = root
                return result

            result = user_function(*args, **kwds)

            if full:
                # Use the old root to store the new key and result.
                oldroot = root
                oldroot[KEY] = key
                oldroot[RESULT] = result
                # Empty the oldest link and make it the new root.
                # Keep a reference to the old key and old result to
                # prevent their ref counts from going to zero during the
                # update. That will prevent potentially arbitrary object
                # clean-up code (i.e. __del__) from running while we're
                # still adjusting the links.
                root = oldroot[NEXT]
                oldkey = root[KEY]
                oldresult = root[RESULT]
                root[KEY] = root[RESULT] = None
                # Now update the cache dictionary.
                del cache[oldkey]
                # Save the potentially reentrant cache[key] assignment
                # for last, after the root and links have been put in
                # a consistent state.
                cache[key] = oldroot
            else:
                # Put result in a new link at the front of the queue.
                last = root[PREV]
                link = [last, root, key, result]
                last[NEXT] = root[PREV] = cache[key] = link
                # Use the cache_len bound method instead of the len() function
                # which could potentially be wrapped in an lru_cache itself.
                full = (cache_len() >= maxsize)
            return result

プログラム全体としては大きく二つに分かれ、以下のuser_functionのコールを境に前半はキャッシュの探索、 後半がキャッシュがヒットしなかった場合のキャッシュの保存になっている。

 result = user_function(*args, **kwds)

キャッシュの探索の部分を見てみると、キーの生成とキャッシュの取得後、キャッシュがヒットした場合には、何やらデータ構造をいじっている様子がわかる。

            key = make_key(args, kwds, typed)
            link = cache_get(key)
            if link is not None:
                # Move the link to the front of the circular queue
                link_prev, link_next, _key, result = link
                link_prev[NEXT] = link_next
                link_next[PREV] = link_prev
                last = root[PREV]
                last[NEXT] = root[PREV] = link
                link[PREV] = last
                link[NEXT] = root
                return result

ネタを明かしてしまうと、lru_cacheはハッシュテーブルに加えて、双方向の連結リストを用いてキャッシュを実装している。

f:id:mitsuo_0114:20190529232121p:plain

イメージ的にはこんな感じ。

                # Move the link to the front of the circular queue
                link_prev, link_next, _key, result = link
                link_prev[NEXT] = link_next
                link_next[PREV] = link_prev

まず、この3行ではキャッシュヒットしたデータを持つノードを連結リストの中から取り除く作業をしている。

双方向連結リストでは、自分の一つ前のノードの「次のノードへのポインタ」に対し、自分の次のノードのポインタを渡してあげ、 自分の次のノードの「前のノードへのポインタ」には、自分の前のノードのポインタを渡してあげる。

これによって、自らへの参照を取り除くことができる。

その後、取り除いたノードを一番最後に挿入する。

                last = root[PREV]
                last[NEXT] = root[PREV] = link
                link[PREV] = last
                link[NEXT] = root

rootはダミーオブジェクトであり、双方向連結リストではよく用いられるもので、 一番最後のオブジェクトの次、一番最初のオブジェクトの前に存在するイメージ。

ここでは、rootの前=一番最後のノードを取得し、この最後のノードと、rootの間に入れることで一番後ろにデータを挿入することになる。

これらの作業によって、キャッシュがヒットした場合、そのヒットしたデータを持つノードの順位を変更することが可能になっている。

次に後半の挿入を見てみる。

挿入はキャッシュがヒットしなかった場合に、user_functionの結果を格納するときにおこる。

一番最後の部分を最初に見ると、以下のように新しくデータを作成したうえで、キャッシュへの保存および、双方向連結リストの一番後ろにデータを持ってきていることがわかる。

                # Put result in a new link at the front of the queue.
                last = root[PREV]
                link = [last, root, key, result]
                last[NEXT] = root[PREV] = cache[key] = link
                # Use the cache_len bound method instead of the len() function
                # which could potentially be wrapped in an lru_cache itself.
                full = (cache_len() >= maxsize)

ここの部分が、新しくノードを作っているところ。4つは順番に前のノード、次のノード、キー、データの意味合い。

                link = [last, root, key, result]

最後に、fullであり、追い出しをする場合には以下。

                elif full:
                    # Use the old root to store the new key and result.
                    oldroot = root
                    oldroot[KEY] = key
                    oldroot[RESULT] = result

                    # Empty the oldest link and make it the new root.
                    # Keep a reference to the old key and old result to
                    # prevent their ref counts from going to zero during the
                    # update. That will prevent potentially arbitrary object
                    # clean-up code (i.e. __del__) from running while we're
                    # still adjusting the links.
                    root = oldroot[NEXT]
                    oldkey = root[KEY]
                    oldresult = root[RESULT]
                    root[KEY] = root[RESULT] = None
                    # Now update the cache dictionary.
                    del cache[oldkey]
                    # Save the potentially reentrant cache[key] assignment
                    # for last, after the root and links have been put in
                    # a consistent state.
                    cache[key] = oldroot

コメントを読むに、gcによってデータが意図せぬところで回収されないように参照カウンタに関して少し気を配っているようだが、 やっていることはデータの上書きである。

まず、現在rootを指しているノードに対して、今回のkeyとresultを挿入する。

                    oldroot = root
                    oldroot[KEY] = key
                    oldroot[RESULT] = result

次に、rootの次のノードを指し示し(参照カウンタをキープしたうえで)、ここのデータをNoneに書き換える。

これによって、先ほどまでrootがさしていたノードの次のノードが新しいrootになる。この辺り変数名が親切じゃない気もする。

                    root = oldroot[NEXT]
                    oldkey = root[KEY]
                    oldresult = root[RESULT]
                    root[KEY] = root[RESULT] = None

その後、キャッシュテーブルからも追い出したデータを削除し、新しく追加されたデータをキャッシュテーブルに追加する。

                    del cache[oldkey]
                    # Save the potentially reentrant cache[key] assignment
                    # for last, after the root and links have been put in
                    # a consistent state.
                    cache[key] = oldroot

なかなか興味深い実装で、ここまで有効にデータ構造を使ってるとは思わなかった。

確かに

「現在の最低順位のデータがどれかを知る」、

「任意の位置のデータを最高順位に持ってきて、それ以外のすべてのデータの順位を一つ下げる」

というあたりの要件を考えるにハッシュテーブルと、最初と最後をいじれる双方向連結リストが必要なんだなと思う。

Androidにもlru_cacheという名前の実装があるので今度実装を見てみたい。