Fix how map keys are handled, adding ClassKey as first class support,…

… adding tests for all possible annotation values in custom keys.

.run/Desktop App.run.xml

@@ -1,6 +1,6 @@
 <component name="ProjectRunConfigurationManager">
   <configuration default="false" name="Desktop App" type="Application" factoryName="Application">
-    <option name="ALTERNATIVE_JRE_PATH" value="azul-17" />
+    <option name="ALTERNATIVE_JRE_PATH" value="zulu-17" />
     <option name="ALTERNATIVE_JRE_PATH_ENABLED" value="true" />
     <option name="MAIN_CLASS_NAME" value="com.r0adkll.kimchi.restaurant.MainKt" />
     <module name="kimchi.sample.desktopApp.main" />

annotations/src/commonMain/kotlin/com/r0adkll/kimchi/annotations/MapKey.kt

@@ -2,6 +2,8 @@
 // SPDX-License-Identifier: Apache-2.0
 package com.r0adkll.kimchi.annotations
 
+import kotlin.reflect.KClass
+
 /**
  * Defines a key that can be used for contributing multibinding elements to a map. Annotations
  * annotated with [MapKey] should have a single parameter of any type. Kimchi will then use this
@@ -27,26 +29,41 @@ package com.r0adkll.kimchi.annotations
  * ```
  */
 @Target(AnnotationTarget.ANNOTATION_CLASS)
-@Retention(AnnotationRetention.RUNTIME)
+@Retention(AnnotationRetention.BINARY)
 annotation class MapKey
 
 @MapKey
 @Target(AnnotationTarget.CLASS)
-@Retention(AnnotationRetention.RUNTIME)
+@Retention(AnnotationRetention.BINARY)
 annotation class StringKey(
   val value: String,
 )
 
 @MapKey
 @Target(AnnotationTarget.CLASS)
-@Retention(AnnotationRetention.RUNTIME)
+@Retention(AnnotationRetention.BINARY)
 annotation class IntKey(
   val value: Int,
 )
 
 @MapKey
 @Target(AnnotationTarget.CLASS)
-@Retention(AnnotationRetention.RUNTIME)
+@Retention(AnnotationRetention.BINARY)
 annotation class LongKey(
   val value: Long,
 )
+
+@MapKey
+@Target(AnnotationTarget.CLASS)
+@Retention(AnnotationRetention.BINARY)
+annotation class ClassKey(
+  val value: KClass<*>,
+)
+
+@MapKey
+annotation class ByteKey(
+  val value: Byte,
+)
+
+@ByteKey(0)
+class Test

compiler-utils/src/main/kotlin/com/r0adkll/kimchi/util/ksp/KspUtil.kt

@@ -6,6 +6,7 @@ import com.google.devtools.ksp.symbol.ClassKind
 import com.google.devtools.ksp.symbol.KSAnnotated
 import com.google.devtools.ksp.symbol.KSAnnotation
 import com.google.devtools.ksp.symbol.KSClassDeclaration
+import com.google.devtools.ksp.symbol.KSDeclaration
 import com.google.devtools.ksp.symbol.KSType
 import com.google.devtools.ksp.symbol.KSTypeAlias
 import com.google.devtools.ksp.symbol.KSTypeReference
@@ -69,3 +70,6 @@ public val KSType.classDeclaration: KSClassDeclaration
 
 public val KSClassDeclaration.isInterface: Boolean
   get() = this.classKind == ClassKind.INTERFACE
+
+public val KSDeclaration.safeRequiredQualifiedName: String
+  get() = qualifiedName!!.asString().replace(".", "_")

compiler-utils/src/testFixtures/kotlin/com/r0adkll/kimchi/ReflectionAssertionBuilder.kt

