text_splitter/splitter/
code.rs

1use std::{cmp::Ordering, ops::Range};
2
3use thiserror::Error;
4use tree_sitter::{Language, LanguageError, Parser, TreeCursor, MIN_COMPATIBLE_LANGUAGE_VERSION};
5
6use crate::{
7    splitter::{SemanticLevel, Splitter},
8    trim::Trim,
9    ChunkConfig, ChunkSizer, TryChunkSizer,
10};
11
12use super::ChunkCharIndex;
13
14/// Indicates there was an error with creating a `CodeSplitter`.
15/// The `Display` implementation will provide a human-readable error message to
16/// help debug the issue that caused the error.
17#[derive(Error, Debug)]
18#[error(transparent)]
19pub struct CodeSplitterError(#[from] CodeSplitterErrorRepr);
20
21/// Private error and free to change across minor version of the crate.
22#[derive(Error, Debug)]
23enum CodeSplitterErrorRepr {
24    #[error(
25        "Language version {0:?} is too old. Expected at least version {min_version}",
26        min_version=MIN_COMPATIBLE_LANGUAGE_VERSION,
27    )]
28    LanguageError(LanguageError),
29}
30
31/// Source code splitter. Recursively splits chunks into the largest
32/// semantic units that fit within the chunk size. Also will attempt to merge
33/// neighboring chunks if they can fit within the given chunk size.
34#[derive(Debug)]
35pub struct CodeSplitter<Sizer>
36where
37    Sizer: TryChunkSizer,
38{
39    /// Method of determining chunk sizes.
40    chunk_config: ChunkConfig<Sizer>,
41    /// Language to use for parsing the code.
42    language: Language,
43}
44
45impl<Sizer> CodeSplitter<Sizer>
46where
47    Sizer: TryChunkSizer,
48{
49    /// Creates a new [`CodeSplitter`].
50    ///
51    /// ```
52    /// use text_splitter::CodeSplitter;
53    ///
54    /// // By default, the chunk sizer is based on characters.
55    /// let splitter = CodeSplitter::new(tree_sitter_rust::LANGUAGE, 512).expect("Invalid language");
56    /// ```
57    ///
58    /// # Errors
59    ///
60    /// Will return an error if the language version is too old to be compatible
61    /// with the current version of the tree-sitter crate.
62    pub fn new(
63        language: impl Into<Language>,
64        chunk_config: impl Into<ChunkConfig<Sizer>>,
65    ) -> Result<Self, CodeSplitterError> {
66        // Verify that this is a valid language so we can rely on that later.
67        let mut parser = Parser::new();
68        let language = language.into();
69        parser
70            .set_language(&language)
71            .map_err(CodeSplitterErrorRepr::LanguageError)?;
72        Ok(Self {
73            chunk_config: chunk_config.into(),
74            language,
75        })
76    }
77}
78
79impl<Sizer: ChunkSizer> CodeSplitter<Sizer> {
80    /// Generate a list of chunks from a given text. Each chunk will be up to the `chunk_capacity`.
81    ///
82    /// ## Method
83    ///
84    /// To preserve as much semantic meaning within a chunk as possible, each chunk is composed of the largest semantic units that can fit in the next given chunk. For each splitter type, there is a defined set of semantic levels. Here is an example of the steps used:
85    ///
86    /// 1. Split the text by a increasing semantic levels.
87    /// 2. Check the first item for each level and select the highest level whose first item still fits within the chunk size.
88    /// 3. Merge as many of these neighboring sections of this level or above into a chunk to maximize chunk length.
89    ///    Boundaries of higher semantic levels are always included when merging, so that the chunk doesn't inadvertantly cross semantic boundaries.
90    ///
91    /// The boundaries used to split the text if using the `chunks` method, in ascending order:
92    ///
93    /// 1. Characters
94    /// 2. [Unicode Grapheme Cluster Boundaries](https://www.unicode.org/reports/tr29/#Grapheme_Cluster_Boundaries)
95    /// 3. [Unicode Word Boundaries](https://www.unicode.org/reports/tr29/#Word_Boundaries)
96    /// 4. [Unicode Sentence Boundaries](https://www.unicode.org/reports/tr29/#Sentence_Boundaries)
97    /// 5. Ascending depth of the syntax tree. So function would have a higher level than a statement inside of the function, and so on.
98    ///
99    /// Splitting doesn't occur below the character level, otherwise you could get partial bytes of a char, which may not be a valid unicode str.
100    ///
101    /// ```
102    /// use text_splitter::CodeSplitter;
103    ///
104    /// let splitter = CodeSplitter::new(tree_sitter_rust::LANGUAGE, 10).expect("Invalid language");
105    /// let text = "Some text\n\nfrom a\ndocument";
106    /// let chunks = splitter.chunks(text).collect::<Vec<_>>();
107    ///
108    /// assert_eq!(vec!["Some text", "from a", "document"], chunks);
109    /// ```
110    pub fn chunks<'splitter, 'text: 'splitter>(
111        &'splitter self,
112        text: &'text str,
113    ) -> impl Iterator<Item = &'text str> + 'splitter {
114        Splitter::<_>::try_chunks(self, text).map(super::unwrap_infallible)
115    }
116
117    /// Returns an iterator over chunks of the text and their byte offsets.
118    /// Each chunk will be up to the `chunk_capacity`.
119    ///
120    /// See [`CodeSplitter::chunks`] for more information.
121    ///
122    /// ```
123    /// use text_splitter::{ChunkCharIndex, CodeSplitter};
124    ///
125    /// let splitter = CodeSplitter::new(tree_sitter_rust::LANGUAGE, 10).expect("Invalid language");
126    /// let text = "Some text\n\nfrom a\ndocument";
127    /// let chunks = splitter.chunk_indices(text).collect::<Vec<_>>();
128    ///
129    /// assert_eq!(vec![(0, "Some text"), (11, "from a"), (18, "document")], chunks);
130    /// ```
131    pub fn chunk_indices<'splitter, 'text: 'splitter>(
132        &'splitter self,
133        text: &'text str,
134    ) -> impl Iterator<Item = (usize, &'text str)> + 'splitter {
135        Splitter::<_>::try_chunk_indices(self, text).map(super::unwrap_infallible)
136    }
137
138    /// Returns an iterator over chunks of the text with their byte and character offsets.
139    /// Each chunk will be up to the `chunk_capacity`.
140    ///
141    /// See [`CodeSplitter::chunks`] for more information.
142    ///
143    /// This will be more expensive than just byte offsets, and for most usage in Rust, just
144    /// having byte offsets is sufficient. But when interfacing with other languages or systems
145    /// that require character offsets, this will track the character offsets for you,
146    /// accounting for any trimming that may have occurred.
147    ///
148    /// ```
149    /// use text_splitter::{ChunkCharIndex, CodeSplitter};
150    ///
151    /// let splitter = CodeSplitter::new(tree_sitter_rust::LANGUAGE, 10).expect("Invalid language");
152    /// let text = "Some text\n\nfrom a\ndocument";
153    /// let chunks = splitter.chunk_char_indices(text).collect::<Vec<_>>();
154    ///
155    /// assert_eq!(vec![ChunkCharIndex {chunk: "Some text", byte_offset: 0, char_offset: 0}, ChunkCharIndex {chunk: "from a", byte_offset: 11, char_offset: 11}, ChunkCharIndex {chunk: "document", byte_offset: 18, char_offset: 18}], chunks);
156    /// ```
157    pub fn chunk_char_indices<'splitter, 'text: 'splitter>(
158        &'splitter self,
159        text: &'text str,
160    ) -> impl Iterator<Item = ChunkCharIndex<'text>> + 'splitter {
161        Splitter::<_>::try_chunk_char_indices(self, text).map(super::unwrap_infallible)
162    }
163}
164
165impl<Sizer: TryChunkSizer> CodeSplitter<Sizer> {
166    /// Fallible version of [`Self::chunks`] for a [`TryChunkSizer`].
167    ///
168    /// Yields the first sizing error, then ends without calling the sizer again.
169    /// Collect into `Result<Vec<_>, _>` to stop at that error and discard partial
170    /// output. See [`TryChunkSizer`] for an example.
171    pub fn try_chunks<'splitter, 'text: 'splitter>(
172        &'splitter self,
173        text: &'text str,
174    ) -> impl Iterator<Item = Result<&'text str, Sizer::Error>> + 'splitter {
175        Splitter::<_>::try_chunks(self, text)
176    }
177
178    /// Fallible version of [`Self::chunk_indices`], including byte offsets.
179    ///
180    /// Yields the first sizing error, then ends without calling the sizer again.
181    pub fn try_chunk_indices<'splitter, 'text: 'splitter>(
182        &'splitter self,
183        text: &'text str,
184    ) -> impl Iterator<Item = Result<(usize, &'text str), Sizer::Error>> + 'splitter {
185        Splitter::<_>::try_chunk_indices(self, text)
186    }
187
188    /// Fallible version of [`Self::chunk_char_indices`], including byte and
189    /// character offsets.
190    ///
191    /// Yields the first sizing error, then ends without calling the sizer again.
192    pub fn try_chunk_char_indices<'splitter, 'text: 'splitter>(
193        &'splitter self,
194        text: &'text str,
195    ) -> impl Iterator<Item = Result<ChunkCharIndex<'text>, Sizer::Error>> + 'splitter {
196        Splitter::<_>::try_chunk_char_indices(self, text)
197    }
198}
199
200impl<Sizer> Splitter<Sizer> for CodeSplitter<Sizer>
201where
202    Sizer: TryChunkSizer,
203{
204    type Level = Depth;
205
206    const TRIM: Trim = Trim::PreserveIndentation;
207
208    fn chunk_config(&self) -> &ChunkConfig<Sizer> {
209        &self.chunk_config
210    }
211
212    fn parse(&self, text: &str) -> Vec<(Self::Level, Range<usize>)> {
213        let mut parser = Parser::new();
214        parser
215            .set_language(&self.language)
216            // We verify at initialization that the language is valid, so this should be safe.
217            .expect("Error loading language");
218        // The only reason the tree would be None is:
219        // - No language was set (we do that)
220        // - There was a timeout or cancellation option set (we don't)
221        // - So it should be safe to unwrap here
222        let tree = parser.parse(text, None).expect("Error parsing source code");
223
224        CursorOffsets::new(tree.walk()).collect()
225    }
226}
227
228/// New type around a usize to capture the depth of a given code node.
229/// Custom type so that we can implement custom ordering, since we want to
230/// sort items of lower depth as higher priority.
231#[derive(Clone, Copy, Debug, Eq, PartialEq)]
232pub struct Depth(usize);
233
234impl PartialOrd for Depth {
235    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
236        Some(self.cmp(other))
237    }
238}
239
240impl Ord for Depth {
241    fn cmp(&self, other: &Self) -> Ordering {
242        other.0.cmp(&self.0)
243    }
244}
245
246/// New type around a tree-sitter cursor to allow for implementing an iterator.
247/// Each call to `next()` will return the next node in the tree in a depth-first
248/// order.
249struct CursorOffsets<'cursor> {
250    cursor: TreeCursor<'cursor>,
251}
252
253impl<'cursor> CursorOffsets<'cursor> {
254    fn new(cursor: TreeCursor<'cursor>) -> Self {
255        Self { cursor }
256    }
257}
258
259impl Iterator for CursorOffsets<'_> {
260    type Item = (Depth, Range<usize>);
261
262    fn next(&mut self) -> Option<Self::Item> {
263        // There are children (can call this initially because we don't want the root node)
264        if self.cursor.goto_first_child() {
265            return Some((
266                Depth(self.cursor.depth() as usize),
267                self.cursor.node().byte_range(),
268            ));
269        }
270
271        loop {
272            // There are sibling elements to grab
273            if self.cursor.goto_next_sibling() {
274                return Some((
275                    Depth(self.cursor.depth() as usize),
276                    self.cursor.node().byte_range(),
277                ));
278            // Start going back up the tree and check for next sibling on next iteration.
279            } else if self.cursor.goto_parent() {
280                continue;
281            }
282
283            // We have no more siblings or parents, so we're done.
284            return None;
285        }
286    }
287}
288
289impl SemanticLevel for Depth {}
290
291#[cfg(test)]
292mod tests {
293    use tree_sitter::{Node, Tree};
294
295    use super::*;
296
297    #[test]
298    fn rust_splitter() {
299        let splitter = CodeSplitter::new(tree_sitter_rust::LANGUAGE, 16).unwrap();
300        let text = "fn main() {\n    let x = 5;\n}";
301        let chunks = splitter.chunks(text).collect::<Vec<_>>();
302
303        assert_eq!(chunks, vec!["fn main()", "{\n    let x = 5;", "}"]);
304    }
305
306    #[test]
307    fn rust_splitter_indices() {
308        let splitter = CodeSplitter::new(tree_sitter_rust::LANGUAGE, 16).unwrap();
309        let text = "fn main() {\n    let x = 5;\n}";
310        let chunks = splitter.chunk_indices(text).collect::<Vec<_>>();
311
312        assert_eq!(
313            chunks,
314            vec![(0, "fn main()"), (10, "{\n    let x = 5;"), (27, "}")]
315        );
316    }
317
318    #[test]
319    fn rust_splitter_char_indices() {
320        let splitter = CodeSplitter::new(tree_sitter_rust::LANGUAGE, 16).unwrap();
321        let text = "fn main() {\n    let x = 5;\n}";
322        let chunks = splitter.chunk_char_indices(text).collect::<Vec<_>>();
323
324        assert_eq!(
325            chunks,
326            vec![
327                ChunkCharIndex {
328                    chunk: "fn main()",
329                    byte_offset: 0,
330                    char_offset: 0
331                },
332                ChunkCharIndex {
333                    chunk: "{\n    let x = 5;",
334                    byte_offset: 10,
335                    char_offset: 10
336                },
337                ChunkCharIndex {
338                    chunk: "}",
339                    byte_offset: 27,
340                    char_offset: 27
341                }
342            ]
343        );
344    }
345
346    #[test]
347    fn depth_partialord() {
348        assert_eq!(Depth(0).partial_cmp(&Depth(1)), Some(Ordering::Greater));
349        assert_eq!(Depth(1).partial_cmp(&Depth(2)), Some(Ordering::Greater));
350        assert_eq!(Depth(1).partial_cmp(&Depth(1)), Some(Ordering::Equal));
351        assert_eq!(Depth(2).partial_cmp(&Depth(1)), Some(Ordering::Less));
352    }
353
354    #[test]
355    fn depth_ord() {
356        assert_eq!(Depth(0).cmp(&Depth(1)), Ordering::Greater);
357        assert_eq!(Depth(1).cmp(&Depth(2)), Ordering::Greater);
358        assert_eq!(Depth(1).cmp(&Depth(1)), Ordering::Equal);
359        assert_eq!(Depth(2).cmp(&Depth(1)), Ordering::Less);
360    }
361
362    #[test]
363    fn depth_sorting() {
364        let mut depths = vec![Depth(0), Depth(1), Depth(2)];
365        depths.sort();
366        assert_eq!(depths, [Depth(2), Depth(1), Depth(0)]);
367    }
368
369    /// Checks that the optimized version of the code produces the same results as the naive version.
370    #[test]
371    fn optimized_code_offsets() {
372        let mut parser = Parser::new();
373        parser
374            .set_language(&tree_sitter_rust::LANGUAGE.into())
375            .expect("Error loading Rust grammar");
376        let source_code = "fn test() {
377    let x = 1;
378}";
379        let tree = parser
380            .parse(source_code, None)
381            .expect("Error parsing source code");
382
383        let offsets = CursorOffsets::new(tree.walk()).collect::<Vec<_>>();
384
385        assert_eq!(offsets, naive_offsets(&tree));
386    }
387
388    #[test]
389    fn multiple_top_siblings() {
390        let mut parser = Parser::new();
391        parser
392            .set_language(&tree_sitter_rust::LANGUAGE.into())
393            .expect("Error loading Rust grammar");
394        let source_code = "
395fn fn1() {}
396fn fn2() {}
397fn fn3() {}
398fn fn4() {}";
399        let tree = parser
400            .parse(source_code, None)
401            .expect("Error parsing source code");
402
403        let offsets = CursorOffsets::new(tree.walk()).collect::<Vec<_>>();
404
405        assert_eq!(offsets, naive_offsets(&tree));
406    }
407
408    fn naive_offsets(tree: &Tree) -> Vec<(Depth, Range<usize>)> {
409        let root_node = tree.root_node();
410        let mut offsets = vec![];
411        recursive_naive_offsets(&mut offsets, root_node, 0);
412        offsets
413    }
414
415    // Basic version to compare an optimized version against. According to the tree-sitter
416    // documentation, this is not efficient for large trees. But because it is the easiest
417    // to reason about it is a good check for correctness.
418    fn recursive_naive_offsets(
419        collection: &mut Vec<(Depth, Range<usize>)>,
420        node: Node<'_>,
421        depth: usize,
422    ) {
423        // We can skip the root node
424        if depth > 0 {
425            collection.push((Depth(depth), node.byte_range()));
426        }
427
428        for child in node.children(&mut node.walk()) {
429            recursive_naive_offsets(collection, child, depth + 1);
430        }
431    }
432}
text_splitter/splitter/code.rs

text_splitter/splitter/
code.rs
