1use crate::shell::HideConsoleWindow;
8use std::path::{Path, PathBuf};
9use std::process::Command;
10
11#[derive(Debug, Clone)]
12pub struct WorktreeEntry {
13 pub path: PathBuf,
14 pub branch: String,
15 pub sanitized_branch: String,
16 pub namespace: Option<String>,
18}
19
20pub fn discover_worktrees(project_dir: &Path) -> Vec<WorktreeEntry> {
21 if project_dir.join(".jj").exists() {
23 discover_jj_workspaces(project_dir)
24 } else if project_dir.join(".git").exists() {
25 discover_git_worktrees(project_dir)
26 } else {
27 vec![]
28 }
29}
30
31fn discover_jj_workspaces(project_dir: &Path) -> Vec<WorktreeEntry> {
34 let output = match Command::new("jj")
35 .args(["--ignore-working-copy", "workspace", "list"])
36 .current_dir(project_dir)
37 .hide_console_window()
38 .output()
39 {
40 Ok(o) if o.status.success() => o.stdout,
41 _ => return vec![],
42 };
43
44 let names = parse_jj_workspace_names(&output);
45 if names.is_empty() {
46 return vec![];
47 }
48
49 let non_default: Vec<&str> = names
51 .iter()
52 .filter(|n| **n != "default")
53 .map(|n| n.as_str())
54 .collect();
55 let mut roots = std::collections::HashMap::with_capacity(non_default.len());
56
57 std::thread::scope(|s| {
58 let handles: Vec<_> = non_default
59 .iter()
60 .map(|name| s.spawn(move || (*name, get_jj_workspace_root(project_dir, name))))
61 .collect();
62
63 for handle in handles {
64 let (name, root) = handle.join().unwrap();
65 roots.insert(name.to_string(), root);
66 }
67 });
68
69 let mut entries = Vec::with_capacity(names.len());
70 for name in &names {
71 let path = if name == "default" {
72 Some(project_dir.to_path_buf())
73 } else {
74 roots.get(name).cloned().unwrap_or(None)
75 };
76
77 let Some(path) = path else {
78 continue;
79 };
80
81 let sanitized = sanitize_branch(name);
82 if sanitized.is_empty() {
83 log::warn!(
84 "Skipping jj workspace '{}' because its sanitized name is empty \
85 (no ASCII alphanumeric characters)",
86 name,
87 );
88 continue;
89 }
90 entries.push(WorktreeEntry {
91 path,
92 branch: name.to_string(),
93 sanitized_branch: sanitized,
94 namespace: None,
95 });
96 }
97
98 entries
99}
100fn parse_jj_workspace_names(stdout: &[u8]) -> Vec<String> {
101 let text = String::from_utf8_lossy(stdout);
102 text.lines()
103 .filter_map(|line| {
104 let line = line.trim();
105 if line.is_empty() {
106 return None;
107 }
108 line.split_once(':')
109 .map(|(name, _)| name.trim().to_string())
110 .filter(|n| !n.is_empty())
111 })
112 .collect()
113}
114
115#[allow(dead_code)]
116fn parse_jj_workspace_list(
117 stdout: &[u8],
118 mut resolve_path: impl FnMut(&str) -> Option<PathBuf>,
119) -> Vec<WorktreeEntry> {
120 let text = String::from_utf8_lossy(stdout);
121 let mut entries = Vec::new();
122
123 for line in text.lines() {
124 let line = line.trim();
125 if line.is_empty() {
126 continue;
127 }
128
129 let Some((name, _)) = line.split_once(':') else {
130 continue;
131 };
132 let name = name.trim();
133 if name.is_empty() {
134 continue;
135 }
136
137 let path = match resolve_path(name) {
138 Some(p) => p,
139 None => continue,
140 };
141
142 let sanitized = sanitize_branch(name);
143 if sanitized.is_empty() {
144 log::warn!(
145 "Skipping jj workspace '{}' because its sanitized name is empty \
146 (no ASCII alphanumeric characters)",
147 name,
148 );
149 continue;
150 }
151 entries.push(WorktreeEntry {
152 path,
153 branch: name.to_string(),
154 sanitized_branch: sanitized,
155 namespace: None,
156 });
157 }
158
159 entries
160}
161
162fn get_jj_workspace_root(project_dir: &Path, name: &str) -> Option<PathBuf> {
163 let output = Command::new("jj")
164 .args(["--ignore-working-copy", "workspace", "root", "--name", name])
165 .current_dir(project_dir)
166 .hide_console_window()
167 .output()
168 .ok()?;
169
170 if !output.status.success() {
171 return None;
172 }
173
174 let path_str = String::from_utf8_lossy(&output.stdout);
175 let trimmed = path_str.trim();
176 if trimmed.is_empty() {
177 return None;
178 }
179
180 Some(PathBuf::from(trimmed))
181}
182
183fn discover_git_worktrees(project_dir: &Path) -> Vec<WorktreeEntry> {
186 let git_dir = project_dir.join(".git");
193 if git_dir.is_dir() {
194 let worktrees_dir = git_dir.join("worktrees");
195 let has_linked = match std::fs::read_dir(&worktrees_dir) {
196 Ok(mut it) => it.next().is_some(),
197 Err(_) => false,
198 };
199 if !has_linked {
200 let path = project_dir
203 .canonicalize()
204 .unwrap_or_else(|_| project_dir.to_path_buf());
205 let branch = std::fs::read_to_string(git_dir.join("HEAD"))
207 .ok()
208 .and_then(|h| {
209 h.trim()
210 .strip_prefix("ref: refs/heads/")
211 .map(str::to_string)
212 });
213 let porcelain = match &branch {
214 Some(b) => format!("worktree {}\nbranch refs/heads/{b}\n\n", path.display()),
215 None => format!("worktree {}\n\n", path.display()),
216 };
217 return parse_git_worktree_output(porcelain.as_bytes());
218 }
219 }
220
221 let output = match Command::new("git")
224 .args(["worktree", "list", "--porcelain"])
225 .current_dir(project_dir)
226 .hide_console_window()
227 .output()
228 {
229 Ok(o) if o.status.success() => o.stdout,
230 _ => return vec![],
231 };
232
233 parse_git_worktree_output(&output)
234}
235
236fn parse_git_worktree_output(stdout: &[u8]) -> Vec<WorktreeEntry> {
237 let text = String::from_utf8_lossy(stdout);
238 let mut entries = Vec::new();
239 let mut current_path = None;
240 let mut current_branch = None;
241
242 for line in text.lines() {
243 if let Some(path) = line.strip_prefix("worktree ") {
244 current_path = Some(PathBuf::from(path.trim()));
245 current_branch = None;
246 } else if let Some(branch) = line.strip_prefix("branch ") {
247 current_branch = Some(
248 branch
249 .trim()
250 .strip_prefix("refs/heads/")
251 .unwrap_or(branch.trim())
252 .to_string(),
253 );
254 }
255 if line.is_empty() {
256 flush_git_entry(&mut entries, &mut current_path, &mut current_branch);
257 }
258 }
259
260 flush_git_entry(&mut entries, &mut current_path, &mut current_branch);
261
262 entries
263}
264
265fn flush_git_entry(
266 entries: &mut Vec<WorktreeEntry>,
267 path: &mut Option<PathBuf>,
268 branch: &mut Option<String>,
269) {
270 if let (Some(p), Some(b)) = (path.take(), branch.take()) {
271 let sanitized = sanitize_branch(&b);
272 if sanitized.is_empty() {
273 log::warn!(
274 "Skipping git worktree at '{}' because branch '{}' sanitizes to empty \
275 (no ASCII alphanumeric characters)",
276 p.display(),
277 b,
278 );
279 return;
280 }
281 entries.push(WorktreeEntry {
282 path: p,
283 branch: b,
284 sanitized_branch: sanitized,
285 namespace: None,
286 });
287 }
288}
289
290fn sanitize_branch(branch: &str) -> String {
293 let sanitized: String = branch
294 .chars()
295 .map(|c| {
296 if c.is_ascii_alphanumeric() || c == '-' {
297 c
298 } else {
299 '-'
300 }
301 })
302 .collect();
303 sanitized.trim_matches('-').to_string()
304}
305
306#[cfg(test)]
307mod tests {
308 use super::*;
309
310 #[test]
311 fn test_sanitize_branch_simple() {
312 assert_eq!(sanitize_branch("feature-a"), "feature-a");
313 }
314
315 #[test]
316 fn test_sanitize_branch_with_slash() {
317 assert_eq!(
318 sanitize_branch("feature/my-endpoint"),
319 "feature-my-endpoint"
320 );
321 }
322
323 #[test]
324 fn test_sanitize_branch_with_underscore() {
325 assert_eq!(sanitize_branch("fix_bug_123"), "fix-bug-123");
326 }
327
328 #[test]
331 fn test_parse_jj_workspace_list_two_workspaces() {
332 let input =
333 b"default: kkqmkqnm 6aa0ec8e main\nfeature-a: rrqxmqnm 8e9b1c2d feature/my-endpoint\n";
334 let entries = parse_jj_workspace_list(input, |name| {
335 Some(PathBuf::from(format!("/home/user/{}-ws", name)))
336 });
337 assert_eq!(entries.len(), 2);
338 assert_eq!(entries[0].path, PathBuf::from("/home/user/default-ws"));
339 assert_eq!(entries[0].branch, "default");
340 assert_eq!(entries[0].sanitized_branch, "default");
341 assert_eq!(entries[1].path, PathBuf::from("/home/user/feature-a-ws"));
342 assert_eq!(entries[1].branch, "feature-a");
343 assert_eq!(entries[1].sanitized_branch, "feature-a");
344 }
345
346 #[test]
347 fn test_parse_jj_workspace_list_no_colon() {
348 let input = b"some invalid line without colon\n";
349 let entries = parse_jj_workspace_list(input, |_| {
350 panic!("should not be called for unparseable lines")
351 });
352 assert_eq!(entries.len(), 0);
353 }
354
355 #[test]
356 fn test_parse_jj_workspace_list_no_trailing_newline() {
357 let input = b"default: kkqmkqnm 6aa0ec8e main";
358 let entries = parse_jj_workspace_list(input, |_| Some(PathBuf::from("/home/user/myapp")));
359 assert_eq!(entries.len(), 1);
360 assert_eq!(entries[0].path, PathBuf::from("/home/user/myapp"));
361 assert_eq!(entries[0].branch, "default");
362 assert_eq!(entries[0].sanitized_branch, "default");
363 }
364
365 #[test]
366 fn test_parse_jj_workspace_list_skips_unresolved() {
367 let input = b"default: abc123\norphan: def456\n";
368 let entries = parse_jj_workspace_list(input, |name| {
369 if name == "default" {
370 Some(PathBuf::from("/home/user/myapp"))
371 } else {
372 None
373 }
374 });
375 assert_eq!(entries.len(), 1);
376 assert_eq!(entries[0].branch, "default");
377 }
378
379 #[test]
382 fn test_parse_git_worktree_output_two_worktrees() {
383 let input = b"worktree /home/user/myapp\nHEAD abc123\nbranch refs/heads/main\n\nworktree /home/user/myapp-feature-a\nHEAD def456\nbranch refs/heads/feature-a\n";
384 let entries = parse_git_worktree_output(input);
385 assert_eq!(entries.len(), 2);
386 assert_eq!(entries[0].path, PathBuf::from("/home/user/myapp"));
387 assert_eq!(entries[0].branch, "main");
388 assert_eq!(entries[0].sanitized_branch, "main");
389 assert_eq!(entries[1].path, PathBuf::from("/home/user/myapp-feature-a"));
390 assert_eq!(entries[1].branch, "feature-a");
391 assert_eq!(entries[1].sanitized_branch, "feature-a");
392 }
393
394 #[test]
395 fn test_parse_git_worktree_output_detached_head() {
396 let input = b"worktree /home/user/myapp\nHEAD abc123\n\n";
397 let entries = parse_git_worktree_output(input);
398 assert_eq!(entries.len(), 0);
399 }
400
401 #[test]
402 fn test_parse_git_worktree_output_no_trailing_blank() {
403 let input = b"worktree /home/user/myapp\nHEAD abc123\nbranch refs/heads/main";
404 let entries = parse_git_worktree_output(input);
405 assert_eq!(entries.len(), 1);
406 assert_eq!(entries[0].branch, "main");
407 }
408
409 #[test]
410 fn test_sanitize_branch_non_ascii() {
411 assert_eq!(sanitize_branch("fix-バグ"), "fix");
412 assert_eq!(sanitize_branch("fix-ü"), "fix");
413 assert_eq!(sanitize_branch("fix-中"), "fix");
414 }
415
416 #[test]
417 fn test_sanitize_branch_empty() {
418 assert_eq!(sanitize_branch("---"), "");
419 assert_eq!(sanitize_branch("///"), "");
420 assert_eq!(sanitize_branch("___"), "");
421 }
422
423 #[test]
424 fn test_parse_git_worktree_output_empty_sanitized() {
425 let input = b"worktree /home/user/myapp\nHEAD abc123\nbranch refs/heads/---\n\n";
426 let entries = parse_git_worktree_output(input);
427 assert_eq!(entries.len(), 0);
428 }
429
430 #[test]
433 fn test_discover_git_worktrees_single_checkout_no_spawn() {
434 let temp = tempfile::tempdir().unwrap();
435 let repo = temp.path().join("my-repo");
436 std::fs::create_dir(&repo).unwrap();
437 let git_dir = repo.join(".git");
438 std::fs::create_dir_all(git_dir.join("refs")).unwrap();
439 std::fs::write(git_dir.join("HEAD"), "ref: refs/heads/main\n").unwrap();
441
442 let entries = discover_git_worktrees(&repo);
443 assert_eq!(entries.len(), 1);
444 assert_eq!(entries[0].path, repo.canonicalize().unwrap());
447 assert_eq!(entries[0].branch, "main");
448 assert_eq!(entries[0].sanitized_branch, "main");
449 }
450
451 #[test]
454 fn test_discover_git_worktrees_empty_worktrees_dir() {
455 let temp = tempfile::tempdir().unwrap();
456 let repo = temp.path().join("my-repo");
457 std::fs::create_dir(&repo).unwrap();
458 let git_dir = repo.join(".git");
459 std::fs::create_dir_all(git_dir.join("worktrees")).unwrap();
460 std::fs::write(git_dir.join("HEAD"), "ref: refs/heads/main\n").unwrap();
461
462 let entries = discover_git_worktrees(&repo);
463 assert_eq!(entries.len(), 1);
464 assert_eq!(entries[0].branch, "main");
465 }
466
467 #[test]
470 fn test_discover_git_worktrees_single_checkout_detached() {
471 let temp = tempfile::tempdir().unwrap();
472 let repo = temp.path().join("my-repo");
473 std::fs::create_dir(&repo).unwrap();
474 let git_dir = repo.join(".git");
475 std::fs::create_dir_all(git_dir.join("refs")).unwrap();
476 std::fs::write(
477 git_dir.join("HEAD"),
478 "0123456789abcdef0123456789abcdef01234567\n",
479 )
480 .unwrap();
481
482 let entries = discover_git_worktrees(&repo);
483 assert_eq!(entries.len(), 0);
484 }
485
486 #[cfg(unix)]
490 #[test]
491 fn test_discover_git_worktrees_canonicalizes_symlink() {
492 use std::os::unix::fs::symlink;
493
494 let temp = tempfile::tempdir().unwrap();
495 let repo = temp.path().join("real-repo");
496 std::fs::create_dir(&repo).unwrap();
497 let git_dir = repo.join(".git");
498 std::fs::create_dir_all(git_dir.join("refs")).unwrap();
499 std::fs::write(git_dir.join("HEAD"), "ref: refs/heads/main\n").unwrap();
500
501 let link = temp.path().join("link-to-repo");
502 symlink(&repo, &link).unwrap();
503
504 let entries = discover_git_worktrees(&link);
505 assert_eq!(entries.len(), 1);
506 assert_eq!(
507 entries[0].path,
508 repo.canonicalize().unwrap(),
509 "fast path must resolve the symlink like git does"
510 );
511 }
512
513 #[test]
514 fn test_parse_jj_workspace_list_empty_sanitized() {
515 let input = b"---: kkqmkqnm 6aa0ec8e main\n";
516 let entries = parse_jj_workspace_list(input, |_| Some(PathBuf::from("/home/user")));
517 assert_eq!(entries.len(), 0);
518 }
519}