haematite-cli 0.21.3

//! Execute labelled records and directed chains at the captured branch version.

use std::collections::{BTreeSet, VecDeque};
use std::num::NonZeroUsize;

use haematite::db::ReadOnlyDatabase;
use haematite_query::Refusal;
use haematite_query::handle::Error;
use haematite_query::page::Record;
use haematite_query::stage::matching::bind;
use serde_json::value::RawValue;

use super::super::parser::{Direction, Node as PatternNode, Plan, Relation, Stage};
use super::super::source::Scope;
use super::{Context, refused};
use crate::serve::access::Access;
use crate::serve::contracts::identifier;
use crate::serve::protocol::Node;
use crate::serve::storage::{Captured, Failure, document_key, edges, reads};

#[derive(Clone)]
pub(super) struct Request {
    pub collection: String,
    binding: String,
    hops: Vec<Hop>,
}

#[derive(Clone)]
struct Hop {
    collection: String,
    binding: String,
    kind: String,
    direction: Direction,
}

impl Request {
    /// Prepare validated labels and bindings for one directed match.
    ///
    /// # Errors
    /// Refuses invalid labels, repeated bindings, or unsupported node and relation shapes.
    pub fn new(nodes: &[PatternNode], relations: &[Relation]) -> Result<Self, Failure> {
        let Some(first) = nodes.first() else {
            return Err(pattern());
        };
        if nodes.len() > 8 || relations.len() + 1 != nodes.len() {
            return Err(pattern());
        }
        let mut bindings = BTreeSet::new();
        for node in nodes {
            identifier(&node.label.value)?;
            if !bindings.insert(&node.binding.value) {
                return Err(Failure::new(
                    "query_match_binding",
                    "match cannot repeat a binding name",
                ));
            }
        }
        let hops = nodes[1..]
            .iter()
            .zip(relations)
            .map(|(far, relation)| {
                identifier(&relation.kind.value)?;
                Ok(Hop {
                    collection: far.label.value.clone(),
                    binding: far.binding.value.clone(),
                    kind: relation.kind.value.clone(),
                    direction: relation.direction,
                })
            })
            .collect::<Result<_, Failure>>()?;
        Ok(Self {
            collection: first.label.value.clone(),
            binding: first.binding.value.clone(),
            hops,
        })
    }

    pub const fn is_hop(&self) -> bool {
        !self.hops.is_empty()
    }
}

/// Check all labels before worker admission, including an unreachable far label.
pub(super) fn check(reader: &ReadOnlyDatabase, at: &Captured, plan: &Plan) -> Result<(), Failure> {
    for stage in &plan.stages {
        if let Stage::Match { nodes, relations } = &stage.value {
            Request::new(nodes, relations)?;
            for node in nodes {
                if reads::collection(reader, at, &node.label.value)?.is_none() {
                    return Err(Failure::new(
                        "collection_missing",
                        format!(
                            "collection {} does not exist at the captured version",
                            node.label.value
                        ),
                    ));
                }
            }
        }
    }
    Ok(())
}

pub(super) struct Prepared<'a> {
    request: Request,
    access: &'a Access,
    reader: &'a ReadOnlyDatabase,
    at: &'a Captured,
    failure: Option<Failure>,
}

impl<'a> Prepared<'a> {
    pub fn new(context: &Context<'a>, at: &'a Captured, request: &Request) -> Self {
        Self {
            request: request.clone(),
            access: context.access,
            reader: context.reader,
            at,
            failure: None,
        }
    }

    /// Preserve the storage refusal code across the worker-stage boundary.
    ///
    /// # Errors
    /// Returns a named protocol, scan, binding, or storage refusal for an incomplete match.
    pub fn expand(
        &mut self,
        record: Record,
        remaining: u64,
    ) -> Result<(VecDeque<Record>, u64), Error> {
        self.expand_stored(record, remaining).map_err(|failure| {
            let error = Error::new(Refusal::QueryRuntime, failure.message.clone());
            self.failure = Some(failure);
            error
        })
    }

    pub const fn take_failure(&mut self) -> Option<Failure> {
        self.failure.take()
    }

    fn expand_stored(
        &self,
        mut record: Record,
        remaining: u64,
    ) -> Result<(VecDeque<Record>, u64), Failure> {
        if record.origin.collection != self.request.collection {
            return Err(Failure::new(
                "query_protocol",
                "match received a record outside its first label",
            ));
        }
        record.origin.score = None;
        if self.request.hops.is_empty() {
            record.value = bind([(self.request.binding.as_str(), record.value.as_ref())])
                .map_err(|error| refused(&error))?;
            return Ok((VecDeque::from([record]), 0));
        }
        self.hop_records(&record, remaining)
    }

