IDEA-170943 Support @Delegate on methods

plugins/groovy/groovy-psi/src/org/jetbrains/plugins/groovy/GroovyBundle.properties

@@ -378,6 +378,7 @@ repetitive.method.name.0=Repetitive method name ''{0}''
 assign.expected='=' expected
 declared.type.0.have.to.extend.script=Declared type ''{0}'' does not extend ''groovy.lang.Script'' class
 base.script.annotation.is.allowed.only.inside.scripts=Annotation @BaseScript can only be used within a script
+delegate.annotation.is.only.for.methods.without.arguments=Annotation @Delegate couldn't be applied to method with arguments
 0.expressions.on.trait.fields.properties.are.not.supported.in.traits={0} expressions on trait fields/properties are not supported in traits
 only.traits.expected.here=Only traits are expected here
 anonymous.classes.cannot.be.created.from.traits=Anonymous classes cannot be created from traits

plugins/groovy/groovy-psi/src/org/jetbrains/plugins/groovy/annotator/checkers/DelegateAnnotationChecker.java

@@ -0,0 +1,46 @@
+/*
+ * Copyright 2000-2017 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 org.jetbrains.plugins.groovy.annotator.checkers;
+
+import com.intellij.lang.annotation.AnnotationHolder;
+import com.intellij.psi.PsiElement;
+import com.intellij.psi.PsiModifierList;
+import org.jetbrains.annotations.NotNull;
+import org.jetbrains.plugins.groovy.GroovyBundle;
+import org.jetbrains.plugins.groovy.lang.psi.api.auxiliary.modifiers.annotation.GrAnnotation;
+import org.jetbrains.plugins.groovy.lang.psi.api.statements.typedef.members.GrMethod;
+import org.jetbrains.plugins.groovy.lang.psi.util.GroovyCommonClassNames;
+
+/**
+ * Only methods without arguments could be annotated with @Delegate
+ */
+public class DelegateAnnotationChecker extends CustomAnnotationChecker {
+  @Override
+  public boolean checkApplicability(@NotNull AnnotationHolder holder, @NotNull GrAnnotation annotation) {
+    if (!GroovyCommonClassNames.GROOVY_LANG_DELEGATE.equals(annotation.getQualifiedName())) return false;
+    PsiElement annoParent = annotation.getParent();
+    PsiElement owner = annoParent instanceof PsiModifierList ? annoParent.getParent() : annoParent;
+    if (!(owner instanceof GrMethod)) {
+      return false;
+    }
+    if (((GrMethod)owner).getParameterList().getParametersCount() > 0) {
+      holder.createErrorAnnotation(annotation, GroovyBundle.message("delegate.annotation.is.only.for.methods.without.arguments"));
+      return true;
+    }
+
+    return false;
+  }
+}

plugins/groovy/groovy-psi/src/org/jetbrains/plugins/groovy/lang/psi/impl/GrAnnotationUtil.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2000-2016 JetBrains s.r.o.
+ * Copyright 2000-2017 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.
@@ -131,8 +131,9 @@ public class GrAnnotationUtil {
     return (PsiElement)owner;
   }
 
