sema: accept an array literal of enum case references as compile-time constant

Author: nkcsgexi

include/swift/AST/Expr.h

@@ -496,8 +496,8 @@ class alignas(8) Expr : public ASTAllocated<Expr> {
     return getSemanticsProvidingExpr()->getKind() == ExprKind::InOut;
   }
 
-  bool printConstExprValue(llvm::raw_ostream *OS) const;
-  bool isSemanticallyConstExpr() const;
+  bool printConstExprValue(llvm::raw_ostream *OS, llvm::function_ref<bool(Expr*)> additionalCheck) const;
+  bool isSemanticallyConstExpr(llvm::function_ref<bool(Expr*)> additionalCheck = nullptr) const;
 
   /// Returns false if this expression needs to be wrapped in parens when
   /// used inside of a any postfix expression, true otherwise.

lib/APIDigester/ModuleAnalyzerNodes.cpp

@@ -2265,7 +2265,7 @@ class ConstExtractor: public ASTWalker {
   void record(Expr *E, Expr *ValueProvider, StringRef ReferecedD = "") {
     std::string content;
     llvm::raw_string_ostream os(content);
-    ValueProvider->printConstExprValue(&os);
+    ValueProvider->printConstExprValue(&os, nullptr);
     assert(!content.empty());
     auto buffered = SCtx.buffer(content);
     switch(ValueProvider->getKind()) {

lib/AST/Expr.cpp

@@ -198,7 +198,8 @@ Expr *Expr::getSemanticsProvidingExpr() {
   return this;
 }
 
-bool Expr::printConstExprValue(llvm::raw_ostream *OS) const {
+bool Expr::printConstExprValue(llvm::raw_ostream *OS,
+                        llvm::function_ref<bool(Expr*)> additionalCheck) const {
   auto print = [&](StringRef text) {
     if (OS) {
       *OS << text;
@@ -242,7 +243,7 @@ bool Expr::printConstExprValue(llvm::raw_ostream *OS) const {
     for (unsigned N = CE->getNumElements(), I = 0; I != N; I ++) {
       auto Ele = CE->getElement(I);
       auto needComma = I + 1 != N;
-      if (!Ele->printConstExprValue(OS)) {
+      if (!Ele->printConstExprValue(OS, additionalCheck)) {
         return false;
       }
       if (needComma)
@@ -257,7 +258,7 @@ bool Expr::printConstExprValue(llvm::raw_ostream *OS) const {
     for (unsigned N = TE->getNumElements(), I = 0; I != N; I ++) {
       auto Ele = TE->getElement(I);
       auto needComma = I + 1 != N;
-      if (!Ele->printConstExprValue(OS)) {
+      if (!Ele->printConstExprValue(OS, additionalCheck)) {
         return false;
       }
       if (needComma)
@@ -267,12 +268,13 @@ bool Expr::printConstExprValue(llvm::raw_ostream *OS) const {
     return true;
   }
   default:
-    return false;
+    return additionalCheck && additionalCheck(const_cast<Expr*>(this));
   }
 }
 
-bool Expr::isSemanticallyConstExpr() const {
-  return printConstExprValue(nullptr);
+bool Expr::isSemanticallyConstExpr(
+    llvm::function_ref<bool(Expr*)> additionalCheck) const {
+  return printConstExprValue(nullptr, additionalCheck);
 }
 
 Expr *Expr::getValueProvidingExpr() {

lib/Sema/CSFix.cpp

@@ -1143,14 +1143,26 @@ NotCompileTimeConst::create(ConstraintSystem &cs, Type paramTy,
 
 bool NotCompileTimeConst::diagnose(const Solution &solution, bool asNote) const {
   auto *locator = getLocator();
-  // Referencing an enum element directly is considered a compile-time literal.
-  if (auto *d = solution.resolveLocatorToDecl(locator).getDecl()) {
-    if (isa<EnumElementDecl>(d)) {
-      if (!d->hasParameterList()) {
-        return true;
+  if (auto *E = getAsExpr(locator->getAnchor())) {
+    auto isAccepted = E->isSemanticallyConstExpr([&](Expr *E) {
+      if (auto *UMC = dyn_cast<UnresolvedMemberChainResultExpr>(E)) {
+        E = UMC->getSubExpr();
       }
-    }
+      auto locator = solution.getConstraintSystem().getConstraintLocator(E);
+      // Referencing an enum element directly is considered a compile-time literal.
+      if (auto *d = solution.resolveLocatorToDecl(locator).getDecl()) {
+        if (isa<EnumElementDecl>(d)) {
+          if (!d->hasParameterList()) {
+            return true;
+          }
+        }
+      }
+      return false;
+    });
+    if (isAccepted)
+      return true;
   }
+
   NotCompileTimeConstFailure failure(solution, locator);
   return failure.diagnose(asNote);
 }

lib/Sema/CSSimplify.cpp

@@ -1961,14 +1961,8 @@ static ConstraintSystem::TypeMatchResult matchCallArguments(
         auto *locator = cs.getConstraintLocator(loc);
         SourceRange range;
         // simplify locator so the anchor is the exact argument.
-        locator = simplifyLocator(cs, locator, range);
-        if (locator->getPath().empty() &&
-            locator->getAnchor().isExpr(ExprKind::UnresolvedMemberChainResult)) {
-          locator =
-            cs.getConstraintLocator(cast<UnresolvedMemberChainResultExpr>(
-              locator->getAnchor().get<Expr*>())->getSubExpr());
-        }
-        cs.recordFix(NotCompileTimeConst::create(cs, paramTy, locator));
+        cs.recordFix(NotCompileTimeConst::create(cs, paramTy,
+          simplifyLocator(cs, locator, range)));
       }
 
       cs.addConstraint(

test/Sema/const_enum_elements_array.swift

@@ -0,0 +1,20 @@
+// RUN: %target-typecheck-verify-swift
+
+enum E {
+  case a
+  case b
+
+  var z: E { .b }
+}
+
+func getAE() -> E { return .a }
+
+func test_without_const(_ : [E]) {}
+func testArr(_ : _const [E]) {}
+
+testArr([])
+testArr([.a])
+testArr([.a, .b])
+
+testArr([getAE()]) // expected-error {{expect a compile-time constant literal}}
+testArr([.a, .b, .a.z]) // expected-error {{expect a compile-time constant literal}}