Make possible store constant values in stubs (KT-23738)

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

@@ -86,11 +86,11 @@ public interface KtNodeTypes {
     KtNodeType CONSTRUCTOR_DELEGATION_REFERENCE = new KtNodeType.KtLeftBoundNodeType("CONSTRUCTOR_DELEGATION_REFERENCE", KtConstructorDelegationReferenceExpression.class);
 
     // TODO: Not sure if we need separate NT for each kind of constants
-    KtNodeType NULL               = new KtNodeType("NULL", KtConstantExpression.class);
+    IElementType NULL               = KtStubElementTypes.NULL;
-    KtNodeType BOOLEAN_CONSTANT   = new KtNodeType("BOOLEAN_CONSTANT", KtConstantExpression.class);
+    IElementType BOOLEAN_CONSTANT   = KtStubElementTypes.BOOLEAN_CONSTANT;
-    KtNodeType FLOAT_CONSTANT     = new KtNodeType("FLOAT_CONSTANT", KtConstantExpression.class);
+    IElementType FLOAT_CONSTANT     = KtStubElementTypes.FLOAT_CONSTANT;
-    KtNodeType CHARACTER_CONSTANT = new KtNodeType("CHARACTER_CONSTANT", KtConstantExpression.class);
+    IElementType CHARACTER_CONSTANT = KtStubElementTypes.CHARACTER_CONSTANT;
-    KtNodeType INTEGER_CONSTANT   = new KtNodeType("INTEGER_CONSTANT", KtConstantExpression.class);
+    IElementType INTEGER_CONSTANT   = KtStubElementTypes.INTEGER_CONSTANT;
 
     KtNodeType STRING_TEMPLATE    = new KtNodeType("STRING_TEMPLATE", KtStringTemplateExpression.class);
     KtNodeType LONG_STRING_TEMPLATE_ENTRY = new KtNodeType("LONG_STRING_TEMPLATE_ENTRY", KtBlockStringTemplateEntry.class);

compiler/psi/src/org/jetbrains/kotlin/parsing/KotlinExpressionParsing.java

@@ -1873,7 +1873,7 @@ public class KotlinExpressionParsing extends AbstractKotlinParsing {
         literal.done(OBJECT_LITERAL);
     }
 
-    private void parseOneTokenExpression(KtNodeType type) {
+    private void parseOneTokenExpression(IElementType type) {
         PsiBuilder.Marker mark = mark();
         advance();
         mark.done(type);

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

@@ -17,15 +17,29 @@
 package org.jetbrains.kotlin.psi;
 
 import com.intellij.lang.ASTNode;
+import com.intellij.psi.PsiElement;
+import com.intellij.util.IncorrectOperationException;
 import org.jetbrains.annotations.NotNull;
+import org.jetbrains.kotlin.psi.stubs.KotlinConstantExpressionStub;
+import org.jetbrains.kotlin.psi.stubs.elements.KtConstantExpressionElementType;
 
-public class KtConstantExpression extends KtExpressionImpl {
+public class KtConstantExpression
+        extends KtElementImplStub<KotlinConstantExpressionStub> implements KtExpression {
     public KtConstantExpression(@NotNull ASTNode node) {
         super(node);
     }
 
+    public KtConstantExpression(@NotNull KotlinConstantExpressionStub stub) {
+        super(stub, KtConstantExpressionElementType.Companion.kindToConstantElementType(stub.kind()));
+    }
+
     @Override
     public <R, D> R accept(@NotNull KtVisitor<R, D> visitor, D data) {
         return visitor.visitConstantExpression(this, data);
     }
+
+    @Override
+    public PsiElement replace(@NotNull PsiElement newElement) throws IncorrectOperationException {
+        return KtExpressionImpl.Companion.replaceExpression(this, newElement, true, super::replace);
+    }
 }

compiler/psi/src/org/jetbrains/kotlin/psi/stubs/KotlinStubVersions.kt

@@ -23,7 +23,7 @@ object KotlinStubVersions {
     // Though only kotlin declarations (no code in the bodies) are stubbed, please do increase this version
     //      if you are not 100% sure it can be avoided.
     // Increasing this version will lead to reindexing of all kotlin source files on the first IDE startup with the new version.
-    const val SOURCE_STUB_VERSION = 130
+    const val SOURCE_STUB_VERSION = 131
 
     // Binary stub version should be increased if stub format (org.jetbrains.kotlin.psi.stubs.impl) is changed
     //      or changes are made to the core stub building code (org.jetbrains.kotlin.idea.decompiler.stubBuilder).

compiler/psi/src/org/jetbrains/kotlin/psi/stubs/StubInterfaces.kt

@@ -124,6 +124,19 @@ interface KotlinTypeParameterStub : KotlinStubWithFqName<KtTypeParameter> {
     fun isOutVariance(): Boolean
 }
 
+enum class ConstantValueKind {
+    NULL,
+    BOOLEAN_CONSTANT,
+    FLOAT_CONSTANT,
+    CHARACTER_CONSTANT,
+    INTEGER_CONSTANT
+}
+
+interface KotlinConstantExpressionStub : StubElement<KtConstantExpression> {
+    fun kind(): ConstantValueKind
+    fun value(): String
+}
+
 interface KotlinTypeProjectionStub : StubElement<KtTypeProjection> {
     fun getProjectionKind(): KtProjectionKind
 }

compiler/psi/src/org/jetbrains/kotlin/psi/stubs/elements/KtConstantExpressionElementType.kt

@@ -0,0 +1,88 @@
+/*
+ * Copyright 2010-2018 JetBrains s.r.o. 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.psi.stubs.elements
+
+import com.intellij.lang.ASTNode
+import com.intellij.psi.stubs.StubElement
+import com.intellij.psi.stubs.StubInputStream
+import com.intellij.psi.stubs.StubOutputStream
+import com.intellij.util.io.StringRef
+import org.jetbrains.annotations.NonNls
+import org.jetbrains.kotlin.psi.KtConstantExpression
+import org.jetbrains.kotlin.psi.stubs.ConstantValueKind
+import org.jetbrains.kotlin.psi.stubs.KotlinConstantExpressionStub
+import org.jetbrains.kotlin.psi.stubs.impl.KotlinConstantExpressionStubImpl
+
+class KtConstantExpressionElementType(@NonNls debugName: String) :
+    KtStubElementType<KotlinConstantExpressionStub, KtConstantExpression>(
+        debugName,
+        KtConstantExpression::class.java,
+        KotlinConstantExpressionStub::class.java
+    ) {
+
+    override fun shouldCreateStub(node: ASTNode): Boolean {
+        val parent = node.treeParent ?: return false
+        if (parent.elementType != KtStubElementTypes.VALUE_ARGUMENT) return false
+
+        return super.shouldCreateStub(node)
+    }
+
+    override fun createStub(psi: KtConstantExpression, parentStub: StubElement<*>?): KotlinConstantExpressionStub {
+        val elementType = psi.node.elementType as? KtConstantExpressionElementType
+            ?: throw IllegalStateException("Stub element type is expected for constant")
+
+        val value = psi.text ?: ""
+
+        return KotlinConstantExpressionStubImpl(
+            parentStub,
+            elementType,
+            constantElementTypeToKind(elementType),
+            StringRef.fromString(value)
+        )
+    }
+
+    override fun serialize(stub: KotlinConstantExpressionStub, dataStream: StubOutputStream) {
+        dataStream.writeInt(stub.kind().ordinal)
+        dataStream.writeName(stub.value())
+    }
+
+    override fun deserialize(dataStream: StubInputStream, parentStub: StubElement<*>?): KotlinConstantExpressionStub {
+        val kindOrdinal = dataStream.readInt()
+        val value = dataStream.readName() ?: StringRef.fromString("")
+
+        val valueKind = ConstantValueKind.values()[kindOrdinal]
+
+        return KotlinConstantExpressionStubImpl(
+            parentStub,
+            kindToConstantElementType(valueKind),
+            valueKind,
+            value
+        )
+    }
+
+    companion object {
+        fun kindToConstantElementType(kind: ConstantValueKind): KtConstantExpressionElementType {
+            return when (kind) {
+                ConstantValueKind.NULL -> KtStubElementTypes.NULL
+                ConstantValueKind.BOOLEAN_CONSTANT -> KtStubElementTypes.BOOLEAN_CONSTANT
+                ConstantValueKind.FLOAT_CONSTANT -> KtStubElementTypes.FLOAT_CONSTANT
+                ConstantValueKind.CHARACTER_CONSTANT -> KtStubElementTypes.CHARACTER_CONSTANT
+                ConstantValueKind.INTEGER_CONSTANT -> KtStubElementTypes.INTEGER_CONSTANT
+            }
+        }
+
+        private fun constantElementTypeToKind(elementType: KtConstantExpressionElementType): ConstantValueKind {
+            return when (elementType) {
+                KtStubElementTypes.NULL -> ConstantValueKind.NULL
+                KtStubElementTypes.BOOLEAN_CONSTANT -> ConstantValueKind.BOOLEAN_CONSTANT
+                KtStubElementTypes.INTEGER_CONSTANT -> ConstantValueKind.INTEGER_CONSTANT
+                KtStubElementTypes.FLOAT_CONSTANT -> ConstantValueKind.FLOAT_CONSTANT
+                KtStubElementTypes.CHARACTER_CONSTANT -> ConstantValueKind.CHARACTER_CONSTANT
+                else -> throw IllegalStateException("Unknown constant node type: $elementType")
+            }
+        }
+    }
+}

compiler/psi/src/org/jetbrains/kotlin/psi/stubs/elements/KtStubElementTypes.java

@@ -129,6 +129,12 @@ public interface KtStubElementTypes {
     KtPlaceHolderStubElementType<KtConstructorCalleeExpression> CONSTRUCTOR_CALLEE =
             new KtPlaceHolderStubElementType<>("CONSTRUCTOR_CALLEE", KtConstructorCalleeExpression.class);
 
+    KtConstantExpressionElementType NULL                = new KtConstantExpressionElementType("NULL");
+    KtConstantExpressionElementType BOOLEAN_CONSTANT    = new KtConstantExpressionElementType("BOOLEAN_CONSTANT");
+    KtConstantExpressionElementType FLOAT_CONSTANT      = new KtConstantExpressionElementType("FLOAT_CONSTANT");
+    KtConstantExpressionElementType CHARACTER_CONSTANT  = new KtConstantExpressionElementType("CHARACTER_CONSTANT");
+    KtConstantExpressionElementType INTEGER_CONSTANT    = new KtConstantExpressionElementType("INTEGER_CONSTANT");
+
     KtScriptElementType SCRIPT = new KtScriptElementType("SCRIPT");
 
     TokenSet DECLARATION_TYPES =

compiler/psi/src/org/jetbrains/kotlin/psi/stubs/impl/KotlinConstantExpressionStubImpl.kt

@@ -0,0 +1,24 @@
+/*
+ * Copyright 2010-2018 JetBrains s.r.o. 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.psi.stubs.impl
+
+import com.intellij.psi.PsiElement
+import com.intellij.psi.stubs.StubElement
+import com.intellij.util.io.StringRef
+import org.jetbrains.kotlin.psi.KtConstantExpression
+import org.jetbrains.kotlin.psi.stubs.ConstantValueKind
+import org.jetbrains.kotlin.psi.stubs.KotlinConstantExpressionStub
+import org.jetbrains.kotlin.psi.stubs.elements.KtConstantExpressionElementType
+
+class KotlinConstantExpressionStubImpl(
+    parent: StubElement<out PsiElement>?,
+    elementType: KtConstantExpressionElementType,
+    private val kind: ConstantValueKind,
+    private val value: StringRef
+) : KotlinStubBaseImpl<KtConstantExpression>(parent, elementType), KotlinConstantExpressionStub {
+    override fun kind(): ConstantValueKind = kind
+    override fun value(): String = StringRef.toString(value)
+}

idea/testData/stubs/AnnotationValues.expected

@@ -59,21 +59,29 @@ PsiJetFileStubImpl[package=test]
               REFERENCE_EXPRESSION[referencedName=Simple]
         VALUE_ARGUMENT_LIST
           VALUE_ARGUMENT
+            INTEGER_CONSTANT[kind=INTEGER_CONSTANT, value=12]
           VALUE_ARGUMENT
+            INTEGER_CONSTANT[kind=INTEGER_CONSTANT, value=12L]
           VALUE_ARGUMENT
+            INTEGER_CONSTANT[kind=INTEGER_CONSTANT, value=12]
           VALUE_ARGUMENT
+            FLOAT_CONSTANT[kind=FLOAT_CONSTANT, value=3.3]
           VALUE_ARGUMENT
             VALUE_ARGUMENT_NAME
               REFERENCE_EXPRESSION[referencedName=f]
+            FLOAT_CONSTANT[kind=FLOAT_CONSTANT, value=3.3F]
           VALUE_ARGUMENT
             VALUE_ARGUMENT_NAME
               REFERENCE_EXPRESSION[referencedName=c]
+            CHARACTER_CONSTANT[kind=CHARACTER_CONSTANT, value='a']
           VALUE_ARGUMENT
             VALUE_ARGUMENT_NAME
               REFERENCE_EXPRESSION[referencedName=b1]
+            BOOLEAN_CONSTANT[kind=BOOLEAN_CONSTANT, value=true]
           VALUE_ARGUMENT
             VALUE_ARGUMENT_NAME
               REFERENCE_EXPRESSION[referencedName=b2]
+            BOOLEAN_CONSTANT[kind=BOOLEAN_CONSTANT, value=false]
   CLASS[fqName=test.StringLiteral, isEnumEntry=false, isInterface=false, isLocal=false, isTopLevel=true, name=StringLiteral, superNames=[]]
     MODIFIER_LIST[annotation]
     PRIMARY_CONSTRUCTOR
@@ -156,8 +164,11 @@ PsiJetFileStubImpl[package=test]
               REFERENCE_EXPRESSION[referencedName=VarArg]
         VALUE_ARGUMENT_LIST
           VALUE_ARGUMENT
+            INTEGER_CONSTANT[kind=INTEGER_CONSTANT, value=1]
           VALUE_ARGUMENT
+            INTEGER_CONSTANT[kind=INTEGER_CONSTANT, value=2]
           VALUE_ARGUMENT
+            INTEGER_CONSTANT[kind=INTEGER_CONSTANT, value=3]
   CLASS[fqName=test.Arrays, isEnumEntry=false, isInterface=false, isLocal=false, isTopLevel=true, name=Arrays, superNames=[]]
     MODIFIER_LIST[annotation]
     PRIMARY_CONSTRUCTOR