1use std::io::Write;
18use std::path::{Path, PathBuf};
19
20use mkit_core::hash::Hash;
21use mkit_core::layout::{self, RepoLayout};
22use mkit_core::object::Object;
23use mkit_core::refs::{self, RefWriteCondition};
24use mkit_core::store::ObjectStore;
25
26use crate::clap_shim;
27use crate::exit;
28use crate::format;
29use clap::Parser;
30use clap::Subcommand;
31
32#[derive(Debug, Parser)]
33#[command(name = "mkit worktree", about = "Manage linked working trees.")]
34struct WorktreeOpts {
35 #[command(subcommand)]
36 sub: WorktreeCmd,
37}
38
39#[derive(Debug, Subcommand)]
40enum WorktreeCmd {
41 Add {
48 path: String,
49 commit_ish: Option<String>,
50 },
51 List {
53 #[arg(long)]
55 porcelain: bool,
56 },
57 Remove {
59 #[arg(long, short)]
62 force: bool,
63 path: String,
64 },
65 Prune {
67 #[arg(long)]
69 dry_run: bool,
70 },
71}
72
73#[must_use]
74pub fn run(args: &[String]) -> u8 {
75 let opts = match clap_shim::parse::<WorktreeOpts>("mkit worktree", args) {
76 Ok(o) => o,
77 Err(code) => return code,
78 };
79 let cwd = match std::env::current_dir() {
80 Ok(c) => c,
81 Err(e) => return super::error(&format!("cwd: {e}"), exit::CONFIG_ERROR),
82 };
83 let layout = match super::resolve_layout(&cwd) {
84 Ok(layout) => layout,
85 Err(code) => return code,
86 };
87
88 match opts.sub {
89 WorktreeCmd::Add { path, commit_ish } => add(&layout, &cwd, &path, commit_ish.as_deref()),
90 WorktreeCmd::List { porcelain } => list(&layout, porcelain),
91 WorktreeCmd::Remove { force, path } => remove(&layout, &cwd, &path, force),
92 WorktreeCmd::Prune { dry_run } => prune(&layout, dry_run),
93 }
94}
95
96enum HeadPlan {
100 NewBranch { branch: String, start: Hash },
102 ExistingBranch { branch: String, tip: Hash },
104 Detached(Hash),
106}
107
108fn add(layout: &RepoLayout, cwd: &Path, path: &str, commit_ish: Option<&str>) -> u8 {
109 let store = match super::open_store_configured(layout) {
110 Ok(s) => s,
111 Err(e) => return super::error(&format!("open store: {e}"), exit::UNAVAILABLE),
112 };
113
114 let target = canonical_or_lexical(&absolutize(cwd, Path::new(path)));
117 if let Err(code) = check_add_target(layout, &target) {
118 return code;
119 }
120 let plan = match plan_head(layout, &store, &target, commit_ish) {
121 Ok(p) => p,
122 Err(code) => return code,
123 };
124
125 let commit_hash = match &plan {
126 HeadPlan::NewBranch { start, .. } => *start,
127 HeadPlan::ExistingBranch { tip, .. } => *tip,
128 HeadPlan::Detached(h) => *h,
129 };
130 let tree_hash = match store.read_object(&commit_hash) {
131 Ok(Object::Commit(c)) => c.tree_hash,
132 Ok(Object::Remix(r)) => r.tree_hash,
133 Ok(_) => {
134 return super::error(
135 &format!(
136 "{} does not resolve to a commit or remix",
137 format::short_hash(&commit_hash, 8)
138 ),
139 exit::DATAERR,
140 );
141 }
142 Err(e) => return super::error(&format!("read commit: {e}"), exit::GENERAL_ERROR),
143 };
144 if let Err(code) = create_worktree(layout, &store, &plan, &target, tree_hash) {
145 return code;
146 }
147
148 let mut stdout = std::io::stdout().lock();
149 match &plan {
150 HeadPlan::NewBranch { branch, .. } => {
151 let _ = writeln!(stdout, "Preparing worktree (new branch '{branch}')");
152 }
153 HeadPlan::ExistingBranch { branch, .. } => {
154 let _ = writeln!(stdout, "Preparing worktree (checking out '{branch}')");
155 }
156 HeadPlan::Detached(h) => {
157 let _ = writeln!(
158 stdout,
159 "Preparing worktree (detached HEAD {})",
160 format::short_hash(h, 8)
161 );
162 }
163 }
164 let _ = writeln!(
165 stdout,
166 "HEAD is now at {} {}",
167 format::short_hash(&commit_hash, 8),
168 super::commit_subject(&store, &commit_hash)
169 );
170 exit::OK
171}
172
173fn check_add_target(layout: &RepoLayout, target: &Path) -> Result<(), u8> {
175 let siblings =
180 super::all_worktree_layouts(layout).map_err(|e| super::error(&e, exit::DATAERR))?;
181 for (tree_root, _) in &siblings {
182 let root = canonical_or_lexical(tree_root);
183 if target.starts_with(&root) {
184 return Err(super::error(
185 &format!(
186 "'{}' is inside the worktree at '{}'; choose a path outside every \
187 existing worktree",
188 target.display(),
189 tree_root.display()
190 ),
191 exit::USAGE,
192 ));
193 }
194 }
195 match std::fs::read_dir(target) {
196 Ok(mut entries) => {
197 if entries.next().is_some() {
198 return Err(super::error(
199 &format!("'{}' already exists and is not empty", target.display()),
200 exit::CANTCREAT,
201 ));
202 }
203 Ok(())
204 }
205 Err(e) if e.kind() == std::io::ErrorKind::NotFound => Ok(()),
206 Err(_) if target.exists() => Err(super::error(
207 &format!(
208 "'{}' already exists and is not a directory",
209 target.display()
210 ),
211 exit::CANTCREAT,
212 )),
213 Err(e) => Err(super::error(
214 &format!("inspect '{}': {e}", target.display()),
215 exit::GENERAL_ERROR,
216 )),
217 }
218}
219
220fn plan_head(
222 layout: &RepoLayout,
223 store: &ObjectStore,
224 target: &Path,
225 commit_ish: Option<&str>,
226) -> Result<HeadPlan, u8> {
227 let plan = match commit_ish {
228 None => {
229 let Some(branch) = branch_name_from_path(target) else {
230 return Err(super::error(
231 &format!(
232 "cannot derive a branch name from '{}'; pass a commit-ish",
233 target.display()
234 ),
235 exit::USAGE,
236 ));
237 };
238 if matches!(refs::read_ref(layout, &branch), Ok(Some(_))) {
239 return Err(super::error(
240 &format!(
241 "branch '{branch}' already exists; pass it explicitly to check it out"
242 ),
243 exit::CANTCREAT,
244 ));
245 }
246 let start = match refs::resolve_head(layout) {
247 Ok(Some(h)) => h,
248 Ok(None) => {
249 return Err(super::error(
250 "cannot add a worktree: the repository has no commits yet",
251 exit::DATAERR,
252 ));
253 }
254 Err(e) => return Err(super::error(&format!("resolve HEAD: {e}"), exit::DATAERR)),
255 };
256 HeadPlan::NewBranch { branch, start }
257 }
258 Some(spec) => match refs::read_ref(layout, spec) {
259 Ok(Some(tip)) => HeadPlan::ExistingBranch {
260 branch: spec.to_string(),
261 tip,
262 },
263 _ => match super::revspec::resolve_revision(store, layout, spec) {
264 Ok(h) => HeadPlan::Detached(h),
265 Err(e) => {
266 return Err(super::error(
267 &format!("no such branch, tag, or commit: {spec} ({e})"),
268 exit::GENERAL_ERROR,
269 ));
270 }
271 },
272 },
273 };
274
275 if let HeadPlan::ExistingBranch { branch, .. } | HeadPlan::NewBranch { branch, .. } = &plan {
278 match super::branch_checked_out_elsewhere(layout, branch) {
279 Ok(Some(at)) => {
280 return Err(super::error(
281 &format!(
282 "branch '{branch}' is already checked out at '{}'",
283 at.display()
284 ),
285 exit::DATAERR,
286 ));
287 }
288 Ok(None) => {}
289 Err(e) => return Err(super::error(&e, exit::DATAERR)),
290 }
291 if matches!(refs::read_head(layout), Ok(refs::Head::Branch(ref cur)) if cur == branch) {
293 return Err(super::error(
294 &format!("branch '{branch}' is already checked out in this worktree"),
295 exit::DATAERR,
296 ));
297 }
298 }
299 Ok(plan)
300}
301
302fn create_worktree(
305 layout: &RepoLayout,
306 store: &ObjectStore,
307 plan: &HeadPlan,
308 target: &Path,
309 tree_hash: Hash,
310) -> Result<(), u8> {
311 let _lock = super::acquire_worktrees_registry_lock(layout)?;
316
317 if let HeadPlan::ExistingBranch { branch, .. } | HeadPlan::NewBranch { branch, .. } = plan {
321 match super::branch_checked_out_elsewhere(layout, branch) {
322 Ok(None) => {}
323 Ok(Some(at)) => {
324 return Err(super::error(
325 &format!(
326 "branch '{branch}' is already checked out at '{}'",
327 at.display()
328 ),
329 exit::DATAERR,
330 ));
331 }
332 Err(e) => return Err(super::error(&e, exit::DATAERR)),
333 }
334 }
335
336 let Some(id) = free_worktree_id(layout, target) else {
337 return Err(super::error(
338 &format!("cannot derive a worktree id from '{}'", target.display()),
339 exit::USAGE,
340 ));
341 };
342 let state_dir = layout.worktree_state_dir_for(&id);
343 let linked = RepoLayout::linked(target, &state_dir, layout.common_dir());
344
345 if let Err(e) = std::fs::create_dir_all(&state_dir) {
349 return Err(super::error(
350 &format!("create state dir: {e}"),
351 exit::CANTCREAT,
352 ));
353 }
354 let steps: [(&str, std::io::Result<()>); 2] = [
355 (
356 "commondir",
357 std::fs::write(state_dir.join(layout::COMMONDIR_FILE_NAME), b"../..\n"),
358 ),
359 (
360 "back-pointer",
361 std::fs::write(
362 state_dir.join(layout::BACKPOINTER_FILE_NAME),
363 format!("{}\n", target.join(mkit_core::MKIT_DIR).display()),
364 ),
365 ),
366 ];
367 for (what, res) in steps {
368 if let Err(e) = res {
369 return Err(super::error(&format!("write {what}: {e}"), exit::CANTCREAT));
370 }
371 }
372 let head_write = match plan {
373 HeadPlan::NewBranch { branch, .. } | HeadPlan::ExistingBranch { branch, .. } => {
374 refs::write_head_branch(&linked, branch)
375 }
376 HeadPlan::Detached(h) => refs::write_head_detached(&linked, h),
377 };
378 if let Err(e) = head_write {
379 return Err(super::error(&format!("write HEAD: {e}"), exit::CANTCREAT));
380 }
381
382 if let HeadPlan::NewBranch { branch, start } = plan
384 && let Err(e) =
385 super::write_ref_recording_history(layout, branch, RefWriteCondition::Missing, start)
386 {
387 return Err(super::error(
388 &format!("create branch '{branch}': {e}"),
389 exit::CANTCREAT,
390 ));
391 }
392
393 if let Err(e) = std::fs::create_dir_all(target) {
395 return Err(super::error(
396 &format!("create '{}': {e}", target.display()),
397 exit::CANTCREAT,
398 ));
399 }
400 if let Err(e) = layout::write_pointer_file(target, &state_dir) {
401 return Err(super::error(
402 &format!("write worktree pointer: {e}"),
403 exit::CANTCREAT,
404 ));
405 }
406 if let Err(e) = super::restore_worktree_and_index(&linked, store, tree_hash) {
407 return Err(super::error(&e, exit::GENERAL_ERROR));
408 }
409 Ok(())
410}
411
412type ListRow = (PathBuf, Option<Hash>, Option<String>, Option<String>);
417
418fn list(layout: &RepoLayout, porcelain: bool) -> u8 {
419 let store = match super::open_store_configured(layout) {
420 Ok(s) => s,
421 Err(e) => return super::error(&format!("open store: {e}"), exit::UNAVAILABLE),
422 };
423 let _ = store; let mut rows: Vec<ListRow> = Vec::new();
426 let siblings = match super::all_worktree_layouts(layout) {
427 Ok(s) => s,
428 Err(e) => return super::error(&e, exit::DATAERR),
429 };
430 for (tree_root, candidate) in &siblings {
431 let head = refs::resolve_head(candidate).ok().flatten();
432 let branch = match refs::read_head(candidate) {
433 Ok(refs::Head::Branch(name)) => Some(name),
434 _ => None,
435 };
436 rows.push((tree_root.clone(), head, branch, None));
437 }
438 match layout::worktrees(layout) {
440 Ok(entries) => {
441 for wt in entries {
442 if let Some(reason) = wt.prunable {
443 let shown = wt.tree_root.unwrap_or_else(|| wt.state_dir.clone());
444 rows.push((shown, None, None, Some(reason)));
445 }
446 }
447 }
448 Err(e) => return super::error(&format!("worktree registry: {e}"), exit::DATAERR),
449 }
450
451 let mut stdout = std::io::stdout().lock();
452 for (path, head, branch, prunable) in rows {
453 if porcelain {
454 let _ = writeln!(stdout, "worktree {}", path.display());
455 if let Some(h) = head {
456 let _ = writeln!(stdout, "HEAD {}", mkit_core::hash::to_hex(&h));
457 }
458 match (&branch, &prunable) {
459 (_, Some(reason)) => {
460 let _ = writeln!(stdout, "prunable {reason}");
461 }
462 (Some(b), None) => {
463 let _ = writeln!(stdout, "branch refs/heads/{b}");
464 }
465 (None, None) => {
466 let _ = writeln!(stdout, "detached");
467 }
468 }
469 let _ = writeln!(stdout);
470 } else {
471 let hash_col = head.map_or_else(|| "-".repeat(8), |h| format::short_hash(&h, 8));
472 let desc = match (&branch, &prunable) {
473 (_, Some(reason)) => format!("(prunable: {reason})"),
474 (Some(b), None) => format!("[{b}]"),
475 (None, None) => "(detached HEAD)".to_owned(),
476 };
477 let _ = writeln!(stdout, "{} {hash_col} {desc}", path.display());
478 }
479 }
480 exit::OK
481}
482
483fn remove(layout: &RepoLayout, cwd: &Path, path: &str, force: bool) -> u8 {
486 let target = canonical_or_lexical(&absolutize(cwd, Path::new(path)));
487
488 let main_root = layout.common_dir().parent().map(canonical_or_lexical);
489 if main_root.as_deref() == Some(target.as_path()) {
490 return super::error("the main working tree cannot be removed", exit::USAGE);
491 }
492 if canonical_or_lexical(cwd).starts_with(&target) {
493 return super::error(
494 "cannot remove the worktree you are currently inside",
495 exit::USAGE,
496 );
497 }
498
499 let entries = match layout::worktrees(layout) {
500 Ok(e) => e,
501 Err(e) => return super::error(&format!("worktree registry: {e}"), exit::DATAERR),
502 };
503 let Some(wt) = entries.into_iter().find(|wt| {
504 wt.tree_root
505 .as_deref()
506 .is_some_and(|root| canonical_or_lexical(root) == target)
507 }) else {
508 return super::error(
509 &format!(
510 "'{}' is not a linked worktree of this repository",
511 target.display()
512 ),
513 exit::USAGE,
514 );
515 };
516
517 if !force && wt.prunable.is_none() {
521 let linked = RepoLayout::linked(&target, &wt.state_dir, layout.common_dir());
522 if let Some(op) = mkit_core::ops::conflict_state::in_progress_op_name(&linked) {
523 return super::error(
524 &format!("worktree has a {op} in progress; resolve it or pass --force"),
525 exit::DATAERR,
526 );
527 }
528 match worktree_is_dirty(&linked) {
529 Ok(Some(why)) => {
530 return super::error(
531 &format!("worktree contains {why}; commit, stash, or pass --force"),
532 exit::DATAERR,
533 );
534 }
535 Ok(None) => {}
536 Err(e) => return super::error(&e, exit::DATAERR),
537 }
538 }
539
540 let _lock = match super::acquire_worktrees_registry_lock(layout) {
541 Ok(l) => l,
542 Err(code) => return code,
543 };
544 let _target_lock = if wt.state_dir.is_dir() {
549 let target_layout = RepoLayout::linked(&target, &wt.state_dir, layout.common_dir());
550 match super::acquire_worktree_lock(&target_layout) {
551 Ok(l) => Some(l),
552 Err(code) => return code,
553 }
554 } else {
555 None
556 };
557 if target.exists()
560 && let Err(e) = std::fs::remove_dir_all(&target)
561 {
562 return super::error(
563 &format!("remove '{}': {e}", target.display()),
564 exit::GENERAL_ERROR,
565 );
566 }
567 if let Err(e) = std::fs::remove_dir_all(&wt.state_dir) {
568 return super::error(
569 &format!("remove state dir '{}': {e}", wt.state_dir.display()),
570 exit::GENERAL_ERROR,
571 );
572 }
573 exit::OK
574}
575
576fn worktree_is_dirty(linked: &RepoLayout) -> Result<Option<String>, String> {
579 let store = super::open_store_configured(linked).map_err(|e| format!("open store: {e}"))?;
580 let head_tree = super::current_head_tree(linked, &store)?;
581 let Some(head_tree) = head_tree else {
582 return Ok(None); };
584 let idx = super::read_or_seed_index_from_head(linked, &store)?;
587 let mut paths = Vec::new();
588 super::collect_worktree_paths(
589 linked.worktree_root(),
590 linked.worktree_root(),
591 "",
592 &mut paths,
593 )
594 .map_err(|e| format!("scan worktree: {e}"))?;
595 for p in paths {
596 let abs = linked.worktree_root().join(&p);
597 if abs.is_dir() {
598 continue;
599 }
600 if !super::index_tracks_path_or_descendant(&idx, &p) {
601 return Ok(Some(format!("untracked file '{p}'")));
602 }
603 }
604 match super::ensure_restore_safe(linked, &store, head_tree) {
609 Ok(()) => Ok(None),
610 Err(why) if why.starts_with("restore would") => Ok(Some("local changes".to_owned())),
611 Err(why) => Err(why),
612 }
613}
614
615fn prune(layout: &RepoLayout, dry_run: bool) -> u8 {
618 let _lock = if dry_run {
624 None
625 } else {
626 match super::acquire_worktrees_registry_lock(layout) {
627 Ok(l) => Some(l),
628 Err(code) => return code,
629 }
630 };
631 let entries = match layout::worktrees(layout) {
632 Ok(e) => e,
633 Err(e) => return super::error(&format!("worktree registry: {e}"), exit::DATAERR),
634 };
635 let mut stdout = std::io::stdout().lock();
636 for wt in entries {
637 let Some(reason) = wt.prunable else { continue };
638 if dry_run {
639 let _ = writeln!(stdout, "would prune worktrees/{}: {reason}", wt.id);
640 continue;
641 }
642 if let Err(e) = std::fs::remove_dir_all(&wt.state_dir) {
643 return super::error(
644 &format!("prune worktrees/{}: {e}", wt.id),
645 exit::GENERAL_ERROR,
646 );
647 }
648 let _ = writeln!(stdout, "pruned worktrees/{}: {reason}", wt.id);
649 }
650 exit::OK
651}
652
653fn absolutize(cwd: &Path, path: &Path) -> PathBuf {
658 if path.is_absolute() {
659 path.to_path_buf()
660 } else {
661 cwd.join(path)
662 }
663}
664
665fn canonical_or_lexical(p: &Path) -> PathBuf {
672 if let Ok(c) = p.canonicalize() {
673 return c;
674 }
675 let mut missing = Vec::new();
676 let mut cur = p;
677 while let Some(parent) = cur.parent() {
678 if let Some(name) = cur.file_name() {
679 missing.push(name.to_owned());
680 }
681 if let Ok(c) = parent.canonicalize() {
682 let mut out = c;
683 for name in missing.iter().rev() {
684 out.push(name);
685 }
686 return out;
687 }
688 cur = parent;
689 }
690 p.to_path_buf()
691}
692
693fn branch_name_from_path(target: &Path) -> Option<String> {
696 let base = target.file_name()?.to_string_lossy();
697 let candidate: String = base
698 .chars()
699 .map(|c| {
700 if c.is_ascii_alphanumeric() || c == '.' || c == '_' || c == '-' {
701 c
702 } else {
703 '-'
704 }
705 })
706 .collect();
707 let candidate = candidate.trim_matches(['-', '.']).to_string();
708 refs::validate_ref_name(&candidate).then_some(candidate)
709}
710
711fn free_worktree_id(layout: &RepoLayout, target: &Path) -> Option<String> {
714 let base = target.file_name()?.to_string_lossy();
715 let sanitized: String = base
716 .chars()
717 .map(|c| {
718 if c.is_ascii_alphanumeric() || c == '.' || c == '_' || c == '-' {
719 c
720 } else {
721 '-'
722 }
723 })
724 .collect();
725 let sanitized = sanitized.trim_matches('-').to_string();
726 if !layout::validate_worktree_id(&sanitized) {
727 return None;
728 }
729 if !layout.worktree_state_dir_for(&sanitized).exists() {
730 return Some(sanitized);
731 }
732 (1..10_000).find_map(|n| {
733 let candidate = format!("{sanitized}-{n}");
734 (layout::validate_worktree_id(&candidate)
735 && !layout.worktree_state_dir_for(&candidate).exists())
736 .then_some(candidate)
737 })
738}