Optimize the plots sections to reduce the number of useless re-renderings (#3341)

* Optimize the plots sections to reduce the number of useless re-renderings

* Revert to section variable name to reduce number of changes

Author: sroy3

webview/src/plots/components/templatePlots/TemplatePlots.tsx

@@ -1,17 +1,11 @@
 import { TemplatePlotGroup } from 'dvc/src/plots/webview/contract'
-import React, {
-  DragEvent,
-  useState,
-  useEffect,
-  useRef,
-  useCallback
-} from 'react'
+import React, { DragEvent, useState, useCallback } from 'react'
 import cx from 'classnames'
 import { useDispatch, useSelector } from 'react-redux'
 import { MessageFromWebviewType } from 'dvc/src/webview/contract'
 import { AddedSection } from './AddedSection'
 import { TemplatePlotsGrid } from './TemplatePlotsGrid'
-import { PlotGroup } from './templatePlotsSlice'
+import { PlotGroup, updateSections } from './templatePlotsSlice'
 import { removeFromPreviousAndAddToNewSection } from './util'
 import { sendMessage } from '../../../shared/vscode'
 import { createIDWithIndex, getIDIndex } from '../../../util/ids'
@@ -30,9 +24,7 @@ export enum NewSectionBlock {
 }
 
 export const TemplatePlots: React.FC = () => {
-  const { plotSections, size } = useSelector(
-    (state: PlotsState) => state.template
-  )
+  const { size, sections } = useSelector((state: PlotsState) => state.template)
   const draggedOverGroup = useSelector(
     (state: PlotsState) => state.dragAndDrop.draggedOverGroup
   )
@@ -43,29 +35,30 @@ export const TemplatePlots: React.FC = () => {
     (state: PlotsState) => state.webview.selectedRevisions
   )
 
-  const [sections, setSections] = useState<PlotGroup[]>(plotSections)
   const [hoveredSection, setHoveredSection] = useState('')
   const nbItemsPerRow = size
-  const shouldSendMessage = useRef(true)
   const dispatch = useDispatch()
 
-  useEffect(() => {
-    shouldSendMessage.current = false
-    setSections(plotSections)
-  }, [plotSections, setSections])
-
-  useEffect(() => {
-    if (sections && shouldSendMessage.current) {
-      sendMessage({
-        payload: sections.map(section => ({
-          group: section.group,
-          paths: section.entries
-        })),
-        type: MessageFromWebviewType.REORDER_PLOTS_TEMPLATES
-      })
-    }
-    shouldSendMessage.current = true
-  }, [sections])
+  const sendReorderMessage = useCallback((sections: PlotGroup[]) => {
+    sendMessage({
+      payload: sections.map(section => ({
+        group: section.group,
+        paths: section.entries
+      })),
+      type: MessageFromWebviewType.REORDER_PLOTS_TEMPLATES
+    })
+  }, [])
+
+  const setSections = useCallback(
+    (sections: PlotGroup[]) => {
+      /* Although the following dispatch duplicates the work the reducer will do when the state returns 
+         from the extension, this is necessary to not see any flickering in the order as the returned state 
+         sometimes takes a while to come back */
+      dispatch(updateSections(sections))
+      sendReorderMessage(sections)
+    },
+    [dispatch, sendReorderMessage]
+  )
 
   const handleDropInSection = useCallback(
     (
@@ -94,16 +87,14 @@ export const TemplatePlots: React.FC = () => {
 
   const setSectionEntries = useCallback(
     (index: number, entries: string[]) => {
-      setSections(sections => {
-        const updatedSections = [...sections]
-        updatedSections[index] = {
-          ...sections[index],
-          entries
-        }
-        return updatedSections
-      })
+      const updatedSections = [...sections]
+      updatedSections[index] = {
+        ...updatedSections[index],
+        entries
+      }
+      setSections(updatedSections)
     },
-    [setSections]
+    [setSections, sections]
   )
 
   if (sectionIsLoading(selectedRevisions)) {
@@ -235,7 +226,6 @@ export const TemplatePlots: React.FC = () => {
               useVirtualizedGrid={useVirtualizedGrid}
               nbItemsPerRow={nbItemsPerRow}
               parentDraggedOver={draggedOverGroup === groupId}
-              entries={section.entries}
             />
           </div>
         )

webview/src/plots/components/templatePlots/TemplatePlotsGrid.tsx

@@ -1,6 +1,6 @@
 import cx from 'classnames'
 import { Section } from 'dvc/src/plots/webview/contract'
-import React, { useEffect, useState, useCallback, useMemo } from 'react'
+import React, { useEffect, useCallback, useMemo } from 'react'
 import { useDispatch, useSelector } from 'react-redux'
 import { changeDisabledDragIds, changeSize } from './templatePlotsSlice'
 import { VirtualizedGrid } from '../../../shared/components/virtualizedGrid/VirtualizedGrid'
@@ -24,7 +24,6 @@ interface TemplatePlotsGridProps {
   setSectionEntries: (groupIndex: number, entries: string[]) => void
   useVirtualizedGrid?: boolean
   nbItemsPerRow: number
-  entries: string[]
   parentDraggedOver?: boolean
 }
 
@@ -36,12 +35,13 @@ export const TemplatePlotsGrid: React.FC<TemplatePlotsGridProps> = ({
   setSectionEntries,
   useVirtualizedGrid,
   nbItemsPerRow,
-  entries,
   parentDraggedOver
 }) => {
   const dispatch = useDispatch()
-  const [order, setOrder] = useState<string[]>([])
   const currentSize = useSelector((state: PlotsState) => state.template.size)
+  const entries = useSelector(
+    (state: PlotsState) => state.template.sections[groupIndex].entries
+  )
 
   const disabledDragPlotIds = useSelector(
     (state: PlotsState) => state.template.disabledDragPlotIds
@@ -64,10 +64,6 @@ export const TemplatePlotsGrid: React.FC<TemplatePlotsGridProps> = ({
     e.stopPropagation()
   }, [])
 
-  useEffect(() => {
-    setOrder(entries)
-  }, [entries])
-
   useEffect(() => {
     const panels = document.querySelectorAll('.vega-bindings')
     return () => {
@@ -88,11 +84,8 @@ export const TemplatePlotsGrid: React.FC<TemplatePlotsGridProps> = ({
     }
   }, [addDisabled, removeDisabled, disableClick])
 
-  const setEntriesOrder = (order: string[]) => {
-    setOrder(order)
-
+  const setEntriesOrder = (order: string[]) =>
     setSectionEntries(groupIndex, order)
-  }
 
   const plotClassName = cx(styles.plot, {
     [styles.multiViewPlot]: multiView
@@ -107,7 +100,7 @@ export const TemplatePlotsGrid: React.FC<TemplatePlotsGridProps> = ({
 
   const items = useMemo(
     () =>
-      order.map((plot: string) => {
+      entries.map((plot: string) => {
         const colSpan =
           (multiView &&
             plotDataStore[Section.TEMPLATE_PLOTS][plot].revisions?.length) ||
@@ -134,7 +127,7 @@ export const TemplatePlotsGrid: React.FC<TemplatePlotsGridProps> = ({
         )
       }),
     [
-      order,
+      entries,
       plotClassName,
       addEventsOnViewReady,
       currentSize,
@@ -145,7 +138,7 @@ export const TemplatePlotsGrid: React.FC<TemplatePlotsGridProps> = ({
 
   return (
     <DragDropContainer
-      order={order}
+      order={entries}
       setOrder={setEntriesOrder}
       items={items}
       group={groupId}

webview/src/plots/components/templatePlots/templatePlotsSlice.ts

@@ -12,17 +12,17 @@ export type PlotGroup = { group: TemplatePlotGroup; entries: string[] }
 export interface TemplatePlotsState extends Omit<TemplatePlotsData, 'plots'> {
   isCollapsed: boolean
   hasData: boolean
-  plotSections: PlotGroup[]
   plotsSnapshots: { [key: string]: string }
+  sections: PlotGroup[]
   disabledDragPlotIds: string[]
 }
 
 export const templatePlotsInitialState: TemplatePlotsState = {
   disabledDragPlotIds: [],
   hasData: false,
   isCollapsed: DEFAULT_SECTION_COLLAPSED[Section.TEMPLATE_PLOTS],
-  plotSections: [],
   plotsSnapshots: {},
+  sections: [],
   size: DEFAULT_SECTION_SIZES[Section.TEMPLATE_PLOTS]
 }
 
@@ -48,6 +48,7 @@ export const templatePlotsSlice = createSlice({
         entries: section.entries.map(entry => entry.id),
         group: section.group
       }))
+
       const plots = action.payload.plots?.flatMap(section => section.entries)
       const plotsIds = plots?.map(plot => plot.id) || []
       const snapShots = addPlotsWithSnapshots(plots, Section.TEMPLATE_PLOTS)
@@ -56,15 +57,29 @@ export const templatePlotsSlice = createSlice({
       return {
         ...state,
         hasData: !!action.payload,
-        plotSections,
         plotsSnapshots: snapShots,
+        sections:
+          JSON.stringify(plotSections) === JSON.stringify(state.sections)
+            ? state.sections
+            : plotSections,
         size: Math.abs(action.payload.size)
       }
+    },
+    updateSections: (state, action: PayloadAction<PlotGroup[]>) => {
+      return {
+        ...state,
+        sections: action.payload
+      }
     }
   }
 })
 
-export const { update, setCollapsed, changeSize, changeDisabledDragIds } =
-  templatePlotsSlice.actions
+export const {
+  update,
+  setCollapsed,
+  changeSize,
+  changeDisabledDragIds,
+  updateSections
+} = templatePlotsSlice.actions
 
 export default templatePlotsSlice.reducer