Limit scope of changes only to AutoDiff code

Author: kovdan01

include/swift/AST/Expr.h

@@ -4482,8 +4482,6 @@ class ClosureExpr : public AbstractClosureExpr {
 class AutoClosureExpr : public AbstractClosureExpr {
   BraceStmt *Body;
 
-  ApplyExpr *getUnwrappedCurryThunkImpl() const;
-
 public:
   enum class Kind : uint8_t {
     // An autoclosure with type () -> Result. Formed from type checking an
@@ -4545,10 +4543,6 @@ class AutoClosureExpr : public AbstractClosureExpr {
   /// - otherwise, returns nullptr for convenience.
   Expr *getUnwrappedCurryThunkExpr() const;
 
-  /// Same as getUnwrappedCurryThunkExpr, but get the called ValueDecl instead
-  /// of the expr.
-  ValueDecl *getUnwrappedCurryThunkCalledValue() const;
-
   // Implement isa/cast/dyncast/etc.
   static bool classof(const Expr *E) {
     return E->getKind() == ExprKind::AutoClosure;

lib/AST/Expr.cpp

@@ -2145,7 +2145,7 @@ Expr *AutoClosureExpr::getSingleExpressionBody() const {
   return cast<ReturnStmt>(Body->getLastElement().get<Stmt *>())->getResult();
 }
 
-ApplyExpr *AutoClosureExpr::getUnwrappedCurryThunkImpl() const {
+Expr *AutoClosureExpr::getUnwrappedCurryThunkExpr() const {
   auto maybeUnwrapOpenExistential = [](Expr *expr) {
     if (auto *openExistential = dyn_cast<OpenExistentialExpr>(expr)) {
       expr = openExistential->getSubExpr()->getSemanticsProvidingExpr();
@@ -2189,7 +2189,7 @@ ApplyExpr *AutoClosureExpr::getUnwrappedCurryThunkImpl() const {
     body = maybeUnwrapConversions(body);
 
     if (auto *outerCall = dyn_cast<ApplyExpr>(body)) {
-      return outerCall;
+      return outerCall->getFn();
     }
 
     assert(false && "Malformed curry thunk?");
@@ -2210,7 +2210,7 @@ ApplyExpr *AutoClosureExpr::getUnwrappedCurryThunkImpl() const {
       if (auto *outerCall = dyn_cast<ApplyExpr>(innerBody)) {
         auto outerFn = maybeUnwrapConversions(outerCall->getFn());
         if (auto *innerCall = dyn_cast<ApplyExpr>(outerFn)) {
-          return innerCall;
+          return innerCall->getFn();
         }
       }
     }
@@ -2223,20 +2223,6 @@ ApplyExpr *AutoClosureExpr::getUnwrappedCurryThunkImpl() const {
   return nullptr;
 }
 
-Expr *AutoClosureExpr::getUnwrappedCurryThunkExpr() const {
-  ApplyExpr *ae = getUnwrappedCurryThunkImpl();
-  if (ae == nullptr)
-    return nullptr;
-  return ae->getFn();
-}
-
-ValueDecl *AutoClosureExpr::getUnwrappedCurryThunkCalledValue() const {
-  ApplyExpr *ae = getUnwrappedCurryThunkImpl();
-  if (ae == nullptr)
-    return nullptr;
-  return ae->getCalledValue(/*skipFunctionConversions=*/true);
-}
-
 FORWARD_SOURCE_LOCS_TO(UnresolvedPatternExpr, subPattern)
 
 TypeExpr::TypeExpr(TypeRepr *Repr)

lib/SILOptimizer/Mandatory/Differentiation.cpp

@@ -485,8 +485,18 @@ SILFunction *DifferentiationTransformer::getUnwrappedCurryThunkFunction(
   if (autoCE == nullptr)
     return nullptr;
 
-  auto *afd =
-      cast<AbstractFunctionDecl>(autoCE->getUnwrappedCurryThunkCalledValue());
+  AbstractFunctionDecl *afd = nullptr;
+  Expr *expr = autoCE->getUnwrappedCurryThunkExpr();
+  switch (autoCE->getThunkKind()) {
+  default:
+    llvm_unreachable("Only single and double curry thunks are expected");
+  case AutoClosureExpr::Kind::SingleCurryThunk:
+  case AutoClosureExpr::Kind::DoubleCurryThunk:
+    afd = cast<AbstractFunctionDecl>(cast<ApplyExpr>(expr)->getCalledValue(
+        /*skipFunctionConversions=*/true));
+    break;
+  }
+  assert(afd);
 
   auto silFnIt = afdToSILFn.find(afd);
   if (silFnIt == afdToSILFn.end()) {