From eaf4e1556204abb4dff38cf192dc5e9b9fd4a201 Mon Sep 17 00:00:00 2001
From: Matt Seddon <37993418+mattseddon@users.noreply.github.com>
Date: Wed, 13 Dec 2023 10:58:07 +1100
Subject: [PATCH] Reduce number of plots resize observers (#5097)

---
 .../components/customPlots/CustomPlots.tsx    |  19 +-
 .../customPlots/CustomPlotsGrid.tsx           |  36 ++++
 .../customPlots/customPlotsSlice.ts           |  28 ++-
 .../templatePlots/TemplatePlots.tsx           | 203 +++++++++++-------
 .../templatePlots/templatePlotsSlice.ts       |  16 ++
 webview/src/plots/hooks/useGetPlot.ts         |  49 ++---
 .../plots/hooks/useObserveGridDimensions.ts   |  50 +++++
 7 files changed, 285 insertions(+), 116 deletions(-)
 create mode 100644 webview/src/plots/components/customPlots/CustomPlotsGrid.tsx
 create mode 100644 webview/src/plots/hooks/useObserveGridDimensions.ts

diff --git a/webview/src/plots/components/customPlots/CustomPlots.tsx b/webview/src/plots/components/customPlots/CustomPlots.tsx
index 60795ab671..68c172ab8c 100644
--- a/webview/src/plots/components/customPlots/CustomPlots.tsx
+++ b/webview/src/plots/components/customPlots/CustomPlots.tsx
@@ -1,11 +1,10 @@
-import React, { DragEvent, useEffect, useState } from 'react'
-import { useSelector } from 'react-redux'
-import { PlotsSection } from 'dvc/src/plots/webview/contract'
+import React, { DragEvent, useEffect, useRef, useState } from 'react'
 import cx from 'classnames'
+import { useSelector } from 'react-redux'
 import { NoPlotsAdded } from './NoPlotsAdded'
+import { CustomPlotsGrid } from './CustomPlotsGrid'
 import styles from '../styles.module.scss'
 import { shouldUseVirtualizedGrid } from '../util'
-import { Grid } from '../Grid'
 import { LoadingSection, sectionIsLoading } from '../LoadingSection'
 import { PlotsState } from '../../store'
 import { changeOrderWithDraggedInfo } from '../../../util/array'
@@ -33,6 +32,8 @@ export const CustomPlots: React.FC<CustomPlotsProps> = ({ plotsIds }) => {
     (state: PlotsState) => state.webview.selectedRevisions
   )
 
+  const gridRef = useRef<HTMLDivElement>(null)
+
   useEffect(() => {
     setOrder(plotsIds)
   }, [plotsIds])
@@ -82,15 +83,15 @@ export const CustomPlots: React.FC<CustomPlotsProps> = ({ plotsIds }) => {
       onDragLeave={() => setOnSection(false)}
       onDragOver={handleDragOver}
       onDrop={handleDropAtTheEnd}
+      ref={gridRef}
     >
-      <Grid
-        setOrder={setPlotsIdsOrder}
+      <CustomPlotsGrid
+        gridRef={gridRef}
         nbItemsPerRow={nbItemsPerRow}
-        useVirtualizedGrid={useVirtualizedGrid}
         order={order}
-        groupId="custom-plots"
         parentDraggedOver={onSection}
-        sectionKey={PlotsSection.CUSTOM_PLOTS}
+        setOrder={setPlotsIdsOrder}
+        useVirtualizedGrid={useVirtualizedGrid}
       />
     </div>
   )
diff --git a/webview/src/plots/components/customPlots/CustomPlotsGrid.tsx b/webview/src/plots/components/customPlots/CustomPlotsGrid.tsx
new file mode 100644
index 0000000000..37b0633988
--- /dev/null
+++ b/webview/src/plots/components/customPlots/CustomPlotsGrid.tsx
@@ -0,0 +1,36 @@
+import React, { RefObject } from 'react'
+import { PlotsSection } from 'dvc/src/plots/webview/contract'
+import { useObserveGridDimensions } from '../../hooks/useObserveGridDimensions'
+import { Grid } from '../Grid'
+
+interface CustomPlotsGridProps {
+  gridRef: RefObject<HTMLDivElement>
+  nbItemsPerRow: number
+  order: string[]
+  parentDraggedOver: boolean
+  useVirtualizedGrid?: boolean
+  setOrder: (order: string[]) => void
+}
+
+export const CustomPlotsGrid: React.FC<CustomPlotsGridProps> = ({
+  gridRef,
+  nbItemsPerRow,
+  parentDraggedOver,
+  order,
+  setOrder,
+  useVirtualizedGrid
+}) => {
+  useObserveGridDimensions(PlotsSection.CUSTOM_PLOTS, gridRef)
+
+  return (
+    <Grid
+      setOrder={setOrder}
+      nbItemsPerRow={nbItemsPerRow}
+      useVirtualizedGrid={useVirtualizedGrid}
+      order={order}
+      groupId="custom-plots"
+      parentDraggedOver={parentDraggedOver}
+      sectionKey={PlotsSection.CUSTOM_PLOTS}
+    />
+  )
+}
diff --git a/webview/src/plots/components/customPlots/customPlotsSlice.ts b/webview/src/plots/components/customPlots/customPlotsSlice.ts
index aca045a30a..440fcb77d0 100644
--- a/webview/src/plots/components/customPlots/customPlotsSlice.ts
+++ b/webview/src/plots/components/customPlots/customPlotsSlice.ts
@@ -10,12 +10,14 @@ import {
 import { addPlotsWithSnapshots, removePlots } from '../plotDataStore'
 
 export interface CustomPlotsState extends Omit<CustomPlotsData, 'plots'> {
-  isCollapsed: boolean
   hasData: boolean
   hasItems: boolean
+  isCollapsed: boolean
+  isInDragAndDropMode: boolean
   plotsIds: string[]
   plotsSnapshots: { [key: string]: string }
-  isInDragAndDropMode: boolean
+  sectionHeight: number
+  sectionWidth: number
 }
 
 export const customPlotsInitialState: CustomPlotsState = {
@@ -29,7 +31,9 @@ export const customPlotsInitialState: CustomPlotsState = {
   nbItemsPerRow:
     DEFAULT_SECTION_NB_ITEMS_PER_ROW_OR_WIDTH[PlotsSection.CUSTOM_PLOTS],
   plotsIds: [],
-  plotsSnapshots: {}
+  plotsSnapshots: {},
+  sectionHeight: 0,
+  sectionWidth: 0
 }
 
 export const customPlotsSlice = createSlice({
@@ -75,16 +79,28 @@ export const customPlotsSlice = createSlice({
         plotsIds: plots?.map(plot => plot.id) || [],
         plotsSnapshots
       }
+    },
+    updateSectionDimensions: (
+      state,
+      action: PayloadAction<{ sectionHeight: number; sectionWidth: number }>
+    ) => {
+      const { sectionHeight, sectionWidth } = action.payload
+      return {
+        ...state,
+        sectionHeight,
+        sectionWidth
+      }
     }
   }
 })
 
 export const {
-  update,
-  setCollapsed,
   changeSize,
+  clearState,
+  setCollapsed,
   toggleDragAndDropMode,
-  clearState
+  update,
+  updateSectionDimensions
 } = customPlotsSlice.actions
 
 export default customPlotsSlice.reducer
diff --git a/webview/src/plots/components/templatePlots/TemplatePlots.tsx b/webview/src/plots/components/templatePlots/TemplatePlots.tsx
index da97fe24ee..e716033ddd 100644
--- a/webview/src/plots/components/templatePlots/TemplatePlots.tsx
+++ b/webview/src/plots/components/templatePlots/TemplatePlots.tsx
@@ -1,5 +1,11 @@
-import { TemplatePlotGroup } from 'dvc/src/plots/webview/contract'
-import React, { DragEvent, useState, useCallback } from 'react'
+import { PlotsSection, TemplatePlotGroup } from 'dvc/src/plots/webview/contract'
+import React, {
+  DragEvent,
+  useState,
+  useCallback,
+  useRef,
+  RefObject
+} from 'react'
 import cx from 'classnames'
 import { useDispatch, useSelector } from 'react-redux'
 import { AddedSection } from './AddedSection'
@@ -11,12 +17,117 @@ import { createIDWithIndex, getIDIndex } from '../../../util/ids'
 import styles from '../styles.module.scss'
 import { shouldUseVirtualizedGrid } from '../util'
 import { PlotsState } from '../../store'
-import { setDraggedOverGroup } from '../../../shared/components/dragDrop/dragDropSlice'
+import {
+  DraggedInfo,
+  setDraggedOverGroup
+} from '../../../shared/components/dragDrop/dragDropSlice'
 import { isSameGroup } from '../../../shared/components/dragDrop/util'
 import { changeOrderWithDraggedInfo } from '../../../util/array'
 import { LoadingSection, sectionIsLoading } from '../LoadingSection'
 import { reorderTemplatePlots } from '../../util/messages'
 import { TooManyPlots } from '../TooManyPlots'
+import { useObserveGridDimensions } from '../../hooks/useObserveGridDimensions'
+
+interface TemplatePlotGroupsProps {
+  draggedRef: DraggedInfo
+  draggedOverGroup: string
+  gridRef: RefObject<HTMLDivElement>
+  handleDropInSection: (
+    draggedId: string,
+    draggedGroup: string,
+    groupId: string,
+    position?: number
+  ) => void
+  handleEnteringSection: (groupId: string) => void
+  nbItemsPerRow: number
+  sections: PlotGroup[]
+  setSectionEntries: (index: number, entries: string[]) => void
+  setSections: (sections: PlotGroup[]) => void
+}
+
+const TemplatePlotGroups: React.FC<TemplatePlotGroupsProps> = ({
+  draggedOverGroup,
+  draggedRef,
+  gridRef,
+  handleDropInSection,
+  handleEnteringSection,
+  nbItemsPerRow,
+  sections,
+  setSectionEntries,
+  setSections
+}) => {
+  useObserveGridDimensions(PlotsSection.TEMPLATE_PLOTS, gridRef)
+
+  return sections.map((section, i) => {
+    const groupId = createIDWithIndex(section.group, i)
+    const useVirtualizedGrid = shouldUseVirtualizedGrid(
+      Object.keys(section.entries).length,
+      nbItemsPerRow
+    )
+
+    const isMultiView = section.group === TemplatePlotGroup.MULTI_VIEW
+
+    const classes = cx(styles.sectionWrapper, {
+      [styles.multiViewPlotsGrid]: isMultiView,
+      [styles.singleViewPlotsGrid]: !isMultiView,
+      [styles.noBigGrid]: !useVirtualizedGrid
+    })
+
+    const handleDropAtTheEnd = () => {
+      handleEnteringSection('')
+      if (!draggedRef) {
+        return
+      }
+
+      if (draggedRef.group === groupId) {
+        const order = section.entries
+        const updatedSections = [...sections]
+
+        const newOrder = changeOrderWithDraggedInfo(order, draggedRef)
+        updatedSections[i] = {
+          ...sections[i],
+          entries: newOrder
+        }
+        setSections(updatedSections)
+      } else if (isSameGroup(draggedRef.group, groupId)) {
+        handleDropInSection(
+          draggedRef.itemId,
+          draggedRef.group,
+          groupId,
+          section.entries.length
+        )
+      }
+    }
+
+    const handleDragOver = (e: DragEvent) => {
+      e.preventDefault()
+      handleEnteringSection(groupId)
+    }
+
+    return (
+      <div
+        key={groupId}
+        id={groupId}
+        data-testid={`plots-section_${groupId}`}
+        className={classes}
+        onDragEnter={() => handleEnteringSection(groupId)}
+        onDragOver={handleDragOver}
+        onDrop={handleDropAtTheEnd}
+      >
+        <TemplatePlotsGrid
+          groupId={groupId}
+          groupIndex={i}
+          onDropInSection={handleDropInSection}
+          multiView={isMultiView}
+          setSectionEntries={setSectionEntries}
+          useVirtualizedGrid={useVirtualizedGrid}
+          nbItemsPerRow={nbItemsPerRow}
+          parentDraggedOver={draggedOverGroup === groupId}
+        />
+      </div>
+    )
+  })
+}
 
 export enum NewSectionBlock {
   TOP = 'drop-section-top',
@@ -37,6 +148,8 @@ export const TemplatePlots: React.FC = () => {
     (state: PlotsState) => state.webview.selectedRevisions
   )
 
+  const gridRef = useRef<HTMLDivElement>(null)
+
   const [hoveredSection, setHoveredSection] = useState('')
   const dispatch = useDispatch()
 
@@ -151,87 +264,29 @@ export const TemplatePlots: React.FC = () => {
   }
 
   return (
-    <>
+    <div ref={gridRef}>
       <AddedSection
         {...newDropSection}
         id={NewSectionBlock.TOP}
         closestSection={firstSection}
       />
-      {sections.map((section, i) => {
-        const groupId = createIDWithIndex(section.group, i)
-        const useVirtualizedGrid = shouldUseVirtualizedGrid(
-          Object.keys(section.entries).length,
-          nbItemsPerRow
-        )
-
-        const isMultiView = section.group === TemplatePlotGroup.MULTI_VIEW
-
-        const classes = cx(styles.sectionWrapper, {
-          [styles.multiViewPlotsGrid]: isMultiView,
-          [styles.singleViewPlotsGrid]: !isMultiView,
-          [styles.noBigGrid]: !useVirtualizedGrid
-        })
-
-        const handleDropAtTheEnd = () => {
-          handleEnteringSection('')
-          if (!draggedRef) {
-            return
-          }
-
-          if (draggedRef.group === groupId) {
-            const order = section.entries
-            const updatedSections = [...sections]
-
-            const newOrder = changeOrderWithDraggedInfo(order, draggedRef)
-            updatedSections[i] = {
-              ...sections[i],
-              entries: newOrder
-            }
-            setSections(updatedSections)
-          } else if (isSameGroup(draggedRef.group, groupId)) {
-            handleDropInSection(
-              draggedRef.itemId,
-              draggedRef.group,
-              groupId,
-              section.entries.length
-            )
-          }
-        }
-
-        const handleDragOver = (e: DragEvent) => {
-          e.preventDefault()
-          handleEnteringSection(groupId)
-        }
-
-        return (
-          <div
-            key={groupId}
-            id={groupId}
-            data-testid={`plots-section_${groupId}`}
-            className={classes}
-            onDragEnter={() => handleEnteringSection(groupId)}
-            onDragOver={handleDragOver}
-            onDrop={handleDropAtTheEnd}
-          >
-            <TemplatePlotsGrid
-              groupId={groupId}
-              groupIndex={i}
-              onDropInSection={handleDropInSection}
-              multiView={isMultiView}
-              setSectionEntries={setSectionEntries}
-              useVirtualizedGrid={useVirtualizedGrid}
-              nbItemsPerRow={nbItemsPerRow}
-              parentDraggedOver={draggedOverGroup === groupId}
-            />
-          </div>
-        )
-      })}
+      <TemplatePlotGroups
+        draggedRef={draggedRef}
+        draggedOverGroup={draggedOverGroup}
+        gridRef={gridRef}
+        handleDropInSection={handleDropInSection}
+        handleEnteringSection={handleEnteringSection}
+        nbItemsPerRow={nbItemsPerRow}
+        sections={sections}
+        setSections={setSections}
+        setSectionEntries={setSectionEntries}
+      />
       <AddedSection
         {...newDropSection}
         id={NewSectionBlock.BOTTOM}
         closestSection={lastSection}
       />
       {shouldShowTooManyPlotsMessage && <TooManyPlots />}
-    </>
+    </div>
   )
 }
diff --git a/webview/src/plots/components/templatePlots/templatePlotsSlice.ts b/webview/src/plots/components/templatePlots/templatePlotsSlice.ts
index bc1e793373..a2ace2778f 100644
--- a/webview/src/plots/components/templatePlots/templatePlotsSlice.ts
+++ b/webview/src/plots/components/templatePlots/templatePlotsSlice.ts
@@ -16,7 +16,9 @@ export interface TemplatePlotsState extends Omit<TemplatePlotsData, 'plots'> {
   hasData: boolean
   hasItems: boolean
   plotsSnapshots: { [key: string]: string }
+  sectionHeight: number
   sections: PlotGroup[]
+  sectionWidth: number
   shouldShowTooManyPlotsMessage: boolean
   isInDragAndDropMode: boolean
 }
@@ -30,6 +32,8 @@ export const templatePlotsInitialState: TemplatePlotsState = {
   nbItemsPerRow:
     DEFAULT_SECTION_NB_ITEMS_PER_ROW_OR_WIDTH[PlotsSection.TEMPLATE_PLOTS],
   plotsSnapshots: {},
+  sectionHeight: 0,
+  sectionWidth: 0,
   sections: [],
   shouldShowTooManyPlotsMessage: false,
   smoothPlotValues: {}
@@ -90,6 +94,17 @@ export const templatePlotsSlice = createSlice({
         smoothPlotValues: action.payload.smoothPlotValues
       }
     },
+    updateSectionDimensions: (
+      state,
+      action: PayloadAction<{ sectionHeight: number; sectionWidth: number }>
+    ) => {
+      const { sectionHeight, sectionWidth } = action.payload
+      return {
+        ...state,
+        sectionHeight,
+        sectionWidth
+      }
+    },
     updateSections: (state, action: PayloadAction<PlotGroup[]>) => {
       return {
         ...state,
@@ -111,6 +126,7 @@ export const {
   changeSize,
   toggleDragAndDropMode,
   updateSections,
+  updateSectionDimensions,
   updateShouldShowTooManyPlotsMessage,
   clearState
 } = templatePlotsSlice.actions
diff --git a/webview/src/plots/hooks/useGetPlot.ts b/webview/src/plots/hooks/useGetPlot.ts
index e1ee462fee..28ecadc0bd 100644
--- a/webview/src/plots/hooks/useGetPlot.ts
+++ b/webview/src/plots/hooks/useGetPlot.ts
@@ -9,49 +9,44 @@ import { fillTemplate } from '../components/vegaLite/util'
 export const useGetPlot = (
   section: PlotsSection,
   id: string,
-  plotRef: RefObject<HTMLButtonElement | HTMLDivElement>,
+  parentRef: RefObject<HTMLButtonElement | HTMLDivElement>,
   plotFocused: boolean
-  // eslint-disable-next-line sonarjs/cognitive-complexity
 ): VisualizationSpec | undefined => {
   const storeSection =
     section === PlotsSection.TEMPLATE_PLOTS ? 'template' : 'custom'
-  const { plotsSnapshots } = useSelector(
-    (state: PlotsState) => state[storeSection]
-  )
+  const {
+    plotsSnapshots,
+    nbItemsPerRow,
+    height: plotHeight,
+    sectionHeight,
+    sectionWidth
+  } = useSelector((state: PlotsState) => state[storeSection])
 
   const [spec, setSpec] = useState<VisualizationSpec | undefined>()
 
-  const [height, setHeight] = useState(0)
-  const [width, setWidth] = useState(0)
-
   useEffect(() => {
-    const resizeObserver = new ResizeObserver(() => {
-      if (!plotRef.current) {
-        return
-      }
-      const { height, width } = plotRef.current.getBoundingClientRect()
-      setHeight(height)
-      setWidth(width)
-    })
-
-    if (plotRef.current) {
-      resizeObserver.observe(plotRef.current)
-    }
-
-    return () => {
-      resizeObserver.disconnect()
+    if (!parentRef.current) {
+      return
     }
-  }, [plotRef])
-
-  useEffect(() => {
     const plot = plotDataStore[section][id]
+    const { height, width } = parentRef.current.getBoundingClientRect()
 
     const spec = fillTemplate(plot, width, height, plotFocused)
     if (!spec) {
       return
     }
     setSpec(spec)
-  }, [height, id, plotFocused, plotsSnapshots, section, width])
+  }, [
+    id,
+    nbItemsPerRow,
+    parentRef,
+    plotFocused,
+    plotHeight,
+    plotsSnapshots,
+    section,
+    sectionHeight,
+    sectionWidth
+  ])
 
   return spec
 }
diff --git a/webview/src/plots/hooks/useObserveGridDimensions.ts b/webview/src/plots/hooks/useObserveGridDimensions.ts
new file mode 100644
index 0000000000..cc307dfbe3
--- /dev/null
+++ b/webview/src/plots/hooks/useObserveGridDimensions.ts
@@ -0,0 +1,50 @@
+import type { RefObject } from 'react'
+import { useEffect } from 'react'
+import { useDispatch } from 'react-redux'
+import { PlotsSection } from 'dvc/src/plots/webview/contract'
+import { updateSectionDimensions as updateTemplateDimensions } from '../components/templatePlots/templatePlotsSlice'
+import { updateSectionDimensions as updateCustomDimensions } from '../components/customPlots/customPlotsSlice'
+
+const updateBySection: {
+  [section: string]:
+    | typeof updateTemplateDimensions
+    | typeof updateCustomDimensions
+} = {
+  [PlotsSection.TEMPLATE_PLOTS]: updateTemplateDimensions,
+  [PlotsSection.CUSTOM_PLOTS]: updateCustomDimensions
+}
+
+export const useObserveGridDimensions = (
+  sectionKey: PlotsSection,
+  ref: RefObject<HTMLElement>
+): void => {
+  const dispatch = useDispatch()
+
+  useEffect(() => {
+    const updateSectionDimensions = updateBySection[sectionKey]
+    const resizeObserver = new ResizeObserver(() => {
+      if (!updateSectionDimensions) {
+        return
+      }
+
+      if (!ref.current) {
+        dispatch(updateSectionDimensions({ sectionHeight: 0, sectionWidth: 0 }))
+        return
+      }
+
+      const { height, width } = ref.current.getBoundingClientRect()
+
+      dispatch(
+        updateSectionDimensions({ sectionHeight: height, sectionWidth: width })
+      )
+    })
+
+    if (ref.current) {
+      resizeObserver.observe(ref.current)
+    }
+
+    return () => {
+      resizeObserver.disconnect()
+    }
+  }, [dispatch, ref, sectionKey])
+}