swiftide_integrations/treesitter/
outline_code_tree_sitter.rs

1//! Add the outline of the code in the given file to the metadata of a node, using tree-sitter.
2use anyhow::Result;
3use async_trait::async_trait;
4
5use swiftide_core::Transformer;
6use swiftide_core::indexing::TextNode;
7
8use crate::treesitter::{CodeOutliner, SupportedLanguages};
9
10/// `OutlineCodeTreeSitter` adds a "Outline" field to the metadata of a node that contains
11/// a summary of the code in the node. It uses the tree-sitter parser to parse the code and
12/// remove any information that is less relevant for tasks that consider the file as a whole.
13#[swiftide_macros::indexing_transformer(metadata_field_name = "Outline", derive(skip_default))]
14pub struct OutlineCodeTreeSitter {
15    outliner: CodeOutliner,
16    minimum_file_size: Option<usize>,
17}
18
19impl OutlineCodeTreeSitter {
20    /// Tries to create a `OutlineCodeTreeSitter` instance for a given programming language.
21    ///
22    /// # Parameters
23    /// - `lang`: The programming language to be used to parse the code. It should implement
24    ///   `TryInto<SupportedLanguages>`.
25    ///
26    /// # Returns
27    /// - `Result<Self>`: Returns an instance of `OutlineCodeTreeSitter` if successful, otherwise
28    ///   returns an error.
29    ///
30    /// # Errors
31    /// - Returns an error if the language is not supported or if the `CodeOutliner` fails to build.
32    pub fn try_for_language(
33        lang: impl TryInto<SupportedLanguages>,
34        minimum_file_size: Option<usize>,
35    ) -> Result<Self> {
36        Ok(Self {
37            outliner: CodeOutliner::builder().try_language(lang)?.build()?,
38            minimum_file_size,
39            client: None,
40            concurrency: None,
41            indexing_defaults: None,
42        })
43    }
44}
45
46#[async_trait]
47impl Transformer for OutlineCodeTreeSitter {
48    type Input = String;
49    type Output = String;
50    /// Adds context to the metadata of a `TextNode` containing code in the "Outline" field.
51    ///
52    /// It uses the `CodeOutliner` to generate the context.
53    ///
54    /// # Parameters
55    /// - `node`: The `TextNode` containing the code of which the context is to be generated.
56    ///
57    /// # Returns
58    /// - `TextNode`: The same `TextNode` instances, with the metadata updated to include the
59    ///   generated context.
60    ///
61    /// # Errors
62    /// - If the code outlining fails, an error is sent downstream.
63    #[tracing::instrument(skip_all, name = "transformers.outline_code_tree_sitter")]
64    async fn transform_node(&self, mut node: TextNode) -> Result<TextNode> {
65        if let Some(minimum_file_size) = self.minimum_file_size
66            && node.chunk.len() < minimum_file_size
67        {
68            return Ok(node);
69        }
70
71        let outline_result = self.outliner.outline(&node.chunk)?;
72        node.metadata.insert(NAME, outline_result);
73        Ok(node)
74    }
75}