haematite-cli 0.21.8

//! Read a query's source and ordered stages without opening any records.

mod group_by;
mod r#match;
#[path = "run/percentile_parse.rs"]
mod percentile;
mod rerank;
mod rg;
mod shape;
mod source;
mod spans;
mod stages;

#[cfg(test)]
mod rerank_tests;
pub(in crate::serve) mod window;

use std::fmt;
use std::num::NonZeroU32;

pub use r#match::{Direction, Node, Relation};
pub use source::{Argument, Root, Source};
pub use stages::{Question, Stage};

use crate::serve::protocol::Metric;

/// A byte range in the original query, excluding surrounding whitespace.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct Span {
    pub start: usize,
    pub end: usize,
}

/// Keep a parsed operation beside its location for query checking.
#[derive(Debug, PartialEq, Eq)]
pub struct Spanned<T> {
    pub value: T,
    pub span: Span,
}

#[derive(Debug, PartialEq, Eq)]
pub struct Plan {
    pub source: Source,
    pub stages: Vec<Spanned<Stage>>,
}

/// The stored-vector search understood by the existing query runner.
#[derive(Debug, PartialEq, Eq)]
pub struct Search {
    pub field: String,
    pub record: String,
    pub measure: Metric,
    pub limit: NonZeroU32,
}

impl Plan {
    #[must_use]
    pub fn plain(text: &str) -> Self {
        Self {
            source: Source::All,
            stages: vec![Spanned {
                value: Stage::Jq(text.to_owned()),
                span: Span {
                    start: 0,
                    end: text.len(),
                },
            }],
        }
    }
}

/// A named refusal and the byte range to mark in the query text.
#[derive(Debug, PartialEq, Eq)]
pub enum Refused {
    NotYet(&'static str, Span),
    Unserved(&'static str, Span),
    Parse(String, Span),
    WholeRecords(String, Span),
    Percentile(String, Span),
    Flag(String, Span),
    Window(&'static str, String, Span),
}

impl std::error::Error for Refused {}

impl Refused {
    #[must_use]
    pub const fn code(&self) -> &'static str {
        match self {
            Self::NotYet(..) | Self::Unserved(..) | Self::WholeRecords(..) => {
                "query_stage_not_available"
            }
            Self::Parse(..) => "query_parse",
            Self::Percentile(..) => "query_percentile_range",
            Self::Flag(..) => "query_flag",
            Self::Window(code, ..) => code,
        }
    }

    #[must_use]
    pub const fn span(&self) -> Span {
        match self {
            Self::NotYet(_, span)
            | Self::Unserved(_, span)
            | Self::Parse(_, span)
            | Self::WholeRecords(_, span)
            | Self::Percentile(_, span)
            | Self::Flag(_, span)
            | Self::Window(_, _, span) => *span,
        }
    }
}

impl fmt::Display for Refused {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Self::NotYet(stage, _) => write!(formatter, "{stage} is not available yet")?,
            Self::Unserved(words, _) => formatter.write_str(words)?,
            Self::Parse(message, _)
            | Self::WholeRecords(message, _)
            | Self::Percentile(message, _)
            | Self::Flag(message, _)
            | Self::Window(_, message, _) => formatter.write_str(message)?,
        }
        let span = self.span();
        write!(formatter, " at bytes {}..{}", span.start, span.end)
    }
}

/// Read a source and stages in exactly the order written.
///
/// # Errors
/// Returns a named refusal with its original byte span for invalid or unavailable syntax.
pub fn pipeline(text: &str) -> Result<Plan, Refused> {
    pipeline_with(text, |_, _| {})
}

/// Report each completed operation while applying the same pipeline rules.
pub(super) fn pipeline_with(
    text: &str,
    mut completed: impl FnMut(&str, Span),
) -> Result<Plan, Refused> {
    let mut input = Input { text, rest: text };
    input.trim();
    let (mut plan, complete) = start_pipeline(&mut input, &mut completed)?;
    if complete {
        return Ok(plan);
    }
    let mut saw_jq = false;
    let mut saw_count = false;
    loop {
        input.trim();
        let start = input.position();
        stages::position(&input, &plan, saw_count, saw_jq)?;
        let value = stages::read(&mut input)?;
        let span = Span {
            start,
            end: input.position(),
        };
        r#match::position(&value, &plan.source, plan.stages.len(), span)?;
        if matches!(value, Stage::Jq(_)) {
            if saw_jq {
                return Err(Refused::NotYet("a second jq stage", span));
            }
            saw_jq = true;
        }
        let name = match &value {
            Stage::Similar { .. } => "similar",
            Stage::Rerank { .. } => "rerank",
            Stage::Limit(_) => "limit",
            Stage::Sort { .. } => "sort",
            Stage::Count => "count",
            Stage::Percentile { .. } => "percentile",
            Stage::GroupBy { .. } => "group_by",
            Stage::Jq(_) => "jq",
            Stage::Shape(_) => "{",
            Stage::Rg { .. } => "rg",
            Stage::Match { .. } => "match",
            Stage::Window(_) => "window",
        };
        completed(name, span);
        saw_count = matches!(value, Stage::Count);
        plan.stages.push(Spanned { value, span });
        input.trim();
        if input.rest.is_empty() {
            return Ok(plan);
        }
        if saw_count && !input.rest.starts_with('|') {
            return Err(Refused::Parse(
                "nothing follows count".to_owned(),
                spans::stage(input.rest, input.position()),
            ));
        }
        input.require("|", "a stage is followed only by `|`")?;
    }
}

