Merge branch 'siriuszhuang-master'

comtihon · Jul 28, 2015 · 9564550 · 9564550
2 parents 0f7c07e + 7364597
commit 9564550
Show file tree

Hide file tree

Showing 3 changed files with 66 additions and 51 deletions.
diff --git a/src/bson.erl b/src/bson.erl
@@ -25,7 +25,7 @@
 % Conceptually a document is a list of label-value pairs (associative array, dictionary, record). However, for read/write-ability, it is implemented as a flat tuple, ie. the list becomes a tuple and the pair braces are elided, so you just have alternating labels and values where each value is associated with the previous label.
 % To distinguish a tagged value such as {uuid, _} (see value() type below) from a document with field name 'uuid' we made sure all valid tagged value types have an odd number of elements (documents have even number of elements). So actually only {bin, uuid, _} is a valid value, {uuid, _} is a document.
 
--type label() :: atom().
+-type label() :: binary().
 
 -spec doc_foldl(fun ((label(), value(), A) -> A), A, document()) -> A.
 %@doc Reduce document by applying given function to each field with result of previous field's application, starting with given initial result.
@@ -64,31 +64,34 @@ flatten([{Label, Value} | Fields]) -> [Label, Value | flatten(Fields)].
 
 -spec lookup(label(), document()) -> maybe (value()).
 %@doc Value of field in document if there
+lookup(Label, Doc) when is_atom(Label) -> lookup(atom_to_binary(Label, utf8), Doc);
 lookup(Label, Doc) ->
-  Parts = string:tokens(atom_to_list(Label), "."),
+  Parts = binary:split(Label, <<".">>, []),
   case length(Parts) of
     1 ->
-      case find(list_to_atom(hd(Parts)), Doc) of
+      case find(hd(Parts), Doc) of
         {Index} -> {element(Index * 2 + 2, Doc)};
         {} -> {} end;
     _ ->
-      case find(list_to_atom(hd(Parts)), Doc) of
-        {Index} -> lookup(list_to_atom(string:join(tl(Parts), ".")), element(Index * 2 + 2, Doc));
+      case find(hd(Parts), Doc) of
+        {Index} -> lookup(hd(tl(Parts)), element(Index * 2 + 2, Doc));
         {} -> {} end
   end.
 
 -spec lookup(label(), document(), value()) -> value().
 %@doc Value of field in document if there or default
+lookup(Label, Doc, Default) when is_atom(Label) ->
+  lookup(atom_to_binary(Label, utf8), Doc, Default);
 lookup(Label, Doc, Default) ->
-  Parts = string:tokens(atom_to_list(Label), "."),
+  Parts = binary:split(Label, <<".">>, []),
   case length(Parts) of
     1 ->
-      case find(list_to_atom(hd(Parts)), Doc) of
+      case find(hd(Parts), Doc) of
         {Index} -> element(Index * 2 + 2, Doc);
         {} -> Default end;
     _ ->
-      case find(list_to_atom(hd(Parts)), Doc) of
-        {Index} -> lookup(list_to_atom(string:join(tl(Parts), ".")), element(Index * 2 + 2, Doc), Default);
+      case find(hd(Parts), Doc) of
+        {Index} -> lookup(hd(tl(Parts)), element(Index * 2 + 2, Doc), Default);
         {} -> Default end
   end.
 
@@ -99,16 +102,26 @@ find(Label, Doc) -> findN(Label, Doc, 0, tuple_size(Doc) div 2).
 -spec findN(label(), document(), integer(), integer()) -> maybe (integer()).
 %@doc Find field index in document from first index (inclusive) to second index (exclusive).
 findN(_Label, _Doc, High, High) -> {};
-findN(Label, Doc, Low, High) -> case element(Low * 2 + 1, Doc) of
-                                  Label -> {Low};
-                                  _ -> findN(Label, Doc, Low + 1, High) end.
+findN(Label, Doc, Low, High) ->
+  case element(Low * 2 + 1, Doc) of
+    Label -> {Low};
+    AtomKey when is_atom(AtomKey) ->
+      case atom_to_binary(AtomKey, utf8) =:= Label of
+         true -> {Low};
+         false -> findN(Label, Doc, Low + 1, High)
+      end;
+    _ -> findN(Label, Doc, Low + 1, High)
+  end.
 
 -spec at(label(), document()) -> value().
 %@doc Value of field in document, error if missing
