[BOLT][NFC] Fix undefined behavior in encodeAnnotationImm

Fix UBSan-reported issue in MCPlusBuilder::encodeAnnotationImm (left shift of a
negative value).

Test Plan:
```
ninja check-bolt
...
PASS: BOLT-Unit :: Core/./CoreTests/AArch64/MCPlusBuilderTester.Annotation/0 (1 of 140)
PASS: BOLT-Unit :: Core/./CoreTests/X86/MCPlusBuilderTester.Annotation/0 (131 of 134)
```

Reviewed By: maksfb, yota9

Differential Revision: https://reviews.llvm.org/D120260

3 files changed, 32 insertions, 6 deletions

diff --git a/bolt/include/bolt/Core/MCPlusBuilder.h b/bolt/include/bolt/Core/MCPlusBuilder.h
index 3876b6e379da..93970e02a4d1 100644
--- a/bolt/include/bolt/Core/MCPlusBuilder.h
+++ b/bolt/include/bolt/Core/MCPlusBuilder.h
@@ -74,9 +74,9 @@ private:
   AllocatorIdTy MaxAllocatorId = 0;
 
   /// We encode Index and Value into a 64-bit immediate operand value.
-  static int64_t encodeAnnotationImm(unsigned Index, int64_t Value) {
-    assert(Index < 256 && "annotation index max value exceeded");
-    assert((Value == (Value << 8) >> 8) && "annotation value out of range");
+  static int64_t encodeAnnotationImm(uint8_t Index, int64_t Value) {
+    if (LLVM_UNLIKELY(Value != extractAnnotationValue(Value)))
+      report_fatal_error("annotation value out of range");
 
     Value &= 0xff'ffff'ffff'ffff;
     Value |= (int64_t)Index << 56;
@@ -85,14 +85,13 @@ private:
   }
 
   /// Extract annotation index from immediate operand value.
-  static unsigned extractAnnotationIndex(int64_t ImmValue) {
+  static uint8_t extractAnnotationIndex(int64_t ImmValue) {
     return ImmValue >> 56;
   }
 
   /// Extract annotation value from immediate operand value.
   static int64_t extractAnnotationValue(int64_t ImmValue) {
-    ImmValue &= 0xff'ffff'ffff'ffff;
-    return (ImmValue << 8) >> 8;
+    return SignExtend64<56>(ImmValue & 0xff'ffff'ffff'ffffULL);
   }
 
   MCInst *getAnnotationInst(const MCInst &Inst) const {
diff --git a/bolt/unittests/Core/CMakeLists.txt b/bolt/unittests/Core/CMakeLists.txt
index 1516eb4a6762..7abb0bea041a 100644
--- a/bolt/unittests/Core/CMakeLists.txt
+++ b/bolt/unittests/Core/CMakeLists.txt
@@ -3,6 +3,7 @@ set(LLVM_LINK_COMPONENTS
   BOLTRewrite
   DebugInfoDWARF
   Object
+  MC
   ${LLVM_TARGETS_TO_BUILD}
   )
 
diff --git a/bolt/unittests/Core/MCPlusBuilder.cpp b/bolt/unittests/Core/MCPlusBuilder.cpp
index 2158f652c071..a1ebde8c1123 100644
--- a/bolt/unittests/Core/MCPlusBuilder.cpp
+++ b/bolt/unittests/Core/MCPlusBuilder.cpp
@@ -110,3 +110,29 @@ TEST_P(MCPlusBuilderTester, AliasSmallerAX) {
 }
 
 #endif // X86_AVAILABLE
+
+TEST_P(MCPlusBuilderTester, Annotation) {
+  MCInst Inst;
+  bool Success = BC->MIB->createTailCall(Inst, BC->Ctx->createNamedTempSymbol(),
+                                         BC->Ctx.get());
+  ASSERT_TRUE(Success);
+  MCSymbol *LPSymbol = BC->Ctx->createNamedTempSymbol("LP");
+  uint64_t Value = INT32_MIN;
+  // Test encodeAnnotationImm using this indirect way
+  BC->MIB->addEHInfo(Inst, MCPlus::MCLandingPad(LPSymbol, Value));
+  // Round-trip encoding-decoding check for negative values
+  Optional<MCPlus::MCLandingPad> EHInfo = BC->MIB->getEHInfo(Inst);
+  ASSERT_TRUE(EHInfo.hasValue());
+  MCPlus::MCLandingPad LP = EHInfo.getValue();
+  uint64_t DecodedValue = LP.second;
+  ASSERT_EQ(Value, DecodedValue);
+
+  // Large int64 should trigger an out of range assertion
+  Value = 0x1FF'FFFF'FFFF'FFFFULL;
+  Inst.clear();
+  Success = BC->MIB->createTailCall(Inst, BC->Ctx->createNamedTempSymbol(),
+                                    BC->Ctx.get());
+  ASSERT_TRUE(Success);
+  ASSERT_DEATH(BC->MIB->addEHInfo(Inst, MCPlus::MCLandingPad(LPSymbol, Value)),
+               "annotation value out of range");
+}