Fix the bug that `MergedObject` may cause panic. #539

CHANGELOG.md

@@ -4,6 +4,10 @@ All notable changes to this project will be documented in this file.
 The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
+## [2.9.4] 2021-06-21
+
+- Fix the bug that `MergedObject` may cause panic. [#539](https://github.com/async-graphql/async-graphql/issues/539#issuecomment-862209442)
+
 ## [2.9.3] 2021-06-17
 
 - Bump upstream crate `bson` from `v1.2.0` to `v2.0.0-beta.1`. [#516](https://github.com/async-graphql/async-graphql/pull/516)

src/registry/mod.rs

@@ -100,6 +100,7 @@ impl<'a> MetaTypeName<'a> {
     }
 }
 
+#[derive(Clone)]
 pub struct MetaInputValue {
     pub name: &'static str,
     pub description: Option<&'static str>,
@@ -117,12 +118,13 @@ type ComputeComplexityFn = fn(
     usize,
 ) -> ServerResult<usize>;
 
+#[derive(Clone)]
 pub enum ComplexityType {
     Const(usize),
     Fn(ComputeComplexityFn),
 }
 
-#[derive(Debug)]
+#[derive(Debug, Clone)]
 pub enum Deprecation {
     NoDeprecated,
     Deprecated { reason: Option<&'static str> },
@@ -149,6 +151,7 @@ impl Deprecation {
     }
 }
 
+#[derive(Clone)]
 pub struct MetaField {
     pub name: String,
     pub description: Option<&'static str>,
@@ -163,6 +166,7 @@ pub struct MetaField {
     pub compute_complexity: Option<ComplexityType>,
 }
 
+#[derive(Clone)]
 pub struct MetaEnumValue {
     pub name: &'static str,
     pub description: Option<&'static str>,
@@ -172,6 +176,7 @@ pub struct MetaEnumValue {
 
 type MetaVisibleFn = fn(&Context<'_>) -> bool;
 
+#[derive(Clone)]
 pub enum MetaType {
     Scalar {
         name: String,
@@ -373,7 +378,8 @@ impl Registry {
     pub fn create_dummy_type<T: Type>(&mut self) -> MetaType {
         T::create_type_info(self);
         self.types
-            .remove(&*T::type_name())
+            .get(&*T::type_name())
+            .cloned()
             .expect("You definitely encountered a bug!")
     }
 
@@ -423,8 +429,8 @@ impl Registry {
         self.types.values().any(|ty| match ty {
             MetaType::Object {
                 keys: Some(keys), ..
-            } => !keys.is_empty(),
-            MetaType::Interface {
+            }
+            | MetaType::Interface {
                 keys: Some(keys), ..
             } => !keys.is_empty(),
             _ => false,
@@ -606,4 +612,123 @@ impl Registry {
             None => {}
         }
     }
+
+    pub fn remove_unused_types(&mut self) {
+        let mut used_types = HashSet::new();
+        let mut unused_types = HashSet::new();
+
+        fn traverse_field<'a>(
+            types: &'a IndexMap<String, MetaType>,
+            used_types: &mut HashSet<&'a str>,
+            field: &'a MetaField,
+        ) {
+            traverse_type(
+                types,
+                used_types,
+                MetaTypeName::concrete_typename(&field.ty),
+            );
+            for arg in field.args.values() {
+                traverse_input_value(types, used_types, arg);
+            }
+        }
+
+        fn traverse_input_value<'a>(
+            types: &'a IndexMap<String, MetaType>,
+            used_types: &mut HashSet<&'a str>,
+            input_value: &'a MetaInputValue,
+        ) {
+            traverse_type(
+                types,
+                used_types,
+                MetaTypeName::concrete_typename(&input_value.ty),
+            );
+        }
+
+        fn traverse_type<'a>(
+            types: &'a IndexMap<String, MetaType>,
+            used_types: &mut HashSet<&'a str>,
+            type_name: &'a str,
+        ) {
+            if used_types.contains(type_name) {
+                return;
+            }
+
+            if let Some(ty) = types.get(type_name) {
+                used_types.insert(type_name);
+                match ty {
+                    MetaType::Object { fields, .. } => {
+                        for field in fields.values() {
+                            traverse_field(types, used_types, field);
+                        }
+                    }
+                    MetaType::Interface {
+                        fields,
+                        possible_types,
+                        ..
+                    } => {
+                        for field in fields.values() {
+                            traverse_field(types, used_types, field);
+                        }
+                        for type_name in possible_types.iter() {
+                            traverse_type(types, used_types, type_name);
+                        }
+                    }
+                    MetaType::Union { possible_types, .. } => {
+                        for type_name in possible_types.iter() {
+                            traverse_type(types, used_types, type_name);
+                        }
+                    }
+                    MetaType::InputObject { input_fields, .. } => {
+                        for field in input_fields.values() {
+                            traverse_input_value(types, used_types, field);
+                        }
+                    }
+                    _ => {}
+                }
+            }
+        }
+
+        for type_name in Some(&self.query_type)
+            .into_iter()
+            .chain(self.mutation_type.iter())
+            .chain(self.subscription_type.iter())
+        {
+            traverse_type(&self.types, &mut used_types, type_name);
+        }
+
+        for ty in self.types.values().filter(|ty| match ty {
+            MetaType::Object {
+                keys: Some(keys), ..
+            }
+            | MetaType::Interface {
+                keys: Some(keys), ..
+            } => !keys.is_empty(),
+            _ => false,
+        }) {
+            traverse_type(&self.types, &mut used_types, ty.name());
+        }
+
+        fn is_system_type(name: &str) -> bool {
+            if name.starts_with("__") {
+                return true;
+            }
+
+            name == "Boolean"
+                || name == "Int"
+                || name == "Float"
+                || name == "String"
+                || name == "ID"
+        }
+
+        for ty in self.types.values() {
+            let name = ty.name();
+            if !is_system_type(name) && !used_types.contains(name) {
+                unused_types.insert(name.to_string());
+            }
+        }
+
+        for type_name in unused_types {
+            self.types.remove(&type_name);
+        }
+    }
 }

src/schema.rs

@@ -322,6 +322,7 @@ where
             Subscription::create_type_info(&mut registry);
         }
 
+        registry.remove_unused_types();
         registry
     }
 

src/types/merged_object.rs

@@ -22,16 +22,6 @@ impl<A: Type, B: Type> Type for MergedObject<A, B> {
             let mut fields = IndexMap::new();
             let mut cc = CacheControl::default();
 
-            if let MetaType::Object {
-                fields: a_fields,
-                cache_control: a_cc,
-                ..
-            } = registry.create_dummy_type::<A>()
-            {
-                fields.extend(a_fields);
-                cc = cc.merge(&a_cc);
-            }
-
             if let MetaType::Object {
                 fields: b_fields,
                 cache_control: b_cc,
@@ -42,6 +32,16 @@ impl<A: Type, B: Type> Type for MergedObject<A, B> {
                 cc = cc.merge(&b_cc);
             }
 
+            if let MetaType::Object {
+                fields: a_fields,
+                cache_control: a_cc,
+                ..
+            } = registry.create_dummy_type::<A>()
+            {
+                fields.extend(a_fields);
+                cc = cc.merge(&a_cc);
+            }
+
             MetaType::Object {
                 name: Self::type_name().to_string(),
                 description: None,

tests/merged_object.rs

@@ -346,3 +346,58 @@ pub async fn test_issue_333() {
         })
     )
 }
+
+#[tokio::test]
+pub async fn test_issue_539() {
+    // https://github.com/async-graphql/async-graphql/issues/539#issuecomment-862209442
+
+    struct Query;
+
+    #[Object]
+    impl Query {
+        async fn value(&self) -> i32 {
+            10
+        }
+    }
+
+    #[derive(SimpleObject)]
+    struct A {
+        a: Option<Box<A>>,
+    }
+
+    #[derive(SimpleObject)]
+    struct B {
+        b: Option<Box<B>>,
+    }
+
+    #[derive(MergedObject)]
+    pub struct Mutation(A, B);
+
+    let schema = Schema::new(
+        Query,
+        Mutation(A { a: None }, B { b: None }),
+        EmptySubscription,
+    );
+    assert_eq!(
+        schema
+            .execute("{ __type(name: \"Mutation\") { fields { name type { name } } } }")
+            .await
+            .into_result()
+            .unwrap()
+            .data,
+        value!({
+            "__type": {
+                "fields": [
+                    {
+                        "name": "a",
+                        "type": { "name": "A" },
+                    },
+                    {
+                        "name": "b",
+                        "type": { "name": "B" },
+                    }
+                ]
+            }
+        })
+    )
+}