sausage/userland/git/src/checkout.rs

//! Worktree writer shared by checkout and future clone/merge operations.
//! Caller must hold the repository .git lock. No refs are changed here.
use crate::Result;
use crate::repo::error;
use gix::bstr::ByteSlice;
use gix::index::{
    State,
    entry::{Mode, Stat},
};
use std::{
    collections::{BTreeMap, BTreeSet},
    fs,
    path::Path,
};

type Files = BTreeMap<String, (gix::ObjectId, Mode)>;

fn tree_state(repo: &gix::Repository, tree: Option<gix::ObjectId>) -> Result<State> {
    match tree {
        Some(id) => State::from_tree(&id, &repo.objects, Default::default())
            .map_err(|e| error("checkout tree", e)),
        None => Ok(State::new(repo.object_hash())),
    }
}

fn files(state: &State) -> Result<Files> {
    state
        .entries()
        .iter()
        .map(|e| {
            let path = e
                .path(state)
                .to_str()
                .map_err(|e| error("checkout path UTF-8", e))?
                .to_owned();
            crate::repo::validate_path(&path)?;
            if path
                .split('/')
                .any(|s| s.eq_ignore_ascii_case(".git") || s.contains('\\'))
            {
                return Err(error("checkout path", path));
            }
            Ok((path, (e.id, e.mode)))
        })
        .collect()
}

fn metadata(path: &Path) -> Result<Option<fs::Metadata>> {
    match fs::symlink_metadata(path) {
        Ok(m) => Ok(Some(m)),
        Err(e)
            if matches!(
                e.kind(),
                std::io::ErrorKind::NotFound | std::io::ErrorKind::NotADirectory
            ) =>
        {
            Ok(None)
        }
        Err(e) => Err(error(
            "checkout metadata",
            format!("{}: {e}", path.display()),
        )),
    }
}

fn regular(mode: Mode) -> bool {
    mode == Mode::FILE || mode == Mode::FILE_EXECUTABLE
}

fn clean(repo: &gix::Repository, path: &Path, expected: (gix::ObjectId, Mode)) -> Result<bool> {
    let Some(meta) = metadata(path)? else {
        // These entries are intentionally not materialized by this writer.
        return Ok(!regular(expected.1));
    };
    if expected.1 == Mode::SYMLINK && meta.file_type().is_symlink() {
        let target = fs::read_link(path).map_err(|e| error("checkout symlink", e))?;
        let target = target
            .to_str()
            .ok_or_else(|| error("checkout symlink", "non-UTF-8 link target"))?;
        return Ok(gix::objs::compute_hash(
            repo.object_hash(),
            gix::objs::Kind::Blob,
            target.as_bytes(),
        )
        .map_err(|e| error("checkout symlink hash", e))?
            == expected.0);
    }
    if !meta.is_file() || !regular(expected.1) {
        return Ok(false);
    }
    #[cfg(unix)]
    {
        use std::os::unix::fs::PermissionsExt;
        if (meta.permissions().mode() & 0o111 != 0) != (expected.1 == Mode::FILE_EXECUTABLE) {
            return Ok(false);
        }
    }
    let bytes = fs::read(path).map_err(|e| error("checkout read", e))?;
    let id = gix::objs::compute_hash(repo.object_hash(), gix::objs::Kind::Blob, &bytes)
        .map_err(|e| error("checkout hash", e))?;
    Ok(id == expected.0)
}

/// Verify every leaf in a directory is tracked and scheduled for removal. Never remove untracked leaves.
fn removable_dir(root: &Path, dir: &Path, removed: &BTreeSet<String>) -> Result<bool> {
    for child in fs::read_dir(dir).map_err(|e| error("checkout directory", e))? {
        let child = child.map_err(|e| error("checkout directory", e))?;
        let path = child.path();
        let meta = fs::symlink_metadata(&path).map_err(|e| error("checkout metadata", e))?;
        if meta.is_dir() {
            if !removable_dir(root, &path, removed)? {
                return Ok(false);
            }
        } else if !path
            .strip_prefix(root)
            .ok()
            .and_then(Path::to_str)
            .is_some_and(|p| removed.contains(p))
        {
            return Ok(false);
        }
    }
    Ok(true)
}

/// Diff trees, preflight all overwrite risks, write raw blobs, remove tracked paths and empty
/// directories, then atomically publish the index last. On I/O failure HEAD/refs remain unchanged;
/// worktree edits already completed are not rolled back. Filters/attributes are not applied.
/// Unrelated staged and unstaged edits survive a non-force switch. Force resets tracked paths,
/// but never overwrites/deletes an untracked file or follows a symlink directory.
pub(crate) fn write_tree_to_worktree(
    repo: &gix::Repository,
    old_tree: Option<gix::ObjectId>,
    new_tree: gix::ObjectId,
    force: bool,
) -> Result<Vec<String>> {
    write_tree(repo, old_tree, new_tree, force, false)
}