-at(Label, Document) -> case lookup(Label, Document) of
-% {} -> erlang:error (missing_field, [Label, Document]);
-                         {} -> null;
-                         {Value} -> Value end.
+at(Label, Document) when is_atom(Label) ->
+  at(atom_to_binary(Label, utf8), Document);
+at(Label, Document) ->
+  case lookup(Label, Document) of
+    {} -> null;
+    {Value} -> Value
+  end.
 
 -spec include([label()], document()) -> document().
 %@doc Project given fields of document
@@ -128,21 +141,23 @@ exclude(Labels, Document) ->
 
 -spec update(label(), value(), document()) -> document().
 %@doc Replace field with new value, adding to end if new
+update(Label, Value, Document) when is_atom(Label) ->
+  update(atom_to_binary(Label, utf8), Value, Document);
 update(Label, Value, Document) ->
-  Parts = string:tokens(atom_to_list(Label), "."),
+  Parts = binary:split(Label, <<".">>, []),
   case length(Parts) of
     1 ->
-      case find(list_to_atom(hd(Parts)), Document) of
+      case find(hd(Parts), Document) of
         {Index} -> setelement(Index * 2 + 2, Document, Value);
         {} ->
           Doc = erlang:append_element(Document, Label),
           erlang:append_element(Doc, Value) end;
     _ ->
-      case find(list_to_atom(hd(Parts)), Document) of
+      case find(hd(Parts), Document) of
         {Index} ->
-          setelement(Index * 2 + 2, Document, update(list_to_atom(string:join(tl(Parts), ".")), Value, element(Index * 2 + 2, Document)));
-        {} -> Doc = erlang:append_element(Document, list_to_atom(hd(Parts))),
-          erlang:append_element(Doc, update(list_to_atom(string:join(tl(Parts), ".")), Value, {})) end
+          setelement(Index * 2 + 2, Document, update(hd(tl(Parts)), Value, element(Index * 2 + 2, Document)));
+        {} -> Doc = erlang:append_element(Document, hd(Parts)),
+          erlang:append_element(Doc, update(hd(tl(Parts)), Value, {})) end
   end.
 
 
@@ -271,4 +286,4 @@ map_to_bson(Map) ->
 -spec proplist_to_bson(proplists:proplist()) -> document().
 proplist_to_bson(Proplist) ->
   L = lists:foldr(fun({A, B}, Acc) -> [A | [B | Acc]] end, [], Proplist),
-  list_to_tuple(L).
+  list_to_tuple(L).
diff --git a/src/bson_binary.erl b/src/bson_binary.erl
@@ -99,7 +99,7 @@ get_document(<<?get_int32(N), Bin/binary>>) ->
 get_fields(<<>>) -> [];
 get_fields(Bin) ->
   {Name, Value, Bin1} = get_field(Bin),
-  [binary_to_atom(Name, utf8), Value | get_fields(Bin1)].
+  [Name, Value | get_fields(Bin1)].
 
 -spec put_array(bson:arr()) -> binary().
 % encoded same as document with labels '0', '1', etc.

diff --git a/test/bson_tests.erl b/test/bson_tests.erl
@@ -25,9 +25,9 @@ bson_test() ->
   {b, {x, 2, y, 3}, a, 1, c, 4.2} = bson:update(c, 4.2, Doc),
   {b, {x, 13, y, 3}, a, 1, c, [mon, tue, wed]} = bson:update('b.x', 13, Doc),
   {b, {x, 2, y, 14}, a, 1, c, [mon, tue, wed]} = bson:update('b.y', 14, Doc),
