Update to use LinearDifferentiableFunctionTypeComponent.

Author: rxwei

include/swift/AST/AutoDiff.h

@@ -52,9 +52,19 @@ enum class NormalDifferentiableFunctionTypeComponent : uint8_t {
   VJP = 2
 };
 
-enum class LinearDifferentiableFunctionTypeComponent : uint8_t {
-  Original = 0,
-  Transpose = 1
+struct LinearDifferentiableFunctionTypeComponent {
+  enum innerty : unsigned {
+    Original = 0,
+    Transpose = 1,
+  } rawValue;
+
+  LinearDifferentiableFunctionTypeComponent() = default;
+  LinearDifferentiableFunctionTypeComponent(innerty rawValue)
+      : rawValue(rawValue) {}
+  explicit LinearDifferentiableFunctionTypeComponent(unsigned rawValue) :
+      LinearDifferentiableFunctionTypeComponent((innerty)rawValue) {}
+  explicit LinearDifferentiableFunctionTypeComponent(StringRef name);
+  operator innerty() const { return rawValue; }
 };
 
 class ParsedAutoDiffParameter {

include/swift/SIL/SILBuilder.h

@@ -535,7 +535,7 @@ class SILBuilder {
   }
 
   LinearFunctionExtractInst *createLinearFunctionExtract(
-      SILLocation Loc, LinearFunctionExtractee Extractee,
+      SILLocation Loc, LinearDifferentiableFunctionTypeComponent Extractee,
       SILValue TheFunction) {
     return insert(new (getModule()) LinearFunctionExtractInst(
         getModule(), getSILDebugLocation(Loc), Extractee, TheFunction));

include/swift/SIL/SILInstruction.h

@@ -8022,42 +8022,33 @@ class LinearFunctionExtractInst
     : public InstructionBase<
           SILInstructionKind::LinearFunctionExtractInst,
           SingleValueInstruction> {
-public:
-  struct Extractee {
-    enum innerty : unsigned {
-      Original = 0,
-      Transpose = 1
-    } rawValue;
-    Extractee() = default;
-    Extractee(innerty rawValue) : rawValue(rawValue) {}
-    explicit Extractee(unsigned rawValue) : Extractee((innerty)rawValue) {}
-    explicit Extractee(StringRef name);
-    operator innerty() const { return rawValue; }
-  };
-
 private:
   /// The extractee.
-  Extractee extractee;
+  LinearDifferentiableFunctionTypeComponent extractee;
   /// The list containing the `@differentiable(linear)` function operand.
   FixedOperandList<1> operands;
 
   static SILType
-  getExtracteeType(SILValue function, Extractee extractee, SILModule &module);
+  getExtracteeType(SILValue function,
+                   LinearDifferentiableFunctionTypeComponent extractee,
+                   SILModule &module);
 
 public:
   explicit LinearFunctionExtractInst(
-      SILModule &module, SILDebugLocation debugLoc, Extractee extractee,
+      SILModule &module, SILDebugLocation debugLoc,
+      LinearDifferentiableFunctionTypeComponent extractee,
       SILValue theFunction);
 
-  Extractee getExtractee() const { return extractee; }
+  LinearDifferentiableFunctionTypeComponent getExtractee() const {
+    return extractee;
+  }
 
   SILValue getFunctionOperand() const { return operands[0].get(); }
   ArrayRef<Operand> getAllOperands() const { return operands.asArray(); }
   MutableArrayRef<Operand> getAllOperands() { return operands.asArray(); }
 };
 
-typedef LinearFunctionExtractInst::Extractee
-    LinearFunctionExtractee;
+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,16 @@ AutoDiffDerivativeFunctionKind(StringRef string) {
   rawValue = *result;
 }
 
+LinearDifferentiableFunctionTypeComponent::
+LinearDifferentiableFunctionTypeComponent(StringRef string) {
+  Optional<innerty> result =
+      llvm::StringSwitch<Optional<innerty>>(string)
+          .Case("original", Original)
+          .Case("transpose", Transpose);
+  assert(result && "Invalid string");
+  rawValue = *result;
+}
+
 // TODO(TF-874): This helper is inefficient and should be removed. Unwrapping at
 // most once (for curried method types) is sufficient.
 static void unwrapCurryLevels(AnyFunctionType *fnTy,

lib/ParseSIL/ParseSIL.cpp

@@ -3064,7 +3064,7 @@ bool SILParser::parseSILInstruction(SILBuilder &B) {
   case SILInstructionKind::LinearFunctionExtractInst: {
     // Parse the rest of the instruction: an extractee, a linear function
     // operand, and a debug location.
-    LinearFunctionExtractee extractee;
+    LinearDifferentiableFunctionTypeComponent extractee;
     StringRef extracteeNames[2] = {"original", "transpose"};
     SILValue functionOperand;
     SourceLoc lastLoc;

lib/SIL/SILInstructions.cpp

@@ -633,15 +633,6 @@ SILType LinearFunctionInst::getLinearFunctionType(
   return SILType::getPrimitiveObjectType(diffTy);
 }
 
-LinearFunctionExtractInst::Extractee::Extractee(StringRef string) {
-  Optional<innerty> result =
-      llvm::StringSwitch<Optional<innerty>>(string)
-          .Case("original", Original)
-          .Case("transpose", Transpose);
-  assert(result && "Invalid string");
-  rawValue = *result;
-}
-
 LinearFunctionInst::LinearFunctionInst(
     SILDebugLocation Loc, IndexSubset *ParameterIndices,
     SILValue OriginalFunction, Optional<SILValue> TransposeFunction,
@@ -729,14 +720,16 @@ DifferentiableFunctionExtractInst::DifferentiableFunctionExtractInst(
       extractee(extractee), operands(this, theFunction) {}
 
 SILType LinearFunctionExtractInst::
-getExtracteeType(SILValue function, Extractee extractee, SILModule &module) {
+getExtracteeType(
+    SILValue function, LinearDifferentiableFunctionTypeComponent extractee,
+    SILModule &module) {
   auto fnTy = function->getType().castTo<SILFunctionType>();
   assert(fnTy->getDifferentiabilityKind() == DifferentiabilityKind::Linear);
   auto originalFnTy = fnTy->getWithoutDifferentiability();
   switch (extractee) {
-  case Extractee::Original:
+  case LinearDifferentiableFunctionTypeComponent::Original:
     return SILType::getPrimitiveObjectType(originalFnTy);
-  case Extractee::Transpose:
+  case LinearDifferentiableFunctionTypeComponent::Transpose:
     auto transposeFnTy = originalFnTy->getAutoDiffTransposeFunctionType(
         fnTy->getDifferentiationParameterIndices(), module.Types,
         LookUpConformanceInModule(module.getSwiftModule()));
@@ -745,8 +738,8 @@ getExtracteeType(SILValue function, Extractee extractee, SILModule &module) {
 }
 
 LinearFunctionExtractInst::LinearFunctionExtractInst(
-    SILModule &module, SILDebugLocation debugLoc, Extractee extractee,
-    SILValue theFunction)
+    SILModule &module, SILDebugLocation debugLoc,
+    LinearDifferentiableFunctionTypeComponent extractee, SILValue theFunction)
     : InstructionBase(debugLoc,
                       getExtracteeType(theFunction, extractee, module)),
       extractee(extractee), operands(this, theFunction) {}