delete isMethodFlag from the indices

Author: Marc Rasi

include/swift/AST/AutoDiff.h

@@ -96,45 +96,36 @@ class AutoDiffParameterIndices : public llvm::FoldingSetNode {
   /// Bits corresponding to parameters in the set are "on", and bits
   /// corresponding to parameters not in the set are "off".
   ///
-  /// Normally, the bits correspond to the function's parameters in order. For
-  /// example,
+  /// For non-method functions, the bits correspond to the function's
+  //// parameters in order. For example,
   ///
   ///   Function type: (A, B, C) -> R
   ///   Bits: [A][B][C]
   ///
-  /// When `isMethodFlag` is set, the bits correspond to the function's
-  /// non-self parameters in order, followed by the function's self parameter.
-  /// For example,
+  /// For methods, the bits correspond to the function's non-self parameters
+  /// in order, followed by the function's self parameter. For example,
   ///
   ///   Function type: (Self) -> (A, B, C) -> R
   ///   Bits: [A][B][C][Self]
   ///
   const llvm::SmallBitVector indices;
 
-  /// Whether the function is a method.
-  ///
-  const bool isMethodFlag;
-
-  AutoDiffParameterIndices(llvm::SmallBitVector indices, bool isMethodFlag)
-      : indices(indices), isMethodFlag(isMethodFlag) {}
+  AutoDiffParameterIndices(llvm::SmallBitVector indices)
+      : indices(indices) {}
 
   static AutoDiffParameterIndices *get(llvm::SmallBitVector indices,
-                                       bool isMethodFlag, ASTContext &C);
+                                       ASTContext &C);
 
 public:
-  bool isMethod() const { return isMethodFlag; }
-
   /// Allocates and initializes an `AutoDiffParameterIndices` corresponding to
   /// the given `string` generated by `getString()`. If the string is invalid,
   /// returns nullptr.
   static AutoDiffParameterIndices *create(ASTContext &C, StringRef string);
 
   /// Returns a textual string description of these indices,
   ///
-  ///   [FM][SU]+
+  ///   [SU]+
   ///
-  /// "F" means that `isMethodFlag` is false
-  /// "M" means that `isMethodFlag` is true
   /// "S" means that the corresponding index is set
   /// "U" means that the corresponding index is unset
   std::string getString() const;
