Fix incorrect tree snapshot encoding/decoding (#198)

Author: 7hong13

yorkie/src/main/kotlin/dev/yorkie/api/ElementConverter.kt

@@ -48,7 +48,7 @@ import dev.yorkie.document.crdt.TextValue
 import dev.yorkie.document.time.ActorID
 import dev.yorkie.document.time.TimeTicket.Companion.InitialTimeTicket
 import dev.yorkie.util.IndexTreeNode
-import dev.yorkie.util.traverse
+import dev.yorkie.util.traverseAll
 
 internal typealias PBJsonElement = dev.yorkie.api.v1.JSONElement
 internal typealias PBJsonElementSimple = dev.yorkie.api.v1.JSONElementSimple
@@ -224,6 +224,7 @@ internal fun List<PBTreeNode>.toCrdtTreeRootNode(): CrdtTreeNode? {
         }.takeIf { it > -1 }?.plus(i + 1) ?: continue
         nodes[parentIndex].prepend(nodes[i])
     }
+    root.updateDescendantSize()
     return CrdtTree(root, InitialTimeTicket).root
 }
 
@@ -315,7 +316,7 @@ internal fun RgaTreeList.toPBRgaNodes(): List<PBRgaNode> {
 internal fun CrdtTreeNode.toPBTreeNodes(): List<PBTreeNode> {
     val treeNode = this
     return buildList {
-        traverse(treeNode) { node, nodeDepth ->
+        traverseAll(treeNode) { node, nodeDepth ->
             val pbTreeNode = treeNode {
                 id = node.id.toPBTreeNodeID()
                 type = node.type

yorkie/src/main/kotlin/dev/yorkie/document/crdt/CrdtTree.kt

@@ -1,5 +1,6 @@
 package dev.yorkie.document.crdt
 
+import androidx.annotation.VisibleForTesting
 import dev.yorkie.document.CrdtTreeNodeIDStruct
 import dev.yorkie.document.CrdtTreePosStruct
 import dev.yorkie.document.JsonSerializable
@@ -15,6 +16,7 @@ import dev.yorkie.util.IndexTreeNodeList
 import dev.yorkie.util.TokenType
 import dev.yorkie.util.TreePos
 import dev.yorkie.util.TreeToken
+import dev.yorkie.util.traverseAll
 import java.util.TreeMap
 
 public typealias TreePosRange = Pair<CrdtTreePos, CrdtTreePos>
@@ -46,14 +48,18 @@ internal data class CrdtTree(
         get() = indexTree.root.toTreeNode()
 
     init {
-        indexTree.traverse { node, _ ->
+        indexTree.traverseAll { node, _ ->
             nodeMapByID[node.id] = node
         }
     }
 
     val size: Int
         get() = indexTree.size
 
+    @VisibleForTesting
+    val nodeSize: Int
+        get() = nodeMapByID.size
+
     /**
      * Applies the given [attributes] of the given [range].
      */
@@ -444,17 +450,6 @@ internal data class CrdtTree(
         return TreeOperationResult(changes, gcPairs)
     }
 
-    private fun traverseAll(
-        node: CrdtTreeNode,
-        depth: Int = 0,
-        action: ((CrdtTreeNode, Int) -> Unit),
-    ) {
-        node.allChildren.forEach { child ->
-            traverseAll(child, depth + 1, action)
-        }
-        action.invoke(node, depth)
-    }
-
     private fun traverseInPosRange(
         fromParent: CrdtTreeNode,
         fromLeft: CrdtTreeNode,

yorkie/src/main/kotlin/dev/yorkie/util/IndexTree.kt

@@ -138,6 +138,13 @@ internal class IndexTree<T : IndexTreeNode<T>>(val root: T) {
         traverse(root, 0, action)
     }
 
+    /**
+     * Traverses the whole tree (include tombstones) with postorder traversal.
+     */
+    fun traverseAll(action: ((T, Int) -> Unit)) {
+        traverseAll(root, 0, action)
+    }
+
     /**
      * Finds the position of the given [index] in the tree.
      */
@@ -370,7 +377,7 @@ internal data class TreePos<T : IndexTreeNode<T>>(
  * document of text-based editors.
  */
 @Suppress("UNCHECKED_CAST")
-internal abstract class IndexTreeNode<T : IndexTreeNode<T>>() {
+internal abstract class IndexTreeNode<T : IndexTreeNode<T>> {
     abstract val type: String
 
     protected abstract val childNodes: IndexTreeNodeList<T>
@@ -422,7 +429,8 @@ internal abstract class IndexTreeNode<T : IndexTreeNode<T>>() {
     var onRemovedListener: OnRemovedListener<T>? = null
 
     /**
-     * Updates the size of the ancestors.
+     * Updates the size of the ancestors. It is used when
+     * the size of the node is changed.
      */
     fun updateAncestorSize() {
         var parent = parent
@@ -434,6 +442,20 @@ internal abstract class IndexTreeNode<T : IndexTreeNode<T>>() {
         }
     }
 
+    /**
+     * Updates the size of the descendants. It is used when
+     * the tree is newly created and the size of the descendants is not calculated.
+     */
+    fun updateDescendantSize(): Int {
+        if (isRemoved) {
+            size = 0
+            return 0
+        }
+
+        size += children.sumOf { it.updateDescendantSize() }
+        return paddedSize
+    }
+
     fun findOffset(node: T): Int {
         check(!isText) {
             "Text node cannot have children"
@@ -474,10 +496,6 @@ internal abstract class IndexTreeNode<T : IndexTreeNode<T>>() {
 
         childNodes.add(0, node)
         node.parent = this as T
-
-        if (!node.isRemoved) {
-            node.updateAncestorSize()
-        }
     }
 
     /**

yorkie/src/main/kotlin/dev/yorkie/util/IndexTreeExtensions.kt

@@ -11,6 +11,17 @@ internal fun <T : IndexTreeNode<T>> traverse(
     action.invoke(node, depth)
 }
 
+internal fun <T : IndexTreeNode<T>> traverseAll(
+    node: T,
+    depth: Int = 0,
+    action: ((T, Int) -> Unit),
+) {
+    node.allChildren.forEach { child ->
+        traverseAll(child, depth + 1, action)
+    }
+    action.invoke(node, depth)
+}
+
 internal fun <T : IndexTreeNode<T>> findCommonAncestor(node1: T, node2: T): T? {
     if (node1 == node2) {
         return node1

yorkie/src/test/kotlin/dev/yorkie/api/ConverterTest.kt

@@ -4,6 +4,7 @@ import com.google.protobuf.kotlin.toByteStringUtf8
 import dev.yorkie.api.v1.jSONElement
 import dev.yorkie.api.v1.operation
 import dev.yorkie.core.PresenceChange
+import dev.yorkie.document.Document
 import dev.yorkie.document.change.Change
 import dev.yorkie.document.change.ChangeID
 import dev.yorkie.document.change.ChangePack
@@ -25,6 +26,9 @@ import dev.yorkie.document.crdt.RgaTreeList
 import dev.yorkie.document.crdt.RgaTreeSplit
 import dev.yorkie.document.crdt.RgaTreeSplitNodeID
 import dev.yorkie.document.crdt.RgaTreeSplitPos
+import dev.yorkie.document.json.JsonObject
+import dev.yorkie.document.json.JsonTree
+import dev.yorkie.document.json.TreeBuilder.element
 import dev.yorkie.document.operation.AddOperation
 import dev.yorkie.document.operation.EditOperation
 import dev.yorkie.document.operation.IncreaseOperation
@@ -44,6 +48,7 @@ import dev.yorkie.document.time.TimeTicket.Companion.MaxTimeTicket
 import dev.yorkie.util.IndexTreeNode.Companion.DEFAULT_ROOT_TYPE
 import java.util.Date
 import kotlin.test.assertEquals
+import kotlinx.coroutines.test.runTest
 import org.junit.Assert.assertThrows
 import org.junit.Test
 
@@ -413,6 +418,31 @@ class ConverterTest {
         }
     }
 
+    @Test
+    fun `should encode and decode tree properly`() = runTest {
+        val document = Document(Document.Key(""))
+        document.updateAsync { root, _ ->
+            root.setNewTree(
+                "t",
+                element("r") {
+                    element("p") { text { "12" } }
+                    element("p") { text { "34" } }
+                },
+            ).editByPath(listOf(0, 1), listOf(1, 1))
+        }.await()
+
+        fun JsonObject.tree() = getAs<JsonTree>("t")
+
+        assertEquals("<r><p>14</p></r>", document.getRoot().tree().toXml())
+        assertEquals(4, document.getRoot().tree().size)
+
+        val tree = document.getRoot().target["t"]
+        val bytes = tree.toByteString()
+        val obj = bytes.toCrdtTree()
+        assertEquals(document.getRoot().tree().target.nodeSize, obj.nodeSize)
+        assertEquals(document.getRoot().tree().size, obj.size)
+    }
+
     private class TestOperation(
         override val parentCreatedAt: TimeTicket,
         override var executedAt: TimeTicket,