From 3c0baf641d922a20de06160f7e44e697ee98baeb Mon Sep 17 00:00:00 2001
From: Sen Wang <senwang@Sens-MBP-M2Pro.local>
Date: Fri, 25 Aug 2023 18:56:39 -0700
Subject: [PATCH] SelectiveView

---
 apps/ui/src/components/Sidebar.tsx   | 32 +++++++++
 apps/ui/src/lib/store/canvasSlice.ts | 98 +++++++++++++++++++++++-----
 2 files changed, 112 insertions(+), 18 deletions(-)

diff --git a/apps/ui/src/components/Sidebar.tsx b/apps/ui/src/components/Sidebar.tsx
index 4f68a8b6..62206e7e 100644
--- a/apps/ui/src/components/Sidebar.tsx
+++ b/apps/ui/src/components/Sidebar.tsx
@@ -1294,10 +1294,24 @@ function PodTreeItem({ id, node2children }) {
   );
 }
 
+function dfs(node, node2children, visited) {
+  if (!visited.has(node)) {
+    visited.add(node);
+    const children = node2children.get(node);
+    children.forEach((n) => {
+      dfs(n, node2children, visited);
+    });
+  }
+}
+
 function TableofPods() {
   const store = useContext(RepoContext);
   if (!store) throw new Error("Missing BearContext.Provider in the tree");
   const nodesMap = useStore(store, (state) => state.getNodesMap());
+  const selectedToc = useStore(store, (state) => state.selectedToc);
+  const setSelectedToc = useStore(store, (state) => state.setSelectedToc);
+  const setSelectiveView = useStore(store, (state) => state.setSelectiveView);
+  const clearSelectedToc = useStore(store, (state) => state.clearSelectedToc);
   let node2children = new Map<string, string[]>();
   let keys = new Set(nodesMap.keys());
 
@@ -1323,6 +1337,23 @@ function TableofPods() {
     }
   }
 
+  const handleSelect = (event, nodeIds: string[]) => {
+    if (nodeIds.length > 1) {
+      const selectedItems = new Set<string>();
+      //select all descendants of a node
+      nodeIds.forEach((n) => {
+        dfs(n, node2children, selectedItems);
+      });
+
+      if (selectedItems !== selectedToc) {
+        setSelectedToc(selectedItems);
+      }
+      setSelectiveView(true);
+    } else {
+      clearSelectedToc();
+      setSelectiveView(false);
+    }
+  };
   return (
     <TreeView
       aria-label="multi-select"
@@ -1331,6 +1362,7 @@ function TableofPods() {
       defaultExpanded={Array.from(node2children.keys()).filter(
         (key) => node2children!.get(key!)!.length > 0
       )}
+      onNodeSelect={handleSelect}
       multiSelect
     >
       {node2children.size > 0 &&
diff --git a/apps/ui/src/lib/store/canvasSlice.ts b/apps/ui/src/lib/store/canvasSlice.ts
index f757df1c..b41988f4 100644
--- a/apps/ui/src/lib/store/canvasSlice.ts
+++ b/apps/ui/src/lib/store/canvasSlice.ts
@@ -283,6 +283,12 @@ export interface CanvasSlice {
   centerSelection: boolean;
   setCenterSelection: (b: boolean) => void;
 
+  isSelectiveView: boolean;
+  setSelectiveView: (b: boolean) => void;
+  selectedToc: Set<string>;
+  clearSelectedToc: () => void;
+  setSelectedToc: (selectedItems: Set<string>) => void;
+
   handlePaste(event: ClipboardEvent, position: XYPosition): void;
   handleCopy(event: ClipboardEvent): void;
 
@@ -382,6 +388,23 @@ export const createCanvasSlice: StateCreator<MyState, [], [], CanvasSlice> = (
     set({ centerSelection: b });
   },
 
+  isSelectiveView: false,
+  setSelectiveView: (b: boolean) => set({ isSelectiveView: b }),
+  selectedToc: new Set(),
+  setSelectedToc(selectedItems) {
+    set(
+      produce((state: MyState) => {
+        state.selectedToc = selectedItems;
+      })
+    );
+  },
+  clearSelectedToc: () => {
+    set(
+      produce((state: MyState) => {
+        state.selectedToc.clear();
+      })
+    );
+  },
   focusedEditor: undefined,
   setFocusedEditor: (id?: string) =>
     set(
@@ -408,29 +431,68 @@ export const createCanvasSlice: StateCreator<MyState, [], [], CanvasSlice> = (
    * This function handles the real updates to the reactflow nodes to render.
    */
   updateView: () => {
+    let nodes: Node[] = [];
     const nodesMap = get().getNodesMap();
-    let selectedPods = get().selectedPods;
-    let nodes = Array.from<Node>(nodesMap.values());
-    nodes = nodes
-      .sort((a: Node, b: Node) => a.data.level - b.data.level)
-      .map((node) => ({
-        ...node,
-        style: {
-          ...node.style,
-          backgroundColor:
-            node.type === "SCOPE" ? level2color[node.data.level] : undefined,
-        },
-        selected: selectedPods.has(node.id),
-        // className: get().dragHighlight === node.id ? "active" : "",
-        className: match(node.id)
-          .with(get().dragHighlight || "", () => "active")
-          .otherwise(() => undefined),
-      }));
+    if (!get().isSelectiveView) {
+      let selectedPods = get().selectedPods;
+      nodes = Array.from<Node>(nodesMap.values());
+      nodes = nodes
+        .sort((a: Node, b: Node) => a.data.level - b.data.level)
+        .map((node) => ({
+          ...node,
+          style: {
+            ...node.style,
+            backgroundColor:
+              node.type === "SCOPE" ? level2color[node.data.level] : undefined,
+          },
+          selected: selectedPods.has(node.id),
+          // className: get().dragHighlight === node.id ? "active" : "",
+          className: match(node.id)
+            .with(get().dragHighlight || "", () => "active")
+            .otherwise(() => undefined),
+        }));
+    } else {
+      nodes = Array.from(
+        [...get().selectedToc].map((id): Node => {
+          return nodesMap.get(id)!;
+        })
+      );
+      nodes = nodes
+        .sort((a: Node, b: Node) => a.data.level - b.data.level)
+        .map((node) => ({
+          ...node,
+          style: {
+            ...node.style,
+            backgroundColor:
+              node.type === "SCOPE" ? level2color[node.data.level] : undefined,
+          },
+          parentNode:
+            node.parentNode === undefined ||
+            get().selectedToc.has(node.parentNode)
+              ? node.parentNode
+              : undefined,
+          position:
+            node.parentNode === undefined ||
+            get().selectedToc.has(node.parentNode)
+              ? node.position
+              : getAbsPos(node, nodesMap),
+          selected: get().selectedToc.has(node.id),
+          // className: get().dragHighlight === node.id ? "active" : "",
+          className: match(node.id)
+            .with(get().dragHighlight || "", () => "active")
+            .otherwise(() => undefined),
+        }));
+    }
 
     set({ nodes });
     // edges view
     const edgesMap = get().getEdgesMap();
-    set({ edges: Array.from<Edge>(edgesMap.values()).filter((e) => e) });
+    const nodeSet = new Set<Node>(Array.from(nodes));
+    set({
+      edges: Array.from<Edge>(edgesMap.values()).filter(
+        (e) => nodeSet.has(e.sourceNode!) && nodeSet.has(e.targetNode!)
+      ),
+    });
   },
 
   addNode: (type, position, parent) => {