-  {b, {x, 2, y, 3, z, 0}, a, 1, c, [mon, tue, wed]} = bson:update('b.z', 0, Doc),
-  {b, {x, 2, y, 3}, a, 1, c, [mon, tue, wed], d, {x, 15}} = bson:update('d.x', 15, Doc),
-  {b, 0, a, 1, c, 2, d, 3} = bson:merge({c, 2, d, 3, b, 0}, Doc),
+  {b, {x, 2, y, 3, <<"z">>, 0}, a, 1, c, [mon, tue, wed]} = bson:update('b.z', 0, Doc),
+  {b, {x, 2, y, 3}, a, 1, c, [mon, tue, wed], <<"d">>, {<<"x">>, 15}} = bson:update('d.x', 15, Doc),
+  {b, 0, a, 1, c, 2, <<"d">>, 3} = bson:merge({c, 2, d, 3, b, 0}, Doc),
   {a, 1, b, 2, c, 3, d, 4} = bson:append({a, 1, b, 2}, {c, 3, d, 4}),
   [{b, {x, 2, y, 3}}, {a, 1}, {c, [mon, tue, wed]}] = bson:fields(Doc).
 
@@ -47,35 +47,35 @@ binary_test() ->
   Bin = <<49, 0, 0, 0, 4, 66, 83, 79, 78, 0, 38, 0, 0, 0, 2, 48, 0, 8, 0, 0, 0, 97, 119, 101, 115, 111, 109, 101, 0, 1, 49, 0, 51, 51, 51, 51, 51, 51, 20, 64, 16, 50, 0, 194, 7, 0, 0, 0, 0>>,
   VBin = <<200, 12, 240, 129, 100, 90, 56, 198, 34, 0, 0>>,
   Time = bson:timenow(),
-  Doc1 = {a, -4.230845,
-    b, <<"hello">>,
-    c, {x, -1, y, 2.2001},
-    d, [23, 45, 200],
-    eeeeeeeee, {bin, bin, VBin},
-    f, {bin, function, VBin},
-    g, {bin, uuid, Bin},
-    h, {bin, md5, VBin},
-    i, {bin, userdefined, Bin},
-    j, bson:objectid(bson:unixtime_to_secs(Time), <<2:24/big, 3:16/big>>, 4),
-    k1, false,
-    k2, true,
-    l, Time,
-    m, undefined,
-    n, {regex, <<"foo">>, <<"bar">>},
-    o1, {javascript, {}, <<"function(x) = x + 1;">>},
-    o2, {javascript, {x, 0, y, <<"foo">>}, <<"function(a) = a + x">>},
-    p, atom,
-    q1, -2000444000,
-    q2, -8000111000222001,
-    r, {mongostamp, 100022, 995332003},
-    s1, 'MIN_KEY',
-    s2, 'MAX_KEY'},
+  Doc1 = {<<"a">>, -4.230845,
+    <<"b">>, <<"hello">>,
+    <<"c">>, {<<"x">>, -1, <<"y">>, 2.2001},
+    <<"d">>, [23, 45, 200],
+    <<"eeeeeeeee">>, {bin, bin, VBin},
+    <<"f">>, {bin, function, VBin},
+    <<"g">>, {bin, uuid, Bin},
+    <<"h">>, {bin, md5, VBin},
+    <<"i">>, {bin, userdefined, Bin},
+    <<"j">>, bson:objectid(bson:unixtime_to_secs(Time), <<2:24/big, 3:16/big>>, 4),
+    <<"k1">>, false,
+    <<"k2">>, true,
+    <<"l">>, Time,
+    <<"m">>, undefined,
+    <<"n">>, {regex, <<"foo">>, <<"bar">>},
+    <<"o1">>, {javascript, {}, <<"function(x) = x + 1;">>},
+    <<"o2">>, {javascript, {<<"x">>, 0, <<"y">>, <<"foo">>}, <<"function(a) = a + x">>},
+    <<"p">>, atom,
+    <<"q1">>, -2000444000,
+    <<"q2">>, -8000111000222001,
+    <<"r">>, {mongostamp, 100022, 995332003},
+    <<"s1">>, 'MIN_KEY',
+    <<"s2">>, 'MAX_KEY'},
   Bin1 = bson_binary:put_document(Doc1),
   {Doc1, <<>>} = bson_binary:get_document(Bin1).
 
 put_document_test() ->
   Doc = {<<"key">>, <<"value">>},
-  ?assertEqual({{key, <<"value">>}, <<>>}, bson_binary:get_document(bson_binary:put_document(Doc))).
+  ?assertEqual({{<<"key">>, <<"value">>}, <<>>}, bson_binary:get_document(bson_binary:put_document(Doc))).
 
 bson_int_too_large_test() ->
   Doc1 = {int, 16#7fffffffffffffff + 1},