Follow-up on primitive value classes

mockito-kotlin/src/main/kotlin/org/mockito/kotlin/ArgumentCaptor.kt

@@ -25,43 +25,84 @@
 
 package org.mockito.kotlin
 
-import org.mockito.kotlin.internal.createInstance
 import org.mockito.ArgumentCaptor
+import org.mockito.kotlin.internal.toKotlinType
+import org.mockito.kotlin.internal.createInstance
+import org.mockito.kotlin.internal.valueClassInnerClass
 import java.lang.reflect.Array
 import kotlin.reflect.KClass
+import kotlin.reflect.KType
+import kotlin.reflect.typeOf
 
 /**
  * Creates a [KArgumentCaptor] for given type.
+ *
+ * Caution: this factory method cannot be used to create a captor for a suspend
+ * function, please refer to [suspendFunctionArgumentCaptor] for that.
+ * This incompatibility is caused by the use of `typeOf<T>()` which is the way
+ * to determine runtime nullability of T, but this function is yet incompatible with
+ * suspend functions at compile time. That incompatibility has been declared since
+ * Kotlin 1.6, and the promised proper support for suspend functions has not been
+ * delivered ever since.
+ * See [Kotlin issue KT-47562](https://youtrack.jetbrains.com/issue/KT-47562/Support-suspend-functional-types-in-typeOf)
+ * for more details.
+ */
+inline fun <reified T : Any?> argumentCaptor(): KArgumentCaptor<T> {
+    return KArgumentCaptor(typeOf<T>())
+}
+
+/**
+ * Creates a [KArgumentCaptor] for given (suspend function).
  */
