use keyhog_core::{Chunk, ChunkMetadata, Source, SourceError};
use std::path::{Path, PathBuf};
use std::process::Command;
pub struct GitDiffSource {
repo_path: PathBuf,
base_ref: String,
head_ref: Option<String>,
}
impl GitDiffSource {
pub fn new(repo_path: PathBuf, base_ref: impl Into<String>) -> Self {
Self {
repo_path,
base_ref: base_ref.into(),
head_ref: None,
}
}
pub fn with_head_ref(mut self, head_ref: impl Into<String>) -> Self {
self.head_ref = Some(head_ref.into());
self
}
}
impl Source for GitDiffSource {
fn name(&self) -> &str {
"git-diff"
}
fn chunks(&self) -> Box<dyn Iterator<Item = Result<Chunk, SourceError>> + '_> {
match stream_added_lines(&self.repo_path, &self.base_ref, self.head_ref.as_deref()) {
Ok(iter) => Box::new(iter),
Err(e) => Box::new(std::iter::once(Err(e))),
}
}
fn as_any(&self) -> &dyn std::any::Any {
self
}
}
fn stream_added_lines(
repo_path: &Path,
base_ref: &str,
head_ref: Option<&str>,
) -> Result<impl Iterator<Item = Result<Chunk, SourceError>>, SourceError> {
let base_ref = super::validate_ref_name(base_ref)?;
let head_ref = head_ref.map(super::validate_ref_name).transpose()?;
let repo_root = super::canonical_repo_root(repo_path)?;
let repo_arg = super::validate_repo_path(&repo_root)?;
super::verify_ref(&repo_arg, &base_ref)?;
let base_commit = super::get_commit_hash(&repo_arg, &base_ref)?;
let head_commit = if let Some(head_ref) = head_ref.as_deref() {
super::verify_ref(&repo_arg, head_ref)?;
Some(super::get_commit_hash(&repo_arg, head_ref)?)
} else {
None
};
let mut command = Command::new(super::git_bin()?);
command.args(["-C", &repo_arg, "diff", "-U0", "--end-of-options"]);
command.arg(&base_commit);
if let Some(head_commit) = head_commit.as_deref() {
command.arg(head_commit);
}
command.stdout(std::process::Stdio::piped());
command.stderr(std::process::Stdio::piped());
let mut child = command.spawn().map_err(SourceError::Io)?;
let stdout = child
.stdout
.take()
.ok_or_else(|| SourceError::Io(std::io::Error::other("missing stdout")))?;
let mut reader = std::io::BufReader::new(stdout);
let metadata_commit = head_commit.unwrap_or_else(|| base_commit.clone());
let author = super::get_commit_author(&repo_arg, &metadata_commit)?;
let date = super::get_commit_date(&repo_arg, &metadata_commit)?;
let mut current_path: Option<String> = None;
let mut current_content = String::new();
let mut in_hunk = false;
let mut done = false;
let mut line_buf: Vec<u8> = Vec::new();
let mut current_base_line: usize = 0;
Ok(std::iter::from_fn(move || {
if done {
return None;
}
loop {
let line = match super::read_capped_line(
&mut reader,
&mut line_buf,
super::MAX_GIT_LINE_BYTES,
) {
Ok(n) if n > 0 => {
let l = String::from_utf8_lossy(&line_buf);
l.trim_end_matches('\n').trim_end_matches('\r').to_string()
}
Err(e) => {
done = true;
return Some(Err(SourceError::Io(e)));
}
Ok(_) => {
done = true;
if let Some(ref path) = current_path {
if !current_content.trim().is_empty() {
return Some(Ok(Chunk {
data: current_content.trim().to_string().into(),
metadata: ChunkMetadata {
base_offset: 0,
base_line: current_base_line,
source_type: "git-diff".into(),
path: Some(path.clone()),
commit: Some(metadata_commit.clone()),
author: Some(author.clone()),
date: Some(date.clone()),
mtime_ns: None,
size_bytes: None,
},
}));
}
}
return None;
}
};
if line.starts_with("diff --git ") {
let prev_path = current_path.take();
let prev_content = std::mem::take(&mut current_content);
let prev_base_line = current_base_line;
in_hunk = false;
current_base_line = 0;
if let Some(path) = prev_path {
if !prev_content.trim().is_empty() {
return Some(Ok(Chunk {
data: prev_content.trim().to_string().into(),
metadata: ChunkMetadata {
base_offset: 0,
base_line: prev_base_line,
source_type: "git-diff".into(),
path: Some(path),
commit: Some(metadata_commit.clone()),
author: Some(author.clone()),
date: Some(date.clone()),
mtime_ns: None,
size_bytes: None,
},
}));
}
}
continue;
}
if line.starts_with("deleted file mode") {
current_path = None;
continue;
}
if line.starts_with("new file mode")
|| line.starts_with("index ")
|| line.starts_with("--- ")
{
continue;
}
if let Some(path_part) = line.strip_prefix("+++ b/") {
current_path = Some(path_part.trim().to_string());
continue;
}
if line.starts_with("@@") && line.contains("@@") {
let new_start = super::parse_hunk_new_start(&line).unwrap_or(1);
let prev_content = std::mem::take(&mut current_content);
let prev_base_line = current_base_line;
current_base_line = new_start.saturating_sub(1);
in_hunk = true;
if let Some(ref path) = current_path {
if !prev_content.trim().is_empty() {
return Some(Ok(Chunk {
data: prev_content.trim().to_string().into(),
metadata: ChunkMetadata {
base_offset: 0,
base_line: prev_base_line,
source_type: "git-diff".into(),
path: Some(path.clone()),
commit: Some(metadata_commit.clone()),
author: Some(author.clone()),
date: Some(date.clone()),
mtime_ns: None,
size_bytes: None,
},
}));
}
}
continue;
}
if in_hunk && line.starts_with('+') && !line.starts_with("+++") {
current_content.push_str(&line[1..]);
current_content.push('\n');
}
if current_content.len() > 10 * 1024 * 1024 {
if let Some(ref path) = current_path {
if !current_content.trim().is_empty() {
let flush_base_line = current_base_line;
current_base_line = current_base_line
.saturating_add(memchr::memchr_iter(b'\n', current_content.as_bytes()).count());
let chunk_content = current_content.trim().to_string();
current_content = String::new();
return Some(Ok(Chunk {
data: chunk_content.into(),
metadata: ChunkMetadata {
base_offset: 0,
base_line: flush_base_line,
source_type: "git-diff".into(),
path: Some(path.clone()),
commit: Some(metadata_commit.clone()),
author: Some(author.clone()),
date: Some(date.clone()),
mtime_ns: None,
size_bytes: None,
},
}));
}
}
}
}
}))
}