macltok

public

Star0

Watch0

126 commits 3 branches 0 tags 0 releases3.0 MiB

use tantivy::{
    Term,
    query::{BooleanQuery, Occur, PhraseQuery, Query, TermQuery},
    schema::{Field, IndexRecordOption},
    tokenizer::{
        AsciiFoldingFilter, LowerCaser, RawTokenizer, TextAnalyzer, Token, TokenStream, Tokenizer,
    },
};

use crate::{Error, ErrorKind};

pub(crate) const TOKENIZER: &str = "macltok";
/// longer runs are hashes, keys or noise, and in chars so cyrillic and greek words fit
const MAX_WORD_CHARS: usize = 64;

#[derive(Clone, Copy, Debug, Eq, PartialEq)]
enum Kind {
    Word,
    /// one char of a han, kana, hangul or thai run, or an emoji
    Single,
    /// two neighbouring chars of such a run, sharing the first one's position
    Pair,
}

#[derive(Clone, Debug, Eq, PartialEq)]
struct Piece {
    from: usize,
    to: usize,
    position: usize,
    kind: Kind,
    /// a single inside a run of two or more, which its pairs already cover in a query
    paired: bool,
}

/// scripts written without spaces, where a char is the only boundary there is
fn unspaced(c: char) -> bool {
    matches!(c,
        '\u{3040}'..='\u{30FF}' | '\u{31F0}'..='\u{31FF}' | '\u{3400}'..='\u{4DBF}'
        | '\u{4E00}'..='\u{9FFF}' | '\u{F900}'..='\u{FAFF}' | '\u{FF66}'..='\u{FF9F}'
        | '\u{1100}'..='\u{11FF}' | '\u{3130}'..='\u{318F}' | '\u{AC00}'..='\u{D7AF}'
        | '\u{0E00}'..='\u{0E7F}' | '\u{20000}'..='\u{2FA1F}')
}

fn emoji(c: char) -> bool {
    matches!(c, '\u{1F300}'..='\u{1FAFF}' | '\u{2600}'..='\u{27BF}')
}

fn pieces(text: &str) -> Vec<Piece> {
    let mut pieces = Vec::new();
    let mut position = 0;
    let mut chars = text.char_indices().peekable();
    while let Some((from, c)) = chars.next() {
        if unspaced(c) {
            let mut run = vec![(from, from + c.len_utf8())];
            while let Some(&(at, next)) = chars.peek() {
                if !unspaced(next) {
                    break;
                }
                run.push((at, at + next.len_utf8()));
                chars.next();
            }
            let paired = run.len() > 1;
            for (index, &(from, to)) in run.iter().enumerate() {
                pieces.push(Piece {
                    from,
                    to,
                    position,
                    kind: Kind::Single,
                    paired,
                });
                if let Some(&(_, next_to)) = run.get(index + 1) {
                    pieces.push(Piece {
                        from,
                        to: next_to,
                        position,
                        kind: Kind::Pair,
                        paired: false,
                    });
                }
                position += 1;
            }
        } else if emoji(c) {
            pieces.push(Piece {
                from,
                to: from + c.len_utf8(),
                position,
                kind: Kind::Single,
                paired: false,
            });
            position += 1;
        } else if c.is_alphanumeric() {
            let mut to = from + c.len_utf8();
            let mut count = 1;
            while let Some(&(at, next)) = chars.peek() {
                if !next.is_alphanumeric() || unspaced(next) {
                    break;
                }
                to = at + next.len_utf8();
                count += 1;
                chars.next();
            }
            if count <= MAX_WORD_CHARS {
                pieces.push(Piece {
                    from,
                    to,
                    position,
                    kind: Kind::Word,
                    paired: false,
                });
                position += 1;
            }
        }
    }
    pieces
}

#[derive(Clone)]
pub(crate) struct MacltokTokenizer;

pub(crate) struct PieceStream {
    tokens: std::vec::IntoIter<Token>,
    token: Token,
}

impl TokenStream for PieceStream {
    fn advance(&mut self) -> bool {
        match self.tokens.next() {
            Some(token) => {
                self.token = token;
                true
            }
            None => false,
        }
    }

    fn token(&self) -> &Token {
        &self.token
    }

    fn token_mut(&mut self) -> &mut Token {
        &mut self.token
    }
}

