Expand the worktree functionality

src/cmd.rs

@@ -66,16 +66,46 @@ pub struct Worktree {
 #[derive(Parser)]
 pub enum WorktreeAction {
     #[clap(about = "Add a new worktree")]
-    Add(WorktreeActionArgs),
+    Add(WorktreeAddArgs),
     #[clap(about = "Add an existing worktree")]
-    Delete(WorktreeActionArgs),
+    Delete(WorktreeDeleteArgs),
+    #[clap(about = "Show state of existing worktrees")]
+    Status(WorktreeStatusArgs),
+    #[clap(about = "Clean all worktrees that do not contain uncommited/unpushed changes")]
+    Clean(WorktreeCleanArgs),
 }
 
 #[derive(Parser)]
-pub struct WorktreeActionArgs {
+pub struct WorktreeAddArgs {
     #[clap(about = "Name of the worktree")]
     pub name: String,
+
+    #[clap(
+        short = 'p',
+        long = "branch-prefix",
+        about = "Prefix for the branch name"
+    )]
+    pub branch_prefix: Option<String>,
+    #[clap(short = 't', long = "track", about = "Remote branch to track")]
+    pub track: Option<String>,
 }
+#[derive(Parser)]
+pub struct WorktreeDeleteArgs {
+    #[clap(about = "Name of the worktree")]
+    pub name: String,
+
+    #[clap(
+        long = "force",
+        about = "Force deletion, even when there are uncommitted/unpushed changes"
+    )]
+    pub force: bool,
+}
+
+#[derive(Parser)]
+pub struct WorktreeStatusArgs {}
+
+#[derive(Parser)]
+pub struct WorktreeCleanArgs {}
 
 pub fn parse() -> Opts {
     Opts::parse()

src/lib.rs

@@ -99,6 +99,16 @@ fn expand_path(path: &Path) -> PathBuf {
     Path::new(&expanded_path).to_path_buf()
 }
 
+fn get_default_branch(repo: &git2::Repository) -> Result<git2::Branch, String> {
+    match repo.find_branch("main", git2::BranchType::Local) {
+        Ok(branch) => Ok(branch),
+        Err(_) => match repo.find_branch("master", git2::BranchType::Local) {
+            Ok(branch) => Ok(branch),
+            Err(_) => Err(String::from("Could not determine default branch")),
+        },
+    }
+}
+
 fn sync_trees(config: Config) {
     for tree in config.trees {
         let repos = tree.repos.unwrap_or_default();
@@ -448,6 +458,7 @@ fn add_table_header(table: &mut Table) {
         .apply_modifier(comfy_table::modifiers::UTF8_ROUND_CORNERS)
         .set_header(vec![
             Cell::new("Repo"),
+            Cell::new("Worktree"),
             Cell::new("Status"),
             Cell::new("Branches"),
             Cell::new("HEAD"),
@@ -455,12 +466,35 @@ fn add_table_header(table: &mut Table) {
         ]);
 }
 
-fn add_repo_status(table: &mut Table, repo_name: &str, repo_handle: &git2::Repository) {
+fn add_worktree_table_header(table: &mut Table) {
-    let repo_status = get_repo_status(repo_handle);
+    table
+        .load_preset(comfy_table::presets::UTF8_FULL)
+        .apply_modifier(comfy_table::modifiers::UTF8_ROUND_CORNERS)
+        .set_header(vec![
+            Cell::new("Worktree"),
+            Cell::new("Status"),
+            Cell::new("Branch"),
+            Cell::new("Remote branch"),
+        ]);
+}
+
+fn add_repo_status(
+    table: &mut Table,
+    repo_name: &str,
+    repo_handle: &git2::Repository,
+    is_worktree: bool,
+) {
+    let repo_status = get_repo_status(repo_handle, is_worktree);
 
     table.add_row(vec![
         repo_name,
+        match is_worktree {
+            true => "\u{2714}",
+            false => "",
+        },
         &match repo_status.changes {
+            None => String::from("-"),
+            Some(changes) => match changes {
                 Some(changes) => {
                     let mut out = Vec::new();
                     if changes.files_new > 0 {
@@ -476,6 +510,7 @@ fn add_repo_status(table: &mut Table, repo_name: &str, repo_handle: &git2::Repos
                 }
                 None => String::from("\u{2714}"),
             },
+        },
         &repo_status
             .branches
             .iter()
@@ -504,10 +539,13 @@ fn add_repo_status(table: &mut Table, repo_name: &str, repo_handle: &git2::Repos
             .collect::<String>()
             .trim()
             .to_string(),
-        &match repo_status.head {
+        &match is_worktree {
+            true => String::from(""),
+            false => match repo_status.head {
                 Some(head) => head,
                 None => String::from("Empty"),
             },
+        },
         &repo_status
             .remotes
             .iter()
@@ -518,6 +556,72 @@ fn add_repo_status(table: &mut Table, repo_name: &str, repo_handle: &git2::Repos
     ]);
 }
 
+fn add_worktree_status(table: &mut Table, worktree_name: &str, repo: &git2::Repository) {
+    let repo_status = get_repo_status(repo, false);
+
+    let head = repo.head().unwrap();
+
+    if !head.is_branch() {
+        print_error("No branch checked out in worktree");
+        process::exit(1);
+    }
+
+    let local_branch_name = head.shorthand().unwrap();
+    let local_branch = repo
+        .find_branch(local_branch_name, git2::BranchType::Local)
+        .unwrap();
+
+    let upstream_output = match local_branch.upstream() {
+        Ok(remote_branch) => {
+            let remote_branch_name = remote_branch.name().unwrap().unwrap().to_string();
+
+            let (ahead, behind) = repo
+                .graph_ahead_behind(
+                    local_branch.get().peel_to_commit().unwrap().id(),
+                    remote_branch.get().peel_to_commit().unwrap().id(),
+                )
+                .unwrap();
+
+            format!(
+                "{}{}\n",
+                &remote_branch_name,
+                &match (ahead, behind) {
+                    (0, 0) => String::from(""),
+                    (d, 0) => format!(" [+{}]", &d),
+                    (0, d) => format!(" [-{}]", &d),
+                    (d1, d2) => format!(" [+{}/-{}]", &d1, &d2),
+                },
+            )
+        }
+        Err(_) => String::from(""),
+    };
+
+    table.add_row(vec![
+        worktree_name,
+        &match repo_status.changes {
+            None => String::from(""),
+            Some(changes) => match changes {
+                Some(changes) => {
+                    let mut out = Vec::new();
+                    if changes.files_new > 0 {
+                        out.push(format!("New: {}\n", changes.files_new))
+                    }
+                    if changes.files_modified > 0 {
+                        out.push(format!("Modified: {}\n", changes.files_modified))
+                    }
+                    if changes.files_deleted > 0 {
+                        out.push(format!("Deleted: {}\n", changes.files_deleted))
+                    }
+                    out.into_iter().collect::<String>().trim().to_string()
+                }
+                None => String::from("\u{2714}"),
+            },
+        },
+        local_branch_name,
+        &upstream_output,
+    ]);
+}
+
 fn show_single_repo_status(path: &Path, is_worktree: bool) {
     let mut table = Table::new();
     add_table_header(&mut table);
@@ -547,7 +651,7 @@ fn show_single_repo_status(path: &Path, is_worktree: bool) {
         },
     };
 
-    add_repo_status(&mut table, &repo_name, &repo_handle.unwrap());
+    add_repo_status(&mut table, &repo_name, &repo_handle.unwrap(), is_worktree);
 
     println!("{}", table);
 }
@@ -588,12 +692,104 @@ fn show_status(config: Config) {
 
             let repo_handle = repo_handle.unwrap();
 
-            add_repo_status(&mut table, &repo.name, &repo_handle);
+            add_repo_status(&mut table, &repo.name, &repo_handle, repo.worktree_setup);
         }
         println!("{}", table);
     }
 }
 
+enum WorktreeRemoveFailureReason {
+    Changes(String),
+    Error(String),
+}
+
+fn remove_worktree(
+    name: &str,
+    worktree_dir: &Path,
+    force: bool,
+    main_repo: &git2::Repository,
+) -> Result<(), WorktreeRemoveFailureReason> {
+    if !worktree_dir.exists() {
+        return Err(WorktreeRemoveFailureReason::Error(format!(
+            "{} does not exist",
+            name
+        )));
+    }
+    let worktree_repo = match open_repo(worktree_dir, false) {
+        Ok(r) => r,
+        Err(e) => {
+            return Err(WorktreeRemoveFailureReason::Error(format!(
+                "Error opening repo: {}",
+                e
+            )));
+        }
+    };
+
+    let head = worktree_repo.head().unwrap();
+    if !head.is_branch() {
+        return Err(WorktreeRemoveFailureReason::Error(String::from(
+            "No branch checked out in worktree",
+        )));
+    }
+
+    let branch_name = head.shorthand().unwrap();
+    if !branch_name.ends_with(name) {
+        return Err(WorktreeRemoveFailureReason::Error(format!(
+            "Branch {} is checked out in worktree, this does not look correct",
+            &branch_name
+        )));
+    }
+
+    let mut branch = worktree_repo
+        .find_branch(branch_name, git2::BranchType::Local)
+        .unwrap();
+
+    if !force {
+        let status = get_repo_status(&worktree_repo, false);
+        if status.changes.is_some() {
+            return Err(WorktreeRemoveFailureReason::Changes(String::from(
+                "Changes found in worktree",
+            )));
+        }
+
+        match branch.upstream() {
+            Ok(remote_branch) => {
+                let (ahead, behind) = worktree_repo
+                    .graph_ahead_behind(
+                        branch.get().peel_to_commit().unwrap().id(),
+                        remote_branch.get().peel_to_commit().unwrap().id(),
+                    )
+                    .unwrap();
+
+                if (ahead, behind) != (0, 0) {
+                    return Err(WorktreeRemoveFailureReason::Changes(format!(
+                        "Branch {} is not in line with remote branch",
+                        name
+                    )));
+                }
+            }
+            Err(_) => {
+                return Err(WorktreeRemoveFailureReason::Changes(format!(
+                    "No remote tracking branch for branch {} found",
+                    name
+                )));
+            }
+        }
+    }
+
+    if let Err(e) = std::fs::remove_dir_all(&worktree_dir) {
+        return Err(WorktreeRemoveFailureReason::Error(format!(
+            "Error deleting {}: {}",
+            &worktree_dir.display(),
+            e
+        )));
+    }
+    main_repo.find_worktree(name).unwrap().prune(None).unwrap();
+    branch.delete().unwrap();
+
+    Ok(())
+}
+
 pub fn run() {
     let opts = cmd::parse();
 
@@ -667,15 +863,13 @@ pub fn run() {
                 }
             };
 
-            match args.action {
-                cmd::WorktreeAction::Add(action_args) => {
             let repo = match open_repo(&dir, true) {
                 Ok(r) => r,
                 Err(e) => {
                     match e.kind {
-                                RepoErrorKind::NotFound => print_error(
+                        RepoErrorKind::NotFound => {
-                                    "Current directory does not contain a worktree setup",
+                            print_error("Current directory does not contain a worktree setup")
-                                ),
+                        }
                         _ => print_error(&format!("Error opening repo: {}", e)),
                     }
                     process::exit(1);
@@ -686,14 +880,63 @@ pub fn run() {
                 .worktrees()
                 .unwrap()
                 .iter()
-                        .map(|e| e.unwrap())
+                .map(|e| e.unwrap().to_string())
-                        .collect::<String>();
+                .collect::<Vec<String>>();
+
+            match args.action {
+                cmd::WorktreeAction::Add(action_args) => {
                     if worktrees.contains(&action_args.name) {
-                        print_error("Worktree directory already exists");
+                        print_error("Worktree already exists");
                         process::exit(1);
                     }
 
-                    match repo.worktree(&action_args.name, &dir.join(&action_args.name), None) {
+                    let branch_name = match action_args.branch_prefix {
+                        Some(prefix) => format!("{}{}", &prefix, &action_args.name),
+                        None => action_args.name.clone(),
+                    };
+
+                    let checkout_commit = match &action_args.track {
+                        Some(upstream_branch_name) => {
+                            match repo.find_branch(upstream_branch_name, git2::BranchType::Remote) {
+                                Ok(branch) => branch.into_reference().peel_to_commit().unwrap(),
+                                Err(_) => {
+                                    print_error(&format!(
+                                        "Remote branch {} not found",
+                                        &upstream_branch_name
+                                    ));
+                                    process::exit(1);
+                                }
+                            }
+                        }
+                        None => get_default_branch(&repo)
+                            .unwrap()
+                            .into_reference()
+                            .peel_to_commit()
+                            .unwrap(),
+                    };
+
+                    let mut target_branch =
+                        match repo.find_branch(&branch_name, git2::BranchType::Local) {
+                            Ok(branchref) => branchref,
+                            Err(_) => repo.branch(&branch_name, &checkout_commit, false).unwrap(),
+                        };
+
+                    if let Some(upstream_branch_name) = action_args.track {
+                        target_branch
+                            .set_upstream(Some(&upstream_branch_name))
+                            .unwrap();
+                    }
+
+                    let worktree = repo.worktree(
+                        &action_args.name,
+                        &dir.join(&action_args.name),
+                        Some(
+                            git2::WorktreeAddOptions::new()
+                                .reference(Some(&target_branch.into_reference())),
+                        ),
+                    );
+
+                    match worktree {
                         Ok(_) => print_success(&format!("Worktree {} created", &action_args.name)),
                         Err(e) => {
                             print_error(&format!("Error creating worktree: {}", e));
@@ -701,64 +944,122 @@ pub fn run() {
                         }
                     };
                 }
+
                 cmd::WorktreeAction::Delete(action_args) => {
                     let worktree_dir = dir.join(&action_args.name);
-                    if !worktree_dir.exists() {
+
-                        print_error(&format!("{} does not exist", &action_args.name));
+                    match remove_worktree(
+                        &action_args.name,
+                        &worktree_dir,
+                        action_args.force,
+                        &repo,
+                    ) {
+                        Ok(_) => print_success(&format!("Worktree {} deleted", &action_args.name)),
+                        Err(error) => {
+                            match error {
+                                WorktreeRemoveFailureReason::Error(msg) => {
+                                    print_error(&msg);
                                     process::exit(1);
                                 }
-                    let repo = match open_repo(&worktree_dir, false) {
+                                WorktreeRemoveFailureReason::Changes(changes) => {
+                                    print_warning(&format!(
+                                        "Changes in worktree: {}. Refusing to delete",
+                                        changes
+                                    ));
+                                }
+                            }
+                            process::exit(1);
+                        }
+                    }
+                }
+                cmd::WorktreeAction::Status(_args) => {
+                    let mut table = Table::new();
+                    add_worktree_table_header(&mut table);
+                    for worktree in &worktrees {
+                        let repo_dir = &dir.join(&worktree);
+                        if repo_dir.exists() {
+                            let repo = match open_repo(repo_dir, false) {
                                 Ok(r) => r,
                                 Err(e) => {
                                     print_error(&format!("Error opening repo: {}", e));
                                     process::exit(1);
                                 }
                             };
-                    let status = get_repo_status(&repo);
+                            add_worktree_status(&mut table, worktree, &repo);
-                    if status.changes.is_some() {
+                        } else {
-                        print_error("Changes found in worktree, refusing to delete!");
+                            print_warning(&format!(
-                        process::exit(1);
+                                "Worktree {} does not have a directory",
-                    }
+                                &worktree
-
-                    let mut branch = repo
-                        .find_branch(&action_args.name, git2::BranchType::Local)
-                        .unwrap();
-                    match branch.upstream() {
-                        Ok(remote_branch) => {
-                            let (ahead, behind) = repo
-                                .graph_ahead_behind(
-                                    branch.get().peel_to_commit().unwrap().id(),
-                                    remote_branch.get().peel_to_commit().unwrap().id(),
-                                )
-                                .unwrap();
-
-                            if (ahead, behind) != (0, 0) {
-                                print_error(&format!("Branch {} is not in line with remote branch, refusing to delete worktree!", &action_args.name));
-                                process::exit(1);
-                            }
-                        }
-                        Err(_) => {
-                            print_error(&format!("No remote tracking branch for branch {} found, refusing to delete worktree!", &action_args.name));
-                            process::exit(1);
-                        }
-                    }
-
-                    match std::fs::remove_dir_all(&worktree_dir) {
-                        Ok(_) => print_success(&format!("Worktree {} deleted", &action_args.name)),
-                        Err(e) => {
-                            print_error(&format!(
-                                "Error deleting {}: {}",
-                                &worktree_dir.display(),
-                                e
                             ));
+                        }
+                    }
+                    for entry in std::fs::read_dir(&dir).unwrap() {
+                        let dirname = path_as_string(
+                            &entry
+                                .unwrap()
+                                .path()
+                                .strip_prefix(&dir)
+                                .unwrap()
+                                .to_path_buf(),
+                        );
+                        if dirname == GIT_MAIN_WORKTREE_DIRECTORY {
+                            continue;
+                        }
+                        if !&worktrees.contains(&dirname) {
+                            print_warning(&format!(
+                                "Found {}, which is not a valid worktree directory!",
+                                &dirname
+                            ));
+                        }
+                    }
+                    println!("{}", table);
+                }
+                cmd::WorktreeAction::Clean(_args) => {
+                    for worktree in &worktrees {
+                        let repo_dir = &dir.join(&worktree);
+                        if repo_dir.exists() {
+                            match remove_worktree(worktree, repo_dir, false, &repo) {
+                                Ok(_) => print_success(&format!("Worktree {} deleted", &worktree)),
+                                Err(error) => match error {
+                                    WorktreeRemoveFailureReason::Changes(changes) => {
+                                        print_warning(&format!(
+                                            "Changes found in {}: {}, skipping",
+                                            &worktree, &changes
+                                        ));
+                                        continue;
+                                    }
+                                    WorktreeRemoveFailureReason::Error(e) => {
+                                        print_error(&e);
                                         process::exit(1);
                                     }
+                                },
                             }
-                    repo.find_worktree(&action_args.name)
+                        } else {
+                            print_warning(&format!(
+                                "Worktree {} does not have a directory",
+                                &worktree
+                            ));
+                        }
+                    }
+                    for entry in std::fs::read_dir(&dir).unwrap() {
+                        let dirname = path_as_string(
+                            &entry
                                 .unwrap()
-                        .prune(None)
+                                .path()
-                        .unwrap();
+                                .strip_prefix(&dir)
-                    branch.delete().unwrap();
+                                .unwrap()
+                                .to_path_buf(),
+                        );
+                        if dirname == GIT_MAIN_WORKTREE_DIRECTORY {
+                            continue;
+                        }
+                        if !&worktrees.contains(&dirname) {
+                            print_warning(&format!(
+                                "Found {}, which is not a valid worktree directory!",
+                                &dirname
+                            ));
+                        }
+                    }
                 }
             }
         }

src/repo.rs

@@ -90,11 +90,14 @@ pub struct RepoStatus {
 
     pub head: Option<String>,
 
-    pub changes: Option<RepoChanges>,
+    // None(_) => Could not get changes (e.g. because it's a worktree setup
+    // Some(None) => No changes
+    // Some(Some(_)) => Changes
+    pub changes: Option<Option<RepoChanges>>,
 
     pub worktrees: usize,
 
-    pub submodules: Vec<(String, SubmoduleStatus)>,
+    pub submodules: Option<Vec<(String, SubmoduleStatus)>>,
 
     pub branches: Vec<(String, Option<(String, RemoteTrackingStatus)>)>,
 }
@@ -253,7 +256,7 @@ pub fn clone_repo(
     }
 }
 
-pub fn get_repo_status(repo: &git2::Repository) -> RepoStatus {
+pub fn get_repo_status(repo: &git2::Repository, is_worktree: bool) -> RepoStatus {
     let operation = match repo.state() {
         git2::RepositoryState::Clean => None,
         state => Some(state),
@@ -268,11 +271,17 @@ pub fn get_repo_status(repo: &git2::Repository) -> RepoStatus {
         .map(|repo_name| repo_name.unwrap().to_string())
         .collect::<Vec<String>>();
 
-    let head = match empty {
+    let head = match is_worktree {
+        true => None,
+        false => match empty {
             true => None,
             false => Some(repo.head().unwrap().shorthand().unwrap().to_string()),
+        },
     };
 
+    let changes = match is_worktree {
+        true => None,
+        false => {
             let statuses = repo
                 .statuses(Some(
                     git2::StatusOptions::new()
@@ -281,8 +290,8 @@ pub fn get_repo_status(repo: &git2::Repository) -> RepoStatus {
                 ))
                 .unwrap();
 
-    let changes = match statuses.is_empty() {
+            match statuses.is_empty() {
-        true => None,
+                true => Some(None),
                 false => {
                     let mut files_new = 0;
                     let mut files_modified = 0;
@@ -298,7 +307,9 @@ pub fn get_repo_status(repo: &git2::Repository) -> RepoStatus {
                                 | git2::Status::WT_TYPECHANGE,
                         ) {
                             files_modified += 1;
-                } else if status_bits.intersects(git2::Status::INDEX_NEW | git2::Status::WT_NEW) {
+                        } else if status_bits
+                            .intersects(git2::Status::INDEX_NEW | git2::Status::WT_NEW)
+                        {
                             files_new += 1;
                         } else if status_bits
                             .intersects(git2::Status::INDEX_DELETED | git2::Status::WT_DELETED)
@@ -311,16 +322,21 @@ pub fn get_repo_status(repo: &git2::Repository) -> RepoStatus {
                             "is_empty() returned true, but no file changes were detected. This is a bug!"
                         );
                     }
-            Some(RepoChanges {
+                    Some(Some(RepoChanges {
                         files_new,
                         files_modified,
                         files_deleted,
-            })
+                    }))
+                }
+            }
         }
     };
 
     let worktrees = repo.worktrees().unwrap().len();
 
+    let submodules = match is_worktree {
+        true => None,
+        false => {
             let mut submodules = Vec::new();
             for submodule in repo.submodules().unwrap() {
                 let submodule_name = submodule.name().unwrap().to_string();
@@ -346,6 +362,9 @@ pub fn get_repo_status(repo: &git2::Repository) -> RepoStatus {
 
                 submodules.push((submodule_name, submodule_status));
             }
+            Some(submodules)
+        }
+    };
 
     let mut branches = Vec::new();
     for (local_branch, _) in repo