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#[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 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 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 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 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 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 pub fn new_plain_english_curated_chars(source: &[char]) -> Self {
86 Self::new_from_chars(Lrc::from(source), &PlainEnglish, &FstDictionary::curated())
87 }
88
89 pub fn new_plain_english_curated(text: &str) -> Self {
92 Self::new(text, &PlainEnglish, &FstDictionary::curated())
93 }
94
95 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 pub fn new_plain_english(text: &str, dictionary: &impl Dictionary) -> Self {
111 Self::new(text, &PlainEnglish, dictionary)
112 }
113
114 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 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 pub fn new_markdown_default_curated(text: &str) -> Self {
137 Self::new_markdown_curated(text, MarkdownOptions::default())
138 }
139
140 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 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 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 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; let mut wi = 0; 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 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 fn condense_indices(&mut self, indices: &[usize], stretch_len: usize) {
239 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 let old = self.tokens.clone();
249 self.tokens.clear();
250
251 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 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 pub fn get_token(&self, index: usize) -> Option<&Token> {
292 self.tokens.get(index)
293 }
294
295 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 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 pub fn fat_tokens(&self) -> impl Iterator<Item = FatToken> + '_ {
344 self.tokens().map(|token| token.to_fat(&self.source))
345 }
346
347 pub fn get_next_word_from_offset(&self, base: usize, offset: isize) -> Option<&Token> {
350 if !self.get_token_offset(base, offset)?.kind.is_whitespace() {
352 return None;
353 }
354 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 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 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 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 for quote in "e_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 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 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 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 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 = ©[cursor];
505
506 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 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 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 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 = ©[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 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 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 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 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 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 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 self.tokens[cursor - 2].span = Span::new(
792 self.tokens[cursor - 2].span.start,
793 self.tokens[cursor].span.end,
794 );
795
796 to_remove.push_back(cursor - 1);
798 to_remove.push_back(cursor);
799 }
800
801 cursor += 1;
803
804 if cursor >= self.tokens.len() {
805 break;
806 }
807 }
808
809 self.tokens.remove_indices(to_remove);
811 }
812
813 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 fn condense_ampersand_pairs(&mut self) {
877 self.condense_delimited_pairs(
878 |kind| kind.is_ampersand(),
879 &[
880 ('b', 'b'), ('b', 'w'), ('g', 't'), ('k', 'r'), ('q', 'a'), ('r', 'b'), ('r', 'd'), ('r', 'r'), ('s', 'p'), ],
890 );
891 }
892
893 fn condense_slash_pairs(&mut self) {
895 self.condense_delimited_pairs(
896 |kind| kind.is_slash(),
897 &[
898 ('a', 'c'), ('b', 'w'), ('c', 'o'), ('d', 'c'), ('d', 'l'), ('i', 'o'), ('j', 'k'), ('n', 'a'), ('r', 'c'), ('s', 'n'), ('y', 'n'), ('y', 'o'), ],
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 assert!(
1223 doc.tokens
1224 .iter()
1225 .all(|t| t.kind.is_word() || t.kind.is_whitespace())
1226 );
1227 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}