public static @Nullable Module findModuleForPsiElement(@NotNull PsiElement element) {
    PsiFile containingFile = element.getContainingFile();
    PsiElement highestPsi = containingFile == null ? element : containingFile;
    if (!highestPsi.isValid()) {
      return null;
    }
    Project project = highestPsi.getProject();
    if (project.isDefault()) return null;
    ProjectFileIndex fileIndex = ProjectFileIndex.getInstance(project);

    if (element instanceof PsiFileSystemItem && (!(element instanceof PsiFile) || element.getContext() == null)) {
      VirtualFile vFile = ((PsiFileSystemItem)element).getVirtualFile();
      if (vFile == null) {
        vFile = containingFile == null ? null : containingFile.getOriginalFile().getVirtualFile();
        if (vFile == null) {
          return element.getUserData(KEY_MODULE);
        }
      }

      if (fileIndex.isInLibrary(vFile)) {
        List<OrderEntry> orderEntries = fileIndex.getOrderEntriesForFile(vFile);
        if (orderEntries.isEmpty()) {
          return null;
        }

        if (orderEntries.size() == 1 && orderEntries.get(0) instanceof LibraryOrSdkOrderEntry) {
          return orderEntries.get(0).getOwnerModule();
        }

        Optional<Module> module = orderEntries
          .stream()
          .filter(entry -> entry instanceof LibraryOrSdkOrderEntry)
          .map(OrderEntry::getOwnerModule)
          .min(ModuleManager.getInstance(project).moduleDependencyComparator());
        //there may be no LibraryOrSdkOrderEntry if the file is located under both module source root and a library root
        if (module.isPresent()) {
          return module.get();
        }
      }

      return fileIndex.getModuleForFile(vFile);
    }
    if (containingFile != null) {
      PsiElement context;
      while ((context = containingFile.getContext()) != null) {
        PsiFile file = context.getContainingFile();
        if (file == null) break;
        containingFile = file;
      }

      if (containingFile.getUserData(KEY_MODULE) != null) {
        return containingFile.getUserData(KEY_MODULE);
      }

      PsiFile originalFile = containingFile.getOriginalFile();
      if (originalFile.getUserData(KEY_MODULE) != null) {
        return originalFile.getUserData(KEY_MODULE);
      }

      VirtualFile virtualFile = originalFile.getVirtualFile();
      if (virtualFile != null) {
        return fileIndex.getModuleForFile(virtualFile);
      }
    }

    return element.getUserData(KEY_MODULE);
  }
public static @Nullable Module findModuleForPsiElement(@NotNull PsiElement element)