From 02ead13ee9310368b46c8009cc5666e738291bc6 Mon Sep 17 00:00:00 2001
From: Romain Lebran <rlebran@netwo.com>
Date: Tue, 30 Jan 2024 08:06:33 +0100
Subject: [PATCH] Fix multiple flatten with all_of

---
 schemars/src/flatten.rs                       | 63 +++++++++++++++-
 .../tests/expected/multiple_enum_flatten.json | 75 +++++++++++++++++++
 schemars/tests/multiple_enum_flatten.rs       | 33 ++++++++
 3 files changed, 168 insertions(+), 3 deletions(-)
 create mode 100644 schemars/tests/expected/multiple_enum_flatten.json
 create mode 100644 schemars/tests/multiple_enum_flatten.rs

diff --git a/schemars/src/flatten.rs b/schemars/src/flatten.rs
index e55cb7fc..68a20681 100644
--- a/schemars/src/flatten.rs
+++ b/schemars/src/flatten.rs
@@ -12,9 +12,66 @@ impl Schema {
         } else if is_null_type(&other) {
             return self;
         }
-        let s1: SchemaObject = self.into();
-        let s2: SchemaObject = other.into();
-        Schema::Object(s1.merge(s2))
+        let mut s1: SchemaObject = self.into();
+        let mut s2: SchemaObject = other.into();
+
+        let s1_subschemas = &mut s1.subschemas;
+        let s1_all_of = s1_subschemas.as_mut().and_then(|sub_sch| sub_sch.all_of.as_mut());
+        let s2_subschemas = &mut s2.subschemas;
+        let s2_all_of = s2_subschemas.as_mut().and_then(|sub_sch| sub_sch.all_of.as_mut());
+
+        match (s1_all_of, s2_all_of) {
+            (Some(s1_all_of), Some(s2_all_of)) => {
+                let mut s1_all_of = std::mem::take(s1_all_of);
+                s1_all_of.append(s2_all_of);
+                Schema::Object(SchemaObject {
+                    subschemas: Some(Box::new(SubschemaValidation {
+                        all_of: Some(s1_all_of),
+                        ..Default::default()
+                    })),
+                    ..Default::default()
+                })
+            }
+            (Some(s1_all_of), None) => {
+                let mut s1_all_of = std::mem::take(s1_all_of);
+                s1_all_of.push(Schema::Object(s2));
+                Schema::Object(SchemaObject {
+                    subschemas: Some(Box::new(SubschemaValidation {
+                        all_of: Some(s1_all_of),
+                        ..Default::default()
+                    })),
+                    ..Default::default()
+                })
+            }
+            (None, Some(s2_all_of)) => {
+                let mut s2_all_of = std::mem::take(s2_all_of);
+                s2_all_of.push(Schema::Object(s1));
+                Schema::Object(SchemaObject {
+                    subschemas: Some(Box::new(SubschemaValidation {
+                        all_of: Some(s2_all_of),
+                        ..Default::default()
+                    })),
+                    ..Default::default()
+                })
+            }
+            (None, None) => {
+                match (s1_subschemas.is_some(), s2_subschemas.is_some()) {
+                    (true, true) => Schema::Object(SchemaObject {
+                        subschemas: Some(Box::new(SubschemaValidation {
+                            all_of: Some(vec![
+                                Schema::Object(s1),
+                                Schema::Object(s2),
+                            ]),
+                            ..Default::default()
+                        })),
+                        ..Default::default()
+                    }),
+                    (true, false) |
+                    (false, true) |
+                    (false, false) => Schema::Object(s1.merge(s2)),
+                }
+            }
+        }
     }
 }
 
diff --git a/schemars/tests/expected/multiple_enum_flatten.json b/schemars/tests/expected/multiple_enum_flatten.json
new file mode 100644
index 00000000..d754f449
--- /dev/null
+++ b/schemars/tests/expected/multiple_enum_flatten.json
@@ -0,0 +1,75 @@
+{
+  "$schema": "http://json-schema.org/draft-07/schema#",
+  "title": "Flat",
+  "allOf": [
+    {
+      "type": "object",
+      "oneOf": [
+        {
+          "type": "object",
+          "required": [
+            "B"
+          ],
+          "properties": {
+            "B": {
+              "type": "boolean"
+            }
+          },
+          "additionalProperties": false
+        },
+        {
+          "type": "object",
+          "required": [
+            "S"
+          ],
+          "properties": {
+            "S": {
+              "type": "string"
+            }
+          },
+          "additionalProperties": false
+        }
+      ],
+      "required": [
+        "f"
+      ],
+      "properties": {
+        "f": {
+          "type": "number",
+          "format": "float"
+        }
+      }
+    },
+    {
+      "oneOf": [
+        {
+          "type": "object",
+          "required": [
+            "U"
+          ],
+          "properties": {
+            "U": {
+              "type": "integer",
+              "format": "uint32",
+              "minimum": 0.0
+            }
+          },
+          "additionalProperties": false
+        },
+        {
+          "type": "object",
+          "required": [
+            "F"
+          ],
+          "properties": {
+            "F": {
+              "type": "number",
+              "format": "double"
+            }
+          },
+          "additionalProperties": false
+        }
+      ]
+    }
+  ]
+}
diff --git a/schemars/tests/multiple_enum_flatten.rs b/schemars/tests/multiple_enum_flatten.rs
new file mode 100644
index 00000000..5e9ae93c
--- /dev/null
+++ b/schemars/tests/multiple_enum_flatten.rs
@@ -0,0 +1,33 @@
+mod util;
+use schemars::JsonSchema;
+use util::*;
+
+#[allow(dead_code)]
+#[derive(JsonSchema)]
+#[schemars(rename = "Flat")]
+struct Flat {
+  f: f32,
+  #[schemars(flatten)]
+  e1: Enum1,
+  #[schemars(flatten)]
+  e2: Enum2,
+}
+
+#[allow(dead_code)]
+#[derive(JsonSchema)]
+enum Enum1 {
+  B(bool),
+  S(String),
+}
+
+#[allow(dead_code)]
+#[derive(JsonSchema)]
+enum Enum2 {
+  U(u32),
+  F(f64)
+}
+
+#[test]
+fn test_flat_schema() -> TestResult {
+  test_default_generated_schema::<Flat>("multiple_enum_flatten")
+}