From c277e5eba6054c1006d934ea539c6839c8ef408b Mon Sep 17 00:00:00 2001
From: Patrick Tasse <patrick.tasse@gmail.com>
Date: Thu, 17 Oct 2024 15:54:57 -0400
Subject: [PATCH] Fix time graph filtering and collapsing

Update the filtering of empty nodes to include visible nodes even if any
of their parent is filtered-out. Recursively replace filtered-out nodes
by their visible children, if any, but only if the filtered-out parent
node is not collapsed.

In TimeGraphOutputComponent update the arrows layer after empty nodes
have been determined, and filter out empty nodes from the rowIds passed
to the arrows layer.

Signed-off-by: Patrick Tasse <patrick.tasse@gmail.com>
---
 .../components/timegraph-output-component.tsx |  6 ++---
 .../components/utils/filter-tree/table.tsx    |  2 +-
 .../src/components/utils/filter-tree/tree.tsx |  2 +-
 .../components/utils/filter-tree/utils.tsx    | 23 +++++++++++--------
 4 files changed, 19 insertions(+), 14 deletions(-)

diff --git a/packages/react-components/src/components/timegraph-output-component.tsx b/packages/react-components/src/components/timegraph-output-component.tsx
index a9451d35..8f999b49 100644
--- a/packages/react-components/src/components/timegraph-output-component.tsx
+++ b/packages/react-components/src/components/timegraph-output-component.tsx
@@ -1056,9 +1056,6 @@ export class TimegraphOutputComponent extends AbstractTreeOutputComponent<Timegr
             additionalProperties
         );
         this.updateMarkersData(timeGraphData.rangeEvents, newRange, nbTimes);
-        if (fetchArrows) {
-            this.arrowLayer.addArrows(timeGraphData.arrows, this.getTimegraphRowIds().rowIds);
-        }
         this.rangeEventsLayer.addRangeEvents(timeGraphData.rangeEvents);
 
         if (document.getElementById(this.props.traceId + this.props.outputDescriptor.id + 'handleSpinner')) {
@@ -1085,6 +1082,9 @@ export class TimegraphOutputComponent extends AbstractTreeOutputComponent<Timegr
             emptyNodes = [];
         }
 
+        if (fetchArrows) {
+            this.arrowLayer.addArrows(timeGraphData.arrows, this.getTimegraphRowIds().rowIds.filter(rowId => !emptyNodes.includes(rowId)));
+        }
         this.setState({ emptyNodes });
 
         // Apply the pending selection here since the row provider had been called before this method.
diff --git a/packages/react-components/src/components/utils/filter-tree/table.tsx b/packages/react-components/src/components/utils/filter-tree/table.tsx
index bdbd9f70..5100d516 100644
--- a/packages/react-components/src/components/utils/filter-tree/table.tsx
+++ b/packages/react-components/src/components/utils/filter-tree/table.tsx
@@ -80,7 +80,7 @@ export class Table extends React.Component<TableProps> {
         let nodes = sortNodes(this.props.nodes, this.props.sortConfig);
 
         if (this.props.hideEmptyNodes) {
-            nodes = filterEmptyNodes(nodes, this.props.emptyNodes);
+            nodes = filterEmptyNodes(nodes, this.props.emptyNodes, this.props.collapsedNodes);
         }
 
         return (
diff --git a/packages/react-components/src/components/utils/filter-tree/tree.tsx b/packages/react-components/src/components/utils/filter-tree/tree.tsx
index fc9fae4f..b88bec2d 100644
--- a/packages/react-components/src/components/utils/filter-tree/tree.tsx
+++ b/packages/react-components/src/components/utils/filter-tree/tree.tsx
@@ -3,7 +3,7 @@ import { TreeNode } from './tree-node';
 import { Message } from './message';
 import { Filter } from './filter';
 import { Table } from './table';
-import { getAllExpandedNodeIds, filterEmptyNodes } from './utils';
+import { getAllExpandedNodeIds } from './utils';
 import { SortConfig, sortNodes } from './sort';
 import ColumnHeader from './column-header';
 import { isEqual } from 'lodash';
diff --git a/packages/react-components/src/components/utils/filter-tree/utils.tsx b/packages/react-components/src/components/utils/filter-tree/utils.tsx
index a7222a1b..afd440e1 100644
--- a/packages/react-components/src/components/utils/filter-tree/utils.tsx
+++ b/packages/react-components/src/components/utils/filter-tree/utils.tsx
@@ -75,26 +75,31 @@ export const validateNumArray = (arr: any | undefined): boolean => {
  * Removes specified nodes from the tree structure.
  * @param nodes The array of root nodes of the tree.
  * @param nodesToRemove An array of node IDs to be removed.
+ * @param collapsedNodes The array of collapsed node IDs.
  * @returns A new array of root nodes with specified nodes removed.
  */
-export function filterEmptyNodes(nodes: TreeNode[], nodesToRemove: number[]): TreeNode[] {
+export function filterEmptyNodes(nodes: TreeNode[], nodesToRemove: number[], collapsedNodes: number[]): TreeNode[] {
     // return nodes;
     return nodes.reduce((acc: TreeNode[], node) => {
-        // If the current node is in the removal list, don't add it to the accumulator
-        if (nodesToRemove.includes(node.id)) {
-            return acc;
-        }
-
         // Create a new node object with the same properties
         const newNode: TreeNode = { ...node };
 
         // Recursively remove nodes from children
         if (newNode.children.length > 0) {
-            newNode.children = filterEmptyNodes(newNode.children, nodesToRemove);
+            newNode.children = filterEmptyNodes(newNode.children, nodesToRemove, collapsedNodes);
+        }
+
+        if (!nodesToRemove.includes(node.id)) {
+            // If the new node is not in the removal list, add it to the accumulator
+            acc.push(newNode);
+        } else if (!collapsedNodes.includes(node.id)) {
+            // If the new node is in the removal list and expanded, add its filtered children to the accumulator
+            newNode.children.forEach(child => {
+                child.parentId = newNode.parentId;
+                acc.push(child);
+            });
         }
 
-        // Add the new node to the accumulator
-        acc.push(newNode);
         return acc;
     }, []);
 }