[AST] Refactor KeyPathExpr constructors

Now that the CSApply just uses components, we can
better split up the key path constructors to either
accept a set of resolved components, or a parsed
root or path.

Author: hamishknight

include/swift/AST/Expr.h

@@ -5562,24 +5562,43 @@ class KeyPathExpr : public Expr {
 private:
   llvm::MutableArrayRef<Component> Components;
 
-public:
-  /// Create a new #keyPath expression.
-  KeyPathExpr(ASTContext &C,
-              SourceLoc keywordLoc, SourceLoc lParenLoc,
-              ArrayRef<Component> components,
-              SourceLoc rParenLoc,
-              bool isImplicit = false);
+  KeyPathExpr(SourceLoc startLoc, Expr *parsedRoot, Expr *parsedPath,
+              SourceLoc endLoc, bool hasLeadingDot, bool isObjC,
+              bool isImplicit);
 
+  /// Create a key path with unresolved root and path expressions.
   KeyPathExpr(SourceLoc backslashLoc, Expr *parsedRoot, Expr *parsedPath,
-              bool hasLeadingDot, bool isImplicit = false)
-      : Expr(ExprKind::KeyPath, isImplicit), StartLoc(backslashLoc),
-        EndLoc(parsedPath ? parsedPath->getEndLoc() : parsedRoot->getEndLoc()),
-        ParsedRoot(parsedRoot), ParsedPath(parsedPath),
-        HasLeadingDot(hasLeadingDot) {
-    assert((parsedRoot || parsedPath) &&
-           "keypath must have either root or path");
-    Bits.KeyPathExpr.IsObjC = false;
-  }
+              bool hasLeadingDot, bool isImplicit);
+
+  /// Create a key path with components.
+  KeyPathExpr(ASTContext &ctx, SourceLoc startLoc,
+              ArrayRef<Component> components, SourceLoc endLoc, bool isObjC,
+              bool isImplicit);
+
+public:
+  /// Create a new parsed Swift key path expression.
+  static KeyPathExpr *createParsed(ASTContext &ctx, SourceLoc backslashLoc,
+                                   Expr *parsedRoot, Expr *parsedPath,
+                                   bool hasLeadingDot);
+
+  /// Create a new parsed #keyPath expression.
+  static KeyPathExpr *createParsedPoundKeyPath(ASTContext &ctx,
+                                               SourceLoc keywordLoc,
+                                               SourceLoc lParenLoc,
+                                               ArrayRef<Component> components,
+                                               SourceLoc rParenLoc);
+
+  /// Create an implicit Swift key path expression with a set of resolved
+  /// components.
+  static KeyPathExpr *createImplicit(ASTContext &ctx, SourceLoc backslashLoc,
+                                     ArrayRef<Component> components,
+                                     SourceLoc endLoc);
+
+  /// Create an implicit Swift key path expression with a root and path
+  /// expression to be resolved.
+  static KeyPathExpr *createImplicit(ASTContext &ctx, SourceLoc backslashLoc,
+                                     Expr *parsedRoot, Expr *parsedPath,
+                                     bool hasLeadingDot);
 
   SourceLoc getLoc() const { return StartLoc; }
   SourceRange getSourceRange() const { return SourceRange(StartLoc, EndLoc); }

lib/AST/Expr.cpp

@@ -1987,17 +1987,64 @@ OpenedArchetypeType *OpenExistentialExpr::getOpenedArchetype() const {
   return type->castTo<OpenedArchetypeType>();
 }
 