impl Tokenizer for MacltokTokenizer {
    type TokenStream<'a> = PieceStream;

    fn token_stream<'a>(&'a mut self, text: &'a str) -> PieceStream {
        let tokens: Vec<_> = pieces(text)
            .into_iter()
            .map(|piece| Token {
                offset_from: piece.from,
                offset_to: piece.to,
                position: piece.position,
                text: text[piece.from..piece.to].to_owned(),
                position_length: 1,
            })
            .collect();
        PieceStream {
            tokens: tokens.into_iter(),
            token: Token::default(),
        }
    }
}

pub(crate) fn analyzer() -> TextAnalyzer {
    TextAnalyzer::builder(MacltokTokenizer)
        .filter(LowerCaser)
        .filter(AsciiFoldingFilter)
        .build()
}

fn fold(text: &str) -> String {
    let mut folder = TextAnalyzer::builder(RawTokenizer::default())
        .filter(LowerCaser)
        .filter(AsciiFoldingFilter)
        .build();
    let mut stream = folder.token_stream(text);
    if stream.advance() {
        stream.token().text.clone()
    } else {
        String::new()
    }
}

/// the byte ranges in `text` whose analysed token is one of `terms`, overlaps joined
pub(crate) fn match_ranges(text: &str, terms: &[String]) -> Vec<(usize, usize)> {
    if terms.is_empty() {
        return Vec::new();
    }
    let mut analyzer = analyzer();
    let mut stream = analyzer.token_stream(text);
    let mut ranges: Vec<(usize, usize)> = Vec::new();
    while stream.advance() {
        let token = stream.token();
        if !terms.iter().any(|term| *term == token.text) {
            continue;
        }
        match ranges.last_mut() {
            Some(last) if token.offset_from < last.1 => last.1 = last.1.max(token.offset_to),
            _ => ranges.push((token.offset_from, token.offset_to)),
        }
    }
    ranges
}

enum Clause {
    Words(String),
    Phrase(String),
}

/// quotes make a phrase, a leading `-` excludes, anything else is a word that must match.
/// no field syntax, so `re: deploy` or a url is just its words
fn clauses(query: &str) -> Vec<(bool, Clause)> {
    let mut clauses = Vec::new();
    let mut rest = query;
    loop {
        rest = rest.trim_start();
        if rest.is_empty() {
            return clauses;
        }
        let (excluded, body) = match rest.strip_prefix('-') {
            Some(body) if body.starts_with(|c: char| !c.is_whitespace()) => (true, body),
            _ => (false, rest),
        };
        if let Some(quoted) = body.strip_prefix('"') {
            let end = quoted.find('"').unwrap_or(quoted.len());
            clauses.push((excluded, Clause::Phrase(quoted[..end].to_owned())));
            rest = quoted.get(end + 1..).unwrap_or("");
        } else {
            let end = body.find(char::is_whitespace).unwrap_or(body.len());
            clauses.push((excluded, Clause::Words(body[..end].to_owned())));
            rest = &body[end..];
        }
    }
}

fn term(field: Field, text: &str) -> Box<dyn Query> {
    Box::new(TermQuery::new(
        Term::from_field_text(field, text),
        IndexRecordOption::WithFreqs,
    ))
}

/// every term the clause needs, and its query, `None` when nothing in it is searchable
fn clause_query(field: Field, clause: &Clause) -> Option<(Box<dyn Query>, Vec<String>)> {
    let text = match clause {
        Clause::Words(text) | Clause::Phrase(text) => text,
    };
    let found = pieces(text);
    if let Clause::Phrase(_) = clause {
        let ordered: Vec<_> = found
            .iter()
            .filter(|piece| piece.kind != Kind::Pair)
            .map(|piece| (piece.position, fold(&text[piece.from..piece.to])))
            .filter(|(_, folded)| !folded.is_empty())
            .collect();
        let terms: Vec<String> = ordered.iter().map(|(_, folded)| folded.clone()).collect();
        return match ordered.len() {
            0 => None,
            1 => Some((term(field, &ordered[0].1), terms)),
            _ => {
                let first = ordered[0].0;
                let offsets = ordered
                    .iter()
                    .map(|(position, folded)| {
                        (position - first, Term::from_field_text(field, folded))
                    })
                    .collect();
                Some((Box::new(PhraseQuery::new_with_offset(offsets)), terms))
            }
        };
    }
    let mut terms: Vec<String> = found
        .iter()
        .filter(|piece| !piece.paired)
        .map(|piece| fold(&text[piece.from..piece.to]))
        .filter(|folded| !folded.is_empty())
        .collect();
    terms.dedup();
    match terms.len() {
        0 => None,
        1 => Some((term(field, &terms[0]), terms)),
        _ => {
            let musts = terms
                .iter()
                .map(|text| (Occur::Must, term(field, text)))
                .collect();
            Some((Box::new(BooleanQuery::new(musts)), terms))
        }
    }
}

