use std::collections::HashSet;
use std::io::{BufRead, Write};
use std::path::{Path, PathBuf};
use std::sync::atomic::{AtomicBool, Ordering};
use clap::Parser;
use mkit_core::hash::{Hash, ZERO};
use mkit_core::ignore::{self, IgnoreList};
use mkit_core::index::{self, EntryStatus, Index, IndexEntry};
use mkit_core::layout::RepoLayout;
use mkit_core::object::{Blob, Object};
use mkit_core::ops::{HunkLineKind, PatchHunk, apply_hunks_subset, enumerate_hunks};
use mkit_core::serialize;
use mkit_core::store::{ObjectSink, ObjectStore};
use mkit_core::worktree;
use rayon::prelude::*;
use crate::clap_shim;
use crate::exit;
#[derive(Debug, Parser)]
#[command(
name = "mkit add",
about = "Stage files (paths, `.`, `-A`, or `-u`) into the index."
)]
#[allow(clippy::struct_excessive_bools)]
struct AddOpts {
#[arg(short = 'A', long)]
all: bool,
#[arg(short = 'u', long)]
update: bool,
#[arg(short = 'f', long)]
force: bool,
#[arg(short = 'p', long)]
patch: bool,
paths: Vec<String>,
}
pub(super) fn stage_tracked_changes(
layout: &RepoLayout,
store: &ObjectStore,
) -> Result<(), String> {
let root = layout.worktree_root();
let mut idx = super::read_or_seed_index_from_head(layout, store)?;
let batch = store.batch();
for entry in &mut idx.entries {
if entry.status == EntryStatus::Removed {
continue;
}
if !index::validate_index_path(&entry.path) {
return Err(format!("invalid index path: {}", entry.path));
}
let abs = root.join(&entry.path);
let meta = match abs.symlink_metadata() {
Ok(meta) => meta,
Err(e)
if matches!(
e.kind(),
std::io::ErrorKind::NotFound | std::io::ErrorKind::NotADirectory
) =>
{
entry.status = EntryStatus::Removed;
entry.object_hash = ZERO;
continue;
}
Err(e) => return Err(format!("metadata {}: {e}", abs.display())),
};
if worktree::stat_matches(entry, &meta) {
continue;
}
let (status, h, stat) = if meta.file_type().is_file() {
let (h, opened_meta) = worktree::hash_file_with_metadata(&batch, &abs)
.map_err(|e| format!("read/store {}: {e}", abs.display()))?;
let stat = worktree::stat_cache_fields(&opened_meta);
(file_status_from_meta(&opened_meta, entry.status), h, stat)
} else if meta.file_type().is_symlink() {
let target = std::fs::read_link(&abs)
.map_err(|e| format!("read link {}: {e}", abs.display()))?;
let target_str = target
.to_str()
.ok_or_else(|| "symlink target is not valid UTF-8".to_string())?;
if !worktree::validate_symlink_target(target_str) {
return Err(format!("invalid symlink target: {target_str}"));
}
let blob = Object::Blob(Blob {
data: target_str.as_bytes().to_vec(),
});
let ser = serialize::serialize(&blob).map_err(|e| format!("serialize: {e}"))?;
let h = batch.put(&ser).map_err(|e| format!("store: {e}"))?;
(EntryStatus::Symlink, h, (0, 0, 0, 0))
} else {
entry.status = EntryStatus::Removed;
entry.object_hash = ZERO;
continue;
};
entry.status = status;
entry.object_hash = h;
entry.mtime_ns = stat.0;
entry.size = stat.1;
entry.ino = stat.2;
entry.ctime_ns = stat.3;
}
batch.commit().map_err(|e| format!("store: {e}"))?;
index::write_index(layout, &idx).map_err(|e| format!("write index: {e}"))
}
#[cfg(unix)]
fn file_status_from_meta(meta: &std::fs::Metadata, _previous: EntryStatus) -> EntryStatus {
use std::os::unix::fs::PermissionsExt;
if meta.permissions().mode() & 0o111 != 0 {
EntryStatus::Executable
} else {
EntryStatus::Blob
}
}
#[cfg(not(unix))]
fn file_status_from_meta(_meta: &std::fs::Metadata, previous: EntryStatus) -> EntryStatus {
if previous == EntryStatus::Executable {
EntryStatus::Executable
} else {
EntryStatus::Blob
}
}
fn worktree_err_exit_code(e: &worktree::WorktreeError) -> u8 {
match e {
worktree::WorktreeError::Io(_) | worktree::WorktreeError::FileTooLarge(_) => exit::NOINPUT,
worktree::WorktreeError::Object(_) | worktree::WorktreeError::Store(_) => exit::CANTCREAT,
worktree::WorktreeError::InvalidSymlinkTarget(_) | worktree::WorktreeError::InvalidUtf8 => {
exit::DATAERR
}
}
}
#[must_use]
pub fn run(args: &[String]) -> u8 {
let opts = match clap_shim::parse::<AddOpts>("mkit add", args) {
Ok(o) => o,
Err(code) => return code,
};
let cwd = match std::env::current_dir() {
Ok(p) => p,
Err(e) => return emit_err(&format!("cwd: {e}"), exit::NOINPUT),
};
let layout = match super::resolve_layout(&cwd) {
Ok(layout) => layout,
Err(code) => return code,
};
let store = match super::open_store_configured(&layout) {
Ok(s) => s,
Err(e) => return emit_err(&format!("not a mkit repo: {e}"), exit::GENERAL_ERROR),
};
let _lock = match super::acquire_worktree_lock(&layout) {
Ok(l) => l,
Err(code) => return code,
};
if opts.patch {
if opts.all || opts.update {
return emit_err(
"-p/--patch cannot be combined with -A/--all or -u/--update",
exit::USAGE,
);
}
if opts.paths.is_empty() {
return emit_err("-p/--patch requires one or more file paths", exit::USAGE);
}
return run_patch(&layout, &store, &opts.paths, opts.force);
}
if opts.all && opts.update {
return emit_err("cannot combine -A/--all with -u/--update", exit::USAGE);
}
if (opts.all || opts.update) && !opts.paths.is_empty() {
return emit_err(
"-A/--all and -u/--update take no path arguments",
exit::USAGE,
);
}
if opts.update {
return match stage_tracked_changes(&layout, &store) {
Ok(()) => exit::OK,
Err(e) => emit_err(&e, exit::GENERAL_ERROR),
};
}
let mut idx = match super::read_or_seed_index_from_head(&layout, &store) {
Ok(i) => i,
Err(e) => return emit_err(&e, exit::GENERAL_ERROR),
};
let batch = store.batch();
if opts.all {
if let Err(code) = add_whole_worktree(&cwd, &batch, &mut idx) {
return code;
}
} else if opts.paths.is_empty() {
return emit_err(
"no paths given (use `.`, -A, -u, or one or more paths)",
exit::USAGE,
);
} else {
let ignores = match ignore::load(&cwd) {
Ok(i) => i,
Err(e) => return emit_err(&format!("read ignore file: {e}"), exit::GENERAL_ERROR),
};
for target in &opts.paths {
if target == "." {
if let Err(code) = add_whole_worktree(&cwd, &batch, &mut idx) {
return code;
}
} else {
let p = Path::new(target);
let abs = if p.is_absolute() {
p.to_path_buf()
} else {
cwd.join(p)
};
if let Err(e) = ensure_within_repo(&cwd, &abs) {
return emit_err(&e, exit::DATAERR);
}
match add_one(&cwd, p, &batch, &mut idx, &ignores, opts.force) {
Ok(_) => {}
Err(code) => return code,
}
}
}
}
if let Err(e) = batch.commit() {
return emit_err(&format!("store: {e}"), exit::CANTCREAT);
}
match index::write_index(&layout, &idx) {
Ok(()) => exit::OK,
Err(e) => emit_err(&format!("write index: {e}"), exit::CANTCREAT),
}
}
fn add_whole_worktree(
root: &Path,
sink: &(dyn ObjectSink + Sync),
idx: &mut Index,
) -> Result<(), u8> {
let ignores = match ignore::load(root) {
Ok(i) => i,
Err(e) => {
return Err(emit_err(
&format!("read ignore file: {e}"),
exit::GENERAL_ERROR,
));
}
};
let mut seen = HashSet::new();
let mut pending = Vec::new();
add_tree(
root,
root,
false,
sink,
idx,
&ignores,
&mut seen,
&mut pending,
)?;
let hashed = hash_pending_batch(&pending, sink);
if let Some(pos) = hashed
.iter()
.position(|h| matches!(h, HashOutcome::Failed(_)))
{
let HashOutcome::Failed(e) = &hashed[pos] else {
unreachable!("position() just matched a Failed variant")
};
return Err(emit_err(&e.message, e.code));
}
for (p, outcome) in pending.into_iter().zip(hashed) {
let HashOutcome::Done(hashed_file) = outcome else {
unreachable!(
"Skipped only occurs once a Failed entry exists, and the check above already returned on any Failed entry"
)
};
stage_hashed(idx, p.rel_str.clone(), hashed_file);
seen.insert(p.rel_str);
}
mark_missing_paths_removed(root, idx, &seen);
Ok(())
}
const HASH_FANOUT_FILES_PER_THREAD: usize = 8;
fn hash_fanout_threshold() -> usize {
HASH_FANOUT_FILES_PER_THREAD.saturating_mul(rayon::current_num_threads())
}
fn hash_one(sink: &dyn ObjectSink, aborted: &AtomicBool, p: &PendingHash) -> HashOutcome {
if aborted.load(Ordering::Relaxed) {
return HashOutcome::Skipped;
}
match hash_pending(sink, p) {
Ok(v) => HashOutcome::Done(v),
Err(e) => {
aborted.store(true, Ordering::Relaxed);
HashOutcome::Failed(e)
}
}
}
fn hash_pending_batch(pending: &[PendingHash], sink: &(dyn ObjectSink + Sync)) -> Vec<HashOutcome> {
let aborted = AtomicBool::new(false);
if pending.len() < hash_fanout_threshold() {
return pending
.iter()
.map(|p| hash_one(sink, &aborted, p))
.collect();
}
pending
.par_iter()
.map(|p| hash_one(sink, &aborted, p))
.collect()
}
enum HashOutcome {
Done(HashedFile),
Failed(HashError),
Skipped,
}
struct PendingHash {
abs: PathBuf,
rel_str: String,
previous_status: EntryStatus,
}
enum Routed {
Done(String),
NeedsHash(PendingHash),
}
fn route_path(
root: &Path,
rel: &Path,
sink: &dyn ObjectSink,
idx: &mut Index,
ignores: &IgnoreList,
force: bool,
) -> Result<Routed, u8> {
let abs = if rel.is_absolute() {
rel.to_path_buf()
} else {
root.join(rel)
};
let meta = abs
.symlink_metadata()
.map_err(|e| emit_err(&format!("metadata {}: {e}", abs.display()), exit::NOINPUT))?;
let rel_str = abs
.strip_prefix(root)
.unwrap_or(rel)
.to_string_lossy()
.replace('\\', "/");
if !index::validate_index_path(&rel_str) {
return Err(emit_err(&format!("invalid path: {rel_str}"), exit::DATAERR));
}
let existing_pos = idx.find_entry(&rel_str);
let previous_status = existing_pos.map_or(EntryStatus::Blob, |i| idx.entries[i].status);
let already_tracked = previous_status != EntryStatus::Removed && existing_pos.is_some();
if !force && !already_tracked && ignores.is_ignored_with_ancestors(&rel_str, meta.is_dir()) {
return Err(emit_err(
&format!("path '{rel_str}' is ignored; use -f to add it anyway"),
exit::USAGE,
));
}
if let Some(existing) = existing_pos
&& worktree::stat_matches(&idx.entries[existing], &meta)
{
return Ok(Routed::Done(rel_str));
}
if meta.file_type().is_file() {
Ok(Routed::NeedsHash(PendingHash {
abs,
rel_str,
previous_status,
}))
} else if meta.file_type().is_symlink() {
let target = std::fs::read_link(&abs)
.map_err(|e| emit_err(&format!("read link {}: {e}", abs.display()), exit::NOINPUT))?;
let target_str = match target.to_str() {
Some(t) => t.to_string(),
None => return Err(emit_err("symlink target is not valid UTF-8", exit::DATAERR)),
};
if !worktree::validate_symlink_target(&target_str) {
return Err(emit_err(
&format!("invalid symlink target: {target_str}"),
exit::DATAERR,
));
}
let blob = Object::Blob(Blob {
data: target_str.into_bytes(),
});
let ser = serialize::serialize(&blob)
.map_err(|e| emit_err(&format!("serialize: {e}"), exit::DATAERR))?;
let h = sink
.put(&ser)
.map_err(|e| emit_err(&format!("store: {e}"), exit::CANTCREAT))?;
let entry = IndexEntry {
path: rel_str.clone(),
status: EntryStatus::Symlink,
object_hash: h,
mtime_ns: 0,
size: 0,
ino: 0,
ctime_ns: 0,
};
idx.remove_directory_conflicts(&entry.path);
idx.upsert_entry(entry);
Ok(Routed::Done(rel_str))
} else {
Err(emit_err(
&format!("not a regular file: {}", abs.display()),
exit::NOINPUT,
))
}
}
type HashedFile = (EntryStatus, Hash, (u64, u64, u64, u64));
struct HashError {
message: String,
code: u8,
}
fn hash_pending(sink: &dyn ObjectSink, p: &PendingHash) -> Result<HashedFile, HashError> {
let (h, opened_meta) =
worktree::hash_file_with_metadata(sink, &p.abs).map_err(|e| HashError {
message: format!("{}: {e}", p.abs.display()),
code: worktree_err_exit_code(&e),
})?;
let stat = worktree::stat_cache_fields(&opened_meta);
let status = file_status_from_meta(&opened_meta, p.previous_status);
Ok((status, h, stat))
}
fn stage_hashed(idx: &mut Index, rel_str: String, hashed: HashedFile) {
let (status, h, stat) = hashed;
let entry = IndexEntry {
path: rel_str,
status,
object_hash: h,
mtime_ns: stat.0,
size: stat.1,
ino: stat.2,
ctime_ns: stat.3,
};
idx.remove_directory_conflicts(&entry.path);
idx.upsert_entry(entry);
}
fn add_one(
root: &Path,
rel: &Path,
sink: &dyn ObjectSink,
idx: &mut Index,
ignores: &IgnoreList,
force: bool,
) -> Result<String, u8> {
match route_path(root, rel, sink, idx, ignores, force)? {
Routed::Done(rel_str) => Ok(rel_str),
Routed::NeedsHash(p) => {
let hashed = hash_pending(sink, &p).map_err(|e| emit_err(&e.message, e.code))?;
stage_hashed(idx, p.rel_str.clone(), hashed);
Ok(p.rel_str)
}
}
}
fn add_tree(
root: &Path,
dir: &Path,
parent_ignored: bool,
sink: &dyn ObjectSink,
idx: &mut Index,
ignores: &IgnoreList,
seen: &mut HashSet<String>,
pending: &mut Vec<PendingHash>,
) -> Result<(), u8> {
let rd = std::fs::read_dir(dir)
.map_err(|e| emit_err(&format!("read dir {}: {e}", dir.display()), exit::NOINPUT))?;
for ent in rd.flatten() {
let p = ent.path();
let meta = p
.symlink_metadata()
.map_err(|e| emit_err(&format!("metadata {}: {e}", p.display()), exit::NOINPUT))?;
let is_dir = meta.file_type().is_dir();
let rel_path = p
.strip_prefix(root)
.unwrap_or(&p)
.to_string_lossy()
.replace('\\', "/");
let entry_ignored = parent_ignored || ignores.is_ignored(&rel_path, is_dir);
if entry_ignored && !super::index_tracks_path_or_descendant(idx, &rel_path) {
continue;
}
if meta.file_type().is_dir() {
add_tree(root, &p, entry_ignored, sink, idx, ignores, seen, pending)?;
} else if meta.file_type().is_file() || meta.file_type().is_symlink() {
match route_path(root, &p, sink, idx, ignores, true)? {
Routed::Done(rel) => {
seen.insert(rel);
}
Routed::NeedsHash(pend) => pending.push(pend),
}
}
}
Ok(())
}
fn mark_missing_paths_removed(root: &Path, idx: &mut Index, seen: &HashSet<String>) {
for entry in &mut idx.entries {
if entry.status != EntryStatus::Removed
&& !seen.contains(&entry.path)
&& matches!(
root.join(&entry.path).symlink_metadata(),
Err(e) if matches!(
e.kind(),
std::io::ErrorKind::NotFound | std::io::ErrorKind::NotADirectory
)
)
{
entry.status = EntryStatus::Removed;
entry.object_hash = ZERO;
}
}
}
struct PatchOutcome {
staged: bool,
quit: bool,
}
fn run_patch(layout: &RepoLayout, store: &ObjectStore, paths: &[String], force: bool) -> u8 {
let root = layout.worktree_root();
let mut idx = match super::read_or_seed_index_from_head(layout, store) {
Ok(i) => i,
Err(e) => return emit_err(&e, exit::GENERAL_ERROR),
};
let ignores = match ignore::load(root) {
Ok(i) => i,
Err(e) => return emit_err(&format!("read ignore file: {e}"), exit::GENERAL_ERROR),
};
let stdin = std::io::stdin();
let mut input = stdin.lock();
let mut any_staged = false;
for target in paths {
match patch_one_file(
root,
Path::new(target),
store,
&mut idx,
&ignores,
force,
&mut input,
) {
Ok(outcome) => {
any_staged |= outcome.staged;
if outcome.quit {
break;
}
}
Err(code) => return code,
}
}
if any_staged && let Err(e) = index::write_index(layout, &idx) {
return emit_err(&format!("write index: {e}"), exit::CANTCREAT);
}
exit::OK
}
fn patch_one_file(
root: &Path,
rel: &Path,
store: &ObjectStore,
idx: &mut Index,
ignores: &IgnoreList,
force: bool,
input: &mut impl BufRead,
) -> Result<PatchOutcome, u8> {
let abs = if rel.is_absolute() {
rel.to_path_buf()
} else {
root.join(rel)
};
let meta = abs
.symlink_metadata()
.map_err(|e| emit_err(&format!("metadata {}: {e}", abs.display()), exit::NOINPUT))?;
let rel_str = abs
.strip_prefix(root)
.unwrap_or(rel)
.to_string_lossy()
.replace('\\', "/");
if !index::validate_index_path(&rel_str) {
return Err(emit_err(&format!("invalid path: {rel_str}"), exit::DATAERR));
}
if let Err(e) = ensure_within_repo(root, &abs) {
return Err(emit_err(&e, exit::DATAERR));
}
if !meta.file_type().is_file() {
return Err(emit_err(
&format!("-p/--patch supports regular files only: {rel_str}"),
exit::USAGE,
));
}
let already_tracked = idx
.find_entry(&rel_str)
.is_some_and(|i| idx.entries[i].status != EntryStatus::Removed);
if !force && !already_tracked && ignores.is_ignored_with_ancestors(&rel_str, false) {
return Err(emit_err(
&format!("path '{rel_str}' is ignored; use -f to add it anyway"),
exit::USAGE,
));
}
let base = match idx.find_entry(&rel_str) {
Some(i) if idx.entries[i].status != EntryStatus::Removed => {
worktree::read_blob(store, &idx.entries[i].object_hash)
.map_err(|e| emit_err(&format!("read staged blob: {e}"), exit::GENERAL_ERROR))?
}
_ => Vec::new(),
};
let previous_status = idx
.find_entry(&rel_str)
.map_or(EntryStatus::Blob, |i| idx.entries[i].status);
let (opened_meta, work_bytes) = worktree::read_regular_file_bounded(&abs)
.map_err(|e| emit_err(&format!("read {}: {e}", abs.display()), exit::NOINPUT))?;
let hunks = match enumerate_hunks(&base, &work_bytes) {
None => {
eprintln!("{rel_str}: binary file — skipped (use `mkit add` to stage whole)");
return Ok(PatchOutcome {
staged: false,
quit: false,
});
}
Some(h) if h.is_empty() => {
eprintln!("{rel_str}: no changes to stage");
return Ok(PatchOutcome {
staged: false,
quit: false,
});
}
Some(h) => h,
};
let (selected, quit) = select_hunks(&rel_str, &hunks, input)?;
if selected.is_empty() {
return Ok(PatchOutcome {
staged: false,
quit,
});
}
let new_bytes = apply_hunks_subset(&base, &hunks, &selected);
let h = worktree::store_file_object(store, &new_bytes)
.map_err(|e| emit_err(&format!("store: {e}"), exit::CANTCREAT))?;
let status = file_status_from_meta(&opened_meta, previous_status);
let entry = IndexEntry {
path: rel_str.clone(),
status,
object_hash: h,
mtime_ns: 0,
size: 0,
ino: 0,
ctime_ns: 0,
};
idx.remove_directory_conflicts(&entry.path);
idx.upsert_entry(entry);
eprintln!(
"{rel_str}: staged {} of {} hunks",
selected.len(),
hunks.len()
);
Ok(PatchOutcome { staged: true, quit })
}
fn select_hunks(
path: &str,
hunks: &[PatchHunk],
input: &mut impl BufRead,
) -> Result<(Vec<usize>, bool), u8> {
let mut stderr = std::io::stderr().lock();
let mut selected = Vec::new();
let mut auto: Option<bool> = None;
let mut i = 0;
while i < hunks.len() {
if let Some(stage_rest) = auto {
if stage_rest {
selected.push(i);
}
i += 1;
continue;
}
render_hunk(&mut stderr, path, i, hunks.len(), &hunks[i]);
let _ = write!(stderr, "Stage this hunk [y,n,q,a,d,?]? ");
let _ = stderr.flush();
let mut line = String::new();
let read = input
.read_line(&mut line)
.map_err(|e| emit_err(&format!("read input: {e}"), exit::NOINPUT))?;
if read == 0 {
return Ok((selected, true));
}
match line.trim().chars().next() {
Some('y') => {
selected.push(i);
i += 1;
}
Some('n') => i += 1,
Some('q') => return Ok((selected, true)),
Some('a') => {
selected.push(i);
auto = Some(true);
i += 1;
}
Some('d') => auto = Some(false),
_ => {
let _ = writeln!(
stderr,
"y - stage this hunk\nn - skip this hunk\nq - quit; stage selected hunks\na - stage this and all later hunks in the file\nd - skip this and all later hunks in the file\n? - print help"
);
}
}
}
Ok((selected, false))
}
fn render_hunk(out: &mut impl Write, path: &str, idx: usize, total: usize, hunk: &PatchHunk) {
let _ = writeln!(out, "--- {path} (hunk {}/{total}) ---", idx + 1);
let _ = writeln!(
out,
"@@ -{} +{} @@",
range_str(hunk.old_start, hunk.old_len),
range_str(hunk.new_start, hunk.new_len)
);
for l in &hunk.lines {
let prefix = match l.kind {
HunkLineKind::Context => b' ',
HunkLineKind::Added => b'+',
HunkLineKind::Removed => b'-',
};
let mut buf = vec![prefix];
buf.extend_from_slice(&l.text);
buf.push(b'\n');
let _ = out.write_all(&buf);
if !l.has_newline {
let _ = writeln!(out, "\\ No newline at end of file");
}
}
}
fn range_str(start: usize, len: usize) -> String {
if len == 1 {
start.to_string()
} else {
format!("{start},{len}")
}
}
fn ensure_within_repo(root: &Path, abs: &Path) -> Result<(), String> {
use std::path::Component;
let parent = abs
.parent()
.ok_or_else(|| format!("invalid path: {}", abs.display()))?;
let real_parent = parent
.canonicalize()
.map_err(|e| format!("path {}: {e}", parent.display()))?;
let real_root = root.canonicalize().map_err(|e| format!("repo root: {e}"))?;
if real_parent != real_root && !real_parent.starts_with(&real_root) {
return Err(format!("path is outside repository: {}", abs.display()));
}
if let Ok(rel) = abs.strip_prefix(root) {
let comps: Vec<Component<'_>> = rel.components().collect();
let parent_count = comps.len().saturating_sub(1); let mut cur = root.to_path_buf();
for comp in &comps[..parent_count] {
if let Component::Normal(name) = comp {
cur.push(name);
if matches!(cur.symlink_metadata(), Ok(m) if m.file_type().is_symlink()) {
return Err(format!(
"path traverses a symbolic link ({}): refusing to stage beyond it",
cur.display()
));
}
}
}
}
Ok(())
}
use super::error as emit_err;