refactor(edgeless): move element tree utilities to std

toeverything · Sep 11, 2024 · a21a0ac · a21a0ac
1 parent ff1da1b
commit a21a0ac
Show file tree

Hide file tree

Showing 17 changed files with 462 additions and 214 deletions.
diff --git a/packages/affine/block-surface/src/element-model/mindmap.ts b/packages/affine/block-surface/src/element-model/mindmap.ts
@@ -1,5 +1,6 @@
 import type {
   BaseElementProps,
+  GfxModel,
   SerializedElement,
   SurfaceBlockModel,
 } from '@blocksuite/block-std/gfx';
@@ -856,24 +857,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(() => {

diff --git a/packages/affine/model/src/blocks/frame/frame-model.ts b/packages/affine/model/src/blocks/frame/frame-model.ts
@@ -1,9 +1,11 @@
 import {
+  getDescendantElementsImpl,
   type GfxBlockElementModel,
   type GfxContainerElement,
   gfxContainerSymbol,
   type GfxElementGeometry,
   type GfxModel,
+  hasDescendantElementImpl,
   type PointTestOptions,
   SurfaceBlockModel,
 } from '@blocksuite/block-std/gfx';
@@ -71,21 +73,27 @@ export class FrameBlockModel
     return [...(this.childElementIds ? Object.keys(this.childElementIds) : [])];
   }
 
-  addChild(element: BlockSuite.EdgelessModel | string): void {
-    const id = typeof element === 'string' ? element : element.id;
+  addChild(element: GfxModel) {
     this.doc.transact(() => {
       if (!this.childElementIds) this.childElementIds = {};
-      this.childElementIds[id] = true;
+      this.childElementIds[element.id] = true;
     });
   }
 
   override containsBound(bound: Bound): boolean {
     return this.elementBound.contains(bound);
   }
 
-  hasDescendant(element: string | GfxModel): boolean {
-    const id = typeof element === 'string' ? element : element.id;
-    return !!this.childElementIds?.[id];
+  getDescendantElements(): GfxModel[] {
+    return getDescendantElementsImpl(this);
+  }
+
+  hasChild(element: GfxModel): boolean {
+    return this.childElementIds ? element.id in this.childElementIds : false;
+  }
+
+  hasDescendantElement(element: GfxModel): boolean {
+    return hasDescendantElementImpl(this, element);
   }
 
   override includesPoint(x: number, y: number, _: PointTestOptions): boolean {
@@ -100,10 +108,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];
     });
   }
 }

diff --git a/packages/affine/model/src/elements/group/group.ts b/packages/affine/model/src/elements/group/group.ts
@@ -1,5 +1,6 @@
 import type {
   BaseElementProps,
+  GfxModel,
   SerializedElement,
 } from '@blocksuite/block-std/gfx';
 import type { Y } from '@blocksuite/store';
@@ -58,13 +59,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 +77,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);
     });
   }
 

diff --git a/packages/blocks/src/root-block/edgeless/components/rects/edgeless-selected-rect.ts b/packages/blocks/src/root-block/edgeless/components/rects/edgeless-selected-rect.ts
@@ -83,7 +83,6 @@ import {
   isImageBlock,
   isNoteBlock,
 } from '../../utils/query.js';
-import { getTopElements } from '../../utils/tree.js';
 import '../auto-complete/edgeless-auto-complete.js';
 import '../connector/connector-handle.js';
 import { HandleDirection } from '../resize/resize-handles.js';
