From f32caada92c3045deefeec72817179cbc8a3d369 Mon Sep 17 00:00:00 2001
From: "Anna.Kozlova" <anna.kozlova@jetbrains.com>
Date: Tue, 20 Nov 2018 10:52:43 +0100
Subject: [PATCH] switch expression: inference for result expression fixed;
 more testdata on overload resolution

---
 .../src/com/intellij/psi/LambdaUtil.java         |  3 ++-
 .../graphInference/PsiPolyExpressionUtil.java    |  9 +++++++++
 .../advHighlighting12/SimpleInferenceCases.java  | 16 ++++++++++++++++
 3 files changed, 27 insertions(+), 1 deletion(-)

diff --git a/java/java-psi-api/src/com/intellij/psi/LambdaUtil.java b/java/java-psi-api/src/com/intellij/psi/LambdaUtil.java
index 6fcca9d9822c..9908ca0b4d6d 100644
--- a/java/java-psi-api/src/com/intellij/psi/LambdaUtil.java
+++ b/java/java-psi-api/src/com/intellij/psi/LambdaUtil.java
@@ -722,6 +722,7 @@ public class LambdaUtil {
                                                     PsiExpressionList.class,
                                                     PsiLambdaExpression.class,
                                                     PsiConditionalExpression.class,
+                                                    PsiSwitchExpression.class,
                                                     PsiCodeBlock.class,
                                                     PsiCall.class);
     while (true) {
@@ -750,7 +751,7 @@ public class LambdaUtil {
         }
       }
 
-      if (parent instanceof PsiConditionalExpression && !PsiPolyExpressionUtil.isPolyExpression((PsiExpression)parent)) {
+      if ((parent instanceof PsiConditionalExpression || parent instanceof PsiSwitchExpression) && !PsiPolyExpressionUtil.isPolyExpression((PsiExpression)parent)) {
         break;
       }
 
diff --git a/java/java-psi-api/src/com/intellij/psi/impl/source/resolve/graphInference/PsiPolyExpressionUtil.java b/java/java-psi-api/src/com/intellij/psi/impl/source/resolve/graphInference/PsiPolyExpressionUtil.java
index dba8832aae0e..3d03148448e8 100644
--- a/java/java-psi-api/src/com/intellij/psi/impl/source/resolve/graphInference/PsiPolyExpressionUtil.java
+++ b/java/java-psi-api/src/com/intellij/psi/impl/source/resolve/graphInference/PsiPolyExpressionUtil.java
@@ -3,6 +3,7 @@ package com.intellij.psi.impl.source.resolve.graphInference;
 
 import com.intellij.psi.*;
 import com.intellij.psi.infos.MethodCandidateInfo;
+import com.intellij.psi.util.PsiTreeUtil;
 import com.intellij.psi.util.PsiUtil;
 import com.intellij.psi.util.TypeConversionUtil;
 import org.jetbrains.annotations.NotNull;
@@ -124,9 +125,17 @@ public class PsiPolyExpressionUtil {
     return context instanceof PsiExpressionList ||
            context instanceof PsiArrayInitializerExpression ||
            context instanceof PsiConditionalExpression && (expr instanceof PsiCallExpression || isPolyExpression((PsiExpression)context)) ||
+           isSwitchExpressionAssignmentOrInvocationContext(expr) ||
            isAssignmentContext(expr, context);
   }
 
+  private static boolean isSwitchExpressionAssignmentOrInvocationContext(PsiExpression expr) {
+    PsiSwitchExpression switchExpression = PsiTreeUtil.getParentOfType(expr, PsiSwitchExpression.class);
+    return switchExpression  != null && 
+           PsiUtil.getSwitchResultExpressions(switchExpression).contains(expr) && 
+           isInAssignmentOrInvocationContext(switchExpression);
+  }
+
   private static boolean isAssignmentContext(PsiExpression expr, PsiElement context) {
     return PsiUtil.isCondition(expr, context) ||
            context instanceof PsiReturnStatement ||
diff --git a/java/java-tests/testData/codeInsight/daemonCodeAnalyzer/advHighlighting12/SimpleInferenceCases.java b/java/java-tests/testData/codeInsight/daemonCodeAnalyzer/advHighlighting12/SimpleInferenceCases.java
index 03373b102e65..bdb71811f137 100644
--- a/java/java-tests/testData/codeInsight/daemonCodeAnalyzer/advHighlighting12/SimpleInferenceCases.java
+++ b/java/java-tests/testData/codeInsight/daemonCodeAnalyzer/advHighlighting12/SimpleInferenceCases.java
@@ -1,11 +1,27 @@
+import java.util.function.*;
 class MyTest {
   <T> T foo(T t) {
     return t;
   }
+  <T> T foo(Supplier<T> t) {
+        return t.get();
+    }
+
+  <T> T foo(IntSupplier t) {
+    return null;
+  }
+
+  static <K> K bar() {
+    return null;
+  }
 
   void m(int i) {
     String s = foo(switch (i) {default -> "str";});
     String s1 = <error descr="Incompatible types. Required String but 'foo' was inferred to T:
 no instance(s) of type variable(s)  exist so that Object conforms to String">foo(switch (i) {case 1 -> new Object(); default -> "str";});</error>
+    String s2 =  foo(() -> switch (i) {
+            default -> "str";
+        });
+    String s3 = foo(() -> switch (i) {default -> bar();});
   }
 }
\ No newline at end of file