use std::path::{Path, PathBuf};
use anyhow::{Context, Result};
use rayon::prelude::*;
use rustc_hash::FxHashSet;
use tracing::info;
use crate::candidate_files::collect_candidate_files;
use crate::config::limits::LIMITS;
use crate::edges;
use crate::fragmentation::process_files_for_fragments;
use crate::git::CatFileBatch;
use crate::graph::{self, Graph};
use crate::tokenizer::count_tokens;
use crate::types::{Fragment, FragmentId};
pub struct ProjectGraph {
pub fragments: Vec<Fragment>,
pub graph: Graph,
pub root_dir: PathBuf,
}
impl ProjectGraph {
pub fn node_count(&self) -> usize {
self.graph.node_count()
}
pub fn edge_count(&self) -> usize {
self.graph.edge_count()
}
}
pub struct ProjectGraphOptions {
pub use_git_batch_reader: bool,
pub skip_expensive_edges: Option<bool>,
}
impl Default for ProjectGraphOptions {
fn default() -> Self {
Self {
use_git_batch_reader: false,
skip_expensive_edges: None,
}
}
}
pub fn build_project_graph(root_dir: &Path) -> Result<ProjectGraph> {
build_project_graph_with_options(root_dir, &ProjectGraphOptions::default())
}
pub fn build_project_graph_with_options(
root_dir: &Path,
options: &ProjectGraphOptions,
) -> Result<ProjectGraph> {
let resolved_root = root_dir
.canonicalize()
.with_context(|| format!("failed to canonicalize root_dir '{}'", root_dir.display()))?;
let included_set: FxHashSet<PathBuf> = FxHashSet::default();
let candidate_files = collect_candidate_files(&resolved_root, &included_set);
info!(
"project_graph: found {} candidate files",
candidate_files.len()
);
let mut seen_frag_ids: FxHashSet<FragmentId> = FxHashSet::default();
let mut all_fragments = if options.use_git_batch_reader {
let mut batch_reader = CatFileBatch::new(&resolved_root)?;
let frags = process_files_for_fragments(
&candidate_files,
&resolved_root,
&[],
&mut seen_frag_ids,
Some(&mut batch_reader),
false,
);
batch_reader.close();
frags
} else {
process_files_for_fragments(
&candidate_files,
&resolved_root,
&[],
&mut seen_frag_ids,
None,
false,
)
};
assign_token_counts(&mut all_fragments);
info!(
"project_graph: {} fragments from {} files",
all_fragments.len(),
candidate_files.len()
);
let skip_expensive = options
.skip_expensive_edges
.unwrap_or_else(|| all_fragments.len() > LIMITS.skip_expensive_threshold);
let capped = edges::collect_capped_edges(
&all_fragments,
Some(resolved_root.as_path()),
skip_expensive,
);
let graph = graph::build_graph_capped(&all_fragments, capped);
Ok(ProjectGraph {
fragments: all_fragments,
graph,
root_dir: resolved_root,
})
}
fn assign_token_counts(fragments: &mut [Fragment]) {
fragments.par_iter_mut().for_each(|frag| {
if frag.token_count == 0 {
frag.token_count = count_tokens(&frag.content) + LIMITS.overhead_per_fragment;
}
});
}
#[cfg(test)]
mod tests {
use super::*;
use std::fs;
use tempfile::TempDir;
fn git(dir: &Path, args: &[&str]) {
let status = crate::git::git_command(dir)
.args(args)
.status()
.unwrap_or_else(|e| panic!("git {args:?}: {e}"));
assert!(status.success(), "git {args:?} failed");
}
fn init_git_repo(dir: &Path) {
git(dir, &["init", "-q", "-b", "main"]);
git(dir, &["config", "user.email", "test@example.com"]);
git(dir, &["config", "user.name", "Test"]);
git(dir, &["config", "commit.gpgsign", "false"]);
}
fn commit_all(dir: &Path) {
git(dir, &["add", "-A"]);
git(dir, &["commit", "-q", "-m", "initial"]);
}
fn write_file(root: &Path, rel: &str, content: &str) {
let path = root.join(rel);
if let Some(parent) = path.parent() {
fs::create_dir_all(parent).expect("create parent");
}
fs::write(&path, content).expect("write file");
}
#[test]
fn build_project_graph_on_tiny_python_project() {
let tmp = TempDir::new().expect("tempdir");
let root = tmp.path();
init_git_repo(root);
write_file(
root,
"alpha.py",
"def alpha():\n return beta()\n\ndef beta():\n return 1\n",
);
write_file(
root,
"consumer.py",
"from alpha import alpha\n\ndef main():\n return alpha()\n",
);
commit_all(root);
let pg = build_project_graph(root).expect("build_project_graph");
assert!(
pg.node_count() >= 3,
"expected fragments, got {}",
pg.node_count()
);
assert_eq!(pg.fragments.len(), pg.node_count());
assert!(pg.root_dir.is_absolute());
}
#[test]
fn build_project_graph_empty_dir_yields_no_fragments() {
let tmp = TempDir::new().expect("tempdir");
let root = tmp.path();
init_git_repo(root);
write_file(root, ".gitkeep", "");
commit_all(root);
let pg = build_project_graph(root).expect("build_project_graph");
assert_eq!(pg.fragments.len(), 0);
assert_eq!(pg.node_count(), 0);
assert_eq!(pg.edge_count(), 0);
}
#[test]
fn build_project_graph_produces_edge_categories() {
let tmp = TempDir::new().expect("tempdir");
let root = tmp.path();
init_git_repo(root);
write_file(
root,
"lib.py",
"def helper(x):\n return x + 1\n\ndef other(y):\n return helper(y) * 2\n",
);
write_file(
root,
"app.py",
"from lib import helper, other\n\ndef run():\n return helper(other(3))\n",
);
commit_all(root);
let pg = build_project_graph(root).expect("build_project_graph");
assert!(pg.fragments.len() >= 2);
assert_eq!(pg.graph.categorized_edge_count(), pg.graph.edge_count());
}
}