Handle tombstoned elements when deleted node is restored through undo

src/document/crdt/root.ts

@@ -82,18 +82,7 @@ export class CRDTRoot {
     this.removedElementSetByCreatedAt = new Set();
     this.elementHasRemovedNodesSetByCreatedAt = new Set();
     this.opsForTest = [];
-
-    this.elementPairMapByCreatedAt.set(
-      this.rootObject.getCreatedAt().toIDString(),
-      { element: this.rootObject },
-    );
-
-    rootObject.getDescendants(
-      (elem: CRDTElement, parent: CRDTContainer): boolean => {
-        this.registerElement(elem, parent);
-        return false;
-      },
-    );
+    this.registerElement(rootObject, undefined);
   }
 
   /**
@@ -160,23 +149,49 @@ export class CRDTRoot {
   }
 
   /**
-   * `registerElement` registers the given element to hash table.
+   * `registerElement` registers the given element and its descendants to hash table.
    */
-  public registerElement(element: CRDTElement, parent: CRDTContainer): void {
+  public registerElement(element: CRDTElement, parent?: CRDTContainer): void {
     this.elementPairMapByCreatedAt.set(element.getCreatedAt().toIDString(), {
       parent,
       element,
     });
+
+    if (element instanceof CRDTContainer) {
+      element.getDescendants((elem, parent) => {
+        this.registerElement(elem, parent);
+        return false;
+      });
+    }
   }
 
   /**
-   * `deregisterElement` deregister the given element from hash table.
+   * `deregisterElement` deregister the given element and its descendants from hash table.
    */
-  public deregisterElement(element: CRDTElement): void {
-    this.elementPairMapByCreatedAt.delete(element.getCreatedAt().toIDString());
-    this.removedElementSetByCreatedAt.delete(
-      element.getCreatedAt().toIDString(),
-    );
+  public deregisterElement(element: CRDTElement): number {
+    let count = 0;
+
+    const callback = (elem: CRDTElement) => {
+      const createdAt = elem.getCreatedAt().toIDString();
+      this.elementPairMapByCreatedAt.delete(createdAt);
+      this.removedElementSetByCreatedAt.delete(createdAt);
+      count++;
+    };
+    const deregisterDescendants = (container: CRDTContainer) => {
+      container.getDescendants((elem) => {
+        callback(elem);
+        if (elem instanceof CRDTContainer) {
+          deregisterDescendants(elem);
+        }
+        return false;
+      });
+    };
+
+    callback(element);
+    if (element instanceof CRDTContainer) {
+      deregisterDescendants(element);
+    }
+    return count;
   }
 
   /**
@@ -263,7 +278,7 @@ export class CRDTRoot {
         ticket.compare(pair.element.getRemovedAt()!) >= 0
       ) {
         pair.parent!.purge(pair.element);
-        count += this.garbageCollectInternal(pair.element);
+        count += this.deregisterElement(pair.element);
       }
     }
 
@@ -283,25 +298,6 @@ export class CRDTRoot {
     return count;
   }
 
-  private garbageCollectInternal(element: CRDTElement): number {
-    let count = 0;
-
-    // eslint-disable-next-line @typescript-eslint/no-unused-vars
-    const callback = (elem: CRDTElement, parent?: CRDTContainer): boolean => {
-      this.deregisterElement(elem);
-      count++;
-      return false;
-    };
-
-    callback(element);
-
-    if (element instanceof CRDTContainer) {
-      element.getDescendants(callback);
-    }
-
-    return count;
-  }
-
   /**
    * `toJSON` returns the JSON encoding of this root object.
    */

src/document/json/array.ts

@@ -458,12 +458,6 @@ export class ArrayProxy {
     const element = buildCRDTElement(context, value, createdAt);
     target.insertAfter(prevCreatedAt, element);
     context.registerElement(element, target);
-    if (element instanceof CRDTContainer) {
-      element.getDescendants((elem, parent) => {
-        context.registerElement(elem, parent);
-        return false;
-      });
-    }
     context.push(
       AddOperation.create(
         target.getCreatedAt(),

src/document/json/object.ts

@@ -155,12 +155,6 @@ export class ObjectProxy {
     if (removed) {
       context.registerRemovedElement(removed);
     }
-    if (element instanceof CRDTContainer) {
-      element.getDescendants((elem, parent) => {
-        context.registerElement(elem, parent);
-        return false;
-      });
-    }
     context.push(
       SetOperation.create(
         key,

src/document/operation/add_operation.ts

@@ -72,12 +72,6 @@ export class AddOperation extends Operation {
     const value = this.value.deepcopy();
     array.insertAfter(this.prevCreatedAt, value);
     root.registerElement(value, array);
-    if (value instanceof CRDTContainer) {
-      value.getDescendants((elem, parent) => {
-        root.registerElement(elem, parent);
-        return false;
-      });
-    }
 
     return {
       opInfos: [

src/document/operation/set_operation.ts

@@ -90,16 +90,19 @@ export class SetOperation extends Operation {
 
     const value = this.value.deepcopy();
     const removed = obj.set(this.key, value, this.getExecutedAt());
+    // NOTE(chacha912): When resetting elements with the pre-existing createdAt
+    // during undo/redo, it's essential to handle previously tombstoned elements.
+    // In non-GC languages, there may be a need to execute both deregister and purge.
+    if (
+      source === OpSource.UndoRedo &&
+      root.findByCreatedAt(value.getCreatedAt())
+    ) {
+      root.deregisterElement(value);
+    }
     root.registerElement(value, obj);
     if (removed) {
       root.registerRemovedElement(removed);
     }
-    if (value instanceof CRDTContainer) {
-      value.getDescendants((elem, parent) => {
-        root.registerElement(elem, parent);
-        return false;
-      });
-    }
 
     return {
       opInfos: [

test/integration/object_test.ts

@@ -868,5 +868,50 @@ describe('Object', function () {
       await client2.sync();
       assert.equal(doc2.toSortedJSON(), '{"shape":{"color":"red"}}');
     });
+
+    it(`Should clean up the references to a previously deleted node when the deleted node is restored through undo`, async function ({
+      task,
+    }) {
+      interface TestDoc {
+        shape: { color: string };
+      }
+      const docKey = toDocKey(`${task.name}-${new Date().getTime()}`);
+      const doc1 = new Document<TestDoc>(docKey);
+      const doc2 = new Document<TestDoc>(docKey);
+
+      const client1 = new Client(testRPCAddr);
+      const client2 = new Client(testRPCAddr);
+      await client1.activate();
+      await client2.activate();
+
+      await client1.attach(doc1, { isRealtimeSync: false });
+      await client2.attach(doc2, { isRealtimeSync: false });
+
+      doc1.update((root) => {
+        root.shape = { color: 'black' };
+      });
+      await client1.sync();
+      await client2.sync();
+      assert.equal(doc1.toSortedJSON(), '{"shape":{"color":"black"}}');
+      assert.equal(doc2.toSortedJSON(), '{"shape":{"color":"black"}}');
+
+      doc2.update((root) => {
+        root.shape = { color: 'yellow' };
+      });
+      await client2.sync();
+      await client1.sync();
+      assert.equal(doc1.toSortedJSON(), '{"shape":{"color":"yellow"}}');
+      assert.equal(doc2.toSortedJSON(), '{"shape":{"color":"yellow"}}');
+
+      doc2.history.undo();
+      await client2.sync();
+      await client1.sync();
+      assert.equal(doc1.toSortedJSON(), '{"shape":{"color":"black"}}');
+      assert.equal(doc2.toSortedJSON(), '{"shape":{"color":"black"}}');
+
+      // NOTE(chacha912): removedElementSetByCreatedAt should only retain
+      // the entry for `{shape: {color: 'yellow'}}`.
+      assert.equal(doc2.getGarbageLen(), 2);
+    });
   });
 });