[groovy] Add inference of @ClosureParams annotation for explicitly-typed closure parameters

GitOrigin-RevId: 5e5091c45bcbddd104ac5c30c15cad70f419e058

plugins/groovy/src/org/jetbrains/plugins/groovy/intentions/style/inference/InferenceProcess.kt

@@ -3,11 +3,14 @@ package org.jetbrains.plugins.groovy.intentions.style.inference
 
 import com.intellij.psi.PsiTypeParameter
 import com.intellij.psi.search.SearchScope
+import org.jetbrains.plugins.groovy.intentions.closure.isClosureType
 import org.jetbrains.plugins.groovy.intentions.style.inference.driver.*
 import org.jetbrains.plugins.groovy.intentions.style.inference.graph.InferenceUnitGraph
 import org.jetbrains.plugins.groovy.intentions.style.inference.graph.createGraphFromInferenceVariables
 import org.jetbrains.plugins.groovy.intentions.style.inference.graph.determineDependencies
+import org.jetbrains.plugins.groovy.lang.psi.api.statements.params.GrParameter
 import org.jetbrains.plugins.groovy.lang.psi.api.statements.typedef.members.GrMethod
+import org.jetbrains.plugins.groovy.lang.psi.util.GroovyCommonClassNames
 
 /**
  * Performs full substitution for non-typed parameters of [method]
@@ -32,6 +35,13 @@ fun runInferenceProcess(method: GrMethod, scope: SearchScope): GrMethod {
   return instantiateTypeParameters(parameterizedDriver, inferredGraph, method, typeUsage)
 }
 
+private val forbiddenAnnotations =
+  setOf(GroovyCommonClassNames.GROOVY_LANG_DELEGATES_TO,
+        GroovyCommonClassNames.GROOVY_TRANSFORM_STC_CLOSURE_PARAMS)
+
+internal fun GrParameter.eligibleForExtendedInference(): Boolean =
+  typeElement == null || (type.isClosureType() && (annotations.map { it.qualifiedName } intersect forbiddenAnnotations).isEmpty())
+
 private fun createDriver(method: GrMethod,
                          scope: SearchScope): InferenceDriver {
   val virtualMethod = createVirtualMethod(method) ?: return EmptyDriver

plugins/groovy/src/org/jetbrains/plugins/groovy/intentions/style/inference/driver/CommonDriver.kt

@@ -54,7 +54,7 @@ class CommonDriver private constructor(private val targetParameters: Set<GrParam
                          scope: SearchScope): InferenceDriver {
       val elementFactory = GroovyPsiElementFactory.getInstance(virtualMethod.project)
       val targetParameters = setUpParameterMapping(method, virtualMethod)
-        .filter { it.key.typeElement == null }
+        .filter { it.key.eligibleForExtendedInference() }
         .map { it.value }
         .toSet()
       val typeParameters = mutableListOf<PsiTypeParameter>()

plugins/groovy/src/org/jetbrains/plugins/groovy/lang/psi/typeEnhancers/InferredClosureParamsEnhancer.kt

@@ -10,6 +10,7 @@ import com.intellij.psi.util.parentOfType
 import org.jetbrains.plugins.groovy.intentions.style.inference.CollectingGroovyInferenceSessionBuilder
 import org.jetbrains.plugins.groovy.intentions.style.inference.MethodParameterAugmenter
 import org.jetbrains.plugins.groovy.intentions.style.inference.driver.closure.compose
+import org.jetbrains.plugins.groovy.intentions.style.inference.eligibleForExtendedInference
 import org.jetbrains.plugins.groovy.lang.psi.api.GrFunctionalExpression
 import org.jetbrains.plugins.groovy.lang.psi.api.GroovyMethodResult
 import org.jetbrains.plugins.groovy.lang.psi.api.statements.expressions.GrCall
@@ -25,7 +26,7 @@ class InferredClosureParamsEnhancer : ClosureParamsEnhancer() {
   override fun getClosureParameterType(closureBlock: GrFunctionalExpression, index: Int): PsiType? {
     val methodCall = closureBlock.parentOfType<GrCall>() ?: return null
     val method = (methodCall.resolveMethod() as? GrMethod)
-    method?.takeIf { it.parameters.any { parameter -> parameter.typeElement == null } } ?: return null
+    method?.takeIf { it.parameters.any { parameter -> parameter.eligibleForExtendedInference() } } ?: return null
     val signatures = computeSignatures(methodCall, closureBlock, method) ?: return null
     val parameters = closureBlock.allParameters
     return signatures.singleOrNull { it.size == parameters.size }?.getOrNull(index)
@@ -58,7 +59,7 @@ class InferredClosureParamsEnhancer : ClosureParamsEnhancer() {
                                             virtualMethod: GrMethod): GrParameter? {
     val method = resolveResult.candidate?.method ?: return null
     val methodParameter = (resolveResult.candidate?.argumentMapping?.targetParameter(
-      ExpressionArgument(closureBlock))?.psi as? GrParameter)?.takeIf { it.typeElement == null } ?: return null
+      ExpressionArgument(closureBlock))?.psi as? GrParameter)?.takeIf { it.eligibleForExtendedInference() } ?: return null
     return virtualMethod.parameters.getOrNull(method.parameterList.getParameterIndex(methodParameter)) ?: return null
   }
 

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

@@ -132,6 +132,7 @@ def method(Box<A> box) {
   }
 
   void testOverloadedInClosure() {
+    RecursionManager.disableAssertOnRecursionPrevention(myFixture.testRootDisposable)
     testHighlighting '''
 def <T> void foo(T t, Closure cl) {}
 

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

@@ -167,6 +167,7 @@ def method(Box<A> box) {
   }
 
   void testOverloadedInClosure() {
+    RecursionManager.disableAssertOnRecursionPrevention(myFixture.testRootDisposable)
     testHighlighting '''
 def <T> void foo(T t, Closure cl) {}
 

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

@@ -86,6 +86,47 @@ def foo(a) {
 }
 
 foo(null as Number)
+''', JAVA_LANG_INTEGER
+  }
+
+
+  void 'test inference with explicit type'() {
+    doTest '''
+def foo(Closure a) {
+  a(1)
+}
+
+foo {<caret>it}
+''', JAVA_LANG_INTEGER
+  }
+
+  void 'test inference with explicit type with dependency on implicit type'() {
+    doTest '''
+def foo(a, Closure b) {
+  b(a)
+}
+
+foo(1) {<caret>it}
+''', JAVA_LANG_INTEGER
+  }
+
+  void 'test do not infer annotation for already annotated parameter'() {
+    doTest '''
+def foo(@DelegatesTo(Integer) Closure b) {
+  b(1)
+}
+
+foo {<caret>it}
+''', null
+  }
+
+  void 'test do infer annotation for already annotated parameter when annotation is trivial'() {
+    doTest '''
+def foo(@NotNull Closure b) {
+  b(1)
+}
+
+foo {<caret>it}
 ''', JAVA_LANG_INTEGER
   }
 }