Add some predicated searching functions to StringRef.

This adds 4 new functions to StringRef, which can be used to
take or drop characters while a certain condition is met, or
until a certain condition is met.  They are:

take_while - Return characters until a condition is not met.
take_until - Return characters until a condition is met.
drop_while - Remove characters until a condition is not met.
drop_until - Remove characters until a condition is met.

Internally, all of these functions delegate to two additional
helper functions which can be used to search for the position
of a character meeting or not meeting a condition, which are:

find_if - Find the first character matching a predicate.
find_if_not - Find the first character not matching a predicate.

Differential Revision: https://reviews.llvm.org/D24842

llvm-svn: 282346

2 files changed, 118 insertions, 0 deletions

diff --git a/llvm/include/llvm/ADT/StringRef.h b/llvm/include/llvm/ADT/StringRef.h
index 44b14d20dee..f6a348e8ab4 100644
--- a/llvm/include/llvm/ADT/StringRef.h
+++ b/llvm/include/llvm/ADT/StringRef.h
@@ -10,6 +10,7 @@
 #ifndef LLVM_ADT_STRINGREF_H
 #define LLVM_ADT_STRINGREF_H
 
+#include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/iterator_range.h"
 #include "llvm/Support/Compiler.h"
 #include <algorithm>
@@ -276,6 +277,32 @@ namespace llvm {
       return npos;
     }
 
+    /// Search for the first character satisfying the predicate \p F
+    ///
+    /// \returns The index of the first character satisfying \p F starting from
+    /// \p From, or npos if not found.
+    LLVM_ATTRIBUTE_ALWAYS_INLINE
+    LLVM_ATTRIBUTE_UNUSED_RESULT
+    size_t find_if(function_ref<bool(char)> F, size_t From = 0) const {
+      StringRef S = drop_front(From);
+      while (!S.empty()) {
+        if (F(S.front()))
+          return size() - S.size();
+        S = S.drop_front();
+      }
+      return npos;
+    }
+
+    /// Search for the first character not satisfying the predicate \p F
+    ///
+    /// \returns The index of the first character not satisfying \p F starting
+    /// from \p From, or npos if not found.
+    LLVM_ATTRIBUTE_ALWAYS_INLINE
+    LLVM_ATTRIBUTE_UNUSED_RESULT
+    size_t find_if_not(function_ref<bool(char)> F, size_t From = 0) const {
+      return find_if([F](char c) { return !F(c); }, From);
+    }
+
     /// Search for the first string \p Str in the string.
     ///
     /// \returns The index of the first occurrence of \p Str, or npos if not
