Add support for top-level Kotlin functions #847

utbot-core/src/main/kotlin/org/utbot/common/KClassUtil.kt

@@ -2,15 +2,9 @@ package org.utbot.common
 
 import java.lang.reflect.InvocationTargetException
 import java.lang.reflect.Method
-import kotlin.reflect.KClass
-
-val Class<*>.nameOfPackage: String get() = `package`?.name?:""
 
 fun Method.invokeCatching(obj: Any?, args: List<Any?>) = try {
  Result.success(invoke(obj, *args.toTypedArray()))
 } catch (e: InvocationTargetException) {
  Result.failure<Nothing>(e.targetException)
-}
-
-val KClass<*>.allNestedClasses: List<KClass<*>>
- get() = listOf(this) + nestedClasses.flatMap { it.allNestedClasses }
+}

utbot-core/src/main/kotlin/org/utbot/common/ReflectionUtil.kt

@@ -111,4 +111,10 @@ val Class<*>.isFinal
  get() = Modifier.isFinal(modifiers)
 
 val Class<*>.isProtected
- get() = Modifier.isProtected(modifiers)
+ get() = Modifier.isProtected(modifiers)
+
+val Class<*>.nameOfPackage: String
+ get() = `package`?.name?:""
+
+val Class<*>.allNestedClasses: List<Class<*>>
+ get() = listOf(this) + this.declaredClasses.flatMap { it.allNestedClasses }

utbot-framework-api/src/main/kotlin/org/utbot/framework/plugin/api/util/IdUtil.kt

@@ -24,11 +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
 
@@ -178,6 +181,14 @@ val ClassId.isDoubleType: Boolean
 val ClassId.isClassType: Boolean
  get() = this == classClassId
 
+/**
+ * Checks if the class is a Kotlin class with kind File (see [Metadata.kind] for more details)
+ */
+val ClassId.isKotlinFile: Boolean
+ get() = jClass.annotations.filterIsInstance<Metadata>().singleOrNull()?.let {
+ KotlinClassHeader.Kind.getById(it.kind) == KotlinClassHeader.Kind.FILE_FACADE
+ } ?: false
+
 val voidClassId = ClassId("void")
 val booleanClassId = ClassId("boolean")
 val byteClassId = ClassId("byte")
@@ -430,6 +441,12 @@ val MethodId.method: Method
  ?: error("Can't find method $signature in ${declaringClass.name}")
  }
 
+/**
+ * See [KCallable.extensionReceiverParameter] for more details
+ */
+val MethodId.extensionReceiverParameterIndex: Int?
+ get() = this.method.kotlinFunction?.extensionReceiverParameter?.index
+
 // TODO: maybe cache it somehow in the future
 val ConstructorId.constructor: Constructor<*>
  get() {
@@ -484,6 +501,7 @@ val Method.displayName: String
 
 val KCallable<*>.declaringClazz: Class<*>
  get() = when (this) {
+ is KFunction<*> -> javaMethod?.declaringClass?.kotlin
  is CallableReference -> owner as? KClass<*>
  else -> instanceParameter?.type?.classifier as? KClass<*>
  }?.java ?: tryConstructor(this) ?: error("Can't get parent class for $this")

utbot-framework-test/src/test/kotlin/org/utbot/examples/codegen/FileWithTopLevelFunctionsTest.kt

@@ -0,0 +1,43 @@
+package org.utbot.examples.codegen
+
+import org.junit.jupiter.api.Test
+import org.utbot.testcheckers.eq
+import org.utbot.tests.infrastructure.UtValueTestCaseChecker
+import kotlin.reflect.KFunction3
+
+@Suppress("UNCHECKED_CAST")
+internal class FileWithTopLevelFunctionsTest : UtValueTestCaseChecker(testClass = FileWithTopLevelFunctionsReflectHelper.clazz.kotlin) {
+ @Test
+ fun topLevelSumTest() {
+ check(
+ ::topLevelSum,
+ eq(1),
+ )
+ }
+
+ @Test
+ fun extensionOnBasicTypeTest() {
+ check(
+ Int::extensionOnBasicType,
+ eq(1),
+ )
+ }
+
+ @Test
+ fun extensionOnCustomClassTest() {
+ check(
+ // NB: cast is important here because we need to treat receiver as an argument to be able to check its content in matchers
+ CustomClass::extensionOnCustomClass as KFunction3<*, CustomClass, CustomClass, Boolean>,
+ eq(2),
+ { receiver, argument, result -> receiver === argument && result == true },
+ { receiver, argument, result -> receiver !== argument && result == false },
+ additionalDependencies = dependenciesForClassExtensions
+ )
+ }
+
+ companion object {
+ // Compilation of extension methods for ref objects produces call to
+ // `kotlin.jvm.internal.Intrinsics::checkNotNullParameter`, so we need to add it to dependencies
+ val dependenciesForClassExtensions = arrayOf<Class<*>>(kotlin.jvm.internal.Intrinsics::class.java)
+ }
+}

utbot-framework/src/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 => here we suppose that
+ // methods for BuiltinClasses are not extensions by default (which should be true as long as we suppose them to be java methods)
+ if (executableId.classId !is BuiltinClassId) {
+ 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)) {

utbot-framework/src/main/kotlin/org/utbot/framework/codegen/model/visitor/CgAbstractRenderer.kt

@@ -131,8 +131,13 @@ internal abstract class CgAbstractRenderer(
  }
  }
 
