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