Change GetMatchingPaths to avoid traversing unnecesscary paths

tensorflow/core/platform/file_system_helper.cc

@@ -55,30 +55,48 @@ void ForEach(int first, int last, const std::function<void(int)>& f) {
 Status GetMatchingPaths(FileSystem* fs, Env* env, const string& pattern,
                         std::vector<string>* results) {
   results->clear();
-  // Find the fixed prefix by looking for the first wildcard.
+  if (pattern.empty()) {
+    return Status::OK();
+  }
+
   string fixed_prefix = pattern.substr(0, pattern.find_first_of("*?[\\"));
   string eval_pattern = pattern;
-  std::vector<string> all_files;
   string dir(io::Dirname(fixed_prefix));
   // If dir is empty then we need to fix up fixed_prefix and eval_pattern to
   // include . as the top level directory.
   if (dir.empty()) {
     dir = ".";
     fixed_prefix = io::JoinPath(dir, fixed_prefix);
-    eval_pattern = io::JoinPath(dir, pattern);
+    eval_pattern = io::JoinPath(dir, eval_pattern);
   }
-
+  bool is_directory = pattern[pattern.size() - 1] == '/';
+#ifdef PLATFORM_WINDOWS
+  is_directory = is_directory || pattern[pattern.size() - 1] == '\\';
+#endif
+  std::vector<string> dirs;
+  if (!is_directory) {
+    dirs.push_back(eval_pattern);
+  }
+  StringPiece tmp_dir(io::Dirname(eval_pattern));
+  while (tmp_dir.size() > dir.size()) {
+    dirs.push_back(string(tmp_dir));
+    tmp_dir = io::Dirname(tmp_dir);
+  }
+  dirs.push_back(dir);
+  std::reverse(dirs.begin(), dirs.end());
   // Setup a BFS to explore everything under dir.
-  std::deque<string> dir_q;
-  dir_q.push_back(dir);
+  std::deque<std::pair<string, int>> dir_q;
+  dir_q.push_back({dirs[0], 0});
   Status ret;  // Status to return.
   // children_dir_status holds is_dir status for children. It can have three
   // possible values: OK for true; FAILED_PRECONDITION for false; CANCELLED
   // if we don't calculate IsDirectory (we might do that because there isn't
   // any point in exploring that child path).
   std::vector<Status> children_dir_status;
   while (!dir_q.empty()) {
-    string current_dir = dir_q.front();
+    string current_dir = dir_q.front().first;
+    int dir_index = dir_q.front().second;
+    dir_index++;
     dir_q.pop_front();
     std::vector<string> children;
     Status s = fs->GetChildren(current_dir, &children);
@@ -91,16 +109,17 @@ Status GetMatchingPaths(FileSystem* fs, Env* env, const string& pattern,
     // This IsDirectory call can be expensive for some FS. Parallelizing it.
     children_dir_status.resize(children.size());
     ForEach(0, children.size(),
-            [fs, &current_dir, &children, &fixed_prefix,
-             &children_dir_status](int i) {
+            [fs, &current_dir, &children, &dirs, dir_index,
+             is_directory, &children_dir_status](int i) {
               const string child_path = io::JoinPath(current_dir, children[i]);
-              // In case the child_path doesn't start with the fixed_prefix then
-              // we don't need to explore this path.
-              if (!absl::StartsWith(child_path, fixed_prefix)) {
+              if (!fs->Match(child_path, dirs[dir_index])) {
                 children_dir_status[i] = Status(tensorflow::error::CANCELLED,
                                                 "Operation not needed");
+              } else if (dir_index != dirs.size() - 1){
+                children_dir_status[i] = fs->IsDirectory(child_path);Status::OK();
               } else {
-                children_dir_status[i] = fs->IsDirectory(child_path);
+                children_dir_status[i] = is_directory ?
+                    fs->IsDirectory(child_path) : Status::OK();
               }
             });
     for (size_t i = 0; i < children.size(); ++i) {
@@ -109,18 +128,13 @@ Status GetMatchingPaths(FileSystem* fs, Env* env, const string& pattern,
       if (children_dir_status[i].code() == tensorflow::error::CANCELLED) {
         continue;
       }
-      // If the child is a directory add it to the queue.
       if (children_dir_status[i].ok()) {
-        dir_q.push_back(child_path);
+        if (dir_index != dirs.size() - 1) {
+          dir_q.push_back({child_path, dir_index});
+        } else {
+          results->push_back(child_path);
+        }
       }
-      all_files.push_back(child_path);
-    }
-  }
-
-  // Match all obtained files to the input pattern.
-  for (const auto& f : all_files) {
-    if (fs->Match(f, eval_pattern)) {
-      results->push_back(f);
     }
   }
   return ret;