-KeyPathExpr::KeyPathExpr(ASTContext &C, SourceLoc keywordLoc,
-                         SourceLoc lParenLoc, ArrayRef<Component> components,
-                         SourceLoc rParenLoc, bool isImplicit)
-    : Expr(ExprKind::KeyPath, isImplicit), StartLoc(keywordLoc),
-      LParenLoc(lParenLoc), EndLoc(rParenLoc),
-      Components(C.AllocateUninitialized<Component>(components.size())) {
-  // Copy components into the AST context.
-  std::uninitialized_copy(components.begin(), components.end(),
-                          Components.begin());
-
-  Bits.KeyPathExpr.IsObjC = true;
+KeyPathExpr::KeyPathExpr(SourceLoc startLoc, Expr *parsedRoot,
+                         Expr *parsedPath, SourceLoc endLoc, bool hasLeadingDot,
+                         bool isObjC, bool isImplicit)
+    : Expr(ExprKind::KeyPath, isImplicit), StartLoc(startLoc), EndLoc(endLoc),
+      ParsedRoot(parsedRoot), ParsedPath(parsedPath),
+      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)
+    : KeyPathExpr(backslashLoc, parsedRoot, parsedPath,
+                  parsedPath ? parsedPath->getEndLoc()
+                             : parsedRoot->getEndLoc(),
+                  hasLeadingDot, /*isObjC*/ false, isImplicit) {
+  assert((parsedRoot || parsedPath) &&
+         "Key path must have either root or path");
+}
+
+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) {
+  assert(!components.empty());
+  Components = ctx.AllocateCopy(components);
+}
+
+KeyPathExpr *KeyPathExpr::createParsedPoundKeyPath(
+    ASTContext &ctx, SourceLoc keywordLoc, SourceLoc lParenLoc,
+    ArrayRef<Component> components, SourceLoc rParenLoc) {
+  return new (ctx) KeyPathExpr(ctx, keywordLoc, components, rParenLoc,
+                               /*isObjC*/ true, /*isImplicit*/ false);
+}
+
+KeyPathExpr *KeyPathExpr::createParsed(ASTContext &ctx, SourceLoc backslashLoc,
+                                       Expr *parsedRoot, Expr *parsedPath,
+                                       bool hasLeadingDot) {
+  return new (ctx) KeyPathExpr(backslashLoc, parsedRoot, parsedPath,
+                               hasLeadingDot, /*isImplicit*/ false);
+}
+
+KeyPathExpr *KeyPathExpr::createImplicit(ASTContext &ctx,
+                                         SourceLoc backslashLoc,
+                                         ArrayRef<Component> components,
+                                         SourceLoc endLoc) {
+  return new (ctx) KeyPathExpr(ctx, backslashLoc, components, endLoc,
+                               /*isObjC*/ false, /*isImplicit*/ true);
+}
+
+KeyPathExpr *KeyPathExpr::createImplicit(ASTContext &ctx,
+                                         SourceLoc backslashLoc,
+                                         Expr *parsedRoot, Expr *parsedPath,
+                                         bool hasLeadingDot) {
+  return new (ctx) KeyPathExpr(backslashLoc, parsedRoot, parsedPath,
+                               hasLeadingDot, /*isImplicit*/ true);
 }
 
 void

lib/Parse/ParseExpr.cpp

@@ -656,17 +656,17 @@ ParserResult<Expr> Parser::parseExprKeyPath() {
     // Add the code completion expression to the path result.
     CodeCompletionExpr *CC = new (Context)
         CodeCompletionExpr(pathResult.getPtrOrNull(), Tok.getLoc());
-    auto keypath = new (Context)
-        KeyPathExpr(backslashLoc, rootResult.getPtrOrNull(), CC, hasLeadingDot);
+    auto *keypath = KeyPathExpr::createParsed(
+        Context, backslashLoc, rootResult.getPtrOrNull(), CC, hasLeadingDot);
     if (CodeCompletion)
       CodeCompletion->completeExprKeyPath(keypath, DotLoc);
     consumeToken(tok::code_complete);
     return makeParserCodeCompletionResult(keypath);
   }
 
-  auto keypath =
-      new (Context) KeyPathExpr(backslashLoc, rootResult.getPtrOrNull(),
-                                pathResult.getPtrOrNull(), hasLeadingDot);
+  auto *keypath = KeyPathExpr::createParsed(
+      Context, backslashLoc, rootResult.getPtrOrNull(),
+      pathResult.getPtrOrNull(), hasLeadingDot);
   return makeParserResult(parseStatus, keypath);
 }
 
