Merge pull request #73 from SimonMarquis/feat/companion-objects

SimonMarquis · Apr 22, 2023 · 10c8e10 · 10c8e10
2 parents caccabe + 9697647
commit 10c8e10
Show file tree

Hide file tree

Showing 9 changed files with 141 additions and 24 deletions.
diff --git a/.editorconfig b/.editorconfig
@@ -15,6 +15,7 @@ ij_kotlin_allow_trailing_comma_on_call_site = true
 ij_kotlin_name_count_to_use_star_import = 2147483647
 ij_kotlin_name_count_to_use_star_import_for_members = 2147483647
 ij_kotlin_packages_to_use_import_on_demand = unset
+ij_kotlin_class_annotation_wrap = normal
 
 [{*.md,*.mdown,*.markdown}]
 indent_size = 2

diff --git a/README.md b/README.md
@@ -50,6 +50,17 @@ And then you'll have access to the `sealedObjectInstances()` extension on the co
 val flags: Set<FeatureFlag> = FeatureFlag::class.sealedObjectInstances()
 ```
 
+> **Note**  
+> You can annotate the `companion object` to access a simpler extension function (no more `::class` prefix).
+> ```kotlin
+> sealed class FeatureFlag {
+>     @SealedObjectInstances companion object
+>     /*...*/
+> }
+>
+> val flags: Set<FeatureFlag> = FeatureFlag.sealedObjectInstances()
+> ```
+
 ## Setup
 
 In the module's build script, apply the `com.google.devtools.ksp` plugin with the matching Kotlin version: [![Maven Central](https://img.shields.io/maven-central/v/com.google.devtools.ksp/com.google.devtools.ksp.gradle.plugin?label=%20&color=success)](https://central.sonatype.com/artifact/com.google.devtools.ksp/com.google.devtools.ksp.gradle.plugin)

diff --git a/app/src/main/kotlin/fr/smarquis/sealed/Debug.kt b/app/src/main/kotlin/fr/smarquis/sealed/Debug.kt
@@ -15,8 +15,8 @@
  */
 package fr.smarquis.sealed
 
-@SealedObjectInstances
 sealed class Debug(override val isEnabled: Boolean = false) : FeatureFlag() {
+    @SealedObjectInstances companion object
     object Logs : Debug(true)
     object Traces : Debug()
     object StrictMode : Debug()

diff --git a/app/src/main/kotlin/fr/smarquis/sealed/FeatureFlag.kt b/app/src/main/kotlin/fr/smarquis/sealed/FeatureFlag.kt
@@ -15,7 +15,7 @@
  */
 package fr.smarquis.sealed
 
-@SealedObjectInstances
 sealed class FeatureFlag {
+    @SealedObjectInstances companion object;
     abstract val isEnabled: Boolean
 }
diff --git a/app/src/main/kotlin/fr/smarquis/sealed/UI.kt b/app/src/main/kotlin/fr/smarquis/sealed/UI.kt
@@ -15,8 +15,8 @@
  */
 package fr.smarquis.sealed
 
-@SealedObjectInstances
 sealed class UI(override val isEnabled: Boolean = false) : FeatureFlag() {
+    @SealedObjectInstances companion object
     object Animations : UI()
     object Framerate : UI()
 }
diff --git a/build.gradle.kts b/build.gradle.kts
@@ -32,7 +32,6 @@ allprojects {
         val licenseHeader = rootProject.file("spotless/spotless.kt")
         format("misc") {
             target("**/*.md", "**/.gitignore")
-            trimTrailingWhitespace()
             endWithNewline()
         }
         kotlin {

diff --git a/processor/src/main/kotlin/fr/smarquis/sealed/SealedObjectInstancesProcessor.kt b/processor/src/main/kotlin/fr/smarquis/sealed/SealedObjectInstancesProcessor.kt
@@ -52,32 +52,54 @@ internal class SealedObjectInstancesProcessor(
     override fun process(resolver: Resolver): List<KSAnnotated> {
         resolver.getSymbolsWithAnnotation(SealedObjectInstances::class.qualifiedName!!)
             .filterIsInstance<KSClassDeclaration>()
-            .filter { SEALED in it.modifiers }
-            .forEach(::process)
+            .groupBy { it.sealedClass() }
+            .filterNotNullValues()
+            .mapValues { it.annotations() }
+            .forEach { it.process() }
         return emptyList()
     }
 
-    private fun process(declaration: KSClassDeclaration) = declaration.annotations()
+    /**
+     * @return the corresponding sealed class declaration for the [this] declaration.
+     */
+    private fun KSClassDeclaration.sealedClass() = when {
+        SEALED in modifiers -> this
+        isCompanionObject -> parentDeclaration as KSClassDeclaration
+        else -> environment.logger.error(
+            message = "Failed to find a corresponding sealed class!",
+            symbol = this,
+        ).let { null }
+    }
+
+    @Suppress("UNCHECKED_CAST")
+    private fun <K : Any, V> Map<K?, V>.filterNotNullValues(): Map<K, V> = filterKeys { it != null } as Map<K, V>
+
+    /**
+     * @return the [List] of [SealedObjectInstances] annotations for [this] entry.
+     */
+    @OptIn(KspExperimental::class)
+    private fun Map.Entry<KSClassDeclaration, List<KSClassDeclaration>>.annotations() = value
+        .flatMap { it.getAnnotationsByType(SealedObjectInstances::class) }
+        .also {
+            if (it.size == it.distinctBy(SealedObjectInstances::name).size) return@also
+            environment.logger.error(
+                message = "Duplicated names: ${it.groupingBy(SealedObjectInstances::name).eachCount()}",
+                symbol = key,
+            )
+        }
+
+    private fun Map.Entry<KSClassDeclaration, List<SealedObjectInstances>>.process() = value
         .groupBy { it.fileName.takeUnless(String::isBlank) }
-        .mapKeys { declaration.createNewFile(it.key) }
+        .mapKeys { key.createNewFile(it.key) }
         .forEach { (file, annotations) ->
             file.writer().use {
-                it.appendHeader(declaration)
+                it.appendHeader(key)
                 annotations.forEach { annotation ->
-                    it.appendMethod(declaration, annotation)
+                    it.appendMethod(key, annotation)
                 }
             }
         }
 
-    @OptIn(KspExperimental::class)
-    private fun KSClassDeclaration.annotations() = getAnnotationsByType(SealedObjectInstances::class).toList().also {
-        if (it.size == it.distinctBy(SealedObjectInstances::name).size) return@also
-        environment.logger.error(
-            message = "Duplicated names: ${it.groupingBy(SealedObjectInstances::name).eachCount()}",
-            symbol = this,
-        )
-    }
-
     private fun KSClassDeclaration.createNewFile(fileName: String?) = runCatching {
         environment.codeGenerator.createNewFile(
             dependencies = Dependencies(
@@ -118,13 +140,28 @@ internal class SealedObjectInstancesProcessor(
             postfix = ")",
         ) { it.qualifiedName!!.asString() }
 
-        // language=kotlin
+        // language=kotlin ~ Extension on the KClass
         """
         /** @return [$rawClassName] of sealed object instances of type [$sealedClassName]. */
         ${visibility.modifier()} fun $receiverType.$methodName()$returnType = $collectionBuilder
         """.trimIndent().let(::appendLine)
+
+        sealed.companionOrNull()?.let {
+            // language=kotlin ~ Extension on the companion object
+            """
+            /** @return [$rawClassName] of sealed object instances of type [$sealedClassName]. */
+            ${visibility.modifier()} fun $sealedClassName$genericReceiverType.$it.$methodName()$returnType = $sealedClassName$genericReceiverType::class.$methodName()
+            """.trimIndent().let(::appendLine)
+        }
     }
 
+    /**
+     * @return the companion object of [this] declaration or `null` if it does not exist.
+     */
+    private fun KSClassDeclaration.companionOrNull() = declarations
+        .filterIsInstance<KSClassDeclaration>()
+        .singleOrNull { it.isCompanionObject }
+
     private fun KSClassDeclaration.getVisibility(
         annotation: SealedObjectInstances,
     ): Visibility = when (val visibility = annotation.visibility) {

diff --git a/processor/src/test/kotlin/fr/smarquis/sealed/CompanionObjectTest.kt b/processor/src/test/kotlin/fr/smarquis/sealed/CompanionObjectTest.kt
@@ -0,0 +1,69 @@
+/*
+ * Copyright (C) 2023 Simon Marquis
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *    https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package fr.smarquis.sealed
+
+import kotlin.test.Test
+import kotlin.test.assertSame
+
+class CompanionObjectTest {
+
+    @SealedObjectInstances
+    sealed class AnnotatedSealedClassWithCompanionObject {
+        companion object
+        object INSTANCE : AnnotatedSealedClassWithCompanionObject()
+    }
+
+    @Test
+    fun testAnnotatedSealedClassWithCompanionObject() = assertSame(
+        expected = AnnotatedSealedClassWithCompanionObject::class.sealedObjectInstances().single(),
+        actual = AnnotatedSealedClassWithCompanionObject.sealedObjectInstances().single(),
+    )
+
+    sealed class SealedClassWithAnnotatedCompanionObject {
+        @SealedObjectInstances companion object
+        object INSTANCE : SealedClassWithAnnotatedCompanionObject()
+    }
+
+    @Test
+    fun testSealedClassWithAnnotatedCompanionObject() = assertSame(
+        expected = SealedClassWithAnnotatedCompanionObject::class.sealedObjectInstances().single(),
+        actual = SealedClassWithAnnotatedCompanionObject.sealedObjectInstances().single(),
+    )
+
+    @SealedObjectInstances(name = "onSealedClass")
+    sealed class SealedClassWithCompanionObjectBothAnnotated {
+        @SealedObjectInstances(name = "onCompanionObject")
+        companion object
+        object INSTANCE : SealedClassWithCompanionObjectBothAnnotated()
+    }
+
+    @Test
+    fun testSealedClassWithCompanionObjectBothAnnotated() = assertSame(
+        expected = SealedClassWithAnnotatedCompanionObject::class.sealedObjectInstances().single(),
+        actual = SealedClassWithAnnotatedCompanionObject.sealedObjectInstances().single(),
+    )
+
+    sealed class SealedClassWithNameCompanionObject {
+        @SealedObjectInstances companion object Foo
+        object INSTANCE : SealedClassWithNameCompanionObject()
+    }
+
+    @Test
+    fun testSealedClassWithNameCompanionObject() = assertSame(
+        expected = SealedClassWithNameCompanionObject::class.sealedObjectInstances().single(),
+        actual = SealedClassWithNameCompanionObject.sealedObjectInstances().single(),
+    )
+}
diff --git a/processor/src/test/kotlin/fr/smarquis/sealed/NestedExtensionReferencesTest.kt b/processor/src/test/kotlin/fr/smarquis/sealed/NestedExtensionReferencesTest.kt
@@ -24,14 +24,14 @@ class NestedExtensionReferencesTest {
     sealed class Sealed {
         object Object : Sealed()
         companion object {
-            val values get() = Sealed::class.sealedObjectInstances()
-            val lazyValues by lazy(Sealed::class::sealedObjectInstances)
+            val values get() = Sealed.sealedObjectInstances()
+            val lazyValues by lazy(Companion::sealedObjectInstances)
         }
     }
 
     @Test
     fun `referencing extension from within the class compiles and returns the same values`() {
-        assertEquals(Sealed::class.sealedObjectInstances(), Sealed.values)
-        assertEquals(Sealed::class.sealedObjectInstances(), Sealed.lazyValues)
+        assertEquals(Sealed.sealedObjectInstances(), Sealed.values)
+        assertEquals(Sealed.sealedObjectInstances(), Sealed.lazyValues)
     }
 }