Adds stashes & worktrees to Graph model

eamodio · eamodio · commit 1c989b05098c · 2024-09-06T18:20:03.000-04:00
diff --git a/src/env/node/git/localGitProvider.ts b/src/env/node/git/localGitProvider.ts
@@ -128,7 +128,7 @@ import type { GitTreeEntry } from '../../../git/models/tree';
 import type { GitUser } from '../../../git/models/user';
 import { isUserMatch } from '../../../git/models/user';
 import type { GitWorktree } from '../../../git/models/worktree';
-import { getWorktreeId, getWorktreesByBranch } from '../../../git/models/worktree';
+import { getWorktreeId, groupWorktreesByBranch } from '../../../git/models/worktree';
 import { parseGitBlame } from '../../../git/parsers/blameParser';
 import { parseGitBranches } from '../../../git/parsers/branchParser';
 import {
@@ -2319,21 +2319,25 @@ export class LocalGitProvider implements GitProvider, Disposable {
 		const refParser = getRefParser();
 		const statsParser = getGraphStatsParser();
 
-		const [refResult, stashResult, branchesResult, remotesResult, currentUserResult, worktreesByBranchResult] =
+		const [refResult, stashResult, branchesResult, remotesResult, currentUserResult, worktreesResult] =
 			await Promise.allSettled([
 				this.git.log(repoPath, undefined, ...refParser.arguments, '-n1', options?.ref ?? 'HEAD'),
 				this.getStash(repoPath),
 				this.getBranches(repoPath),
 				this.getRemotes(repoPath),
 				this.getCurrentUser(repoPath),
-				getWorktreesByBranch(this.container.git.getRepository(repoPath), { includeMainWorktree: true }),
+				this.container.git
+					.getWorktrees(repoPath)
+					.then(w => [w, groupWorktreesByBranch(w, { includeMainWorktree: true })]) satisfies Promise<
+					[GitWorktree[], Map<string, GitWorktree>]
+				>,
 			]);
 
 		const branches = getSettledValue(branchesResult)?.values;
 		const branchMap = branches != null ? new Map(branches.map(r => [r.name, r])) : new Map<string, GitBranch>();
 		const headBranch = branches?.find(b => b.current);
 		const headRefUpstreamName = headBranch?.upstream?.name;
-		const worktreesByBranch = getSettledValue(worktreesByBranchResult);
+		const [worktrees, worktreesByBranch] = getSettledValue(worktreesResult) ?? [[], new Map<string, GitWorktree>()];
 
 		let branchIdOfMainWorktree: string | undefined;
 		if (worktreesByBranch != null) {
@@ -2351,13 +2355,15 @@ export class LocalGitProvider implements GitProvider, Disposable {
 
 		const downstreamMap = new Map<string, string[]>();
 
+		let stashes: Map<string, GitStashCommit> | undefined;
 		let stdin: string | undefined;
 
 		// TODO@eamodio this is insanity -- there *HAS* to be a better way to get git log to return stashes
 		const stash = getSettledValue(stashResult);
 		if (stash?.commits.size) {
+			stashes = new Map(stash.commits);
 			stdin = join(
-				map(stash.commits.values(), c => c.sha.substring(0, 9)),
+				map(stashes.values(), c => c.sha.substring(0, 9)),
 				'\n',
 			);
 		}
@@ -2421,6 +2427,8 @@ export class LocalGitProvider implements GitProvider, Disposable {
 									branches: branchMap,
 									remotes: remoteMap,
 									downstreams: downstreamMap,
+									stashes: stashes,
+									worktrees: worktrees,
 									worktreesByBranch: worktreesByBranch,
 									rows: [],
 								};
@@ -2444,6 +2452,8 @@ export class LocalGitProvider implements GitProvider, Disposable {
 						branches: branchMap,
 						remotes: remoteMap,
 						downstreams: downstreamMap,
+						stashes: stashes,
+						worktrees: worktrees,
 						worktreesByBranch: worktreesByBranch,
 						rows: [],
 					};
@@ -2838,6 +2848,8 @@ export class LocalGitProvider implements GitProvider, Disposable {
 				branches: branchMap,
 				remotes: remoteMap,
 				downstreams: downstreamMap,
+				stashes: stashes,
+				worktrees: worktrees,
 				worktreesByBranch: worktreesByBranch,
 				rows: rows,
 				id: sha,
diff --git a/src/git/models/graph.ts b/src/git/models/graph.ts
@@ -10,6 +10,7 @@ import type {
 import type { GkProviderId } from '../../gk/models/repositoryIdentities';
 import type { Brand, Unbrand } from '../../system/brand';
 import type { GitBranch } from './branch';
+import type { GitStashCommit } from './commit';
 import type { GitRemote } from './remote';
 import type { GitWorktree } from './worktree';
 
@@ -50,6 +51,8 @@ export interface GitGraph {
 	readonly branches: Map<string, GitBranch>;
 	readonly remotes: Map<string, GitRemote>;
 	readonly downstreams: Map<string, string[]>;
+	readonly stashes: Map<string, GitStashCommit> | undefined;
+	readonly worktrees: GitWorktree[] | undefined;
 	readonly worktreesByBranch: Map<string, GitWorktree> | undefined;
 
 	/** The rows for the set of commits requested */
diff --git a/src/git/models/worktree.ts b/src/git/models/worktree.ts
@@ -345,12 +345,10 @@ export async function getWorktreesByBranch(
 	if (repos == null) return worktreesByBranch;
 
 	async function addWorktrees(repo: Repository) {
-		const worktrees = await repo.getWorktrees();
-		for (const wt of worktrees) {
-			if (wt.branch == null || (!options?.includeMainWorktree && wt.main)) continue;
-
-			worktreesByBranch.set(wt.branch.id, wt);
-		}
+		groupWorktreesByBranch(await repo.getWorktrees(), {
+			includeMainWorktree: options?.includeMainWorktree,
+			worktreesByBranch: worktreesByBranch,
+		});
 	}
 
 	if (!Array.isArray(repos)) {
@@ -362,6 +360,22 @@ export async function getWorktreesByBranch(
 	return worktreesByBranch;
 }
 
+export function groupWorktreesByBranch(
+	worktrees: GitWorktree[],
+	options?: { includeMainWorktree?: boolean; worktreesByBranch?: Map<string, GitWorktree> },
+) {
+	const worktreesByBranch = options?.worktreesByBranch ?? new Map<string, GitWorktree>();
+	if (worktrees == null) return worktreesByBranch;
+
+	for (const wt of worktrees) {
+		if (wt.branch == null || (!options?.includeMainWorktree && wt.main)) continue;
+
+		worktreesByBranch.set(wt.branch.id, wt);
+	}
+
+	return worktreesByBranch;
+}
+
 export function getOpenedWorktreesByBranch(
 	worktreesByBranch: Map<string, GitWorktree> | undefined,
 ): Set<string> | undefined {
diff --git a/src/plus/integrations/providers/github/githubGitProvider.ts b/src/plus/integrations/providers/github/githubGitProvider.ts
@@ -42,7 +42,7 @@ import { GitUri } from '../../../../git/gitUri';
 import type { GitBlame, GitBlameAuthor, GitBlameLine, GitBlameLines } from '../../../../git/models/blame';
 import type { BranchSortOptions } from '../../../../git/models/branch';
 import { getBranchId, getBranchNameWithoutRemote, GitBranch, sortBranches } from '../../../../git/models/branch';
-import type { GitCommitLine } from '../../../../git/models/commit';
+import type { GitCommitLine, GitStashCommit } from '../../../../git/models/commit';
 import { getChangedFilesCount, GitCommit, GitCommitIdentity } from '../../../../git/models/commit';
 import { deletedOrMissing, uncommitted } from '../../../../git/models/constants';
 import { GitContributor } from '../../../../git/models/contributor';
@@ -84,6 +84,7 @@ import { getTagId, GitTag, sortTags } from '../../../../git/models/tag';
 import type { GitTreeEntry } from '../../../../git/models/tree';
 import type { GitUser } from '../../../../git/models/user';
 import { isUserMatch } from '../../../../git/models/user';
+import type { GitWorktree } from '../../../../git/models/worktree';
 import { getRemoteProviderMatcher, loadRemoteProviders } from '../../../../git/remotes/remoteProviders';
 import type { GitSearch, GitSearchResultData, GitSearchResults } from '../../../../git/search';
 import { getSearchQueryComparisonKey, parseSearchQuery } from '../../../../git/search';
@@ -1329,6 +1330,9 @@ export class GitHubGitProvider implements GitProvider, Disposable {
 			getSettledValue(currentUserResult),
 			avatars,
 			ids,
+			undefined,
+			undefined,
+			undefined,
 			{ ...options, useAvatars: useAvatars },
 		);
 	}
@@ -1346,6 +1350,9 @@ export class GitHubGitProvider implements GitProvider, Disposable {
 		currentUser: GitUser | undefined,
 		avatars: Map<string, string>,
 		ids: Set<string>,
+		stashes: Map<string, GitStashCommit> | undefined,
+		worktrees: GitWorktree[] | undefined,
+		worktreesByBranch: Map<string, GitWorktree> | undefined,
 		options?: {
 			branch?: string;
 			include?: { stats?: boolean };
@@ -1365,7 +1372,9 @@ export class GitHubGitProvider implements GitProvider, Disposable {
 				branches: branchMap,
 				remotes: remoteMap,
 				downstreams: downstreamMap,
-				worktreesByBranch: undefined,
+				stashes: stashes,
+				worktrees: worktrees,
+				worktreesByBranch: worktreesByBranch,
 				rows: [],
 			};
 		}
@@ -1380,7 +1389,9 @@ export class GitHubGitProvider implements GitProvider, Disposable {
 				branches: branchMap,
 				remotes: remoteMap,
 				downstreams: downstreamMap,
-				worktreesByBranch: undefined,
+				stashes: stashes,
+				worktrees: worktrees,
+				worktreesByBranch: worktreesByBranch,
 				rows: [],
 			};
 		}
@@ -1629,7 +1640,9 @@ export class GitHubGitProvider implements GitProvider, Disposable {
 			branches: branchMap,
 			remotes: remoteMap,
 			downstreams: downstreamMap,
-			worktreesByBranch: undefined,
+			stashes: stashes,
+			worktrees: worktrees,
+			worktreesByBranch: worktreesByBranch,
 			rows: rows,
 			id: options?.ref,
 
@@ -1653,6 +1666,9 @@ export class GitHubGitProvider implements GitProvider, Disposable {
 					currentUser,
 					avatars,
 					ids,
+					stashes,
+					worktrees,
+					worktreesByBranch,
 					options,
 				);
 			},
