Entitree Flex

import type { Edge, Node } from '@xyflow/vue';
import type { TreeNodeData } from './initialElements';
import { Position } from '@xyflow/vue';
import { layoutFromMap } from 'entitree-flex';
 
export type EntitreeNodeData = {
  label: string;
  direction: string;
  isRoot?: boolean;
  isSpouse?: boolean;
  isSibling?: boolean;
  children?: string[];
  siblings?: string[];
  spouses?: string[];
};
 
export type EntitreeNode = Node<EntitreeNodeData, 'custom'>;
 
const nodeWidth = 150;
const nodeHeight = 36;
 
const Orientation = {
  Vertical: 'vertical',
  Horizontal: 'horizontal',
} as const;
 
const entitreeSettings = {
  clone: true, // returns a copy of the input, if your application does not allow editing the original object
  enableFlex: true, // has slightly better performance if turned off (node.width, node.height will not be read)
  firstDegreeSpacing: 100, // spacing between nodes belonging to the same source, e.g. children with same parent
  nextAfterAccessor: 'spouses', // the side node prop used to go sideways, AFTER the current node
  nextAfterSpacing: 100, // the spacing of the "side" nodes AFTER the current node
  nextBeforeAccessor: 'siblings', // the side node prop used to go sideways, BEFORE the current node
  nextBeforeSpacing: 100, // the spacing of the "side" nodes BEFORE the current node
  nodeHeight, // default node height in px
  nodeWidth, // default node width in px
  orientation: Orientation.Vertical, // "vertical" = parents top / children bottom, "horizontal" = parents left
  rootX: 0, // set root position if other than 0
  rootY: 0, // set root position if other than 0
  secondDegreeSpacing: 100, // spacing between nodes not belonging to the same parent, e.g. "cousin" nodes
  sourcesAccessor: 'parents', // the prop used as the array of ancestor ids
  sourceTargetSpacing: 100, // the "vertical" spacing between nodes in vertical orientation, horizontal otherwise
  targetsAccessor: 'children', // the prop used as the array of children ids
};
 
const { Top, Bottom, Left, Right } = Position;
 
export function layoutElements(tree: Record<string, TreeNodeData>, rootId: string | number, direction: 'TB' | 'LR' = 'TB') {
  const isTreeHorizontal = direction === 'LR';
 
  const { nodes: entitreeNodes, rels: entitreeEdges } = layoutFromMap(rootId, tree, {
    ...entitreeSettings,
    orientation: isTreeHorizontal ? Orientation.Horizontal : Orientation.Vertical,
  });
 
  const nodes: EntitreeNode[] = [];
  const edges: Edge[] = [];
 
  entitreeEdges.forEach((edge) => {
    const sourceNode = edge.source.id;
    const targetNode = edge.target.id;
 
    // check if the target node is a spouse or a sibling to pick which handles to connect
    const isTargetSpouse = !!edge.target.isSpouse;
    const isTargetSibling = !!edge.target.isSibling;
 
    let sourceHandle: Position;
    let targetHandle: Position;
 
    if (isTargetSpouse) {
      sourceHandle = isTreeHorizontal ? Bottom : Right;
      targetHandle = isTreeHorizontal ? Top : Left;
    }
    else if (isTargetSibling) {
      sourceHandle = isTreeHorizontal ? Top : Left;
      targetHandle = isTreeHorizontal ? Bottom : Right;
    }
    else {
      sourceHandle = isTreeHorizontal ? Right : Bottom;
      targetHandle = isTreeHorizontal ? Left : Top;
    }
 
    edges.push({
      id: `e${sourceNode}${targetNode}`,
      source: sourceNode,
      target: targetNode,
      type: 'smoothstep',
      animated: true,
      sourceHandle,
      targetHandle,
    });
  });
 
  entitreeNodes.forEach((node) => {
    const isSpouse = !!node.isSpouse;
    const isSibling = !!node.isSibling;
 
    let sourcePosition: Position;
    let targetPosition: Position;
 
    if (isSpouse) {
      sourcePosition = isTreeHorizontal ? Bottom : Right;
      targetPosition = isTreeHorizontal ? Top : Left;
    }
    else if (isSibling) {
      sourcePosition = isTreeHorizontal ? Top : Left;
      targetPosition = isTreeHorizontal ? Bottom : Right;
    }
    else {
      sourcePosition = isTreeHorizontal ? Right : Bottom;
      targetPosition = isTreeHorizontal ? Left : Top;
    }
 
    nodes.push({
      id: node.id,
      type: 'custom',
      position: { x: node.x, y: node.y },
      width: nodeWidth,
      height: nodeHeight,
      sourcePosition,
      targetPosition,
      // `isRoot` (and isSpouse/isSibling/children/…) come from the spread entitree node
      data: { label: node.name, direction, ...node },
    });
  });
 
  return { nodes, edges };
}