Change tree iterate

pkg/acl/tree/change.go

@@ -25,6 +25,7 @@ type Change struct {
 	IsSnapshot      bool
 	DecryptedChange []byte // TODO: check if we need it
 	ParsedModel     interface{}
+	visited         bool
 
 	Content *aclpb.Change
 	Sign    []byte

pkg/acl/tree/tree_test.go

@@ -105,6 +105,7 @@ func TestTree_Add(t *testing.T) {
 		t.Log(time.Since(st))
 		assert.Equal(t, []string{"9999"}, tr.Heads())
 	})
+	// TODO: add my tests
 }
 
 func TestTree_Hash(t *testing.T) {
@@ -157,9 +158,10 @@ func TestTree_Iterate(t *testing.T) {
 			newChange("1", "0", "0"),
 			newChange("1.1", "0", "1"),
 			newChange("1.2", "0", "1"),
+			newChange("1.4", "0", "1.2"),
 			newChange("1.3", "0", "1"),
 			newChange("1.3.1", "0", "1.3"),
-			newChange("1.2+3", "0", "1.2", "1.3.1"),
+			newChange("1.2+3", "0", "1.4", "1.3.1"),
 			newChange("1.2+3.1", "0", "1.2+3"),
 			newChange("10", "0", "1.2+3.1", "1.1"),
 			newChange("last", "0", "10"),
@@ -174,7 +176,7 @@ func TestTree_Iterate(t *testing.T) {
 			res = append(res, c.Id)
 			return true
 		})
-		assert.Equal(t, []string{"0", "1", "1.1", "1.2", "1.3", "1.3.1", "1.2+3", "1.2+3.1", "10", "last"}, res)
+		assert.Equal(t, []string{"0", "1", "1.1", "1.2", "1.4", "1.3", "1.3.1", "1.2+3", "1.2+3.1", "10", "last"}, res)
 	})
 }
 
@@ -211,4 +213,17 @@ func BenchmarkTree_Add(b *testing.B) {
 			tr.AddFast(getChanges()...)
 		}
 	})
+	// prepare linear tree
+	tr := new(Tree)
+	tr.AddFast(newSnapshot("0", ""))
+	for j := 0; j < 10000; j++ {
+		tr.Add(newChange(fmt.Sprint(j+1), "0", fmt.Sprint(j)))
+	}
+	b.Run("add linear", func(b *testing.B) {
+		for i := 0; i < b.N; i++ {
+			tr.Iterate("0", func(c *Change) (isContinue bool) {
+				return true
+			})
+		}
+	})
 }

pkg/acl/tree/treeiterator.go

@@ -1,10 +1,15 @@
 package tree
 
-import "sync"
+import (
+	"sync"
+)
 
 var itPool = &sync.Pool{
 	New: func() interface{} {
-		return &iterator{}
+		return &iterator{
+			stack:  make([]*Change, 0, 100),
+			resBuf: make([]*Change, 0, 100),
+		}
 	},
 }
 
@@ -17,11 +22,9 @@ func freeIterator(i *iterator) {
 }
 
 type iterator struct {
-	compBuf    []*Change
-	queue      []*Change
-	doneMap    map[*Change]struct{}
-	breakpoint *Change
-	f          func(c *Change) bool
+	resBuf []*Change
+	stack  []*Change
+	f      func(c *Change) bool
 }
 
 func (i *iterator) iterateSkip(start *Change, skipBefore *Change, f func(c *Change) (isContinue bool)) {
@@ -35,124 +38,47 @@ func (i *iterator) iterateSkip(start *Change, skipBefore *Change, f func(c *Chan
 	})
 }
 
