2 files changed, 31 insertions, 10 deletions

diff --git a/llvm/lib/ProfileData/SampleProfReader.cpp b/llvm/lib/ProfileData/SampleProfReader.cpp
index d184969e046..899343f72f7 100644
--- a/llvm/lib/ProfileData/SampleProfReader.cpp
+++ b/llvm/lib/ProfileData/SampleProfReader.cpp
@@ -100,6 +100,12 @@ static bool ParseHead(const StringRef &Input, StringRef &FName,
   return true;
 }
 
+
+/// \brief Returns true if line offset \p L is legal (only has 16 bits).
+static bool isOffsetLegal(unsigned L) {
+  return (L & 0xffff) == L;
+}
+
 /// \brief Parse \p Input as line sample.
 ///
 /// \param Input input line.
@@ -124,7 +130,7 @@ static bool ParseLine(const StringRef &Input, bool &IsCallsite, uint32_t &Depth,
   StringRef Loc = Input.substr(Depth, n1 - Depth);
   size_t n2 = Loc.find('.');
   if (n2 == StringRef::npos) {
-    if (Loc.getAsInteger(10, LineOffset))
+    if (Loc.getAsInteger(10, LineOffset) || !isOffsetLegal(LineOffset))
       return false;
     Discriminator = 0;
   } else {
@@ -308,6 +314,10 @@ SampleProfileReaderBinary::readProfile(FunctionSamples &FProfile) {
     if (std::error_code EC = LineOffset.getError())
       return EC;
 
+    if (!isOffsetLegal(*LineOffset)) {
+      return std::error_code();
+    }
+
     auto Discriminator = readNumber<uint64_t>();
     if (std::error_code EC = Discriminator.getError())
       return EC;
diff --git a/llvm/lib/Transforms/IPO/SampleProfile.cpp b/llvm/lib/Transforms/IPO/SampleProfile.cpp
index e0691515c95..b797321b2f8 100644
--- a/llvm/lib/Transforms/IPO/SampleProfile.cpp
+++ b/llvm/lib/Transforms/IPO/SampleProfile.cpp
@@ -122,6 +122,7 @@ protected:
   void buildEdges(Function &F);
   bool propagateThroughEdges(Function &F);
   void computeDominanceAndLoopInfo(Function &F);
+  unsigned getOffset(unsigned L, unsigned H) const;
 
   /// \brief Map basic blocks to their computed weights.
   ///
@@ -174,6 +175,17 @@ protected:
 };
 }
 
+/// \brief Returns the offset of lineno \p L to head_lineno \p H
+///
+/// \param L  Lineno
+/// \param H  Header lineno of the function
+///
+/// \returns offset to the header lineno. 16 bits are used to represent offset.
+/// We assume that a single function will not exceed 65535 LOC.
+unsigned SampleProfileLoader::getOffset(unsigned L, unsigned H) const {
+  return (L - H) & 0xffff;
+}
+
 /// \brief Print the weight of edge \p E on stream \p OS.
 ///
 /// \param OS  Stream to emit the output to.
@@ -229,11 +241,9 @@ SampleProfileLoader::getInstWeight(const Instruction &Inst) const {
   const DILocation *DIL = DLoc;
   unsigned Lineno = DLoc.getLine();
   unsigned HeaderLineno = DIL->getScope()->getSubprogram()->getLine();
-  if (Lineno < HeaderLineno)
-    return std::error_code();
 
-  ErrorOr<uint64_t> R =
-      FS->findSamplesAt(Lineno - HeaderLineno, DIL->getDiscriminator());
+  ErrorOr<uint64_t> R = FS->findSamplesAt(getOffset(Lineno, HeaderLineno),
+                                          DIL->getDiscriminator());
   if (R)
     DEBUG(dbgs() << "    " << Lineno << "." << DIL->getDiscriminator() << ":"
                  << Inst << " (line offset: " << Lineno - HeaderLineno << "."
@@ -308,7 +318,7 @@ SampleProfileLoader::findCalleeFunctionSamples(const CallInst &Inst) const {
     return nullptr;
   }
   DISubprogram *SP = DIL->getScope()->getSubprogram();
-  if (!SP || DIL->getLine() < SP->getLine())
+  if (!SP)
     return nullptr;
 
   Function *CalleeFunc = Inst.getCalledFunction();
@@ -321,8 +331,9 @@ SampleProfileLoader::findCalleeFunctionSamples(const CallInst &Inst) const {
   if (FS == nullptr)
     return nullptr;
 
-  return FS->findFunctionSamplesAt(CallsiteLocation(
-      DIL->getLine() - SP->getLine(), DIL->getDiscriminator(), CalleeName));
+  return FS->findFunctionSamplesAt(
+      CallsiteLocation(getOffset(DIL->getLine(), SP->getLine()),
+                       DIL->getDiscriminator(), CalleeName));
 }
 
 /// \brief Get the FunctionSamples for an instruction.
@@ -345,10 +356,10 @@ SampleProfileLoader::findFunctionSamples(const Instruction &Inst) const {
   for (const DILocation *DIL = Inst.getDebugLoc(); DIL;
        DIL = DIL->getInlinedAt()) {
     DISubprogram *SP = DIL->getScope()->getSubprogram();
-    if (!SP || DIL->getLine() < SP->getLine())
+    if (!SP)
       return nullptr;
     if (!CalleeName.empty()) {
-      S.push_back(CallsiteLocation(DIL->getLine() - SP->getLine(),
+      S.push_back(CallsiteLocation(getOffset(DIL->getLine(), SP->getLine()),
                                    DIL->getDiscriminator(), CalleeName));
     }
     CalleeName = SP->getLinkageName();