From d88764638372cbd5055f98f13469810f100f7d87 Mon Sep 17 00:00:00 2001
From: Mike Fortman <michael.fortman@datastax.com>
Date: Wed, 7 May 2025 09:12:22 -0500
Subject: [PATCH] refactor: performance improvements for canvas
 controls/toolbar (#7930)

* performance improvements for canvas controls/toolbar

* [autofix.ci] apply automated fixes

---------

Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>
---
 .../core/canvasControlsComponent/index.tsx    | 19 ++---
 .../core/flowToolbarComponent/index.tsx       | 71 ++-----------------
 .../PageComponent/MemoizedComponents.tsx      | 61 ++++++++++++++++
 .../components/PageComponent/index.tsx        | 54 ++++----------
 .../components/nodeToolbarComponent/index.tsx | 16 -----
 5 files changed, 90 insertions(+), 131 deletions(-)
 create mode 100644 src/frontend/src/pages/FlowPage/components/PageComponent/MemoizedComponents.tsx

diff --git a/src/frontend/src/components/core/canvasControlsComponent/index.tsx b/src/frontend/src/components/core/canvasControlsComponent/index.tsx
index 53cf7658c..1c3a9869c 100644
--- a/src/frontend/src/components/core/canvasControlsComponent/index.tsx
+++ b/src/frontend/src/components/core/canvasControlsComponent/index.tsx
@@ -13,7 +13,8 @@ import {
   type ReactFlowState,
 } from "@xyflow/react";
 import { cloneDeep } from "lodash";
-import { useEffect } from "react";
+import { useCallback, useEffect } from "react";
+import { useShallow } from "zustand/react/shallow";
 import { shallow } from "zustand/shallow";
 
 type CustomControlButtonProps = {
@@ -70,20 +71,22 @@ const CanvasControls = ({ children }) => {
     shallow,
   );
   const saveFlow = useSaveFlow();
-  const currentFlow = useFlowStore((state) => state.currentFlow);
+  const isLocked = useFlowStore(
+    useShallow((state) => state.currentFlow?.locked),
+  );
   const setCurrentFlow = useFlowStore((state) => state.setCurrentFlow);
   const autoSaving = useFlowsManagerStore((state) => state.autoSaving);
 
   useEffect(() => {
-    const isLocked = currentFlow?.locked;
     store.setState({
       nodesDraggable: !isLocked,
       nodesConnectable: !isLocked,
       elementsSelectable: !isLocked,
     });
-  }, [currentFlow?.locked]);
+  }, [isLocked]);
 
-  const handleSaveFlow = () => {
+  const handleSaveFlow = useCallback(() => {
+    const currentFlow = useFlowStore.getState().currentFlow;
     if (!currentFlow) return;
     const newFlow = cloneDeep(currentFlow);
     newFlow.locked = isInteractive;
@@ -92,16 +95,16 @@ const CanvasControls = ({ children }) => {
     } else {
       setCurrentFlow(newFlow);
     }
-  };
+  }, [isInteractive, autoSaving, saveFlow, setCurrentFlow]);
 
