WalkPath: allow to skip dirs and to abort a walk

GitOrigin-RevId: 6be49d1c098b0d17b63ae1f26dbefb6904320d6b

tdutils/td/utils/port/path.cpp

@@ -79,11 +79,16 @@ Status mkpath(CSlice path, int32 mode) {
 }
 
 Status rmrf(CSlice path) {
-  return walk_path(path, [](CSlice path, bool is_dir) {
+  return walk_path(path, [](CSlice path, WalkPath::Type type) {
-    if (is_dir) {
+    switch (type) {
-      rmdir(path).ignore();
+      case WalkPath::Type::EnterDir:
-    } else {
+        break;
-      unlink(path).ignore();
+      case WalkPath::Type::ExitDir:
+        rmdir(path).ignore();
+        break;
+      case WalkPath::Type::NotDir:
+        unlink(path).ignore();
+        break;
     }
   });
 }
@@ -246,18 +251,16 @@ Result<string> mkdtemp(CSlice dir, Slice prefix) {
 }
 
 namespace detail {
-Status walk_path_dir(string &path, FileFd fd,
+using WalkFunction = std::function<WalkPath::Action(CSlice name, WalkPath::Type type)>;
-                     const std::function<void(CSlice name, bool is_directory)> &func) TD_WARN_UNUSED_RESULT;
+Result<bool> walk_path_dir(string &path, FileFd fd, const WalkFunction &func) TD_WARN_UNUSED_RESULT;
 
-Status walk_path_dir(string &path,
+Result<bool> walk_path_dir(string &path, const WalkFunction &func) TD_WARN_UNUSED_RESULT;
-                     const std::function<void(CSlice name, bool is_directory)> &func) TD_WARN_UNUSED_RESULT;
 
-Status walk_path_file(string &path,
+Result<bool> walk_path_file(string &path, const WalkFunction &func) TD_WARN_UNUSED_RESULT;
-                      const std::function<void(CSlice name, bool is_directory)> &func) TD_WARN_UNUSED_RESULT;
 
-Status walk_path(string &path, const std::function<void(CSlice name, bool is_directory)> &func) TD_WARN_UNUSED_RESULT;
+Result<bool> walk_path(string &path, const WalkFunction &func) TD_WARN_UNUSED_RESULT;
 
-Status walk_path_subdir(string &path, DIR *dir, const std::function<void(CSlice name, bool is_directory)> &func) {
+Result<bool> walk_path_subdir(string &path, DIR *dir, const WalkFunction &func) {
   while (true) {
     errno = 0;
     auto *entry = readdir(dir);
@@ -266,7 +269,7 @@ Status walk_path_subdir(string &path, DIR *dir, const std::function<void(CSlice
       return Status::PosixError(readdir_errno, "readdir");
     }
     if (entry == nullptr) {
-      return Status::OK();
+      return true;
     }
     Slice name = Slice(static_cast<const char *>(entry->d_name));
     if (name == "." || name == "..") {
@@ -280,7 +283,7 @@ Status walk_path_subdir(string &path, DIR *dir, const std::function<void(CSlice
     SCOPE_EXIT {
       path.resize(size);
     };
-    Status status;
+    Result<bool> status;
 #ifdef DT_DIR
     if (entry->d_type == DT_UNKNOWN) {
       status = walk_path(path, func);
@@ -293,22 +296,40 @@ Status walk_path_subdir(string &path, DIR *dir, const std::function<void(CSlice
 #warning "Slow walk_path"
     status = walk_path(path, func);
 #endif
-    if (status.is_error()) {
+    TRY_RESULT(is_ok, std::move(status));
-      return status;
+    if (!is_ok) {
+      return false;
     }
   }
 }
 
-Status walk_path_dir(string &path, DIR *subdir, const std::function<void(CSlice name, bool is_directory)> &func) {
+Result<bool> walk_path_dir(string &path, DIR *subdir, const WalkFunction &func) {
   SCOPE_EXIT {
     closedir(subdir);
   };
-  TRY_STATUS(walk_path_subdir(path, subdir, func));
+  switch (func(path, WalkPath::Type::EnterDir)) {
-  func(path, true);
+    case WalkPath::Action::Abort:
-  return Status::OK();
+      return false;
+    case WalkPath::Action::SkipDir:
+      return true;
+    case WalkPath::Action::Continue:
+      break;
+  }
+  TRY_RESULT(is_ok, walk_path_subdir(path, subdir, func));
+  if (!is_ok) {
+    return false;
+  }
+  switch (func(path, WalkPath::Type::ExitDir)) {
+    case WalkPath::Action::Abort:
+      return false;
+    case WalkPath::Action::SkipDir:
+    case WalkPath::Action::Continue:
+      break;
+  }
+  return true;
 }
 
-Status walk_path_dir(string &path, FileFd fd, const std::function<void(CSlice name, bool is_directory)> &func) {
+Result<bool> walk_path_dir(string &path, FileFd fd, const WalkFunction &func) {
   auto native_fd = fd.move_as_native_fd();
   auto *subdir = fdopendir(native_fd.fd());
   if (subdir == nullptr) {
@@ -318,7 +339,7 @@ Status walk_path_dir(string &path, FileFd fd, const std::function<void(CSlice na
   return walk_path_dir(path, subdir, func);
 }
 
-Status walk_path_dir(string &path, const std::function<void(CSlice name, bool is_directory)> &func) {
+Result<bool> walk_path_dir(string &path, const WalkFunction &func) {
   auto *subdir = opendir(path.c_str());
   if (subdir == nullptr) {
     return OS_ERROR(PSLICE() << tag("opendir", path));
@@ -326,12 +347,18 @@ Status walk_path_dir(string &path, const std::function<void(CSlice name, bool is
   return walk_path_dir(path, subdir, func);
 }
 
-Status walk_path_file(string &path, const std::function<void(CSlice name, bool is_directory)> &func) {
+Result<bool> walk_path_file(string &path, const WalkFunction &func) {
-  func(path, false);
+  switch (func(path, WalkPath::Type::NotDir)) {
-  return Status::OK();
+    case WalkPath::Action::Abort:
+      return false;
+    case WalkPath::Action::SkipDir:
+    case WalkPath::Action::Continue:
+      break;
+  }
+  return true;
 }
 
-Status walk_path(string &path, const std::function<void(CSlice name, bool is_directory)> &func) {
+Result<bool> walk_path(string &path, const WalkFunction &func) {
   TRY_RESULT(fd, FileFd::open(path, FileFd::Read));
   auto stat = fd.stat();
   bool is_dir = stat.is_dir_;
@@ -345,15 +372,16 @@ Status walk_path(string &path, const std::function<void(CSlice name, bool is_dir
     return walk_path_file(path, func);
   }
 
-  return Status::OK();
+  return true;
 }
 }  // namespace detail
 
-Status walk_path(CSlice path, const std::function<void(CSlice name, bool is_directory)> &func) {
+Status WalkPath::do_run(CSlice path, const detail::WalkFunction &func) {
   string curr_path;
   curr_path.reserve(PATH_MAX + 10);
   curr_path = path.c_str();
-  return detail::walk_path(curr_path, func);
+  TRY_STATUS(detail::walk_path(curr_path, func));
+  return Status::OK();
 }
 
 #endif
@@ -518,9 +546,17 @@ Result<std::pair<FileFd, string>> mkstemp(CSlice dir) {
   return Status::Error(PSLICE() << "Can't create temporary file \"" << file_pattern << '"');
 }
 
-static Status walk_path_dir(const std::wstring &dir_name,
+static Result<bool> walk_path_dir(const std::wstring &dir_name,
-                            const std::function<void(CSlice name, bool is_directory)> &func) {
+                                  const std::function<WalkPath::Action(CSlice name, WalkPath::Type type)> &func) {
   std::wstring name = dir_name + L"\\*";
+  switch (func(entry_name, WalkPath::Type::EnterDir)) {
+    case WalkPath::Action::Abort:
+      return false;
+    case WalkPath::Action::SkipDir:
+      return true;
+    case WalkPath::Action::Continue:
+      break;
+  }
 
   WIN32_FIND_DATA file_data;
   auto handle = FindFirstFileExW(name.c_str(), FindExInfoStandard, &file_data, FindExSearchNameMatch, nullptr, 0);
@@ -536,9 +572,18 @@ static Status walk_path_dir(const std::wstring &dir_name,
     TRY_RESULT(entry_name, from_wstring(full_name));
     if (file_data.cFileName[0] != '.') {
       if ((file_data.dwFileAttributes & FILE_ATTRIBUTE_DIRECTORY) != 0) {
-        TRY_STATUS(walk_path_dir(full_name, func));
+        TRY_RESULT(is_ok, walk_path_dir(full_name, func));
+        if (!is_ok) {
+          return false;
+        }
       } else {
-        func(entry_name, false);
+        switch (func(entry_name, WalkPath::Type::NotDir)) {
+          case WalkPath::Action::Abort:
+            return false;
+          case WalkPath::Action::SkipDir:
+          case WalkPath::Action::Continue:
+            break;
+        }
       }
     }
     auto status = FindNextFileW(handle, &file_data);
@@ -551,18 +596,25 @@ static Status walk_path_dir(const std::wstring &dir_name,
     }
   }
   TRY_RESULT(entry_name, from_wstring(dir_name));
-  func(entry_name, true);
+  switch (func(entry_name, WalkPath::Type::ExitDir)) {
-  return Status::OK();
+    case WalkPath::Action::Abort:
+      return false;
+    case WalkPath::Action::SkipDir:
+    case WalkPath::Action::Continue:
+      break;
+  }
+  return true;
 }
 
-Status walk_path(CSlice path, const std::function<void(CSlice name, bool is_directory)> &func) {
+Status WalkPath::do_run(CSlice path, const std::function<Action(CSlice name, Type)> &func) {
   TRY_RESULT(wpath, to_wstring(path));
   Slice path_slice = path;
   while (!path_slice.empty() && (path_slice.back() == '/' || path_slice.back() == '\\')) {
     path_slice.remove_suffix(1);
     wpath.pop_back();
   }
-  return walk_path_dir(wpath, func);
+  TRY_STATUS(walk_path_dir(wpath, func));
+  return Status::OK();
 }
 
 #endif

tdutils/td/utils/port/path.h

@@ -40,6 +40,32 @@ Result<std::pair<FileFd, string>> mkstemp(CSlice dir) TD_WARN_UNUSED_RESULT;
 
 Result<string> mkdtemp(CSlice dir, Slice prefix) TD_WARN_UNUSED_RESULT;
 
-Status walk_path(CSlice path, const std::function<void(CSlice name, bool is_directory)> &func) TD_WARN_UNUSED_RESULT;
+class WalkPath {
+ public:
+  enum class Action { Continue, Abort, SkipDir };
+  enum class Type { EnterDir, ExitDir, NotDir };
+
+  template <class F, class R = decltype(std::declval<F>()("", Type::ExitDir))>
+  static std::enable_if_t<std::is_same<R, Action>::value, Status> run(CSlice path, F &&func) TD_WARN_UNUSED_RESULT {
+    return do_run(path, func);
+  }
+  template <class F, class R = decltype(std::declval<F>()("", Type::ExitDir))>
+  static std::enable_if_t<!std::is_same<R, Action>::value, Status> run(CSlice path, F &&func) TD_WARN_UNUSED_RESULT {
+    return do_run(path, [&](CSlice name, Type type) {
+      func(name, type);
+      return Action::Continue;
+    });
+  }
+
+ private:
+  static TD_WARN_UNUSED_RESULT Status do_run(CSlice path,
+                                             const std::function<WalkPath::Action(CSlice name, Type type)> &func);
+};
+
+// deprecated interface
+template <class F>
+TD_WARN_UNUSED_RESULT Status walk_path(CSlice path, F &&func) {
+  return WalkPath::run(path, func);
+}
 
 }  // namespace td

tdutils/test/port.cpp

@@ -10,6 +10,7 @@
 #include "td/utils/port/path.h"
 #include "td/utils/Slice.h"
 #include "td/utils/tests.h"
+#include "td/utils/misc.h"
 
 using namespace td;
 
@@ -17,7 +18,7 @@ TEST(Port, files) {
   CSlice main_dir = "test_dir";
   rmrf(main_dir).ignore();
   ASSERT_TRUE(FileFd::open(main_dir, FileFd::Write).is_error());
-  ASSERT_TRUE(walk_path(main_dir, [](CSlice name, bool is_directory) { UNREACHABLE(); }).is_error());
+  ASSERT_TRUE(walk_path(main_dir, [](CSlice name, WalkPath::Type type) { UNREACHABLE(); }).is_error());
   mkdir(main_dir).ensure();
   mkdir(PSLICE() << main_dir << TD_DIR_SLASH << "A").ensure();
   mkdir(PSLICE() << main_dir << TD_DIR_SLASH << "B").ensure();
@@ -34,16 +35,39 @@ TEST(Port, files) {
   int cnt = 0;
   const int ITER_COUNT = 1000;
   for (int i = 0; i < ITER_COUNT; i++) {
-    walk_path(main_dir,
+    walk_path(main_dir, [&](CSlice name, WalkPath::Type type) {
-              [&](CSlice name, bool is_directory) {
+      if (type == WalkPath::Type::NotDir) {
-                if (!is_directory) {
+        ASSERT_TRUE(name == fd_path || name == fd2_path);
-                  ASSERT_TRUE(name == fd_path || name == fd2_path);
+      }
-                }
+      cnt++;
-                cnt++;
+    }).ensure();
-              })
-        .ensure();
   }
-  ASSERT_EQ(7 * ITER_COUNT, cnt);
+  ASSERT_EQ((5 * 2 + 2) * ITER_COUNT, cnt);
+  bool was_abort = false;
+  walk_path(main_dir, [&](CSlice name, WalkPath::Type type) {
+    CHECK(!was_abort);
+    if (type == WalkPath::Type::EnterDir && ends_with(name, PSLICE() << "/B")) {
+      was_abort = true;
+      return WalkPath::Action::Abort;
+    }
+    return WalkPath::Action::Continue;
+  }).ensure();
+  CHECK(was_abort);
+
+  cnt = 0;
+  bool is_first_dir = true;
+  walk_path(main_dir, [&](CSlice name, WalkPath::Type type) {
+    cnt++;
+    if (type == WalkPath::Type::EnterDir) {
+      if (is_first_dir) {
+        is_first_dir = false;
+      } else {
+        return WalkPath::Action::SkipDir;
+      }
+    }
+    return WalkPath::Action::Continue;
+  }).ensure();
+  ASSERT_EQ(6, cnt);
 
   ASSERT_EQ(0u, fd.get_size());
   ASSERT_EQ(12u, fd.write("Hello world!").move_as_ok());