[AutoDiff upstream] [AST] Add 'AutoDiffIndexSubset' data structure.

`AutoDiffIndexSubset` is a fixed-size bit vector that is used for efficiently representing a subset of indices in automatic differentiation, specifically for representing a subset of parameters and results of a function to differentiate with respect to. It is uniqued in `ASTContext`.

This patch adds definition and unit tests for `AutoDiffIndexSubset` along with new files `AutoDiff.h` and `AutoDiff.cpp` into the 'AST' target, with no changes to the compiler's behavior. More data structures used for AutoDiff will be added to these files.

----------------------------

This is part of the ongoing effort to merge the experimental [differentiable programming feature](https://forums.swift.org/t/differentiable-programming-mega-proposal/28547) (informally referred to as "AutoDiff") to the 'master' branch for code reviews and better maintenance.

Upstreaming task: [TF-879](https://bugs.swift.org/browse/TF-879)

Author: rxwei

include/swift/AST/ASTContext.h

@@ -109,6 +109,7 @@ namespace swift {
   class TypeAliasDecl;
   class VarDecl;
   class UnifiedStatsReporter;
+  class AutoDiffIndexSubset;
 
   enum class KnownProtocolKind : uint8_t;
 

include/swift/AST/AutoDiff.h

@@ -0,0 +1,273 @@
+//===--- AutoDiff.h - Swift Differentiable Programming --------------------===//
+//
+// This source file is part of the Swift.org open source project
+//
+// Copyright (c) 2014 - 2017 Apple Inc. and the Swift project authors
+// Licensed under Apache License v2.0 with Runtime Library Exception
+//
+// See https://swift.org/LICENSE.txt for license information
+// See https://swift.org/CONTRIBUTORS.txt for the list of Swift project authors
+//
+//===----------------------------------------------------------------------===//
+//
+//  This file defines AST support for the experimental differentiable
+//  programming feature.
+//
+//===----------------------------------------------------------------------===//
+
+
+#ifndef SWIFT_AST_AUTODIFF_H
+#define SWIFT_AST_AUTODIFF_H
+
+#include "ASTContext.h"
+#include "llvm/ADT/SmallBitVector.h"
+#include "swift/Basic/Range.h"
+
+namespace swift {
+
+/// 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);
+  }
+  
+  /// Returns an iterator range of indices in the index subset.
+  iterator_range<iterator> getIndices() const {
+    return make_range(begin(), end());
+  }
+
+  /// Returns the number of indices in the index subset.
+  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/ASTContext.cpp

@@ -17,6 +17,7 @@
 #include "swift/AST/ASTContext.h"
 #include "ForeignRepresentationInfo.h"
 #include "SubstitutionMapStorage.h"
+#include "swift/AST/AutoDiff.h"
 #include "swift/AST/ClangModuleLoader.h"
 #include "swift/AST/ConcreteDeclRef.h"
 #include "swift/AST/DiagnosticEngine.h"
@@ -427,6 +428,9 @@ FOR_KNOWN_FOUNDATION_TYPES(CACHE_FOUNDATION_DECL)
   llvm::FoldingSet<DeclName::CompoundDeclName> CompoundNames;
   llvm::DenseMap<UUID, OpenedArchetypeType *> OpenedExistentialArchetypes;
 
+  /// For uniquifying `AutoDiffIndexSubset` allocations.
+  llvm::FoldingSet<AutoDiffIndexSubset> AutoDiffIndexSubsets;
+
   /// A cache of information about whether particular nominal types
   /// are representable in a foreign language.
   llvm::DenseMap<NominalTypeDecl *, ForeignRepresentationInfo>
@@ -4616,3 +4620,24 @@ void VarDecl::setOriginalWrappedProperty(VarDecl *originalProperty) {
   assert(ctx.getImpl().OriginalWrappedProperties.count(this) == 0);
   ctx.getImpl().OriginalWrappedProperties[this] = originalProperty;
 }
+
+AutoDiffIndexSubset *
+AutoDiffIndexSubset::get(ASTContext &ctx, const SmallBitVector &indices) {
+  auto &foldingSet = ctx.getImpl().AutoDiffIndexSubsets;
+  llvm::FoldingSetNodeID id;
+  unsigned capacity = indices.size();
+  id.AddInteger(capacity);
+  for (unsigned index : indices.set_bits())
+    id.AddInteger(index);
+  void *insertPos = nullptr;
+  auto *existing = foldingSet.FindNodeOrInsertPos(id, insertPos);
+  if (existing)
+    return existing;
+  auto sizeToAlloc = sizeof(AutoDiffIndexSubset) +
+    getNumBitWordsNeededForCapacity(capacity);
+  auto *buf = reinterpret_cast<AutoDiffIndexSubset *>(
+    ctx.Allocate(sizeToAlloc, alignof(AutoDiffIndexSubset)));
+  auto *newNode = new (buf) AutoDiffIndexSubset(indices);
+  foldingSet.InsertNode(newNode, insertPos);
+  return newNode;
+}