diamonds: new expressions from return statements

java/java-tests/testData/codeInsight/daemonCodeAnalyzer/advHighlighting7/DiamondMisc.java

@@ -3,4 +3,43 @@ import java.util.List;
 
 class Test {
   List<String> queue = new ArrayList<>();
-}
+}
+
+class DD {
+    P1<P<String>> l = new L<String>() {
+        @Override
+        void f() {
+        }
+    };
+
+    P1<P<String>> l1 = new L<>();
+
+    P1<P<String>> foo() {
+        return new L<>();
+    }
+
+    String s = "";
+}
+
+class L<K> extends P1<P<K>> {
+    void f() {
+    }
+}
+
+class P1<P1T> extends P<P1T> {
+}
+
+class P<PT> {
+}
+
+
+class Test1 {
+  void bar() {
+    foo(new FF<>());
+  }
+
+  void foo(F<F<String>> p) {}
+}
+
+class FF<X> extends F<X>{}
+class F<T> {}

java/openapi/src/com/intellij/psi/PsiDiamondType.java

@@ -130,14 +130,30 @@ public class PsiDiamondType extends PsiType {
         }
       } else if (parent instanceof PsiVariable) {
         return getComponentTypes((PsiVariable)parent);
+      } else if (parent instanceof PsiReturnStatement) {
+        final PsiMethod containingMethod = PsiTreeUtil.getParentOfType(parent, PsiMethod.class);
+        if (containingMethod != null) {
+          final PsiType returnType = containingMethod.getReturnType();
+          if (returnType != null) {
+            final PsiExpression returnValue = ((PsiReturnStatement)parent).getReturnValue();
+            if (returnValue instanceof PsiNewExpression) {
+              return getComponentTypes(returnType, returnValue);
+            }
+            return getComponentTypes(returnType);
+          }
+        }
       }
     }
     return NULL_TYPES;
   }
 
   private static PsiType[] getComponentTypes(PsiVariable declaredElement) {
-    PsiType lType = declaredElement.getType();
+    final PsiType lType = declaredElement.getType();
     final PsiExpression initializer = declaredElement.getInitializer();
+    return getComponentTypes(lType, initializer);
+  }
+
+  private static PsiType[] getComponentTypes(PsiType lType, PsiExpression initializer) {
     if (initializer instanceof PsiNewExpression) {
       final PsiNewExpression newExpression = (PsiNewExpression)initializer;
       final PsiJavaCodeReferenceElement classReference = newExpression.getClassOrAnonymousClassReference();
@@ -145,7 +161,7 @@ public class PsiDiamondType extends PsiType {
         final String text = classReference.getReferenceName();
         if (text != null) {
           final PsiClass psiClass =
-            JavaPsiFacade.getInstance(declaredElement.getProject()).getResolveHelper().resolveReferencedClass(text, initializer);
+            JavaPsiFacade.getInstance(initializer.getProject()).getResolveHelper().resolveReferencedClass(text, initializer);
           final PsiType substitute = substitute(psiClass, lType);
           if (substitute != null) {
             lType = substitute;