new inference: session container (I)

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

@@ -65,6 +65,7 @@ public class InferenceSession {
   private final Set<InferenceVariable> myInferenceVariables = new LinkedHashSet<InferenceVariable>();
   private final List<ConstraintFormula> myConstraints = new ArrayList<ConstraintFormula>();
   private final Set<ConstraintFormula> myConstraintsCopy = new HashSet<ConstraintFormula>();
+  private final InferenceSessionContainer myInferenceSessionContainer = new InferenceSessionContainer();
 
   private PsiSubstitutor mySiteSubstitutor;
   private final PsiManager myManager;
@@ -79,12 +80,6 @@ public class InferenceSession {
   private final PsiElement myContext;
 
   private PsiSubstitutor myInferenceSubstitution = PsiSubstitutor.EMPTY;
-  private final Map<PsiElement, InferenceSession> myNestedSessions = new HashMap<PsiElement, InferenceSession>();
-  public void registerNestedSession(InferenceSession session) {
-    propagateVariables(session.getInferenceVariables());
-    myNestedSessions.put(session.getContext(), session);
-    myNestedSessions.putAll(session.myNestedSessions);
-  }
 
   public InferenceSession(PsiTypeParameter[] typeParams,
                           PsiType[] leftTypes, 
@@ -363,7 +358,7 @@ public class InferenceSession {
             }
           }
         }
-        final InferenceSession nestedCallSession = findNestedCallSession(arg);
+        final InferenceSession nestedCallSession = myInferenceSessionContainer.findNestedCallSession(arg, this);
         final PsiType parameterType =
           nestedCallSession.substituteWithInferenceVariables(getParameterType(parameters, i, siteSubstitutor, varargs));
         if (!isPertinentToApplicability(arg, parentMethod)) {
@@ -1210,7 +1205,7 @@ public class InferenceSession {
       final PsiExpression expression = ((ExpressionCompatibilityConstraint)formula).getExpression();
       final PsiCall callExpression = PsiTreeUtil.getParentOfType(expression, PsiCall.class, false);
       if (callExpression != null) {
-        final InferenceSession session = myNestedSessions.get(callExpression);
+        final InferenceSession session = myInferenceSessionContainer.findNestedCallSession(callExpression, null);
         if (session != null) {
           formula.apply(session.myInferenceSubstitution, true);
           collectVarsToResolve(varsToResolve, (InputOutputConstraintFormula)formula);
@@ -1682,12 +1677,8 @@ public class InferenceSession {
     return myInferenceSubstitution.substitute(type);
   }
 
-  public InferenceSession findNestedCallSession(PsiExpression arg) {
-    InferenceSession session = myNestedSessions.get(PsiTreeUtil.getParentOfType(arg, PsiCall.class));
-    if (session == null) {
-      session = this;
-    }
-    return session;
+  public InferenceSessionContainer getInferenceSessionContainer() {
+    return myInferenceSessionContainer;
   }
 
   public PsiType startWithFreshVars(PsiType type) {

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

@@ -0,0 +1,44 @@
+/*
+ * Copyright 2000-2015 JetBrains s.r.o.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.intellij.psi.impl.source.resolve.graphInference;
+
+import com.intellij.psi.PsiCall;
+import com.intellij.psi.PsiElement;
+import com.intellij.psi.util.PsiTreeUtil;
+import org.jetbrains.annotations.Contract;
+import org.jetbrains.annotations.Nullable;
+
+import java.util.HashMap;
+import java.util.Map;
+
+public class InferenceSessionContainer {
+  private final Map<PsiElement, InferenceSession> myNestedSessions = new HashMap<PsiElement, InferenceSession>();
+
+  public InferenceSessionContainer() {
+  }
+
+  public void registerNestedSession(InferenceSession targetSession, InferenceSession session) {
+    targetSession.propagateVariables(session.getInferenceVariables());
+    myNestedSessions.put(session.getContext(), session);
+    myNestedSessions.putAll(session.getInferenceSessionContainer().myNestedSessions);
+  }
+
+  @Contract("_, !null -> !null")
+  public InferenceSession findNestedCallSession(PsiElement arg, @Nullable InferenceSession defaultSession) {
+    InferenceSession session = myNestedSessions.get(PsiTreeUtil.getParentOfType(arg, PsiCall.class));
+    return session == null ? defaultSession : session;
+  }
+}

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

@@ -129,7 +129,7 @@ public class ExpressionCompatibilityConstraint extends InputOutputConstraintForm
           }
           callSession.registerReturnTypeConstraints(siteSubstitutor.substitute(returnType), myT);
           if (callSession.repeatInferencePhases(true)) {
-            session.registerNestedSession(callSession);
+            session.getInferenceSessionContainer().registerNestedSession(session, callSession);
           } else {
             return false;
           }

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

@@ -3,7 +3,6 @@ package com.intellij.psi.impl.source.resolve.graphInference.constraints;
 import com.intellij.psi.*;
 import com.intellij.psi.impl.source.resolve.graphInference.FunctionalInterfaceParameterizationUtil;
 import com.intellij.psi.impl.source.resolve.graphInference.InferenceSession;
-import com.intellij.psi.util.PsiTreeUtil;
 import com.intellij.psi.util.PsiUtil;
 
 import java.util.List;
@@ -63,7 +62,7 @@ public class LambdaExpressionCompatibilityConstraint implements ConstraintFormul
         if (!myExpression.isValueCompatible()) {
           return false;
         }
-        InferenceSession callsession = session.findNestedCallSession(myExpression);
+        InferenceSession callsession = session.getInferenceSessionContainer().findNestedCallSession(myExpression, session);
         returnType = callsession.substituteWithInferenceVariables(substitutor.substitute(returnType));
         if (!callsession.isProperType(returnType)) {
           for (PsiExpression returnExpression : returnExpressions) {