Refactor experiment table branch data

extension/src/experiments/model/index.ts

@@ -26,7 +26,8 @@ import {
   isQueued,
   isRunning,
   GitRemoteStatus,
-  RunningExperiment
+  RunningExperiment,
+  WORKSPACE_BRANCH
 } from '../webview/contract'
 import { reorderListSubset } from '../../util/array'
 import {
@@ -413,7 +414,7 @@ export class ExperimentsModel extends ModelWithPersistence {
     const commitsBySha = this.applyFiltersToCommits()
 
     const rows: Commit[] = [
-      { branch: undefined, ...this.addDetails(this.workspace) }
+      { branch: WORKSPACE_BRANCH, ...this.addDetails(this.workspace) }
     ]
     for (const { branch, sha } of this.rowOrder) {
       const commit = commitsBySha[sha]
@@ -520,6 +521,10 @@ export class ExperimentsModel extends ModelWithPersistence {
     return [this.currentBranch as string, ...this.selectedBranches]
   }
 
+  public getSelectedBranches() {
+    return this.selectedBranches
+  }
+
   public getAvailableBranchesToShow() {
     return this.availableBranchesToShow
   }

extension/src/experiments/webview/contract.ts

@@ -53,7 +53,7 @@ export type Experiment = {
   starred?: boolean
   status?: ExperimentStatus
   timestamp?: string | null
-  branch?: string | undefined
+  branch?: string | typeof WORKSPACE_BRANCH
 }
 
 export const isRunning = (status: ExperimentStatus | undefined): boolean =>
@@ -92,6 +92,8 @@ export type Column = {
   width?: number
 }
 
+export const WORKSPACE_BRANCH = null
+
 export type TableData = {
   changes: string[]
   cliError: string | null
@@ -107,6 +109,7 @@ export type TableData = {
   hasRunningWorkspaceExperiment: boolean
   isShowingMoreCommits: Record<string, boolean>
   rows: Commit[]
+  selectedBranches: string[]
   selectedForPlotsCount: number
   sorts: SortDefinition[]
 }

extension/src/experiments/webview/messages.ts

@@ -276,6 +276,7 @@ export class WebviewMessages {
         this.experiments.hasRunningWorkspaceExperiment(),
       isShowingMoreCommits: this.experiments.getIsShowingMoreCommits(),
       rows: this.experiments.getRowData(),
+      selectedBranches: this.experiments.getSelectedBranches(),
       selectedForPlotsCount: this.experiments.getSelectedRevisions().length,
       sorts: this.experiments.getSorts()
     }

extension/src/test/fixtures/expShow/base/rows.ts

@@ -1,7 +1,8 @@
 import { join } from '../../../util/path'
 import {
   Commit,
-  GitRemoteStatus
+  GitRemoteStatus,
+  WORKSPACE_BRANCH
 } from '../../../../experiments/webview/contract'
 import { copyOriginalColors } from '../../../../experiments/model/status/colors'
 import { shortenForLabel } from '../../../../util/string'
@@ -20,7 +21,7 @@ const colorsList = copyOriginalColors()
 
 const rowsFixture: Commit[] = [
   {
-    branch: undefined,
+    branch: WORKSPACE_BRANCH,
     commit: undefined,
     description: undefined,
     deps: {

extension/src/test/fixtures/expShow/base/tableData.ts

@@ -3,6 +3,7 @@ import rowsFixture from './rows'
 import columnsFixture from './columns'
 
 const tableDataFixture: TableData = {
+  selectedBranches: [],
   cliError: null,
   changes: [],
   columnOrder: [],

extension/src/test/fixtures/expShow/dataTypes/rows.ts

@@ -1,9 +1,12 @@
 import { EXPERIMENT_WORKSPACE_ID } from '../../../../cli/dvc/contract'
-import { Commit } from '../../../../experiments/webview/contract'
+import {
+  Commit,
+  WORKSPACE_BRANCH
+} from '../../../../experiments/webview/contract'
 
 const data: Commit[] = [
   {
-    branch: undefined,
+    branch: WORKSPACE_BRANCH,
     displayColor: undefined,
     id: EXPERIMENT_WORKSPACE_ID,
     label: EXPERIMENT_WORKSPACE_ID,

extension/src/test/fixtures/expShow/deeplyNested/rows.ts

@@ -1,8 +1,9 @@
 import { EXPERIMENT_WORKSPACE_ID } from '../../../../cli/dvc/contract'
+import { WORKSPACE_BRANCH } from '../../../../experiments/webview/contract'
 
 export const data = [
   {
-    branch: undefined,
+    branch: WORKSPACE_BRANCH,
     id: EXPERIMENT_WORKSPACE_ID,
     label: EXPERIMENT_WORKSPACE_ID,
     params: {

extension/src/test/fixtures/expShow/survival/rows.ts

@@ -1,10 +1,13 @@
 import { join } from '../../../util/path'
-import { Commit } from '../../../../experiments/webview/contract'
+import {
+  Commit,
+  WORKSPACE_BRANCH
+} from '../../../../experiments/webview/contract'
 import { EXPERIMENT_WORKSPACE_ID } from '../../../../cli/dvc/contract'
 
 const data: Commit[] = [
   {
-    branch: undefined,
+    branch: WORKSPACE_BRANCH,
     id: EXPERIMENT_WORKSPACE_ID,
     label: EXPERIMENT_WORKSPACE_ID,
     metrics: {

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

@@ -1728,7 +1728,6 @@ describe('App', () => {
 
       const multipleBranches = {
         ...tableDataFixture,
-        branches,
         hasData: true,
         rows: [
           workspace as Commit,
@@ -1747,7 +1746,8 @@ describe('App', () => {
             branch: branches[2],
             subRows: undefined
           }))
-        ]
+        ],
+        selectedBranches: branches.slice(1)
       }
 
       renderTable(multipleBranches)

webview/src/experiments/components/App.tsx

@@ -8,6 +8,7 @@ import { TableData } from 'dvc/src/experiments/webview/contract'
 import Experiments from './Experiments'
 import {
   update,
+  updateSelectedBranches,
   updateChanges,
   updateCliError,
   updateColumnOrder,
@@ -87,6 +88,9 @@ export const App: React.FC<Record<string, unknown>> = () => {
               case 'rows':
                 dispatch(updateRows(data.data.rows))
                 continue
+              case 'selectedBranches':
+                dispatch(updateSelectedBranches(data.data.selectedBranches))
+                continue
               case 'selectedForPlotsCount':
                 dispatch(
                   updateSelectedForPlotsCount(data.data.selectedForPlotsCount)

webview/src/experiments/components/table/Indicators.tsx

@@ -83,9 +83,8 @@ export const Indicators = () => {
     (state: ExperimentsState) => state.tableData.selectedForPlotsCount
   )
 
-  const branchesSelected = useSelector(
-    (state: ExperimentsState) =>
-      Math.max(state.tableData.branches.filter(Boolean).length - 1, 0) // We always have one branch by default (the current one which is not selected) and undefined for the workspace
+  const branchesSelected = useSelector((state: ExperimentsState) =>
+    Math.max(state.tableData.selectedBranches.length, 0)
   )
   const { hasBranchesToSelect } = useSelector(
     (state: ExperimentsState) => state.tableData

webview/src/experiments/components/table/RowSelectionContext.tsx

@@ -61,7 +61,6 @@ export const RowSelectionProvider: React.FC<{ children: React.ReactNode }> = ({
           row: {
             original: { id, branch }
           }
-          // eslint-disable-next-line @typescript-eslint/no-unsafe-return
         }) => getCompositeId(id, branch)
       ),
       ...selectionHistory

webview/src/experiments/components/table/Table.test.tsx

@@ -261,7 +261,6 @@ describe('Table', () => {
 
       const tableDataWithColumnSetting: TableDataState = {
         ...sortingTableDataFixture,
-        branches: ['main'],
         columnWidths
       }
       render(
@@ -302,7 +301,6 @@ describe('Table', () => {
 
       const tableDataWithColumnSetting: TableDataState = {
         ...sortingTableDataFixture,
-        branches: ['main'],
         columnWidths
       }
       render(

webview/src/experiments/components/table/body/RowContextMenu.tsx

@@ -3,6 +3,7 @@ import { useSelector } from 'react-redux'
 import { MessageFromWebviewType } from 'dvc/src/webview/contract'
 import {
   ExperimentStatus,
+  WORKSPACE_BRANCH,
   isQueued,
   isRunning,
   isRunningInQueue
@@ -292,7 +293,7 @@ const getSingleSelectMenuOptions = (
 
 const getContextMenuOptions = (
   id: string,
-  branch: string | undefined,
+  branch: string | undefined | typeof WORKSPACE_BRANCH,
   isWorkspace: boolean,
   projectHasCheckpoints: boolean,
   hasRunningWorkspaceExperiment: boolean,

webview/src/experiments/components/table/body/TableContent.test.tsx

@@ -808,12 +808,12 @@ describe('TableContent', () => {
     })
   } as unknown as Table<Experiment>
 
-  const renderTableContent = (rowsInstance = instance, branches = ['main']) => {
+  const renderTableContent = (rowsInstance = instance) => {
     return render(
       <Provider
         store={configureStore({
           preloadedState: {
-            tableData: { ...tableData, branches }
+            tableData
           },
           reducer: experimentsReducers
         })}
@@ -845,7 +845,7 @@ describe('TableContent', () => {
         ]
       })
     } as unknown as Table<Experiment>
-    renderTableContent(multipleBranchesInstance, ['main', 'new-branch'])
+    renderTableContent(multipleBranchesInstance)
 
     expect(screen.getAllByTestId('branch-name').length).toBe(2)
     expect(screen.getByText('main')).toBeInTheDocument()

webview/src/experiments/components/table/body/TableContent.tsx

@@ -1,10 +1,8 @@
 import React, { Fragment, RefObject, useCallback, useContext } from 'react'
-import { useSelector } from 'react-redux'
-import { EXPERIMENT_WORKSPACE_ID } from 'dvc/src/cli/dvc/contract'
 import { TableBody } from './TableBody'
+import { collectBranchWithRows } from './util'
 import { BranchDivider } from './branchDivider/BranchDivider'
 import { RowSelectionContext } from '../RowSelectionContext'
-import { ExperimentsState } from '../../../store'
 import { InstanceProp, RowProp } from '../../../util/interfaces'
 
 interface TableContentProps extends InstanceProp {
@@ -19,7 +17,6 @@ export const TableContent: React.FC<TableContentProps> = ({
 }) => {
   const { rows, flatRows } = instance.getRowModel()
   const { batchSelection, lastSelectedRow } = useContext(RowSelectionContext)
-  const { branches } = useSelector((state: ExperimentsState) => state.tableData)
 
   const batchRowSelection = useCallback(
     ({ row: { id } }: RowProp) => {
@@ -52,11 +49,9 @@ export const TableContent: React.FC<TableContentProps> = ({
 
   return (
     <>
-      {branches.map(branch => {
-        const branchRows = rows.filter(row => row.original.branch === branch)
-
+      {collectBranchWithRows(rows).map(([branch, branchRows]) => {
         return (
-          <Fragment key={`${branch || EXPERIMENT_WORKSPACE_ID}`}>
+          <Fragment key={branch}>
             {branchRows.map((row, i) => {
               const isFirstBranchRow = branch && i === 0
               return (

webview/src/experiments/components/table/body/util.ts

@@ -0,0 +1,32 @@
+import { Row } from '@tanstack/react-table'
+import {
+  Experiment,
+  WORKSPACE_BRANCH
+} from 'dvc/src/experiments/webview/contract'
+
+export const collectBranchWithRows = (
+  rows: Row<Experiment>[]
+): [string | typeof WORKSPACE_BRANCH, Row<Experiment>[]][] => {
+  const branchesWithRows: [
+    string | typeof WORKSPACE_BRANCH,
+    Row<Experiment>[]
+  ][] = []
+
+  let branchRows = []
+
+  for (let i = 0; i < rows.length; i++) {
+    const row = rows[i]
+    const branch = row.original.branch
+    if (branch === undefined) {
+      continue
+    }
+    const next = rows[i + 1]
+    branchRows.push(row)
+    if (!next || next.original.branch !== branch) {
+      branchesWithRows.push([branch, branchRows])
+      branchRows = []
+    }
+  }
+
+  return branchesWithRows
+}

webview/src/experiments/state/tableDataSlice.ts

@@ -10,11 +10,9 @@ import { keepReferenceIfEqual } from '../../util/objects'
 
 export interface TableDataState extends TableData {
   hasData?: boolean
-  branches: (string | undefined)[]
 }
 
 export const tableDataInitialState: TableDataState = {
-  branches: ['current'],
   changes: [],
   cliError: null,
   columnOrder: [],
@@ -30,6 +28,7 @@ export const tableDataInitialState: TableDataState = {
   hasRunningWorkspaceExperiment: false,
   isShowingMoreCommits: {},
   rows: [],
+  selectedBranches: [],
   selectedForPlotsCount: 0,
   sorts: []
 }
@@ -112,15 +111,9 @@ export const tableDataSlice = createSlice({
         state.rows,
         action.payload
       ) as Experiment[]
-
-      const branches: (string | undefined)[] = []
-      for (const { branch } of state.rows) {
-        if (!branches.includes(branch)) {
-          branches.push(branch)
-        }
-      }
-
-      state.branches = branches
+    },
+    updateSelectedBranches: (state, action: PayloadAction<string[]>) => {
+      state.selectedBranches = action.payload
     },
     updateSelectedForPlotsCount: (state, action: PayloadAction<number>) => {
       state.selectedForPlotsCount = action.payload
@@ -139,8 +132,8 @@ export const {
   updateChanges,
   updateCliError,
   updateColumnOrder,
-  updateColumnWidths,
   updateColumns,
+  updateColumnWidths,
   updateFilters,
   updateHasBranchesToSelect,
   updateHasCheckpoints,
@@ -150,6 +143,7 @@ export const {
   updateHasRunningWorkspaceExperiment,
   updateIsShowingMoreCommits,
   updateRows,
+  updateSelectedBranches,
   updateSelectedForPlotsCount,
   updateSorts
 } = tableDataSlice.actions

webview/src/experiments/util/rows.ts

@@ -1 +1,4 @@
-export const getCompositeId = (id: string, branch = '') => `${id}-${branch}`
+export const getCompositeId = (
+  id: string,
+  branch: string | undefined | null = ''
+) => `${id}-${branch}`