From c5280f0411b3c86d615fca59afe5bca68e93f24f Mon Sep 17 00:00:00 2001
From: "Vladislav.Soroka" <Vladislav.Soroka@jetbrains.com>
Date: Tue, 6 Dec 2016 17:49:46 +0300
Subject: [PATCH] gradle: resolve closure type to delegate based on return
 method type

---
 .../gradle/service/resolve/GradleMiscContributor.kt | 13 ++++++++++++-
 1 file changed, 12 insertions(+), 1 deletion(-)

diff --git a/plugins/gradle/src/org/jetbrains/plugins/gradle/service/resolve/GradleMiscContributor.kt b/plugins/gradle/src/org/jetbrains/plugins/gradle/service/resolve/GradleMiscContributor.kt
index 27e991d14a20..3caa7c745bfc 100644
--- a/plugins/gradle/src/org/jetbrains/plugins/gradle/service/resolve/GradleMiscContributor.kt
+++ b/plugins/gradle/src/org/jetbrains/plugins/gradle/service/resolve/GradleMiscContributor.kt
@@ -18,12 +18,14 @@ package org.jetbrains.plugins.gradle.service.resolve
 import com.intellij.patterns.PsiJavaPatterns.psiElement
 import com.intellij.psi.PsiClass
 import com.intellij.psi.PsiElement
+import com.intellij.psi.PsiType
 import com.intellij.psi.ResolveState
 import com.intellij.psi.scope.PsiScopeProcessor
 import groovy.lang.Closure
 import org.jetbrains.plugins.gradle.service.resolve.GradleCommonClassNames.*
 import org.jetbrains.plugins.groovy.lang.psi.GroovyFile
 import org.jetbrains.plugins.groovy.lang.psi.api.statements.blocks.GrClosableBlock
+import org.jetbrains.plugins.groovy.lang.psi.api.statements.expressions.GrMethodCall
 import org.jetbrains.plugins.groovy.lang.psi.api.statements.expressions.path.GrMethodCallExpression
 import org.jetbrains.plugins.groovy.lang.psi.impl.GroovyPsiManager
 import org.jetbrains.plugins.groovy.lang.psi.impl.statements.expressions.TypesUtil
@@ -65,8 +67,8 @@ class GradleMiscContributor : GradleMethodContextContributor {
 //      return DelegatesToInfo(TypesUtil.createType("org.gradle.api.publish.PublicationContainer", closure), Closure.DELEGATE_FIRST)
 //    }
 
+    val parent = closure.parent
     if (domainCollectionWithTypeClosure.accepts(closure)) {
-      val parent = closure.parent
       if (parent is GrMethodCallExpression) {
         val psiElement = parent.argumentList.allArguments.singleOrNull()?.reference?.resolve()
         if (psiElement is PsiClass) {
@@ -74,6 +76,15 @@ class GradleMiscContributor : GradleMethodContextContributor {
         }
       }
     }
+
+    // resolve closure type to delegate based on return method type, e.g.
+    // FlatDirectoryArtifactRepository flatDir(Closure configureClosure)
+    if(parent is GrMethodCall) {
+      val psiType = parent.invokedExpression.type
+      if(psiType != null && psiType != PsiType.VOID) {
+        return DelegatesToInfo(psiType, Closure.DELEGATE_FIRST)
+      }
+    }
     return null
   }