feat(editor): selection as store extension (#9605)

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

@@ -1,2 +1,3 @@
 export * from './extension';
+export * from './selection';
 export * from './store-extension';

blocksuite/framework/store/src/extension/selection/base.ts

@@ -0,0 +1,44 @@
+import { BlockSuiteError, ErrorCode } from '@blocksuite/global/exceptions';
+
+import type { SelectionConstructor } from './types';
+
+export type BaseSelectionOptions = {
+  blockId: string;
+};
+
+export abstract class BaseSelection {
+  static readonly group: string;
+
+  static readonly type: string;
+
+  readonly blockId: string;
+
+  get group(): string {
+    return (this.constructor as SelectionConstructor).group;
+  }
+
+  get type(): string {
+    return (this.constructor as SelectionConstructor).type as string;
+  }
+
+  constructor({ blockId }: BaseSelectionOptions) {
+    this.blockId = blockId;
+  }
+
+  static fromJSON(_: Record<string, unknown>): BaseSelection {
+    throw new BlockSuiteError(
+      ErrorCode.SelectionError,
+      'You must override this method'
+    );
+  }
+
+  abstract equals(other: BaseSelection): boolean;
+
+  is<T extends SelectionConstructor>(
+    type: T
+  ): this is T extends SelectionConstructor<infer U> ? U : never {
+    return this.type === type.type;
+  }
+
+  abstract toJSON(): Record<string, unknown>;
+}

blocksuite/framework/store/src/extension/selection/identifier.ts

@@ -0,0 +1,17 @@
+import { createIdentifier } from '@blocksuite/global/di';
+
+import type { ExtensionType } from '../extension';
+import type { SelectionConstructor } from './types';
+
+export const SelectionIdentifier =
+  createIdentifier<SelectionConstructor>('Selection');
+
+export function SelectionExtension(
+  selectionCtor: SelectionConstructor
+): ExtensionType {
+  return {
+    setup: di => {
+      di.addImpl(SelectionIdentifier(selectionCtor.type), () => selectionCtor);
+    },
+  };
+}

blocksuite/framework/store/src/extension/selection/index.ts

@@ -0,0 +1,4 @@
+export * from './base';
+export * from './identifier';
+export * from './selection-extension';
+export * from './types';

blocksuite/framework/store/src/extension/selection/selection-extension.ts

@@ -0,0 +1,187 @@
+import { BlockSuiteError, ErrorCode } from '@blocksuite/global/exceptions';
+import { Slot } from '@blocksuite/global/utils';
+import { computed, signal } from '@preact/signals-core';
+
+import type { Store } from '../../model';
+import { nanoid } from '../../utils/id-generator';
+import type { StackItem } from '../../yjs';
+import { StoreExtension } from '../store-extension';
+import type { BaseSelection } from './base';
+import { SelectionIdentifier } from './identifier';
+import type { SelectionConstructor } from './types';
+
+export class StoreSelectionExtension extends StoreExtension {
+  static override readonly key = 'selection';
+
+  private readonly _id = `${this.store.id}:${nanoid()}`;
+  private _selectionConstructors: Record<string, SelectionConstructor> = {};
+  private readonly _selections = signal<BaseSelection[]>([]);
+  private readonly _remoteSelections = signal<Map<number, BaseSelection[]>>(
+    new Map()
+  );
+
+  private readonly _itemAdded = (event: { stackItem: StackItem }) => {
+    event.stackItem.meta.set('selection-state', this._selections.value);
+  };
+
+  private readonly _itemPopped = (event: { stackItem: StackItem }) => {
+    const selection = event.stackItem.meta.get('selection-state');
+    if (selection) {
+      this.set(selection as BaseSelection[]);
+    }
+  };
+
+  private readonly _jsonToSelection = (json: Record<string, unknown>) => {
+    const ctor = this._selectionConstructors[json.type as string];
+    if (!ctor) {
+      throw new BlockSuiteError(
+        ErrorCode.SelectionError,
+        `Unknown selection type: ${json.type}`
+      );
+    }
+    return ctor.fromJSON(json);
+  };
+
+  slots = {
+    changed: new Slot<BaseSelection[]>(),
+    remoteChanged: new Slot<Map<number, BaseSelection[]>>(),
+  };
+
+  constructor(store: Store) {
+    super(store);
+
+    this.store.provider.getAll(SelectionIdentifier).forEach(ctor => {
+      [ctor].flat().forEach(ctor => {
+        this._selectionConstructors[ctor.type] = ctor;
+      });
+    });
+
+    this.store.awarenessStore.awareness.on(
+      'change',
+      (change: { updated: number[]; added: number[]; removed: number[] }) => {
+        const all = change.updated.concat(change.added).concat(change.removed);
+        const localClientID = this.store.awarenessStore.awareness.clientID;
+        const exceptLocal = all.filter(id => id !== localClientID);
+        const hasLocal = all.includes(localClientID);
+        if (hasLocal) {
+          const localSelectionJson =
+            this.store.awarenessStore.getLocalSelection(this._id);
+          const localSelection = localSelectionJson.map(json => {
+            return this._jsonToSelection(json);
+          });
+          this._selections.value = localSelection;
+        }
+
+        // Only consider remote selections from other clients
+        if (exceptLocal.length > 0) {
+          const map = new Map<number, BaseSelection[]>();
+          this.store.awarenessStore.getStates().forEach((state, id) => {
+            if (id === this.store.awarenessStore.awareness.clientID) return;
+            // selection id starts with the same block collection id from others clients would be considered as remote selections
+            const selection = Object.entries(state.selectionV2)
+              .filter(([key]) => key.startsWith(this.store.id))
+              .flatMap(([_, selection]) => selection);
+
+            const selections = selection
+              .map(json => {
+                try {
+                  return this._jsonToSelection(json);
+                } catch (error) {
+                  console.error(
+                    'Parse remote selection failed:',
+                    id,
+                    json,
+                    error
+                  );
+                  return null;
+                }
+              })
+              .filter((sel): sel is BaseSelection => !!sel);
+
+            map.set(id, selections);
+          });
+          this._remoteSelections.value = map;
+          this.slots.remoteChanged.emit(map);
+        }
+      }
+    );
+
+    this.store.history.on('stack-item-added', this._itemAdded);
+    this.store.history.on('stack-item-popped', this._itemPopped);
+  }
+
+  get value() {
+    return this._selections.value;
+  }
+
+  get remoteSelections() {
+    return this._remoteSelections.value;
+  }
+
+  clear(types?: string[]) {
+    if (types) {
+      const values = this.value.filter(
+        selection => !types.includes(selection.type)
+      );
+      this.set(values);
+    } else {
+      this.set([]);
+    }
+  }
+
+  create<T extends SelectionConstructor>(
+    Type: T,
+    ...args: ConstructorParameters<T>
+  ): InstanceType<T> {
+    return new Type(...args) as InstanceType<T>;
+  }
+
+  getGroup(group: string) {
+    return this.value.filter(s => s.group === group);
+  }
+
+  filter<T extends SelectionConstructor>(type: T) {
+    return this.filter$(type).value;
+  }
+
+  filter$<T extends SelectionConstructor>(type: T) {
+    return computed(() =>
+      this.value.filter((sel): sel is InstanceType<T> => sel.is(type))
+    );
+  }
+
+  find<T extends SelectionConstructor>(type: T) {
+    return this.find$(type).value;
+  }
+
+  find$<T extends SelectionConstructor>(type: T) {
+    return computed(() =>
+      this.value.find((sel): sel is InstanceType<T> => sel.is(type))
+    );
+  }
+
+  set(selections: BaseSelection[]) {
+    this.store.awarenessStore.setLocalSelection(
+      this._id,
+      selections.map(s => s.toJSON())
+    );
+    this.slots.changed.emit(selections);
+  }
+
+  setGroup(group: string, selections: BaseSelection[]) {
+    const current = this.value.filter(s => s.group !== group);
+    this.set([...current, ...selections]);
+  }
+
+  update(fn: (currentSelections: BaseSelection[]) => BaseSelection[]) {
+    const selections = fn(this.value);
+    this.set(selections);
+  }
+
+  fromJSON(json: Record<string, unknown>[]) {
+    const selections = json.map(json => {
+      return this._jsonToSelection(json);
+    });
+    return this.set(selections);
+  }
+}

blocksuite/framework/store/src/extension/selection/types.ts

@@ -0,0 +1,9 @@
+import type { BaseSelection } from './base';
+
+export interface SelectionConstructor<T extends BaseSelection = BaseSelection> {
+  type: string;
+  group: string;
+
+  new (...args: any[]): T;
+  fromJSON(json: Record<string, unknown>): T;
+}

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

@@ -4,6 +4,7 @@ import { type Disposable, Slot } from '@blocksuite/global/utils';
 import { signal } from '@preact/signals-core';
 
 import type { ExtensionType } from '../../extension/extension.js';
+import { StoreSelectionExtension } from '../../extension/index.js';
 import type { Schema } from '../../schema/index.js';
 import {
   Block,
@@ -27,29 +28,33 @@ export type StoreOptions = {
   extensions?: ExtensionType[];
 };
 
+const internalExtensions = [StoreSelectionExtension];
+
 export class Store {
+  readonly userExtensions: ExtensionType[];
+
   private readonly _provider: ServiceProvider;
 
   private readonly _runQuery = (block: Block) => {
     runQuery(this._query, block);
   };
 
-  protected readonly _doc: Doc;
+  private readonly _doc: Doc;
 
-  protected readonly _blocks = signal<Record<string, Block>>({});
+  private readonly _blocks = signal<Record<string, Block>>({});
 
-  protected readonly _crud: DocCRUD;
+  private readonly _crud: DocCRUD;
 
-  protected readonly _disposeBlockUpdated: Disposable;
+  private readonly _disposeBlockUpdated: Disposable;
 
-  protected readonly _query: Query = {
+  private readonly _query: Query = {
     match: [],
     mode: 'loose',
   };
 
-  protected _readonly = signal(false);
+  private readonly _readonly = signal(false);
 
-  protected readonly _schema: Schema;
+  private readonly _schema: Schema;
 
   readonly slots: Doc['slots'] & {
     /** This is always triggered after `doc.load` is called. */
@@ -295,7 +300,12 @@ export class Store {
     const container = new Container();
     container.addImpl(StoreIdentifier, () => this);
 
+    internalExtensions.forEach(ext => {
+      ext.setup(container);
+    });
+
     const userExtensions = extensions ?? [];
+    this.userExtensions = userExtensions;
     userExtensions.forEach(extension => {
       extension.setup(container);
     });