Specialised set difference function for dtree.

1 files changed, 7 insertions, 3 deletions

diff --git a/src/dtree.erl b/src/dtree.erl
index 8eff0644..0e0785bf 100644
--- a/src/dtree.erl
+++ b/src/dtree.erl
@@ -91,8 +91,7 @@ take(PKs, SK, {P, S}) ->
         none         -> {[], {P, S}};
         {value, PKS} -> TakenPKS = gb_sets:from_list(PKs),
                         PKSInter = gb_sets:intersection(PKS, TakenPKS),
-                        PKSDiff =
-                            rabbit_misc:gb_sets_difference(PKS, TakenPKS),
+                        PKSDiff  = gb_sets_difference_unsafe(PKS, TakenPKS),
                         {KVs, P1} = take2(PKSInter, SK, P),
                         {KVs, {P1, case gb_sets:is_empty(PKSDiff) of
                                        true  -> gb_trees:delete(SK, S);
@@ -153,9 +152,14 @@ take_all2(PKS, P) ->
 prune(SKS, PKS, S) ->
     gb_sets:fold(fun (SK0, S0) ->
                          PKS1 = gb_trees:get(SK0, S0),
-                         PKS2 = rabbit_misc:gb_sets_difference(PKS1, PKS),
+                         PKS2 = gb_sets_difference_unsafe(PKS1, PKS),
                          case gb_sets:is_empty(PKS2) of
                              true  -> gb_trees:delete(SK0, S0);
                              false -> gb_trees:update(SK0, PKS2, S0)
                          end
                  end, S, SKS).
+
+%% This function assumes that all the elements we're deleting from the
+%% first set are present.
+gb_sets_difference_unsafe(S1, S2) ->
+    lists:foldl(fun gb_sets:delete/2, S1, gb_sets:to_list(S2)).