1 files changed, 21 insertions, 17 deletions

diff --git a/llvm/lib/ProfileData/InstrProf.cpp b/llvm/lib/ProfileData/InstrProf.cpp
index 86cd76c9a20..54c37e30dba 100644
--- a/llvm/lib/ProfileData/InstrProf.cpp
+++ b/llvm/lib/ProfileData/InstrProf.cpp
@@ -442,6 +442,24 @@ static std::unique_ptr<ValueProfData> allocValueProfData(uint32_t TotalSize) {
                                             ValueProfData());
 }
 
+instrprof_error ValueProfData::checkIntegrity() {
+  if (NumValueKinds > IPVK_Last + 1)
+    return instrprof_error::malformed;
+  // Total size needs to be mulltiple of quadword size.
+  if (TotalSize % sizeof(uint64_t))
+    return instrprof_error::malformed;
+
+  ValueProfRecord *VR = getFirstValueProfRecord(this);
+  for (uint32_t K = 0; K < this->NumValueKinds; K++) {
+    if (VR->Kind > IPVK_Last)
+      return instrprof_error::malformed;
+    VR = getValueProfRecordNext(VR);
+    if ((char *)VR - (char *)this > (ptrdiff_t)TotalSize)
+      return instrprof_error::malformed;
+  }
+  return instrprof_error::success;
+}
+
 ErrorOr<std::unique_ptr<ValueProfData>>
 ValueProfData::getValueProfData(const unsigned char *D,
                                 const unsigned char *const BufferEnd,
@@ -452,31 +470,17 @@ ValueProfData::getValueProfData(const unsigned char *D,
 
   const unsigned char *Header = D;
   uint32_t TotalSize = swapToHostOrder<uint32_t>(Header, Endianness);
-  uint32_t NumValueKinds = swapToHostOrder<uint32_t>(Header, Endianness);
-
   if (D + TotalSize > BufferEnd)
     return instrprof_error::too_large;
-  if (NumValueKinds > IPVK_Last + 1)
-    return instrprof_error::malformed;
-  // Total size needs to be mulltiple of quadword size.
-  if (TotalSize % sizeof(uint64_t))
-    return instrprof_error::malformed;
 
   std::unique_ptr<ValueProfData> VPD = allocValueProfData(TotalSize);
-
   memcpy(VPD.get(), D, TotalSize);
   // Byte swap.
   VPD->swapBytesToHost(Endianness);
 
-  // Data integrity check:
-  ValueProfRecord *VR = getFirstValueProfRecord(VPD.get());
-  for (uint32_t K = 0; K < VPD->NumValueKinds; K++) {
-    if (VR->Kind > IPVK_Last)
-      return instrprof_error::malformed;
-    VR = getValueProfRecordNext(VR);
-    if ((char *)VR - (char *)VPD.get() > (ptrdiff_t)TotalSize)
-      return instrprof_error::malformed;
-  }
+  instrprof_error EC = VPD->checkIntegrity();
+  if (EC != instrprof_error::success)
+    return EC;
 
   return std::move(VPD);
 }