use std::cell::RefCell;
use std::path::Path;
use super::repo;
pub const GIT_MAIN_WORKTREE_DIRECTORY: &str = ".git-main-working-tree";
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn invalid_worktree_names() {
assert!(add_worktree(Path::new("/tmp/"), "/leadingslash", None, false).is_err());
assert!(add_worktree(Path::new("/tmp/"), "trailingslash/", None, false).is_err());
assert!(add_worktree(Path::new("/tmp/"), "//", None, false).is_err());
assert!(add_worktree(Path::new("/tmp/"), "test//test", None, false).is_err());
assert!(add_worktree(Path::new("/tmp/"), "test test", None, false).is_err());
assert!(add_worktree(Path::new("/tmp/"), "test\ttest", None, false).is_err());
}
}
struct Init;
struct WithLocalBranchName<'a> {
local_branch_name: String,
local_branch: RefCell<Option<Option<repo::Branch<'a>>>>,
}
struct WithLocalTargetSelected<'a> {
local_branch_name: String,
local_branch: Option<repo::Branch<'a>>,
target_commit: Option<Box<repo::Commit<'a>>>,
}
struct WithRemoteTrackingBranch<'a> {
local_branch_name: String,
local_branch: Option<repo::Branch<'a>>,
target_commit: Option<Box<repo::Commit<'a>>>,
remote_tracking_branch: Option<(String, String)>,
prefix: Option<String>,
}
struct Worktree<'a, S: WorktreeState> {
repo: &'a repo::RepoHandle,
extra: S,
}
impl<'a> WithLocalBranchName<'a> {
fn new(name: String) -> Self {
Self {
local_branch_name: name,
local_branch: RefCell::new(None),
}
}
}
trait WorktreeState {}
impl WorktreeState for Init {}
impl<'a> WorktreeState for WithLocalBranchName<'a> {}
impl<'a> WorktreeState for WithLocalTargetSelected<'a> {}
impl<'a> WorktreeState for WithRemoteTrackingBranch<'a> {}
impl<'a> Worktree<'a, Init> {
fn new(repo: &'a repo::RepoHandle) -> Self {
Self {
repo,
extra: Init {},
}
}
fn set_local_branch_name(self, name: &str) -> Worktree<'a, WithLocalBranchName<'a>> {
Worktree::<WithLocalBranchName> {
repo: self.repo,
extra: WithLocalBranchName::new(name.to_string()),
}
}
}
impl<'a, 'b> Worktree<'a, WithLocalBranchName<'b>>
where
'a: 'b,
{
fn check_local_branch(&self) {
let mut branchref = self.extra.local_branch.borrow_mut();
if branchref.is_none() {
let branch = self.repo.find_local_branch(&self.extra.local_branch_name);
*branchref = Some(if let Ok(branch) = branch {
Some(branch)
} else {
None
});
}
}
fn local_branch_already_exists(&self) -> bool {
if let Some(branch) = &*self.extra.local_branch.borrow() {
return branch.is_some();
}
self.check_local_branch();
(*self.extra.local_branch.borrow())
.as_ref()
.unwrap()
.is_some()
}
fn select_commit(
self,
commit: Option<Box<repo::Commit<'b>>>,
) -> Worktree<'a, WithLocalTargetSelected<'b>> {
self.check_local_branch();
Worktree::<'a, WithLocalTargetSelected> {
repo: self.repo,
extra: WithLocalTargetSelected::<'b> {
local_branch_name: self.extra.local_branch_name,
local_branch: self.extra.local_branch.into_inner().unwrap(),
target_commit: commit,
},
}
}
}
impl<'a> Worktree<'a, WithLocalTargetSelected<'a>> {
fn set_remote_tracking_branch(
self,
branch: Option<(&str, &str)>,
prefix: Option<&str>,
) -> Worktree<'a, WithRemoteTrackingBranch<'a>> {
Worktree::<WithRemoteTrackingBranch> {
repo: self.repo,
extra: WithRemoteTrackingBranch {
local_branch_name: self.extra.local_branch_name,
local_branch: self.extra.local_branch,
target_commit: self.extra.target_commit,
remote_tracking_branch: branch.map(|(s1, s2)| (s1.to_string(), s2.to_string())),
prefix: prefix.map(|prefix| prefix.to_string()),
},
}
}
}
impl<'a> Worktree<'a, WithRemoteTrackingBranch<'a>> {
fn create(self, directory: &Path) -> Result<Option<Vec<String>>, String> {
let mut warnings: Vec<String> = vec![];
let mut branch = if let Some(branch) = self.extra.local_branch {
branch
} else {
self.repo.create_branch(
&self.extra.local_branch_name,
&*self.extra.target_commit.unwrap(),
)?
};
if let Some((remote_name, remote_branch_name)) = self.extra.remote_tracking_branch {
let remote_branch_with_prefix = if let Some(ref prefix) = self.extra.prefix {
if let Ok(remote_branch) = self
.repo
.find_remote_branch(&remote_name, &format!("{prefix}/{remote_branch_name}"))
{
Some(remote_branch)
} else {
None
}
} else {
None
};
let remote_branch_without_prefix = if let Ok(remote_branch) = self
.repo
.find_remote_branch(&remote_name, &remote_branch_name)
{
Some(remote_branch)
} else {
None
};
let remote_branch = if let Some(ref _prefix) = self.extra.prefix {
remote_branch_with_prefix
} else {
remote_branch_without_prefix
};
match remote_branch {
Some(remote_branch) => {
if branch.commit()?.id().hex_string()
!= remote_branch.commit()?.id().hex_string()
{
warnings.push(format!("The local branch \"{}\" and the remote branch \"{}/{}\" differ. Make sure to push/pull afterwards!", &self.extra.local_branch_name, &remote_name, &remote_branch_name));
}
branch.set_upstream(&remote_name, &remote_branch.basename()?)?;
}
None => {
let mut remote = match self.repo.find_remote(&remote_name)? {
Some(remote) => remote,
None => return Err(format!("Remote \"{remote_name}\" not found")),
};
if !remote.is_pushable()? {
return Err(format!(
"Cannot push to non-pushable remote \"{remote_name}\""
));
}
if let Some(prefix) = self.extra.prefix {
remote.push(
&self.extra.local_branch_name,
&format!("{}/{}", prefix, remote_branch_name),
self.repo,
)?;
branch.set_upstream(
&remote_name,
&format!("{}/{}", prefix, remote_branch_name),
)?;
} else {
remote.push(
&self.extra.local_branch_name,
&remote_branch_name,
self.repo,
)?;
branch.set_upstream(&remote_name, &remote_branch_name)?;
}
}
}
}
if self.extra.local_branch_name.contains('/') {
let path = Path::new(&self.extra.local_branch_name);
if let Some(base) = path.parent() {
std::fs::create_dir_all(
directory
.join(GIT_MAIN_WORKTREE_DIRECTORY)
.join("worktrees")
.join(base),
)
.map_err(|error| error.to_string())?;
std::fs::create_dir_all(base).map_err(|error| error.to_string())?;
}
}
self.repo.new_worktree(
&self.extra.local_branch_name,
&directory.join(&self.extra.local_branch_name),
&branch,
)?;
Ok(if warnings.is_empty() {
None
} else {
Some(warnings)
})
}
}
fn validate_worktree_name(name: &str) -> Result<(), String> {
if name.starts_with('/') || name.ends_with('/') {
return Err(format!(
"Invalid worktree name: {}. It cannot start or end with a slash",
name
));
}
if name.contains("//") {
return Err(format!(
"Invalid worktree name: {}. It cannot contain two consecutive slashes",
name
));
}
if name.contains(char::is_whitespace) {
return Err(format!(
"Invalid worktree name: {}. It cannot contain whitespace",
name
));
}
Ok(())
}
pub fn add_worktree(
directory: &Path,
name: &str,
track: Option<(&str, &str)>,
no_track: bool,
) -> Result<Option<Vec<String>>, String> {
let mut warnings: Vec<String> = vec![];
validate_worktree_name(name)?;
let repo = repo::RepoHandle::open(directory, true).map_err(|error| match error.kind {
repo::RepoErrorKind::NotFound => {
String::from("Current directory does not contain a worktree setup")
}
_ => format!("Error opening repo: {}", error),
})?;
let remotes = &repo.remotes()?;
let config = repo::read_worktree_root_config(directory)?;
if repo.find_worktree(name).is_ok() {
return Err(format!("Worktree {} already exists", &name));
}
let track_config = config.and_then(|config| config.track);
let prefix = track_config
.as_ref()
.and_then(|track| track.default_remote_prefix.as_ref());
let enable_tracking = track_config.as_ref().map_or(false, |track| track.default);
let default_remote = track_config
.as_ref()
.map(|track| track.default_remote.clone());
let default_branch_head = repo.default_branch()?.commit_owned()?;
let worktree = Worktree::<Init>::new(&repo).set_local_branch_name(name);
let get_remote_head = |remote_name: &str,
remote_branch_name: &str|
-> Result<Option<Box<repo::Commit>>, String> {
if let Ok(remote_branch) = repo.find_remote_branch(remote_name, remote_branch_name) {
Ok(Some(Box::new(remote_branch.commit_owned()?)))
} else {
Ok(None)
}
};
let worktree = if worktree.local_branch_already_exists() {
worktree.select_commit(None)
} else if let Some((remote_name, remote_branch_name)) = if no_track { None } else { track } {
if let Ok(remote_branch) = repo.find_remote_branch(remote_name, remote_branch_name) {
worktree.select_commit(Some(Box::new(remote_branch.commit_owned()?)))
} else {
worktree.select_commit(Some(Box::new(default_branch_head)))
}
} else {
match remotes.len() {
0 => worktree.select_commit(Some(Box::new(default_branch_head))),
1 => {
let remote_name = &remotes[0];
let commit: Option<Box<repo::Commit>> = ({
if let Some(prefix) = prefix {
get_remote_head(remote_name, &format!("{prefix}/{name}"))?
} else {
None
}
})
.or(get_remote_head(remote_name, name)?)
.or_else(|| Some(Box::new(default_branch_head)));
worktree.select_commit(commit)
}
_ => {
let commit = if let Some(ref default_remote) = default_remote {
if let Some(ref prefix) = prefix {
if let Ok(remote_branch) = repo
.find_remote_branch(default_remote, &format!("{prefix}/{name}"))
{
Some(Box::new(remote_branch.commit_owned()?))
} else {
None
}
} else {
None
}
.or({
if let Ok(remote_branch) =
repo.find_remote_branch(default_remote, name)
{
Some(Box::new(remote_branch.commit_owned()?))
} else {
None
}
})
} else {
None
}.or({
let mut commits = vec![];
for remote_name in remotes.iter() {
let remote_head: Option<Box<repo::Commit>> = ({
if let Some(ref prefix) = prefix {
if let Ok(remote_branch) = repo.find_remote_branch(
remote_name,
&format!("{prefix}/{name}"),
) {
Some(Box::new(remote_branch.commit_owned()?))
} else {
None
}
} else {
None
}
})
.or({
if let Ok(remote_branch) =
repo.find_remote_branch(remote_name, name)
{
Some(Box::new(remote_branch.commit_owned()?))
} else {
None
}
})
.or(None);
commits.push(remote_head);
}
let mut commits = commits
.into_iter()
.flatten()
.collect::<Vec<Box<repo::Commit>>>();
if commits.is_empty() {
Some(Box::new(default_branch_head))
} else if commits.len() == 1 {
Some(commits.swap_remove(0))
} else if commits.windows(2).any(|window| {
let c1 = &window[0];
let c2 = &window[1];
(*c1).id().hex_string() != (*c2).id().hex_string()
}) {
warnings.push(
"Branch exists on multiple remotes, but they deviate. Selecting default branch instead".to_string()
);
Some(Box::new(default_branch_head))
} else {
Some(commits.swap_remove(0))
}
});
worktree.select_commit(commit)
}
}
};
let worktree = if no_track {
worktree.set_remote_tracking_branch(None, prefix.map(|s| s.as_str()))
} else if let Some((remote_name, remote_branch_name)) = track {
worktree.set_remote_tracking_branch(
Some((remote_name, remote_branch_name)),
None, )
} else if !enable_tracking {
worktree.set_remote_tracking_branch(None, prefix.map(|s| s.as_str()))
} else {
match remotes.len() {
0 => worktree.set_remote_tracking_branch(None, prefix.map(|s| s.as_str())),
1 => worktree
.set_remote_tracking_branch(Some((&remotes[0], name)), prefix.map(|s| s.as_str())),
_ => {
if let Some(default_remote) = default_remote {
worktree.set_remote_tracking_branch(
Some((&default_remote, name)),
prefix.map(|s| s.as_str()),
)
} else {
worktree.set_remote_tracking_branch(None, prefix.map(|s| s.as_str()))
}
}
}
};
worktree.create(directory)?;
Ok(if warnings.is_empty() {
None
} else {
Some(warnings)
})
}