xyflow/packages/system/src/utils/store.ts

import {
  NodeBase,
  CoordinateExtent,
  Dimensions,
  NodeDimensionUpdate,
  NodeOrigin,
  PanZoomInstance,
  Transform,
  XYPosition,
  XYZPosition,
  ConnectionLookup,
  EdgeBase,
  EdgeLookup,
  InternalNodeBase,
} from '../types';
import { getDimensions, getHandleBounds } from './dom';
import { isNumeric } from './general';
import { getNodePositionWithOrigin } from './graph';

type ParentNodes = Set<string>;

export function updateAbsolutePositions<NodeType extends NodeBase>(
  nodeLookup: Map<string, InternalNodeBase<NodeType>>,
  nodeOrigin: NodeOrigin = [0, 0],
  parentNodes?: ParentNodes
) {
  for (const [, node] of nodeLookup) {
    if (node.parentNode && !nodeLookup.has(node.parentNode)) {
      throw new Error(`Parent node ${node.parentNode} not found`);
    }

    if (node.parentNode || parentNodes?.has(node.id)) {
      const parentNode = node.parentNode ? nodeLookup.get(node.parentNode) : null;
      const { x, y, z } = calculateXYZPosition(
        node,
        nodeLookup,
        {
          ...node.position,
          z: node.internals.z ?? 0,
        },
        parentNode?.origin || nodeOrigin
      );

      const positionChanged = x !== node.internals.positionAbsolute?.x || y !== node.internals.positionAbsolute?.y;
      node.internals.positionAbsolute = positionChanged
        ? {
            x,
            y,
          }
        : node.internals.positionAbsolute;

      node.internals.z = z;

      if (parentNodes?.has(node.id)) {
        node.internals.isParent = true;
      }
    }

    nodeLookup.set(node.id, node);
  }
}

type UpdateNodesOptions<NodeType extends NodeBase> = {
  nodeOrigin?: NodeOrigin;
  elevateNodesOnSelect?: boolean;
  defaults?: Partial<NodeType>;
};

export function adoptUserProvidedNodes<NodeType extends NodeBase>(
  nodes: NodeType[],
  nodeLookup: Map<string, InternalNodeBase<NodeType>>,
  options: UpdateNodesOptions<NodeType> = {
    nodeOrigin: [0, 0] as NodeOrigin,
    elevateNodesOnSelect: true,
    defaults: {},
  }
) {
  const tmpLookup = new Map(nodeLookup);
  nodeLookup.clear();
  const parentNodes: ParentNodes = new Set();
  const selectedNodeZ: number = options?.elevateNodesOnSelect ? 1000 : 0;

  nodes.forEach((n) => {
    const currentStoreNode = tmpLookup.get(n.id);
    if (n.parentNode) {
      parentNodes.add(n.parentNode);
    }
    if (n === currentStoreNode?.internals?.userProvidedNode) {
      nodeLookup.set(n.id, currentStoreNode);
      return currentStoreNode;
    }

    const node: InternalNodeBase<NodeType> = {
      ...options.defaults,
      ...n,
      computed: {
        width: n.computed?.width,
        height: n.computed?.height,
      },
      internals: {
        positionAbsolute: n.position,
        handleBounds: currentStoreNode?.internals?.handleBounds,
        z: (isNumeric(n.zIndex) ? n.zIndex : 0) + (n.selected ? selectedNodeZ : 0),
        userProvidedNode: n,
        isParent: false,
      },
    };

    nodeLookup.set(node.id, node);
  });

  if (parentNodes.size > 0) {
    updateAbsolutePositions(nodeLookup, options.nodeOrigin, parentNodes);
  }
}

