2 files changed, 52 insertions, 7 deletions

diff --git a/llvm/include/llvm/ADT/PointerEmbeddedInt.h b/llvm/include/llvm/ADT/PointerEmbeddedInt.h
index 9562819f283..546597f4577 100644
--- a/llvm/include/llvm/ADT/PointerEmbeddedInt.h
+++ b/llvm/include/llvm/ADT/PointerEmbeddedInt.h
@@ -11,6 +11,7 @@
 #define LLVM_ADT_POINTEREMBEDDEDINT_H
 
 #include "llvm/ADT/DenseMapInfo.h"
+#include "llvm/Support/MathExtras.h"
 #include "llvm/Support/PointerLikeTypeTraits.h"
 #include <climits>
 
@@ -30,6 +31,8 @@ template <typename IntT, int Bits = sizeof(IntT) * CHAR_BIT>
 class PointerEmbeddedInt {
   uintptr_t Value;
 
+  // Note: This '<' is correct; using '<=' would result in some shifts
+  // overflowing their storage types.
   static_assert(Bits < sizeof(uintptr_t) * CHAR_BIT,
                 "Cannot embed more bits than we have in a pointer!");
 
@@ -42,26 +45,34 @@ class PointerEmbeddedInt {
     Mask = static_cast<uintptr_t>(-1) << Bits
   };
 
+  static constexpr const struct RawValueTag {} RawValue = RawValueTag();
+
   friend class PointerLikeTypeTraits<PointerEmbeddedInt>;
 
-  explicit PointerEmbeddedInt(uintptr_t Value) : Value(Value) {}
+  explicit PointerEmbeddedInt(uintptr_t Value, RawValueTag) : Value(Value) {}
 
 public:
   PointerEmbeddedInt() : Value(0) {}
 
-  PointerEmbeddedInt(IntT I) : Value(static_cast<uintptr_t>(I) << Shift) {
-    assert((I & Mask) == 0 && "Integer has bits outside those preserved!");
+  PointerEmbeddedInt(IntT I) {
+    *this = I;
   }
 
   PointerEmbeddedInt &operator=(IntT I) {
-    assert((I & Mask) == 0 && "Integer has bits outside those preserved!");
+    assert((std::is_signed<IntT>::value ? llvm::isInt<Bits>(I)
+                                        : llvm::isUInt<Bits>(I)) &&
+           "Integer has bits outside those preserved!");
     Value = static_cast<uintptr_t>(I) << Shift;
     return *this;
   }
 
   // Note that this implicit conversion additionally allows all of the basic
   // comparison operators to work transparently, etc.
-  operator IntT() const { return static_cast<IntT>(Value >> Shift); }
+  operator IntT() const {
+    if (std::is_signed<IntT>::value)
+      return static_cast<IntT>(static_cast<intptr_t>(Value) >> Shift);
+    return static_cast<IntT>(Value >> Shift);
+  }
 };
 
 // Provide pointer like traits to support use with pointer unions and sum
@@ -75,10 +86,10 @@ public:
     return reinterpret_cast<void *>(P.Value);
   }
   static inline T getFromVoidPointer(void *P) {
-    return T(reinterpret_cast<uintptr_t>(P));
+    return T(reinterpret_cast<uintptr_t>(P), T::RawValue);
   }
   static inline T getFromVoidPointer(const void *P) {
-    return T(reinterpret_cast<uintptr_t>(P));
+    return T(reinterpret_cast<uintptr_t>(P), T::RawValue);
   }
 
   enum { NumLowBitsAvailable = T::Shift };
diff --git a/llvm/unittests/ADT/PointerEmbeddedIntTest.cpp b/llvm/unittests/ADT/PointerEmbeddedIntTest.cpp
index b10365a2f61..9c27f8ee655 100644
--- a/llvm/unittests/ADT/PointerEmbeddedIntTest.cpp
+++ b/llvm/unittests/ADT/PointerEmbeddedIntTest.cpp
@@ -43,4 +43,38 @@ TEST(PointerEmbeddedIntTest, Basic) {
   EXPECT_FALSE(42 >= J);
 }
 
+TEST(PointerEmbeddedIntTest, intptr_t) {
+  PointerEmbeddedInt<intptr_t, CHAR_BIT> IPos = 42, INeg = -42;
+  EXPECT_EQ(42, IPos);
+  EXPECT_EQ(-42, INeg);
+
+  PointerEmbeddedInt<uintptr_t, CHAR_BIT> U = 42, USaturated = 255;
+  EXPECT_EQ(42U, U);
+  EXPECT_EQ(255U, USaturated);
+
+  PointerEmbeddedInt<intptr_t, std::numeric_limits<intptr_t>::digits>
+      IMax = std::numeric_limits<intptr_t>::max() >> 1,
+      IMin = std::numeric_limits<intptr_t>::min() >> 1;
+  EXPECT_EQ(std::numeric_limits<intptr_t>::max() >> 1, IMax);
+  EXPECT_EQ(std::numeric_limits<intptr_t>::min() >> 1, IMin);
+
+  PointerEmbeddedInt<uintptr_t, std::numeric_limits<uintptr_t>::digits - 1>
+      UMax = std::numeric_limits<uintptr_t>::max() >> 1,
+      UMin = std::numeric_limits<uintptr_t>::min() >> 1;
+  EXPECT_EQ(std::numeric_limits<uintptr_t>::max() >> 1, UMax);
+  EXPECT_EQ(std::numeric_limits<uintptr_t>::min() >> 1, UMin);
+}
+
+TEST(PointerEmbeddedIntTest, PointerLikeTypeTraits) {
+  PointerEmbeddedInt<int, CHAR_BIT> I = 42;
+  using ITraits = PointerLikeTypeTraits<decltype(I)>;
+  EXPECT_EQ(42, ITraits::getFromVoidPointer(ITraits::getAsVoidPointer(I)));
+
+  PointerEmbeddedInt<uintptr_t, std::numeric_limits<uintptr_t>::digits - 1>
+      Max = std::numeric_limits<uintptr_t>::max() >> 1;
+  using MaxTraits = PointerLikeTypeTraits<decltype(Max)>;
+  EXPECT_EQ(std::numeric_limits<uintptr_t>::max() >> 1,
+            MaxTraits::getFromVoidPointer(MaxTraits::getAsVoidPointer(Max)));
+}
+
 } // end anonymous namespace