Use trigrams to speed up SpecialCaseList.

[krasin]

Summary

it's often the case when the rules in the SpecialCaseList
are of the form hel.o*bar. That gives us a chance to build
trigram index to quickly discard 99% of inputs without
running a full regex. A similar idea was used in Google Code Search
as described in the blog post:
https://swtch.com/~rsc/regexp/regexp4.html

The check is defeated, if there's at least one regex
more complicated than that. In this case, all inputs
will go through the regex. That said, the real-world
rules are often simple or can be simplied. That considerably
speeds up compiling Chromium with CFI and UBSan.

As measured on Chromium's content_message_generator.cc:

before, CFI: 44 s
after, CFI: 23 s
after, CFI, no blacklist: 23 s (~1% slower, but 3 runs were unable to show the difference)
after, regular compilation to bitcode: 23 s

[krasin]

Add test cases for trigrams into SpecialCaseList unit test.

[krasin]

Simplify code; ensure build() is called.

[krasin]

Remove debug leftovers.

[krasin]

Get rid of build()

[krasin]

Speed improvements.

[pcc]

Nice!

Do you have before/after performance numbers for a large Chromium TU? Maybe also measure without the blacklist for a baseline.

lib/Support/Regex.cpp

199

Does this belong in Regex? It feels more like an implementation detail of TrigramIndex.

lib/Support/SpecialCaseList.cpp

34

Maybe use a SmallVector instead of std::vector since the size is bounded?

This code also needs more comments. Please at least leave comments for TrigramIndex and its fields.

47

Maybe fold this into the use? I found it a little confusing to use a variable named "I" for something other than a loop counter.

50

Maybe use a range for loop here?

67

I wonder how important it is to test Was here. It looks like things would still work if you allow trigrams to have multiple identical indices, and as you mention in the comment above it isn't worth worrying about popular trigrams.

82

Remove braces

[krasin]

Peter, thank you for your comments. I have added some benchmarks to the description and will address the rest of your points tomorrow.

[krasin]

Address the comments.

[krasin]

Please, take another look.

lib/Support/Regex.cpp

199

It belongs to Regex, because it uses RegexAdvancedMetachars constant that is very similar to RegexMetachars. If one is updated, another must be updated as well.

Also, isLiteralERE and isSimpleWildcard are very similar in meaning. If one lives in Regex, why not another?

lib/Support/SpecialCaseList.cpp

34

SmallVector: done
Comments to TrigramIndex and its fields: done.

47

Inlined and added a comment.

67

Actually, pretty important. Consider a test case like the following:

blablablablabla
blabla

If we don't use Was, the index will look the following:

"bla" -> [0, 0, 0, 0]
"lab" -> [0, 0, 0, 0]
"abl" -> [0, 0, 0, 0]

and counts: [13, 0]. One of the counts is zero => check defeated => all queries go through a regex => SLOW.

In the current scheme the index will be:
"bla" -> [0, 1]
"lab" -> [0, 1]
"abl" -> [0, 1]
and counts: [13, 4]. That gives a way to filter out all queries without these trigrams.

A second issue is that when we count trigrams at eval time, duplicates will increase per-rule counts (while we still have non-inflated limits) => regexps are called much more often => SLOW

Unfortunately, while I have a test case like that, it does not really have a way to know if the check is defeated or not. I only test for correctness. It means that speed regressions like I have explained above are still possible without breaking any tests.

