calculate target type in the same way for statement/expressions lambdas

java/java-psi-api/src/com/intellij/psi/util/PsiTypesUtil.java

@@ -209,7 +209,7 @@ public class PsiTypesUtil {
     else if (parent instanceof PsiReturnStatement) {
       final PsiLambdaExpression lambdaExpression = PsiTreeUtil.getParentOfType(parent, PsiLambdaExpression.class);
       if (lambdaExpression != null) {
-        return LambdaUtil.getFunctionalInterfaceReturnType(lambdaExpression.getFunctionalInterfaceType());
+        return null;
       }
       else {
         PsiMethod method = PsiTreeUtil.getParentOfType(parent, PsiMethod.class);

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

@@ -525,11 +525,22 @@ public class InferenceSession {
       return getTargetType((PsiExpression)parent);
     }
     else if (parent instanceof PsiLambdaExpression) {
-      if (PsiUtil.skipParenthesizedExprUp(parent.getParent()) instanceof PsiExpressionList) {
-        final PsiType typeTypeByParentCall = getTargetType((PsiLambdaExpression)parent);
-        return LambdaUtil.getFunctionalInterfaceReturnType(FunctionalInterfaceParameterizationUtil.getGroundTargetType(typeTypeByParentCall, (PsiLambdaExpression)parent));
+      return getTargetTypeByContainingLambda((PsiLambdaExpression)parent);
+    }
+    else if (parent instanceof PsiReturnStatement) {
+      return getTargetTypeByContainingLambda(PsiTreeUtil.getParentOfType(parent, PsiLambdaExpression.class));
+    }
+    return null;
+  }
+
+  private static PsiType getTargetTypeByContainingLambda(PsiLambdaExpression lambdaExpression) {
+    if (lambdaExpression != null) {
+      if (PsiUtil.skipParenthesizedExprUp(lambdaExpression.getParent()) instanceof PsiExpressionList) {
+        final PsiType typeTypeByParentCall = getTargetType(lambdaExpression);
+        return LambdaUtil.getFunctionalInterfaceReturnType(
+          FunctionalInterfaceParameterizationUtil.getGroundTargetType(typeTypeByParentCall, lambdaExpression));
       }
-      return LambdaUtil.getFunctionalInterfaceReturnType(((PsiLambdaExpression)parent).getFunctionalInterfaceType());
+      return LambdaUtil.getFunctionalInterfaceReturnType(lambdaExpression.getFunctionalInterfaceType());
     }
     return null;
   }

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

@@ -0,0 +1,17 @@
+
+import java.util.List;
+import java.util.Optional;
+import java.util.function.Function;
+
+class SameCalls<ST> {
+  <B> List<B> bar(B a) {return null;}
+  <R> Optional<R> foo(Function<String, Optional<R>> computable) {return null;}
+
+  List<String> ff(SameCalls<String> sc){
+    return sc.foo((x) -> {
+      return Optional.of(bar(x));
+    }).get();
+  }
+
+
+}

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

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