@@ -352,6 +379,9 @@ namespace llvm {
     LLVM_ATTRIBUTE_ALWAYS_INLINE
     bool contains(StringRef Other) const { return find(Other) != npos; }
 
+    LLVM_ATTRIBUTE_ALWAYS_INLINE
+    bool contains(char C) const { return find_first_of(C) != npos; }
+
     /// @}
     /// @name Helpful Algorithms
     /// @{
@@ -496,6 +526,22 @@ namespace llvm {
       return drop_front(size() - N);
     }
 
+    /// Return the longest prefix of 'this' such that every character
+    /// in the prefix satisfies the given predicate.
+    LLVM_ATTRIBUTE_ALWAYS_INLINE
+    LLVM_ATTRIBUTE_UNUSED_RESULT
+    StringRef take_while(function_ref<bool(char)> F) const {
+      return substr(0, find_if_not(F));
+    }
+
+    /// Return the longest prefix of 'this' such that no character in
+    /// the prefix satisfies the given predicate.
+    LLVM_ATTRIBUTE_ALWAYS_INLINE
+    LLVM_ATTRIBUTE_UNUSED_RESULT
+    StringRef take_until(function_ref<bool(char)> F) const {
+      return substr(0, find_if(F));
+    }
+
     /// Return a StringRef equal to 'this' but with the first \p N elements
     /// dropped.
     LLVM_ATTRIBUTE_ALWAYS_INLINE
@@ -514,6 +560,22 @@ namespace llvm {
       return substr(0, size()-N);
     }
 
+    /// Return a StringRef equal to 'this', but with all characters satisfying
+    /// the given predicate dropped from the beginning of the string.
+    LLVM_ATTRIBUTE_ALWAYS_INLINE
+    LLVM_ATTRIBUTE_UNUSED_RESULT
+    StringRef drop_while(function_ref<bool(char)> F) const {
+      return substr(find_if_not(F));
+    }
+
+    /// Return a StringRef equal to 'this', but with all characters not
+    /// satisfying the given predicate dropped from the beginning of the string.
+    LLVM_ATTRIBUTE_ALWAYS_INLINE
+    LLVM_ATTRIBUTE_UNUSED_RESULT
+    StringRef drop_until(function_ref<bool(char)> F) const {
+      return substr(find_if(F));
+    }
+
     /// Returns true if this StringRef has the given prefix and removes that
     /// prefix.
     LLVM_ATTRIBUTE_ALWAYS_INLINE
diff --git a/llvm/unittests/ADT/StringRefTest.cpp b/llvm/unittests/ADT/StringRefTest.cpp
index 39ed71b77e9..e1f54967572 100644
--- a/llvm/unittests/ADT/StringRefTest.cpp
+++ b/llvm/unittests/ADT/StringRefTest.cpp
@@ -864,4 +864,60 @@ TEST(StringRefTest, Take) {
   EXPECT_TRUE(Taken.empty());
 }
 
+TEST(StringRefTest, FindIf) {
+  StringRef Punct("Test.String");
+  StringRef NoPunct("ABCDEFG");
+  StringRef Empty;
+
+  auto IsPunct = [](char c) { return ::ispunct(c); };
+  auto IsAlpha = [](char c) { return ::isalpha(c); };
+  EXPECT_EQ(4, Punct.find_if(IsPunct));
+  EXPECT_EQ(StringRef::npos, NoPunct.find_if(IsPunct));
+  EXPECT_EQ(StringRef::npos, Empty.find_if(IsPunct));
+
+  EXPECT_EQ(4, Punct.find_if_not(IsAlpha));
+  EXPECT_EQ(StringRef::npos, NoPunct.find_if_not(IsAlpha));
+  EXPECT_EQ(StringRef::npos, Empty.find_if_not(IsAlpha));
+}
+
+TEST(StringRefTest, TakeWhileUntil) {
+  StringRef Test("String With 1 Number");
+
+  StringRef Taken = Test.take_while([](char c) { return ::isdigit(c); });
+  EXPECT_EQ("", Taken);
+
+  Taken = Test.take_until([](char c) { return ::isdigit(c); });
+  EXPECT_EQ("String With ", Taken);
+
+  Taken = Test.take_while([](char c) { return true; });
+  EXPECT_EQ(Test, Taken);
+
+  Taken = Test.take_until([](char c) { return true; });
+  EXPECT_EQ("", Taken);
+
+  Test = "";
+  Taken = Test.take_while([](char c) { return true; });
+  EXPECT_EQ("", Taken);
+}
+
+TEST(StringRefTest, DropWhileUntil) {
+  StringRef Test("String With 1 Number");
+
+  StringRef Taken = Test.drop_while([](char c) { return ::isdigit(c); });
+  EXPECT_EQ(Test, Taken);
+
+  Taken = Test.drop_until([](char c) { return ::isdigit(c); });
+  EXPECT_EQ("1 Number", Taken);
+
+  Taken = Test.drop_while([](char c) { return true; });
+  EXPECT_EQ("", Taken);
+
+  Taken = Test.drop_until([](char c) { return true; });
+  EXPECT_EQ(Test, Taken);
+
+  StringRef EmptyString = "";
+  Taken = EmptyString.drop_while([](char c) { return true; });
+  EXPECT_EQ("", Taken);
+}
+
 } // end anonymous namespace