haematite-cli 0.21.8

use haematite_query::stage::shape::Selection;
use serde_json::Value;

use super::{Input, Refused, Span, first};

#[derive(Debug, PartialEq, Eq)]
pub enum Source {
    All,
    Shaped(Vec<Root>),
    Collection {
        name: String,
        arguments: Vec<Argument>,
    },
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Root {
    pub name: String,
    pub arguments: Vec<Argument>,
    pub selection: Selection,
    pub span: Span,
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Argument {
    pub field: String,
    pub value: Value,
}

pub(super) fn read(input: &mut Input<'_>) -> Result<Source, Refused> {
    if input.word() == "match" {
        return Ok(Source::All);
    }
    input.reserved()?;
    match input.word() {
        "percentile" => {
            return Err(input.parse("`percentile` is a stage and must follow a collection"));
        }
        "similar" | "sort" | "count" | "limit" | "window" | "rerank" => {
            return Err(input.parse(format!(
                "`{}` is a stage and must follow a collection",
                input.word()
            )));
        }
        "rg" | "grep" => {
            return Err(input.parse("`rg` is a stage and must follow a collection"));
        }
        "group_by" => {
            return Err(input.parse("`group_by` is a stage and must follow a collection"));
        }
        _ => {}
    }
    collection(input)
}

/// Root braces identify a collection, even when its name is a stage word.
pub(super) fn collection(input: &mut Input<'_>) -> Result<Source, Refused> {
    let name = input.word().to_owned();
    if !name.starts_with(first) {
        return Err(input.parse("a query that is not jq starts with a collection"));
    }
    input.advance(name.len());
    input.trim();
    if !input.take("(") {
        return Ok(Source::Collection {
            name,
            arguments: Vec::new(),
        });
    }
    if name == "collection" {
        input.trim();
        let start = input.position();
        let name = input.string("a collection name")?;
        if name.is_empty() || name.contains('\0') {
            return Err(Refused::Parse(
                "a collection name has 1 or more characters and no NUL".to_owned(),
                Span {
                    start,
                    end: input.position(),
                },
            ));
        }
        input.require(")", "`collection(` takes one quoted name and a `)`")?;
        return Ok(Source::Collection {
            name,
            arguments: Vec::new(),
        });
    }
    let mut arguments = Vec::new();
    loop {
        input.trim();
        let length = input
            .rest
            .find(|character| !first(character))
            .unwrap_or(input.rest.len());
        let field = input.rest.split_at(length).0.to_owned();
        if field.is_empty() {
            return Err(input.parse("a collection argument names a top-level field before `:`"));
        }
        input.advance(length);
        input.require(
            ":",
            "a collection argument separates its field and JSON value with `:`",
        )?;
        input.trim();
        let value = value(input)?;
        arguments.push(Argument { field, value });
        input.trim();
        if input.take(")") {
            return Ok(Source::Collection { name, arguments });
        }
        if input.word() != "and" {
            return Err(input.parse("collection arguments are joined by `and` and closed by `)`"));
        }
        input.advance(3);
    }
}

/// Find argument boundaries outside JSON strings, arrays and objects.
fn value(input: &mut Input<'_>) -> Result<Value, Refused> {
    let mut quoted = false;
    let mut escaped = false;
    let mut depth = 0_usize;
    let mut end = input.rest.len();
    for (index, character) in input.rest.char_indices() {
        if quoted {
            if escaped {
                escaped = false;
            } else if character == '\\' {
                escaped = true;
            } else if character == '"' {
                quoted = false;
            }
            continue;
        }
        let (before, after) = input.rest.split_at(index);
        match character {
            '"' => quoted = true,
            '[' | '{' => depth += 1,
            ']' | '}' => depth = depth.saturating_sub(1),
            ')' if depth == 0 => {
                end = index;
                break;
            }
            'a' if depth == 0
                && before.ends_with(char::is_whitespace)
                && after
                    .strip_prefix("and")
                    .is_some_and(|rest| rest.starts_with(char::is_whitespace)) =>
            {
                end = index;
                break;
            }
            _ => {}
        }
    }
    let literal = input.rest.split_at(end).0.trim_end();
    let span = Span {
        start: input.position(),
        end: input.position() + literal.len(),
    };
    let value = serde_json::from_str(literal).map_err(|error| {
        Refused::Parse(
            format!("a collection argument needs a JSON value: {error}"),
            span,
        )
    })?;
    input.advance(end);
    Ok(value)
}

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