@@ -85,7 +85,8 @@ public inline fun Assertion.Builder<KClass<*>>.withFunction(
   name: String,
   crossinline block: Assertion.Builder<KFunction<*>>.() -> Unit,
 ): Assertion.Builder<KClass<*>> {
-  return with({ functions.find { it.name == name } }) {
+  println("Finding function for: $name")
+  return with(name, { functions.find { it.name == name } }) {
     isNotNull().block()
   }
 }

compiler/src/main/kotlin/com/r0adkll/kimchi/util/kotlinpoet/Bindings.kt

@@ -7,6 +7,7 @@ import com.google.devtools.ksp.symbol.KSClassDeclaration
 import com.r0adkll.kimchi.annotations.ContributesMultibinding
 import com.r0adkll.kimchi.util.buildFun
 import com.r0adkll.kimchi.util.buildProperty
+import com.r0adkll.kimchi.util.ksp.MapKeyValue
 import com.r0adkll.kimchi.util.ksp.findBindingTypeFor
 import com.r0adkll.kimchi.util.ksp.findMapKey
 import com.r0adkll.kimchi.util.ksp.findQualifier
@@ -17,7 +18,6 @@ import com.squareup.kotlinpoet.FunSpec
 import com.squareup.kotlinpoet.PropertySpec
 import com.squareup.kotlinpoet.TypeName
 import com.squareup.kotlinpoet.TypeSpec
-import com.squareup.kotlinpoet.asTypeName
 import com.squareup.kotlinpoet.ksp.addOriginatingKSFile
 import com.squareup.kotlinpoet.ksp.toAnnotationSpec
 import com.squareup.kotlinpoet.ksp.toClassName
@@ -134,12 +134,12 @@ private fun TypeSpec.Builder.addProvidesFunction(
 private fun TypeSpec.Builder.addMappingProvidesFunction(
   boundClass: KSClassDeclaration,
   boundType: KSClassDeclaration,
-  mapKey: Any,
+  mapKey: MapKeyValue,
   isBindable: Boolean,
 ) {
   addFunction(
-    FunSpec.buildFun("provide${boundType.simpleName.asString()}_$mapKey") {
-      returns(pairTypeOf(mapKey::class.asTypeName(), boundType.toClassName()))
+    FunSpec.buildFun("provide${boundType.simpleName.asString()}_${mapKey.functionSuffix()}") {
+      returns(pairTypeOf(mapKey.type(), boundType.toClassName()))
 
       addAnnotation(Provides::class)
       addAnnotation(IntoMap::class)
@@ -149,19 +149,13 @@ private fun TypeSpec.Builder.addMappingProvidesFunction(
       }
 
       if (isBindable) {
+        val (format, value) = mapKey.value()
         addParameter("value", boundClass.toClassName())
-        if (mapKey is String) {
-          addStatement("return (%S to value)", mapKey)
-        } else {
-          addStatement("return (%L to value)", mapKey)
-        }
+        addStatement("return ($format to value)", value)
       } else {
+        val (format, value) = mapKey.value()
         val valueTemplate = if (boundClass.classKind == ClassKind.OBJECT) "%T" else "%T()"
-        if (mapKey is String) {
-          addStatement("return (%S to $valueTemplate)", mapKey, boundClass.toClassName())
-        } else {
-          addStatement("return (%L to $valueTemplate)", mapKey, boundClass.toClassName())
-        }
+        addStatement("return ($format to $valueTemplate)", value, boundClass.toClassName())
       }
     },
   )

compiler/src/main/kotlin/com/r0adkll/kimchi/util/ksp/MapKeys.kt

@@ -2,18 +2,26 @@
 // SPDX-License-Identifier: Apache-2.0
 package com.r0adkll.kimchi.util.ksp
 
+import com.google.devtools.ksp.symbol.ClassKind
 import com.google.devtools.ksp.symbol.KSClassDeclaration
+import com.google.devtools.ksp.symbol.KSType
+import com.google.devtools.ksp.symbol.KSValueArgument
 import com.r0adkll.kimchi.annotations.MapKey
+import com.r0adkll.kimchi.util.KimchiException
 import com.squareup.kotlinpoet.ParameterizedTypeName
 import com.squareup.kotlinpoet.ParameterizedTypeName.Companion.parameterizedBy
+import com.squareup.kotlinpoet.STAR
 import com.squareup.kotlinpoet.TypeName
 import com.squareup.kotlinpoet.asTypeName
+import com.squareup.kotlinpoet.ksp.toClassName
+import com.squareup.kotlinpoet.ksp.toTypeName
+import kotlin.reflect.KClass
 
 /**
  * Find the map key value included in [com.r0adkll.kimchi.annotations.ContributesMultibinding]
  * usage to use as the key when generating the bindings on the graph
  */
-fun KSClassDeclaration.findMapKey(): Any? {
+fun KSClassDeclaration.findMapKey(): MapKeyValue? {
   val mapKeyAnnotation = annotations.find { annotation ->
     annotation.isAnnotationOf(MapKey::class)
   } ?: return null
@@ -22,11 +30,124 @@ fun KSClassDeclaration.findMapKey(): Any? {
     .firstOrNull()
     ?: error("MapKey's must define a single argument to use as the value of the key.")
 
-  return mapKeyArgument.value
-    ?: error("MapKey is not provided with a single value that we can find")
+  return MapKeyValue(mapKeyArgument)
 }
 
 fun pairTypeOf(vararg typeNames: TypeName): ParameterizedTypeName {
   return Pair::class.asTypeName()
     .parameterizedBy(*typeNames)
 }
+
+data class MapKeyValue(private val valueArgument: KSValueArgument) {
+
+  fun functionSuffix(): String {
+    val value = valueArgument.value ?: throw KimchiException(
+      "Unable to determine the MapKey value: ${valueArgument.name?.asString()} = ${valueArgument.value}",
+    )
+
+    return when (value) {
+      is Byte -> "${value}b"
+      is Short -> "${value}s"
+
+      is Int,
+      is Long,
+      is Char,
+      is String,
+      is Boolean,
+      -> value.toString()
+
+      is Float,
+      is Double,
+      -> value.toString().replace(".", "_")
+
+      is KSType -> value.declaration.simpleName.asString()
+      is KSClassDeclaration -> value.safeRequiredQualifiedName
+      else -> throw KimchiException(
+        "Unable to determine the MapKey name: ${valueArgument.name?.asString()} = ${valueArgument.value}",
+      )
+    }
+  }
+
+  fun type(): TypeName {
+    val value = valueArgument.value ?: throw KimchiException(
+      "Unable to determine the MapKey type: ${valueArgument.name?.asString()} = ${valueArgument.value}",
+    )
+    return when (value) {
+      is Byte -> Byte::class.asTypeName()
+      is Short -> Short::class.asTypeName()
+      is Int -> Int::class.asTypeName()
+      is Long -> Long::class.asTypeName()
+      is Float -> Float::class.asTypeName()
+      is Double -> Double::class.asTypeName()
+      is Char -> Char::class.asTypeName()
+      is String -> String::class.asTypeName()
+      is Boolean -> Boolean::class.asTypeName()
+
+      is KSType -> KClass::class.asTypeName().parameterizedBy(STAR)
+      is KSClassDeclaration -> when (value.classKind) {
+        ClassKind.ENUM_CLASS -> value.toClassName()
+        ClassKind.ENUM_ENTRY -> (value.parent as? KSClassDeclaration)?.toClassName()
+          ?: throw KimchiException(
+            "Unsupported map key detected",
+            value,
+          )
+        else -> throw KimchiException(
+          "Unsupported map key detected",
+          value,
+        )
+      }
+      else -> throw KimchiException(
+        "Unable to determine the map key type",
+        valueArgument,
+      )
+    }
+  }
+
+  fun value(): Pair<String, Any> {
+    val value = valueArgument.value ?: throw KimchiException(
+      "Unable to determine the MapKey value: ${valueArgument.name?.asString()} = ${valueArgument.value}",
+    )
+
+    val format = when (value) {
+      is Byte,
+      is Short,
+      is Int,
+      is Long,
+      is Float,
+      is Double,
+      is Boolean,
+      -> "%L"
+
+      is Char -> "'%L'"
+      is String -> "%S"
+
+      is KSType -> "%T::class"
+      is KSClassDeclaration -> "%L"
+      else -> throw KimchiException(
+        "Unable to determine the MapKey value: ${valueArgument.name?.asString()} = ${valueArgument.value}",
+      )
+    }
+
+    val argValue = when (value) {
+      is Byte -> "$value.toByte()"
+      is Short -> "$value.toShort()"
+      is Float -> "$value.toFloat()"
+
+      is Int,
+      is Long,
+      is Double,
+      is Boolean,
+      is Char,
+      is String,
+      -> value
+
+      is KSType -> value.toTypeName()
+      is KSClassDeclaration -> value
+      else -> throw KimchiException(
+        "Unable to determine the MapKey value: ${valueArgument.name?.asString()} = ${valueArgument.value}",
+      )
+    }
+
+    return format to argValue
+  }
+}

compiler/src/test/kotlin/com/r0adkll/kimchi/TestSources.kt

@@ -69,6 +69,68 @@ val SingleIn = """
 val JvmCompilationResult.singleIn: KClass<*>
   get() = kotlinClass("kimchi.SingleIn")
 
+@Language("kotlin")
+val ByteKey = """
+  package kimchi
+  import com.r0adkll.kimchi.annotations.MapKey
+  @MapKey
+  annotation class ByteKey(val value: Byte)
+""".trimIndent()
+
+@Language("kotlin")
+val ShortKey = """
+  package kimchi
+  import com.r0adkll.kimchi.annotations.MapKey
+  @MapKey
+  annotation class ShortKey(val value: Short)
+""".trimIndent()
+
+@Language("kotlin")
+val FloatKey = """
+  package kimchi
+  import com.r0adkll.kimchi.annotations.MapKey
+  @MapKey
+  annotation class FloatKey(val value: Float)
+""".trimIndent()
+
+@Language("kotlin")
+val DoubleKey = """
+  package kimchi
+  import com.r0adkll.kimchi.annotations.MapKey
+  @MapKey
+  annotation class DoubleKey(val value: Double)
+""".trimIndent()
+
+@Language("kotlin")
+val CharKey = """
+  package kimchi
+  import com.r0adkll.kimchi.annotations.MapKey
+  @MapKey
+  annotation class CharKey(val value: Char)
+""".trimIndent()
+
+@Language("kotlin")
+val BooleanKey = """
+  package kimchi
+  import com.r0adkll.kimchi.annotations.MapKey
+  @MapKey
+  annotation class BooleanKey(val value: Boolean)
+""".trimIndent()
+
+enum class Keynum {
+  First, Second, Third
+}
+
+@Language("kotlin")
+val EnumKey = """
+  package kimchi
+  import com.r0adkll.kimchi.annotations.MapKey
+  import com.r0adkll.kimchi.Keynum
+
+  @MapKey
+  annotation class EnumKey(val value: Keynum)
+""".trimIndent()
+
 /**
  * Add source files here to include in every compiler test
  */
@@ -78,6 +140,13 @@ val commonTestSources = arrayOf(
   TestQualifier,
   TestComponent,
   SingleIn,
+  ByteKey,
+  ShortKey,
+  FloatKey,
+  DoubleKey,
+  CharKey,
+  BooleanKey,
+  EnumKey,
 )
 
 /**