From e8dd9ee39997177d8e417d4d381d74665717c0b5 Mon Sep 17 00:00:00 2001
From: Jesse Wilson <jwilson@squareup.com>
Date: Mon, 28 Oct 2024 18:04:35 -0400
Subject: [PATCH] Don't crash if burstValues() has only one argument

---
 .../burst/kotlin/BurstKotlinPluginTest.kt     | 26 +++++++++++++++
 .../kotlin/app/cash/burst/kotlin/Argument.kt  | 33 +++++++++----------
 2 files changed, 42 insertions(+), 17 deletions(-)

diff --git a/burst-kotlin-plugin-tests/src/test/kotlin/app/cash/burst/kotlin/BurstKotlinPluginTest.kt b/burst-kotlin-plugin-tests/src/test/kotlin/app/cash/burst/kotlin/BurstKotlinPluginTest.kt
index 04485ce..dbb0642 100644
--- a/burst-kotlin-plugin-tests/src/test/kotlin/app/cash/burst/kotlin/BurstKotlinPluginTest.kt
+++ b/burst-kotlin-plugin-tests/src/test/kotlin/app/cash/burst/kotlin/BurstKotlinPluginTest.kt
@@ -19,6 +19,7 @@ import assertk.assertThat
 import assertk.assertions.contains
 import assertk.assertions.containsExactly
 import assertk.assertions.containsExactlyInAnyOrder
+import assertk.assertions.isEmpty
 import assertk.assertions.isFalse
 import assertk.assertions.isTrue
 import com.tschuchort.compiletesting.JvmCompilationResult
@@ -308,6 +309,31 @@ class BurstKotlinPluginTest {
     baseLog.clear()
   }
 
+  @Test
+  fun burstValuesJustOne() {
+    val result = compile(
+      sourceFile = SourceFile.kotlin(
+        "CoffeeTest.kt",
+        """
+        import app.cash.burst.Burst
+        import app.cash.burst.burstValues
+        import kotlin.test.Test
+
+        @Burst
+        class CoffeeTest {
+          @Test
+          fun test(volume: Int = burstValues(12)) {
+          }
+        }
+        """,
+      ),
+    )
+    assertEquals(KotlinCompilation.ExitCode.OK, result.exitCode, result.messages)
+
+    val baseClass = result.classLoader.loadClass("CoffeeTest")
+    assertThat(baseClass.testSuffixes).isEmpty()
+  }
+
   @Test
   fun burstValuesHasReasonableSymbolName() {
     val result = compile(
diff --git a/burst-kotlin-plugin/src/main/kotlin/app/cash/burst/kotlin/Argument.kt b/burst-kotlin-plugin/src/main/kotlin/app/cash/burst/kotlin/Argument.kt
index ad23f45..3057617 100644
--- a/burst-kotlin-plugin/src/main/kotlin/app/cash/burst/kotlin/Argument.kt
+++ b/burst-kotlin-plugin/src/main/kotlin/app/cash/burst/kotlin/Argument.kt
@@ -19,7 +19,6 @@ import org.jetbrains.kotlin.backend.common.extensions.IrPluginContext
 import org.jetbrains.kotlin.ir.IrElement
 import org.jetbrains.kotlin.ir.backend.js.utils.valueArguments
 import org.jetbrains.kotlin.ir.declarations.IrClass
-import org.jetbrains.kotlin.ir.declarations.IrDeclarationParent
 import org.jetbrains.kotlin.ir.declarations.IrEnumEntry
 import org.jetbrains.kotlin.ir.declarations.IrValueParameter
 import org.jetbrains.kotlin.ir.expressions.IrCall
@@ -87,13 +86,16 @@ private class BooleanArgument(
   }
 }
 
+@UnsafeDuringIrConstructionAPI
 private class BurstValuesArgument(
-  private val declarationParent: IrDeclarationParent,
-  override val isDefault: Boolean,
-  override val name: String,
-  val value: IrExpression,
+  private val parameter: IrValueParameter,
+  private val value: IrExpression,
+  index: Int,
 ) : Argument {
-  override fun expression() = value.deepCopyWithSymbols(declarationParent)
+  override val isDefault = index == 0
+  override val name = value.suggestedName() ?: index.toString()
+
+  override fun expression() = value.deepCopyWithSymbols(parameter.parent)
 
   override fun <R, D> accept(visitor: IrElementVisitor<R, D>, data: D): R {
     return value.accept(visitor, data)
@@ -136,24 +138,21 @@ private fun burstValuesArguments(
   burstApisCall: IrCall,
 ): List<Argument> {
   return buildList {
-    val defaultExpression = burstApisCall.valueArguments[0] ?: unexpectedParameter(parameter)
     add(
       BurstValuesArgument(
-        declarationParent = parameter.parent,
-        isDefault = true,
-        name = defaultExpression.suggestedName() ?: "0",
-        value = defaultExpression,
+        parameter = parameter,
+        value = burstApisCall.valueArguments[0] ?: unexpectedParameter(parameter),
+        index = size,
       ),
     )
 
-    for ((index, element) in (burstApisCall.valueArguments[1] as IrVararg).elements.withIndex()) {
-      val varargExpression = element as? IrExpression ?: unexpectedParameter(parameter)
+    val varargs = burstApisCall.valueArguments[1] as? IrVararg ?: return@buildList
+    for (element in varargs.elements) {
       add(
         BurstValuesArgument(
-          declarationParent = parameter.parent,
-          isDefault = false,
-          name = varargExpression.suggestedName() ?: (index + 1).toString(),
-          value = varargExpression,
+          parameter = parameter,
+          value = element as? IrExpression ?: unexpectedParameter(parameter),
+          index = size,
         ),
       )
     }