Pick up changes from IndexSubset.

Author: bgogul

include/swift/AST/Attr.h

@@ -1758,7 +1758,7 @@ class DifferentiableAttr final
   /// specified.
   FuncDecl *VJPFunction = nullptr;
   /// The differentiation parameters' indices, resolved by the type checker.
-  AutoDiffIndexSubset *ParameterIndices = nullptr;
+  IndexSubset *ParameterIndices = nullptr;
   /// The trailing where clause (optional).
   TrailingWhereClause *WhereClause = nullptr;
   /// The generic signature for autodiff associated functions. Resolved by the
@@ -1777,7 +1777,7 @@ class DifferentiableAttr final
 
   explicit DifferentiableAttr(ASTContext &context, bool implicit,
                               SourceLoc atLoc, SourceRange baseRange,
-                              bool linear, AutoDiffIndexSubset *indices,
+                              bool linear, IndexSubset *indices,
                               Optional<DeclNameWithLoc> jvp,
                               Optional<DeclNameWithLoc> vjp,
                               GenericSignature derivativeGenericSignature);
@@ -1793,7 +1793,7 @@ class DifferentiableAttr final
 
   static DifferentiableAttr *create(ASTContext &context, bool implicit,
                                     SourceLoc atLoc, SourceRange baseRange,
-                                    bool linear, AutoDiffIndexSubset *indices,
+                                    bool linear, IndexSubset *indices,
                                     Optional<DeclNameWithLoc> jvp,
                                     Optional<DeclNameWithLoc> vjp,
                                     GenericSignature derivativeGenSig);
@@ -1808,10 +1808,10 @@ class DifferentiableAttr final
   /// registered VJP.
   Optional<DeclNameWithLoc> getVJP() const { return VJP; }
 
