Adding code to walk trees

tree/tree.go

@@ -0,0 +1,116 @@
+package tree
+
+type Tree interface {
+	Parent() Tree
+	Child(i int) Tree
+	NumChildren() int
+}
+
+type WalkDirection int
+
+const (
+	Forward WalkDirection = iota
+	Reverse
+)
+
+type WalkOrder int
+
+const (
+	PreOrder WalkOrder = iota
+	PostOrder
+)
+
+// Visitor is the visitor function for a tree walk.
+// if cont is false on return, the walk terminates. If skipChildren is true
+// on return the children and their descendants of the current node are
+// skipped.
+type Visitor func(t Tree, depth int) (continu, skipChildren bool)
+
+// Walk walks the tree of which `tree` is a member node. This function walks the node
+// and it's children, but also flows up to the parent if this is not the root of the tree.
+// It's like continuing a tree walk from the root of the tree that was interrupted at the
+// specified node.
+//
+// `dir` specifies whether the walk of children is done from the last to first,
+// or first to last. As well WalkOrder specifies whether the parent is printed before or after children.
+// skip: if true, skip the current node and it's children are not walked
+func Walk(tree Tree, visitor Visitor, dir WalkDirection, order WalkOrder, depth int, skip bool) {
+	walk(tree, visitor, dir, order, depth, skip)
+
+	// Now continue the walk of the tree from the sibling before/after this node.
+	walkSiblings(tree, visitor, dir, order, depth)
+}
+
+func walk(tree Tree, visitor Visitor, dir WalkDirection, order WalkOrder, depth int, skip bool) {
+	if tree == nil {
+		return
+	}
+
+	i := 0
+	inc := 1
+	end := tree.NumChildren()
+
+	if dir == Reverse {
+		i = end - 1
+		inc = -1
+		end = -1
+	}
+
+	// Visit this node and it's decendants, if desired
+	if !skip {
+		if order == PreOrder {
+			visitor(tree, depth)
+		}
+
+		for ; i != end; i += inc {
+			ch := tree.Child(i)
+			walk(ch, visitor, dir, order, depth+1, false)
+		}
+
+		if order == PostOrder {
+			visitor(tree, depth)
+		}
+	}
+
+}
+
+// Walk the siblings of `tree` which has depth `depth`.
+func walkSiblings(tree Tree, visitor Visitor, dir WalkDirection, order WalkOrder, depth int) {
+	if tree.Parent() == nil {
+		return
+	}
+
+	i := 0
+	inc := 1
+	end := tree.Parent().NumChildren()
+
+	if dir == Reverse {
+		i = end - 1
+		inc = -1
+		end = -1
+	}
+
+	ignore := true
+
+	for ; i != end; i += inc {
+		ch := tree.Parent().Child(i)
+
+		if ch == tree {
+			ignore = false
+			continue
+		}
+		if ignore {
+			if ch == tree {
+				ignore = false
+			}
+			continue
+		}
+
+		walk(ch, visitor, dir, order, depth, false)
+	}
+
+	if tree.Parent() != nil && order == PostOrder {
+		visitor(tree.Parent(), depth-1)
+	}
+	walkSiblings(tree.Parent(), visitor, dir, order, depth-1)
+}

tree/walk_test.go

