Skip to main content

harper_core/
document.rs

1use std::cmp::Ordering;
2use std::collections::VecDeque;
3use std::fmt::Display;
4
5use harper_brill::{Chunker, Tagger, brill_tagger, burn_chunker};
6use itertools::Itertools;
7
8use crate::expr::{Expr, ExprExt, FirstMatchOf, Repeating, SequenceExpr};
9use crate::parsers::{Markdown, MarkdownOptions, Parser, PlainEnglish};
10use crate::punctuation::Punctuation;
11use crate::spell::{Dictionary, FstDictionary};
12use crate::vec_ext::VecExt;
13use crate::{CharStringExt, FatStringToken, FatToken, Lrc, Token, TokenKind, TokenStringExt};
14use crate::{OrdinalSuffix, Span};
15
16/// A document containing some amount of lexed and parsed English text.
17#[derive(Debug, Clone)]
18pub struct Document {
19    source: Lrc<[char]>,
20    tokens: Vec<Token>,
21}
22
23impl Default for Document {
24    fn default() -> Self {
25        Self::new("", &PlainEnglish, &FstDictionary::curated())
26    }
27}
28
29impl Document {
30    /// Locate all the tokens that intersect a provided span.
31    ///
32    /// Desperately needs optimization.
33    pub fn token_indices_intersecting(&self, span: Span<char>) -> Vec<usize> {
34        self.tokens()
35            .enumerate()
36            .filter_map(|(idx, tok)| tok.span.overlaps_with(span).then_some(idx))
37            .collect()
38    }
39
40    /// Locate all the tokens that intersect a provided span and convert them to [`FatToken`]s.
41    ///
42    /// Desperately needs optimization.
43    pub fn fat_tokens_intersecting(&self, span: Span<char>) -> Vec<FatToken> {
44        let indices = self.token_indices_intersecting(span);
45
46        indices
47            .into_iter()
48            .map(|i| self.tokens[i].to_fat(&self.source))
49            .collect()
50    }
51
52    /// Lexes and parses text to produce a document using a provided language
53    /// parser and dictionary.
54    pub fn new(text: &str, parser: &impl Parser, dictionary: &impl Dictionary) -> Self {
55        let source: Lrc<_> = text.chars().collect();
56
57        Self::new_from_chars(source, parser, dictionary)
58    }
59
60    /// Lexes and parses text to produce a document using a provided language
61    /// parser and the included curated dictionary.
62    pub fn new_curated(text: &str, parser: &impl Parser) -> Self {
63        let source: Lrc<_> = text.chars().collect();
64
65        Self::new_from_chars(source, parser, &FstDictionary::curated())
66    }
67
68    /// Lexes and parses text to produce a document using a provided language
69    /// parser and dictionary.
70    pub fn new_from_chars(
71        source: Lrc<[char]>,
72        parser: &impl Parser,
73        dictionary: &impl Dictionary,
74    ) -> Self {
75        let tokens = parser.parse(&source);
76
77        let mut document = Self { source, tokens };
78        document.parse(dictionary);
79
80        document
81    }
82
83    /// Create a new document from character data using the built-in [`PlainEnglish`]
84    /// parser and curated dictionary. This avoids string-to-char conversions.
85    pub fn new_plain_english_curated_chars(source: &[char]) -> Self {
86        Self::new_from_chars(Lrc::from(source), &PlainEnglish, &FstDictionary::curated())
87    }
88
89    /// Parse text to produce a document using the built-in [`PlainEnglish`]
90    /// parser and curated dictionary.
91    pub fn new_plain_english_curated(text: &str) -> Self {
92        Self::new(text, &PlainEnglish, &FstDictionary::curated())
93    }
94
95    /// Create a new document simply by tokenizing the provided input and applying fix-ups. The
96    /// contained words will not contain any metadata.
97    ///
98    /// This avoids running potentially expensive metadata generation code, so this is more
99    /// efficient if you don't need that information.
100    pub(crate) fn new_basic_tokenize(text: &str, parser: &impl Parser) -> Self {
101        let source: Lrc<_> = text.chars().collect();
102        let tokens = parser.parse(&source);
103        let mut document = Self { source, tokens };
104        document.apply_fixups();
105        document
106    }
107
108    /// Parse text to produce a document using the built-in [`PlainEnglish`]
109    /// parser and a provided dictionary.
110    pub fn new_plain_english(text: &str, dictionary: &impl Dictionary) -> Self {
111        Self::new(text, &PlainEnglish, dictionary)
112    }
113
114    /// Parse text to produce a document using the built-in [`Markdown`] parser
115    /// and curated dictionary.
116    pub fn new_markdown_curated(text: &str, markdown_options: MarkdownOptions) -> Self {
117        Self::new(
118            text,
119            &Markdown::new(markdown_options),
120            &FstDictionary::curated(),
121        )
122    }
123
124    /// Create a new document from character data using the built-in [`Markdown`] parser
125    /// and curated dictionary. This avoids string-to-char conversions.
126    pub fn new_markdown_default_curated_chars(chars: &[char]) -> Self {
127        Self::new_from_chars(
128            chars.to_vec().into(),
129            &Markdown::default(),
130            &FstDictionary::curated(),
131        )
132    }
133
134    /// Parse text to produce a document using the built-in [`Markdown`] parser
135    /// and curated dictionary with the default Markdown configuration.
136    pub fn new_markdown_default_curated(text: &str) -> Self {
137        Self::new_markdown_curated(text, MarkdownOptions::default())
138    }
139
140    /// Parse text to produce a document using the built-in [`PlainEnglish`]
141    /// parser and the curated dictionary.
142    pub fn new_markdown(
143        text: &str,
144        markdown_options: MarkdownOptions,
145        dictionary: &impl Dictionary,
146    ) -> Self {
147        Self::new(text, &Markdown::new(markdown_options), dictionary)
148    }
149
150    /// Parse text to produce a document using the built-in [`PlainEnglish`]
151    /// parser and the curated dictionary with the default Markdown configuration.
152    pub fn new_markdown_default(text: &str, dictionary: &impl Dictionary) -> Self {
153        Self::new_markdown(text, MarkdownOptions::default(), dictionary)
154    }
155
156    fn apply_fixups(&mut self) {
157        self.condense_spaces();
158        self.condense_newlines();
159        self.newlines_to_breaks();
160        self.condense_dotted_initialisms();
161        self.condense_number_suffixes();
162        self.condense_ellipsis();
163        self.condense_dotted_truncations();
164        self.condense_common_top_level_domains();
165        self.condense_filename_extensions();
166        self.condense_tldr();
167        self.condense_ampersand_pairs();
168        self.condense_slash_pairs();
169        self.match_quotes();
170    }
171
172    /// Re-parse important language constructs.
173    ///
174    /// Should be run after every change to the underlying [`Self::source`].
175    fn parse(&mut self, dictionary: &impl Dictionary) {
176        self.apply_fixups();
177
178        let chunker = burn_chunker();
179        let tagger = brill_tagger();
180
181        for sent in self.tokens.iter_sentences_mut() {
182            let token_strings: Vec<_> = sent
183                .iter()
184                .filter(|t| !t.kind.is_whitespace())
185                .map(|t| t.get_str(&self.source))
186                .collect();
187
188            let token_tags = tagger.tag_sentence(&token_strings);
189            let np_flags = chunker.chunk_sentence(&token_strings, &token_tags);
190
191            // Annotate DictWord metadata
192            let word_sources: Vec<_> = sent
193                .iter()
194                .filter(|t| matches!(t.kind, TokenKind::Word(_)))
195                .map(|t| t.get_ch(&self.source))
196                .collect();
197
198            let mut ti = 0; // Index for token_tags/np_flags (all non-whitespace tokens)
199            let mut wi = 0; // Index for word_sources (only word tokens)
200            for token in sent.iter_mut() {
201                if let TokenKind::Word(meta) = &mut token.kind {
202                    let word_source = word_sources[wi];
203                    let mut found_meta = dictionary
204                        .get_word_metadata(word_source)
205                        .map(|c| c.into_owned());
206
207                    if let Some(inner) = &mut found_meta {
208                        inner.pos_tag = token_tags[ti].or_else(|| inner.infer_pos_tag());
209                        inner.np_member = Some(np_flags[ti]);
210                    }
211
212                    *meta = found_meta;
213                    ti += 1;
214                    wi += 1;
215                } else if !token.kind.is_whitespace() {
216                    ti += 1;
217                }
218            }
219        }
220    }
221
222    /// Convert all sets of newlines greater than 2 to paragraph breaks.
223    fn newlines_to_breaks(&mut self) {
224        for token in &mut self.tokens {
225            if let TokenKind::Newline(n) = token.kind
226                && n >= 2
227            {
228                token.kind = TokenKind::ParagraphBreak;
229            }
230        }
231    }
232
233    /// Given a list of indices, this function removes the subsequent
234    /// `stretch_len - 1` elements after each index.
235    ///
236    /// Will extend token spans to include removed elements.
237    /// Assumes condensed tokens are contiguous in source text.
238    fn condense_indices(&mut self, indices: &[usize], stretch_len: usize) {
239        // Update spans
240        for idx in indices {
241            let end_tok = self.tokens[idx + stretch_len - 1].clone();
242            let start_tok = &mut self.tokens[*idx];
243
244            start_tok.span.end = end_tok.span.end;
245        }
246
247        // Trim
248        let old = self.tokens.clone();
249        self.tokens.clear();
250
251        // Keep first chunk.
252        self.tokens
253            .extend_from_slice(&old[0..indices.first().copied().unwrap_or(indices.len())]);
254
255        let mut iter = indices.iter().peekable();
256
257        while let (Some(a_idx), b) = (iter.next(), iter.peek()) {
258            self.tokens.push(old[*a_idx].clone());
259
260            if let Some(b_idx) = b {
261                self.tokens
262                    .extend_from_slice(&old[a_idx + stretch_len..**b_idx]);
263            }
264        }
265
266        // Keep last chunk.
267        self.tokens.extend_from_slice(
268            &old[indices
269                .last()
270                .map(|v| v + stretch_len)
271                .unwrap_or(indices.len())..],
272        );
273    }
274
275    pub fn get_token_at_char_index(&self, char_index: usize) -> Option<&Token> {
276        let index = self
277            .tokens
278            .binary_search_by(|t| {
279                if t.span.overlaps_with(Span::new_with_len(char_index, 1)) {
280                    Ordering::Equal
281                } else {
282                    t.span.start.cmp(&char_index)
283                }
284            })
285            .ok()?;
286
287        Some(&self.tokens[index])
288    }
289
290    /// Defensively attempt to grab a specific token.
291    pub fn get_token(&self, index: usize) -> Option<&Token> {
292        self.tokens.get(index)
293    }
294
295    /// Get a token at a signed offset from a base index, or None if out of bounds.
296    pub fn get_token_offset(&self, base: usize, offset: isize) -> Option<&Token> {
297        match base.checked_add_signed(offset) {
298            None => None,
299            Some(idx) => self.get_token(idx),
300        }
301    }
302
303    /// Get an iterator over all the tokens contained in the document.
304    pub fn tokens(&self) -> impl Iterator<Item = &Token> + '_ {
305        self.tokens.iter()
306    }
307
308    pub fn iter_nominal_phrases(&self) -> impl Iterator<Item = &[Token]> {
309        fn is_np_member(t: &Token) -> bool {
310            t.kind
311                .as_word()
312                .and_then(|x| x.as_ref())
313                .and_then(|w| w.np_member)
314                .unwrap_or(false)
315        }
316
317        fn trim(slice: &[Token]) -> &[Token] {
318            let mut start = 0;
319            let mut end = slice.len();
320            while start < end && slice[start].kind.is_whitespace() {
321                start += 1;
322            }
323            while end > start && slice[end - 1].kind.is_whitespace() {
324                end -= 1;
325            }
326            &slice[start..end]
327        }
328
329        self.tokens
330            .as_slice()
331            .split(|t| !(is_np_member(t) || t.kind.is_whitespace()))
332            .filter_map(|s| {
333                let s = trim(s);
334                if s.iter().any(is_np_member) {
335                    Some(s)
336                } else {
337                    None
338                }
339            })
340    }
341
342    /// Get an iterator over all the tokens contained in the document.
343    pub fn fat_tokens(&self) -> impl Iterator<Item = FatToken> + '_ {
344        self.tokens().map(|token| token.to_fat(&self.source))
345    }
346
347    /// Get the next or previous word token relative to a base index, if separated by whitespace.
348    /// Returns None if the next/previous token is not a word or does not exist.
349    pub fn get_next_word_from_offset(&self, base: usize, offset: isize) -> Option<&Token> {
350        // Look for whitespace at the expected offset
351        if !self.get_token_offset(base, offset)?.kind.is_whitespace() {
352            return None;
353        }
354        // Now look beyond the whitespace for a word token
355        let word_token = self.get_token_offset(base, offset + offset.signum());
356        let word_token = word_token?;
357        word_token.kind.is_word().then_some(word_token)
358    }
359
360    /// Get an iterator over all the tokens contained in the document.
361    pub fn fat_string_tokens(&self) -> impl Iterator<Item = FatStringToken> + '_ {
362        self.fat_tokens().map(|t| t.into())
363    }
364
365    pub fn get_span_content(&self, span: &Span<char>) -> &[char] {
366        span.get_content(&self.source)
367    }
368
369    pub fn get_span_content_str(&self, span: &Span<char>) -> String {
370        String::from_iter(self.get_span_content(span))
371    }
372
373    pub fn get_full_string(&self) -> String {
374        self.get_span_content_str(&Span::new(0, self.source.len()))
375    }
376
377    pub fn get_full_content(&self) -> &[char] {
378        &self.source
379    }
380
381    pub fn get_source(&self) -> &[char] {
382        &self.source
383    }
384
385    pub fn get_tokens(&self) -> &[Token] {
386        &self.tokens
387    }
388
389    /// Searches for quotation marks and fills the
390    /// [`Punctuation::Quote::twin_loc`] field. This is on a best-effort
391    /// basis.
392    ///
393    /// Current algorithm is based on https://leancrew.com/all-this/2025/03/a-mac-smart-quote-curiosity
394    fn match_quotes(&mut self) {
395        let mut pg_indices: Vec<_> = vec![0];
396        pg_indices.extend(self.iter_paragraph_break_indices());
397        pg_indices.push(self.tokens.len());
398
399        // Avoid allocation in loop
400        let mut quote_indices = Vec::new();
401        let mut open_quote_indices = Vec::new();
402
403        for (start, end) in pg_indices.into_iter().tuple_windows() {
404            let pg = &mut self.tokens[start..end];
405
406            quote_indices.clear();
407            quote_indices.extend(pg.iter_quote_indices());
408            open_quote_indices.clear();
409
410            // Find open quotes first.
411            for quote in &quote_indices {
412                let is_open = *quote == 0
413                    || pg[0..*quote].iter_word_likes().next().is_none()
414                    || pg[quote - 1].kind.is_whitespace()
415                    || matches!(
416                        pg[quote - 1].kind.as_punctuation(),
417                        Some(Punctuation::LessThan)
418                            | Some(Punctuation::OpenRound)
419                            | Some(Punctuation::OpenSquare)
420                            | Some(Punctuation::OpenCurly)
421                            | Some(Punctuation::EmDash)
422                            | Some(Punctuation::EnDash)
423                            | Some(Punctuation::Apostrophe)
424                            | Some(Punctuation::OpenSingle)
425                    );
426
427                if is_open {
428                    open_quote_indices.push(*quote);
429                }
430            }
431
432            while let Some(open_idx) = open_quote_indices.pop() {
433                let Some(close_idx) = pg[open_idx + 1..].iter_quote_indices().next() else {
434                    continue;
435                };
436
437                if pg[close_idx + open_idx + 1]
438                    .kind
439                    .as_quote()
440                    .unwrap()
441                    .twin_loc
442                    .is_some()
443                {
444                    continue;
445                }
446
447                pg[open_idx].kind.as_mut_quote().unwrap().twin_loc =
448                    Some(close_idx + open_idx + start + 1);
449                pg[close_idx + open_idx + 1]
450                    .kind
451                    .as_mut_quote()
452                    .unwrap()
453                    .twin_loc = Some(open_idx + start);
454            }
455        }
456    }
457
458    /// Searches for number suffixes and condenses them down into single tokens
459    fn condense_number_suffixes(&mut self) {
460        if self.tokens.len() < 2 {
461            return;
462        }
463
464        let mut replace_starts = Vec::new();
465
466        for idx in 0..self.tokens.len() - 1 {
467            let b = &self.tokens[idx + 1];
468            let a = &self.tokens[idx];
469
470            // TODO: Allow spaces between `a` and `b`
471
472            if let (TokenKind::Number(..), TokenKind::Word(..)) = (&a.kind, &b.kind)
473                && let Some(found_suffix) =
474                    OrdinalSuffix::from_chars(self.get_span_content(&b.span))
475            {
476                self.tokens[idx].kind.as_mut_number().unwrap().suffix = Some(found_suffix);
477                replace_starts.push(idx);
478            }
479        }
480
481        self.condense_indices(&replace_starts, 2);
482    }
483
484    /// Searches for multiple sequential space tokens and condenses them down
485    /// into one.
486    fn condense_spaces(&mut self) {
487        let mut cursor = 0;
488        let copy = self.tokens.clone();
489
490        let mut remove_these = VecDeque::new();
491
492        while cursor < self.tokens.len() {
493            // Locate a stretch of one or more newline tokens.
494            let start_tok = &mut self.tokens[cursor];
495
496            if let TokenKind::Space(start_count) = &mut start_tok.kind {
497                loop {
498                    cursor += 1;
499
500                    if cursor >= copy.len() {
501                        break;
502                    }
503
504                    let child_tok = &copy[cursor];
505
506                    // Only condense adjacent spans
507                    if start_tok.span.end != child_tok.span.start {
508                        break;
509                    }
510
511                    if let TokenKind::Space(n) = child_tok.kind {
512                        *start_count += n;
513                        start_tok.span.end = child_tok.span.end;
514                        remove_these.push_back(cursor);
515                        cursor += 1;
516                    } else {
517                        break;
518                    };
519                }
520            }
521
522            cursor += 1;
523        }
524
525        self.tokens.remove_indices(remove_these);
526    }
527
528    thread_local! {
529        static DOTTED_TRUNCATION_EXPR: Lrc<FirstMatchOf> = Document::uncached_dotted_truncation_expr();
530    }
531
532    fn uncached_dotted_truncation_expr() -> Lrc<FirstMatchOf> {
533        Lrc::new(FirstMatchOf::new(vec![
534            Box::new(SequenceExpr::word_set(&["esp", "etc", "vs"]).then_period()),
535            Box::new(
536                SequenceExpr::aco("et")
537                    .then_whitespace()
538                    .t_aco("al")
539                    .then_period(),
540            ),
541        ]))
542    }
543
544    /// Assumes that the first matched token is the canonical one to be condensed into.
545    /// Takes a callback that can be used to retroactively edit the canonical token afterwards.
546    fn condense_expr<F>(&mut self, expr: &impl Expr, edit: F)
547    where
548        F: Fn(&mut Token),
549    {
550        let matches = expr.iter_matches_in_doc(self).collect::<Vec<_>>();
551
552        let mut remove_indices = VecDeque::with_capacity(matches.len());
553
554        for m in matches {
555            remove_indices.extend(m.start + 1..m.end);
556            self.tokens[m.start].span = self.tokens[m.into_iter()].span().unwrap();
557            edit(&mut self.tokens[m.start]);
558        }
559
560        self.tokens.remove_indices(remove_indices);
561    }
562
563    fn condense_dotted_truncations(&mut self) {
564        self.condense_expr(&Self::DOTTED_TRUNCATION_EXPR.with(|v| v.clone()), |_| {})
565    }
566
567    /// Searches for multiple sequential newline tokens and condenses them down
568    /// into one.
569    fn condense_newlines(&mut self) {
570        let mut cursor = 0;
571        let copy = self.tokens.clone();
572
573        let mut remove_these = VecDeque::new();
574
575        while cursor < self.tokens.len() {
576            // Locate a stretch of one or more newline tokens.
577            let start_tok = &mut self.tokens[cursor];
578
579            if let TokenKind::Newline(start_count) = &mut start_tok.kind {
580                loop {
581                    cursor += 1;
582
583                    if cursor >= copy.len() {
584                        break;
585                    }
586
587                    let child_tok = &copy[cursor];
588                    if let TokenKind::Newline(n) = child_tok.kind {
589                        *start_count += n;
590                        start_tok.span.end = child_tok.span.end;
591                        remove_these.push_back(cursor);
592                        cursor += 1;
593                    } else {
594                        break;
595                    };
596                }
597            }
598
599            cursor += 1;
600        }
601
602        self.tokens.remove_indices(remove_these);
603    }
604
605    /// Condenses words like "i.e.", "e.g." and "N.S.A." down to single words
606    /// using a state machine.
607    fn condense_dotted_initialisms(&mut self) {
608        if self.tokens.len() < 2 {
609            return;
610        }
611
612        let mut to_remove = VecDeque::new();
613
614        let mut cursor = 1;
615
616        let mut initialism_start = None;
617
618        loop {
619            let a = &self.tokens[cursor - 1];
620            let b = &self.tokens[cursor];
621
622            let is_initialism_chunk = a.kind.is_word() && a.span.len() == 1 && b.kind.is_period();
623
624            if is_initialism_chunk {
625                if initialism_start.is_none() {
626                    initialism_start = Some(cursor - 1);
627                } else {
628                    to_remove.push_back(cursor - 1);
629                }
630
631                to_remove.push_back(cursor);
632                cursor += 1;
633            } else {
634                if let Some(start) = initialism_start {
635                    let end = self.tokens[cursor - 2].span.end;
636                    let start_tok: &mut Token = &mut self.tokens[start];
637                    start_tok.span.end = end;
638                }
639
640                initialism_start = None;
641            }
642
643            cursor += 1;
644
645            if cursor >= self.tokens.len() - 1 {
646                break;
647            }
648        }
649
650        self.tokens.remove_indices(to_remove);
651    }
652
653    /// Condenses likely filename extensions down to single tokens.
654    fn condense_filename_extensions(&mut self) {
655        if self.tokens.len() < 2 {
656            return;
657        }
658
659        let mut to_remove = VecDeque::new();
660
661        let mut cursor = 1;
662
663        let mut ext_start = None;
664
665        loop {
666            // left context, dot, extension, right context
667            let l = self.get_token_offset(cursor, -2);
668            let d = &self.tokens[cursor - 1];
669            let x = &self.tokens[cursor];
670            let r = self.get_token_offset(cursor, 1);
671
672            let is_ext_chunk = d.kind.is_period()
673                && x.kind.is_word()
674                && x.span.len() <= 3
675                && ((l.is_none_or(|t| t.kind.is_whitespace())
676                    && r.is_none_or(|t| t.kind.is_whitespace()))
677                    || (l.is_some_and(|t| t.kind.is_open_round())
678                        && r.is_some_and(|t| t.kind.is_close_round())))
679                && {
680                    let ext_chars = x.get_ch(&self.source);
681                    ext_chars.iter().all(|c| c.is_ascii_lowercase())
682                        || ext_chars.iter().all(|c| c.is_ascii_uppercase())
683                };
684
685            if is_ext_chunk {
686                if ext_start.is_none() {
687                    ext_start = Some(cursor - 1);
688                    self.tokens[cursor - 1].kind = TokenKind::Unlintable;
689                } else {
690                    to_remove.push_back(cursor - 1);
691                }
692
693                to_remove.push_back(cursor);
694                cursor += 1;
695            } else {
696                if let Some(start) = ext_start {
697                    let end = self.tokens[cursor - 2].span.end;
698                    let start_tok: &mut Token = &mut self.tokens[start];
699                    start_tok.span.end = end;
700                }
701
702                ext_start = None;
703            }
704
705            cursor += 1;
706
707            if cursor >= self.tokens.len() {
708                break;
709            }
710        }
711
712        self.tokens.remove_indices(to_remove);
713    }
714
715    /// Condenses common top-level domains (for example: `.blog`, `.com`) down to single tokens.
716    fn condense_common_top_level_domains(&mut self) {
717        const COMMON_TOP_LEVEL_DOMAINS: &[&str; 106] = &[
718            "ai", "app", "blog", "co", "com", "dev", "edu", "gov", "info", "io", "me", "mil",
719            "net", "org", "shop", "tech", "uk", "us", "xyz", "jp", "de", "fr", "br", "it", "ru",
720            "es", "pl", "ca", "au", "cn", "in", "nl", "eu", "ch", "id", "at", "kr", "cz", "mx",
721            "be", "tv", "se", "tr", "tw", "al", "ua", "ir", "vn", "cl", "sk", "ly", "cc", "to",
722            "no", "fi", "pt", "dk", "ar", "hu", "tk", "gr", "il", "news", "ro", "my", "biz", "ie",
723            "za", "nz", "sg", "ee", "th", "pe", "bg", "hk", "rs", "lt", "link", "ph", "club", "si",
724            "site", "mobi", "by", "cat", "wiki", "la", "ga", "xxx", "cf", "hr", "ng", "jobs",
725            "online", "kz", "ug", "gq", "ae", "is", "lv", "pro", "fm", "tips", "ms", "sa", "int",
726        ];
727
728        if self.tokens.len() < 2 {
729            return;
730        }
731
732        let mut to_remove = VecDeque::new();
733        for cursor in 1..self.tokens.len() {
734            // left context, dot, tld, right context
735            let l = self.get_token_offset(cursor, -2);
736            let d = &self.tokens[cursor - 1];
737            let tld = &self.tokens[cursor];
738            let r = self.get_token_offset(cursor, 1);
739
740            let is_tld_chunk = d.kind.is_period()
741                && tld.kind.is_word()
742                && tld
743                    .get_ch(&self.source)
744                    .iter()
745                    .all(|c| c.is_ascii_alphabetic())
746                && tld
747                    .get_ch(&self.source)
748                    .eq_any_ignore_ascii_case_str(COMMON_TOP_LEVEL_DOMAINS)
749                && ((l.is_none_or(|t| t.kind.is_whitespace())
750                    && r.is_none_or(|t| t.kind.is_whitespace()))
751                    || (l.is_some_and(|t| t.kind.is_open_round())
752                        && r.is_some_and(|t| t.kind.is_close_round())));
753
754            if is_tld_chunk {
755                self.tokens[cursor - 1].kind = TokenKind::Unlintable;
756                self.tokens[cursor - 1].span.end = self.tokens[cursor].span.end;
757                to_remove.push_back(cursor);
758            }
759        }
760
761        self.tokens.remove_indices(to_remove);
762    }
763
764    /// Condenses "tl;dr" down to a single word token.
765    fn condense_tldr(&mut self) {
766        if self.tokens.len() < 3 {
767            return;
768        }
769
770        let mut to_remove = VecDeque::new();
771        let mut cursor = 2;
772
773        loop {
774            let tl = &self.tokens[cursor - 2];
775            let simicolon = &self.tokens[cursor - 1];
776            let dr = &self.tokens[cursor];
777
778            let is_tldr_chunk = tl.kind.is_word()
779                && tl.span.len() == 2
780                && tl.get_ch(&self.source).eq_ch(&['t', 'l'])
781                && simicolon.kind.is_semicolon()
782                && dr.kind.is_word()
783                && dr.span.len() >= 2
784                && dr.span.len() <= 3
785                && dr
786                    .get_ch(&self.source)
787                    .eq_any_ignore_ascii_case_chars(&[&['d', 'r'], &['d', 'r', 's']]);
788
789            if is_tldr_chunk {
790                // Update the first token to be the full "tl;dr" as a word
791                self.tokens[cursor - 2].span = Span::new(
792                    self.tokens[cursor - 2].span.start,
793                    self.tokens[cursor].span.end,
794                );
795
796                // Mark the semicolon and "dr" tokens for removal
797                to_remove.push_back(cursor - 1);
798                to_remove.push_back(cursor);
799            }
800
801            // Skip ahead since we've processed these tokens
802            cursor += 1;
803
804            if cursor >= self.tokens.len() {
805                break;
806            }
807        }
808
809        // Remove the marked tokens in reverse order to maintain correct indices
810        self.tokens.remove_indices(to_remove);
811    }
812
813    /// Allows condensing of delimited pairs of tokens into a single token.
814    ///
815    /// # Arguments
816    ///
817    /// * `is_delimiter` - A function that returns `true` if the token is a delimiter.
818    /// * `valid_pairs` - A slice of tuples representing the valid pairs of tokens to condense.
819    ///
820    fn condense_delimited_pairs<F>(&mut self, is_delimiter: F, valid_pairs: &[(char, char)])
821    where
822        F: Fn(&TokenKind) -> bool,
823    {
824        if self.tokens.len() < 3 {
825            return;
826        }
827
828        let mut to_remove = VecDeque::new();
829        let mut cursor = 2;
830
831        loop {
832            let l1 = &self.tokens[cursor - 2];
833            let delim = &self.tokens[cursor - 1];
834            let l2 = &self.tokens[cursor];
835
836            let is_delimited_chunk = l1.kind.is_word()
837                && l1.span.len() == 1
838                && is_delimiter(&delim.kind)
839                && l2.kind.is_word()
840                && l2.span.len() == 1;
841
842            if is_delimited_chunk {
843                let (l1, l2) = (
844                    l1.get_ch(&self.source).first(),
845                    l2.get_ch(&self.source).first(),
846                );
847
848                let is_valid_pair = match (l1, l2) {
849                    (Some(l1), Some(l2)) => {
850                        let pair = (l1.to_ascii_lowercase(), l2.to_ascii_lowercase());
851                        valid_pairs.contains(&pair)
852                    }
853                    _ => false,
854                };
855
856                if is_valid_pair {
857                    self.tokens[cursor - 2].span = Span::new(
858                        self.tokens[cursor - 2].span.start,
859                        self.tokens[cursor].span.end,
860                    );
861                    to_remove.push_back(cursor - 1);
862                    to_remove.push_back(cursor);
863                }
864            }
865
866            cursor += 1;
867            if cursor >= self.tokens.len() {
868                break;
869            }
870        }
871
872        self.tokens.remove_indices(to_remove);
873    }
874
875    // Condenses "ampersand pairs" such as "R&D" or "Q&A" into single tokens.
876    fn condense_ampersand_pairs(&mut self) {
877        self.condense_delimited_pairs(
878            |kind| kind.is_ampersand(),
879            &[
880                ('b', 'b'), // bed & breakfast
881                ('b', 'w'), // black & white
882                ('g', 't'), // gin & tonic
883                ('k', 'r'), // Kernighan & Ritchie
884                ('q', 'a'), // question & answer
885                ('r', 'b'), // rhythm & blues
886                ('r', 'd'), // research & development
887                ('r', 'r'), // rest & relaxation
888                ('s', 'p'), // Standard & Poor's
889            ],
890        );
891    }
892
893    // Condenses "slash pairs" such as "I/O" into single tokens.
894    fn condense_slash_pairs(&mut self) {
895        self.condense_delimited_pairs(
896            |kind| kind.is_slash(),
897            &[
898                ('a', 'c'), // aircon; alternating current
899                ('b', 'w'), // black and white
900                ('c', 'o'), // care of
901                ('d', 'c'), // direct current
902                ('d', 'l'), // download
903                ('i', 'o'), // input/output
904                ('j', 'k'), // just kidding
905                ('n', 'a'), // not applicable
906                ('r', 'c'), // radio control
907                ('s', 'n'), // serial number
908                ('y', 'n'), // yes/no
909                ('y', 'o'), // years old
910            ],
911        );
912    }
913
914    fn uncached_ellipsis_pattern() -> Lrc<Repeating> {
915        let period = SequenceExpr::default().then_period();
916        Lrc::new(Repeating::new(Box::new(period), 2))
917    }
918
919    thread_local! {
920        static ELLIPSIS_EXPR: Lrc<Repeating> = Document::uncached_ellipsis_pattern();
921    }
922
923    fn condense_ellipsis(&mut self) {
924        let expr = Self::ELLIPSIS_EXPR.with(|v| v.clone());
925        self.condense_expr(&expr, |tok| {
926            tok.kind = TokenKind::Punctuation(Punctuation::Ellipsis)
927        });
928    }
929}
930
931impl TokenStringExt for Document {
932    fn tokens(&self) -> &[Token] {
933        &self.tokens
934    }
935
936    fn tokens_mut(&mut self) -> &mut [Token] {
937        &mut self.tokens
938    }
939}
940
941impl Display for Document {
942    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
943        for token in &self.tokens {
944            write!(f, "{}", self.get_span_content_str(&token.span))?;
945        }
946
947        Ok(())
948    }
949}
950
951#[cfg(test)]
952mod tests {
953    use itertools::Itertools;
954
955    use super::Document;
956    use crate::TokenStringExt;
957    use crate::{Span, parsers::MarkdownOptions};
958
959    fn assert_condensed_contractions(text: &str, final_tok_count: usize) {
960        let document = Document::new_plain_english_curated(text);
961
962        assert_eq!(document.tokens.len(), final_tok_count);
963
964        let document = Document::new_markdown_curated(text, MarkdownOptions::default());
965
966        assert_eq!(document.tokens.len(), final_tok_count);
967    }
968
969    #[test]
970    fn simple_contraction() {
971        assert_condensed_contractions("isn't", 1);
972    }
973
974    #[test]
975    fn simple_contraction2() {
976        assert_condensed_contractions("wasn't", 1);
977    }
978
979    #[test]
980    fn simple_contraction3() {
981        assert_condensed_contractions("There's", 1);
982    }
983
984    #[test]
985    fn simple_contraction4() {
986        assert_condensed_contractions("doesn't", 1);
987    }
988
989    #[test]
990    fn medium_contraction() {
991        assert_condensed_contractions("isn't wasn't", 3);
992    }
993
994    #[test]
995    fn medium_contraction2() {
996        assert_condensed_contractions("There's no way", 5);
997    }
998
999    #[test]
1000    fn selects_token_at_char_index() {
1001        let text = "There were three little pigs. They built three little homes.";
1002        let document = Document::new_plain_english_curated(text);
1003
1004        let got = document.get_token_at_char_index(19).unwrap();
1005
1006        assert!(got.kind.is_word());
1007        assert_eq!(got.span, Span::new(17, 23));
1008    }
1009
1010    fn assert_token_count(source: &str, count: usize) {
1011        let document = Document::new_plain_english_curated(source);
1012
1013        dbg!(document.tokens().map(|t| t.kind.clone()).collect_vec());
1014        assert_eq!(document.tokens.len(), count);
1015    }
1016
1017    #[test]
1018    fn condenses_number_suffixes() {
1019        assert_token_count("1st", 1);
1020        assert_token_count("This is the 2nd test", 9);
1021        assert_token_count("This is the 3rd test", 9);
1022        assert_token_count(
1023            "It works even with weird capitalization like this: 600nD",
1024            18,
1025        );
1026    }
1027
1028    #[test]
1029    fn condenses_ie() {
1030        assert_token_count("There is a thing (i.e. that one)", 15);
1031        assert_token_count("We are trying to condense \"i.e.\"", 13);
1032        assert_token_count(r#"Condenses words like "i.e.", "e.g." and "N.S.A.""#, 20);
1033    }
1034
1035    #[test]
1036    fn condenses_eg() {
1037        assert_token_count("We are trying to condense \"e.g.\"", 13);
1038        assert_token_count(r#"Condenses words like "i.e.", "e.g." and "N.S.A.""#, 20);
1039    }
1040
1041    #[test]
1042    fn condenses_nsa() {
1043        assert_token_count(r#"Condenses words like "i.e.", "e.g." and "N.S.A.""#, 20);
1044    }
1045
1046    #[test]
1047    fn parses_ellipsis() {
1048        assert_token_count("...", 1);
1049    }
1050
1051    #[test]
1052    fn parses_long_ellipsis() {
1053        assert_token_count(".....", 1);
1054    }
1055
1056    #[test]
1057    fn parses_short_ellipsis() {
1058        assert_token_count("..", 1);
1059    }
1060
1061    #[test]
1062    fn selects_token_at_offset() {
1063        let doc = Document::new_plain_english_curated("Foo bar baz");
1064
1065        let tok = doc.get_token_offset(1, -1).unwrap();
1066
1067        assert_eq!(tok.span, Span::new(0, 3));
1068    }
1069
1070    #[test]
1071    fn cant_select_token_before_start() {
1072        let doc = Document::new_plain_english_curated("Foo bar baz");
1073
1074        let tok = doc.get_token_offset(0, -1);
1075
1076        assert!(tok.is_none());
1077    }
1078
1079    #[test]
1080    fn select_next_word_pos_offset() {
1081        let doc = Document::new_plain_english_curated("Foo bar baz");
1082
1083        let bar = doc.get_next_word_from_offset(0, 1).unwrap();
1084        let bar = doc.get_span_content(&bar.span);
1085        assert_eq!(bar, ['b', 'a', 'r']);
1086    }
1087
1088    #[test]
1089    fn select_next_word_neg_offset() {
1090        let doc = Document::new_plain_english_curated("Foo bar baz");
1091
1092        let bar = doc.get_next_word_from_offset(2, -1).unwrap();
1093        let bar = doc.get_span_content(&bar.span);
1094        assert_eq!(bar, ['F', 'o', 'o']);
1095    }
1096
1097    #[test]
1098    fn cant_select_next_word_not_from_whitespace() {
1099        let doc = Document::new_plain_english_curated("Foo bar baz");
1100
1101        let tok = doc.get_next_word_from_offset(0, 2);
1102
1103        assert!(tok.is_none());
1104    }
1105
1106    #[test]
1107    fn cant_select_next_word_before_start() {
1108        let doc = Document::new_plain_english_curated("Foo bar baz");
1109
1110        let tok = doc.get_next_word_from_offset(0, -1);
1111
1112        assert!(tok.is_none());
1113    }
1114
1115    #[test]
1116    fn cant_select_next_word_with_punctuation_instead_of_whitespace() {
1117        let doc = Document::new_plain_english_curated("Foo, bar, baz");
1118
1119        let tok = doc.get_next_word_from_offset(0, 1);
1120
1121        assert!(tok.is_none());
1122    }
1123
1124    #[test]
1125    fn cant_select_next_word_with_punctuation_after_whitespace() {
1126        let doc = Document::new_plain_english_curated("Foo \"bar\", baz");
1127
1128        let tok = doc.get_next_word_from_offset(0, 1);
1129
1130        assert!(tok.is_none());
1131    }
1132
1133    #[test]
1134    fn condenses_filename_extensions() {
1135        let doc = Document::new_plain_english_curated(".c and .exe and .js");
1136        assert!(doc.tokens[0].kind.is_unlintable());
1137        assert!(doc.tokens[4].kind.is_unlintable());
1138        assert!(doc.tokens[8].kind.is_unlintable());
1139    }
1140
1141    #[test]
1142    fn condense_filename_extension_ok_at_start_and_end() {
1143        let doc = Document::new_plain_english_curated(".c and .EXE");
1144        assert!(doc.tokens.len() == 5);
1145        assert!(doc.tokens[0].kind.is_unlintable());
1146        assert!(doc.tokens[4].kind.is_unlintable());
1147    }
1148
1149    #[test]
1150    fn doesnt_condense_filename_extensions_with_mixed_case() {
1151        let doc = Document::new_plain_english_curated(".c and .Exe");
1152        assert!(doc.tokens.len() == 6);
1153        assert!(doc.tokens[0].kind.is_unlintable());
1154        assert!(doc.tokens[4].kind.is_punctuation());
1155        assert!(doc.tokens[5].kind.is_word());
1156    }
1157
1158    #[test]
1159    fn doesnt_condense_filename_extensions_with_non_letters() {
1160        let doc = Document::new_plain_english_curated(".COM and .C0M");
1161        assert!(doc.tokens.len() == 6);
1162        assert!(doc.tokens[0].kind.is_unlintable());
1163        assert!(doc.tokens[4].kind.is_punctuation());
1164        assert!(doc.tokens[5].kind.is_word());
1165    }
1166
1167    #[test]
1168    fn doesnt_condense_filename_extensions_longer_than_three() {
1169        let doc = Document::new_plain_english_curated(".dll and .dlls");
1170        assert!(doc.tokens.len() == 6);
1171        assert!(doc.tokens[0].kind.is_unlintable());
1172        assert!(doc.tokens[4].kind.is_punctuation());
1173        assert!(doc.tokens[5].kind.is_word());
1174    }
1175
1176    #[test]
1177    fn condense_filename_extension_in_parens() {
1178        let doc = Document::new_plain_english_curated(
1179            "true for the manual installation when trying to run the executable(.exe) after a manual download",
1180        );
1181        assert!(doc.tokens.len() > 23);
1182        assert!(doc.tokens[21].kind.is_open_round());
1183        assert!(doc.tokens[22].kind.is_unlintable());
1184        assert!(doc.tokens[23].kind.is_close_round());
1185    }
1186
1187    #[test]
1188    fn condense_tldr_uppercase() {
1189        let doc = Document::new_plain_english_curated("TL;DR");
1190        assert!(doc.tokens.len() == 1);
1191        assert!(doc.tokens[0].kind.is_word());
1192        assert!(doc.tokens[0].span.len() == 5);
1193    }
1194
1195    #[test]
1196    fn condense_tldr_lowercase() {
1197        let doc = Document::new_plain_english_curated("tl;dr");
1198        assert!(doc.tokens.len() == 1);
1199        assert!(doc.tokens[0].kind.is_word());
1200    }
1201
1202    #[test]
1203    fn condense_tldr_mixed_case_1() {
1204        let doc = Document::new_plain_english_curated("tl;DR");
1205        assert!(doc.tokens.len() == 1);
1206        assert!(doc.tokens[0].kind.is_word());
1207    }
1208
1209    #[test]
1210    fn condense_tldr_mixed_case_2() {
1211        let doc = Document::new_plain_english_curated("TL;Dr");
1212        assert!(doc.tokens.len() == 1);
1213        assert!(doc.tokens[0].kind.is_word());
1214    }
1215
1216    #[test]
1217    fn condense_tldr_pural() {
1218        let doc = Document::new_plain_english_curated(
1219            "managing the flow between components to produce relevant TL;DRs of current news articles",
1220        );
1221        // no token is a punctuation token - only words with whitespace between
1222        assert!(
1223            doc.tokens
1224                .iter()
1225                .all(|t| t.kind.is_word() || t.kind.is_whitespace())
1226        );
1227        // one of the word tokens contains a ';' character
1228        let tldrs = doc
1229            .tokens
1230            .iter()
1231            .filter(|t| t.get_ch(&doc.source).contains(&';'))
1232            .collect_vec();
1233        assert!(tldrs.len() == 1);
1234        assert!(tldrs[0].get_str(&doc.source) == "TL;DRs");
1235    }
1236
1237    #[test]
1238    fn condense_common_top_level_domains() {
1239        let doc = Document::new_plain_english_curated(".blog and .com and .NET");
1240        assert!(doc.tokens.len() == 9);
1241        assert!(doc.tokens[0].kind.is_unlintable());
1242        assert!(doc.tokens[4].kind.is_unlintable());
1243        assert!(doc.tokens[8].kind.is_unlintable());
1244    }
1245
1246    #[test]
1247    fn condense_common_top_level_domains_in_parens() {
1248        let doc = Document::new_plain_english_curated("(.blog)");
1249        assert!(doc.tokens.len() == 3);
1250        assert!(doc.tokens[0].kind.is_open_round());
1251        assert!(doc.tokens[1].kind.is_unlintable());
1252        assert!(doc.tokens[2].kind.is_close_round());
1253    }
1254
1255    #[test]
1256    fn doesnt_condense_unknown_top_level_domains() {
1257        let doc = Document::new_plain_english_curated(".harper");
1258        assert!(doc.tokens.len() == 2);
1259        assert!(doc.tokens[0].kind.is_punctuation());
1260        assert!(doc.tokens[1].kind.is_word());
1261    }
1262
1263    #[test]
1264    fn condense_r_and_d_caps() {
1265        let doc = Document::new_plain_english_curated("R&D");
1266        assert!(doc.tokens.len() == 1);
1267        assert!(doc.tokens[0].kind.is_word());
1268    }
1269
1270    #[test]
1271    fn condense_r_and_d_mixed_case() {
1272        let doc = Document::new_plain_english_curated("R&d");
1273        assert!(doc.tokens.len() == 1);
1274        assert!(doc.tokens[0].kind.is_word());
1275    }
1276
1277    #[test]
1278    fn condense_r_and_d_lowercase() {
1279        let doc = Document::new_plain_english_curated("r&d");
1280        assert!(doc.tokens.len() == 1);
1281        assert!(doc.tokens[0].kind.is_word());
1282    }
1283
1284    #[test]
1285    fn dont_condense_r_and_d_with_spaces() {
1286        let doc = Document::new_plain_english_curated("R & D");
1287        assert!(doc.tokens.len() == 5);
1288        assert!(doc.tokens[0].kind.is_word());
1289        assert!(doc.tokens[1].kind.is_whitespace());
1290        assert!(doc.tokens[2].kind.is_ampersand());
1291        assert!(doc.tokens[3].kind.is_whitespace());
1292        assert!(doc.tokens[4].kind.is_word());
1293    }
1294
1295    #[test]
1296    fn condense_q_and_a() {
1297        let doc =
1298            Document::new_plain_english_curated("A Q&A platform software for teams at any scales.");
1299        assert!(doc.tokens.len() >= 3);
1300        assert!(doc.tokens[2].kind.is_word());
1301        assert!(doc.tokens[2].get_str(&doc.source) == "Q&A");
1302    }
1303
1304    #[test]
1305    fn dont_allow_mixed_r_and_d_with_q_and_a() {
1306        let doc = Document::new_plain_english_curated("R&A or Q&D");
1307        assert!(doc.tokens.len() == 9);
1308        assert!(doc.tokens[1].kind.is_ampersand() || doc.tokens[7].kind.is_ampersand());
1309    }
1310
1311    #[test]
1312    fn condense_io() {
1313        let doc = Document::new_plain_english_curated("I/O");
1314        assert!(doc.tokens.len() == 1);
1315        assert!(doc.tokens[0].kind.is_word());
1316    }
1317
1318    #[test]
1319    fn finds_unmatched_quotes_in_document() {
1320        let raw = r#"
1321This is a paragraph with a single word "quoted."
1322
1323This is a second paragraph with no quotes.
1324
1325This is a third paragraph with a single erroneous "quote.
1326
1327This is a final paragraph with a weird "quote and a not-weird "quote".
1328            "#;
1329
1330        let doc = Document::new_markdown_default_curated(raw);
1331
1332        let quote_twins: Vec<_> = doc
1333            .iter_quotes()
1334            .map(|t| t.kind.as_quote().unwrap().twin_loc)
1335            .collect();
1336
1337        assert_eq!(
1338            quote_twins,
1339            vec![Some(19), Some(16), None, None, Some(89), Some(87)]
1340        )
1341    }
1342
1343    #[test]
1344    fn issue_1901() {
1345        let raw = r#"
1346"A quoted line"
1347"A quote without a closing mark
1348"Another quoted lined"
1349"The last quoted line"
1350            "#;
1351
1352        let doc = Document::new_markdown_default_curated(raw);
1353
1354        let quote_twins: Vec<_> = doc
1355            .iter_quotes()
1356            .map(|t| t.kind.as_quote().unwrap().twin_loc)
1357            .collect();
1358
1359        assert_eq!(
1360            quote_twins,
1361            vec![
1362                Some(6),
1363                Some(0),
1364                None,
1365                Some(27),
1366                Some(21),
1367                Some(37),
1368                Some(29)
1369            ]
1370        )
1371    }
1372}