[AutoDiff] Refactor 'AutoDiffIndexSubset' to take a 'SmallBitVector' in the base factory method. (#25058)

Change the base implementation of `AutoDiffIndexSubset::get` to take a `const SmallBitVector &` instead of an `ArrayRef<unsigned>` because bit vectors are far more efficient.

Author: rxwei

include/swift/AST/AutoDiff.h

@@ -276,11 +276,11 @@ class AutoDiffIndexSubset : public llvm::FoldingSetNode {
     return {const_cast<BitWord *>(getBitWordsData()), getNumBitWords()};
   }
 
-  explicit AutoDiffIndexSubset(unsigned capacity, ArrayRef<unsigned> indices)
-      : capacity(capacity),
+  explicit AutoDiffIndexSubset(const SmallBitVector &indices)
+      : capacity((unsigned)indices.size()),
         numBitWords(getNumBitWordsNeededForCapacity(capacity)) {
     std::uninitialized_fill_n(getBitWordsData(), numBitWords, 0);
-    for (auto i : indices) {
+    for (auto i : indices.set_bits()) {
       unsigned bitWordIndex, offset;
       std::tie(bitWordIndex, offset) = getBitWordIndexAndOffset(i);
       getBitWord(bitWordIndex) |= (1 << offset);
@@ -294,22 +294,28 @@ class AutoDiffIndexSubset : public llvm::FoldingSetNode {
 
   // Defined in ASTContext.h.
   static AutoDiffIndexSubset *get(ASTContext &ctx,
-                                  unsigned capacity,
-                                  ArrayRef<unsigned> indices);
+                                  const SmallBitVector &indices);
 
-  static AutoDiffIndexSubset *getDefault(ASTContext &ctx,
-                                         unsigned capacity,
+  static AutoDiffIndexSubset *get(ASTContext &ctx, unsigned capacity,
+                                  ArrayRef<unsigned> indices) {
+    SmallBitVector indicesBitVec(capacity, false);
+    for (auto index : indices)
+      indicesBitVec.set(index);
+    return AutoDiffIndexSubset::get(ctx, indicesBitVec);
+  }
+
+  static AutoDiffIndexSubset *getDefault(ASTContext &ctx, unsigned capacity,
                                          bool includeAll = false) {
-    if (includeAll)
-      return getFromRange(ctx, capacity, IntRange<>(capacity));
-    return get(ctx, capacity, {});
+    return get(ctx, SmallBitVector(capacity, includeAll));
   }
 
-  static AutoDiffIndexSubset *getFromRange(ASTContext &ctx,
-                                           unsigned capacity,
-                                           IntRange<> range) {
-    return get(ctx, capacity,
-               SmallVector<unsigned, 8>(range.begin(), range.end()));
+  static AutoDiffIndexSubset *getFromRange(ASTContext &ctx, unsigned capacity,
+                                           unsigned start, unsigned end) {
+    assert(start < capacity);
+    assert(end <= capacity);
+    SmallBitVector bitVec(capacity);
+    bitVec.set(start, end);
+    return get(ctx, bitVec);
   }
 
   unsigned getNumBitWords() const {

lib/AST/ASTContext.cpp

@@ -4521,21 +4521,13 @@ AutoDiffParameterIndices::get(llvm::SmallBitVector indices, ASTContext &C) {
 }
 
 AutoDiffIndexSubset *
-AutoDiffIndexSubset::get(ASTContext &ctx, unsigned capacity,
-                         ArrayRef<unsigned> indices) {
+AutoDiffIndexSubset::get(ASTContext &ctx, const SmallBitVector &indices) {
   auto &foldingSet = ctx.getImpl().AutoDiffIndexSubsets;
   llvm::FoldingSetNodeID id;
+  unsigned capacity = indices.size();
   id.AddInteger(capacity);
-#ifndef NDEBUG
-  int last = -1;
-#endif
-  for (unsigned index : indices) {
-#ifndef NDEBUG
-    assert((int)index > last && "Indices must be ascending");
-    last = (int)index;
-#endif
+  for (unsigned index : indices.set_bits())
     id.AddInteger(index);
-  }
   void *insertPos = nullptr;
   auto *existing = foldingSet.FindNodeOrInsertPos(id, insertPos);
   if (existing)
@@ -4544,7 +4536,7 @@ AutoDiffIndexSubset::get(ASTContext &ctx, unsigned capacity,
       getNumBitWordsNeededForCapacity(capacity);
   auto *buf = reinterpret_cast<AutoDiffIndexSubset *>(
       ctx.Allocate(sizeToAlloc, alignof(AutoDiffIndexSubset)));
-  auto *newNode = new (buf) AutoDiffIndexSubset(capacity, indices);
+  auto *newNode = new (buf) AutoDiffIndexSubset(indices);
   foldingSet.InsertNode(newNode, insertPos);
   return newNode;
 }

lib/AST/AutoDiff.cpp

@@ -385,28 +385,27 @@ AutoDiffIndexSubset *AutoDiffIndexSubset::adding(unsigned index,
   assert(index < getCapacity());
   if (contains(index))
     return const_cast<AutoDiffIndexSubset *>(this);
-  SmallVector<unsigned, 8> newIndices;
-  newIndices.reserve(capacity + 1);
+  SmallBitVector newIndices(capacity);
   bool inserted = false;
   for (auto curIndex : getIndices()) {
     if (!inserted && curIndex > index) {
-      newIndices.push_back(index);
+      newIndices.set(index);
       inserted = true;
     }
-    newIndices.push_back(curIndex);
+    newIndices.set(curIndex);
   }
-  return get(ctx, capacity, newIndices);
+  return get(ctx, newIndices);
 }
 
 AutoDiffIndexSubset *AutoDiffIndexSubset::extendingCapacity(
     ASTContext &ctx, unsigned newCapacity) const {
   assert(newCapacity >= capacity);
   if (newCapacity == capacity)
     return const_cast<AutoDiffIndexSubset *>(this);
-  SmallVector<unsigned, 8> indices;
+  SmallBitVector indices(newCapacity);
   for (auto index : getIndices())
-    indices.push_back(index);
-  return AutoDiffIndexSubset::get(ctx, newCapacity, indices);
+    indices.set(index);
+  return AutoDiffIndexSubset::get(ctx, indices);
 }
 
 int AutoDiffIndexSubset::findNext(int startIndex) const {