Support/FileSystem: Add directory_iterator implementation.

llvm-svn: 120989

include/llvm/Support/FileSystem.h

@@ -544,19 +544,29 @@ class directory_entry {
 
 public:
   explicit directory_entry(const Twine &path, file_status st = file_status(),
-                                       file_status symlink_st = file_status());
+                                       file_status symlink_st = file_status())
+    : Path(path.str())
+    , Status(st)
+    , SymlinkStatus(symlink_st) {}
+
+  directory_entry() {}
 
   void assign(const Twine &path, file_status st = file_status(),
-                          file_status symlink_st = file_status());
+                          file_status symlink_st = file_status()) {
+    Path = path.str();
+    Status = st;
+    SymlinkStatus = symlink_st;
+  }
+
   void replace_filename(const Twine &filename, file_status st = file_status(),
                               file_status symlink_st = file_status());
 
-  const SmallVectorImpl<char> &path() const;
+  StringRef path() const { return Path; }
   error_code status(file_status &result) const;
   error_code symlink_status(file_status &result) const;
 
-  bool operator==(const directory_entry& rhs) const;
-  bool operator!=(const directory_entry& rhs) const;
+  bool operator==(const directory_entry& rhs) const { return Path == rhs.Path; }
+  bool operator!=(const directory_entry& rhs) const { return !(*this == rhs); }
   bool operator< (const directory_entry& rhs) const;
   bool operator<=(const directory_entry& rhs) const;
   bool operator> (const directory_entry& rhs) const;
@@ -567,16 +577,41 @@ public:
 /// operator++ because we need an error_code. If it's really needed we can make
 /// it call report_fatal_error on error.
 class directory_iterator {
-  // implementation directory iterator status
+  intptr_t IterationHandle;
+  directory_entry CurrentEntry;
+
+  // Platform implementations implement these functions to handle iteration.
+  friend error_code directory_iterator_construct(directory_iterator& it,
+                                                 const StringRef &path);
+  friend error_code directory_iterator_increment(directory_iterator& it);
+  friend error_code directory_iterator_destruct(directory_iterator& it);
 
 public:
-  explicit directory_iterator(const Twine &path, error_code &ec);
+  explicit directory_iterator(const Twine &path, error_code &ec)
+  : IterationHandle(0) {
+    SmallString<128> path_storage;
+    ec = directory_iterator_construct(*this, path.toStringRef(path_storage));
+  }
+
+  /// Construct end iterator.
+  directory_iterator() : IterationHandle(0) {}
+
+  ~directory_iterator() {
+    directory_iterator_destruct(*this);
+  }
+
   // No operator++ because we need error_code.
-  directory_iterator &increment(error_code &ec);
+  directory_iterator &increment(error_code &ec) {
+    ec = directory_iterator_increment(*this);
+    return *this;
+  }
 
-  const directory_entry &operator*() const;
-  const directory_entry *operator->() const;
+  const directory_entry &operator*() const { return CurrentEntry; }
+  const directory_entry *operator->() const { return &CurrentEntry; };
 
+  bool operator!=(const directory_iterator &RHS) const {
+    return CurrentEntry != RHS.CurrentEntry;
+  }
   // Other members as required by
   // C++ Std, 24.1.1 Input iterators [input.iterators]
 };

lib/Support/PathV2.cpp

@@ -688,6 +688,16 @@ error_code create_directories(const Twine &path, bool &existed) {
   return create_directory(p, existed);
 }
 
+void directory_entry::replace_filename(const Twine &filename, file_status st,
+                                       file_status symlink_st) {
+  SmallString<128> path(Path.begin(), Path.end());
+  path::remove_filename(path);
+  path::append(path, filename);
+  Path = path.str();
+  Status = st;
+  SymlinkStatus = symlink_st;
+}
+
 } // end namespace fs
 } // end namespace sys
 } // end namespace llvm

lib/Support/Windows/PathV2.inc

@@ -121,6 +121,15 @@ namespace {
   typedef ScopedHandle<HCRYPTPROV, HCRYPTPROV(INVALID_HANDLE_VALUE),
                        BOOL (WINAPI*)(HCRYPTPROV), CryptReleaseContext>
     ScopedCryptContext;
+  bool is_separator(const wchar_t value) {
+    switch (value) {
+    case L'\\':
+    case L'/':
+      return true;
+    default:
+      return false;
+    }
+  }
 }
 
 namespace llvm {
@@ -598,6 +607,74 @@ retry_create_file:
   result_fd = fd;
   return success;
 }
