java: propagate (completion) inference policy to nested sessions (IDEA-208607)

GitOrigin-RevId: 63b587a222e027c92ac3eebb1a5aeb79b669b17e

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

@@ -56,7 +56,7 @@ public class InferenceSession {
   private PsiSubstitutor myRestoreNameSubstitution = PsiSubstitutor.EMPTY;
   private MethodCandidateInfo myCurrentMethod;
 
-  public InferenceSession(InitialInferenceState initialState) {
+  public InferenceSession(InitialInferenceState initialState, ParameterTypeInferencePolicy policy) {
     myContext = initialState.getContext();
     myManager = myContext.getManager();
 
@@ -69,7 +69,7 @@ public class InferenceSession {
     }
     myInferenceSessionContainer = initialState.getInferenceSessionContainer();
     myErased = initialState.isErased();
-    myPolicy = DefaultParameterTypeInferencePolicy.INSTANCE;
+    myPolicy = policy;
   }
 
   public InferenceSession(PsiTypeParameter[] typeParams,
@@ -126,6 +126,11 @@ public class InferenceSession {
     myCurrentMethod = currentMethod;
   }
 
+  @NotNull
+  public ParameterTypeInferencePolicy getInferencePolicy() {
+    return myPolicy;
+  }
+
   public static PsiType createTypeParameterTypeWithUpperBound(@NotNull PsiType upperBound, @NotNull PsiElement place) {
     final PsiElementFactory elementFactory = JavaPsiFacade.getElementFactory(place.getProject());
 

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

@@ -128,7 +128,7 @@ public class InferenceSessionContainer {
     final CompoundInitialState compoundInitialState = createState(parentSession);
     InitialInferenceState initialInferenceState = compoundInitialState.getInitialState(parent);
     if (initialInferenceState != null) {
-      final InferenceSession childSession = new InferenceSession(initialInferenceState);
+      InferenceSession childSession = new InferenceSession(initialInferenceState, parentSession.getInferencePolicy());
       return childSession.collectAdditionalAndInfer(parameters, arguments, currentMethod, compoundInitialState.getInitialSubstitutor());
     }
 
@@ -165,7 +165,7 @@ public class InferenceSessionContainer {
 
               //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);
+              InferenceSession sessionInsideLambda = new InferenceSession(initialInferenceState, parentSession.getInferencePolicy());
               sessionInsideLambda.collectAdditionalAndInfer(methodParameters, argumentList.getExpressions(), ((MethodCandidateInfo)result), compoundInitialState.getInitialSubstitutor());
               return inferNested(parameters, arguments, parent, currentMethod, sessionInsideLambda);
             }

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

@@ -35,7 +35,7 @@ public class PsiGraphInferenceHelper implements PsiInferenceHelper {
       );
     }
     else {
-      final InferenceSession inferenceSession = new InferenceSession(new PsiTypeParameter[]{typeParameter}, partialSubstitutor, myManager, null);
+      InferenceSession inferenceSession = new InferenceSession(new PsiTypeParameter[]{typeParameter}, partialSubstitutor, myManager, null, policy);
       inferenceSession.initExpressionConstraints(parameters, arguments, null, false);
       substitutor = inferenceSession.infer();
     }

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

@@ -155,7 +155,7 @@ public class ExpressionCompatibilityConstraint extends InputOutputConstraintForm
 
       if (typeParams != null) {
         PsiSubstitutor siteSubstitutor = InferenceSession.chooseSiteSubstitutor(currentMethod, resolveResult, method);
-        final InferenceSession callSession = new InferenceSession(typeParams, siteSubstitutor, expression.getManager(), expression);
+        InferenceSession callSession = new InferenceSession(typeParams, siteSubstitutor, expression.getManager(), expression, session.getInferencePolicy());
         callSession.propagateVariables(session);
         if (method != null) {
           final PsiExpression[] args = argumentList.getExpressions();

java/java-tests/testData/codeInsight/completion/smartType/NestedCollectorsCounting-out.java

@@ -0,0 +1,10 @@
+import java.util.stream.*;
+
+class X {
+
+  {
+    Stream<String> stringStream = Stream.of("1", "2", "2");
+    stringStream.collect(Collectors.groupingBy(s -> s, Collectors.counting())<caret>);
+  }
+
+}

java/java-tests/testData/codeInsight/completion/smartType/NestedCollectorsCounting.java

@@ -0,0 +1,10 @@
+import java.util.stream.*;
+
+class X {
+
+  {
+    Stream<String> stringStream = Stream.of("1", "2", "2");
+    stringStream.collect(Collectors.groupingBy(s -> s, Collectors.count<caret>));
+  }
+
+}

java/java-tests/testSrc/com/intellij/java/codeInsight/completion/SmartType18CompletionTest.java

@@ -293,5 +293,7 @@ public void testConvertToObjectStream() {
     myFixture.assertPreferredCompletionItems(0, "String.class");
   }
 
+  public void testNestedCollectorsCounting() { doTest(false); }
+
   public void testFilterInaccessibleConstructors() { doAntiTest(); }
 }