1use std::path::{Path, PathBuf};
2
3use anyhow::{Context, Result};
4use rayon::prelude::*;
5use rustc_hash::FxHashSet;
6use tracing::info;
7
8use crate::candidate_files::collect_candidate_files;
9use crate::config::limits::LIMITS;
10use crate::edges;
11use crate::fragmentation::process_files_for_fragments;
12use crate::git::CatFileBatch;
13use crate::graph::{self, Graph};
14use crate::tokenizer::count_tokens;
15use crate::types::{Fragment, FragmentId};
16
17pub struct ProjectGraph {
18 pub fragments: Vec<Fragment>,
19 pub graph: Graph,
20 pub root_dir: PathBuf,
21}
22
23impl ProjectGraph {
24 pub fn node_count(&self) -> usize {
25 self.graph.node_count()
26 }
27
28 pub fn edge_count(&self) -> usize {
29 self.graph.edge_count()
30 }
31}
32
33pub struct ProjectGraphOptions {
34 pub use_git_batch_reader: bool,
35 pub skip_expensive_edges: Option<bool>,
36}
37
38impl Default for ProjectGraphOptions {
39 fn default() -> Self {
40 Self {
41 use_git_batch_reader: false,
42 skip_expensive_edges: None,
43 }
44 }
45}
46
47pub fn build_project_graph(root_dir: &Path) -> Result<ProjectGraph> {
48 build_project_graph_with_options(root_dir, &ProjectGraphOptions::default())
49}
50
51pub fn build_project_graph_with_options(
52 root_dir: &Path,
53 options: &ProjectGraphOptions,
54) -> Result<ProjectGraph> {
55 let resolved_root = root_dir
56 .canonicalize()
57 .with_context(|| format!("failed to canonicalize root_dir '{}'", root_dir.display()))?;
58
59 let included_set: FxHashSet<PathBuf> = FxHashSet::default();
60 let candidate_files = collect_candidate_files(&resolved_root, &included_set);
61
62 info!(
63 "project_graph: found {} candidate files",
64 candidate_files.len()
65 );
66
67 let mut seen_frag_ids: FxHashSet<FragmentId> = FxHashSet::default();
68 let mut all_fragments = if options.use_git_batch_reader {
69 let mut batch_reader = CatFileBatch::new(&resolved_root)?;
70 let frags = process_files_for_fragments(
71 &candidate_files,
72 &resolved_root,
73 &[],
74 &mut seen_frag_ids,
75 Some(&mut batch_reader),
76 false,
77 );
78 batch_reader.close();
79 frags
80 } else {
81 process_files_for_fragments(
82 &candidate_files,
83 &resolved_root,
84 &[],
85 &mut seen_frag_ids,
86 None,
87 false,
88 )
89 };
90
91 assign_token_counts(&mut all_fragments);
92
93 info!(
94 "project_graph: {} fragments from {} files",
95 all_fragments.len(),
96 candidate_files.len()
97 );
98
99 let skip_expensive = options
100 .skip_expensive_edges
101 .unwrap_or_else(|| all_fragments.len() > LIMITS.skip_expensive_threshold);
102
103 let capped = edges::collect_capped_edges(
104 &all_fragments,
105 Some(resolved_root.as_path()),
106 skip_expensive,
107 crate::deadline::Deadline::none(),
108 );
109
110 let graph = graph::build_graph_capped(&all_fragments, capped);
111
112 Ok(ProjectGraph {
113 fragments: all_fragments,
114 graph,
115 root_dir: resolved_root,
116 })
117}
118
119fn assign_token_counts(fragments: &mut [Fragment]) {
120 fragments.par_iter_mut().for_each(|frag| {
121 if frag.token_count == 0 {
122 frag.token_count = count_tokens(&frag.content) + LIMITS.overhead_per_fragment;
123 }
124 });
125}
126
127#[cfg(test)]
128mod tests {
129 use super::*;
130 use std::fs;
131 use tempfile::TempDir;
132
133 fn git(dir: &Path, args: &[&str]) {
134 let status = crate::git::git_command(dir)
135 .args(args)
136 .status()
137 .unwrap_or_else(|e| panic!("git {args:?}: {e}"));
138 assert!(status.success(), "git {args:?} failed");
139 }
140
141 fn init_git_repo(dir: &Path) {
142 git(dir, &["init", "-q", "-b", "main"]);
143 git(dir, &["config", "user.email", "test@example.com"]);
144 git(dir, &["config", "user.name", "Test"]);
145 git(dir, &["config", "commit.gpgsign", "false"]);
146 }
147
148 fn commit_all(dir: &Path) {
149 git(dir, &["add", "-A"]);
150 git(dir, &["commit", "-q", "-m", "initial"]);
151 }
152
153 fn write_file(root: &Path, rel: &str, content: &str) {
154 let path = root.join(rel);
155 if let Some(parent) = path.parent() {
156 fs::create_dir_all(parent).expect("create parent");
157 }
158 fs::write(&path, content).expect("write file");
159 }
160
161 #[test]
162 fn build_project_graph_on_tiny_python_project() {
163 let tmp = TempDir::new().expect("tempdir");
164 let root = tmp.path();
165 init_git_repo(root);
166 write_file(
167 root,
168 "alpha.py",
169 "def alpha():\n return beta()\n\ndef beta():\n return 1\n",
170 );
171 write_file(
172 root,
173 "consumer.py",
174 "from alpha import alpha\n\ndef main():\n return alpha()\n",
175 );
176 commit_all(root);
177
178 let pg = build_project_graph(root).expect("build_project_graph");
179 assert!(
180 pg.node_count() >= 3,
181 "expected fragments, got {}",
182 pg.node_count()
183 );
184 assert_eq!(pg.fragments.len(), pg.node_count());
185 assert!(pg.root_dir.is_absolute());
186 }
187
188 #[test]
189 fn build_project_graph_empty_dir_yields_no_fragments() {
190 let tmp = TempDir::new().expect("tempdir");
191 let root = tmp.path();
192 init_git_repo(root);
193 write_file(root, ".gitkeep", "");
194 commit_all(root);
195
196 let pg = build_project_graph(root).expect("build_project_graph");
197 assert_eq!(pg.fragments.len(), 0);
198 assert_eq!(pg.node_count(), 0);
199 assert_eq!(pg.edge_count(), 0);
200 }
201
202 #[test]
203 fn build_project_graph_produces_edge_categories() {
204 let tmp = TempDir::new().expect("tempdir");
205 let root = tmp.path();
206 init_git_repo(root);
207 write_file(
208 root,
209 "lib.py",
210 "def helper(x):\n return x + 1\n\ndef other(y):\n return helper(y) * 2\n",
211 );
212 write_file(
213 root,
214 "app.py",
215 "from lib import helper, other\n\ndef run():\n return helper(other(3))\n",
216 );
217 commit_all(root);
218
219 let pg = build_project_graph(root).expect("build_project_graph");
220 assert!(pg.fragments.len() >= 2);
221 assert_eq!(pg.graph.categorized_edge_count(), pg.graph.edge_count());
222 }
223}