-  public static List<PsiClass> getClassArrayValue(@NotNull PsiAnnotation annotation, @NotNull String attributeName) {
-    PsiAnnotationMemberValue value = annotation.findAttributeValue(attributeName);
+  public static List<PsiClass> getClassArrayValue(@NotNull PsiAnnotation annotation, @NotNull String attributeName, boolean declared) {
+    PsiAnnotationMemberValue value =
+      declared ? annotation.findDeclaredAttributeValue(attributeName) : annotation.findAttributeValue(attributeName);
     if (value instanceof PsiArrayInitializerMemberValue) {
       return ContainerUtil.mapNotNull(((PsiArrayInitializerMemberValue)value).getInitializers(), GrAnnotationUtil::getPsiClass);
     }
@@ -144,8 +145,9 @@ public class GrAnnotationUtil {
     return Collections.emptyList();
   }
 
-  public static List<String> getStringArrayValue(@NotNull PsiAnnotation annotation, @NotNull String attributeName) {
-    PsiAnnotationMemberValue value = annotation.findAttributeValue(attributeName);
+  public static List<String> getStringArrayValue(@NotNull PsiAnnotation annotation, @NotNull String attributeName, boolean declared) {
+    PsiAnnotationMemberValue value =
+      declared ? annotation.findDeclaredAttributeValue(attributeName) : annotation.findAttributeValue(attributeName);
     if (value instanceof PsiArrayInitializerMemberValue) {
       return ContainerUtil.mapNotNull(((PsiArrayInitializerMemberValue)value).getInitializers(), memberValue -> {
         String string = getString(memberValue);

plugins/groovy/groovy-psi/src/org/jetbrains/plugins/groovy/lang/psi/util/GrTraitUtil.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2000-2016 JetBrains s.r.o.
+ * Copyright 2000-2017 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.
@@ -82,7 +82,7 @@ public class GrTraitUtil {
           PsiAnnotation annotation = AnnotationUtil.findAnnotation(clazz, "groovy.transform.SelfType");
           if (annotation != null) {
             result.addAll(
-              GrAnnotationUtil.getClassArrayValue(annotation, "value")
+              GrAnnotationUtil.getClassArrayValue(annotation, "value", false)
             );
           }
         }

plugins/groovy/groovy-psi/src/org/jetbrains/plugins/groovy/transformations/impl/DelegateTransformationSupport.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2000-2016 JetBrains s.r.o.
+ * Copyright 2000-2017 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.
@@ -45,19 +45,28 @@ import org.jetbrains.plugins.groovy.transformations.TransformationContext;
 import java.util.*;
 
 public class DelegateTransformationSupport implements AstTransformationSupport {
-
   @Override
   public void applyTransformation(@NotNull TransformationContext context) {
+    Map<PsiType, PsiAnnotation> declaredTypes = ContainerUtil.newLinkedHashMap();
     for (GrField field : context.getFields()) {
       final PsiAnnotation annotation = PsiImplUtil.getAnnotation(field, GroovyCommonClassNames.GROOVY_LANG_DELEGATE);
       if (annotation == null) continue;
+      declaredTypes.putIfAbsent(field.getDeclaredType(), annotation);
 
-      final PsiType type = field.getDeclaredType();
-      if (!(type instanceof PsiClassType)) continue;
+    }
+    for (PsiMethod method : context.getMethods()) {
+      final PsiAnnotation annotation = PsiImplUtil.getAnnotation(method, GroovyCommonClassNames.GROOVY_LANG_DELEGATE);
+      if (annotation == null) continue;
+      if (method.getParameterList().getParametersCount() > 0) continue;
+      declaredTypes.putIfAbsent(method.getReturnType(), annotation);
+    }
+
+    declaredTypes.forEach((type, annotation) -> {
+      if (!(type instanceof PsiClassType)) return;
 
       final PsiClassType.ClassResolveResult delegateResult = ((PsiClassType)type).resolveGenerics();
       final PsiClass delegate = delegateResult.getElement();
-      if (delegate == null) continue;
+      if (delegate == null) return;
 
       DelegateProcessor processor = new DelegateProcessor(context, delegate, annotation);
       delegate.processDeclarations(
@@ -67,7 +76,7 @@ public class DelegateTransformationSupport implements AstTransformationSupport {
         context.getCodeClass()
       );
 
-      if (!processor.myInterfaces) continue;
+      if (!processor.myInterfaces) return;
 
       Set<PsiClass> visited = ContainerUtil.newHashSet();
       Queue<Pair<PsiClass, PsiSubstitutor>> queue = ContainerUtil.newLinkedList(Pair.create(delegate, delegateResult.getSubstitutor()));
@@ -89,7 +98,7 @@ public class DelegateTransformationSupport implements AstTransformationSupport {
           }
         }
       }
-    }
+    });
   }
 
   private static class DelegateProcessor extends GrScopeProcessorWithHints {
@@ -197,17 +206,17 @@ public class DelegateTransformationSupport implements AstTransformationSupport {
         result = Conditions.and(result, method -> PsiImplUtil.getAnnotation(method, CommonClassNames.JAVA_LANG_DEPRECATED) == null);
       }
 
-      List<String> excludes = GrAnnotationUtil.getStringArrayValue(annotation, "excludes");
+      List<String> excludes = GrAnnotationUtil.getStringArrayValue(annotation, "excludes", true);
       if (!excludes.isEmpty()) {
         return Conditions.and(result, method -> !excludes.contains(method.getName()));
       }
 
-      List<String> includes = GrAnnotationUtil.getStringArrayValue(annotation, "includes");
+      List<String> includes = GrAnnotationUtil.getStringArrayValue(annotation, "includes", true);
       if (!includes.isEmpty()) {
         return Conditions.and(result, method -> includes.contains(method.getName()));
       }
 
-      List<PsiClass> excludeTypes = GrAnnotationUtil.getClassArrayValue(annotation, "excludeTypes");
+      List<PsiClass> excludeTypes = GrAnnotationUtil.getClassArrayValue(annotation, "excludeTypes", true);
       if (!excludeTypes.isEmpty()) {
         return Conditions.and(result, method -> {
           for (PsiClass excludeProvider : excludeTypes) {
@@ -219,7 +228,7 @@ public class DelegateTransformationSupport implements AstTransformationSupport {
         });
       }
 
-      List<PsiClass> includeTypes = GrAnnotationUtil.getClassArrayValue(annotation, "includeTypes");
+      List<PsiClass> includeTypes = GrAnnotationUtil.getClassArrayValue(annotation, "includeTypes", true);
       if (!includeTypes.isEmpty()) {
         return Conditions.and(result, method -> {
           for (PsiClass includeProvider : includeTypes) {

plugins/groovy/src/META-INF/plugin.xml

@@ -251,6 +251,7 @@
     <customAnnotationChecker implementation="org.jetbrains.plugins.groovy.annotator.checkers.GrAliasAnnotationChecker"/>
     <customAnnotationChecker implementation="org.jetbrains.plugins.groovy.annotator.checkers.AnnotationCollectorChecker" order="first"/>
     <customAnnotationChecker implementation="org.jetbrains.plugins.groovy.annotator.checkers.DelegatesToAnnotationChecker"/>
+    <customAnnotationChecker implementation="org.jetbrains.plugins.groovy.annotator.checkers.DelegateAnnotationChecker"/>
     <customAnnotationChecker implementation="org.jetbrains.plugins.groovy.griffon.GriffonPropertyListenerAnnotationChecker"/>
     <customAnnotationChecker implementation="org.jetbrains.plugins.groovy.annotator.checkers.BaseScriptAnnotationChecker"/>
 

plugins/groovy/test/org/jetbrains/plugins/groovy/lang/highlighting/GroovyHighlightingTest.groovy

@@ -349,6 +349,23 @@ class Baz implements I {
     myFixture.testHighlighting(false, false, false)
   }
 
+  void testMethodDelegateError() {
+    myFixture.configureByText('a.groovy','''
+class A {
+  def foo(){}
+}
+
+class B {
+  <error>@Delegate</error> A getA(int i){return new A()}
+}
+
+new B().fo<caret>o()
+''')
+
+    fixture.checkHighlighting()
+  }
+
+
   void testPrimitiveTypeParams() {
     myFixture.configureByText('a.groovy', '''\
 List<<error descr="Primitive type parameters are not allowed in type parameter list">int</error>> list = new ArrayList<int><EOLError descr="'(' expected"></EOLError>

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

@@ -1,5 +1,5 @@
 /*
- * Copyright 2000-2016 JetBrains s.r.o.
+ * Copyright 2000-2017 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.
@@ -51,6 +51,20 @@ class B {
   @Delegate A a
 }
 
+new B().fo<caret>o()
+''')
+  }
+
+  void testSimple2() {
+    doTest('''
+class A {
+  def foo(){}
+}
+
+class B {
+  @Delegate A getA(){return new A()}
+}
+
 new B().fo<caret>o()
 ''')
   }
@@ -111,6 +125,55 @@ class B {
   @Delegate A1 a1
 }
 
+new B().fo<caret>o()
+''')
+
+    def prototype = resolved.prototype as PsiMethod
+    def cc = prototype.containingClass
+    assertEquals 'A2', cc.name
+  }
+
+  void testSelectFirst3() {
+    def resolved = doTest('''
+class A1 {
+  def foo(){}
+}
+
+class A2 {
+  def foo(){}
+}
+
+
+class B {
+  @Delegate A2 bar()
+  @Delegate A1 bar2()
+}
+
+new B().fo<caret>o()
+''')
+
+    def prototype = resolved.prototype as PsiMethod
+    def cc = prototype.containingClass
+    assertEquals 'A2', cc.name
+  }
+
+  void testSelectFirst4() {
+    def resolved = doTest('''
+class A1 {
+  def foo(){}
+}
+
+class A2 {
+  def foo(){}
+}
+
+
+class B {
+  @Delegate A1 bar2()
+  @Delegate A2 bar // fields are processed before methods
+  
+}
+
 new B().fo<caret>o()
 ''')
 
@@ -279,6 +342,15 @@ class MyClass {
     @Delegate
     HashMap<String, Integer> map = new HashMap<String, Integer>()
 }
+''')
+  }
+
+  void 'test delegate method with generics'() {
+    assertAllMethodsImplemented('a.groovy', '''
+class MyClass {
+    @Delegate
+    HashMap<String, Integer> getMap() {return new HashMap<>}
+}
 ''')
   }
 }