Merge pull request swiftlang#74145 from tbkka/tbkka-remotemirror-mpe-fixes

Expand the work from swiftlang#73491 to support more MPE layouts.

This is also switches the MPE layout code to exclusively use the new approach. The key observation: existing reflection metadata seems to already provide enough information in all cases, so we can abandon an earlier effort to add spare bitmask data.

There is some risk that the old code may have handled some enums that aren't (yet) fully supported by the new code. However, I did a bunch of experiments and found that our existing test cases have decent coverage of the major capabilities and I've also added a new test case specifically to exercise enum layouts being accessed resiliently. So if there are regressions from this change, they should be minor and easily fixed.

Resolves rdar://129281368

Author: tbkka

include/swift/RemoteInspection/BitMask.h

@@ -74,7 +74,16 @@ class BitMask {
 
   BitMask(unsigned sizeInBytes, uint64_t sourceMask): size(sizeInBytes) {
     mask = (uint8_t *)calloc(1, sizeInBytes);
-    memcpy(mask, &sourceMask, sizeInBytes);
+    if (!mask) {
+      assert(false && "Failed to allocate Bitmask");
+      size = 0;
+      return;
+    }
+    size_t toCopy = sizeInBytes;
+    if (toCopy > sizeof(sourceMask)) {
+      toCopy = sizeof(sourceMask);
+    }
+    memcpy(mask, &sourceMask, toCopy);
   }
 
   // Construct a bitmask of the appropriate number of bytes

stdlib/public/RemoteInspection/TypeLowering.cpp

@@ -251,7 +251,7 @@ BuiltinTypeInfo::BuiltinTypeInfo(unsigned Size, unsigned Alignment,
 
 // Builtin.Int<N> is mangled as 'Bi' N '_'
 // Returns 0 if this isn't an Int
-static unsigned isIntType(std::string name) {
+static unsigned intTypeBitSize(std::string name) {
   llvm::StringRef nameRef(name);
   if (nameRef.starts_with("Bi") && nameRef.endswith("_")) {
     llvm::StringRef naturalRef = nameRef.drop_front(2).drop_back();
@@ -274,11 +274,12 @@ bool BuiltinTypeInfo::readExtraInhabitantIndex(
   }
   // If it has extra inhabitants, it could be an integer type with extra
   // inhabitants (such as a bool) or a pointer.
-  unsigned intSize = isIntType(Name);
+  unsigned intSize = intTypeBitSize(Name);
   if (intSize > 0) {
     // This is an integer type
 
-    // If it cannot have extra inhabitants, return early...
+    // If extra inhabitants are impossible, return early...
+    // (assert in debug builds)
     assert(intSize < getSize() * 8
 	   && "Standard-sized int cannot have extra inhabitants");
     if (intSize > 64 || getSize() > 8 || intSize >= getSize() * 8) {
@@ -319,14 +320,15 @@ bool BuiltinTypeInfo::readExtraInhabitantIndex(
 }
 
 BitMask BuiltinTypeInfo::getSpareBits(TypeConverter &TC, bool &hasAddrOnly) const {
-  unsigned intSize = isIntType(Name);
+  unsigned intSize = intTypeBitSize(Name);
   if (intSize > 0) {
     // Odd-sized integers export spare bits
     // In particular: bool fields are Int1 and export 7 spare bits
     auto mask = BitMask::oneMask(getSize());
     mask.keepOnlyMostSignificantBits(getSize() * 8 - intSize);
     return mask;
   } else if (
+    Name == "ypXp" || // Any.Type
     Name == "yyXf" // 'yyXf' =  @thin () -> Void function
   ) {
     // Builtin types that expose pointer spare bits
@@ -590,7 +592,9 @@ class NoPayloadEnumTypeInfo: public EnumTypeInfo {
   }
 
   BitMask getSpareBits(TypeConverter &TC, bool &hasAddrOnly) const override {
-    return BitMask::zeroMask(getSize());
+    auto mask = BitMask(getSize(), maskForCount(getNumCases()));
+    mask.complement();
+    return mask;
   }
 
   bool projectEnumValue(remote::MemoryReader &reader,
@@ -660,7 +664,17 @@ class SinglePayloadEnumTypeInfo: public EnumTypeInfo {
   }
 
   BitMask getSpareBits(TypeConverter &TC, bool &hasAddrOnly) const override {
-    return BitMask::zeroMask(getSize());
+    FieldInfo PayloadCase = getCases()[0];
+    size_t payloadSize = PayloadCase.TI.getSize();
+    if (getSize() <= payloadSize) {
+      return BitMask::zeroMask(getSize());
+    }
+    size_t tagSize = getSize() - payloadSize;
+    auto mask = BitMask::oneMask(getSize());
+    mask.keepOnlyMostSignificantBits(tagSize * 8); // Clear payload bits
+    auto tagMaskUsedBits = BitMask(getSize(), maskForCount(getNumCases()));
+    mask.andNotMask(tagMaskUsedBits, payloadSize); // Clear used tag bits
+    return mask;
   }
 
   // Think of a single-payload enum as being encoded in "pages".
@@ -1973,6 +1987,15 @@ class EnumTypeInfoBuilder {
     default: Kind = EnumKind::MultiPayloadEnum; break;
     }
 
+    // Sanity:  Ignore any enum that claims to have a size more than 1MiB
+    // This avoids allocating lots of memory for spare bit mask calculations
+    // when clients try to interpret random chunks of memory as type descriptions.
+    if (Size > (1024ULL * 1024)) {
+      unsigned Stride = ((Size + Alignment - 1) & ~(Alignment - 1));
+      return TC.makeTypeInfo<UnsupportedEnumTypeInfo>(
+	Size, Alignment, Stride, NumExtraInhabitants, BitwiseTakable, Kind, Cases);
+    }
+
     if (Cases.size() == 1) {
       if (EffectivePayloadCases == 0) {
         // Zero-sized enum with only one empty case
@@ -2058,11 +2081,10 @@ class EnumTypeInfoBuilder {
     //
 
     // Do we have a fixed layout?
-    // TODO: Test whether a missing FixedDescriptor is actually relevant.
     auto FixedDescriptor = TC.getBuilder().getBuiltinTypeDescriptor(TR);
     if (!FixedDescriptor || GenericPayloadCases > 0) {
       // This is a "dynamic multi-payload enum".  For example,
-      // this occurs with:
+      // this occurs with generics such as:
       // ```
       // class ClassWithEnum<T> {
       //   enum E {
@@ -2072,6 +2094,12 @@ class EnumTypeInfoBuilder {
       //   var e: E?
       // }
       // ```
+      // and when we have a resilient inner enum, such as:
+      // ```
+      // enum E2 {
+      //   case y(E1_resilient)
+      //   case z(Int)
+      // }
       auto tagCounts = getEnumTagCounts(Size, EffectiveNoPayloadCases,
                                         EffectivePayloadCases);
       Size += tagCounts.numTagBytes;
@@ -2103,7 +2131,6 @@ class EnumTypeInfoBuilder {
     unsigned Stride = ((Size + Alignment - 1) & ~(Alignment - 1));
     if (Stride == 0)
       Stride = 1;
-    auto PayloadSize = EnumTypeInfo::getPayloadSizeForCases(Cases);
 
     // Compute the spare bit mask and determine if we have any address-only fields
     auto localSpareBitMask = BitMask::oneMask(Size);
@@ -2115,44 +2142,11 @@ class EnumTypeInfoBuilder {
       }
     }
 
-    // See if we have MPE bit mask information from the compiler...
-    // TODO: drop this?
-
-    // Uncomment the following line to dump the MPE section every time we come through here...
-    //TC.getBuilder().dumpMultiPayloadEnumSection(std::cerr); // DEBUG helper
-
-    auto MPEDescriptor = TC.getBuilder().getMultiPayloadEnumDescriptor(TR);
-    if (MPEDescriptor && MPEDescriptor->usesPayloadSpareBits()) {
-      // We found compiler-provided spare bit data...
-      auto PayloadSpareBitMaskByteCount = MPEDescriptor->getPayloadSpareBitMaskByteCount();
-      auto PayloadSpareBitMaskByteOffset = MPEDescriptor->getPayloadSpareBitMaskByteOffset();
-      auto SpareBitMask = MPEDescriptor->getPayloadSpareBits();
-      BitMask compilerSpareBitMask(PayloadSize, SpareBitMask,
-                            PayloadSpareBitMaskByteCount, PayloadSpareBitMaskByteOffset);
-      
-      if (compilerSpareBitMask.isZero() || hasAddrOnly) {
-        // If there are no spare bits, use the "simple" tag-only implementation.
-        return TC.makeTypeInfo<TaggedMultiPayloadEnumTypeInfo>(
-          Size, Alignment, Stride, NumExtraInhabitants,
-          BitwiseTakable, Cases, EffectivePayloadCases);
-      }
-
-#if 0  // TODO: This should be !defined(NDEBUG)
-      // Verify that compiler provided and local spare bit info agree...
-      // TODO: If we could make this actually work, then we wouldn't need the
-      // bulky compiler-provided info, would we?
-      assert(localSpareBitMask == compilerSpareBitMask);
-#endif
-
-      // Use compiler-provided spare bit information
-      return TC.makeTypeInfo<MultiPayloadEnumTypeInfo>(
-        Size, Alignment, Stride, NumExtraInhabitants,
-        BitwiseTakable, Cases, compilerSpareBitMask,
-        EffectivePayloadCases);
-    }
-
     if (localSpareBitMask.isZero() || hasAddrOnly) {
-      // Simple case that does not use spare bits
+      // Simple tag-only layout does not use spare bits.
+      // Either:
+      // * There are no spare bits, or
+      // * We can't copy it to strip spare bits.
       return TC.makeTypeInfo<TaggedMultiPayloadEnumTypeInfo>(
         Size, Alignment, Stride, NumExtraInhabitants,
         BitwiseTakable, Cases, EffectivePayloadCases);

validation-test/Reflection/Inputs/reflect_Enum_values_resilient_enums.swift

@@ -0,0 +1,10 @@
+
+public enum E1_resilient {
+case a
+case b
+}
+
+public enum E2_resilient {
+case c(E1_resilient)
+case d(E1_resilient)
+}

validation-test/Reflection/reflect_Enum_values_resilient.swift

@@ -0,0 +1,75 @@
+// RUN: %empty-directory(%t)
+
+// RUN: %target-build-swift-dylib(%t/%target-library-name(resilient_enums)) -enable-library-evolution %S/Inputs/reflect_Enum_values_resilient_enums.swift -emit-module -emit-module-path %t/resilient_enums.swiftmodule -module-name resilient_enums
+// RUN: %target-codesign %t/%target-library-name(resilient_enums)
+
+// RUN: %target-build-swift -lswiftSwiftReflectionTest %s -L %t -I %t -lresilient_enums -o %t/reflect_Enum_values_resilient %target-rpath(%t)
+// RUN: %target-codesign %t/reflect_Enum_values_resilient
+
+// RUN: %target-run %target-swift-reflection-test %t/reflect_Enum_values_resilient | tee /dev/stderr | %FileCheck %s --dump-input=fail
+
+// REQUIRES: executable_test
+// REQUIRES: objc_interop, OS=macosx
+// UNSUPPORTED: use_os_stdlib
+
+import SwiftReflectionTest
+
+import resilient_enums
+
+// Non-resilient enum wrapping a resilient enum
+// This doesn't use spare bits of the inner enum
+enum E2 {
+case y(E1_resilient)
+case z(E1_resilient)
+}
+
+// Contrast:
+// E2_resilient is a resilient enum wrapping a resilient enum
+// This does use spare bits of the inner enum
+
+
+// CHECK: Reflecting an enum value.
+// CHECK-NEXT: Type reference:
+// CHECK-NEXT: (enum reflect_Enum_values_resilient.E2)
+// CHECK-NEXT: Value: .y(.a)
+
+reflect(enumValue: E2.y(.a))
+
+// CHECK: Reflecting an enum value.
+// CHECK-NEXT: Type reference:
+// CHECK-NEXT: (enum reflect_Enum_values_resilient.E2)
+// CHECK-NEXT: Value: .z(.b)
+
+reflect(enumValue: E2.z(.b))
+
+// CHECK: Reflecting an enum value.
+// CHECK-NEXT: Type reference:
+// CHECK-NEXT: (enum resilient_enums.E1_resilient)
+// CHECK-NEXT: Value: .a
+
+reflect(enumValue: E1_resilient.a)
+
+// CHECK: Reflecting an enum value.
+// CHECK-NEXT: Type reference:
+// CHECK-NEXT: (enum resilient_enums.E1_resilient)
+// CHECK-NEXT: Value: .b
+
+reflect(enumValue: E1_resilient.b)
+
+// CHECK: Reflecting an enum value.
+// CHECK-NEXT: Type reference:
+// CHECK-NEXT: (enum resilient_enums.E2_resilient)
+// CHECK-NEXT: Value: .c(.a)
+
+reflect(enumValue: E2_resilient.c(.a))
+
+// CHECK: Reflecting an enum value.
+// CHECK-NEXT: Type reference:
+// CHECK-NEXT: (enum resilient_enums.E2_resilient)
+// CHECK-NEXT: Value: .d(.b)
+
+reflect(enumValue: E2_resilient.d(.b))
+
+doneReflecting()
+
+// CHECK: Done.