-  AutoDiffIndexSubset *getParameterIndices() const {
+  IndexSubset *getParameterIndices() const {
     return ParameterIndices;
   }
-  void setParameterIndices(AutoDiffIndexSubset *pi) {
+  void setParameterIndices(IndexSubset *pi) {
     ParameterIndices = pi;
   }
 

include/swift/AST/AutoDiff.h

@@ -87,247 +87,6 @@ class ParsedAutoDiffParameter {
   }
 };
 
-/// An efficient index subset data structure, uniqued in `ASTContext`.
-/// Stores a bit vector representing set indices and a total capacity.
-class AutoDiffIndexSubset : public llvm::FoldingSetNode {
-public:
-  typedef uint64_t BitWord;
-
-  static constexpr unsigned bitWordSize = sizeof(BitWord);
-  static constexpr unsigned numBitsPerBitWord = bitWordSize * 8;
-
-  static std::pair<unsigned, unsigned>
-  getBitWordIndexAndOffset(unsigned index) {
-    auto bitWordIndex = index / numBitsPerBitWord;
-    auto bitWordOffset = index % numBitsPerBitWord;
-    return {bitWordIndex, bitWordOffset};
-  }
-
-  static unsigned getNumBitWordsNeededForCapacity(unsigned capacity) {
-    if (capacity == 0) return 0;
-    return capacity / numBitsPerBitWord + 1;
-  }
-
-private:
-  /// The total capacity of the index subset, which is `1` less than the largest
-  /// index.
-  unsigned capacity;
-  /// The number of bit words in the index subset.
-  unsigned numBitWords;
-
-  BitWord *getBitWordsData() {
-    return reinterpret_cast<BitWord *>(this + 1);
-  }
-
-  const BitWord *getBitWordsData() const {
-    return reinterpret_cast<const BitWord *>(this + 1);
-  }
-
-  ArrayRef<BitWord> getBitWords() const {
-    return {getBitWordsData(), getNumBitWords()};
-  }
-
-  BitWord getBitWord(unsigned i) const {
-    return getBitWordsData()[i];
-  }
-
-  BitWord &getBitWord(unsigned i) {
-    return getBitWordsData()[i];
-  }
-
-  MutableArrayRef<BitWord> getMutableBitWords() {
-    return {const_cast<BitWord *>(getBitWordsData()), getNumBitWords()};
-  }
-
-  explicit AutoDiffIndexSubset(const SmallBitVector &indices)
-      : capacity((unsigned)indices.size()),
-        numBitWords(getNumBitWordsNeededForCapacity(capacity)) {
-    std::uninitialized_fill_n(getBitWordsData(), numBitWords, 0);
-    for (auto i : indices.set_bits()) {
-      unsigned bitWordIndex, offset;
-      std::tie(bitWordIndex, offset) = getBitWordIndexAndOffset(i);
-      getBitWord(bitWordIndex) |= (1 << offset);
-    }
-  }
-
-public:
-  AutoDiffIndexSubset() = delete;
-  AutoDiffIndexSubset(const AutoDiffIndexSubset &) = delete;
-  AutoDiffIndexSubset &operator=(const AutoDiffIndexSubset &) = delete;
-
-  // Defined in ASTContext.cpp.
-  static AutoDiffIndexSubset *get(ASTContext &ctx,
-                                  const SmallBitVector &indices);
-
-  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) {
-    return get(ctx, SmallBitVector(capacity, includeAll));
-  }
-
-  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);
-  }
-
-  /// Creates an index subset corresponding to the given string generated by
-  /// `getString()`. If the string is invalid, returns nullptr.
-  static AutoDiffIndexSubset *getFromString(ASTContext &ctx, StringRef string);
-
-  /// Returns the number of bit words used to store the index subset.
-  // Note: Use `getCapacity()` to get the total index subset capacity.
-  // This is public only for unit testing
-  // (in unittests/AST/SILAutoDiffIndices.cpp).
-  unsigned getNumBitWords() const {
-    return numBitWords;
-  }
-
-  /// Returns the capacity of the index subset.
-  unsigned getCapacity() const {
-    return capacity;
-  }
-
-  /// Returns a textual string description of these indices.
-  ///
-  /// It has the format `[SU]+`, where the total number of characters is equal
-  /// to the capacity, and where "S" means that the corresponding index is
-  /// contained and "U" means that the corresponding index is not.
-  std::string getString() const;
-
-  class iterator;
-
-  iterator begin() const {
-    return iterator(this);
-  }
-
-  iterator end() const {
-    return iterator(this, (int)capacity);
-  }
-
-  iterator_range<iterator> getIndices() const {
-    return make_range(begin(), end());
-  }
-
-  unsigned getNumIndices() const {
-    return (unsigned)std::distance(begin(), end());
-  }
-
-  SmallBitVector getBitVector() const {
-    SmallBitVector indicesBitVec(capacity, false);
-    for (auto index : getIndices())
-      indicesBitVec.set(index);
-    return indicesBitVec;
-  }
-
-  bool contains(unsigned index) const {
-    unsigned bitWordIndex, offset;
-    std::tie(bitWordIndex, offset) = getBitWordIndexAndOffset(index);
-    return getBitWord(bitWordIndex) & (1 << offset);
-  }
-
-  bool isEmpty() const {
-    return llvm::all_of(getBitWords(), [](BitWord bw) { return !(bool)bw; });
-  }
-
-  bool equals(AutoDiffIndexSubset *other) const {
-    return capacity == other->getCapacity() &&
-        getBitWords().equals(other->getBitWords());
-  }
-
-  bool isSubsetOf(AutoDiffIndexSubset *other) const;
-  bool isSupersetOf(AutoDiffIndexSubset *other) const;
-
-  AutoDiffIndexSubset *adding(unsigned index, ASTContext &ctx) const;
-  AutoDiffIndexSubset *extendingCapacity(ASTContext &ctx,
-                                         unsigned newCapacity) const;
-
-  void Profile(llvm::FoldingSetNodeID &id) const {
-    id.AddInteger(capacity);
-    for (auto index : getIndices())
-      id.AddInteger(index);
-  }
-
-  void print(llvm::raw_ostream &s = llvm::outs()) const {
-    s << '{';
-    interleave(range(capacity), [this, &s](unsigned i) { s << contains(i); },
-               [&s] { s << ", "; });
-    s << '}';
-  }
-
-  void dump(llvm::raw_ostream &s = llvm::errs()) const {
-    s << "(autodiff_index_subset capacity=" << capacity << " indices=(";
-    interleave(getIndices(), [&s](unsigned i) { s << i; },
-               [&s] { s << ", "; });
-    s << "))";
-  }
-
-  int findNext(int startIndex) const;
-  int findFirst() const { return findNext(-1); }
-  int findPrevious(int endIndex) const;
-  int findLast() const { return findPrevious(capacity); }
-
-  class iterator {
-  public:
-    typedef unsigned value_type;
-    typedef unsigned difference_type;
-    typedef unsigned * pointer;
-    typedef unsigned & reference;
-    typedef std::forward_iterator_tag iterator_category;
-
-  private:
-    const AutoDiffIndexSubset *parent;
-    int current = 0;
-
-    void advance() {
-      assert(current != -1 && "Trying to advance past end.");
-      current = parent->findNext(current);
-    }
-
-  public:
-    iterator(const AutoDiffIndexSubset *parent, int current)
-        : parent(parent), current(current) {}
-    explicit iterator(const AutoDiffIndexSubset *parent)
-        : iterator(parent, parent->findFirst()) {}
-    iterator(const iterator &) = default;
-
-    iterator operator++(int) {
-      auto prev = *this;
-      advance();
-      return prev;
-    }
-
-    iterator &operator++() {
-      advance();
-      return *this;
-    }
-
-    unsigned operator*() const { return current; }
-
-    bool operator==(const iterator &other) const {
-      assert(parent == other.parent &&
-             "Comparing iterators from different AutoDiffIndexSubsets");
-      return current == other.current;
-    }
-
-    bool operator!=(const iterator &other) const {
-      assert(parent == other.parent &&
-             "Comparing iterators from different AutoDiffIndexSubsets");
-      return current != other.current;
-    }
-  };
-};
-
 } // end namespace swift
 
 #endif // SWIFT_AST_AUTODIFF_H

