diff --git a/crates/code2prompt-core/src/configuration.rs b/crates/code2prompt-core/src/configuration.rs index 16c2c04..9790251 100644 --- a/crates/code2prompt-core/src/configuration.rs +++ b/crates/code2prompt-core/src/configuration.rs @@ -7,7 +7,7 @@ use crate::tokenizer::TokenizerType; use crate::{sort::FileSortMethod, tokenizer::TokenFormat}; use derive_builder::Builder; use serde::{Deserialize, Serialize}; -use std::collections::HashMap; +use std::collections::{HashMap, HashSet}; use std::path::PathBuf; /// Configuration object defining preferences and filters for code prompt generation. @@ -97,6 +97,13 @@ pub struct Code2PromptConfig { /// If set, contains two branch names for which code2prompt will generate a git diff. pub diff_branches: Option<(String, String)>, + /// Populated internally from `diff_branches` before traversal: the set of + /// relative file paths that changed between the two branches. When + /// present, `discover_files` prunes both the source tree and the + /// collected file content down to just these paths, mirroring how + /// `include_patterns` filters both. Not user-configurable directly. + pub diff_files: Option>, + /// If set, contains two branch names for which code2prompt will retrieve the git log. pub log_branches: Option<(String, String)>, diff --git a/crates/code2prompt-core/src/git.rs b/crates/code2prompt-core/src/git.rs index d359b9b..2f6971b 100644 --- a/crates/code2prompt-core/src/git.rs +++ b/crates/code2prompt-core/src/git.rs @@ -8,7 +8,8 @@ use anyhow::{Context, Result}; use git2::{DiffOptions, Repository}; use log::info; -use std::path::Path; +use std::collections::HashSet; +use std::path::{Path, PathBuf}; /// Extract git diff for staged changes in repository. /// @@ -135,6 +136,68 @@ pub fn get_git_diff_between_branches( Ok(String::from_utf8_lossy(&diff_text).into_owned()) } +/// Computes the set of relative file paths that changed between two branches. +/// +/// This performs the same tree-to-tree diff as [`get_git_diff_between_branches`], +/// but instead of rendering a textual patch it collects the paths touched by +/// each delta (both sides, so a rename contributes its old and new path). +/// This set is used to prune the source tree and file content down to only +/// the files that actually changed when `--git-diff-branch` is active. +/// +/// # Arguments +/// +/// * `repo_path` - A reference to the path of the git repository +/// * `branch1` - The name of the first branch +/// * `branch2` - The name of the second branch +/// +/// # Returns +/// +/// * `Result>` - The relative paths that changed between the two branches +pub fn get_git_diff_file_paths( + repo_path: &Path, + branch1: &str, + branch2: &str, +) -> Result> { + info!("Opening repository at path: {:?}", repo_path); + let repo = Repository::open(repo_path).context("Failed to open repository")?; + + for branch in [branch1, branch2].iter() { + if !branch_exists(&repo, branch) { + return Err(anyhow::anyhow!("Branch {} doesn't exist!", branch)); + } + } + + let branch1_commit = repo.revparse_single(branch1)?.peel_to_commit()?; + let branch2_commit = repo.revparse_single(branch2)?.peel_to_commit()?; + + let branch1_tree = branch1_commit.tree()?; + let branch2_tree = branch2_commit.tree()?; + + let diff = repo + .diff_tree_to_tree( + Some(&branch1_tree), + Some(&branch2_tree), + Some(DiffOptions::new().ignore_whitespace(true)), + ) + .context("Failed to generate diff between branches")?; + + let mut paths = HashSet::new(); + for delta in diff.deltas() { + if let Some(path) = delta.old_file().path() { + paths.insert(path.to_path_buf()); + } + if let Some(path) = delta.new_file().path() { + paths.insert(path.to_path_buf()); + } + } + + info!( + "Collected {} changed file path(s) between branches", + paths.len() + ); + Ok(paths) +} + /// Retrieves the git log between two branches for the repository at the provided path /// /// # Arguments diff --git a/crates/code2prompt-core/src/path.rs b/crates/code2prompt-core/src/path.rs index 06e765d..9d3af3f 100644 --- a/crates/code2prompt-core/src/path.rs +++ b/crates/code2prompt-core/src/path.rs @@ -120,12 +120,20 @@ fn discover_files( let path = entry.path(); if let Ok(relative_path) = path.strip_prefix(&canonical_root_path) { // Use SelectionEngine if available, otherwise fall back to pattern matching - let entry_match = if let Some(engine) = selection_engine.as_mut() { + let mut entry_match = if let Some(engine) = selection_engine.as_mut() { engine.is_selected(relative_path) } else { should_include_file(relative_path, &include_globset, &exclude_globset) }; + // When `--git-diff-branch` is active, further restrict matches to + // only the files that changed between the two branches. This + // mirrors how include/exclude patterns prune both the tree and + // the file content list. + if let Some(diff_files) = &config.diff_files { + entry_match = entry_match && diff_files.contains(relative_path); + } + // Directory Tree let include_in_tree = config.full_directory_tree || entry_match; diff --git a/crates/code2prompt-core/src/session.rs b/crates/code2prompt-core/src/session.rs index 32b7d11..a880e42 100644 --- a/crates/code2prompt-core/src/session.rs +++ b/crates/code2prompt-core/src/session.rs @@ -10,7 +10,9 @@ use std::sync::Arc; use crate::analysis::CodebaseAnalysis; use crate::configuration::Code2PromptConfig; use crate::entity_map::FileCodeMap; -use crate::git::{get_git_diff, get_git_diff_between_branches, get_git_log}; +use crate::git::{ + get_git_diff, get_git_diff_between_branches, get_git_diff_file_paths, get_git_log, +}; use crate::path::{FileEntry, display_name, traverse_directory, wrap_code_block}; use crate::selection::SelectionEngine; use crate::template::{OutputFormat, handlebars_setup, render_template}; @@ -257,6 +259,22 @@ impl Code2PromptSession { /// Loads the codebase data (source tree and file list) into the session. pub fn load_codebase(&mut self) -> Result<()> { + // When comparing two branches, restrict both the source tree and the + // file content to only the files that actually changed, mirroring + // how `--include` filters both the tree and content. If the diff + // can't be computed (e.g. a branch doesn't exist), fall back to the + // unfiltered codebase; `load_git_diff_between_branches` will surface + // the underlying error to the user. + if let Some((branch1, branch2)) = self.config.diff_branches.clone() { + match get_git_diff_file_paths(&self.config.path, &branch1, &branch2) { + Ok(diff_files) => self.config.diff_files = Some(diff_files), + Err(e) => log::warn!( + "Could not compute changed files between '{branch1}' and '{branch2}', \ + showing unfiltered codebase: {e}" + ), + } + } + let (tree, files) = traverse_directory(&self.config, Some(&mut self.selection_engine)) .with_context(|| "Failed to traverse directory")?; diff --git a/crates/code2prompt-core/tests/session_integration_test.rs b/crates/code2prompt-core/tests/session_integration_test.rs index d5ea98e..376e88e 100644 --- a/crates/code2prompt-core/tests/session_integration_test.rs +++ b/crates/code2prompt-core/tests/session_integration_test.rs @@ -3,6 +3,7 @@ use code2prompt_core::configuration::Code2PromptConfig; use code2prompt_core::session::Code2PromptSession; use std::fs; +use std::path::Path; use tempfile::TempDir; #[cfg(test)] @@ -178,4 +179,102 @@ mod tests { assert_eq!(selected_files.len(), 1); assert_eq!(selected_files[0], main_rs_relative); } + + /// Regression test for https://github.com/mufeedvh/code2prompt/issues/176 + /// + /// With `--git-diff-branch`, the source tree (and file content) must be + /// pruned down to only the files that actually changed between the two + /// branches, the same way `--include` filters both the tree and content. + #[test] + fn test_git_diff_branch_prunes_source_tree_to_changed_files() { + use git2::{Repository, RepositoryInitOptions, Signature}; + + let temp_dir = TempDir::new().unwrap(); + let repo_path = temp_dir.path(); + + let mut binding = RepositoryInitOptions::new(); + let init_options = binding.initial_head("master"); + let repo = Repository::init_opts(repo_path, init_options) + .expect("Failed to initialize repository"); + + fs::create_dir_all(repo_path.join("src")).unwrap(); + fs::write(repo_path.join("src/changed.rs"), "fn changed() {}").unwrap(); + fs::write(repo_path.join("src/unchanged.rs"), "fn unchanged() {}").unwrap(); + + let signature = Signature::now("Test", "test@example.com").unwrap(); + + // Commit both files on master + let mut index = repo.index().unwrap(); + index.add_path(Path::new("src/changed.rs")).unwrap(); + index.add_path(Path::new("src/unchanged.rs")).unwrap(); + index.write().unwrap(); + let tree_id = index.write_tree().unwrap(); + let tree = repo.find_tree(tree_id).unwrap(); + let master_commit = repo + .commit( + Some("HEAD"), + &signature, + &signature, + "Initial commit", + &tree, + &[], + ) + .expect("Failed to commit"); + + // Create a feature branch and only modify one of the two files there + repo.branch("feature", &repo.find_commit(master_commit).unwrap(), false) + .expect("Failed to create new branch"); + repo.set_head("refs/heads/feature").unwrap(); + repo.checkout_head(None).unwrap(); + + fs::write( + repo_path.join("src/changed.rs"), + "fn changed() { /* updated */ }", + ) + .unwrap(); + + let mut index = repo.index().unwrap(); + index.add_path(Path::new("src/changed.rs")).unwrap(); + index.write().unwrap(); + let tree_id = index.write_tree().unwrap(); + let tree = repo.find_tree(tree_id).unwrap(); + repo.commit( + Some("HEAD"), + &signature, + &signature, + "Update changed.rs", + &tree, + &[&repo.find_commit(master_commit).unwrap()], + ) + .expect("Failed to commit on feature branch"); + + let config = Code2PromptConfig::builder() + .path(repo_path.to_path_buf()) + .diff_branches(Some(("master".to_string(), "feature".to_string()))) + .build() + .unwrap(); + + let mut session = Code2PromptSession::new(config); + session.load_codebase().expect("Failed to load codebase"); + + let source_tree = session.data.source_tree.clone().unwrap_or_default(); + assert!( + source_tree.contains("changed.rs"), + "expected diffed file in source tree, got:\n{source_tree}" + ); + assert!( + !source_tree.contains("unchanged.rs"), + "expected unchanged file to be pruned from source tree, got:\n{source_tree}" + ); + + let files = session.data.files.clone().unwrap_or_default(); + assert!( + files.iter().any(|f| f.path.contains("changed.rs")), + "expected diffed file in file content list" + ); + assert!( + !files.iter().any(|f| f.path.contains("unchanged.rs")), + "expected unchanged file to be excluded from file content list" + ); + } }