From c31ff5cd8395be5eba93c5285ad3dd3bc4bd12ac Mon Sep 17 00:00:00 2001
From: Yanyan-Wang <yanyanwang93@gmail.com>
Date: Thu, 2 Nov 2023 15:57:34 +0800
Subject: [PATCH] feat: pin node with menu

---
 .../components/PinNodeWithMenu/Component.tsx  | 34 ++++++------
 .../src/components/common/handlePinNode.tsx   | 55 ++++++++-----------
 2 files changed, 41 insertions(+), 48 deletions(-)

diff --git a/packages/gi-assets-basic/src/components/PinNodeWithMenu/Component.tsx b/packages/gi-assets-basic/src/components/PinNodeWithMenu/Component.tsx
index d79ac73ec..973ea4a8f 100644
--- a/packages/gi-assets-basic/src/components/PinNodeWithMenu/Component.tsx
+++ b/packages/gi-assets-basic/src/components/PinNodeWithMenu/Component.tsx
@@ -5,8 +5,8 @@ import $i18n from '../../i18n';
 import { handlePinNode, handleUnPinNode } from '../common/handlePinNode';
 
 // import { INode } from '@antv/g6';
-type INode = any;
 
+const hasPinned = nodeModel => nodeModel?.data.__style?.badgeShapes?.find(badge => badge.name === 'pin');
 export interface PinNodeMenuItemProps {
   contextmenu: any;
   controlledValues?: {
@@ -17,17 +17,17 @@ export interface PinNodeMenuItemProps {
 
 const PinNodeMenuItem: React.FunctionComponent<PinNodeMenuItemProps> = props => {
   const { contextmenu, controlledValues } = props;
-  const { graph, layout, restartForceSimulation, updateHistory } = useContext();
-  const isForce = layout.type === 'graphin-force';
+  const { graph, layout, updateHistory } = useContext();
+  const isForce = layout.type === 'graphin-force' || layout.type === 'force' || layout.type === 'd3force';
 
   const [pinned, setPinned] = useState(false);
 
   const targetNode = contextmenu.item;
 
   const handleLockNode = (propId?: string, action?: 'pin' | 'unpin') => {
-    const target = contextmenu.item;
+    let { id } = contextmenu.item;
     // 仅支持对节点的操作
-    const invalidFromContextMenu = !target || target.destroyed || target.getType?.() !== 'node';
+    const invalidFromContextMenu = !graph.getNodeData(id);
     if (invalidFromContextMenu && !propId) {
       handleUpateHistory(
         undefined,
@@ -39,9 +39,7 @@ const PinNodeMenuItem: React.FunctionComponent<PinNodeMenuItemProps> = props =>
     }
     contextmenu.onClose();
 
-    let item;
     let pinAction;
-    let id;
     if (propId) {
       if (!graph.findById(propId)) {
         handleUpateHistory(
@@ -53,19 +51,23 @@ const PinNodeMenuItem: React.FunctionComponent<PinNodeMenuItemProps> = props =>
         return;
       }
       id = propId;
-      item = graph.findById(propId);
       pinAction = action || 'pin';
     } else {
-      item = target;
-      const model = target.getModel();
-      id = model.id;
-      pinAction = model.pinned ? 'unpin' : 'pin';
-      setPinned(model.pinned);
+      const model = graph.getNodeData(id);
+      const pinned = hasPinned(model);
+      pinAction = pinned ? 'unpin' : 'pin';
+      setPinned(pinned);
     }
     if (pinAction === 'unpin') {
-      handleUnPinNode(item, graph, restartForceSimulation, isForce);
+      handleUnPinNode(id, graph);
     } else {
-      handlePinNode(item, graph, restartForceSimulation, { dragNodeMass: 100000, isForce });
+      handlePinNode(id, graph, { dragNodeMass: 1000 });
+    }
+    if (isForce) {
+      graph.layout({
+        animated: true,
+        presetLayout: {},
+      });
     }
     handleUpateHistory(id, pinAction, true);
   };
@@ -109,7 +111,7 @@ const PinNodeMenuItem: React.FunctionComponent<PinNodeMenuItemProps> = props =>
 
   return (
     <Menu.Item key="lock-node" eventKey="lock-node" onClick={() => handleLockNode()}>
-      {targetNode.pinned
+      {hasPinned(targetNode)
         ? $i18n.get({ id: 'basic.components.PinNodeWithMenu.Component.Unfix', dm: '解除固定' })
         : $i18n.get({ id: 'basic.components.PinNodeWithMenu.Component.FixedNode', dm: '固定节点' })}
     </Menu.Item>
diff --git a/packages/gi-assets-basic/src/components/common/handlePinNode.tsx b/packages/gi-assets-basic/src/components/common/handlePinNode.tsx
index 1d6bba05f..c81c68b0c 100644
--- a/packages/gi-assets-basic/src/components/common/handlePinNode.tsx
+++ b/packages/gi-assets-basic/src/components/common/handlePinNode.tsx
@@ -2,8 +2,15 @@ import { IGraph, icons } from '@antv/gi-sdk';
 
 export const handlePinNode = (itemId, graph, params) => {
   const { x, y, color, dragNodeMass } = params || {};
+  const nodeModel = graph.getNodeData(itemId);
+  if (!nodeModel) return;
 
   try {
+    const { x: ox, y: oy, __style } = nodeModel.data;
+    const pinPosition = {
+      x: isNaN(x) ? ox : x,
+      y: isNaN(y) ? oy : y,
+    };
     const pinBadge = {
       text: icons.pushpin,
       position: 'leftBottom',
@@ -15,17 +22,14 @@ export const handlePinNode = (itemId, graph, params) => {
       [
         {
           id: itemId,
-          data: {
-            x,
-            y,
-          },
+          data: pinPosition,
         },
       ],
       true,
       true,
     );
 
-    const { badgeShapes = [] } = graph.getNodeData(itemId)?.data || {};
+    const { badgeShapes = [] } = __style || {};
     if (!badgeShapes.find(badge => badge.name === 'pin')) {
       badgeShapes.push(pinBadge);
       graph.updateData('node', {
@@ -43,34 +47,21 @@ export const handlePinNode = (itemId, graph, params) => {
   }
 };
 
-export const handleUnPinNode = (target, graph: IGraph, restartForceSimulation, isForce) => {
+export const handleUnPinNode = (itemId, graph: IGraph) => {
   try {
-    const model = target.getModel();
-    // 目前仅支持GraphinNode 节点
-    const isGraphinNode = model.type === 'graphin-circle';
-    if (!isGraphinNode) {
-      return;
-    }
-    const style = model.style || { badges: [] };
-    const badges = [...style.badges];
-    const index = badges.findIndex(({ value }) => value === icons.pushpin);
-    badges.splice(index, 1);
-    // 更改样式
-    graph.updateItem(model.id, {
-      layout: {
-        ...model.layout,
-        force: { mass: null },
-      },
-      mass: 1,
-      pinned: false,
-      style: {
-        badges,
-      },
-    });
-
-    // 重启力导
-    if (isForce) {
-      restartForceSimulation([model], graph);
+    const { badgeShapes = [] } = graph.getNodeData(itemId)?.data?.__style || {};
+    const index = badgeShapes.findIndex(({ name }) => name === 'pin');
+    if (index > -1) {
+      badgeShapes.splice(index, 1);
+      graph.updateData('node', {
+        id: itemId,
+        data: {
+          mass: undefined,
+          __style: {
+            badgeShapes,
+          },
+        },
+      });
     }
   } catch (error) {
     console.log(error);