Fix a crash where we mixed up void and Unit (#174)

We need to use a helper to move statements from
one function to another, otherwise early returns break.

Co-authored-by: Jesse Wilson <jwilson@squareup.com>

Author: swankjesse

TROUBLESHOOTING.md

@@ -0,0 +1,20 @@
+Burst Troubleshooting
+=====================
+
+NoSuchMethodError
+-----------------
+
+```
+java.lang.NoSuchMethodError: 'void com.example.AbstractFooTest.intercept(app.cash.burst.TestFunction)'
+    at com.example.StandardFooTest.intercept(StandardFooTest.kt:100)
+```
+
+You’ll get this if the Burst Gradle Plugin wasn’t applied to a class that needs it. Update the
+corresponding `build.gradle.kts` to apply it:
+
+```kotlin
+plugins {
+  id("app.cash.burst")
+  ...
+}
+```

burst-kotlin-plugin-tests/src/test/kotlin/app/cash/burst/kotlin/TestInterceptorKotlinPluginTest.kt

@@ -1480,4 +1480,59 @@ class TestInterceptorKotlinPluginTest {
       "running BaseTest.happyPath",
     )
   }
+
+  /**
+   * Our initial implementation failed bytecode validation when a function returned 'void' when the
+   * interceptor required it to return 'Unit'.
+   */
+  @Test
+  fun earlyReturn() {
+    val log = BurstTester(
+      packageName = "com.example",
+    ).compileAndRun(
+      SourceFile.kotlin(
+        "Main.kt",
+        """
+        package com.example
+
+        import app.cash.burst.InterceptTest
+        import app.cash.burst.TestFunction
+        import app.cash.burst.TestInterceptor
+        import kotlin.test.AfterTest
+        import kotlin.test.BeforeTest
+        import kotlin.test.Test
+
+        class SampleTest {
+          @InterceptTest
+          val interceptor = object : TestInterceptor {
+            override fun intercept(testFunction: TestFunction) {
+              log("intercepting")
+              testFunction()
+              log("intercepted")
+            }
+          }
+
+          @Test
+          fun happyPath() {
+            if (true) {
+              log("early return")
+              return
+            }
+            log("no early return")
+          }
+        }
+
+        fun main(vararg args: String) {
+          SampleTest().happyPath()
+        }
+        """,
+      ),
+    )
+
+    assertThat(log).containsExactly(
+      "intercepting",
+      "early return",
+      "intercepted",
+    )
+  }
 }

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

@@ -50,7 +50,6 @@ import org.jetbrains.kotlin.ir.util.defaultType
 import org.jetbrains.kotlin.ir.util.packageFqName
 import org.jetbrains.kotlin.ir.util.parentClassOrNull
 import org.jetbrains.kotlin.ir.util.patchDeclarationParents
-import org.jetbrains.kotlin.ir.util.statements
 import org.jetbrains.kotlin.ir.util.superClass
 import org.jetbrains.kotlin.name.Name
 
@@ -364,7 +363,7 @@ internal class InterceptorInjector(
     // If there's no function body to rewrite, we're probably looking at a superclass from another
     // module. The rewrite won't be emitted anywhere, but we still need to generate its symbols for
     // subclasses to call.
-    val originalBody = original.body ?: return
+    if (original.body == null) return
 
     if (original.modality != Modality.FINAL && originalParent.modality != Modality.FINAL) {
       unexpectedOpenFunction(original)
@@ -380,9 +379,7 @@ internal class InterceptorInjector(
         arguments[0] = irString(packageName)
         arguments[1] = irString(className)
         arguments[2] = irString(original.name.asString())
-        arguments[3] = localLambda(original.symbol) {
-          +originalBody.statements
-        }
+        arguments[3] = moveBodyToLocalLambda(original)
       }
 
       +irCall(

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

@@ -15,6 +15,7 @@
  */
 package app.cash.burst.kotlin
 
+import org.jetbrains.kotlin.backend.common.ir.moveBodyTo
 import org.jetbrains.kotlin.backend.common.lower.DeclarationIrBuilder
 import org.jetbrains.kotlin.backend.common.lower.irCatch
 import org.jetbrains.kotlin.cli.common.messages.CompilerMessageLocationWithRange
@@ -227,6 +228,26 @@ fun IrBlockBodyBuilder.localLambda(
   )
 }
 
+/** Moves the body of [original] to a newly-created local lambda. */
+fun IrBlockBodyBuilder.moveBodyToLocalLambda(
+  original: IrSimpleFunction,
+): IrFunctionExpressionImpl {
+  return IrFunctionExpressionImpl(
+    startOffset = startOffset,
+    endOffset = endOffset,
+    type = context.irBuiltIns.functionN(0).typeWith(context.irBuiltIns.unitType),
+    function = context.irFactory.buildFun {
+      this.name = Name.special("<anonymous>")
+      this.returnType = context.irBuiltIns.unitType
+      this.origin = IrDeclarationOrigin.Companion.LOCAL_FUNCTION_FOR_LAMBDA
+      this.visibility = DescriptorVisibilities.LOCAL
+    }.apply {
+      body = original.moveBodyTo(this, mapOf())
+    },
+    origin = IrStatementOrigin.Companion.LAMBDA,
+  )
+}
+
 internal fun IrBlockBuilder.irAccumulateFailure(
   burstApis: BurstApis,
   failure: IrVariable,