function calculateXYZPosition<NodeType extends NodeBase>(
  node: NodeType,
  nodeLookup: Map<string, InternalNodeBase<NodeType>>,
  result: XYZPosition,
  nodeOrigin: NodeOrigin
): XYZPosition {
  if (!node.parentNode) {
    return result;
  }

  const parentNode = nodeLookup.get(node.parentNode)!;
  const { position: parentNodePosition } = getNodePositionWithOrigin(parentNode, parentNode?.origin || nodeOrigin);

  return calculateXYZPosition(
    parentNode,
    nodeLookup,
    {
      x: (result.x ?? 0) + parentNodePosition.x,
      y: (result.y ?? 0) + parentNodePosition.y,
      z: (parentNode.internals.z ?? 0) > (result.z ?? 0) ? parentNode.internals.z ?? 0 : result.z ?? 0,
    },
    parentNode.origin || nodeOrigin
  );
}

export function updateNodeDimensions<NodeType extends NodeBase>(
  updates: Map<string, NodeDimensionUpdate>,
  nodeLookup: Map<string, InternalNodeBase<NodeType>>,
  domNode: HTMLElement | null,
  nodeOrigin?: NodeOrigin,
  onUpdate?: (id: string, dimensions: Dimensions) => void
): { hasUpdate: boolean } {
  const viewportNode = domNode?.querySelector('.xyflow__viewport');
  let hasUpdate = false;

  if (!viewportNode) {
    return { hasUpdate };
  }

  const style = window.getComputedStyle(viewportNode);
  const { m22: zoom } = new window.DOMMatrixReadOnly(style.transform);

  updates.forEach((update) => {
    const node = nodeLookup.get(update.id);

    if (node) {
      const dimensions = getDimensions(update.nodeElement);
      const doUpdate = !!(
        dimensions.width &&
        dimensions.height &&
        (node.computed?.width !== dimensions.width || node.computed?.height !== dimensions.height || update.forceUpdate)
      );

      if (doUpdate) {
        hasUpdate = true;
        onUpdate?.(node.id, dimensions);

        const newNode = {
          ...node,
          computed: {
            ...node.computed,
            ...dimensions,
          },
          internals: {
            ...node.internals,
            handleBounds: {
              source: getHandleBounds('.source', update.nodeElement, zoom, node.origin || nodeOrigin),
              target: getHandleBounds('.target', update.nodeElement, zoom, node.origin || nodeOrigin),
            },
          },
        };
        nodeLookup.set(node.id, newNode);
      }
    }
  });

  return { hasUpdate };
}

export function panBy({
  delta,
  panZoom,
  transform,
  translateExtent,
  width,
  height,
}: {
  delta: XYPosition;
  panZoom: PanZoomInstance | null;
  transform: Transform;
  translateExtent: CoordinateExtent;
  width: number;
  height: number;
}) {
  if (!panZoom || (!delta.x && !delta.y)) {
    return false;
  }

  const nextViewport = panZoom.setViewportConstrained(
    {
      x: transform[0] + delta.x,
      y: transform[1] + delta.y,
      zoom: transform[2],
    },
    [
      [0, 0],
      [width, height],
    ],
    translateExtent
  );

  const transformChanged =
    !!nextViewport &&
    (nextViewport.x !== transform[0] || nextViewport.y !== transform[1] || nextViewport.k !== transform[2]);

  return transformChanged;
}

export function updateConnectionLookup(connectionLookup: ConnectionLookup, edgeLookup: EdgeLookup, edges: EdgeBase[]) {
  connectionLookup.clear();
  edgeLookup.clear();

  for (const edge of edges) {
    const { source, target, sourceHandle = null, targetHandle = null } = edge;

    const sourceKey = `${source}-source-${sourceHandle}`;
    const targetKey = `${target}-target-${targetHandle}`;

    const prevSource = connectionLookup.get(sourceKey) || new Map();
    const prevTarget = connectionLookup.get(targetKey) || new Map();
    const connection = { edgeId: edge.id, source, target, sourceHandle, targetHandle };

    edgeLookup.set(edge.id, edge);
    connectionLookup.set(sourceKey, prevSource.set(`${target}-${targetHandle}`, connection));
    connectionLookup.set(targetKey, prevTarget.set(`${source}-${sourceHandle}`, connection));
  }
}