From 650c434c1a7fdda73c6cab704c43a661c47fbddc Mon Sep 17 00:00:00 2001
From: Graham Esau <gesau@hotmail.co.uk>
Date: Tue, 6 Aug 2019 00:52:06 +0100
Subject: [PATCH] Ensure all schema types are set to nullable when appropriate

---
 schemars/src/gen.rs         | 33 ++++++++++++++++++++++++++++++--
 schemars/src/make_schema.rs | 18 +++++++++++-------
 schemars/tests/schema.json  | 38 +++++++++++++++++++++----------------
 3 files changed, 64 insertions(+), 25 deletions(-)

diff --git a/schemars/src/gen.rs b/schemars/src/gen.rs
index 0227d06..f8d7b8b 100644
--- a/schemars/src/gen.rs
+++ b/schemars/src/gen.rs
@@ -89,7 +89,7 @@ impl SchemaGenerator {
     }
 
     pub fn subschema_for<T: ?Sized + MakeSchema>(&mut self) -> Schema {
-        if !T::generates_ref_schema() {
+        if !T::is_referenceable() {
             return T::make_schema(self);
         }
 
@@ -100,7 +100,6 @@ impl SchemaGenerator {
             .map(|(n, _)| n.clone())
             .unwrap_or_else(|| {
                 let name = self.make_unique_name::<T>();
-                self.names.insert(name.clone());
                 self.insert_new_subschema_for::<T>(type_id, name.clone());
                 name
             });
@@ -113,6 +112,7 @@ impl SchemaGenerator {
         type_id: SchemaTypeId,
         name: String,
     ) {
+        self.names.insert(name.clone());
         let dummy = Schema::Bool(false);
         // insert into definitions BEFORE calling make_schema to avoid infinite recursion
         self.definitions.insert(type_id.clone(), (name, dummy));
@@ -153,6 +153,35 @@ impl SchemaGenerator {
         schema
     }
 
+    pub(crate) fn try_get_schema_object<'a>(
+        &'a self,
+        mut schema: &'a Schema,
+    ) -> Option<SchemaObject> {
+        loop {
+            match schema {
+                Schema::Object(o) => return Some(o.clone()),
+                Schema::Bool(true) => return Some(Default::default()),
+                Schema::Bool(false) => {
+                    return Some(SchemaObject {
+                        not: Some(Schema::Bool(true).into()),
+                        ..Default::default()
+                    })
+                }
+                Schema::Ref(r) => {
+                    let definitions_path_len = self.settings().definitions_path.len();
+                    let name = r.reference.get(definitions_path_len..)?;
+                    // FIXME this is pretty inefficient
+                    schema = self
+                        .definitions
+                        .values()
+                        .filter(|(n, _)| n == name)
+                        .map(|(_, s)| s)
+                        .next()?;
+                }
+            }
+        }
+    }
+
     fn make_unique_name<T: ?Sized + MakeSchema>(&mut self) -> String {
         let base_name = T::schema_name();
         if self.names.contains(&base_name) {
diff --git a/schemars/src/make_schema.rs b/schemars/src/make_schema.rs
index 2b5a3bb..68ad085 100644
--- a/schemars/src/make_schema.rs
+++ b/schemars/src/make_schema.rs
@@ -20,7 +20,7 @@ pub trait MakeSchema {
         core::any::type_name::<Self>().replace(|c: char| !c.is_ascii_alphanumeric(), "_")
     }
 
-    fn generates_ref_schema() -> bool {
+    fn is_referenceable() -> bool {
         true
     }
 
@@ -29,7 +29,7 @@ pub trait MakeSchema {
 
 macro_rules! no_ref_schema {
     () => {
-        fn generates_ref_schema() -> bool {
+        fn is_referenceable() -> bool {
             false
         }
     };
@@ -261,13 +261,15 @@ map_impl!(<K: Eq + core::hash::Hash, V, H: core::hash::BuildHasher> MakeSchema f
 ////////// OPTION //////////
 
 impl<T: MakeSchema> MakeSchema for Option<T> {
-    no_ref_schema!();
+    fn is_referenceable() -> bool {
+        T::is_referenceable()
+    }
 
     fn make_schema(gen: &mut SchemaGenerator) -> Schema {
         let settings = gen.settings();
         let make_any_of = settings.option_any_of_null;
         let set_nullable = settings.option_nullable;
-        let mut schema = match gen.subschema_for::<T>() {
+        let schema = match gen.subschema_for::<T>() {
             Schema::Bool(true) => true.into(),
             Schema::Bool(false) => <()>::make_schema(gen),
             schema => {
@@ -283,9 +285,11 @@ impl<T: MakeSchema> MakeSchema for Option<T> {
             }
         };
         if set_nullable {
-            // FIXME still need to handle ref schemas here
-            if let Schema::Object(ref mut o) = schema {
-                o.extensions.insert("nullable".to_owned(), true.into());
+            let deref = gen.try_get_schema_object(&schema);
+            debug_assert!(deref.is_some(), "Could not get schema object: {:?}", schema);
+            if let Some(mut schema) = deref {
+                schema.extensions.insert("nullable".to_owned(), json!(true));
+                return Schema::Object(schema);
             }
         };
         schema
diff --git a/schemars/tests/schema.json b/schemars/tests/schema.json
index 2da5bfa..111d775 100644
--- a/schemars/tests/schema.json
+++ b/schemars/tests/schema.json
@@ -13,6 +13,26 @@
     }
   ],
   "definitions": {
+    "core__option__Option_schemars__schema__SingleOrVec_schemars__schema__InstanceType__": {
+      "anyOf": [
+        {
+          "$ref": "#/definitions/schemars__schema__SingleOrVec_schemars__schema__InstanceType_"
+        },
+        {
+          "type": "null"
+        }
+      ]
+    },
+    "core__option__Option_schemars__schema__SingleOrVec_schemars__schema__Schema__": {
+      "anyOf": [
+        {
+          "$ref": "#/definitions/schemars__schema__SingleOrVec_schemars__schema__Schema_"
+        },
+        {
+          "type": "null"
+        }
+      ]
+    },
     "schemars__schema__InstanceType": {
       "enum": [
         "null",
@@ -117,14 +137,7 @@
           "additionalProperties": true
         },
         "items": {
-          "anyOf": [
-            {
-              "$ref": "#/definitions/schemars__schema__SingleOrVec_schemars__schema__Schema_"
-            },
-            {
-              "type": "null"
-            }
-          ]
+          "$ref": "#/definitions/core__option__Option_schemars__schema__SingleOrVec_schemars__schema__Schema__"
         },
         "not": {
           "anyOf": [
@@ -179,14 +192,7 @@
           ]
         },
         "type": {
-          "anyOf": [
-            {
-              "$ref": "#/definitions/schemars__schema__SingleOrVec_schemars__schema__InstanceType_"
-            },
-            {
-              "type": "null"
-            }
-          ]
+          "$ref": "#/definitions/core__option__Option_schemars__schema__SingleOrVec_schemars__schema__InstanceType__"
         }
       }
     },