Manage seen heads

commonspace/object/tree/objecttree/changediffer.go

@@ -2,6 +2,7 @@ package objecttree
 
 import (
 	"go.uber.org/zap"
+	"golang.org/x/exp/slices"
 
 	"github.com/anyproto/any-sync/util/slice"
 )
@@ -41,11 +42,15 @@ func NewChangeDiffer(tree ReadableObjectTree, hasChanges hasChangesFunc) (*Chang
 	return diff, nil
 }
 
-func (d *ChangeDiffer) RemoveBefore(ids []string) (removed []string, notFound []string) {
+func (d *ChangeDiffer) RemoveBefore(ids []string, seenHeads []string) (removed []string, notFound []string, heads []string) {
-	var attached []*Change
+	var (
+		attached    []*Change
+		attachedIds []string
+	)
 	for _, id := range ids {
 		if ch, ok := d.attached[id]; ok {
 			attached = append(attached, ch)
+			attachedIds = append(attachedIds, id)
 			continue
 		}
 		// check if we have it at the bottom
@@ -53,7 +58,11 @@ func (d *ChangeDiffer) RemoveBefore(ids []string) (removed []string, notFound []
 			notFound = append(notFound, id)
 		}
 	}
+	heads = make([]string, 0, len(seenHeads))
+	heads = append(heads, seenHeads...)
 	d.dfsPrev(attached, func(ch *Change) (isContinue bool) {
+		heads = append(heads, ch.Id)
+		heads = append(heads, ch.PreviousIds...)
 		removed = append(removed, ch.Id)
 		return true
 	}, func(changes []*Change) {
@@ -66,6 +75,12 @@ func (d *ChangeDiffer) RemoveBefore(ids []string) (removed []string, notFound []
 			})
 		}
 	})
+	slices.Sort(heads)
+	heads = slice.RemoveRepeatedSorted(heads)
+	heads = append(heads, attachedIds...)
+	heads = append(heads, seenHeads...)
+	heads = slice.RemoveUniqueElementsSorted(heads)
+	heads = slice.DiscardDuplicatesSorted(heads)
 	return
 }
 
@@ -162,15 +177,21 @@ func NewDiffManager(initHeads, curHeads []string, treeBuilder treeBuilderFunc, o
 }
 
 func (d *DiffManager) Init() {
-	removed, _ := d.differ.RemoveBefore(d.seenHeads)
+	removed, _, seenHeads := d.differ.RemoveBefore(d.seenHeads, []string{})
+	d.seenHeads = seenHeads
 	d.onRemove(removed)
 }
 
+func (d *DiffManager) SeenHeads() []string {
+	return d.seenHeads
+}
+
 func (d *DiffManager) Remove(ids []string) {
-	removed, notFound := d.differ.RemoveBefore(ids)
+	removed, notFound, seenHeads := d.differ.RemoveBefore(ids, d.seenHeads)
 	for _, id := range notFound {
 		d.notFound[id] = struct{}{}
 	}
+	d.seenHeads = seenHeads
 	d.onRemove(removed)
 }
 
@@ -203,6 +224,7 @@ func (d *DiffManager) Update(objTree ObjectTree) {
 	d.differ.Add(toAdd...)
 	d.heads = make([]string, 0, len(d.heads))
 	d.heads = append(d.heads, objTree.Heads()...)
-	removed, _ := d.differ.RemoveBefore(toRemove)
+	removed, _, seenHeads := d.differ.RemoveBefore(toRemove, d.seenHeads)
+	d.seenHeads = seenHeads
 	d.onRemove(removed)
 }

commonspace/object/tree/objecttree/changediffer_test.go

@@ -23,9 +23,10 @@ func TestChangeDiffer_Add(t *testing.T) {
 			return false
 		})
 		differ.Add(changes...)
-		res, notFound := differ.RemoveBefore([]string{"7"})
+		res, notFound, seenHeads := differ.RemoveBefore([]string{"7"}, []string{"0"})
 		require.Len(t, notFound, 0)
 		require.Equal(t, len(changes), len(res))
