text_splitter/
chunk_size.rs

1use std::{
2    borrow::Cow,
3    cell::{Ref, RefMut},
4    cmp::Ordering,
5    convert::Infallible,
6    fmt,
7    ops::{Deref, Range, RangeFrom, RangeFull, RangeInclusive, RangeTo, RangeToInclusive},
8    rc::Rc,
9    sync::Arc,
10};
11
12use ahash::AHashMap;
13use itertools::Itertools;
14use thiserror::Error;
15
16mod characters;
17#[cfg(feature = "tokenizers")]
18mod huggingface;
19#[cfg(feature = "tiktoken-rs")]
20mod tiktoken;
21
22use crate::trim::Trim;
23pub use characters::Characters;
24
25const PROBE_OVERSHOOT_DIVISOR: usize = 4;
26const PROBE_START_BYTE_FACTOR: usize = 8;
27
28/// Indicates there was an error with the chunk capacity configuration.
29/// The `Display` implementation will provide a human-readable error message to
30/// help debug the issue that caused the error.
31#[derive(Error, Debug)]
32#[error(transparent)]
33pub struct ChunkCapacityError(#[from] ChunkCapacityErrorRepr);
34
35/// Private error and free to change across minor version of the crate.
36#[derive(Error, Debug)]
37enum ChunkCapacityErrorRepr {
38    #[error("Max chunk size must be greater than or equal to the desired chunk size")]
39    MaxLessThanDesired,
40}
41
42/// Describes the valid chunk size(s) that can be generated.
43///
44/// The `desired` size is the target size for the chunk. In most cases, this
45/// will also serve as the maximum size of the chunk. It is always possible
46/// that a chunk may be returned that is less than the `desired` value, as
47/// adding the next piece of text may have made it larger than the `desired`
48/// capacity.
49///
50/// The `max` size is the maximum possible chunk size that can be generated.
51/// By setting this to a larger value than `desired`, it means that the chunk
52/// should be as close to `desired` as possible, but can be larger if it means
53/// staying at a larger semantic level.
54///
55/// The splitter will consume text until at maxumum somewhere between `desired`
56/// and `max`, if they differ, but never above `max`.
57///
58/// If you need to ensure a fixed size, set `desired` and `max` to the same
59/// value. For example, if you are trying to maximize the context window for an
60/// embedding.
61///
62/// If you are loosely targeting a size, but have some extra room, for example
63/// in a RAG use case where you roughly want a certain part of a document, you
64/// can set `max` to your absolute maxumum, and the splitter can stay at a
65/// higher semantic level when determining the chunk.
66#[derive(Copy, Clone, Debug, PartialEq)]
67pub struct ChunkCapacity {
68    pub(crate) desired: usize,
69    pub(crate) max: usize,
70}
71
72impl ChunkCapacity {
73    /// Create a new `ChunkCapacity` with the same `desired` and `max` size.
74    #[must_use]
75    pub fn new(size: usize) -> Self {
76        Self {
77            desired: size,
78            max: size,
79        }
80    }
81
82    /// The `desired` size is the target size for the chunk. In most cases, this
83    /// will also serve as the maximum size of the chunk. It is always possible
84    /// that a chunk may be returned that is less than the `desired` value, as
85    /// adding the next piece of text may have made it larger than the `desired`
86    /// capacity.
87    #[must_use]
88    pub fn desired(&self) -> usize {
89        self.desired
90    }
91
92    /// The `max` size is the maximum possible chunk size that can be generated.
93    /// By setting this to a larger value than `desired`, it means that the chunk
94    /// should be as close to `desired` as possible, but can be larger if it means
95    /// staying at a larger semantic level.
96    #[must_use]
97    pub fn max(&self) -> usize {
98        self.max
99    }
100
101    /// If you need to ensure a fixed size, set `desired` and `max` to the same
102    /// value. For example, if you are trying to maximize the context window for an
103    /// embedding.
104    ///
105    /// If you are loosely targeting a size, but have some extra room, for example
106    /// in a RAG use case where you roughly want a certain part of a document, you
107    /// can set `max` to your absolute maxumum, and the splitter can stay at a
108    /// higher semantic level when determining the chunk.
109    ///
110    /// # Errors
111    ///
112    /// If the `max` size is less than the `desired` size, an error is returned.
113    pub fn with_max(mut self, max: usize) -> Result<Self, ChunkCapacityError> {
114        if max < self.desired {
115            Err(ChunkCapacityError(
116                ChunkCapacityErrorRepr::MaxLessThanDesired,
117            ))
118        } else {
119            self.max = max;
120            Ok(self)
121        }
122    }
123
124    /// Validate if a given chunk fits within the capacity
125    ///
126    /// - `Ordering::Less` indicates more could be added
127    /// - `Ordering::Equal` indicates the chunk is within the capacity range
128    /// - `Ordering::Greater` indicates the chunk is larger than the capacity
129    #[must_use]
130    pub fn fits(&self, chunk_size: usize) -> Ordering {
131        if chunk_size < self.desired {
132            Ordering::Less
133        } else if chunk_size > self.max {
134            Ordering::Greater
135        } else {
136            Ordering::Equal
137        }
138    }
139}
140
141impl From<usize> for ChunkCapacity {
142    fn from(size: usize) -> Self {
143        ChunkCapacity::new(size)
144    }
145}
146
147impl From<Range<usize>> for ChunkCapacity {
148    fn from(range: Range<usize>) -> Self {
149        ChunkCapacity::new(range.start)
150            .with_max(range.end.saturating_sub(1).max(range.start))
151            .expect("invalid range")
152    }
153}
154
155impl From<RangeFrom<usize>> for ChunkCapacity {
156    fn from(range: RangeFrom<usize>) -> Self {
157        ChunkCapacity::new(range.start)
158            .with_max(usize::MAX)
159            .expect("invalid range")
160    }
161}
162
163impl From<RangeFull> for ChunkCapacity {
164    fn from(_: RangeFull) -> Self {
165        ChunkCapacity::new(usize::MIN)
166            .with_max(usize::MAX)
167            .expect("invalid range")
168    }
169}
170
171impl From<RangeInclusive<usize>> for ChunkCapacity {
172    fn from(range: RangeInclusive<usize>) -> Self {
173        ChunkCapacity::new(*range.start())
174            .with_max(*range.end())
175            .expect("invalid range")
176    }
177}
178
179impl From<RangeTo<usize>> for ChunkCapacity {
180    fn from(range: RangeTo<usize>) -> Self {
181        ChunkCapacity::new(usize::MIN)
182            .with_max(range.end.saturating_sub(1))
183            .expect("invalid range")
184    }
185}
186
187impl From<RangeToInclusive<usize>> for ChunkCapacity {
188    fn from(range: RangeToInclusive<usize>) -> Self {
189        ChunkCapacity::new(usize::MIN)
190            .with_max(range.end)
191            .expect("invalid range")
192    }
193}
194
195/// Determines the size of a given chunk.
196pub trait ChunkSizer {
197    /// Determine the size of a given chunk to use for validation
198    fn size(&self, chunk: &str) -> usize;
199}
200
201/// Determines the size of a chunk, returning an error if sizing fails.
202///
203/// All [`ChunkSizer`] implementations automatically implement this trait with
204/// [`Infallible`] as their error type. This adapter preserves their behavior,
205/// including any panics; it does not catch panics or convert them into errors.
206/// A type implementing [`ChunkSizer`] cannot also provide a separate
207/// implementation of this trait. Use a newtype for different fallible behavior.
208///
209/// ```
210/// use text_splitter::{ChunkConfig, TextSplitter, TryChunkSizer};
211///
212/// struct CheckedCharacters;
213///
214/// impl TryChunkSizer for CheckedCharacters {
215///     type Error = &'static str;
216///
217///     fn try_size(&self, chunk: &str) -> Result<usize, Self::Error> {
218///         if chunk.contains('\0') {
219///             Err("NUL is not supported")
220///         } else {
221///             Ok(chunk.chars().count())
222///         }
223///     }
224/// }
225///
226/// let splitter = TextSplitter::new(ChunkConfig::new(10).with_sizer(CheckedCharacters));
227/// let chunks = splitter.try_chunks("hello\0world").collect::<Result<Vec<_>, _>>();
228/// assert_eq!(chunks, Err("NUL is not supported"));
229/// ```
230pub trait TryChunkSizer {
231    /// The error returned when sizing fails.
232    type Error;
233
234    /// Determine the size of a chunk.
235    ///
236    /// # Errors
237    ///
238    /// Returns the sizer's error if the chunk cannot be sized.
239    fn try_size(&self, chunk: &str) -> Result<usize, Self::Error>;
240}
241
242impl<T: ChunkSizer + ?Sized> TryChunkSizer for T {
243    type Error = Infallible;
244
245    fn try_size(&self, chunk: &str) -> Result<usize, Self::Error> {
246        Ok(self.size(chunk))
247    }
248}
249
250impl<T> ChunkSizer for &T
251where
252    T: ChunkSizer,
253{
254    fn size(&self, chunk: &str) -> usize {
255        (*self).size(chunk)
256    }
257}
258
259impl<T> ChunkSizer for Ref<'_, T>
260where
261    T: ChunkSizer,
262{
263    fn size(&self, chunk: &str) -> usize {
264        self.deref().size(chunk)
265    }
266}
267
268impl<T> ChunkSizer for RefMut<'_, T>
269where
270    T: ChunkSizer,
271{
272    fn size(&self, chunk: &str) -> usize {
273        self.deref().size(chunk)
274    }
275}
276
277impl<T> ChunkSizer for Box<T>
278where
279    T: ChunkSizer,
280{
281    fn size(&self, chunk: &str) -> usize {
282        self.deref().size(chunk)
283    }
284}
285
286impl<T> ChunkSizer for Cow<'_, T>
287where
288    T: ChunkSizer + ToOwned + ?Sized,
289    <T as ToOwned>::Owned: ChunkSizer,
290{
291    fn size(&self, chunk: &str) -> usize {
292        self.as_ref().size(chunk)
293    }
294}
295
296impl<T> ChunkSizer for Rc<T>
297where
298    T: ChunkSizer,
299{
300    fn size(&self, chunk: &str) -> usize {
301        self.deref().size(chunk)
302    }
303}
304
305impl<T> ChunkSizer for Arc<T>
306where
307    T: ChunkSizer,
308{
309    fn size(&self, chunk: &str) -> usize {
310        self.as_ref().size(chunk)
311    }
312}
313
314/// Indicates there was an error with the chunk configuration.
315/// The `Display` implementation will provide a human-readable error message to
316/// help debug the issue that caused the error.
317#[derive(Error, Debug)]
318#[error(transparent)]
319pub struct ChunkConfigError(#[from] ChunkConfigErrorRepr);
320
321/// Private error and free to change across minor version of the crate.
322#[derive(Error, Debug)]
323enum ChunkConfigErrorRepr {
324    #[error("The overlap is larger than or equal to the desired chunk capacity")]
325    OverlapLargerThanCapacity,
326}
327
328/// Configuration for how chunks should be created
329#[derive(Debug)]
330pub struct ChunkConfig<Sizer>
331where
332    Sizer: TryChunkSizer,
333{
334    /// The chunk capacity to use for filling chunks
335    pub(crate) capacity: ChunkCapacity,
336    /// The amount of overlap between chunks. Defaults to 0.
337    pub(crate) overlap: usize,
338    /// The chunk sizer to use for determining the size of each chunk
339    pub(crate) sizer: Sizer,
340    /// Whether whitespace will be trimmed from the beginning and end of each chunk
341    pub(crate) trim: bool,
342}
343
344impl ChunkConfig<Characters> {
345    /// Create a basic configuration for chunking with only the required value a chunk capacity.
346    ///
347    /// By default, chunk sizes will be calculated based on the number of characters in each chunk.
348    /// You can set a custom chunk sizer by calling [`Self::with_sizer`].
349    ///
350    /// By default, chunks will be trimmed. If you want to preserve whitespace,
351    /// call [`Self::with_trim`] and set it to `false`.
352    #[must_use]
353    pub fn new(capacity: impl Into<ChunkCapacity>) -> Self {
354        Self {
355            capacity: capacity.into(),
356            overlap: 0,
357            sizer: Characters,
358            trim: true,
359        }
360    }
361}
362
363impl<Sizer> ChunkConfig<Sizer>
364where
365    Sizer: TryChunkSizer,
366{
367    /// Retrieve a reference to the chunk capacity for this configuration.
368    pub fn capacity(&self) -> &ChunkCapacity {
369        &self.capacity
370    }
371
372    /// Retrieve the amount of overlap between chunks.
373    pub fn overlap(&self) -> usize {
374        self.overlap
375    }
376
377    /// Set the amount of overlap between chunks.
378    ///
379    /// # Errors
380    ///
381    /// Will return an error if the overlap is larger than or equal to the chunk capacity.
382    pub fn with_overlap(mut self, overlap: usize) -> Result<Self, ChunkConfigError> {
383        if overlap >= self.capacity.desired {
384            Err(ChunkConfigError(
385                ChunkConfigErrorRepr::OverlapLargerThanCapacity,
386            ))
387        } else {
388            self.overlap = overlap;
389            Ok(self)
390        }
391    }
392
393    /// Retrieve a reference to the chunk sizer for this configuration.
394    pub fn sizer(&self) -> &Sizer {
395        &self.sizer
396    }
397
398    /// Set a custom chunk sizer to use for determining the size of each chunk
399    ///
400    /// ```
401    /// use text_splitter::{Characters, ChunkConfig};
402    ///
403    /// let config = ChunkConfig::new(512).with_sizer(Characters);
404    /// ```
405    #[must_use]
406    pub fn with_sizer<S: TryChunkSizer>(self, sizer: S) -> ChunkConfig<S> {
407        ChunkConfig {
408            capacity: self.capacity,
409            overlap: self.overlap,
410            sizer,
411            trim: self.trim,
412        }
413    }
414
415    /// Whether chunkd should have whitespace trimmed from the beginning and end or not.
416    pub fn trim(&self) -> bool {
417        self.trim
418    }
419
420    /// Specify whether chunks should have whitespace trimmed from the
421    /// beginning and end or not.
422    ///
423    /// If `false` (default), joining all chunks should return the original
424    /// string.
425    /// If `true`, all chunks will have whitespace removed from beginning and end.
426    ///
427    /// ```
428    /// use text_splitter::ChunkConfig;
429    ///
430    /// let config = ChunkConfig::new(512).with_trim(false);
431    /// ```
432    #[must_use]
433    pub fn with_trim(mut self, trim: bool) -> Self {
434        self.trim = trim;
435        self
436    }
437}
438
439impl<T> From<T> for ChunkConfig<Characters>
440where
441    T: Into<ChunkCapacity>,
442{
443    fn from(capacity: T) -> Self {
444        Self::new(capacity)
445    }
446}
447
448/// A memoized chunk sizer that caches the size of chunks.
449/// Very helpful when the same chunk is being validated multiple times, which
450/// happens often, and can be expensive to compute, such as with tokenizers.
451#[derive(Debug)]
452pub struct MemoizedChunkSizer<'sizer, Sizer>
453where
454    Sizer: TryChunkSizer,
455{
456    /// Cache of chunk sizes per byte offset range for base capacity
457    size_cache: AHashMap<Range<usize>, usize>,
458    /// The sizer used for caluclating chunk sizes
459    sizer: &'sizer Sizer,
460}
461
462impl<'sizer, Sizer> MemoizedChunkSizer<'sizer, Sizer>
463where
464    Sizer: TryChunkSizer,
465{
466    /// Wrap any chunk sizer for memoization
467    pub fn new(sizer: &'sizer Sizer) -> Self {
468        Self {
469            size_cache: AHashMap::new(),
470            sizer,
471        }
472    }
473
474    /// Determine the size of a given chunk to use for validation,
475    /// returning a cached value if it exists, and storing the result if not.
476    pub fn chunk_size(
477        &mut self,
478        offset: usize,
479        chunk: &str,
480        trim: Trim,
481    ) -> Result<usize, Sizer::Error> {
482        let (offset, chunk) = trim.trim(offset, chunk);
483        match self.size_cache.entry(offset..(offset + chunk.len())) {
484            std::collections::hash_map::Entry::Occupied(entry) => Ok(*entry.get()),
485            std::collections::hash_map::Entry::Vacant(entry) => {
486                let size = self.sizer.try_size(chunk)?;
487                entry.insert(size);
488                Ok(size)
489            }
490        }
491    }
492
493    /// Find the best level to start splitting the text
494    pub fn find_correct_level<'text, L, Boundaries>(
495        &mut self,
496        offset: usize,
497        capacity: &ChunkCapacity,
498        levels_with_first_chunk: impl Iterator<Item = (L, &'text str, Option<L>)>,
499        mut lower_boundaries_for: impl FnMut(Option<L>, usize) -> Boundaries,
500        trim: Trim,
501    ) -> Result<(Option<L>, Option<usize>), Sizer::Error>
502    where
503        L: Copy + fmt::Debug,
504        Boundaries: Iterator<Item = usize>,
505    {
506        let mut semantic_level = None;
507        let mut max_offset = None;
508
509        // We assume that larger levels are also longer. We can skip lower levels if going to a higher level would result in a shorter text
510        let levels_with_first_chunk = levels_with_first_chunk.coalesce(
511            |(a_level, a_str, a_lower_level), (b_level, b_str, b_lower_level)| {
512                if a_str.len() >= b_str.len() {
513                    Ok((b_level, b_str, b_lower_level))
514                } else {
515                    Err((
516                        (a_level, a_str, a_lower_level),
517                        (b_level, b_str, b_lower_level),
518                    ))
519                }
520            },
521        );
522
523        for (level, str, lower_level) in levels_with_first_chunk {
524            // Skip tokenizing levels that we know are too small anyway.
525            let len = str.len();
526            if len > capacity.max {
527                let mut lower_boundaries = lower_boundaries_for(lower_level, offset + len);
528                let fits = self.chunk_fits_with_boundaries(
529                    offset,
530                    str,
531                    capacity,
532                    &mut lower_boundaries,
533                    trim,
534                )?;
535                // If this no longer fits, we use the level we are at.
536                if fits.is_gt() {
537                    // Keep the old search window while avoiding the full-size call
538                    // when lower semantic boundaries prove this section is too large.
539                    max_offset = Some(offset + len);
540                    break;
541                }
542            }
543            // Otherwise break up the text with the next level
544            semantic_level = Some(level);
545        }
546
547        Ok((semantic_level, max_offset))
548    }
549
550    fn chunk_fits_with_boundaries(
551        &mut self,
552        offset: usize,
553        chunk: &str,
554        capacity: &ChunkCapacity,
555        lower_boundaries: &mut impl Iterator<Item = usize>,
556        trim: Trim,
557    ) -> Result<Ordering, Sizer::Error> {
558        let chunk_end = offset + chunk.len();
559        let probe_max = probe_max_size(capacity);
560        let probe_start = offset
561            .saturating_add(capacity.max.saturating_mul(PROBE_START_BYTE_FACTOR).max(1))
562            .min(chunk_end);
563        if probe_start == chunk_end {
564            let chunk_size = self.chunk_size(offset, chunk, trim)?;
565            return Ok(capacity.fits(chunk_size));
566        }
567
568        let mut step: usize = 1;
569        let mut next_boundary =
570            lower_boundaries.find(|&boundary| boundary > offset && boundary >= probe_start);
571
572        while let Some(boundary) = next_boundary {
573            if boundary > chunk_end {
574                break;
575            }
576
577            let prefix = chunk
578                .get(..(boundary - offset))
579                .expect("valid character boundary");
580            let chunk_size = self.chunk_size(offset, prefix, trim)?;
581            let fits = capacity.fits(chunk_size);
582            if chunk_size > probe_max || boundary == chunk_end {
583                return Ok(fits);
584            }
585
586            let skip = step.saturating_sub(1);
587            step = step.saturating_mul(2);
588            next_boundary = lower_boundaries.nth(skip);
589        }
590
591        let chunk_size = self.chunk_size(offset, chunk, trim)?;
592        Ok(capacity.fits(chunk_size))
593    }
594
595    /// Clear the cached values. Once we've moved the cursor,
596    /// we don't need to keep the old values around.
597    pub fn clear_cache(&mut self) {
598        self.size_cache.clear();
599    }
600}
601
602fn probe_max_size(capacity: &ChunkCapacity) -> usize {
603    capacity
604        .max
605        .saturating_add((capacity.max / PROBE_OVERSHOOT_DIVISOR).max(1))
606}
607
608#[cfg(test)]
609mod tests {
610    use std::{
611        cell::RefCell,
612        sync::atomic::{self, AtomicUsize},
613    };
614
615    use crate::trim::Trim;
616
617    use super::*;
618
619    #[test]
620    fn check_chunk_capacity() {
621        let chunk = "12345";
622
623        assert_eq!(
624            ChunkCapacity::from(4).fits(Characters.size(chunk)),
625            Ordering::Greater
626        );
627        assert_eq!(
628            ChunkCapacity::from(5).fits(Characters.size(chunk)),
629            Ordering::Equal
630        );
631        assert_eq!(
632            ChunkCapacity::from(6).fits(Characters.size(chunk)),
633            Ordering::Less
634        );
635    }
636
637    #[test]
638    fn check_chunk_capacity_for_range() {
639        let chunk = "12345";
640
641        assert_eq!(
642            ChunkCapacity::from(0..0).fits(Characters.size(chunk)),
643            Ordering::Greater
644        );
645        assert_eq!(
646            ChunkCapacity::from(0..5).fits(Characters.size(chunk)),
647            Ordering::Greater
648        );
649        assert_eq!(
650            ChunkCapacity::from(5..6).fits(Characters.size(chunk)),
651            Ordering::Equal
652        );
653        assert_eq!(
654            ChunkCapacity::from(6..100).fits(Characters.size(chunk)),
655            Ordering::Less
656        );
657    }
658
659    #[test]
660    fn check_chunk_capacity_for_range_from() {
661        let chunk = "12345";
662
663        assert_eq!(
664            ChunkCapacity::from(0..).fits(Characters.size(chunk)),
665            Ordering::Equal
666        );
667        assert_eq!(
668            ChunkCapacity::from(5..).fits(Characters.size(chunk)),
669            Ordering::Equal
670        );
671        assert_eq!(
672            ChunkCapacity::from(6..).fits(Characters.size(chunk)),
673            Ordering::Less
674        );
675    }
676
677    #[test]
678    fn check_chunk_capacity_for_range_full() {
679        let chunk = "12345";
680
681        assert_eq!(
682            ChunkCapacity::from(..).fits(Characters.size(chunk)),
683            Ordering::Equal
684        );
685    }
686
687    #[test]
688    fn check_chunk_capacity_for_range_inclusive() {
689        let chunk = "12345";
690
691        assert_eq!(
692            ChunkCapacity::from(0..=4).fits(Characters.size(chunk)),
693            Ordering::Greater
694        );
695        assert_eq!(
696            ChunkCapacity::from(5..=6).fits(Characters.size(chunk)),
697            Ordering::Equal
698        );
699        assert_eq!(
700            ChunkCapacity::from(4..=5).fits(Characters.size(chunk)),
701            Ordering::Equal
702        );
703        assert_eq!(
704            ChunkCapacity::from(6..=100).fits(Characters.size(chunk)),
705            Ordering::Less
706        );
707    }
708
709    #[test]
710    fn check_chunk_capacity_for_range_to() {
711        let chunk = "12345";
712
713        assert_eq!(
714            ChunkCapacity::from(..0).fits(Characters.size(chunk)),
715            Ordering::Greater
716        );
717        assert_eq!(
718            ChunkCapacity::from(..5).fits(Characters.size(chunk)),
719            Ordering::Greater
720        );
721        assert_eq!(
722            ChunkCapacity::from(..6).fits(Characters.size(chunk)),
723            Ordering::Equal
724        );
725    }
726
727    #[test]
728    fn check_chunk_capacity_for_range_to_inclusive() {
729        let chunk = "12345";
730
731        assert_eq!(
732            ChunkCapacity::from(..=4).fits(Characters.size(chunk)),
733            Ordering::Greater
734        );
735        assert_eq!(
736            ChunkCapacity::from(..=5).fits(Characters.size(chunk)),
737            Ordering::Equal
738        );
739        assert_eq!(
740            ChunkCapacity::from(..=6).fits(Characters.size(chunk)),
741            Ordering::Equal
742        );
743    }
744
745    #[derive(Default)]
746    struct CountingSizer {
747        calls: AtomicUsize,
748    }
749
750    impl ChunkSizer for CountingSizer {
751        // Return character version, but count calls
752        fn size(&self, chunk: &str) -> usize {
753            self.calls.fetch_add(1, atomic::Ordering::SeqCst);
754            Characters.size(chunk)
755        }
756    }
757
758    #[test]
759    fn memoized_sizer_only_calculates_once_per_text() {
760        let sizer = CountingSizer::default();
761        let mut memoized_sizer = MemoizedChunkSizer::new(&sizer);
762        let text = "1234567890";
763        for _ in 0..10 {
764            memoized_sizer.chunk_size(0, text, Trim::All).unwrap();
765        }
766
767        assert_eq!(memoized_sizer.sizer.calls.load(atomic::Ordering::SeqCst), 1);
768    }
769
770    #[test]
771    fn memoized_sizer_calculates_once_per_different_text() {
772        let sizer = CountingSizer::default();
773        let mut memoized_sizer = MemoizedChunkSizer::new(&sizer);
774        let text = "1234567890";
775        for i in 0..10 {
776            memoized_sizer
777                .chunk_size(0, text.get(0..i).unwrap(), Trim::All)
778                .unwrap();
779        }
780
781        assert_eq!(
782            memoized_sizer.sizer.calls.load(atomic::Ordering::SeqCst),
783            10
784        );
785    }
786
787    #[test]
788    fn can_clear_cache_on_memoized_sizer() {
789        let sizer = CountingSizer::default();
790        let mut memoized_sizer = MemoizedChunkSizer::new(&sizer);
791        let text = "1234567890";
792        for _ in 0..10 {
793            memoized_sizer.chunk_size(0, text, Trim::All).unwrap();
794            memoized_sizer.clear_cache();
795        }
796
797        assert_eq!(
798            memoized_sizer.sizer.calls.load(atomic::Ordering::SeqCst),
799            10
800        );
801    }
802
803    #[test]
804    fn memoized_sizer_caches_successes_but_not_errors() {
805        struct RetrySizer(AtomicUsize);
806
807        impl TryChunkSizer for RetrySizer {
808            type Error = &'static str;
809
810            fn try_size(&self, chunk: &str) -> Result<usize, Self::Error> {
811                if self.0.fetch_add(1, atomic::Ordering::SeqCst) == 0 {
812                    Err("first call failed")
813                } else {
814                    Ok(Characters.size(chunk))
815                }
816            }
817        }
818
819        let sizer = RetrySizer(AtomicUsize::new(0));
820        let mut memoized = MemoizedChunkSizer::new(&sizer);
821        assert_eq!(
822            memoized.chunk_size(0, "abc", Trim::All),
823            Err("first call failed")
824        );
825        assert!(memoized.size_cache.is_empty());
826        assert_eq!(memoized.chunk_size(0, "abc", Trim::All), Ok(3));
827        assert_eq!(memoized.chunk_size(0, "abc", Trim::All), Ok(3));
828        assert_eq!(sizer.0.load(atomic::Ordering::SeqCst), 2);
829    }
830
831    #[test]
832    fn boundary_probe_is_skipped_when_probe_start_is_chunk_end() {
833        let sizer = CountingSizer::default();
834        let mut memoized_sizer = MemoizedChunkSizer::new(&sizer);
835        let mut lower_boundaries = [usize::MAX].into_iter();
836
837        let fits = memoized_sizer.chunk_fits_with_boundaries(
838            0,
839            "12345678901",
840            &ChunkCapacity::new(10),
841            &mut lower_boundaries,
842            Trim::All,
843        );
844
845        assert_eq!(fits, Ok(Ordering::Greater));
846        assert_eq!(sizer.calls.load(atomic::Ordering::SeqCst), 1);
847        assert_eq!(lower_boundaries.next(), Some(usize::MAX));
848    }
849
850    #[test]
851    fn boundary_probe_stops_when_next_boundary_exceeds_chunk_end() {
852        let sizer = CountingSizer::default();
853        let mut memoized_sizer = MemoizedChunkSizer::new(&sizer);
854        let chunk = "1234567890".repeat(9);
855        let mut lower_boundaries = [usize::MAX].into_iter();
856
857        let fits = memoized_sizer.chunk_fits_with_boundaries(
858            0,
859            &chunk,
860            &ChunkCapacity::new(10),
861            &mut lower_boundaries,
862            Trim::All,
863        );
864
865        assert_eq!(fits, Ok(Ordering::Greater));
866        assert_eq!(sizer.calls.load(atomic::Ordering::SeqCst), 1);
867        assert_eq!(lower_boundaries.next(), None);
868    }
869
870    #[test]
871    fn basic_chunk_config() {
872        let config = ChunkConfig::new(10);
873        assert_eq!(config.capacity, 10.into());
874        assert_eq!(config.sizer, Characters);
875        assert!(config.trim());
876    }
877
878    #[test]
879    fn disable_trimming() {
880        let config = ChunkConfig::new(10).with_trim(false);
881        assert!(!config.trim());
882    }
883
884    #[test]
885    fn new_sizer() {
886        #[derive(Debug, PartialEq)]
887        struct BasicSizer;
888
889        impl ChunkSizer for BasicSizer {
890            fn size(&self, _chunk: &str) -> usize {
891                unimplemented!()
892            }
893        }
894
895        let config = ChunkConfig::new(10).with_sizer(BasicSizer);
896        assert_eq!(config.capacity, 10.into());
897        assert_eq!(config.sizer, BasicSizer);
898        assert!(config.trim());
899    }
900
901    #[test]
902    fn chunk_capacity_max_and_desired_equal() {
903        let capacity = ChunkCapacity::new(10);
904        assert_eq!(capacity.desired(), 10);
905        assert_eq!(capacity.max(), 10);
906    }
907
908    #[test]
909    fn chunk_capacity_can_adjust_max() {
910        let capacity = ChunkCapacity::new(10).with_max(20).unwrap();
911        assert_eq!(capacity.desired(), 10);
912        assert_eq!(capacity.max(), 20);
913    }
914
915    #[test]
916    fn chunk_capacity_max_cant_be_less_than_desired() {
917        let capacity = ChunkCapacity::new(10);
918        let err = capacity.with_max(5).unwrap_err();
919        assert_eq!(
920            err.to_string(),
921            "Max chunk size must be greater than or equal to the desired chunk size"
922        );
923        assert_eq!(capacity.desired(), 10);
924        assert_eq!(capacity.max(), 10);
925    }
926
927    #[test]
928    fn set_chunk_overlap() {
929        let config = ChunkConfig::new(10).with_overlap(5).unwrap();
930        assert_eq!(config.overlap(), 5);
931    }
932
933    #[test]
934    fn cant_set_overlap_larger_than_capacity() {
935        let chunk_config = ChunkConfig::new(5);
936        let err = chunk_config.with_overlap(10).unwrap_err();
937        assert_eq!(
938            err.to_string(),
939            "The overlap is larger than or equal to the desired chunk capacity"
940        );
941    }
942
943    #[test]
944    fn cant_set_overlap_larger_than_desired() {
945        let chunk_config = ChunkConfig::new(5..15);
946        let err = chunk_config.with_overlap(10).unwrap_err();
947        assert_eq!(
948            err.to_string(),
949            "The overlap is larger than or equal to the desired chunk capacity"
950        );
951    }
952
953    #[test]
954    fn chunk_size_reference() {
955        let config = ChunkConfig::new(1).with_sizer(&Characters);
956        config.sizer().size("chunk");
957    }
958
959    #[test]
960    fn chunk_size_cow() {
961        let sizer: Cow<'_, Characters> = Cow::Owned(Characters);
962        let config = ChunkConfig::new(1).with_sizer(sizer);
963        config.sizer().size("chunk");
964
965        let sizer = Cow::Borrowed(&Characters);
966        let config = ChunkConfig::new(1).with_sizer(sizer);
967        config.sizer().size("chunk");
968    }
969
970    #[test]
971    fn chunk_size_arc() {
972        let sizer = Arc::new(Characters);
973        let config = ChunkConfig::new(1).with_sizer(sizer);
974        config.sizer().size("chunk");
975    }
976
977    #[test]
978    fn chunk_size_ref() {
979        let sizer = RefCell::new(Characters);
980        let config = ChunkConfig::new(1).with_sizer(sizer.borrow());
981        config.sizer().size("chunk");
982    }
983
984    #[test]
985    fn chunk_size_ref_mut() {
986        let sizer = RefCell::new(Characters);
987        let config = ChunkConfig::new(1).with_sizer(sizer.borrow_mut());
988        config.sizer().size("chunk");
989    }
990
991    #[test]
992    fn chunk_size_box() {
993        let sizer = Box::new(Characters);
994        let config = ChunkConfig::new(1).with_sizer(sizer);
995        config.sizer().size("chunk");
996    }
997
998    #[test]
999    fn chunk_size_rc() {
1000        let sizer = Rc::new(Characters);
1001        let config = ChunkConfig::new(1).with_sizer(sizer);
1002        config.sizer().size("chunk");
1003    }
1004}
text_splitter/chunk_size.rs

text_splitter/
chunk_size.rs
