From ac4072eda48125e63cdc67fcdf37aa8a047de957 Mon Sep 17 00:00:00 2001
From: Tagir Valeev <tagir.valeev@jetbrains.com>
Date: Fri, 26 Jan 2024 10:38:00 +0100
Subject: [PATCH] [java-injection] Support annotations on methods returning
 template processor (IDEA-333006)

GitOrigin-RevId: e08169d48cbd23db8cfaf2ee4466037bb70a2219
---
 .../intelliLang/JavaLanguageInjectionSupportTest.kt   | 11 +++++++----
 .../plugins/intelliLang/util/AnnotationUtilEx.java    |  8 ++++++--
 2 files changed, 13 insertions(+), 6 deletions(-)

diff --git a/plugins/IntelliLang/IntelliLang-tests/test/org/intellij/plugins/intelliLang/JavaLanguageInjectionSupportTest.kt b/plugins/IntelliLang/IntelliLang-tests/test/org/intellij/plugins/intelliLang/JavaLanguageInjectionSupportTest.kt
index 93c7bb79c533..9463f75e1fd8 100644
--- a/plugins/IntelliLang/IntelliLang-tests/test/org/intellij/plugins/intelliLang/JavaLanguageInjectionSupportTest.kt
+++ b/plugins/IntelliLang/IntelliLang-tests/test/org/intellij/plugins/intelliLang/JavaLanguageInjectionSupportTest.kt
@@ -256,16 +256,19 @@ class JavaLanguageInjectionSupportTest : AbstractLanguageInjectionTestCase() {
 
       class Hello {
       	void test(@Subst("ClassName") String name, String message) {
-        	@Language("JAVA")
-          String s = STR.""${'"'}
+          String s = javaProcessor().""${'"'}
                 class \{name} {
                     void main() {
                       System.out.println("\{message}");
                     }
                 }""${'"'};
       	}
-      }
-    """.trimIndent())
+            
+        @Language("JAVA")
+        private static StringTemplate.Processor<String, RuntimeException> javaProcessor() {
+          return STR;
+        }
+    }""".trimIndent())
     myFixture.checkHighlighting()
     injectionTestFixture.assertInjectedContent("""
       class ClassName {
diff --git a/plugins/IntelliLang/java-support/org/intellij/plugins/intelliLang/util/AnnotationUtilEx.java b/plugins/IntelliLang/java-support/org/intellij/plugins/intelliLang/util/AnnotationUtilEx.java
index f2b715578e81..12828a69f401 100644
--- a/plugins/IntelliLang/java-support/org/intellij/plugins/intelliLang/util/AnnotationUtilEx.java
+++ b/plugins/IntelliLang/java-support/org/intellij/plugins/intelliLang/util/AnnotationUtilEx.java
@@ -156,8 +156,12 @@ public final class AnnotationUtilEx {
         PsiElement target = ref.resolve();
         if (target instanceof PsiVariable var) {
           visitor.visitVariable(var);
-        } else if (target instanceof PsiMethod method) {
-          visitor.visitMethodReturnStatement(parent, method);
+        }
+      }
+      else if (templ.getProcessor() instanceof PsiCallExpression call && call != element) {
+        PsiMethod method = call.resolveMethod();
+        if (method != null) {
+          visitor.visitMethodReturnStatement(call, method);
         }
       }
     }