Added ad-hoc KtValueArgument.stringTemplateExpression

Use only stub stringTemplateExpression if it is present and don't use ast nodes as it could be unbounded

#KT-47034 Fixed

compiler/frontend.java/src/org/jetbrains/kotlin/fileClasses/JvmFileClassUtil.kt

@@ -30,6 +30,7 @@ import org.jetbrains.kotlin.psi.*
 import org.jetbrains.kotlin.resolve.DescriptorToSourceUtils
 import org.jetbrains.kotlin.resolve.jvm.JvmClassName
 import org.jetbrains.kotlin.serialization.deserialization.descriptors.DeserializedMemberDescriptor
+import org.jetbrains.kotlin.utils.addToStdlib.safeAs
 
 object JvmFileClassUtil {
     val JVM_NAME: FqName = FqName("kotlin.jvm.JvmName")
@@ -99,9 +100,13 @@ object JvmFileClassUtil {
         }
 
     fun getLiteralStringFromAnnotation(annotation: KtAnnotationEntry): String? {
-        val argumentExpression = annotation.valueArguments.firstOrNull()?.getArgumentExpression() ?: return null
-        val stringTemplate = argumentExpression as? KtStringTemplateExpression ?: return null
-        val singleEntry = stringTemplate.entries.singleOrNull() as? KtLiteralStringTemplateEntry ?: return null
+        val stringTemplateExpression = annotation.valueArguments.firstOrNull()?.run {
+            when (this) {
+                is KtValueArgument -> stringTemplateExpression
+                else -> getArgumentExpression().safeAs<KtStringTemplateExpression>()
+            }
+        } ?: return null
+        val singleEntry = stringTemplateExpression.entries.singleOrNull() as? KtLiteralStringTemplateEntry ?: return null
         return singleEntry.text
     }
 }

compiler/psi/src/org/jetbrains/kotlin/psi/KtValueArgument.java

@@ -59,6 +59,10 @@ public class KtValueArgument extends KtElementImplStub<KotlinValueArgumentStub<?
             KtStubElementTypes.STRING_TEMPLATE
     );
 
+    private static final TokenSet STRING_TEMPLATE_EXPRESSIONS_TYPES = TokenSet.create(
+            KtStubElementTypes.STRING_TEMPLATE
+    );
+
     @Override
     @Nullable @IfNotParsed
     public KtExpression getArgumentExpression() {
@@ -73,6 +77,19 @@ public class KtValueArgument extends KtElementImplStub<KotlinValueArgumentStub<?
         return findChildByClass(KtExpression.class);
     }
 
