inference: don't fix the lambda return type for nested lambda processing (IDEA-169161)

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

@@ -112,7 +112,7 @@ public class InferenceSessionContainer {
           }
 
           if (session != null) {
-            final PsiSubstitutor childSubstitutor = inferNested(typeParameters, parameters, arguments, partialSubstitutor, (PsiCall)parent, policy, properties, session);
+            final PsiSubstitutor childSubstitutor = inferNested(parameters, arguments, (PsiCall)parent, properties, session);
             if (childSubstitutor != null) return childSubstitutor;
           }
           else if (topLevelCall instanceof PsiMethodCallExpression) {
@@ -127,14 +127,11 @@ public class InferenceSessionContainer {
     return inferenceSession.infer(parameters, arguments, parent);
   }
 
-  private static PsiSubstitutor inferNested(final PsiTypeParameter[] typeParameters,
-                                            @NotNull final PsiParameter[] parameters,
+  private static PsiSubstitutor inferNested(@NotNull final PsiParameter[] parameters,
                                             @NotNull final PsiExpression[] arguments,
-                                            final PsiSubstitutor partialSubstitutor,
                                             @NotNull final PsiCall parent,
-                                            @NotNull final ParameterTypeInferencePolicy policy,
                                             @NotNull final MethodCandidateInfo.CurrentCandidateProperties properties,
-                                            final InferenceSession parentSession) {
+                                            @NotNull final InferenceSession parentSession) {
     final CompoundInitialState compoundInitialState = createState(parentSession);
     InitialInferenceState initialInferenceState = compoundInitialState.getInitialState(parent);
     if (initialInferenceState != null) {
@@ -174,25 +171,12 @@ public class InferenceSessionContainer {
                 if (methodParameters.length == 0) {
                   break;
                 }
-                final PsiType parameterType = PsiTypesUtil.getParameterType(methodParameters, idx, true);
-                final PsiType parameterTypeInTermsOfSession = initialInferenceState.getInferenceSubstitutor().substitute(parameterType);
-                final PsiType lambdaTargetType = compoundInitialState.getInitialSubstitutor().substitute(parameterTypeInTermsOfSession);
-                if (call.equals(PsiTreeUtil.getParentOfType(parent, PsiCall.class, true))) {
-                  return LambdaUtil.performWithLambdaTargetType((PsiLambdaExpression)gParent, partialSubstitutor.substitute(lambdaTargetType),
-                                                                () -> {
-                                                                    //parent was mentioned in the top inference session
-                                                                    //just proceed with the target type
-                                                                    final InferenceSession inferenceSession = new InferenceSession(typeParameters, partialSubstitutor, parent.getManager(), parent, policy);
-                                                                    inferenceSession.initExpressionConstraints(parameters, arguments, parent);
-                                                                    return inferenceSession.infer(parameters, arguments, parent);
-                                                                });
-                }
 
                 //one of the grand parents were found in the top inference session
                 //start from it as it is the top level call
                 final InferenceSession sessionInsideLambda = new InferenceSession(initialInferenceState);
                 sessionInsideLambda.collectAdditionalAndInfer(methodParameters, argumentList.getExpressions(), ((MethodCandidateInfo)result).createProperties(), compoundInitialState.getInitialSubstitutor());
-                return inferNested(typeParameters, parameters, arguments, partialSubstitutor, parent, policy, properties, sessionInsideLambda);
+                return inferNested(parameters, arguments, parent, properties, sessionInsideLambda);
               }
             }
             else {

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

@@ -0,0 +1,117 @@
+
+import java.util.function.Function;
+
+
+import static java.util.function.Function.identity;
+
+class Example {
+
+  @FunctionalInterface
+  public interface Functor<A, F extends Functor> {
+    <B> Functor<B, F> fmap(Function<? super A, ? extends B> fn);
+  }
+
+  @FunctionalInterface
+  public interface NaturalTransformation<A, FA extends Functor<A, ?>, GA extends Functor<A, ?>> extends Function<FA, GA> {
+    static <A, FA extends Functor<A, ?>> NaturalTransformation<A, FA, FA> endo(
+      NaturalTransformation<A, FA, FA> naturalTransformation) {
+      return naturalTransformation;
+    }
+
+    static void main(String[] args) {
+      NaturalTransformation<Integer, NatF<Integer>, NatF<Integer>> happy =
+        natF -> natF.match(identity(), s -> NatF.s(s.carrier() + 1)); //no problems
+
+      // compiler exception reported on the following line
+      NaturalTransformation<Integer, NatF<Integer>, NatF<Integer>> unhappy =
+        NaturalTransformation.endo(natF -> natF.match(identity(), s -> NatF.s(s.carrier() + 1)));
+    }
+  }
+
+  public interface Coproduct2<A, B> {
+    <R> R match(Function<? super A, ? extends R> aFn, Function<? super B, ? extends R> bFn);
+  }
+
+  public static abstract class NatF<A> implements Functor<A, NatF>, Coproduct2<NatF.Z<A>, NatF.S<A>> {
+
+    private NatF() {
+    }
+
+    public static <A> NatF<A> z() {
+      return new Z<>();
+    }
+
+    public static <A> NatF<A> s(A a) {
+      return new S<>(a);
+    }
+
+    public static final class Z<A> extends NatF<A> {
+      private Z() {
+      }
+
+      @Override
+      @SuppressWarnings("unchecked")
+      public <B> Z<B> fmap(Function<? super A, ? extends B> fn) {
+        return (Z<B>) this;
+      }
+
+      @Override
+      public <R> R match(Function<? super Z<A>, ? extends R> aFn,
+                         Function<? super S<A>, ? extends R> bFn) {
+        return aFn.apply(this);
+      }
+    }
+
+    public static final class S<A> extends NatF<A> {
+      private final A a;
+
+      private S(A a) {
+        this.a = a;
+      }
+
+      public A carrier() {
+        return a;
+      }
+
+      @Override
+      public <B> S<B> fmap(Function<? super A, ? extends B> fn) {
+        return new S<>(fn.apply(carrier()));
+      }
+
+      @Override
+      public <R> R match(Function<? super Z<A>, ? extends R> aFn,
+                         Function<? super S<A>, ? extends R> bFn) {
+        return bFn.apply(this);
+      }
+    }
+  }
+}
+
+class ExampleSimplified {
+  static void main(String[] args) {
+    Function<Z<Integer>, Z<Integer>> zzFunction = t -> t;
+    Function<NatF<Integer>, NatF<Integer>> unhappy = endo(natF -> natF.match(zzFunction, s -> s1(s.carrier())));
+  }
+
+  static <FA> Function<FA, FA> endo(Function<FA, FA> function) {
+    return function;
+  }
+
+  public static <A1> NatF<A1> s1(A1 a) {
+    return null;
+  }
+
+  public static abstract class NatF<A> {
+    abstract <R> R match(Function<Z<A>, ? extends R> aFn,
+                         Function<S<A>, ? extends R> bFn);
+  }
+
+  public static abstract class Z<A> extends NatF<A> {}
+
+  public static abstract class S<A> extends NatF<A> {
+
+    public A carrier() {
+      return null;
+    }
+  }
+}

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

@@ -153,6 +153,7 @@ public class NewLambdaHighlightingTest extends LightDaemonAnalyzerTestCase {
   public void testCollectLambdaAdditionalConstraintsByGroundType() { doTest(); }
   public void testInferTypeParametersFromFunctionalInterfaceInputs() { doTest(); }
   public void testGroundTargetTypeWhenAbstractMethodInSuperclass() { doTest(); }
+  public void testNestedLambdasWithInferenceOfReturnTypeInTheLatestLambda() { doTest(); }
 
   private void doTest() {
     IdeaTestUtil.setTestVersion(JavaSdkVersion.JDK_1_8, getModule(), getTestRootDisposable());