Fix algorithms and add tests

pkg/acl/tree/change.go

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

pkg/acl/tree/tree.go

@@ -24,10 +24,11 @@ type Tree struct {
 	// missed id -> list of dependency ids
 	waitList       map[string][]string
 	invalidChanges map[string]struct{}
+	possibleRoots  []*Change
 
 	// bufs
-	iterCompBuf []*Change
-	iterQueue   []*Change
+	visitedBuf []*Change
+	stackBuf   []*Change
 
 	duplicateEvents int
 }
@@ -59,14 +60,13 @@ func (t *Tree) AddFast(changes ...*Change) {
 func (t *Tree) AddMergedHead(c *Change) error {
 	// check that it was not inserted previously
 	if _, ok := t.attached[c.Id]; ok {
-		return fmt.Errorf("change already exists")
+		return fmt.Errorf("change already exists") // TODO: named error
 	} else if _, ok := t.unAttached[c.Id]; ok {
 		return fmt.Errorf("change already exists")
 	}
-	t.add(c)
 
 	// check that it was attached after adding
-	if _, ok := t.attached[c.Id]; !ok {
+	if !t.add(c) {
 		return fmt.Errorf("change is not attached")
 	}
 
@@ -172,6 +172,7 @@ func (t *Tree) add(c *Change) (attached bool) {
 		t.unAttached = make(map[string]*Change)
 		t.waitList = make(map[string][]string)
 		t.invalidChanges = make(map[string]struct{})
+		t.possibleRoots = make([]*Change, 0, 10)
 		return true
 	}
 	if len(c.PreviousIds) > 1 {
@@ -216,6 +217,9 @@ func (t *Tree) attach(c *Change, newEl bool) {
 	if !newEl {
 		delete(t.unAttached, c.Id)
 	}
+	if c.IsSnapshot {
+		t.possibleRoots = append(t.possibleRoots, c)
+	}
 
 	// add next to all prev changes
 	for _, id := range c.PreviousIds {
@@ -266,25 +270,31 @@ func (t *Tree) after(id1, id2 string) (found bool) {
 	return
 }
 
-func (t *Tree) dfs(startChange string) (uniqMap map[string]*Change) {
-	stack := make([]*Change, 0, 10)
-	stack = append(stack, t.attached[startChange])
-	uniqMap = map[string]*Change{}
+func (t *Tree) dfsPrev(stack []*Change, visit func(ch *Change), afterVisit func([]*Change)) {
+	t.visitedBuf = t.visitedBuf[:0]
 
 	for len(stack) > 0 {
 		ch := stack[len(stack)-1]
 		stack = stack[:len(stack)-1]
-		if _, exists := uniqMap[ch.Id]; exists {
+		if ch.visited {
 			continue
 		}
 
-		uniqMap[ch.Id] = ch
+		ch.visited = true
+		t.visitedBuf = append(t.visitedBuf, ch)
 
-		for _, prev := range ch.PreviousIds {
-			stack = append(stack, t.attached[prev])
+		for _, prevId := range ch.PreviousIds {
+			prevCh := t.attached[prevId]
+			if !prevCh.visited {
+				stack = append(stack, prevCh)
+			}
 		}
+		visit(ch)
+	}
+	afterVisit(t.visitedBuf)
+	for _, ch := range t.visitedBuf {
+		ch.visited = false
 	}
-	return uniqMap
 }
 
 func (t *Tree) updateHeads() {

pkg/acl/tree/tree_test.go

@@ -186,22 +186,24 @@ func TestTree_CheckRootReduce(t *testing.T) {
 		tr.Add(
 			newSnapshot("0", ""),
 			newSnapshot("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.4", "1.3.1"),
-			newChange("1.2+3.1", "0", "1.2+3"),
-			newSnapshot("10", "0", "1.2+3.1", "1.1"),
+			newSnapshot("1.2+3", "1", "1.2", "1.3.1"),
+			newChange("1.2+3.1", "1", "1.2+3"),
+			newChange("1.2+3.2", "1", "1.2+3"),
+			newSnapshot("10", "1.2+3", "1.2+3.1", "1.2+3.2"),
 			newChange("last", "10", "10"),
 		)
 		t.Run("check root", func(t *testing.T) {
 			total := tr.checkRoot(tr.attached["10"])
 			assert.Equal(t, 1, total)
 
+			total = tr.checkRoot(tr.attached["1.2+3"])
+			assert.Equal(t, 4, total)
+
 			total = tr.checkRoot(tr.attached["1"])
-			assert.Equal(t, 9, total)
+			assert.Equal(t, 8, total)
 		})
 		t.Run("reduce", func(t *testing.T) {
 			tr.reduceTree()
@@ -304,6 +306,9 @@ func BenchmarkTree_Add(b *testing.B) {
 			tr.AddFast(getChanges()...)
 		}
 	})
+}
+
+func BenchmarkTree_IterateLinear(b *testing.B) {
 	// prepare linear tree
 	tr := new(Tree)
 	tr.AddFast(newSnapshot("0", ""))

pkg/acl/tree/treeiterator.go

@@ -46,19 +46,28 @@ func (i *iterator) topSort(start *Change) {
 		ch := stack[len(stack)-1]
 		stack = stack[:len(stack)-1]
 
-		if ch.visited {
+		// this looks a bit clumsy, but the idea is that we will go through the change again as soon as we finished
+		// going through its branches
+		if ch.branchesFinished {
 			i.resBuf = append(i.resBuf, ch)
+			ch.branchesFinished = false
+			continue
+		}
+
+		// in theory, it may be the case that we add the change two times
+		// but probably due to the way how we build the tree, we won't need it
+		if ch.visited {
 			continue
 		}
 
-		ch.visited = true
 		stack = append(stack, ch)
+		ch.visited = true
+		ch.branchesFinished = true
 
 		for j := 0; j < len(ch.Next); j++ {
-			if ch.Next[j].visited {
-				continue
+			if !ch.Next[j].visited {
+				stack = append(stack, ch.Next[j])
 			}
-			stack = append(stack, ch.Next[j])
 		}
 	}
 	for _, ch := range i.resBuf {