swiftide_integrations/treesitter/
metadata_refs_defs_code.rs

1//! Adds references and definitions found in code as metadata to chunks
2//!
3//! Uses tree-sitter to do the extractions. It tries to only get unique definitions and references,
4//! and only references that are not local.
5//!
6//! See the [`crate::treesitter::CodeParser`] tests for some examples.
7//!
8//! # Example
9//!
10//! ```no_run
11//! # use swiftide_core::indexing::TextNode;
12//! # use swiftide_integrations::treesitter::transformers::metadata_refs_defs_code::*;
13//! # use swiftide_core::Transformer;
14//! # #[tokio::main]
15//! # async fn main() -> Result<(), Box<dyn std::error::Error>> {
16//! let transformer = MetadataRefsDefsCode::try_from_language("rust").unwrap();
17//! let code = r#"
18//!   fn main() {
19//!     println!("Hello, World!");
20//!   }
21//! "#;
22//! let mut node = TextNode::new(code.to_string());
23//!
24//! node = transformer.transform_node(node).await.unwrap();
25//!
26//! assert_eq!(
27//!     node.metadata.get(NAME_REFERENCES).unwrap().as_str().unwrap(),
28//!     "println"
29//! );
30//! assert_eq!(
31//!     node.metadata.get(NAME_DEFINITIONS).unwrap().as_str().unwrap(),
32//!     "main"
33//! );
34//! # Ok(())
35//! # }
36//! ```
37use std::sync::Arc;
38
39use swiftide_core::{Transformer, indexing::TextNode};
40
41use crate::treesitter::{CodeParser, SupportedLanguages};
42use anyhow::{Context as _, Result};
43use async_trait::async_trait;
44
45pub const NAME_REFERENCES: &str = "References (code)";
46pub const NAME_DEFINITIONS: &str = "Definitions (code)";
47
48/// `MetadataRefsDefsCode` is responsible for extracting references and definitions.
49#[swiftide_macros::indexing_transformer(derive(skip_default))]
50pub struct MetadataRefsDefsCode {
51    code_parser: Arc<CodeParser>,
52}
53
54impl MetadataRefsDefsCode {
55    /// Tries to build a new `MetadataRefsDefsCode` transformer
56    ///
57    /// # Errors
58    ///
59    /// Language is not supported by tree-sitter
60    pub fn try_from_language(language: impl TryInto<SupportedLanguages>) -> Result<Self> {
61        let language: SupportedLanguages = language
62            .try_into()
63            .ok()
64            .context("Treesitter language not supported")?;
65
66        MetadataRefsDefsCode::builder()
67            .code_parser(CodeParser::from_language(language))
68            .build()
69    }
70}
71
72#[async_trait]
73impl Transformer for MetadataRefsDefsCode {
74    type Input = String;
75    type Output = String;
76    /// Extracts references and definitions from code and
77    /// adds them as metadata to the node if present
78    async fn transform_node(&self, mut node: TextNode) -> Result<TextNode> {
79        let refs_defs = self
80            .code_parser
81            .parse(&node.chunk)?
82            .references_and_definitions()?;
83
84        if !refs_defs.references.is_empty() {
85            node.metadata
86                .insert(NAME_REFERENCES.to_string(), refs_defs.references.join(","));
87        }
88
89        if !refs_defs.definitions.is_empty() {
90            node.metadata.insert(
91                NAME_DEFINITIONS.to_string(),
92                refs_defs.definitions.join(","),
93            );
94        }
95        Ok(node)
96    }
97}
98
99#[cfg(test)]
100mod test {
101
102    use super::*;
103    use test_case::test_case;
104
105    #[test_case("rust", "fn main() { println!(\"Hello, World!\"); }", "println", "main"; "rust")]
106    #[test_case("ruby", "def main; puts 'Hello, World!'; end", "puts", "main"; "ruby")]
107    #[test_case("python", "def main(): print('Hello, World!')", "print", "main"; "python")]
108    #[test_case("javascript", "function main() { console.log('Hello, World!'); }", "log", "main"; "javascript")]
109    #[test_case("typescript", "function main() { console.log('Hello, World!'); }", "log", "main"; "typescript")]
110    #[test_case("java", "public class Main { public static void main(String[] args) { System.out.println(\"Hello, World!\"); } }", "println", "Main,main"; "java")]
111    #[tokio::test]
112    async fn assert_refs_defs_from_code(
113        lang: &str,
114        code: &str,
115        expected_references: &str,
116        expected_definitions: &str,
117    ) {
118        let transformer = MetadataRefsDefsCode::try_from_language(lang).unwrap();
119        let node = TextNode::new(code);
120
121        let node = transformer.transform_node(node).await.unwrap();
122
123        let references = node
124            .metadata
125            .get(NAME_REFERENCES)
126            .unwrap()
127            .as_str()
128            .unwrap()
129            .to_string();
130        let definitions = node
131            .metadata
132            .get(NAME_DEFINITIONS)
133            .unwrap()
134            .as_str()
135            .unwrap()
136            .to_string();
137
138        assert_eq!(references, expected_references);
139        assert_eq!(definitions, expected_definitions);
140    }
141}
swiftide_integrations/treesitter/metadata_refs_defs_code.rs

swiftide_integrations/treesitter/
metadata_refs_defs_code.rs
