feat(editor): type safe draft model and transformer (#10486)

blocksuite/framework/store/src/adapter/base.ts

@@ -1,6 +1,11 @@
 import { BlockSuiteError } from '@blocksuite/global/exceptions';
 
-import type { DraftModel, Store } from '../model/index.js';
+import {
+  BlockModel,
+  type DraftModel,
+  type Store,
+  toDraftModel,
+} from '../model/index.js';
 import type { AssetsManager } from '../transformer/assets.js';
 import type { Slice, Transformer } from '../transformer/index.js';
 import type {
@@ -72,9 +77,11 @@ export abstract class BaseAdapter<AdapterTarget = unknown> {
     this.job = job;
   }
 
-  async fromBlock(model: DraftModel) {
+  async fromBlock(model: BlockModel | DraftModel) {
     try {
-      const blockSnapshot = this.job.blockToSnapshot(model);
+      const draftModel =
+        model instanceof BlockModel ? toDraftModel(model) : model;
+      const blockSnapshot = this.job.blockToSnapshot(draftModel);
       if (!blockSnapshot) return;
       return await this.fromBlockSnapshot({
         snapshot: blockSnapshot,

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

@@ -4,12 +4,16 @@ type PropsInDraft = 'version' | 'flavour' | 'role' | 'id' | 'keys' | 'text';
 
 type ModelProps<Model> = Model extends BlockModel<infer U> ? U : never;
 
+const draftModelSymbol = Symbol('draftModel');
+
 export type DraftModel<Model extends BlockModel = BlockModel> = Pick<
   Model,
   PropsInDraft
 > & {
   children: DraftModel[];
-} & ModelProps<Model>;
+} & ModelProps<Model> & {
+    [draftModelSymbol]: true;
+  };
 
 export function toDraftModel<Model extends BlockModel = BlockModel>(
   origin: Model

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

@@ -16,7 +16,6 @@ import {
   type BlockModel,
   type BlockOptions,
   type BlockProps,
-  type DraftModel,
 } from '../block/index.js';
 import type { Doc } from '../doc.js';
 import { DocCRUD } from './crud.js';
@@ -100,10 +99,10 @@ export class Store {
   };
 
   updateBlock: {
-    <T extends Partial<BlockProps>>(model: BlockModel, props: T): void;
-    (model: BlockModel, callback: () => void): void;
+    <T extends Partial<BlockProps>>(model: BlockModel | string, props: T): void;
+    (model: BlockModel | string, callback: () => void): void;
   } = (
-    model: BlockModel,
+    modelOrId: BlockModel | string,
     callBackOrProps: (() => void) | Partial<BlockProps>
   ) => {
     if (this.readonly) {
@@ -113,6 +112,17 @@ export class Store {
 
     const isCallback = typeof callBackOrProps === 'function';
 
+    const model =
+      typeof modelOrId === 'string'
+        ? this.getBlock(modelOrId)?.model
+        : modelOrId;
+    if (!model) {
+      throw new BlockSuiteError(
+        ErrorCode.ModelCRUDError,
+        `updating block: ${modelOrId} not found`
+      );
+    }
+
     if (!isCallback) {
       const parent = this.getParent(model);
       this.schema.validate(
@@ -549,7 +559,7 @@ export class Store {
   }
 
   deleteBlock(
-    model: DraftModel,
+    model: BlockModel | string,
     options: {
       bringChildrenTo?: BlockModel;
       deleteChildren?: boolean;
@@ -575,7 +585,10 @@ export class Store {
     };
 
     this.transact(() => {
-      this._crud.deleteBlock(model.id, opts);
+      this._crud.deleteBlock(
+        typeof model === 'string' ? model : model.id,
+        opts
+      );
     });
   }
 

blocksuite/framework/store/src/transformer/middleware.ts

@@ -1,6 +1,6 @@
 import type { Slot } from '@blocksuite/global/utils';
 
-import type { DraftModel, Store } from '../model/index.js';
+import type { BlockModel, DraftModel, Store } from '../model/index.js';
 import type { AssetsManager } from './assets.js';
 import type { Slice } from './slice.js';
 import type {
@@ -48,7 +48,7 @@ export type BeforeExportPayload =
       type: 'info';
     };
 
-export type FinalPayload =
+export type AfterExportPayload =
   | {
       snapshot: BlockSnapshot;
       type: 'block';
@@ -71,11 +71,34 @@ export type FinalPayload =
       type: 'info';
     };
 
+export type AfterImportPayload =
+  | {
+      snapshot: BlockSnapshot;
+      type: 'block';
+      model: BlockModel;
+      parent?: string;
+      index?: number;
+    }
+  | {
+      snapshot: DocSnapshot;
+      type: 'page';
+      page: Store;
+    }
+  | {
+      snapshot: SliceSnapshot;
+      type: 'slice';
+      slice: Slice;
+    }
+  | {
+      snapshot: CollectionInfoSnapshot;
+      type: 'info';
+    };
+
 export type TransformerSlots = {
   beforeImport: Slot<BeforeImportPayload>;
-  afterImport: Slot<FinalPayload>;
+  afterImport: Slot<AfterImportPayload>;
   beforeExport: Slot<BeforeExportPayload>;
-  afterExport: Slot<FinalPayload>;
+  afterExport: Slot<AfterExportPayload>;
 };
 
 type TransformerMiddlewareOptions = {

blocksuite/framework/store/src/transformer/slice.ts

@@ -1,4 +1,9 @@
-import type { DraftModel, Store } from '../model/index';
+import {
+  BlockModel,
+  type DraftModel,
+  type Store,
+  toDraftModel,
+} from '../model/index';
 
 type SliceData = {
   content: DraftModel[];
@@ -21,9 +26,15 @@ export class Slice {
 
   constructor(readonly data: SliceData) {}
 
-  static fromModels(doc: Store, models: DraftModel[]) {
+  static fromModels(doc: Store, models: DraftModel[] | BlockModel[]) {
+    const draftModels = models.map(model => {
+      if (model instanceof BlockModel) {
+        return toDraftModel(model);
+      }
+      return model;
+    });
     return new Slice({
-      content: models,
+      content: draftModels,
       workspaceId: doc.workspace.id,
       pageId: doc.id,
     });

blocksuite/framework/store/src/transformer/transformer.ts

@@ -1,19 +1,21 @@
 import { BlockSuiteError, ErrorCode } from '@blocksuite/global/exceptions';
 import { nextTick, Slot } from '@blocksuite/global/utils';
 
-import type {
+import {
   BlockModel,
-  BlockSchemaType,
-  DraftModel,
-  Store,
+  type BlockSchemaType,
+  type DraftModel,
+  type Store,
+  toDraftModel,
 } from '../model/index.js';
 import type { Schema } from '../schema/index.js';
 import { AssetsManager } from './assets.js';
 import { BaseBlockTransformer } from './base.js';
 import type {
+  AfterExportPayload,
+  AfterImportPayload,
   BeforeExportPayload,
   BeforeImportPayload,
-  FinalPayload,
   TransformerMiddleware,
   TransformerSlots,
 } from './middleware.js';
@@ -66,14 +68,19 @@ export class Transformer {
 
   private readonly _slots: TransformerSlots = {
     beforeImport: new Slot<BeforeImportPayload>(),
-    afterImport: new Slot<FinalPayload>(),
+    afterImport: new Slot<AfterImportPayload>(),
     beforeExport: new Slot<BeforeExportPayload>(),
-    afterExport: new Slot<FinalPayload>(),
+    afterExport: new Slot<AfterExportPayload>(),
   };
 
-  blockToSnapshot = (model: DraftModel): BlockSnapshot | undefined => {
+  blockToSnapshot = (
+    model: DraftModel | BlockModel
+  ): BlockSnapshot | undefined => {
     try {
-      const snapshot = this._blockToSnapshot(model);
+      const draftModel =
+        model instanceof BlockModel ? toDraftModel(model) : model;
+
+      const snapshot = this._blockToSnapshot(draftModel);
 
       if (!snapshot) {
         return;
@@ -103,7 +110,7 @@ export class Transformer {
           'Root block not found in doc'
         );
       }
-      const blocks = this.blockToSnapshot(rootModel);
+      const blocks = this.blockToSnapshot(toDraftModel(rootModel));
       if (!blocks) {
         return;
       }
@@ -286,7 +293,8 @@ export class Transformer {
 
       const contentBlocks = blockTree.children
         .map(tree => doc.getBlockById(tree.draft.id))
-        .filter(Boolean) as DraftModel[];
+        .filter((x): x is BlockModel => x !== null)
+        .map(model => toDraftModel(model));
 
       const slice = new Slice({
         content: contentBlocks,