diff --git a/ext/containers/rbtree_map/rbtree.c b/ext/containers/rbtree_map/rbtree.c index 24cf1e4d..bd609a86 100644 --- a/ext/containers/rbtree_map/rbtree.c +++ b/ext/containers/rbtree_map/rbtree.c @@ -412,8 +412,9 @@ static VALUE rbtree_max_key(VALUE self) { static VALUE rbtree_delete(VALUE self, VALUE key) { VALUE deleted_value; rbtree *tree = get_tree_from_self(self); - if(!tree->root) + if(get(tree, tree->root, key) == Qnil) { return Qnil; + } tree->root = delete(tree, tree->root, key, &deleted_value); if(tree->root) diff --git a/lib/containers/rb_tree_map.rb b/lib/containers/rb_tree_map.rb index c7bf63c8..9386bc21 100644 --- a/lib/containers/rb_tree_map.rb +++ b/lib/containers/rb_tree_map.rb @@ -129,7 +129,7 @@ def max_key # map.min_key #=> "GA" def delete(key) result = nil - if @root + if @root && get(key) @root, result = delete_recursive(@root, key) @root.color = :black if @root end diff --git a/spec/rb_tree_map_spec.rb b/spec/rb_tree_map_spec.rb index 32040aaa..63ba6990 100644 --- a/spec/rb_tree_map_spec.rb +++ b/spec/rb_tree_map_spec.rb @@ -34,6 +34,7 @@ shared_examples "non-empty rbtree" do before(:each) do @num_items = 1000 + @non_existing_element = @num_items+1 @random_array = Array.new(@num_items) { rand(@num_items) } @random_array.each { |x| @tree[x] = x } end @@ -51,7 +52,7 @@ end it "should not #has_key? keys it doesn't have" do - expect(@tree.has_key?(100000)).to be false + expect(@tree.has_key?(@non_existing_element)).to be false end it "should #has_key? keys it does have" do @@ -88,6 +89,10 @@ counter += 1 end end + + it "should allow deleting non-existent keys" do + expect(@tree.delete(@non_existing_element)).to be_nil + end end describe "empty rbtreemap" do