KT-81448: Support wildcards for power-assert function selection

libraries/tools/kotlin-power-assert/src/common/kotlin/org/jetbrains/kotlin/powerassert/gradle/PowerAssertGradleExtension.kt

@@ -21,7 +21,9 @@ package org.jetbrains.kotlin.powerassert.gradle
 
 import org.gradle.api.model.ObjectFactory
 import org.gradle.api.provider.SetProperty
+import org.intellij.lang.annotations.Language
 import org.jetbrains.kotlin.gradle.ExperimentalKotlinGradlePluginApi
+import java.util.regex.Pattern
 import javax.inject.Inject
 
 @ExperimentalKotlinGradlePluginApi
@@ -34,6 +36,19 @@ abstract class PowerAssertGradleExtension @Inject constructor(
      */
     val functions: SetProperty<String> = objectFactory.setProperty(String::class.java).convention(setOf("kotlin.assert"))
 
+    /**
+     * Defines regexes that are used to match functions to transform in addition to those in [functions].
+     * Regexes are applied to the fully-qualified path of the function, as used in [functions].
+     * Any function whose fully-qualified path entirely matches a regex in this set will be transformed.
+     *
+     * Some examples of common patterns include
+     * * `kotlin\.test\.assert.*` (kotlin-test)
+     * * `org\.junit\.jupiter\.api\.Assertions\.assert.*` (Junit platform)
+     * * `my\.test\.framework.+\.assert.*` (custom framework)
+     */
+    // Java's Pattern is used here instead of Kotlin's Regex for Groovy compatability
+    val functionRegexes: SetProperty<Pattern> = objectFactory.setProperty(Pattern::class.java).convention(setOf())
+
     /**
      * Defines the Kotlin SourceSets by name which will be transformed by the Power-Assert compiler plugin.
      * When the provider returns `null` - which is the default - all test SourceSets will be transformed.

libraries/tools/kotlin-power-assert/src/common/kotlin/org/jetbrains/kotlin/powerassert/gradle/PowerAssertGradlePlugin.kt

@@ -30,6 +30,7 @@ class PowerAssertGradlePlugin : KotlinCompilerPluginSupportPlugin {
         private const val POWER_ASSERT_ARTIFACT_NAME = "kotlin-power-assert-compiler-plugin-embeddable"
 
         private const val FUNCTION_ARG_NAME = "function"
+        private const val FUNCTION_REGEX_ARG_NAME = "functionRegex"
     }
 
     override fun apply(target: Project) {
@@ -52,11 +53,16 @@ class PowerAssertGradlePlugin : KotlinCompilerPluginSupportPlugin {
     ): Provider<List<SubpluginOption>> {
         val project = kotlinCompilation.target.project
         val extension = project.extensions.getByType(PowerAssertGradleExtension::class.java)
-        return extension.functions.map { functions ->
-            functions.map {
-                SubpluginOption(key = FUNCTION_ARG_NAME, value = it)
-            }
+
+        val functionOptions = extension.functions.map {
+            it.map { SubpluginOption(key = FUNCTION_ARG_NAME, value = it) }
+        }
+
+        val patternOptions = extension.functionRegexes.map {
+            it.map { SubpluginOption(key = FUNCTION_REGEX_ARG_NAME, value = "${it.flags()}:${it.pattern()}") }
         }
+
+        return functionOptions.zip(patternOptions) { a, b -> a + b }
     }
 
     override fun getCompilerPluginId(): String = "org.jetbrains.kotlin.powerassert"

plugins/power-assert/power-assert.backend/src/org/jetbrains/kotlin/powerassert/PowerAssertCallTransformer.kt

@@ -59,7 +59,7 @@ class PowerAssertCallTransformer(
     override fun visitCall(expression: IrCall): IrExpression {
         val function = expression.symbol.owner
         val fqName = function.kotlinFqName
-        if (function.parameters.isEmpty() || configuration.functions.none { fqName == it }) {
+        if (function.parameters.isEmpty() || !configuration.isTransformEnabledFor(fqName)) {
             return super.visitCall(expression)
         }
 
@@ -293,3 +293,7 @@ val IrFunction.callableId: CallableId
             CallableId(parent.kotlinFqName, name)
         }
     }
+
+private fun PowerAssertConfiguration.isTransformEnabledFor(fqName: FqName): Boolean {
+    return functions.any { fqName == it } || functionRegexes.any { it.matches(fqName.asString()) }
+}

plugins/power-assert/power-assert.backend/src/org/jetbrains/kotlin/powerassert/PowerAssertConfiguration.kt

@@ -15,6 +15,7 @@ import org.jetbrains.kotlin.name.FqName
 class PowerAssertConfiguration(
     private val configuration: CompilerConfiguration,
     val functions: Set<FqName>,
+    val functionRegexes: Set<Regex>
 ) {
     val constTracker: EvaluatedConstTracker? get() = configuration[CommonConfigurationKeys.EVALUATED_CONST_TRACKER]
     val messageCollector: MessageCollector get() = configuration.messageCollector

plugins/power-assert/power-assert.cli/src/org/jetbrains/kotlin/powerassert/PowerAssertCommandLineProcessor.kt

@@ -24,6 +24,7 @@ import org.jetbrains.kotlin.compiler.plugin.CliOption
 import org.jetbrains.kotlin.compiler.plugin.CommandLineProcessor
 import org.jetbrains.kotlin.config.CompilerConfiguration
 import org.jetbrains.kotlin.powerassert.PowerAssertPluginNames.PLUGIN_ID
+import java.util.regex.Pattern
 
 class PowerAssertCommandLineProcessor : CommandLineProcessor {
     override val pluginId: String get() = PLUGIN_ID
@@ -36,6 +37,13 @@ class PowerAssertCommandLineProcessor : CommandLineProcessor {
             required = false, // TODO required for Kotlin/JS
             allowMultipleOccurrences = true,
         ),
+        CliOption(
+            optionName = "functionRegex",
+            valueDescription = "regex matched against a function full-qualified name. Format is '\$flagsInt:\$pattern'.",
+            description = "regex matched against the fully qualified path of function to intercept",
+            required = false, // TODO required for Kotlin/JS
+            allowMultipleOccurrences = true,
+        ),
     )
 
     override fun processOption(
@@ -45,6 +53,11 @@ class PowerAssertCommandLineProcessor : CommandLineProcessor {
     ) {
         return when (option.optionName) {
             "function" -> configuration.add(KEY_FUNCTIONS, value)
+            "functionRegex" -> {
+                val flags = value.substringBefore(':', "").toIntOrNull() ?: 0
+                val pattern = value.substringAfter(':')
+                configuration.add(KEY_FUNCTION_REGEXES, Pattern.compile(pattern, flags).toRegex())
+            }
             else -> error("Unexpected config option ${option.optionName}")
         }
     }

plugins/power-assert/power-assert.cli/src/org/jetbrains/kotlin/powerassert/PowerAssertCompilerPluginRegistrar.kt

@@ -27,26 +27,30 @@ import org.jetbrains.kotlin.name.FqName
 import org.jetbrains.kotlin.powerassert.PowerAssertPluginNames.PLUGIN_ID
 
 val KEY_FUNCTIONS = CompilerConfigurationKey<List<String>>("fully-qualified function names")
+val KEY_FUNCTION_REGEXES = CompilerConfigurationKey<List<Regex>>("function fqn regexes")
 
 class PowerAssertCompilerPluginRegistrar(
     private val functions: Set<FqName>,
+    private val functionRegexs: Set<Regex>,
 ) : CompilerPluginRegistrar() {
     @Suppress("unused")
-    constructor() : this(emptySet()) // Used by service loader
+    constructor() : this(emptySet(), emptySet()) // Used by service loader
 
     override val pluginId: String get() = PLUGIN_ID
 
     override val supportsK2: Boolean = true
 
     override fun ExtensionStorage.registerExtensions(configuration: CompilerConfiguration) {
         val functions = configuration[KEY_FUNCTIONS]?.map { FqName(it) } ?: functions
-        if (functions.isEmpty()) return
+        val functionRegexs = configuration[KEY_FUNCTION_REGEXES] ?: functionRegexs
+        if (functions.isEmpty() && functionRegexs.isEmpty()) return
 
         IrGenerationExtension.registerExtension(
             PowerAssertIrGenerationExtension(
                 PowerAssertConfiguration(
                     configuration,
-                    functions.toSet()
+                    functions.toSet(),
+                    functionRegexs.toSet(),
                 )
             )
         )

plugins/power-assert/testData/codegen/regex/junitAssertionsRegex.box.txt

@@ -0,0 +1,16 @@
+assertTrue: ---
+assertTrue(booleanValue)
+           |
+           false
+ ==> expected: <true> but was: <false>---
+assertEquals: ---
+assertEquals(a, b)
+             |  |
+             |  5
+             3
+ ==> expected: <3> but was: <5>---
+assertFalse: ---
+assertFalse(booleanValue)
+            |
+            true
+ ==> expected: <false> but was: <true>---

plugins/power-assert/testData/codegen/regex/junitAssertionsRegex.kt

@@ -0,0 +1,27 @@
+// FUNCTION_REGEX: org\.junit\.jupiter\.api\.Assertions\.assert.*
+// WITH_JUNIT5
+
+import org.junit.jupiter.api.Assertions
+
+fun box(): String = runAll(
+    "assertTrue" to { test1() },
+    "assertEquals" to { test2() },
+    "assertFalse" to { test3() },
+)
+
+fun test1() {
+    val booleanValue = false
+    Assertions.assertTrue(booleanValue)
+}
+
+
+fun test2() {
+    val a = 3
+    val b = 5
+    Assertions.assertEquals(a, b)
+}
+
+fun test3() {
+    val booleanValue = true
+    Assertions.assertFalse(booleanValue)
+}

plugins/power-assert/testData/codegen/regex/kotlinTestRegex.box.txt

@@ -0,0 +1,16 @@
+assertTrue: ---
+assertTrue(booleanValue)
+           |
+           false
+---
+assertEquals: ---
+assertEquals(a, b)
+             |  |
+             |  5
+             3
+. Expected <3>, actual <5>.---
+assertFalse: ---
+assertFalse(booleanValue)
+            |
+            true
+---

plugins/power-assert/testData/codegen/regex/kotlinTestRegex.kt

@@ -0,0 +1,26 @@
+// FUNCTION_REGEX: kotlin\.test\.assert.*
+
+import kotlin.test.*
+
+fun box(): String = runAll(
+    "assertTrue" to { test1() },
+    "assertEquals" to { test2() },
+    "assertFalse" to { test3() },
+)
+
+fun test1() {
+    val booleanValue = false
+    assertTrue(booleanValue)
+}
+
+
+fun test2() {
+    val a = 3
+    val b = 5
+    assertEquals(a, b)
+}
+
+fun test3() {
+    val booleanValue = true
+    assertFalse(booleanValue)
+}

plugins/power-assert/testFixtures/org/jetbrains/kotlin/powerassert/PowerAssertConfigurationDirectives.kt

@@ -12,6 +12,10 @@ object PowerAssertConfigurationDirectives : SimpleDirectivesContainer() {
         description = "Functions targeted by Power-Assert transformation",
         multiLine = true,
     )
+    val FUNCTION_REGEX by stringDirective(
+        description = "Regexes for functions targeted by Power-Assert transformation",
+        multiLine = true,
+    )
 
     val WITH_JUNIT5 by directive("Add JUnit5 to classpath")
 }

plugins/power-assert/testFixtures/org/jetbrains/kotlin/powerassert/PowerAssertTests.kt

@@ -74,8 +74,10 @@ class PowerAssertEnvironmentConfigurator(testServices: TestServices) : Environme
         val functions = moduleStructure.allDirectives[PowerAssertConfigurationDirectives.FUNCTION]
             .ifEmpty { listOf("kotlin.assert") }
             .mapTo(mutableSetOf()) { FqName(it) }
+        val functionRegexs = moduleStructure.allDirectives[PowerAssertConfigurationDirectives.FUNCTION_REGEX]
+            .mapTo(mutableSetOf()) { Regex(it) }
 
-        IrGenerationExtension.registerExtension(PowerAssertIrGenerationExtension(PowerAssertConfiguration(configuration, functions)))
+        IrGenerationExtension.registerExtension(PowerAssertIrGenerationExtension(PowerAssertConfiguration(configuration, functions, functionRegexs)))
     }
 }