Add a mapping from revision to branches while using only one exp show call

extension/src/cli/dvc/contract.ts

@@ -1,3 +1,4 @@
+import { CommitData } from '../../experiments/webview/contract'
 import { Plot } from '../../plots/webview/contract'
 
 export const MIN_CLI_VERSION = '2.58.1'
@@ -98,13 +99,17 @@ export type ExpWithError = {
   rev: string
   name?: string
   branch: string | undefined
+  commit?: CommitData
+  description?: string
 } & DvcError
 
 type ExpWithData = {
   rev: string
   name?: string
   branch: string | undefined
   data: ExpData
+  commit?: CommitData
+  description?: string
 }
 
 export type ExpState = ExpWithData | ExpWithError

extension/src/cli/git/reader.ts

@@ -49,15 +49,19 @@ export class GitReader extends GitCli {
     return !output
   }
 
-  public async getCommitMessages(cwd: string, sha: string): Promise<string> {
+  public async getCommitMessages(
+    cwd: string,
+    revision: string,
+    revisions: string
+  ): Promise<string> {
     const options = getOptions(
       cwd,
       Command.LOG,
-      sha,
+      revision,
       Flag.PRETTY_FORMAT_COMMIT_MESSAGE,
       Flag.SEPARATE_WITH_NULL,
       Flag.NUMBER,
-      '1'
+      revisions
     )
     try {
       return await this.executeProcess(options)

extension/src/data/index.ts

@@ -10,7 +10,14 @@ import { DeferredDisposable } from '../class/deferred'
 import { isSameOrChild } from '../fileSystem'
 
 export abstract class BaseData<
-  T extends { data: PlotsOutputOrError; revs: string[] } | ExpShowOutput
+  T extends
+    | { data: PlotsOutputOrError; revs: string[] }
+    | {
+        currentBranch: string
+        expShow: ExpShowOutput
+        order: { branch: string; sha: string }[]
+        gitLog: string
+      }
 > extends DeferredDisposable {
   public readonly onDidUpdate: Event<T>
   public readonly onDidChangeDvcYaml: Event<void>

extension/src/experiments/data/index.ts

@@ -7,14 +7,19 @@ import {
 import { getRelativePattern } from '../../fileSystem/relativePattern'
 import { createFileSystemWatcher } from '../../fileSystem/watcher'
 import { AvailableCommands, InternalCommands } from '../../commands/internal'
-import { EXPERIMENT_WORKSPACE_ID, ExpShowOutput } from '../../cli/dvc/contract'
+import { ExpShowOutput } from '../../cli/dvc/contract'
 import { BaseData } from '../../data'
-import { DOT_DVC, ExperimentFlag } from '../../cli/dvc/constants'
-import { gitPath } from '../../cli/git/constants'
+import { Args, DOT_DVC, ExperimentFlag } from '../../cli/dvc/constants'
+import { COMMITS_SEPARATOR, gitPath } from '../../cli/git/constants'
 import { getGitPath } from '../../fileSystem'
 import { ExperimentsModel } from '../model'
 
-export class ExperimentsData extends BaseData<ExpShowOutput> {
+export class ExperimentsData extends BaseData<{
+  currentBranch: string
+  expShow: ExpShowOutput
+  order: { branch: string; sha: string }[]
+  gitLog: string
+}> {
   private readonly experiments: ExperimentsModel
 
   constructor(
@@ -47,43 +52,58 @@ export class ExperimentsData extends BaseData<ExpShowOutput> {
     )
 
     void this.updateAvailableBranchesToSelect(allBranches)
-    const data: ExpShowOutput = []
 
     const { branches, currentBranch } = await this.getBranchesToShowWithCurrent(
       allBranches
     )
+    let gitLog = ''
+    const order: { branch: string; sha: string }[] = []
+    const args: Args = []
 
-    await Promise.all(
-      branches.map(async branch => {
-        const branchFlags = [
-          ExperimentFlag.REV,
-          branch,
-          ExperimentFlag.NUM_COMMIT,
-          this.experiments.getNbOfCommitsToShow(branch).toString()
-        ]
-
-        const output = (await this.expShow(
-          branchFlags,
-          branch
-        )) as ExpShowOutput
-
-        if (branch !== currentBranch) {
-          const workspaceIndex = output.findIndex(
-            exp => exp.rev === EXPERIMENT_WORKSPACE_ID
-          )
-          output.splice(workspaceIndex, 1)
-        }
-        data.push(...output)
-      })
+    for (const branch of branches) {
+      gitLog = await this.collectGitLogAndOrder(gitLog, branch, order, args)
+    }
+
+    const expShow = await this.internalCommands.executeCommand<ExpShowOutput>(
+      AvailableCommands.EXP_SHOW,
+      this.dvcRoot,
+      ...args
     )
 
-    this.collectFiles(data)
+    this.collectFiles({ expShow })
 
-    return this.notifyChanged(data)
+    return this.notifyChanged({ currentBranch, expShow, gitLog, order })
   }
 
-  protected collectFiles(data: ExpShowOutput) {
-    this.collectedFiles = collectFiles(data, this.collectedFiles)
+  protected collectFiles({ expShow }: { expShow: ExpShowOutput }) {
+    this.collectedFiles = collectFiles(expShow, this.collectedFiles)
+  }
+
+  private async collectGitLogAndOrder(
+    gitLog: string,
+    branch: string,
+    order: { branch: string; sha: string }[],
+    args: Args
+  ) {
+    const nbOfCommitsToShow = this.experiments.getNbOfCommitsToShow(branch)
+
+    const branchGitLog = await this.internalCommands.executeCommand(
+      AvailableCommands.GIT_GET_COMMIT_MESSAGES,
+      this.dvcRoot,
+      branch,
+      String(nbOfCommitsToShow)
+    )
+    gitLog = [gitLog, branchGitLog].join(COMMITS_SEPARATOR)
+
+    for (const commit of branchGitLog.split(COMMITS_SEPARATOR)) {
+      const [sha] = commit.split('\n')
+      order.push({ branch, sha })
+      if (args.includes(sha)) {
+        continue
+      }
+      args.push(ExperimentFlag.REV, sha)
+    }
+    return gitLog
   }
 
   private async getBranchesToShowWithCurrent(allBranches: string[]) {
@@ -94,24 +114,16 @@ export class ExperimentsData extends BaseData<ExpShowOutput> {
 
     this.experiments.pruneBranchesToShow(allBranches)
 
-    const branches = this.experiments.getBranchesToShow()
+    const branches = [
+      currentBranch,
+      ...this.experiments
+        .getBranchesToShow()
+        .filter(branch => branch !== currentBranch)
+    ]
 
-    if (!branches.includes(currentBranch)) {
-      branches.push(currentBranch)
-      this.experiments.setBranchesToShow(branches)
-    }
     return { branches, currentBranch }
   }
 
-  private async expShow(flags: (ExperimentFlag | string)[], branch?: string) {
-    const data = await this.internalCommands.executeCommand<ExpShowOutput>(
-      AvailableCommands.EXP_SHOW,
-      this.dvcRoot,
-      ...flags
-    )
-    return data.map(exp => ({ ...exp, branch }))
-  }
-
   private async updateAvailableBranchesToSelect(branches?: string[]) {
     const allBranches =
       branches ||

extension/src/experiments/index.ts

@@ -10,6 +10,7 @@ import omit from 'lodash.omit'
 import { addStarredToColumns } from './columns/like'
 import { setContextForEditorTitleIcons } from './context'
 import { ExperimentsModel } from './model'
+import { combineOutputs } from './util'
 import {
   pickExperiment,
   pickExperiments,
@@ -33,7 +34,7 @@ import { DecorationProvider } from './model/decorationProvider'
 import { starredFilter } from './model/filterBy/constants'
 import { ResourceLocator } from '../resourceLocator'
 import { AvailableCommands, InternalCommands } from '../commands/internal'
-import { EXPERIMENT_WORKSPACE_ID, ExpShowOutput } from '../cli/dvc/contract'
+import { ExpShowOutput } from '../cli/dvc/contract'
 import { ViewKey } from '../webview/constants'
 import { BaseRepository } from '../webview/repository'
 import { Title } from '../vscode/title'
@@ -43,7 +44,6 @@ import { Toast } from '../vscode/toast'
 import { ConfigKey } from '../vscode/config'
 import { checkSignalFile, pollSignalFileForProcess } from '../fileSystem'
 import { DVCLIVE_ONLY_RUNNING_SIGNAL_FILE } from '../cli/dvc/constants'
-import { COMMITS_SEPARATOR } from '../cli/git/constants'
 
 export const ExperimentsScale = {
   ...omit(ColumnType, 'TIMESTAMP'),
@@ -169,13 +169,24 @@ export class Experiments extends BaseRepository<TableData> {
     return this.data.managedUpdate()
   }
 
-  public async setState(data: ExpShowOutput) {
+  public async setState({
+    currentBranch,
+    expShow,
+    gitLog,
+    order
+  }: {
+    currentBranch: string
+    expShow: ExpShowOutput
+    order: { branch: string; sha: string }[]
+    gitLog: string
+  }) {
+    const data = combineOutputs(currentBranch, expShow, gitLog, order)
+
     const hadCheckpoints = this.hasCheckpoints()
     const dvcLiveOnly = await this.checkSignalFile()
-    const commitsOutput = await this.getCommitOutput(data)
     await Promise.all([
       this.columns.transformAndSet(data),
-      this.experiments.transformAndSet(data, dvcLiveOnly, commitsOutput)
+      this.experiments.transformAndSet(data, dvcLiveOnly)
     ])
 
     if (hadCheckpoints !== this.hasCheckpoints()) {
@@ -538,25 +549,6 @@ export class Experiments extends BaseRepository<TableData> {
     return this.webviewMessages.sendWebviewMessage()
   }
 
-  private async getCommitOutput(data: ExpShowOutput | undefined) {
-    if (!data || data.length === 0) {
-      return
-    }
-    let output = ''
-    for (const commit of data) {
-      if (commit.rev === EXPERIMENT_WORKSPACE_ID) {
-        continue
-      }
-      output += await this.internalCommands.executeCommand(
-        AvailableCommands.GIT_GET_COMMIT_MESSAGES,
-        this.dvcRoot,
-        commit.rev
-      )
-      output += COMMITS_SEPARATOR
-    }
-    return output
-  }
-
   private createWebviewMessageHandler() {
     const webviewMessages = new WebviewMessages(
       this.dvcRoot,

extension/src/experiments/model/collect.test.ts

@@ -1,14 +1,9 @@
 import { collectExperiments } from './collect'
-import { COMMITS_SEPARATOR } from '../../cli/git/constants'
 import { generateTestExpShowOutput } from '../../test/util/experiments'
 
 describe('collectExperiments', () => {
   it('should return an empty array if no commits are present', () => {
-    const { commits } = collectExperiments(
-      generateTestExpShowOutput({}),
-      false,
-      ''
-    )
+    const { commits } = collectExperiments(generateTestExpShowOutput({}), false)
     expect(commits).toStrictEqual([])
   })
 
@@ -26,19 +21,19 @@ describe('collectExperiments', () => {
   )
 
   it('should define a workspace', () => {
-    const { workspace } = collectExperiments(expShowWithTwoCommits, false, '')
+    const { workspace } = collectExperiments(expShowWithTwoCommits, false)
 
     expect(workspace).toBeDefined()
   })
 
   it('should find two branches from a repo with two branches', () => {
-    const { commits } = collectExperiments(expShowWithTwoCommits, false, '')
+    const { commits } = collectExperiments(expShowWithTwoCommits, false)
 
     expect(commits.length).toStrictEqual(2)
   })
 
   it('should list commits in the same order as they are collected', () => {
-    const { commits } = collectExperiments(expShowWithTwoCommits, false, '')
+    const { commits } = collectExperiments(expShowWithTwoCommits, false)
     const [commitA, commitB] = commits
 
     expect(commitA.id).toStrictEqual('branchA')
@@ -48,44 +43,19 @@ describe('collectExperiments', () => {
   it('should find two experiments on commitA', () => {
     const { experimentsByCommit } = collectExperiments(
       expShowWithTwoCommits,
-      false,
-      ''
+      false
     )
     expect(experimentsByCommit.get('branchA')?.length).toStrictEqual(2)
   })
 
   it('should find no experiments on branchB', () => {
     const { experimentsByCommit } = collectExperiments(
       expShowWithTwoCommits,
-      false,
-      ''
+      false
     )
     expect(experimentsByCommit.get('branchB')).toBeUndefined()
   })
 
-  it('should add data from git to commits if git log output is provided', () => {
-    const { commits } = collectExperiments(
-      expShowWithTwoCommits,
-      false,
-      `61bed4ce8913eca7f73ca754d65bc5daad1520e2\nJohn Smith\n3 days ago\nrefNames:tag: v.1.1\nmessage:add new feature${COMMITS_SEPARATOR}351e42ace3cb6a3a853c65bef285e60748cc6341\nrenovate[bot]\n5 weeks ago\nrefNames:\nmessage:update various dependencies\n* update dvc\n* update dvclive`
-    )
-    const [branch1, branch2] = commits
-    expect(branch1.description).toStrictEqual('add new feature')
-    expect(branch2.description).toStrictEqual('update various dependencies ...')
-    expect(branch1.commit).toStrictEqual({
-      author: 'John Smith',
-      date: '3 days ago',
-      message: 'add new feature',
-      tags: ['v.1.1']
-    })
-    expect(branch2.commit).toStrictEqual({
-      author: 'renovate[bot]',
-      date: '5 weeks ago',
-      message: 'update various dependencies\n* update dvc\n* update dvclive',
-      tags: []
-    })
-  })
-
   it('should not collect the same experiment twice', () => {
     const main = {
       experiments: [
@@ -114,8 +84,7 @@ describe('collectExperiments', () => {
 
     const { experimentsByCommit, commits } = collectExperiments(
       expShowWithDuplicateCommits,
-      false,
-      ''
+      false
     )
 
     expect(commits.length).toStrictEqual(3)