@@ -1081,20 +1080,21 @@ export class EdgelessSelectedRect extends WithDisposable(LitElement) {
 
   #adjustFrame(frame: FrameBlockModel, bound: Bound) {
     const frameManager = this.edgeless.service.frame;
+    const treeManager = this.surface.model.tree;
 
     const oldChildren = frameManager.getChildElementsInFrame(frame);
 
     this.edgeless.service.updateElement(frame.id, {
       xywh: bound.serialize(),
     });
 
-    const newChildren = getTopElements(
-      frameManager.getElementsInFrameBound(frame)
-    ).concat(
-      oldChildren.filter(oldChild => {
-        return frame.intersectsBound(oldChild.elementBound);
-      })
-    );
+    const newChildren = treeManager
+      .getTopElements(frameManager.getElementsInFrameBound(frame))
+      .concat(
+        oldChildren.filter(oldChild => {
+          return frame.intersectsBound(oldChild.elementBound);
+        })
+      );
 
     frameManager.removeAllChildrenFromFrame(frame);
     frameManager.addElementsToFrame(frame, newChildren);

diff --git a/packages/blocks/src/root-block/edgeless/frame-manager.ts b/packages/blocks/src/root-block/edgeless/frame-manager.ts
@@ -7,7 +7,11 @@ import {
   renderableInEdgeless,
 } from '@blocksuite/affine-block-surface';
 import { GroupElementModel } from '@blocksuite/affine-model';
-import { isGfxContainerElm } from '@blocksuite/block-std/gfx';
+import {
+  getTopElements,
+  type GfxModel,
+  isGfxContainerElm,
+} from '@blocksuite/block-std/gfx';
 import { BlockSuiteError, ErrorCode } from '@blocksuite/global/exceptions';
 import {
   Bound,
@@ -22,7 +26,6 @@ import type { EdgelessRootService } from '../../index.js';
 import { GfxBlockModel } from './block-model.js';
 import { edgelessElementsBound } from './utils/bound-utils.js';
 import { isFrameBlock } from './utils/query.js';
-import { getTopElements } from './utils/tree.js';
 
 const MIN_FRAME_WIDTH = 800;
 const MIN_FRAME_HEIGHT = 640;
@@ -33,7 +36,7 @@ export class FrameOverlay extends Overlay {
 
   private _frame: FrameBlockModel | null = null;
 
-  private _innerElements: BlockSuite.EdgelessModel[] = [];
+  private _innerElements: GfxModel[] = [];
 
   private get _frameManager() {
     return this._edgelessService.frame;
@@ -117,7 +120,7 @@ export class EdgelessFrameManager {
   }
 
   constructor(private _rootService: EdgelessRootService) {
-    this._watchElementAddedOrDeleted();
+    this._watchElementAdded();
   }
 
   private _addFrameBlock(bound: Bound) {
@@ -145,13 +148,19 @@ export class EdgelessFrameManager {
     return frameModel;
   }
 
-  private _watchElementAddedOrDeleted() {
+  private _watchElementAdded() {
     this._disposable.add(
       this._rootService.surface.elementAdded.on(({ id, local }) => {
         const element = this._rootService.surface.getElementById(id);
         if (element && local) {
           const frame = this.getFrameFromPoint(element.elementBound.center);
 
+          // if the container created with a frame, skip it.
+          // |<-- Container |< -- frame -->| Container -->|
+          if (isGfxContainerElm(element) && frame && element.hasChild(frame)) {
+            return;
+          }
+
           // TODO(@L-Sun): refactor this in a tree manager
           if (element instanceof GroupElementModel) {
             if (frame && element.hasChild(frame)) return;
@@ -208,10 +217,7 @@ export class EdgelessFrameManager {
   /**
    * Reset parent of elements to the frame
    */
-  addElementsToFrame(
-    frame: FrameBlockModel,
-    elements: BlockSuite.EdgelessModel[]
-  ) {
+  addElementsToFrame(frame: FrameBlockModel, elements: GfxModel[]) {
     if (frame.childElementIds === undefined) {
       elements = [...elements, ...this.getChildElementsInFrame(frame)];
       frame.childElementIds = {};
@@ -223,25 +229,8 @@ export class EdgelessFrameManager {
 
     if (elements.length === 0) return;
 
-    this._rootService.doc.transact(() => {
-      elements.forEach(element => {
-        // TODO(@L-Sun): refactor this. This branch is avoid circle, but it's better to handle in a tree manager
-        if (isGfxContainerElm(element) && element.childIds.includes(frame.id)) {
-          if (isFrameBlock(element)) {
-            this.removeParentFrame(frame);
-          } else if (element instanceof GroupElementModel) {
-            // eslint-disable-next-line unicorn/prefer-dom-node-remove
-            element.removeChild(frame.id);
-          }
-        }
-
-        const parentFrame = this.getParentFrame(element);
-        if (parentFrame) {
-          // eslint-disable-next-line unicorn/prefer-dom-node-remove
-          parentFrame.removeChild(element);
-        }
-        frame.addChild(element);
-      });
+    elements.forEach(element => {
+      frame.addChild(element);
     });
   }
 
@@ -263,7 +252,7 @@ export class EdgelessFrameManager {
     return frameModel;
   }
 
-  createFrameOnElements(elements: BlockSuite.EdgelessModel[]) {
+  createFrameOnElements(elements: GfxModel[]) {
     let bound = edgelessElementsBound(
       this._rootService.selection.selectedElements
     );
@@ -318,7 +307,7 @@ export class EdgelessFrameManager {
    * 1. The frame doesn't have `childElements`, return all elements in the frame bound but not owned by another frame.
    * 2. Return all child elements of the frame if `childElements` exists.
    */
-  getChildElementsInFrame(frame: FrameBlockModel): BlockSuite.EdgelessModel[] {
+  getChildElementsInFrame(frame: FrameBlockModel): GfxModel[] {
     if (frame.childElementIds === undefined) {
       return this.getElementsInFrameBound(frame).filter(
         element => this.getParentFrame(element) !== null
@@ -338,7 +327,7 @@ export class EdgelessFrameManager {
    */
   getElementsInFrameBound(frame: FrameBlockModel, fullyContained = true) {
     const bound = Bound.deserialize(frame.xywh);
-    const elements: BlockSuite.EdgelessModel[] = this._rootService.gfx.grid
+    const elements: GfxModel[] = this._rootService.gfx.grid
       .search(bound, fullyContained)
       .filter(element => element !== frame);
 
@@ -358,10 +347,9 @@ export class EdgelessFrameManager {
     return null;
   }
 
-  getParentFrame(element: BlockSuite.EdgelessModel) {
-    return this.frames.find(frame => {
-      return frame.childIds.includes(element.id);
-    });
+  getParentFrame(element: GfxModel) {
+    const container = element.container;
+    return container && isFrameBlock(container) ? container : null;
   }
 
   removeAllChildrenFromFrame(frame: FrameBlockModel) {
@@ -370,7 +358,7 @@ export class EdgelessFrameManager {
     });
   }
 
-  removeParentFrame(element: BlockSuite.EdgelessModel) {
+  removeParentFrame(element: GfxModel) {
     const parentFrame = this.getParentFrame(element);
     if (parentFrame) {
       // eslint-disable-next-line unicorn/prefer-dom-node-remove

diff --git a/packages/blocks/src/root-block/edgeless/tools/default-tool.ts b/packages/blocks/src/root-block/edgeless/tools/default-tool.ts
@@ -5,7 +5,6 @@ import type {
   NoteBlockModel,
 } from '@blocksuite/affine-model';
 import type { PointerEventState } from '@blocksuite/block-std';
-import type { PointTestOptions } from '@blocksuite/block-std/gfx';
 import type { IVec } from '@blocksuite/global/utils';
 
 import {
@@ -26,6 +25,10 @@ import {
   handleNativeRangeAtPoint,
   resetNativeSelection,
 } from '@blocksuite/affine-shared/utils';
+import {
+  isGfxContainerElm,
+  type PointTestOptions,
+} from '@blocksuite/block-std/gfx';
 import {
   Bound,
   DisposableGroup,
@@ -54,7 +57,6 @@ import {
   mountShapeTextEditor,
   mountTextElementEditor,
 } from '../utils/text.js';
-import { getAllDescendantElements, getTopElements } from '../utils/tree.js';
 import { EdgelessToolController } from './edgeless-tool.js';
 
 export enum DefaultModeDragType {
@@ -223,7 +225,10 @@ export class DefaultToolController extends EdgelessToolController<DefaultTool> {
     const h = Math.abs(startY - curY);
     const bound = new Bound(x, y, w, h);
 
-    const elements = getTopElements(service.gfx.getElementsByBound(bound));
+    const treeManager = this._surface.model.tree;
+    const elements = treeManager.getTopElements(
+      service.gfx.getElementsByBound(bound)
+    );
 
     const set = new Set(
       tools.shiftKey ? [...elements, ...selection.selectedElements] : elements
@@ -923,7 +928,8 @@ export class DefaultToolController extends EdgelessToolController<DefaultTool> {
 
     {
       const frameManager = this._service.frame;
-      const toBeMovedTopElements = getTopElements(this._toBeMoved);
+      const treeManager = this._surface.model.tree;
+      const toBeMovedTopElements = treeManager.getTopElements(this._toBeMoved);
       if (this._hoveredFrame) {
         frameManager.addElementsToFrame(
           this._hoveredFrame,
@@ -1021,17 +1027,18 @@ export class DefaultToolController extends EdgelessToolController<DefaultTool> {
 
     const elements = this.edgelessSelectionManager.selectedElements;
     const toBeMoved = new Set(elements);
+    const tree = this._surface.model.tree;
     elements.forEach(element => {
       if (element.group instanceof MindmapElementModel && elements.length > 1) {
-        getAllDescendantElements(element.group).forEach(ele =>
-          toBeMoved.add(ele)
-        );
-      } else {
-        getAllDescendantElements(element).forEach(ele => {
+        tree
+          .getDescendantElements(element.group)
+          .forEach(ele => toBeMoved.add(ele));
+      } else if (isGfxContainerElm(element)) {
+        tree.getDescendantElements(element).forEach(ele => {
           if (ele.group instanceof MindmapElementModel) {
-            getAllDescendantElements(ele.group).forEach(_el =>
-              toBeMoved.add(_el)
-            );
+            tree
+              .getDescendantElements(ele.group)
+              .forEach(_el => toBeMoved.add(_el));
           }
           toBeMoved.add(ele);
         });