refactor(edgeless): element tree manager

packages/affine/model/src/blocks/frame/frame-model.ts

@@ -1,11 +1,15 @@
+import type {
+  GfxBlockElementModel,
+  GfxContainerElement,
+  GfxElementGeometry,
+  GfxModel,
+  PointTestOptions,
+} from '@blocksuite/block-std/gfx';
+
 import {
-  type GfxBlockElementModel,
-  type GfxContainerElement,
+  descendantElementsImpl,
   gfxContainerSymbol,
-  type GfxElementGeometry,
-  type GfxModel,
-  type PointTestOptions,
-  SurfaceBlockModel,
+  hasDescendantElementImpl,
 } from '@blocksuite/block-std/gfx';
 import { Bound, type SerializedXYWH } from '@blocksuite/global/utils';
 import { BlockModel, defineBlockSchema, type Text } from '@blocksuite/store';
@@ -49,17 +53,14 @@ export class FrameBlockModel
   [gfxContainerSymbol] = true as const;
 
   get childElements() {
-    const surface = this.doc
-      .getBlocks()
-      .find(model => model instanceof SurfaceBlockModel);
-    if (!surface) return [];
+    if (!this.surface) return [];
 
-    const elements: BlockSuite.EdgelessModel[] = [];
+    const elements: GfxModel[] = [];
 
     for (const key of this.childIds) {
       const element =
-        surface.getElementById(key) ||
-        (surface.doc.getBlockById(key) as GfxBlockElementModel);
+        this.surface.getElementById(key) ||
+        (this.surface.doc.getBlockById(key) as GfxBlockElementModel);
 
       element && elements.push(element);
     }
@@ -71,10 +72,13 @@ export class FrameBlockModel
     return [...(this.childElementIds ? Object.keys(this.childElementIds) : [])];
   }
 
-  addChild(element: BlockSuite.EdgelessModel | string): void {
-    const id = typeof element === 'string' ? element : element.id;
+  get descendantElements(): GfxModel[] {
+    return descendantElementsImpl(this);
+  }
+
+  addChild(element: GfxModel) {
     this.doc.transact(() => {
-      this.childElementIds = { ...this.childElementIds, [id]: true };
+      this.childElementIds = { ...this.childElementIds, [element.id]: true };
     });
   }
 
@@ -99,9 +103,12 @@ export class FrameBlockModel
     return this.elementBound.contains(bound);
   }
 
-  hasDescendant(element: string | GfxModel): boolean {
-    const id = typeof element === 'string' ? element : element.id;
-    return !!this.childElementIds?.[id];
+  hasChild(element: GfxModel): boolean {
+    return this.childElementIds ? element.id in this.childElementIds : false;
+  }
+
+  hasDescendant(element: GfxModel): boolean {
+    return hasDescendantElementImpl(this, element);
   }
 
   override includesPoint(x: number, y: number, _: PointTestOptions): boolean {
@@ -116,10 +123,9 @@ export class FrameBlockModel
     );
   }
 
-  removeChild(element: BlockSuite.EdgelessModel | string): void {
-    const id = typeof element === 'string' ? element : element.id;
+  removeChild(element: GfxModel): void {
     this.doc.transact(() => {
-      this.childElementIds && delete this.childElementIds[id];
+      this.childElementIds && delete this.childElementIds[element.id];
     });
   }
 }

packages/affine/model/src/elements/group/group.ts

@@ -1,7 +1,9 @@
 import type {
   BaseElementProps,
+  GfxModel,
   SerializedElement,
 } from '@blocksuite/block-std/gfx';
