[AutoDiff] NFC: Change DifferentiableFunctionExtractee to a top-level type.

include/swift/AST/AutoDiff.h

@@ -45,11 +45,58 @@ enum class DifferentiabilityKind: uint8_t {
   Linear = 2
 };
 
-// TODO(TF-904): Replace `DifferentiableFunctionExtractInst::Extractee`.
-enum class NormalDifferentiableFunctionTypeComponent : uint8_t {
-  Original = 0,
-  JVP = 1,
-  VJP = 2
+/// The kind of an linear map.
+struct AutoDiffLinearMapKind {
+  enum innerty : uint8_t {
+    // The differential function.
+    Differential = 0,
+    // The pullback function.
+    Pullback = 1
+  } rawValue;
+
+  AutoDiffLinearMapKind() = default;
+  AutoDiffLinearMapKind(innerty rawValue) : rawValue(rawValue) {}
+  operator innerty() const { return rawValue; }
+};
+
+/// The kind of a derivative function.
+struct AutoDiffDerivativeFunctionKind {
+  enum innerty : uint8_t {
+   // The Jacobian-vector products function.
+   JVP = 0,
+   // The vector-Jacobian products function.
+   VJP = 1
+  } rawValue;
+
+  AutoDiffDerivativeFunctionKind() = default;
+  AutoDiffDerivativeFunctionKind(innerty rawValue) : rawValue(rawValue) {}
+  AutoDiffDerivativeFunctionKind(AutoDiffLinearMapKind linMapKind)
+      : rawValue(static_cast<innerty>(linMapKind.rawValue)) {}
+  explicit AutoDiffDerivativeFunctionKind(StringRef string);
+  operator innerty() const { return rawValue; }
+  AutoDiffLinearMapKind getLinearMapKind() {
+    return (AutoDiffLinearMapKind::innerty)rawValue;
+  }
+};
+
+struct NormalDifferentiableFunctionTypeComponent {
+  enum innerty : unsigned {
+    Original = 0,
+    JVP = 1,
+    VJP = 2
+  } rawValue;
+
+  NormalDifferentiableFunctionTypeComponent() = default;
+  NormalDifferentiableFunctionTypeComponent(innerty rawValue)
+      : rawValue(rawValue) {}
+  NormalDifferentiableFunctionTypeComponent(
+      AutoDiffDerivativeFunctionKind kind);
+  explicit NormalDifferentiableFunctionTypeComponent(unsigned rawValue) :
+      NormalDifferentiableFunctionTypeComponent((innerty)rawValue) {}
+  explicit NormalDifferentiableFunctionTypeComponent(StringRef name);
+  operator innerty() const { return rawValue; }
+
+  Optional<AutoDiffDerivativeFunctionKind> getAsDerivativeFunctionKind() const;
 };
 
 struct LinearDifferentiableFunctionTypeComponent {
@@ -196,40 +243,6 @@ inline llvm::raw_ostream &operator<<(llvm::raw_ostream &s,
   return s;
 }
 
-/// The kind of an linear map.
-struct AutoDiffLinearMapKind {
-  enum innerty : uint8_t {
-    // The differential function.
-    Differential = 0,
-    // The pullback function.
-    Pullback = 1
-  } rawValue;
-
-  AutoDiffLinearMapKind() = default;
-  AutoDiffLinearMapKind(innerty rawValue) : rawValue(rawValue) {}
-  operator innerty() const { return rawValue; }
-};
-
-/// The kind of a derivative function.
-struct AutoDiffDerivativeFunctionKind {
-  enum innerty : uint8_t {
-   // The Jacobian-vector products function.
-   JVP = 0,
-   // The vector-Jacobian products function.
-   VJP = 1
-  } rawValue;
-
-  AutoDiffDerivativeFunctionKind() = default;
-  AutoDiffDerivativeFunctionKind(innerty rawValue) : rawValue(rawValue) {}
-  AutoDiffDerivativeFunctionKind(AutoDiffLinearMapKind linMapKind)
-      : rawValue(static_cast<innerty>(linMapKind.rawValue)) {}
-  explicit AutoDiffDerivativeFunctionKind(StringRef string);
-  operator innerty() const { return rawValue; }
-  AutoDiffLinearMapKind getLinearMapKind() {
-    return (AutoDiffLinearMapKind::innerty)rawValue;
-  }
-};
-
 /// Identifies an autodiff derivative function configuration:
 /// - Parameter indices.
 /// - Result indices.

include/swift/SIL/SILBuilder.h

@@ -528,7 +528,7 @@ class SILBuilder {
   }
 
   DifferentiableFunctionExtractInst *createDifferentiableFunctionExtract(
-      SILLocation Loc, DifferentiableFunctionExtractee Extractee,
+      SILLocation Loc, NormalDifferentiableFunctionTypeComponent Extractee,
       SILValue TheFunction) {
     return insert(new (getModule()) DifferentiableFunctionExtractInst(
         getModule(), getSILDebugLocation(Loc), Extractee, TheFunction));
@@ -546,7 +546,7 @@ class SILBuilder {
                                               SILValue TheFunction) {
     return insert(new (getModule()) DifferentiableFunctionExtractInst(
         getModule(), getSILDebugLocation(Loc),
-        DifferentiableFunctionExtractee::Original, TheFunction));
+        NormalDifferentiableFunctionTypeComponent::Original, TheFunction));
   }
 
   BuiltinInst *createBuiltin(SILLocation Loc, Identifier Name, SILType ResultTy,

include/swift/SIL/SILInstruction.h

@@ -7967,42 +7967,29 @@ class DifferentiableFunctionExtractInst
     : public InstructionBase<
           SILInstructionKind::DifferentiableFunctionExtractInst,
           SingleValueInstruction> {
-public:
-  struct Extractee {
-    enum innerty : unsigned {
-      Original = 0,
-      JVP = 1,
-      VJP = 2
-    } rawValue;
-    Extractee() = default;
-    Extractee(innerty rawValue) : rawValue(rawValue) {}
-    explicit Extractee(unsigned rawValue) : Extractee((innerty)rawValue) {}
-    Extractee(AutoDiffDerivativeFunctionKind kind);
-    explicit Extractee(StringRef name);
-    operator innerty() const { return rawValue; }
-
-    Optional<AutoDiffDerivativeFunctionKind>
-    getExtracteeAsDerivativeFunction() const;
-  };
-
 private:
   /// The extractee.
-  Extractee extractee;
+  NormalDifferentiableFunctionTypeComponent extractee;
   /// The list containing the `@differentiable` function operand.
   FixedOperandList<1> operands;
 
   static SILType
-  getExtracteeType(SILValue function, Extractee extractee, SILModule &module);
+  getExtracteeType(
+      SILValue function, NormalDifferentiableFunctionTypeComponent extractee,
+      SILModule &module);
 
 public:
   explicit DifferentiableFunctionExtractInst(
-      SILModule &module, SILDebugLocation debugLoc, Extractee extractee,
+      SILModule &module, SILDebugLocation debugLoc,
+      NormalDifferentiableFunctionTypeComponent extractee,
       SILValue theFunction);
 
-  Extractee getExtractee() const { return extractee; }
+  NormalDifferentiableFunctionTypeComponent getExtractee() const {
+      return extractee;
+  }
 
   AutoDiffDerivativeFunctionKind getDerivativeFunctionKind() const {
-    auto kind = extractee.getExtracteeAsDerivativeFunction();
+    auto kind = extractee.getAsDerivativeFunctionKind();
     assert(kind);
     return *kind;
   }
@@ -8012,9 +7999,6 @@ class DifferentiableFunctionExtractInst
   MutableArrayRef<Operand> getAllOperands() { return operands.asArray(); }
 };
 
-typedef DifferentiableFunctionExtractInst::Extractee
-    DifferentiableFunctionExtractee;
-
 /// `linear_function_extract` - given an `@differentiable(linear)` function
 /// representing a bundle of the original function and the transpose function,
 /// extract the specified function.
@@ -8047,8 +8031,6 @@ class LinearFunctionExtractInst
   ArrayRef<Operand> getAllOperands() const { return operands.asArray(); }
   MutableArrayRef<Operand> getAllOperands() { return operands.asArray(); }
 };
-
-typedef LinearDifferentiableFunctionTypeComponent LinearFunctionExtractee;
 // SWIFT_ENABLE_TENSORFLOW END
 
 // This is defined out of line to work around the fact that this depends on

lib/AST/AutoDiff.cpp

@@ -32,6 +32,33 @@ AutoDiffDerivativeFunctionKind(StringRef string) {
   rawValue = *result;
 }
 
+NormalDifferentiableFunctionTypeComponent::
+NormalDifferentiableFunctionTypeComponent(AutoDiffDerivativeFunctionKind kind) {
+  switch (kind) {
+  case AutoDiffDerivativeFunctionKind::JVP: rawValue = JVP; return;
+  case AutoDiffDerivativeFunctionKind::VJP: rawValue = VJP; return;
+  }
+}
+
+NormalDifferentiableFunctionTypeComponent::
+NormalDifferentiableFunctionTypeComponent(StringRef string) {
+  Optional<innerty> result = llvm::StringSwitch<Optional<innerty>>(string)
+      .Case("original", Original)
+      .Case("jvp", JVP)
+      .Case("vjp", VJP);
+  assert(result && "Invalid string");
+  rawValue = *result;
+}
+
+Optional<AutoDiffDerivativeFunctionKind>
+NormalDifferentiableFunctionTypeComponent::getAsDerivativeFunctionKind() const {
+  switch (rawValue) {
+  case Original: return None;
+  case JVP: return {AutoDiffDerivativeFunctionKind::JVP};
+  case VJP: return {AutoDiffDerivativeFunctionKind::VJP};
+  }
+}
+
 LinearDifferentiableFunctionTypeComponent::
 LinearDifferentiableFunctionTypeComponent(StringRef string) {
   Optional<innerty> result =

lib/IRGen/GenDiffFunc.cpp

@@ -41,34 +41,34 @@ class DifferentiableFuncFieldInfo final
     : public RecordField<DifferentiableFuncFieldInfo> {
 public:
   DifferentiableFuncFieldInfo(
-      DifferentiableFunctionExtractee component, const TypeInfo &type,
+      NormalDifferentiableFunctionTypeComponent component, const TypeInfo &type,
       IndexSubset *parameterIndices)
       : RecordField(type), component(component),
         parameterIndices(parameterIndices) {}
 
   /// The field index.
-  const DifferentiableFunctionExtractee component;
+  const NormalDifferentiableFunctionTypeComponent component;
 
   /// The parameter indices.
   IndexSubset *parameterIndices;
 
   std::string getFieldName() const {
     switch (component) {
-    case DifferentiableFunctionExtractee::Original:
+    case NormalDifferentiableFunctionTypeComponent::Original:
       return "original";
-    case DifferentiableFunctionExtractee::JVP:
+    case NormalDifferentiableFunctionTypeComponent::JVP:
       return "jvp";
-    case DifferentiableFunctionExtractee::VJP:
+    case NormalDifferentiableFunctionTypeComponent::VJP:
       return "vjp";
     }
   }
 
   SILType getType(IRGenModule &IGM, SILType t) const {
     auto fnTy = t.castTo<SILFunctionType>();
     auto origFnTy = fnTy->getWithoutDifferentiability();
-    if (component == DifferentiableFunctionExtractee::Original)
+    if (component == NormalDifferentiableFunctionTypeComponent::Original)
       return SILType::getPrimitiveObjectType(origFnTy);
-    auto kind = *component.getExtracteeAsDerivativeFunction();
+    auto kind = *component.getAsDerivativeFunctionKind();
     auto assocTy = origFnTy->getAutoDiffDerivativeFunctionType(
         parameterIndices, /*resultIndex*/ 0, kind,
         IGM.getSILTypes(), LookUpConformanceInModule(IGM.getSwiftModule()));
@@ -79,8 +79,8 @@ class DifferentiableFuncFieldInfo final
 class DifferentiableFuncTypeInfo final
     : public RecordTypeInfo<DifferentiableFuncTypeInfo, LoadableTypeInfo,
                             DifferentiableFuncFieldInfo> {
-  using super =
-      RecordTypeInfo<DifferentiableFuncTypeInfo, LoadableTypeInfo, DifferentiableFuncFieldInfo>;
+  using super = RecordTypeInfo<DifferentiableFuncTypeInfo, LoadableTypeInfo,
+                               DifferentiableFuncFieldInfo>;
 
 public:
   DifferentiableFuncTypeInfo(
@@ -117,7 +117,7 @@ class DifferentiableFuncTypeInfo final
 
 class DifferentiableFuncTypeBuilder
     : public RecordTypeBuilder<DifferentiableFuncTypeBuilder, DifferentiableFuncFieldInfo,
-                               DifferentiableFunctionExtractee> {
+                               NormalDifferentiableFunctionTypeComponent> {
 
   SILFunctionType *originalType;
   IndexSubset *parameterIndices;
@@ -151,15 +151,15 @@ class DifferentiableFuncTypeBuilder
   }
 
   DifferentiableFuncFieldInfo getFieldInfo(
-      unsigned index, DifferentiableFunctionExtractee component,
+      unsigned index, NormalDifferentiableFunctionTypeComponent component,
       const TypeInfo &fieldTI) {
     return DifferentiableFuncFieldInfo(component, fieldTI, parameterIndices);
   }
 
-  SILType getType(DifferentiableFunctionExtractee component) {
-    if (component == DifferentiableFunctionExtractee::Original)
+  SILType getType(NormalDifferentiableFunctionTypeComponent component) {
+    if (component == NormalDifferentiableFunctionTypeComponent::Original)
       return SILType::getPrimitiveObjectType(originalType->getCanonicalType());
-    auto kind = *component.getExtracteeAsDerivativeFunction();
+    auto kind = *component.getAsDerivativeFunctionKind();
     auto assocTy = originalType->getAutoDiffDerivativeFunctionType(
         parameterIndices, /*resultIndex*/ 0, kind, IGM.getSILTypes(),
         LookUpConformanceInModule(IGM.getSwiftModule()));
@@ -320,9 +320,9 @@ class LinearFuncTypeBuilder
 const TypeInfo *
 TypeConverter::convertNormalDifferentiableFunctionType(SILFunctionType *type) {
   DifferentiableFuncTypeBuilder builder(IGM, type);
-  return builder.layout({DifferentiableFunctionExtractee::Original,
-                         DifferentiableFunctionExtractee::JVP,
-                         DifferentiableFunctionExtractee::VJP});
+  return builder.layout({NormalDifferentiableFunctionTypeComponent::Original,
+                         NormalDifferentiableFunctionTypeComponent::JVP,
+                         NormalDifferentiableFunctionTypeComponent::VJP});
 }
 
 const TypeInfo *

lib/ParseSIL/ParseSIL.cpp

@@ -3041,7 +3041,7 @@ bool SILParser::parseSILInstruction(SILBuilder &B) {
   case SILInstructionKind::DifferentiableFunctionExtractInst: {
     // Parse the rest of the instruction: an extractee, a differentiable
     // function operand, and a debug location.
-    DifferentiableFunctionExtractee extractee;
+    NormalDifferentiableFunctionTypeComponent extractee;
     StringRef extracteeNames[3] = {"original", "jvp", "vjp"};
     SILValue functionOperand;
     SourceLoc lastLoc;

lib/SIL/SILInstructions.cpp

@@ -664,45 +664,16 @@ LinearFunctionInst *LinearFunctionInst::create(
       HasOwnership);
 }
 
-DifferentiableFunctionExtractInst::Extractee::Extractee(
-    AutoDiffDerivativeFunctionKind kind) {
-  switch (kind) {
-  case AutoDiffDerivativeFunctionKind::JVP:
-    rawValue = JVP;
-    return;
-  case AutoDiffDerivativeFunctionKind::VJP:
-    rawValue = VJP;
-    return;
-  }
-}
-
-DifferentiableFunctionExtractInst::Extractee::Extractee(StringRef string) {
-  Optional<innerty> result = llvm::StringSwitch<Optional<innerty>>(string)
-      .Case("original", Original)
-      .Case("jvp", JVP)
-      .Case("vjp", VJP);
-  assert(result && "Invalid string");
-  rawValue = *result;
-}
-
-Optional<AutoDiffDerivativeFunctionKind>
-DifferentiableFunctionExtractInst::Extractee::
-getExtracteeAsDerivativeFunction() const {
-  switch (rawValue) {
-  case Original: return None;
-  case JVP: return {AutoDiffDerivativeFunctionKind::JVP};
-  case VJP: return {AutoDiffDerivativeFunctionKind::VJP};
-  }
-}
-
 SILType DifferentiableFunctionExtractInst::
-getExtracteeType(SILValue function, Extractee extractee, SILModule &module) {
+getExtracteeType(
+    SILValue function, NormalDifferentiableFunctionTypeComponent extractee,
+    SILModule &module) {
   auto fnTy = function->getType().castTo<SILFunctionType>();
   assert(fnTy->getDifferentiabilityKind() == DifferentiabilityKind::Normal);
   auto originalFnTy = fnTy->getWithoutDifferentiability();
-  auto kindOpt = extractee.getExtracteeAsDerivativeFunction();
+  auto kindOpt = extractee.getAsDerivativeFunctionKind();
   if (!kindOpt) {
-    assert(extractee == Extractee::Original);
+    assert(extractee == NormalDifferentiableFunctionTypeComponent::Original);
     return SILType::getPrimitiveObjectType(originalFnTy);
   }
   auto resultFnTy = originalFnTy->getAutoDiffDerivativeFunctionType(
@@ -713,8 +684,8 @@ getExtracteeType(SILValue function, Extractee extractee, SILModule &module) {
 }
 
 DifferentiableFunctionExtractInst::DifferentiableFunctionExtractInst(
-    SILModule &module, SILDebugLocation debugLoc, Extractee extractee,
-    SILValue theFunction)
+    SILModule &module, SILDebugLocation debugLoc,
+    NormalDifferentiableFunctionTypeComponent extractee, SILValue theFunction)
     : InstructionBase(debugLoc,
                       getExtracteeType(theFunction, extractee, module)),
       extractee(extractee), operands(this, theFunction) {}