@@ -0,0 +1,228 @@
+package tree
+
+import (
+	"testing"
+)
+
+type Node struct {
+	name     string
+	parent   *Node
+	children []*Node
+}
+
+func (n Node) Parent() Tree {
+	if n.parent == nil {
+		return nil
+	}
+	return n.parent
+}
+
+func (n Node) Child(i int) Tree {
+	return n.children[i]
+}
+
+func (n Node) NumChildren() int {
+	return len(n.children)
+}
+
+func (n *Node) String() string {
+	return n.name
+}
+
+/*
+
+  a
+		b
+			d
+			e
+				f
+					m
+		c
+			g
+				i
+				j
+			h
+				k
+					n
+				l
+					o
+
+*/
+
+type TestData struct {
+	root  *Node
+	nodes map[string]*Node
+	depth map[string]int
+}
+
+func makeTestData() (d *TestData) {
+	d = &TestData{nodes: map[string]*Node{}, depth: map[string]int{}}
+
+	mkNode := func(nm string, parent string) *Node {
+		n := &Node{name: nm}
+		d.nodes[nm] = n
+
+		if parent != "" {
+			n.parent = d.nodes[parent]
+			d.nodes[parent].children = append(d.nodes[parent].children, n)
+		}
+
+		return n
+	}
+
+	d.root = mkNode("a", "")
+	mkNode("b", "a")
+	mkNode("c", "a")
+	mkNode("d", "b")
+	mkNode("e", "b")
+	mkNode("f", "e")
+	mkNode("m", "f")
+	mkNode("g", "c")
+	mkNode("h", "c")
+	mkNode("i", "g")
+	mkNode("j", "g")
+	mkNode("k", "h")
+	mkNode("l", "h")
+	mkNode("n", "k")
+	mkNode("o", "l")
+
+	d.depth = map[string]int{"a": 0, "b": 1, "d": 2, "e": 2, "f": 3, "m": 4, "c": 1, "g": 2, "i": 3, "j": 3, "h": 2, "k": 3, "n": 4, "l": 3, "o": 4}
+
+	return d
+}
+
+func makeSimpleVisitor(t *testing.T, data *TestData, expectedOrder []string, ndx *int) Visitor {
+	return func(tree Tree, depth int) (continu, skipChildren bool) {
+		n := tree.(*Node)
+
+		if *ndx >= len(expectedOrder) {
+			t.Fatalf("Visitor was called for node %s after the end of the expected nodes (called too many times)", n.name)
+		}
+
+		if expectedOrder[*ndx] != n.name {
+			t.Fatalf("Expected %s but got node %s", expectedOrder[*ndx], n.name)
+		}
+		if data.depth[n.name] != depth {
+			t.Fatalf("Expected depth %d but got depth %d at node %s", data.depth[n.name], depth, n.name)
+		}
+		(*ndx)++
+		return true, false
+	}
+
+}
+
+func testWalk(t *testing.T, expectedOrder []string, treeNode string, dir WalkDirection, order WalkOrder, depth int, skip bool) {
+	data := makeTestData()
+
+	ndx := 0
+
+	tree := data.nodes[treeNode]
+
+	visitor := makeSimpleVisitor(t, data, expectedOrder, &ndx)
+
+	Walk(tree, visitor, dir, order, depth, skip)
+
+	if ndx < len(expectedOrder) {
+		t.Fatalf("Not enough nodes visited. Walk stopped at %s", expectedOrder[ndx-1])
+	}
+
+}
+
+func TestWalk(t *testing.T) {
+
+	tests := []struct {
+		name          string
+		expectedOrder []string
+		tree          string
+		dir           WalkDirection
+		order         WalkOrder
+		depth         int
+		skip          bool
+	}{
+		{
+			"PreOrderForwardWalkFromRoot",
+			[]string{"a", "b", "d", "e", "f", "m", "c", "g", "i", "j", "h", "k", "n", "l", "o"},
+			"a",
+			Forward, PreOrder, 0, false,
+		},
+		{
+			"PostOrderForwardWalkFromRoot",
+			[]string{"d", "m", "f", "e", "b", "i", "j", "g", "n", "k", "o", "l", "h", "c", "a"},
+			"a",
+			Forward, PostOrder, 0, false,
+		},
+		{
+			"PreOrderReverseWalkFromRoot",
+			[]string{"a", "c", "h", "l", "o", "k", "n", "g", "j", "i", "b", "e", "f", "m", "d"},
+			"a",
+			Reverse, PreOrder, 0, false,
+		},
+		{
+			"PostOrderReverseWalkFromRoot",
+			[]string{"o", "l", "n", "k", "h", "j", "i", "g", "c", "m", "f", "e", "d", "b", "a"},
+			"a",
+			Reverse, PostOrder, 0, false,
+		},
+		{
+			"PreOrderForwardWalkFromDepth1Node",
+			[]string{"b", "d", "e", "f", "m", "c", "g", "i", "j", "h", "k", "n", "l", "o"},
+			"b",
+			Forward, PreOrder, 1, false,
+		},
+		{
+			"PreOrderForwardWalkFromLeafNode",
+			[]string{"m", "c", "g", "i", "j", "h", "k", "n", "l", "o"},
+			"m",
+			Forward, PreOrder, 4, false,
+		},
+		{
+			"PreOrderForwardWalkFromDepth1NodeSkip",
+			[]string{"c", "g", "i", "j", "h", "k", "n", "l", "o"},
+			"b",
+			Forward, PreOrder, 1, true,
+		},
+		{
+			"PostOrderForwardWalkFromDepth1Node",
+			[]string{"d", "m", "f", "e", "b", "i", "j", "g", "n", "k", "o", "l", "h", "c", "a"},
+			"b",
+			Forward, PostOrder, 1, false,
+		},
+		{
+			"PreOrderReverseWalkFromDepth1Node",
+			[]string{"b", "e", "f", "m", "d"},
+			"b",
+			Reverse, PreOrder, 1, false,
+		},
+		{
+			"PreOrderReverseWalkFromLeafNode",
+			[]string{"m", "d"},
+			"m",
+			Reverse, PreOrder, 4, false,
+		},
+		{
+			"PostOrderReverseWalkFromDepth1Node",
+			[]string{"m", "f", "e", "d", "b", "a"},
+			"b",
+			Reverse, PostOrder, 1, false,
+		},
+		{
+			"PostOrderReverseWalkFromDepth1NodeSkip",
+			[]string{"m", "f", "e", "d", "b", "a"},
+			"c",
+			Reverse, PostOrder, 1, true,
+		},
+		{
+			"PostOrderReverseWalkFromDepth2NodeSkip",
+			[]string{"c", "m", "f", "e", "d", "b", "a"},
+			"g",
+			Reverse, PostOrder, 2, true,
+		},
+	}
+
+	for _, tc := range tests {
+		t.Run(tc.name, func(t *testing.T) {
+			testWalk(t, tc.expectedOrder, tc.tree, tc.dir, tc.order, tc.depth, tc.skip)
+		})
+	}
+
+}