[Parser] Replace postfix ellipsis with a prefix 'repeat' keyword for pack

expansion expressions, and create them in the parser instead of during
pre-checking in the constraint system.

Author: hborla

include/swift/AST/Expr.h

@@ -3595,21 +3595,22 @@ class PackElementExpr final : public Expr {
 /// call argument lists, the elements of a tuple value, and the source
 /// of a for-in loop.
 class PackExpansionExpr final : public Expr {
+  SourceLoc RepeatLoc;
   Expr *PatternExpr;
-  SourceLoc DotsLoc;
   GenericEnvironment *Environment;
 
-  PackExpansionExpr(Expr *patternExpr,
-                    SourceLoc dotsLoc,
+  PackExpansionExpr(SourceLoc repeatLoc,
+                    Expr *patternExpr,
                     GenericEnvironment *environment,
                     bool implicit, Type type)
     : Expr(ExprKind::PackExpansion, implicit, type),
-      PatternExpr(patternExpr), DotsLoc(dotsLoc), Environment(environment) {}
+      RepeatLoc(repeatLoc), PatternExpr(patternExpr),
+      Environment(environment) {}
 
 public:
   static PackExpansionExpr *create(ASTContext &ctx,
+                                   SourceLoc repeatLoc,
                                    Expr *patternExpr,
-                                   SourceLoc dotsLoc,
                                    GenericEnvironment *environment,
                                    bool implicit = false,
                                    Type type = Type());
@@ -3626,12 +3627,16 @@ class PackExpansionExpr final : public Expr {
     return Environment;
   }
 
+  void setGenericEnvironment(GenericEnvironment *env) {
+    Environment = env;
+  }
+
   SourceLoc getStartLoc() const {
-    return PatternExpr->getStartLoc();
+    return RepeatLoc;
   }
 
   SourceLoc getEndLoc() const {
-    return DotsLoc;
+    return PatternExpr->getEndLoc();
   }
 
   static bool classof(const Expr *E) {

lib/AST/Expr.cpp

@@ -1250,10 +1250,10 @@ VarargExpansionExpr *VarargExpansionExpr::createArrayExpansion(ASTContext &ctx,
 }
 
 PackExpansionExpr *
-PackExpansionExpr::create(ASTContext &ctx, Expr *patternExpr,
-                          SourceLoc dotsLoc, GenericEnvironment *environment,
+PackExpansionExpr::create(ASTContext &ctx, SourceLoc repeatLoc,
+                          Expr *patternExpr, GenericEnvironment *environment,
                           bool implicit, Type type) {
-  return new (ctx) PackExpansionExpr(patternExpr, dotsLoc, environment,
+  return new (ctx) PackExpansionExpr(repeatLoc, patternExpr, environment,
                                      implicit, type);
 }
 

lib/Parse/ParseExpr.cpp

@@ -511,6 +511,19 @@ ParserResult<Expr> Parser::parseExprUnary(Diag<> Message, bool isExprBasic) {
   // First check to see if we have the start of a regex literal `/.../`.
   tryLexRegexLiteral(/*forUnappliedOperator*/ false);
 
+  // 'repeat' as an expression prefix is a pack expansion expression.
+  if (Tok.is(tok::kw_repeat)) {
+    SourceLoc repeatLoc = consumeToken();
+    auto patternExpr = parseExpr(Message);
+    if (patternExpr.isNull())
+      return patternExpr;
+
+    auto *expansion =
+        PackExpansionExpr::create(Context, repeatLoc, patternExpr.get(),
+                                  /*genericEnv*/ nullptr);
+    return makeParserResult(expansion);
+  }
+
   switch (Tok.getKind()) {
   default:
     // If the next token is not an operator, just parse this as expr-postfix.

lib/Parse/ParsePattern.cpp

@@ -163,7 +163,8 @@ bool Parser::startsParameterName(bool isClosure) {
     if (!Tok.isContextualKeyword("isolated") &&
         !Tok.isContextualKeyword("some") &&
         !Tok.isContextualKeyword("any") &&
-        !Tok.isContextualKeyword("each"))
+        !Tok.isContextualKeyword("each") &&
+        !Tok.is(tok::kw_repeat))
       return true;
 
     // "isolated" can be an argument label, but it's also a contextual keyword,

lib/Parse/ParseStmt.cpp

@@ -46,7 +46,6 @@ bool Parser::isStartOfStmt() {
   case tok::kw_guard:
   case tok::kw_while:
   case tok::kw_do:
-  case tok::kw_repeat:
   case tok::kw_for:
   case tok::kw_break:
   case tok::kw_continue:
@@ -62,6 +61,13 @@ bool Parser::isStartOfStmt() {
   case tok::pound_sourceLocation:
     return true;
 
+  case tok::kw_repeat:
+    // 'repeat' followed by anything other than a brace stmt
+    // is a pack expansion expression.
+    // FIXME: 'repeat' followed by '{' could be a pack expansion
+    // with a closure pattern.
+    return peekToken().is(tok::l_brace);
+
   case tok::pound_line:
     // #line at the start of a line is a directive, when within, it is an expr.
     return Tok.isAtStartOfLine();

lib/Sema/PreCheckExpr.cpp

@@ -399,57 +399,19 @@ static BinaryExpr *getCompositionExpr(Expr *expr) {
   return nullptr;
 }
 
-static Expr *getPackExpansion(DeclContext *dc, Expr *expr, SourceLoc opLoc) {
-  // FIXME: The parser should create PackExpansionExprs directly, pack
-  // elements should be discovered via PackExpansionExpr::getExpandedPacks,
-  // and the generic environment should be created lazily when solving
+static void
+setPackElementEnvironment(DeclContext *dc, PackExpansionExpr *expansion) {
+  // FIXME: The generic environment should be created lazily when solving
   // PackElementOf constraints.
-  struct PackReferenceFinder : public ASTWalker {
-    DeclContext *dc;
-    GenericEnvironment *environment;
-
-    PackReferenceFinder(DeclContext *dc)
-      : dc(dc), environment(nullptr) {}
-
-    void createElementEnvironment() {
-      if (!environment) {
-        auto &ctx = dc->getASTContext();
-        auto sig = ctx.getOpenedElementSignature(
-            dc->getGenericSignatureOfContext().getCanonicalSignature());
-        auto *contextEnv = dc->getGenericEnvironmentOfContext();
-        auto contextSubs = contextEnv->getForwardingSubstitutionMap();
-        environment =
-            GenericEnvironment::forOpenedElement(sig, UUID::fromTime(),
-                                                 contextSubs);
-      }
-    }
-
-    virtual PreWalkResult<Expr *> walkToExprPre(Expr *E) override {
-      auto *packElement = dyn_cast<PackElementExpr>(E);
-      if (!packElement)
-        return Action::Continue(E);
-
-      createElementEnvironment();
-      return Action::Continue(packElement);
-    }
-
-    virtual PreWalkAction walkToTypeReprPre(TypeRepr *T) override {
-      if (isa<PackReferenceTypeRepr>(T)) {
-        createElementEnvironment();
-      }
-
-      return Action::Continue();
-    }
-  } packReferenceFinder(dc);
-
-  auto *pattern = expr->walk(packReferenceFinder);
-
-  if (packReferenceFinder.environment != nullptr) {
-    return PackExpansionExpr::create(dc->getASTContext(), pattern,
-                                     opLoc, packReferenceFinder.environment);
-  }
-
-  return nullptr;
+  auto &ctx = dc->getASTContext();
+  auto sig = ctx.getOpenedElementSignature(
+      dc->getGenericSignatureOfContext().getCanonicalSignature());
+  auto *contextEnv = dc->getGenericEnvironmentOfContext();
+  auto contextSubs = contextEnv->getForwardingSubstitutionMap();
+  auto *environment =
+      GenericEnvironment::forOpenedElement(sig, UUID::fromTime(),
+                                           contextSubs);
+  expansion->setGenericEnvironment(environment);
 }
 
 /// Bind an UnresolvedDeclRefExpr by performing name lookup and
@@ -1202,12 +1164,10 @@ namespace {
         return Action::Continue(result);
       }
 
-      // Rewrite postfix unary '...' expressions containing pack
-      // references to PackExpansionExpr.
-      if (auto *operand = isPostfixEllipsisOperator(expr)) {
-        if (auto *expansion = getPackExpansion(DC, operand, expr->getLoc())) {
-          return Action::Continue(expansion);
-        }
+      // Set the generic environment of a PackExpansionExpr.
+      if (auto *expansion = dyn_cast<PackExpansionExpr>(expr)) {
+        setPackElementEnvironment(DC, expansion);
+        return Action::Continue(expansion);
       }
 
       // Type check the type parameters in an UnresolvedSpecializeExpr.

test/Constraints/pack-expansion-expressions.swift

@@ -3,44 +3,44 @@
 // REQUIRES: asserts
 
 func tuplify<T...>(_ t: T...) -> (T...) {
-  return ((each t)...)
+  return (repeat each t)
 }
 
 func prepend<First, Rest...>(value: First, to rest: Rest...) -> (First, Rest...) {
-  return (value, (each rest)...)
+  return (value, repeat each rest)
 }
 
 func concatenate<T..., U...>(_ first: T..., with second: U...) -> (T..., U...) {
-  return ((each first)..., (each second)...)
+  return (repeat each first, repeat each second)
 }
 
 func zip<T..., U...>(_ first: T..., with second: U...) -> ((T, U)...) {
-  return ((each first, each second)...)
+  return (repeat (each first, each second))
 }
 
 func forward<U...>(_ u: U...) -> (U...) {
-  return tuplify((each u)...)
+  return tuplify(repeat each u)
 }
 
 func forwardAndMap<U..., V...>(us u: U..., vs v: V...) -> ([(U, V)]...) {
-  return tuplify([(each u, each v)]...)
+  return tuplify(repeat [(each u, each v)])
 }
 
 func variadicMap<T..., Result...>(_ t: T..., transform: ((T) -> Result)...) -> (Result...) {
-  return ((each transform)(each t)...)
+  return (repeat (each transform)(each t))
 }
 
 func coerceExpansion<T...>(_ value: T...) {
   func promoteToOptional<Wrapped...>(_: Wrapped?...) {}
 
-  promoteToOptional((each value)...)
+  promoteToOptional(repeat each value)
 }
 
 func localValuePack<T...>(_ t: T...) -> (T..., T...) {
-  let local = (each t)...
-  let localAnnotated: T... = (each t)...
+  let local = repeat each t
+  let localAnnotated: T... = repeat each t
 
-  return ((each local)..., (each localAnnotated)...)
+  return (repeat each local, repeat each localAnnotated)
 }
 
 protocol P {
@@ -52,21 +52,21 @@ protocol P {
 }
 
 func outerArchetype<T..., U>(t: T..., u: U) where T: P {
-  let _: (T.A, U)... = ((each t).value, u)...
+  let _: (T.A, U)... = repeat ((each t).value, u)
 }
 
 func sameElement<T..., U>(t: T..., u: U) where T: P, T == U {
-  let _: T... = (each t).f(u)...
+  let _: T... = repeat (each t).f(u)
 }
 
 func forEachEach<C..., U>(c: C..., function: (U) -> Void)
     where C: Collection, C.Element == U {
-  _ = (each c).forEach(function)...
+  _ = repeat (each c).forEach(function)
 }
 
 func typeReprPacks<T...>(_ t: T...) where T: ExpressibleByIntegerLiteral {
-  _ = Array<each T>()...
-  _ = (1 as each T)...
+  _ = repeat Array<each T>()
+  _ = repeat 1 as each T
 
   _ = Array<each T>() // expected-error {{pack reference 'T' can only appear in pack expansion or generic requirement}}
   _ = 1 as each T // expected-error {{pack reference 'T' can only appear in pack expansion or generic requirement}}

test/Constraints/variadic_generic_constraints.swift

@@ -69,10 +69,10 @@ let _ = zip(t: 1, 2, 3, u: "hi", "hello", "greetings")  // ok
 let _ = zip(t: 1, u: "hi", "hello", "greetings")  // expected-error {{type of expression is ambiguous without more context}}
 
 func goodCallToZip<T..., U...>(t: T..., u: U...) where ((T, U)...): Any {
-  _ = zip(t: (each t)..., u: (each u)...)
+  _ = zip(t: repeat each t, u: repeat each u)
 }
 
 func badCallToZip<T..., U...>(t: T..., u: U...) {
-  _ = zip(t: (each t)..., u: (each u)...)
+  _ = zip(t: repeat each t, u: repeat each u)
   // expected-error@-1 {{global function 'zip(t:u:)' requires the type packs 'T' and 'U' have the same shape}}
 }

test/Constraints/variadic_generic_functions.swift

@@ -28,7 +28,7 @@ func invalidPacks() {
 
 func call() {
   func multipleParameters<T...>(xs: T..., ys: T...) -> (T...) {
-    return ((each xs)...)
+    return (repeat each xs)
   }
   multipleParameters()
 
@@ -37,7 +37,7 @@ func call() {
   multipleParameters(xs: "", 5.0, ys: 5.0, "") // expected-error {{type of expression is ambiguous without more context}}
 
   func multipleSequences<T..., U...>(xs: T..., ys: U...) -> (T...) {
-    return ((each ys)...)
+    return (repeat each ys)
     // expected-error@-1 {{cannot convert return expression of type '(U...)' to return type '(T...)'}}
   }