graph, graphql: Use Arc<ObjctType> to reference object types

Rather than use a string name, use the actual object type to identify
types. It's not possible to do this with plain references, for example,
because we pass a reference to a SelectionSet to graph::spawn_blocking, so
we do the next best thing and use an Arc. Unfortunately, because
`graphql_parser` doesn't wrap its object types in an Arc, that means we
need to keep a copy of all of them in ApiSchema.

Author: lutter

graph/src/data/schema.rs

@@ -1,3 +1,4 @@
+use crate::cheap_clone::CheapClone;
 use crate::components::store::{EntityType, SubgraphStore};
 use crate::data::graphql::ext::{DirectiveExt, DirectiveFinder, DocumentExt, TypeExt, ValueExt};
 use crate::data::store::ValueType;
@@ -353,6 +354,7 @@ pub struct ApiSchema {
     // Root types for the api schema.
     pub query_type: Arc<ObjectType>,
     pub subscription_type: Option<Arc<ObjectType>>,
+    object_types: HashMap<String, Arc<ObjectType>>,
 }
 
 impl ApiSchema {
@@ -376,10 +378,19 @@ impl ApiSchema {
             .cloned()
             .map(Arc::new);
 
+        let object_types = HashMap::from_iter(
+            api_schema
+                .document
+                .get_object_type_definitions()
+                .into_iter()
+                .map(|obj_type| (obj_type.name.clone(), Arc::new(obj_type.clone()))),
+        );
+
         Ok(Self {
             schema: api_schema,
             query_type: Arc::new(query_type),
             subscription_type,
+            object_types,
         })
     }
 
@@ -403,6 +414,17 @@ impl ApiSchema {
     pub fn interfaces_for_type(&self, type_name: &EntityType) -> Option<&Vec<InterfaceType>> {
         self.schema.interfaces_for_type(type_name)
     }
+
+    /// Return an `Arc` around the `ObjectType` from our internal cache
+    ///
+    /// # Panics
+    /// If `obj_type` is not part of this schema, this function panics
+    pub fn object_type(&self, obj_type: &ObjectType) -> Arc<ObjectType> {
+        self.object_types
+            .get(&obj_type.name)
+            .expect("ApiSchema.object_type is only used with existing types")
+            .cheap_clone()
+    }
 }
 
 fn add_introspection_schema(schema: &mut Document) {

graphql/src/execution/ast.rs

@@ -3,7 +3,7 @@ use std::{collections::HashSet, ops::Deref};
 use graph::{
     components::store::EntityType,
     data::graphql::{DocumentExt, ObjectOrInterface},
-    prelude::{anyhow, q, r, s, QueryExecutionError, Schema, ValueMap},
+    prelude::{anyhow, q, r, s, ApiSchema, QueryExecutionError, ValueMap},
 };
 use graphql_parser::Pos;
 
@@ -23,13 +23,16 @@ use crate::schema::ast::ObjectType;
 pub struct SelectionSet {
     // Map object types to the list of fields that should be selected for
     // them
-    items: Vec<(String, Vec<Field>)>,
+    items: Vec<(ObjectType, Vec<Field>)>,
 }
 
 impl SelectionSet {
     /// Create a new `SelectionSet` that can handle the given types
-    pub fn new(types: Vec<String>) -> Self {
-        let items = types.into_iter().map(|name| (name, Vec::new())).collect();
+    pub fn new(types: Vec<ObjectType>) -> Self {
+        let items = types
+            .into_iter()
+            .map(|obj_type| (obj_type, Vec::new()))
+            .collect();
         SelectionSet { items }
     }
 
@@ -76,29 +79,28 @@ impl SelectionSet {
     }
 
     /// Iterate over all types and the fields for those types
-    pub fn fields(&self) -> impl Iterator<Item = (&str, impl Iterator<Item = &Field>)> {
+    pub fn fields(&self) -> impl Iterator<Item = (&ObjectType, impl Iterator<Item = &Field>)> {
         self.items
             .iter()
-            .map(|(name, fields)| (name.as_str(), fields.iter()))
+            .map(|(obj_type, fields)| (obj_type, fields.iter()))
     }
 
     /// Iterate over all types and the fields that are not leaf fields, i.e.
     /// whose selection sets are not empty
-    pub fn interior_fields(&self) -> impl Iterator<Item = (&str, impl Iterator<Item = &Field>)> {
-        self.items.iter().map(|(name, fields)| {
-            (
-                name.as_str(),
-                fields.iter().filter(|field| !field.is_leaf()),
-            )
-        })
+    pub fn interior_fields(
+        &self,
+    ) -> impl Iterator<Item = (&ObjectType, impl Iterator<Item = &Field>)> {
+        self.items
+            .iter()
+            .map(|(obj_type, fields)| (obj_type, fields.iter().filter(|field| !field.is_leaf())))
     }
 
     /// Iterate over all fields for the given object type
-    pub fn fields_for(&self, obj_type: &s::ObjectType) -> impl Iterator<Item = &Field> {
+    pub fn fields_for(&self, obj_type: &ObjectType) -> impl Iterator<Item = &Field> {
         let item = self
             .items
             .iter()
-            .find(|(name, _)| name == &obj_type.name)
+            .find(|(our_type, _)| our_type == obj_type)
             .expect("there is an entry for the type");
         item.1.iter()
     }
@@ -254,14 +256,14 @@ impl ValueMap for Field {
 /// object types that implement them, and possibly narrowing further when
 /// expanding fragments with type conitions
 #[derive(Debug, Clone, PartialEq)]
-pub enum ObjectTypeSet {
+pub(crate) enum ObjectTypeSet {
     Any,
-    Only(HashSet<String>),
+    Only(HashSet<ObjectType>),
 }
 
 impl ObjectTypeSet {
     pub fn convert(
-        schema: &Schema,
+        schema: &ApiSchema,
         type_cond: Option<&q::TypeCondition>,
     ) -> Result<ObjectTypeSet, QueryExecutionError> {
         match type_cond {
@@ -270,64 +272,60 @@ impl ObjectTypeSet {
         }
     }
 
-    pub fn from_name(schema: &Schema, name: &str) -> Result<ObjectTypeSet, QueryExecutionError> {
-        let set = resolve_object_types(schema, name)?
-            .into_iter()
-            .map(|ty| ty.name().to_string())
-            .collect();
+    pub fn from_name(schema: &ApiSchema, name: &str) -> Result<ObjectTypeSet, QueryExecutionError> {
+        let set = resolve_object_types(schema, name)?;
         Ok(ObjectTypeSet::Only(set))
     }
 
-    fn matches_name(&self, name: &str) -> bool {
+    fn contains(&self, obj_type: &ObjectType) -> bool {
         match self {
             ObjectTypeSet::Any => true,
-            ObjectTypeSet::Only(set) => set.contains(name),
+            ObjectTypeSet::Only(set) => set.contains(obj_type),
         }
     }
 
     pub fn intersect(self, other: &ObjectTypeSet) -> ObjectTypeSet {
         match self {
             ObjectTypeSet::Any => other.clone(),
-            ObjectTypeSet::Only(set) => ObjectTypeSet::Only(
-                set.into_iter()
-                    .filter(|ty| other.matches_name(ty))
-                    .collect(),
-            ),
+            ObjectTypeSet::Only(set) => {
+                ObjectTypeSet::Only(set.into_iter().filter(|ty| other.contains(ty)).collect())
+            }
         }
     }
 
     /// Return a list of the object type names that are in this type set and
     /// are also implementations of `current_type`
     pub fn type_names(
         &self,
-        schema: &Schema,
+        schema: &ApiSchema,
         current_type: ObjectOrInterface<'_>,
-    ) -> Result<Vec<String>, QueryExecutionError> {
+    ) -> Result<Vec<ObjectType>, QueryExecutionError> {
         Ok(resolve_object_types(schema, current_type.name())?
             .into_iter()
-            .map(|obj| obj.name().to_string())
-            .filter(|name| match self {
+            .filter(|obj_type| match self {
                 ObjectTypeSet::Any => true,
-                ObjectTypeSet::Only(set) => set.contains(name.as_str()),
+                ObjectTypeSet::Only(set) => set.contains(obj_type),
             })
-            .collect::<Vec<String>>())
+            .collect())
     }
 }
 
 /// Look up the type `name` from the schema and resolve interfaces
 /// and unions until we are left with a set of concrete object types
-pub(crate) fn resolve_object_types<'a>(
-    schema: &'a Schema,
+pub(crate) fn resolve_object_types(
+    schema: &ApiSchema,
     name: &str,
-) -> Result<HashSet<ObjectType<'a>>, QueryExecutionError> {
+) -> Result<HashSet<ObjectType>, QueryExecutionError> {
     let mut set = HashSet::new();
     match schema
+        .schema
         .document
         .get_named_type(name)
         .ok_or_else(|| QueryExecutionError::AbstractTypeError(name.to_string()))?
     {
         s::TypeDefinition::Interface(intf) => {
             for obj_ty in &schema.types_for_interface()[&EntityType::new(intf.name.to_string())] {
+                let obj_ty = schema.object_type(obj_ty);
                 set.insert(obj_ty.into());
             }
         }
@@ -337,6 +335,7 @@ pub(crate) fn resolve_object_types<'a>(
             }
         }
         s::TypeDefinition::Object(ty) => {
+            let ty = schema.object_type(ty);
             set.insert(ty.into());
         }
         s::TypeDefinition::Scalar(_)

graphql/src/execution/execution.rs

@@ -222,7 +222,7 @@ pub(crate) fn get_field<'a>(
     name: &str,
 ) -> Option<s::Field> {
     if name == "__schema" || name == "__type" {
-        let object_type = *INTROSPECTION_QUERY_TYPE;
+        let object_type = &*INTROSPECTION_QUERY_TYPE;
         sast::get_field(object_type, name).cloned()
     } else {
         sast::get_field(object_type, name).cloned()
@@ -265,7 +265,7 @@ where
 pub(crate) fn execute_root_selection_set_uncached(
     ctx: &ExecutionContext<impl Resolver>,
     selection_set: &a::SelectionSet,
-    root_type: &s::ObjectType,
+    root_type: &sast::ObjectType,
 ) -> Result<Object, Vec<QueryExecutionError>> {
     // Split the top-level fields into introspection fields and
     // regular data fields
@@ -314,7 +314,7 @@ pub(crate) fn execute_root_selection_set_uncached(
 pub(crate) async fn execute_root_selection_set<R: Resolver>(
     ctx: Arc<ExecutionContext<R>>,
     selection_set: Arc<a::SelectionSet>,
-    root_type: Arc<s::ObjectType>,
+    root_type: sast::ObjectType,
     block_ptr: Option<BlockPtr>,
 ) -> Arc<QueryResult> {
     // Cache the cache key to not have to calculate it twice - once for lookup
@@ -458,7 +458,7 @@ pub(crate) async fn execute_root_selection_set<R: Resolver>(
 fn execute_selection_set<'a>(
     ctx: &'a ExecutionContext<impl Resolver>,
     selection_set: &'a a::SelectionSet,
-    object_type: &s::ObjectType,
+    object_type: &sast::ObjectType,
     prefetched_value: Option<r::Value>,
 ) -> Result<r::Value, Vec<QueryExecutionError>> {
     Ok(r::Value::Object(execute_selection_set_to_map(
@@ -472,7 +472,7 @@ fn execute_selection_set<'a>(
 fn execute_selection_set_to_map<'a>(
     ctx: &'a ExecutionContext<impl Resolver>,
     selection_set: &'a a::SelectionSet,
-    object_type: &s::ObjectType,
+    object_type: &sast::ObjectType,
     prefetched_value: Option<r::Value>,
 ) -> Result<Object, Vec<QueryExecutionError>> {
     let mut prefetched_object = match prefetched_value {
@@ -806,12 +806,15 @@ fn complete_value(
                 }
 
                 // Complete object types recursively
-                s::TypeDefinition::Object(object_type) => execute_selection_set(
-                    ctx,
-                    &field.selection_set,
-                    object_type,
-                    Some(resolved_value),
-                ),
+                s::TypeDefinition::Object(object_type) => {
+                    let object_type = ctx.query.schema.object_type(object_type).into();
+                    execute_selection_set(
+                        ctx,
+                        &field.selection_set,
+                        &object_type,
+                        Some(resolved_value),
+                    )
+                }
 
                 // Resolve interface types using the resolved value and complete the value recursively
                 s::TypeDefinition::Interface(_) => {
@@ -820,7 +823,7 @@ fn complete_value(
                     execute_selection_set(
                         ctx,
                         &field.selection_set,
-                        object_type,
+                        &object_type,
                         Some(resolved_value),
                     )
                 }
@@ -832,7 +835,7 @@ fn complete_value(
                     execute_selection_set(
                         ctx,
                         &field.selection_set,
-                        object_type,
+                        &object_type,
                         Some(resolved_value),
                     )
                 }
@@ -850,14 +853,16 @@ fn resolve_abstract_type<'a>(
     ctx: &'a ExecutionContext<impl Resolver>,
     abstract_type: &s::TypeDefinition,
     object_value: &r::Value,
-) -> Result<&'a s::ObjectType, Vec<QueryExecutionError>> {
+) -> Result<sast::ObjectType, Vec<QueryExecutionError>> {
     // Let the resolver handle the type resolution, return an error if the resolution
     // yields nothing
-    ctx.resolver
+    let obj_type = ctx
+        .resolver
         .resolve_abstract_type(ctx.query.schema.document(), abstract_type, object_value)
         .ok_or_else(|| {
             vec![QueryExecutionError::AbstractTypeError(
                 sast::get_type_name(abstract_type).to_string(),
             )]
-        })
+        })?;
+    Ok(ctx.query.schema.object_type(obj_type).into())
 }

graphql/src/execution/query.rs

@@ -51,8 +51,8 @@ struct SelectedFields<'a>(&'a a::SelectionSet);
 impl<'a> std::fmt::Display for SelectedFields<'a> {
     fn fmt(&self, fmt: &mut std::fmt::Formatter) -> Result<(), std::fmt::Error> {
         let mut first = true;
-        for (name, fields) in self.0.fields() {
-            write!(fmt, "{}:", name)?;
+        for (obj_type, fields) in self.0.fields() {
+            write!(fmt, "{}:", obj_type.name)?;
             for field in fields {
                 if first {
                     write!(fmt, "{}", field.response_key())?;
@@ -207,9 +207,9 @@ impl Query {
     {
         let mut bcs: Vec<(BlockConstraint, (a::SelectionSet, ErrorPolicy))> = Vec::new();
 
-        let root_type = self.schema.query_type.as_ref();
+        let root_type = sast::ObjectType::from(self.schema.query_type.cheap_clone());
         let mut prev_bc: Option<BlockConstraint> = None;
-        for field in self.selection_set.fields_for(root_type) {
+        for field in self.selection_set.fields_for(&root_type) {
             let bc = match field.argument_value("block") {
                 Some(bc) => BlockConstraint::try_from_value(bc).map_err(|_| {
                     vec![QueryExecutionError::InvalidArgumentError(
@@ -811,7 +811,7 @@ impl Transform {
         } else {
             let field_type = parent_type.field(&name).expect("field names are valid");
             let ty = field_type.field_type.get_base_type();
-            let type_set = a::ObjectTypeSet::from_name(&self.schema.schema, ty)?;
+            let type_set = a::ObjectTypeSet::from_name(&self.schema, ty)?;
             let ty = self.schema.document().object_or_interface(ty).unwrap();
             self.expand_selection_set(selection_set, &type_set, ty)?
         };
@@ -841,7 +841,7 @@ impl Transform {
         let mut visited_fragments = HashSet::new();
 
         // All the types that could possibly be returned by this selection set
-        let types = type_set.type_names(&self.schema.schema, ty)?;
+        let types = type_set.type_names(&self.schema, ty)?;
         let mut newset = a::SelectionSet::new(types);
 
         for sel in items {
@@ -920,8 +920,7 @@ impl Transform {
             None => ty,
         };
         if !skip {
-            let type_set =
-                a::ObjectTypeSet::convert(&self.schema.schema, frag_cond)?.intersect(type_set);
+            let type_set = a::ObjectTypeSet::convert(&self.schema, frag_cond)?.intersect(type_set);
             let selection_set = self.expand_selection_set(selection_set, &type_set, ty)?;
             newset.merge(selection_set, directives);
         }