/// the query and the terms to highlight. one with nothing left to find is `Invalid`,
/// rather than everything in scope
pub(crate) fn build_query(
    field: Field,
    query: &str,
) -> Result<(Box<dyn Query>, Vec<String>), Error> {
    let mut subqueries = Vec::new();
    let mut highlights = Vec::new();
    let mut positive = false;
    for (excluded, clause) in clauses(query) {
        let Some((subquery, terms)) = clause_query(field, &clause) else {
            continue;
        };
        if excluded {
            subqueries.push((Occur::MustNot, subquery));
        } else {
            positive = true;
            highlights.extend(terms);
            subqueries.push((Occur::Must, subquery));
        }
    }
    if !positive {
        return Err(Error::new(
            ErrorKind::Invalid,
            "search needs at least one word that isn't excluded",
        ));
    }
    highlights.sort();
    highlights.dedup();
    Ok((Box::new(BooleanQuery::new(subqueries)), highlights))
}

#[cfg(test)]
#[path = "text_tests.rs"]
mod tests;

1	use tantivy::{
2	Term,
3	query::{BooleanQuery, Occur, PhraseQuery, Query, TermQuery},
4	schema::{Field, IndexRecordOption},
5	tokenizer::{
6	AsciiFoldingFilter, LowerCaser, RawTokenizer, TextAnalyzer, Token, TokenStream, Tokenizer,
7	},
8	};
9
10	use crate::{Error, ErrorKind};
11
12	pub(crate) const TOKENIZER: &str = "macltok";
13	/// longer runs are hashes, keys or noise, and in chars so cyrillic and greek words fit
14	const MAX_WORD_CHARS: usize = 64;
15
16	#[derive(Clone, Copy, Debug, Eq, PartialEq)]
17	enum Kind {
18	Word,
19	/// one char of a han, kana, hangul or thai run, or an emoji
20	Single,
21	/// two neighbouring chars of such a run, sharing the first one's position
22	Pair,
23	}
24
25	#[derive(Clone, Debug, Eq, PartialEq)]
26	struct Piece {
27	from: usize,
28	to: usize,
29	position: usize,
30	kind: Kind,
31	/// a single inside a run of two or more, which its pairs already cover in a query
32	paired: bool,
33	}
34
35	/// scripts written without spaces, where a char is the only boundary there is
36	fn unspaced(c: char) -> bool {
37	matches!(c,
38	'\u{3040}'..='\u{30FF}' \| '\u{31F0}'..='\u{31FF}' \| '\u{3400}'..='\u{4DBF}'
39	\| '\u{4E00}'..='\u{9FFF}' \| '\u{F900}'..='\u{FAFF}' \| '\u{FF66}'..='\u{FF9F}'
40	\| '\u{1100}'..='\u{11FF}' \| '\u{3130}'..='\u{318F}' \| '\u{AC00}'..='\u{D7AF}'
41	\| '\u{0E00}'..='\u{0E7F}' \| '\u{20000}'..='\u{2FA1F}')
42	}
43
44	fn emoji(c: char) -> bool {
45	matches!(c, '\u{1F300}'..='\u{1FAFF}' \| '\u{2600}'..='\u{27BF}')
46	}
47
48	fn pieces(text: &str) -> Vec<Piece> {
49	let mut pieces = Vec::new();
50	let mut position = 0;
51	let mut chars = text.char_indices().peekable();
52	while let Some((from, c)) = chars.next() {
53	if unspaced(c) {
54	let mut run = vec![(from, from + c.len_utf8())];
55	while let Some(&(at, next)) = chars.peek() {
56	if !unspaced(next) {
57	break;
58	}
59	run.push((at, at + next.len_utf8()));
60	chars.next();
61	}
62	let paired = run.len() > 1;
63	for (index, &(from, to)) in run.iter().enumerate() {
64	pieces.push(Piece {
65	from,
66	to,
67	position,
68	kind: Kind::Single,
69	paired,
70	});
71	if let Some(&(_, next_to)) = run.get(index + 1) {
72	pieces.push(Piece {
73	from,
74	to: next_to,
75	position,
76	kind: Kind::Pair,
77	paired: false,
78	});
79	}
80	position += 1;
81	}
82	} else if emoji(c) {
83	pieces.push(Piece {
84	from,
85	to: from + c.len_utf8(),
86	position,
87	kind: Kind::Single,
88	paired: false,
89	});
90	position += 1;
91	} else if c.is_alphanumeric() {
92	let mut to = from + c.len_utf8();
93	let mut count = 1;
94	while let Some(&(at, next)) = chars.peek() {
95	if !next.is_alphanumeric() \|\| unspaced(next) {
96	break;
97	}
98	to = at + next.len_utf8();
99	count += 1;
100	chars.next();
101	}
102	if count <= MAX_WORD_CHARS {
103	pieces.push(Piece {
104	from,
105	to,
106	position,
107	kind: Kind::Word,
108	paired: false,
109	});
110	position += 1;
111	}
112	}
113	}
114	pieces
115	}
116
117	#[derive(Clone)]
118	pub(crate) struct MacltokTokenizer;
119
120	pub(crate) struct PieceStream {
121	tokens: std::vec::IntoIter<Token>,
122	token: Token,
123	}
124
125	impl TokenStream for PieceStream {
126	fn advance(&mut self) -> bool {
127	match self.tokens.next() {
128	Some(token) => {
129	self.token = token;
130	true
131	}
132	None => false,
133	}
134	}
135
136	fn token(&self) -> &Token {
137	&self.token
138	}
139
140	fn token_mut(&mut self) -> &mut Token {
141	&mut self.token
142	}
143	}
144
145	impl Tokenizer for MacltokTokenizer {
146	type TokenStream<'a> = PieceStream;
147
148	fn token_stream<'a>(&'a mut self, text: &'a str) -> PieceStream {
149	let tokens: Vec<_> = pieces(text)
150	.into_iter()
151	.map(\|piece\| Token {
152	offset_from: piece.from,
153	offset_to: piece.to,
154	position: piece.position,
155	text: text[piece.from..piece.to].to_owned(),
156	position_length: 1,
157	})
158	.collect();
159	PieceStream {
160	tokens: tokens.into_iter(),
161	token: Token::default(),
162	}
163	}
164	}
165
166	pub(crate) fn analyzer() -> TextAnalyzer {
167	TextAnalyzer::builder(MacltokTokenizer)
168	.filter(LowerCaser)
169	.filter(AsciiFoldingFilter)
170	.build()
171	}
172
173	fn fold(text: &str) -> String {
174	let mut folder = TextAnalyzer::builder(RawTokenizer::default())
175	.filter(LowerCaser)
176	.filter(AsciiFoldingFilter)
177	.build();
178	let mut stream = folder.token_stream(text);
179	if stream.advance() {
180	stream.token().text.clone()
181	} else {
182	String::new()
183	}
184	}
185
186	/// the byte ranges in `text` whose analysed token is one of `terms`, overlaps joined
187	pub(crate) fn match_ranges(text: &str, terms: &[String]) -> Vec<(usize, usize)> {
188	if terms.is_empty() {
189	return Vec::new();
190	}
191	let mut analyzer = analyzer();
192	let mut stream = analyzer.token_stream(text);
193	let mut ranges: Vec<(usize, usize)> = Vec::new();
194	while stream.advance() {
195	let token = stream.token();
196	if !terms.iter().any(\|term\| *term == token.text) {
197	continue;
198	}
199	match ranges.last_mut() {
200	Some(last) if token.offset_from < last.1 => last.1 = last.1.max(token.offset_to),
201	_ => ranges.push((token.offset_from, token.offset_to)),
202	}
203	}
204	ranges
205	}
206
207	enum Clause {
208	Words(String),
209	Phrase(String),
210	}
211
212	/// quotes make a phrase, a leading `-` excludes, anything else is a word that must match.
213	/// no field syntax, so `re: deploy` or a url is just its words
214	fn clauses(query: &str) -> Vec<(bool, Clause)> {
215	let mut clauses = Vec::new();
216	let mut rest = query;
217	loop {
218	rest = rest.trim_start();
219	if rest.is_empty() {
220	return clauses;
221	}
222	let (excluded, body) = match rest.strip_prefix('-') {
223	Some(body) if body.starts_with(\|c: char\| !c.is_whitespace()) => (true, body),
224	_ => (false, rest),
225	};
226	if let Some(quoted) = body.strip_prefix('"') {
227	let end = quoted.find('"').unwrap_or(quoted.len());
228	clauses.push((excluded, Clause::Phrase(quoted[..end].to_owned())));
229	rest = quoted.get(end + 1..).unwrap_or("");
230	} else {
231	let end = body.find(char::is_whitespace).unwrap_or(body.len());
232	clauses.push((excluded, Clause::Words(body[..end].to_owned())));
233	rest = &body[end..];
234	}
235	}
236	}
237
238	fn term(field: Field, text: &str) -> Box<dyn Query> {
239	Box::new(TermQuery::new(
240	Term::from_field_text(field, text),
241	IndexRecordOption::WithFreqs,
242	))
243	}
244
245	/// every term the clause needs, and its query, `None` when nothing in it is searchable
246	fn clause_query(field: Field, clause: &Clause) -> Option<(Box<dyn Query>, Vec<String>)> {
247	let text = match clause {
248	Clause::Words(text) \| Clause::Phrase(text) => text,
249	};
250	let found = pieces(text);
251	if let Clause::Phrase(_) = clause {
252	let ordered: Vec<_> = found
253	.iter()
254	.filter(\|piece\| piece.kind != Kind::Pair)
255	.map(\|piece\| (piece.position, fold(&text[piece.from..piece.to])))
256	.filter(\|(_, folded)\| !folded.is_empty())
257	.collect();
258	let terms: Vec<String> = ordered.iter().map(\|(_, folded)\| folded.clone()).collect();
259	return match ordered.len() {
260	0 => None,
261	1 => Some((term(field, &ordered[0].1), terms)),
262	_ => {
263	let first = ordered[0].0;
264	let offsets = ordered
265	.iter()
266	.map(\|(position, folded)\| {
267	(position - first, Term::from_field_text(field, folded))
268	})
269	.collect();
270	Some((Box::new(PhraseQuery::new_with_offset(offsets)), terms))
271	}
272	};
273	}
274	let mut terms: Vec<String> = found
275	.iter()
276	.filter(\|piece\| !piece.paired)
277	.map(\|piece\| fold(&text[piece.from..piece.to]))
278	.filter(\|folded\| !folded.is_empty())
279	.collect();
280	terms.dedup();
281	match terms.len() {
282	0 => None,
283	1 => Some((term(field, &terms[0]), terms)),
284	_ => {
285	let musts = terms
286	.iter()
287	.map(\|text\| (Occur::Must, term(field, text)))
288	.collect();
289	Some((Box::new(BooleanQuery::new(musts)), terms))
290	}
291	}
292	}
293
294	/// the query and the terms to highlight. one with nothing left to find is `Invalid`,
295	/// rather than everything in scope
296	pub(crate) fn build_query(
297	field: Field,
298	query: &str,
299	) -> Result<(Box<dyn Query>, Vec<String>), Error> {
300	let mut subqueries = Vec::new();
301	let mut highlights = Vec::new();
302	let mut positive = false;
303	for (excluded, clause) in clauses(query) {
304	let Some((subquery, terms)) = clause_query(field, &clause) else {
305	continue;
306	};
307	if excluded {
308	subqueries.push((Occur::MustNot, subquery));
309	} else {
310	positive = true;
311	highlights.extend(terms);
312	subqueries.push((Occur::Must, subquery));
313	}
314	}
315	if !positive {
316	return Err(Error::new(
317	ErrorKind::Invalid,
318	"search needs at least one word that isn't excluded",
319	));
320	}
321	highlights.sort();
322	highlights.dedup();
323	Ok((Box::new(BooleanQuery::new(subqueries)), highlights))
324	}
325
326	#[cfg(test)]
327	#[path = "text_tests.rs"]
328	mod tests;