Merge pull request #27364 from CodaFi/raw-smackdown

Clean up enum raw expression validation a bit

Author: CodaFi

include/swift/AST/Decl.h

@@ -6329,9 +6329,7 @@ class EnumElementDecl : public DeclContext, public ValueDecl {
 
   /// The raw value literal for the enum element, or null.
   LiteralExpr *RawValueExpr;
-  /// The type-checked raw value expression.
-  Expr *TypeCheckedRawValueExpr = nullptr;
-  
+
 public:
   EnumElementDecl(SourceLoc IdentifierLoc, DeclName Name,
                   ParameterList *Params,
@@ -6364,13 +6362,6 @@ class EnumElementDecl : public DeclContext, public ValueDecl {
   bool hasRawValueExpr() const { return RawValueExpr; }
   LiteralExpr *getRawValueExpr() const { return RawValueExpr; }
   void setRawValueExpr(LiteralExpr *e) { RawValueExpr = e; }
-  
-  Expr *getTypeCheckedRawValueExpr() const {
-    return TypeCheckedRawValueExpr;
-  }
-  void setTypeCheckedRawValueExpr(Expr *e) {
-    TypeCheckedRawValueExpr = e;
-  }
 
   /// Return the containing EnumDecl.
   EnumDecl *getParentEnum() const {

lib/AST/ASTWalker.cpp

@@ -415,15 +415,7 @@ class Traversal : public ASTVisitor<Traversal, Expr*, Stmt*,
       visit(PL);
     }
 
-    // The getRawValueExpr should remain the untouched original LiteralExpr for
-    // serialization and validation purposes. We only traverse the type-checked
-    // form, unless we haven't populated it yet.
-    if (auto *rawValueExpr = ED->getTypeCheckedRawValueExpr()) {
-      if (auto newRawValueExpr = doIt(rawValueExpr))
-        ED->setTypeCheckedRawValueExpr(newRawValueExpr);
-      else
-        return true;
-    } else if (auto *rawLiteralExpr = ED->getRawValueExpr()) {
+    if (auto *rawLiteralExpr = ED->getRawValueExpr()) {
       Expr *newRawExpr = doIt(rawLiteralExpr);
       if (auto newRawLiteralExpr = dyn_cast<LiteralExpr>(newRawExpr))
         ED->setRawValueExpr(newRawLiteralExpr);

lib/Sema/DerivedConformanceRawRepresentable.cpp

@@ -83,14 +83,6 @@ deriveBodyRawRepresentable_raw(AbstractFunctionDecl *toRawDecl, void *) {
   assert(rawTy);
   rawTy = toRawDecl->mapTypeIntoContext(rawTy);
 
-#ifndef NDEBUG
-  for (auto elt : enumDecl->getAllElements()) {
-    assert(elt->getTypeCheckedRawValueExpr() &&
-           "Enum element has no literal - missing a call to checkEnumRawValues()");
-    assert(elt->getTypeCheckedRawValueExpr()->getType()->isEqual(rawTy));
-  }
-#endif
-
   if (enumDecl->isObjC()) {
     // Special case: ObjC enums are represented by their raw value, so just use
     // a bitcast.
@@ -297,14 +289,6 @@ deriveBodyRawRepresentable_init(AbstractFunctionDecl *initDecl, void *) {
   assert(rawTy);
   rawTy = initDecl->mapTypeIntoContext(rawTy);
 
-#ifndef NDEBUG
-  for (auto elt : enumDecl->getAllElements()) {
-    assert(elt->getTypeCheckedRawValueExpr() &&
-           "Enum element has no literal - missing a call to checkEnumRawValues()");
-    assert(elt->getTypeCheckedRawValueExpr()->getType()->isEqual(rawTy));
-  }
-#endif
-
   bool isStringEnum =
     (rawTy->getNominalOrBoundGenericNominal() == C.getStringDecl());
   llvm::SmallVector<Expr *, 16> stringExprs;

lib/Sema/TypeCheckDecl.cpp

@@ -1433,7 +1433,6 @@ static LiteralExpr *getAutomaticRawValueExpr(TypeChecker &TC,
 
 static void checkEnumRawValues(TypeChecker &TC, EnumDecl *ED) {
   Type rawTy = ED->getRawType();
-
   if (!rawTy) {
     return;
   }
@@ -1490,8 +1489,10 @@ static void checkEnumRawValues(TypeChecker &TC, EnumDecl *ED) {
 
   for (auto elt : ED->getAllElements()) {
     // Skip if the raw value expr has already been checked.
-    if (elt->getTypeCheckedRawValueExpr())
+    if (elt->hasRawValueExpr() && elt->getRawValueExpr()->getType()) {
+      prevValue = elt->getRawValueExpr();
       continue;
+    }
 
     // Make sure the element is checked out before we poke at it.
     // FIXME: Make isInvalid work with interface types
@@ -1509,15 +1510,7 @@ static void checkEnumRawValues(TypeChecker &TC, EnumDecl *ED) {
       continue;
     }
 
-    // Check the raw value expr, if we have one.
-    if (auto *rawValue = elt->getRawValueExpr()) {
-      Expr *typeCheckedExpr = rawValue;
-      auto resultTy = TC.typeCheckExpression(typeCheckedExpr, ED,
-                                             TypeLoc::withoutLoc(rawTy),
-                                             CTP_EnumCaseRawValue);
-      if (resultTy) {
-        elt->setTypeCheckedRawValueExpr(typeCheckedExpr);
-      }
+    if (elt->hasRawValueExpr()) {
       lastExplicitValueElt = elt;
     } else {
       // If the enum element has no explicit raw value, try to
@@ -1529,43 +1522,20 @@ static void checkEnumRawValues(TypeChecker &TC, EnumDecl *ED) {
         break;
       }
       elt->setRawValueExpr(nextValue);
-      Expr *typeChecked = nextValue;
-      auto resultTy = TC.typeCheckExpression(
-          typeChecked, ED, TypeLoc::withoutLoc(rawTy), CTP_EnumCaseRawValue);
-      if (resultTy)
-        elt->setTypeCheckedRawValueExpr(typeChecked);
     }
     prevValue = elt->getRawValueExpr();
     assert(prevValue && "continued without setting raw value of enum case");
 
     // If we didn't find a valid initializer (maybe the initial value was
     // incompatible with the raw value type) mark the entry as being erroneous.
-    if (!elt->getTypeCheckedRawValueExpr()) {
+    TC.checkRawValueExpr(ED, elt);
+    if (!prevValue->getType() || prevValue->getType()->hasError()) {
       elt->setInvalid();
       continue;
     }
 
-    TC.checkEnumElementErrorHandling(elt);
-
-    // Find the type checked version of the LiteralExpr used for the raw value.
-    // this is unfortunate, but is needed because we're digging into the
-    // literals to get information about them, instead of accepting general
-    // expressions.
-    LiteralExpr *rawValue = elt->getRawValueExpr();
-    if (!rawValue->getType()) {
-      elt->getTypeCheckedRawValueExpr()->forEachChildExpr([&](Expr *E)->Expr* {
-        if (E->getKind() == rawValue->getKind())
-          rawValue = cast<LiteralExpr>(E);
-        return E;
-      });
-      elt->setRawValueExpr(rawValue);
-    }
-
-    prevValue = rawValue;
-    assert(prevValue && "continued without setting raw value of enum case");
-
     // Check that the raw value is unique.
-    RawValueKey key(rawValue);
+    RawValueKey key{prevValue};
     RawValueSource source{elt, lastExplicitValueElt};
 
     auto insertIterPair = uniqueRawValues.insert({key, source});
@@ -3030,6 +3000,13 @@ class DeclChecker : public DeclVisitor<DeclChecker> {
       TC.checkDefaultArguments(PL, EED);
     }
 
+    // Yell if our parent doesn't have a raw type but we have a raw value.
+    if (!EED->getParentEnum()->hasRawType() && EED->hasRawValueExpr()) {
+      TC.diagnose(EED->getRawValueExpr()->getLoc(),
+                  diag::enum_raw_value_without_raw_type);
+      EED->setInvalid();
+    }
+    
     checkAccessControl(TC, EED);
   }
 
@@ -4112,22 +4089,6 @@ void TypeChecker::validateDecl(ValueDecl *D) {
                              TypeResolverContext::EnumElementDecl);
     }
 
-    // If we have a raw value, make sure there's a raw type as well.
-    if (auto *rawValue = EED->getRawValueExpr()) {
-      if (!ED->hasRawType()) {
-        diagnose(rawValue->getLoc(),diag::enum_raw_value_without_raw_type);
-        // Recover by setting the raw type as this element's type.
-        Expr *typeCheckedExpr = rawValue;
-        if (!typeCheckExpression(typeCheckedExpr, ED)) {
-          EED->setTypeCheckedRawValueExpr(typeCheckedExpr);
-          checkEnumElementErrorHandling(EED);
-        }
-      } else {
-        // Wait until the second pass, when all the raw value expressions
-        // can be checked together.
-      }
-    }
-
     // Now that we have an argument type we can set the element's declared
     // type.
     EED->computeType();

lib/Sema/TypeCheckError.cpp

@@ -1663,9 +1663,9 @@ void TypeChecker::checkInitializerErrorHandling(Initializer *initCtx,
 /// perhaps accidentally, and (2) allows the verifier to assert that
 /// all calls have been checked.
 void TypeChecker::checkEnumElementErrorHandling(EnumElementDecl *elt) {
-  if (auto init = elt->getTypeCheckedRawValueExpr()) {
+  if (auto *rawValue = elt->getRawValueExpr()) {
     CheckErrorCoverage checker(*this, Context::forEnumElementInitializer(elt));
-    init->walk(checker);
+    rawValue->walk(checker);
   }
 }
 

lib/Sema/TypeCheckStmt.cpp

@@ -1911,6 +1911,20 @@ void TypeChecker::checkDefaultArguments(ParameterList *params,
   }
 }
 
+void TypeChecker::checkRawValueExpr(EnumDecl *ED, EnumElementDecl *EED) {
+  Type rawTy = ED->getRawType();
+  Expr *rawValue = EED->getRawValueExpr();
+  assert((rawTy && rawValue) && "Cannot check missing raw value!");
+
+  if (ED->getGenericEnvironmentOfContext() != nullptr)
+    rawTy = ED->mapTypeIntoContext(rawTy);
+
+  if (typeCheckExpression(rawValue, ED, TypeLoc::withoutLoc(rawTy),
+                          CTP_EnumCaseRawValue)) {
+    checkEnumElementErrorHandling(EED);
+  }
+}
+
 static Type getFunctionBuilderType(FuncDecl *FD) {
   Type builderType = FD->getFunctionBuilderType();
 

lib/Sema/TypeChecker.h

@@ -1008,6 +1008,8 @@ class TypeChecker final : public LazyResolver {
 
   /// Check the default arguments that occur within this value decl.
   void checkDefaultArguments(ParameterList *params, ValueDecl *VD);
+  /// Check the raw value expression in this enum element.
+  void checkRawValueExpr(EnumDecl *parent, EnumElementDecl *Elt);
 
   virtual void resolveDeclSignature(ValueDecl *VD) override {
     validateDecl(VD);