Support coroutine tests in Burst

We were previously broken with @burst tests that
called runTest, when executing on Kotlin/JS.

Author: squarejesse

burst-gradle-plugin/src/test/kotlin/app/cash/burst/gradle/BurstGradlePluginTest.kt

@@ -17,6 +17,7 @@
 package app.cash.burst.gradle
 
 import assertk.assertThat
+import assertk.assertions.contains
 import assertk.assertions.containsExactlyInAnyOrder
 import assertk.assertions.isEqualTo
 import assertk.assertions.isFalse
@@ -69,13 +70,29 @@ class BurstGradlePluginTest {
     // Each test class is executed normally with nothing skipped.
     with(tester.readTestSuite("CoffeeTest_Regular", testTaskName)) {
       assertThat(testCases.map { it.name }).containsExactlyInAnyOrder(
-        "test_Milk[$platformName]",
-        "test_None[$platformName]",
-        "test_Oat[$platformName]",
+        "basicTest_Milk[$platformName]",
+        "basicTest_None[$platformName]",
+        "basicTest_Oat[$platformName]",
+        "coroutinesTest_Milk[$platformName]",
+        "coroutinesTest_None[$platformName]",
+        "coroutinesTest_Oat[$platformName]",
       )
 
-      val sampleSpecialization = testCases.single { it.name == "test_Milk[$platformName]" }
+      val sampleSpecialization = testCases.single { it.name == "basicTest_Milk[$platformName]" }
       assertThat(sampleSpecialization.skipped).isFalse()
+
+      assertThat(systemOut).contains(
+        """
+        |set up Regular
+        |running Regular Oat in coffeeCoroutine
+        |
+        """.trimMargin(),
+        """
+        |set up Regular
+        |running Regular Oat
+        |
+        """.trimMargin(),
+      )
     }
 
     if (checkKlibMetadata) {
@@ -86,10 +103,14 @@ class BurstGradlePluginTest {
       val coffeeTestMetadata = klibMetadata.classes.first { it.name == "CoffeeTest" }
       assertThat(coffeeTestMetadata.functions.map { it.name }).containsExactlyInAnyOrder(
         "setUp",
-        "test",
-        "test_Milk",
-        "test_None",
-        "test_Oat",
+        "basicTest",
+        "basicTest_Milk",
+        "basicTest_None",
+        "basicTest_Oat",
+        "coroutinesTest",
+        "coroutinesTest_Milk",
+        "coroutinesTest_None",
+        "coroutinesTest_Oat",
       )
     }
   }

burst-gradle-plugin/src/test/projects/multiplatform/lib/build.gradle.kts

@@ -17,6 +17,8 @@ kotlin {
     commonTest {
       dependencies {
         implementation(kotlin("test"))
+        implementation(libs.kotlinx.coroutines.core)
+        implementation(libs.kotlinx.coroutines.test)
       }
     }
   }

burst-gradle-plugin/src/test/projects/multiplatform/lib/src/commonTest/kotlin/CoffeeTest.kt

@@ -1,6 +1,12 @@
 import app.cash.burst.Burst
+import kotlin.coroutines.coroutineContext
 import kotlin.test.BeforeTest
 import kotlin.test.Test
+import kotlin.time.Duration.Companion.milliseconds
+import kotlinx.coroutines.CoroutineName
+import kotlinx.coroutines.async
+import kotlinx.coroutines.delay
+import kotlinx.coroutines.test.runTest
 
 @Burst
 class CoffeeTest(
@@ -12,9 +18,18 @@ class CoffeeTest(
   }
 
   @Test
-  fun test(dairy: Dairy) {
+  fun basicTest(dairy: Dairy) {
     println("running $espresso $dairy")
   }
+
+  @Test
+  fun coroutinesTest(dairy: Dairy) = runTest(CoroutineName("coffeeCoroutine")) {
+    val deferred = async {
+      println("running $espresso $dairy in ${coroutineContext[CoroutineName]?.name}")
+    }
+    delay(1000.milliseconds)
+    deferred.await()
+  }
 }
 
 enum class Espresso { Decaf, Regular, Double }

burst-kotlin-plugin/src/main/kotlin/app/cash/burst/kotlin/Argument.kt

@@ -43,9 +43,6 @@ internal sealed interface Argument {
   /** True if this argument matches the default parameter value. */
   val isDefault: Boolean
 
-  /** Where to assign this argument to in a call. */
-  val indexInParameters: Int
-
   /** A string that's safe to use in a declaration name. */
   val name: String
 
@@ -64,9 +61,6 @@ private class EnumValueArgument(
 ) : Argument {
   override val name = value.name.identifier
 
-  override val indexInParameters: Int
-    get() = original.indexInParameters
-
   override fun expression() =
     IrGetEnumValueImpl(original.startOffset, original.endOffset, type, value.symbol)
 
@@ -83,9 +77,6 @@ private class BooleanArgument(
 ) : Argument {
   override val name = value.toString()
 
-  override val indexInParameters: Int
-    get() = original.indexInParameters
-
   override fun expression() =
     IrConstImpl.boolean(original.startOffset, original.endOffset, booleanType, value)
 
@@ -101,9 +92,6 @@ private class NullArgument(
 ) : Argument {
   override val name = "null"
 
-  override val indexInParameters: Int
-    get() = original.indexInParameters
-
   override fun expression() = IrConstImpl.constNull(original.startOffset, original.endOffset, type)
 
   override fun <R, D> accept(visitor: IrVisitor<R, D>, data: D): R {
@@ -118,8 +106,6 @@ private class BurstValuesArgument(
   index: Int,
 ) : Argument {
   override val isDefault = index == 0
-  override val indexInParameters: Int
-    get() = parameter.indexInParameters
   override val name = value.suggestedName() ?: index.toString()
 
   override fun expression() = value.deepCopyWithSymbols(parameter.parent)

burst-kotlin-plugin/src/main/kotlin/app/cash/burst/kotlin/BurstApis.kt

@@ -13,6 +13,8 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
+@file:OptIn(UnsafeDuringIrConstructionAPI::class)
+
 package app.cash.burst.kotlin
 
 import org.jetbrains.kotlin.backend.common.extensions.IrPluginContext
@@ -24,7 +26,9 @@ import org.jetbrains.kotlin.ir.symbols.IrClassSymbol
 import org.jetbrains.kotlin.ir.symbols.IrFunctionSymbol
 import org.jetbrains.kotlin.ir.symbols.IrPropertySymbol
 import org.jetbrains.kotlin.ir.symbols.IrSimpleFunctionSymbol
+import org.jetbrains.kotlin.ir.symbols.UnsafeDuringIrConstructionAPI
 import org.jetbrains.kotlin.ir.types.IrType
+import org.jetbrains.kotlin.ir.types.classFqName
 import org.jetbrains.kotlin.ir.types.classOrNull
 import org.jetbrains.kotlin.ir.types.defaultType
 import org.jetbrains.kotlin.ir.util.hasAnnotation
@@ -37,7 +41,8 @@ internal class BurstApis private constructor(
   private val testClassSymbols: List<IrClassSymbol>,
   val beforeTestSymbols: List<IrClassSymbol>,
   val afterTestSymbols: List<IrClassSymbol>,
-  val runTestSymbols: List<IrFunctionSymbol>,
+  /** Null if `kotlinx.coroutines.test` isn't in this build. */
+  val runTestSymbol: IrFunctionSymbol?,
 ) {
   val burstValues: IrFunctionSymbol = pluginContext.referenceFunctions(burstValuesId).single()
 
@@ -46,6 +51,7 @@ internal class BurstApis private constructor(
     testInterceptorClassId = burstFqPackage.classId("TestInterceptor"),
   )!!
 
+  /** Null if `app.cash.burst.coroutines` isn't in this build. */
   val coroutinesTestInterceptorApis: TestInterceptorApis? = pluginContext.testInterceptorApis(
     testFunctionClassId = burstCoroutinesFqPackage.classId("CoroutineTestFunction"),
     testInterceptorClassId = burstCoroutinesFqPackage.classId("CoroutineTestInterceptor"),
@@ -75,13 +81,8 @@ internal class BurstApis private constructor(
       .firstOrNull { it in afterTestSymbols }
   }
 
-  fun isEitherTestInterceptor(property: IrProperty): Boolean {
-    return testInterceptorApis.isTestInterceptor(property) ||
-      coroutinesTestInterceptorApis?.isTestInterceptor(property) == true
-  }
-
   fun isRunTest(irCall: IrCall): Boolean {
-    return irCall.symbol in runTestSymbols
+    return runTestSymbol != null && irCall.symbol == runTestSymbol
   }
 
   companion object {
@@ -114,14 +115,18 @@ internal class BurstApis private constructor(
         pluginContext.referenceClass(kotlinAfterTestClassId),
       )
 
-      val runTestSymbols = pluginContext.referenceFunctions(runTestId).toList()
+      val runTestSymbol = pluginContext.referenceFunctions(runTestId).singleOrNull {
+        it.owner.parameters.size == 3 &&
+          it.owner.parameters[0].type.classFqName == coroutineContextId.asSingleFqName() &&
+          it.owner.parameters[1].type.classFqName == durationId.asSingleFqName()
+      }
 
       return BurstApis(
         pluginContext = pluginContext,
         testClassSymbols = testClassSymbols,
         beforeTestSymbols = beforeTestSymbols,
         afterTestSymbols = afterTestSymbols,
-        runTestSymbols = runTestSymbols,
+        runTestSymbol = runTestSymbol,
       )
     }
   }
@@ -168,6 +173,12 @@ private fun IrPluginContext.testInterceptorApis(
 private val kotlinPackage = FqPackageName("kotlin")
 private val throwableAddSuppressedId = kotlinPackage.callableId("addSuppressed")
 
+private val kotlinCoroutinePackage = FqPackageName("kotlin.coroutines")
+private val coroutineContextId = kotlinCoroutinePackage.callableId("CoroutineContext")
+
+private val kotlinTimeFqPackage = FqPackageName("kotlin.time")
+private val durationId = kotlinTimeFqPackage.classId("Duration")
+
 private val burstFqPackage = FqPackageName("app.cash.burst")
 private val burstCoroutinesFqPackage = FqPackageName("app.cash.burst.coroutines")
 private val burstAnnotationId = burstFqPackage.classId("Burst")

burst-kotlin-plugin/src/main/kotlin/app/cash/burst/kotlin/BurstIrGenerationExtension.kt

@@ -33,6 +33,7 @@ class BurstIrGenerationExtension(
   override fun generate(moduleFragment: IrModuleFragment, pluginContext: IrPluginContext) {
     // Skip the rewrite if the Burst APIs aren't loaded. We don't expect to find @Burst anywhere.
     val burstApis = BurstApis.maybeCreate(pluginContext) ?: return
+    val testFunctionReader = TestFunctionReader(burstApis)
 
     val transformer = object : IrElementTransformerVoidWithContext() {
       override fun visitClassNew(declaration: IrClass): IrStatement {
@@ -61,16 +62,15 @@ class BurstIrGenerationExtension(
         val originalFunctions = classDeclaration.functions.toList()
 
         for (function in originalFunctions) {
-          val testAnnotationClassSymbol = burstApis.findTestAnnotation(function) ?: continue
+          val testFunction = testFunctionReader.readOrNull(function) ?: continue
           if (!classHasAtBurst && !function.hasAtBurst) continue
 
           try {
             val specializer = FunctionSpecializer(
               pluginContext = pluginContext,
               burstApis = burstApis,
               originalParent = classDeclaration,
-              original = function,
-              testAnnotationClassSymbol = testAnnotationClassSymbol,
+              original = testFunction,
             )
             specializer.generateSpecializations()
           } catch (e: BurstCompilationException) {

burst-kotlin-plugin/src/main/kotlin/app/cash/burst/kotlin/ClassSpecializer.kt

@@ -148,8 +148,9 @@ internal class ClassSpecializer(
           context = pluginContext,
           symbol = superConstructor.symbol,
         ) {
+          arguments.clear()
           for (argument in specialization.arguments) {
-            arguments[argument.indexInParameters] = argument.expression()
+            arguments += argument.expression()
           }
         }
         statements += irInstanceInitializerCall(
@@ -176,8 +177,9 @@ internal class ClassSpecializer(
           context = pluginContext,
           symbol = superConstructor.symbol,
         ) {
+          arguments.clear()
           for (argument in specialization.arguments) {
-            arguments[argument.indexInParameters] = argument.expression()
+            arguments += argument.expression()
           }
         }
       }