feat: support create milestone from yDoc (#1781)

packages/y-indexeddb/src/__tests__/index.spec.ts

@@ -7,9 +7,17 @@ import { __unstableSchemas, AffineSchemas } from '@blocksuite/blocks/models';
 import { assertExists, uuidv4, Workspace } from '@blocksuite/store';
 import { openDB } from 'idb';
 import { beforeEach, describe, expect, test, vi } from 'vitest';
+import { applyUpdate, Doc, encodeStateAsUpdate } from 'yjs';
 
 import type { WorkspacePersist } from '../index';
-import { createIndexedDBProvider, dbVersion, setMergeCount } from '../index';
+import {
+  createIndexedDBProvider,
+  dbVersion,
+  getMilestones,
+  markMilestone,
+  revertUpdate,
+  setMergeCount,
+} from '../index';
 
 async function getUpdates(id: string): Promise<ArrayBuffer[]> {
   const db = await openDB('affine-local', dbVersion);
@@ -146,3 +154,66 @@ describe('indexeddb provider', () => {
     }
   });
 });
+
+describe('milestone', () => {
+  test('milestone', async () => {
+    const doc = new Doc();
+    const map = doc.getMap('map');
+    const array = doc.getArray('array');
+    map.set('1', 1);
+    array.push([1]);
+    await markMilestone('1', doc, 'test1');
+    const milestones = await getMilestones('1');
+    assertExists(milestones);
+    expect(milestones).toBeDefined();
+    expect(Object.keys(milestones).length).toBe(1);
+    expect(milestones.test1).toBeInstanceOf(Uint8Array);
+    const snapshot = new Doc();
+    applyUpdate(snapshot, milestones.test1);
+    {
+      const map = snapshot.getMap('map');
+      expect(map.get('1')).toBe(1);
+    }
+    map.set('1', 2);
+    {
+      const map = snapshot.getMap('map');
+      expect(map.get('1')).toBe(1);
+    }
+    revertUpdate(doc, milestones.test1, {
+      map: 'Map',
+      array: 'Array',
+    });
+    {
+      const map = doc.getMap('map');
+      expect(map.get('1')).toBe(1);
+    }
+
+    const fn = vi.fn(() => true);
+    doc.gcFilter = fn;
+    expect(fn).toBeCalledTimes(0);
+
+    for (let i = 0; i < 1e5; i++) {
+      map.set(`${i}`, i + 1);
+    }
+    for (let i = 0; i < 1e5; i++) {
+      map.delete(`${i}`);
+    }
+    for (let i = 0; i < 1e5; i++) {
+      map.set(`${i}`, i - 1);
+    }
+
+    expect(fn).toBeCalled();
+
+    const doc2 = new Doc();
+    applyUpdate(doc2, encodeStateAsUpdate(doc));
+
+    revertUpdate(doc2, milestones.test1, {
+      map: 'Map',
+      array: 'Array',
+    });
+    {
+      const map = doc2.getMap('map');
+      expect(map.get('1')).toBe(1);
+    }
+  });
+});

packages/y-indexeddb/src/index.ts

@@ -5,13 +5,57 @@ import {
   diffUpdate,
   Doc,
   encodeStateAsUpdate,
+  encodeStateVector,
   mergeUpdates,
+  UndoManager,
 } from 'yjs';
 
 const indexeddbOrigin = Symbol('indexeddb-provider-origin');
+const snapshotOrigin = Symbol('snapshot-origin');
 
 let mergeCount = 500;
 
+type Metadata = Record<string, 'Text' | 'Map' | 'Array'>;
+
+export function revertUpdate(
+  doc: Doc,
+  snapshotUpdate: Uint8Array,
+  metadata: Metadata
+) {
+  const snapshotDoc = new Doc();
+  applyUpdate(snapshotDoc, snapshotUpdate, snapshotOrigin);
+
+  const currentStateVector = encodeStateVector(doc);
+  const snapshotStateVector = encodeStateVector(snapshotDoc);
+
+  const changesSinceSnapshotUpdate = encodeStateAsUpdate(
+    doc,
+    snapshotStateVector
+  );
+  const undoManager = new UndoManager(
+    [...snapshotDoc.share.keys()].map(key => {
+      if (metadata[key] === 'Text') {
+        return snapshotDoc.getText(key);
+      } else if (metadata[key] === 'Map') {
+        return snapshotDoc.getMap(key);
+      } else if (metadata[key] === 'Array') {
+        return snapshotDoc.getArray(key);
+      }
+      throw new Error('Unknown type');
+    }),
+    {
+      trackedOrigins: new Set([snapshotOrigin]),
+    }
+  );
+  applyUpdate(snapshotDoc, changesSinceSnapshotUpdate, snapshotOrigin);
+  undoManager.undo();
+  const revertChangesSinceSnapshotUpdate = encodeStateAsUpdate(
+    snapshotDoc,
+    currentStateVector
+  );
+  applyUpdate(doc, revertChangesSinceSnapshotUpdate, snapshotOrigin);
+}
+
 export class EarlyDisconnectError extends Error {
   constructor() {
     super('Early disconnect');
@@ -69,6 +113,52 @@ export interface OldYjsDB extends DBSchema {
   };
 }
 
+export const markMilestone = async (
+  id: string,
+  doc: Doc,
+  name: string,
+  dbName = 'affine-local'
+): Promise<void> => {
+  const dbPromise = openDB<BlockSuiteBinaryDB>(dbName, dbVersion, {
+    upgrade: upgradeDB,
+  });
+  const db = await dbPromise;
+  const store = db
+    .transaction('milestone', 'readwrite')
+    .objectStore('milestone');
+  const milestone = await store.get('id');
+  const binary = encodeStateAsUpdate(doc);
+  if (!milestone) {
+    await store.put({
+      id,
+      milestone: {
+        [name]: binary,
+      },
+    });
+  } else {
+    milestone.milestone[name] = binary;
+    await store.put(milestone);
+  }
+};
+
+export const getMilestones = async (
+  id: string,
+  dbName = 'affine-local'
+): Promise<null | WorkspaceMilestone['milestone']> => {
+  const dbPromise = openDB<BlockSuiteBinaryDB>(dbName, dbVersion, {
+    upgrade: upgradeDB,
+  });
+  const db = await dbPromise;
+  const store = db
+    .transaction('milestone', 'readonly')
+    .objectStore('milestone');
+  const milestone = await store.get(id);
+  if (!milestone) {
+    return null;
+  }
+  return milestone.milestone;
+};
+
 export const createIndexedDBProvider = (
   id: string,
   doc: Doc,