Merge pull request #108 from Workiva/rt_get

Add get methods.

Author: dustinhiatt-wf

mock/rangetree.go

@@ -50,3 +50,12 @@ func (m *RangeTree) InsertAtDimension(dimension uint64, index,
 func (m *RangeTree) Apply(interval rangetree.Interval, fn func(rangetree.Entry) bool) {
 	m.Called(interval, fn)
 }
+
+func (m *RangeTree) Get(entries ...rangetree.Entry) rangetree.Entries {
+	ifc := m.Called(entries).Get(0)
+	if ifc == nil {
+		return nil
+	}
+
+	return ifc.(rangetree.Entries)
+}

rangetree/immutable.go

@@ -234,6 +234,35 @@ func (irt *immutableRangeTree) Query(interval Interval) Entries {
 	return entries
 }
 
+func (irt *immutableRangeTree) get(entry Entry) Entry {
+	on := irt.top
+	for i := uint64(1); i <= irt.dimensions; i++ {
+		n, _ := on.get(entry.ValueAtDimension(i))
+		if n == nil {
+			return nil
+		}
+		if i == irt.dimensions {
+			return n.entry
+		}
+		on = n.orderedNodes
+	}
+
+	return nil
+}
+
+// Get returns any entries that exist at the addresses provided by the
+// given entries.  Entries are returned in the order in which they are
+// received.  If an entry cannot be found, a nil is returned in its
+// place.
+func (irt *immutableRangeTree) Get(entries ...Entry) Entries {
+	result := make(Entries, 0, len(entries))
+	for _, entry := range entries {
+		result = append(result, irt.get(entry))
+	}
+
+	return result
+}
+
 // Len returns the number of items in this tree.
 func (irt *immutableRangeTree) Len() uint64 {
 	return irt.number

rangetree/immutable_test.go

@@ -514,6 +514,16 @@ func TestImmutableInsertInvalidNumber(t *testing.T) {
 	assert.Equal(t, tree, tree1)
 }
 
+func TestImmutableGet(t *testing.T) {
+	tree, entries := constructMultiDimensionalImmutableTree(2)
+
+	result := tree.Get(entries...)
+	assert.Equal(t, entries, result)
+
+	result = tree.Get(constructMockEntry(10000, 5000, 5000))
+	assert.Equal(t, Entries{nil}, result)
+}
+
 func BenchmarkImmutableInsertFirstDimension(b *testing.B) {
 	numItems := int64(100000)
 

rangetree/interface.go

@@ -67,6 +67,11 @@ type RangeTree interface {
 	// cancel iteration.  Altering the entry in such a way that its location
 	// changes will result in undefined behavior.
 	Apply(interval Interval, fn func(Entry) bool)
+	// Get returns any entries that exist at the addresses provided by the
+	// given entries.  Entries are returned in the order in which they are
+	// received.  If an entry cannot be found, a nil is returned in its
+	// place.
+	Get(entries ...Entry) Entries
 	// InsertAtDimension will increment items at and above the given index
 	// by the number provided.  Provide a negative number to to decrement.
 	// Returned are two lists.  The first list is a list of entries that

rangetree/orderedtree.go

@@ -120,6 +120,35 @@ func (ot *orderedTree) delete(entry Entry) *node {
 	return node
 }
 
+func (ot *orderedTree) get(entry Entry) Entry {
+	on := ot.top
+	for i := uint64(1); i <= ot.dimensions; i++ {
+		n, _ := on.get(entry.ValueAtDimension(i))
+		if n == nil {
+			return nil
+		}
+		if i == ot.dimensions {
+			return n.entry
+		}
+		on = n.orderedNodes
+	}
+
+	return nil
+}
+
+// Get returns any entries that exist at the addresses provided by the
+// given entries.  Entries are returned in the order in which they are
+// received.  If an entry cannot be found, a nil is returned in its
+// place.
+func (ot *orderedTree) Get(entries ...Entry) Entries {
+	result := make(Entries, 0, len(entries))
+	for _, entry := range entries {
+		result = append(result, ot.get(entry))
+	}
+
+	return result
+}
+
 // 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,

rangetree/orderedtree_test.go

@@ -346,6 +346,16 @@ func TestOverwrites(t *testing.T) {
 	assert.Equal(t, Entries{entry}, overwritten)
 }
 
+func TestGet(t *testing.T) {
+	tree, entries := constructMultiDimensionalOrderedTree(2)
+
+	result := tree.Get(entries...)
+	assert.Equal(t, entries, result)
+
+	result = tree.Get(constructMockEntry(10000, 5000, 5000))
+	assert.Equal(t, Entries{nil}, result)
+}
+
 func TestTreeApply(t *testing.T) {
 	tree, entries := constructMultiDimensionalOrderedTree(2)
 
@@ -610,3 +620,25 @@ func BenchmarkDeleteSecondDimension(b *testing.B) {
 		tree.InsertAtDimension(2, 0, -1)
 	}
 }
+
+func BenchmarkGetMultiDimensions(b *testing.B) {
+	numItemsX := 10000
+	numItemsY := 100
+
+	tree := newOrderedTree(2)
+	entries := make(Entries, 0, numItemsY*numItemsX)
+
+	for i := 0; i < numItemsX; i++ {
+		for j := 0; j < numItemsY; j++ {
+			e := constructMockEntry(uint64(j*numItemsY+i), int64(i), int64(j))
+			entries = append(entries, e)
+		}
+	}
+
+	tree.Add(entries...)
+	b.ResetTimer()
+
+	for i := 0; i < b.N; i++ {
+		tree.Get(entries[i%len(entries)])
+	}
+}