add new test and make test coverage to 100% :)

iter_test.go

@@ -0,0 +1,47 @@
+package rbtree
+
+import (
+	"testing"
+
+	"github.com/fanyang01/tree/common"
+	"github.com/stretchr/testify/assert"
+)
+
+func TestIter(t *testing.T) {
+	n := 1 << 10
+	tr := New(common.CompareInt)
+
+	assert.Nil(t, tr.PostorderFirst())
+	assert.Nil(t, tr.PreorderFirst())
+
+	for i := 0; i < n; i++ {
+		tr.Insert(i)
+	}
+
+	assert.Nil(t, tr.Prev(tr.First()))
+	assert.Nil(t, tr.Next(tr.Last()))
+
+	for i, x := 0, tr.First(); i < n; i++ {
+		assert.NotNil(t, x)
+		assert.Equal(t, i, x.Value())
+		x = tr.Next(x)
+	}
+	for i, x := n-1, tr.Last(); i >= 0; i-- {
+		assert.NotNil(t, x)
+		assert.Equal(t, i, x.Value())
+		x = tr.Prev(x)
+	}
+
+	tr.Clean()
+	for i := 0; i < n; i++ {
+		tr.Insert(r.Intn(n))
+	}
+
+	for x := tr.PostorderFirst(); x != nil; x = tr.PostorderNext(x) {
+		assert.NotNil(t, tr.PostorderFirstNode(x))
+	}
+
+	for x := tr.PreorderFirst(); x != nil; x = tr.PreorderNext(x) {
+		assert.NotNil(t, tr.PreorderLastNode(x))
+	}
+}

rbtree_test.go → tree_test.go

@@ -11,9 +11,10 @@ import (
 
 var r = rand.New(rand.NewSource(time.Now().UnixNano()))
 
-func TestAssert(t *testing.T) {
+func TestTree(t *testing.T) {
 	n := 1 << 16
 	tr := New(common.CompareInt)
+
 	assert.True(t, tr.IsEmpty())
 
 	for i := 0; i < n; i++ {
@@ -27,20 +28,6 @@ func TestAssert(t *testing.T) {
 	_, ok := tr.Insert(0)
 	assert.False(t, ok)
 
-	assert.Nil(t, tr.Prev(tr.First()))
-	assert.Nil(t, tr.Next(tr.Last()))
-
-	for i, x := 0, tr.First(); i < n; i++ {
-		assert.NotNil(t, x)
-		assert.Equal(t, i, x.Value().(int))
-		x = tr.Next(x)
-	}
-	for i, x := n-1, tr.Last(); i >= 0; i-- {
-		assert.NotNil(t, x)
-		assert.Equal(t, i, x.Value().(int))
-		x = tr.Prev(x)
-	}
-
 	for i := n - 1; i >= n/2; i-- {
 		tr.Delete(tr.Search(i))
 		assert.Nil(t, tr.Search(i))
@@ -74,6 +61,12 @@ func TestAssert(t *testing.T) {
 	assert.Nil(t, tr.First())
 	assert.Nil(t, tr.Last())
 
+	tr.Insert(1)
+	assert.Nil(t, tr.Root().Parent())
+	assert.Nil(t, tr.Root().Left())
+	assert.Nil(t, tr.Root().Right())
+	tr.Clean()
+
 	for i := 0; i < n; i++ {
 		random := r.Intn(n)
 		tr.Insert(random)

walk.go

@@ -11,7 +11,7 @@ type Visitor interface {
 type VisitFunc func(n *Node) bool
 
 // Visit implements the Visitor interface
-func (f WalkerFunc) Visit(n *Node) (w Visitor) {
+func (f VisitFunc) Visit(n *Node) (w Visitor) {
 	if ok := f(n); ok {
 		return f
 	}
@@ -71,8 +71,9 @@ func (t *Tree) WalkPreorder(v Visitor) {
 
 // WalkSubPreorder traverses subtree rooted at x in pre-order, x self is also visited.
 func (t *Tree) WalkSubPreorder(v Visitor, x *Node) {
-	for n := t.PreorderLastNode(x); x != n; x = t.PreorderNext(n) {
-		v = v.Visit(n)
+	var n *Node
+	for n = t.PreorderLastNode(x); x != n; x = t.PreorderNext(x) {
+		v = v.Visit(x)
 		if v == nil {
 			return
 		}

walk_test.go

@@ -0,0 +1,78 @@
+package rbtree
+
+import (
+	"testing"
+
+	"github.com/fanyang01/tree/common"
+	"github.com/stretchr/testify/assert"
+)
+
+func TestWalk(t *testing.T) {
+	n := 1 << 10
+	tr := New(common.CompareInt)
+	for i := 0; i < n; i++ {
+		tr.Insert(i)
+	}
+
+	var i int
+	fn := func(x *Node) bool {
+		assert.Equal(t, i, x.Value())
+		i++
+		return true
+	}
+	tr.Walk(VisitFunc(fn))
+
+	i = n - 1
+	fn = func(x *Node) bool {
+		assert.Equal(t, i, x.Value())
+		i--
+		return true
+	}
+	tr.WalkReverse(VisitFunc(fn))
+
+	fn = func(x *Node) bool {
+		left, right := x.Left(), x.Right()
+		if left != nil {
+			assert.True(t, x.Value().(int) > left.Value().(int))
+		}
+		if right != nil {
+			assert.True(t, x.Value().(int) < right.Value().(int))
+		}
+		return true
+	}
+	tr.WalkPostorder(VisitFunc(fn))
+	tr.WalkPreorder(VisitFunc(fn))
+
+	size := 0
+	fn = func(x *Node) bool {
+		size++
+		return true
+	}
+	tr.WalkPostorder(VisitFunc(fn))
+	assert.Equal(t, n, size)
+
+	size = 0
+	tr.WalkSubPostorder(VisitFunc(fn), tr.Root())
+	assert.Equal(t, n, size)
+
+	size = 0
+	tr.WalkPreorder(VisitFunc(fn))
+	assert.Equal(t, n, size)
+
+	size = 0
+	tr.WalkSubPreorder(VisitFunc(fn), tr.Root())
+	assert.Equal(t, n, size)
+
+	size = 0
+	fn = func(x *Node) bool {
+		size++
+		return false
+	}
+	tr.Walk(VisitFunc(fn))
+	tr.WalkReverse(VisitFunc(fn))
+	tr.WalkPostorder(VisitFunc(fn))
+	tr.WalkPreorder(VisitFunc(fn))
+	tr.WalkSubPostorder(VisitFunc(fn), tr.Root())
+	tr.WalkSubPreorder(VisitFunc(fn), tr.Root())
+	assert.Equal(t, 6, size)
+}