From 94d3c6adf91ea97742359c7192be1547062e08ba Mon Sep 17 00:00:00 2001
From: Stephanie Roy <stephanie@iterative.ai>
Date: Tue, 21 Feb 2023 11:40:46 -0500
Subject: [PATCH 1/7] Optimize rendering time of plots by making sure snap
 points do not trigger re-rendering

---
 webview/src/plots/components/ZoomablePlot.tsx | 34 +++++++++---
 .../checkpointPlots/CheckpointPlot.tsx        | 14 +++--
 .../templatePlots/TemplatePlotsGrid.tsx       | 14 +++--
 webview/src/plots/hooks/useResize.ts          | 52 -------------------
 webview/src/plots/hooks/useSnapPoints.ts      | 17 ++++++
 5 files changed, 55 insertions(+), 76 deletions(-)
 delete mode 100644 webview/src/plots/hooks/useResize.ts
 create mode 100644 webview/src/plots/hooks/useSnapPoints.ts

diff --git a/webview/src/plots/components/ZoomablePlot.tsx b/webview/src/plots/components/ZoomablePlot.tsx
index 0ce5636d12..22c4c0ede6 100644
--- a/webview/src/plots/components/ZoomablePlot.tsx
+++ b/webview/src/plots/components/ZoomablePlot.tsx
@@ -1,5 +1,8 @@
+import { AnyAction } from '@reduxjs/toolkit'
 import cx from 'classnames'
-import React, { useEffect, useRef, useState } from 'react'
+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 } from 'react-redux'
 import { PlainObject, VisualizationSpec } from 'react-vega'
 import { Renderers } from 'vega'
@@ -9,6 +12,7 @@ import styles from './styles.module.scss'
 import { Resizer } from './Resizer'
 import { config } from './constants'
 import { GripIcon } from '../../shared/components/dragDrop/GripIcon'
