Optimize plot re-rendering (#3337)

* Optimize rendering time of plots by making sure snap points do not trigger re-rendering

* Use a type for snappoints

* Optimize rendering of plots

* Do not drill down vega plot content

* Fix checkpoint plots

* Add linting rule for TODOs

* Made a few variables clearer

Author: sroy3

.eslintrc.js

@@ -148,6 +148,7 @@ module.exports = {
     // https://github.com/typescript-eslint/typescript-eslint/issues/2540#issuecomment-692505191
     'no-use-before-define': 'off',
     'no-void': ['error', { allowAsStatement: true }],
+    'no-warning-comments': 'error',
     quotes: ['error', 'single', { avoidEscape: true }],
     'react-hooks/exhaustive-deps': 'error',
     'react-hooks/rules-of-hooks': 'error',

webview/src/plots/components/App.test.tsx

@@ -27,8 +27,7 @@ import {
   Revision,
   Section,
   TemplatePlotGroup,
-  TemplatePlotsData,
-  TemplatePlotSection
+  TemplatePlotsData
 } from 'dvc/src/plots/webview/contract'
 import {
   MessageFromWebviewType,
@@ -41,8 +40,12 @@ import { VisualizationSpec } from 'react-vega'
 import { App } from './App'
 import { NewSectionBlock } from './templatePlots/TemplatePlots'
 import { SectionDescription } from './PlotsContainer'
-import { CheckpointPlotsById, plotDataStore } from './plotDataStore'
-import { setMaxPlotSize } from './webviewSlice'
+import {
+  CheckpointPlotsById,
+  plotDataStore,
+  TemplatePlotsById
+} from './plotDataStore'
+import { setSnapPoints } from './webviewSlice'
 import { plotsReducers, plotsStore } from '../store'
 import { vsCodeApi } from '../../shared/api'
 import {
@@ -126,7 +129,7 @@ describe('App', () => {
 
   const setWrapperSize = (store: typeof plotsStore) =>
     act(() => {
-      store.dispatch(setMaxPlotSize(1000))
+      store.dispatch(setSnapPoints(1000))
     })
 
   const templatePlot = templatePlotsFixture.plots[0].entries[0]
@@ -201,8 +204,8 @@ describe('App', () => {
     jest
       .spyOn(HTMLElement.prototype, 'clientHeight', 'get')
       .mockImplementation(() => heightToSuppressVegaError)
-    plotDataStore.checkpoint = {} as CheckpointPlotsById
-    plotDataStore.template = [] as TemplatePlotSection[]
+    plotDataStore[Section.CHECKPOINT_PLOTS] = {} as CheckpointPlotsById
+    plotDataStore[Section.TEMPLATE_PLOTS] = {} as TemplatePlotsById
   })
 
   afterEach(() => {
@@ -1427,7 +1430,7 @@ describe('App', () => {
 
     const resizeScreen = (width: number, store: typeof plotsStore) => {
       act(() => {
-        store.dispatch(setMaxPlotSize(width))
+        store.dispatch(setSnapPoints(width))
       })
       act(() => {
         global.innerWidth = width

webview/src/plots/components/Plots.tsx

@@ -6,7 +6,7 @@ import { CheckpointPlotsWrapper } from './checkpointPlots/CheckpointPlotsWrapper
 import { TemplatePlotsWrapper } from './templatePlots/TemplatePlotsWrapper'
 import { ComparisonTableWrapper } from './comparisonTable/ComparisonTableWrapper'
 import { Ribbon } from './ribbon/Ribbon'
-import { setMaxPlotSize, setZoomedInPlot } from './webviewSlice'
+import { setSnapPoints, setZoomedInPlot } from './webviewSlice'
 import { EmptyState } from '../../shared/components/emptyState/EmptyState'
 import { Modal } from '../../shared/components/modal/Modal'
 import { WebviewWrapper } from '../../shared/components/webviewWrapper/WebviewWrapper'
@@ -38,7 +38,7 @@ const PlotsContent = () => {
     const onResize = () => {
       wrapperRef.current &&
         dispatch(
-          setMaxPlotSize(wrapperRef.current.getBoundingClientRect().width - 100)
+          setSnapPoints(wrapperRef.current.getBoundingClientRect().width - 100)
         )
     }
     window.addEventListener('resize', onResize)

webview/src/plots/components/ZoomablePlot.tsx

@@ -1,45 +1,54 @@
+import { AnyAction } from '@reduxjs/toolkit'
 import cx from 'classnames'
-import React, { useEffect, useRef, useState } from 'react'
-import { useDispatch } from 'react-redux'
-import { PlainObject, VisualizationSpec } from 'react-vega'
+import { MessageFromWebviewType } from 'dvc/src/webview/contract'
+import { Section } from 'dvc/src/plots/webview/contract'
+import React, { useCallback, useEffect, useRef, useState } from 'react'
+import { useDispatch, useSelector } from 'react-redux'
+import { VisualizationSpec } from 'react-vega'
 import { Renderers } from 'vega'
 import VegaLite, { VegaLiteProps } from 'react-vega/lib/VegaLite'
 import { setZoomedInPlot } from './webviewSlice'
 import styles from './styles.module.scss'
 import { Resizer } from './Resizer'
 import { config } from './constants'
+import { PlotsState } from '../store'
+import { useGetPlot } from '../hooks/useGetPlot'
 import { GripIcon } from '../../shared/components/dragDrop/GripIcon'
+import { sendMessage } from '../../shared/vscode'
 
 interface ZoomablePlotProps {
-  spec: VisualizationSpec
-  data?: PlainObject
+  spec?: VisualizationSpec
   id: string
   onViewReady?: () => void
-  toggleDrag: (enabled: boolean) => void
-  onResize: (diff: number) => void
-  snapPoints: number[]
+  toggleDrag: (enabled: boolean, id: string) => void
+  changeSize: (size: number) => AnyAction
   currentSnapPoint: number
-  size: number
+  section: Section
+  shouldNotResize?: boolean
 }
 
 export const ZoomablePlot: React.FC<ZoomablePlotProps> = ({
-  spec,
-  data,
+  spec: createdSpec,
   id,
   onViewReady,
   toggleDrag,
-  onResize,
-  snapPoints,
+  changeSize,
   currentSnapPoint,
-  size
+  section,
+  shouldNotResize
 }) => {
+  const snapPoints = useSelector(
+    (state: PlotsState) => state.webview.snapPoints
+  )
+  const { data, content: spec } = useGetPlot(section, id, createdSpec)
   const dispatch = useDispatch()
   const previousSpecsAndData = useRef(JSON.stringify({ data, spec }))
   const currentPlotProps = useRef<VegaLiteProps>()
   const clickDisabled = useRef(false)
   const enableClickTimeout = useRef(0)
   const [isExpanding, setIsExpanding] = useState(false)
   const newSpecsAndData = JSON.stringify({ data, spec })
+  const size = snapPoints[currentSnapPoint - 1]
 
   const plotProps: VegaLiteProps = {
     actions: false,
@@ -69,20 +78,34 @@ export const ZoomablePlot: React.FC<ZoomablePlotProps> = ({
   const handleOnClick = () =>
     !clickDisabled.current && dispatch(setZoomedInPlot({ id, plot: plotProps }))
 
+  const onResize = useCallback(
+    (newSnapPoint: number) => {
+      dispatch(changeSize(newSnapPoint))
+      sendMessage({
+        payload: { section, size: newSnapPoint },
+        type: MessageFromWebviewType.RESIZE_PLOTS
+      })
+    },
+    [dispatch, changeSize, section]
+  )
+
   const commonResizerProps = {
     onGrab: () => {
       clickDisabled.current = true
-      toggleDrag(false)
+      toggleDrag(false, id)
     },
     onRelease: () => {
-      toggleDrag(true)
+      toggleDrag(true, id)
       enableClickTimeout.current = window.setTimeout(
         () => (clickDisabled.current = false),
         0
       )
     },
     onResize
   }
+  if (!data && !spec) {
+    return null
+  }
   return (
     <button
       className={cx(styles.zoomablePlot, {
@@ -94,7 +117,7 @@ export const ZoomablePlot: React.FC<ZoomablePlotProps> = ({
       {currentPlotProps.current && (
         <VegaLite {...plotProps} onNewView={onViewReady} />
       )}
-      {snapPoints.length > 0 && (
+      {!shouldNotResize && (
         <Resizer
           className={styles.plotVerticalResizer}
           {...commonResizerProps}

webview/src/plots/components/checkpointPlots/CheckpointPlot.tsx

@@ -8,7 +8,6 @@ import styles from '../styles.module.scss'
 import { withScale } from '../../../util/styles'
 import { plotDataStore } from '../plotDataStore'
 import { PlotsState } from '../../store'
-import { useResize } from '../../hooks/useResize'
 
 interface CheckpointPlotProps {
   id: string
@@ -23,12 +22,9 @@ export const CheckpointPlot: React.FC<CheckpointPlotProps> = ({
   const plotSnapshot = useSelector(
     (state: PlotsState) => state.checkpoint.plotsSnapshots[id]
   )
-  const [plot, setPlot] = useState(plotDataStore.checkpoint[id])
+  const [plot, setPlot] = useState(plotDataStore[Section.CHECKPOINT_PLOTS][id])
   const currentSize = useSelector((state: PlotsState) => state.checkpoint.size)
-  const { onResize: handleResize, snapPoints } = useResize(
-    Section.CHECKPOINT_PLOTS,
-    changeSize
-  )
+
   const spec = useMemo(() => {
     const title = plot?.title
     if (!title) {
@@ -38,15 +34,13 @@ export const CheckpointPlot: React.FC<CheckpointPlotProps> = ({
   }, [plot?.title, colors])
 
   useEffect(() => {
-    setPlot(plotDataStore.checkpoint[id])
+    setPlot(plotDataStore[Section.CHECKPOINT_PLOTS][id])
   }, [plotSnapshot, id])
 
   if (!plot) {
     return null
   }
 
-  const { values } = plot
-
   const key = `plot-${id}`
 
   const toggleDrag = (enabled: boolean) => {
@@ -57,13 +51,11 @@ export const CheckpointPlot: React.FC<CheckpointPlotProps> = ({
     <div className={styles.plot} data-testid={key} id={id} style={withScale(1)}>
       <ZoomablePlot
         spec={spec}
-        data={{ values }}
-        id={key}
+        id={id}
         toggleDrag={toggleDrag}
-        onResize={handleResize}
-        snapPoints={snapPoints}
+        changeSize={changeSize}
         currentSnapPoint={currentSize}
-        size={snapPoints[currentSize - 1]}
+        section={Section.CHECKPOINT_PLOTS}
       />
     </div>
   )

webview/src/plots/components/checkpointPlots/checkpointPlotsSlice.ts

@@ -5,10 +5,7 @@ import {
   DEFAULT_SECTION_SIZES,
   Section
 } from 'dvc/src/plots/webview/contract'
-import {
-  addCheckpointPlotsWithSnapshots,
-  removeCheckpointPlots
-} from '../plotDataStore'
+import { addPlotsWithSnapshots, removePlots } from '../plotDataStore'
 export interface CheckpointPlotsState
   extends Omit<CheckpointPlotsData, 'plots'> {
   isCollapsed: boolean
@@ -48,8 +45,8 @@ export const checkpointPlotsSlice = createSlice({
       }
       const { plots, ...statePayload } = action.payload
       const plotsIds = plots?.map(plot => plot.id) || []
-      const snapShots = addCheckpointPlotsWithSnapshots(plots)
-      removeCheckpointPlots(plotsIds)
+      const snapShots = addPlotsWithSnapshots(plots, Section.CHECKPOINT_PLOTS)
+      removePlots(plotsIds, Section.CHECKPOINT_PLOTS)
       return {
         ...state,
         ...statePayload,

webview/src/plots/components/plotDataStore.ts

@@ -1,30 +1,34 @@
 import {
   CheckpointPlotData,
-  TemplatePlotSection
+  Section,
+  TemplatePlotEntry
 } from 'dvc/src/plots/webview/contract'
 
 export type CheckpointPlotsById = { [key: string]: CheckpointPlotData }
+export type TemplatePlotsById = { [key: string]: TemplatePlotEntry }
 
 export const plotDataStore = {
-  checkpoint: {} as CheckpointPlotsById,
-  template: [] as TemplatePlotSection[]
+  [Section.CHECKPOINT_PLOTS]: {} as CheckpointPlotsById,
+  [Section.TEMPLATE_PLOTS]: {} as TemplatePlotsById,
+  [Section.COMPARISON_TABLE]: {} as CheckpointPlotsById // This category is unused but exists only to make typings easier
 }
 
-export const addCheckpointPlotsWithSnapshots = (
-  plots: CheckpointPlotData[]
+export const addPlotsWithSnapshots = (
+  plots: (CheckpointPlotData | TemplatePlotEntry)[],
+  section: Section
 ) => {
   const snapShots: { [key: string]: string } = {}
   for (const plot of plots || []) {
-    plotDataStore.checkpoint[plot.id] = plot
+    plotDataStore[section][plot.id] = plot
     snapShots[plot.id] = JSON.stringify(plot)
   }
   return snapShots
 }
 
-export const removeCheckpointPlots = (currentIds: string[]) => {
-  for (const plotId of Object.keys(plotDataStore.checkpoint)) {
+export const removePlots = (currentIds: string[], section: Section) => {
+  for (const plotId of Object.keys(plotDataStore[section])) {
     if (!currentIds.includes(plotId)) {
-      delete plotDataStore.checkpoint[plotId]
+      delete plotDataStore[section][plotId]
     }
   }
 }

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

@@ -1,10 +1,10 @@
 import React, { DragEvent } from 'react'
 import { useSelector } from 'react-redux'
 import cx from 'classnames'
-import { TemplatePlotSection } from 'dvc/src/plots/webview/contract'
+import { PlotGroup } from './templatePlotsSlice'
 import styles from '../styles.module.scss'
-import { getIDWithoutIndex } from '../../../util/ids'
 import { PlotsState } from '../../store'
+import { getIDWithoutIndex } from '../../../util/ids'
 import { Icon } from '../../../shared/components/Icon'
 import { GraphLine } from '../../../shared/components/icons'
 import { useDeferedDragLeave } from '../../../shared/hooks/useDeferedDragLeave'
@@ -14,7 +14,7 @@ interface AddedSectionProps {
   hoveredSection: string
   setHoveredSection: (section: string) => void
   onDrop: (e: DragEvent<HTMLElement>) => void
-  closestSection: TemplatePlotSection
+  closestSection: PlotGroup
   acceptedGroups: string[]
 }