From e444dd371aa54db7bc258fea5400b4e5f5186fc8 Mon Sep 17 00:00:00 2001
From: Jeff Phillips <jephilli@redhat.com>
Date: Thu, 8 Aug 2024 14:42:16 -0400
Subject: [PATCH] fix(DagreGroupLayout): Keep node order when groups are
 collapsed (#231)

---
 .../module/src/layouts/DagreGroupsLayout.ts   | 26 +++++++++++++------
 1 file changed, 18 insertions(+), 8 deletions(-)

diff --git a/packages/module/src/layouts/DagreGroupsLayout.ts b/packages/module/src/layouts/DagreGroupsLayout.ts
index 261589b6..d0cedb7d 100644
--- a/packages/module/src/layouts/DagreGroupsLayout.ts
+++ b/packages/module/src/layouts/DagreGroupsLayout.ts
@@ -1,5 +1,5 @@
 import * as dagre from '@dagrejs/dagre';
-import { Edge, Graph, GRAPH_LAYOUT_END_EVENT, Layout, Node } from '../types';
+import { Edge, Graph, GRAPH_LAYOUT_END_EVENT, isNode, Layout, Node } from '../types';
 import { BaseLayout, LAYOUT_DEFAULTS } from './BaseLayout';
 import { LayoutLink } from './LayoutLink';
 import { LayoutNode } from './LayoutNode';
@@ -120,7 +120,7 @@ export class DagreGroupsLayout extends BaseLayout implements Layout {
           (n) =>
             n.element.getParent()?.getId() === parentGroup?.id ||
             (!parentGroup && n.element.getParent()?.getId() === graph.getId())
-        );
+        ) as DagreNode[];
         const layerEdges = this.edges.filter(
           (edge) =>
             (layerGroups.find((n) => n.id === edge.sourceNode.id) ||
@@ -129,20 +129,30 @@ export class DagreGroupsLayout extends BaseLayout implements Layout {
               layerNodes.find((n) => n.id === edge.targetNode.id))
         );
 
+        const dagreNodes: DagreNode[] = [];
+
         // Layout any child groups first
         layerGroups.forEach((group) => {
           doLayout(group);
 
           // Add the child group node (now with the correct dimensions) to the graph
           const dagreNode = new DagreNode(group.element, group.padding);
-          const updateNode = dagreNode.getUpdatableNode();
-          dagreGraph.setNode(group.id, updateNode);
+          dagreNodes.push(dagreNode);
         });
 
-        layerNodes?.forEach((node) => {
-          const updateNode = (node as DagreNode).getUpdatableNode();
-          dagreGraph.setNode(node.id, updateNode);
-        });
+        dagreNodes.push(...layerNodes);
+
+        // Set the nodes in the graph in the same order give to maintain ordering when groups are collapsed
+        this.graph
+          .getController()
+          .getElements()
+          .filter((e) => isNode(e))
+          .forEach((node) => {
+            const updateNode = dagreNodes.find((dagreNode) => dagreNode.id === node.getId());
+            if (updateNode) {
+              dagreGraph.setNode(updateNode.id, updateNode.getUpdatableNode());
+            }
+          });
 
         layerEdges?.forEach((dagreEdge) => {
           dagreGraph.setEdge(dagreEdge.source.id, dagreEdge.target.id, dagreEdge);