lru.erl

%% =====================================================================
%% Redistribution and use in source and binary forms, with or without
%% modification, are permitted provided that the following conditions
%% are met:
%%
%% 1. Redistributions of source code must retain the above
%% copyright notice, this list of conditions and the following
%% disclaimer.
%%
%% 2. Redistributions in binary form must reproduce the above
%% copyright notice, this list of conditions and the following
%% disclaimer in the documentation and/or other materials provided
%% with the distribution.
%%
%% 3. The name of the author may not be used to endorse or promote
%% products derived from this software without specific prior
%% written permission.
%%
%% THIS SOFTWARE IS PROVIDED BY THE AUTHOR ``AS IS'' AND ANY EXPRESS
%% OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
%% WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
%% ARE DISCLAIMED. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY
%% DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
%% DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE
%% GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
%% INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY,
%% WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING
%% NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
%% SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
%%
%% @author Hunter Morris <hunter.morris@smarkets.com>
%% @copyright 2009 Smarkets Limited.
%% @end
%% =====================================================================

%% TODO: Runtime selection of map type

%% @doc Erlang LRU
%%
%% This module implements a least-recently-used data structure.  It is
%% effectively a list of values where the head of the list is the most
%% recently used value.  When a value is used, it is moved to the head
%% of the list.  The last element in the list is the
%% least-recently-used value. This structure is often used in a cache
%% to decide which values to evict when the cache becomes full.
%%
%% This module uses a dict to implement the LRU efficiently for large
%% values, but can be configured to use gb_trees instead if the
%% compiler constant LRU_WITH_GB_TREES is defined.
%% @end

-module(lru).
-author('Hunter Morris <hunter.morris@smarkets.com>').

-export([empty/0]).
-export([null/1, size/1, to_list/1, member/2]).
-export([hit/2, delete/2, pop/1, last/1]).

-record(lru, {
          head,
          last,
          map
         }).

-record(lru_item, {
          prev,
          next
         }).

%% @spec () -> lru()
%%
%% @type lru() = {'lru', term(), term(), map()}
%% @type map() = gb_tree() | dict()
%%
%% @doc Returns an empty LRU
%% O(1)

empty() ->
    #lru{map=map_empty()}.

%% @spec (lru()) -> bool()
%%
%% @doc Returns true if the LRU is empty
%% O(1)

