1use std::{
2 fs,
3 path::{Path, PathBuf},
4 sync::{
5 Arc, Mutex,
6 atomic::{AtomicU64, Ordering},
7 },
8};
9
10use anyhow::{Context, Result, anyhow};
11use tokio::{sync::watch, task};
12use tree_sitter::{Parser, Query, Tree};
13
14use crate::fetch::{HierarchyQuery, HierarchyResponse, WorkspaceSymbolMatch};
15
16mod index;
17
18use index::ProjectIndex;
19
20#[derive(Clone, Copy, Debug, Eq, PartialEq)]
21pub enum TreeSitterLanguage {
22 Rust,
23 C,
24 Cpp,
25 Python,
26}
27
28impl TreeSitterLanguage {
29 pub fn detect(workspace_root: &Path) -> Option<Self> {
30 if workspace_root.join("Cargo.toml").is_file()
31 || contains_source_with_extension(workspace_root, &["rs"])
32 {
33 return Some(Self::Rust);
34 }
35 if contains_source_with_extension(workspace_root, &["cc", "cpp", "cxx", "hpp"]) {
36 return Some(Self::Cpp);
37 }
38 if workspace_root.join("compile_commands.json").is_file()
39 || workspace_root.join("CMakeLists.txt").is_file()
40 || contains_source_with_extension(workspace_root, &["c", "h"])
41 {
42 return Some(Self::C);
43 }
44 if workspace_root.join("pyproject.toml").is_file()
45 || workspace_root.join("setup.py").is_file()
46 || workspace_root.join("requirements.txt").is_file()
47 || contains_source_with_extension(workspace_root, &["py"])
48 {
49 return Some(Self::Python);
50 }
51
52 None
53 }
54
55 pub fn name(self) -> &'static str {
56 match self {
57 Self::Rust => "Rust",
58 Self::C => "C",
59 Self::Cpp => "C++",
60 Self::Python => "Python",
61 }
62 }
63
64 pub(super) fn grammar(self) -> tree_sitter::Language {
65 match self {
66 Self::Rust => tree_sitter_rust::LANGUAGE.into(),
67 Self::C => tree_sitter_c::LANGUAGE.into(),
68 Self::Cpp => tree_sitter_cpp::LANGUAGE.into(),
69 Self::Python => tree_sitter_python::LANGUAGE.into(),
70 }
71 }
72
73 pub(super) fn tags_query(self) -> &'static str {
74 match self {
75 Self::Rust => tree_sitter_rust::TAGS_QUERY,
76 Self::C => tree_sitter_c::TAGS_QUERY,
77 Self::Cpp => tree_sitter_cpp::TAGS_QUERY,
78 Self::Python => tree_sitter_python::TAGS_QUERY,
79 }
80 }
81
82 pub(super) fn call_query(self) -> &'static str {
83 match self {
84 Self::Rust | Self::Python => self.tags_query(),
85 Self::C | Self::Cpp => {
86 r#"
87 (call_expression
88 function: (identifier) @name) @reference.call
89 (call_expression
90 function: (field_expression
91 field: (field_identifier) @name)) @reference.call
92 "#
93 }
94 }
95 }
96
97 pub(super) fn accepts_path(self, path: &Path) -> bool {
98 let extension = path.extension().and_then(|extension| extension.to_str());
99 match self {
100 Self::Rust => extension == Some("rs"),
101 Self::C => matches!(extension, Some("c" | "h")),
102 Self::Cpp => matches!(extension, Some("cc" | "cpp" | "cxx" | "h" | "hpp")),
103 Self::Python => extension == Some("py"),
104 }
105 }
106}
107
108pub struct TreeSitterProvider {
115 workspace_root: PathBuf,
116 language: TreeSitterLanguage,
117 parser: Parser,
118 symbol_query: Query,
119 shared: Arc<SharedIndex>,
120}
121
122struct SharedIndex {
123 workspace_root: PathBuf,
124 language: TreeSitterLanguage,
125 state: Mutex<IndexState>,
126 next_build_id: AtomicU64,
127 #[cfg(test)]
128 build_count: std::sync::atomic::AtomicUsize,
129 #[cfg(test)]
130 pause_build: std::sync::atomic::AtomicBool,
131}
132
133type IndexBuildResult = std::result::Result<Arc<ProjectIndex>, Arc<str>>;
134
135enum IndexState {
138 Empty,
139 Building {
140 build_id: u64,
141 receiver: watch::Receiver<Option<IndexBuildResult>>,
142 },
143 Ready(Arc<ProjectIndex>),
144}
145
146#[derive(Clone)]
147pub struct WorkspaceSymbolClient {
148 shared: Arc<SharedIndex>,
149}
150
151#[derive(Clone)]
152pub struct HierarchyClient {
153 shared: Arc<SharedIndex>,
154}
155
156impl std::fmt::Debug for TreeSitterProvider {
157 fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
158 formatter
159 .debug_struct("TreeSitterProvider")
160 .field("workspace_root", &self.workspace_root)
161 .field("language", &self.language)
162 .finish_non_exhaustive()
163 }
164}
165
166impl TreeSitterProvider {
167 pub fn start(workspace_root: impl Into<PathBuf>, language: TreeSitterLanguage) -> Result<Self> {
168 let workspace_root = workspace_root.into();
169 let grammar = language.grammar();
170 let mut parser = Parser::new();
171 parser
172 .set_language(&grammar)
173 .with_context(|| format!("failed to initialize {} grammar", language.name()))?;
174 let symbol_query = Query::new(&grammar, language.tags_query())
175 .with_context(|| format!("failed to initialize {} symbol query", language.name()))?;
176
177 Ok(Self {
178 shared: Arc::new(SharedIndex {
179 workspace_root: workspace_root.clone(),
180 language,
181 state: Mutex::new(IndexState::Empty),
182 next_build_id: AtomicU64::new(1),
183 #[cfg(test)]
184 build_count: std::sync::atomic::AtomicUsize::new(0),
185 #[cfg(test)]
186 pause_build: std::sync::atomic::AtomicBool::new(false),
187 }),
188 workspace_root,
189 language,
190 parser,
191 symbol_query,
192 })
193 }
194
195 pub fn workspace_root(&self) -> &Path {
196 &self.workspace_root
197 }
198
199 pub fn language(&self) -> TreeSitterLanguage {
200 self.language
201 }
202
203 pub fn parse(&mut self, source: &str) -> Result<Tree> {
204 self.parser
205 .parse(source, None)
206 .with_context(|| format!("{} parser returned no syntax tree", self.language.name()))
207 }
208
209 pub fn symbol_capture_names(&self) -> &[&str] {
210 self.symbol_query.capture_names()
211 }
212
213 pub fn workspace_symbol_client(&self) -> WorkspaceSymbolClient {
214 WorkspaceSymbolClient {
215 shared: Arc::clone(&self.shared),
216 }
217 }
218
219 pub fn hierarchy_client(&self) -> HierarchyClient {
220 HierarchyClient {
221 shared: Arc::clone(&self.shared),
222 }
223 }
224}
225
226impl std::fmt::Debug for WorkspaceSymbolClient {
227 fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
228 formatter
229 .debug_struct("TreeSitterWorkspaceSymbolClient")
230 .finish_non_exhaustive()
231 }
232}
233
234impl WorkspaceSymbolClient {
235 pub async fn query(&self, _query: &str) -> Result<Vec<WorkspaceSymbolMatch>> {
236 Ok(load_index(&self.shared).await?.workspace_symbols())
237 }
238}
239
240impl std::fmt::Debug for HierarchyClient {
241 fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
242 formatter
243 .debug_struct("TreeSitterHierarchyClient")
244 .finish_non_exhaustive()
245 }
246}
247
248impl HierarchyClient {
249 pub async fn query(&self, query: HierarchyQuery) -> Result<HierarchyResponse> {
250 load_index(&self.shared).await?.hierarchy(query)
251 }
252}
253
254async fn load_index(shared: &Arc<SharedIndex>) -> Result<Arc<ProjectIndex>> {
255 let (receiver, build) = {
256 let mut state = shared
257 .state
258 .lock()
259 .expect("Tree-sitter index state mutex poisoned");
260 match &*state {
261 IndexState::Ready(index) => return Ok(Arc::clone(index)),
262 IndexState::Building { receiver, .. } => (receiver.clone(), None),
263 IndexState::Empty => {
264 let build_id = shared.next_build_id.fetch_add(1, Ordering::Relaxed);
265 let (sender, receiver) = watch::channel(None);
266 *state = IndexState::Building {
267 build_id,
268 receiver: receiver.clone(),
269 };
270 (receiver, Some((build_id, sender)))
271 }
272 }
273 };
274
275 if let Some((build_id, sender)) = build {
276 spawn_index_build(Arc::clone(shared), build_id, sender);
277 }
278
279 wait_for_index(receiver).await
280}
281
282fn spawn_index_build(
283 shared: Arc<SharedIndex>,
284 build_id: u64,
285 sender: watch::Sender<Option<IndexBuildResult>>,
286) {
287 let _build_task = tokio::spawn(async move {
288 let workspace_root = shared.workspace_root.clone();
289 let language = shared.language;
290 #[cfg(test)]
291 let build_shared = Arc::clone(&shared);
292 let result = task::spawn_blocking(move || {
293 #[cfg(test)]
294 {
295 use std::{thread, time::Duration};
296
297 build_shared.build_count.fetch_add(1, Ordering::SeqCst);
298 while build_shared.pause_build.load(Ordering::SeqCst) {
299 thread::sleep(Duration::from_millis(1));
300 }
301 }
302 ProjectIndex::build(&workspace_root, language)
303 })
304 .await
305 .context("Tree-sitter indexing task failed")
306 .and_then(|result| result)
307 .map(Arc::new)
308 .map_err(|error| Arc::<str>::from(format!("{error:#}")));
309
310 {
311 let mut state = shared
312 .state
313 .lock()
314 .expect("Tree-sitter index state mutex poisoned");
315 if matches!(
316 &*state,
317 IndexState::Building {
318 build_id: active_build_id,
319 ..
320 } if *active_build_id == build_id
321 ) {
322 *state = match &result {
323 Ok(index) => IndexState::Ready(Arc::clone(index)),
324 Err(_) => IndexState::Empty,
325 };
326 }
327 }
328
329 sender.send_replace(Some(result));
330 });
331}
332
333async fn wait_for_index(
334 mut receiver: watch::Receiver<Option<IndexBuildResult>>,
335) -> Result<Arc<ProjectIndex>> {
336 loop {
337 if let Some(result) = receiver.borrow().clone() {
338 return result.map_err(|error| anyhow!(error.to_string()));
339 }
340 receiver
341 .changed()
342 .await
343 .context("Tree-sitter indexing task ended without a result")?;
344 }
345}
346
347fn contains_source_with_extension(workspace_root: &Path, extensions: &[&str]) -> bool {
348 fs::read_dir(workspace_root).is_ok_and(|entries| {
349 entries.filter_map(Result::ok).any(|entry| {
350 entry
351 .path()
352 .extension()
353 .and_then(|extension| extension.to_str())
354 .is_some_and(|extension| extensions.contains(&extension))
355 })
356 })
357}
358
359#[cfg(test)]
360mod tests {
361 use std::{
362 fs,
363 sync::atomic::Ordering,
364 time::{SystemTime, UNIX_EPOCH},
365 };
366
367 use super::{TreeSitterLanguage, TreeSitterProvider};
368 use crate::{
369 fetch::{FetchSource, HierarchyQuery},
370 state::{HierarchyDirection, HierarchyKind, SourceLocation, SymbolIdentity},
371 };
372
373 #[test]
374 fn detects_supported_workspace_languages() {
375 let workspace = temporary_workspace("detect");
376 fs::write(workspace.join("main.cpp"), "int main() {}\n").unwrap();
377
378 assert_eq!(
379 TreeSitterLanguage::detect(&workspace),
380 Some(TreeSitterLanguage::Cpp)
381 );
382
383 fs::remove_dir_all(workspace).unwrap();
384 }
385
386 #[test]
387 fn initializes_and_parses_each_supported_grammar() {
388 let cases = [
389 (TreeSitterLanguage::Rust, "fn main() {}"),
390 (TreeSitterLanguage::C, "int main(void) { return 0; }"),
391 (TreeSitterLanguage::Cpp, "int main() { return 0; }"),
392 (TreeSitterLanguage::Python, "def main():\n pass\n"),
393 ];
394
395 for (language, source) in cases {
396 let mut provider = TreeSitterProvider::start(".", language).unwrap();
397 let tree = provider.parse(source).unwrap();
398 assert!(
399 !tree.root_node().has_error(),
400 "{} parse failed",
401 language.name()
402 );
403 assert!(!provider.symbol_capture_names().is_empty());
404 }
405 }
406
407 #[test]
408 fn normalizes_tree_sitter_byte_columns_to_utf16() {
409 assert_eq!(super::index::utf16_column("é😀name", 6, 6).unwrap(), 3);
410 }
411
412 #[tokio::test]
413 async fn indexes_rust_symbols_and_bidirectional_static_calls_once() {
414 let workspace = temporary_workspace("rust-index");
415 fs::write(
416 workspace.join("lib.rs"),
417 r#"
418 struct Worker;
419 trait Job {
420 fn execute(&self) {}
421 }
422 impl Worker {
423 fn run(&self) {
424 helper();
425 self.finish();
426 }
427 fn finish(&self) {}
428 }
429 fn helper() {}
430 "#,
431 )
432 .unwrap();
433 fs::create_dir(workspace.join("target")).unwrap();
434 fs::write(workspace.join("target/ignored.rs"), "fn ignored() {}\n").unwrap();
435 let provider = TreeSitterProvider::start(&workspace, TreeSitterLanguage::Rust).unwrap();
436 let symbol_client = provider.workspace_symbol_client();
437 let hierarchy_client = provider.hierarchy_client();
438
439 let symbols = symbol_client.query("").await.unwrap();
440 let names = symbols
441 .iter()
442 .map(|symbol| symbol.name.as_str())
443 .collect::<Vec<_>>();
444 assert!(names.contains(&"Worker::run"));
445 assert!(names.contains(&"Worker::finish"));
446 assert!(names.contains(&"Job::execute"));
447 assert!(names.contains(&"helper"));
448 assert!(!names.contains(&"ignored"));
449
450 let run = identity(&symbols, "Worker::run", HierarchyKind::Call);
451 let outgoing = hierarchy_client
452 .query(HierarchyQuery {
453 symbol: run.clone(),
454 direction: HierarchyDirection::Outgoing,
455 })
456 .await
457 .unwrap();
458 assert_eq!(outgoing.source, FetchSource::TreeSitter);
459 assert_eq!(
460 outgoing
461 .children
462 .iter()
463 .map(|child| child.symbol.as_str())
464 .collect::<Vec<_>>(),
465 ["helper", "Worker::finish"]
466 );
467
468 let helper = identity(&symbols, "helper", HierarchyKind::Call);
469 let incoming = hierarchy_client
470 .query(HierarchyQuery {
471 symbol: helper,
472 direction: HierarchyDirection::Incoming,
473 })
474 .await
475 .unwrap();
476 assert_eq!(incoming.children, [run]);
477 fs::remove_dir_all(workspace).unwrap();
478 }
479
480 #[tokio::test]
481 async fn indexes_python_methods_calls_and_type_inheritance() {
482 let workspace = temporary_workspace("python-index");
483 fs::write(
484 workspace.join("main.py"),
485 "class Base:\n pass\n\nclass Child(Base):\n def run(self):\n helper()\n\ndef helper():\n pass\n",
486 )
487 .unwrap();
488 let provider = TreeSitterProvider::start(&workspace, TreeSitterLanguage::Python).unwrap();
489 let symbols = provider.workspace_symbol_client().query("").await.unwrap();
490 let hierarchy = provider.hierarchy_client();
491
492 let run = identity(&symbols, "Child.run", HierarchyKind::Call);
493 let outgoing = hierarchy
494 .query(HierarchyQuery {
495 symbol: run,
496 direction: HierarchyDirection::Outgoing,
497 })
498 .await
499 .unwrap();
500 assert_eq!(
501 outgoing
502 .children
503 .iter()
504 .map(|child| child.symbol.as_str())
505 .collect::<Vec<_>>(),
506 ["helper"]
507 );
508
509 let base = identity(&symbols, "Base", HierarchyKind::Type);
510 let child = identity(&symbols, "Child", HierarchyKind::Type);
511 let subtypes = hierarchy
512 .query(HierarchyQuery {
513 symbol: base.clone(),
514 direction: HierarchyDirection::Outgoing,
515 })
516 .await
517 .unwrap();
518 assert_eq!(subtypes.children.as_slice(), std::slice::from_ref(&child));
519 let supertypes = hierarchy
520 .query(HierarchyQuery {
521 symbol: child,
522 direction: HierarchyDirection::Incoming,
523 })
524 .await
525 .unwrap();
526 assert_eq!(supertypes.children, [base]);
527 fs::remove_dir_all(workspace).unwrap();
528 }
529
530 #[tokio::test]
531 async fn python_self_calls_prefer_the_method_on_the_current_class() {
532 let workspace = temporary_workspace("python-self-call");
533 fs::write(
534 workspace.join("main.py"),
535 "class First:\n def finish(self):\n pass\n\nclass Second:\n def run(self):\n self.finish()\n\n def finish(self):\n pass\n",
536 )
537 .unwrap();
538 let provider = TreeSitterProvider::start(&workspace, TreeSitterLanguage::Python).unwrap();
539 let symbols = provider.workspace_symbol_client().query("").await.unwrap();
540 let run = identity(&symbols, "Second.run", HierarchyKind::Call);
541
542 let outgoing = provider
543 .hierarchy_client()
544 .query(HierarchyQuery {
545 symbol: run,
546 direction: HierarchyDirection::Outgoing,
547 })
548 .await
549 .unwrap();
550
551 assert_eq!(
552 outgoing
553 .children
554 .iter()
555 .map(|child| child.symbol.as_str())
556 .collect::<Vec<_>>(),
557 ["Second.finish"]
558 );
559 fs::remove_dir_all(workspace).unwrap();
560 }
561
562 #[tokio::test]
563 async fn indexes_c_and_cpp_calls_plus_cpp_inheritance() {
564 let c_workspace = temporary_workspace("c-index");
565 fs::write(
566 c_workspace.join("main.c"),
567 "void helper(void);\nvoid helper(void) {}\nvoid run(void) { helper(); }\n",
568 )
569 .unwrap();
570 let c_provider = TreeSitterProvider::start(&c_workspace, TreeSitterLanguage::C).unwrap();
571 assert_call_edge(&c_provider, "run", "helper").await;
572
573 let cpp_workspace = temporary_workspace("cpp-index");
574 fs::write(
575 cpp_workspace.join("main.cpp"),
576 "class Base {};\nclass Child : public Base { public: void run(); };\nvoid helper() {}\nvoid Child::run() { helper(); }\n",
577 )
578 .unwrap();
579 let cpp_provider =
580 TreeSitterProvider::start(&cpp_workspace, TreeSitterLanguage::Cpp).unwrap();
581 assert_call_edge(&cpp_provider, "Child::run", "helper").await;
582 let symbols = cpp_provider
583 .workspace_symbol_client()
584 .query("")
585 .await
586 .unwrap();
587 let base = identity(&symbols, "Base", HierarchyKind::Type);
588 let child = identity(&symbols, "Child", HierarchyKind::Type);
589 let response = cpp_provider
590 .hierarchy_client()
591 .query(HierarchyQuery {
592 symbol: base,
593 direction: HierarchyDirection::Outgoing,
594 })
595 .await
596 .unwrap();
597 assert_eq!(response.children, [child]);
598
599 fs::remove_dir_all(c_workspace).unwrap();
600 fs::remove_dir_all(cpp_workspace).unwrap();
601 }
602
603 #[tokio::test]
604 async fn leaves_ambiguous_targets_unbound_and_requires_an_exact_root() {
605 let workspace = temporary_workspace("ambiguous-index");
606 fs::write(workspace.join("first.rs"), "pub fn helper() {}\n").unwrap();
607 fs::write(workspace.join("second.rs"), "pub fn helper() {}\n").unwrap();
608 fs::write(workspace.join("main.rs"), "fn run() { helper(); }\n").unwrap();
609 let provider = TreeSitterProvider::start(&workspace, TreeSitterLanguage::Rust).unwrap();
610 let symbols = provider.workspace_symbol_client().query("").await.unwrap();
611 let run = identity(&symbols, "run", HierarchyKind::Call);
612
613 let response = provider
614 .hierarchy_client()
615 .query(HierarchyQuery {
616 symbol: run,
617 direction: HierarchyDirection::Outgoing,
618 })
619 .await
620 .unwrap();
621 assert!(response.children.is_empty());
622
623 let error = provider
624 .hierarchy_client()
625 .query(HierarchyQuery {
626 symbol: SymbolIdentity {
627 symbol: "helper".to_owned(),
628 kind: HierarchyKind::Call,
629 location: None,
630 },
631 direction: HierarchyDirection::Incoming,
632 })
633 .await
634 .unwrap_err();
635 assert!(error.to_string().contains("ambiguous"));
636 fs::remove_dir_all(workspace).unwrap();
637 }
638
639 #[tokio::test]
640 async fn cancelling_the_first_query_does_not_restart_project_indexing() {
641 let workspace = temporary_workspace("cancelled-index-query");
642 fs::write(workspace.join("lib.rs"), "fn main() {}\n").unwrap();
643 let provider = TreeSitterProvider::start(&workspace, TreeSitterLanguage::Rust).unwrap();
644 provider.shared.pause_build.store(true, Ordering::SeqCst);
645 let first_client = provider.workspace_symbol_client();
646 let first_query = tokio::spawn(async move { first_client.query("").await });
647
648 let mut build_started = false;
649 for _ in 0..10_000 {
650 if provider.shared.build_count.load(Ordering::SeqCst) == 1 {
651 build_started = true;
652 break;
653 }
654 tokio::task::yield_now().await;
655 }
656 if !build_started {
657 provider.shared.pause_build.store(false, Ordering::SeqCst);
658 panic!("background index build did not start");
659 }
660
661 first_query.abort();
662 assert!(first_query.await.unwrap_err().is_cancelled());
663 provider.shared.pause_build.store(false, Ordering::SeqCst);
664 let symbols = provider.workspace_symbol_client().query("").await.unwrap();
665
666 assert!(symbols.iter().any(|symbol| symbol.name == "main"));
667 assert_eq!(provider.shared.build_count.load(Ordering::SeqCst), 1);
668 fs::remove_dir_all(workspace).unwrap();
669 }
670
671 async fn assert_call_edge(provider: &TreeSitterProvider, caller: &str, callee: &str) {
672 let symbols = provider.workspace_symbol_client().query("").await.unwrap();
673 let caller = identity(&symbols, caller, HierarchyKind::Call);
674 let response = provider
675 .hierarchy_client()
676 .query(HierarchyQuery {
677 symbol: caller,
678 direction: HierarchyDirection::Outgoing,
679 })
680 .await
681 .unwrap();
682 assert_eq!(
683 response
684 .children
685 .iter()
686 .map(|child| child.symbol.as_str())
687 .collect::<Vec<_>>(),
688 [callee]
689 );
690 }
691
692 fn identity(
693 symbols: &[crate::fetch::WorkspaceSymbolMatch],
694 name: &str,
695 kind: HierarchyKind,
696 ) -> SymbolIdentity {
697 let symbol = symbols
698 .iter()
699 .find(|symbol| symbol.name == name)
700 .unwrap_or_else(|| panic!("missing indexed symbol {name:?}"));
701 let position = symbol.range.unwrap().start;
702 SymbolIdentity {
703 symbol: symbol.name.clone(),
704 kind,
705 location: Some(SourceLocation {
706 uri: symbol.uri.to_string(),
707 line: Some(position.line),
708 character: Some(position.character),
709 }),
710 }
711 }
712
713 fn temporary_workspace(name: &str) -> std::path::PathBuf {
714 let unique = SystemTime::now()
715 .duration_since(UNIX_EPOCH)
716 .unwrap()
717 .as_nanos();
718 let path = std::env::temp_dir().join(format!("cgraph-{name}-{unique}"));
719 fs::create_dir(&path).unwrap();
720 path
721 }
722}