pickle.erl

%%%% Copyright (C) 2005-2008 Wager Labs, SA
%%%%
%%%% THE WORK (AS DEFINED BELOW) IS PROVIDED UNDER THE TERMS OF THIS 
%%%% CREATIVE COMMONS PUBLIC LICENSE ("CCPL" OR "LICENSE"). THE WORK IS 
%%%% PROTECTED BY COPYRIGHT AND/OR OTHER APPLICABLE LAW. ANY USE OF 
%%%% THE WORK OTHER THAN AS AUTHORIZED UNDER THIS LICENSE OR COPYRIGHT 
%%%% LAW IS PROHIBITED.
%%%%
%%%% BY EXERCISING ANY RIGHTS TO THE WORK PROVIDED HERE, YOU ACCEPT 
%%%% AND AGREE TO BE BOUND BY THE TERMS OF THIS LICENSE. TO THE EXTENT 
%%%% THIS LICENSE MAY BE CONSIDERED TO BE A CONTRACT, THE LICENSOR GRANTS 
%%%% YOU THE RIGHTS CONTAINED HERE IN CONSIDERATION OF YOUR ACCEPTANCE 
%%%% OF SUCH TERMS AND CONDITIONS.
%%%%
%%%% Please see LICENSE for full legal details and the following URL
%%%% for a human-readable explanation:
%%%%
%%%% http://creativecommons.org/licenses/by-nc-sa/3.0/us/
%%%%

-module(pickle).

%%%
%%% Serialization tools
%%%

-export([pickle/2, unpickle/2]).
-export([byte/0, short/0, sshort/0, int/0, 
         sint/0, long/0, slong/0, price/0]).
-export([list/2, choice/2, optional/1, wrap/2,
         tuple/1, record/2, binary/1, wstring/0]).
-export([string/0]).

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

%%% Pickle and unpickle. We accumulate into a list.

pickle({Pickler, _}, Value) ->
    lists:reverse(Pickler([], Value)).

unpickle({_, Pickler}, Bin) ->
    element(1, Pickler(Bin)).

%%% Byte

byte() -> 
    {fun write_byte/2, fun read_byte/1}.

write_byte(Acc, Byte) -> 
    [<<Byte:8>>|Acc].

read_byte(Bin) -> 
    <<Byte:8, Rest/binary>> = Bin,
    {Byte, Rest}.

%%% Unsigned short

short() -> 
    {fun write_short/2, fun read_short/1}.

write_short(Acc, Word) -> 
    [<<Word:16>>|Acc].

read_short(Bin) -> 
    <<Word:16, Rest/binary>> = Bin,
    {Word, Rest}.

%%% Signed short

sshort() -> 
    {fun write_sshort/2, fun read_sshort/1}.

write_sshort(Acc, Word) -> 
    [<<Word:16/signed>>|Acc].

read_sshort(Bin) -> 
    <<Word:16/signed, Rest/binary>> = Bin,
    {Word, Rest}.

%%% Unsigned int

int() -> 
    {fun write_int/2, fun read_int/1}.

write_int(Acc, Word) -> 
    [<<Word:32>>|Acc].

read_int(Bin) -> 
    <<Word:32, Rest/binary>> = Bin,
    {Word, Rest}.

%%% Signed int

sint() -> 
    {fun write_sint/2, fun read_sint/1}.

write_sint(Acc, Word) -> 
    [<<Word:32/signed>>|Acc].

read_sint(Bin) -> 
    <<Word:32/signed, Rest/binary>> = Bin,
    {Word, Rest}.

%%% Unsigned long

long() -> 
    {fun write_long/2, fun read_long/1}.

write_long(Acc, Word) -> 
    [<<Word:64>>|Acc].

read_long(Bin) -> 
    <<Word:64, Rest/binary>> = Bin,
    {Word, Rest}.

%%% Signed long

slong() -> 
    {fun write_slong/2, fun read_slong/1}.

write_slong(Acc, Word) -> 
    [<<Word:64/signed>>|Acc].

read_slong(Bin) -> 
    <<Word:64/signed, Rest/binary>> = Bin,
    {Word, Rest}.

%%% Price with 5 decimal points

