From 0bb3a4ccb717154c73b04e648bb3e61b7de38529 Mon Sep 17 00:00:00 2001 From: Nicolay Mitropolsky Date: Thu, 14 May 2020 19:46:25 +0300 Subject: [PATCH] 201: Uast: KotlinUastCodeGenerationPlugin on the level needed for ReactorSupport (KT-38062) --- .../KotlinUastCodeGenerationPlugin.kt.201 | 322 +++++- .../expressions/KotlinULiteralExpression.kt | 20 +- .../expressions/KotlinUReturnExpression.kt | 8 +- .../kotlin/internal/KotlinFakeUElement.kt | 18 + .../tests/KotlinUastGenerationTest.kt.201 | 938 ++++++++++++++++++ 5 files changed, 1273 insertions(+), 33 deletions(-) create mode 100644 plugins/uast-kotlin/src/org/jetbrains/uast/kotlin/internal/KotlinFakeUElement.kt create mode 100644 plugins/uast-kotlin/tests/KotlinUastGenerationTest.kt.201 diff --git a/plugins/uast-kotlin-idea/src/org/jetbrains/uast/kotlin/generate/KotlinUastCodeGenerationPlugin.kt.201 b/plugins/uast-kotlin-idea/src/org/jetbrains/uast/kotlin/generate/KotlinUastCodeGenerationPlugin.kt.201 index c9ca537310a..58b2c58d163 100644 --- a/plugins/uast-kotlin-idea/src/org/jetbrains/uast/kotlin/generate/KotlinUastCodeGenerationPlugin.kt.201 +++ b/plugins/uast-kotlin-idea/src/org/jetbrains/uast/kotlin/generate/KotlinUastCodeGenerationPlugin.kt.201 @@ -6,15 +6,25 @@ package org.jetbrains.uast.kotlin.generate import com.intellij.lang.Language +import com.intellij.openapi.diagnostic.logger import com.intellij.openapi.project.Project import com.intellij.openapi.util.text.StringUtil +import com.intellij.psi.PsiClassType import com.intellij.psi.PsiElement import com.intellij.psi.PsiType +import com.intellij.psi.util.PsiTreeUtil import org.jetbrains.kotlin.idea.KotlinLanguage +import org.jetbrains.kotlin.idea.caches.resolve.getResolutionFacade +import org.jetbrains.kotlin.idea.core.KotlinNameSuggester import org.jetbrains.kotlin.idea.core.ShortenReferences +import org.jetbrains.kotlin.idea.refactoring.fqName.fqName +import org.jetbrains.kotlin.idea.util.getResolutionScope +import org.jetbrains.kotlin.idea.util.resolveToKotlinType +import org.jetbrains.kotlin.incremental.components.NoLookupLocation +import org.jetbrains.kotlin.name.Name import org.jetbrains.kotlin.psi.* -import org.jetbrains.kotlin.psi.psiUtil.getPossiblyQualifiedCallExpression -import org.jetbrains.kotlin.psi.psiUtil.quoteIfNeeded +import org.jetbrains.kotlin.psi.psiUtil.* +import org.jetbrains.kotlin.resolve.scopes.utils.findClassifier import org.jetbrains.kotlin.utils.addToStdlib.cast import org.jetbrains.kotlin.utils.addToStdlib.safeAs import org.jetbrains.uast.* @@ -22,6 +32,8 @@ import org.jetbrains.uast.generate.UParameterInfo import org.jetbrains.uast.generate.UastCodeGenerationPlugin import org.jetbrains.uast.generate.UastElementFactory import org.jetbrains.uast.kotlin.* +import org.jetbrains.uast.kotlin.internal.KotlinFakeUElement +import org.jetbrains.uast.kotlin.internal.toSourcePsiFakeAware class KotlinUastCodeGenerationPlugin : UastCodeGenerationPlugin { override val language: Language @@ -31,7 +43,7 @@ class KotlinUastCodeGenerationPlugin : UastCodeGenerationPlugin { KotlinUastElementFactory(project) override fun replace(oldElement: UElement, newElement: T, elementType: Class): T? { - val oldPsi = oldElement.sourcePsi ?: return null + val oldPsi = oldElement.toSourcePsiFakeAware().singleOrNull() ?: return null val newPsi = newElement.sourcePsi?.let { when { it is KtCallExpression && it.parent is KtQualifiedExpression -> it.parent @@ -42,9 +54,23 @@ class KotlinUastCodeGenerationPlugin : UastCodeGenerationPlugin { val psiFactory = KtPsiFactory(oldPsi.project) val oldParentPsi = oldPsi.parent val (updOldPsi, updNewPsi) = when { - oldParentPsi is KtStringTemplateExpression && oldParentPsi.entries.size == 1 -> oldParentPsi to newPsi - oldPsi is KtStringTemplateEntry && newPsi !is KtStringTemplateEntry && newPsi is KtExpression -> oldPsi to psiFactory.createBlockStringTemplateEntry(newPsi) - else -> oldPsi to newPsi + oldParentPsi is KtStringTemplateExpression && oldParentPsi.entries.size == 1 -> + oldParentPsi to newPsi + + oldPsi is KtStringTemplateEntry && newPsi !is KtStringTemplateEntry && newPsi is KtExpression -> + oldPsi to psiFactory.createBlockStringTemplateEntry(newPsi) + + oldPsi is KtBlockExpression && newPsi is KtBlockExpression -> { + if (!hasBraces(oldPsi) && hasBraces(newPsi)) { + oldPsi to psiFactory.createLambdaExpression("none", newPsi.statements.joinToString("\n") { "println()" }).bodyExpression!!.also { + it.statements.zip(newPsi.statements).forEach { it.first.replace(it.second) } + } + } else + oldPsi to newPsi + } + + else -> + oldPsi to newPsi } return when (val replaced = updOldPsi.replace(updNewPsi)?.safeAs()?.let { ShortenReferences.DEFAULT.process(it) }) { @@ -54,10 +80,18 @@ class KotlinUastCodeGenerationPlugin : UastCodeGenerationPlugin { } } +private fun hasBraces(oldPsi: KtBlockExpression): Boolean = oldPsi.lBrace != null && oldPsi.rBrace != null + class KotlinUastElementFactory(project: Project) : UastElementFactory { private val psiFactory = KtPsiFactory(project) + @Deprecated("use version with context parameter") override fun createQualifiedReference(qualifiedName: String, context: UElement?): UQualifiedReferenceExpression? { + logger().error("Please switch caller to the version with a context parameter") + return createQualifiedReference(qualifiedName, context?.sourcePsi) + } + + /*override*/ fun createQualifiedReference(qualifiedName: String, context: PsiElement?): UQualifiedReferenceExpression? { return psiFactory.createExpression(qualifiedName).let { when (it) { is KtDotQualifiedExpression -> KotlinUQualifiedReferenceExpression(it, null) @@ -68,21 +102,43 @@ class KotlinUastElementFactory(project: Project) : UastElementFactory { } } - override fun createCallExpression(receiver: UExpression?, methodName: String, parameters: List, expectedReturnType: PsiType?, kind: UastCallKind, context: PsiElement?): UCallExpression? { + override fun createCallExpression( + receiver: UExpression?, + methodName: String, + parameters: List, + expectedReturnType: PsiType?, + kind: UastCallKind, + context: PsiElement? + ): UCallExpression? { if (kind != UastCallKind.METHOD_CALL) return null + val typeParams = (context as? KtElement)?.let { kontext -> + val resolutionFacade = kontext.getResolutionFacade() + (expectedReturnType as? PsiClassType)?.parameters?.map { it.resolveToKotlinType(resolutionFacade) } + } + val name = methodName.quoteIfNeeded() val methodCall = psiFactory.createExpression( - if (receiver != null) "a.$name()" else "$name()" + buildString { + if (receiver != null) + append("a.") + append(name) + if (typeParams != null) { + append(typeParams.joinToString(", ", "<", ">") { type -> + type.fqName?.asString() ?: "" + }) + } + append("()") + } ).getPossiblyQualifiedCallExpression() ?: return null if (receiver != null) { - methodCall.parent.safeAs()?.receiverExpression?.replace(receiver.sourcePsi!!) + methodCall.parent.safeAs()?.receiverExpression?.replace(wrapULiteral(receiver).sourcePsi!!) } val valueArgumentList = methodCall.valueArgumentList for (parameter in parameters) { - valueArgumentList?.addArgument(psiFactory.createArgument(parameter.sourcePsi as? KtExpression)) + valueArgumentList?.addArgument(psiFactory.createArgument(wrapULiteral(parameter).sourcePsi as KtExpression)) } return KotlinUFunctionCallExpression(methodCall, null) @@ -93,65 +149,281 @@ class KotlinUastElementFactory(project: Project) : UastElementFactory { return KotlinStringULiteralExpression(psiFactory.createExpression(StringUtil.wrapWithDoubleQuote(text)), null) } + /*override*/ fun createNullLiteral(context: PsiElement?): ULiteralExpression { + return psiFactory.createExpression("null").toUElementOfType()!! + } + + /*override*/ fun createIntLiteral(value: Int, context: PsiElement?): ULiteralExpression { + return psiFactory.createExpression(value.toString()).toUElementOfType()!! + } + + private fun KtExpression.ensureBlockExpressionBraces(): KtExpression { + if (this !is KtBlockExpression || hasBraces(this)) return this + val blockExpression = psiFactory.createBlock(this.statements.joinToString("\n") { "println()" }) + for ((placeholder, statement) in blockExpression.statements.zip(this.statements)) { + placeholder.replace(statement) + } + return blockExpression + } + + @Deprecated("use version with context parameter") override fun createIfExpression(condition: UExpression, thenBranch: UExpression, elseBranch: UExpression?): UIfExpression? { + logger().error("Please switch caller to the version with a context parameter") + return createIfExpression(condition, thenBranch, elseBranch, null) + } + + /*override*/ fun createIfExpression( + condition: UExpression, + thenBranch: UExpression, + elseBranch: UExpression?, + context: PsiElement? + ): UIfExpression? { val conditionPsi = condition.sourcePsi as? KtExpression ?: return null val thenBranchPsi = thenBranch.sourcePsi as? KtExpression ?: return null val elseBranchPsi = elseBranch?.sourcePsi as? KtExpression - return KotlinUIfExpression(psiFactory.createIf(conditionPsi, thenBranchPsi, elseBranchPsi), null) + return KotlinUIfExpression(psiFactory.createIf(conditionPsi, thenBranchPsi.ensureBlockExpressionBraces(), elseBranchPsi?.ensureBlockExpressionBraces()), null) } + @Deprecated("use version with context parameter") override fun createParenthesizedExpression(expression: UExpression): UParenthesizedExpression? { + logger().error("Please switch caller to the version with a context parameter") + return createParenthesizedExpression(expression, null) + } + + /*override*/ fun createParenthesizedExpression(expression: UExpression, context: PsiElement?): UParenthesizedExpression? { val source = expression.sourcePsi ?: return null val parenthesized = psiFactory.createExpression("(${source.text})") as? KtParenthesizedExpression ?: return null return KotlinUParenthesizedExpression(parenthesized, null) } + @Deprecated("use version with context parameter") override fun createSimpleReference(name: String): USimpleNameReferenceExpression? { + logger().error("Please switch caller to the version with a context parameter") + return createSimpleReference(name, null) + } + + /*override*/ fun createSimpleReference(name: String, context: PsiElement?): USimpleNameReferenceExpression? { return KotlinUSimpleReferenceExpression(psiFactory.createSimpleName(name), null) } + @Deprecated("use version with context parameter") override fun createSimpleReference(variable: UVariable): USimpleNameReferenceExpression? { - return createSimpleReference(variable.name ?: return null) + logger().error("Please switch caller to the version with a context parameter") + return createSimpleReference(variable, null) } + /*override*/ fun createSimpleReference(variable: UVariable, context: PsiElement?): USimpleNameReferenceExpression? { + return createSimpleReference(variable.name ?: return null, context) + } + + @Deprecated("use version with context parameter") override fun createReturnExpresion(expression: UExpression?, inLambda: Boolean): UReturnExpression? { - val returnExpression = psiFactory.createExpression("return") as KtReturnExpression - expression?.sourcePsi?.let { returnExpression.add(it) } + logger().error("Please switch caller to the version with a context parameter") + return createReturnExpresion(expression, inLambda, null) + } + + /*override*/ fun createReturnExpresion(expression: UExpression?, inLambda: Boolean, context: PsiElement?): UReturnExpression? { + val label = if (inLambda && context != null) getParentLambdaLabelName(context)?.let { "@$it" } ?: "" else "" + val returnExpression = psiFactory.createExpression("return$label 1") as KtReturnExpression + val sourcePsi = expression?.sourcePsi + if (sourcePsi != null) { + returnExpression.returnedExpression!!.replace(sourcePsi) + } else { + returnExpression.returnedExpression?.delete() + } return KotlinUReturnExpression(returnExpression, null) } - override fun createBinaryExpression(leftOperand: UExpression, rightOperand: UExpression, operator: UastBinaryOperator): UBinaryExpression? { + private fun getParentLambdaLabelName(context: PsiElement): String? { + val lambdaExpression = context.getParentOfType(false) ?: return null + lambdaExpression.parent.safeAs()?.let { return it.getLabelName() } + val callExpression = lambdaExpression.getStrictParentOfType() ?: return null + callExpression.valueArguments.find { + it.getArgumentExpression()?.unpackFunctionLiteral(allowParentheses = false) === lambdaExpression + } ?: return null + return callExpression.getCallNameExpression()?.text + } + + @Deprecated("use version with context parameter") + override fun createBinaryExpression( + leftOperand: UExpression, + rightOperand: UExpression, + operator: UastBinaryOperator + ): UBinaryExpression? { + logger().error("Please switch caller to the version with a context parameter") + return createBinaryExpression(leftOperand, rightOperand, operator, null) + } + + /*override*/ fun createBinaryExpression( + leftOperand: UExpression, + rightOperand: UExpression, + operator: UastBinaryOperator, + context: PsiElement? + ): UBinaryExpression? { + val binaryExpression = joinBinaryExpression(leftOperand, rightOperand, operator) ?: return null + return KotlinUBinaryExpression(binaryExpression, null) + } + + private fun joinBinaryExpression( + leftOperand: UExpression, + rightOperand: UExpression, + operator: UastBinaryOperator + ): KtBinaryExpression? { val leftPsi = leftOperand.sourcePsi ?: return null val rightPsi = rightOperand.sourcePsi ?: return null val binaryExpression = psiFactory.createExpression("a ${operator.text} b") as? KtBinaryExpression ?: return null binaryExpression.left?.replace(leftPsi) binaryExpression.right?.replace(rightPsi) - return KotlinUBinaryExpression(binaryExpression, null) + return binaryExpression } - override fun createFlatBinaryExpression(leftOperand: UExpression, rightOperand: UExpression, operator: UastBinaryOperator): UPolyadicExpression? { - return createBinaryExpression(leftOperand, rightOperand, operator) + @Deprecated("use version with context parameter") + override fun createFlatBinaryExpression( + leftOperand: UExpression, + rightOperand: UExpression, + operator: UastBinaryOperator + ): UPolyadicExpression? { + logger().error("Please switch caller to the version with a context parameter") + return createFlatBinaryExpression(leftOperand, rightOperand, operator, null) } + /*override*/ fun createFlatBinaryExpression( + leftOperand: UExpression, + rightOperand: UExpression, + operator: UastBinaryOperator, + context: PsiElement? + ): UPolyadicExpression? { + + fun unwrapParentheses(exp: KtExpression?) { + if (exp !is KtParenthesizedExpression) return + if (!KtPsiUtil.areParenthesesUseless(exp)) return + exp.expression?.let { exp.replace(it) } + } + + val binaryExpression = joinBinaryExpression(leftOperand, rightOperand, operator) ?: return null + unwrapParentheses(binaryExpression.left) + unwrapParentheses(binaryExpression.right) + + return psiFactory.createExpression(binaryExpression.text).toUElementOfType()!! + } + + @Deprecated("use version with context parameter") override fun createBlockExpression(expressions: List): UBlockExpression? { - if (expressions.any { it.sourcePsi == null}) return null - val block = psiFactory.createBlock(expressions.joinToString(separator = "\n") { it.sourcePsi?.text ?: "" }) + logger().error("Please switch caller to the version with a context parameter") + return createBlockExpression(expressions, null) + } + + /*override*/ fun createBlockExpression(expressions: List, context: PsiElement?): UBlockExpression? { + val sourceExpressions = expressions.flatMap { it.toSourcePsiFakeAware() } + val block = psiFactory.createBlock( + sourceExpressions.joinToString(separator = "\n") { "println()" } + ) + for ((placeholder, psiElement) in block.statements.zip(sourceExpressions)) { + placeholder.replace(psiElement) + } return KotlinUBlockExpression(block, null) } + @Deprecated("use version with context parameter") override fun createDeclarationExpression(declarations: List): UDeclarationsExpression? { - return object : KotlinUDeclarationsExpression(null) { + logger().error("Please switch caller to the version with a context parameter") + return createDeclarationExpression(declarations, null) + } + + /*override*/ fun createDeclarationExpression(declarations: List, context: PsiElement?): UDeclarationsExpression? { + return object : KotlinUDeclarationsExpression(null), KotlinFakeUElement { override var declarations: List = declarations + override fun unwrapToSourcePsi(): List = declarations.flatMap { it.toSourcePsiFakeAware() } } } - override fun createLambdaExpression(parameters: List, body: UExpression): ULambdaExpression? { - TODO("Not yet implemented") + /*override*/ fun createLambdaExpression( + parameters: List, + body: UExpression, + context: PsiElement? + ): ULambdaExpression? { + val resolutionFacade = (context as? KtElement)?.getResolutionFacade() + val validator = (context as? KtElement)?.let { usedNamesFilter(it) } ?: { true } + + val newLambdaStatements = if (body is UBlockExpression) { + body.expressions.flatMap { member -> + when { + member is UReturnExpression -> member.returnExpression?.toSourcePsiFakeAware().orEmpty() + else -> member.toSourcePsiFakeAware() + } + } + } else + listOf(body.sourcePsi!!) + + val ktLambdaExpression = psiFactory.createLambdaExpression( + parameters.joinToString(", ") { p -> + val ktype = resolutionFacade?.let { p.type?.resolveToKotlinType(it) } + StringBuilder().apply { + append(p.suggestedName ?: ktype?.let { KotlinNameSuggester.suggestNamesByType(it, validator).firstOrNull() }) + ?: KotlinNameSuggester.suggestNameByName("v", validator) + ktype?.fqName?.toString()?.let { append(": ").append(it) } + } + }, + newLambdaStatements.joinToString("\n") { "placeholder" } + ) + + for ((old, new) in ktLambdaExpression.bodyExpression!!.statements.zip(newLambdaStatements)) { + old.replace(new) + } + + return ktLambdaExpression.toUElementOfType()!! } - override fun createLocalVariable(suggestedName: String?, type: PsiType?, initializer: UExpression, immutable: Boolean): ULocalVariable? { - TODO("Not yet implemented") + @Deprecated("use version with context parameter") + override fun createLambdaExpression(parameters: List, body: UExpression): ULambdaExpression? { + logger().error("Please switch caller to the version with a context parameter") + return createLambdaExpression(parameters, body, null) } + + /*override*/ fun createLocalVariable( + suggestedName: String?, + type: PsiType?, + initializer: UExpression, + immutable: Boolean, + context: PsiElement? + ): ULocalVariable? { + val resolutionFacade = (context as? KtElement)?.getResolutionFacade() + val validator = (context as? KtElement)?.let { usedNamesFilter(it) } ?: { true } + val ktype = resolutionFacade?.let { type?.resolveToKotlinType(it) } + + val function = psiFactory.createFunction( + buildString { + append("fun foo() { ") + append(if (immutable) "val" else "var") + append(" ") + append(suggestedName ?: ktype?.let { KotlinNameSuggester.suggestNamesByType(it, validator).firstOrNull() }) + ?: KotlinNameSuggester.suggestNameByName("v", validator) + ktype?.fqName?.toString()?.let { append(": ").append(it) } + append(" = null") + append("}") + } + ) + + val ktVariable = PsiTreeUtil.findChildOfType(function, KtVariableDeclaration::class.java)!! + val newVariable = ktVariable.initializer!!.replace(initializer.sourcePsi!!).parent + return newVariable.toUElementOfType() as ULocalVariable + } + + @Deprecated("use version with context parameter") + override fun createLocalVariable( + suggestedName: String?, + type: PsiType?, + initializer: UExpression, + immutable: Boolean + ): ULocalVariable? { + logger().error("Please switch caller to the version with a context parameter") + return createLocalVariable(suggestedName, type, initializer, immutable, null) + } +} + +private fun usedNamesFilter(context: KtElement): (String) -> Boolean { + val scope = context.getResolutionScope() + return { name: String -> scope.findClassifier(Name.identifier(name), NoLookupLocation.FROM_IDE) == null } } \ No newline at end of file diff --git a/plugins/uast-kotlin/src/org/jetbrains/uast/kotlin/expressions/KotlinULiteralExpression.kt b/plugins/uast-kotlin/src/org/jetbrains/uast/kotlin/expressions/KotlinULiteralExpression.kt index 4c71b4b9f29..710a899425f 100644 --- a/plugins/uast-kotlin/src/org/jetbrains/uast/kotlin/expressions/KotlinULiteralExpression.kt +++ b/plugins/uast-kotlin/src/org/jetbrains/uast/kotlin/expressions/KotlinULiteralExpression.kt @@ -23,22 +23,26 @@ import org.jetbrains.kotlin.psi.KtConstantExpression import org.jetbrains.kotlin.psi.KtEscapeStringTemplateEntry import org.jetbrains.uast.UElement import org.jetbrains.uast.ULiteralExpression +import org.jetbrains.uast.kotlin.internal.KotlinFakeUElement +import org.jetbrains.uast.wrapULiteral class KotlinULiteralExpression( - override val sourcePsi: KtConstantExpression, - givenParent: UElement? -) : KotlinAbstractUExpression(givenParent), ULiteralExpression, KotlinUElementWithType, KotlinEvaluatableUElement { + override val sourcePsi: KtConstantExpression, + givenParent: UElement? +) : KotlinAbstractUExpression(givenParent), ULiteralExpression, KotlinUElementWithType, KotlinEvaluatableUElement, KotlinFakeUElement { override val isNull: Boolean get() = sourcePsi.unwrapBlockOrParenthesis().node?.elementType == KtNodeTypes.NULL override val value by lz { evaluate() } + + override fun unwrapToSourcePsi(): List = listOfNotNull(wrapULiteral(this).sourcePsi) } class KotlinStringULiteralExpression( - override val sourcePsi: PsiElement, - givenParent: UElement?, - val text: String -) : KotlinAbstractUExpression(givenParent), ULiteralExpression, KotlinUElementWithType{ + override val sourcePsi: PsiElement, + givenParent: UElement?, + val text: String +) : KotlinAbstractUExpression(givenParent), ULiteralExpression, KotlinUElementWithType, KotlinFakeUElement { constructor(psi: PsiElement, uastParent: UElement?) : this(psi, uastParent, if (psi is KtEscapeStringTemplateEntry) psi.unescapedValue else psi.text) @@ -48,4 +52,6 @@ class KotlinStringULiteralExpression( override fun evaluate() = value override fun getExpressionType(): PsiType? = PsiType.getJavaLangString(sourcePsi.manager, sourcePsi.resolveScope) + + override fun unwrapToSourcePsi(): List = listOfNotNull(wrapULiteral(this).sourcePsi) } diff --git a/plugins/uast-kotlin/src/org/jetbrains/uast/kotlin/expressions/KotlinUReturnExpression.kt b/plugins/uast-kotlin/src/org/jetbrains/uast/kotlin/expressions/KotlinUReturnExpression.kt index 5a3387faf4a..dd398d1be99 100644 --- a/plugins/uast-kotlin/src/org/jetbrains/uast/kotlin/expressions/KotlinUReturnExpression.kt +++ b/plugins/uast-kotlin/src/org/jetbrains/uast/kotlin/expressions/KotlinUReturnExpression.kt @@ -21,6 +21,8 @@ import org.jetbrains.kotlin.psi.KtReturnExpression import org.jetbrains.uast.UElement import org.jetbrains.uast.UExpression import org.jetbrains.uast.UReturnExpression +import org.jetbrains.uast.kotlin.internal.KotlinFakeUElement +import org.jetbrains.uast.kotlin.internal.toSourcePsiFakeAware class KotlinUReturnExpression( override val sourcePsi: KtReturnExpression, @@ -31,11 +33,15 @@ class KotlinUReturnExpression( class KotlinUImplicitReturnExpression( givenParent: UElement? -) : KotlinAbstractUExpression(givenParent), UReturnExpression, KotlinUElementWithType { +) : KotlinAbstractUExpression(givenParent), UReturnExpression, KotlinUElementWithType, KotlinFakeUElement { override val psi: PsiElement? get() = null override lateinit var returnExpression: UExpression internal set + override fun unwrapToSourcePsi(): List { + return returnExpression.toSourcePsiFakeAware() + } + } \ No newline at end of file diff --git a/plugins/uast-kotlin/src/org/jetbrains/uast/kotlin/internal/KotlinFakeUElement.kt b/plugins/uast-kotlin/src/org/jetbrains/uast/kotlin/internal/KotlinFakeUElement.kt new file mode 100644 index 00000000000..aa4a176dbf2 --- /dev/null +++ b/plugins/uast-kotlin/src/org/jetbrains/uast/kotlin/internal/KotlinFakeUElement.kt @@ -0,0 +1,18 @@ +/* + * Copyright 2010-2020 JetBrains s.r.o. and Kotlin Programming Language contributors. + * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file. + */ + +package org.jetbrains.uast.kotlin.internal + +import com.intellij.psi.PsiElement +import org.jetbrains.uast.UElement + +interface KotlinFakeUElement { + fun unwrapToSourcePsi(): List +} + +fun UElement.toSourcePsiFakeAware(): List { + if (this is KotlinFakeUElement) return this.unwrapToSourcePsi() + return listOfNotNull(this.sourcePsi) +} \ No newline at end of file diff --git a/plugins/uast-kotlin/tests/KotlinUastGenerationTest.kt.201 b/plugins/uast-kotlin/tests/KotlinUastGenerationTest.kt.201 new file mode 100644 index 00000000000..77716adb13d --- /dev/null +++ b/plugins/uast-kotlin/tests/KotlinUastGenerationTest.kt.201 @@ -0,0 +1,938 @@ +/* + * Copyright 2010-2020 JetBrains s.r.o. and Kotlin Programming Language contributors. + * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file. + */ + +package org.jetbrains.uast.test.kotlin + +import com.intellij.openapi.command.WriteCommandAction +import com.intellij.openapi.util.Key +import com.intellij.openapi.util.Ref +import com.intellij.psi.* +import com.intellij.testFramework.LightProjectDescriptor +import com.intellij.testFramework.UsefulTestCase +import junit.framework.TestCase +import org.jetbrains.kotlin.idea.KotlinLanguage +import org.jetbrains.kotlin.idea.test.KotlinLightCodeInsightFixtureTestCase +import org.jetbrains.kotlin.idea.test.KotlinWithJdkAndRuntimeLightProjectDescriptor +import org.jetbrains.kotlin.psi.* +import org.jetbrains.kotlin.utils.addToStdlib.cast +import org.jetbrains.uast.* +import org.jetbrains.uast.expressions.UInjectionHost +import org.jetbrains.uast.generate.UParameterInfo +import org.jetbrains.uast.generate.UastCodeGenerationPlugin +import org.jetbrains.uast.generate.refreshed +import org.jetbrains.uast.generate.replace +import org.jetbrains.uast.kotlin.generate.KotlinUastElementFactory +import org.jetbrains.uast.test.env.kotlin.findElementByTextFromPsi +import org.jetbrains.uast.test.env.kotlin.findUElementByTextFromPsi +import kotlin.test.fail as kfail + +class KotlinUastGenerationTest : KotlinLightCodeInsightFixtureTestCase() { + + override fun getProjectDescriptor(): LightProjectDescriptor = + KotlinWithJdkAndRuntimeLightProjectDescriptor.INSTANCE + + + private val psiFactory + get() = KtPsiFactory(project) + private val generatePlugin: UastCodeGenerationPlugin + get() = UastCodeGenerationPlugin.byLanguage(KotlinLanguage.INSTANCE)!! + private val uastElementFactory + get() = generatePlugin.getElementFactory(myFixture.project) as KotlinUastElementFactory + + fun `test logical and operation with simple operands`() { + val left = psiFactory.createExpression("true").toUElementOfType() + ?: kfail("Cannot create left UExpression") + val right = psiFactory.createExpression("false").toUElementOfType() + ?: kfail("Cannot create right UExpression") + + val expression = uastElementFactory.createBinaryExpression(left, right, UastBinaryOperator.LOGICAL_AND, dummyContextFile()) + ?: kfail("Cannot create expression") + + TestCase.assertEquals("true && false", expression.sourcePsi?.text) + } + + fun `test logical and operation with simple operands with parenthesis`() { + val left = psiFactory.createExpression("(true)").toUElementOfType() + ?: kfail("Cannot create left UExpression") + val right = psiFactory.createExpression("(false)").toUElementOfType() + ?: kfail("Cannot create right UExpression") + + val expression = uastElementFactory.createFlatBinaryExpression(left, right, UastBinaryOperator.LOGICAL_AND, dummyContextFile()) + ?: kfail("Cannot create expression") + + TestCase.assertEquals("true && false", expression.sourcePsi?.text) + TestCase.assertEquals(""" + UBinaryExpression (operator = &&) + ULiteralExpression (value = true) + ULiteralExpression (value = false) + """.trimIndent(), expression.putIntoFunctionBody().asRecursiveLogString().trim()) + } + + fun `test logical and operation with simple operands with parenthesis polyadic`() { + val left = psiFactory.createExpression("(true && false)").toUElementOfType() + ?: kfail("Cannot create left UExpression") + val right = psiFactory.createExpression("(false)").toUElementOfType() + ?: kfail("Cannot create right UExpression") + + val expression = uastElementFactory.createFlatBinaryExpression(left, right, UastBinaryOperator.LOGICAL_AND, dummyContextFile()) + ?: kfail("Cannot create expression") + + TestCase.assertEquals("true && false && false", expression.sourcePsi?.text) + TestCase.assertEquals(""" + UBinaryExpression (operator = &&) + UBinaryExpression (operator = &&) + ULiteralExpression (value = null) + ULiteralExpression (value = null) + ULiteralExpression (value = null) + """.trimIndent(), expression.asRecursiveLogString().trim()) + } + + fun `test simple reference creating from variable`() { + val context = dummyContextFile() + val variable = uastElementFactory.createLocalVariable( + "a", PsiType.INT, uastElementFactory.createNullLiteral(context), false, context + ) ?: kfail("cannot create variable") + + val reference = uastElementFactory.createSimpleReference(variable, context) ?: kfail("cannot create reference") + TestCase.assertEquals("a", reference.identifier) + } + + fun `test simple reference by name`() { + val reference = uastElementFactory.createSimpleReference("a", dummyContextFile()) ?: kfail("cannot create reference") + TestCase.assertEquals("a", reference.identifier) + } + + fun `test parenthesised expression`() { + val expression = psiFactory.createExpression("a + b").toUElementOfType() + ?: kfail("cannot create expression") + val parenthesizedExpression = uastElementFactory.createParenthesizedExpression(expression, dummyContextFile()) + ?: kfail("cannot create parenthesized expression") + + TestCase.assertEquals("(a + b)", parenthesizedExpression.sourcePsi?.text) + } + + fun `test return expression`() { + val expression = psiFactory.createExpression("a + b").toUElementOfType() + ?: kfail("Cannot find plugin") + + val returnExpression = uastElementFactory.createReturnExpresion(expression, false, dummyContextFile()) ?: kfail("cannot create return expression") + TestCase.assertEquals("a + b", returnExpression.returnExpression?.asRenderString()) + TestCase.assertEquals("return a + b", returnExpression.sourcePsi?.text) + } + + fun `test variable declaration without type`() { + val expression = psiFactory.createExpression("1 + 2").toUElementOfType() + ?: kfail("cannot create variable declaration") + + val declaration = uastElementFactory.createLocalVariable("a", null, expression, false, dummyContextFile()) ?: kfail("cannot create variable") + + TestCase.assertEquals("var a = 1 + 2", declaration.sourcePsi?.text) + } + + fun `test variable declaration with type`() { + val expression = psiFactory.createExpression("b").toUElementOfType() + ?: kfail("cannot create variable declaration") + + val declaration = uastElementFactory.createLocalVariable("a", PsiType.DOUBLE, expression, false, dummyContextFile()) ?: kfail("cannot create variable") + + TestCase.assertEquals("var a: kotlin.Double = b", declaration.sourcePsi?.text) + } + + fun `test final variable declaration`() { + val expression = psiFactory.createExpression("b").toUElementOfType() + ?: kfail("cannot create variable declaration") + + val declaration = uastElementFactory.createLocalVariable("a", PsiType.DOUBLE, expression, true, dummyContextFile()) + ?: kfail("cannot create variable") + + TestCase.assertEquals("val a: kotlin.Double = b", declaration.sourcePsi?.text) + } + + fun `test final variable declaration with unique name`() { + val expression = psiFactory.createExpression("b").toUElementOfType() + ?: kfail("cannot create variable declaration") + + val declaration = uastElementFactory.createLocalVariable("a", PsiType.DOUBLE, expression, true, dummyContextFile()) + ?: kfail("cannot create variable") + + TestCase.assertEquals("val a: kotlin.Double = b", declaration.sourcePsi?.text) + TestCase.assertEquals(""" + ULocalVariable (name = a) + USimpleNameReferenceExpression (identifier = b) + """.trimIndent(), declaration.asRecursiveLogString().trim()) + } + + fun `test block expression`() { + val statement1 = psiFactory.createExpression("System.out.println()").toUElementOfType() + ?: kfail("cannot create statement") + val statement2 = psiFactory.createExpression("System.out.println(2)").toUElementOfType() + ?: kfail("cannot create statement") + + val block = uastElementFactory.createBlockExpression(listOf(statement1, statement2), dummyContextFile()) ?: kfail("cannot create block") + + TestCase.assertEquals(""" + { + System.out.println() + System.out.println(2) + } + """.trimIndent(), block.sourcePsi?.text + ) + } + + fun `test lambda expression`() { + val statement = psiFactory.createExpression("System.out.println()").toUElementOfType() + ?: kfail("cannot create statement") + + val lambda = uastElementFactory.createLambdaExpression( + listOf( + UParameterInfo(PsiType.INT, "a"), + UParameterInfo(null, "b") + ), + statement, + dummyContextFile() + ) ?: kfail("cannot create lambda") + + TestCase.assertEquals("{ a: kotlin.Int, b -> System.out.println() }", lambda.sourcePsi?.text) + TestCase.assertEquals(""" + ULambdaExpression + UParameter (name = a) + UAnnotation (fqName = org.jetbrains.annotations.NotNull) + UParameter (name = b) + UAnnotation (fqName = null) + UBlockExpression + UQualifiedReferenceExpression + UQualifiedReferenceExpression + USimpleNameReferenceExpression (identifier = System) + USimpleNameReferenceExpression (identifier = out) + UCallExpression (kind = UastCallKind(name='method_call'), argCount = 0)) + UIdentifier (Identifier (println)) + USimpleNameReferenceExpression (identifier = println, resolvesTo = null) + """.trimIndent(), lambda.putIntoFunctionBody().asRecursiveLogString().trim()) + } + + private fun UExpression.putIntoFunctionBody(): UExpression { + val file = myFixture.configureByText("dummyFile.kt", "fun foo() { TODO() }") as KtFile + val ktFunction = file.declarations.single { it.name == "foo" } as KtFunction + val uMethod = ktFunction.toUElementOfType()!! + return runWriteCommand { + uMethod.uastBody.cast().expressions.single().replace(this)!! + } + } + + private fun T.putIntoVarInitializer(): T { + val file = myFixture.configureByText("dummyFile.kt", "val foo = TODO()") as KtFile + val ktFunction = file.declarations.single { it.name == "foo" } as KtProperty + val uMethod = ktFunction.toUElementOfType()!! + return runWriteCommand { + @Suppress("UNCHECKED_CAST") + generatePlugin.replace(uMethod.uastInitializer!!, this, UExpression::class.java) as T + } + } + + private fun runWriteCommand(uExpression: () -> T): T { + val result = Ref() + WriteCommandAction.runWriteCommandAction(project) { + result.set(uExpression()) + } + return result.get() + } + + fun `test lambda expression with explicit types`() { + val statement = psiFactory.createExpression("System.out.println()").toUElementOfType() + ?: kfail("cannot create statement") + + val lambda = uastElementFactory.createLambdaExpression( + listOf( + UParameterInfo(PsiType.INT, "a"), + UParameterInfo(PsiType.DOUBLE, "b") + ), + statement, + dummyContextFile() + ) ?: kfail("cannot create lambda") + + TestCase.assertEquals("{ a: kotlin.Int, b: kotlin.Double -> System.out.println() }", lambda.sourcePsi?.text) + TestCase.assertEquals(""" + ULambdaExpression + UParameter (name = a) + UAnnotation (fqName = org.jetbrains.annotations.NotNull) + UParameter (name = b) + UAnnotation (fqName = org.jetbrains.annotations.NotNull) + UBlockExpression + UQualifiedReferenceExpression + UQualifiedReferenceExpression + USimpleNameReferenceExpression (identifier = System) + USimpleNameReferenceExpression (identifier = out) + UCallExpression (kind = UastCallKind(name='method_call'), argCount = 0)) + UIdentifier (Identifier (println)) + USimpleNameReferenceExpression (identifier = println, resolvesTo = null) + """.trimIndent(), lambda.putIntoFunctionBody().asRecursiveLogString().trim()) + } + + fun `test lambda expression with simplified block body with context`() { + val r = psiFactory.createExpression("return \"10\"").toUElementOfType() + ?: kfail("cannot create return") + + val block = uastElementFactory.createBlockExpression(listOf(r), dummyContextFile()) ?: kfail("cannot create block") + + val lambda = uastElementFactory.createLambdaExpression(listOf(UParameterInfo(null, "a")), block, dummyContextFile()) + ?: kfail("cannot create lambda") + TestCase.assertEquals("""{ a -> "10" }""".trimMargin(), lambda.sourcePsi?.text) + TestCase.assertEquals(""" + ULambdaExpression + UParameter (name = a) + UAnnotation (fqName = org.jetbrains.annotations.NotNull) + UBlockExpression + UReturnExpression + ULiteralExpression (value = "10") + """.trimIndent(), lambda.putIntoVarInitializer().asRecursiveLogString().trim()) + } + + fun `test function argument replacement`() { + + val file = myFixture.configureByText( + "test.kt", """ + fun f(a: Any){} + + fun main(){ + f(a) + } + """.trimIndent() + ) + + val expression = file.findUElementByTextFromPsi("f(a)") + val newArgument = psiFactory.createExpression("b").toUElementOfType() + ?: kfail("cannot create reference") + + WriteCommandAction.runWriteCommandAction(project) { + TestCase.assertNotNull(expression.valueArguments[0].replace(newArgument)) + } + + val updated = expression.refreshed() ?: kfail("cannot update expression") + TestCase.assertEquals("f", updated.methodName) + TestCase.assertTrue(updated.valueArguments[0] is USimpleNameReferenceExpression) + TestCase.assertEquals("b", (updated.valueArguments[0] as USimpleNameReferenceExpression).identifier) + TestCase.assertEquals(""" + UCallExpression (kind = UastCallKind(name='method_call'), argCount = 1)) + UIdentifier (Identifier (f)) + USimpleNameReferenceExpression (identifier = f, resolvesTo = null) + USimpleNameReferenceExpression (identifier = b) + """.trimIndent(), updated.asRecursiveLogString().trim()) + } + + fun `test suggested name`() { + val expression = psiFactory.createExpression("f(a) + 1").toUElementOfType() + ?: kfail("cannot create expression") + val variable = uastElementFactory.createLocalVariable(null, PsiType.INT, expression, true, dummyContextFile()) + ?: kfail("cannot create variable") + + TestCase.assertEquals("val i: kotlin.Int = f(a) + 1", variable.sourcePsi?.text) + TestCase.assertEquals(""" + ULocalVariable (name = i) + UBinaryExpression (operator = +) + UCallExpression (kind = UastCallKind(name='method_call'), argCount = 1)) + UIdentifier (Identifier (f)) + USimpleNameReferenceExpression (identifier = , resolvesTo = null) + USimpleNameReferenceExpression (identifier = a) + ULiteralExpression (value = null) + """.trimIndent(), variable.asRecursiveLogString().trim()) + } + + fun `test method call generation with receiver`() { + val receiver = psiFactory.createExpression(""""10"""").toUElementOfType() + ?: kfail("cannot create receiver") + val arg1 = psiFactory.createExpression("1").toUElementOfType() + ?: kfail("cannot create arg1") + val arg2 = psiFactory.createExpression("2").toUElementOfType() + ?: kfail("cannot create arg2") + val methodCall = uastElementFactory.createCallExpression( + receiver, + "substring", + listOf(arg1, arg2), + null, + UastCallKind.METHOD_CALL + ) ?: kfail("cannot create call") + + TestCase.assertEquals(""""10".substring(1,2)""", methodCall.uastParent?.sourcePsi?.text) + TestCase.assertEquals(""" + UQualifiedReferenceExpression + ULiteralExpression (value = "10") + UCallExpression (kind = UastCallKind(name='method_call'), argCount = 2)) + UIdentifier (Identifier (substring)) + USimpleNameReferenceExpression (identifier = , resolvesTo = null) + ULiteralExpression (value = null) + ULiteralExpression (value = null) + """.trimIndent(), methodCall.uastParent?.asRecursiveLogString()?.trim() + ) + } + + fun `test method call generation without receiver`() { + val arg1 = psiFactory.createExpression("1").toUElementOfType() + ?: kfail("cannot create arg1") + val arg2 = psiFactory.createExpression("2").toUElementOfType() + ?: kfail("cannot create arg2") + val methodCall = uastElementFactory.createCallExpression( + null, + "substring", + listOf(arg1, arg2), + null, + UastCallKind.METHOD_CALL + ) ?: kfail("cannot create call") + + TestCase.assertEquals("""substring(1,2)""", methodCall.sourcePsi?.text) + } + + fun `test method call generation with generics restoring`() { + val arrays = psiFactory.createExpression("java.util.Arrays").toUElementOfType() + ?: kfail("cannot create receiver") + val methodCall = uastElementFactory.createCallExpression( + arrays, + "asList", + listOf(), + createTypeFromText("java.util.List", null), + UastCallKind.METHOD_CALL, + dummyContextFile() + ) ?: kfail("cannot create call") + TestCase.assertEquals("java.util.Arrays.asList()", methodCall.uastParent?.sourcePsi?.text) + } + + fun `test method call generation with generics restoring 2 parameters`() { + val collections = psiFactory.createExpression("java.util.Collections").toUElementOfType() + ?: kfail("cannot create receiver") + TestCase.assertEquals("java.util.Collections", collections.asRenderString()) + val methodCall = uastElementFactory.createCallExpression( + collections, + "emptyMap", + listOf(), + createTypeFromText( + "java.util.Map", + null + ), + UastCallKind.METHOD_CALL, + dummyContextFile() + ) ?: kfail("cannot create call") + TestCase.assertEquals("emptyMap()", methodCall.sourcePsi?.text) + TestCase.assertEquals("java.util.Collections.emptyMap()", methodCall.sourcePsi?.parent?.text) + TestCase.assertEquals( + """ + UQualifiedReferenceExpression + UQualifiedReferenceExpression + UQualifiedReferenceExpression + USimpleNameReferenceExpression (identifier = java) + USimpleNameReferenceExpression (identifier = util) + USimpleNameReferenceExpression (identifier = Collections) + UCallExpression (kind = UastCallKind(name='method_call'), argCount = 0)) + UIdentifier (Identifier (emptyMap)) + USimpleNameReferenceExpression (identifier = , resolvesTo = null) + """.trimIndent(), methodCall.uastParent?.asRecursiveLogString()?.trim() + ) + } + + private fun dummyContextFile(): KtFile = myFixture.configureByText("file.kt", "fun foo() {}") as KtFile + + fun `test method call generation with generics restoring 1 parameter with 1 existing`() { + val a = psiFactory.createExpression("A").toUElementOfType() + ?: kfail("cannot create a receiver") + val param = psiFactory.createExpression("\"a\"").toUElementOfType() + ?: kfail("cannot create a parameter") + val methodCall = uastElementFactory.createCallExpression( + a, + "kek", + listOf(param), + createTypeFromText( + "java.util.Map", + null + ), + UastCallKind.METHOD_CALL, + dummyContextFile() + ) ?: kfail("cannot create call") + + TestCase.assertEquals("A.kek(\"a\")", methodCall.sourcePsi?.parent?.text) + TestCase.assertEquals(""" + UQualifiedReferenceExpression + USimpleNameReferenceExpression (identifier = A) + UCallExpression (kind = UastCallKind(name='method_call'), argCount = 1)) + UIdentifier (Identifier (kek)) + USimpleNameReferenceExpression (identifier = , resolvesTo = null) + ULiteralExpression (value = "a") + """.trimIndent(), methodCall.uastParent?.asRecursiveLogString()?.trim()) + } + + + //not implemented (currently we dont perform resolve in code generating) + fun `ignore method call generation with generics restoring 1 parameter with 1 unused `() { + val aClassFile = myFixture.configureByText("A.kt", + """ + object A { + fun kek(a: T1): Map { + return TODO(); + } + } + """.trimIndent() + ) + val a = psiFactory.createExpression("A").toUElementOfType() + ?: kfail("cannot create a receiver") + val param = psiFactory.createExpression("\"a\"").toUElementOfType() + ?: kfail("cannot create a parameter") + val methodCall = uastElementFactory.createCallExpression( + a, + "kek", + listOf(param), + createTypeFromText( + "java.util.Map", + null + ), + UastCallKind.METHOD_CALL, + aClassFile + ) ?: kfail("cannot create call") + + TestCase.assertEquals("A.kek(\"a\")", methodCall.sourcePsi?.text) + } + + + fun `test method call generation with generics with context`() { + val file = myFixture.configureByText("file.kt", """ + class A { + fun method(): List { TODO() } + } + + fun main(){ + val a = A() + println(a) + } + """.trimIndent() + ) as KtFile + + val reference = file.findUElementByTextFromPsi("println(a)") + .findElementByTextFromPsi("a") + + val callExpression = uastElementFactory.createCallExpression( + reference, + "method", + emptyList(), + createTypeFromText( + "java.util.List", + null + ), + UastCallKind.METHOD_CALL, + context = reference.sourcePsi + ) ?: kfail("cannot create method call") + + TestCase.assertEquals("a.method()", callExpression.uastParent?.sourcePsi?.text) + TestCase.assertEquals(""" + UQualifiedReferenceExpression + USimpleNameReferenceExpression (identifier = a) + UCallExpression (kind = UastCallKind(name='method_call'), argCount = 0)) + UIdentifier (Identifier (method)) + USimpleNameReferenceExpression (identifier = , resolvesTo = null) + """.trimIndent(), callExpression.uastParent?.asRecursiveLogString()?.trim() + ) + + } + + + fun `test replace lambda implicit return value`() { + + val file = myFixture.configureByText( + "file.kt", """ + fun main(){ + val a: (Int) -> String = { + println(it) + println(2) + "abc" + } + } + """.trimIndent() + ) as KtFile + + val uLambdaExpression = file.findUElementByTextFromPsi("\"abc\"") + .getParentOfType() ?: kfail("cant get lambda") + + val expressions = uLambdaExpression.body.cast().expressions + UsefulTestCase.assertSize(3, expressions) + + val uReturnExpression = expressions.last() as UReturnExpression + val newStringLiteral = uastElementFactory.createStringLiteralExpression("def", file) ?: kfail("cannot create method call") + + val defReturn = runWriteCommand { uReturnExpression.replace(newStringLiteral) ?: kfail("cant replace") } + val uLambdaExpression2 = defReturn.getParentOfType() ?: kfail("cant get lambda") + + TestCase.assertEquals("{\n println(it)\n println(2)\n \"def\"\n }", uLambdaExpression2.sourcePsi?.text) + TestCase.assertEquals( + """ + ULambdaExpression + UParameter (name = it) + UBlockExpression + UCallExpression (kind = UastCallKind(name='method_call'), argCount = 1)) + UIdentifier (Identifier (println)) + USimpleNameReferenceExpression (identifier = println, resolvesTo = null) + USimpleNameReferenceExpression (identifier = it) + UCallExpression (kind = UastCallKind(name='method_call'), argCount = 1)) + UIdentifier (Identifier (println)) + USimpleNameReferenceExpression (identifier = println, resolvesTo = null) + ULiteralExpression (value = 2) + UReturnExpression + ULiteralExpression (value = "def") + """.trimIndent(), uLambdaExpression2.asRecursiveLogString().trim() + ) + + } + + + private class UserDataChecker { + + private val storedData = Any() + private val KEY = Key.create("testKey") + + private lateinit var uniqueStringLiteralText: String + + fun store(uElement: UInjectionHost) { + val psiElement = uElement.sourcePsi as KtStringTemplateExpression + uniqueStringLiteralText = psiElement.text + psiElement.putCopyableUserData(KEY, storedData) + } + + fun checkUserDataAlive(uElement: UElement) { + val psiElements = uElement.let { SyntaxTraverser.psiTraverser(it.sourcePsi) } + .filter(KtStringTemplateExpression::class.java) + .filter { it.text == uniqueStringLiteralText }.toList() + + UsefulTestCase.assertNotEmpty(psiElements) + UsefulTestCase.assertTrue("uElement still should keep the userdata", psiElements.any { storedData === it!!.getCopyableUserData(KEY) }) + } + + } + + fun `test add intermediate returns to lambda`() { + + val file = myFixture.configureByText( + "file.kt", """ + fun main(){ + val a: (Int) -> String = lname@{ + println(it) + println(2) + "abc" + } + } + """.trimIndent() + ) as KtFile + + val aliveChecker = UserDataChecker() + val uLambdaExpression = file.findUElementByTextFromPsi("\"abc\"") + .also(aliveChecker::store) + .getParentOfType() ?: kfail("cant get lambda") + + val oldBlockExpression = uLambdaExpression.body.cast() + UsefulTestCase.assertSize(3, oldBlockExpression.expressions) + + val conditionalExit = with(uastElementFactory) { + createIfExpression( + createBinaryExpression( + createSimpleReference("it", uLambdaExpression.sourcePsi)!!, + createIntLiteral(3, uLambdaExpression.sourcePsi), + UastBinaryOperator.GREATER, + uLambdaExpression.sourcePsi + )!!, + createReturnExpresion( + createStringLiteralExpression("exit", uLambdaExpression.sourcePsi)!!, true, + uLambdaExpression.sourcePsi + )!!, + null, + uLambdaExpression.sourcePsi + )!! + } + + val newBlockExpression = uastElementFactory.createBlockExpression( + listOf(conditionalExit) + oldBlockExpression.expressions, + uLambdaExpression.sourcePsi + )!! + + aliveChecker.checkUserDataAlive(newBlockExpression) + + val uLambdaExpression2 = runWriteCommand { + oldBlockExpression.replace(newBlockExpression) ?: kfail("cant replace") + }.getParentOfType() ?: kfail("cant get lambda") + + aliveChecker.checkUserDataAlive(uLambdaExpression2) + + TestCase.assertEquals( + """ + lname@{ + if (it > 3) return@lname "exit" + println(it) + println(2) + "abc" + } + """.trimIndent(), uLambdaExpression2.sourcePsi?.parent?.text + ) + TestCase.assertEquals( + """ + ULambdaExpression + UParameter (name = it) + UBlockExpression + UIfExpression + UBinaryExpression (operator = >) + USimpleNameReferenceExpression (identifier = it) + ULiteralExpression (value = 3) + UReturnExpression + ULiteralExpression (value = "exit") + UCallExpression (kind = UastCallKind(name='method_call'), argCount = 1)) + UIdentifier (Identifier (println)) + USimpleNameReferenceExpression (identifier = println, resolvesTo = null) + USimpleNameReferenceExpression (identifier = it) + UCallExpression (kind = UastCallKind(name='method_call'), argCount = 1)) + UIdentifier (Identifier (println)) + USimpleNameReferenceExpression (identifier = println, resolvesTo = null) + ULiteralExpression (value = 2) + UReturnExpression + ULiteralExpression (value = "abc") + """.trimIndent(), uLambdaExpression2.asRecursiveLogString().trim() + ) + + } + + fun `test converting lambda to if`() { + + val file = myFixture.configureByText( + "file.kt", """ + fun foo(call: (Int) -> String): String = call.invoke(2) + + fun main() { + foo { + println(it) + println(2) + "abc" + } + } + } + """.trimIndent() + ) as KtFile + + val aliveChecker = UserDataChecker() + val uLambdaExpression = file.findUElementByTextFromPsi("\"abc\"") + .also { aliveChecker.store(it) } + .getParentOfType() ?: kfail("cant get lambda") + + val oldBlockExpression = uLambdaExpression.body.cast() + aliveChecker.checkUserDataAlive(oldBlockExpression) + + val newLambda = with(uastElementFactory) { + createLambdaExpression( + listOf(UParameterInfo(null, "it")), + createIfExpression( + createBinaryExpression( + createSimpleReference("it", uLambdaExpression.sourcePsi)!!, + createIntLiteral(3, uLambdaExpression.sourcePsi), + UastBinaryOperator.GREATER, + uLambdaExpression.sourcePsi + )!!, + oldBlockExpression, + createReturnExpresion( + createStringLiteralExpression("exit", uLambdaExpression.sourcePsi)!!, true, + uLambdaExpression.sourcePsi + )!!, + uLambdaExpression.sourcePsi + )!!.also { + aliveChecker.checkUserDataAlive(it) + }, + uLambdaExpression.sourcePsi + )!! + } + aliveChecker.checkUserDataAlive(newLambda) + + val uLambdaExpression2 = runWriteCommand { + uLambdaExpression.replace(newLambda) ?: kfail("cant replace") + } + + TestCase.assertEquals( + """ + { it -> + if (it > 3) { + println(it) + println(2) + "abc" + } else return@foo "exit" + } + """.trimIndent(), uLambdaExpression2.sourcePsi?.parent?.text + ) + TestCase.assertEquals( + """ + ULambdaExpression + UParameter (name = it) + UAnnotation (fqName = null) + UBlockExpression + UReturnExpression + UIfExpression + UBinaryExpression (operator = >) + USimpleNameReferenceExpression (identifier = it) + ULiteralExpression (value = 3) + UBlockExpression + UCallExpression (kind = UastCallKind(name='method_call'), argCount = 1)) + UIdentifier (Identifier (println)) + USimpleNameReferenceExpression (identifier = println, resolvesTo = null) + USimpleNameReferenceExpression (identifier = it) + UCallExpression (kind = UastCallKind(name='method_call'), argCount = 1)) + UIdentifier (Identifier (println)) + USimpleNameReferenceExpression (identifier = println, resolvesTo = null) + ULiteralExpression (value = 2) + ULiteralExpression (value = "abc") + UReturnExpression + ULiteralExpression (value = "exit") + """.trimIndent(), uLambdaExpression2.asRecursiveLogString().trim() + ) + aliveChecker.checkUserDataAlive(uLambdaExpression2) + + } + + fun `test removing unnecessary type parameters while replace`() { + val aClassFile = myFixture.configureByText( + "A.kt", + """ + class A { + fun method():List = TODO() + } + """.trimIndent() + ) + + val reference = psiFactory.createExpression("a") + .toUElementOfType() ?: kfail("cannot create reference expression") + val callExpression = uastElementFactory.createCallExpression( + reference, + "method", + emptyList(), + createTypeFromText( + "java.util.List", + null + ), + UastCallKind.METHOD_CALL, + context = aClassFile + ) ?: kfail("cannot create method call") + + val listAssigment = myFixture.addFileToProject("temp.kt", """ + fun foo(kek: List) { + val list: List = kek + } + """.trimIndent()).findUElementByTextFromPsi("val list: List = kek") + + WriteCommandAction.runWriteCommandAction(project) { + val methodCall = listAssigment.uastInitializer?.replace(callExpression) ?: kfail("cannot replace!") + // originally result expected be `a.method()` but we expect to clean up type arguments in other plase + TestCase.assertEquals("a.method()", methodCall.sourcePsi?.parent?.text) + } + + } + + fun `test create if`() { + val condition = psiFactory.createExpression("true").toUElementOfType() + ?: kfail("cannot create condition") + val thenBranch = psiFactory.createBlock("{a(b);}").toUElementOfType() + ?: kfail("cannot create then branch") + val elseBranch = psiFactory.createExpression("c++").toUElementOfType() + ?: kfail("cannot create else branch") + val ifExpression = uastElementFactory.createIfExpression(condition, thenBranch, elseBranch, dummyContextFile()) + ?: kfail("cannot create if expression") + TestCase.assertEquals("if (true) {\n { a(b); }\n } else c++", ifExpression.sourcePsi?.text) + } + + fun `test qualified reference`() { + val reference = uastElementFactory.createQualifiedReference("java.util.List", myFixture.file) + TestCase.assertEquals("java.util.List", reference?.sourcePsi?.text) + } + + fun `test build lambda from returning a variable`() { + val context = dummyContextFile() + val localVariable = uastElementFactory.createLocalVariable("a", null, uastElementFactory.createNullLiteral(context), true, context) + ?: kfail("cannot create variable") + val declarationExpression = + uastElementFactory.createDeclarationExpression(listOf(localVariable), context) ?: kfail("cannot create declaration expression") + val returnExpression = uastElementFactory.createReturnExpresion( + uastElementFactory.createSimpleReference(localVariable, context), false, context + ) ?: kfail("cannot create return expression") + val block = uastElementFactory.createBlockExpression(listOf(declarationExpression, returnExpression), context) + ?: kfail("cannot create block expression") + + TestCase.assertEquals(""" + UBlockExpression + UDeclarationsExpression + ULocalVariable (name = a) + ULiteralExpression (value = null) + UReturnExpression + USimpleNameReferenceExpression (identifier = a) + """.trimIndent(), block.asRecursiveLogString().trim()) + + + val lambda = uastElementFactory.createLambdaExpression(listOf(), block, context) ?: kfail("cannot create lambda expression") + TestCase.assertEquals("{ val a = null\na }", lambda.sourcePsi?.text) + + TestCase.assertEquals(""" + ULambdaExpression + UBlockExpression + UDeclarationsExpression + ULocalVariable (name = a) + ULiteralExpression (value = null) + UReturnExpression + USimpleNameReferenceExpression (identifier = a) + """.trimIndent(), lambda.putIntoVarInitializer().asRecursiveLogString().trim()) + } + + fun `test expand oneline lambda`() { + + val context = dummyContextFile() + val parameters = listOf(UParameterInfo(PsiType.INT, "a")) + val oneLineLambda = with(uastElementFactory) { + createLambdaExpression( + parameters, + createBinaryExpression( + createSimpleReference("a", context)!!, + createSimpleReference("a", context)!!, + UastBinaryOperator.PLUS, context + )!!, context + )!! + }.putIntoVarInitializer() + + val lambdaReturn = (oneLineLambda.body as UBlockExpression).expressions.single() + + val lambda = with(uastElementFactory) { + createLambdaExpression( + parameters, + createBlockExpression( + listOf( + createCallExpression( + null, + "println", + listOf(createSimpleReference("a", context)!!), + PsiType.VOID, + UastCallKind.METHOD_CALL, + context + )!!, + lambdaReturn + ), + context + )!!, context + )!! + } + + TestCase.assertEquals("{ a: kotlin.Int -> println(a)\na + a }", lambda.sourcePsi?.text) + + TestCase.assertEquals(""" + ULambdaExpression + UParameter (name = a) + UAnnotation (fqName = org.jetbrains.annotations.NotNull) + UBlockExpression + UCallExpression (kind = UastCallKind(name='method_call'), argCount = 1)) + UIdentifier (Identifier (println)) + USimpleNameReferenceExpression (identifier = println, resolvesTo = null) + USimpleNameReferenceExpression (identifier = a) + UReturnExpression + UBinaryExpression (operator = +) + USimpleNameReferenceExpression (identifier = a) + USimpleNameReferenceExpression (identifier = a) + """.trimIndent(), lambda.putIntoVarInitializer().asRecursiveLogString().trim()) + } + + private fun createTypeFromText(s: String, newClass: PsiElement?): PsiType? { + return JavaPsiFacade.getElementFactory(myFixture.project).createTypeFromText(s, newClass) + } + +}