+    @Nullable
+    public KtStringTemplateExpression getStringTemplateExpression() {
+        KotlinPlaceHolderStub<? extends KtValueArgument> stub = getStub();
+        KtExpression expression;
+        if (stub != null) {
+            KtExpression[] stringTemplateExpressions = stub.getChildrenByType(STRING_TEMPLATE_EXPRESSIONS_TYPES, KtExpression.EMPTY_ARRAY);
+            expression = stringTemplateExpressions.length != 0 ? stringTemplateExpressions[0] : null;
+        } else {
+            expression = findChildByClass(KtExpression.class);
+        }
+        return expression instanceof KtStringTemplateExpression ? (KtStringTemplateExpression)expression : null;
+    }
+
     @Override
     @Nullable
     public KtValueArgumentName getArgumentName() {

idea/idea-test-framework/test/org/jetbrains/kotlin/idea/test/AstAccessControl.kt

@@ -6,20 +6,40 @@
 package org.jetbrains.kotlin.idea.test
 
 import com.intellij.openapi.Disposable
+import com.intellij.openapi.application.ApplicationManager
 import com.intellij.openapi.project.Project
 import com.intellij.openapi.vfs.VfsUtilCore
 import com.intellij.openapi.vfs.VirtualFile
 import com.intellij.openapi.vfs.VirtualFileFilter
+import com.intellij.psi.PsiFile
 import com.intellij.psi.PsiManager
 import com.intellij.psi.impl.PsiManagerImpl
 import com.intellij.testFramework.fixtures.CodeInsightTestFixture
 import org.jetbrains.kotlin.idea.KotlinFileType
+import org.jetbrains.kotlin.idea.test.KotlinLightCodeInsightFixtureTestCaseBase.*
+import org.jetbrains.kotlin.psi.KtFile
 import org.jetbrains.kotlin.test.InTextDirectivesUtils
 import kotlin.test.fail
 
 object AstAccessControl {
     val ALLOW_AST_ACCESS_DIRECTIVE: String = "ALLOW_AST_ACCESS"
 
+    @JvmStatic
+    fun dropPsiAndTestWithControlledAccessToAst(
+        shouldFail: Boolean,
+        file: KtFile,
+        disposable: Disposable,
+        testBody: () -> Unit
+    ) {
+        testWithControlledAccessToAst(shouldFail, file.project, disposable) {
+            // clean up AST
+            ApplicationManager.getApplication().runWriteAction { file.onContentReload() }
+            assertNotNull("file is parsed from AST", file.stub)
+
+            testBody()
+        }
+    }
+
     // Please provide at least one test that fails ast switch check (shouldFail should be true for at least one test)
     // This kind of inconvenience is justified by the fact that the check can be invalidated by slight misconfiguration of the test
     // leading to all tests passing

idea/tests/org/jetbrains/kotlin/fileClasses/JvmFileClassUtilTest.kt

@@ -0,0 +1,70 @@
+/*
+ * Copyright 2010-2021 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.kotlin.fileClasses
+
+import com.intellij.psi.PsiElement
+import com.intellij.psi.stubs.StubElement
+import com.intellij.psi.tree.IElementType
+import org.jetbrains.kotlin.idea.test.AstAccessControl.dropPsiAndTestWithControlledAccessToAst
+import org.jetbrains.kotlin.idea.test.KotlinLightCodeInsightFixtureTestCaseBase
+import org.jetbrains.kotlin.psi.KtAnnotationEntry
+import org.jetbrains.kotlin.psi.KtFile
+import org.jetbrains.kotlin.psi.stubs.elements.KtStubElementTypes
+
+class JvmFileClassUtilTest : KotlinLightCodeInsightFixtureTestCaseBase() {
+    fun testStubAccessWithACorruptedJvmName() {
+        doTestJvmNameStubAccess(
+            """
+            <!ILLEGAL_JVM_NAME!>@JvmName(<!NO_VALUE_FOR_PARAMETER!>)<!><!>
+            fun foo() {}
+        """.trimIndent(), null
+        )
+    }
+
+    fun testStubAccessWithAValidJvmName() {
+        doTestJvmNameStubAccess(
+            """
+            @JvmName("bar")
+            fun foo() {}
+        """.trimIndent(), "bar"
+        )
+    }
+
+    private fun doTestJvmNameStubAccess(content: String, expected: String?) {
+        val ktFile = myFixture.configureByText("jvmName.kt", content) as KtFile
+        assertNull("file is parsed from AST", ktFile.stub)
+
+        dropPsiAndTestWithControlledAccessToAst(true, ktFile, testRootDisposable) {
+            val annotationEntries =
+                ktFile.findDescendantStubChildrenByType<KtAnnotationEntry>(KtStubElementTypes.ANNOTATION_ENTRY)
+            assertTrue(annotationEntries.all { it.stub != null })
+            assertEquals(1, annotationEntries.size)
+
+            with(JvmFileClassUtil.getLiteralStringFromAnnotation(annotationEntries.first())) {
+                expected?.run { assertEquals(expected, this) } ?: run { assertNull(this) }
+            }
+        }
+    }
+
+    inline fun <reified T : PsiElement> KtFile.findDescendantStubChildrenByType(elementType: IElementType): List<T> {
+        val array = emptyArray<T>()
+        val result = mutableListOf<T>()
+        stub?.forEachDescendantStubChildren { stubElement ->
+            stubElement.getChildrenByType(elementType, array).let { array ->
+                array.forEach { result += it }
+            }
+        }
+        return result
+    }
+
+    fun StubElement<*>.forEachDescendantStubChildren(action: (StubElement<*>) -> Unit) {
+        action(this)
+        childrenStubs.forEach {
+            it.forEachDescendantStubChildren(action)
+        }
+    }
+
+}

idea/tests/org/jetbrains/kotlin/idea/stubs/KotlinStubsTest.java

@@ -5,7 +5,6 @@
 
 package org.jetbrains.kotlin.idea.stubs;
 
-import com.intellij.openapi.application.ApplicationManager;
 import com.intellij.psi.PsiFile;
 import com.intellij.testFramework.LightProjectDescriptor;
 import com.intellij.testFramework.fixtures.LightCodeInsightFixtureTestCase;
@@ -22,6 +21,8 @@ import org.junit.runner.RunWith;
 
 import java.util.List;
 
+import static org.jetbrains.kotlin.idea.test.AstAccessControl.dropPsiAndTestWithControlledAccessToAst;
+
 @RunWith(JUnit3WithIdeaConfigurationRunner.class)
 public class KotlinStubsTest extends LightCodeInsightFixtureTestCase {
     @NotNull
@@ -56,12 +57,12 @@ public class KotlinStubsTest extends LightCodeInsightFixtureTestCase {
         assertEquals(1, astDeclarations.size());
         assertTrue(astDeclarations.get(0) instanceof KtScript);
 
-        // clean up AST
-        ApplicationManager.getApplication().runWriteAction(() -> ktFile.onContentReload());
+        dropPsiAndTestWithControlledAccessToAst(true, ktFile, getTestRootDisposable(), () -> {
+            List<KtDeclaration> stubDeclarations = ktFile.getDeclarations();
+            assertEquals(1, stubDeclarations.size());
+            assertTrue(stubDeclarations.get(0) instanceof KtScript);
 
-        assertNotNull("file is parsed from AST", ktFile.getStub());
-        List<KtDeclaration> stubDeclarations = ktFile.getDeclarations();
-        assertEquals(1, stubDeclarations.size());
-        assertTrue(stubDeclarations.get(0) instanceof KtScript);
+            return null;
+        });
     }
 }