price() -> 
    {fun write_price/2, fun read_price/1}.

write_price(Acc, Word) -> 
    Price = trunc(Word * 10000),
    [<<Price:32>>|Acc].

read_price(Bin) -> 
    <<Word:32, Rest/binary>> = Bin,
    Price = Word / 10000,
    {Price, Rest}.

%%% List. We supply a pickler for list length 
%%% as well as a pickler for list elements.

list(Len, Elem) ->
    {fun(Acc, List) -> write_list(Len, Elem, Acc, List) end, 
     fun(Bin) -> read_list(Len, Elem, Bin) end }.

write_list({Len, _}, {Elem, _}, Acc, List) ->
    Acc1 = Len(Acc, length(List)),
    Fun = fun(A, Acc2) -> Elem(Acc2, A) end,
    lists:foldr(Fun, Acc1, List).

read_list({_, Len}, {_, Elem}, Bin) ->
    {N, Bin1} = Len(Bin),
    read_list(N, [], Elem, Bin1).

read_list(0, Acc, _, Bin) -> {Acc, Bin};
read_list(N, Acc, Elem, Bin) ->
    {E, Bin1} = Elem(Bin),
    read_list(N - 1, [E|Acc], Elem, Bin1).

%%% Alternative selection. This could probably use some
%%% deeper thinking. We take a pickler for the tag
%%% as well as a tuple of two functions. The first one
%%% returns the tag value and a pickler based on the supplied
%%% value. The second one selects a pickler based on a tag value.

choice(Tag, Choice) ->
    {fun(Acc, Value) -> write_choice(Tag, Choice, Acc, Value) end,
     fun(Bin) -> read_choice(Tag, Choice, Bin) end }.

write_choice({Tag, _}, {Choice, _}, Acc, Value) 
  when is_function(Tag), 
       is_function(Choice) ->
    {T, {Pickler, _}} = Choice(Value),
    Acc1 = Tag(Acc, T),
    Pickler(Acc1, Value).

read_choice({_, Tag}, {_, Choice}, Bin) 
  when is_function(Tag), 
       is_function(Choice) ->
    {T, Bin1} = Tag(Bin),
    {_, Pickler} = Choice(T),
    Pickler(Bin1).

%%% Optional value. Use 'none' to indicate no value.

optional(Pickler) ->
    {fun(Acc, Value) -> write_optional(Pickler, Acc, Value) end,
     fun(Bin) -> read_optional(Pickler, Bin) end}.

write_optional(_, Acc, none) ->
    [<<0>>|Acc];

write_optional({Pickler, _}, Acc, Value) ->
    Pickler([<<1>>|Acc], Value).

read_optional({_, Pickler}, Bin) ->
    <<Opt:8, Bin1/binary>> = Bin,
    case Opt of 
        0 -> {none, Bin1};
        _ -> Pickler(Bin1)
    end.

%%% Wrapper. Take a pickler and a wrapper tuple of two functions
%%% where the first one is used to convert the value before 
%%% pickling and the second one after unpickling.

wrap(Wrap, Pickler) ->
    {fun(Acc, Value) -> write_wrap(Wrap, Pickler, Acc, Value) end,
     fun(Bin) -> read_wrap(Wrap, Pickler, Bin) end}.

write_wrap({Wrap, _}, {Pickler, _}, Acc, Value) ->
    Pickler(Acc, Wrap(Value)).

read_wrap({_, Wrap}, {_, Pickler}, Bin) ->
    {Value, Bin1} = Pickler(Bin),
    {Wrap(Value), Bin1}.

%%% Erlang does not support enumerations but I want to have
%%% {cow, sheep, horse} as well as [{cow, 10}, {sheep, 100}]
%%% and be able to marshal these back and forth. Enumerated 
%%% values start from 1 for the tuple case.

prep_enum_tuple(Enum)
  when is_tuple(Enum) ->
    prep_enum_tuple(Enum, size(Enum), [], []).

prep_enum_tuple(_, 0, Acc1, Acc2) ->
    {Acc1, Acc2};

prep_enum_tuple(Enum, N, Acc1, Acc2) ->
    prep_enum_tuple(Enum, N - 1, 
                    [{element(N, Enum), N}|Acc1],
                    [{N, element(N, Enum)}|Acc2]).

