haematite-cli 0.21.1

use std::collections::BTreeMap;

use haematite_query::page::{Allowance, Page, Resume};
use serde_json::value::RawValue;
use serde_json::{Value, json};

use super::answer::{Head, answer, engine, fitted, named, rows, rows_bytes};
use super::plan::{Executed, Stored, present, records, scope};
use super::{
    Asked, Binding, Context, Seal, bind, frame_bytes, plan, refused, start_failure, widest_next,
    worker_engine,
};
use crate::serve::query::admission::Running;
use crate::serve::query::parser::{Plan, Refused};
use crate::serve::query::{Continuation, Settings};
use crate::serve::service::Service;
use crate::serve::storage::{Captured, Failure};

pub(super) fn worked(
    context: &Context<'_>,
    asked: &Asked<'_>,
    settings: &Settings,
    parsed: Result<Plan, Refused>,
    parameters: BTreeMap<String, Box<RawValue>>,
    service: Option<&Service>,
) -> Result<Box<RawValue>, Failure> {
    let request = asked.request;
    let at = asked.at;
    let scan = asked.scan;
    let alternate = parsed.as_ref().ok().map(|plan| plan::jq(plan).to_owned());
    let frame = frame_bytes(asked.budget)?;
    let mut running = Running::start(
        context.slots,
        Some(settings),
        frame,
        super::shape::program(&request.query, &parsed),
        alternate,
        parameters,
    )
    .map_err(|error| start_failure(error, parsed.as_ref().err()))?;
    let executed = scope(
        running.ready(),
        parsed.ok(),
        settings.memory_bytes,
        context,
        at,
        asked.now,
        scan,
    )?;
    let scope = &executed.scope;
    let binding = bind(request, &worker_engine(running.ready()), scan)?;
    if let Some(token) = asked.token {
        token.check_binding(&binding)?;
    }
    present(context.reader, at, scope)?;
    let engine = engine(running.ready());
    let resume = asked.token.map(|token| &token.resume);
    let read_before = resume.map_or(0, |resume| resume.records_read);
    let mut source = records(context, at, &executed, resume, scan, service)?;
    source.prepare_match(context, at, executed.matching.as_ref());
    let compared = source.compared();
    source.prepare(context, at, &executed.reranks, service, scan)?;
    let summary = compared.map_or(
        Value::Null,
        |compared| json!({"exact": true, "compared": compared}),
    );
    let widest = widest_next(
        at,
        named(scope),
        asked.now,
        &binding,
        context.seal,
        executed.limits().len(),
    )?;
    let mut head = Head {
        reading: executed.reading,
        stages: &executed.names,
        collection: named(scope),
        read: u64::MAX,
        skipped: u64::MAX,
        search: &summary,
        unscored: source.unscored().map(|_| u64::MAX),
        next: Some(&widest),
    };
    let envelope = answer(at, &json!([]), &engine, &head);
    let allowance = Allowance {
        rows: request.page.rows,
        bytes: rows_bytes(asked.budget, &envelope)?,
        // Comparing and reading share this page's scan allowance.
        scan_records: read_before.saturating_add(scan.saturating_sub(compared.unwrap_or(0))),
    };
    let page = read_page(&mut running, &mut source, &executed, &allowance, resume)?;
    let next = sealed_next(page.next, at, asked.now, binding, context.seal)?;
    running.stop().map_err(|error| refused(&error))?;
    let rows = rows(page.rows, request.rows)?;
    head.read = page.records_read;
    head.skipped = page.skipped;
    head.unscored = source.unscored();
    head.next = next.as_deref();
    fitted(answer(at, &json!([]), &engine, &head), &rows, asked.budget)
}

fn read_page(
    running: &mut Running<'_>,
    source: &mut Stored<'_>,
    executed: &Executed,
    allowance: &Allowance,
    resume: Option<&Resume>,
) -> Result<Page, Failure> {
    let worker = running.worker().map_err(|error| refused(&error))?;
    haematite_query::stage::read_page(worker, source, &executed.stages, allowance, resume).map_err(
        |error| {
            source
                .match_failure()
                .unwrap_or_else(|| source.failure(&error))
        },
    )
}

fn sealed_next(
    next: Option<Resume>,
    at: &Captured,
    now: &RawValue,
    binding: Binding,
    seal: &Seal,
) -> Result<Option<String>, Failure> {
    next.map(|resume| {
        Continuation::new(
            at.branch.clone(),
            at.version,
            resume,
            now.to_owned(),
            binding,
        )?
        .encode(seal)
    })
    .transpose()
}