Fix determination whether the class is final on light AST

GitOrigin-RevId: 6fb426544cb5b3accc2118388124fe1be012b2b9

java/java-analysis-impl/src/com/intellij/codeInspection/dataFlow/inference/ContractInferenceIndex.kt

@@ -37,7 +37,8 @@ private fun indexFile(tree: LighterAST): Map<Int, MethodData> {
   return visitor.result
 }
 
-internal data class ClassData(val hasSuper : Boolean, val hasPureInitializer : Boolean, val fieldModifiers : Map<String, LighterASTNode?>)
+internal data class ClassData(val hasSuper : Boolean, val hasPureInitializer : Boolean, val isFinal : Boolean, 
+                              val fieldModifiers : Map<String, LighterASTNode?>)
 
 private class InferenceVisitor(val tree : LighterAST) : RecursiveLighterASTNodeWalkingVisitor(tree) {
   var methodIndex = 0
@@ -61,10 +62,19 @@ private class InferenceVisitor(val tree : LighterAST) : RecursiveLighterASTNodeW
 
   private fun calcClassData(aClass: LighterASTNode) : ClassData {
     var hasSuper = aClass.tokenType == ANONYMOUS_CLASS
+    var isFinal = aClass.tokenType == ANONYMOUS_CLASS
     val fieldModifiers = HashMap<String, LighterASTNode?>()
     val initializers = ArrayList<LighterASTNode>()
     for (child in tree.getChildren(aClass)) {
       when(child.tokenType) {
+        JavaTokenType.RECORD_KEYWORD, JavaTokenType.ENUM_KEYWORD -> isFinal = true
+        MODIFIER_LIST -> {
+          isFinal = LightTreeUtil.firstChildOfType(tree, child, JavaTokenType.FINAL_KEYWORD) != null
+        }
+        ENUM_CONSTANT -> {
+          // We rely that enum constants go after ENUM_KEYWORD
+          isFinal = isFinal && LightTreeUtil.firstChildOfType(tree, child, ENUM_CONSTANT_INITIALIZER) == null
+        }
         EXTENDS_LIST -> {
           if (LightTreeUtil.firstChildOfType(tree, child, JAVA_CODE_REFERENCE) != null) {
             hasSuper = true
@@ -106,10 +116,11 @@ private class InferenceVisitor(val tree : LighterAST) : RecursiveLighterASTNodeW
         if (!pureInitializer) break
       }
     }
-    return ClassData(hasSuper, pureInitializer, fieldModifiers)
+    return ClassData(hasSuper, pureInitializer, isFinal, fieldModifiers)
   }
   
-  private fun getInferenceMode(method: LighterASTNode): JavaSourceInference.InferenceMode {
+  private fun getInferenceMode(method: LighterASTNode, clsData: ClassData?): JavaSourceInference.InferenceMode {
+    if (clsData?.isFinal == true) return JavaSourceInference.InferenceMode.ENABLED
     // PsiUtil#canBeOverridden logic on LighterAST
     val ctor = LightTreeUtil.firstChildOfType(tree, method, TYPE) == null
     if (ctor) return JavaSourceInference.InferenceMode.ENABLED
@@ -118,19 +129,14 @@ private class InferenceVisitor(val tree : LighterAST) : RecursiveLighterASTNodeW
     val isFinal = LightTreeUtil.firstChildOfType(tree, modifiers, JavaTokenType.FINAL_KEYWORD) != null
     val isPrivate = LightTreeUtil.firstChildOfType(tree, modifiers, JavaTokenType.PRIVATE_KEYWORD) != null
     if (isStatic || isFinal || isPrivate) return JavaSourceInference.InferenceMode.ENABLED
-    val aClass = tree.getParent(method)
-    if (aClass?.tokenType == ANONYMOUS_CLASS) return JavaSourceInference.InferenceMode.ENABLED
-    val classModifiers = LightTreeUtil.firstChildOfType(tree, aClass, MODIFIER_LIST)
-    if (LightTreeUtil.firstChildOfType(tree, classModifiers, JavaTokenType.FINAL_KEYWORD) != null) {
-      return JavaSourceInference.InferenceMode.ENABLED
-    }
     return JavaSourceInference.InferenceMode.PARAMETERS
   }
 
   private fun calcData(method: LighterASTNode): MethodData? {
     val body = LightTreeUtil.firstChildOfType(tree, method, CODE_BLOCK) ?: return null
     val parameterNames = getParameterNames(tree, method)
-    val inferenceMode = getInferenceMode(method)
+    val clsData = classData[tree.getParent(method)]
+    val inferenceMode = getInferenceMode(method, clsData)
     if (inferenceMode == JavaSourceInference.InferenceMode.PARAMETERS && parameterNames.isEmpty()) {
       return null
     }
@@ -139,7 +145,6 @@ private class InferenceVisitor(val tree : LighterAST) : RecursiveLighterASTNodeW
     if (inferenceMode == JavaSourceInference.InferenceMode.PARAMETERS) {
       return createData(body, emptyList(), null, null, notNullParams)
     }
-    val clsData = classData[tree.getParent(method)]
     val fieldMap = clsData?.fieldModifiers ?: emptyMap()
     // Constructor which has super classes may implicitly call impure super constructor, so don't infer purity for subclasses
     val ctor = LightTreeUtil.firstChildOfType(tree, method, TYPE) == null

java/java-tests/testSrc/com/intellij/java/codeInspection/PurityInferenceFromSourceTest.groovy

@@ -463,6 +463,32 @@ class Super {
     assert !JavaMethodContractUtil.isPure(clazz.methods[0])
     assert JavaMethodContractUtil.isPure(clazz.superClass.methods[0])
   }
+  
+  void "test enum method"() {
+    def clazz = myFixture.addClass """
+enum X {
+  A;
+  
+  Iterable<?> onXyz() {
+    return java.util.Collections.emptyList();
+  }
+}
+"""
+    assert JavaMethodContractUtil.isPure(clazz.methods[0])
+  }
+  
+  void "test enum method with subclass"() {
+    def clazz = myFixture.addClass """
+enum X {
+  A, B {};
+  
+  Iterable<?> onXyz() {
+    return java.util.Collections.emptyList();
+  }
+}
+"""
+    assert !JavaMethodContractUtil.isPure(clazz.methods[0])
+  }
 
   private void assertPure(String classBody) {
     assertMutationSignature(classBody, MutationSignature.pure())