-inline fun <reified T : Any> argumentCaptor(): KArgumentCaptor<T> {
+inline fun <reified T : Function<*>> suspendFunctionArgumentCaptor(): KArgumentCaptor<T> {
     return KArgumentCaptor(T::class)
 }
 
 /**
  * Creates 2 [KArgumentCaptor]s for given types.
+ *
+ * Caution: this factory method cannot be used to create a captor for a suspend
+ * function, please refer to [suspendFunctionArgumentCaptor] for that.
+ * This incompatibility is caused by the use of `typeOf<T>()` which is the way
+ * to determine runtime nullability of T, but this function is yet incompatible with
+ * suspend functions at compile time. That incompatibility has been declared since
+ * Kotlin 1.6, and the promised proper support for suspend functions has not been
+ * delivered ever since.
+ * See [Kotlin issue KT-47562](https://youtrack.jetbrains.com/issue/KT-47562/Support-suspend-functional-types-in-typeOf)
+ * for more details.
  */
 inline fun <reified A : Any, reified B : Any> argumentCaptor(
-    a: KClass<A> = A::class,
-    b: KClass<B> = B::class
+    @Suppress("unused") a: KClass<A> = A::class,
+    @Suppress("unused") b: KClass<B> = B::class
 ): Pair<KArgumentCaptor<A>, KArgumentCaptor<B>> {
     return Pair(
-        KArgumentCaptor(a),
-        KArgumentCaptor(b)
+        KArgumentCaptor(typeOf<A>()),
+        KArgumentCaptor(typeOf<B>())
     )
 }
 
 /**
  * Creates 3 [KArgumentCaptor]s for given types.
+ *
+ * Caution: this factory method cannot be used to create a captor for a suspend
+ * function, please refer to [suspendFunctionArgumentCaptor] for that.
+ * This incompatibility is caused by the use of `typeOf<T>()` which is the way
+ * to determine runtime nullability of T, but this function is yet incompatible with
+ * suspend functions at compile time. That incompatibility has been declared since
+ * Kotlin 1.6, and the promised proper support for suspend functions has not been
+ * delivered ever since.
+ * See [Kotlin issue KT-47562](https://youtrack.jetbrains.com/issue/KT-47562/Support-suspend-functional-types-in-typeOf)
+ * for more details.
  */
 inline fun <reified A : Any, reified B : Any, reified C : Any> argumentCaptor(
-    a: KClass<A> = A::class,
-    b: KClass<B> = B::class,
-    c: KClass<C> = C::class
+    @Suppress("unused") a: KClass<A> = A::class,
+    @Suppress("unused") b: KClass<B> = B::class,
+    @Suppress("unused") c: KClass<C> = C::class
 ): Triple<KArgumentCaptor<A>, KArgumentCaptor<B>, KArgumentCaptor<C>> {
     return Triple(
-        KArgumentCaptor(a),
-        KArgumentCaptor(b),
-        KArgumentCaptor(c)
+        KArgumentCaptor(typeOf<A>()),
+        KArgumentCaptor(typeOf<B>()),
+        KArgumentCaptor(typeOf<C>())
     )
 }
 
@@ -71,21 +112,46 @@ class ArgumentCaptorHolder4<out A, out B, out C, out D>(
     val third: C,
     val fourth: D
 ) {
-
     operator fun component1() = first
     operator fun component2() = second
     operator fun component3() = third
     operator fun component4() = fourth
 }
 
+/**
+ * Creates 4 [KArgumentCaptor]s for given types.
+ *
+ * Caution: this factory method cannot be used to create a captor for a suspend
+ * function, please refer to [suspendFunctionArgumentCaptor] for that.
+ * This incompatibility is caused by the use of `typeOf<T>()` which is the way
+ * to determine runtime nullability of T, but this function is yet incompatible with
+ * suspend functions at compile time. That incompatibility has been declared since
+ * Kotlin 1.6, and the promised proper support for suspend functions has not been
+ * delivered ever since.
+ * See [Kotlin issue KT-47562](https://youtrack.jetbrains.com/issue/KT-47562/Support-suspend-functional-types-in-typeOf)
+ * for more details.
+ */
+inline fun <reified A : Any, reified B : Any, reified C : Any, reified D : Any> argumentCaptor(
+    @Suppress("unused") a: KClass<A> = A::class,
+    @Suppress("unused") b: KClass<B> = B::class,
+    @Suppress("unused") c: KClass<C> = C::class,
+    @Suppress("unused") d: KClass<D> = D::class
+): ArgumentCaptorHolder4<KArgumentCaptor<A>, KArgumentCaptor<B>, KArgumentCaptor<C>, KArgumentCaptor<D>> {
+    return ArgumentCaptorHolder4(
+        KArgumentCaptor(typeOf<A>()),
+        KArgumentCaptor(typeOf<B>()),
+        KArgumentCaptor(typeOf<C>()),
+        KArgumentCaptor(typeOf<D>())
+    )
+}
+
 class ArgumentCaptorHolder5<out A, out B, out C, out D, out E>(
     val first: A,
     val second: B,
     val third: C,
     val fourth: D,
     val fifth: E
 ) {
-
     operator fun component1() = first
     operator fun component2() = second
     operator fun component3() = third
@@ -94,38 +160,31 @@ class ArgumentCaptorHolder5<out A, out B, out C, out D, out E>(
 }
 
 /**
- * Creates 4 [KArgumentCaptor]s for given types.
- */
-inline fun <reified A : Any, reified B : Any, reified C : Any, reified D : Any> argumentCaptor(
-    a: KClass<A> = A::class,
-    b: KClass<B> = B::class,
-    c: KClass<C> = C::class,
-    d: KClass<D> = D::class
-): ArgumentCaptorHolder4<KArgumentCaptor<A>, KArgumentCaptor<B>, KArgumentCaptor<C>, KArgumentCaptor<D>> {
-    return ArgumentCaptorHolder4(
-        KArgumentCaptor(a),
-        KArgumentCaptor(b),
-        KArgumentCaptor(c),
-        KArgumentCaptor(d)
-    )
-}
-
-/**
- * Creates 4 [KArgumentCaptor]s for given types.
+ * Creates 5 [KArgumentCaptor]s for given types.
+ *
+ * Caution: this factory method cannot be used to create a captor for a suspend
+ * function, please refer to [suspendFunctionArgumentCaptor] for that.
+ * This incompatibility is caused by the use of `typeOf<T>()` which is the way
+ * to determine runtime nullability of T, but this function is yet incompatible with
+ * suspend functions at compile time. That incompatibility has been declared since
+ * Kotlin 1.6, and the promised proper support for suspend functions has not been
+ * delivered ever since.
+ * See [Kotlin issue KT-47562](https://youtrack.jetbrains.com/issue/KT-47562/Support-suspend-functional-types-in-typeOf)
+ * for more details.
  */
 inline fun <reified A : Any, reified B : Any, reified C : Any, reified D : Any, reified E : Any> argumentCaptor(
-    a: KClass<A> = A::class,
-    b: KClass<B> = B::class,
-    c: KClass<C> = C::class,
-    d: KClass<D> = D::class,
-    e: KClass<E> = E::class
+    @Suppress("unused") a: KClass<A> = A::class,
+    @Suppress("unused") b: KClass<B> = B::class,
+    @Suppress("unused") c: KClass<C> = C::class,
+    @Suppress("unused") d: KClass<D> = D::class,
+    @Suppress("unused") e: KClass<E> = E::class
 ): ArgumentCaptorHolder5<KArgumentCaptor<A>, KArgumentCaptor<B>, KArgumentCaptor<C>, KArgumentCaptor<D>, KArgumentCaptor<E>> {
     return ArgumentCaptorHolder5(
-        KArgumentCaptor(a),
-        KArgumentCaptor(b),
-        KArgumentCaptor(c),
-        KArgumentCaptor(d),
-        KArgumentCaptor(e)
+        KArgumentCaptor(typeOf<A>()),
+        KArgumentCaptor(typeOf<B>()),
+        KArgumentCaptor(typeOf<C>()),
+        KArgumentCaptor(typeOf<D>()),
+        KArgumentCaptor(typeOf<E>())
     )
 }
 
@@ -140,7 +199,7 @@ inline fun <reified T : Any> argumentCaptor(f: KArgumentCaptor<T>.() -> Unit): K
  * Creates a [KArgumentCaptor] for given nullable type.
  */
 inline fun <reified T : Any> nullableArgumentCaptor(): KArgumentCaptor<T?> {
-    return KArgumentCaptor(T::class)
+    return KArgumentCaptor(typeOf<T>())
 }
 
 /**
@@ -157,58 +216,61 @@ inline fun <reified T : Any> capture(captor: ArgumentCaptor<T>): T {
     return captor.capture() ?: createInstance()
 }
 
-class KArgumentCaptor<out T : Any?> (
-    private val tClass: KClass<*>
+class KArgumentCaptor<out T : Any?>(
+    private val clazz: KClass<*>,
+    private val isMarkedNullable: Boolean = false
 ) {
+    constructor(kType: KType):this(
+        kType.classifier as KClass<*>,
+        kType.isMarkedNullable
+    )
+
     private val captor: ArgumentCaptor<Any?> =
-        if (tClass.isValue) {
-            val boxImpl =
-                tClass.java.declaredMethods
-                    .single { it.name == "box-impl" && it.parameterCount == 1 }
-            boxImpl.parameters[0].type // is the boxed type of the value type
+        if (clazz.isValue && !isMarkedNullable) {
+            clazz.valueClassInnerClass()
         } else {
-            tClass.java
+            clazz
         }.let {
-            ArgumentCaptor.forClass(it)
+            ArgumentCaptor.forClass(it.java)
         }
 
     /**
      * The first captured value of the argument.
      * @throws IndexOutOfBoundsException if the value is not available.
      */
     val firstValue: T
-        get() = toKotlinType(captor.firstValue)
+        get() = captor.firstValue.toKotlinType(clazz)
 
     /**
      * The second captured value of the argument.
      * @throws IndexOutOfBoundsException if the value is not available.
      */
     val secondValue: T
-        get() = toKotlinType(captor.secondValue)
+        get() = captor.secondValue.toKotlinType(clazz)
 
     /**
      * The third captured value of the argument.
      * @throws IndexOutOfBoundsException if the value is not available.
      */
     val thirdValue: T
-        get() = toKotlinType(captor.thirdValue)
+        get() = captor.thirdValue.toKotlinType(clazz)
 
     /**
      * The last captured value of the argument.
      * @throws IndexOutOfBoundsException if the value is not available.
      */
     val lastValue: T
-        get() = toKotlinType(captor.lastValue)
+        get() = captor.lastValue.toKotlinType(clazz)
 
     /**
      * The *only* captured value of the argument,
      * or throws an exception if no value or more than one value was captured.
      */
     val singleValue: T
-        get() = toKotlinType(captor.singleValue)
+        get() = captor.singleValue.toKotlinType(clazz)
 
     val allValues: List<T>
-        get() = captor.allValues.map(::toKotlinType)
+        get() = captor.allValues.map { it.toKotlinType(clazz) }
 
     @Suppress("UNCHECKED_CAST")
     fun capture(): T {
@@ -219,28 +281,14 @@ class KArgumentCaptor<out T : Any?> (
         // In Java, `captor.capture` returns null and so the method is called with `[null]`
         // In Kotlin, we have to create `[null]` explicitly.
         // This code-path is applied for non-vararg array arguments as well, but it seems to work fine.
-        return captor.capture() as T ?: if (tClass.java.isArray) {
+        return captor.capture().toKotlinType(clazz) ?: if (clazz.java.isArray) {
             singleElementArray()
         } else {
-            createInstance(tClass)
+            createInstance(clazz)
         } as T
     }
 
-    private fun singleElementArray(): Any? = Array.newInstance(tClass.java.componentType, 1)
-
-    @Suppress("UNCHECKED_CAST")
-    private fun toKotlinType(rawCapturedValue: Any?) : T {
-        return if(tClass.isValue) {
-            rawCapturedValue
-                ?.let {
-                    val boxImpl =
-                        tClass.java.declaredMethods.single { it.name == "box-impl" && it.parameterCount == 1 }
-                    boxImpl.invoke(null, it)
-                } as T
-        } else {
-            rawCapturedValue as T
-        }
-    }
+    private fun singleElementArray(): Any? = Array.newInstance(clazz.java.componentType, 1)
 }
 
 val <T> ArgumentCaptor<T>.firstValue: T