+import { sendMessage } from '../../shared/vscode'
 
 interface ZoomablePlotProps {
   spec: VisualizationSpec
@@ -16,10 +20,11 @@ interface ZoomablePlotProps {
   id: string
   onViewReady?: () => void
   toggleDrag: (enabled: boolean) => void
-  onResize: (diff: number) => void
-  snapPoints: number[]
+  changeSize: (size: number) => AnyAction
   currentSnapPoint: number
-  size: number
+  section: Section
+  snapPoints: [number, number, number, number]
+  shouldNotResize?: boolean
 }
 
 export const ZoomablePlot: React.FC<ZoomablePlotProps> = ({
@@ -28,10 +33,11 @@ export const ZoomablePlot: React.FC<ZoomablePlotProps> = ({
   id,
   onViewReady,
   toggleDrag,
-  onResize,
-  snapPoints,
+  changeSize,
   currentSnapPoint,
-  size
+  section,
+  snapPoints,
+  shouldNotResize
 }) => {
   const dispatch = useDispatch()
   const previousSpecsAndData = useRef(JSON.stringify({ data, spec }))
@@ -40,6 +46,7 @@ export const ZoomablePlot: React.FC<ZoomablePlotProps> = ({
   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,6 +76,17 @@ 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
@@ -94,7 +112,7 @@ export const ZoomablePlot: React.FC<ZoomablePlotProps> = ({
       {currentPlotProps.current && (
         <VegaLite {...plotProps} onNewView={onViewReady} />
       )}
-      {snapPoints.length > 0 && (
+      {!shouldNotResize && (
         <Resizer
           className={styles.plotVerticalResizer}
           {...commonResizerProps}
diff --git a/webview/src/plots/components/checkpointPlots/CheckpointPlot.tsx b/webview/src/plots/components/checkpointPlots/CheckpointPlot.tsx
index d3a545d9c2..1f490acf12 100644
--- a/webview/src/plots/components/checkpointPlots/CheckpointPlot.tsx
+++ b/webview/src/plots/components/checkpointPlots/CheckpointPlot.tsx
@@ -8,7 +8,7 @@ import styles from '../styles.module.scss'
 import { withScale } from '../../../util/styles'
 import { plotDataStore } from '../plotDataStore'
 import { PlotsState } from '../../store'
-import { useResize } from '../../hooks/useResize'
+import { useSnapPoints } from '../../hooks/useSnapPoints'
 
 interface CheckpointPlotProps {
   id: string
@@ -25,10 +25,8 @@ export const CheckpointPlot: React.FC<CheckpointPlotProps> = ({
   )
   const [plot, setPlot] = useState(plotDataStore.checkpoint[id])
   const currentSize = useSelector((state: PlotsState) => state.checkpoint.size)
-  const { onResize: handleResize, snapPoints } = useResize(
-    Section.CHECKPOINT_PLOTS,
-    changeSize
-  )
+  const snapPoints = useSnapPoints()
+
   const spec = useMemo(() => {
     const title = plot?.title
     if (!title) {
@@ -60,10 +58,10 @@ export const CheckpointPlot: React.FC<CheckpointPlotProps> = ({
         data={{ values }}
         id={key}
         toggleDrag={toggleDrag}
-        onResize={handleResize}
-        snapPoints={snapPoints}
+        changeSize={changeSize}
         currentSnapPoint={currentSize}
-        size={snapPoints[currentSize - 1]}
+        section={Section.CHECKPOINT_PLOTS}
+        snapPoints={snapPoints}
       />
     </div>
   )
diff --git a/webview/src/plots/components/templatePlots/TemplatePlotsGrid.tsx b/webview/src/plots/components/templatePlots/TemplatePlotsGrid.tsx
index 6a962f083e..e8b6a1c263 100644
--- a/webview/src/plots/components/templatePlots/TemplatePlotsGrid.tsx
+++ b/webview/src/plots/components/templatePlots/TemplatePlotsGrid.tsx
@@ -16,7 +16,7 @@ import { DropTarget } from '../DropTarget'
 import styles from '../styles.module.scss'
 import { ZoomablePlot } from '../ZoomablePlot'
 import { PlotsState } from '../../store'
-import { useResize } from '../../hooks/useResize'
+import { useSnapPoints } from '../../hooks/useSnapPoints'
 
 interface TemplatePlotsGridProps {
   entries: TemplatePlotEntry[]
@@ -48,15 +48,12 @@ export const TemplatePlotsGrid: React.FC<TemplatePlotsGridProps> = ({
 }) => {
   const dispatch = useDispatch()
   const [order, setOrder] = useState<string[]>([])
+  const snapPoints = useSnapPoints()
 
   const disabledDragPlotIds = useSelector(
     (state: PlotsState) => state.template.disabledDragPlotIds
   )
   const currentSize = useSelector((state: PlotsState) => state.template.size)
-  const { onResize: handleResize, snapPoints } = useResize(
-    Section.TEMPLATE_PLOTS,
-    changeSize
-  )
 
   const addDisabled = useCallback(
     (e: Event) => {
@@ -139,10 +136,11 @@ export const TemplatePlotsGrid: React.FC<TemplatePlotsGridProps> = ({
           spec={{ ...content, ...autoSize } as VisualizationSpec}
           onViewReady={addEventsOnViewReady}
           toggleDrag={toggleDrag}
-          onResize={handleResize}
-          snapPoints={multiView ? [] : snapPoints}
+          changeSize={changeSize}
           currentSnapPoint={currentSize}
-          size={snapPoints[currentSize - 1]}
+          shouldNotResize={multiView}
+          section={Section.TEMPLATE_PLOTS}
+          snapPoints={snapPoints}
         />
       </div>
     )
diff --git a/webview/src/plots/hooks/useResize.ts b/webview/src/plots/hooks/useResize.ts
deleted file mode 100644
index 4e54326479..0000000000
--- a/webview/src/plots/hooks/useResize.ts
+++ /dev/null
@@ -1,52 +0,0 @@
-import { AnyAction } from '@reduxjs/toolkit'
-import { Section } from 'dvc/src/plots/webview/contract'
-import { MessageFromWebviewType } from 'dvc/src/webview/contract'
-import { useCallback, useEffect, useState } from 'react'
-import { useDispatch, useSelector } from 'react-redux'
-import { sendMessage } from '../../shared/vscode'
-import { PlotsState } from '../store'
-
-interface SnapPoints {
-  1: number
-  2: number
-  3: number
-  4: number
-}
-
-const initialSnapPoints = {
-  1: 0,
-  2: 0,
-  3: 0,
-  4: 0
-}
-
-export const useResize = (
-  section: Section,
-  changeSize: (size: number) => AnyAction
-) => {
-  const dispatch = useDispatch()
-  const maxSize = useSelector((state: PlotsState) => state.webview.maxPlotSize)
-  const [snapPoints, setSnapPoints] = useState<SnapPoints>(initialSnapPoints)
-
-  useEffect(() => {
-    setSnapPoints({
-      1: maxSize,
-      2: maxSize / 2,
-      3: maxSize / 3,
-      4: maxSize / 4
-    })
-  }, [maxSize])
-
-  const onResize = useCallback(
-    (newSnapPoint: number) => {
-      dispatch(changeSize(newSnapPoint))
-      sendMessage({
-        payload: { section, size: newSnapPoint },
-        type: MessageFromWebviewType.RESIZE_PLOTS
-      })
-    },
-    [dispatch, changeSize, section]
-  )
-
-  return { onResize, snapPoints: Object.values(snapPoints) }
-}
diff --git a/webview/src/plots/hooks/useSnapPoints.ts b/webview/src/plots/hooks/useSnapPoints.ts
new file mode 100644
index 0000000000..055cb51a57
--- /dev/null
+++ b/webview/src/plots/hooks/useSnapPoints.ts
@@ -0,0 +1,17 @@
+import { useEffect, useState } from 'react'
+import { useSelector } from 'react-redux'
+import { PlotsState } from '../store'
+
+const initialSnapPoints: [number, number, number, number] = [0, 0, 0, 0]
+
+export const useSnapPoints = () => {
+  const maxSize = useSelector((state: PlotsState) => state.webview.maxPlotSize)
+  const [snapPoints, setSnapPoints] =
+    useState<[number, number, number, number]>(initialSnapPoints)
+
+  useEffect(() => {
+    setSnapPoints([maxSize, maxSize / 2, maxSize / 3, maxSize / 4])
+  }, [maxSize])
+
+  return snapPoints
+}

From 449241e0429cf88fe24311bafe23023d0c2d314e Mon Sep 17 00:00:00 2001
From: Stephanie Roy <stephanie@iterative.ai>
Date: Tue, 21 Feb 2023 13:34:41 -0500
Subject: [PATCH 2/7] Use a type for snappoints

---
 webview/src/plots/components/ZoomablePlot.tsx | 3 ++-
 webview/src/plots/hooks/useSnapPoints.ts      | 5 ++---
 2 files changed, 4 insertions(+), 4 deletions(-)

diff --git a/webview/src/plots/components/ZoomablePlot.tsx b/webview/src/plots/components/ZoomablePlot.tsx
index 22c4c0ede6..fef065baea 100644
--- a/webview/src/plots/components/ZoomablePlot.tsx
+++ b/webview/src/plots/components/ZoomablePlot.tsx
@@ -13,6 +13,7 @@ import { Resizer } from './Resizer'
 import { config } from './constants'
 import { GripIcon } from '../../shared/components/dragDrop/GripIcon'
 import { sendMessage } from '../../shared/vscode'
+import { SnapPoints } from '../hooks/useSnapPoints'
 
 interface ZoomablePlotProps {
   spec: VisualizationSpec
@@ -23,7 +24,7 @@ interface ZoomablePlotProps {
   changeSize: (size: number) => AnyAction
   currentSnapPoint: number
   section: Section
-  snapPoints: [number, number, number, number]
+  snapPoints: SnapPoints
   shouldNotResize?: boolean
 }
 
diff --git a/webview/src/plots/hooks/useSnapPoints.ts b/webview/src/plots/hooks/useSnapPoints.ts
index 055cb51a57..53ebaf3633 100644
--- a/webview/src/plots/hooks/useSnapPoints.ts
+++ b/webview/src/plots/hooks/useSnapPoints.ts
@@ -2,12 +2,11 @@ import { useEffect, useState } from 'react'
 import { useSelector } from 'react-redux'
 import { PlotsState } from '../store'
 
-const initialSnapPoints: [number, number, number, number] = [0, 0, 0, 0]
+export type SnapPoints = [number, number, number, number]
 
 export const useSnapPoints = () => {
   const maxSize = useSelector((state: PlotsState) => state.webview.maxPlotSize)
-  const [snapPoints, setSnapPoints] =
-    useState<[number, number, number, number]>(initialSnapPoints)
+  const [snapPoints, setSnapPoints] = useState<SnapPoints>([0, 0, 0, 0])
 
   useEffect(() => {
     setSnapPoints([maxSize, maxSize / 2, maxSize / 3, maxSize / 4])

From 6c65589a3d87bfaf6d4318944f6eccaf4f23ae08 Mon Sep 17 00:00:00 2001
From: Stephanie Roy <stephanie@iterative.ai>
Date: Tue, 21 Feb 2023 14:29:34 -0500
Subject: [PATCH 3/7] Optimize rendering of plots

---
 .../templatePlots/TemplatePlotsGrid.tsx       | 38 ++++----------
 .../templatePlots/TemplatePlotsGridItem.tsx   | 50 +++++++++++++++++++
 webview/src/util/wdyr.ts                      |  2 +-
 3 files changed, 62 insertions(+), 28 deletions(-)
 create mode 100644 webview/src/plots/components/templatePlots/TemplatePlotsGridItem.tsx

diff --git a/webview/src/plots/components/templatePlots/TemplatePlotsGrid.tsx b/webview/src/plots/components/templatePlots/TemplatePlotsGrid.tsx
index e8b6a1c263..ea1e1aec0a 100644
--- a/webview/src/plots/components/templatePlots/TemplatePlotsGrid.tsx
+++ b/webview/src/plots/components/templatePlots/TemplatePlotsGrid.tsx
@@ -1,22 +1,21 @@
 import cx from 'classnames'
-import { Section, TemplatePlotEntry } from 'dvc/src/plots/webview/contract'
+import { TemplatePlotEntry } from 'dvc/src/plots/webview/contract'
 import { reorderObjectList } from 'dvc/src/util/array'
 import React, { useEffect, useState, useCallback } from 'react'
 import { useDispatch, useSelector } from 'react-redux'
-import { VisualizationSpec } from 'react-vega'
 import { changeDisabledDragIds, changeSize } from './templatePlotsSlice'
+import { TemplatePlotsGridItem } from './TemplatePlotsGridItem'
 import { VirtualizedGrid } from '../../../shared/components/virtualizedGrid/VirtualizedGrid'
 import {
   DragDropContainer,
   OnDrop,
   WrapperProps
 } from '../../../shared/components/dragDrop/DragDropContainer'
-import { withScale } from '../../../util/styles'
 import { DropTarget } from '../DropTarget'
 import styles from '../styles.module.scss'
-import { ZoomablePlot } from '../ZoomablePlot'
 import { PlotsState } from '../../store'
 import { useSnapPoints } from '../../hooks/useSnapPoints'
+import { withScale } from '../../../util/styles'
 
 interface TemplatePlotsGridProps {
   entries: TemplatePlotEntry[]
@@ -30,11 +29,6 @@ interface TemplatePlotsGridProps {
   parentDraggedOver?: boolean
 }
 
-const autoSize = {
-  height: 'container',
-  width: 'container'
-}
-
 export const TemplatePlotsGrid: React.FC<TemplatePlotsGridProps> = ({
   entries,
   groupId,
@@ -53,7 +47,6 @@ export const TemplatePlotsGrid: React.FC<TemplatePlotsGridProps> = ({
   const disabledDragPlotIds = useSelector(
     (state: PlotsState) => state.template.disabledDragPlotIds
   )
-  const currentSize = useSelector((state: PlotsState) => state.template.size)
 
   const addDisabled = useCallback(
     (e: Event) => {
@@ -87,14 +80,14 @@ export const TemplatePlotsGrid: React.FC<TemplatePlotsGridProps> = ({
     }
   }, [addDisabled, removeDisabled, disableClick])
 
-  const addEventsOnViewReady = () => {
+  const addEventsOnViewReady = useCallback(() => {
     const panels = document.querySelectorAll('.vega-bindings')
     for (const panel of Object.values(panels)) {
       panel.addEventListener('mouseenter', addDisabled)
       panel.addEventListener('mouseleave', removeDisabled)
       panel.addEventListener('click', disableClick)
     }
-  }
+  }, [addDisabled, removeDisabled, disableClick])
 
   const setEntriesOrder = (order: string[]) => {
     setOrder(order)
@@ -116,31 +109,22 @@ export const TemplatePlotsGrid: React.FC<TemplatePlotsGridProps> = ({
   })
 
   const items = reorderedItems.map((plot: TemplatePlotEntry) => {
-    const { id, content, multiView, revisions } = plot
+    const { id, revisions } = plot
     const nbRevisions = (multiView && revisions?.length) || 1
 
-    const toggleDrag = (enabled: boolean) => {
-      dispatch(changeDisabledDragIds(enabled ? [] : [id]))
-    }
-
     return (
       <div
         key={id}
+        id={id}
         className={plotClassName}
         data-testid={`plot_${id}`}
-        id={id}
         style={withScale(nbRevisions)}
       >
-        <ZoomablePlot
-          id={id}
-          spec={{ ...content, ...autoSize } as VisualizationSpec}
-          onViewReady={addEventsOnViewReady}
-          toggleDrag={toggleDrag}
-          changeSize={changeSize}
-          currentSnapPoint={currentSize}
-          shouldNotResize={multiView}
-          section={Section.TEMPLATE_PLOTS}
+        <TemplatePlotsGridItem
+          plot={plot}
+          addEventsOnViewReady={addEventsOnViewReady}
           snapPoints={snapPoints}
+          changeSize={changeSize}
         />
       </div>
     )
diff --git a/webview/src/plots/components/templatePlots/TemplatePlotsGridItem.tsx b/webview/src/plots/components/templatePlots/TemplatePlotsGridItem.tsx
new file mode 100644
index 0000000000..2a6b7bbbd4
--- /dev/null
+++ b/webview/src/plots/components/templatePlots/TemplatePlotsGridItem.tsx
@@ -0,0 +1,50 @@
+import { AnyAction } from '@reduxjs/toolkit'
+import { Section, TemplatePlotEntry } from 'dvc/src/plots/webview/contract'
+import React, { memo, useMemo } from 'react'
+import { useDispatch, useSelector } from 'react-redux'
+import { VisualizationSpec } from 'react-vega'
+import { changeDisabledDragIds } from './templatePlotsSlice'
+import { ZoomablePlot } from '../ZoomablePlot'
+import { SnapPoints } from '../../hooks/useSnapPoints'
+import { PlotsState } from '../../store'
+
+interface TemplatePlotsGridItemProps {
+  plot: TemplatePlotEntry
+  addEventsOnViewReady: () => void
+  snapPoints: SnapPoints
+  changeSize: (size: number) => AnyAction
+}
+
+const TemplatePlotsGridItemComponent: React.FC<TemplatePlotsGridItemProps> = ({
+  plot,
+  addEventsOnViewReady,
+  snapPoints,
+  changeSize
+}) => {
+  const dispatch = useDispatch()
+  const currentSize = useSelector((state: PlotsState) => state.template.size)
+  const { id, content, multiView } = plot
+  const toggleDrag = (enabled: boolean) => {
+    dispatch(changeDisabledDragIds(enabled ? [] : [id]))
+  }
+  const spec = useMemo(
+    () => ({ ...content, height: 'container', width: 'container' }),
+    [content]
+  ) as VisualizationSpec
+
+  return (
+    <ZoomablePlot
+      id={id}
+      spec={spec}
+      onViewReady={addEventsOnViewReady}
+      toggleDrag={toggleDrag}
+      changeSize={changeSize}
+      currentSnapPoint={currentSize}
+      shouldNotResize={multiView}
+      section={Section.TEMPLATE_PLOTS}
+      snapPoints={snapPoints}
+    />
+  )
+}
+
+export const TemplatePlotsGridItem = memo(TemplatePlotsGridItemComponent)
diff --git a/webview/src/util/wdyr.ts b/webview/src/util/wdyr.ts
index 801a885d4c..d65f830622 100644
--- a/webview/src/util/wdyr.ts
+++ b/webview/src/util/wdyr.ts
@@ -5,6 +5,6 @@ import React from 'react'
 if (process.env.NODE_ENV === 'development') {
   const whyDidYouRender = require('@welldone-software/why-did-you-render')
   whyDidYouRender(React, {
-    trackAllPureComponents: true
+    trackAllPureComponents: false
   })
 }

From fa821a22f92bb4f7c6a8513371c943ae561e774e Mon Sep 17 00:00:00 2001
From: Stephanie Roy <stephanie@iterative.ai>
Date: Wed, 22 Feb 2023 11:25:43 -0500
Subject: [PATCH 4/7] Do not drill down vega plot content

---
 webview/src/plots/components/App.test.tsx     |  19 +--
 webview/src/plots/components/Plots.tsx        |   4 +-
 webview/src/plots/components/ZoomablePlot.tsx |  29 +++--
 .../checkpointPlots/CheckpointPlot.tsx        |  10 +-
 .../checkpointPlots/checkpointPlotsSlice.ts   |   9 +-
 webview/src/plots/components/plotDataStore.ts |  22 ++--
 .../components/templatePlots/AddedSection.tsx |   6 +-
 .../templatePlots/TemplatePlots.tsx           | 120 +++++++++---------
 .../templatePlots/TemplatePlotsGrid.tsx       |  95 ++++++++------
 .../templatePlots/TemplatePlotsGridItem.tsx   |  50 --------
 .../templatePlots/templatePlotsSlice.ts       |  24 +++-
 .../plots/components/templatePlots/util.ts    |  28 ++--
 webview/src/plots/components/webviewSlice.ts  |  19 ++-
 webview/src/plots/hooks/useGetPlot.ts         |  51 ++++++++
 webview/src/plots/hooks/useSnapPoints.ts      |  16 ---
 15 files changed, 256 insertions(+), 246 deletions(-)
 delete mode 100644 webview/src/plots/components/templatePlots/TemplatePlotsGridItem.tsx
 create mode 100644 webview/src/plots/hooks/useGetPlot.ts
 delete mode 100644 webview/src/plots/hooks/useSnapPoints.ts

diff --git a/webview/src/plots/components/App.test.tsx b/webview/src/plots/components/App.test.tsx
index 68ac919ed4..d0d562c877 100644
--- a/webview/src/plots/components/App.test.tsx
+++ b/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
diff --git a/webview/src/plots/components/Plots.tsx b/webview/src/plots/components/Plots.tsx
index 1d5bc127af..8f990648df 100644
--- a/webview/src/plots/components/Plots.tsx
+++ b/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)
diff --git a/webview/src/plots/components/ZoomablePlot.tsx b/webview/src/plots/components/ZoomablePlot.tsx
index fef065baea..ef90539cb1 100644
--- a/webview/src/plots/components/ZoomablePlot.tsx
+++ b/webview/src/plots/components/ZoomablePlot.tsx
@@ -3,43 +3,44 @@ import cx from 'classnames'
 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 } from 'react-redux'
-import { PlainObject, VisualizationSpec } from 'react-vega'
+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'
-import { SnapPoints } from '../hooks/useSnapPoints'
 
 interface ZoomablePlotProps {
-  spec: VisualizationSpec
-  data?: PlainObject
+  spec?: VisualizationSpec
   id: string
   onViewReady?: () => void
-  toggleDrag: (enabled: boolean) => void
+  toggleDrag: (enabled: boolean, id: string) => void
   changeSize: (size: number) => AnyAction
   currentSnapPoint: number
   section: Section
-  snapPoints: SnapPoints
   shouldNotResize?: boolean
 }
 
 export const ZoomablePlot: React.FC<ZoomablePlotProps> = ({
-  spec,
-  data,
+  spec: createdSpec,
   id,
   onViewReady,
   toggleDrag,
   changeSize,
   currentSnapPoint,
   section,
-  snapPoints,
   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>()
@@ -60,6 +61,7 @@ export const ZoomablePlot: React.FC<ZoomablePlotProps> = ({
   currentPlotProps.current = plotProps
 
   useEffect(() => {
+    // TODO Review this as this should be handled by the useGetPlot
     if (previousSpecsAndData.current !== newSpecsAndData) {
       dispatch(
         setZoomedInPlot({ id, plot: currentPlotProps.current, refresh: true })
@@ -91,10 +93,10 @@ export const ZoomablePlot: React.FC<ZoomablePlotProps> = ({
   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
@@ -102,6 +104,9 @@ export const ZoomablePlot: React.FC<ZoomablePlotProps> = ({
     },
     onResize
   }
+  if (!data && !spec) {
+    return null
+  }
   return (
     <button
       className={cx(styles.zoomablePlot, {
diff --git a/webview/src/plots/components/checkpointPlots/CheckpointPlot.tsx b/webview/src/plots/components/checkpointPlots/CheckpointPlot.tsx
index 1f490acf12..f6a1addf87 100644
--- a/webview/src/plots/components/checkpointPlots/CheckpointPlot.tsx
+++ b/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 { useSnapPoints } from '../../hooks/useSnapPoints'
 
 interface CheckpointPlotProps {
   id: string
@@ -23,9 +22,8 @@ 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 snapPoints = useSnapPoints()
 
   const spec = useMemo(() => {
     const title = plot?.title
@@ -36,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) => {
@@ -55,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}
         toggleDrag={toggleDrag}
         changeSize={changeSize}
         currentSnapPoint={currentSize}
         section={Section.CHECKPOINT_PLOTS}
-        snapPoints={snapPoints}
       />
     </div>
   )
diff --git a/webview/src/plots/components/checkpointPlots/checkpointPlotsSlice.ts b/webview/src/plots/components/checkpointPlots/checkpointPlotsSlice.ts
index 2ea0890849..88dba386e2 100644
--- a/webview/src/plots/components/checkpointPlots/checkpointPlotsSlice.ts
+++ b/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,
diff --git a/webview/src/plots/components/plotDataStore.ts b/webview/src/plots/components/plotDataStore.ts
index ca675ba78e..76e106ab0e 100644
--- a/webview/src/plots/components/plotDataStore.ts
+++ b/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]
     }
   }
 }
diff --git a/webview/src/plots/components/templatePlots/AddedSection.tsx b/webview/src/plots/components/templatePlots/AddedSection.tsx
index 0ec00bccad..786f5b1524 100644
--- a/webview/src/plots/components/templatePlots/AddedSection.tsx
+++ b/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[]
 }
 
diff --git a/webview/src/plots/components/templatePlots/TemplatePlots.tsx b/webview/src/plots/components/templatePlots/TemplatePlots.tsx
index 0664b4c7d7..a1b4af4fef 100644
--- a/webview/src/plots/components/templatePlots/TemplatePlots.tsx
+++ b/webview/src/plots/components/templatePlots/TemplatePlots.tsx
@@ -1,22 +1,23 @@
-import {
-  TemplatePlotEntry,
-  TemplatePlotGroup,
-  TemplatePlotSection
-} from 'dvc/src/plots/webview/contract'
-import React, { DragEvent, useState, useEffect, useRef } from 'react'
+import { TemplatePlotGroup } from 'dvc/src/plots/webview/contract'
+import React, {
+  DragEvent,
+  useState,
+  useEffect,
+  useRef,
+  useCallback
+} from 'react'
 import cx from 'classnames'
 import { useDispatch, useSelector } from 'react-redux'
 import { MessageFromWebviewType } from 'dvc/src/webview/contract'
-import { reorderObjectList } from 'dvc/src/util/array'
 import { AddedSection } from './AddedSection'
 import { TemplatePlotsGrid } from './TemplatePlotsGrid'
+import { PlotGroup } from './templatePlotsSlice'
 import { removeFromPreviousAndAddToNewSection } from './util'
 import { sendMessage } from '../../../shared/vscode'
 import { createIDWithIndex, getIDIndex } from '../../../util/ids'
 import styles from '../styles.module.scss'
 import { shouldUseVirtualizedGrid } from '../util'
 import { PlotsState } from '../../store'
-import { plotDataStore } from '../plotDataStore'
 import { setDraggedOverGroup } from '../../../shared/components/dragDrop/dragDropSlice'
 import { EmptyState } from '../../../shared/components/emptyState/EmptyState'
 import { isSameGroup } from '../../../shared/components/dragDrop/DragDropContainer'
@@ -29,7 +30,7 @@ export enum NewSectionBlock {
 }
 
 export const TemplatePlots: React.FC = () => {
-  const { plotsSnapshot, size } = useSelector(
+  const { plotSections, size } = useSelector(
     (state: PlotsState) => state.template
   )
   const draggedOverGroup = useSelector(
@@ -42,7 +43,7 @@ export const TemplatePlots: React.FC = () => {
     (state: PlotsState) => state.webview.selectedRevisions
   )
 
-  const [sections, setSections] = useState<TemplatePlotSection[]>([])
+  const [sections, setSections] = useState<PlotGroup[]>(plotSections)
   const [hoveredSection, setHoveredSection] = useState('')
   const nbItemsPerRow = size
   const shouldSendMessage = useRef(true)
@@ -50,15 +51,15 @@ export const TemplatePlots: React.FC = () => {
 
   useEffect(() => {
     shouldSendMessage.current = false
-    setSections(plotDataStore.template)
-  }, [plotsSnapshot, setSections])
+    setSections(plotSections)
+  }, [plotSections, setSections])
 
   useEffect(() => {
     if (sections && shouldSendMessage.current) {
       sendMessage({
         payload: sections.map(section => ({
           group: section.group,
-          paths: section.entries.map(({ id }) => id)
+          paths: section.entries
         })),
         type: MessageFromWebviewType.REORDER_PLOTS_TEMPLATES
       })
@@ -66,6 +67,45 @@ export const TemplatePlots: React.FC = () => {
     shouldSendMessage.current = true
   }, [sections])
 
+  const handleDropInSection = useCallback(
+    (
+      draggedId: string,
+      draggedGroup: string,
+      groupId: string,
+      position?: number
+    ) => {
+      if (draggedGroup === groupId) {
+        return
+      }
+      const oldGroupId = getIDIndex(draggedGroup)
+      const newGroupId = getIDIndex(groupId)
+      const updatedSections = removeFromPreviousAndAddToNewSection(
+        sections,
+        oldGroupId,
+        draggedId,
+        newGroupId,
+        position
+      )
+
+      setSections(updatedSections)
+    },
+    [setSections, sections]
+  )
+
+  const setSectionEntries = useCallback(
+    (index: number, entries: string[]) => {
+      setSections(sections => {
+        const updatedSections = [...sections]
+        updatedSections[index] = {
+          ...sections[index],
+          entries
+        }
+        return updatedSections
+      })
+    },
+    [setSections]
+  )
+
   if (sectionIsLoading(selectedRevisions)) {
     return <LoadingSection />
   }
@@ -73,17 +113,6 @@ export const TemplatePlots: React.FC = () => {
     return <EmptyState isFullScreen={false}>No Plots to Display</EmptyState>
   }
 
-  const setSectionEntries = (index: number, entries: TemplatePlotEntry[]) => {
-    setSections(sections => {
-      const updatedSections = [...sections]
-      updatedSections[index] = {
-        ...sections[index],
-        entries
-      }
-      return updatedSections
-    })
-  }
-
   const firstSection = sections[0]
   const lastSection = sections.slice(-1)[0]
 
@@ -104,14 +133,11 @@ export const TemplatePlots: React.FC = () => {
       draggedId
     )
 
-    const { group, entries } = sections[draggedSectionId]
+    const { group } = sections[draggedSectionId]
 
     setHoveredSection('')
-    const entry = entries.find(
-      ({ id }) => id === draggedId
-    ) as TemplatePlotEntry
     const newSection = {
-      entries: [entry],
+      entries: [draggedId],
       group
     }
 
@@ -126,32 +152,6 @@ export const TemplatePlots: React.FC = () => {
     }
   }
 
-  const handleDropInSection = (
-    draggedId: string,
-    draggedGroup: string,
-    groupId: string,
-    position?: number
-  ) => {
-    if (draggedGroup === groupId) {
-      return
-    }
-    const oldGroupId = getIDIndex(draggedGroup)
-    const newGroupId = getIDIndex(groupId)
-    const entry = sections[oldGroupId].entries.find(
-      ({ id }) => id === draggedId
-    ) as TemplatePlotEntry
-    const updatedSections = removeFromPreviousAndAddToNewSection(
-      sections,
-      oldGroupId,
-      draggedId,
-      newGroupId,
-      entry,
-      position
-    )
-
-    setSections(updatedSections)
-  }
-
   const handleEnteringSection = (groupId: string) => {
     dispatch(setDraggedOverGroup(groupId))
   }
@@ -192,17 +192,13 @@ export const TemplatePlots: React.FC = () => {
           }
 
           if (draggedRef.group === groupId) {
-            const order = section.entries.map(s => s.id)
+            const order = section.entries
             const updatedSections = [...sections]
 
             const newOrder = changeOrderWithDraggedInfo(order, draggedRef)
             updatedSections[i] = {
               ...sections[i],
-              entries: reorderObjectList<TemplatePlotEntry>(
-                newOrder,
-                section.entries,
-                'id'
-              )
+              entries: newOrder
             }
             setSections(updatedSections)
           } else if (isSameGroup(draggedRef.group, groupId)) {
@@ -231,7 +227,6 @@ export const TemplatePlots: React.FC = () => {
             onDrop={handleDropAtTheEnd}
           >
             <TemplatePlotsGrid
-              entries={section.entries}
               groupId={groupId}
               groupIndex={i}
               onDropInSection={handleDropInSection}
@@ -240,6 +235,7 @@ export const TemplatePlots: React.FC = () => {
               useVirtualizedGrid={useVirtualizedGrid}
               nbItemsPerRow={nbItemsPerRow}
               parentDraggedOver={draggedOverGroup === groupId}
+              entries={section.entries}
             />
           </div>
         )
diff --git a/webview/src/plots/components/templatePlots/TemplatePlotsGrid.tsx b/webview/src/plots/components/templatePlots/TemplatePlotsGrid.tsx
index ea1e1aec0a..3a2efc9522 100644
--- a/webview/src/plots/components/templatePlots/TemplatePlotsGrid.tsx
+++ b/webview/src/plots/components/templatePlots/TemplatePlotsGrid.tsx
@@ -1,10 +1,8 @@
 import cx from 'classnames'
-import { TemplatePlotEntry } from 'dvc/src/plots/webview/contract'
-import { reorderObjectList } from 'dvc/src/util/array'
-import React, { useEffect, useState, useCallback } from 'react'
+import { Section } from 'dvc/src/plots/webview/contract'
+import React, { useEffect, useState, useCallback, useMemo } from 'react'
 import { useDispatch, useSelector } from 'react-redux'
 import { changeDisabledDragIds, changeSize } from './templatePlotsSlice'
-import { TemplatePlotsGridItem } from './TemplatePlotsGridItem'
 import { VirtualizedGrid } from '../../../shared/components/virtualizedGrid/VirtualizedGrid'
 import {
   DragDropContainer,
@@ -14,23 +12,23 @@ import {
 import { DropTarget } from '../DropTarget'
 import styles from '../styles.module.scss'
 import { PlotsState } from '../../store'
-import { useSnapPoints } from '../../hooks/useSnapPoints'
 import { withScale } from '../../../util/styles'
+import { ZoomablePlot } from '../ZoomablePlot'
+import { plotDataStore } from '../plotDataStore'
 
 interface TemplatePlotsGridProps {
-  entries: TemplatePlotEntry[]
   groupId: string
   groupIndex: number
   onDropInSection: OnDrop
   multiView: boolean
-  setSectionEntries: (groupIndex: number, entries: TemplatePlotEntry[]) => void
+  setSectionEntries: (groupIndex: number, entries: string[]) => void
   useVirtualizedGrid?: boolean
   nbItemsPerRow: number
+  entries: string[]
   parentDraggedOver?: boolean
 }
 
 export const TemplatePlotsGrid: React.FC<TemplatePlotsGridProps> = ({
-  entries,
   groupId,
   groupIndex,
   onDropInSection,
@@ -38,11 +36,12 @@ export const TemplatePlotsGrid: React.FC<TemplatePlotsGridProps> = ({
   setSectionEntries,
   useVirtualizedGrid,
   nbItemsPerRow,
+  entries,
   parentDraggedOver
 }) => {
   const dispatch = useDispatch()
   const [order, setOrder] = useState<string[]>([])
-  const snapPoints = useSnapPoints()
+  const currentSize = useSelector((state: PlotsState) => state.template.size)
 
   const disabledDragPlotIds = useSelector(
     (state: PlotsState) => state.template.disabledDragPlotIds
@@ -66,7 +65,7 @@ export const TemplatePlotsGrid: React.FC<TemplatePlotsGridProps> = ({
   }, [])
 
   useEffect(() => {
-    setOrder(entries.map(({ id }) => id))
+    setOrder(entries)
   }, [entries])
 
   useEffect(() => {
@@ -92,43 +91,57 @@ export const TemplatePlotsGrid: React.FC<TemplatePlotsGridProps> = ({
   const setEntriesOrder = (order: string[]) => {
     setOrder(order)
 
-    setSectionEntries(
-      groupIndex,
-      reorderObjectList<TemplatePlotEntry>(order, entries, 'id')
-    )
+    setSectionEntries(groupIndex, order)
   }
 
-  const reorderedItems = reorderObjectList<TemplatePlotEntry>(
-    order,
-    entries,
-    'id'
-  )
-
   const plotClassName = cx(styles.plot, {
     [styles.multiViewPlot]: multiView
   })
 
-  const items = reorderedItems.map((plot: TemplatePlotEntry) => {
-    const { id, revisions } = plot
-    const nbRevisions = (multiView && revisions?.length) || 1
-
-    return (
-      <div
-        key={id}
-        id={id}
-        className={plotClassName}
-        data-testid={`plot_${id}`}
-        style={withScale(nbRevisions)}
-      >
-        <TemplatePlotsGridItem
-          plot={plot}
-          addEventsOnViewReady={addEventsOnViewReady}
-          snapPoints={snapPoints}
-          changeSize={changeSize}
-        />
-      </div>
-    )
-  })
+  const toggleDrag = useCallback(
+    (enabled: boolean, id: string) => {
+      dispatch(changeDisabledDragIds(enabled ? [] : [id]))
+    },
+    [dispatch]
+  )
+
+  const items = useMemo(
+    () =>
+      order.map((plot: string) => {
+        const nbRevisions =
+          (multiView &&
+            plotDataStore[Section.TEMPLATE_PLOTS][plot].revisions?.length) ||
+          1
+
+        return (
+          <div
+            key={plot}
+            id={plot}
+            className={plotClassName}
+            data-testid={`plot_${plot}`}
+            style={withScale(nbRevisions)}
+          >
+            <ZoomablePlot
+              id={plot}
+              onViewReady={addEventsOnViewReady}
+              toggleDrag={toggleDrag}
+              changeSize={changeSize}
+              currentSnapPoint={currentSize}
+              shouldNotResize={multiView}
+              section={Section.TEMPLATE_PLOTS}
+            />
+          </div>
+        )
+      }),
+    [
+      order,
+      plotClassName,
+      addEventsOnViewReady,
+      currentSize,
+      multiView,
+      toggleDrag
+    ]
+  )
 
   return (
     <DragDropContainer
diff --git a/webview/src/plots/components/templatePlots/TemplatePlotsGridItem.tsx b/webview/src/plots/components/templatePlots/TemplatePlotsGridItem.tsx
deleted file mode 100644
index 2a6b7bbbd4..0000000000
--- a/webview/src/plots/components/templatePlots/TemplatePlotsGridItem.tsx
+++ /dev/null
@@ -1,50 +0,0 @@
-import { AnyAction } from '@reduxjs/toolkit'
-import { Section, TemplatePlotEntry } from 'dvc/src/plots/webview/contract'
-import React, { memo, useMemo } from 'react'
-import { useDispatch, useSelector } from 'react-redux'
-import { VisualizationSpec } from 'react-vega'
-import { changeDisabledDragIds } from './templatePlotsSlice'
-import { ZoomablePlot } from '../ZoomablePlot'
-import { SnapPoints } from '../../hooks/useSnapPoints'
-import { PlotsState } from '../../store'
-
-interface TemplatePlotsGridItemProps {
-  plot: TemplatePlotEntry
-  addEventsOnViewReady: () => void
-  snapPoints: SnapPoints
-  changeSize: (size: number) => AnyAction
-}
-
-const TemplatePlotsGridItemComponent: React.FC<TemplatePlotsGridItemProps> = ({
-  plot,
-  addEventsOnViewReady,
-  snapPoints,
-  changeSize
-}) => {
-  const dispatch = useDispatch()
-  const currentSize = useSelector((state: PlotsState) => state.template.size)
-  const { id, content, multiView } = plot
-  const toggleDrag = (enabled: boolean) => {
-    dispatch(changeDisabledDragIds(enabled ? [] : [id]))
-  }
-  const spec = useMemo(
-    () => ({ ...content, height: 'container', width: 'container' }),
-    [content]
-  ) as VisualizationSpec
-
-  return (
-    <ZoomablePlot
-      id={id}
-      spec={spec}
-      onViewReady={addEventsOnViewReady}
-      toggleDrag={toggleDrag}
-      changeSize={changeSize}
-      currentSnapPoint={currentSize}
-      shouldNotResize={multiView}
-      section={Section.TEMPLATE_PLOTS}
-      snapPoints={snapPoints}
-    />
-  )
-}
-
-export const TemplatePlotsGridItem = memo(TemplatePlotsGridItemComponent)
diff --git a/webview/src/plots/components/templatePlots/templatePlotsSlice.ts b/webview/src/plots/components/templatePlots/templatePlotsSlice.ts
index a142c60c33..eef678f2bc 100644
--- a/webview/src/plots/components/templatePlots/templatePlotsSlice.ts
+++ b/webview/src/plots/components/templatePlots/templatePlotsSlice.ts
@@ -3,14 +3,17 @@ import {
   DEFAULT_SECTION_COLLAPSED,
   DEFAULT_SECTION_SIZES,
   Section,
+  TemplatePlotGroup,
   TemplatePlotsData
 } from 'dvc/src/plots/webview/contract'
-import { plotDataStore } from '../plotDataStore'
+import { addPlotsWithSnapshots, removePlots } from '../plotDataStore'
 
+export type PlotGroup = { group: TemplatePlotGroup; entries: string[] }
 export interface TemplatePlotsState extends Omit<TemplatePlotsData, 'plots'> {
   isCollapsed: boolean
   hasData: boolean
-  plotsSnapshot: string
+  plotSections: PlotGroup[]
+  plotsSnapshots: { [key: string]: string }
   disabledDragPlotIds: string[]
 }
 
@@ -18,7 +21,8 @@ export const templatePlotsInitialState: TemplatePlotsState = {
   disabledDragPlotIds: [],
   hasData: false,
   isCollapsed: DEFAULT_SECTION_COLLAPSED[Section.TEMPLATE_PLOTS],
-  plotsSnapshot: '',
+  plotSections: [],
+  plotsSnapshots: {},
   size: DEFAULT_SECTION_SIZES[Section.TEMPLATE_PLOTS]
 }
 
@@ -36,17 +40,27 @@ export const templatePlotsSlice = createSlice({
       state.isCollapsed = action.payload
     },
     update: (state, action: PayloadAction<TemplatePlotsData>) => {
-      plotDataStore.template = action.payload?.plots
       if (!action.payload) {
         return templatePlotsInitialState
       }
 
+      const plotSections = action.payload.plots?.map(section => ({
+        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)
+      removePlots(plotsIds, Section.TEMPLATE_PLOTS)
+
       return {
         ...state,
         hasData: !!action.payload,
-        plotsSnapshot: JSON.stringify(action.payload.plots),
+        plotSections,
+        plotsSnapshots: snapShots,
         size: Math.abs(action.payload.size)
       }
+      // TODO: add sections in here with the entries ids
     }
   }
 })
diff --git a/webview/src/plots/components/templatePlots/util.ts b/webview/src/plots/components/templatePlots/util.ts
index c072ee03e6..3beb6414fd 100644
--- a/webview/src/plots/components/templatePlots/util.ts
+++ b/webview/src/plots/components/templatePlots/util.ts
@@ -1,51 +1,43 @@
-import {
-  TemplatePlotSection,
-  TemplatePlotEntry
-} from 'dvc/src/plots/webview/contract'
+import { PlotGroup } from './templatePlotsSlice'
 
-const remove = (section: TemplatePlotSection, entryId: string) => {
-  const entries = section.entries.filter(({ id }) => id !== entryId)
+const remove = (section: PlotGroup, entryId: string) => {
+  const entries = section.entries.filter(id => id !== entryId)
   return {
     entries,
     group: section.group
   }
 }
 
-const add = (
-  section: TemplatePlotSection,
-  entry: TemplatePlotEntry,
-  position?: number
-) => {
+const add = (section: PlotGroup, entryId: string, position?: number) => {
   const entries = [...section.entries]
   entries.splice(
     position === undefined ? entries.length - 1 : position,
     0,
-    entry
+    entryId
   )
   return { entries, group: section.group }
 }
 
-const cleanup = (section: TemplatePlotSection) =>
+const cleanup = (section: PlotGroup) =>
   section.entries.length > 0 ? section : null
 
 export const removeFromPreviousAndAddToNewSection = (
-  sections: TemplatePlotSection[],
+  sections: PlotGroup[],
   oldSectionIndex: number,
   entryId: string,
   newGroupIndex?: number,
-  entry?: TemplatePlotEntry,
   position?: number
 ) => {
   const newSections = sections.map((section, i) => {
     if (i === oldSectionIndex) {
       return remove(section, entryId)
-    } else if (i === newGroupIndex && entry) {
-      return add(section, entry, position)
+    } else if (i === newGroupIndex) {
+      return add(section, entryId, position)
     }
     return section
   })
 
   return newSections
     .map(section => cleanup(section))
-    .filter(Boolean) as TemplatePlotSection[]
+    .filter(Boolean) as PlotGroup[]
 }
diff --git a/webview/src/plots/components/webviewSlice.ts b/webview/src/plots/components/webviewSlice.ts
index 43bf51cfa5..ab767233e4 100644
--- a/webview/src/plots/components/webviewSlice.ts
+++ b/webview/src/plots/components/webviewSlice.ts
@@ -8,21 +8,23 @@ type ZoomedInPlotState = {
   id: string
   refresh?: boolean
 }
+
+type SnapPoints = [number, number, number, number]
 export interface WebviewState {
   hasData: boolean
   hasPlots: boolean
   hasUnselectedPlots: boolean
   selectedRevisions: Revision[]
   zoomedInPlot: ZoomedInPlotState | undefined
-  maxPlotSize: number
+  snapPoints: SnapPoints
 }
 
 export const webviewInitialState: WebviewState = {
   hasData: false,
   hasPlots: false,
   hasUnselectedPlots: false,
-  maxPlotSize: 375,
   selectedRevisions: [],
+  snapPoints: [0, 0, 0, 0],
   zoomedInPlot: {
     id: '',
     plot: undefined
@@ -36,11 +38,16 @@ export const webviewSlice = createSlice({
     initialize: (state: { hasData: boolean }) => {
       state.hasData = true
     },
-    setMaxPlotSize: (
-      state: { maxPlotSize: number },
+    setSnapPoints: (
+      state: { snapPoints: SnapPoints },
       action: PayloadAction<number>
     ) => {
-      state.maxPlotSize = action.payload
+      state.snapPoints = [
+        action.payload,
+        action.payload / 2,
+        action.payload / 3,
+        action.payload / 4
+      ]
     },
     setZoomedInPlot: (
       // eslint-disable-next-line @typescript-eslint/no-explicit-any
@@ -86,7 +93,7 @@ export const {
   updateHasUnselectedPlots,
   updateSelectedRevisions,
   setZoomedInPlot,
-  setMaxPlotSize
+  setSnapPoints
 } = webviewSlice.actions
 
 export default webviewSlice.reducer
diff --git a/webview/src/plots/hooks/useGetPlot.ts b/webview/src/plots/hooks/useGetPlot.ts
new file mode 100644
index 0000000000..16e6a5b9b6
--- /dev/null
+++ b/webview/src/plots/hooks/useGetPlot.ts
@@ -0,0 +1,51 @@
+import {
+  CheckpointPlotData,
+  CheckpointPlotValues,
+  Section,
+  TemplatePlotEntry
+} from 'dvc/src/plots/webview/contract'
+import { useCallback, useEffect, useState } from 'react'
+import { useSelector } from 'react-redux'
+import { VisualizationSpec } from 'react-vega'
+import { plotDataStore } from '../components/plotDataStore'
+import { PlotsState } from '../store'
+
+export const useGetPlot = (
+  section: Section,
+  id: string,
+  spec?: VisualizationSpec
+) => {
+  const isCheckpointPlot = section === Section.CHECKPOINT_PLOTS
+  const storeSection = isCheckpointPlot ? 'checkpoint' : 'template'
+  const snapshot = useSelector(
+    (state: PlotsState) => state[storeSection].plotsSnapshots
+  )
+  const [data, setData] = useState<CheckpointPlotValues | undefined>(undefined)
+  const [content, setContent] = useState<VisualizationSpec | undefined>(spec)
+
+  const setPlotData = useCallback(() => {
+    const plot = plotDataStore[section][id]
+    if (!plot) {
+      return
+    }
+
+    if (isCheckpointPlot) {
+      setData((plot as CheckpointPlotData).values)
+      setContent(spec)
+      return
+    }
+
+    setData(undefined)
+    setContent({
+      ...(plot as TemplatePlotEntry).content,
+      height: 'container',
+      width: 'container'
+    } as VisualizationSpec)
+  }, [id, isCheckpointPlot, setData, setContent, section, spec])
+
+  useEffect(() => {
+    setPlotData()
+  }, [snapshot, setPlotData])
+
+  return { content, data }
+}
diff --git a/webview/src/plots/hooks/useSnapPoints.ts b/webview/src/plots/hooks/useSnapPoints.ts
deleted file mode 100644
index 53ebaf3633..0000000000
--- a/webview/src/plots/hooks/useSnapPoints.ts
+++ /dev/null
@@ -1,16 +0,0 @@
-import { useEffect, useState } from 'react'
-import { useSelector } from 'react-redux'
-import { PlotsState } from '../store'
-
-export type SnapPoints = [number, number, number, number]
-
-export const useSnapPoints = () => {
-  const maxSize = useSelector((state: PlotsState) => state.webview.maxPlotSize)
-  const [snapPoints, setSnapPoints] = useState<SnapPoints>([0, 0, 0, 0])
-
-  useEffect(() => {
-    setSnapPoints([maxSize, maxSize / 2, maxSize / 3, maxSize / 4])
-  }, [maxSize])
-
-  return snapPoints
-}

From 2ef17568dff4794fd7c615ca08f654a88c4056b4 Mon Sep 17 00:00:00 2001
From: Stephanie Roy <stephanie@iterative.ai>
Date: Wed, 22 Feb 2023 16:06:14 -0500
Subject: [PATCH 5/7] Fix checkpoint plots

---
 .../plots/components/checkpointPlots/CheckpointPlot.tsx    | 2 +-
 webview/src/plots/hooks/useGetPlot.ts                      | 7 +++----
 2 files changed, 4 insertions(+), 5 deletions(-)

diff --git a/webview/src/plots/components/checkpointPlots/CheckpointPlot.tsx b/webview/src/plots/components/checkpointPlots/CheckpointPlot.tsx
index f6a1addf87..38f3163b2f 100644
--- a/webview/src/plots/components/checkpointPlots/CheckpointPlot.tsx
+++ b/webview/src/plots/components/checkpointPlots/CheckpointPlot.tsx
@@ -51,7 +51,7 @@ export const CheckpointPlot: React.FC<CheckpointPlotProps> = ({
     <div className={styles.plot} data-testid={key} id={id} style={withScale(1)}>
       <ZoomablePlot
         spec={spec}
-        id={key}
+        id={id}
         toggleDrag={toggleDrag}
         changeSize={changeSize}
         currentSnapPoint={currentSize}
diff --git a/webview/src/plots/hooks/useGetPlot.ts b/webview/src/plots/hooks/useGetPlot.ts
index 16e6a5b9b6..74f7275cda 100644
--- a/webview/src/plots/hooks/useGetPlot.ts
+++ b/webview/src/plots/hooks/useGetPlot.ts
@@ -1,12 +1,11 @@
 import {
   CheckpointPlotData,
-  CheckpointPlotValues,
   Section,
   TemplatePlotEntry
 } from 'dvc/src/plots/webview/contract'
 import { useCallback, useEffect, useState } from 'react'
 import { useSelector } from 'react-redux'
-import { VisualizationSpec } from 'react-vega'
+import { PlainObject, VisualizationSpec } from 'react-vega'
 import { plotDataStore } from '../components/plotDataStore'
 import { PlotsState } from '../store'
 
@@ -20,7 +19,7 @@ export const useGetPlot = (
   const snapshot = useSelector(
     (state: PlotsState) => state[storeSection].plotsSnapshots
   )
-  const [data, setData] = useState<CheckpointPlotValues | undefined>(undefined)
+  const [data, setData] = useState<PlainObject | undefined>(undefined)
   const [content, setContent] = useState<VisualizationSpec | undefined>(spec)
 
   const setPlotData = useCallback(() => {
@@ -30,7 +29,7 @@ export const useGetPlot = (
     }
 
     if (isCheckpointPlot) {
-      setData((plot as CheckpointPlotData).values)
+      setData({ values: (plot as CheckpointPlotData).values })
       setContent(spec)
       return
     }

From 3d26208b0dd95a56a126bc2dd53b856f22dfc16b Mon Sep 17 00:00:00 2001
From: Stephanie Roy <stephanie@iterative.ai>
Date: Thu, 23 Feb 2023 09:32:30 -0500
Subject: [PATCH 6/7] Add linting rule for TODOs

---
 .eslintrc.js                                                     | 1 +
 webview/src/plots/components/ZoomablePlot.tsx                    | 1 -
 webview/src/plots/components/templatePlots/templatePlotsSlice.ts | 1 -
 3 files changed, 1 insertion(+), 2 deletions(-)

diff --git a/.eslintrc.js b/.eslintrc.js
index 2e42a0d122..251a6a0ff0 100644
--- a/.eslintrc.js
+++ b/.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',
diff --git a/webview/src/plots/components/ZoomablePlot.tsx b/webview/src/plots/components/ZoomablePlot.tsx
index ef90539cb1..1c0c6f7e4b 100644
--- a/webview/src/plots/components/ZoomablePlot.tsx
+++ b/webview/src/plots/components/ZoomablePlot.tsx
@@ -61,7 +61,6 @@ export const ZoomablePlot: React.FC<ZoomablePlotProps> = ({
   currentPlotProps.current = plotProps
 
   useEffect(() => {
-    // TODO Review this as this should be handled by the useGetPlot
     if (previousSpecsAndData.current !== newSpecsAndData) {
       dispatch(
         setZoomedInPlot({ id, plot: currentPlotProps.current, refresh: true })
diff --git a/webview/src/plots/components/templatePlots/templatePlotsSlice.ts b/webview/src/plots/components/templatePlots/templatePlotsSlice.ts
index eef678f2bc..97c1c5c066 100644
--- a/webview/src/plots/components/templatePlots/templatePlotsSlice.ts
+++ b/webview/src/plots/components/templatePlots/templatePlotsSlice.ts
@@ -60,7 +60,6 @@ export const templatePlotsSlice = createSlice({
         plotsSnapshots: snapShots,
         size: Math.abs(action.payload.size)
       }
-      // TODO: add sections in here with the entries ids
     }
   }
 })

From 12d254676ae52d9e49985a1626271094d1303ce2 Mon Sep 17 00:00:00 2001
From: Stephanie Roy <stephanie@iterative.ai>
Date: Thu, 23 Feb 2023 09:35:51 -0500
Subject: [PATCH 7/7] Made a few variables clearer

---
 .../src/plots/components/templatePlots/TemplatePlotsGrid.tsx  | 4 ++--
 webview/src/plots/components/webviewSlice.ts                  | 1 +
 2 files changed, 3 insertions(+), 2 deletions(-)

diff --git a/webview/src/plots/components/templatePlots/TemplatePlotsGrid.tsx b/webview/src/plots/components/templatePlots/TemplatePlotsGrid.tsx
index 3a2efc9522..87f2373e34 100644
--- a/webview/src/plots/components/templatePlots/TemplatePlotsGrid.tsx
+++ b/webview/src/plots/components/templatePlots/TemplatePlotsGrid.tsx
@@ -108,7 +108,7 @@ export const TemplatePlotsGrid: React.FC<TemplatePlotsGridProps> = ({
   const items = useMemo(
     () =>
       order.map((plot: string) => {
-        const nbRevisions =
+        const colSpan =
           (multiView &&
             plotDataStore[Section.TEMPLATE_PLOTS][plot].revisions?.length) ||
           1
@@ -119,7 +119,7 @@ export const TemplatePlotsGrid: React.FC<TemplatePlotsGridProps> = ({
             id={plot}
             className={plotClassName}
             data-testid={`plot_${plot}`}
-            style={withScale(nbRevisions)}
+            style={withScale(colSpan)}
           >
             <ZoomablePlot
               id={plot}
diff --git a/webview/src/plots/components/webviewSlice.ts b/webview/src/plots/components/webviewSlice.ts
index ab767233e4..17e7e435b5 100644
--- a/webview/src/plots/components/webviewSlice.ts
+++ b/webview/src/plots/components/webviewSlice.ts
@@ -42,6 +42,7 @@ export const webviewSlice = createSlice({
       state: { snapPoints: SnapPoints },
       action: PayloadAction<number>
     ) => {
+      // Action payload here is the max size of the plots webview. When changed, we re-calculate the snap points.
       state.snapPoints = [
         action.payload,
         action.payload / 2,