Fix incorrect tree snapshot encoding/decoding

Sources/API/Converter.swift

@@ -911,7 +911,7 @@ extension Converter {
             return pbTreeNodes
         }
 
-        traverse(node: node) { node, depth in
+        traverseAll(node: node) { node, depth in
             var pbTreeNode = PbTreeNode()
             pbTreeNode.id = toTreeNodeID(node.id)
             pbTreeNode.type = node.type
@@ -993,6 +993,8 @@ extension Converter {
             try parent?.prepend(contentsOf: [nodes[index]])
         }
 
+        root.updateDescendantsSize()
+
         // build CRDTTree from the root to construct the links between nodes.
         return CRDTTree(root: root, createdAt: TimeTicket.initial).root
     }

Sources/Document/CRDT/CRDTTree.swift

@@ -566,7 +566,7 @@ class CRDTTree: CRDTElement {
         self.indexTree = IndexTree(root: root)
         self.nodeMapByID = LLRBTree()
 
-        self.indexTree.traverse { node, _ in
+        self.indexTree.traverseAll { node, _ in
             self.nodeMapByID.put(node.id, node)
         }
     }
@@ -938,6 +938,13 @@ class CRDTTree: CRDTElement {
         self.indexTree.size
     }
 
+    /**
+     * `nodeSize` returns the size of the LLRBTree.
+     */
+    var nodeSize: Int {
+        self.nodeMapByID.size
+    }
+
     /**
      * toXML returns the XML encoding of this tree.
      */

Sources/Document/Json/JSONTree.swift

@@ -304,6 +304,17 @@ public class JSONTree {
         return tree.size
     }
 
+    /**
+     * `getNodeSize` returns the node size of this tree.
+     */
+    public func getNodeSize() throws -> Int {
+        guard self.context != nil, let tree else {
+            throw YorkieError.unexpected(message: "it is not initialized yet")
+        }
+
+        return tree.nodeSize
+    }
+
     /**
      * `getIndexTree` returns the index tree of this tree.
      */

Sources/Util/IndexTree.swift

@@ -136,7 +136,8 @@ protocol IndexTreeNode: AnyObject {
 
 extension IndexTreeNode {
     /**
-     * `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.
      */
     func updateAncestorsSize() {
         var parent = self.parent
@@ -148,6 +149,27 @@ extension IndexTreeNode {
         }
     }
 
+    /**
+     * `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.
+     */
+    @discardableResult
+    func updateDescendantsSize() -> Int {
+        if self.isRemoved {
+            self.size = 0
+            return 0
+        }
+
+        var sum = 0
+        for child in self.children {
+            sum += child.updateDescendantsSize()
+        }
+
+        self.size += sum
+
+        return self.paddedSize
+    }
+
     /**
      * `isText` returns true if the node is a text node.
      */
@@ -245,7 +267,8 @@ extension IndexTreeNode {
     }
 
     /**
-     * `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.
      */
     func prepend(contentsOf newNode: [Self]) throws {
         guard self.isText == false else {
@@ -256,10 +279,6 @@ extension IndexTreeNode {
 
         for node in newNode {
             node.parent = self
-
-            if node.isRemoved == false {
-                node.updateAncestorsSize()
-            }
         }
     }
 

Tests/Unit/API/V1/ConverterTests.swift

@@ -365,4 +365,36 @@ class ConverterTests: XCTestCase {
 
         XCTAssert(!(samplePresence == converted))
     }
+
+    func test_should_encode_and_decode_tree_properly() async throws {
+        let doc = Document(key: "test-doc")
+
+        try await doc.update { root, _ in
+            root.tree = JSONTree(initialRoot: JSONTreeElementNode(type: "r", children: [
+                JSONTreeElementNode(type: "p", children: [JSONTreeTextNode(value: "12")]),
+                JSONTreeElementNode(type: "p", children: [JSONTreeTextNode(value: "34")])
+            ]))
+
+            try (root.tree as? JSONTree)?.editByPath([0, 1], [1, 1])
+        }
+
+        let xml = await(doc.getRoot().tree as? JSONTree)?.toXML()
+        let size = try await(doc.getRoot().tree as? JSONTree)?.getSize()
+
+        XCTAssertEqual(xml, "<r><p>14</p></r>")
+        XCTAssertEqual(size, 4)
+
+        let bytes = try await Converter.objectToBytes(obj: doc.getRootObject())
+        let obj = try Converter.bytesToObject(bytes: bytes)
+
+        let nodeSize = try await(doc.getRoot().tree as? JSONTree)?.getNodeSize()
+        let nodeSize2 = (obj.get(key: "tree") as? CRDTTree)?.nodeSize
+
+        XCTAssertEqual(nodeSize, nodeSize2)
+
+        let size1 = try await(doc.getRoot().tree as? JSONTree)?.getSize()
+        let size2 = (obj.get(key: "tree") as? CRDTTree)?.size
+
+        XCTAssertEqual(size1, size2)
+    }
 }