1 files changed, 10 insertions, 3 deletions
diff --git a/src/couchdb/couch_key_tree.erl b/src/couchdb/couch_key_tree.erl
index 647a3455..830820f3 100644
--- a/src/couchdb/couch_key_tree.erl
+++ b/src/couchdb/couch_key_tree.erl
@@ -278,13 +278,20 @@ count_leafs_simple([{_Key, _Value, SubTree} | RestTree]) ->
 map(_Fun, []) ->
     [];
 map(Fun, [{Pos, Tree}|Rest]) ->
-    [NewTree] = map_simple(Fun, Pos, [Tree]),
-    [{Pos, NewTree} | map(Fun, Rest)].
+    case erlang:fun_info(Fun, arity) of
+    {arity, 2} ->
+        [NewTree] = map_simple(fun(A,B,_C) -> Fun(A,B) end, Pos, [Tree]),
+        [{Pos, NewTree} | map(Fun, Rest)];
+    {arity, 3} ->
+        [NewTree] = map_simple(Fun, Pos, [Tree]),
+        [{Pos, NewTree} | map(Fun, Rest)]
+    end.
 
 map_simple(_Fun, _Pos, []) ->
     [];
 map_simple(Fun, Pos, [{Key, Value, SubTree} | RestTree]) ->
-    Value2 = Fun({Pos, Key}, Value),
+    Value2 = Fun({Pos, Key}, Value, 
+            if SubTree == [] -> leaf; true -> branch end),
     [{Key, Value2, map_simple(Fun, Pos + 1, SubTree)} | map_simple(Fun, Pos, RestTree)].