Make treeset take an interface, make the trees satisfy that interface

sets/treeset/enumerable.go

@@ -6,7 +6,7 @@ package treeset
 
 import (
 	"github.com/emirpasic/gods/containers"
-	rbt "github.com/emirpasic/gods/trees/redblacktree"
+	rbt "github.com/spewspews/gods/trees/redblacktree"
 )
 
 func assertEnumerableImplementation() {
@@ -24,7 +24,7 @@ func (set *Set) Each(f func(index int, value interface{})) {
 // Map invokes the given function once for each element and returns a
 // container containing the values returned by the given function.
 func (set *Set) Map(f func(index int, value interface{}) interface{}) *Set {
-	newSet := &Set{tree: rbt.NewWith(set.tree.Comparator)}
+	newSet := &Set{tree: set.tree.New()}
 	iterator := set.Iterator()
 	for iterator.Next() {
 		newSet.Add(f(iterator.Index(), iterator.Value()))
@@ -34,7 +34,7 @@ func (set *Set) Map(f func(index int, value interface{}) interface{}) *Set {
 
 // Select returns a new container containing all elements for which the given function returns a true value.
 func (set *Set) Select(f func(index int, value interface{}) bool) *Set {
-	newSet := &Set{tree: rbt.NewWith(set.tree.Comparator)}
+	newSet := &Set{tree: rbt.NewWith(set.tree.Comparator())}
 	iterator := set.Iterator()
 	for iterator.Next() {
 		if f(iterator.Index(), iterator.Value()) {

sets/treeset/iterator.go

@@ -6,7 +6,7 @@ package treeset
 
 import (
 	"github.com/emirpasic/gods/containers"
-	rbt "github.com/emirpasic/gods/trees/redblacktree"
+	"github.com/spewspews/gods/trees"
 )
 
 func assertIteratorImplementation() {
@@ -16,8 +16,8 @@ func assertIteratorImplementation() {
 // Iterator returns a stateful iterator whose values can be fetched by an index.
 type Iterator struct {
 	index    int
-	iterator rbt.Iterator
-	tree     *rbt.Tree
+	iterator containers.ReverseIteratorWithKey
+	tree     trees.Tree
 }
 
 // Iterator holding the iterator's state

sets/treeset/treeset.go

@@ -11,10 +11,12 @@ package treeset
 
 import (
 	"fmt"
+	"strings"
+
 	"github.com/emirpasic/gods/sets"
-	rbt "github.com/emirpasic/gods/trees/redblacktree"
 	"github.com/emirpasic/gods/utils"
-	"strings"
+	"github.com/spewspews/gods/trees"
+	rbt "github.com/spewspews/gods/trees/redblacktree"
 )
 
 func assertSetImplementation() {
@@ -23,7 +25,7 @@ func assertSetImplementation() {
 
 // Set holds elements in a red-black tree
 type Set struct {
-	tree *rbt.Tree
+	tree trees.Tree
 }
 
 var itemExists = struct{}{}
@@ -43,6 +45,11 @@ func NewWithStringComparator() *Set {
 	return &Set{tree: rbt.NewWithStringComparator()}
 }
 
+// NewWithTree instantiates a new empty set with given tree
+func NewWithTree(tree trees.Tree) (set *Set) {
+	return &Set{tree: tree}
+}
+
 // Add adds the items (one or more) to the set.
 func (set *Set) Add(items ...interface{}) {
 	for _, item := range items {

sets/treeset/treeset_test.go

@@ -7,6 +7,8 @@ package treeset
 import (
 	"fmt"
 	"testing"
+
+	"github.com/spewspews/gods/trees/avltree"
 )
 
 func TestSetAdd(t *testing.T) {
@@ -27,6 +29,24 @@ func TestSetAdd(t *testing.T) {
 	}
 }
 
+func TestSetAVLAdd(t *testing.T) {
+	set := NewWithTree(avltree.NewWithIntComparator())
+	set.Add()
+	set.Add(1)
+	set.Add(2)
+	set.Add(2, 3)
+	set.Add()
+	if actualValue := set.Empty(); actualValue != false {
+		t.Errorf("Got %v expected %v", actualValue, false)
+	}
+	if actualValue := set.Size(); actualValue != 3 {
+		t.Errorf("Got %v expected %v", actualValue, 3)
+	}
+	if actualValue, expectedValue := fmt.Sprintf("%d%d%d", set.Values()...), "123"; actualValue != expectedValue {
+		t.Errorf("Got %v expected %v", actualValue, expectedValue)
+	}
+}
+
 func TestSetContains(t *testing.T) {
 	set := NewWithIntComparator()
 	set.Add(3, 1, 2)

trees/avltree/avltree.go

@@ -12,8 +12,8 @@ import (
 	"io/ioutil"
 	"log"
 
-	"github.com/emirpasic/gods/trees"
 	"github.com/emirpasic/gods/utils"
+	"github.com/spewspews/gods/trees"
 )
 
 func assertTreeImplementation() {
@@ -26,7 +26,7 @@ var dbgLog = log.New(ioutil.Discard, "avltree: ", log.LstdFlags)
 type Tree struct {
 	Root       *Node
 	size       int
-	Comparator utils.Comparator
+	comparator utils.Comparator
 }
 
 // Node is a single element within the tree
@@ -40,17 +40,27 @@ type Node struct {
 
 // NewWith instantiates an AVL tree with the custom comparator.
 func NewWith(comparator utils.Comparator) *Tree {
-	return &Tree{Comparator: comparator}
+	return &Tree{comparator: comparator}
 }
 
 // NewWithIntComparator instantiates an AVL tree with the IntComparator, i.e. keys are of type int.
 func NewWithIntComparator() *Tree {
-	return &Tree{Comparator: utils.IntComparator}
+	return &Tree{comparator: utils.IntComparator}
 }
 
 // NewWithStringComparator instantiates an AVL tree with the StringComparator, i.e. keys are of type string.
 func NewWithStringComparator() *Tree {
-	return &Tree{Comparator: utils.StringComparator}
+	return &Tree{comparator: utils.StringComparator}
+}
+
+// Comparator returns the comparator function for the tree.
+func (t *Tree) Comparator() utils.Comparator {
+	return t.comparator
+}
+
+// New returns a new empty tree with the same comparator.
+func (t *Tree) New() trees.Tree {
+	return &Tree{comparator: t.comparator}
 }
 
 // Size returns the number of elements stored in the tree.
@@ -75,7 +85,7 @@ func (t *Tree) Clear() {
 func (t *Tree) Get(key interface{}) (value interface{}, found bool) {
 	n := t.Root
 	for n != nil {
-		cmp := t.Comparator(key, n.Key)
+		cmp := t.comparator(key, n.Key)
 		switch {
 		case cmp == 0:
 			return n.Value, true
@@ -100,7 +110,7 @@ func (t *Tree) Floor(key interface{}) (floor *Node, found bool) {
 	found = false
 	n := t.Root
 	for n != nil {
-		c := t.Comparator(key, n.Key)
+		c := t.comparator(key, n.Key)
 		switch {
 		case c == 0:
 			return n, true
@@ -129,7 +139,7 @@ func (t *Tree) Ceiling(key interface{}) (floor *Node, found bool) {
 	found = false
 	n := t.Root
 	for n != nil {
-		c := t.Comparator(key, n.Key)
+		c := t.comparator(key, n.Key)
 		switch {
 		case c == 0:
 			return n, true
@@ -158,7 +168,7 @@ func (t *Tree) Put(key interface{}, value interface{}) {
 			return true
 		}
 
-		c := t.Comparator(key, q.Key)
+		c := t.comparator(key, q.Key)
 		if c == 0 {
 			q.Key = key
 			q.Value = value
@@ -192,7 +202,7 @@ func (t *Tree) Remove(key interface{}) {
 			return false
 		}
 
-		c := t.Comparator(key, q.Key)
+		c := t.comparator(key, q.Key)
 		if c == 0 {
 			t.size--
 			if q.c[1] == nil {

trees/avltree/iterator.go

@@ -24,8 +24,8 @@ const (
 )
 
 // Iterator returns a stateful iterator whose elements are key/value pairs.
-func (tree *Tree) Iterator() Iterator {
-	return Iterator{tree: tree, node: nil, position: begin}
+func (tree *Tree) Iterator() containers.ReverseIteratorWithKey {
+	return &Iterator{tree: tree, node: nil, position: begin}
 }
 
 // Next moves the iterator to the next element and returns true if there was a next element in the container.

trees/redblacktree/iterator.go

@@ -24,8 +24,8 @@ const (
 )
 
 // Iterator returns a stateful iterator whose elements are key/value pairs.
-func (tree *Tree) Iterator() Iterator {
-	return Iterator{tree: tree, node: nil, position: begin}
+func (tree *Tree) Iterator() containers.ReverseIteratorWithKey {
+	return &Iterator{tree: tree, node: nil, position: begin}
 }
 
 // Next moves the iterator to the next element and returns true if there was a next element in the container.
@@ -55,7 +55,7 @@ func (iterator *Iterator) Next() bool {
 		node := iterator.node
 		for iterator.node.Parent != nil {
 			iterator.node = iterator.node.Parent
-			if iterator.tree.Comparator(node.Key, iterator.node.Key) <= 0 {
+			if iterator.tree.comparator(node.Key, iterator.node.Key) <= 0 {
 				goto between
 			}
 		}
@@ -97,7 +97,7 @@ func (iterator *Iterator) Prev() bool {
 		node := iterator.node
 		for iterator.node.Parent != nil {
 			iterator.node = iterator.node.Parent
-			if iterator.tree.Comparator(node.Key, iterator.node.Key) >= 0 {
+			if iterator.tree.comparator(node.Key, iterator.node.Key) >= 0 {
 				goto between
 			}
 		}

trees/redblacktree/redblacktree.go

@@ -13,8 +13,9 @@ package redblacktree
 
 import (
 	"fmt"
-	"github.com/emirpasic/gods/trees"
+
 	"github.com/emirpasic/gods/utils"
+	"github.com/spewspews/gods/trees"
 )
 
 func assertTreeImplementation() {
@@ -31,7 +32,7 @@ const (
 type Tree struct {
 	Root       *Node
 	size       int
-	Comparator utils.Comparator
+	comparator utils.Comparator
 }
 
 // Node is a single element within the tree
@@ -46,17 +47,27 @@ type Node struct {
 
 // NewWith instantiates a red-black tree with the custom comparator.
 func NewWith(comparator utils.Comparator) *Tree {
-	return &Tree{Comparator: comparator}
+	return &Tree{comparator: comparator}
 }
 
 // NewWithIntComparator instantiates a red-black tree with the IntComparator, i.e. keys are of type int.
 func NewWithIntComparator() *Tree {
-	return &Tree{Comparator: utils.IntComparator}
+	return &Tree{comparator: utils.IntComparator}
 }
 
 // NewWithStringComparator instantiates a red-black tree with the StringComparator, i.e. keys are of type string.
 func NewWithStringComparator() *Tree {
-	return &Tree{Comparator: utils.StringComparator}
+	return &Tree{comparator: utils.StringComparator}
+}
+
+// Comparator returns the comparator function for the tree.
+func (t *Tree) Comparator() utils.Comparator {
+	return t.comparator
+}
+
+// New returns a new empty tree with the same comparator.
+func (t *Tree) New() trees.Tree {
+	return &Tree{comparator: t.comparator}
 }
 
 // Put inserts node into the tree.
@@ -69,7 +80,7 @@ func (tree *Tree) Put(key interface{}, value interface{}) {
 		node := tree.Root
 		loop := true
 		for loop {
-			compare := tree.Comparator(key, node.Key)
+			compare := tree.comparator(key, node.Key)
 			switch {
 			case compare == 0:
 				node.Key = key
@@ -204,7 +215,7 @@ func (tree *Tree) Floor(key interface{}) (floor *Node, found bool) {
 	found = false
 	node := tree.Root
 	for node != nil {
-		compare := tree.Comparator(key, node.Key)
+		compare := tree.comparator(key, node.Key)
 		switch {
 		case compare == 0:
 			return node, true
@@ -233,7 +244,7 @@ func (tree *Tree) Ceiling(key interface{}) (ceiling *Node, found bool) {
 	found = false
 	node := tree.Root
 	for node != nil {
-		compare := tree.Comparator(key, node.Key)
+		compare := tree.comparator(key, node.Key)
 		switch {
 		case compare == 0:
 			return node, true
@@ -300,7 +311,7 @@ func output(node *Node, prefix string, isTail bool, str *string) {
 func (tree *Tree) lookup(key interface{}) *Node {
 	node := tree.Root
 	for node != nil {
-		compare := tree.Comparator(key, node.Key)
+		compare := tree.comparator(key, node.Key)
 		switch {
 		case compare == 0:
 			return node

trees/trees.go

@@ -9,10 +9,21 @@
 // Reference: https://en.wikipedia.org/wiki/Tree_%28data_structure%29
 package trees
 
-import "github.com/emirpasic/gods/containers"
+import (
+	"github.com/emirpasic/gods/containers"
+	"github.com/emirpasic/gods/utils"
+)
 
 // Tree interface that all trees implement
 type Tree interface {
+	Comparator() utils.Comparator
+	New() Tree
+	Iterator() containers.ReverseIteratorWithKey
+	Put(key interface{}, value interface{})
+	Remove(key interface{})
+	Get(key interface{}) (interface{}, bool)
+	Keys() []interface{}
+
 	containers.Container
 	// Empty() bool
 	// Size() int