From c3d1bb52f8fb3a1348a27b013bbd179165fa0f2d Mon Sep 17 00:00:00 2001
From: Saul Shanabrook <s.shanabrook@gmail.com>
Date: Wed, 23 Oct 2024 10:50:08 -0400
Subject: [PATCH] Preserve filtered egraph on change

closes https://github.com/egraphs-good/egraph-visualizer/issues/2
---
 package.json       |  2 +-
 src/Visualizer.tsx | 23 ++++++++++++++++-------
 2 files changed, 17 insertions(+), 8 deletions(-)

diff --git a/package.json b/package.json
index d37af2c..04e36d1 100644
--- a/package.json
+++ b/package.json
@@ -1,6 +1,6 @@
 {
   "name": "egraph-visualizer",
-  "version": "2.0.4",
+  "version": "2.0.5",
   "repository": {
     "type": "git",
     "url": "git+https://github.com/egraphs-good/egraph-visualizer.git"
diff --git a/src/Visualizer.tsx b/src/Visualizer.tsx
index 9bcd202..f721751 100644
--- a/src/Visualizer.tsx
+++ b/src/Visualizer.tsx
@@ -333,28 +333,31 @@ function LayoutFlow({
   outerElem,
   innerElem,
   aspectRatio,
+  firstEgraph,
 }: {
   egraph: string;
   outerElem: HTMLDivElement;
   innerElem: HTMLDivElement;
   aspectRatio: number;
+  firstEgraph: string;
 }) {
   const [useInteractiveLayout, setUseInteractiveLayout] = useState(false);
   const [mergeEdges, setMergeEdges] = useState(false);
   const previousLayoutRef = useRef<PreviousLayout | null>(null);
-  // e-class ID we have currently selected, store egraph string as well so we know if this selection is outdated
-  const [selectedNodeWithEGraph, setSelectedNodeWithEGraph] = useState<(SelectedNode & { egraph: string }) | null>(null);
+  // e-class ID we have currently selected, store the first egraph string as well so we know if this selection is outdated,
+  // if our whole list of egraphs changes, but keep the selection if we have simply added a new egraph on
+  const [selectedNodeWithEGraph, setSelectedNodeWithEGraph] = useState<(SelectedNode & { firstEgraph: string }) | null>(null);
   const selectedNode = useMemo(() => {
-    if (selectedNodeWithEGraph && selectedNodeWithEGraph.egraph === egraph) {
+    if (selectedNodeWithEGraph && selectedNodeWithEGraph.firstEgraph === firstEgraph) {
       return selectedNodeWithEGraph;
     }
     return null;
-  }, [selectedNodeWithEGraph, egraph]);
+  }, [selectedNodeWithEGraph, firstEgraph]);
   const setSelectedNode = useCallback(
     (node: { type: "class" | "node"; id: string } | null) => {
-      setSelectedNodeWithEGraph(node ? { ...node, egraph } : null);
+      setSelectedNodeWithEGraph(node ? { ...node, firstEgraph } : null);
     },
-    [setSelectedNodeWithEGraph, egraph]
+    [setSelectedNodeWithEGraph, firstEgraph]
   );
 
   const getNodeSize = useCallback(
@@ -469,7 +472,13 @@ export function Visualizer({ egraphs, height = null, resize = false }: { egraphs
       </div>
       <SelectSider length={egraphs.length} onSelect={onSelect} selected={actualSelected} />
       {outerElem && innerElem && aspectRatio && (
-        <LayoutFlow aspectRatio={aspectRatio} egraph={egraphs[actualSelected]} outerElem={outerElem} innerElem={innerElem} />
+        <LayoutFlow
+          aspectRatio={aspectRatio}
+          firstEgraph={egraphs[0]}
+          egraph={egraphs[actualSelected]}
+          outerElem={outerElem}
+          innerElem={innerElem}
+        />
       )}
     </div>
   );