fn start_pipeline(
    input: &mut Input<'_>,
    completed: &mut impl FnMut(&str, Span),
) -> Result<(Plan, bool), Refused> {
    let source_start = input.position();
    let source_word = input.word().to_owned();
    let source_word = if input.rest.starts_with('{') {
        "{".to_owned()
    } else {
        source_word
    };
    let mut initial_shape = None;
    let source = if input.rest.starts_with('\'') {
        Source::All
    } else if input.rest.starts_with('{') {
        let (source, span, selection) = shape::source(input.rest, input.position())?;
        input.advance(span.end - span.start);
        initial_shape = selection.map(|selection| Spanned {
            value: Stage::Shape(selection),
            span,
        });
        source
    } else {
        source::read(input)?
    };
    let mut plan = Plan {
        source,
        stages: Vec::new(),
    };
    if let Some(stage) = initial_shape {
        plan.stages.push(stage);
    }
    if plan.source != Source::All {
        let consumed = input
            .text
            .split_at(source_start)
            .1
            .split_at(input.position() - source_start)
            .0;
        let end = source_start + consumed.trim_end().len();
        completed(
            &source_word,
            Span {
                start: source_start,
                end,
            },
        );
        input.trim();
        if input.rest.is_empty() {
            return Ok((plan, true));
        }
        input.require("|", "a collection is followed only by `|`")?;
    }
    Ok((plan, false))
}

struct Input<'a> {
    text: &'a str,
    rest: &'a str,
}

impl Input<'_> {
    const fn position(&self) -> usize {
        self.text.len() - self.rest.len()
    }

    fn trim(&mut self) {
        self.rest = self.rest.trim_start();
    }

    const fn advance(&mut self, bytes: usize) {
        self.rest = self.rest.split_at(bytes).1;
    }

    fn take(&mut self, token: &str) -> bool {
        if let Some(rest) = self.rest.strip_prefix(token) {
            self.rest = rest;
            true
        } else {
            false
        }
    }

    fn require(&mut self, token: &str, message: &str) -> Result<(), Refused> {
        self.trim();
        if self.take(token) {
            Ok(())
        } else {
            Err(self.parse(message))
        }
    }

    fn word(&self) -> &str {
        let length = self
            .rest
            .find(|character| !bare(character))
            .unwrap_or(self.rest.len());
        self.rest.split_at(length).0
    }

    fn token_span(&self) -> Span {
        let length = self.word().len();
        Span {
            start: self.position(),
            end: self.position()
                + if length == 0 {
                    self.rest.chars().next().map_or(0, char::len_utf8)
                } else {
                    length
                },
        }
    }

    fn parse(&self, message: impl Into<String>) -> Refused {
        Refused::Parse(message.into(), self.token_span())
    }

    fn not_yet(&self, name: &'static str) -> Refused {
        Refused::NotYet(name, self.token_span())
    }

    fn reserved(&self) -> Result<(), Refused> {
        let name = if self.rest.starts_with('{') {
            "{"
        } else {
            self.word()
        };
        super::words::reserved(name).map_or(Ok(()), |word| Err(self.not_yet(word)))
    }

    fn string(&mut self, what: &str) -> Result<String, Refused> {
        self.trim();
        if !self.rest.starts_with('"') {
            return Err(self.parse(format!("{what} is written in double quotes")));
        }
        let mut escaped = false;
        for (index, character) in self.rest.char_indices().skip(1) {
            if escaped {
                escaped = false;
            } else if character == '\\' {
                escaped = true;
            } else if character == '"' {
                let (literal, rest) = self.rest.split_at(index + 1);
                let value = serde_json::from_str(literal).map_err(|error| {
                    Refused::Parse(
                        format!("{what} is not a string: {error}"),
                        Span {
                            start: self.position(),
                            end: self.position() + literal.len(),
                        },
                    )
                })?;
                self.rest = rest;
                return Ok(value);
            }
        }
        Err(self.parse(format!("{what}'s quotes are not closed")))
    }
}

const fn first(character: char) -> bool {
    character.is_ascii_alphanumeric() || character == '_'
}

const fn bare(character: char) -> bool {
    first(character) || matches!(character, '.' | '-')
}

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

#[cfg(test)]
#[path = "parser_pipeline_tests.rs"]
mod pipeline_tests;

#[cfg(test)]
#[path = "parser/shape_tests.rs"]
mod shape_tests;