refactor(edgeless): move element tree utilities to std

toeverything · Sep 9, 2024 · 818cecd · 818cecd
1 parent 3f2041c
commit 818cecd
Show file tree

Hide file tree

Showing 15 changed files with 353 additions and 184 deletions.
diff --git a/packages/affine/model/src/blocks/frame/frame-model.ts b/packages/affine/model/src/blocks/frame/frame-model.ts
@@ -71,11 +71,10 @@ 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(elementId: string): void {
     this.doc.transact(() => {
       if (!this.childElementIds) this.childElementIds = {};
-      this.childElementIds[id] = true;
+      this.childElementIds[elementId] = true;
     });
   }
 
@@ -100,10 +99,9 @@ export class FrameBlockModel
     );
   }
 
-  removeChild(element: BlockSuite.EdgelessModel | string): void {
-    const id = typeof element === 'string' ? element : element.id;
+  removeChild(elementId: string): void {
     this.doc.transact(() => {
-      this.childElementIds && delete this.childElementIds[id];
+      this.childElementIds && delete this.childElementIds[elementId];
     });
   }
 }

diff --git a/packages/affine/model/src/elements/group/group.ts b/packages/affine/model/src/elements/group/group.ts
@@ -58,13 +58,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(elementId: string) {
     this.surface.doc.transact(() => {
-      this.children.set(id, true);
+      this.children.set(elementId, true);
     });
   }
 
@@ -80,13 +76,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(elementId: string) {
     if (!this.children) {
       return;
     }
     this.surface.doc.transact(() => {
-      this.children.delete(id);
+      this.children.delete(elementId);
     });
   }
 

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,7 @@ import {
   renderableInEdgeless,
 } from '@blocksuite/affine-block-surface';
 import { GroupElementModel } from '@blocksuite/affine-model';
-import { isGfxContainerElm } from '@blocksuite/block-std/gfx';
+import { type GfxModel, isGfxContainerElm } from '@blocksuite/block-std/gfx';
 import { BlockSuiteError, ErrorCode } from '@blocksuite/global/exceptions';
 import {
   Bound,
@@ -22,7 +22,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 +32,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;
@@ -67,14 +66,16 @@ export class FrameOverlay extends Overlay {
     this._frame = frame;
 
     if (highlightInnerElements) {
+      const treeManager = this._edgelessService.surface.tree;
+
       const innerElements = new Set(
-        getTopElements(
-          this._frameManager.getElementsInFrameBound(frame)
-        ).concat(
-          this._frameManager.getChildElementsInFrame(frame).filter(child => {
-            return frame.intersectsBound(child.elementBound);
-          })
-        )
+        treeManager
+          .getTopElements(this._frameManager.getElementsInFrameBound(frame))
+          .concat(
+            this._frameManager.getChildElementsInFrame(frame).filter(child => {
+              return frame.intersectsBound(child.elementBound);
+            })
+          )
       );
 
       this._innerElements = [...innerElements];
@@ -116,8 +117,12 @@ export class EdgelessFrameManager {
     return this._rootService.frames;
   }
 
+  get tree() {
+    return this._rootService.surface.tree;
+  }
+
   constructor(private _rootService: EdgelessRootService) {
-    this._watchElementAddedOrDeleted();
+    this._watchElementAdded();
   }
 
   private _addFrameBlock(bound: Bound) {
@@ -145,13 +150,23 @@ 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 &&
+            this.tree.hasChild(element, frame)
+          ) {
+            return;
+          }
+
           // TODO(@L-Sun): refactor this in a tree manager
           if (element instanceof GroupElementModel) {
             if (frame && element.hasChild(frame)) return;
@@ -208,10 +223,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 +235,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 => {
+      this.tree.addChildToContainer(frame, element);
     });
   }
 
@@ -250,7 +245,7 @@ export class EdgelessFrameManager {
 
     this.addElementsToFrame(
       frameModel,
-      getTopElements(this.getElementsInFrameBound(frameModel))
+      this.tree.getTopElements(this.getElementsInFrameBound(frameModel))
     );
 
     this._rootService.doc.captureSync();
@@ -263,7 +258,7 @@ export class EdgelessFrameManager {
     return frameModel;
   }
 
-  createFrameOnElements(elements: BlockSuite.EdgelessModel[]) {
+  createFrameOnElements(elements: GfxModel[]) {
     let bound = edgelessElementsBound(
       this._rootService.selection.selectedElements
     );
@@ -279,7 +274,7 @@ export class EdgelessFrameManager {
 
     const frameModel = this._addFrameBlock(bound);
 
-    this.addElementsToFrame(frameModel, getTopElements(elements));
+    this.addElementsToFrame(frameModel, this.tree.getTopElements(elements));
 
     this._rootService.doc.captureSync();
 
@@ -318,7 +313,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 +333,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 +353,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 = this.tree.getContainer(element);
+    return container && isFrameBlock(container) ? container : null;
   }
 
   removeAllChildrenFromFrame(frame: FrameBlockModel) {
@@ -370,11 +364,10 @@ 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
-      parentFrame.removeChild(element);
+      this.tree.removeChildFromContainer(element);
     }
   }
 }

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);
         });

diff --git a/packages/blocks/src/root-block/edgeless/tools/frame-tool.ts b/packages/blocks/src/root-block/edgeless/tools/frame-tool.ts
@@ -3,10 +3,10 @@ import type { PointerEventState } from '@blocksuite/block-std';
 import type { IPoint, IVec } from '@blocksuite/global/utils';
 
 import { TelemetryProvider } from '@blocksuite/affine-shared/services';
+import { getTopElements } from '@blocksuite/block-std/gfx';
 import { Bound, noop, Vec } from '@blocksuite/global/utils';
 import { DocCollection } from '@blocksuite/store';
 
-import { getTopElements } from '../utils/tree.js';
 import { EdgelessToolController } from './edgeless-tool.js';
 
 type FrameTool = {