Fix incorrect tree snapshot encoding/decoding (#836)

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: Yourim Cha <[email protected]>
Co-authored-by: raararaara <[email protected]>

src/api/converter.ts

@@ -105,7 +105,7 @@ import {
   CRDTTreeNode,
   CRDTTreeNodeID,
 } from '@yorkie-js-sdk/src/document/crdt/tree';
-import { traverse } from '../util/index_tree';
+import { traverseAll } from '../util/index_tree';
 import { TreeStyleOperation } from '../document/operation/tree_style_operation';
 import { RHT } from '../document/crdt/rht';
 
@@ -601,7 +601,7 @@ function toTreeNodes(node: CRDTTreeNode): Array<PbTreeNode> {
   }
 
   const pbTreeNodes: Array<PbTreeNode> = [];
-  traverse(node, (n, depth) => {
+  traverseAll(node, (n, depth) => {
     const pbTreeNode = new PbTreeNode({
       id: toTreeNodeID(n.id),
       type: n.type,
@@ -1047,6 +1047,8 @@ function fromTreeNodes(
     parent!.prepend(nodes[i]);
   }
 
+  root.updateDescendantsSize();
+
   // build CRDTTree from the root to construct the links between nodes.
   return CRDTTree.create(root, InitialTimeTicket).getRoot();
 }

src/document/crdt/tree.ts

@@ -774,7 +774,7 @@ export class CRDTTree extends CRDTElement implements GCParent {
     this.indexTree = new IndexTree<CRDTTreeNode>(root);
     this.nodeMapByID = new LLRBTree(CRDTTreeNodeID.createComparator());
 
-    this.indexTree.traverse((node) => {
+    this.indexTree.traverseAll((node) => {
       this.nodeMapByID.put(node.id, node);
     });
   }
@@ -1274,6 +1274,13 @@ export class CRDTTree extends CRDTElement implements GCParent {
     return this.indexTree.size;
   }
 
+  /**
+   * `getNodeSize` returns the size of the LLRBTree.
+   */
+  public getNodeSize(): number {
+    return this.nodeMapByID.size();
+  }
+
   /**
    * `getIndexTree` returns the index tree.
    */

src/document/json/tree.ts

@@ -253,6 +253,17 @@ export class Tree {
     return this.tree.getSize();
   }
 
+  /**
+   * `getNodeSize` returns the node size of this tree.
+   */
+  public getNodeSize(): number {
+    if (!this.context || !this.tree) {
+      throw new Error('it is not initialized yet');
+    }
+
+    return this.tree.getNodeSize();
+  }
+
   /**
    * `getIndexTree` returns the index tree of this tree.
    */

src/util/index_tree.ts

@@ -134,7 +134,8 @@ export abstract class IndexTreeNode<T extends IndexTreeNode<T>> {
   }
 
   /**
-   * `updateAncestorsSize` updates the size of the ancestors.
+   * `updateAncestorsSize` updates the size of the ancestors. It is used when
+   * the size of the node is changed.
    */
   updateAncestorsSize(): void {
     let parent: T | undefined = this.parent;
@@ -146,6 +147,26 @@ export abstract class IndexTreeNode<T extends IndexTreeNode<T>> {
     }
   }
 
+  /**
+   * `updateDescendantsSize` updates the size of the descendants. It is used when
+   * the tree is newly created and the size of the descendants is not calculated.
+   */
+  updateDescendantsSize(): number {
+    if (this.isRemoved) {
+      this.size = 0;
+      return 0;
+    }
+
+    let sum = 0;
+    for (const child of this._children) {
+      sum += child.updateDescendantsSize();
+    }
+
+    this.size += sum;
+
+    return this.paddedSize;
+  }
+
   /**
    * `isText` returns true if the node is a text node.
    */
@@ -276,7 +297,8 @@ export abstract class IndexTreeNode<T extends IndexTreeNode<T>> {
   }
 
   /**
-   * `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 snapshot.
    */
   prepend(...newNode: Array<T>): void {
     if (this.isText) {
@@ -286,10 +308,6 @@ export abstract class IndexTreeNode<T extends IndexTreeNode<T>> {
     this._children.unshift(...newNode);
     for (const node of newNode) {
       node.parent = this as any;
-
-      if (!node.isRemoved) {
-        node.updateAncestorsSize();
-      }
     }
   }
 

test/unit/api/converter_test.ts

@@ -17,7 +17,7 @@
 import { describe, it, assert } from 'vitest';
 import { Document } from '@yorkie-js-sdk/src/document/document';
 import { converter } from '@yorkie-js-sdk/src/api/converter';
-import { Counter, Text } from '@yorkie-js-sdk/src/yorkie';
+import { Counter, Text, Tree } from '@yorkie-js-sdk/src/yorkie';
 import { CounterType } from '@yorkie-js-sdk/src/document/crdt/counter';
 
 describe('Converter', function () {
@@ -87,4 +87,37 @@ describe('Converter', function () {
     assert.equal(bytes.length, 12);
     assert.equal(converter.toHexString(bytes), hexString);
   });
+
+  it('should encode and decode tree properly', function () {
+    const doc = new Document<{
+      tree: Tree;
+    }>('test-doc');
+
+    doc.update((root) => {
+      root.tree = new Tree({
+        type: 'r',
+        children: [
+          { type: 'p', children: [{ type: 'text', value: '12' }] },
+          { type: 'p', children: [{ type: 'text', value: '34' }] },
+        ],
+      });
+
+      root.tree.editByPath([0, 1], [1, 1]);
+    });
+    assert.equal(doc.getRoot().tree.toXML(), /*html*/ `<r><p>14</p></r>`);
+    assert.equal(doc.getRoot().tree.getSize(), 4);
+
+    const bytes = converter.objectToBytes(doc.getRootObject());
+    const obj = converter.bytesToObject(bytes);
+
+    assert.equal(
+      doc.getRoot().tree.getNodeSize(),
+      (obj.get('tree') as unknown as Tree).getNodeSize(),
+    );
+
+    assert.equal(
+      doc.getRoot().tree.getSize(),
+      (obj.get('tree') as unknown as Tree).getSize(),
+    );
+  });
 });