From a17b102756ad4b469d63e972a4a9cd57f4af2fe5 Mon Sep 17 00:00:00 2001
From: mcrakhman <rakhmanov.m@gmail.com>
Date: Sat, 20 Aug 2022 23:55:18 +0200
Subject: [PATCH] Few tree fixes

---
 pkg/acl/tree/doctree.go |  5 ++++-
 pkg/acl/tree/tree.go    | 29 ++++++++++++++++++++++++++---
 2 files changed, 30 insertions(+), 4 deletions(-)

diff --git a/pkg/acl/tree/doctree.go b/pkg/acl/tree/doctree.go
index f7ae2e4c..d0d373d4 100644
--- a/pkg/acl/tree/doctree.go
+++ b/pkg/acl/tree/doctree.go
@@ -238,7 +238,10 @@ func (d *docTree) AddContent(ctx context.Context, aclList list.ACLList, content
 		// clearing tree, because we already fixed everything in the last snapshot
 		d.tree = &Tree{}
 	}
-	d.tree.AddFast(ch) // TODO: Add head
+	err = d.tree.AddMergedHead(ch)
+	if err != nil {
+		panic("error in adding head")
+	}
 	rawCh := &aclpb.RawChange{
 		Payload:   fullMarshalledChange,
 		Signature: ch.Signature(),
diff --git a/pkg/acl/tree/tree.go b/pkg/acl/tree/tree.go
index 1836719c..687561fa 100644
--- a/pkg/acl/tree/tree.go
+++ b/pkg/acl/tree/tree.go
@@ -56,6 +56,31 @@ func (t *Tree) AddFast(changes ...*Change) {
 	t.updateHeads()
 }
 
+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")
+	} 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 {
+		return fmt.Errorf("change is not attached")
+	}
+
+	// check that previous heads have new change as next
+	for _, prevHead := range t.headIds {
+		head := t.attached[prevHead]
+		if len(head.Next) != 1 || head.Next[0].Id != c.Id {
+			return fmt.Errorf("this is not a new head")
+		}
+	}
+	t.headIds = []string{c.Id}
+	return nil
+}
+
 func (t *Tree) Add(changes ...*Change) (mode Mode) {
 	var beforeHeadIds = t.headIds
 	var attached bool
@@ -263,7 +288,7 @@ func (t *Tree) dfs(startChange string) (uniqMap map[string]*Change) {
 }
 
 func (t *Tree) updateHeads() {
-	var newHeadIds, newMetaHeadIds []string
+	var newHeadIds []string
 	t.iterate(t.root, func(c *Change) (isContinue bool) {
 		if len(c.Next) == 0 {
 			newHeadIds = append(newHeadIds, c.Id)
@@ -271,9 +296,7 @@ func (t *Tree) updateHeads() {
 		return true
 	})
 	t.headIds = newHeadIds
-	t.metaHeadIds = newMetaHeadIds
 	sort.Strings(t.headIds)
-	sort.Strings(t.metaHeadIds)
 }
 
 func (t *Tree) iterate(start *Change, f func(c *Change) (isContinue bool)) {