From c961aee31d1b841cfe8914536e891d6322751267 Mon Sep 17 00:00:00 2001
From: Gogul Balakrishnan <bgogul@google.com>
Date: Mon, 21 Oct 2019 23:46:41 -0700
Subject: [PATCH] Fix memory error in the implementation of IndexSubset.

---
 include/swift/AST/IndexSubset.h | 4 ++++
 lib/AST/ASTContext.cpp          | 2 +-
 2 files changed, 5 insertions(+), 1 deletion(-)

diff --git a/include/swift/AST/IndexSubset.h b/include/swift/AST/IndexSubset.h
index 3bdb021ea4e0f..e2b8a98370ee6 100644
--- a/include/swift/AST/IndexSubset.h
+++ b/include/swift/AST/IndexSubset.h
@@ -57,6 +57,10 @@ class IndexSubset : public llvm::FoldingSetNode {
   /// The number of bit words in the index subset.
   unsigned numBitWords;
 
+  static unsigned getNumBytesNeededForCapacity(unsigned capacity) {
+    return getNumBitWordsNeededForCapacity(capacity) * bitWordSize;
+  }
+
   BitWord *getBitWordsData() {
     return reinterpret_cast<BitWord *>(this + 1);
   }
diff --git a/lib/AST/ASTContext.cpp b/lib/AST/ASTContext.cpp
index 0bc13d4aa62b9..9288d42065638 100644
--- a/lib/AST/ASTContext.cpp
+++ b/lib/AST/ASTContext.cpp
@@ -4822,7 +4822,7 @@ IndexSubset::get(ASTContext &ctx, const SmallBitVector &indices) {
   if (existing)
     return existing;
   auto sizeToAlloc = sizeof(IndexSubset) +
-      getNumBitWordsNeededForCapacity(capacity);
+      getNumBytesNeededForCapacity(capacity);
   auto *buf = reinterpret_cast<IndexSubset *>(
       ctx.Allocate(sizeToAlloc, alignof(IndexSubset)));
   auto *newNode = new (buf) IndexSubset(indices);