From e79593cf345541ba14ee91bda5a8149cbc2c5da2 Mon Sep 17 00:00:00 2001
From: Paul de Vrieze <paul.devrieze@gmail.com>
Date: Mon, 20 Feb 2023 18:58:32 +0000
Subject: [PATCH] Tidy up the code a bit (and use the cast method on the
 serializer to avoid casting warnings). Use a map to allow faster lookup of
 serializers to class (to allow for overlapping hierarchies).

---
 .../modules/PolymorphicModuleBuilder.kt       | 24 ++++++++++---------
 .../features/PolymorphicSealedChildTest.kt    |  2 +-
 2 files changed, 14 insertions(+), 12 deletions(-)

diff --git a/core/commonMain/src/kotlinx/serialization/modules/PolymorphicModuleBuilder.kt b/core/commonMain/src/kotlinx/serialization/modules/PolymorphicModuleBuilder.kt
index 5a83e9233d..a6a7ada907 100644
--- a/core/commonMain/src/kotlinx/serialization/modules/PolymorphicModuleBuilder.kt
+++ b/core/commonMain/src/kotlinx/serialization/modules/PolymorphicModuleBuilder.kt
@@ -19,15 +19,23 @@ public class PolymorphicModuleBuilder<in Base : Any> @PublishedApi internal cons
     private val baseClass: KClass<Base>,
     private val baseSerializer: KSerializer<Base>? = null
 ) {
-    private val subclasses: MutableList<Pair<KClass<out Base>, KSerializer<out Base>>> = mutableListOf()
+    private val subclasses: MutableMap<KClass<out Base>, KSerializer<out Base>> = mutableMapOf()
     private var defaultSerializerProvider: ((Base) -> SerializationStrategy<Base>?)? = null
     private var defaultDeserializerProvider: ((String?) -> DeserializationStrategy<Base>?)? = null
 
+
+    /**
+     * Registers the child serializers for the sealed [subclass] [serializer] in the resulting module under the [base class][Base].
+     */
+    public inline fun <reified T : Base> subclassesOf(): Unit =
+        subclassesOf(serializer<T>())
+
+
     /**
-     * Registers the subclasses of the given class as subclasses of the outer class. Currently this requires `baseClass`
+     * Registers the subclasses of the given class as subclasses of the outer class. This currently requires `baseClass`
      * to be sealed.
      */
-    public fun <T: Base> subclassesOf(baseClass: KClass<T>, serializer: KSerializer<T>) {
+    public fun <T: Base> subclassesOf(serializer: KSerializer<T>) {
         require(serializer is SealedClassSerializer) {
             "subClassesOf only supports automatic adding of subclasses of sealed types."
         }
@@ -42,7 +50,7 @@ public class PolymorphicModuleBuilder<in Base : Any> @PublishedApi internal cons
      * Registers a [subclass] [serializer] in the resulting module under the [base class][Base].
      */
     public fun <T : Base> subclass(subclass: KClass<T>, serializer: KSerializer<T>) {
-        subclasses.add(subclass to serializer)
+        subclasses[subclass] = serializer
     }
 
     /**
@@ -132,14 +140,8 @@ public inline fun <Base : Any, reified T : Base> PolymorphicModuleBuilder<Base>.
 public inline fun <Base : Any, reified T : Base> PolymorphicModuleBuilder<Base>.subclass(clazz: KClass<T>): Unit =
     subclass(clazz, serializer())
 
-/**
- * Registers the child serializers for the sealed [subclass] [serializer] in the resulting module under the [base class][Base].
- */
-public inline fun <Base : Any, reified T : Base> PolymorphicModuleBuilder<Base>.subclassesOf(serializer: KSerializer<T>): Unit =
-    subclassesOf(T::class, serializer)
-
 /**
  * Registers the child serializers for the sealed class [T] in the resulting module under the [base class][Base].
  */
 public inline fun <Base : Any, reified T : Base> PolymorphicModuleBuilder<Base>.subclassesOf(clazz: KClass<T>): Unit =
-    subclassesOf(clazz, serializer())
+    subclassesOf(clazz.serializer())
diff --git a/formats/json-tests/commonTest/src/kotlinx/serialization/features/PolymorphicSealedChildTest.kt b/formats/json-tests/commonTest/src/kotlinx/serialization/features/PolymorphicSealedChildTest.kt
index 1957026c2c..47a3788dca 100644
--- a/formats/json-tests/commonTest/src/kotlinx/serialization/features/PolymorphicSealedChildTest.kt
+++ b/formats/json-tests/commonTest/src/kotlinx/serialization/features/PolymorphicSealedChildTest.kt
@@ -34,7 +34,7 @@ class PolymorphicSealedChildTest {
 
     val sealedModule = SerializersModule {
         polymorphic(FooBase::class) {
-            subclassesOf(Foo.serializer())
+            subclassesOf<Foo>()
         }
     }