lib/AST/Attr.cpp

@@ -20,6 +20,7 @@
 #include "swift/AST/Decl.h"
 #include "swift/AST/Expr.h"
 #include "swift/AST/GenericEnvironment.h"
+#include "swift/AST/IndexSubset.h"
 #include "swift/AST/Module.h"
 #include "swift/AST/TypeRepr.h"
 #include "swift/AST/Types.h"
@@ -353,13 +354,13 @@ static void printShortFormAvailable(ArrayRef<const DeclAttribute *> Attrs,
 }
 
 static std::string getDifferentiationParametersClauseString(
-    const AbstractFunctionDecl *function, AutoDiffIndexSubset *indices,
+    const AbstractFunctionDecl *function, IndexSubset *indices,
     ArrayRef<ParsedAutoDiffParameter> parsedParams) {
   bool isInstanceMethod = function && function->isInstanceMember();
   std::string result;
   llvm::raw_string_ostream printer(result);
 
-  // Use parameter indices from `AutoDiffIndexSubset`, if specified.
+  // Use parameter indices from `IndexSubset`, if specified.
   if (indices) {
     auto parameters = indices->getBitVector();
     auto parameterCount = parameters.count();
@@ -1340,7 +1341,7 @@ DifferentiableAttr::DifferentiableAttr(ASTContext &context, bool implicit,
 DifferentiableAttr::DifferentiableAttr(ASTContext &context, bool implicit,
                                        SourceLoc atLoc, SourceRange baseRange,
                                        bool linear,
-                                       AutoDiffIndexSubset *indices,
+                                       IndexSubset *indices,
                                        Optional<DeclNameWithLoc> jvp,
                                        Optional<DeclNameWithLoc> vjp,
                                        GenericSignature derivativeGenSig)
@@ -1368,7 +1369,7 @@ DifferentiableAttr::create(ASTContext &context, bool implicit,
 DifferentiableAttr *
 DifferentiableAttr::create(ASTContext &context, bool implicit,
                            SourceLoc atLoc, SourceRange baseRange,
-                           bool linear, AutoDiffIndexSubset *indices,
+                           bool linear, IndexSubset *indices,
                            Optional<DeclNameWithLoc> jvp,
                            Optional<DeclNameWithLoc> vjp,
                            GenericSignature derivativeGenSig) {