+		require.Equal(t, []string{"7"}, seenHeads)
 	})
 	t.Run("remove in two parts", func(t *testing.T) {
 		changes := []*Change{
@@ -42,12 +43,14 @@ func TestChangeDiffer_Add(t *testing.T) {
 			return false
 		})
 		differ.Add(changes...)
-		res, notFound := differ.RemoveBefore([]string{"4"})
+		res, notFound, seenHeads := differ.RemoveBefore([]string{"4"}, []string{"0"})
 		require.Len(t, notFound, 0)
 		require.Equal(t, 2, len(res))
-		res, notFound = differ.RemoveBefore([]string{"7"})
+		require.Equal(t, []string{"4"}, seenHeads)
+		res, notFound, seenHeads = differ.RemoveBefore([]string{"7"}, seenHeads)
 		require.Len(t, notFound, 0)
 		require.Equal(t, 6, len(res))
+		require.Equal(t, []string{"7"}, seenHeads)
 	})
 	t.Run("add and remove", func(t *testing.T) {
 		changes := []*Change{
@@ -60,9 +63,10 @@ func TestChangeDiffer_Add(t *testing.T) {
 			return false
 		})
 		differ.Add(changes...)
-		res, notFound := differ.RemoveBefore([]string{"3"})
+		res, notFound, seenHeads := differ.RemoveBefore([]string{"3"}, []string{"0"})
 		require.Len(t, notFound, 0)
 		require.Equal(t, len(changes), len(res))
+		require.Equal(t, []string{"3"}, seenHeads)
 		changes = []*Change{
 			newChange("4", "0", "0"),
 			newChange("5", "0", "4"),
@@ -70,15 +74,17 @@ func TestChangeDiffer_Add(t *testing.T) {
 			newChange("7", "0", "3", "6"),
 		}
 		differ.Add(changes...)
-		res, notFound = differ.RemoveBefore([]string{"7"})
+		res, notFound, seenHeads = differ.RemoveBefore([]string{"7"}, seenHeads)
 		require.Len(t, notFound, 0)
 		require.Equal(t, len(changes), len(res))
+		require.Equal(t, []string{"7"}, seenHeads)
 	})
 	t.Run("remove not found", func(t *testing.T) {
 		differ, _ := NewChangeDiffer(nil, func(ids ...string) bool {
 			return false
 		})
-		_, notFound := differ.RemoveBefore([]string{"3", "4", "5"})
+		_, notFound, seenHeads := differ.RemoveBefore([]string{"3", "4", "5"}, []string{"0"})
+		require.Equal(t, seenHeads, []string{"0"})
 		require.Len(t, notFound, 3)
 	})
 	t.Run("exists in storage", func(t *testing.T) {
@@ -91,8 +97,9 @@ func TestChangeDiffer_Add(t *testing.T) {
 			}
 			return true
 		})
-		res, notFound := differ.RemoveBefore([]string{"3", "4", "5"})
+		res, notFound, seenHeads := differ.RemoveBefore([]string{"3", "4", "5"}, []string{"7"})
 		require.Len(t, res, 0)
 		require.Len(t, notFound, 0)
+		require.Equal(t, seenHeads, []string{"7"})
 	})
 }

util/slice/slice.go

@@ -159,6 +159,50 @@ func DiscardDuplicatesSorted[T comparable](sorted []T) []T {
 	return sorted[:cnt]
 }
 
+func RemoveRepeatedSorted[T constraints.Ordered](nums []T) []T {
+	if len(nums) == 0 {
+		return nums
+	}
+	writeIndex := 0
+	i := 0
+	for i < len(nums) {
+		count := 1
+		for i+1 < len(nums) && nums[i] == nums[i+1] {
+			count++
+			i++
+		}
+		if count == 1 {
+			nums[writeIndex] = nums[i]
+			writeIndex++
+		}
+		i++
+	}
+	return nums[:writeIndex]
+}
+
+func RemoveUniqueElementsSorted[T constraints.Ordered](nums []T) []T {
+	if len(nums) == 0 {
+		return nums
+	}
+	writeIndex := 0
+	i := 0
+	for i < len(nums) {
+		count := 1
+		for i+1 < len(nums) && nums[i] == nums[i+1] {
+			count++
+			i++
+		}
+		if count > 1 {
+			for j := 0; j < count; j++ {
+				nums[writeIndex] = nums[i]
+				writeIndex++
+			}
+		}
+		i++
+	}
+	return nums[:writeIndex]
+}
+
 func DiscardDuplicatesSortedFunc[T any](sorted []T, equal func(T, T) bool) []T {
 	if len(sorted) <= 1 {
 		return sorted