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