1 files changed, 41 insertions, 8 deletions
diff --git a/lib/Basic/DiagnosticIDs.cpp b/lib/Basic/DiagnosticIDs.cpp
index 1237fd5db1..4716a449a7 100644
--- a/lib/Basic/DiagnosticIDs.cpp
+++ b/lib/Basic/DiagnosticIDs.cpp
@@ -107,16 +107,49 @@ static const StaticDiagInfoRec *GetDiagInfo(unsigned DiagID) {
   }
 #endif
 
-  // Search the diagnostic table with a binary search.
-  StaticDiagInfoRec Find = { static_cast<unsigned short>(DiagID),
-                             0, 0, 0, 0, 0, 0, 0, 0, 0, 0};
-
-  const StaticDiagInfoRec *Found =
-    std::lower_bound(StaticDiagInfo, StaticDiagInfo + StaticDiagInfoSize, Find);
-  if (Found == StaticDiagInfo + StaticDiagInfoSize ||
-      Found->DiagID != DiagID)
+  // Out of bounds diag. Can't be in the table.
+  using namespace diag;
+  if (DiagID >= DIAG_UPPER_LIMIT)
     return 0;
 
+  // Compute the index of the requested diagnostic in the static table.
+  // 1. Add the number of diagnostics in each category preceeding the
+  //    diagnostic and of the category the diagnostic is in. This gives us
+  //    the offset of the category in the table.
+  // 2. Subtract the number of IDs in each category from our ID. This gives us
+  //    the offset of the diagnostic in the category.
+  // This is cheaper than a binary search on the table as it doesn't touch
+  // memory at all.
+  unsigned Offset = 0;
+  unsigned ID = DiagID;
+#define DIAG_START_COMMON 0 // Sentinel value.
+#define CATEGORY(NAME, PREV) \
+  if (DiagID > DIAG_START_##NAME) { \
+    Offset += NUM_BUILTIN_##PREV##_DIAGNOSTICS - DIAG_START_##PREV; \
+    ID -= DIAG_START_##NAME - DIAG_START_##PREV + 1; \
+  }
+CATEGORY(DRIVER, COMMON)
+CATEGORY(FRONTEND, DRIVER)
+CATEGORY(SERIALIZATION, FRONTEND)
+CATEGORY(LEX, SERIALIZATION)
+CATEGORY(PARSE, LEX)
+CATEGORY(AST, PARSE)
+CATEGORY(COMMENT, AST)
+CATEGORY(SEMA, COMMENT)
+CATEGORY(ANALYSIS, SEMA)
+#undef CATEGORY
+#undef DIAG_START_COMMON
+
+  // Avoid out of bounds reads.
+  if (ID + Offset >= StaticDiagInfoSize)
+    return 0;
+
+  const StaticDiagInfoRec *Found = &StaticDiagInfo[ID + Offset];
+  // If the diag id doesn't match we found a different diag, abort. This can
+  // happen when this function is called with an ID that points into a hole in
+  // the diagID space.
+  if (Found->DiagID != DiagID)
+    return 0;
   return Found;
 }