new inference: cache lambda functional type during expression constraints reduction (IDEA-121315)

java/java-psi-api/src/com/intellij/psi/LambdaUtil.java

@@ -33,6 +33,7 @@ import java.util.*;
  * Date: 7/17/12
  */
 public class LambdaUtil {
+  public static ThreadLocal<Map<PsiElement, PsiType>> ourFunctionTypes = new ThreadLocal<Map<PsiElement, PsiType>>();
   private static final Logger LOG = Logger.getInstance("#" + LambdaUtil.class.getName());
   @NonNls public static final String JAVA_LANG_FUNCTIONAL_INTERFACE = "java.lang.FunctionalInterface";
 
@@ -318,6 +319,13 @@ public class LambdaUtil {
               final int finalLambdaIdx = adjustLambdaIdx(lambdaIdx, (PsiMethod)resolve, parameters);
               if (finalLambdaIdx < parameters.length) {
                 if (!tryToSubstitute) return getNormalizedType(parameters[finalLambdaIdx]);
+                final Map<PsiElement, PsiType> map = ourFunctionTypes.get();
+                if (map != null) {
+                  final PsiType type = map.get(expression);
+                  if (type != null) {
+                    return type;
+                  }
+                }
                 return PsiResolveHelper.ourGraphGuard.doPreventingRecursion(expression, true, new Computable<PsiType>() {
                   @Override
                   public PsiType compute() {

java/java-psi-impl/src/com/intellij/psi/impl/source/resolve/graphInference/InferenceSession.java

@@ -823,15 +823,19 @@ public class InferenceSession {
         MethodCandidateInfo.updateSubstitutor(argumentList, substitutor);
       }
 
-      for (ConstraintFormula additionalConstraint : subset) {
-        additionalConstraint.apply(substitutor);
-      }
+      try {
+        for (ConstraintFormula additionalConstraint : subset) {
+          additionalConstraint.apply(substitutor);
+        }
 
-      myConstraints.addAll(subset);
-      if (!repeatInferencePhases(true)) {
-        return false;
+        myConstraints.addAll(subset);
+        if (!repeatInferencePhases(true)) {
+          return false;
+        }
+      }
+      finally {
+        LambdaUtil.ourFunctionTypes.set(null);
       }
-
     }
     return true;
   }

java/java-psi-impl/src/com/intellij/psi/impl/source/resolve/graphInference/constraints/InputOutputConstraintFormula.java

@@ -22,9 +22,7 @@ import com.intellij.psi.impl.source.resolve.graphInference.InferenceVariable;
 import com.intellij.psi.util.PsiUtil;
 import org.jetbrains.annotations.Nullable;
 
-import java.util.Collections;
-import java.util.HashSet;
-import java.util.Set;
+import java.util.*;
 
 /**
  * User: anna
@@ -110,5 +108,11 @@ public abstract class InputOutputConstraintFormula implements ConstraintFormula
   @Override
   public void apply(PsiSubstitutor substitutor) {
     setT(substitutor.substitute(getT()));
+    Map<PsiElement, PsiType> map = LambdaUtil.ourFunctionTypes.get();
+    if (map == null) {
+      map = new HashMap<PsiElement, PsiType>();
+      LambdaUtil.ourFunctionTypes.set(map);
+    }
+    map.put(getExpression(), getT());
   }
 }

java/java-tests/testData/codeInsight/daemonCodeAnalyzer/lambda/newLambda/IDEA121315.java

@@ -0,0 +1,22 @@
+import java.util.Map;
+import java.util.stream.Collectors;
+import java.util.stream.Stream;
+
+public class IDEA121315 {
+  class Issue {
+    Long getId() {
+      return 1l;
+    }
+  }
+
+  <T> T id(T i) {
+    return i;
+  }
+
+  void foo(Stream<Issue> map){
+    Map<Long, Issue> id2Issue = map.collect(Collectors.toMap(null, p -> id(p)));
+    Map<Long, Issue> id2Issue1 = map.collect(Collectors.toMap(null, p -> p));
+    Map<Long, Issue> id2Issue2 = map.collect(Collectors.toMap(null, this::id));
+
+  }
+}

java/java-tests/testSrc/com/intellij/codeInsight/daemon/lambda/NewLambdaHighlightingTest.java

@@ -146,6 +146,10 @@ public class NewLambdaHighlightingTest extends LightDaemonAnalyzerTestCase {
     doTest();
   }
 
+  public void testIDEA121315() throws Exception {
+    doTest();
+  }
+
   private void doTest() {
     doTest(false);
   }