private static @NotNull Set<PsiClass> findDirectSubClassesInFile(@NotNull PsiClass sealedClass) {
    Set<PsiClass> subClasses = new HashSet<>();

    if (sealedClass.isEnum()) {
      for (PsiField field : sealedClass.getFields()) {
        if (field instanceof PsiEnumConstant) {
          ContainerUtil.addIfNotNull(subClasses, ((PsiEnumConstant)field).getInitializingClass());
        }
      }
      return subClasses;
    }

    sealedClass.getContainingFile().accept(new JavaElementVisitor() {
      @Override
      public void visitJavaFile(@NotNull PsiJavaFile file) {
        for (PsiClass psiClass : file.getClasses()) {
          visitClass(psiClass);
        }
      }

      @Override
      public void visitClass(@NotNull PsiClass psiClass) {
        for (PsiClass inner : psiClass.getInnerClasses()) {
          visitClass(inner);
        }
        if (psiClass.isInheritor(sealedClass, false)) {
          subClasses.add(psiClass);
        }
      }
    });
    return subClasses;
  }
private static @NotNull Set<PsiClass> findDirectSubClassesInFile(@NotNull PsiClass sealedClass)