From b6d44294798f54bde04be4f2c2d02137cac9df26 Mon Sep 17 00:00:00 2001
From: Aaron Ballman <aaron@aaronballman.com>
Date: Thu, 11 Jan 2018 18:47:15 +0000
Subject: [PATCH] Use size_t to represent the size of a StringMapEntry length
 and alignment rather than unsigned.

Patch by Matt Davis.

llvm-svn: 322305
---
 include/llvm/ADT/StringMap.h    | 20 ++++++++---------
 unittests/ADT/StringMapTest.cpp | 40 +++++++++++++++++++++++++++++++++
 2 files changed, 49 insertions(+), 11 deletions(-)

diff --git a/include/llvm/ADT/StringMap.h b/include/llvm/ADT/StringMap.h
index 6c2830b4491..d34d5ed7e60 100644
--- a/include/llvm/ADT/StringMap.h
+++ b/include/llvm/ADT/StringMap.h
@@ -37,12 +37,12 @@ template<typename ValueTy> class StringMapKeyIterator;
 
 /// StringMapEntryBase - Shared base class of StringMapEntry instances.
 class StringMapEntryBase {
-  unsigned StrLen;
+  size_t StrLen;
 
 public:
-  explicit StringMapEntryBase(unsigned Len) : StrLen(Len) {}
+  explicit StringMapEntryBase(size_t Len) : StrLen(Len) {}
 
-  unsigned getKeyLength() const { return StrLen; }
+  size_t getKeyLength() const { return StrLen; }
 };
 
 /// StringMapImpl - This is the base class of StringMap that is shared among
@@ -127,10 +127,10 @@ class StringMapEntry : public StringMapEntryBase {
 public:
   ValueTy second;
 
-  explicit StringMapEntry(unsigned strLen)
+  explicit StringMapEntry(size_t strLen)
     : StringMapEntryBase(strLen), second() {}
   template <typename... InitTy>
-  StringMapEntry(unsigned strLen, InitTy &&... InitVals)
+  StringMapEntry(size_t strLen, InitTy &&... InitVals)
       : StringMapEntryBase(strLen), second(std::forward<InitTy>(InitVals)...) {}
   StringMapEntry(StringMapEntry &E) = delete;
 
@@ -155,13 +155,12 @@ public:
   template <typename AllocatorTy, typename... InitTy>
   static StringMapEntry *Create(StringRef Key, AllocatorTy &Allocator,
                                 InitTy &&... InitVals) {
-    unsigned KeyLength = Key.size();
+    size_t KeyLength = Key.size();
 
     // Allocate a new item with space for the string at the end and a null
     // terminator.
-    unsigned AllocSize = static_cast<unsigned>(sizeof(StringMapEntry))+
-      KeyLength+1;
-    unsigned Alignment = alignof(StringMapEntry);
+    size_t AllocSize = sizeof(StringMapEntry) + KeyLength + 1;
+    size_t Alignment = alignof(StringMapEntry);
 
     StringMapEntry *NewItem =
       static_cast<StringMapEntry*>(Allocator.Allocate(AllocSize,Alignment));
@@ -203,8 +202,7 @@ public:
   template<typename AllocatorTy>
   void Destroy(AllocatorTy &Allocator) {
     // Free memory referenced by the item.
-    unsigned AllocSize =
-        static_cast<unsigned>(sizeof(StringMapEntry)) + getKeyLength() + 1;
+    size_t AllocSize = sizeof(StringMapEntry) + getKeyLength() + 1;
     this->~StringMapEntry();
     Allocator.Deallocate(static_cast<void *>(this), AllocSize);
   }
diff --git a/unittests/ADT/StringMapTest.cpp b/unittests/ADT/StringMapTest.cpp
index b5c63695ff3..6e0ea0e48ff 100644
--- a/unittests/ADT/StringMapTest.cpp
+++ b/unittests/ADT/StringMapTest.cpp
@@ -12,6 +12,7 @@
 #include "llvm/ADT/Twine.h"
 #include "llvm/Support/DataTypes.h"
 #include "gtest/gtest.h"
+#include <limits>
 #include <tuple>
 using namespace llvm;
 
@@ -492,4 +493,43 @@ TEST(StringMapCustomTest, EmplaceTest) {
   EXPECT_EQ(42, Map["abcd"].Data);
 }
 
+// Test that StringMapEntryBase can handle size_t wide sizes.
+TEST(StringMapCustomTest, StringMapEntryBaseSize) {
+  size_t LargeValue;
+
+  // Test that the entry can represent max-unsigned.
+  if (sizeof(size_t) <= sizeof(unsigned))
+    LargeValue = std::numeric_limits<unsigned>::max();
+  else
+    LargeValue = std::numeric_limits<unsigned>::max() + 1ULL;
+  StringMapEntryBase LargeBase(LargeValue);
+  EXPECT_EQ(LargeValue, LargeBase.getKeyLength());
+
+  // Test that the entry can hold at least max size_t.
+  LargeValue = std::numeric_limits<size_t>::max();
+  StringMapEntryBase LargerBase(LargeValue);
+  LargeValue = std::numeric_limits<size_t>::max();
+  EXPECT_EQ(LargeValue, LargerBase.getKeyLength());
+}
+
+// Test that StringMapEntry can handle size_t wide sizes.
+TEST(StringMapCustomTest, StringMapEntrySize) {
+  size_t LargeValue;
+
+  // Test that the entry can represent max-unsigned.
+  if (sizeof(size_t) <= sizeof(unsigned))
+    LargeValue = std::numeric_limits<unsigned>::max();
+  else
+    LargeValue = std::numeric_limits<unsigned>::max() + 1ULL;
+  StringMapEntry<int> LargeEntry(LargeValue);
+  StringRef Key = LargeEntry.getKey();
+  EXPECT_EQ(LargeValue, Key.size());
+
+  // Test that the entry can hold at least max size_t.
+  LargeValue = std::numeric_limits<size_t>::max();
+  StringMapEntry<int> LargerEntry(LargeValue);
+  Key = LargerEntry.getKey();
+  EXPECT_EQ(LargeValue, Key.size());
+}
+
 } // end anonymous namespace