From 41d910fb2d99d15248dd94104ff41a4da549205c Mon Sep 17 00:00:00 2001
From: Tagir Valeev <Tagir.Valeev@jetbrains.com>
Date: Mon, 24 Apr 2017 13:44:36 +0700
Subject: [PATCH] Fix race conditions in ConstantExpressionVisitor.

When the same condition is concurrently evaluated in two threads, it was possible to get it not evaluated in one of threads:
1. Shared PsiElement user data could be cleared by one thread before used by another
2. Subsequent getCached() calls in visitElement and elementFinished may return null and non-null if value was cached in-between by another thread

This should fix flacky DataFlowRangeAnalysisTest.testLongRangeKnownMethods test as it reason was "Long.MIN_VALUE + 1" is not always calculated to constant expression resulting in different CFG for the same code
---
 .../psi/impl/ConstantExpressionVisitor.java   | 24 ++++++++-----------
 .../impl/JavaConstantExpressionEvaluator.java |  7 ++++--
 2 files changed, 15 insertions(+), 16 deletions(-)

diff --git a/java/java-psi-impl/src/com/intellij/psi/impl/ConstantExpressionVisitor.java b/java/java-psi-impl/src/com/intellij/psi/impl/ConstantExpressionVisitor.java
index a9b3bf46df9f..32dbc20496cf 100644
--- a/java/java-psi-impl/src/com/intellij/psi/impl/ConstantExpressionVisitor.java
+++ b/java/java-psi-impl/src/com/intellij/psi/impl/ConstantExpressionVisitor.java
@@ -15,15 +15,16 @@
  */
 package com.intellij.psi.impl;
 
-import com.intellij.openapi.util.Key;
 import com.intellij.psi.*;
 import com.intellij.psi.tree.IElementType;
 import com.intellij.psi.util.ClassUtil;
 import com.intellij.psi.util.ConstantEvaluationOverflowException;
 import com.intellij.psi.util.ConstantExpressionUtil;
+import com.intellij.util.containers.HashMap;
 import com.intellij.util.containers.StringInterner;
 import gnu.trove.THashSet;
 
+import java.util.Map;
 import java.util.Set;
 import java.util.concurrent.ConcurrentMap;
 
@@ -32,6 +33,7 @@ class ConstantExpressionVisitor extends JavaElementVisitor implements PsiConstan
   private final StringInterner myInterner = new StringInterner();
 
   private Set<PsiVariable> myVisitedVars;
+  private Map<PsiElement, Object> myCachedValues = new HashMap<>();
   private final boolean myThrowExceptionOnOverflow;
 
   private Object myResult;
@@ -50,20 +52,13 @@ class ConstantExpressionVisitor extends JavaElementVisitor implements PsiConstan
     store(element, myResult);
     return myResult;
   }
-  private static final Key<Object> VALUE = Key.create("VALUE");
-  private static Object getStoredValue(PsiElement element) {
-    if (element == null) {
-      return null;
-    }
-    try {
-      return element.getUserData(VALUE);
-    }
-    finally {
-      element.putUserData(VALUE, null);
-    }
+
+  private Object getStoredValue(PsiElement element) {
+    return myCachedValues.remove(element);
   }
-  static void store(PsiElement element, Object value) {
-    element.putUserData(VALUE, value);
+
+  void store(PsiElement element, Object value) {
+    myCachedValues.put(element, value);
   }
 
   @Override
@@ -543,6 +538,7 @@ class ConstantExpressionVisitor extends JavaElementVisitor implements PsiConstan
       String name = ClassUtil.getJVMClassName((PsiClass)element);
       try {
         Class aClass = Class.forName(name);
+        //noinspection unchecked
         myResult = Enum.valueOf(aClass, constant);
       }
       catch (Throwable ignore) { }
diff --git a/java/java-psi-impl/src/com/intellij/psi/impl/JavaConstantExpressionEvaluator.java b/java/java-psi-impl/src/com/intellij/psi/impl/JavaConstantExpressionEvaluator.java
index 4d82d0195fc2..a9d78fda273b 100644
--- a/java/java-psi-impl/src/com/intellij/psi/impl/JavaConstantExpressionEvaluator.java
+++ b/java/java-psi-impl/src/com/intellij/psi/impl/JavaConstantExpressionEvaluator.java
@@ -64,6 +64,9 @@ public class JavaConstantExpressionEvaluator extends JavaRecursiveElementWalking
       Object result = myConstantExpressionVisitor.handle(element);
       cache(element, result);
     }
+    else {
+      myConstantExpressionVisitor.store(element, value == NO_VALUE ? null : value);
+    }
   }
 
   @Override
@@ -74,7 +77,7 @@ public class JavaConstantExpressionEvaluator extends JavaRecursiveElementWalking
       // will cache back in elementFinished()
     }
     else {
-      ConstantExpressionVisitor.store(element, value == NO_VALUE ? null : value);
+      myConstantExpressionVisitor.store(element, value == NO_VALUE ? null : value);
     }
   }
 
@@ -122,7 +125,7 @@ public class JavaConstantExpressionEvaluator extends JavaRecursiveElementWalking
         PsiElement operand = ((PsiPrefixExpression)expression).getOperand();
         if (operand == null) return null;
         Object value = evaluator.myConstantExpressionVisitor.handle(operand);
-        ConstantExpressionVisitor.store(operand, value);
+        evaluator.myConstantExpressionVisitor.store(operand, value);
       }
       return evaluator.myConstantExpressionVisitor.handle(expression);
     }