Skip to main content

text_splitter/
splitter.rs

1use std::{cmp::Ordering, convert::Infallible, fmt, iter::once, ops::Range};
2
3use either::Either;
4use itertools::Itertools;
5use strum::IntoEnumIterator;
6
7use self::fallback::FallbackLevel;
8use crate::{
9    chunk_size::MemoizedChunkSizer, trim::Trim, ChunkCapacity, ChunkConfig, TryChunkSizer,
10};
11
12#[cfg(feature = "code")]
13mod code;
14mod fallback;
15#[cfg(feature = "markdown")]
16mod markdown;
17mod text;
18
19#[cfg(feature = "code")]
20pub use code::{CodeSplitter, CodeSplitterError};
21#[cfg(feature = "markdown")]
22pub use markdown::MarkdownSplitter;
23pub use text::TextSplitter;
24
25fn unwrap_infallible<T>(result: Result<T, Infallible>) -> T {
26    match result {
27        Ok(value) => value,
28        Err(never) => match never {},
29    }
30}
31
32/// Shared interface for splitters that can generate chunks of text based on the
33/// associated semantic level.
34trait Splitter<Sizer>
35where
36    Sizer: TryChunkSizer,
37{
38    type Level: SemanticLevel;
39
40    /// Trimming behavior to use when trimming chunks
41    const TRIM: Trim = Trim::All;
42
43    /// Retrieve the splitter chunk configuration
44    fn chunk_config(&self) -> &ChunkConfig<Sizer>;
45
46    /// Generate a list of offsets for each semantic level within the text.
47    fn parse(&self, text: &str) -> Vec<(Self::Level, Range<usize>)>;
48
49    /// Returns an iterator over chunks of the text and their byte offsets.
50    /// Each chunk will be up to the max size of the `ChunkConfig`.
51    fn try_chunk_indices<'splitter, 'text: 'splitter>(
52        &'splitter self,
53        text: &'text str,
54    ) -> impl Iterator<Item = Result<(usize, &'text str), Sizer::Error>> + 'splitter
55    where
56        Sizer: 'splitter,
57    {
58        TextChunks::<Sizer, Self::Level>::new(
59            self.chunk_config(),
60            text,
61            self.parse(text),
62            Self::TRIM,
63        )
64    }
65
66    /// Returns an iterator over chunks of the text and their byte and character offsets.
67    /// Each chunk will be up to the max size of the `ChunkConfig`.
68    ///
69    /// This will be more expensive than just byte offsets, and for most usage in Rust, just having byte offsets is sufficient.
70    /// But when interfacing with other languages or systems that require character offsets, this will track the character offsets
71    /// for you, accounting for any trimming that may have occurred.
72    fn try_chunk_char_indices<'splitter, 'text: 'splitter>(
73        &'splitter self,
74        text: &'text str,
75    ) -> impl Iterator<Item = Result<ChunkCharIndex<'text>, Sizer::Error>> + 'splitter
76    where
77        Sizer: 'splitter,
78    {
79        TextChunksWithCharIndices::<Sizer, Self::Level>::new(
80            self.chunk_config(),
81            text,
82            self.parse(text),
83            Self::TRIM,
84        )
85    }
86
87    /// Generate a list of chunks from a given text.
88    /// Each chunk will be up to the max size of the `ChunkConfig`.
89    fn try_chunks<'splitter, 'text: 'splitter>(
90        &'splitter self,
91        text: &'text str,
92    ) -> impl Iterator<Item = Result<&'text str, Sizer::Error>> + 'splitter
93    where
94        Sizer: 'splitter,
95    {
96        self.try_chunk_indices(text)
97            .map(|chunk| chunk.map(|(_, t)| t))
98    }
99}
100
101/// Custom-defined levels of semantic splitting for custom document types.
102trait SemanticLevel: Copy + fmt::Debug + Ord + PartialOrd + 'static {
103    /// Given a level, split the text into sections based on the level.
104    /// Level ranges are also provided of items that are equal to or greater than the current level.
105    /// Default implementation assumes that all level ranges should be treated
106    /// as their own item.
107    fn sections(
108        text: &str,
109        level_ranges: impl Iterator<Item = (Self, Range<usize>)>,
110    ) -> impl Iterator<Item = (usize, &str)> {
111        let mut cursor = 0;
112        let mut final_match = false;
113        level_ranges
114            .batching(move |it| {
115                loop {
116                    match it.next() {
117                        // If we've hit the end, actually return None
118                        None if final_match => return None,
119                        // First time we hit None, return the final section of the text
120                        None => {
121                            final_match = true;
122                            return text.get(cursor..).map(|t| Either::Left(once((cursor, t))));
123                        }
124                        // Return text preceding match + the match
125                        Some((_, range)) => {
126                            if range.start < cursor {
127                                continue;
128                            }
129                            let offset = cursor;
130                            let prev_section = text
131                                .get(offset..range.start)
132                                .expect("invalid character sequence");
133                            let separator = text
134                                .get(range.start..range.end)
135                                .expect("invalid character sequence");
136                            cursor = range.end;
137                            return Some(Either::Right(
138                                [(offset, prev_section), (range.start, separator)].into_iter(),
139                            ));
140                        }
141                    }
142                }
143            })
144            .flatten()
145            .filter(|(_, s)| !s.is_empty())
146    }
147}
148
149/// Captures information about document structure for a given text, and their
150/// various semantic levels
151#[derive(Debug)]
152struct SemanticSplitRanges<Level>
153where
154    Level: SemanticLevel,
155{
156    /// Current cursor in the ranges list, so that we can skip over items we've
157    /// already processed.
158    cursor: usize,
159    /// Range of each semantic item and its precalculated semantic level
160    ranges: Vec<(Level, Range<usize>)>,
161}
162
163impl<Level> SemanticSplitRanges<Level>
164where
165    Level: SemanticLevel,
166{
167    fn new(mut ranges: Vec<(Level, Range<usize>)>) -> Self {
168        // Sort by start. If start is equal, sort by end in reverse order, so larger ranges come first.
169        ranges.sort_unstable_by(|(_, a), (_, b)| {
170            a.start.cmp(&b.start).then_with(|| b.end.cmp(&a.end))
171        });
172        Self { cursor: 0, ranges }
173    }
174
175    /// Retrieve ranges for all sections of a given level after an offset
176    fn ranges_after_offset(
177        &self,
178        offset: usize,
179    ) -> impl Iterator<Item = (Level, Range<usize>)> + '_ {
180        self.ranges[self.cursor..]
181            .iter()
182            .filter(move |(_, sep)| sep.start >= offset)
183            .map(|(l, r)| (*l, r.start..r.end))
184    }
185    /// Retrieve ranges for all sections of a given level after an offset
186    fn level_ranges_after_offset(
187        &self,
188        offset: usize,
189        level: Level,
190    ) -> impl Iterator<Item = (Level, Range<usize>)> + '_ {
191        // Find the first item of this level. Allows us to skip larger items of a higher level that surround this one.
192        // Otherwise all lower levels would only return the first item of the higher level that wraps it.
193        let first_item = self
194            .ranges_after_offset(offset)
195            .position(|(l, _)| l == level)
196            .and_then(|i| {
197                self.ranges_after_offset(offset)
198                    .skip(i)
199                    .coalesce(|(a_level, a_range), (b_level, b_range)| {
200                        // If we are at the first item, if two neighboring elements have the same level and start, take the shorter one
201                        if a_level == b_level && a_range.start == b_range.start && i == 0 {
202                            Ok((b_level, b_range))
203                        } else {
204                            Err(((a_level, a_range), (b_level, b_range)))
205                        }
206                    })
207                    // Just take the first of these items
208                    .next()
209            });
210        // let first_item = self.ranges_after_offset(offset).find(|(l, _)| l == &level);
211        self.ranges_after_offset(offset)
212            .filter(move |(l, _)| l >= &level)
213            .skip_while(move |(l, r)| {
214                first_item.as_ref().is_some_and(|(_, fir)| {
215                    (l > &level && r.contains(&fir.start))
216                        || (l == &level && r.start == fir.start && r.end > fir.end)
217                })
218            })
219    }
220
221    /// Return a unique, sorted list of all line break levels present before the next max level, added
222    /// to all of the base semantic levels, in order from smallest to largest
223    fn levels_in_remaining_text(&self, offset: usize) -> impl Iterator<Item = Level> + '_ {
224        self.ranges_after_offset(offset)
225            .map(|(l, _)| l)
226            .sorted()
227            .dedup()
228    }
229
230    /// Split a given text into iterator over each semantic chunk
231    fn semantic_chunks<'splitter, 'text: 'splitter>(
232        &'splitter self,
233        offset: usize,
234        text: &'text str,
235        semantic_level: Level,
236    ) -> impl Iterator<Item = (usize, &'text str)> + 'splitter {
237        Level::sections(
238            text,
239            self.level_ranges_after_offset(offset, semantic_level)
240                .map(move |(l, sep)| (l, sep.start - offset..sep.end - offset)),
241        )
242        .map(move |(i, str)| (offset + i, str))
243    }
244
245    /// Clear out ranges we have moved past so future iterations are faster
246    fn update_cursor(&mut self, cursor: usize) {
247        self.cursor += self.ranges[self.cursor..]
248            .iter()
249            .position(|(_, range)| range.start >= cursor)
250            .unwrap_or_else(|| self.ranges.len() - self.cursor);
251    }
252}
253
254/// Returns chunks of text with their byte offsets as an iterator.
255#[derive(Debug)]
256struct TextChunks<'text, 'sizer, Sizer, Level>
257where
258    Sizer: TryChunkSizer,
259    Level: SemanticLevel,
260{
261    /// Overal capacity of the chunk
262    capacity: ChunkCapacity,
263    /// How to validate chunk sizes
264    chunk_sizer: MemoizedChunkSizer<'sizer, Sizer>,
265    /// Average number of sections in a chunk for each level
266    chunk_stats: ChunkStats,
267    /// Current byte offset in the `text`
268    cursor: usize,
269    /// Reusable container for next sections to avoid extra allocations
270    next_sections: Vec<(usize, &'text str)>,
271    /// Overlap capacity
272    overlap: ChunkCapacity,
273    /// Previous item's end byte offset
274    prev_item_end: usize,
275    /// Splitter used for determining semantic levels.
276    semantic_split: SemanticSplitRanges<Level>,
277    /// Original text to iterate over and generate chunks from
278    text: &'text str,
279    /// The trimming method to apply
280    trim: Trim,
281    /// Whether iteration has ended, including after a sizing error.
282    finished: bool,
283}
284
285impl<'sizer, 'text: 'sizer, Sizer, Level> TextChunks<'text, 'sizer, Sizer, Level>
286where
287    Sizer: TryChunkSizer,
288    Level: SemanticLevel,
289{
290    /// Generate new [`TextChunks`] iterator for a given text.
291    /// Starts with an offset of 0
292    fn new(
293        chunk_config: &'sizer ChunkConfig<Sizer>,
294        text: &'text str,
295        offsets: Vec<(Level, Range<usize>)>,
296        trim: Trim,
297    ) -> Self {
298        let ChunkConfig {
299            capacity,
300            overlap,
301            sizer,
302            trim: trim_enabled,
303        } = chunk_config;
304        Self {
305            capacity: *capacity,
306            chunk_sizer: MemoizedChunkSizer::new(sizer),
307            chunk_stats: ChunkStats::new(),
308            cursor: 0,
309            next_sections: Vec::new(),
310            overlap: (*overlap).into(),
311            prev_item_end: 0,
312            semantic_split: SemanticSplitRanges::new(offsets),
313            text,
314            trim: if *trim_enabled { trim } else { Trim::None },
315            finished: false,
316        }
317    }
318
319    /// Generate the next chunk, applying trimming settings.
320    /// Returns final byte offset and str.
321    fn next_chunk(&mut self) -> Result<(usize, &'text str), Sizer::Error> {
322        self.semantic_split.update_cursor(self.cursor);
323        let low = self.update_next_sections()?;
324        let (start, end) = self.binary_search_next_chunk(low)?;
325        // Chunk selection uses offsets from slices of this same text, so both
326        // endpoints are valid UTF-8 boundaries.
327        let chunk = &self.text[start..end];
328        self.chunk_stats.update_max_chunk_size(end - start);
329
330        // Reset caches so we can reuse the memory allocation
331        self.chunk_sizer.clear_cache();
332        // Optionally move cursor back if overlap is desired
333        self.update_cursor(end)?;
334
335        // Trim whitespace if user requested it
336        Ok(self.trim.trim(start, chunk))
337    }
338
339    /// Use binary search to find the next chunk that fits within the chunk size
340    fn binary_search_next_chunk(&mut self, mut low: usize) -> Result<(usize, usize), Sizer::Error> {
341        let start = self.cursor;
342        let mut end = self.cursor;
343        let mut equals_found = false;
344        let mut high = self.next_sections.len().saturating_sub(1);
345        let mut successful_index = None;
346        let mut successful_chunk_size = None;
347
348        while low <= high {
349            let mid = low + (high - low) / 2;
350            let (offset, str) = self.next_sections[mid];
351            let text_end = offset + str.len();
352            let chunk = &self.text[start..text_end];
353            let chunk_size = self.chunk_sizer.chunk_size(start, chunk, self.trim)?;
354            let fits = self.capacity.fits(chunk_size);
355
356            match fits {
357                Ordering::Less => {
358                    // We got further than the last one, so update end
359                    if text_end > end {
360                        end = text_end;
361                        successful_index = Some(mid);
362                        successful_chunk_size = Some(chunk_size);
363                    }
364                }
365                Ordering::Equal => {
366                    // If we found a smaller equals use it. Or if this is the first equals we found
367                    if text_end < end || !equals_found {
368                        end = text_end;
369                        successful_index = Some(mid);
370                        successful_chunk_size = Some(chunk_size);
371                    }
372                    equals_found = true;
373                }
374                Ordering::Greater => {
375                    // If we're too big on our smallest run, we must return at least one section
376                    if mid == 0 && start == end {
377                        end = text_end;
378                        successful_index = Some(mid);
379                        successful_chunk_size = Some(chunk_size);
380                    }
381                }
382            }
383
384            // Adjust search area
385            if fits.is_lt() {
386                low = mid + 1;
387            } else if mid > 0 {
388                high = mid - 1;
389            } else {
390                // Nothing to adjust
391                break;
392            }
393        }
394
395        if let (Some(successful_index), Some(chunk_size)) =
396            (successful_index, successful_chunk_size)
397        {
398            let mut range = successful_index..self.next_sections.len();
399            // We've already checked the successful index
400            range.next();
401
402            for index in range {
403                let (offset, str) = self.next_sections[index];
404                let text_end = offset + str.len();
405                let chunk = &self.text[start..text_end];
406                let size = self.chunk_sizer.chunk_size(start, chunk, self.trim)?;
407                if size <= chunk_size {
408                    if text_end > end {
409                        end = text_end;
410                    }
411                } else {
412                    break;
413                }
414            }
415        }
416
417        Ok((start, end))
418    }
419
420    /// Use binary search to find the sections that fit within the overlap size.
421    /// If no overlap deisired, return end.
422    fn update_cursor(&mut self, end: usize) -> Result<(), Sizer::Error> {
423        if self.overlap.max == 0 {
424            self.cursor = end;
425            return Ok(());
426        }
427
428        // Binary search for overlap
429        let mut start = end;
430        let mut low = 0;
431        // Find closest index that would work
432        let mut high = match self
433            .next_sections
434            .binary_search_by_key(&end, |(offset, str)| offset + str.len())
435        {
436            Ok(i) | Err(i) => i,
437        };
438
439        while low <= high {
440            let mid = low + (high - low) / 2;
441            let (offset, _) = self.next_sections[mid];
442            let chunk_size = self.chunk_sizer.chunk_size(
443                offset,
444                self.text.get(offset..end).expect("Invalid range"),
445                self.trim,
446            )?;
447            let fits = self.overlap.fits(chunk_size);
448
449            // We got further than the last one, so update start
450            if fits.is_le() && offset < start && offset > self.cursor {
451                start = offset;
452            }
453
454            // Adjust search area
455            if fits.is_lt() && mid > 0 {
456                high = mid - 1;
457            } else {
458                low = mid + 1;
459            }
460        }
461
462        self.cursor = start;
463        Ok(())
464    }
465
466    /// Find the ideal next sections, breaking it up until we find the largest chunk.
467    /// Increasing length of chunk until we find biggest size to minimize validation time
468    /// on huge chunks
469    #[expect(clippy::too_many_lines)]
470    fn update_next_sections(&mut self) -> Result<usize, Sizer::Error> {
471        // First thing, clear out the list, but reuse the allocated memory
472        self.next_sections.clear();
473
474        let remaining_text = self.text.get(self.cursor..).unwrap();
475        let mut lower_level = None;
476
477        let (semantic_level, mut max_offset) = self.chunk_sizer.find_correct_level(
478            self.cursor,
479            &self.capacity,
480            self.semantic_split
481                .levels_in_remaining_text(self.cursor)
482                .filter_map(|level| {
483                    let first_chunk = self
484                        .semantic_split
485                        .semantic_chunks(self.cursor, remaining_text, level)
486                        .next();
487
488                    let result = first_chunk.map(|(_, str)| {
489                        let candidate_lower_level = lower_level;
490                        (level, str, candidate_lower_level)
491                    });
492
493                    lower_level = Some(level);
494                    result
495                }),
496            |lower_level, chunk_end| {
497                lower_level.map_or_else(
498                    || Either::Left(std::iter::empty()),
499                    |lower_level| {
500                        Either::Right(
501                            self.semantic_split
502                                .semantic_chunks(self.cursor, remaining_text, lower_level)
503                                .map(|(offset, text)| offset + text.len())
504                                .take_while(move |end| *end <= chunk_end),
505                        )
506                    },
507                )
508            },
509            self.trim,
510        )?;
511
512        let sections = if let Some(semantic_level) = semantic_level {
513            Either::Left(self.semantic_split.semantic_chunks(
514                self.cursor,
515                remaining_text,
516                semantic_level,
517            ))
518        } else {
519            let (semantic_level, fallback_max_offset) = self.chunk_sizer.find_correct_level(
520                self.cursor,
521                &self.capacity,
522                FallbackLevel::iter().filter_map(|level| {
523                    level
524                        .sections(remaining_text)
525                        .next()
526                        .map(|(_, str)| (level, str, level.boundary_level_for_probe()))
527                }),
528                |lower_level, chunk_end| {
529                    lower_level.map_or_else(
530                        || Either::Left(std::iter::empty()),
531                        |lower_level| {
532                            Either::Right(
533                                lower_level
534                                    .sections(remaining_text)
535                                    .map(|(offset, text)| self.cursor + offset + text.len())
536                                    .take_while(move |end| *end <= chunk_end),
537                            )
538                        },
539                    )
540                },
541                self.trim,
542            )?;
543
544            max_offset = match (fallback_max_offset, max_offset) {
545                (Some(fallback), Some(max)) => Some(fallback.min(max)),
546                (fallback, max) => fallback.or(max),
547            };
548
549            let fallback_level = semantic_level.unwrap_or(FallbackLevel::Char);
550
551            Either::Right(
552                fallback_level
553                    .sections(remaining_text)
554                    .map(|(offset, text)| (self.cursor + offset, text)),
555            )
556        };
557
558        let mut sections = sections
559            .take_while(move |(offset, _)| max_offset.is_none_or(|max| *offset <= max))
560            .filter(|(_, str)| !str.is_empty());
561
562        // Start filling up the next sections. Since calculating the size of the chunk gets more expensive
563        // the farther we go, we conservatively check for a smaller range to do the later binary search in.
564        let mut low = 0;
565        let mut prev_equals: Option<usize> = None;
566        let max = self.capacity.max;
567        let mut target_offset = self.chunk_stats.max_chunk_size.unwrap_or(max);
568
569        loop {
570            let prev_num = self.next_sections.len();
571            for (offset, str) in sections.by_ref() {
572                self.next_sections.push((offset, str));
573                if offset + str.len() > (self.cursor.saturating_add(target_offset)) {
574                    break;
575                }
576            }
577            let new_num = self.next_sections.len();
578            // If we've iterated through the whole iterator, break here.
579            if new_num - prev_num == 0 {
580                break;
581            }
582
583            // Check if the last item fits
584            if let Some(&(offset, str)) = self.next_sections.last() {
585                let text_end = offset + str.len();
586                if (text_end - self.cursor) < target_offset {
587                    break;
588                }
589                let chunk_size = self.chunk_sizer.chunk_size(
590                    offset,
591                    self.text.get(self.cursor..text_end).expect("Invalid range"),
592                    self.trim,
593                )?;
594                let fits = self.capacity.fits(chunk_size);
595
596                if fits.is_le() {
597                    let final_offset = offset + str.len() - self.cursor;
598                    let size = chunk_size.max(1);
599                    let diff = (max - size).max(1);
600                    let avg_size = final_offset.div_ceil(size);
601
602                    target_offset = final_offset
603                        .saturating_add(diff.saturating_mul(avg_size))
604                        .saturating_add(final_offset.div_ceil(10));
605                }
606
607                match fits {
608                    Ordering::Less => {
609                        // We know we can go higher
610                        low = new_num.saturating_sub(1);
611                    }
612                    Ordering::Equal => {
613                        // Don't update low because it could be a range
614                        // If we've seen a previous equals, we can break if the size is bigger already
615                        if let Some(prev) = prev_equals {
616                            if prev < chunk_size {
617                                break;
618                            }
619                        }
620                        prev_equals = Some(chunk_size);
621                    }
622                    Ordering::Greater => {
623                        break;
624                    }
625                }
626            }
627        }
628
629        Ok(low)
630    }
631}
632
633impl<'sizer, 'text: 'sizer, Sizer, Level> Iterator for TextChunks<'text, 'sizer, Sizer, Level>
634where
635    Sizer: TryChunkSizer,
636    Level: SemanticLevel,
637{
638    type Item = Result<(usize, &'text str), Sizer::Error>;
639
640    fn next(&mut self) -> Option<Self::Item> {
641        if self.finished {
642            return None;
643        }
644        loop {
645            // Make sure we haven't reached the end
646            if self.cursor >= self.text.len() {
647                self.finished = true;
648                return None;
649            }
650
651            match self.next_chunk() {
652                Err(error) => {
653                    self.finished = true;
654                    return Some(Err(error));
655                }
656                // Make sure we didn't get an empty chunk. Should only happen in
657                // cases where we trim.
658                Ok((_, "")) => {}
659                Ok(c) => {
660                    let item_end = c.0 + c.1.len();
661                    // Skip because we've emitted a chunk whose content we've already emitted
662                    if item_end <= self.prev_item_end {
663                        continue;
664                    }
665                    self.prev_item_end = item_end;
666                    return Some(Ok(c));
667                }
668            }
669        }
670    }
671}
672
673/// Represents a chunk of text along with its byte and character offsets.
674#[derive(Debug, Clone, Copy, PartialEq, Eq)]
675pub struct ChunkCharIndex<'text> {
676    /// The text of the generated chunk.
677    pub chunk: &'text str,
678    /// The byte offset of the chunk within the original text.
679    pub byte_offset: usize,
680    /// The character offset of the chunk within the original text.
681    pub char_offset: usize,
682}
683
684/// Returns chunks of text with their byte and character offsets as an iterator.
685#[derive(Debug)]
686struct TextChunksWithCharIndices<'text, 'sizer, Sizer, Level>
687where
688    Sizer: TryChunkSizer,
689    Level: SemanticLevel,
690{
691    /// The text being chunked.
692    text: &'text str,
693    /// The main iterator over chunks of text.
694    text_chunks: TextChunks<'text, 'sizer, Sizer, Level>,
695    /// The byte offset of the previous chunk within the original text.
696    byte_offset: usize,
697    /// The character offset of the previous chunk within the original text.
698    char_offset: usize,
699}
700
701impl<'sizer, 'text: 'sizer, Sizer, Level> TextChunksWithCharIndices<'text, 'sizer, Sizer, Level>
702where
703    Sizer: TryChunkSizer,
704    Level: SemanticLevel,
705{
706    /// Generate new [`TextChunksWithCharIndices`] iterator for a given text.
707    /// Starts with an offset of 0
708    fn new(
709        chunk_config: &'sizer ChunkConfig<Sizer>,
710        text: &'text str,
711        offsets: Vec<(Level, Range<usize>)>,
712        trim: Trim,
713    ) -> Self {
714        Self {
715            text,
716            text_chunks: TextChunks::new(chunk_config, text, offsets, trim),
717            byte_offset: 0,
718            char_offset: 0,
719        }
720    }
721}
722
723impl<'sizer, 'text: 'sizer, Sizer, Level> Iterator
724    for TextChunksWithCharIndices<'text, 'sizer, Sizer, Level>
725where
726    Sizer: TryChunkSizer,
727    Level: SemanticLevel,
728{
729    type Item = Result<ChunkCharIndex<'text>, Sizer::Error>;
730
731    fn next(&mut self) -> Option<Self::Item> {
732        self.text_chunks.next().map(|result| {
733            let (byte_offset, chunk) = result?;
734            let preceding_text = self
735                .text
736                .get(self.byte_offset..byte_offset)
737                .expect("Invalid byte sequence");
738            self.byte_offset = byte_offset;
739            self.char_offset += preceding_text.chars().count();
740            Ok(ChunkCharIndex {
741                chunk,
742                byte_offset,
743                char_offset: self.char_offset,
744            })
745        })
746    }
747}
748
749/// Keeps track of the average size of chunks as we go
750#[derive(Debug, Default)]
751struct ChunkStats {
752    /// The size of the biggest chunk we've seen, if we have seen at least one
753    max_chunk_size: Option<usize>,
754}
755
756impl ChunkStats {
757    fn new() -> Self {
758        Self::default()
759    }
760
761    /// Update statistics after the chunk has been produced
762    fn update_max_chunk_size(&mut self, size: usize) {
763        self.max_chunk_size = self.max_chunk_size.map(|s| s.max(size)).or(Some(size));
764    }
765}
766
767#[cfg(test)]
768mod tests {
769    use super::*;
770
771    #[test]
772    fn chunk_stats_empty() {
773        let stats = ChunkStats::new();
774        assert_eq!(stats.max_chunk_size, None);
775    }
776
777    #[test]
778    fn chunk_stats_one() {
779        let mut stats = ChunkStats::new();
780        stats.update_max_chunk_size(10);
781        assert_eq!(stats.max_chunk_size, Some(10));
782    }
783
784    #[test]
785    fn chunk_stats_multiple() {
786        let mut stats = ChunkStats::new();
787        stats.update_max_chunk_size(10);
788        stats.update_max_chunk_size(20);
789        stats.update_max_chunk_size(30);
790        assert_eq!(stats.max_chunk_size, Some(30));
791    }
792
793    impl SemanticLevel for usize {}
794
795    #[test]
796    fn semantic_ranges_are_sorted() {
797        let ranges = SemanticSplitRanges::new(vec![(0, 0..1), (1, 0..2), (0, 1..2), (2, 0..4)]);
798
799        assert_eq!(
800            ranges.ranges,
801            vec![(2, 0..4), (1, 0..2), (0, 0..1), (0, 1..2)]
802        );
803    }
804
805    #[test]
806    fn semantic_ranges_skip_previous_ranges() {
807        let mut ranges = SemanticSplitRanges::new(vec![(0, 0..1), (1, 0..2), (0, 1..2), (2, 0..4)]);
808
809        ranges.update_cursor(1);
810
811        assert_eq!(
812            ranges.ranges_after_offset(0).collect::<Vec<_>>(),
813            vec![(0, 1..2)]
814        );
815    }
816}