From 475270e10404ed58e5003e8292eb7047a563b22a Mon Sep 17 00:00:00 2001
From: AlixBa <alix.baggio@outlook.fr>
Date: Mon, 5 Jul 2021 12:43:42 +0200
Subject: [PATCH] Allow union types to reuse the same member

closes #826
---
 .../scala/caliban/tools/SchemaWriter.scala    | 33 ++++++++++++++-----
 .../caliban/tools/SchemaWriterSpec.scala      | 33 +++++++++++++------
 2 files changed, 48 insertions(+), 18 deletions(-)

diff --git a/tools/src/main/scala/caliban/tools/SchemaWriter.scala b/tools/src/main/scala/caliban/tools/SchemaWriter.scala
index 05433d4d2..036528285 100644
--- a/tools/src/main/scala/caliban/tools/SchemaWriter.scala
+++ b/tools/src/main/scala/caliban/tools/SchemaWriter.scala
@@ -24,7 +24,7 @@ object SchemaWriter {
       .map(union => (union, union.memberTypes.flatMap(schema.objectTypeDefinition)))
       .toMap
 
-    val unions = unionTypes.map { case (union, objects) => writeUnion(union, objects) }.mkString("\n")
+    val unions = writeUnions(unionTypes)
 
     val objects = schema.objectTypeDefinitions
       .filterNot(obj =>
@@ -146,17 +146,34 @@ object SchemaWriter {
           }
        """
 
-  def writeUnion(typedef: UnionTypeDefinition, objects: List[ObjectTypeDefinition])(implicit
+  def writeUnions(unions: Map[UnionTypeDefinition, List[ObjectTypeDefinition]])(implicit
     scalarMappings: ScalarMappings
   ): String =
-    s"""${writeDescription(typedef.description)}sealed trait ${typedef.name} extends scala.Product with scala.Serializable
-
-          object ${typedef.name} {
-            ${objects
-      .map(o => s"${writeObject(o)} extends ${typedef.name}")
-      .mkString("\n")}
+    if (unions.nonEmpty) {
+      val invertedUnions = unions.foldLeft(Map.empty[ObjectTypeDefinition, List[UnionTypeDefinition]]) {
+        case (map, (unionType, objectTypes)) =>
+          objectTypes.foldLeft(map) { case (map, objectType) =>
+            map.updated(objectType, map.getOrElse(objectType, List.empty) :+ unionType)
           }
+      }
+
+      val codegenUnions = unions.keys.map { union =>
+        s"""${writeDescription(
+          union.description
+        )}sealed trait ${union.name} extends scala.Product with scala.Serializable"""
+      }
+        .mkString("\n")
+
+      val codegenObjects = invertedUnions.map { case (obj, unions) =>
+        s"${writeObject(obj)} extends ${unions.map(_.name).mkString(" with ")}"
+      }
+        .mkString("\n")
+
+      s"""$codegenUnions
+        
+        $codegenObjects
        """
+    } else ""
 
   def writeField(field: FieldDefinition, of: ObjectTypeDefinition)(implicit scalarMappings: ScalarMappings): String =
     if (field.args.nonEmpty) {
diff --git a/tools/src/test/scala/caliban/tools/SchemaWriterSpec.scala b/tools/src/test/scala/caliban/tools/SchemaWriterSpec.scala
index d05333cc9..e9ad64e5c 100644
--- a/tools/src/test/scala/caliban/tools/SchemaWriterSpec.scala
+++ b/tools/src/test/scala/caliban/tools/SchemaWriterSpec.scala
@@ -221,10 +221,19 @@ object SchemaWriterSpec extends DefaultRunnableSpec {
              Captain or Pilot
              \"\"\"
           """
+        val role2  =
+          s"""
+              \"\"\"
+             role2
+             Captain or Pilot
+             \"\"\"
+          """
         val schema =
           s"""
              $role
              union Role = Captain | Pilot
+             $role2
+             union Role2 = Captain | Pilot
              
              type Captain {
                "ship" shipName: String!
@@ -237,23 +246,27 @@ object SchemaWriterSpec extends DefaultRunnableSpec {
 
         assertM(gen(schema))(
           equalTo {
-            val role =
+            val role  =
               s"""\"\"\"role
+Captain or Pilot\"\"\""""
+            val role2 =
+              s"""\"\"\"role2
 Captain or Pilot\"\"\""""
             s"""import caliban.schema.Annotations._
 
 object Types {
 
   @GQLDescription($role)
-  sealed trait Role extends scala.Product with scala.Serializable
-
-  object Role {
-    case class Captain(
-      @GQLDescription("ship")
-      shipName: String
-    )                                  extends Role
-    case class Pilot(shipName: String) extends Role
-  }
+  sealed trait Role  extends scala.Product with scala.Serializable
+  @GQLDescription($role2)
+  sealed trait Role2 extends scala.Product with scala.Serializable
+
+  case class Captain(
+    @GQLDescription("ship")
+    shipName: String
+  )                                  extends Role
+      with Role2
+  case class Pilot(shipName: String) extends Role with Role2
 
 }
 """