And I will need a separate header for the TrigramIndex to make a more granular unit test possible. I would be willing to extract it, if you think that is a good idea (as usual, I don't feel the conventions in the LLVM projects).

[pcc]

lib/Support/Regex.cpp

199

I doubt that we'd actually be changing the definition of what an ERE is, the specification comes from POSIX and people are relying on the current semantics. On the other hand, the set of metacharacters for isSimpleWildcard is a property of how powerful TrigramIndex is, and I'd certainly expect it to change if we made TrigramIndex more powerful.

Besides, only SpecialCaseList mangles * into .*, which would seem to be critical to how we define this character list (i.e. you wouldn't be able to extract the trigram "abc" from the pattern "abc*" without this mangling).

lib/Support/SpecialCaseList.cpp

31

"a heuristic"

52

This should specify that these are indices into Counts.

67

Okay, I understand why this is necessary now; I hadn't had the insight that duplicate indices would effectively square the per-rule count.

As for your first point: it would seem that the index would be defeated under the current implementation with something like this:

blabla
blablabla
blablablabla
blablablablabla
blablablablablabla

or perhaps more realistically, something like:

class1.h
class2.h
class3.h
class4.h
class.h

It isn't clear how likely this would be in practice, though, and I suspect that you could mitigate against the more realistic scenarios by sorting the patterns by length beforehand.

Regarding testing: yes, we generally tend to extract algorithmically tricky classes and write unit tests for them (see the tests for CFI and devirtualization internals in llvm/unittests/Transforms/IPO).

[krasin]

Extract TrigramIndex{.h,.cpp,Test.cpp}.

[krasin]

nits.

[krasin]

Please, take another look.

lib/Support/Regex.cpp

199

On the other hand, the set of metacharacters for isSimpleWildcard is a property of how powerful TrigramIndex is, and I'd certainly expect it to change if we made TrigramIndex more powerful.

This is a good argument. I buy it. Done.

lib/Support/SpecialCaseList.cpp

52

Sorry, I have lost the context here after moving the code around. Probably addressed, but I am not sure.

67

Actually, in

class1.h
class2.h
class3.h
class4.h

we have trigrams ss1, ss2, ss3, ss4 which will help us to keep it working.

Realistically, the current implementation works for CFI blacklist and almost works for UBSanVptr blacklist in Chromium. The latter fails due to using \+ in the regex, and I hope to address escaping in the next CL.

[krasin]

lib/Support/SpecialCaseList.cpp

67

By "fails" I mean that the corresponding TrigramIndex is defeated. The blacklist still works correctly, just slower.

[pcc]

lib/Support/SpecialCaseList.cpp

67

Unless I'm mistaken, after adding class[1-4].h, the trigrams would look like this:

"cla" -> [0, 1, 2, 3]
"las" -> [0, 1, 2, 3]
"ass" -> [0, 1, 2, 3]
"ss1" -> [0]
"ss2" -> [1]
"ss3" -> [2]
"ss4" -> [3]

When we add "class.h", we will try to add only the trigrams "cla", "las" and "ass", all of which would fail due to exceeding the bound, no?

[krasin]

lib/Support/SpecialCaseList.cpp

67

This is correct, and regexp will be executed for the query. But the check will still correctly predict that we don't need to do that for queries like "foo" or "bar". So, while the check is not exact, it's not defeated either.

[krasin]

lib/Support/SpecialCaseList.cpp

67

actually, not. The check will not execute query for class.h. I will add a test case to demonstrate that.

[krasin]

Add one more test case.

[pcc]

Maybe I'm not making myself clear. This test doesn't pass on my machine.

TEST_F(TrigramIndexTest, foo) {
  std::unique_ptr<TrigramIndex> TI = makeTrigramIndex(
      {"class1.h", "class2.h", "class3.h", "class4.h", "class.h"});
  EXPECT_FALSE(TI->isDefeated());
}

[krasin]

lib/Support/SpecialCaseList.cpp

67

What happens during class.h matching is the following: it considers "cla", "las", "ass" and makes counts for all 4 rules equal to 3. But they need to be at least 4 to trigger a check.

An example of a query that will go through is claclaass.h. In this case, the count will become 4, and isDefinitelyOut will return false.

[krasin]

In D27188#610048, @pcc wrote:

Maybe I'm not making myself clear. This test doesn't pass on my machine.

TEST_F(TrigramIndexTest, foo) {
  std::unique_ptr<TrigramIndex> TI = makeTrigramIndex(
      {"class1.h", "class2.h", "class3.h", "class4.h", "class.h"});
  EXPECT_FALSE(TI->isDefeated());
}

Ah, yes, I incorrectly read your example. I thought class.h was a query, not a rule. Yes, that's correct, and I will add this test case as well (although, I have a similar one already).

[krasin]

Add a test case proposed by Peter.

[pcc]

LGTM with nits

include/llvm/Support/TrigramIndex.h

1 ↗

(On Diff #79845)

80 cols

30–35 ↗

(On Diff #79845)

LLVM headers come before system headers

42 ↗

(On Diff #79845)

There's no need for a constructor on this class; you can let Counts be default initialized and use in-class initializers for the other members.

lib/Support/CMakeLists.txt

97–98 ↗

(On Diff #79845)

Sort alphabetically

lib/Support/SpecialCaseList.cpp

52

This pertained to the comment "// Index holds a list of rules indices for each trigram.". It doesn't look like you've changed this. Note that you can click the double left arrow in the top left to see comments in their original context.

67

To be clear, I think it's fine if the implementation isn't perfect; this is certainly a strict improvement over the status quo. I just wanted to point out one somewhat realistic case where this might fail.

lib/Support/TrigramIndex.cpp

1–15 ↗

(On Diff #79845)

This comment needs to be updated.

17–22 ↗

(On Diff #79845)

LLVM headers come before system headers.

24 ↗

(On Diff #79845)

We normally write using namespace llvm; instead of enclosing the entire file in a namespace.

unittests/Support/TrigramIndexTest.cpp

1 ↗

(On Diff #79845)

80 cols

10–15 ↗

(On Diff #79845)

LLVM headers come before system headers

[krasin]

Address final nits. Also sync.

[pcc]

Still LGTM

include/llvm/Support/TrigramIndex.h

65–66 ↗

(On Diff #79851)

Does std::unordered_map<unsigned, SmallVector<size_t, 4>> Index{256}; work?

lib/Support/SpecialCaseList.cpp

210

auto &Entry

[krasin]

Shorter in-class initialization (thx Peter!)

[krasin]

Format nits.

[krasin]

include/llvm/Support/TrigramIndex.h

65–66 ↗

(On Diff #79851)

Yes, it does. Thank you!

include/llvm/Support/Regex.h

public:
     /// string.
     std::string sub(StringRef Repl, StringRef String,
                     std::string *Error = nullptr);
 
     /// \brief If this function returns true, ^Str$ is an extended regular
     /// expression that matches Str and only Str.
     static bool isLiteralERE(StringRef Str);
 
+    /// \brief If this function returns true, Str is a simple wildcard with
+    /// regular chars and optional star and dot wildcards like foo*b.r. This is a
+    /// common case for sanitizers blacklists.
+    static bool isSimpleWildcard(StringRef Str);
+
     /// \brief Turn String into a regex by escaping its special characters.
     static std::string escape(StringRef String);
 
   private:
     struct llvm_regex *preg;
     int error;
   };
 }
 
 #endif // LLVM_SUPPORT_REGEX_H

lib/Support/Regex.cpp

std::string Regex::sub(StringRef Repl, StringRef String,
   // And finally the suffix.
   Res += StringRef(Matches[0].end(), String.end() - Matches[0].end());
 
   return Res;
 }
 
 // These are the special characters matched in functions like "p_ere_exp".
 static const char RegexMetachars[] = "()^$|*+?.[]\\{}";
+static const char RegexAdvancedMetachars[] = "()^$|+?[]\\{}";
 
 bool Regex::isLiteralERE(StringRef Str) {
   // Check for regex metacharacters.  This list was derived from our regex
   // implementation in regcomp.c and double checked against the POSIX extended
   // regular expression specification.
   return Str.find_first_of(RegexMetachars) == StringRef::npos;
 }
 
+bool Regex::isSimpleWildcard(StringRef Str) {
+  // Check for regex metacharacters other than '*' and '.'.
+  return Str.find_first_of(RegexAdvancedMetachars) == StringRef::npos;
+}
+
 std::string Regex::escape(StringRef String) {
   std::string RegexStr;
   for (unsigned i = 0, e = String.size(); i != e; ++i) {
     if (strchr(RegexMetachars, String[i]))
       RegexStr += '\\';
     RegexStr += String[i];
   }
 
   return RegexStr;
 }

lib/Support/SpecialCaseList.cpp

 //===----------------------------------------------------------------------===//
 
 #include "llvm/Support/SpecialCaseList.h"
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/ADT/StringExtras.h"
 #include "llvm/ADT/StringSet.h"
 #include "llvm/Support/MemoryBuffer.h"
 #include "llvm/Support/Regex.h"
+#include <unordered_map>
+#include <set>
 #include <string>
 #include <system_error>
 #include <utility>
 
 namespace llvm {
 
+// TrigramIndex implements a heuristics that allows to filter ~99% incoming
+// queries when all regular expressions in the SpecialCaseList are simple
+// wildcards with '*' and '.'. If rules are more complicated, the check is
+// defeated and it will always pass the queries to a full regex.
+// The basic idea is that in order for a wildcard to match a query, the query
+// needs to have all trigrams which exist in the wildcard. We create a trigram
+// index (trigram -> list of rules with it) and then count trigrams in the query
+// for each rule. If the count for one of the rules reaches the expected value,
+// the check passes the query to a regexp. If none of the rules got enough
+// trigrams, the check tells that the query is definitely not matched by any
+// of the rules, and no regex matching is needed.
+// A similar idea was used in Google Code Search as described in the blog post:
+// https://swtch.com/~rsc/regexp/regexp4.html
+struct TrigramIndex {
+  // If true, the rules are too complicated for the check to work, and full
+  // regex matching is needed for every rule.
+  bool Defeated;
+  // The minimum number of trigrams which should match for a rule to have a
+  // chance to match the query. The number of elements equals the number of
+  // regex rules in the SpecialCaseList.
+  std::vector<unsigned> Counts;
+  // Index holds a list of rules indices for each trigram.
+  // If a trigram is too common (>4 rules with it), we stop tracking it,
+  // which increases the probability for a need to match using regex, but
+  // decreases the costs in the regular case.
+  std::unordered_map<unsigned, SmallVector<size_t, 4>> Index;
+
+  TrigramIndex(): Index(256) {}
+
+  void insert(std::string Regex) {
+    if (Defeated) return;
+    if (!Regex::isSimpleWildcard(Regex)) {
+      Defeated = true;
+      return;
+    }
+
+    std::set<unsigned> Was;
+    unsigned Cnt = 0;
+    unsigned Tri = 0;
+    unsigned Len = 0;
+    for (unsigned Char : Regex) {
+      if (Char == '.' || Char == '*') {
+        Tri = 0;
+        Len = 0;
+        continue;
+      }
+      Tri = ((Tri << 8) + Char) & 0xFFFFFF;
+      Len++;
+      if (Len < 3)
+        continue;
+      // We don't want the index to grow too much for the popular trigrams,
+      // as they are weak signals. It's ok to still require them for the
+      // rules we have already processed. It's just a small additional
+      // computational cost.
+      if (Index[Tri].size() >= 4)
+        continue;
+      Cnt++;
+      if (!Was.count(Tri)) {
+        // Adding the current rule to the index.
+        Index[Tri].push_back(Counts.size());
+        Was.insert(Tri);
+      }
+    }
+    if (!Cnt) {
+      // This rule does not have remarkable trigrams to rely on.
+      // We have to always call the full regex chain.
+      Defeated = true;
+      return;
+    }
+    Counts.push_back(Cnt);
+  }
+
+  bool isDefinitelyOut(StringRef Query) const {
+    if (Defeated)
+      return false;
+    std::vector<unsigned> CurCounts(Counts.size());
+    unsigned Tri = 0;
+    for (size_t I = 0; I < Query.size(); I++) {
+      Tri = ((Tri << 8) + Query[I]) & 0xFFFFFF;
+      if (I < 2)
+        continue;
+      const auto &II = Index.find(Tri);
+      if (II == Index.end())
+        continue;
+      for (size_t J : II->second) {
+        CurCounts[J]++;
+        // If we have reached a desired limit, we have to look at the query
+        // more closely by running a full regex.
+        if (CurCounts[J] >= Counts[J])
+          return false;
+      }
+    }
+    return true;
+  }
+};
+
 /// Represents a set of regular expressions.  Regular expressions which are
 /// "literal" (i.e. no regex metacharacters) are stored in Strings, while all
 /// others are represented as a single pipe-separated regex in RegEx.  The
 /// reason for doing so is efficiency; StringSet is much faster at matching
 /// literal strings than Regex.
 struct SpecialCaseList::Entry {
   StringSet<> Strings;
+  TrigramIndex Trigrams;
   std::unique_ptr<Regex> RegEx;
 
   bool match(StringRef Query) const {
-    return Strings.count(Query) || (RegEx && RegEx->match(Query));
+    if (Strings.count(Query))
+      return true;
+    if (Trigrams.isDefinitelyOut(Query))
+      return false;
+    return RegEx && RegEx->match(Query);
   }
 };
 
 SpecialCaseList::SpecialCaseList() : Entries(), Regexps(), IsCompiled(false) {}
 
 std::unique_ptr<SpecialCaseList>
 SpecialCaseList::create(const std::vector<std::string> &Paths,
                         std::string &Error) {
if (SplitLine.second.empty()) {
       return false;
     }
 
     std::pair<StringRef, StringRef> SplitRegexp = SplitLine.second.split("=");
     std::string Regexp = SplitRegexp.first;
     StringRef Category = SplitRegexp.second;
 
     // See if we can store Regexp in Strings.
+    auto& Entry = Entries[Prefix][Category];
     if (Regex::isLiteralERE(Regexp)) {
-      Entries[Prefix][Category].Strings.insert(Regexp);
+      Entry.Strings.insert(Regexp);
       continue;
     }
+    Entry.Trigrams.insert(Regexp);
 
     // Replace * with .*
     for (size_t pos = 0; (pos = Regexp.find("*", pos)) != std::string::npos;
          pos += strlen(".*")) {
       Regexp.replace(pos, strlen("*"), ".*");
     }
 
     // Check that the regexp is valid.

unittests/Support/SpecialCaseListTest.cpp

TEST_F(SpecialCaseListTest, MultipleBlacklists) {
   EXPECT_FALSE(SCL->inSection("src", "ban"));
   EXPECT_TRUE(SCL->inSection("src", "ban", "init"));
   EXPECT_TRUE(SCL->inSection("src", "tomfoolery"));
   EXPECT_TRUE(SCL->inSection("src", "tomfoglery"));
   for (auto &Path : Files)
     sys::fs::remove(Path);
 }
 
+TEST_F(SpecialCaseListTest, NoTrigramsInRules) {
+  std::unique_ptr<SpecialCaseList> SCL = makeSpecialCaseList("fun:b.r\n"
+                                                             "fun:za*az\n");
+  EXPECT_TRUE(SCL->inSection("fun", "bar"));
+  EXPECT_FALSE(SCL->inSection("fun", "baz"));
+  EXPECT_TRUE(SCL->inSection("fun", "zakaz"));
+  EXPECT_FALSE(SCL->inSection("fun", "zaraza"));
+}
+
+TEST_F(SpecialCaseListTest, NoTrigramsInARule) {
+  std::unique_ptr<SpecialCaseList> SCL = makeSpecialCaseList("fun:*bar*\n"
+                                                             "fun:za*az\n");
+  EXPECT_TRUE(SCL->inSection("fun", "abara"));
+  EXPECT_FALSE(SCL->inSection("fun", "bor"));
+  EXPECT_TRUE(SCL->inSection("fun", "zakaz"));
+  EXPECT_FALSE(SCL->inSection("fun", "zaraza"));
+}
+
+TEST_F(SpecialCaseListTest, RepetitiveRule) {
+  std::unique_ptr<SpecialCaseList> SCL = makeSpecialCaseList("fun:*bar*bar*bar*bar*\n"
+                                                             "fun:bar*\n");
+  EXPECT_TRUE(SCL->inSection("fun", "bara"));
+  EXPECT_FALSE(SCL->inSection("fun", "abara"));
+  EXPECT_TRUE(SCL->inSection("fun", "barbarbarbar"));
+  EXPECT_TRUE(SCL->inSection("fun", "abarbarbarbar"));
+  EXPECT_FALSE(SCL->inSection("fun", "abarbarbar"));
+}
+
+TEST_F(SpecialCaseListTest, SpecialSymbolRule) {
+  std::unique_ptr<SpecialCaseList> SCL = makeSpecialCaseList("src:*c\\+\\+abi*\n");
+  EXPECT_TRUE(SCL->inSection("src", "c++abi"));
+  EXPECT_FALSE(SCL->inSection("src", "c\\+\\+abi"));
+}
+
+TEST_F(SpecialCaseListTest, PopularTrigram) {
+  std::unique_ptr<SpecialCaseList> SCL = makeSpecialCaseList("fun:*aaaaaa*\n"
+                                                             "fun:*aaaaa*\n"
+                                                             "fun:*aaaa*\n"
+                                                             "fun:*aaa*\n");
+  EXPECT_TRUE(SCL->inSection("fun", "aaa"));
+  EXPECT_TRUE(SCL->inSection("fun", "aaaa"));
+  EXPECT_TRUE(SCL->inSection("fun", "aaaabbbaaa"));
+}
+
 }