Lex: Check buckets on header map construction

If the number of buckets is not a power of two, immediately recognize
the header map as corrupt, rather than waiting for the first lookup.  I
converted the later check to an assert.

git-svn-id: https://llvm.org/svn/llvm-project/cfe/trunk@261448 91177308-0d34-0410-b5e6-96231b3b80d8

lib/Lex/HeaderMap.cpp

@@ -19,6 +19,7 @@
 #include "llvm/Support/DataTypes.h"
 #include "llvm/Support/MathExtras.h"
 #include "llvm/Support/MemoryBuffer.h"
+#include "llvm/Support/SwapByteOrder.h"
 #include <cstdio>
 #include <memory>
 using namespace clang;
@@ -82,6 +83,15 @@ bool HeaderMapImpl::checkHeader(const llvm::MemoryBuffer &File,
   if (Header->Reserved != 0)
     return false;
 
+  // Check the number of buckets.
+  auto NumBuckets = NeedsByteSwap
+                        ? llvm::sys::getSwappedBytes(Header->NumBuckets)
+                        : Header->NumBuckets;
+
+  // If the number of buckets is not a power of two, the headermap is corrupt.
+  if (NumBuckets & (NumBuckets - 1))
+    return false;
+
   // Okay, everything looks good.
   return true;
 }
@@ -191,10 +201,8 @@ StringRef HeaderMapImpl::lookupFilename(StringRef Filename,
   const HMapHeader &Hdr = getHeader();
   unsigned NumBuckets = getEndianAdjustedWord(Hdr.NumBuckets);
 
-  // If the number of buckets is not a power of two, the headermap is corrupt.
+  // Don't probe infinitely.  This should be checked before constructing.
-  // Don't probe infinitely.
+  assert(!(NumBuckets & (NumBuckets - 1)) && "Expected power of 2");
-  if (NumBuckets & (NumBuckets-1))
-    return StringRef();
 
   // Linearly probe the hash table.
   for (unsigned Bucket = HashHMapKey(Filename);; ++Bucket) {

unittests/Lex/HeaderMapTest.cpp

@@ -91,4 +91,13 @@ TEST(HeaderMapTest, checkHeaderValidButEmpty) {
   ASSERT_TRUE(NeedsSwap);
 }
 
+TEST(HeaderMapTest, checkHeader3Buckets) {
+  MapFile<3, 1> File;
+  ASSERT_EQ(3 * sizeof(HMapBucket), sizeof(File.Buckets));
+
+  File.init();
+  bool NeedsSwap;
+  ASSERT_FALSE(HeaderMapImpl::checkHeader(*File.getBuffer(), NeedsSwap));
+}
+
 } // end namespace