From 43d19fa513559280e65abc5502071afdbfdea929 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Fran=C3=A7ois=20Saint-Jacques?= <fsaintjacques@gmail.com>
Date: Tue, 24 Sep 2019 10:22:35 -0400
Subject: [PATCH] Address comments

---
 cpp/src/arrow/filesystem/path_tree.cc      |  2 +-
 cpp/src/arrow/filesystem/path_tree.h       | 30 +++++++++++------
 cpp/src/arrow/filesystem/path_tree_test.cc | 38 +++++++++++++++++++++-
 3 files changed, 58 insertions(+), 12 deletions(-)

diff --git a/cpp/src/arrow/filesystem/path_tree.cc b/cpp/src/arrow/filesystem/path_tree.cc
index b234e35c99796..682d3ba10215c 100644
--- a/cpp/src/arrow/filesystem/path_tree.cc
+++ b/cpp/src/arrow/filesystem/path_tree.cc
@@ -123,7 +123,7 @@ std::ostream& operator<<(std::ostream& os, const std::shared_ptr<PathTree>& tree
 }
 
 bool operator==(const std::shared_ptr<PathTree>& lhs,
-                             const std::shared_ptr<PathTree>& rhs) {
+                const std::shared_ptr<PathTree>& rhs) {
   if (lhs == NULLPTR && rhs == NULLPTR) {
     return true;
   } else if (lhs != NULLPTR && rhs != NULLPTR) {
diff --git a/cpp/src/arrow/filesystem/path_tree.h b/cpp/src/arrow/filesystem/path_tree.h
index 6add0fd7e6402..50ec1f9470450 100644
--- a/cpp/src/arrow/filesystem/path_tree.h
+++ b/cpp/src/arrow/filesystem/path_tree.h
@@ -25,6 +25,8 @@
 #include <utility>
 #include <vector>
 
+#include "arrow/status.h"
+
 namespace arrow {
 namespace fs {
 
@@ -57,21 +59,29 @@ class ARROW_EXPORT PathTree {
 
   /// \brief Visit with eager pruning.
   template <typename Visitor, typename Matcher>
-  void Visit(Visitor v, Matcher m) const {
-    if (!m(stats_)) {
-      return;
+  Status Visit(Visitor&& v, Matcher&& m) const {
+    bool match = false;
+    ARROW_RETURN_NOT_OK(m(stats_, &match));
+    if (!match) {
+      return Status::OK();
     }
 
-    v(stats_);
+    ARROW_RETURN_NOT_OK(v(stats_));
+
+    for (const auto& t : subtrees_) {
+      ARROW_RETURN_NOT_OK(t->Visit(v, m));
+    }
 
-    auto recurse = [&v, &m](const std::shared_ptr<PathTree>& t) { t->Visit(v, m); };
-    std::for_each(subtrees_.cbegin(), subtrees_.cend(), recurse);
+    return Status::OK();
   }
 
   template <typename Visitor>
-  void Visit(Visitor v) const {
-    auto always_match = [](const std::shared_ptr<PathTree>& t) { return true; };
-    return Visit(std::move(v), always_match);
+  Status Visit(Visitor&& v) const {
+    auto always_match = [](const FileStats& t, bool* match) {
+      *match = true;
+      return Status::OK();
+    };
+    return Visit(v, always_match);
   }
 
   bool operator==(const PathTree& other) const {
@@ -82,7 +92,7 @@ class ARROW_EXPORT PathTree {
   FileStats stats_;
   std::vector<std::shared_ptr<PathTree>> subtrees_;
 
-  // The AddChild method is conveninent to create trees in a top-down fashion,
+  // The AddChild method is convenient to create trees in a top-down fashion,
   // e.g. the Make factory constructor.
   void AddChild(std::shared_ptr<PathTree> child) {
     subtrees_.push_back(std::move(child));
diff --git a/cpp/src/arrow/filesystem/path_tree_test.cc b/cpp/src/arrow/filesystem/path_tree_test.cc
index 86127aba12d5e..fb38ad45835ef 100644
--- a/cpp/src/arrow/filesystem/path_tree_test.cc
+++ b/cpp/src/arrow/filesystem/path_tree_test.cc
@@ -28,6 +28,8 @@
 #include "arrow/filesystem/test_util.h"
 #include "arrow/testing/gtest_util.h"
 
+using testing::ContainerEq;
+
 namespace arrow {
 namespace fs {
 
@@ -45,7 +47,7 @@ void AssertMakePathTree(std::vector<FileStats> stats,
   std::vector<std::shared_ptr<PathTree>> actual;
 
   ASSERT_OK(PathTree::Make(stats, &actual));
-  EXPECT_THAT(actual, testing::ContainerEq(expected));
+  EXPECT_THAT(actual, ContainerEq(expected));
 }
 
 TEST(TestPathTree, Basic) {
@@ -59,6 +61,9 @@ TEST(TestPathTree, Basic) {
   AssertMakePathTree({Dir("AA"), File("AA/BB/bb")},
                      {PT(Dir("AA"), {PT(File("AA/BB/bb"))})});
 
+  // Ancestors should link to parent irregardless of the ordering
+  AssertMakePathTree({File("AA/aa"), Dir("AA")}, {PT(Dir("AA"), {PT(File("AA/aa"))})});
+
   // Multiple roots are supported.
   AssertMakePathTree({File("aa"), File("bb")}, {PT(File("aa")), PT(File("bb"))});
   AssertMakePathTree(
@@ -139,5 +144,36 @@ TEST(TestPathTree, HourlyETL) {
   AssertMakePathTree(stats, forest);
 }
 
+TEST(TestPathTree, Visit) {
+  std::shared_ptr<PathTree> tree;
+  ASSERT_OK(PathTree::Make({Dir("A"), File("A/a")}, &tree));
+
+  // Should propagate failure
+  auto visit_noop = [](const FileStats&) { return Status::OK(); };
+  ASSERT_OK(tree->Visit(visit_noop));
+  auto visit_fail = [](const FileStats&) { return Status::Invalid(""); };
+  ASSERT_RAISES(Invalid, tree->Visit(visit_fail));
+  auto match_fail = [](const FileStats&, bool* match) { return Status::Invalid(""); };
+  ASSERT_RAISES(Invalid, tree->Visit(visit_noop, match_fail));
+
+  // Ensure basic visit of all nodes
+  std::vector<FileStats> collect;
+  auto visit = [&collect](const FileStats& f) {
+    collect.push_back(f);
+    return Status::OK();
+  };
+  ASSERT_OK(tree->Visit(visit));
+  EXPECT_THAT(collect, ContainerEq(std::vector<FileStats>{Dir("A"), File("A/a")}));
+
+  // Matcher should be evaluated on all nodes
+  collect.resize(0);
+  auto match_dir = [](const FileStats& s, bool* m) {
+    *m = s.IsDirectory();
+    return Status::OK();
+  };
+  ASSERT_OK(tree->Visit(visit, match_dir));
+  EXPECT_THAT(collect, ContainerEq(std::vector<FileStats>{Dir("A")}));
+}
+
 }  // namespace fs
 }  // namespace arrow