Merge pull request #30084 from slavapestov/sr-75-prep

More Sema cleanups to prepare for curry thunks

Author: slavapestov

include/swift/AST/Expr.h

@@ -284,10 +284,10 @@ class alignas(8) Expr {
     Discriminator : 16
   );
 
-  SWIFT_INLINE_BITFIELD(AutoClosureExpr, AbstractClosureExpr, 1,
-    /// True if this autoclosure was built for a function conversion, and
-    /// not an actual @autoclosure parameter.
-    IsThunk : 1
+  SWIFT_INLINE_BITFIELD(AutoClosureExpr, AbstractClosureExpr, 2,
+    /// If the autoclosure was built for a curry thunk, the thunk kind is
+    /// stored here.
+    Kind : 2
   );
 
   SWIFT_INLINE_BITFIELD(ClosureExpr, AbstractClosureExpr, 1,
@@ -3790,22 +3790,36 @@ class AutoClosureExpr : public AbstractClosureExpr {
   BraceStmt *Body;
 
 public:
+  enum class Kind : uint8_t {
+    // An autoclosure with type () -> Result. Formed from type checking an
+    // @autoclosure argument in a function call.
+    None = 0,
+
+    // An autoclosure with type (Args...) -> Result. Formed from type checking a
+    // partial application.
+    SingleCurryThunk = 1,
+
+    // An autoclosure with type (Self) -> (Args...) -> Result. Formed from type
+    // checking a partial application.
+    DoubleCurryThunk = 2
+  };
+
   AutoClosureExpr(Expr *Body, Type ResultTy, unsigned Discriminator,
                   DeclContext *Parent)
       : AbstractClosureExpr(ExprKind::AutoClosure, ResultTy, /*Implicit=*/true,
                             Discriminator, Parent) {
     if (Body != nullptr)
       setBody(Body);
 
-    Bits.AutoClosureExpr.IsThunk = false;
+    Bits.AutoClosureExpr.Kind = 0;
   }
 
-  bool isThunk() const {
-    return Bits.AutoClosureExpr.IsThunk;
+  Kind getThunkKind() const {
+    return Kind(Bits.AutoClosureExpr.Kind);
   }
 
-  void setIsThunk(bool isThunk) {
-    Bits.AutoClosureExpr.IsThunk = isThunk;
+  void setThunkKind(Kind K) {
+    Bits.AutoClosureExpr.Kind = uint8_t(K);
   }
 
   SourceRange getSourceRange() const;
@@ -3824,6 +3838,16 @@ class AutoClosureExpr : public AbstractClosureExpr {
   /// The body of an autoclosure always consists of a single expression.
   Expr *getSingleExpressionBody() const;
 
+  /// Unwraps a curry thunk. Basically, this gives you what the old AST looked
+  /// like, before Sema started building curry thunks. This is really only
+  /// meant for legacy usages.
+  ///
+  /// The behavior is as follows, based on the kind:
+  /// - for double curry thunks, returns the original DeclRefExpr.
+  /// - for single curry thunks, returns the ApplyExpr for the self call.
+  /// - otherwise, returns nullptr for convenience.
+  Expr *getUnwrappedCurryThunkExpr() const;
+
   // Implement isa/cast/dyncast/etc.
   static bool classof(const Expr *E) {
     return E->getKind() == ExprKind::AutoClosure;

lib/AST/Expr.cpp

@@ -1843,6 +1843,46 @@ Expr *AutoClosureExpr::getSingleExpressionBody() const {
   return cast<ReturnStmt>(Body->getFirstElement().get<Stmt *>())->getResult();
 }
 
+Expr *AutoClosureExpr::getUnwrappedCurryThunkExpr() const {
+  switch (getThunkKind()) {
+  case AutoClosureExpr::Kind::None:
+    break;
+
+  case AutoClosureExpr::Kind::SingleCurryThunk: {
+    auto *body = getSingleExpressionBody();
+    body = body->getSemanticsProvidingExpr();
+    if (auto *outerCall = dyn_cast<ApplyExpr>(body)) {
+      return outerCall->getFn();
+    }
+
+    assert(false && "Malformed curry thunk?");
+    break;
+  }
+
+  case AutoClosureExpr::Kind::DoubleCurryThunk: {
+    auto *body = getSingleExpressionBody();
+    if (auto *innerClosure = dyn_cast<AutoClosureExpr>(body)) {
+      assert(innerClosure->getThunkKind() ==
+               AutoClosureExpr::Kind::SingleCurryThunk);
+      auto *innerBody = innerClosure->getSingleExpressionBody();
+      innerBody = innerBody->getSemanticsProvidingExpr();
+      if (auto *outerCall = dyn_cast<ApplyExpr>(innerBody)) {
+        if (auto *innerCall = dyn_cast<ApplyExpr>(outerCall->getFn())) {
+          if (auto *declRef = dyn_cast<DeclRefExpr>(innerCall->getFn())) {
+            return declRef;
+          }
+        }
+      }
+    }
+
+    assert(false && "Malformed curry thunk?");
+    break;
+  }
+  }
+
+  return nullptr;
+}
+
 FORWARD_SOURCE_LOCS_TO(UnresolvedPatternExpr, subPattern)
 
 TypeExpr::TypeExpr(TypeLoc TyLoc)

lib/IDE/SourceEntityWalker.cpp

@@ -256,6 +256,21 @@ std::pair<bool, Expr *> SemaAnnotator::walkToExprPre(Expr *E) {
   if (auto *CtorRefE = dyn_cast<ConstructorRefCallExpr>(E))
     CtorRefs.push_back(CtorRefE);
 
+  if (auto *ACE = dyn_cast<AutoClosureExpr>(E)) {
+    if (auto *SubExpr = ACE->getUnwrappedCurryThunkExpr()) {
+      if (auto *DRE = dyn_cast<DeclRefExpr>(SubExpr)) {
+        if (!passReference(DRE->getDecl(), DRE->getType(),
+                                  DRE->getNameLoc(),
+                          ReferenceMetaData(getReferenceKind(Parent.getAsExpr(), DRE),
+                                            OpAccess))) {
+          return { false, nullptr };
+        }
+
+        return { true, E };
+      }
+    }
+  }
+
   if (!isa<InOutExpr>(E) &&
       !isa<LoadExpr>(E) &&
       !isa<OpenExistentialExpr>(E) &&

lib/IDE/Utils.cpp

@@ -989,28 +989,9 @@ std::pair<Type, ConcreteDeclRef> swift::ide::getReferencedDecl(Expr *expr) {
     expr = selfApplyExpr->getFn();
 
   // Look through curry thunks.
-  if (auto *closure = dyn_cast<AutoClosureExpr>(expr)) {
-    if (closure->isThunk()) {
-      auto *body = closure->getSingleExpressionBody();
-      if (isa<AutoClosureExpr>(body) &&
-          closure->getParameters()->size() == 1)
-       expr = closure->getSingleExpressionBody();
-    }
-  }
-
-  if (auto *closure = dyn_cast<AutoClosureExpr>(expr)) {
-    if (closure->isThunk()) {
-      auto *body = closure->getSingleExpressionBody();
-      body = body->getSemanticsProvidingExpr();
-      if (auto *outerCall = dyn_cast<ApplyExpr>(body)) {
-        if (auto *innerCall = dyn_cast<ApplyExpr>(outerCall->getFn())) {
-          if (auto *declRef = dyn_cast<DeclRefExpr>(innerCall->getFn())) {
-            expr = declRef;
-          }
-        }
-      }
-    }
-  }
+  if (auto *closure = dyn_cast<AutoClosureExpr>(expr))
+    if (auto *unwrappedThunk = closure->getUnwrappedCurryThunkExpr())
+      expr = unwrappedThunk;
 
   // If this is an IUO result, unwrap the optional type.
   auto refDecl = expr->getReferencedDecl();