N-ary capture (#3)

README.md

@@ -50,4 +50,42 @@ An example of this is the code used to generate this README.md.
 Capturing source from blocks allows sample code to be run and
 tested during generation.
 
-View the source here: [readme/src/main/kotlin/Main.kt](readme/src/main/kotlin/Main.kt)
+View the source here: [readme/src/main/kotlin/Main.kt](readme/src/main/kotlin/Main.kt)
+
+## Parameterized Blocks
+
+The default `CapturedBlock` interface doesn't accept any
+arguments to `invoke` and is only generic on the return type. This
+means the captured source block must depend only on state from the
+enclosing scope. To write source capturing versions of builder blocks
+or common higher-order functions like `map` and `filter` you will
+need to define your own capture interface that extends `Capturable`.
+
+
+```kotlin
+/* must be a fun interface to support SAM conversion from blocks */
+fun interface CustomCapturable<T, R> : Capturable<CustomCapturable<T, R>> {
+    /* invoke is not special. this could be any single abstract method */
+    operator fun invoke(arg: T): R
+
+    /* withSource is called by the plugin to add source information */
+    override fun withSource(source: String): CustomCapturable<T, R> =
+        object : CustomCapturable<T, R> by this { override fun source(): String = source }
+}
+```
+
+Once you have declared your own `Capturable` you can use it
+in a similar way to `CapturedBlock` from above.
+
+```kotlin
+fun <T> List<T>.mapped(block: CustomCapturable<T, T>): String {
+    return "$this.map { ${block.source()} } = ${map { block(it) }}"
+}
+
+check(
+    listOf(1, 2, 3).mapped { x -> x*2 } ==
+    "[1, 2, 3].map { x -> x*2 } = [2, 4, 6]"
+)
+```
+
+If it is present, the block's argument list is considered part of its source text.

kapshot-plugin-kotlin/src/main/kotlin/io/koalaql/kapshot/plugin/CaptureTransformer.kt

@@ -13,13 +13,15 @@ import org.jetbrains.kotlin.ir.expressions.IrTypeOperator
 import org.jetbrains.kotlin.ir.expressions.IrTypeOperatorCall
 import org.jetbrains.kotlin.ir.symbols.IrSimpleFunctionSymbol
 import org.jetbrains.kotlin.ir.types.IrSimpleType
+import org.jetbrains.kotlin.ir.types.IrType
+import org.jetbrains.kotlin.ir.util.dumpKotlinLike
 import org.jetbrains.kotlin.ir.util.kotlinFqName
 import java.io.File
 
 class CaptureTransformer(
     private val context: IrPluginContext,
     private val addSourceToBlock: IrSimpleFunctionSymbol,
-    private val captureBlockFqn: String = "io.koalaql.kapshot.CapturedBlock"
+    private val capturableFqn: String = "io.koalaql.kapshot.Capturable"
 ): IrElementTransformerVoidWithContext() {
     private fun transformSam(expression: IrTypeOperatorCall): IrExpression {
         val symbol = currentScope!!.scope.scopeOwnerSymbol
@@ -49,7 +51,7 @@ class CaptureTransformer(
 
             val trimmed = fileText.substring(startOffset, endOffset).trimIndent().trim()
 
-            addSourceCall.putTypeArgument(0, context.irBuiltIns.stringType)
+            addSourceCall.putTypeArgument(0, expression.type)
 
             /* super call here rather than directly using expression is required to support nesting. otherwise we don't transform the subtree */
             addSourceCall.putValueArgument(0, super.visitTypeOperator(expression))
@@ -59,12 +61,21 @@ class CaptureTransformer(
         }
     }
 
+    private fun typeIsFqn(type: IrType, fqn: String): Boolean {
+        if (type !is IrSimpleType) return false
+
+        return when (val owner = type.classifier.owner) {
+            is IrClass -> owner.kotlinFqName.asString() == fqn
+            else -> false
+        }
+    }
+
     override fun visitTypeOperator(expression: IrTypeOperatorCall): IrExpression {
         if (expression.operator == IrTypeOperator.SAM_CONVERSION) {
             when (val type = expression.type) {
                 is IrSimpleType -> {
                     when (val owner = type.classifier.owner) {
-                        is IrClass -> if (owner.kotlinFqName.asString() == captureBlockFqn) {
+                        is IrClass -> if (owner.superTypes.any { typeIsFqn(it, capturableFqn) }) {
                             return transformSam(expression)
                         }
                     }

kapshot-plugin-kotlin/src/main/kotlin/io/koalaql/kapshot/plugin/GenerationExtension.kt

@@ -2,12 +2,15 @@ package io.koalaql.kapshot.plugin
 
 import org.jetbrains.kotlin.backend.common.extensions.IrGenerationExtension
 import org.jetbrains.kotlin.backend.common.extensions.IrPluginContext
+import org.jetbrains.kotlin.cli.common.messages.CompilerMessageSeverity
+import org.jetbrains.kotlin.cli.common.messages.MessageCollector
 import org.jetbrains.kotlin.ir.declarations.IrModuleFragment
 import org.jetbrains.kotlin.name.CallableId
 import org.jetbrains.kotlin.name.FqName
 import org.jetbrains.kotlin.name.Name
 
 class GenerationExtension(
+    private val messages: MessageCollector
 ) : IrGenerationExtension {
     override fun generate(moduleFragment: IrModuleFragment, pluginContext: IrPluginContext) {
         val addSourceToBlock = pluginContext

kapshot-plugin-kotlin/src/main/kotlin/io/koalaql/kapshot/plugin/Registrar.kt

@@ -2,6 +2,7 @@ package io.koalaql.kapshot.plugin
 
 import com.google.auto.service.AutoService
 import org.jetbrains.kotlin.backend.common.extensions.IrGenerationExtension
+import org.jetbrains.kotlin.cli.common.CLIConfigurationKeys
 import org.jetbrains.kotlin.compiler.plugin.CompilerPluginRegistrar
 import org.jetbrains.kotlin.config.CompilerConfiguration
 import org.jetbrains.kotlin.compiler.plugin.ExperimentalCompilerApi
@@ -12,6 +13,8 @@ class Registrar: CompilerPluginRegistrar() {
     override val supportsK2: Boolean get() = false
 
     override fun ExtensionStorage.registerExtensions(configuration: CompilerConfiguration) {
-        IrGenerationExtension.registerExtension(GenerationExtension())
+        IrGenerationExtension.registerExtension(GenerationExtension(
+            configuration.getNotNull(CLIConfigurationKeys.MESSAGE_COLLECTOR_KEY)
+        ))
     }
 }

kapshot-runtime/src/main/kotlin/io/koalaql/kapshot/Capturable.kt

@@ -0,0 +1,7 @@
+package io.koalaql.kapshot
+
+interface Capturable<T : Capturable<T>> {
+    fun source(): String = error("there is no source code for this block")
+
+    fun withSource(source: String): T
+}

kapshot-runtime/src/main/kotlin/io/koalaql/kapshot/CapturedBlock.kt

@@ -1,7 +1,11 @@
 package io.koalaql.kapshot
 
-fun interface CapturedBlock<T> {
+fun interface CapturedBlock<T>: Capturable<CapturedBlock<T>> {
     operator fun invoke(): T
 
-    fun source(): String = error("there is no source code for this block")
+    override fun withSource(source: String): CapturedBlock<T> {
+        return object : CapturedBlock<T> by this {
+            override fun source(): String = source
+        }
+    }
 }

kapshot-runtime/src/main/kotlin/io/koalaql/kapshot/CapturedBlocks.kt

@@ -1,6 +1,4 @@
 package io.koalaql.kapshot
 
-fun <T> addSourceToBlock(block: CapturedBlock<T>, source: String): CapturedBlock<T> = object : CapturedBlock<T> {
-    override fun invoke(): T = block()
-    override fun source(): String = source
-}
+fun <T : Capturable<T>> addSourceToBlock(block: Capturable<T>, source: String): T =
+    block.withSource(source)

readme/src/main/kotlin/Main.kt

@@ -1,3 +1,4 @@
+import io.koalaql.kapshot.Capturable
 import io.koalaql.kapshot.CapturedBlock
 import kotlin.io.path.Path
 import kotlin.io.path.writeText
@@ -7,7 +8,15 @@ fun execSource(block: CapturedBlock<*>): String {
     return block.source()
 }
 
+fun interface CustomCapturable<T, R> : Capturable<CustomCapturable<T, R>> {
+    operator fun invoke(arg: T): R
+
+    override fun withSource(source: String): CustomCapturable<T, R> =
+        object : CustomCapturable<T, R> by this { override fun source(): String = source }
+}
+
 fun generateMarkdown(): String {
+    val capturableClass = Capturable::class
     val blockClass = CapturedBlock::class
     val importStatement = "import ${blockClass.qualifiedName}"
 
@@ -73,6 +82,52 @@ Capturing source from blocks allows sample code to be run and
 tested during generation.
 
 View the source here: [readme/src/main/kotlin/Main.kt](readme/src/main/kotlin/Main.kt)
+
+## Parameterized Blocks
+
+The default `${blockClass.simpleName}` interface doesn't accept any
+arguments to `invoke` and is only generic on the return type. This
+means the captured source block must depend only on state from the
+enclosing scope. To write source capturing versions of builder blocks
+or common higher-order functions like `map` and `filter` you will
+need to define your own capture interface that extends `${capturableClass.simpleName}`.
+
+${""/*
+The code below can't be captured from a block due to limitations on where
+interfaces can be declared. Perhaps a future version of this plugin could
+add a sourceOf<T>() operator for class/interface definitions
+*/}
+```kotlin
+/* must be a fun interface to support SAM conversion from blocks */
+fun interface CustomCapturable<T, R> : Capturable<CustomCapturable<T, R>> {
+    /* invoke is not special. this could be any single abstract method */
+    operator fun invoke(arg: T): R
+
+    /* withSource is called by the plugin to add source information */
+    override fun withSource(source: String): CustomCapturable<T, R> =
+        object : CustomCapturable<T, R> by this { override fun source(): String = source }
+}
+```
+
+Once you have declared your own `${capturableClass.simpleName}` you can use it
+in a similar way to `${blockClass.simpleName}` from above.
+
+```kotlin
+${
+    execSource {
+        fun <T> List<T>.mapped(block: CustomCapturable<T, T>): String {
+            return "$this.map { ${block.source()} } = ${map { block(it) }}"
+        }
+
+        check(
+            listOf(1, 2, 3).mapped { x -> x*2 } ==
+            "[1, 2, 3].map { x -> x*2 } = [2, 4, 6]"
+        )
+    }
+}
+```
+
+If it is present, the block's argument list is considered part of its source text. 
     """.trim()
 }
 

testing/build.gradle.kts

@@ -2,12 +2,18 @@ plugins {
     id("io.koalaql.kapshot-plugin")
 
     kotlin("jvm") version "1.8.0"
+
+    application
 }
 
 repositories {
     mavenCentral()
 }
 
+application {
+    mainClass.set("MainKt")
+}
+
 dependencies {
     testImplementation(kotlin("test"))
 }

testing/src/test/kotlin/CapturedBlockTests.kt

@@ -1,3 +1,4 @@
+import io.koalaql.kapshot.Capturable
 import io.koalaql.kapshot.CapturedBlock
 import org.junit.jupiter.api.Test
 import kotlin.test.assertEquals
@@ -83,4 +84,24 @@ i""",
             2 + 2
         }}|""", "\uD83C\uDF85|2 + 2|")
     }
+
+    fun interface CapturedTest1<T, R>: Capturable<CapturedTest1<T, R>> {
+        operator fun invoke(arg: T): R
+
+        override fun withSource(source: String): CapturedTest1<T, R> = object : CapturedTest1<T, R> by this {
+            override fun source(): String = source
+        }
+    }
+
+    @Test
+    fun `user defined unary capture block`() {
+        fun List<Int>.sourceyMap(block: CapturedTest1<Int, Int>): String {
+            return "$this.map { ${block.source()} } = ${map { block(it) }}"
+        }
+
+        assertEquals(
+            "[1, 2, 3].map { it*2 } = [2, 4, 6]",
+            listOf(1, 2, 3).sourceyMap { it*2 }
+        )
+    }
 }