use std::collections::{BTreeMap, BTreeSet};
const MAX_DOCUMENT_BYTES: usize = 1024 * 1024;
const MAX_QUERY_BYTES: usize = 4096;
const MAX_TOKENS: usize = 65_536;
const MAX_QUERY_TERMS: usize = 1024;
const MAX_QUERY_DEPTH: usize = 64;
pub const INDEX_FORMAT_VERSION: &str = "1";
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Token {
pub text: String,
pub position: usize,
pub start: usize,
pub end: usize,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum Query {
Term { text: String, prefix: bool },
Not(Box<Query>),
And(Box<Query>, Box<Query>),
Or(Box<Query>, Box<Query>),
Phrase(Box<Query>, Box<Query>),
}
pub fn tokenize(config: &str, input: &str) -> Result<Vec<Token>, String> {
validate_config(config)?;
if input.len() > MAX_DOCUMENT_BYTES {
return Err("document exceeds 1048576 bytes".into());
}
let mut output = Vec::new();
let mut start = None;
for (offset, ch) in input.char_indices() {
if ch.is_alphanumeric() || ch == '_' {
start.get_or_insert(offset);
} else if let Some(token_start) = start.take() {
push_token(config, input, token_start, offset, &mut output)?;
}
}
if let Some(token_start) = start {
push_token(config, input, token_start, input.len(), &mut output)?;
}
Ok(output)
}
fn push_token(
config: &str,
input: &str,
start: usize,
end: usize,
output: &mut Vec<Token>,
) -> Result<(), String> {
if output.len() >= MAX_TOKENS {
return Err("document exceeds 65536 tokens".into());
}
output.push(Token {
text: normalize(config, &input[start..end]),
position: output.len() + 1,
start,
end,
});
Ok(())
}
fn validate_config(config: &str) -> Result<(), String> {
if matches!(config.to_ascii_lowercase().as_str(), "simple" | "english") {
Ok(())
} else {
Err(format!("unsupported text search configuration '{config}'"))
}
}
fn normalize(config: &str, token: &str) -> String {
let mut token = token.to_lowercase();
if config.eq_ignore_ascii_case("english") && token.is_ascii() {
for suffix in ["ing", "ed", "es", "s"] {
if token.len() > suffix.len() + 2 && token.ends_with(suffix) {
token.truncate(token.len() - suffix.len());
break;
}
}
}
token
}
pub fn to_tsvector(config: &str, document: &str) -> Result<String, String> {
let mut terms = BTreeMap::<String, Vec<usize>>::new();
for token in tokenize(config, document)? {
terms.entry(token.text).or_default().push(token.position);
}
Ok(terms
.into_iter()
.map(|(term, positions)| {
format!(
"'{}':{}",
term.replace('\'', "''"),
positions
.into_iter()
.map(|position| position.to_string())
.collect::<Vec<_>>()
.join(",")
)
})
.collect::<Vec<_>>()
.join(" "))
}
pub fn parse_tsquery(config: &str, input: &str) -> Result<Query, String> {
validate_config(config)?;
validate_query_input(input)?;
let tokens = lex_query(config, input)?;
if tokens.is_empty() {
return Err("TSQUERY is empty".into());
}
if tokens
.iter()
.filter(|token| matches!(token, QueryToken::Term(_, _)))
.count()
> MAX_QUERY_TERMS
{
return Err("TSQUERY exceeds 1024 terms".into());
}
let mut parser = QueryParser { tokens, index: 0 };
let query = parser.parse_or(0)?;
if parser.index != parser.tokens.len() {
return Err("TSQUERY contains an unexpected token".into());
}
Ok(query)
}
pub fn plainto_tsquery(config: &str, input: &str) -> Result<Query, String> {
validate_query_input(input)?;
let tokens = tokenize(config, input)?;
if tokens.len() > MAX_QUERY_TERMS {
return Err("TSQUERY exceeds 1024 terms".into());
}
terms_to_and(tokens.into_iter().map(|token| token.text))
}
pub fn websearch_to_tsquery(config: &str, input: &str) -> Result<Query, String> {
validate_config(config)?;
validate_query_input(input)?;
let mut groups = Vec::<Vec<Query>>::new();
let mut current = Vec::new();
let mut chars = input.char_indices().peekable();
while let Some((offset, ch)) = chars.next() {
if ch.is_whitespace() {
continue;
}
if matches!(ch, 'O' | 'o')
&& input[offset..]
.get(..2)
.is_some_and(|value| value.eq_ignore_ascii_case("OR"))
&& input[offset + 2..]
.chars()
.next()
.is_none_or(char::is_whitespace)
{
chars.next();
if current.is_empty() {
return Err("TSQUERY has an empty OR branch".into());
}
groups.push(std::mem::take(&mut current));
continue;
}
let negated = ch == '-';
let first = if negated {
chars
.next()
.ok_or_else(|| "TSQUERY has a dangling '-'".to_string())?
} else {
(offset, ch)
};
let query = if first.1 == '"' {
let start = first.0 + first.1.len_utf8();
let mut end = None;
for (position, next) in chars.by_ref() {
if next == '"' {
end = Some(position);
break;
}
}
let end = end.ok_or_else(|| "TSQUERY has an unterminated quote".to_string())?;
let terms = tokenize(config, &input[start..end])?;
terms_to_phrase(terms.into_iter().map(|token| token.text))?
} else {
let start = first.0;
let mut end = input.len();
while let Some(&(position, next)) = chars.peek() {
if next.is_whitespace() {
end = position;
break;
}
chars.next();
}
let terms = tokenize(config, &input[start..end])?;
terms_to_and(terms.into_iter().map(|token| token.text))?
};
current.push(if negated {
Query::Not(Box::new(query))
} else {
query
});
}
if !current.is_empty() {
groups.push(current);
}
if groups.is_empty() {
return Err("TSQUERY is empty".into());
}
let query = groups
.into_iter()
.map(|group| and_queries(group.into_iter()))
.collect::<Result<Vec<_>, _>>()?
.into_iter()
.reduce(|left, right| Query::Or(Box::new(left), Box::new(right)))
.ok_or_else(|| "TSQUERY is empty".to_string())?;
if query_term_count(&query) > MAX_QUERY_TERMS {
return Err("TSQUERY exceeds 1024 terms".into());
}
Ok(query)
}
fn validate_query_input(input: &str) -> Result<(), String> {
if input.len() > MAX_QUERY_BYTES {
Err("TSQUERY exceeds 4096 bytes".into())
} else {
Ok(())
}
}
fn query_term_count(query: &Query) -> usize {
match query {
Query::Term { .. } => 1,
Query::Not(query) => query_term_count(query),
Query::And(left, right) | Query::Or(left, right) | Query::Phrase(left, right) => {
query_term_count(left) + query_term_count(right)
}
}
}
fn terms_to_and(terms: impl IntoIterator<Item = String>) -> Result<Query, String> {
and_queries(terms.into_iter().map(|text| Query::Term {
text,
prefix: false,
}))
}
fn terms_to_phrase(terms: impl IntoIterator<Item = String>) -> Result<Query, String> {
terms
.into_iter()
.map(|text| Query::Term {
text,
prefix: false,
})
.reduce(|left, right| Query::Phrase(Box::new(left), Box::new(right)))
.ok_or_else(|| "TSQUERY phrase is empty".into())
}
fn and_queries(queries: impl Iterator<Item = Query>) -> Result<Query, String> {
queries
.reduce(|left, right| Query::And(Box::new(left), Box::new(right)))
.ok_or_else(|| "TSQUERY is empty".into())
}
#[derive(Debug, Clone, PartialEq, Eq)]
enum QueryToken {
Term(String, bool),
And,
Or,
Not,
Phrase,
Left,
Right,
}
fn lex_query(config: &str, input: &str) -> Result<Vec<QueryToken>, String> {
let mut output = Vec::new();
let mut index = 0;
while index < input.len() {
let ch = input[index..].chars().next().unwrap();
if ch.is_whitespace() {
index += ch.len_utf8();
continue;
}
let token = match ch {
'&' => QueryToken::And,
'|' => QueryToken::Or,
'!' => QueryToken::Not,
'(' => QueryToken::Left,
')' => QueryToken::Right,
'<' if input[index..].starts_with("<->") => {
index += 2;
QueryToken::Phrase
}
_ if ch.is_alphanumeric() || ch == '_' => {
let start = index;
index += ch.len_utf8();
while index < input.len() {
let next = input[index..].chars().next().unwrap();
if !(next.is_alphanumeric() || next == '_') {
break;
}
index += next.len_utf8();
}
let end = index;
let prefix = input[index..].starts_with(":*");
if prefix {
index += 2;
}
output.push(QueryToken::Term(
normalize(config, &input[start..end]),
prefix,
));
continue;
}
_ => return Err(format!("TSQUERY contains unsupported character '{ch}'")),
};
output.push(token);
index += ch.len_utf8();
}
Ok(output)
}
struct QueryParser {
tokens: Vec<QueryToken>,
index: usize,
}
impl QueryParser {
fn parse_or(&mut self, depth: usize) -> Result<Query, String> {
let mut query = self.parse_and(depth + 1)?;
while self.take(&QueryToken::Or) {
query = Query::Or(Box::new(query), Box::new(self.parse_and(depth + 1)?));
}
Ok(query)
}
fn parse_and(&mut self, depth: usize) -> Result<Query, String> {
let mut query = self.parse_phrase(depth + 1)?;
while self.take(&QueryToken::And) {
query = Query::And(Box::new(query), Box::new(self.parse_phrase(depth + 1)?));
}
Ok(query)
}
fn parse_phrase(&mut self, depth: usize) -> Result<Query, String> {
let mut query = self.parse_unary(depth + 1)?;
while self.take(&QueryToken::Phrase) {
query = Query::Phrase(Box::new(query), Box::new(self.parse_unary(depth + 1)?));
}
Ok(query)
}
fn parse_unary(&mut self, depth: usize) -> Result<Query, String> {
if depth > MAX_QUERY_DEPTH {
return Err("TSQUERY exceeds nesting depth 64".into());
}
if self.take(&QueryToken::Not) {
return Ok(Query::Not(Box::new(self.parse_unary(depth + 1)?)));
}
match self.tokens.get(self.index).cloned() {
Some(QueryToken::Term(text, prefix)) => {
self.index += 1;
Ok(Query::Term { text, prefix })
}
Some(QueryToken::Left) => {
self.index += 1;
let query = self.parse_or(depth + 1)?;
if !self.take(&QueryToken::Right) {
return Err("TSQUERY is missing ')'".into());
}
Ok(query)
}
_ => Err("TSQUERY expects a term".into()),
}
}
fn take(&mut self, expected: &QueryToken) -> bool {
if self.tokens.get(self.index) == Some(expected) {
self.index += 1;
true
} else {
false
}
}
}
pub fn format_query(query: &Query) -> String {
format_query_precedence(query, 0)
}
fn format_query_precedence(query: &Query, parent: u8) -> String {
let (precedence, value) = match query {
Query::Term { text, prefix } => (5, format!("{text}{}", if *prefix { ":*" } else { "" })),
Query::Not(query) => (4, format!("!{}", format_query_precedence(query, 4))),
Query::Phrase(left, right) => (
3,
format!(
"{} <-> {}",
format_query_precedence(left, 3),
format_query_precedence(right, 3)
),
),
Query::And(left, right) => (
2,
format!(
"{} & {}",
format_query_precedence(left, 2),
format_query_precedence(right, 2)
),
),
Query::Or(left, right) => (
1,
format!(
"{} | {}",
format_query_precedence(left, 1),
format_query_precedence(right, 1)
),
),
};
if precedence < parent {
format!("({value})")
} else {
value
}
}
pub fn matches_query(tokens: &[Token], query: &Query) -> bool {
!positions(tokens, query).is_empty()
}
fn positions(tokens: &[Token], query: &Query) -> BTreeSet<usize> {
match query {
Query::Term { text, prefix } => tokens
.iter()
.filter(|token| {
if *prefix {
token.text.starts_with(text)
} else {
token.text == *text
}
})
.map(|token| token.position)
.collect(),
Query::Not(query) => positions(tokens, query)
.is_empty()
.then_some(0)
.into_iter()
.collect(),
Query::And(left, right) => {
let left_positions = positions(tokens, left);
if left_positions.is_empty() || positions(tokens, right).is_empty() {
BTreeSet::new()
} else {
left_positions
}
}
Query::Or(left, right) => positions(tokens, left)
.union(&positions(tokens, right))
.copied()
.collect(),
Query::Phrase(left, right) => {
let left = positions(tokens, left);
positions(tokens, right)
.into_iter()
.filter(|position| {
position
.checked_sub(1)
.is_some_and(|previous| left.contains(&previous))
})
.collect()
}
}
}
pub fn index_terms(query: &Query, output: &mut BTreeSet<String>) -> bool {
match query {
Query::Term {
text,
prefix: false,
} => {
output.insert(text.clone());
true
}
Query::Term { prefix: true, .. } | Query::Not(_) => false,
Query::And(left, right) | Query::Or(left, right) | Query::Phrase(left, right) => {
index_terms(left, output) && index_terms(right, output)
}
}
}
pub fn rank(tokens: &[Token], query: &Query) -> f64 {
if !matches_query(tokens, query) || tokens.is_empty() {
return 0.0;
}
let matched = tokens
.iter()
.filter(|token| positive_token_match(query, &token.text))
.count();
matched as f64 / tokens.len() as f64
}
fn positive_token_match(query: &Query, token: &str) -> bool {
match query {
Query::Term { text, prefix } => {
if *prefix {
token.starts_with(text)
} else {
token == text
}
}
Query::Not(_) => false,
Query::And(left, right) | Query::Or(left, right) | Query::Phrase(left, right) => {
positive_token_match(left, token) || positive_token_match(right, token)
}
}
}
pub fn headline(config: &str, document: &str, query: &Query) -> Result<String, String> {
let tokens = tokenize(config, document)?;
let mut output = String::with_capacity(document.len());
let mut cursor = 0;
for token in tokens {
output.push_str(&document[cursor..token.start]);
if positive_token_match(query, &token.text) {
output.push_str("<b>");
output.push_str(&document[token.start..token.end]);
output.push_str("</b>");
} else {
output.push_str(&document[token.start..token.end]);
}
cursor = token.end;
}
output.push_str(&document[cursor..]);
if output.len() > MAX_DOCUMENT_BYTES {
return Err("headline exceeds 1048576 bytes".into());
}
Ok(output)
}
pub fn parse_tsvector(input: &str) -> Result<Vec<Token>, String> {
if input.len() > MAX_DOCUMENT_BYTES {
return Err("TSVECTOR exceeds 1048576 bytes".into());
}
let mut tokens = Vec::new();
for entry in input.split_whitespace() {
let (term, positions) = entry
.rsplit_once(':')
.ok_or_else(|| "invalid TSVECTOR entry".to_string())?;
let term = term
.strip_prefix('\'')
.and_then(|value| value.strip_suffix('\''))
.ok_or_else(|| "invalid TSVECTOR term".to_string())?
.replace("''", "'");
for position in positions.split(',') {
if tokens.len() >= MAX_TOKENS {
return Err("TSVECTOR exceeds 65536 positions".into());
}
let position = position
.parse::<usize>()
.map_err(|_| "invalid TSVECTOR position".to_string())?;
tokens.push(Token {
text: term.clone(),
position,
start: 0,
end: 0,
});
}
}
tokens.sort_by_key(|token| token.position);
Ok(tokens)
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn phrase_tracks_the_rightmost_position() {
let tokens = tokenize("simple", "one two three").unwrap();
let query = parse_tsquery("simple", "one <-> two <-> three").unwrap();
assert!(matches_query(&tokens, &query));
}
#[test]
fn vector_round_trip_preserves_rank() {
let vector = to_tsvector("simple", "the quick brown fox").unwrap();
let query = plainto_tsquery("simple", "quick fox").unwrap();
assert_eq!(rank(&parse_tsvector(&vector).unwrap(), &query), 0.5);
}
#[test]
fn public_resource_limits_fail_closed() {
assert!(tokenize("simple", &"x".repeat(MAX_DOCUMENT_BYTES + 1)).is_err());
assert!(parse_tsquery("simple", &"x".repeat(MAX_QUERY_BYTES + 1)).is_err());
}
}