1 files changed, 39 insertions, 4 deletions
diff --git a/llvm/lib/ProfileData/InstrProfWriter.cpp b/llvm/lib/ProfileData/InstrProfWriter.cpp
index d07668322354..987c0b175d3c 100644
--- a/llvm/lib/ProfileData/InstrProfWriter.cpp
+++ b/llvm/lib/ProfileData/InstrProfWriter.cpp
@@ -285,7 +285,7 @@ static void setSummary(IndexedInstrProf::Summary *TheSummary,
     TheSummary->setEntry(I, Res[I]);
 }
 
-void InstrProfWriter::writeImpl(ProfOStream &OS) {
+Error InstrProfWriter::writeImpl(ProfOStream &OS) {
   using namespace IndexedInstrProf;
 
   OnDiskChainedHashTableGenerator<InstrProfRecordWriterTrait> Generator;
@@ -376,12 +376,19 @@ void InstrProfWriter::writeImpl(ProfOStream &OS) {
        (int)CSSummarySize}};
 
   OS.patch(PatchItems, sizeof(PatchItems) / sizeof(*PatchItems));
+
+  for (const auto &I : FunctionData)
+    for (const auto &F : I.getValue())
+      if (Error E = validateRecord(F.second))
+        return E;
+
+  return Error::success();
 }
 
-void InstrProfWriter::write(raw_fd_ostream &OS) {
+Error InstrProfWriter::write(raw_fd_ostream &OS) {
   // Write the hash table.
   ProfOStream POS(OS);
-  writeImpl(POS);
+  return writeImpl(POS);
 }
 
 std::unique_ptr<MemoryBuffer> InstrProfWriter::writeBuffer() {
@@ -389,7 +396,8 @@ std::unique_ptr<MemoryBuffer> InstrProfWriter::writeBuffer() {
   raw_string_ostream OS(Data);
   ProfOStream POS(OS);
   // Write the hash table.
-  writeImpl(POS);
+  if (Error E = writeImpl(POS))
+    return nullptr;
   // Return this in an aligned memory buffer.
   return MemoryBuffer::getMemBufferCopy(Data);
 }
@@ -399,6 +407,27 @@ static const char *ValueProfKindStr[] = {
 #include "llvm/ProfileData/InstrProfData.inc"
 };
 
+Error InstrProfWriter::validateRecord(const InstrProfRecord &Func) {
+  for (uint32_t VK = 0; VK <= IPVK_Last; VK++) {
+    uint32_t NS = Func.getNumValueSites(VK);
+    if (!NS)
+      continue;
+    for (uint32_t S = 0; S < NS; S++) {
+      uint32_t ND = Func.getNumValueDataForSite(VK, S);
+      std::unique_ptr<InstrProfValueData[]> VD = Func.getValueForSite(VK, S);
+      bool WasZero = false;
+      for (uint32_t I = 0; I < ND; I++)
+        if ((VK != IPVK_IndirectCallTarget) && (VD[I].Value == 0)) {
+          if (WasZero)
+            return make_error<InstrProfError>(instrprof_error::invalid_prof);
+          WasZero = true;
+        }
+    }
+  }
+
+  return Error::success();
+}
+
 void InstrProfWriter::writeRecordInText(StringRef Name, uint64_t Hash,
                                         const InstrProfRecord &Func,
                                         InstrProfSymtab &Symtab,
@@ -473,5 +502,11 @@ Error InstrProfWriter::writeText(raw_fd_ostream &OS) {
     writeRecordInText(Name, Func.first, Func.second, Symtab, OS);
   }
 
+  for (const auto &record : OrderedFuncData) {
+    const FuncPair &Func = record.second;
+    if (Error E = validateRecord(Func.second))
+      return E;
+  }
+
   return Error::success();
 }