use std::collections::HashMap;
use serde_json::Value as JsonValue;
use thiserror::Error;
use crate::graphql::{
DirectiveError, DirectiveEvaluator, FragmentError, FragmentResolver,
types::{FieldSelection, FragmentDefinition},
};
#[derive(Debug, Error)]
#[non_exhaustive]
pub enum SelectionError {
#[error(transparent)]
Fragment(#[from] FragmentError),
#[error(transparent)]
Directive(#[from] DirectiveError),
}
impl From<SelectionError> for crate::error::FraiseQLError {
fn from(err: SelectionError) -> Self {
let path = match &err {
SelectionError::Fragment(_) => "fragments",
SelectionError::Directive(_) => "directives",
};
Self::Validation {
message: err.to_string(),
path: Some(path.to_string()),
}
}
}
pub fn resolve(
selections: &[FieldSelection],
fragments: &[FragmentDefinition],
max_depth: u32,
) -> Result<Vec<FieldSelection>, SelectionError> {
Ok(FragmentResolver::new(fragments)
.with_max_depth(max_depth)
.resolve_spreads(selections)?)
}
#[allow(clippy::implicit_hasher)]
pub fn filter(
selections: &[FieldSelection],
variables: &HashMap<String, JsonValue>,
) -> Result<Vec<FieldSelection>, SelectionError> {
Ok(merge_by_response_key(DirectiveEvaluator::filter_selections(
selections, variables,
)?))
}
fn merge_by_response_key(selections: Vec<FieldSelection>) -> Vec<FieldSelection> {
let mut merged: Vec<FieldSelection> = Vec::with_capacity(selections.len());
for selection in selections {
let same_field = merged.iter_mut().find(|m| {
m.response_key() == selection.response_key()
&& m.name == selection.name
&& m.arguments == selection.arguments
});
match same_field {
Some(first) => first.nested_fields.extend(selection.nested_fields),
None => merged.push(selection),
}
}
for field in &mut merged {
if !field.nested_fields.is_empty() {
field.nested_fields = merge_by_response_key(std::mem::take(&mut field.nested_fields));
}
}
merged
}
#[allow(clippy::implicit_hasher)]
pub fn resolve_and_filter(
selections: &[FieldSelection],
fragments: &[FragmentDefinition],
variables: &HashMap<String, JsonValue>,
max_depth: u32,
) -> Result<Vec<FieldSelection>, SelectionError> {
filter(&resolve(selections, fragments, max_depth)?, variables)
}
#[must_use]
pub fn variables_map(variables: Option<&JsonValue>) -> HashMap<String, JsonValue> {
match variables {
Some(JsonValue::Object(map)) => map.iter().map(|(k, v)| (k.clone(), v.clone())).collect(),
_ => HashMap::new(),
}
}
#[derive(Debug)]
pub struct CollectedFields {
written: Vec<FieldSelection>,
collected: std::sync::Mutex<Vec<(Outcomes, std::sync::Arc<Collected>)>>,
}
#[derive(Debug)]
pub struct Collected {
fields: std::sync::Arc<[FieldSelection]>,
checked: [std::sync::OnceLock<()>; 2],
projection: std::sync::OnceLock<(ProjectionKey, KeptProjection)>,
}
#[derive(Debug, Clone)]
pub struct KeptProjection {
pub hint: std::sync::Arc<crate::schema::SqlProjectionHint>,
pub decimals_as_text: bool,
pub built: bool,
}
#[derive(Debug, Clone, Copy)]
pub enum FieldsCheck {
Document = 0,
Locale = 1,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ProjectionKey {
pub request_chain: Option<Vec<String>>,
masked: Vec<(String, String)>,
}
impl ProjectionKey {
#[must_use]
pub fn new(
request_chain: Option<Vec<String>>,
masked: impl IntoIterator<Item = (String, String)>,
) -> Self {
let mut masked: Vec<_> = masked.into_iter().collect();
masked.sort_unstable();
masked.dedup();
Self {
request_chain,
masked,
}
}
#[must_use]
pub fn masks(&self, type_name: &str, field: &str) -> bool {
self.masked
.binary_search_by(|(t, f)| (t.as_str(), f.as_str()).cmp(&(type_name, field)))
.is_ok()
}
#[must_use]
pub fn masked(&self) -> &[(String, String)] {
&self.masked
}
}
impl Collected {
fn new(fields: Vec<FieldSelection>) -> std::sync::Arc<Self> {
std::sync::Arc::new(Self {
fields: fields.into(),
checked: [std::sync::OnceLock::new(), std::sync::OnceLock::new()],
projection: std::sync::OnceLock::new(),
})
}
pub fn projection_once<E>(
&self,
key: &ProjectionKey,
build: impl FnOnce(&[FieldSelection], &ProjectionKey) -> Result<KeptProjection, E>,
) -> Result<KeptProjection, E> {
if let Some((kept_key, kept)) = self.projection.get() {
if kept_key == key {
return Ok(kept.clone());
}
return build(&self.fields, key);
}
let built = build(&self.fields, key)?;
let (kept_key, kept) = self.projection.get_or_init(|| (key.clone(), built.clone()));
Ok(if kept_key == key { kept.clone() } else { built })
}
#[must_use]
pub const fn fields(&self) -> &std::sync::Arc<[FieldSelection]> {
&self.fields
}
pub fn check_once<E>(
&self,
which: FieldsCheck,
check: impl FnOnce(&[FieldSelection]) -> Result<(), E>,
) -> Result<(), E> {
let passed = &self.checked[which as usize];
if passed.get().is_none() {
check(&self.fields)?;
let _ = passed.set(());
}
Ok(())
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
struct Outcomes {
bits: u64,
count: u32,
}
impl CollectedFields {
const KEPT: usize = 16;
#[must_use]
pub const fn new(written: Vec<FieldSelection>) -> Self {
Self {
written,
collected: std::sync::Mutex::new(Vec::new()),
}
}
#[must_use]
pub fn written(&self) -> &[FieldSelection] {
&self.written
}
#[allow(clippy::implicit_hasher)]
pub fn collect(
&self,
variables: &HashMap<String, JsonValue>,
) -> Result<std::sync::Arc<Collected>, SelectionError> {
let mut outcomes = Some(Outcomes { bits: 0, count: 0 });
record_outcomes(&self.written, variables, &mut outcomes)?;
let Some(outcomes) = outcomes else {
return Ok(Collected::new(filter(&self.written, variables)?));
};
let mut collected =
self.collected.lock().unwrap_or_else(std::sync::PoisonError::into_inner);
if let Some((_, entry)) = collected.iter().find(|(kept, _)| *kept == outcomes) {
return Ok(entry.clone());
}
let entry = Collected::new(filter(&self.written, variables)?);
if collected.len() < Self::KEPT {
collected.push((outcomes, entry.clone()));
}
Ok(entry)
}
}
fn record_outcomes(
selections: &[FieldSelection],
variables: &HashMap<String, JsonValue>,
outcomes: &mut Option<Outcomes>,
) -> Result<(), DirectiveError> {
for selection in selections {
if !selection.directives.is_empty() {
let kept = DirectiveEvaluator::evaluate_directives(selection, variables)?;
if let Some(o) = outcomes {
if o.count == u64::BITS {
*outcomes = None;
} else {
o.bits |= u64::from(kept) << o.count;
o.count += 1;
}
}
if !kept {
continue;
}
}
record_outcomes(&selection.nested_fields, variables, outcomes)?;
}
Ok(())
}
#[cfg(test)]
#[path = "selection_set_tests.rs"]
mod selection_set_tests;