java inference: support nested constraints inside code blocks (IDEA-237446)

GitOrigin-RevId: e865db7a88560615b1abca31b3ceefeb1ae76573

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

@@ -152,10 +152,7 @@ public class InferenceSessionContainer {
     //find the nearest parent which appears in the map and start inference with a provided target type for a nested lambda
     while (true) {
       if (gParent instanceof PsiReturnStatement) { //process code block lambda
-        final PsiElement returnContainer = gParent.getParent();
-        if (returnContainer instanceof PsiCodeBlock) {
-          gParent = returnContainer.getParent();
-        }
+        gParent = PsiTreeUtil.getParentOfType(gParent, PsiLambdaExpression.class);
       }
       if (gParent instanceof PsiConditionalExpression) {
         gParent = PsiUtil.skipParenthesizedExprUp(gParent.getParent());

java/java-tests/testData/codeInsight/daemonCodeAnalyzer/lambda/additionalConstraints/ConstraintsInsideNestedBlocks.java

@@ -0,0 +1,28 @@
+
+import java.util.List;
+import java.util.function.Function;
+
+class MyTest<O> {
+  private void test(MyTest<List<String>> listMono) {
+    expand(id -> listMono.flatMap(l -> {
+      if (l.size() == 0) {
+        return null;
+      }
+      else {
+        return MyTest.empty();
+      }
+    }));
+  }
+  void expand(Function<? super List<String>, ? extends MyTest<? extends List<String>>> expander) {
+
+  }
+
+  public final <R> MyTest<R> flatMap(Function<? super List<String>, ? extends MyTest<? extends R>> transformer) {
+    return null;
+  }
+
+
+  public static <T> MyTest<T> empty() {
+    return null;
+  }
+}

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

@@ -39,6 +39,10 @@ public class NewInferenceCollectingAdditionalConstraintsTest extends LightDaemon
     doTest();
   }
 
+  public void testConstraintsInsideNestedBlocks() {
+    doTest();
+  }
+
   public void testNestedLambdaExpressionsWithTopStandaloneMethod() {
     doTest();
   }