codepod-io · senwang86 · Aug 26, 2023
diff --git 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
@@ -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) => {