prep_enum_list(Enum) 
  when is_list(Enum) ->
                                                % expect a list of {tag, #value} pairs
    Inv = fun({Key, Val}) -> {Val, Key} end,
    InvEnum = lists:map(Inv, Enum),
    {Enum, InvEnum}.

wrap_enum(Enum) 
  when is_tuple(Enum) ->
    wrap_enum_1(prep_enum_tuple(Enum));

wrap_enum(Enum) 
  when is_list(Enum) ->
    wrap_enum_1(prep_enum_list(Enum)).

wrap_enum_1({List1, List2}) ->
    F = fun(A, B) -> A < B end,
    %% gb_trees needs an ordered list
    Dict1 = lists:sort(F, List1),
    Dict2 = lists:sort(F, List2),
    Tree1 = gb_trees:from_orddict(Dict1),
    Tree2 = gb_trees:from_orddict(Dict2),
    {fun(Key) -> gb_trees:get(Key, Tree1) end,
     fun(Key) -> gb_trees:get(Key, Tree2) end}.       

enum(Enum, Pickler) ->
    wrap(wrap_enum(Enum), Pickler).

%%% Tuple. Uses a tuple of picklers of the same size.

tuple(Picklers) 
  when is_tuple(Picklers) ->
    wrap({fun tuple_to_list/1, 
          fun list_to_tuple/1}, 
         tuple_0(tuple_to_list(Picklers))).

%%% Record. We rely on Erlang records being tuples
%%% and just add the record tag as the first element
%%% when unpickling the record.

record(Tag, Picklers) 
  when is_tuple(Picklers) ->
    wrap({fun(Record) -> record_to_list(Tag, Record) end,
          fun(List) -> list_to_record(Tag, List) end}, 
         tuple_0(tuple_to_list(Picklers))).

write_tuple_0([], Acc, _) ->
    Acc;

write_tuple_0([{Pickler, _}|Rest], Acc, [Value|Tuple]) ->
    write_tuple_0(Rest, Pickler(Acc, Value), Tuple).

read_tuple_0(Picklers, Bin) ->
    read_tuple_0(Picklers, Bin, []).

read_tuple_0([], Bin, Acc) ->
    {lists:reverse(Acc), Bin};

read_tuple_0([{_, Pickler}|Rest], Bin, Acc) ->
    {Value, Bin1} = Pickler(Bin),
    read_tuple_0(Rest, Bin1, [Value|Acc]).

%%% It's convenient to be able to convert the tuple
%%% to a list first as there's no erlang:prepend_element/2.

tuple_0(Picklers) 
  when is_list(Picklers) ->
    {fun(Acc, Value) -> write_tuple_0(Picklers, Acc, Value) end,
     fun(Bin) -> read_tuple_0(Picklers, Bin) end}.

record_to_list(Tag, Record) 
  when is_atom(Tag) ->
    lists:nthtail(1, tuple_to_list(Record)).

list_to_record(Tag, List) 
  when is_atom(Tag), 
       is_list(List) ->
    list_to_tuple([Tag|List]).

%%% Binary

binary(Size) -> 
    {fun (Acc, Bin) -> write_binary(Size, Acc, Bin) end, 
     fun (Bin) -> read_binary(Size, Bin) end}.

write_binary({Size, _}, Acc, undefined) -> 
    Size(Acc, 0);

write_binary({Size, _}, Acc, Bin) -> 
    Acc1 = Size(Acc, size(Bin)),
    [Bin|Acc1].

read_binary({_, Size}, Bin) -> 
    {N, Bin1} = Size(Bin),
    <<Value:N/binary, Bin2/binary>> = Bin1,
    {Value, Bin2}.

%%% Wide string, little-endian style

wstring() ->
    {fun write_wstring/2, 
     fun read_wstring/1}.

write_wstring(Acc, []) ->
    [<<0:16>>|Acc];

write_wstring(Acc, [H|T]) ->
    write_wstring([<<H:16>>|Acc], T).

read_wstring(Bin) ->
    read_wstring(Bin, []).

read_wstring(<<0:16, Bin/binary>>, Acc) ->
    {lists:reverse(Acc), Bin};

read_wstring(<<X:16, Bin/binary>>, Acc) ->
    read_wstring(Bin, [X|Acc]).

string() ->    
    binary(byte()).

%%%
%%% Unit test
%%%

-define(pickle(Value, Pickler),
        fun() ->
                ?assertEqual(Value, 
                             unpickle(Pickler,
                                      list_to_binary(pickle(Pickler, 
                                                            Value))))
        end()).

byte_test() ->
    X = 16#ff,
    ?pickle(X, byte()).

short_test() ->
    X = 16#ffff,
    ?pickle(X, short()).

sshort_test() ->
    X = -1,
    ?pickle(X, sshort()).

int_test() ->
    X = 16#ffffffff,
    ?pickle(X, int()).

sint_test() ->
    X = -1,
    ?pickle(X, sint()).

long_test() ->
    X = 16#aabbccddeeff0011,
    ?pickle(X, long()).

slong_test() ->
    X = -1,
    ?pickle(X, slong()).

list_test() ->
    X = "Wazzup!",
    ?pickle(X, list(int(), byte())).

%%% A choice of serializing either a list or a long.

value2tag(Action) 
  when is_list(Action) ->
    {0, list(byte(), byte())};

value2tag(_) ->
    {1, long()}.

tag2value(0) ->
    list(byte(), byte());

tag2value(1) ->
    long().

selector() ->
    {fun value2tag/1, fun tag2value/1}.

selector_test() ->
    X1 = "Just testing",
    X2 = 16#ffff,
    ?pickle(X1, choice(byte(), selector())),
    ?pickle(X2, choice(byte(), selector())).

%%% Optional value

optional_test() ->
    X1 = none,
    X2 = 55,
    ?pickle(X1, optional(byte())),
    ?pickle(X2, optional(byte())).

%%% Tuple given as a tuple and a list of key/value pairs.

tuple_enum_test() ->
    %% tuple enum
    Enum = {cow, sheep, horse},
    {FROM, TO} = wrap_enum(Enum),
    ?assertEqual(1, FROM(cow)),
    ?assertEqual(2, FROM(sheep)),
    ?assertEqual(3, FROM(horse)),
    ?assertEqual(cow, TO(1)),
    ?assertEqual(sheep, TO(2)),
    ?assertEqual(horse, TO(3)).

list_enum_test() ->
    %% list enum
    Enum = [{cow, 20}, {sheep, 30}, {horse, 40}],
    {FROM, TO} = wrap_enum(Enum),
    ?assertEqual(20, FROM(cow)),
    ?assertEqual(30, FROM(sheep)),
    ?assertEqual(40, FROM(horse)),
    ?assertEqual(cow, TO(20)),
    ?assertEqual(sheep, TO(30)),
    ?assertEqual(horse, TO(40)).

enum_test() ->
    Enum1 = {cow, sheep, horse},
    Enum2 = [{cow, 20}, {sheep, 30}, {horse, 40}],
    ?pickle(cow, enum(Enum1, byte())),
    ?pickle(sheep, enum(Enum2, byte())).

tuple_test() ->
    Tuple = {"Joel", 16#ff00, none},
    Spec = {list(byte(),byte()), short(), optional(byte())},
    ?pickle(Tuple, tuple(Spec)).

%%% Nested records.

-record(foo, { a, b }).
-record(bar, { c, d }).
-record(baz, { e, f }).

nested_record_test() ->
    R1 = #foo { 
      a = 10, 
      b = #bar {
        c = 20, 
        d = #baz {
          e = 30,
          f = "Enough nesting!"
         }
       }
     },
    Pickler = record(foo, {
                       byte(),
                       record(bar, {
                                int(),
                                record(baz, {
                                         sshort(),
                                         list(byte(), byte())
                                        })
                               })
                      }),
    ?pickle(R1, Pickler).

%%% Binary

binary_test() ->
    X = <<1, 2, 3, 4>>,
    ?pickle(X, binary(int())).

%%% Wide string

wstring_test() ->
    X = "This is a wide string",
    ?pickle(X, wstring()).