Push subscript label allocation for KeyPathComponents into the root constructor.

include/swift/AST/Expr.h

@@ -4612,24 +4612,24 @@ class KeyPathExpr : public Expr {
     Type ComponentType;
     SourceLoc Loc;
 
-    explicit Component(DeclNameOrRef decl,
+    explicit Component(ASTContext *ctxForCopyingLabels,
+                       DeclNameOrRef decl,
                        Expr *indexExpr,
                        ArrayRef<Identifier> subscriptLabels,
                        Kind kind,
                        Type type,
-                       SourceLoc loc)
-      : Decl(decl), SubscriptIndexExprAndKind(indexExpr, kind),
-        SubscriptLabels(subscriptLabels),
-        ComponentType(type), Loc(loc)
-    {}
+                       SourceLoc loc);
 
   public:
-    Component() : Component({}, nullptr, {}, Kind::Invalid, Type(), SourceLoc()) {}
+    Component()
+      : Component(nullptr, {}, nullptr, {}, Kind::Invalid, Type(), SourceLoc())
+    {}
 
     /// Create an unresolved component for a property.
     static Component forUnresolvedProperty(DeclName UnresolvedName,
                                            SourceLoc Loc) {
-      return Component(UnresolvedName, nullptr, {},
+      return Component(nullptr,
+                       UnresolvedName, nullptr, {},
                        Kind::UnresolvedProperty,
                        Type(),
                        Loc);
@@ -4648,25 +4648,28 @@ class KeyPathExpr : public Expr {
     ///
     /// You shouldn't add new uses of this overload; use the one that takes a
     /// list of index arguments.
-    static Component forUnresolvedSubscriptWithPrebuiltIndexExpr(Expr *index,
+    static Component forUnresolvedSubscriptWithPrebuiltIndexExpr(
+                                         ASTContext &context,
+                                         Expr *index,
                                          ArrayRef<Identifier> subscriptLabels,
                                          SourceLoc loc) {
 
-      return Component({}, index, subscriptLabels, Kind::UnresolvedSubscript,
+      return Component(&context,
+                       {}, index, subscriptLabels, Kind::UnresolvedSubscript,
                        Type(), loc);
     }
 
     /// Create an unresolved optional force `!` component.
     static Component forUnresolvedOptionalForce(SourceLoc BangLoc) {
-      return Component({}, nullptr, {},
+      return Component(nullptr, {}, nullptr, {},
                        Kind::OptionalForce,
                        Type(),
                        BangLoc);
     }
 
     /// Create an unresolved optional chain `?` component.
     static Component forUnresolvedOptionalChain(SourceLoc QuestionLoc) {
-      return Component({}, nullptr, {},
+      return Component(nullptr, {}, nullptr, {},
                        Kind::OptionalChain,
                        Type(),
                        QuestionLoc);
@@ -4676,7 +4679,7 @@ class KeyPathExpr : public Expr {
     static Component forProperty(ConcreteDeclRef property,
                                  Type propertyType,
                                  SourceLoc loc) {
-      return Component(property, nullptr, {},
+      return Component(nullptr, property, nullptr, {},
                        Kind::Property,
                        propertyType,
                        loc);
@@ -4699,21 +4702,19 @@ class KeyPathExpr : public Expr {
     /// list of index arguments.
     static Component forSubscriptWithPrebuiltIndexExpr(
        ConcreteDeclRef subscript, Expr *index, ArrayRef<Identifier> labels,
-       Type elementType, SourceLoc loc) {
-      return Component(subscript, index, {}, Kind::Subscript, elementType, loc);
-    }
+       Type elementType, SourceLoc loc);
 
     /// Create an optional-forcing `!` component.
     static Component forOptionalForce(Type forcedType, SourceLoc bangLoc) {
-      return Component({}, nullptr, {},
+      return Component(nullptr, {}, nullptr, {},
                        Kind::OptionalForce, forcedType,
                        bangLoc);
     }
 
     /// Create an optional-chaining `?` component.
     static Component forOptionalChain(Type unwrappedType,
                                       SourceLoc questionLoc) {
-      return Component({}, nullptr, {},
+      return Component(nullptr, {}, nullptr, {},
                        Kind::OptionalChain, unwrappedType,
                        questionLoc);
     }
@@ -4722,7 +4723,7 @@ class KeyPathExpr : public Expr {
     /// syntax but may appear when the non-optional result of an optional chain
     /// is implicitly wrapped.
     static Component forOptionalWrap(Type wrappedType) {
-      return Component({}, nullptr, {},
+      return Component(nullptr, {}, nullptr, {},
                        Kind::OptionalWrap, wrappedType,
                        SourceLoc());
     }

lib/AST/ASTWalker.cpp

@@ -957,6 +957,7 @@ class Traversal : public ASTVisitor<Traversal, Expr*, Stmt*,
                 component.getLoc())
             : KeyPathExpr::Component
                          ::forUnresolvedSubscriptWithPrebuiltIndexExpr(
+                E->getType()->getASTContext(),
                 newIndex, component.getSubscriptLabels(), component.getLoc());
         }
         break;

lib/AST/Expr.cpp

@@ -2133,7 +2133,7 @@ KeyPathExpr::Component::forSubscript(ASTContext &ctx,
                                    indexArgLabelsScratch,
                                    indexArgLabelLocsScratch);
   return forSubscriptWithPrebuiltIndexExpr(subscript, index,
-                                           ctx.AllocateCopy(indexArgLabels),
+                                           indexArgLabels,
                                            elementType,
                                            lSquareLoc);
 }
@@ -2153,7 +2153,29 @@ KeyPathExpr::Component::forUnresolvedSubscript(ASTContext &ctx,
                                    trailingClosure, /*implicit*/ false,
                                    indexArgLabelsScratch,
                                    indexArgLabelLocsScratch);
-  return forUnresolvedSubscriptWithPrebuiltIndexExpr(index,
-                                               ctx.AllocateCopy(indexArgLabels),
+  return forUnresolvedSubscriptWithPrebuiltIndexExpr(ctx, index,
+                                               indexArgLabels,
                                                lSquareLoc);
 }
+
+KeyPathExpr::Component::Component(ASTContext *ctxForCopyingLabels,
+                     DeclNameOrRef decl,
+                     Expr *indexExpr,
+                     ArrayRef<Identifier> subscriptLabels,
+                     Kind kind,
+                     Type type,
+                     SourceLoc loc)
+    : Decl(decl), SubscriptIndexExprAndKind(indexExpr, kind),
+      SubscriptLabels(subscriptLabels.empty()
+                       ? subscriptLabels
+                       : ctxForCopyingLabels->AllocateCopy(subscriptLabels)),
+      ComponentType(type), Loc(loc)
+  {}
+
+KeyPathExpr::Component
+KeyPathExpr::Component::forSubscriptWithPrebuiltIndexExpr(
+       ConcreteDeclRef subscript, Expr *index, ArrayRef<Identifier> labels,
+       Type elementType, SourceLoc loc) {
+  return Component(&elementType->getASTContext(),
+                   subscript, index, {}, Kind::Subscript, elementType, loc);
+}

lib/Sema/TypeCheckConstraints.cpp

@@ -1465,6 +1465,7 @@ void PreCheckExpression::resolveKeyPathExpr(KeyPathExpr *KPE) {
         // .[0] or just plain [0]
         components.push_back(
             KeyPathExpr::Component::forUnresolvedSubscriptWithPrebuiltIndexExpr(
+                TC.Context,
                 SE->getIndex(), SE->getArgumentLabels(), SE->getLoc()));
 
         expr = SE->getBase();