[llvm] [ADT] Add predicate based match support to StringSwitch (PR #188046)

[via llvm-commits]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-llvm-adt

Author: Oliver Hunt (ojhunt)

<details>
<summary>Changes</summary>

This introduces `Predicate` and `IfNotPredicate` case selection to StringSwitch to allow use cases like

```
StringSwitch<...>(..)
  .Case("foo", FooTok)
  .Predicate(isAlpha, IdentifierTok)
...
```

This is mostly useful for improving conciseness and clarity when processing generated strings, diagnostics, and similar.

---
Full diff: https://github.com/llvm/llvm-project/pull/188046.diff


2 Files Affected:

- (modified) llvm/include/llvm/ADT/StringSwitch.h (+32) 
- (modified) llvm/unittests/ADT/StringSwitchTest.cpp (+69) 


``````````diff
diff --git a/llvm/include/llvm/ADT/StringSwitch.h b/llvm/include/llvm/ADT/StringSwitch.h
index b144d51364855..3072f3ddbb8b1 100644
--- a/llvm/include/llvm/ADT/StringSwitch.h
+++ b/llvm/include/llvm/ADT/StringSwitch.h
@@ -123,6 +123,38 @@ class StringSwitch {
     return *this;
   }
 
+  // A StringSwitch case that is selected if the specified predicate
+  // returns true for every character in a subject string.
+  StringSwitch &Predicate(function_ref<bool(char)> Pred, T Value) {
+    if (!Result && Str.find_if_not(Pred) == StringRef::npos)
+      Result = std::move(Value);
+    return *this;
+  }
+
+  // A StringSwitch case that is selected if the specified predicate
+  // returns true for the subject string.
+  StringSwitch &Predicate(function_ref<bool(StringRef)> Pred, T Value) {
+    if (!Result && Pred(Str))
+      Result = std::move(Value);
+    return *this;
+  }
+
+  // A StringSwitch case that is selected if the specified predicate
+  // returns false for any character in a subject string.
+  StringSwitch &IfNotPredicate(function_ref<bool(char)> Pred, T Value) {
+    if (!Result && Str.find_if(Pred) == StringRef::npos)
+      Result = std::move(Value);
+    return *this;
+  }
+
+  // A StringSwitch case that is selected if the specified predicate
+  // returns false for the subject string.
+  StringSwitch &IfNotPredicate(function_ref<bool(StringRef)> Pred, T Value) {
+    if (!Result && !Pred(Str))
+      Result = std::move(Value);
+    return *this;
+  }
+
   [[nodiscard]] R Default(T Value) {
     if (Result)
       return std::move(*Result);
diff --git a/llvm/unittests/ADT/StringSwitchTest.cpp b/llvm/unittests/ADT/StringSwitchTest.cpp
index 75d50f4dd1b5b..6975e004534dd 100644
--- a/llvm/unittests/ADT/StringSwitchTest.cpp
+++ b/llvm/unittests/ADT/StringSwitchTest.cpp
@@ -257,6 +257,75 @@ TEST(StringSwitchTest, StringSwitchMultipleMatches) {
   EXPECT_EQ(1, Translate("b"));
 }
 
+TEST(StringPredicateTest, PredicateChar) {
+  auto IsAlpha = [](char Ch) {
+    return (Ch >= 'a' && Ch <= 'z') || (Ch >= 'A' && Ch <= 'Z');
+  };
+  auto PredicateMatch = [&](StringRef S) {
+    return llvm::StringSwitch<int>(S)
+        .Case("A", 0)
+        .Predicate(IsAlpha, 1)
+        .Case("B", 2)
+        .Case("C D", 3)
+        .Default(4);
+  };
+  EXPECT_EQ(0, PredicateMatch("A"));
+  EXPECT_EQ(1, PredicateMatch("abcdef"));
+  EXPECT_EQ(1, PredicateMatch("B"));
+  EXPECT_EQ(3, PredicateMatch("C D"));
+  EXPECT_EQ(4, PredicateMatch("1234"));
+}
+
+TEST(StringPredicateTest, PredicateStr) {
+  auto MatchFn = [](StringRef Str) { return Str == "abc" || Str == "def"; };
+  auto PredicateMatch = [&](StringRef S) {
+    return llvm::StringSwitch<int>(S)
+        .Case("abc", 0)
+        .Predicate(MatchFn, 1)
+        .Case("def", 2)
+        .Default(3);
+  };
+  EXPECT_EQ(0, PredicateMatch("abc"));
+  EXPECT_EQ(1, PredicateMatch("def"));
+  EXPECT_EQ(3, PredicateMatch("ghi"));
+}
+
+TEST(StringPredicateTest, IfNotPredicateChar) {
+  auto IsAlpha = [](char Ch) {
+    return (Ch >= 'a' && Ch <= 'z') || (Ch >= 'A' && Ch <= 'Z');
+  };
+  auto PredicateMatch = [&](StringRef Str) {
+    return llvm::StringSwitch<int>(Str)
+        .Case("AA", 0)
+        .Case("11", 1)
+        .IfNotPredicate(IsAlpha, 2)
+        .Case("BB", 3)
+        .Case("22", 4)
+        .Default(5);
+  };
+  EXPECT_EQ(0, PredicateMatch("AA"));
+  EXPECT_EQ(1, PredicateMatch("11"));
+  EXPECT_EQ(2, PredicateMatch("22"));
+  EXPECT_EQ(2, PredicateMatch("22"));
+  EXPECT_EQ(3, PredicateMatch("BB"));
+  EXPECT_EQ(5, PredicateMatch("ABC123"));
+  EXPECT_EQ(5, PredicateMatch("123ABC"));
+}
+
+TEST(StringPredicateTest, IfNotPredicateStr) {
+  auto MatchFn = [](StringRef Str) { return Str == "abc" || Str == "def"; };
+  auto PredicateMatch = [&](StringRef S) {
+    return llvm::StringSwitch<int>(S)
+        .Case("abc", 0)
+        .IfNotPredicate(MatchFn, 1)
+        .Case("def", 2)
+        .Default(3);
+  };
+  EXPECT_EQ(0, PredicateMatch("abc"));
+  EXPECT_EQ(1, PredicateMatch("ghi"));
+  EXPECT_EQ(2, PredicateMatch("def"));
+}
+
 TEST(StringSwitchTest, DefaultUnreachable) {
   auto Translate = [](StringRef S) {
     return llvm::StringSwitch<int>(S)

``````````

</details>


https://github.com/llvm/llvm-project/pull/188046