Add filter (#16)

Author: johnfercher

README.md

@@ -41,6 +41,7 @@ go get -u github.com/johnfercher/go-tree
 * [GetStructure](https://pkg.go.dev/github.com/johnfercher/go-tree/tree#Node.GetStructure)
 * [IsLeaf](https://pkg.go.dev/github.com/johnfercher/go-tree/tree#Node.IsLeaf)
 * [IsRoot](https://pkg.go.dev/github.com/johnfercher/go-tree/tree#Node.IsLeaf)
+* [Filter](https://pkg.go.dev/github.com/johnfercher/go-tree/tree#Node.Filter)
 
 ### Tree
 * [New](https://pkg.go.dev/github.com/johnfercher/go-tree/tree#New)
@@ -50,6 +51,7 @@ go get -u github.com/johnfercher/go-tree
 * [Get](https://pkg.go.dev/github.com/johnfercher/go-tree/tree#Tree.Get)
 * [GetRoot](https://pkg.go.dev/github.com/johnfercher/go-tree/tree#Tree.GetRoot)
 * [GetStructure](https://pkg.go.dev/github.com/johnfercher/go-tree/tree#Tree.GetStructure)
+* [Filter](https://pkg.go.dev/github.com/johnfercher/go-tree/tree#Tree.Filter)
 
 ## Example
 

node/node.go

@@ -72,12 +72,20 @@ func (n *Node[T]) Backtrack() []*Node[T] {
 // GetStructure retrieves the node structure.
 func (n *Node[T]) GetStructure() []string {
 	var structure []string
+	var current string
+
 	if n.previous == nil {
-		structure = append(structure, fmt.Sprintf("(NULL) -> (%d)", n.id))
+		current = fmt.Sprintf("(NULL) -> (%d)", n.id)
 	} else {
-		structure = append(structure, fmt.Sprintf("(%d) -> (%d)", n.previous.id, n.id))
+		current = fmt.Sprintf("(%d) -> (%d)", n.previous.id, n.id)
+	}
+
+	if n.nexts != nil {
+		current += ", "
 	}
 
+	structure = append(structure, current)
+
 	for _, next := range n.nexts {
 		innerStructure := next.GetStructure()
 		structure = append(structure, innerStructure...)
@@ -91,3 +99,21 @@ func (n *Node[T]) AddNext(node *Node[T]) {
 	node.previous = n
 	n.nexts = append(n.nexts, node)
 }
+
+// Filter remove all sub-nodes that doesn´t respect a rule.
+func (n *Node[T]) Filter(filterFunc func(obj T) bool) (*Node[T], bool) {
+	if !filterFunc(n.GetData()) {
+		return nil, false
+	}
+
+	newNode := New(n.GetData()).WithID(n.GetID())
+
+	for _, next := range n.nexts {
+		innerNode, ok := next.Filter(filterFunc)
+		if ok {
+			newNode.AddNext(innerNode)
+		}
+	}
+
+	return newNode, true
+}

node/node_test.go

@@ -343,3 +343,27 @@ func TestNode_GetStructure(t *testing.T) {
 		assert.NotEmpty(t, str)
 	}
 }
+
+func TestNode_Filter(t *testing.T) {
+	// Arrange
+	n0 := node.New(0).WithID(0)
+	n1 := node.New(1).WithID(1)
+	n2 := node.New(2).WithID(2)
+	n3 := node.New(3).WithID(3)
+
+	n0.AddNext(n1)
+	n0.AddNext(n2)
+	n1.AddNext(n3)
+
+	// Act
+	newN0, ok := n0.Filter(func(obj int) bool {
+		return obj%2 == 0
+	})
+
+	// Assert
+	assert.True(t, ok)
+	assert.Equal(t, 0, newN0.GetID())
+
+	nexts := newN0.GetNexts()
+	assert.Equal(t, 2, nexts[0].GetID())
+}

tree/tree.go

@@ -73,6 +73,23 @@ func (t *Tree[T]) GetStructure() ([]string, bool) {
 	return t.root.GetStructure(), true
 }
 
+// Filter remove all sub-nodes that doesn´t respect a rule.
+func (t *Tree[T]) Filter(filterFunc func(obj T) bool) (*Tree[T], bool) {
+	if t.root == nil {
+		return nil, false
+	}
+
+	newRoot, ok := t.root.Filter(filterFunc)
+	if !ok {
+		return nil, false
+	}
+
+	newTree := New[T]()
+	newTree.AddRoot(newRoot)
+
+	return newTree, true
+}
+
 func (t *Tree[T]) add(parentID int, parentNode *node.Node[T], newNode *node.Node[T]) bool {
 	if parentID == parentNode.GetID() {
 		parentNode.AddNext(newNode)

tree/tree_test.go

@@ -230,3 +230,56 @@ func TestTree_GetStructure_WhenThereIsRoot_ShouldReturnTrue(t *testing.T) {
 		assert.NotEmpty(t, str)
 	}
 }
+
+func TestTree_Filter_WhenThereIsNoRoot_ShouldNotWork(t *testing.T) {
+	// Arrange
+	tr := tree.New[int]()
+
+	// Act
+	newTree, ok := tr.Filter(func(obj int) bool {
+		return obj%2 == 0
+	})
+
+	// Assert
+	assert.False(t, ok)
+	assert.Nil(t, newTree)
+}
+
+func TestTree_Filter_WhenThereIsRootButRulesDoesntApply_ShouldNotWork(t *testing.T) {
+	// Arrange
+	tr := tree.New[int]()
+	tr.AddRoot(node.New(1).WithID(1))
+
+	// Act
+	newTree, ok := tr.Filter(func(obj int) bool {
+		return obj%2 == 0
+	})
+
+	// Assert
+	assert.False(t, ok)
+	assert.Nil(t, newTree)
+}
+
+func TestTree_Filter_WhenEverythingIsOk_ShouldWork(t *testing.T) {
+	// Arrange
+	tr := tree.New[int]()
+
+	tr.AddRoot(node.New(0).WithID(0))
+	tr.Add(0, node.New(1).WithID(1))
+	tr.Add(0, node.New(2).WithID(2))
+	tr.Add(1, node.New(3).WithID(3))
+
+	// Act
+	newTree, ok := tr.Filter(func(obj int) bool {
+		return obj%2 == 0
+	})
+
+	newN0, _ := newTree.GetRoot()
+
+	// Assert
+	assert.True(t, ok)
+	assert.Equal(t, 0, newN0.GetID())
+
+	nexts := newN0.GetNexts()
+	assert.Equal(t, 2, nexts[0].GetID())
+}