null(#lru{head=undefined}) ->
    true;
null(_) ->
    false.

%% @spec (lru()) -> integer()
%%
%% @doc Returns the number of elements in the LRU
%% O(1)

size(#lru{map=Map}) ->
    map_size(Map).

%% @spec (term(), lru()) -> lru()
%%
%% @doc Inserts a value into an LRU.  If the value is already in the
%% LRU, it becomes the head of the list.
%% O(log n)

hit(Value, #lru{map=Map}=Lru) ->
    hit1(Value, Lru, map_lookup(Value, Map)).

hit1(_, Lru, {ok, #lru_item{prev=undefined}}) ->
    Lru;
hit1(Value, #lru{head=Head, map=Map}, {ok, #lru_item{prev=Prev, next=undefined}}) ->
    #lru{map=map_insert(Value, #lru_item{next=Head}, update_next(Prev, undefined, update_prev(Head, Value, Map))),
         head=Value,
         last=Prev};
hit1(Value, #lru{head=Head, last=Last, map=Map}, {ok, #lru_item{prev=Prev, next=Next}}) ->
    Map1 = map_insert(Value, #lru_item{next=Head},
                      update_next(Prev, Next,
                                  update_prev(Next, Prev,
                                              update_prev(Head, Value, Map)))),
    #lru{head=Value, last=Last, map=Map1};
hit1(Value, Lru, error) ->
    insert(Value, Lru).

%% @spec (lru()) -> [term()]
%%
%% @doc Returns a list of the members of the LRU in newest-first
%% order.
%% O(n(log n))

to_list(#lru{head=Head, map=Map}) ->
    case map_null(Map) of
        true ->
            [];
        false ->
            [Head|to_list1(Head, Map)]
    end.

to_list1(Current, Map) ->
    case map_lookup(Current, Map) of
        {ok, #lru_item{next=undefined}} ->
            [];
        {ok, #lru_item{next=Value}} ->
            [Value|to_list1(Value, Map)]
    end.

%% @spec (term(), lru()) -> bool()
%%
%% @doc Returns true if the given element is in the LRU.
%% O(log n)

member(Value, #lru{map=Map}) ->
    map_member(Value, Map).

%% @spec (term(), lru()) -> lru()
%%
%% @doc Removes an element from the LRU if it exists.
%% O(log n)

delete(_, #lru{head=undefined}=L) ->
    L;
delete(Value, #lru{head=Value, last=Last, map=Map}=L) ->
    {#lru_item{next=Next}, Map1} = map_update_lookup(fun const_undefined/2, Value, Map),
    L#lru{head=Next, last=case Next of undefined -> undefined; _ -> Last end, map=Map1};
delete(Value, #lru{last=Value}=L) ->
    pop(L);
delete(Value, #lru{map=Map}=L) ->
        case map_update_lookup(fun const_undefined/2, Value, Map) of
            {#lru_item{prev=Prev, next=Next}, Map1} ->
                L#lru{map=update_next(Prev, Next, update_prev(Next, Prev, Map1))};
            _ ->
                L
        end.

%% @spec (lru()) -> term()
%%
%% @doc Returns the last element of the LRU.  Crashes if the LRU is
%% empty.
%% O(1)

last(#lru{last=undefined}) ->
    throw(empty);
last(#lru{last=Last}) ->
    Last.

%% @spec (lru()) -> lru()
%%
%% @doc Removes the last element of the LRU.  Crashes if the LRU is
%% empty.
%% O(log n)

pop(#lru{last=undefined}) ->
    throw(empty);
pop(#lru{last=LastValue, map=Map}=L) ->
    case map_lookup(LastValue, Map) of
        {ok, #lru_item{prev=undefined}} ->
            empty();
        {ok, #lru_item{prev=NewLast}} ->
            Map1 = map_delete(LastValue, Map),
            NewMap = update_next(NewLast, undefined, Map1),
            L#lru{last=NewLast,
                  map=NewMap}
    end.

%% Item manipulation
update_prev(ToUpdate, NewValue, Map) ->
    F = fun(Item) when is_record(Item, lru_item) -> Item#lru_item{prev=NewValue} end,
    map_update(F, ToUpdate, Map).

update_next(ToUpdate, NewValue, Map) ->
    F = fun(Item) when is_record(Item, lru_item) -> Item#lru_item{next=NewValue} end,
    map_update(F, ToUpdate, Map).

%% O(log n)
insert(Value, #lru{last=undefined}) ->
    #lru{head=Value, last=Value, map=map_singleton(Value, #lru_item{})};
insert(Value, #lru{last=Last, head=Head, map=Map}) ->
    Map1 = map_insert(Value, #lru_item{next=Head}, update_prev(Head, Value, Map)),
    #lru{head=Value, last=Last, map=Map1}.

%% Map utility functions
-ifdef(LRU_WITH_GB_TREES).

-compile({inline, [{map_empty, 0},
                   {map_lookup, 2},
                   {map_size, 1},
                   {map_insert, 3},
                   {map_singleton, 2},
                   {map_null, 1},
                   {map_member, 2},
                   {map_delete, 2},
                   {const_undefined, 2}]}).

map_empty() ->
    gb_trees:empty().

map_lookup(Key, Map) ->
    case gb_trees:lookup(Key, Map) of
        {value, V} ->
            {ok, V};
        none ->
            error
    end.

map_size(Map) ->
    gb_trees:size(Map).

map_update(Fun, Key, Map) when is_function(Fun, 1) ->
    case gb_trees:lookup(Key, Map) of
        none ->
            Map;
        {value, Value} ->
            gb_trees:update(Key, Fun(Value), Map)
    end.

map_insert(Key, Value, Map) ->
    gb_trees:enter(Key, Value, Map).

map_singleton(Key, Value) ->
    gb_trees:from_orddict([{Key, Value}]).

map_null(Map) ->
    gb_trees:is_empty(Map).

map_member(Key, Map) ->
    gb_trees:is_defined(Key, Map).

map_delete(Key, Map) ->
    gb_trees:delete(Key, Map).

map_update_lookup(F, K, Map) ->
    case map_lookup(K, Map) of
        error ->
            {undefined, Map};
        {ok, V} ->
            case F(K, V) of
                undefined ->
                    {V, map_delete(K, Map)};
                V1 ->
                    {V1, map_insert(K, V1, Map)}
            end
    end.

-else. %% dict module as map is default

-compile({inline, [{map_empty, 0},
                   {map_lookup, 2},
                   {map_size, 1},
                   {map_insert, 3},
                   {map_singleton, 2},
                   {map_null, 1},
                   {map_member, 2},
                   {map_delete, 2},
                   {const_undefined, 2}]}).

map_empty() ->
    dict:new().

map_lookup(Value, Map) ->
    dict:find(Value, Map).

map_size(Map) ->
    dict:size(Map).

%% TODO: How badly does the exception impact performance?
map_update(Fun, Key, Map) when is_function(Fun, 1) ->
    try dict:update(Key, Fun, Map)
    catch error:_ -> Map
    end.

map_insert(Key, Value, Map) ->
    dict:store(Key, Value, Map).

map_singleton(Key, Value) ->
    dict:from_list([{Key, Value}]).

map_null(Map) ->
    dict:size(Map) =:= 0.

map_member(Key, Map) ->
    dict:is_key(Key, Map).

map_delete(Key, Map) ->
    dict:erase(Key, Map).

map_update_lookup(F, K, Map) ->
    case map_lookup(K, Map) of
        error ->
            {undefined, Map};
        {ok, V} ->
            case F(K, V) of
                undefined ->
                    {V, map_delete(K, Map)};
                V1 ->
                    {V1, map_insert(K, V1, Map)}
            end
    end.
-endif.

const_undefined(_, _) ->
    undefined.

%% Unit tests
-ifdef(TEST).

-include_lib("eunit/include/eunit.hrl").

unit_test_() ->
    Foo = hit(foo, empty()),
    Foo2 = hit(foo, Foo),
    Foo3 = hit(foo, Foo2),
    Foo4 = hit(foo, Foo3),
    FooList = [foo],
    FooLess = delete(foo, Foo4),
    FooBar = hit(foo, hit(bar, empty())),

    Bar = hit10(0, 50, fun test_name/1, empty()),
    BarList = [cthulhu, hastur, yig, 'shub-niggurath', nug, yeb, tsathoggua, 'yog-sothoth', aforgomon, xexanoth],

    Baz = max_hit10(3, 0, 50, fun test_name/1, empty()),
    BazList = [cthulhu, hastur, yig],

    [?_assert(to_list(empty()) =:= []),
     ?_assert(lru:size(empty()) =:= 0),
     ?_assert(empty() =:= empty()),
     ?_assert(null(empty())),
     ?_assert(lru:size(Foo) =:= 1),
     ?_assert(lru:size(FooLess) =:= 0),
     ?_assert(lru:size(FooBar) =:= 2),
     ?_assert(null(empty())),
     ?_assert(null(delete(foo, hit(foo, empty())))),
     ?_assert(not null(hit(foo, empty()))),
     ?_assert(lru:size(Foo) =:= lru:size(Foo2)),
     ?_assert(lru:size(Foo) =:= lru:size(Foo3)),
     ?_assert(lru:size(Foo) =:= lru:size(Foo4)),
     ?_assert(lru:size(Foo2) =:= lru:size(Foo3)),
     ?_assert(lru:size(Foo2) =:= lru:size(Foo4)),
     ?_assert(lru:size(Foo3) =:= lru:size(Foo4)),
     ?_assert(member(foo, Foo)),
     ?_assert(member(foo, Foo2)),
     ?_assert(member(foo, Foo3)),
     ?_assert(member(foo, Foo4)),
     ?_assert(not member(foo, FooLess)),
     ?_assert(to_list(Foo) =:= FooList),
     ?_assert(to_list(Foo2) =:= FooList),
     ?_assert(to_list(Foo3) =:= FooList),
     ?_assert(to_list(Foo4) =:= FooList),
     ?_assert(lru:size(FooLess) =:= 0),
     ?_assert(to_list(Foo) =:= FooList),
     ?_assert(to_list(Bar) =:= BarList),
     ?_assert(to_list(Baz) =:= BazList),
     ?_assert(last(Baz) =:= yig),
     ?_assert(last(Foo) =:= foo),
     ?_assert(last(Bar) =:= xexanoth),
     ?_assert(last(FooBar) =:= bar),
     ?_assertThrow(empty, last(FooLess)),
     ?_assertThrow(empty, last(empty())),
     ?_assertThrow(empty, pop(FooLess)),
     ?_assertThrow(empty, pop(empty())),
     ?_assert(empty() =:= pop(Foo)),
     ?_assert(empty() =:= delete(nonexistent, empty())),
     ?_assert([cthulhu, hastur] =:= to_list(delete(yig, Baz))),
     ?_assert([hastur, yig] =:= to_list(delete(cthulhu, Baz))),
     ?_assert([cthulhu, hastur, yig] =:= to_list(delete(nothing, Baz))),
     ?_assert([cthulhu, yig] =:= to_list(delete(hastur, Baz))),
     ?_assert([hastur, cthulhu, yig] =:= to_list(hit(hastur, Baz)))].

test_name(I) when I rem 10 =:= 0 ->
    xexanoth;
test_name(I) when I rem 10 =:= 1 ->
    aforgomon;
test_name(I) when I rem 10 =:= 2 ->
    'yog-sothoth';
test_name(I) when I rem 10 =:= 3 ->
    tsathoggua;
test_name(I) when I rem 10 =:= 4 ->
    yeb;
test_name(I) when I rem 10 =:= 5 ->
    nug;
test_name(I) when I rem 10 =:= 6 ->
    'shub-niggurath';
test_name(I) when I rem 10 =:= 7 ->
    yig;
test_name(I) when I rem 10 =:= 8 ->
    hastur;
test_name(I) when I rem 10 =:= 9 ->
    cthulhu.

hit10(I, N, F, Acc) when I < N ->
    hit10(I + 1, N, F, hit(F(I), Acc));
hit10(_, _, _, Acc) ->
    Acc.

max_hit10(Max, I, N, F, Acc) when I < N ->
    Acc1 = hit(F(I), Acc),
    Acc2 = case lru:size(Acc1) of
               Sz when Sz > Max ->
                   pop(Acc1);
               _ ->
                   Acc1
           end,
    max_hit10(Max, I + 1, N, F, Acc2);
max_hit10(_, _, _, _, Acc) ->
    Acc.

-endif.