+
+error_code directory_iterator_construct(directory_iterator& it,
+                                        const StringRef &path) {
+  SmallVector<wchar_t, 128> path_utf16;
+
+  if (error_code ec = UTF8ToUTF16(path,
+                                  path_utf16))
+    return ec;
+
+  // Convert path to the format that Windows is happy with.
+  if (path_utf16.size() > 0 &&
+      !is_separator(path_utf16[path.size() - 1]) &&
+      path_utf16[path.size() - 1] != L':') {
+    path_utf16.push_back(L'\\');
+    path_utf16.push_back(L'*');
+  } else {
+    path_utf16.push_back(L'*');
+  }
+
+  //  Get the first directory entry.
+  WIN32_FIND_DATAW FirstFind;
+  ScopedFindHandle FindHandle(::FindFirstFileW(path_utf16.c_str(), &FirstFind));
+  if (!FindHandle)
+    return windows_error(::GetLastError());
+
+  // Construct the current directory entry.
+  SmallString<128> directory_entry_path_utf8;
+  if (error_code ec = UTF16ToUTF8(FirstFind.cFileName,
+                                  ::wcslen(FirstFind.cFileName),
+                                  directory_entry_path_utf8))
+    return ec;
+
+  it.IterationHandle = intptr_t(FindHandle.take());
+  it.CurrentEntry = directory_entry(path);
+  it.CurrentEntry.replace_filename(Twine(directory_entry_path_utf8));
+
+  return success;
+}
+
+error_code directory_iterator_destruct(directory_iterator& it) {
+  if (it.IterationHandle != 0)
+    // Closes the handle if it's valid.
+    ScopedFindHandle close(HANDLE(it.IterationHandle));
+  it.IterationHandle = 0;
+  it.CurrentEntry = directory_entry();
+  return success;
+}
+
+error_code directory_iterator_increment(directory_iterator& it) {
+  WIN32_FIND_DATAW FindData;
+  if (!::FindNextFileW(HANDLE(it.IterationHandle), &FindData)) {
+    error_code ec = windows_error(::GetLastError());
+    // Check for end.
+    if (ec == windows_error::no_more_files)
+      return directory_iterator_destruct(it);
+    return ec;
+  }
+
+  SmallString<128> directory_entry_path_utf8;
+  if (error_code ec = UTF16ToUTF8(FindData.cFileName,
+                                  ::wcslen(FindData.cFileName),
+                                  directory_entry_path_utf8))
+    return ec;
+
+  it.CurrentEntry.replace_filename(Twine(directory_entry_path_utf8));
+  return success;
+}
+
 } // end namespace fs
 } // end namespace sys
 } // end namespace llvm

lib/Support/Windows/Windows.h

@@ -94,7 +94,11 @@ public:
     return Handle == InvalidHandle ? 0 : unspecified_bool_true;
   }
 
-  typedef ScopedHandle<HANDLE, INVALID_HANDLE_VALUE,
-                       BOOL (WINAPI*)(HANDLE), ::FindClose>
-    ScopedFindHandle;
+  bool operator!() const {
+    return Handle == InvalidHandle;
+  }
 };
+
+typedef ScopedHandle<HANDLE, INVALID_HANDLE_VALUE,
+                      BOOL (WINAPI*)(HANDLE), ::FindClose>
+  ScopedFindHandle;

unittests/Support/Path.cpp

@@ -9,6 +9,7 @@
 
 #include "llvm/Support/FileSystem.h"
 #include "llvm/Support/PathV2.h"
+#include "llvm/Support/ErrorHandling.h"
 
 #include "gtest/gtest.h"
 
@@ -168,6 +169,18 @@ TEST(Support, Path) {
   // Make sure Temp1 doesn't exist.
   ASSERT_FALSE(fs::exists(Twine(TempPath), TempFileExists));
   EXPECT_FALSE(TempFileExists);
+
+  // I've yet to do directory iteration on Unix.
+#ifdef LLVM_ON_WIN32
+  error_code ec;
+  for (fs::directory_iterator i(".", ec), e; i != e; i.increment(ec)) {
+    if (ec) {
+      errs() << ec.message() << '\n';
+      errs().flush();
+      report_fatal_error("Directory iteration failed!");
+    }
+  }
+#endif
 }
 
 } // anonymous namespace