Made the rangetree.delete implementations return deleted entries.

Author: seanstrickland-wf

rangetree/interface.go

@@ -49,13 +49,16 @@ type Interval interface {
 type RangeTree interface {
 	// Add will add the provided entries to the tree.  Any entries that
 	// were overwritten will be returned in the order in which they
-	// were overwritten.  If a cell's addition does not overwrite, a nil
-	// is returned for that cell for its index in the provided cells.
+	// were overwritten.  If an entry's addition does not overwrite, a nil
+	// is returned for that entry's index in the provided cells.
 	Add(entries ...Entry) Entries
 	// Len returns the number of entries in the tree.
 	Len() uint64
 	// Delete will remove the provided entries from the tree.
-	Delete(entries ...Entry)
+	// Any entries that were deleted will be returned in the order in
+	// which they were deleted.  If an entry does not exist to be deleted,
+	// a nil is returned for that entry's index in the provided cells.
+	Delete(entries ...Entry) Entries
 	// Query will return a list of entries that fall within
 	// the provided interval.
 	Query(interval Interval) Entries

rangetree/orderedtree.go

@@ -64,33 +64,31 @@ func (ot *orderedTree) add(entry Entry) *node {
 }
 
 // Add will add the provided entries to the tree.  This method
-// returns a list of cells that were overwritten in the order
-// in which cells were received.  If a cell doesn't overwrite
+// returns a list of entries that were overwritten in the order
+// in which entries were received.  If an entry doesn't overwrite
 // anything, a nil will be returned for that entry in the returned
 // slice.
 func (ot *orderedTree) Add(entries ...Entry) Entries {
 	if len(entries) == 0 {
 		return nil
 	}
 
-	overwrittens := make(Entries, 0, len(entries))
-	for _, entry := range entries {
+	overwrittens := make(Entries, len(entries))
+	for i, entry := range entries {
 		if entry == nil {
 			continue
 		}
 
 		overwritten := ot.add(entry)
 		if overwritten != nil {
-			overwrittens = append(overwrittens, overwritten.entry)
-		} else {
-			overwrittens = append(overwrittens, nil)
+			overwrittens[i] = overwritten.entry
 		}
 	}
 
 	return overwrittens
 }
 
-func (ot *orderedTree) delete(entry Entry) {
+func (ot *orderedTree) delete(entry Entry) *node {
 	ot.resetPath()
 	var index int
 	var node *node
@@ -100,7 +98,7 @@ func (ot *orderedTree) delete(entry Entry) {
 		value := entry.ValueAtDimension(i)
 		node, index = list.get(value)
 		if node == nil { // there's nothing to delete
-			return
+			return nil
 		}
 
 		nb := &nodeBundle{list: list, index: index}
@@ -118,13 +116,32 @@ func (ot *orderedTree) delete(entry Entry) {
 			break
 		}
 	}
+
+	return node
 }
 
-// Delete will remove the entries from the tree.
-func (ot *orderedTree) Delete(entries ...Entry) {
-	for _, entry := range entries {
-		ot.delete(entry)
+// Delete will remove the provided entries from the tree.
+// Any entries that were deleted will be returned in the order in
+// which they were deleted.  If an entry does not exist to be deleted,
+// a nil is returned for that entry's index in the provided cells.
+func (ot *orderedTree) Delete(entries ...Entry) Entries {
+	if len(entries) == 0 {
+		return nil
 	}
+
+	deletedEntries := make(Entries, len(entries))
+	for i, entry := range entries {
+		if entry == nil {
+			continue
+		}
+
+		deleted := ot.delete(entry)
+		if deleted != nil {
+			deletedEntries[i] = deleted.entry
+		}
+	}
+
+	return deletedEntries
 }
 
 // Len returns the number of items in the tree.

rangetree/orderedtree_test.go

@@ -130,6 +130,32 @@ func TestOTAddLargeNumbersMultiDimension(t *testing.T) {
 	assert.Len(t, result, int(numItems))
 }
 
+func TestOTAddReturnsOverwritten(t *testing.T) {
+	tree := newOrderedTree(2)
+
+	starts := []uint64{0, 4, 2, 1, 3}
+
+	entries := make(Entries, 0, len(starts))
+	for _, start := range starts {
+		entries = append(entries, constructMockEntry(start, int64(start), int64(start)))
+	}
+
+	overwritten := tree.Add(entries...)
+
+	assert.Equal(t, Entries{nil, nil, nil, nil, nil}, overwritten)
+
+	oldEntry := entries[2]
+	newEntry := constructMockEntry(10, oldEntry.ValueAtDimension(1),
+		oldEntry.ValueAtDimension(2))
+	overwritten = tree.Add(newEntry)
+
+	assert.Equal(t, Entries{oldEntry}, overwritten)
+
+	result := tree.Query(constructMockInterval(dimension{0, 5}, dimension{0, 5}))
+	assert.Len(t, result, 5)
+	assert.Equal(t, uint64(5), tree.Len())
+}
+
 func BenchmarkOTAddItemsMultiDimensions(b *testing.B) {
 	numItems := b.N
 	entries := make(Entries, 0, numItems)
@@ -256,6 +282,29 @@ func TestOTDeleteEmptyTreeMultiDimensions(t *testing.T) {
 	assert.Equal(t, uint64(0), tree.Len())
 }
 
+func TestOTDeleteReturnsDeleted(t *testing.T) {
+	tree := newOrderedTree(2)
+
+	entries := NewEntries()
+	starts := []uint64{0, 4, 2, 1, 3}
+	for _, start := range starts {
+		entries = append(entries, constructMockEntry(start, int64(start), int64(start)))
+	}
+
+	tree.Add(entries...)
+
+	deleted := tree.Delete(entries[2], constructMockEntry(10, 10, 10))
+
+	assert.Equal(t, Entries{entries[2], nil}, deleted)
+
+	result := tree.Query(constructMockInterval(dimension{0, 11}, dimension{0, 11}))
+
+	assert.Len(t, result, 4)
+	assert.Equal(t, uint64(4), tree.Len())
+
+	assert.NotContains(t, result, entries[2])
+}
+
 func BenchmarkOTDeleteItemsMultiDimensions(b *testing.B) {
 	numItems := uint64(1000)
 	entries := make(Entries, 0, numItems)

rangetree/skiplist/skiplist.go

@@ -192,14 +192,15 @@ func (rt *skipListRT) add(entry rangetree.Entry) rangetree.Entry {
 	panic(`Ran out of dimensions before for loop completed.`)
 }
 
-// Add will add the provided entries to the tree.  Any entries that
-// were overwritten will be returned in the order in which they
-// were overwritten.  If an entry's addition does not overwrite, a nil
-// is returned for that cell for its index in the provided entries.
+// Add will add the provided entries to the tree.  This method
+// returns a list of entries that were overwritten in the order
+// in which entries were received.  If an entry doesn't overwrite
+// anything, a nil will be returned for that entry in the returned
+// slice.
 func (rt *skipListRT) Add(entries ...rangetree.Entry) rangetree.Entries {
-	overwritten := make(rangetree.Entries, 0, len(entries))
-	for _, e := range entries {
-		overwritten = append(overwritten, rt.add(e))
+	overwritten := make(rangetree.Entries, len(entries))
+	for i, e := range entries {
+		overwritten[i] = rt.add(e)
 	}
 
 	return overwritten
@@ -284,10 +285,16 @@ func (rt *skipListRT) delete(entry rangetree.Entry) rangetree.Entry {
 }
 
 // Delete will remove the provided entries from the tree.
-func (rt *skipListRT) Delete(entries ...rangetree.Entry) {
-	for _, e := range entries {
-		rt.delete(e)
+// Any entries that were deleted will be returned in the order in
+// which they were deleted.  If an entry does not exist to be deleted,
+// a nil is returned for that entry's index in the provided cells.
+func (rt *skipListRT) Delete(entries ...rangetree.Entry) rangetree.Entries {
+	deletedEntries := make(rangetree.Entries, len(entries))
+	for i, e := range entries {
+		deletedEntries[i] = rt.delete(e)
 	}
+
+	return deletedEntries
 }
 
 func (rt *skipListRT) apply(sl *skip.SkipList, dimension uint64,

rangetree/skiplist/skiplist_test.go

@@ -103,7 +103,8 @@ func TestRTSingleDimensionDelete(t *testing.T) {
 	m2 := newMockEntry(2)
 	rt.Add(m1, m2)
 
-	rt.Delete(m1, m2)
+	deleted := rt.Delete(m1, newMockEntry(10), m2)
+	assert.Equal(t, rangetree.Entries{m1, nil, m2}, deleted)
 	assert.Equal(t, uint64(0), rt.Len())
 	assert.Equal(t, rangetree.Entries{nil, nil}, rt.Get(m1, m2))
 }
@@ -114,7 +115,8 @@ func TestRTMultiDimensionDelete(t *testing.T) {
 	m2 := newMockEntry(4, 6)
 	rt.Add(m1, m2)
 
-	rt.Delete(m1, m2)
+	deleted := rt.Delete(m1, newMockEntry(10, 10), m2)
+	assert.Equal(t, rangetree.Entries{m1, nil, m2}, deleted)
 	assert.Equal(t, uint64(0), rt.Len())
 	assert.Equal(t, rangetree.Entries{nil, nil}, rt.Get(m1, m2))
 }