+func (i *iterator) topSort(start *Change) {
+	stack := i.stack
+	stack = append(stack, start)
+
+	for len(stack) > 0 {
+		ch := stack[len(stack)-1]
+		stack = stack[:len(stack)-1]
+
+		if ch.visited {
+			i.resBuf = append(i.resBuf, ch)
+			continue
+		}
+
+		ch.visited = true
+		stack = append(stack, ch)
+
+		for j := 0; j < len(ch.Next); j++ {
+			if ch.Next[j].visited {
+				continue
+			}
+			stack = append(stack, ch.Next[j])
+		}
+	}
+	for _, ch := range i.resBuf {
+		ch.visited = false
+	}
+}
+
 func (i *iterator) iterate(start *Change, f func(c *Change) (isContinue bool)) {
 	if start == nil {
 		return
 	}
 	// reset
-	i.queue = i.queue[:0]
-	i.compBuf = i.compBuf[:0]
-	i.doneMap = make(map[*Change]struct{})
-	i.queue = append(i.queue, start)
-	i.breakpoint = nil
+	i.resBuf = i.resBuf[:0]
+	i.stack = i.stack[:0]
 	i.f = f
 
-	for len(i.queue) > 0 {
-		c := i.queue[0]
-		i.queue = i.queue[1:]
-		nl := len(c.Next)
-		if nl == 1 {
-			if !i.iterateLin(c) {
-				return
-			}
-			if i.breakpoint != nil {
-				i.toQueue(i.breakpoint)
-				i.breakpoint = nil
-			}
-		} else {
-			_, done := i.doneMap[c]
-			if !done {
-				if !f(c) {
-					return
-				}
-				i.doneMap[c] = struct{}{}
-			}
-			if nl != 0 {
-				for _, next := range c.Next {
-					i.toQueue(next)
-				}
-			}
-		}
-	}
-}
-
-func (i *iterator) iterateLin(c *Change) bool {
-	for len(c.Next) == 1 {
-		_, done := i.doneMap[c]
-		if !done {
-			if !i.f(c) {
-				return false
-			}
-			i.doneMap[c] = struct{}{}
-		}
-
-		c = c.Next[0]
-		if len(c.PreviousIds) > 1 {
-			break
-		}
-	}
-	if len(c.Next) == 0 && len(c.PreviousIds) <= 1 {
-		if !i.f(c) {
-			return false
-		}
-		i.doneMap[c] = struct{}{}
-	} else {
-		i.breakpoint = c
-	}
-
-	return true
-}
-
-func (i *iterator) comp(c1, c2 *Change) uint8 {
-	if c1.Id == c2.Id {
-		return 0
-	}
-	i.compBuf = i.compBuf[:0]
-	i.compBuf = append(i.compBuf, c1.Next...)
-	var uniq = make(map[*Change]struct{})
-	var appendUniqueToBuf = func(next []*Change) {
-		for _, n := range next {
-			if _, ok := uniq[n]; !ok {
-				i.compBuf = append(i.compBuf, n)
-				uniq[n] = struct{}{}
-			}
-		}
-	}
-	var used int
-	for len(i.compBuf)-used > 0 {
-		l := len(i.compBuf) - used
-		for _, n := range i.compBuf[used:] {
-			delete(uniq, n)
-			if n.Id == c2.Id {
-				return 1
-			} else {
-				appendUniqueToBuf(n.Next)
-			}
-		}
-		used += l
-	}
-	return 2
-}
-
-func (i *iterator) toQueue(c *Change) {
-	var pos = -1
-For:
-	for idx, qc := range i.queue {
-		switch i.comp(c, qc) {
-		// exists
-		case 0:
+	i.topSort(start)
+	for idx := len(i.resBuf) - 1; idx >= 0; idx-- {
+		if !f(i.resBuf[idx]) {
 			return
-		//
-		case 1:
-			pos = idx
-			break For
 		}
 	}
-	if pos == -1 {
-		i.queue = append(i.queue, c)
-	} else if pos == 0 {
-		i.queue = append([]*Change{c}, i.queue...)
-	} else {
-		i.queue = append(i.queue[:pos], append([]*Change{c}, i.queue[pos:]...)...)
-	}
 }