The "blind" difference is unsafe in 'prune'.

1 files changed, 4 insertions, 6 deletions

diff --git a/src/dtree.erl b/src/dtree.erl
index 0e0785bf..b792c681 100644
--- a/src/dtree.erl
+++ b/src/dtree.erl
@@ -91,7 +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  = gb_sets_difference_unsafe(PKS, TakenPKS),
+                        PKSDiff  = gb_sets_difference(PKS, PKSInter),
                         {KVs, P1} = take2(PKSInter, SK, P),
                         {KVs, {P1, case gb_sets:is_empty(PKSDiff) of
                                        true  -> gb_trees:delete(SK, S);
@@ -152,14 +152,12 @@ take_all2(PKS, P) ->
 prune(SKS, PKS, S) ->
     gb_sets:fold(fun (SK0, S0) ->
                          PKS1 = gb_trees:get(SK0, S0),
-                         PKS2 = gb_sets_difference_unsafe(PKS1, PKS),
+                         PKS2 = gb_sets_difference(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)).
+gb_sets_difference(S1, S2) ->
+    lists:foldl(fun gb_sets:delete_any/2, S1, gb_sets:to_list(S2)).