Enhance type inference in Document.subscribe

This commit improves the type inference for the types used in
Document.subscribe. Additionally, it removes the code that forces
casting in the test code to align with these improvements.

---------

Co-authored-by: Hackerwins <susukang98@gmail.com>

Author: chacha912

src/document/crdt/tree.ts

@@ -87,16 +87,37 @@ export enum TreeChangeType {
 /**
  * `TreeChange` represents the change in the tree.
  */
-export interface TreeChange {
-  actor: ActorID;
-  type: TreeChangeType;
-  from: number;
-  to: number;
-  fromPath: Array<number>;
-  toPath: Array<number>;
-  value?: Array<TreeNode> | { [key: string]: any } | Array<string>;
-  splitLevel?: number;
-}
+export type TreeChange =
+  | {
+      actor: ActorID;
+      type: TreeChangeType.Content;
+      from: number;
+      to: number;
+      fromPath: Array<number>;
+      toPath: Array<number>;
+      value?: Array<TreeNode>;
+      splitLevel?: number;
+    }
+  | {
+      actor: ActorID;
+      type: TreeChangeType.Style;
+      from: number;
+      to: number;
+      fromPath: Array<number>;
+      toPath: Array<number>;
+      value: { [key: string]: string };
+      splitLevel?: number;
+    }
+  | {
+      actor: ActorID;
+      type: TreeChangeType.RemoveStyle;
+      from: number;
+      to: number;
+      fromPath: Array<number>;
+      toPath: Array<number>;
+      value?: Array<string>;
+      splitLevel?: number;
+    };
 
 /**
  * `CRDTTreePos` represent a position in the tree. It is used to identify a
@@ -876,7 +897,6 @@ export class CRDTTree extends CRDTGCElement {
     const [toParent, toLeft] = this.findNodesAndSplitText(range[1], editedAt);
 
     const changes: Array<TreeChange> = [];
-    const value = attributesToRemove ? attributesToRemove : undefined;
     this.traverseInPosRange(
       fromParent,
       fromLeft,
@@ -893,13 +913,13 @@ export class CRDTTree extends CRDTGCElement {
           }
 
           changes.push({
+            actor: editedAt.getActorID()!,
             type: TreeChangeType.RemoveStyle,
             from: this.toIndex(fromParent, fromLeft),
             to: this.toIndex(toParent, toLeft),
             fromPath: this.toPath(fromParent, fromLeft),
             toPath: this.toPath(toParent, toLeft),
-            actor: editedAt.getActorID()!,
-            value,
+            value: attributesToRemove,
           });
         }
       },

src/document/document.ts

@@ -369,6 +369,28 @@ export interface PresenceChangedEvent<P extends Indexable>
   value: { clientID: ActorID; presence: P };
 }
 
+type DocEventCallbackMap<P extends Indexable> = {
+  default: NextFn<
+    | SnapshotEvent
+    | LocalChangeEvent<OperationInfo, P>
+    | RemoteChangeEvent<OperationInfo, P>
+  >;
+  presence: NextFn<
+    | InitializedEvent<P>
+    | WatchedEvent<P>
+    | UnwatchedEvent<P>
+    | PresenceChangedEvent<P>
+  >;
+  'my-presence': NextFn<InitializedEvent<P> | PresenceChangedEvent<P>>;
+  others: NextFn<WatchedEvent<P> | UnwatchedEvent<P> | PresenceChangedEvent<P>>;
+  connection: NextFn<ConnectionChangedEvent>;
+  sync: NextFn<SyncStatusChangedEvent>;
+  all: NextFn<TransactionEvent<P>>;
+};
+export type DocEventTopic = keyof DocEventCallbackMap<never>;
+export type DocEventCallback<P extends Indexable> =
+  DocEventCallbackMap<P>[DocEventTopic];
+
 /**
  * Indexable key, value
  * @public
@@ -710,7 +732,7 @@ export class Document<T, P extends Indexable = Indexable> {
    * The callback will be called when the document is changed.
    */
   public subscribe(
-    nextOrObserver: Observer<DocEvent> | NextFn<DocEvent>,
+    next: DocEventCallbackMap<P>['default'],
     error?: ErrorFn,
     complete?: CompleteFn,
   ): Unsubscribe;
@@ -721,7 +743,7 @@ export class Document<T, P extends Indexable = Indexable> {
    */
   public subscribe(
     type: 'presence',
-    next: NextFn<DocEvent<P>>,
+    next: DocEventCallbackMap<P>['presence'],
     error?: ErrorFn,
     complete?: CompleteFn,
   ): Unsubscribe;
@@ -731,7 +753,7 @@ export class Document<T, P extends Indexable = Indexable> {
    */
   public subscribe(
     type: 'my-presence',
-    next: NextFn<DocEvent<P>>,
+    next: DocEventCallbackMap<P>['my-presence'],
     error?: ErrorFn,
     complete?: CompleteFn,
   ): Unsubscribe;
@@ -742,7 +764,7 @@ export class Document<T, P extends Indexable = Indexable> {
    */
   public subscribe(
     type: 'others',
-    next: NextFn<DocEvent<P>>,
+    next: DocEventCallbackMap<P>['others'],
     error?: ErrorFn,
     complete?: CompleteFn,
   ): Unsubscribe;
@@ -752,7 +774,7 @@ export class Document<T, P extends Indexable = Indexable> {
    */
   public subscribe(
     type: 'connection',
-    next: NextFn<DocEvent<P>>,
+    next: DocEventCallbackMap<P>['connection'],
     error?: ErrorFn,
     complete?: CompleteFn,
   ): Unsubscribe;
@@ -762,7 +784,7 @@ export class Document<T, P extends Indexable = Indexable> {
    */
   public subscribe(
     type: 'sync',
-    next: NextFn<DocEvent<P>>,
+    next: DocEventCallbackMap<P>['sync'],
     error?: ErrorFn,
     complete?: CompleteFn,
   ): Unsubscribe;
@@ -775,7 +797,9 @@ export class Document<T, P extends Indexable = Indexable> {
     TOperationInfo extends OperationInfoOf<T, TPath>,
   >(
     targetPath: TPath,
-    next: NextFn<DocEvent<P, TOperationInfo>>,
+    next: NextFn<
+      LocalChangeEvent<TOperationInfo, P> | RemoteChangeEvent<TOperationInfo, P>
+    >,
     error?: ErrorFn,
     complete?: CompleteFn,
   ): Unsubscribe;
@@ -784,7 +808,7 @@ export class Document<T, P extends Indexable = Indexable> {
    */
   public subscribe(
     type: 'all',
-    next: NextFn<TransactionEvent<P>>,
+    next: DocEventCallbackMap<P>['all'],
     error?: ErrorFn,
     complete?: CompleteFn,
   ): Unsubscribe;
@@ -795,11 +819,13 @@ export class Document<T, P extends Indexable = Indexable> {
     TPath extends PathOf<T>,
     TOperationInfo extends OperationInfoOf<T, TPath>,
   >(
-    arg1: TPath | string | Observer<DocEvent<P>> | NextFn<DocEvent<P>>,
+    arg1: TPath | DocEventTopic | DocEventCallbackMap<P>['default'],
     arg2?:
-      | NextFn<DocEvent<P, TOperationInfo>>
-      | NextFn<DocEvent<P>>
-      | NextFn<TransactionEvent<P>>
+      | NextFn<
+          | LocalChangeEvent<TOperationInfo, P>
+          | RemoteChangeEvent<TOperationInfo, P>
+        >
+      | DocEventCallback<P>
       | ErrorFn,
     arg3?: ErrorFn | CompleteFn,
     arg4?: CompleteFn,
@@ -809,7 +835,7 @@ export class Document<T, P extends Indexable = Indexable> {
         throw new Error('Second argument must be a callback function');
       }
       if (arg1 === 'presence') {
-        const callback = arg2 as NextFn<DocEvent<P>>;
+        const callback = arg2 as DocEventCallbackMap<P>['presence'];
         return this.eventStream.subscribe(
           (event) => {
             for (const docEvent of event) {
@@ -830,21 +856,19 @@ export class Document<T, P extends Indexable = Indexable> {
         );
       }
       if (arg1 === 'my-presence') {
-        const callback = arg2 as NextFn<DocEvent<P>>;
+        const callback = arg2 as DocEventCallbackMap<P>['my-presence'];
         return this.eventStream.subscribe(
           (event) => {
             for (const docEvent of event) {
               if (
                 docEvent.type !== DocEventType.Initialized &&
-                docEvent.type !== DocEventType.Watched &&
-                docEvent.type !== DocEventType.Unwatched &&
                 docEvent.type !== DocEventType.PresenceChanged
               ) {
                 continue;
               }
 
               if (
-                docEvent.type !== DocEventType.Initialized &&
+                docEvent.type === DocEventType.PresenceChanged &&
                 docEvent.value.clientID !== this.changeID.getActorID()
               ) {
                 continue;
@@ -858,7 +882,7 @@ export class Document<T, P extends Indexable = Indexable> {
         );
       }
       if (arg1 === 'others') {
-        const callback = arg2 as NextFn<DocEvent<P>>;
+        const callback = arg2 as DocEventCallbackMap<P>['others'];
         return this.eventStream.subscribe(
           (event) => {
             for (const docEvent of event) {
@@ -880,7 +904,7 @@ export class Document<T, P extends Indexable = Indexable> {
         );
       }
       if (arg1 === 'connection') {
-        const callback = arg2 as NextFn<DocEvent<P>>;
+        const callback = arg2 as DocEventCallbackMap<P>['connection'];
         return this.eventStream.subscribe(
           (event) => {
             for (const docEvent of event) {
@@ -895,7 +919,7 @@ export class Document<T, P extends Indexable = Indexable> {
         );
       }
       if (arg1 === 'sync') {
-        const callback = arg2 as NextFn<DocEvent<P>>;
+        const callback = arg2 as DocEventCallbackMap<P>['sync'];
         return this.eventStream.subscribe(
           (event) => {
             for (const docEvent of event) {
@@ -910,31 +934,28 @@ export class Document<T, P extends Indexable = Indexable> {
         );
       }
       if (arg1 === 'all') {
-        const callback = arg2 as NextFn<TransactionEvent<P>>;
+        const callback = arg2 as DocEventCallbackMap<P>['all'];
         return this.eventStream.subscribe(callback, arg3, arg4);
       }
       const target = arg1;
-      const callback = arg2 as NextFn<DocEvent<P>>;
+      const callback = arg2 as NextFn<
+        | LocalChangeEvent<TOperationInfo, P>
+        | RemoteChangeEvent<TOperationInfo, P>
+      >;
       return this.eventStream.subscribe(
         (event) => {
           for (const docEvent of event) {
             if (
-              docEvent.type !== DocEventType.Snapshot &&
               docEvent.type !== DocEventType.LocalChange &&
               docEvent.type !== DocEventType.RemoteChange
             ) {
               continue;
             }
 
-            if (docEvent.type === DocEventType.Snapshot) {
-              target === '$' && callback(docEvent);
-              continue;
-            }
-
-            const targetOps: Array<OperationInfo> = [];
+            const targetOps: Array<TOperationInfo> = [];
             for (const op of docEvent.value.operations) {
               if (this.isSameElementOrChildOf(op.path, target)) {
-                targetOps.push(op);
+                targetOps.push(op as TOperationInfo);
               }
             }
             targetOps.length &&
@@ -949,7 +970,7 @@ export class Document<T, P extends Indexable = Indexable> {
       );
     }
     if (typeof arg1 === 'function') {
-      const callback = arg1 as NextFn<DocEvent<P>>;
+      const callback = arg1 as DocEventCallbackMap<P>['default'];
       const error = arg2 as ErrorFn;
       const complete = arg3 as CompleteFn;
       return this.eventStream.subscribe(

src/document/operation/operation.ts

@@ -149,7 +149,7 @@ export type TreeEditOpInfo = {
   path: string;
   from: number;
   to: number;
-  value: TreeNode;
+  value?: Array<TreeNode>;
   splitLevel: number;
   fromPath: Array<number>;
   toPath: Array<number>;

test/integration/client_test.ts

@@ -140,15 +140,15 @@ describe.sequential('Client', function () {
 
     const unsub1 = {
       syncEvent: d1.subscribe('sync', (event) => {
-        eventCollectorSync1.add(event.value as DocumentSyncStatus);
+        eventCollectorSync1.add(event.value);
       }),
       doc: d1.subscribe((event) => {
         eventCollectorD1.add(event.type);
       }),
     };
     const unsub2 = {
       syncEvent: d2.subscribe('sync', (event) => {
-        eventCollectorSync2.add(event.value as DocumentSyncStatus);
+        eventCollectorSync2.add(event.value);
       }),
       doc: d2.subscribe((event) => {
         eventCollectorD2.add(event.type);
@@ -236,7 +236,7 @@ describe.sequential('Client', function () {
     // 02. c2 changes the sync mode to realtime sync mode.
     const eventCollector = new EventCollector();
     const unsub1 = d2.subscribe('sync', (event) => {
-      eventCollector.add(event.value as DocumentSyncStatus);
+      eventCollector.add(event.value);
     });
     await c2.changeSyncMode(d2, SyncMode.Realtime);
     await eventCollector.waitFor(DocumentSyncStatus.Synced); // sync occurs when resuming
@@ -414,7 +414,7 @@ describe.sequential('Client', function () {
 
     const eventCollector = new EventCollector();
     const unsub1 = d2.subscribe('sync', (event) => {
-      eventCollector.add(event.value as DocumentSyncStatus);
+      eventCollector.add(event.value);
     });
 
     // 01. c2 attach the doc with realtime sync mode at first.
@@ -491,7 +491,7 @@ describe.sequential('Client', function () {
     //     and sync with push-only mode: CP(2, 2) -> CP(3, 2)
     const eventCollector = new EventCollector();
     const unsub = d1.subscribe('sync', (event) => {
-      eventCollector.add(event.value as DocumentSyncStatus);
+      eventCollector.add(event.value);
     });
     d1.update((root) => {
       root.counter.increase(1);
@@ -569,7 +569,7 @@ describe.sequential('Client', function () {
     assert.equal(d1.getRoot().tree.toXML(), '<doc><p>12</p><p>34</p></doc>');
     assert.equal(d2.getRoot().tree.toXML(), '<doc><p>12</p><p>34</p></doc>');
 
-    d1.update((root: any) => {
+    d1.update((root) => {
       root.tree.edit(2, 2, { type: 'text', value: 'a' });
     });
     await c1.sync();
@@ -595,7 +595,7 @@ describe.sequential('Client', function () {
 
     c2.changeSyncMode(d2, SyncMode.Realtime);
 
-    d2.update((root: any) => {
+    d2.update((root) => {
       root.tree.edit(2, 2, { type: 'text', value: 'b' });
     });
     await eventCollectorD1.waitAndVerifyNthEvent(3, DocEventType.RemoteChange);