@@ -690,8 +690,8 @@ ParserResult<Expr> Parser::parseExprKeyPathObjC() {
   auto handleCodeCompletion = [&](SourceLoc DotLoc) -> ParserResult<Expr> {
     KeyPathExpr *expr = nullptr;
     if (!components.empty()) {
-      expr = new (Context)
-          KeyPathExpr(Context, keywordLoc, lParenLoc, components, Tok.getLoc());
+      expr = KeyPathExpr::createParsedPoundKeyPath(
+          Context, keywordLoc, lParenLoc, components, Tok.getLoc());
     }
 
     if (CodeCompletion)
@@ -764,7 +764,7 @@ ParserResult<Expr> Parser::parseExprKeyPathObjC() {
   }
 
   // We're done: create the key-path expression.
-  return makeParserResult<Expr>(new (Context) KeyPathExpr(
+  return makeParserResult<Expr>(KeyPathExpr::createParsedPoundKeyPath(
       Context, keywordLoc, lParenLoc, components, rParenLoc));
 }
 

lib/Sema/CSApply.cpp

@@ -2223,16 +2223,8 @@ namespace {
       auto *anchor = getAsExpr(memberLoc->getAnchor());
 
       auto makeKeyPath = [&](ArrayRef<Component> components) -> Expr * {
-        // Create a new implicit key path. We pass in anchor as the parsed path
-        // to set the end loc, but don't want to keep it as the parsed path,
-        // so clear it out after.
-        auto *kp = new (ctx) KeyPathExpr(/*backslashLoc*/ dotLoc,
-                                         /*parsedRoot*/ nullptr,
-                                         /*parsedPath*/ anchor,
-                                         /*hasLeadingDot*/ false,
-                                         /*isImplicit*/ true);
-        kp->setParsedPath(nullptr);
-        kp->setComponents(ctx, components);
+        auto *kp = KeyPathExpr::createImplicit(ctx, /*backslashLoc*/ dotLoc,
+                                               components, anchor->getEndLoc());
         kp->setType(keyPathTy);
         cs.cacheExprTypes(kp);
 

lib/Sema/TypeCheckStorage.cpp

@@ -918,15 +918,17 @@ static Expr *buildStorageReference(AccessorDecl *accessor,
     Expr *propertyKeyPath = new (ctx) KeyPathDotExpr(SourceLoc());
     propertyKeyPath = UnresolvedDotExpr::createImplicit(ctx, propertyKeyPath,
         enclosingSelfAccess->accessedProperty->getName());
-    propertyKeyPath = new (ctx) KeyPathExpr(
-        SourceLoc(), nullptr, propertyKeyPath, /*hasLeadingDot=*/true);
+    propertyKeyPath = KeyPathExpr::createImplicit(
+        ctx, /*backslashLoc*/ SourceLoc(), /*parsedRoot*/ nullptr,
+        propertyKeyPath, /*hasLeadingDot*/ true);
 
     // Key path referring to the backing storage property.
     Expr *storageKeyPath = new (ctx) KeyPathDotExpr(SourceLoc());
     storageKeyPath = UnresolvedDotExpr::createImplicit(ctx, storageKeyPath,
                                                        storage->getName());
-    storageKeyPath = new (ctx) KeyPathExpr(SourceLoc(), nullptr, storageKeyPath,
-                                           /*hasLeadingDot=*/true);
+    storageKeyPath = KeyPathExpr::createImplicit(
+        ctx, /*backslashLoc*/ SourceLoc(), /*parsedRoot*/ nullptr,
+        storageKeyPath, /*hasLeadingDot*/ true);
     Expr *args[3] = {selfDRE, propertyKeyPath, storageKeyPath};
 
     auto *subscriptDecl = enclosingSelfAccess->subscript;