/// Abort discards tracked changes without refusing dirty files. Paths no longer in the
/// index are untracked, even if HEAD once contained them. Preserve any such obstruction,
/// restoring the HEAD index but leaving a warning/local difference rather than losing bytes.
pub(crate) fn abort_to_tree(repo: &gix::Repository, tree: gix::ObjectId) -> Result<Vec<String>> {
    write_tree(repo, None, tree, true, true)
}
fn write_tree(
    repo: &gix::Repository,
    old_tree: Option<gix::ObjectId>,
    new_tree: gix::ObjectId,
    force: bool,
    preserve_obstructions: bool,
) -> Result<Vec<String>> {
    let root = repo
        .workdir()
        .ok_or_else(|| error("checkout", "bare repository"))?;
    let old = files(&tree_state(repo, old_tree)?)?;
    let mut state = tree_state(repo, Some(new_tree))?;
    let new = files(&state)?;
    let index_path = repo.git_dir().join("index");
    let existing = if index_path
        .try_exists()
        .map_err(|e| error("checkout index", e))?
    {
        Some(
            gix::index::File::at(&index_path, repo.object_hash(), false, Default::default())
                .map_err(|e| error("checkout index", e))?,
        )
    } else {
        None
    };
    if existing.as_ref().is_some_and(|i| i.is_sparse()) {
        return Err(error(
            "checkout index",
            "sparse indexes are unsupported; expand the index before checkout",
        ));
    }
    let indexed = existing
        .as_ref()
        .map(|i| files(i))
        .transpose()?
        .unwrap_or_default();
    let mut affected: BTreeSet<String> = old
        .keys()
        .chain(new.keys())
        .filter(|p| force || old.get(*p) != new.get(*p))
        .cloned()
        .collect();
    if force {
        affected.extend(indexed.keys().cloned());
    } else {
        // A staged addition can conflict structurally even when missing from the worktree.
        // Include index ancestors/descendants of changed paths, so it cannot silently become a D/F conflict.
        let changed = affected.clone();
        affected.extend(
            indexed
                .keys()
                .filter(|p| {
                    let prefix = format!("{p}/");
                    changed
                        .range(prefix.clone()..)
                        .next()
                        .is_some_and(|q| q.starts_with(&prefix))
                        || Path::new(p)
                            .ancestors()
                            .skip(1)
                            .filter_map(Path::to_str)
                            .any(|q| changed.contains(q))
                })
                .cloned(),
        );
    }
    let conflicts: BTreeSet<_> = existing
        .as_ref()
        .map(|i| {
            i.entries()
                .iter()
                .filter(|e| e.stage_raw() != 0)
                .map(|e| e.path(i).to_str_lossy().into_owned())
                .collect()
        })
        .unwrap_or_default();
    let tracked: BTreeSet<_> = old.keys().chain(indexed.keys()).cloned().collect();
    let removed: BTreeSet<_> = affected
        .iter()
        .filter(|p| !new.get(*p).is_some_and(|e| regular(e.1)))
        .cloned()
        .collect();
    let mut dirty = BTreeSet::new();
    let mut blocked = BTreeSet::new();
    let mut warnings = Vec::new();
    // Validate paths in every source before touching anything, including index entries.
    for path in &affected {
        let dest = root.join(path);
        // A symlink/file parent may be replaced only if tracked and scheduled for deletion.
        let mut parent = Path::new(path).parent();
        while let Some(p) = parent.filter(|p| !p.as_os_str().is_empty()) {
            if let Some(m) = metadata(&root.join(p))? {
                if m.file_type().is_symlink()
                    || (!m.is_dir()
                        && !removed.contains(p.to_str().unwrap_or(""))
                        && new.get(path).is_some_and(|e| regular(e.1)))
                {
                    blocked.insert(path.clone());
                }
            }
            parent = p.parent();
        }
        // Inspect the actual worktree: a removed tracked ancestor may have been replaced
        // by a directory containing untracked files. A real file parent yields None via
        // metadata's NotADirectory handling, without hiding existing descendants.
        let meta = metadata(&dest)?;
        if let Some(m) = &meta {
            if m.is_dir() && new.get(path).is_some_and(|e| regular(e.1)) {
                if !removable_dir(root, &dest, &removed)? {
                    blocked.insert(path.clone());
                }
            } else if !m.is_dir() && !tracked.contains(path) {
                blocked.insert(path.clone());
            }
        }
        if !force {
            if indexed.get(path) != old.get(path) && existing.is_some() {
                dirty.insert(path.clone());
            }
            if conflicts.contains(path) {
                dirty.insert(path.clone());
            }
            if let Some(expected) = old.get(path) {
                if !clean(repo, &dest, *expected)? {
                    dirty.insert(path.clone());
                }
            }
        }
    }
    if !blocked.is_empty() {
        let paths = blocked.iter().cloned().collect::<Vec<_>>().join(", ");
        if preserve_obstructions {
            warnings.push(format!("abort preserved untracked/obstructing paths; HEAD index restored but worktree differs: {paths}"));
        } else {
            return Err(error(
                "checkout untracked/obstructing paths (also protected with force)",
                paths,
            ));
        }
    }
    if !dirty.is_empty() {
        return Err(error(
            "checkout would overwrite uncommitted changes",
            dirty.into_iter().collect::<Vec<_>>().join(", "),
        ));
    }
    // Read every required blob before beginning mutations (missing/corrupt objects must not partially checkout).
    let mut blobs = BTreeMap::new();
    for path in &affected {
        if blocked.contains(path) {
            continue;
        }
        if let Some((id, mode)) = new.get(path) {
            if regular(*mode) {
                blobs.insert(
                    path.clone(),
                    repo.find_blob(*id)
                        .map_err(|e| error("checkout blob", e))?
                        .data
                        .to_vec(),
                );
            } else {
                warnings.push(format!(
                    "{path}: symlink/submodule unsupported; skipped worktree materialization"
                ));
            }
        }
    }
    // Remove old tracked leaves first, deepest first; never recursively delete a directory.
    let mut deletions: Vec<_> = affected
        .iter()
        .filter(|p| tracked.contains(*p) && !blocked.contains(*p))
        .collect();
    deletions.sort_by_key(|p| std::cmp::Reverse(p.len()));
    for path in deletions {
        let dest = root.join(path);
        // Ancestors may have been removed already (only affects structurally conflicting staged entries).
        match fs::symlink_metadata(&dest) {
            Ok(m) if !m.is_dir() => {
                fs::remove_file(&dest).map_err(|e| error("checkout remove", e))?
            }
            Ok(_) => {}
            Err(e)
                if matches!(
                    e.kind(),
                    std::io::ErrorKind::NotFound | std::io::ErrorKind::NotADirectory
                ) => {}
            Err(e) => return Err(error("checkout remove", e)),
        }
        prune_empty(root, dest.parent());
    }
    for (path, bytes) in blobs {
        let dest = root.join(&path);
        if dest.is_dir() {
            remove_empty_dirs(&dest)?;
        }
        if let Some(parent) = dest.parent() {
            fs::create_dir_all(parent).map_err(|e| error("checkout mkdir", e))?;
        }
        fs::write(&dest, bytes).map_err(|e| error("checkout write", e))?;
        #[cfg(unix)]
        {
            use std::os::unix::fs::PermissionsExt;
            let mode = if new[&path].1 == Mode::FILE_EXECUTABLE {
                0o755
            } else {
                0o644
            };
            fs::set_permissions(&dest, fs::Permissions::from_mode(mode))
                .map_err(|e| error("checkout chmod", e))?;
        }
    }
    #[cfg(not(unix))]
    if affected
        .iter()
        .any(|p| new.get(p).is_some_and(|e| e.1 == Mode::FILE_EXECUTABLE))
    {
        warnings.push(
            "executable bits preserved in index; filesystem chmod unavailable on this target"
                .into(),
        );
    }
    // Preserve unrelated staging, including staged additions, deletions, and conflict stages.
    if let Some(index) = &existing {
        state.remove_entries(|_, path, _| !affected.contains(path.to_str().unwrap_or("")));
        for entry in index.entries() {
            if !affected.contains(entry.path(index).to_str().unwrap_or("")) {
                state.dangerously_push_entry(
                    entry.stat,
                    entry.id,
                    entry.flags,
                    entry.mode,
                    entry.path(index),
                );
            }
        }
        state.sort_entries();
        state.remove_tree(); // Unrelated staging may differ from the target tree.
    }
    for (entry, path) in state.entries_mut_with_paths() {
        if affected.contains(path.to_str().unwrap_or(""))
            && !blocked.contains(path.to_str().unwrap_or(""))
            && regular(entry.mode)
        {
            let stat = gix::index::fs::Metadata::from_path_no_follow(
                &root.join(path.to_str().map_err(|e| error("checkout path", e))?),
            )
            .map_err(|e| error("checkout stat", e))?;
            entry.stat = Stat::from_fs(&stat).map_err(|e| error("checkout stat", e))?;
        }
    }
    gix::index::File::from_state(state, index_path)
        .write(Default::default())
        .map_err(|e| error("checkout index write", e))?;
    Ok(warnings)
}

fn prune_empty(root: &Path, mut parent: Option<&Path>) {
    while let Some(p) = parent.filter(|p| *p != root && p.starts_with(root)) {
        if fs::remove_dir(p).is_err() {
            break;
        }
        parent = p.parent();
    }
}

fn remove_empty_dirs(path: &Path) -> Result<()> {
    for child in fs::read_dir(path).map_err(|e| error("checkout directory", e))? {
        let child = child.map_err(|e| error("checkout directory", e))?;
        if child
            .file_type()
            .map_err(|e| error("checkout directory", e))?
            .is_dir()
        {
            remove_empty_dirs(&child.path())?;
        }
    }
    fs::remove_dir(path).map_err(|e| error("checkout directory not empty", e))
}