Re-apply "IRGen: Make case numbering consistent for enums with empty payloads"

Now with a small fix.

Author: slavapestov

lib/IRGen/GenDecl.cpp

@@ -2588,7 +2588,14 @@ StringRef IRGenModule::mangleType(CanType type, SmallVectorImpl<char> &buffer) {
   return StringRef(buffer.data(), buffer.size());
 }
 
-/// Is the given declaration resilient?
+/// Do we have to use resilient access patterns when working with this
+/// declaration?
+///
+/// IRGen is primarily concerned with resilient handling of the following:
+/// - For structs, a struct's size might change
+/// - For enums, new cases can be added
+/// - For classes, the superclass might change the size or number
+///   of stored properties
 bool IRGenModule::isResilient(Decl *D, ResilienceScope scope) {
   auto NTD = dyn_cast<NominalTypeDecl>(D);
   if (!NTD)
@@ -2604,6 +2611,23 @@ bool IRGenModule::isResilient(Decl *D, ResilienceScope scope) {
   llvm_unreachable("Bad resilience scope");
 }
 
+// The most general resilience scope where the given declaration is visible.
+ResilienceScope IRGenModule::getResilienceScopeForAccess(NominalTypeDecl *decl) {
+  if (decl->getModuleContext() == SILMod->getSwiftModule() &&
+      decl->getFormalAccess() != Accessibility::Public)
+    return ResilienceScope::Component;
+  return ResilienceScope::Universal;
+}
+
+// The most general resilience scope which has knowledge of the declaration's
+// layout. Calling isResilient() with this scope will always return false.
+ResilienceScope IRGenModule::getResilienceScopeForLayout(NominalTypeDecl *decl) {
+  if (isResilient(decl, ResilienceScope::Universal))
+    return ResilienceScope::Component;
+
+  return getResilienceScopeForAccess(decl);
+}
+
 /// Fetch the witness table access function for a protocol conformance.
 llvm::Function *
 IRGenModule::getAddrOfWitnessTableAccessFunction(

lib/IRGen/GenEnum.cpp

@@ -2707,6 +2707,7 @@ namespace {
 
     CopyDestroyStrategy CopyDestroyKind;
     ReferenceCounting Refcounting;
+    bool AllowFixedLayoutOptimizations;
 
     static EnumPayloadSchema getPayloadSchema(ArrayRef<Element> payloads) {
       // TODO: We might be able to form a nicer schema if the payload elements
@@ -2726,6 +2727,7 @@ namespace {
     MultiPayloadEnumImplStrategy(IRGenModule &IGM,
                                  TypeInfoKind tik,
                                  IsFixedSize_t alwaysFixedSize,
+                                 bool allowFixedLayoutOptimizations,
                                  unsigned NumElements,
                                  std::vector<Element> &&WithPayload,
                                  std::vector<Element> &&WithNoPayload)
@@ -2734,7 +2736,8 @@ namespace {
                                     std::move(WithPayload),
                                     std::move(WithNoPayload),
                                     getPayloadSchema(WithPayload)),
-                                    CopyDestroyKind(Normal)
+        CopyDestroyKind(Normal),
+        AllowFixedLayoutOptimizations(allowFixedLayoutOptimizations)
     {
       assert(ElementsWithPayload.size() > 1);
 
@@ -2794,7 +2797,7 @@ namespace {
       // we might need the payload size if from another module the enum has
       // a dynamic size, which can happen if the enum contains a resilient
       // payload.
-      return !AlwaysFixedSize;
+      return !AllowFixedLayoutOptimizations;
     }
 
     unsigned getPayloadSizeForMetadata() const override {
@@ -4595,20 +4598,26 @@ EnumImplStrategy *EnumImplStrategy::get(TypeConverter &TC,
   unsigned numElements = 0;
   TypeInfoKind tik = Loadable;
   IsFixedSize_t alwaysFixedSize = IsFixedSize;
+  bool allowFixedLayoutOptimizations = true;
   std::vector<Element> elementsWithPayload;
   std::vector<Element> elementsWithNoPayload;
 
-  // The most general resilience scope that can have knowledge of this
-  // enum's layout. If all payload types have a fixed size in this
-  // resilience scope, we can make further assumptions to optimize
-  // layout.
-  ResilienceScope scope = ResilienceScope::Universal;
+  // Resilient enums are manipulated as opaque values, except we still
+  // make the following assumptions:
+  // 1) Physical case indices won't change
+  // 2) The indirect-ness of cases won't change
+  // 3) Payload types won't change in a non-resilient way
+  bool isResilient = TC.IGM.isResilient(theEnum, ResilienceScope::Component);
+  
+  // The most general resilience scope where the enum type is visible.
+  // Case numbering must not depend on any information that is not static
+  // in this resilience scope.
+  ResilienceScope accessScope = TC.IGM.getResilienceScopeForAccess(theEnum);
 
-  // TODO: Replace this with 'public or internal with @availability' check
-  // once that is in place
-  if (theEnum->getFormalAccess() != Accessibility::Public ||
-      TC.IGM.isResilient(theEnum, ResilienceScope::Universal))
-    scope = ResilienceScope::Component;
+  // The most general resilience scope where the enum's layout is known.
+  // Fixed-size optimizations can be applied if all payload types are
+  // fixed-size from this resilience scope.
+  ResilienceScope layoutScope = TC.IGM.getResilienceScopeForLayout(theEnum);
 
   for (auto elt : theEnum->getAllElements()) {
     numElements++;
@@ -4638,14 +4647,20 @@ EnumImplStrategy *EnumImplStrategy::get(TypeConverter &TC,
       = TC.tryGetCompleteTypeInfo(origArgLoweredTy.getSwiftRValueType());
     assert(origArgTI && "didn't complete type info?!");
 
-    // If the unsubstituted argument contains a generic parameter type, or if
-    // the substituted argument is not universally fixed-size, we need to
-    // constrain our layout optimizations to what the runtime can reproduce.
-    if (!origArgTI->isFixedSize(scope))
-      alwaysFixedSize = IsNotFixedSize;
-
-    auto loadableOrigArgTI = dyn_cast<LoadableTypeInfo>(origArgTI);
-    if (loadableOrigArgTI && loadableOrigArgTI->isKnownEmpty()) {
+    // If the unsubstituted argument contains a generic parameter type, or
+    // is not fixed-size in all resilience domains that have knowledge of
+    // this enum's layout, we need to constrain our layout optimizations to
+    // what the runtime can reproduce.
+    if (!isResilient &&
+        !origArgTI->isFixedSize(layoutScope))
+      allowFixedLayoutOptimizations = false;
+
+    // If the payload is empty, turn the case into a no-payload case, but
+    // only if case numbering remains unchanged from all resilience domains
+    // that can see the enum.
+    if (origArgTI->isFixedSize(accessScope) &&
+        isa<LoadableTypeInfo>(origArgTI) &&
+        cast<LoadableTypeInfo>(origArgTI)->isKnownEmpty()) {
       elementsWithNoPayload.push_back({elt, nullptr, nullptr});
     } else {
       // *Now* apply the substitutions and get the type info for the instance's
@@ -4655,16 +4670,22 @@ EnumImplStrategy *EnumImplStrategy::get(TypeConverter &TC,
       auto *substArgTI = &TC.IGM.getTypeInfo(fieldTy);
 
       elementsWithPayload.push_back({elt, substArgTI, origArgTI});
-      if (!substArgTI->isFixedSize())
-        tik = Opaque;
-      else if (!substArgTI->isLoadable() && tik > Fixed)
-        tik = Fixed;
 
-      // If the substituted argument contains a type that is not universally
-      // fixed-size, we need to constrain our layout optimizations to what
-      // the runtime can reproduce.
-      if (!substArgTI->isFixedSize(scope))
-        alwaysFixedSize = IsNotFixedSize;
+      if (!isResilient) {
+        if (!substArgTI->isFixedSize(ResilienceScope::Component))
+          tik = Opaque;
+        else if (!substArgTI->isLoadable() && tik > Fixed)
+          tik = Fixed;
+
+        // If the substituted argument contains a type that is not fixed-size
+        // in all resilience domains that have knowledge of this enum's layout,
+        // we need to constrain our layout optimizations to what the runtime
+        // can reproduce.
+        if (!substArgTI->isFixedSize(layoutScope)) {
+          alwaysFixedSize = IsNotFixedSize;
+          allowFixedLayoutOptimizations = false;
+        }
+      }
     }
   }
 
@@ -4673,12 +4694,7 @@ EnumImplStrategy *EnumImplStrategy::get(TypeConverter &TC,
            + elementsWithNoPayload.size()
          && "not all elements accounted for");
 
-  // Resilient enums are manipulated as opaque values, except we still
-  // make the following assumptions:
-  // 1) Physical case indices won't change
-  // 2) The indirect-ness of cases won't change
-  // 3) Payload types won't change in a non-resilient way
-  if (TC.IGM.isResilient(theEnum, ResilienceScope::Component)) {
+  if (isResilient) {
     return new ResilientEnumImplStrategy(TC.IGM,
                                          numElements,
                                          std::move(elementsWithPayload),
@@ -4702,6 +4718,7 @@ EnumImplStrategy *EnumImplStrategy::get(TypeConverter &TC,
                                          std::move(elementsWithNoPayload));
   if (elementsWithPayload.size() > 1)
     return new MultiPayloadEnumImplStrategy(TC.IGM, tik, alwaysFixedSize,
+                                            allowFixedLayoutOptimizations,
                                             numElements,
                                             std::move(elementsWithPayload),
                                             std::move(elementsWithNoPayload));
@@ -5216,7 +5233,7 @@ MultiPayloadEnumImplStrategy::completeFixedLayout(TypeConverter &TC,
     // The runtime currently does not track spare bits, so we can't use them
     // if the type is layout-dependent. (Even when the runtime does, it will
     // likely only track a subset of the spare bits.)
-    if (!AlwaysFixedSize || TIK < Loadable) {
+    if (!AllowFixedLayoutOptimizations || TIK < Loadable) {
       if (CommonSpareBits.size() < payloadBits)
         CommonSpareBits.extendWithClearBits(payloadBits);
       continue;

lib/IRGen/IRGenModule.h

@@ -504,6 +504,8 @@ class IRGenModule {
   }
 
   bool isResilient(Decl *decl, ResilienceScope scope);
+  ResilienceScope getResilienceScopeForAccess(NominalTypeDecl *decl);
+  ResilienceScope getResilienceScopeForLayout(NominalTypeDecl *decl);
 
   SpareBitVector getSpareBitsForType(llvm::Type *scalarTy, Size size);
 

test/Inputs/resilient_enum.swift

@@ -40,8 +40,6 @@ public struct Color {
   case Bespoke(Color, Color)
 }
 
-// Fixed-layout enum with resilient members
-
 // Resilient enum
 public enum Medium {
   // Empty cases
@@ -60,6 +58,21 @@ public indirect enum IndirectApproach {
   case Angle(Double)
 }
 
+// Resilient enum with resilient empty payload case
+public struct EmptyStruct {
+  public init() {}
+}
+
+public enum ResilientEnumWithEmptyCase {
+  case A // should always be case 1
+  case B // should always be case 2
+  case Empty(EmptyStruct) // should always be case 0
+}
+
+public func getResilientEnumWithEmptyCase() -> [ResilientEnumWithEmptyCase] {
+  return [.A, .B, .Empty(EmptyStruct())]
+}
+
 // Specific enum implementations for executable tests
 public enum ResilientEmptyEnum {
   case X

test/Interpreter/enum_resilience.swift

@@ -379,4 +379,25 @@ ResilientEnumTestSuite.test("DynamicLayoutMultiPayload2") {
   expectEqual(b, [0, 1, 2, 3])
 }
 
+// Make sure case numbers round-trip if payload has zero size
+
+ResilientEnumTestSuite.test("ResilientEnumWithEmptyCase") {
+  let a: [ResilientEnumWithEmptyCase] = getResilientEnumWithEmptyCase()
+
+  let b: [Int] = a.map {
+    switch $0 {
+    case .A:
+      return 0
+    case .B:
+      return 1
+    case .Empty:
+      return 2
+    default:
+      return -1
+    }
+  }
+
+  expectEqual(b, [0, 1, 2])
+}
+
 runAllTests()