@@ -153,11 +144,14 @@ class AutoDiffParameterIndices : public llvm::FoldingSetNode {
   ///   if "Self" and "C" are in the set,
   ///   ==> pushes {Self, C} to `paramTypes`.
   ///
+  /// Pass `isMethod = true` when the function is a method.
+  ///
   /// Pass `selfUncurried = true` when the function type is for a method whose
   /// self parameter has been uncurried as in (A, B, C, Self) -> R.
   ///
   void getSubsetParameterTypes(AnyFunctionType *functionType,
                                SmallVectorImpl<Type> &paramTypes,
+                               bool isMethod,
                                bool selfUncurried = false) const;
 
   /// Returns a bitvector for the SILFunction parameters corresponding to the
@@ -179,14 +173,16 @@ class AutoDiffParameterIndices : public llvm::FoldingSetNode {
   ///   ==> returns 1110
   ///   (because the lowered SIL type is (A, B, C, D) -> R)
   ///
+  /// Pass `isMethod = true` when the function is a method.
+  ///
   /// Pass `selfUncurried = true` when the function type is for a method whose
   /// self parameter has been uncurried as in (A, B, C, Self) -> R.
   ///
   llvm::SmallBitVector getLowered(AnyFunctionType *functionType,
+                                  bool isMethod,
                                   bool selfUncurried = false) const;
 
   void Profile(llvm::FoldingSetNodeID &ID) const {
-    ID.AddBoolean(isMethodFlag);
     ID.AddInteger(indices.size());
     for (unsigned setBit : indices.set_bits())
       ID.AddInteger(setBit);
@@ -196,7 +192,7 @@ class AutoDiffParameterIndices : public llvm::FoldingSetNode {
 /// Builder for `AutoDiffParameterIndices`.
 class AutoDiffParameterIndicesBuilder {
   llvm::SmallBitVector indices;
-  bool isMethodFlag;
+  bool isMethod;
 
   unsigned getNumNonSelfParameters() const;
 
@@ -210,14 +206,14 @@ class AutoDiffParameterIndicesBuilder {
   /// one if it has already been allocated in the `ASTContext`.
   AutoDiffParameterIndices *build(ASTContext &C) const;
 
-  /// Adds the indexed parameter to the set. When `isMethodFlag` is not set,
+  /// Adds the indexed parameter to the set. When `isMethod` is not set,
   /// the indices index into the first parameter list. For example,
   ///
   ///   functionType = (A, B, C) -> R
   ///   paramIndex = 0
   ///   ==> adds "A" to the set.
   ///
-  /// When `isMethodFlag` is set, the indices index into the first non-self
+  /// When `isMethod` is set, the indices index into the first non-self
   /// parameter list. For example,
   ///
   ///   functionType = (Self) -> (A, B, C) -> R
@@ -237,7 +233,7 @@ class AutoDiffParameterIndicesBuilder {
   ///
   void setAllNonSelfParameters();
 
-  /// Adds the self parameter to the set. `isMethodFlag` must be set. For
+  /// Adds the self parameter to the set. `isMethod` must be set. For
   /// example,
   ///   functionType = (Self) -> (A, B, C) -> R
   ///   ==> adds "Self" to the set

include/swift/AST/Types.h

@@ -3127,11 +3127,13 @@ class AnyFunctionType : public TypeBase {
   AnyFunctionType *getAutoDiffAssociatedFunctionType(
       AutoDiffParameterIndices *indices, unsigned resultIndex,
       unsigned differentiationOrder, AutoDiffAssociatedFunctionKind kind,
-      LookupConformanceFn lookupConformance, bool selfUncurried = false);
+      LookupConformanceFn lookupConformance, bool isMethod,
+      bool selfUncurried = false);
 
   AnyFunctionType *
   getAutoDiffAdjointFunctionType(AutoDiffParameterIndices *indices,
-                                 const TupleType *primalResultTy);
+                                 const TupleType *primalResultTy,
+                                 bool isMethod);
 
   /// \brief True if this type allows an implicit conversion from a function
   /// argument expression of type T to a function of type () -> T.

lib/AST/ASTContext.cpp

@@ -5278,12 +5278,10 @@ Optional<TangentSpace> ASTContext::getTangentSpace(CanType type,
 }
 
 AutoDiffParameterIndices *
-AutoDiffParameterIndices::get(llvm::SmallBitVector indices, bool isMethodFlag,
-                              ASTContext &C) {
+AutoDiffParameterIndices::get(llvm::SmallBitVector indices, ASTContext &C) {
   auto &foldingSet = C.getImpl().AutoDiffParameterIndicesSet;
 
   llvm::FoldingSetNodeID id;
-  id.AddBoolean(isMethodFlag);
   id.AddInteger(indices.size());
   for (unsigned setBit : indices.set_bits())
     id.AddInteger(setBit);
@@ -5298,7 +5296,7 @@ AutoDiffParameterIndices::get(llvm::SmallBitVector indices, bool isMethodFlag,
   // SmallBitVector decides to allocate some heap space.
   void *mem = C.Allocate(sizeof(AutoDiffParameterIndices),
                          alignof(AutoDiffParameterIndices));
-  auto *newNode = ::new (mem) AutoDiffParameterIndices(indices, isMethodFlag);
+  auto *newNode = ::new (mem) AutoDiffParameterIndices(indices);
   foldingSet.InsertNode(newNode, insertPos);
 
   return newNode;

lib/AST/AutoDiff.cpp

@@ -110,39 +110,28 @@ static AnyFunctionType *unwrapSelfParameter(AnyFunctionType *functionType,
 /// returns nullptr.
 AutoDiffParameterIndices *
 AutoDiffParameterIndices::create(ASTContext &C, StringRef string) {
-  if (string.size() < 2)
+  if (string.size() < 1)
     return nullptr;
 
-  bool isMethod = false;
   llvm::SmallBitVector indices(string.size() - 1);
-  if (string[0] == 'M')
-    isMethod = true;
-  else if (string[0] != 'F')
-    return nullptr;
   for (unsigned i : range(indices.size())) {
     if (string[i + 1] == 'S')
       indices.set(i);
     else if (string[i + 1] != 'U')
       return nullptr;
   }
 
-  return get(indices, isMethod, C);
+  return get(indices, C);
 }
 
 /// Returns a textual string description of these indices,
 ///
-///   [FM][SU]+
+///   [SU]+
 ///
-/// "F" means that `isMethodFlag` is false
-/// "M" means that `isMethodFlag` is true
 /// "S" means that the corresponding index is set
 /// "U" means that the corresponding index is unset
 std::string AutoDiffParameterIndices::getString() const {
   std::string result;
-  if (isMethodFlag)
-    result += "M";
-  else
-    result += "F";
   for (unsigned i : range(indices.size())) {
     if (indices[i])
       result += "S";
@@ -163,21 +152,23 @@ std::string AutoDiffParameterIndices::getString() const {
 ///   if "Self" and "C" are in the set,
 ///   ==> pushes {Self, C} to `paramTypes`.
 ///
+/// Pass `isMethod = true` when the function is a method.
+///
 /// Pass `selfUncurried = true` when the function type is for a method whose
 /// self parameter has been uncurried as in (A, B, C, Self) -> R.
 ///
 void AutoDiffParameterIndices::getSubsetParameterTypes(
     AnyFunctionType *functionType, SmallVectorImpl<Type> &paramTypes,
-    bool selfUncurried) const {
-  if (selfUncurried && isMethodFlag) {
-    if (isMethodFlag && indices[indices.size() - 1])
+    bool isMethod, bool selfUncurried) const {
+  if (selfUncurried && isMethod) {
+    if (isMethod && indices[indices.size() - 1])
       paramTypes.push_back(functionType->getParams()[functionType->getNumParams() - 1].getPlainType());
     for (unsigned paramIndex : range(functionType->getNumParams() - 1))
       if (indices[paramIndex])
         paramTypes.push_back(functionType->getParams()[paramIndex].getPlainType());
   } else {
-    AnyFunctionType *unwrapped = unwrapSelfParameter(functionType, isMethodFlag);
-    if (isMethodFlag && indices[indices.size() - 1])
+    AnyFunctionType *unwrapped = unwrapSelfParameter(functionType, isMethod);
+    if (isMethod && indices[indices.size() - 1])
       paramTypes.push_back(functionType->getParams()[0].getPlainType());
     for (unsigned paramIndex : range(unwrapped->getNumParams()))
       if (indices[paramIndex])
@@ -213,16 +204,18 @@ static unsigned countNumFlattenedElementTypes(Type type) {
 ///   ==> returns 1110
 ///   (because the lowered SIL type is (A, B, C, D) -> R)
 ///
+/// Pass `isMethod = true` when the function is a method.
+///
 /// Pass `selfUncurried = true` when the function type is a for method whose
 /// self parameter has been uncurried as in (A, B, C, Self) -> R.
 ///
 llvm::SmallBitVector
 AutoDiffParameterIndices::getLowered(AnyFunctionType *functionType,
-                                     bool selfUncurried) const {
+                                     bool isMethod, bool selfUncurried) const {
   // Calculate the lowered sizes of all the parameters.
   AnyFunctionType *unwrapped = selfUncurried
       ? functionType
-      : unwrapSelfParameter(functionType, isMethodFlag);
+      : unwrapSelfParameter(functionType, isMethod);
   SmallVector<unsigned, 8> paramLoweredSizes;
   unsigned totalLoweredSize = 0;
   auto addLoweredParamInfo = [&](Type type) {
@@ -232,7 +225,7 @@ AutoDiffParameterIndices::getLowered(AnyFunctionType *functionType,
   };
   for (auto &param : unwrapped->getParams())
     addLoweredParamInfo(param.getPlainType());
-  if (isMethodFlag && !selfUncurried)
+  if (isMethod && !selfUncurried)
     addLoweredParamInfo(functionType->getParams()[0].getPlainType());
 
   // Construct the result by setting each range of bits that corresponds to each
@@ -256,29 +249,29 @@ static unsigned getNumAutoDiffParameterIndices(AnyFunctionType *functionType,
 }
 
 unsigned AutoDiffParameterIndicesBuilder::getNumNonSelfParameters() const {
-  return indices.size() - (isMethodFlag ? 1 : 0);
+  return indices.size() - (isMethod ? 1 : 0);
 }
 
 AutoDiffParameterIndicesBuilder::AutoDiffParameterIndicesBuilder(
     AnyFunctionType *functionType, bool isMethod, bool setAllParams) :
     indices(getNumAutoDiffParameterIndices(functionType, isMethod),
             setAllParams),
-    isMethodFlag(isMethod) {
+    isMethod(isMethod) {
 }
 
 AutoDiffParameterIndices *
 AutoDiffParameterIndicesBuilder::build(ASTContext &C) const {
-  return AutoDiffParameterIndices::get(indices, isMethodFlag, C);
+  return AutoDiffParameterIndices::get(indices, C);
 }
 
-/// Adds the indexed parameter to the set. When `isMethodFlag` is not set, the
+/// Adds the indexed parameter to the set. When `isMethod` is not set, the
 /// indices index into the first parameter list. For example,
 ///
 ///   functionType = (A, B, C) -> R
 ///   paramIndex = 0
 ///   ==> adds "A" to the set.
 ///
-/// When `isMethodFlag` is set, the indices index into the first non-self
+/// When `isMethod` is set, the indices index into the first non-self
 /// parameter list. For example,
 ///
 ///   functionType = (Self) -> (A, B, C) -> R
@@ -303,14 +296,14 @@ void AutoDiffParameterIndicesBuilder::setAllNonSelfParameters() {
   indices.set(0, getNumNonSelfParameters());
 }
 
-/// Adds the self parameter to the set. `isMethodFlag` must be set. For
+/// Adds the self parameter to the set. `isMethod` must be set. For
 /// example,
 ///
 ///   functionType = (Self) -> (A, B, C) -> R
 ///   ==> adds "Self" to the set
 ///
 void AutoDiffParameterIndicesBuilder::setSelfParameter() {
-  assert(isMethodFlag &&
+  assert(isMethod &&
          "trying to add self param to non-method parameter indices");
   indices.set(indices.size() - 1);
 }

lib/AST/Builtins.cpp

@@ -1077,7 +1077,7 @@ static ValueDecl *getAutoDiffGetAssociatedFunction(
       // stored in the function type.
       auto *vjpType = origFnTy->getAutoDiffAssociatedFunctionType(
           paramIndices, /*resultIndex*/ 0, /*differentiationOrder*/ 1,
-          kind, /*lookupConformance*/ nullptr);
+          kind, /*lookupConformance*/ nullptr, /*isMethod*/ false);
       vjpType = vjpType->withExtInfo(vjpType->getExtInfo().withNoEscape(false));
       return vjpType;
     }

lib/AST/Type.cpp

@@ -4102,13 +4102,14 @@ makeFunctionType(AnyFunctionType *copy, ArrayRef<AnyFunctionType::Param> params,
 }
 
 AnyFunctionType *AnyFunctionType::getAutoDiffAdjointFunctionType(
-    AutoDiffParameterIndices *indices, const TupleType *primalResultTy) {
+    AutoDiffParameterIndices *indices, const TupleType *primalResultTy,
+    bool isMethod) {
   assert(!indices->isEmpty() && "there must be at least one wrt index");
 
   // Compute the return type of the adjoint.
   SmallVector<TupleTypeElt, 8> retElts;
   SmallVector<Type, 8> wrtParamTypes;
-  indices->getSubsetParameterTypes(this, wrtParamTypes);
+  indices->getSubsetParameterTypes(this, wrtParamTypes, isMethod);
   for (auto wrtParamType : wrtParamTypes)
     retElts.push_back(wrtParamType);
 
@@ -4121,7 +4122,7 @@ AnyFunctionType *AnyFunctionType::getAutoDiffAdjointFunctionType(
   // If this is a method, unwrap the function type so that we can see the
   // non-self parameters.
   AnyFunctionType *unwrapped = this;
-  if (indices->isMethod())
+  if (isMethod)
     unwrapped = unwrapped->getResult()->castTo<AnyFunctionType>();
 
   // Compute the adjoint parameters.
@@ -4149,7 +4150,7 @@ AnyFunctionType *AnyFunctionType::getAutoDiffAdjointFunctionType(
 
   // If this is a method, wrap the adjoint type in an additional "(Self) ->"
   // curry level.
-  if (indices->isMethod())
+  if (isMethod)
     adjoint = makeFunctionType(this, getParams(), adjoint);
 
   return adjoint;
@@ -4158,7 +4159,7 @@ AnyFunctionType *AnyFunctionType::getAutoDiffAdjointFunctionType(
 AnyFunctionType *AnyFunctionType::getAutoDiffAssociatedFunctionType(
     AutoDiffParameterIndices *indices, unsigned resultIndex,
     unsigned differentiationOrder, AutoDiffAssociatedFunctionKind kind,
-    LookupConformanceFn lookupConformance, bool selfUncurried) {
+    LookupConformanceFn lookupConformance, bool isMethod, bool selfUncurried) {
   // JVP: (T...) -> ((R...),
   //                 (T.TangentVector...) -> (R.TangentVector...))
   // VJP: (T...) -> ((R...),
@@ -4221,12 +4222,12 @@ AnyFunctionType *AnyFunctionType::getAutoDiffAssociatedFunctionType(
   };
 
   SmallVector<Type, 8> wrtParamTypes;
-  indices->getSubsetParameterTypes(this, wrtParamTypes, selfUncurried);
+  indices->getSubsetParameterTypes(this, wrtParamTypes, isMethod, selfUncurried);
 
   // If this is a method, unwrap the function type so that we can see the
   // non-self parameters and the final result.
   AnyFunctionType *unwrapped = this;
-  if (indices->isMethod() && !selfUncurried)
+  if (isMethod && !selfUncurried)
     unwrapped = unwrapped->getResult()->castTo<AnyFunctionType>();
   Type originalResult = unwrapped->getResult();
 
@@ -4298,7 +4299,7 @@ AnyFunctionType *AnyFunctionType::getAutoDiffAssociatedFunctionType(
 
   // If this is a method, wrap the associated function type in an additional
   // "(Self) ->" curry level.
-  if (indices->isMethod() && !selfUncurried)
+  if (isMethod && !selfUncurried)
     associatedFunction =
         makeFunctionType(this, getParams(), associatedFunction);