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

33

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.

46

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.

49

Maybe use a range for loop here?

66

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.

81

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

33

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

46

Inlined and added a comment.

66

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

30

"a heuristic"

51

This should specify that these are indices into Counts.

66

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

51

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

66

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

66

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

[pcc]

lib/Support/SpecialCaseList.cpp

66

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

66

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

66

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

66

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

51

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.

66

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

192

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 <map>
 #include <string>
 #include <system_error>
 #include <utility>
 
 namespace llvm {
 
+struct TrigramIndex {
+  bool Defeated;
+  bool IsBuilt;
+  std::vector<std::string> Rules;
+  std::vector<int> Counts;
+  std::map<std::string, std::vector<int>> Index;
+
+  void insert(std::string Regex) {
+    if (Defeated) return;
+    if (!Regex::isSimpleWildcard(Regex)) {
+      Defeated = true;
+      fprintf(stderr, "TrigramIndex is defeated (loc3)\n");
+      return;
+    }
+    Rules.push_back(Regex);
+  }
+
+  void build() {
+    assert(!IsBuilt && "build() should only be called once");
+    IsBuilt = true;
+    Counts.resize(Rules.size());
+    for (size_t I = 0; I < Rules.size(); I++) {
+      const auto& Rule = Rules[I];
+      int Cnt = 0;
+      StringSet<> Was;
+      for (size_t J = 0; J < Rule.size() - 2; J++) {
+        std::string Tri = Rule.substr(J, 3);
+        if (Tri.find_first_of(".*") != std::string::npos)
+          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)) {
+          Index[Tri].push_back(I);
+          Was.insert(Tri);
+        }
+      }
+      if (!Cnt) {
+        // This rule does not have remarkable trigrams to rely on.
+        // We have to always call the full regexp chain.
+        fprintf(stderr, "TrigramIndex is defeated (loc2)\n");
+        Defeated = true;
+        return;
+      }
+      Counts[I] = Cnt;
+    }
+  }
+
+  bool isDefinitelyOut(StringRef Query) const {
+    assert(IsBuilt && "TrigramIndex::build() was not called!");
+    //fprintf(stderr, "TrigramIndex:isDefinitelyOut(%s)\n", Query.str().c_str());
+    if (Defeated) {
+      //fprintf(stderr, "TrigramIndex is defeated\n");
+      return false;
+    }
+    std::vector<int> CurCounts(Rules.size());
+    for (int I = 0; I < static_cast<int>(Query.size()) - 2; I++) {
+      std::string Tri = Query.substr(I, 3);
+      if (Tri.find_first_of(".*") != std::string::npos || !Index.count(Tri))
+        continue;
+      for (int J : Index.find(Tri)->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]) {
+          // fprintf(stderr, "TrigramIndex:isDefinitelyOut(%s) is FALSE\n", Query.str().c_str());
+          return false;
+        }
+      }
+    }
+    // fprintf(stderr, "TrigramIndex:isDefinitelyOut(%s) is TRUE\n", Query.str().c_str());
+    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;
-    }
+    } else
+      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.
for (StringMap<StringMap<std::string>>::const_iterator I = Regexps.begin(),
        I != E; ++I) {
     for (StringMap<std::string>::const_iterator II = I->second.begin(),
                                                 IE = I->second.end();
          II != IE; ++II) {
       Entries[I->getKey()][II->getKey()].RegEx.reset(new Regex(II->getValue()));
     }
   }
   Regexps.clear();
+
+  // Build trigram indices
+  for (auto& I : Entries)
+    for (auto& II : I.getValue())
+      II.getValue().Trigrams.build();
   IsCompiled = true;
 }
 
 SpecialCaseList::~SpecialCaseList() {}
 
 bool SpecialCaseList::inSection(StringRef Section, StringRef Query,
                                 StringRef Category) const {
   assert(IsCompiled && "SpecialCaseList::compile() was not called!");

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"));
+}
+
 }