Clean up enum raw expression validation a bit

The distinction between the type checked raw value expression and the regular raw value expression was never important.  Downstream clients were ignoring the type checked form and pulling the text out of the supposed "plain" form.  Drop the distinction and simply don't set back into the raw value expr if we don't have to.

Pushing this through naturally enables some cleanup in checkEnumRawValues.  Factor out type checking the literal value into an helper on the typechecker and pull a common diagnostic into the decl checker.

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

@@ -1432,7 +1432,6 @@ static LiteralExpr *getAutomaticRawValueExpr(TypeChecker &TC,
 
 static void checkEnumRawValues(TypeChecker &TC, EnumDecl *ED) {
   Type rawTy = ED->getRawType();
-
   if (!rawTy) {
     return;
   }
@@ -1489,8 +1488,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
@@ -1508,15 +1509,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
@@ -1528,43 +1521,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});
@@ -3053,6 +3023,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);
   }
 
@@ -4135,22 +4112,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

@@ -1010,6 +1010,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);