+ /**
+ * Returns true if one can call methods of this class without specifying a caller (for example if ClassId represents this instance)
+ */
+ protected abstract val ClassId.methodsAreAccessibleAsTopLevel: Boolean
+
  private val MethodId.accessibleByName: Boolean
- get() = (context.shouldOptimizeImports && this in context.importedStaticMethods) || classId == context.generatedClass
+ get() = (context.shouldOptimizeImports && this in context.importedStaticMethods) || classId.methodsAreAccessibleAsTopLevel
 
  override fun visit(element: CgElement) {
  val error =
@@ -654,8 +659,10 @@ internal abstract class CgAbstractRenderer(
  }
 
  override fun visit(element: CgStaticFieldAccess) {
- print(element.declaringClass.asString())
- print(".")
+ if (!element.declaringClass.methodsAreAccessibleAsTopLevel) {
+ print(element.declaringClass.asString())
+ print(".")
+ }
  print(element.fieldName)
  }
 
@@ -707,7 +714,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 {

utbot-framework/src/main/kotlin/org/utbot/framework/codegen/model/visitor/CgJavaRenderer.kt

@@ -60,6 +60,9 @@ internal class CgJavaRenderer(context: CgRendererContext, printer: CgPrinter = C
 
  override val langPackage: String = "java.lang"
 
+ override val ClassId.methodsAreAccessibleAsTopLevel: Boolean
+ get() = this == context.generatedClass
+
  override fun visit(element: AbstractCgClass<*>) {
  for (annotation in element.annotations) {
  annotation.accept(this)

utbot-framework/src/main/kotlin/org/utbot/framework/codegen/model/visitor/CgKotlinRenderer.kt

@@ -53,6 +53,7 @@ import org.utbot.framework.plugin.api.TypeParameters
 import org.utbot.framework.plugin.api.WildcardTypeParameter
 import org.utbot.framework.plugin.api.util.id
 import org.utbot.framework.plugin.api.util.isArray
+import org.utbot.framework.plugin.api.util.isKotlinFile
 import org.utbot.framework.plugin.api.util.isPrimitive
 import org.utbot.framework.plugin.api.util.isPrimitiveWrapper
 import org.utbot.framework.plugin.api.util.kClass
@@ -72,6 +73,10 @@ internal class CgKotlinRenderer(context: CgRendererContext, printer: CgPrinter =
 
  override val langPackage: String = "kotlin"
 
+ override val ClassId.methodsAreAccessibleAsTopLevel: Boolean
+ // NB: the order of operands is important as `isKotlinFile` uses reflection and thus can't be called on context.generatedClass
+ get() = (this == context.generatedClass) || isKotlinFile
+
  override fun visit(element: AbstractCgClass<*>) {
  for (annotation in element.annotations) {
  annotation.accept(this)

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?>) {
 

utbot-framework/src/main/kotlin/org/utbot/framework/process/EngineMain.kt

@@ -21,6 +21,7 @@ import org.utbot.framework.plugin.api.util.UtContext
 import org.utbot.framework.plugin.api.util.executableId
 import org.utbot.framework.plugin.api.util.id
 import org.utbot.framework.plugin.api.util.jClass
+import org.utbot.framework.plugin.api.util.method
 import org.utbot.framework.plugin.services.JdkInfo
 import org.utbot.framework.process.generated.*
 import org.utbot.framework.util.ConflictTriggers
@@ -36,7 +37,7 @@ import org.utbot.summary.summarize
 import java.io.File
 import java.net.URLClassLoader
 import java.nio.file.Paths
-import kotlin.reflect.full.functions
+import kotlin.reflect.jvm.kotlinFunction
 import kotlin.time.Duration.Companion.seconds
 
 private val messageFromMainTimeoutMillis = 120.seconds
@@ -158,8 +159,8 @@ private fun EngineProcessModel.setup(
  synchronizer.measureExecutionForTermination(findMethodsInClassMatchingSelected) { params ->
  val classId = kryoHelper.readObject<ClassId>(params.classId)
  val selectedSignatures = params.signatures.map { Signature(it.name, it.parametersTypes) }
- FindMethodsInClassMatchingSelectedResult(kryoHelper.writeObject(classId.jClass.kotlin.allNestedClasses.flatMap { clazz ->
- clazz.functions.sortedWith(compareBy { selectedSignatures.indexOf(it.signature()) })
+ FindMethodsInClassMatchingSelectedResult(kryoHelper.writeObject(classId.jClass.allNestedClasses.flatMap { clazz ->
+ clazz.id.allMethods.mapNotNull { it.method.kotlinFunction }.sortedWith(compareBy { selectedSignatures.indexOf(it.signature()) })
  .filter { it.signature().normalized() in selectedSignatures }
  .map { it.executableId }
  }))
@@ -168,7 +169,7 @@ private fun EngineProcessModel.setup(
  val classId = kryoHelper.readObject<ClassId>(params.classId)
  val bySignature = kryoHelper.readObject<Map<Signature, List<String>>>(params.bySignature)
  FindMethodParamNamesResult(kryoHelper.writeObject(
- classId.jClass.kotlin.allNestedClasses.flatMap { it.functions }
+ classId.jClass.allNestedClasses.flatMap { clazz -> clazz.id.allMethods.mapNotNull { it.method.kotlinFunction } }
  .mapNotNull { method -> bySignature[method.signature()]?.let { params -> method.executableId to params } }
  .toMap()
  ))

utbot-intellij/src/main/kotlin/org/utbot/intellij/plugin/process/EngineProcess.kt

@@ -238,7 +238,16 @@ class EngineProcess(parent: Lifetime, val project: Project) {
  }
 
  private fun MemberInfo.paramNames(): List<String> =
- (this.member as PsiMethod).parameterList.parameters.map { it.name }
+ (this.member as PsiMethod).parameterList.parameters.map {
+ if (it.name.startsWith("\$this"))
+ // If member is Kotlin extension function, name of first argument isn't good for further usage,
+ // so we better choose name based on type of receiver.
+ //
+ // There seems no API to check whether parameter is an extension receiver by PSI
+ it.type.presentableText
+ else
+ it.name
+ }
 
  fun generate(
  mockInstalled: Boolean,

utbot-intellij/src/main/kotlin/org/utbot/intellij/plugin/ui/actions/GenerateTestsAction.kt

@@ -18,6 +18,7 @@ import com.intellij.openapi.vfs.VirtualFile
 import com.intellij.psi.*
 import com.intellij.psi.util.PsiTreeUtil
 import com.intellij.refactoring.util.classMembers.MemberInfo
+import org.jetbrains.kotlin.asJava.findFacadeClass
 import org.jetbrains.kotlin.idea.core.getPackage
 import org.jetbrains.kotlin.idea.core.util.toPsiDirectory
 import org.jetbrains.kotlin.idea.core.util.toPsiFile
@@ -26,6 +27,7 @@ import org.utbot.intellij.plugin.util.extractFirstLevelMembers
 import org.utbot.intellij.plugin.util.isVisible
 import java.util.*
 import org.jetbrains.kotlin.j2k.getContainingClass
+import org.jetbrains.kotlin.psi.KtFile
 import org.jetbrains.kotlin.utils.addIfNotNull
 import org.utbot.framework.plugin.api.util.LockFile
 import org.utbot.intellij.plugin.models.packageName
@@ -218,7 +220,7 @@ class GenerateTestsAction : AnAction(), UpdateInBackground {
  }
 
  private fun getAllClasses(directory: PsiDirectory): Set<PsiClass> {
- val allClasses = directory.files.flatMap { getClassesFromFile(it) }.toMutableSet()
+ val allClasses = directory.files.flatMap { PsiElementHandler.makePsiElementHandler(it).getClassesFromFile(it) }.toMutableSet()
  for (subDir in directory.subdirectories) allClasses += getAllClasses(subDir)
  return allClasses
  }
@@ -231,15 +233,10 @@ class GenerateTestsAction : AnAction(), UpdateInBackground {
  if (!dirsArePackages) {
  return emptySet()
  }
- val allClasses = psiFiles.flatMap { getClassesFromFile(it) }.toMutableSet()
+ val allClasses = psiFiles.flatMap { PsiElementHandler.makePsiElementHandler(it).getClassesFromFile(it) }.toMutableSet()
+ allClasses.addAll(psiFiles.mapNotNull { (it as? KtFile)?.findFacadeClass() })
  for (psiDir in psiDirectories) allClasses += getAllClasses(psiDir)
 
  return allClasses
  }
-
- private fun getClassesFromFile(psiFile: PsiFile): List<PsiClass> {
- val psiElementHandler = PsiElementHandler.makePsiElementHandler(psiFile)
- return PsiTreeUtil.getChildrenOfTypeAsList(psiFile, psiElementHandler.classClass)
- .map { psiElementHandler.toPsi(it, PsiClass::class.java) }
- }
 }

utbot-intellij/src/main/kotlin/org/utbot/intellij/plugin/ui/utils/KotlinPsiElementHandler.kt

@@ -2,6 +2,9 @@ package org.utbot.intellij.plugin.ui.utils
 
 import com.intellij.psi.PsiClass
 import com.intellij.psi.PsiElement
+import com.intellij.psi.PsiFile
+import com.intellij.psi.util.findParentOfType
+import org.jetbrains.kotlin.asJava.findFacadeClass
 import org.jetbrains.kotlin.idea.testIntegration.KotlinCreateTestIntention
 import org.jetbrains.kotlin.psi.KtClass
 import org.jetbrains.kotlin.psi.KtClassOrObject
@@ -24,13 +27,27 @@ class KotlinPsiElementHandler(
  return element.toUElement()?.javaPsi as? T ?: error("Could not cast $element to $clazz")
  }
 
- override fun isCreateTestActionAvailable(element: PsiElement): Boolean =
- getTarget(element)?.let { KotlinCreateTestIntention().applicabilityRange(it) != null } ?: false
+ override fun getClassesFromFile(psiFile: PsiFile): List<PsiClass> {
+ return listOfNotNull((psiFile as? KtFile)?.findFacadeClass()) + super.getClassesFromFile(psiFile)
+ }
+
+ override fun isCreateTestActionAvailable(element: PsiElement): Boolean {
+ getTarget(element)?.let {
+ return KotlinCreateTestIntention().applicabilityRange(it) != null
+ }
+ return (element.containingFile as? KtFile)?.findFacadeClass() != null
+ }
 
  private fun getTarget(element: PsiElement?): KtNamedDeclaration? =
  element?.parentsWithSelf
  ?.firstOrNull { it is KtClassOrObject || it is KtNamedDeclaration && it.parent is KtFile } as? KtNamedDeclaration
 
- override fun containingClass(element: PsiElement): PsiClass? =
- element.parentsWithSelf.firstOrNull { it is KtClassOrObject }?.let { toPsi(it, PsiClass::class.java) }
+ override fun containingClass(element: PsiElement): PsiClass? {
+ element.findParentOfType<KtClassOrObject>(strict=false)?.let {
+ return toPsi(it, PsiClass::class.java)
+ }
+ return element.findParentOfType<KtFile>(strict=false)?.findFacadeClass()?.let {
+ toPsi(it, PsiClass::class.java)
+ }
+ }
 }