private static @NotNull PsiElement getElementToHighlight(@NotNull PsiElement element) {
    PsiJavaCodeReferenceElement ref;
    if (element instanceof PsiNewExpression newExpression) {
      ref = newExpression.getClassReference();
    }
    else if (element instanceof PsiMethodCallExpression callExpression) {
      ref = callExpression.getMethodExpression();
    }
    else {
      return element;
    }
    if (ref != null) {
      PsiElement name = ref.getReferenceNameElement();
      return name != null ? name : ref;
    }
    return element;
  }
private static @NotNull PsiElement getElementToHighlight(@NotNull PsiElement element)