refactor(core): move block collection to affine and implement as doc (#9514)

packages/frontend/core/src/modules/doc/services/docs.ts

@@ -2,6 +2,7 @@ import { DebugLogger } from '@affine/debug';
 import { Unreachable } from '@affine/env/constant';
 import type { DocMode } from '@blocksuite/affine/blocks';
 import type { DeltaInsert } from '@blocksuite/affine/inline';
+import { Text } from '@blocksuite/affine/store';
 import type { AffineTextAttributes } from '@blocksuite/affine-shared/types';
 import { LiveData, ObjectPool, Service } from '@toeverything/infra';
 import { omitBy } from 'lodash-es';
@@ -121,7 +122,7 @@ export class DocsService extends Service {
     const { doc, release } = this.open(targetDocId);
     doc.setPriorityLoad(10);
     await doc.waitForSyncReady();
-    const text = new doc.blockSuiteDoc.Text([
+    const text = new Text([
       {
         insert: ' ',
         attributes: {

packages/frontend/core/src/modules/workspace/impl/doc.ts

@@ -0,0 +1,383 @@
+import { type Disposable, Slot } from '@blocksuite/affine/global/utils';
+import {
+  type AwarenessStore,
+  Blocks,
+  type BlockSuiteDoc,
+  type Doc,
+  type GetDocOptions,
+  type Query,
+  type Workspace,
+  type YBlock,
+} from '@blocksuite/affine/store';
+import { signal } from '@preact/signals-core';
+import * as Y from 'yjs';
+
+type DocOptions = {
+  id: string;
+  collection: Workspace;
+  doc: BlockSuiteDoc;
+  awarenessStore: AwarenessStore;
+};
+
+export class DocImpl implements Doc {
+  private _awarenessUpdateDisposable: Disposable | null = null;
+
+  private readonly _canRedo = signal(false);
+
+  private readonly _canUndo = signal(false);
+
+  private readonly _collection: Workspace;
+
+  private readonly _docMap = {
+    undefined: new Map<string, Blocks>(),
+    true: new Map<string, Blocks>(),
+    false: new Map<string, Blocks>(),
+  };
+
+  // doc/space container.
+  private readonly _handleYEvents = (events: Y.YEvent<YBlock | Y.Text>[]) => {
+    events.forEach(event => this._handleYEvent(event));
+  };
+
+  private _history!: Y.UndoManager;
+
+  private readonly _historyObserver = () => {
+    this._updateCanUndoRedoSignals();
+    this.slots.historyUpdated.emit();
+  };
+
+  private readonly _initSubDoc = () => {
+    let subDoc = this.rootDoc.spaces.get(this.id);
+    if (!subDoc) {
+      subDoc = new Y.Doc({
+        guid: this.id,
+      });
+      this.rootDoc.spaces.set(this.id, subDoc);
+      this._loaded = true;
+      this._onLoadSlot.emit();
+    } else {
+      this._loaded = false;
+      this.rootDoc.on('subdocs', this._onSubdocEvent);
+    }
+
+    return subDoc;
+  };
+
+  private _loaded!: boolean;
+
+  private readonly _onLoadSlot = new Slot();
+
+  private readonly _onSubdocEvent = ({
+    loaded,
+  }: {
+    loaded: Set<Y.Doc>;
+  }): void => {
+    const result = Array.from(loaded).find(
+      doc => doc.guid === this._ySpaceDoc.guid
+    );
+    if (!result) {
+      return;
+    }
+    this.rootDoc.off('subdocs', this._onSubdocEvent);
+    this._loaded = true;
+    this._onLoadSlot.emit();
+  };
+
+  /** Indicate whether the block tree is ready */
+  private _ready = false;
+
+  private _shouldTransact = true;
+
+  private readonly _updateCanUndoRedoSignals = () => {
+    const canRedo = this.readonly ? false : this._history.canRedo();
+    const canUndo = this.readonly ? false : this._history.canUndo();
+    if (this._canRedo.peek() !== canRedo) {
+      this._canRedo.value = canRedo;
+    }
+    if (this._canUndo.peek() !== canUndo) {
+      this._canUndo.value = canUndo;
+    }
+  };
+
+  protected readonly _yBlocks: Y.Map<YBlock>;
+
+  /**
+   * @internal Used for convenient access to the underlying Yjs map,
+   * can be used interchangeably with ySpace
+   */
+  protected readonly _ySpaceDoc: Y.Doc;
+
+  readonly awarenessStore: AwarenessStore;
+
+  readonly id: string;
+
+  readonly rootDoc: BlockSuiteDoc;
+
+  readonly slots = {
+    historyUpdated: new Slot(),
+    yBlockUpdated: new Slot<
+      | {
+          type: 'add';
+          id: string;
+        }
+      | {
+          type: 'delete';
+          id: string;
+        }
+    >(),
+  };
+
+  get blobSync() {
+    return this.collection.blobSync;
+  }
+
+  get canRedo() {
+    return this._canRedo.peek();
+  }
+
+  get canUndo() {
+    return this._canUndo.peek();
+  }
+
+  get collection() {
+    return this._collection;
+  }
+
+  get docSync() {
+    return this.collection.docSync;
+  }
+
+  get history() {
+    return this._history;
+  }
+
+  get isEmpty() {
+    return this._yBlocks.size === 0;
+  }
+
+  get loaded() {
+    return this._loaded;
+  }
+
+  get meta() {
+    return this.collection.meta.getDocMeta(this.id);
+  }
+
+  get readonly(): boolean {
+    return this.awarenessStore.isReadonly(this);
+  }
+
+  get ready() {
+    return this._ready;
+  }
+
+  get schema() {
+    return this.collection.schema;
+  }
+
+  get spaceDoc() {
+    return this._ySpaceDoc;
+  }
+
+  get yBlocks() {
+    return this._yBlocks;
+  }
+
+  constructor({ id, collection, doc, awarenessStore }: DocOptions) {
+    this.id = id;
+    this.rootDoc = doc;
+    this.awarenessStore = awarenessStore;
+
+    this._ySpaceDoc = this._initSubDoc();
+
+    this._yBlocks = this._ySpaceDoc.getMap('blocks');
+    this._collection = collection;
+  }
+
+  private _getReadonlyKey(readonly?: boolean): 'true' | 'false' | 'undefined' {
+    return (readonly?.toString() as 'true' | 'false') ?? 'undefined';
+  }
+
+  private _handleVersion() {
+    // Initialization from empty yDoc, indicating that the document is new.
+    if (!this.collection.meta.hasVersion) {
+      this.collection.meta.writeVersion(this.collection);
+    }
+  }
+
+  private _handleYBlockAdd(id: string) {
+    this.slots.yBlockUpdated.emit({ type: 'add', id });
+  }
+
+  private _handleYBlockDelete(id: string) {
+    this.slots.yBlockUpdated.emit({ type: 'delete', id });
+  }
+
+  private _handleYEvent(event: Y.YEvent<YBlock | Y.Text | Y.Array<unknown>>) {
+    // event on top-level block store
+    if (event.target !== this._yBlocks) {
+      return;
+    }
+    event.keys.forEach((value, id) => {
+      try {
+        if (value.action === 'add') {
+          this._handleYBlockAdd(id);
+          return;
+        }
+        if (value.action === 'delete') {
+          this._handleYBlockDelete(id);
+          return;
+        }
+      } catch (e) {
+        console.error('An error occurred while handling Yjs event:');
+        console.error(e);
+      }
+    });
+  }
+
+  private _initYBlocks() {
+    const { _yBlocks } = this;
+    _yBlocks.observeDeep(this._handleYEvents);
+    this._history = new Y.UndoManager([_yBlocks], {
+      trackedOrigins: new Set([this._ySpaceDoc.clientID]),
+    });
+
+    this._history.on('stack-cleared', this._historyObserver);
+    this._history.on('stack-item-added', this._historyObserver);
+    this._history.on('stack-item-popped', this._historyObserver);
+    this._history.on('stack-item-updated', this._historyObserver);
+  }
+
+  /** Capture current operations to undo stack synchronously. */
+  captureSync() {
+    this._history.stopCapturing();
+  }
+
+  clear() {
+    this._yBlocks.clear();
+  }
+
+  clearQuery(query: Query, readonly?: boolean) {
+    const readonlyKey = this._getReadonlyKey(readonly);
+
+    this._docMap[readonlyKey].delete(JSON.stringify(query));
+  }
+
+  private _destroy() {
+    this._ySpaceDoc.destroy();
+    this._onLoadSlot.dispose();
+    this._loaded = false;
+  }
+
+  dispose() {
+    this.slots.historyUpdated.dispose();
+    this._awarenessUpdateDisposable?.dispose();
+
+    if (this.ready) {
+      this._yBlocks.unobserveDeep(this._handleYEvents);
+      this._yBlocks.clear();
+    }
+  }
+
+  getDoc({ readonly, query }: GetDocOptions = {}) {
+    const readonlyKey = this._getReadonlyKey(readonly);
+
+    const key = JSON.stringify(query);
+
+    if (this._docMap[readonlyKey].has(key)) {
+      return this._docMap[readonlyKey].get(key) as Blocks;
+    }
+
+    const doc = new Blocks({
+      blockCollection: this,
+      schema: this.collection.schema,
+      readonly,
+      query,
+    });
+
+    this._docMap[readonlyKey].set(key, doc);
+
+    return doc;
+  }
+
+  load(initFn?: () => void): this {
+    if (this.ready) {
+      return this;
+    }
+
+    this._ySpaceDoc.load();
+
+    if ((this.collection.meta.docs?.length ?? 0) <= 1) {
+      this._handleVersion();
+    }
+
+    this._initYBlocks();
+
+    this._yBlocks.forEach((_, id) => {
+      this._handleYBlockAdd(id);
+    });
+
+    this._awarenessUpdateDisposable = this.awarenessStore.slots.update.on(
+      () => {
+        // change readonly state will affect the undo/redo state
+        this._updateCanUndoRedoSignals();
+      }
+    );
+
+    initFn?.();
+
+    this._ready = true;
+
+    return this;
+  }
+
+  redo() {
+    if (this.readonly) {
+      console.error('cannot modify data in readonly mode');
+      return;
+    }
+    this._history.redo();
+  }
+
+  undo() {
+    if (this.readonly) {
+      console.error('cannot modify data in readonly mode');
+      return;
+    }
+    this._history.undo();
+  }
+
+  remove() {
+    this._destroy();
+    this.rootDoc.spaces.delete(this.id);
+  }
+
+  resetHistory() {
+    this._history.clear();
+  }
+
+  /**
+   * If `shouldTransact` is `false`, the transaction will not be push to the history stack.
+   */
+  transact(fn: () => void, shouldTransact: boolean = this._shouldTransact) {
+    this._ySpaceDoc.transact(
+      () => {
+        try {
+          fn();
+        } catch (e) {
+          console.error(
+            `An error occurred while Y.doc ${this._ySpaceDoc.guid} transacting:`
+          );
+          console.error(e);
+        }
+      },
+      shouldTransact ? this.rootDoc.clientID : null
+    );
+  }
+
+  withoutTransact(callback: () => void) {
+    this._shouldTransact = false;
+    callback();
+    this._shouldTransact = true;
+  }
+}

packages/frontend/core/src/modules/workspace/impl/workspace.ts

@@ -6,10 +6,10 @@ import type { BlockSuiteFlags } from '@blocksuite/affine/global/types';
 import { NoopLogger, Slot } from '@blocksuite/affine/global/utils';
 import {
   AwarenessStore,
-  BlockCollection,
   type Blocks,
   BlockSuiteDoc,
   type CreateDocOptions,
+  type Doc,
   DocCollectionMeta,
   type GetDocOptions,
   type IdGenerator,
@@ -27,6 +27,8 @@ import {
 } from '@blocksuite/affine/sync';
 import { Awareness } from 'y-protocols/awareness.js';
 
+import { DocImpl } from './doc';
+
 type WorkspaceOptions = {
   id?: string;
   schema: Schema;
@@ -62,7 +64,7 @@ export class WorkspaceImpl implements Workspace {
 
   readonly blobSync: BlobEngine;
 
-  readonly blockCollections = new Map<string, BlockCollection>();
+  readonly blockCollections = new Map<string, Doc>();
 
   readonly doc: BlockSuiteDoc;
 
@@ -114,12 +116,11 @@ export class WorkspaceImpl implements Workspace {
 
   private _bindDocMetaEvents() {
     this.meta.docMetaAdded.on(docId => {
-      const doc = new BlockCollection({
+      const doc = new DocImpl({
         id: docId,
         collection: this,
         doc: this.doc,
         awarenessStore: this.awarenessStore,
-        idGenerator: this.idGenerator,
       });
       this.blockCollections.set(doc.id, doc);
     });
@@ -185,8 +186,8 @@ export class WorkspaceImpl implements Workspace {
     this.awarenessSync.disconnect();
   }
 
-  getBlockCollection(docId: string): BlockCollection | null {
-    const space = this.docs.get(docId) as BlockCollection | undefined;
+  getBlockCollection(docId: string): Doc | null {
+    const space = this.docs.get(docId) as Doc | undefined;
     return space ?? null;
   }