From 7aa623772658f6958e4c637287f0f76e4b735c2a 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    | 67 +++++++++++++++++--
 .../caliban/tools/SchemaWriterSpec.scala      | 39 ++++++++---
 2 files changed, 93 insertions(+), 13 deletions(-)

diff --git a/tools/src/main/scala/caliban/tools/SchemaWriter.scala b/tools/src/main/scala/caliban/tools/SchemaWriter.scala
index 05433d4d2..5965bcd2c 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,12 +146,71 @@ 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
+    if (unions.nonEmpty) {
+      val flattened = unions.toList.flatMap { case (unionType, objectTypes) => objectTypes.map(_ -> unionType) }
+
+      val (unionsWithoutReusedMembers, reusedUnionMembers) = flattened
+        .foldLeft(
+          (
+            Map.empty[UnionTypeDefinition, List[ObjectTypeDefinition]],
+            Map.empty[ObjectTypeDefinition, List[UnionTypeDefinition]]
+          )
+        ) {
+          case (
+                (unionsWithoutReusedMembers, reusedUnionMembers),
+                (objectType, unionType)
+              ) =>
+            val isReused = reusedUnionMembers.contains(objectType) ||
+              flattened.exists { case (_objectType, _unionType) =>
+                _unionType.name != unionType.name && _objectType.name == objectType.name
+              }
+
+            if (isReused) {
+              (
+                unionsWithoutReusedMembers,
+                reusedUnionMembers.updated(
+                  objectType,
+                  reusedUnionMembers.getOrElse(objectType, List.empty) :+ unionType
+                )
+              )
+            } else {
+              (
+                unionsWithoutReusedMembers.updated(
+                  unionType,
+                  unionsWithoutReusedMembers.getOrElse(unionType, List.empty) :+ objectType
+                ),
+                reusedUnionMembers
+              )
+            }
+        }
+
+      s"""${unions.keys.map(writeUnionSealedTrait).mkString("\n")}
+        
+        ${unionsWithoutReusedMembers.map { case (union, objects) => writeNotReusedMembers(union, objects) }
+        .mkString("\n")}
+        
+        ${reusedUnionMembers.map { case (objectType, unions) => writeReusedUnionMember(objectType, unions) }
+        .mkString("\n")}
+         """
+    } else ""
+
+  def writeUnionSealedTrait(union: UnionTypeDefinition): String =
+    s"""${writeDescription(
+      union.description
+    )}sealed trait ${union.name} extends scala.Product with scala.Serializable"""
+
+  def writeReusedUnionMember(typedef: ObjectTypeDefinition, unions: List[UnionTypeDefinition])(implicit
+    scalarMappings: ScalarMappings
+  ): String =
+    s"${writeObject(typedef)} extends ${unions.map(_.name).mkString(" with ")}"
 
-          object ${typedef.name} {
+  def writeNotReusedMembers(typedef: UnionTypeDefinition, objects: List[ObjectTypeDefinition])(implicit
+    scalarMappings: ScalarMappings
+  ): String =
+    s"""object ${typedef.name} {
             ${objects
       .map(o => s"${writeObject(o)} extends ${typedef.name}")
       .mkString("\n")}
diff --git a/tools/src/test/scala/caliban/tools/SchemaWriterSpec.scala b/tools/src/test/scala/caliban/tools/SchemaWriterSpec.scala
index d05333cc9..72e7c6a48 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 or Stewart
+             \"\"\"
+          """
         val schema =
           s"""
              $role
              union Role = Captain | Pilot
+             $role2
+             union Role2 = Captain | Pilot | Stewart
              
              type Captain {
                "ship" shipName: String!
@@ -233,28 +242,40 @@ object SchemaWriterSpec extends DefaultRunnableSpec {
              type Pilot {
                shipName: String!
              }
+             
+             type Stewart {
+               shipName: String!
+             }
             """.stripMargin
 
         assertM(gen(schema))(
           equalTo {
-            val role =
+            val role  =
               s"""\"\"\"role
 Captain or Pilot\"\"\""""
+            val role2 =
+              s"""\"\"\"role2
+Captain or Pilot or Stewart\"\"\""""
             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
+
+  object Role2 {
+    case class Stewart(shipName: String) extends Role2
   }
 
+  case class Captain(
+    @GQLDescription("ship")
+    shipName: String
+  )                                  extends Role
+      with Role2
+  case class Pilot(shipName: String) extends Role with Role2
+
 }
 """
           }