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
32trait Splitter<Sizer>
35where
36 Sizer: TryChunkSizer,
37{
38 type Level: SemanticLevel;
39
40 const TRIM: Trim = Trim::All;
42
43 fn chunk_config(&self) -> &ChunkConfig<Sizer>;
45
46 fn parse(&self, text: &str) -> Vec<(Self::Level, Range<usize>)>;
48
49 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 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 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
101trait SemanticLevel: Copy + fmt::Debug + Ord + PartialOrd + 'static {
103 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 None if final_match => return None,
119 None => {
121 final_match = true;
122 return text.get(cursor..).map(|t| Either::Left(once((cursor, t))));
123 }
124 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#[derive(Debug)]
152struct SemanticSplitRanges<Level>
153where
154 Level: SemanticLevel,
155{
156 cursor: usize,
159 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 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 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 fn level_ranges_after_offset(
187 &self,
188 offset: usize,
189 level: Level,
190 ) -> impl Iterator<Item = (Level, Range<usize>)> + '_ {
191 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 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 .next()
209 });
210 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 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 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 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#[derive(Debug)]
256struct TextChunks<'text, 'sizer, Sizer, Level>
257where
258 Sizer: TryChunkSizer,
259 Level: SemanticLevel,
260{
261 capacity: ChunkCapacity,
263 chunk_sizer: MemoizedChunkSizer<'sizer, Sizer>,
265 chunk_stats: ChunkStats,
267 cursor: usize,
269 next_sections: Vec<(usize, &'text str)>,
271 overlap: ChunkCapacity,
273 prev_item_end: usize,
275 semantic_split: SemanticSplitRanges<Level>,
277 text: &'text str,
279 trim: Trim,
281 finished: bool,
283}
284
285impl<'sizer, 'text: 'sizer, Sizer, Level> TextChunks<'text, 'sizer, Sizer, Level>
286where
287 Sizer: TryChunkSizer,
288 Level: SemanticLevel,
289{
290 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 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 let chunk = &self.text[start..end];
328 self.chunk_stats.update_max_chunk_size(end - start);
329
330 self.chunk_sizer.clear_cache();
332 self.update_cursor(end)?;
334
335 Ok(self.trim.trim(start, chunk))
337 }
338
339 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 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 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 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 if fits.is_lt() {
386 low = mid + 1;
387 } else if mid > 0 {
388 high = mid - 1;
389 } else {
390 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 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 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 let mut start = end;
430 let mut low = 0;
431 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 if fits.is_le() && offset < start && offset > self.cursor {
451 start = offset;
452 }
453
454 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 #[expect(clippy::too_many_lines)]
470 fn update_next_sections(&mut self) -> Result<usize, Sizer::Error> {
471 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 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 new_num - prev_num == 0 {
580 break;
581 }
582
583 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 low = new_num.saturating_sub(1);
611 }
612 Ordering::Equal => {
613 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 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 Ok((_, "")) => {}
659 Ok(c) => {
660 let item_end = c.0 + c.1.len();
661 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#[derive(Debug, Clone, Copy, PartialEq, Eq)]
675pub struct ChunkCharIndex<'text> {
676 pub chunk: &'text str,
678 pub byte_offset: usize,
680 pub char_offset: usize,
682}
683
684#[derive(Debug)]
686struct TextChunksWithCharIndices<'text, 'sizer, Sizer, Level>
687where
688 Sizer: TryChunkSizer,
689 Level: SemanticLevel,
690{
691 text: &'text str,
693 text_chunks: TextChunks<'text, 'sizer, Sizer, Level>,
695 byte_offset: usize,
697 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 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#[derive(Debug, Default)]
751struct ChunkStats {
752 max_chunk_size: Option<usize>,
754}
755
756impl ChunkStats {
757 fn new() -> Self {
758 Self::default()
759 }
760
761 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}