use std::{
cell::OnceCell,
collections::{HashMap, HashSet},
};
use ahash::AHashMap;
use serde_json::{Map, Value};
use crate::{
graphql::{DirectiveEvaluator, FieldSelection},
schema::{CompiledSchema, IntrospectionBuilder, IntrospectionType, TypeKind},
};
const MAX_COMPLETION_DEPTH: usize = 64;
#[derive(Debug, Clone, PartialEq, Eq)]
enum TypeRef {
NonNull(Box<TypeRef>),
List(Box<TypeRef>),
Enum(String),
Named(String),
}
impl TypeRef {
fn from_introspection(t: &IntrospectionType) -> Option<Self> {
match t.kind {
TypeKind::NonNull => {
Some(Self::NonNull(Box::new(Self::from_introspection(t.of_type.as_deref()?)?)))
},
TypeKind::List => {
Some(Self::List(Box::new(Self::from_introspection(t.of_type.as_deref()?)?)))
},
TypeKind::Enum => t.name.clone().map(Self::Enum),
_ => t.name.clone().map(Self::Named),
}
}
}
#[derive(Debug, Clone, Default)]
pub struct OutputTypes {
fields: AHashMap<String, AHashMap<String, TypeRef>>,
possible: AHashMap<String, HashSet<String>>,
enums: AHashMap<String, EnumMembers>,
}
#[derive(Debug, Clone)]
struct EnumMembers {
strings: AHashMap<String, String>,
ints: AHashMap<i64, String>,
definition: crate::schema::EnumDefinition,
}
impl EnumMembers {
fn of(e: &crate::schema::EnumDefinition) -> Self {
let mut strings = AHashMap::new();
let mut ints = AHashMap::new();
for v in &e.values {
match &v.value {
Some(crate::schema::StoredEnumValue::Int(n)) => {
ints.insert(*n, v.name.clone());
},
Some(crate::schema::StoredEnumValue::String(s)) => {
strings.insert(s.clone(), v.name.clone());
},
None => {
strings.insert(v.name.clone(), v.name.clone());
},
}
}
Self {
strings,
ints,
definition: e.clone(),
}
}
fn member_of(&self, stored: &Value) -> Option<&String> {
match stored {
Value::String(s) => self.strings.get(s.as_str()),
Value::Number(n) => n.as_i64().and_then(|n| self.ints.get(&n)),
_ => None,
}
}
}
impl OutputTypes {
#[must_use]
pub fn from_schema(schema: &CompiledSchema) -> Self {
let introspection = IntrospectionBuilder::build(schema);
let mut types = Self::default();
for t in &introspection.types {
let Some(name) = &t.name else {
continue;
};
if let Some(fields) = &t.fields {
let fields = fields
.iter()
.filter_map(|f| {
TypeRef::from_introspection(&f.field_type).map(|r| (f.name.clone(), r))
})
.collect();
types.fields.insert(name.clone(), fields);
}
if let Some(possible) = &t.possible_types {
types
.possible
.insert(name.clone(), possible.iter().map(|p| p.name.clone()).collect());
}
}
for e in &schema.enums {
types.enums.insert(e.name.clone(), EnumMembers::of(e));
}
#[cfg(feature = "federation")]
types.add_federation(schema);
types
}
#[cfg(feature = "federation")]
fn add_federation(&mut self, schema: &CompiledSchema) {
let Some(metadata) = schema.federation_metadata() else {
return;
};
let entity = TypeRef::Named("_Entity".to_string());
self.fields.entry("Query".to_string()).or_default().insert(
"_entities".to_string(),
TypeRef::NonNull(Box::new(TypeRef::List(Box::new(entity)))),
);
self.possible
.insert("_Entity".to_string(), metadata.types.iter().map(|t| t.name.clone()).collect());
}
fn type_fields(&self, type_name: &str) -> Option<(&str, &AHashMap<String, TypeRef>)> {
self.fields
.get_key_value(type_name)
.map(|(name, fields)| (name.as_str(), fields))
}
pub fn complete(
&self,
response: &mut Value,
root_type: &str,
selections: &[FieldSelection],
variables: Option<&Value>,
) {
let walk = Walk::new(self, variables);
let mut errors = Vec::new();
let Some(data) = response.get_mut("data").and_then(Value::as_object_mut) else {
return;
};
let mut null_data = false;
let mut path = Vec::new();
if let Some((root_type, fields)) = self.type_fields(root_type) {
walk.for_each_field(selections, root_type, &mut |sel| {
let Some(field_ref) = fields.get(&sel.name) else {
return;
};
let key = sel.response_key();
let Some(value) = data.get_mut(key) else {
return;
};
path.push(Step::Key(key));
let label = Label::Field(root_type, &sel.name);
if walk.position(value, field_ref, sel, label, &mut path, &mut errors, 0).is_err() {
null_data = true;
}
path.pop();
});
}
if null_data {
response["data"] = Value::Null;
}
if errors.is_empty() {
return;
}
match response.get_mut("errors").and_then(Value::as_array_mut) {
Some(existing) => existing.extend(errors),
None => {
if let Some(object) = response.as_object_mut() {
object.insert("errors".to_string(), Value::Array(errors));
}
},
}
}
}
impl OutputTypes {
#[must_use]
pub fn complete_value(
&self,
mut value: Value,
type_name: &str,
selections: &[FieldSelection],
root_key: &str,
) -> (Value, Vec<Value>) {
let walk = Walk::new(self, None);
let sel = FieldSelection {
name: root_key.to_string(),
alias: None,
arguments: Vec::new(),
nested_fields: selections.to_vec(),
directives: Vec::new(),
};
let mut errors = Vec::new();
let mut path = vec![Step::Key(root_key)];
let field_ref = TypeRef::Named(type_name.to_string());
if walk
.position(
&mut value,
&field_ref,
&sel,
Label::Root(root_key),
&mut path,
&mut errors,
0,
)
.is_err()
{
value = Value::Null;
}
(value, errors)
}
}
struct Propagate;
#[derive(Clone, Copy)]
enum Step<'a> {
Key(&'a str),
Index(usize),
}
#[derive(Clone, Copy)]
enum Label<'a> {
Field(&'a str, &'a str),
Root(&'a str),
}
impl std::fmt::Display for Label<'_> {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
Self::Field(owner, field) => write!(f, "{owner}.{field}"),
Self::Root(key) => f.write_str(key),
}
}
}
fn violation(label: Label<'_>, path: &[Step<'_>]) -> Value {
serde_json::json!({
"message": format!("Cannot return null for non-nullable field {label}."),
"path": response_path(path),
})
}
fn not_a_member(
label: Label<'_>,
path: &[Step<'_>],
stored: &Value,
members: &EnumMembers,
) -> Value {
serde_json::json!({
"message": members.definition.not_a_member_message(stored, &label),
"path": response_path(path),
})
}
fn response_path(path: &[Step<'_>]) -> Vec<Value> {
path.iter()
.map(|step| match *step {
Step::Key(key) => Value::String(key.to_string()),
Step::Index(index) => Value::from(index),
})
.collect()
}
struct Walk<'a> {
types: &'a OutputTypes,
variables: Option<&'a Value>,
directive_variables: OnceCell<HashMap<String, Value>>,
}
impl<'a> Walk<'a> {
const fn new(types: &'a OutputTypes, variables: Option<&'a Value>) -> Self {
Self {
types,
variables,
directive_variables: OnceCell::new(),
}
}
fn included(&self, sel: &FieldSelection) -> bool {
if sel.directives.is_empty() {
return true;
}
let variables = self.directive_variables.get_or_init(|| {
self.variables
.and_then(Value::as_object)
.map(|map| map.iter().map(|(k, v)| (k.clone(), v.clone())).collect())
.unwrap_or_default()
});
DirectiveEvaluator::evaluate_directives(sel, variables).unwrap_or(true)
}
fn for_each_field<'s>(
&self,
selections: &'s [FieldSelection],
type_name: &str,
visit: &mut impl FnMut(&'s FieldSelection),
) {
for sel in selections {
if !self.included(sel) {
continue;
}
if let Some(condition) = sel.name.strip_prefix("...on ") {
if self.satisfies(type_name, condition.trim()) {
self.for_each_field(&sel.nested_fields, type_name, visit);
}
continue;
}
if sel.name.starts_with("...") || sel.name.starts_with("__") {
continue;
}
visit(sel);
}
}
fn satisfies(&self, type_name: &str, condition: &str) -> bool {
condition == type_name
|| self.types.possible.get(condition).is_some_and(|p| p.contains(type_name))
}
#[allow(clippy::too_many_arguments)] fn position<'s>(
&self,
value: &mut Value,
field_ref: &TypeRef,
sel: &'s FieldSelection,
label: Label<'s>,
path: &mut Vec<Step<'s>>,
errors: &mut Vec<Value>,
depth: usize,
) -> Result<(), Propagate>
where
'a: 's,
{
match field_ref {
TypeRef::NonNull(inner) => {
self.nullable(value, inner, sel, label, path, errors, depth)?;
if value.is_null() {
errors.push(violation(label, path));
return Err(Propagate);
}
Ok(())
},
other => {
if self.nullable(value, other, sel, label, path, errors, depth).is_err() {
*value = Value::Null;
}
Ok(())
},
}
}
#[allow(clippy::too_many_arguments)] fn nullable<'s>(
&self,
value: &mut Value,
field_ref: &TypeRef,
sel: &'s FieldSelection,
label: Label<'s>,
path: &mut Vec<Step<'s>>,
errors: &mut Vec<Value>,
depth: usize,
) -> Result<(), Propagate>
where
'a: 's,
{
if value.is_null() || depth > MAX_COMPLETION_DEPTH {
return Ok(());
}
match field_ref {
TypeRef::NonNull(_) => self.position(value, field_ref, sel, label, path, errors, depth),
TypeRef::Enum(name) => {
let Some(members) = self.types.enums.get(name.as_str()) else {
return Ok(());
};
let Some(member) = members.member_of(value) else {
errors.push(not_a_member(label, path, value, members));
return Err(Propagate);
};
if value.as_str() != Some(member.as_str()) {
*value = Value::String(member.clone());
}
Ok(())
},
TypeRef::List(inner) => {
let Value::Array(items) = value else {
return Ok(());
};
let mut propagated = false;
for (index, item) in items.iter_mut().enumerate() {
path.push(Step::Index(index));
if self.position(item, inner, sel, label, path, errors, depth + 1).is_err() {
propagated = true;
}
path.pop();
}
if propagated { Err(Propagate) } else { Ok(()) }
},
TypeRef::Named(name) => match value {
Value::Object(map) => self.object(map, name, sel, path, errors, depth),
_ => Ok(()),
},
}
}
fn object<'s>(
&self,
map: &mut Map<String, Value>,
declared: &str,
sel: &'s FieldSelection,
path: &mut Vec<Step<'s>>,
errors: &mut Vec<Value>,
depth: usize,
) -> Result<(), Propagate>
where
'a: 's,
{
let types = self.types;
let Some((concrete, fields)) = map
.get("__typename")
.and_then(Value::as_str)
.and_then(|t| types.type_fields(t))
.or_else(|| types.type_fields(declared))
else {
return Ok(());
};
let mut propagated = false;
self.for_each_field(&sel.nested_fields, concrete, &mut |child| {
let Some(field_ref) = fields.get(&child.name) else {
return;
};
let key = child.response_key();
path.push(Step::Key(key));
let label = Label::Field(concrete, &child.name);
let completed = match map.get_mut(key) {
Some(value) => {
self.position(value, field_ref, child, label, path, errors, depth + 1)
},
None => self.position(
&mut Value::Null,
field_ref,
child,
label,
path,
errors,
depth + 1,
),
};
path.pop();
if completed.is_err() {
propagated = true;
}
});
if propagated { Err(Propagate) } else { Ok(()) }
}
}
mod rows;
#[cfg(test)]
mod tests;