    fn hop_records(
        &self,
        record: &Record,
        remaining: u64,
    ) -> Result<(VecDeque<Record>, u64), Failure> {
        let node = Node {
            collection: record.origin.collection.clone(),
            document_id: record.origin.id.clone(),
        };
        let mut bindings = vec![(self.request.binding.clone(), record.value.clone())];
        let mut walked = Walk {
            rows: VecDeque::new(),
            read: 0,
            remaining,
        };
        self.walk(
            record,
            &node,
            &self.request.hops,
            &mut bindings,
            &mut walked,
        )?;
        Ok((walked.rows, walked.read))
    }

    fn walk(
        &self,
        record: &Record,
        node: &Node,
        hops: &[Hop],
        bindings: &mut Vec<(String, Box<RawValue>)>,
        walked: &mut Walk,
    ) -> Result<(), Failure> {
        let Some((hop, rest)) = hops.split_first() else {
            let value = bind(
                bindings
                    .iter()
                    .map(|(name, value)| (name.as_str(), value.as_ref())),
            )
            .map_err(|error| refused(&error))?;
            walked.rows.push_back(Record {
                value,
                origin: record.origin.clone(),
            });
            return Ok(());
        };
        let mut after = None;
        loop {
            let page = match hop.direction {
                Direction::Outgoing => edges::out(
                    self.reader,
                    self.at,
                    node,
                    Some(&hop.kind),
                    after.as_deref(),
                    NonZeroUsize::MIN,
                    None,
                )?,
                Direction::Incoming => edges::incoming(
                    self.reader,
                    self.at,
                    node,
                    Some(&hop.kind),
                    after.as_deref(),
                    NonZeroUsize::MIN,
                    None,
                )?,
            };
            for edge in page.entries {
                let from_allowed = self
                    .access
                    .can_read_document(&edge.from.collection, &edge.from.document_id);
                let to_allowed = self
                    .access
                    .can_read_document(&edge.to.collection, &edge.to.document_id);
                if !from_allowed || !to_allowed {
                    continue;
                }
                charge(&mut walked.read, walked.remaining)?;
                let far = match hop.direction {
                    Direction::Outgoing => edge.to,
                    Direction::Incoming => edge.from,
                };
                if far.collection != hop.collection {
                    continue;
                }
                charge(&mut walked.read, walked.remaining)?;
                bindings.push((hop.binding.clone(), self.far_record(&far)?));
                self.walk(record, &far, rest, bindings, walked)?;
                bindings.pop().ok_or_else(|| {
                    Failure::new(
                        "query_protocol",
                        "match lost the binding for a completed hop",
                    )
                })?;
            }
            let Some(next) = page.next_after else { break };
            if after.as_ref() == Some(&next) {
                return Err(Failure::new(
                    "query_protocol",
                    "match edge page did not advance",
                ));
            }
            after = Some(next);
        }
        Ok(())
    }

    fn far_record(&self, far: &Node) -> Result<Box<RawValue>, Failure> {
        let bytes = reads::get_bytes(
            self.reader,
            self.at,
            &document_key(&far.collection, &far.document_id),
        )?
        .ok_or_else(|| {
            Failure::new(
                "edge_endpoint_missing",
                "match edge has no far record at the captured version",
            )
        })?;
        serde_json::from_slice(&bytes).map_err(Failure::storage)
    }
}

/// Derive the source scope from the grammar's first label.
pub(super) fn mapped(
    nodes: &[PatternNode],
    relations: &[Relation],
    scope: &mut Scope,
) -> Result<(Request, haematite_query::stage::Stage), Failure> {
    let request = Request::new(nodes, relations)?;
    *scope = Scope::Collection(request.collection.clone());
    let stage = haematite_query::stage::Stage::Match {
        hop: request.is_hop(),
    };
    Ok((request, stage))
}

struct Walk {
    rows: VecDeque<Record>,
    read: u64,
    remaining: u64,
}

fn pattern() -> Failure {
    Failure::new(
        "query_match_pattern",
        "match needs one to eight nodes joined by a directed relation at every hop",
    )
}

fn charge(read: &mut u64, remaining: u64) -> Result<(), Failure> {
    let next = read.checked_add(1).ok_or_else(scan_limit)?;
    if next > remaining {
        return Err(scan_limit());
    }
    *read = next;
    Ok(())
}

fn scan_limit() -> Failure {
    Failure::new(
        "query_scan_limit",
        "match record and edge reads exceed the scan allowance",
    )
}

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