Fix codegen for kotlin extension functions

UnitTestBot · volivan239 · Oct 28, 2022 · Oct 7, 2022 · Oct 10, 2022 · Oct 14, 2022
commit 81880d61969eea4303811fd27928c9241a15e770
diff --git a/utbot-framework-api/src/main/kotlin/org/utbot/framework/plugin/api/util/IdUtil.kt b/utbot-framework-api/src/main/kotlin/org/utbot/framework/plugin/api/util/IdUtil.kt
@@ -24,12 +24,14 @@ import kotlin.reflect.KCallable
 import kotlin.reflect.KClass
 import kotlin.reflect.KFunction
 import kotlin.reflect.KProperty
+import kotlin.reflect.full.extensionReceiverParameter
 import kotlin.reflect.full.instanceParameter
 import kotlin.reflect.jvm.internal.impl.load.kotlin.header.KotlinClassHeader
 import kotlin.reflect.jvm.javaConstructor
 import kotlin.reflect.jvm.javaField
 import kotlin.reflect.jvm.javaGetter
 import kotlin.reflect.jvm.javaMethod
+import kotlin.reflect.jvm.kotlinFunction
 
 // ClassId utils
 
@@ -439,6 +441,9 @@ val MethodId.method: Method
                 ?: error("Can't find method $signature in ${declaringClass.name}")
     }
 
+val MethodId.extensionReceiverParameterIndex: Int?
+    get() = this.method.kotlinFunction?.extensionReceiverParameter?.index
+
 // TODO: maybe cache it somehow in the future
 val ConstructorId.constructor: Constructor<*>
     get() {

diff --git a/...main/kotlin/org/utbot/framework/codegen/model/constructor/tree/CgCallableAccessManager.kt b/...main/kotlin/org/utbot/framework/codegen/model/constructor/tree/CgCallableAccessManager.kt
@@ -49,6 +49,8 @@ import org.utbot.framework.plugin.api.MethodId
 import org.utbot.framework.plugin.api.UtExplicitlyThrownException
 import org.utbot.framework.plugin.api.util.isStatic
 import org.utbot.framework.plugin.api.util.exceptions
+import org.utbot.framework.plugin.api.util.extensionReceiverParameterIndex
+import org.utbot.framework.plugin.api.util.humanReadableName
 import org.utbot.framework.plugin.api.util.id
 import org.utbot.framework.plugin.api.util.isArray
 import org.utbot.framework.plugin.api.util.isSubtypeOf
@@ -110,7 +112,7 @@ internal class CgCallableAccessManagerImpl(val context: CgContext) : CgCallableA
     override operator fun CgIncompleteMethodCall.invoke(vararg args: Any?): CgMethodCall {
         val resolvedArgs = args.resolve()
         val methodCall = if (method.canBeCalledWith(caller, resolvedArgs)) {
-            CgMethodCall(caller, method, resolvedArgs.guardedForDirectCallOf(method))
+            CgMethodCall(caller, method, resolvedArgs.guardedForDirectCallOf(method)).takeCallerFromArgumentsIfNeeded()
         } else {
             method.callWithReflection(caller, resolvedArgs)
         }
@@ -194,6 +196,29 @@ internal class CgCallableAccessManagerImpl(val context: CgContext) : CgCallableA
             else -> false
         }
 
+    /**
+     * For Kotlin extension functions, real caller is one of the arguments in JVM method (and declaration class is omitted),
+     * thus we should move it from arguments to caller
+     *
+     * For example, if we have `Int.f(a: Int)` declared in `Main.kt`, the JVM method signature will be `MainKt.f(Int, Int)`
+     * and in Kotlin we should render this not like `MainKt.f(a, b)` but like `a.f(b)`
+     */
+    private fun CgMethodCall.takeCallerFromArgumentsIfNeeded(): CgMethodCall {
+        if (codegenLanguage == CodegenLanguage.KOTLIN) {
+            // TODO: reflection calls for util and some of mockito methods produce exceptions => currently runCatching is needed
+            //  (but their reflection may be supported, alternatively maybe get rid of reflection somehow here)
+            runCatching {
+                executableId.extensionReceiverParameterIndex?.let { receiverIndex ->
+                    require(caller == null) { "${executableId.humanReadableName} is an extension function but it already has a non-static caller provided" }
+                    val args = arguments.toMutableList()
+                    return CgMethodCall(args.removeAt(receiverIndex), executableId, args, typeParameters)
+                }
+            }
+        }
+
+        return this
+    }
+
     private infix fun CgExpression.canBeArgOf(type: ClassId): Boolean {
         // TODO: SAT-1210 support generics so that we wouldn't need to check specific cases such as this one
         if (this is CgExecutableCall && (executableId == any || executableId == anyOfClass)) {

diff --git a/...framework/src/main/kotlin/org/utbot/framework/codegen/model/visitor/CgAbstractRenderer.kt b/...framework/src/main/kotlin/org/utbot/framework/codegen/model/visitor/CgAbstractRenderer.kt
@@ -711,7 +711,10 @@ internal abstract class CgAbstractRenderer(
         if (caller != null) {
             // 'this' can be omitted, otherwise render caller
             if (caller !is CgThisInstance) {
+                // TODO: we need parentheses for calls like (-1).inv(), do something smarter here
+                if (caller !is CgVariable) print("(")
                 caller.accept(this)
+                if (caller !is CgVariable) print(")")
                 renderAccess(caller)
             }
         } else {

diff --git a/utbot-framework/src/main/kotlin/org/utbot/framework/plugin/api/SignatureUtil.kt b/utbot-framework/src/main/kotlin/org/utbot/framework/plugin/api/SignatureUtil.kt
@@ -4,8 +4,9 @@ import kotlin.reflect.KFunction
 import kotlin.reflect.KParameter
 import kotlin.reflect.jvm.javaType
 
+// Note that rules for obtaining signature here should correlate with PsiMethod.signature()
 fun KFunction<*>.signature() =
-    Signature(this.name, this.parameters.filter { it.kind == KParameter.Kind.VALUE }.map { it.type.javaType.typeName })
+    Signature(this.name, this.parameters.filter { it.kind != KParameter.Kind.INSTANCE }.map { it.type.javaType.typeName })
 
 data class Signature(val name: String, val parameterTypes: List<String?>) {
 

diff --git a/utbot-intellij/src/main/kotlin/org/utbot/intellij/plugin/util/SignaturesHelper.kt b/utbot-intellij/src/main/kotlin/org/utbot/intellij/plugin/util/SignaturesHelper.kt
@@ -7,6 +7,7 @@ import org.utbot.framework.plugin.api.Signature
 fun MemberInfo.signature(): Signature =
         (this.member as PsiMethod).signature()
 
+// Note that rules for obtaining signature here should correlate with KFunction<*>.signature()
 private fun PsiMethod.signature() =
     Signature(this.name, this.parameterList.parameters.map {
         it.type.canonicalText