refactor(editor): introduce store container to make implement doc easier (#12146)

<!-- This is an auto-generated comment: release notes by coderabbit.ai -->
## Summary by CodeRabbit

- **New Features**
  - Introduced a new store management system for handling document stores, improving efficiency and flexibility when working with document data.

- **Refactor**
  - Updated internal store handling to use a centralized store container, simplifying store retrieval and removal across various components.
  - Renamed and updated several store-related method signatures for consistency and clarity.
  - Replaced editor extension loading logic with a new local implementation for better modularity.

- **Chores**
  - Improved and streamlined the export of store-related modules for better maintainability.
  - Removed obsolete and redundant code related to previous store management approaches.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->

blocksuite/affine/blocks/embed-doc/src/embed-synced-doc-block/embed-synced-doc-block.ts

@@ -34,7 +34,7 @@ import {
   LifeCycleWatcher,
 } from '@blocksuite/std';
 import { GfxControllerIdentifier, GfxExtension } from '@blocksuite/std/gfx';
-import { type GetBlocksOptions, type Query, Text } from '@blocksuite/store';
+import { type GetStoreOptions, type Query, Text } from '@blocksuite/store';
 import { computed, signal } from '@preact/signals-core';
 import { html, nothing, type PropertyValues } from 'lit';
 import { query, state } from 'lit/decorators.js';
@@ -403,7 +403,7 @@ export class EmbedSyncedDocBlockComponent extends EmbedBlockComponent<EmbedSynce
   }
 
   get syncedDoc() {
-    const options: GetBlocksOptions = { readonly: true };
+    const options: GetStoreOptions = { readonly: true };
     if (this.isPageMode) options.query = this._pageFilter;
     const doc = this.std.workspace.getDoc(this.model.props.pageId);
     return doc?.getStore(options) ?? null;

blocksuite/affine/blocks/surface-ref/src/portal/note.ts

@@ -56,7 +56,7 @@ export class SurfaceRefNotePortal extends WithDisposable(ShadowlessElement) {
 
     const doc = this.model.doc;
     this._disposables.add(() => {
-      doc.doc.clearQuery(query, true);
+      doc.doc.removeStore({ query, readonly: true });
     });
   }
 

blocksuite/affine/fragments/frame-panel/src/card/frame-preview.ts

@@ -120,7 +120,7 @@ export class FramePreview extends WithDisposable(ShadowlessElement) {
     this._previewDoc =
       doc?.getStore({ readonly: true, query: this._docFilter }) ?? null;
     this.disposables.add(() => {
-      this._originalDoc.doc.clearQuery(this._docFilter);
+      this._originalDoc.doc.removeStore({ query: this._docFilter });
     });
   }
 

blocksuite/framework/store/src/model/doc.ts

@@ -2,12 +2,15 @@ import type * as Y from 'yjs';
 
 import type { AwarenessStore } from '../yjs/awareness.js';
 import type { YBlock } from './block/types.js';
-import type { Query } from './store/query.js';
 import type { Store, StoreOptions } from './store/store.js';
 import type { Workspace } from './workspace.js';
 import type { DocMeta } from './workspace-meta.js';
 
-export type GetBlocksOptions = Omit<StoreOptions, 'schema' | 'doc'>;
+export type GetStoreOptions = Omit<StoreOptions, 'schema' | 'doc'>;
+export type RemoveStoreOptions = Pick<
+  StoreOptions,
+  'query' | 'id' | 'readonly'
+>;
 
 export interface Doc {
   readonly id: string;
@@ -19,8 +22,8 @@ export interface Doc {
   dispose(): void;
 
   clear(): void;
-  getStore(options?: GetBlocksOptions): Store;
-  clearQuery(query: Query, readonly?: boolean): void;
+  getStore(options?: GetStoreOptions): Store;
+  removeStore(options: RemoveStoreOptions): void;
 
   get loaded(): boolean;
   get awarenessStore(): AwarenessStore;

blocksuite/framework/store/src/model/index.ts

@@ -1,5 +1,6 @@
 export * from './block/index.js';
 export * from './doc.js';
 export * from './store/index.js';
+export * from './store-container.js';
 export * from './workspace.js';
 export * from './workspace-meta.js';

blocksuite/framework/store/src/model/store-container.ts

@@ -0,0 +1,74 @@
+import type { Doc, GetStoreOptions, RemoveStoreOptions } from './doc';
+import { type Query, Store } from './store';
+
+export class StoreContainer {
+  private readonly _storeMap = new Map<string, Store>();
+
+  constructor(readonly doc: Doc) {}
+
+  getStore = ({
+    readonly,
+    query,
+    provider,
+    extensions,
+    id,
+  }: GetStoreOptions = {}) => {
+    let idOrOptions: string | { readonly?: boolean; query?: Query };
+    if (readonly || query) {
+      idOrOptions = { readonly, query };
+    } else if (!id) {
+      idOrOptions = this.doc.workspace.idGenerator();
+    } else {
+      idOrOptions = id;
+    }
+    const key = this._getQueryKey(idOrOptions);
+
+    if (this._storeMap.has(key)) {
+      return this._storeMap.get(key) as Store;
+    }
+
+    const doc = new Store({
+      doc: this.doc,
+      readonly,
+      query,
+      provider,
+      extensions,
+    });
+
+    this._storeMap.set(key, doc);
+
+    return doc;
+  };
+
+  removeStore = ({ readonly, query, id }: RemoveStoreOptions) => {
+    let idOrOptions: string | { readonly?: boolean; query?: Query };
+    if (readonly || query) {
+      idOrOptions = { readonly, query };
+    } else if (!id) {
+      return;
+    } else {
+      idOrOptions = id;
+    }
+    const key = this._getQueryKey(idOrOptions);
+    this._storeMap.delete(key);
+  };
+
+  private readonly _getQueryKey = (
+    idOrOptions: string | { readonly?: boolean; query?: Query }
+  ) => {
+    if (typeof idOrOptions === 'string') {
+      return idOrOptions;
+    }
+    const { readonly, query } = idOrOptions;
+    const readonlyKey = this._getReadonlyKey(readonly);
+    const key = JSON.stringify({
+      readonlyKey,
+      query,
+    });
+    return key;
+  };
+
+  private _getReadonlyKey(readonly?: boolean): 'true' | 'false' {
+    return (readonly?.toString() as 'true' | 'false') ?? 'false';
+  }
+}

blocksuite/framework/store/src/test/test-doc.ts

@@ -1,9 +1,12 @@
 import * as Y from 'yjs';
 
 import type { YBlock } from '../model/block/types.js';
-import type { Doc, GetBlocksOptions, Workspace } from '../model/index.js';
-import type { Query } from '../model/store/query.js';
-import { Store } from '../model/store/store.js';
+import {
+  type Doc,
+  type GetStoreOptions,
+  StoreContainer,
+  type Workspace,
+} from '../model/index.js';
 import type { AwarenessStore } from '../yjs/index.js';
 import type { TestWorkspace } from './test-workspace.js';
 
@@ -17,7 +20,7 @@ type DocOptions = {
 export class TestDoc implements Doc {
   private readonly _collection: Workspace;
 
-  private readonly _storeMap = new Map<string, Store>();
+  private readonly _storeContainer: StoreContainer;
 
   private readonly _initSubDoc = () => {
     let subDoc = this.rootDoc.getMap('spaces').get(this.id);
@@ -110,19 +113,15 @@ export class TestDoc implements Doc {
 
     this._yBlocks = this._ySpaceDoc.getMap('blocks');
     this._collection = collection;
-  }
-
-  private _getReadonlyKey(readonly?: boolean): 'true' | 'false' {
-    return (readonly?.toString() as 'true' | 'false') ?? 'false';
+    this._storeContainer = new StoreContainer(this);
   }
 
   clear() {
     this._yBlocks.clear();
   }
 
-  clearQuery(query: Query, readonly?: boolean) {
-    const key = this._getQueryKey({ readonly, query });
-    this._storeMap.delete(key);
+  get removeStore() {
+    return this._storeContainer.removeStore;
   }
 
   private _destroy() {
@@ -136,55 +135,34 @@ export class TestDoc implements Doc {
     }
   }
 
-  private readonly _getQueryKey = (
-    idOrOptions: string | { readonly?: boolean; query?: Query }
-  ) => {
-    if (typeof idOrOptions === 'string') {
-      return idOrOptions;
-    }
-    const { readonly, query } = idOrOptions;
-    const readonlyKey = this._getReadonlyKey(readonly);
-    const key = JSON.stringify({
-      readonlyKey,
-      query,
-    });
-    return key;
-  };
-
   getStore({
     readonly,
     query,
     provider,
     extensions,
     id,
-  }: GetBlocksOptions = {}) {
-    let idOrOptions: string | { readonly?: boolean; query?: Query };
+  }: GetStoreOptions = {}) {
+    const storeExtensions = (
+      this.workspace as TestWorkspace
+    ).storeExtensions.concat(extensions ?? []);
+
+    let storeId: string | undefined;
+
     if (id) {
-      idOrOptions = id;
-    } else if (readonly === undefined && query === undefined) {
-      idOrOptions = this.spaceDoc.guid;
+      storeId = id;
+    } else if (readonly !== undefined || query) {
+      storeId = id;
     } else {
-      idOrOptions = { readonly, query };
-    }
-    const key = this._getQueryKey(idOrOptions);
-
-    if (this._storeMap.has(key)) {
-      return this._storeMap.get(key)!;
+      storeId = this.spaceDoc.guid;
     }
 
-    const doc = new Store({
-      doc: this,
+    return this._storeContainer.getStore({
+      id: storeId,
       readonly,
       query,
       provider,
-      extensions: (this.workspace as TestWorkspace).storeExtensions.concat(
-        extensions ?? []
-      ),
+      extensions: storeExtensions,
     });
-
-    this._storeMap.set(key, doc);
-
-    return doc;
   }
 
   load(initFn?: () => void): this {

packages/frontend/core/src/blocksuite/ai/components/text-renderer.ts

@@ -268,7 +268,7 @@ export class TextRenderer extends WithDisposable(ShadowlessElement) {
               query: this._query,
             });
             this.disposables.add(() => {
-              doc.doc.clearQuery(this._query);
+              doc.doc.removeStore({ query: this._query });
             });
             this._doc.readonly = true;
             this.requestUpdate();

packages/frontend/core/src/blocksuite/block-suite-editor/lit-adaper.tsx

@@ -7,6 +7,7 @@ import {
   type PageEditor,
 } from '@affine/core/blocksuite/editors';
 import type { AffineEditorViewOptions } from '@affine/core/blocksuite/manager/editor-view';
+import { getViewManager } from '@affine/core/blocksuite/manager/migrating-view';
 import { useEnableAI } from '@affine/core/components/hooks/affine/use-enable-ai';
 import type { DocCustomPropertyInfo } from '@affine/core/modules/db';
 import type {
@@ -21,8 +22,9 @@ import { WorkspaceService } from '@affine/core/modules/workspace';
 import track from '@affine/track';
 import type { DocTitle } from '@blocksuite/affine/fragments/doc-title';
 import type { DocMode } from '@blocksuite/affine/model';
-import type { Store } from '@blocksuite/affine/store';
+import type { ExtensionType, Store } from '@blocksuite/affine/store';
 import {
+  type FrameworkProvider,
   useFramework,
   useLiveData,
   useService,
@@ -42,7 +44,6 @@ import {
   type DefaultOpenProperty,
   DocPropertiesTable,
 } from '../../components/doc-properties';
-import { enableEditorExtension } from '../extensions/entry/enable-editor';
 import { BiDirectionalLinkPanel } from './bi-directional-link-panel';
 import { BlocksuiteEditorJournalDocTitle } from './journal-doc-title';
 import { StarterBar } from './starter-bar';
@@ -301,3 +302,20 @@ export const BlocksuiteEdgelessEditor = forwardRef<
     </div>
   );
 });
+
+function enableEditorExtension(
+  framework: FrameworkProvider,
+  mode: 'edgeless' | 'page',
+  enableAI: boolean,
+  options: AffineEditorViewOptions
+): ExtensionType[] {
+  const manager = getViewManager(framework, enableAI, options);
+  if (BUILD_CONFIG.isMobileEdition) {
+    if (mode === 'page') {
+      return manager.get('mobile-page');
+    }
+
+    return manager.get('mobile-edgeless');
+  }
+  return manager.get(mode);
+}

packages/frontend/core/src/blocksuite/extensions/entry/enable-editor.ts

@@ -1,22 +0,0 @@
-import type { AffineEditorViewOptions } from '@affine/core/blocksuite/manager/editor-view';
-import type { ExtensionType } from '@blocksuite/affine/store';
-import { type FrameworkProvider } from '@toeverything/infra';
-
-import { getViewManager } from '../../manager/migrating-view';
-
-export function enableEditorExtension(
-  framework: FrameworkProvider,
-  mode: 'edgeless' | 'page',
-  enableAI: boolean,
-  options: AffineEditorViewOptions
-): ExtensionType[] {
-  const manager = getViewManager(framework, enableAI, options);
-  if (BUILD_CONFIG.isMobileEdition) {
-    if (mode === 'page') {
-      return manager.get('mobile-page');
-    }
-
-    return manager.get('mobile-edgeless');
-  }
-  return manager.get(mode);
-}

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

@@ -3,9 +3,8 @@ import {
   AwarenessStore,
   type Doc,
   type ExtensionType,
-  type GetBlocksOptions,
-  type Query,
-  Store,
+  type GetStoreOptions,
+  StoreContainer,
   type Workspace,
   type YBlock,
 } from '@blocksuite/affine/store';
@@ -21,7 +20,7 @@ type DocOptions = {
 export class DocImpl implements Doc {
   private readonly _collection: Workspace;
 
-  private readonly _storeMap = new Map<string, Store>();
+  private readonly _storeContainer: StoreContainer;
 
   private readonly _initSpaceDoc = () => {
     {
@@ -105,19 +104,15 @@ export class DocImpl implements Doc {
 
     this._yBlocks = this._ySpaceDoc.getMap('blocks');
     this._collection = collection;
-  }
-
-  private _getReadonlyKey(readonly?: boolean): 'true' | 'false' {
-    return (readonly?.toString() as 'true' | 'false') ?? 'false';
+    this._storeContainer = new StoreContainer(this);
   }
 
   clear() {
     this._yBlocks.clear();
   }
 
-  clearQuery(query: Query, readonly?: boolean) {
-    const key = this._getQueryKey({ readonly, query });
-    this._storeMap.delete(key);
+  get removeStore() {
+    return this._storeContainer.removeStore;
   }
 
   private _destroy() {
@@ -134,58 +129,26 @@ export class DocImpl implements Doc {
     }
   }
 
-  private readonly _getQueryKey = (
-    idOrOptions: string | { readonly?: boolean; query?: Query }
-  ) => {
-    if (typeof idOrOptions === 'string') {
-      return idOrOptions;
-    }
-    const { readonly, query } = idOrOptions;
-    const readonlyKey = this._getReadonlyKey(readonly);
-    const key = JSON.stringify({
-      readonlyKey,
-      query,
-    });
-    return key;
-  };
-
   getStore({
     readonly,
     query,
     provider,
     extensions,
     id,
-  }: GetBlocksOptions = {}) {
-    let idOrOptions: string | { readonly?: boolean; query?: Query };
-    if (readonly || query) {
-      idOrOptions = { readonly, query };
-    } else if (!id) {
-      idOrOptions = this.workspace.idGenerator();
-    } else {
-      idOrOptions = id;
-    }
-    const key = this._getQueryKey(idOrOptions);
-
-    if (this._storeMap.has(key)) {
-      return this._storeMap.get(key) as Store;
-    }
-
+  }: GetStoreOptions = {}) {
     const storeExtensions = getStoreManager().get('store');
-    const extensionSet = new Set(
-      storeExtensions.concat(extensions ?? []).concat(this.storeExtensions)
-    );
+    const exts = storeExtensions
+      .concat(extensions ?? [])
+      .concat(this.storeExtensions);
+    const extensionSet = new Set(exts);
 
-    const doc = new Store({
-      doc: this,
+    return this._storeContainer.getStore({
+      id,
       readonly,
       query,
       provider,
       extensions: Array.from(extensionSet),
     });
-
-    this._storeMap.set(key, doc);
-
-    return doc;
   }
 
   load(initFn?: () => void): this {