-  const onToggleInteractivity = () => {
+  const onToggleInteractivity = useCallback(() => {
     store.setState({
       nodesDraggable: !isInteractive,
       nodesConnectable: !isInteractive,
       elementsSelectable: !isInteractive,
     });
     handleSaveFlow();
-  };
+  }, [isInteractive, store, handleSaveFlow]);
 
   return (
     <Panel
diff --git a/src/frontend/src/components/core/flowToolbarComponent/index.tsx b/src/frontend/src/components/core/flowToolbarComponent/index.tsx
index 4a74f5e6e..5bdc0d568 100644
--- a/src/frontend/src/components/core/flowToolbarComponent/index.tsx
+++ b/src/frontend/src/components/core/flowToolbarComponent/index.tsx
@@ -1,7 +1,7 @@
 import ShadTooltip from "@/components/common/shadTooltipComponent";
 import { track } from "@/customization/utils/analytics";
 import { Panel } from "@xyflow/react";
-import { useEffect, useMemo, useState } from "react";
+import { memo, useEffect, useMemo, useState } from "react";
 import { useHotkeys } from "react-hotkeys-hook";
 import ShareModal from "../../../modals/shareModal";
 import useFlowStore from "../../../stores/flowStore";
@@ -11,7 +11,7 @@ import { classNames, cn, isThereModal } from "../../../utils/utils";
 import ForwardedIconComponent from "../../common/genericIconComponent";
 import FlowToolbarOptions from "./components/flow-toolbar-options";
 
-export default function FlowToolbar(): JSX.Element {
+const FlowToolbar = memo(function FlowToolbar(): JSX.Element {
   const preventDefault = true;
   const [open, setOpen] = useState<boolean>(false);
   const [openCodeModal, setOpenCodeModal] = useState<boolean>(false);
@@ -41,75 +41,12 @@ export default function FlowToolbar(): JSX.Element {
   useHotkeys(api, handleAPIWShortcut, { preventDefault });
   useHotkeys(flow, handleShareWShortcut, { preventDefault });
 
-  const hasIO = useFlowStore((state) => state.hasIO);
-  const hasStore = useStoreStore((state) => state.hasStore);
-  const validApiKey = useStoreStore((state) => state.validApiKey);
-  const hasApiKey = useStoreStore((state) => state.hasApiKey);
-  const currentFlow = useFlowStore((state) => state.currentFlow);
-
   useEffect(() => {
     if (open) {
       track("Playground Button Clicked");
     }
   }, [open]);
 
-  const ModalMemo = useMemo(
-    () => (
-      <ShareModal
-        is_component={false}
-        component={currentFlow!}
-        disabled={!hasApiKey || !validApiKey || !hasStore}
-        open={openShareModal}
-        setOpen={setOpenShareModal}
-      >
-        <ShadTooltip
-          content={
-            !hasApiKey || !validApiKey || !hasStore
-              ? "Store API Key Required"
-              : ""
-          }
-          side="bottom"
-          align="end"
-        >
-          <button
-            disabled={!hasApiKey || !validApiKey || !hasStore}
-            className={classNames(
-              "relative inline-flex h-8 w-full items-center justify-center gap-1.5 rounded px-3 py-1.5 text-sm font-semibold text-foreground transition-all duration-150 ease-in-out",
-              !hasApiKey || !validApiKey || !hasStore
-                ? "cursor-not-allowed text-muted-foreground"
-                : "hover:bg-accent",
-            )}
-            data-testid="shared-button-flow"
-            onClick={() => {
-              setOpenShareModal(true);
-            }}
-          >
-            <>
-              <ForwardedIconComponent
-                name="Share2"
-                className={classNames(
-                  "h-4 w-4",
-                  !hasApiKey || !validApiKey || !hasStore
-                    ? "extra-side-bar-save-disable"
-                    : "",
-                )}
-              />
-              <span className="hidden md:block">Share</span>
-            </>
-          </button>
-        </ShadTooltip>
-      </ShareModal>
-    ),
-    [
-      hasApiKey,
-      validApiKey,
-      currentFlow,
-      hasStore,
-      openShareModal,
-      setOpenShareModal,
-    ],
-  );
-
   return (
     <>
       <Panel className="!m-2" position="top-right">
@@ -123,4 +60,6 @@ export default function FlowToolbar(): JSX.Element {
       </Panel>
     </>
   );
-}
+});
+
+export default FlowToolbar;
diff --git a/src/frontend/src/pages/FlowPage/components/PageComponent/MemoizedComponents.tsx b/src/frontend/src/pages/FlowPage/components/PageComponent/MemoizedComponents.tsx
new file mode 100644
index 000000000..7d7c0110e
--- /dev/null
+++ b/src/frontend/src/pages/FlowPage/components/PageComponent/MemoizedComponents.tsx
@@ -0,0 +1,61 @@
+import ForwardedIconComponent from "@/components/common/genericIconComponent";
+import CanvasControls, {
+  CustomControlButton,
+} from "@/components/core/canvasControlsComponent";
+import { SidebarTrigger } from "@/components/ui/sidebar";
+import { cn } from "@/utils/utils";
+import { Background, Panel } from "@xyflow/react";
+import { memo } from "react";
+
+export const MemoizedBackground = memo(() => (
+  <Background size={2} gap={20} className="" />
+));
+
+interface MemoizedCanvasControlsProps {
+  setIsAddingNote: (value: boolean) => void;
+  position: { x: number; y: number };
+  shadowBoxWidth: number;
+  shadowBoxHeight: number;
+}
+
+export const MemoizedCanvasControls = memo(
+  ({
+    setIsAddingNote,
+    position,
+    shadowBoxWidth,
+    shadowBoxHeight,
+  }: MemoizedCanvasControlsProps) => (
+    <CanvasControls>
+      <CustomControlButton
+        iconName="sticky-note"
+        tooltipText="Add Note"
+        onClick={() => {
+          setIsAddingNote(true);
+          const shadowBox = document.getElementById("shadow-box");
+          if (shadowBox) {
+            shadowBox.style.display = "block";
+            shadowBox.style.left = `${position.x - shadowBoxWidth / 2}px`;
+            shadowBox.style.top = `${position.y - shadowBoxHeight / 2}px`;
+          }
+        }}
+        iconClasses="text-primary"
+        testId="add_note"
+      />
+    </CanvasControls>
+  ),
+);
+
+export const MemoizedSidebarTrigger = memo(() => (
+  <Panel
+    className={cn(
+      "react-flow__controls !m-2 flex gap-1.5 rounded-md border border-secondary-hover bg-background fill-foreground stroke-foreground p-1.5 text-primary shadow transition-all duration-300 [&>button]:border-0 [&>button]:bg-background hover:[&>button]:bg-accent",
+      "pointer-events-auto opacity-100 group-data-[open=true]/sidebar-wrapper:pointer-events-none group-data-[open=true]/sidebar-wrapper:-translate-x-full group-data-[open=true]/sidebar-wrapper:opacity-0",
+    )}
+    position="top-left"
+  >
+    <SidebarTrigger className="h-fit w-fit px-3 py-1.5">
+      <ForwardedIconComponent name="PanelRightClose" className="h-4 w-4" />
+      <span className="text-foreground">Components</span>
+    </SidebarTrigger>
+  </Panel>
+));
diff --git a/src/frontend/src/pages/FlowPage/components/PageComponent/index.tsx b/src/frontend/src/pages/FlowPage/components/PageComponent/index.tsx
index 0d4e7a1d2..55c976c3e 100644
--- a/src/frontend/src/pages/FlowPage/components/PageComponent/index.tsx
+++ b/src/frontend/src/pages/FlowPage/components/PageComponent/index.tsx
@@ -1,12 +1,6 @@
 import { DefaultEdge } from "@/CustomEdges";
 import NoteNode from "@/CustomNodes/NoteNode";
-
-import ForwardedIconComponent from "@/components/common/genericIconComponent";
-import CanvasControls, {
-  CustomControlButton,
-} from "@/components/core/canvasControlsComponent";
 import FlowToolbar from "@/components/core/flowToolbarComponent";
-import { SidebarTrigger } from "@/components/ui/sidebar";
 import {
   COLOR_OPTIONS,
   NOTE_NODE_MIN_HEIGHT,
@@ -21,12 +15,10 @@ import { useAddComponent } from "@/hooks/use-add-component";
 import { nodeColorsName } from "@/utils/styleUtils";
 import { cn, isSupportedNodeTypes } from "@/utils/utils";
 import {
-  Background,
   Connection,
   Edge,
   OnNodeDrag,
   OnSelectionChangeParams,
-  Panel,
   ReactFlow,
   reconnectEdge,
   SelectionDragHandler,
@@ -67,6 +59,11 @@ import {
 import ConnectionLineComponent from "../ConnectionLineComponent";
 import SelectionMenu from "../SelectionMenuComponent";
 import UpdateAllComponents from "../UpdateAllComponents";
+import {
+  MemoizedBackground,
+  MemoizedCanvasControls,
+  MemoizedSidebarTrigger,
+} from "./MemoizedComponents";
 import getRandomName from "./utils/get-random-name";
 import isWrappedWithClass from "./utils/is-wrapped-with-class";
 
@@ -591,44 +588,19 @@ export default function Page({
             onPaneClick={onPaneClick}
             onEdgeClick={handleEdgeClick}
           >
-            <Background size={2} gap={20} className="" />
+            <MemoizedBackground />
             {!view && (
               <>
-                <CanvasControls>
-                  <CustomControlButton
-                    iconName="sticky-note"
-                    tooltipText="Add Note"
-                    onClick={() => {
-                      setIsAddingNote(true);
-                      const shadowBox = document.getElementById("shadow-box");
-                      if (shadowBox) {
-                        shadowBox.style.display = "block";
-                        shadowBox.style.left = `${position.current.x - shadowBoxWidth / 2}px`;
-                        shadowBox.style.top = `${position.current.y - shadowBoxHeight / 2}px`;
-                      }
-                    }}
-                    iconClasses="text-primary"
-                    testId="add_note"
-                  />
-                </CanvasControls>
+                <MemoizedCanvasControls
+                  setIsAddingNote={setIsAddingNote}
+                  position={position.current}
+                  shadowBoxWidth={shadowBoxWidth}
+                  shadowBoxHeight={shadowBoxHeight}
+                />
                 <FlowToolbar />
               </>
             )}
-            <Panel
-              className={cn(
-                "react-flow__controls !m-2 flex gap-1.5 rounded-md border border-secondary-hover bg-background fill-foreground stroke-foreground p-1.5 text-primary shadow transition-all duration-300 [&>button]:border-0 [&>button]:bg-background hover:[&>button]:bg-accent",
-                "pointer-events-auto opacity-100 group-data-[open=true]/sidebar-wrapper:pointer-events-none group-data-[open=true]/sidebar-wrapper:-translate-x-full group-data-[open=true]/sidebar-wrapper:opacity-0",
-              )}
-              position="top-left"
-            >
-              <SidebarTrigger className="h-fit w-fit px-3 py-1.5">
-                <ForwardedIconComponent
-                  name="PanelRightClose"
-                  className="h-4 w-4"
-                />
-                <span className="text-foreground">Components</span>
-              </SidebarTrigger>
-            </Panel>
+            <MemoizedSidebarTrigger />
             <div className={cn(componentsToUpdate.length === 0 && "hidden")}>
               <UpdateAllComponents />
             </div>
diff --git a/src/frontend/src/pages/FlowPage/components/nodeToolbarComponent/index.tsx b/src/frontend/src/pages/FlowPage/components/nodeToolbarComponent/index.tsx
index 93c8c3535..62133fbd1 100644
--- a/src/frontend/src/pages/FlowPage/components/nodeToolbarComponent/index.tsx
+++ b/src/frontend/src/pages/FlowPage/components/nodeToolbarComponent/index.tsx
@@ -73,7 +73,6 @@ const NodeToolbarComponent = memo(
     const currentFlowId = useFlowsManagerStore((state) => state.currentFlowId);
     const [openModal, setOpenModal] = useState(false);
     const frozen = data.node?.frozen ?? false;
-    const currentFlow = useFlowStore((state) => state.currentFlow);
     const updateNodeInternals = useUpdateNodeInternals();
 
     const paste = useFlowStore((state) => state.paste);
@@ -93,21 +92,6 @@ const NodeToolbarComponent = memo(
       },
     });
 
-    const flowDataNodes = useMemo(
-      () => currentFlow?.data?.nodes,
-      [currentFlow],
-    );
-
-    const node = useMemo(
-      () => flowDataNodes?.find((n) => n.id === data.id),
-      [flowDataNodes, data.id],
-    );
-
-    const index = useMemo(
-      () => flowDataNodes?.indexOf(node!)!,
-      [flowDataNodes, node],
-    );
-
     const postToolModeValue = usePostTemplateValue({
       node: data.node!,
       nodeId: data.id,