+import type { IVec, PointLocation } from '@blocksuite/global/utils';
 import type { Y } from '@blocksuite/store';
 
 import {
@@ -10,13 +12,7 @@ import {
   local,
   observe,
 } from '@blocksuite/block-std/gfx';
-import {
-  Bound,
-  type IVec,
-  keys,
-  linePolygonIntersects,
-  type PointLocation,
-} from '@blocksuite/global/utils';
+import { Bound, keys, linePolygonIntersects } from '@blocksuite/global/utils';
 import { DocCollection } from '@blocksuite/store';
 
 type GroupElementProps = BaseElementProps & {
@@ -58,13 +54,9 @@ export class GroupElementModel extends GfxGroupLikeElementModel<GroupElementProp
     return props as GroupElementProps;
   }
 
-  addChild(element: BlockSuite.EdgelessModel | string) {
-    const id = typeof element === 'string' ? element : element.id;
-    if (!this.children) {
-      return;
-    }
+  override addChild(element: GfxModel) {
     this.surface.doc.transact(() => {
-      this.children.set(id, true);
+      this.children.set(element.id, true);
     });
   }
 
@@ -80,13 +72,12 @@ export class GroupElementModel extends GfxGroupLikeElementModel<GroupElementProp
     return linePolygonIntersects(start, end, bound.points);
   }
 
-  removeChild(element: BlockSuite.EdgelessModel | string) {
-    const id = typeof element === 'string' ? element : element.id;
+  removeChild(element: GfxModel) {
     if (!this.children) {
       return;
     }
     this.surface.doc.transact(() => {
-      this.children.delete(id);
+      this.children.delete(element.id);
     });
   }
 

packages/affine/model/src/elements/mindmap/mindmap.ts

@@ -1,5 +1,6 @@
 import type {
   BaseElementProps,
+  GfxModel,
   SerializedElement,
 } from '@blocksuite/block-std/gfx';
 import type { SerializedXYWH, XYWH } from '@blocksuite/global/utils';
@@ -16,6 +17,7 @@ import {
   deserializeXYWH,
   keys,
   last,
+  noop,
   pick,
 } from '@blocksuite/global/utils';
 import { DocCollection, type Y } from '@blocksuite/store';
@@ -227,6 +229,14 @@ export class MindmapElementModel extends GfxGroupLikeElementModel<MindmapElement
     return { outdated: true, cacheKey };
   }
 
+  /**
+   * @deprecated
+   * you should not call this method directly
+   */
+  addChild(_element: GfxModel) {
+    noop();
+  }
+
   protected addConnector(
     from: MindmapNode,
     to: MindmapNode,
@@ -635,24 +645,24 @@ export class MindmapElementModel extends GfxGroupLikeElementModel<MindmapElement
     this.requestBuildTree();
   }
 
-  removeChild(id: string) {
-    if (!this._nodeMap.has(id)) {
+  removeChild(element: GfxModel) {
+    if (!this._nodeMap.has(element.id)) {
       return;
     }
 
     const surface = this.surface;
     const removedDescendants: string[] = [];
-    const remove = (element: MindmapNode) => {
-      element.children?.forEach(child => {
+    const remove = (node: MindmapNode) => {
+      node.children?.forEach(child => {
         remove(child);
       });
 
-      this.children?.delete(element.id);
-      removedDescendants.push(element.id);
+      this.children?.delete(node.id);
+      removedDescendants.push(node.id);
     };
 
     surface.doc.transact(() => {
-      remove(this._nodeMap.get(id)!);
+      remove(this._nodeMap.get(element.id)!);
     });
 
     queueMicrotask(() => {

packages/blocks/src/root-block/edgeless/components/auto-complete/auto-complete-panel.ts

@@ -1,7 +1,6 @@
 import type {
   Connection,
   ConnectorElementModel,
-  NoteBlockModel,
   ShapeElementModel,
 } from '@blocksuite/affine-model';
 import type { XYWH } from '@blocksuite/global/utils';
@@ -25,6 +24,7 @@ import {
   FontWeight,
   getShapeName,
   GroupElementModel,
+  NoteBlockModel,
   ShapeStyle,
   TextElementModel,
 } from '@blocksuite/affine-model';
@@ -184,11 +184,13 @@ export class EdgelessAutoCompletePanel extends WithDisposable(LitElement) {
       },
       doc.root?.id
     );
+    const note = doc.getBlock(id)?.model;
+    assertInstanceOf(note, NoteBlockModel);
     doc.addBlock('affine:paragraph', { type: 'text' }, id);
     const group = this.currentSource.group;
 
     if (group instanceof GroupElementModel) {
-      group.addChild(id);
+      group.addChild(note);
     }
     this.connector.target = {
       id,
@@ -245,11 +247,15 @@ export class EdgelessAutoCompletePanel extends WithDisposable(LitElement) {
         y: bound.y,
       });
       if (!textId) return;
+
+      const textElement = edgelessService.getElementById(textId);
+      if (!textElement) return;
+
       edgelessService.updateElement(this.connector.id, {
         target: { id: textId, position },
       });
       if (this.currentSource.group instanceof GroupElementModel) {
-        this.currentSource.group.addChild(textId);
+        this.currentSource.group.addChild(textElement);
       }
 
       this.edgeless.service.selection.set({
@@ -275,7 +281,7 @@ export class EdgelessAutoCompletePanel extends WithDisposable(LitElement) {
         target: { id: textId, position },
       });
       if (this.currentSource.group instanceof GroupElementModel) {
-        this.currentSource.group.addChild(textId);
+        this.currentSource.group.addChild(textElement);
       }
 
       this.edgeless.service.selection.set({

packages/blocks/src/root-block/edgeless/components/auto-complete/utils.ts

@@ -1,3 +1,4 @@
+import type { GfxModel } from '@blocksuite/block-std/gfx';
 import type { XYWH } from '@blocksuite/global/utils';
 
 import {
@@ -16,7 +17,8 @@ import {
   type ShapeName,
   type ShapeStyle,
 } from '@blocksuite/affine-model';
-import { assertExists, Bound } from '@blocksuite/global/utils';
+import { BlockSuiteError, ErrorCode } from '@blocksuite/global/exceptions';
+import { assertType, Bound } from '@blocksuite/global/utils';
 import { DocCollection } from '@blocksuite/store';
 
 import type { EdgelessRootBlockComponent } from '../../edgeless-root-block.js';
@@ -277,12 +279,15 @@ export function createEdgelessElement(
   let id;
   const { service } = edgeless;
 
+  let element: GfxModel | null = null;
+
   if (isShape(current)) {
     id = service.addElement(current.type, {
       ...current.serialize(),
       text: new DocCollection.Y.Text(),
       xywh: bound.serialize(),
     });
+    element = service.getElementById(id);
   } else {
     const { doc } = edgeless;
     id = doc.addBlock(
@@ -295,16 +300,32 @@ export function createEdgelessElement(
       },
       edgeless.model.id
     );
-    const note = doc.getBlockById(id) as NoteBlockModel;
-    assertExists(note);
+    const note = doc.getBlock(id)?.model;
+    if (!note) {
+      throw new BlockSuiteError(
+        ErrorCode.GfxBlockElementError,
+        'Note block is not found after creation'
+      );
+    }
+    assertType<NoteBlockModel>(note);
     doc.updateBlock(note, () => {
       note.edgeless.collapse = true;
     });
     doc.addBlock('affine:paragraph', {}, note.id);
+
+    element = note;
+  }
+
+  if (!element) {
+    throw new BlockSuiteError(
+      ErrorCode.GfxBlockElementError,
+      'Element is not found after creation'
+    );
   }
+
   const group = current.group;
   if (group instanceof GroupElementModel) {
-    group.addChild(id);
+    group.addChild(element);
   }
   return id;
 }
@@ -320,9 +341,10 @@ export function createShapeElement(
     radius: getShapeRadius(targetType),
     text: new DocCollection.Y.Text(),
   });
+  const element = service.getElementById(id);
   const group = current.group;
-  if (group instanceof GroupElementModel) {
-    group.addChild(id);
+  if (group instanceof GroupElementModel && element) {
+    group.addChild(element);
   }
   return id;
 }

packages/blocks/src/root-block/edgeless/components/rects/edgeless-selected-rect.ts

@@ -33,6 +33,7 @@ import {
   requestThrottledConnectedFrame,
   stopPropagation,
 } from '@blocksuite/affine-shared/utils';
+import { getTopElements } from '@blocksuite/block-std/gfx';
 import {
   assertType,
   Bound,
@@ -81,7 +82,6 @@ import {
   isImageBlock,
   isNoteBlock,
 } from '../../utils/query.js';
-import { getTopElements } from '../../utils/tree.js';
 import { HandleDirection } from '../resize/resize-handles.js';
 import { ResizeHandles, type ResizeMode } from '../resize/resize-handles.js';
 import { HandleResizeManager } from '../resize/resize-manager.js';