Fix incorrect tree snapshot encoding/decoding (#881)

The following issues exist when encoding and encoding snapshots of Tree:

- Removed nodes are missing
- Update incorrect size

This commit fixes the incorrect tree snapshot encoding/decoding logic.

---------

Co-authored-by: Youngteac Hong <[email protected]>

api/converter/converter_test.go

@@ -268,4 +268,29 @@ func TestConverter(t *testing.T) {
 		clone := converter.FromPresenceChange(pbChange)
 		assert.Equal(t, change, clone)
 	})
+
+	t.Run("properly encode and decode tree test", func(t *testing.T) {
+		doc := document.New(helper.TestDocKey(t))
+		assert.NoError(t, doc.Update(func(root *json.Object, p *presence.Presence) error {
+			root.SetNewTree("t", &json.TreeNode{
+				Type: "r",
+				Children: []json.TreeNode{
+					{Type: "p", Children: []json.TreeNode{{Type: "text", Value: "12"}}},
+					{Type: "p", Children: []json.TreeNode{{Type: "text", Value: "34"}}},
+				},
+			})
+			assert.Equal(t, "<r><p>12</p><p>34</p></r>", root.GetTree("t").ToXML())
+			root.GetTree("t").EditByPath([]int{0, 1}, []int{1, 1}, nil, 0)
+			return nil
+		}))
+		assert.Equal(t, "<r><p>14</p></r>", doc.Root().GetTree("t").ToXML())
+
+		bytes, err := converter.ObjectToBytes(doc.RootObject())
+		assert.NoError(t, err)
+		obj, err := converter.BytesToObject(bytes)
+		assert.NoError(t, err)
+
+		assert.Equal(t, obj.Get("t").(*crdt.Tree).NodeLen(), doc.Root().GetTree("t").NodeLen())
+		assert.Equal(t, obj.Get("t").(*crdt.Tree).Root().Len(), doc.Root().GetTree("t").Len())
+	})
 }

api/converter/from_pb.go

@@ -596,6 +596,8 @@ func FromTreeNodes(pbNodes []*api.TreeNode) (*crdt.TreeNode, error) {
 		}
 	}
 
+	root.Index.UpdateDescendantsSize()
+
 	// build crdt.Tree from root to construct the links between nodes.
 	return crdt.NewTree(root, nil).Root(), nil
 }

pkg/document/crdt/tree.go

@@ -641,6 +641,11 @@ func (t *Tree) Root() *TreeNode {
 	return t.IndexTree.Root().Value
 }
 
+// NodeLen returns the size of the LLRBTree.
+func (t *Tree) NodeLen() int {
+	return t.NodeMapByID.Len()
+}
+
 // ToXML returns the XML encoding of this tree.
 func (t *Tree) ToXML() string {
 	return ToXML(t.Root())

pkg/index/tree.go

@@ -325,7 +325,8 @@ func (n *Node[V]) SetChildren(children []*Node[V]) error {
 	return nil
 }
 
-// UpdateAncestorsSize updates the size of ancestors.
+// UpdateAncestorsSize updates the size of ancestors. It is used when
+// the size of the node is changed.
 func (n *Node[V]) UpdateAncestorsSize() {
 	parent := n.Parent
 	sign := 1
@@ -340,6 +341,24 @@ func (n *Node[V]) UpdateAncestorsSize() {
 	}
 }
 
+// UpdateDescendantsSize updates the size of descendants. It is used when
+// the tree is newly created and the size of the descendants is not calculated.
+func (n *Node[V]) UpdateDescendantsSize() int {
+	if n.Value.IsRemoved() {
+		n.Length = 0
+		return 0
+	}
+
+	sum := 0
+	for _, child := range n.Children(true) {
+		sum += child.UpdateDescendantsSize()
+	}
+
+	n.Length += sum
+
+	return n.PaddedLength()
+}
+
 // PaddedLength returns the length of the node with padding.
 func (n *Node[V]) PaddedLength() int {
 	length := n.Length
@@ -496,7 +515,8 @@ func (n *Node[V]) insertAtInternal(newNode *Node[V], offset int) error {
 	return nil
 }
 
-// Prepend prepends the given nodes to the children.
+// Prepend prepends the given nodes to the children. It is only used
+// for creating a new node from shapshot.
 func (n *Node[V]) Prepend(children ...*Node[V]) error {
 	if n.IsText() {
 		return ErrInvalidMethodCallForTextNode
@@ -505,10 +525,6 @@ func (n *Node[V]) Prepend(children ...*Node[V]) error {
 	n.children = append(children, n.children...)
 	for _, node := range children {
 		node.Parent = n
-
-		if !node.Value.IsRemoved() {
-			node.UpdateAncestorsSize()
-		}
 	}
 
 	return nil

pkg/llrb/llrb.go

@@ -136,6 +136,11 @@ func (t *Tree[K, V]) Floor(key K) (K, V) {
 	return zeroK, zeroV
 }
 
+// Len returns the length of the tree.
+func (t *Tree[K, V]) Len() int {
+	return t.size
+}
+
 func (t *Tree[K, V]) put(node *Node[K, V], key K, value V) *Node[K, V] {
 	if node == nil {
 		t.size++