fix(transforms): fix mistake when connect to a collapsed node

packages/g6/__tests__/bugs/tree-update-collapsed-node.spec.ts

@@ -0,0 +1,41 @@
+import { idOf, treeToGraphData } from '@/src';
+import { createGraph } from '@@/utils';
+
+describe('bug: tree-update-collapsed-node', () => {
+  it('should be collapsed', async () => {
+    const graph = createGraph({
+      animation: false,
+      x: 50,
+      y: 50,
+      data: treeToGraphData({
+        id: 'root',
+        children: [
+          { id: '1-1', style: { collapsed: true }, children: [{ id: '1-1-1' }] },
+          { id: '1-2', children: [{ id: '1-2-1' }] },
+        ],
+      }),
+      layout: {
+        type: 'indented',
+      },
+    });
+
+    await graph.render();
+
+    await expect(graph).toMatchSnapshot(__filename);
+
+    // set parent of 1-2 to 1-1
+    const edges = graph
+      .getEdgeData()
+      .filter((edge) => edge.target === '1-2')
+      .map(idOf);
+    graph.removeEdgeData(edges);
+    graph.updateNodeData([
+      { id: 'root', children: ['1-1'] },
+      { id: '1-1', children: ['1-1-1', '1-2'] },
+    ]);
+    graph.addEdgeData([{ source: '1-1', target: '1-2' }]);
+    await graph.render();
+
+    await expect(graph).toMatchSnapshot(__filename, 'update collapsed node');
+  });
+});

packages/g6/src/runtime/element.ts

@@ -267,6 +267,7 @@ export class ElementController {
     if (!data) return null;
 
     const { dataChanges, drawData } = data;
+    this.markDestroyElement(drawData);
     // 计算样式 / Calculate style
     this.computeStyle();
     // 创建渲染任务 / Create render task
@@ -476,6 +477,22 @@ export class ElementController {
     });
   }
 
+  /**
+   * <zh/> 标记销毁元素
+   *
+   * <en/> mark destroy element
+   * @param data - <zh/> 绘制数据 | <en/> draw data
+   */
+  private markDestroyElement(data: DrawData) {
+    Object.values(data.remove).forEach((elementData) => {
+      elementData.forEach((datum) => {
+        const id = idOf(datum);
+        const element = this.getElement(id);
+        if (element) markToBeDestroyed(element);
+      });
+    });
+  }
+
   private destroyElement(elementType: ElementType, datum: ElementDatum, context: DrawContext) {
     const { stage = 'exit' } = context;
     const id = idOf(datum);
@@ -538,11 +555,7 @@ export class ElementController {
 
     const preprocess = this.computeChangesAndDrawData({ stage: 'collapse', animation })!;
 
-    preprocess.drawData.remove.nodes.forEach((datum) => {
-      const id = idOf(datum);
-      const element = this.getElement(id);
-      if (element) markToBeDestroyed(element);
-    });
+    this.markDestroyElement(preprocess.drawData);
 
     // 进行预布局，计算出所有元素的位置
     // Perform pre-layout to calculate the position of all elements
@@ -552,7 +565,7 @@ export class ElementController {
     // 重新计算数据 / Recalculate data
     const { drawData } = this.computeChangesAndDrawData({ stage: 'collapse', animation })!;
     const { add, remove, update } = drawData;
-
+    this.markDestroyElement(drawData);
     const context = { animation, stage: 'collapse', data: drawData } as const;
 
     this.destroyElements(remove, context);
@@ -647,7 +660,7 @@ export class ElementController {
     });
 
     const { dataChanges, drawData } = this.computeChangesAndDrawData({ stage: 'collapse', animation })!;
-
+    this.markDestroyElement(drawData);
     const { update, remove } = drawData;
     const context = { animation, stage: 'collapse', data: drawData } as const;
 

packages/g6/src/transforms/collapse-expand-node.ts

@@ -46,8 +46,9 @@ export class CollapseExpandNode extends BaseTransform {
   }
 
   public beforeDraw(input: DrawData): DrawData {
+    const { graph, model } = this.context;
     const {
-      add: { nodes: nodesToAdd },
+      add: { nodes: nodesToAdd, edges: edgesToAdd },
       update: { nodes: nodesToUpdate },
     } = input;
     const nodesToCollapse = new Map<ID, NodeData>();
@@ -57,6 +58,14 @@ export class CollapseExpandNode extends BaseTransform {
       if (isCollapsed(node)) nodesToCollapse.set(id, node);
     });
 
+    // 如果创建了一条连接到收起的节点的边，则将其添加到待展开列表
+    // If an edge is created that connects to a collapsed node, add it to the list to be expanded
+    edgesToAdd.forEach((edge) => {
+      if (graph.getElementType(edge.source) !== 'node') return;
+      const source = graph.getNodeData(edge.source);
+      if (isCollapsed(source)) nodesToCollapse.set(edge.source, source);
+    });
+
     nodesToUpdate.forEach((node, id) => {
       const nodeElement = this.getElement(id);
       if (!nodeElement) return;
@@ -74,13 +83,13 @@ export class CollapseExpandNode extends BaseTransform {
     nodesToCollapse.forEach((node, id) => {
       // 将子节点添加到待删除列表，并删除关联的边
       // Add child nodes to the list to be deleted，and delete the associated edges
-      const descendants = this.context.model!.getDescendantsData(id);
+      const descendants = model.getDescendantsData(id);
       descendants.forEach((descendant) => {
         const id = idOf(descendant);
         if (handledNodes.has(id)) return;
 
         reassignTo(input, 'remove', 'node', descendant);
-        const relatedEdges = this.context.model.getRelatedEdgesData(id);
+        const relatedEdges = model.getRelatedEdgesData(id);
         relatedEdges.forEach((edge) => {
           reassignTo(input, 'remove', 'edge', edge);
         });
@@ -90,7 +99,7 @@ export class CollapseExpandNode extends BaseTransform {
     });
 
     nodesToExpand.forEach((node, id) => {
-      const ancestors = this.context.model.getAncestorsData(id, TREE_KEY);
+      const ancestors = model.getAncestorsData(id, TREE_KEY);
 
       // 如果祖先节点是收起的，添加到移除列表
       // If the ancestor node is collapsed, add it to the removal list