[couchdb-thrift-protocol] 02/19: Add `decode_message` for binary format

[da...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

davisp pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/couchdb-thrift-protocol.git

commit 1a7bc50f257b3b06755ab08035138b17e751da5d
Author: Takeru Ohta <ph...@gmail.com>
AuthorDate: Wed Oct 18 06:55:54 2017 +0900

    Add `decode_message` for binary format
---
 src/thrift_protocol.erl        |   8 ++-
 src/thrift_protocol_binary.erl | 122 +++++++++++++++++++++++++++++------------
 src/thrift_protocol_byte.erl   |  44 +++++++++++++++
 3 files changed, 139 insertions(+), 35 deletions(-)

diff --git a/src/thrift_protocol.erl b/src/thrift_protocol.erl
index 9a5b5cf..3b59cfb 100644
--- a/src/thrift_protocol.erl
+++ b/src/thrift_protocol.erl
@@ -2,7 +2,7 @@
 
 -include("thrift_protocol.hrl").
 
--export([encode_message/2]).
+-export([encode_message/2, decode_message/2]).
 -export([data_type/1]).
 
 -export_type([message/0, message_type/0, struct/0]).
@@ -57,6 +57,12 @@ encode_message(Message, binary) ->
 encode_message(Message, compact) ->
     thrift_protocol_compact:encode_message(Message).
 
+-spec decode_message(binary(), format()) -> {message(), binary()}.
+decode_message(Bin, binary) ->
+    thrift_protocol_binary:decode_message(Bin);
+decode_message(Bin, compact) ->
+    thrift_protocol_compact:decode_message(Bin).
+
 -spec data_type(data()) -> data_type().
 data_type(X) when is_boolean(X) ->
     boolean;
diff --git a/src/thrift_protocol_binary.erl b/src/thrift_protocol_binary.erl
index c4edaa9..0bccf20 100644
--- a/src/thrift_protocol_binary.erl
+++ b/src/thrift_protocol_binary.erl
@@ -1,23 +1,94 @@
 -module(thrift_protocol_binary).
 
 -include("thrift_protocol.hrl").
--include("constants.hrl").
 
--export([encode_message/1]).
+-export([encode_message/1, decode_message/1]).
 
 -define(VERSION, 1).
 
+-spec decode_message(binary()) -> {thrift_protocol:message(), binary()}.
+decode_message(<<1:1, ?VERSION:15, _:8, 0:5, TypeByte:3, NameLen:32, Name:NameLen/binary, SeqId:32/signed, Rest0/binary>>) ->
+    Type = thrift_protocol_byte:to_message_type(TypeByte),
+    {Body, Rest1} = decode_struct(Rest0, #{}),
+    Message =
+        #thrift_protocol_message{
+           method_name = Name,
+           message_type = Type,
+           sequence_id = SeqId,
+           body = Body
+          },
+    {Message, Rest1}.
+
+-spec decode_struct(binary(), Fields) -> {thrift_protocol:struct(), binary()} when
+      Fields :: #{thrift_protocol:field_id() => thrift_protocol:data()}.
+decode_struct(<<0:32, Rest/binary>>, Fields) ->
+    {#thrift_protocol_struct{fields = Fields}, Rest};
+decode_struct(<<TypeByte:8, Id:16/signed, Rest0/binary>>, Fields0) ->
+    Type = thrift_protocol_struct:to_data_type(TypeByte),
+    {Data, Rest1} = decode_data(Type, Rest0),
+    Fields1 = maps:put(Id, Data, Fields0),
+    decode_struct(Rest1, Fields1).
+
+-spec decode_elements(binary(), thrift_protocol:data_type(), non_neg_integer(), [thrift_protocol:data()]) ->
+                             {[thrift_protocol:data()], binary()}.
+decode_elements(Bin, _Type, 0, Elements) ->
+    {lists:reverse(Elements), Bin};
+decode_elements(Bin, Type, Size, Elements) ->
+    {Element, Rest} = decode_data(Type, Bin),
+    decode_elements(Rest, Type, Size - 1, [Element | Elements]).
+
+-spec decode_pairs(binary(), KeyType, ValueType, Size, Acc) -> {Acc, binary()} when
+      KeyType :: thrift_protocol:data_type(),
+      ValueType :: thrift_protocol:data_type(),
+      Size :: non_neg_integer(),
+      Acc :: #{thrift_protocol:data() => thrift_protocol:data()}.
+decode_pairs(Bin, _, _, 0, Pairs) ->
+    {Pairs, Bin};
+decode_pairs(Bin, KeyType, ValueType, Size, Pairs0) ->
+    {Key, Rest0} = decode_data(KeyType, Bin),
+    {Value, Rest1} = decode_data(ValueType, Rest0),
+    Pairs1 = maps:put(Key, Value, Pairs0),
+    decode_pairs(Rest1, KeyType, ValueType, Size - 1, Pairs1).
+
+-spec decode_data(thrift_protocol:data_type(), binary()) -> {thrift_protocol:data(), binary()}.
+decode_data(boolean, <<0, Rest/binary>>) -> {false, Rest};
+decode_data(boolean, <<1, Rest/binary>>) -> {true, Rest};
+decode_data(i8, <<N:8/signed, Rest/binary>>) -> {{i8, N}, Rest};
+decode_data(i16, <<N:16/signed, Rest/binary>>) -> {{i16, N}, Rest};
+decode_data(i32, <<N:32/signed, Rest/binary>>) -> {{i32, N}, Rest};
+decode_data(i64, <<N:64/signed, Rest/binary>>) -> {{i64, N}, Rest};
+decode_data(float, <<N/float, Rest/binary>>) -> {N, Rest};
+decode_data(binary, <<Len:32, Bin:Len/binary, Rest/binary>>) -> {Bin, Rest};
+decode_data(struct, Bin) -> decode_struct(Bin, #{});
+decode_data(map, <<K:8, V:8, Size:32, Bin/binary>>) ->
+    KeyType = thrift_protocol_byte:to_data_type(K),
+    ValueType = thrift_protocol_byte:to_data_type(V),
+    {Pairs, Rest} = decode_pairs(Bin, KeyType, ValueType, Size, #{}),
+    {#thrift_protocol_map{key_type = KeyType, value_type = ValueType, elements = Pairs}, Rest};
+decode_data(set, <<TypeByte:8, Size:32, Bin/binary>>) ->
+    Type = thrift_protocol_byte:to_data_type(TypeByte),
+    {Elements, Rest} = decode_elements(Bin, Type, Size, []),
+    {#thrift_protocol_set{element_type = Type, elements = Elements}, Rest};
+decode_data(list, <<TypeByte:8, Size:32, Bin/binary>>) ->
+    Type = thrift_protocol_byte:to_data_type(TypeByte),
+    {Elements, Rest} = decode_elements(Bin, Type, Size, []),
+    {#thrift_protocol_list{element_type = Type, elements = Elements}, Rest}.
+
 -spec encode_message(thrift_protocol:message()) -> iodata().
 encode_message(Message) ->
     #thrift_protocol_message{method_name = Name, sequence_id = SeqId, body = Body} = Message,
-    Type =
-        case Message#thrift_protocol_message.message_type of
-            call      -> ?MESSAGE_TYPE_CALL;
-            reply     -> ?MESSAGE_TYPE_REPLY;
-            exception -> ?MESSAGE_TYPE_EXCEPTION;
-            oneway    -> ?MESSAGE_TYPE_ONEWAY
-        end,
-    [<<1:1, ?VERSION:15, 0:8, 0:5, Type:3, (byte_size(Name)):32, Name/binary, SeqId:32>> | encode_data(Body)].
+    Type = thrift_protocol_byte:from_message_type(Message#thrift_protocol_message.message_type),
+    [<<1:1, ?VERSION:15, 0:8, 0:5, Type:3, (byte_size(Name)):32, Name/binary, SeqId:32>> | encode_struct(Body)].
+
+-spec encode_struct(thrift_protocol:struct()) -> iodata().
+encode_struct(#thrift_protocol_struct{fields = Fields}) ->
+    maps:fold(
+      fun (Id, Data, Acc) ->
+              Type = thrift_protocol_byte:from_data_type(thrift_protocol:data_type(Data)),
+              [<<Type:8, Id:16>>, encode_data(Data) | Acc]
+      end,
+      [<<0:32>>],
+      Fields).
 
 -spec encode_data(thrift_protocol:data()) -> iodata().
 encode_data(true) ->
@@ -36,14 +107,8 @@ encode_data(N) when is_float(N)->
     <<N/float>>;
 encode_data(B) when is_binary(B) ->
     [<<(byte_size(B)):32>> | B];
-encode_data(#thrift_protocol_struct{fields = Fields}) ->
-    maps:fold(
-      fun (Id, Data, Acc) ->
-              Type = data_type_to_byte(thrift_protocol:data_type(Data)),
-              [<<Type:8, Id:32>>, encode_data(Data) | Acc]
-      end,
-      [<<0:32>>],
-      Fields);
+encode_data(X = #thrift_protocol_struct{}) ->
+    encode_struct(X);
 encode_data(#thrift_protocol_map{key_type = KeyType, value_type = ValueType, elements = Elements}) ->
     IoData =
         maps:fold(
@@ -54,31 +119,20 @@ encode_data(#thrift_protocol_map{key_type = KeyType, value_type = ValueType, ele
           end,
           [],
           Elements),
-    [<<(data_type_to_byte(KeyType)):8, (data_type_to_byte(ValueType)):8, (maps:size(Elements)):32>> | IoData];
+    [<<(thrift_protocol_byte:from_data_type(KeyType)):8,
+       (thrift_protocol_byte:from_data_type(ValueType)):8,
+       (maps:size(Elements)):32>> | IoData];
 encode_data(#thrift_protocol_set{element_type = Type, elements = Elements}) ->
     Size = length(Elements),
-    [<<(data_type_to_byte(Type)):8, Size:32>> |
+    [<<(thrift_protocol_byte:from_data_type(Type)):8, Size:32>> |
      [begin
           Type = thrift_protocol:data_type(E),
           encode_data(E)
       end || E <- Elements]];
 encode_data(#thrift_protocol_list{element_type = Type, elements = Elements}) ->
     Size = length(Elements),
-    [<<(data_type_to_byte(Type)):8, Size:32>> |
+    [<<(thrift_protocol_byte:from_data_type(Type)):8, Size:32>> |
      [begin
           Type = thrift_protocol:data_type(E),
           encode_data(E)
       end || E <- Elements]].
-
--spec data_type_to_byte(thrift_protocol:data_type()) -> byte().
-data_type_to_byte(boolean) -> ?DATA_TYPE_BOOLEAN;
-data_type_to_byte(i8)      -> ?DATA_TYPE_I8;
-data_type_to_byte(i16)     -> ?DATA_TYPE_I16;
-data_type_to_byte(i32)     -> ?DATA_TYPE_I32;
-data_type_to_byte(i64)     -> ?DATA_TYPE_I64;
-data_type_to_byte(float)   -> ?DATA_TYPE_FLOAT;
-data_type_to_byte(binary)  -> ?DATA_TYPE_BINARY;
-data_type_to_byte(struct)  -> ?DATA_TYPE_STRUCT;
-data_type_to_byte(map)     -> ?DATA_TYPE_MAP;
-data_type_to_byte(set)     -> ?DATA_TYPE_SET;
-data_type_to_byte(list)    -> ?DATA_TYPE_LIST.
diff --git a/src/thrift_protocol_byte.erl b/src/thrift_protocol_byte.erl
new file mode 100644
index 0000000..711d8df
--- /dev/null
+++ b/src/thrift_protocol_byte.erl
@@ -0,0 +1,44 @@
+-module(thrift_protocol_byte).
+
+-include("constants.hrl").
+
+-export([from_message_type/1, from_data_type/1]).
+-export([to_message_type/1, to_data_type/1]).
+
+-spec to_message_type(byte()) -> thrift_protocol:message_type().
+to_message_type(?MESSAGE_TYPE_CALL)      -> call;
+to_message_type(?MESSAGE_TYPE_REPLY)     -> reply;
+to_message_type(?MESSAGE_TYPE_EXCEPTION) -> exception;
+to_message_type(?MESSAGE_TYPE_ONEWAY)    -> oneway.
+
+-spec to_data_type(byte()) -> thrift_protocol:data_type().
+to_data_type(?DATA_TYPE_BOOLEAN) -> boolean;
+to_data_type(?DATA_TYPE_I8)      -> i8;
+to_data_type(?DATA_TYPE_I16)     -> i16;
+to_data_type(?DATA_TYPE_I32)     -> i32;
+to_data_type(?DATA_TYPE_I64)     -> i64;
+to_data_type(?DATA_TYPE_FLOAT)   -> float;
+to_data_type(?DATA_TYPE_BINARY)  -> binary;
+to_data_type(?DATA_TYPE_STRUCT)  -> struct;
+to_data_type(?DATA_TYPE_MAP)     -> map;
+to_data_type(?DATA_TYPE_SET)     -> set;
+to_data_type(?DATA_TYPE_LIST)    -> list.
+
+-spec from_message_type(thrift_protocol:message_type()) -> byte().
+from_message_type(call)      -> ?MESSAGE_TYPE_CALL;
+from_message_type(reply)     -> ?MESSAGE_TYPE_REPLY;
+from_message_type(exception) -> ?MESSAGE_TYPE_EXCEPTION;
+from_message_type(oneway)    -> ?MESSAGE_TYPE_ONEWAY.
+
+-spec from_data_type(thrift_protocol:data_type()) -> byte().
+from_data_type(boolean) -> ?DATA_TYPE_BOOLEAN;
+from_data_type(i8)      -> ?DATA_TYPE_I8;
+from_data_type(i16)     -> ?DATA_TYPE_I16;
+from_data_type(i32)     -> ?DATA_TYPE_I32;
+from_data_type(i64)     -> ?DATA_TYPE_I64;
+from_data_type(float)   -> ?DATA_TYPE_FLOAT;
+from_data_type(binary)  -> ?DATA_TYPE_BINARY;
+from_data_type(struct)  -> ?DATA_TYPE_STRUCT;
+from_data_type(map)     -> ?DATA_TYPE_MAP;
+from_data_type(set)     -> ?DATA_TYPE_SET;
+from_data_type(list)    -> ?DATA_TYPE_LIST.