Sema: Use AutoClosureExpr again in KeyPathExpr-as-function expansion

Remove the preallocated closure discriminator from KeyPathExpr and go back
to expanding them using an AutoClosureExpr inside of a CaptureListExpr now
that that's supported. This allows the discriminator to be assigned during
type checking without disturbing the indexing of explicit closure literals.

Author: jckarter

include/swift/AST/Expr.h

@@ -5389,10 +5389,6 @@ class KeyPathExpr : public Expr {
   /// a contextual root type.
   bool HasLeadingDot = false;
 
-  /// When we parse a key path literal, we claim a closure discriminator for it, since it may be used as
-  /// a closure value in function type context.
-  unsigned ClosureDiscriminator;
-
 public:
   /// A single stored component, which will be one of:
   /// - an unresolved DeclNameRef, which has to be type-checked
@@ -5724,12 +5720,11 @@ class KeyPathExpr : public Expr {
 
   KeyPathExpr(SourceLoc startLoc, Expr *parsedRoot, Expr *parsedPath,
               SourceLoc endLoc, bool hasLeadingDot, bool isObjC,
-              bool isImplicit, unsigned closureDiscriminator);
+              bool isImplicit);
 
   /// Create a key path with unresolved root and path expressions.
   KeyPathExpr(SourceLoc backslashLoc, Expr *parsedRoot, Expr *parsedPath,
-              bool hasLeadingDot, bool isImplicit,
-              unsigned closureDiscriminator);
+              bool hasLeadingDot, bool isImplicit);
 
   /// Create a key path with components.
   KeyPathExpr(ASTContext &ctx, SourceLoc startLoc,
@@ -5740,8 +5735,7 @@ class KeyPathExpr : public Expr {
   /// Create a new parsed Swift key path expression.
   static KeyPathExpr *createParsed(ASTContext &ctx, SourceLoc backslashLoc,
      Expr *parsedRoot, Expr *parsedPath,
-     bool hasLeadingDot,
-     unsigned closureDiscriminator = AbstractClosureExpr::InvalidDiscriminator);
+     bool hasLeadingDot);
 
   /// Create a new parsed #keyPath expression.
   static KeyPathExpr *createParsedPoundKeyPath(ASTContext &ctx,
@@ -5836,9 +5830,6 @@ class KeyPathExpr : public Expr {
   /// True if this key path expression has a leading dot.
   bool expectsContextualRoot() const { return HasLeadingDot; }
 
-  /// Return the discriminator to use if this key path becomes a closure.
-  unsigned getClosureDiscriminator() const { return ClosureDiscriminator; }
-
   static bool classof(const Expr *E) {
     return E->getKind() == ExprKind::KeyPath;
   }

lib/AST/Expr.cpp

@@ -2252,24 +2252,21 @@ OpenedArchetypeType *OpenExistentialExpr::getOpenedArchetype() const {
 
 KeyPathExpr::KeyPathExpr(SourceLoc startLoc, Expr *parsedRoot,
                          Expr *parsedPath, SourceLoc endLoc, bool hasLeadingDot,
-                         bool isObjC, bool isImplicit,
-                         unsigned closureDiscriminator)
+                         bool isObjC, bool isImplicit)
     : Expr(ExprKind::KeyPath, isImplicit), StartLoc(startLoc), EndLoc(endLoc),
       ParsedRoot(parsedRoot), ParsedPath(parsedPath),
-      HasLeadingDot(hasLeadingDot), ClosureDiscriminator(closureDiscriminator) {
+      HasLeadingDot(hasLeadingDot) {
   assert(!(isObjC && (parsedRoot || parsedPath)) &&
          "Obj-C key paths should only have components");
   Bits.KeyPathExpr.IsObjC = isObjC;
 }
 
 KeyPathExpr::KeyPathExpr(SourceLoc backslashLoc, Expr *parsedRoot,
-                         Expr *parsedPath, bool hasLeadingDot, bool isImplicit,
-                         unsigned closureDiscriminator)
+                         Expr *parsedPath, bool hasLeadingDot, bool isImplicit)
     : KeyPathExpr(backslashLoc, parsedRoot, parsedPath,
                   parsedPath ? parsedPath->getEndLoc()
                              : parsedRoot->getEndLoc(),
-                  hasLeadingDot, /*isObjC*/ false, isImplicit,
-                  closureDiscriminator) {
+                  hasLeadingDot, /*isObjC*/ false, isImplicit) {
   assert((parsedRoot || parsedPath) &&
          "Key path must have either root or path");
 }
@@ -2278,8 +2275,7 @@ KeyPathExpr::KeyPathExpr(ASTContext &ctx, SourceLoc startLoc,
                          ArrayRef<Component> components, SourceLoc endLoc,
                          bool isObjC, bool isImplicit)
     : KeyPathExpr(startLoc, /*parsedRoot*/ nullptr, /*parsedPath*/ nullptr,
-                  endLoc, /*hasLeadingDot*/ false, isObjC, isImplicit,
-                  /*closure discriminator*/ AbstractClosureExpr::InvalidDiscriminator) {
+                  endLoc, /*hasLeadingDot*/ false, isObjC, isImplicit) {
   assert(!components.empty());
   Components = ctx.AllocateCopy(components);
 }
@@ -2293,11 +2289,9 @@ KeyPathExpr *KeyPathExpr::createParsedPoundKeyPath(
 
 KeyPathExpr *KeyPathExpr::createParsed(ASTContext &ctx, SourceLoc backslashLoc,
                                        Expr *parsedRoot, Expr *parsedPath,
-                                       bool hasLeadingDot,
-                                       unsigned closureDiscriminator) {
+                                       bool hasLeadingDot) {
   return new (ctx) KeyPathExpr(backslashLoc, parsedRoot, parsedPath,
-                               hasLeadingDot, /*isImplicit*/ false,
-                               closureDiscriminator);
+                               hasLeadingDot, /*isImplicit*/ false);
 }
 
 KeyPathExpr *KeyPathExpr::createImplicit(ASTContext &ctx,
@@ -2313,8 +2307,7 @@ KeyPathExpr *KeyPathExpr::createImplicit(ASTContext &ctx,
                                          Expr *parsedRoot, Expr *parsedPath,
                                          bool hasLeadingDot) {
   return new (ctx) KeyPathExpr(backslashLoc, parsedRoot, parsedPath,
-           hasLeadingDot, /*isImplicit*/ true,
-           /*closureDiscriminator*/ AbstractClosureExpr::InvalidDiscriminator);
+           hasLeadingDot, /*isImplicit*/ true);
 }
 
 void

lib/Parse/ParseExpr.cpp

@@ -690,8 +690,7 @@ ParserResult<Expr> Parser::parseExprKeyPath() {
 
   auto *keypath = KeyPathExpr::createParsed(
       Context, backslashLoc, rootResult.getPtrOrNull(),
-      pathResult.getPtrOrNull(), hasLeadingDot,
-      CurLocalContext->claimNextClosureDiscriminator());
+      pathResult.getPtrOrNull(), hasLeadingDot);
   return makeParserResult(parseStatus, keypath);
 }
 

lib/Sema/CSApply.cpp

@@ -4978,24 +4978,14 @@ namespace {
       //     return "{ [$kp$ = \(E)] in $0[keyPath: $kp$] }"
 
       auto &ctx = cs.getASTContext();
-      auto discriminator = E->getClosureDiscriminator();
+      auto discriminator = AutoClosureExpr::InvalidDiscriminator;
 
       FunctionType::ExtInfo closureInfo;
       auto closureTy =
           FunctionType::get({FunctionType::Param(baseTy)}, leafTy, closureInfo);
-
-      ClosureExpr *closure = new (ctx)
-        ClosureExpr(DeclAttributes(),
-                    SourceRange(),
-                    /*captured self*/ nullptr,
-                    /*params*/ nullptr,
-                    SourceLoc(),
-                    SourceLoc(),
-                    SourceLoc(),
-                    SourceLoc(),
-                    /*explicit result type*/ nullptr,
-                    discriminator,
-                    dc);
+      auto closure = new (ctx)
+          AutoClosureExpr(/*set body later*/nullptr, leafTy,
+                          discriminator, dc);
 
       auto param = new (ctx) ParamDecl(
           SourceLoc(),
@@ -5054,14 +5044,10 @@ namespace {
                                  E->getStartLoc(), outerParamRef, E->getEndLoc(),
                                  leafTy, /*implicit=*/true);
       cs.cacheType(application);
-      
-      auto returnStmt = new (ctx) ReturnStmt(SourceLoc(), application);
-      auto bodyStmt = BraceStmt::create(ctx,
-                                SourceLoc(), ASTNode(returnStmt), SourceLoc());
 
       // Finish up the inner closure.
       closure->setParameterList(ParameterList::create(ctx, {param}));
-      closure->setBody(bodyStmt, /*singleExpr*/ true);
+      closure->setBody(application);
       closure->setType(closureTy);
       cs.cacheType(closure);
 

test/SILGen/closure_literal_reabstraction.swift

@@ -16,7 +16,7 @@ func reabstractCaptureListExprArgument() {
 }
 
 // CHECK-LABEL: sil {{.*}} @{{.*}}reabstractKeyPathFunctionArgument
-// CHECK:         [[CLOSURE_FN:%.*]] = function_ref {{.*}}U_
+// CHECK:         [[CLOSURE_FN:%.*]] = function_ref {{.*}}u_
 // CHECK:         [[KP_ARG:%.*]] = copy_value {{.*}} $KeyPath
 // CHECK:         [[CLOSURE:%.*]] = partial_apply {{.*}}[[CLOSURE_FN]]([[KP_ARG]])
 // CHECK:         [[CLOSURE_NE:%.*]] = convert_escape_to_noescape {{.*}} [[CLOSURE]]