resolve & type inference for assignments (+=, -= etc)

plugins/groovy/src/org/jetbrains/plugins/groovy/lang/psi/api/statements/expressions/GrAssignmentExpression.java

@@ -16,22 +16,30 @@
 
 package org.jetbrains.plugins.groovy.lang.psi.api.statements.expressions;
 
+import com.intellij.psi.PsiElement;
+import com.intellij.psi.PsiPolyVariantReference;
 import com.intellij.psi.tree.IElementType;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
+import org.jetbrains.plugins.groovy.lang.psi.api.GroovyResolveResult;
 
 /**
  * @author ilyas
  */
-public interface GrAssignmentExpression extends GrExpression {
+public interface GrAssignmentExpression extends GrExpression, PsiPolyVariantReference {
 
   boolean isTupleAssignment();
 
   @NotNull
-  public GrExpression getLValue();
+  GrExpression getLValue();
 
   @Nullable
-  public GrExpression getRValue();
+  GrExpression getRValue();
 
   IElementType getOperationToken();
+
+  @NotNull
+  GroovyResolveResult[] multiResolve(boolean incompleteCode);
+
+  PsiElement getOpToken();
 }

plugins/groovy/src/org/jetbrains/plugins/groovy/lang/psi/dataFlow/types/TypeDfaInstance.java

@@ -72,16 +72,12 @@ public class TypeDfaInstance implements DfaInstance<Map<String, PsiType>> {
     if (element instanceof GrReferenceExpression && ((GrReferenceExpression) element).getQualifierExpression() == null) {
       final PsiElement parent = element.getParent();
       if (parent instanceof GrAssignmentExpression) {
-        final GrExpression initializer = ((GrAssignmentExpression)parent).getRValue();
-        if (initializer != null) {
-          return new Computable<PsiType>() {
-            @Nullable
-            public PsiType compute() {
-              return initializer.getType();
-            }
-          };
-        }
-        return null;
+        return new Computable<PsiType>() {
+          @Nullable
+          public PsiType compute() {
+            return ((GrAssignmentExpression)parent).getType();
+          }
+        };
       }
 
       if (parent instanceof GrTupleExpression) {

plugins/groovy/src/org/jetbrains/plugins/groovy/lang/psi/impl/statements/expressions/GrAssignmentExpressionImpl.java

@@ -17,21 +17,29 @@
 package org.jetbrains.plugins.groovy.lang.psi.impl.statements.expressions;
 
 import com.intellij.lang.ASTNode;
-import com.intellij.psi.PsiElement;
-import com.intellij.psi.PsiType;
-import com.intellij.psi.PsiVariable;
-import com.intellij.psi.ResolveState;
+import com.intellij.openapi.util.TextRange;
+import com.intellij.psi.*;
+import com.intellij.psi.impl.source.resolve.ResolveCache;
 import com.intellij.psi.scope.PsiScopeProcessor;
 import com.intellij.psi.tree.IElementType;
+import com.intellij.util.ArrayUtil;
+import com.intellij.util.Function;
+import com.intellij.util.IncorrectOperationException;
+import com.intellij.util.NullableFunction;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
+import org.jetbrains.plugins.groovy.lang.lexer.GroovyTokenTypes;
 import org.jetbrains.plugins.groovy.lang.lexer.TokenSets;
 import org.jetbrains.plugins.groovy.lang.psi.GroovyElementVisitor;
+import org.jetbrains.plugins.groovy.lang.psi.api.GroovyResolveResult;
 import org.jetbrains.plugins.groovy.lang.psi.api.statements.GrTupleDeclaration;
 import org.jetbrains.plugins.groovy.lang.psi.api.statements.expressions.GrAssignmentExpression;
 import org.jetbrains.plugins.groovy.lang.psi.api.statements.expressions.GrExpression;
 import org.jetbrains.plugins.groovy.lang.psi.api.statements.expressions.GrReferenceExpression;
 import org.jetbrains.plugins.groovy.lang.psi.api.statements.typedef.members.GrAccessorMethod;
+import org.jetbrains.plugins.groovy.lang.psi.impl.GroovyPsiManager;
+import org.jetbrains.plugins.groovy.lang.psi.impl.PsiImplUtil;
+import org.jetbrains.plugins.groovy.lang.psi.util.PsiUtil;
 import org.jetbrains.plugins.groovy.lang.resolve.processors.ResolverProcessor;
 
 /**
@@ -39,6 +47,31 @@ import org.jetbrains.plugins.groovy.lang.resolve.processors.ResolverProcessor;
  */
 public class GrAssignmentExpressionImpl extends GrExpressionImpl implements GrAssignmentExpression {
 
+  private static final Function<GrAssignmentExpressionImpl, PsiType> TYPE_CALCULATOR =
+    new NullableFunction<GrAssignmentExpressionImpl, PsiType>() {
+      @Override
+      public PsiType fun(GrAssignmentExpressionImpl assignment) {
+        final GroovyResolveResult[] results = assignment.multiResolve(false);
+
+        if (results.length == 0) {
+          final GrExpression rValue = assignment.getRValue();
+          return rValue == null ? null : rValue.getType();
+        }
+
+
+        PsiType returnType = null;
+        final PsiManager manager = assignment.getManager();
+        for (GroovyResolveResult result : results) {
+          final PsiElement element = result.getElement();
+          if (element instanceof PsiMethod) {
+            final PsiType substituted = result.getSubstitutor().substitute(PsiUtil.getSmartReturnType((PsiMethod)element));
+            returnType = TypesUtil.getLeastUpperBoundNullable(returnType, substituted, manager);
+          }
+        }
+        return returnType;
+      }
+    };
+
   public GrAssignmentExpressionImpl(@NotNull ASTNode node) {
     super(node);
   }
@@ -66,11 +99,15 @@ public class GrAssignmentExpressionImpl extends GrExpressionImpl implements GrAs
   }
 
   public IElementType getOperationToken() {
-    return findNotNullChildByType(TokenSets.ASSIGN_OP_SET).getNode().getElementType();
+    return getOpToken().getNode().getElementType();
+  }
+
+  public PsiElement getOpToken() {
+    return findNotNullChildByType(TokenSets.ASSIGN_OP_SET);
   }
 
   public PsiType getType() {
-    return getLValue().getType();
+    return GroovyPsiManager.getInstance(getProject()).getType(this, TYPE_CALCULATOR);
   }
 
   public boolean processDeclarations(@NotNull PsiScopeProcessor processor, @NotNull ResolveState state, PsiElement lastParent, @NotNull PsiElement place) {
@@ -102,4 +139,87 @@ public class GrAssignmentExpressionImpl extends GrExpressionImpl implements GrAs
   public void accept(GroovyElementVisitor visitor) {
     visitor.visitAssignmentExpression(this);
   }
+
+  @NotNull
+  @Override
+  public GroovyResolveResult[] multiResolve(boolean incompleteCode) {
+    return (GroovyResolveResult[])getManager().getResolveCache().resolveWithCaching(this, RESOLVER, false, incompleteCode);
+  }
+
+  @Override
+  public PsiElement getElement() {
+    return this;
+  }
+
+  @Override
+  public TextRange getRangeInElement() {
+    final PsiElement token = getOpToken();
+    assert token != null;
+    final int offset = token.getStartOffsetInParent();
+    return new TextRange(offset, offset + token.getTextLength());
+  }
+
+  @Override
+  public PsiElement resolve() {
+    return PsiImplUtil.extractUniqueElement(multiResolve(false));
+  }
+
+  @NotNull
+  @Override
+  public String getCanonicalText() {
+    return getText();
+  }
+
+  @Override
+  public PsiElement handleElementRename(String newElementName) throws IncorrectOperationException {
+    throw new IncorrectOperationException("assignment expression cannot be renamed");
+  }
+
+  @Override
+  public PsiElement bindToElement(@NotNull PsiElement element) throws IncorrectOperationException {
+    throw new IncorrectOperationException("assignment expression cannot be bound to anything");
+  }
+
+  @Override
+  public boolean isReferenceTo(PsiElement element) {
+    return getManager().areElementsEquivalent(resolve(), element);
+  }
+
+  @NotNull
+  @Override
+  public Object[] getVariants() {
+    return ArrayUtil.EMPTY_OBJECT_ARRAY;
+  }
+
+  @Override
+  public boolean isSoft() {
+    return false;
+  }
+
+  @Override
+  public PsiReference getReference() {
+    final IElementType operationToken = getOperationToken();
+    if (operationToken == GroovyTokenTypes.mASSIGN) return null;
+
+    return this;
+  }
+
+  private static final ResolveCache.PolyVariantResolver<GrAssignmentExpressionImpl> RESOLVER =
+    new ResolveCache.PolyVariantResolver<GrAssignmentExpressionImpl>() {
+      @Override
+      public GroovyResolveResult[] resolve(GrAssignmentExpressionImpl assignmentExpression, boolean incompleteCode) {
+        final IElementType opType = assignmentExpression.getOperationToken();
+        if (opType == null || opType == GroovyTokenTypes.mASSIGN) return GroovyResolveResult.EMPTY_ARRAY;
+
+        final PsiType lType = assignmentExpression.getLValue().getType();
+        if (lType == null) return GroovyResolveResult.EMPTY_ARRAY;
+
+        final GrExpression rightOperand = assignmentExpression.getRValue();
+        PsiType rType = rightOperand == null ? null : rightOperand.getType();
+
+        final IElementType operatorToken = TokenSets.ASSIGNMENTS_TO_OPERATORS.get(opType);
+        return TypesUtil.getOverloadedOperatorCandidates(lType, operatorToken, assignmentExpression, new PsiType[]{rType});
+      }
+    };
+
 }

plugins/groovy/test/org/jetbrains/plugins/groovy/lang/resolve/ResolveMethodTest.groovy

@@ -648,4 +648,10 @@ class Zoo {
   public void resoleAsType() {
     assertInstanceOf resolve("A.groovy"), GrMethod
   }
+
+  public void testPlusAssignment() {
+    final PsiElement resolved = resolve("A.groovy")
+    assertInstanceOf resolved, GrMethod
+    assertEquals("plus", resolved.name)
+  }
 }

plugins/groovy/test/org/jetbrains/plugins/groovy/lang/resolve/TypeInferenceTest.java

@@ -181,4 +181,16 @@ public class TypeInferenceTest extends GroovyResolveTestCase {
   public void testInferWithClosureType() {
     assertTypeEquals("java.util.Date", "A.groovy");
   }
+
+  public void testPlusEquals1() {
+    assertTypeEquals("Test", "A.groovy");
+  }
+
+  public void testPlusEquals2() {
+    assertTypeEquals("java.lang.String", "A.groovy");
+  }
+
+  public void testPlusEquals3() {
+    assertTypeEquals("java.lang.String", "A.groovy");
+  }
 }

plugins/groovy/testdata/resolve/inference/plusEquals1/A.groovy

@@ -0,0 +1,9 @@
+class Test {
+    def plus(def a) {
+        "a"
+    }
+}
+def test = new Test()
+te<ref>st += 2
+print test
+

plugins/groovy/testdata/resolve/inference/plusEquals2/A.groovy

@@ -0,0 +1,9 @@
+class Test {
+    def plus(def a) {
+        "a"
+    }
+}
+def test = new Test()
+test += 2
+print te<ref>st
+

plugins/groovy/testdata/resolve/inference/plusEquals3/A.groovy

@@ -0,0 +1,9 @@
+class Test {
+    def plus(def a) {
+        "a"
+    }
+}
+def test = new Test()
+def aa = test += 2
+print test
+print a<ref>a

plugins/groovy/testdata/resolve/method/plusAssignment/A.groovy

@@ -0,0 +1,7 @@
+class Test {
+  def plus(def a) {
+    new Test()
+  }
+}
+def test=new Test()
+test+<ref>=2