macltok

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;