[SR-5036]: Segfault using KeyPath with NSObject

include/swift/ABI/KeyPath.h

@@ -104,11 +104,11 @@ class KeyPathComponentHeader {
   }
 
   constexpr static KeyPathComponentHeader
-  forStructComponentWithUnresolvedOffset() {
+  forStructComponentWithUnresolvedFieldOffset() {
     return KeyPathComponentHeader(
       (_SwiftKeyPathComponentHeader_StructTag
       << _SwiftKeyPathComponentHeader_DiscriminatorShift)
-      | _SwiftKeyPathComponentHeader_UnresolvedOffsetPayload);
+      | _SwiftKeyPathComponentHeader_UnresolvedFieldOffsetPayload);
   }
 
   constexpr static KeyPathComponentHeader
@@ -128,11 +128,19 @@ class KeyPathComponentHeader {
   }
 
   constexpr static KeyPathComponentHeader
-  forClassComponentWithUnresolvedOffset() {
+  forClassComponentWithUnresolvedFieldOffset() {
     return KeyPathComponentHeader(
-      (_SwiftKeyPathComponentHeader_StructTag
+      (_SwiftKeyPathComponentHeader_ClassTag
+      << _SwiftKeyPathComponentHeader_DiscriminatorShift)
+      | _SwiftKeyPathComponentHeader_UnresolvedFieldOffsetPayload);
+  }
+
+  constexpr static KeyPathComponentHeader
+  forClassComponentWithUnresolvedIndirectOffset() {
+    return KeyPathComponentHeader(
+      (_SwiftKeyPathComponentHeader_ClassTag
       << _SwiftKeyPathComponentHeader_DiscriminatorShift)
-      | _SwiftKeyPathComponentHeader_UnresolvedOffsetPayload);
+      | _SwiftKeyPathComponentHeader_UnresolvedIndirectOffsetPayload);
   }
 
   constexpr static KeyPathComponentHeader

lib/IRGen/GenClass.cpp

@@ -181,14 +181,30 @@ namespace {
     // If not, we can have to access stored properties through the field
     // offset vector in the instantiated type metadata.
     bool ClassHasConcreteLayout = true;
-
+    
   public:
-    ClassLayoutBuilder(IRGenModule &IGM, SILType classType)
+    ClassLayoutBuilder(IRGenModule &IGM, SILType classType,
+                       ReferenceCounting refcounting)
       : StructLayoutBuilder(IGM)
     {
       // Start by adding a heap header.
-      addHeapHeader();
-
+      switch (refcounting) {
+      case swift::irgen::ReferenceCounting::Native:
+        // For native classes, place a full object header.
+        addHeapHeader();
+        break;
+      case swift::irgen::ReferenceCounting::ObjC:
+        // For ObjC-inheriting classes, we don't reliably know the size of the
+        // base class, but NSObject only has an `isa` pointer at most.
+        addNSObjectHeader();
+        break;
+      case swift::irgen::ReferenceCounting::Block:
+      case swift::irgen::ReferenceCounting::Unknown:
+      case swift::irgen::ReferenceCounting::Bridge:
+      case swift::irgen::ReferenceCounting::Error:
+        llvm_unreachable("not a class refcounting kind");
+      }
+
       // Next, add the fields for the given class.
       auto theClass = classType.getClassOrBoundGenericClass();
       assert(theClass);
@@ -382,7 +398,7 @@ void ClassTypeInfo::generateLayout(IRGenModule &IGM, SILType classType) const {
          "already generated layout");
 
   // Add the heap header.
-  ClassLayoutBuilder builder(IGM, classType);
+  ClassLayoutBuilder builder(IGM, classType, Refcount);
 
   // generateLayout can call itself recursively in order to compute a layout
   // for the abstract type.  If classType shares an exemplar types with the
@@ -409,7 +425,8 @@ void ClassTypeInfo::generateLayout(IRGenModule &IGM, SILType classType) const {
 const StructLayout &
 ClassTypeInfo::getLayout(IRGenModule &IGM, SILType classType) const {
   // Return the cached layout if available.
-  if (Layout) return *Layout;
+  if (Layout)
+    return *Layout;
 
   generateLayout(IGM, classType);
   return *Layout;
@@ -489,7 +506,13 @@ irgen::tryEmitConstantClassFragilePhysicalMemberOffset(IRGenModule &IGM,
   }
 }
 
-
+FieldAccess
+irgen::getClassFieldAccess(IRGenModule &IGM, SILType baseType, VarDecl *field) {
+  auto &baseClassTI = IGM.getTypeInfo(baseType).as<ClassTypeInfo>();
+  auto &classLayout = baseClassTI.getClassLayout(IGM, baseType);
+  unsigned fieldIndex = classLayout.getFieldIndex(field);
+  return classLayout.AllFieldAccesses[fieldIndex];
+}
 
 OwnedAddress irgen::projectPhysicalClassMemberAddress(IRGenFunction &IGF,
                                                       llvm::Value *base,

lib/IRGen/GenClass.h

@@ -48,6 +48,7 @@ namespace irgen {
   enum class ReferenceCounting : unsigned char;
   enum class IsaEncoding : unsigned char;
   enum class ClassDeallocationKind : unsigned char;
+  enum class FieldAccess : uint8_t;
 
   OwnedAddress projectPhysicalClassMemberAddress(IRGenFunction &IGF,
                                                  llvm::Value *base,
@@ -138,6 +139,10 @@ namespace irgen {
   tryEmitConstantClassFragilePhysicalMemberOffset(IRGenModule &IGM,
                                                   SILType baseType,
                                                   VarDecl *field);
+
+  FieldAccess getClassFieldAccess(IRGenModule &IGM,
+                                  SILType baseType,
+                                  VarDecl *field);
 
   /// What reference counting mechanism does a class-like type use?
   ReferenceCounting getReferenceCountingForType(IRGenModule &IGM,

lib/IRGen/GenKeyPath.cpp

@@ -25,6 +25,7 @@
 #include "IRGenFunction.h"
 #include "IRGenModule.h"
 #include "ProtocolInfo.h"
+#include "StructLayout.h"
 #include "llvm/ADT/SetVector.h"
 #include "swift/SIL/SILInstruction.h"
 #include "swift/SIL/SILLocation.h"
@@ -209,47 +210,98 @@ IRGenModule::getAddrOfKeyPathPattern(KeyPathPattern *pattern,
     auto &component = pattern->getComponents()[i];
     switch (auto kind = component.getKind()) {
     case KeyPathPatternComponent::Kind::StoredProperty: {
-      // Try to get a constant offset if we can.
       auto property = cast<VarDecl>(component.getStoredPropertyDecl());
-      llvm::Constant *offset;
-      bool isResolved;
-      std::tie(offset, isResolved)
-        = getPropertyOffsetOrIndirectOffset(loweredBaseTy, property);
-      offset = llvm::ConstantExpr::getTruncOrBitCast(offset, Int32Ty);
-      bool isStruct = (bool)loweredBaseTy.getStructOrBoundGenericStruct();
 
-      // If the projection is a statically known integer, try to pack it into
-      // the key path payload.
-      if (isResolved) {
+      auto addFixedOffset = [&](bool isStruct, llvm::Constant *offset) {
+        offset = llvm::ConstantExpr::getTruncOrBitCast(offset, Int32Ty);
         if (auto offsetInt = dyn_cast_or_null<llvm::ConstantInt>(offset)) {
           auto offsetValue = offsetInt->getValue().getZExtValue();
           if (KeyPathComponentHeader::offsetCanBeInline(offsetValue)) {
             auto header = isStruct
               ? KeyPathComponentHeader::forStructComponentWithInlineOffset(offsetValue)
               : KeyPathComponentHeader::forClassComponentWithInlineOffset(offsetValue);
             fields.addInt32(header.getData());
-            break;
+            return;
           }
         }
-
         auto header = isStruct
           ? KeyPathComponentHeader::forStructComponentWithOutOfLineOffset()
           : KeyPathComponentHeader::forClassComponentWithOutOfLineOffset();
         fields.addInt32(header.getData());
-
         fields.add(offset);
-      } else {
-        // Otherwise, stash the offset of the field offset within the metadata
-        // object, so we can pull it out at instantiation time.
-        // TODO: We'll also need a way to handle resilient field offsets, once
-        // field offset vectors no longer cover all fields in the type.
-        KeyPathComponentHeader header = isStruct
-          ? KeyPathComponentHeader::forStructComponentWithUnresolvedOffset()
-          : KeyPathComponentHeader::forClassComponentWithUnresolvedOffset();
+      };
+
+      // For a struct stored property, we may know the fixed offset of the field,
+      // or we may need to fetch it out of the type's metadata at instantiation
+      // time.
+      if (loweredBaseTy.getStructOrBoundGenericStruct()) {
+        if (auto offset = emitPhysicalStructMemberFixedOffset(*this,
+                                                              loweredBaseTy,
+                                                              property)) {
+          // We have a known constant fixed offset.
+          addFixedOffset(/*struct*/ true, offset);
+          break;
+        }
+
+        // If the offset isn't fixed, try instead to get the field offset out
+        // of the type metadata at instantiation time.
+        auto fieldOffset = emitPhysicalStructMemberOffsetOfFieldOffset(
+                                                *this, loweredBaseTy, property);
+        fieldOffset = llvm::ConstantExpr::getTruncOrBitCast(fieldOffset,
+                                                            Int32Ty);
+        auto header = KeyPathComponentHeader::forStructComponentWithUnresolvedFieldOffset();
         fields.addInt32(header.getData());
-        fields.add(offset);
+        fields.add(fieldOffset);
+        break;
+      }
+
+      // For a class, we may know the fixed offset of a field at compile time,
+      // or we may need to fetch it at instantiation time. Depending on the
+      // ObjC-ness and resilience of the class hierarchy, there might be a few
+      // different ways we need to go about this.
+      if (loweredBaseTy.getClassOrBoundGenericClass()) {
+        switch (getClassFieldAccess(*this, loweredBaseTy, property)) {
+        case FieldAccess::ConstantDirect: {
+          // Known constant fixed offset.
+          auto offset = tryEmitConstantClassFragilePhysicalMemberOffset(*this,
+                                                                  loweredBaseTy,
+                                                                  property);
+          assert(offset && "no constant offset for ConstantDirect field?!");
+          addFixedOffset(/*struct*/ false, offset);
+          break;
+        }
+        case FieldAccess::NonConstantDirect: {
+          // A constant offset that's determined at class realization time.
+          // We have to load the offset from a global ivar.
+          auto header =
+            KeyPathComponentHeader::forClassComponentWithUnresolvedIndirectOffset();
+          fields.addInt32(header.getData());
+          auto offsetVar = getAddrOfFieldOffset(property, /*indirect*/ false,
+                                                NotForDefinition);
+          fields.add(cast<llvm::Constant>(offsetVar.getAddress()));
+          break;
+        }
+        case FieldAccess::ConstantIndirect: {
+          // An offset that depends on the instance's generic parameterization,
+          // but whose field offset is at a known vtable offset.
+          auto header =
+            KeyPathComponentHeader::forClassComponentWithUnresolvedFieldOffset();
+          fields.addInt32(header.getData());
+          auto fieldOffset =
+            getClassFieldOffset(*this, loweredBaseTy.getClassOrBoundGenericClass(),
+                                property);
+          fields.addInt32(fieldOffset.getValue());
+          break;
+        }
+        case FieldAccess::NonConstantIndirect:
+          // An offset that depends on the instance's generic parameterization,
+          // whose vtable offset is also unknown.
+          // TODO: This doesn't happen until class resilience is enabled.
+          llvm_unreachable("not implemented");
+        }
+        break;
       }
-      break;
+      llvm_unreachable("not struct or class");
     }
     case KeyPathPatternComponent::Kind::GettableProperty:
     case KeyPathPatternComponent::Kind::SettableProperty: {

lib/IRGen/StructLayout.cpp

@@ -189,6 +189,13 @@ void StructLayoutBuilder::addHeapHeader() {
   StructFields.push_back(IGM.RefCountedStructTy);
 }
 
+void StructLayoutBuilder::addNSObjectHeader() {
+  assert(StructFields.empty() && "adding heap header at a non-zero offset");
+  CurSize = IGM.getPointerSize();
+  CurAlignment = IGM.getPointerAlignment();
+  StructFields.push_back(IGM.ObjCClassPtrTy);
+}
+
 bool StructLayoutBuilder::addFields(llvm::MutableArrayRef<ElementLayout> elts,
                                     LayoutStrategy strategy) {
   // Track whether we've added any storage to our layout.

lib/IRGen/StructLayout.h

@@ -233,9 +233,12 @@ class StructLayoutBuilder {
   StructLayoutBuilder(IRGenModule &IGM) : IGM(IGM) {}
 
   /// Add a swift heap header to the layout.  This must be the first
-  /// call to the layout.
+  /// thing added to the layout.
   void addHeapHeader();
-
+  /// Add the NSObject object header to the layout. This must be the first
+  /// thing added to the layout.
+  void addNSObjectHeader();
+
   /// Add a number of fields to the layout.  The field layouts need
   /// only have the TypeInfo set; the rest will be filled out.
   ///

stdlib/public/SwiftShims/KeyPath.h

@@ -55,9 +55,12 @@ static const __swift_uint32_t _SwiftKeyPathComponentHeader_OptionalTag
   = 3;
 
 static const __swift_uint32_t _SwiftKeyPathComponentHeader_MaximumOffsetPayload
+  = 0x1FFFFFFCU;
+
+static const __swift_uint32_t _SwiftKeyPathComponentHeader_UnresolvedIndirectOffsetPayload
   = 0x1FFFFFFDU;
-
-static const __swift_uint32_t _SwiftKeyPathComponentHeader_UnresolvedOffsetPayload
+  
+static const __swift_uint32_t _SwiftKeyPathComponentHeader_UnresolvedFieldOffsetPayload
   = 0x1FFFFFFEU;
 
 static const __swift_uint32_t _SwiftKeyPathComponentHeader_OutOfLineOffsetPayload

stdlib/public/core/KeyPath.swift

@@ -614,8 +614,11 @@ internal struct RawKeyPathComponent {
     static var outOfLineOffsetPayload: UInt32 {
       return _SwiftKeyPathComponentHeader_OutOfLineOffsetPayload
     }
-    static var unresolvedOffsetPayload: UInt32 {
-      return _SwiftKeyPathComponentHeader_UnresolvedOffsetPayload
+    static var unresolvedFieldOffsetPayload: UInt32 {
+      return _SwiftKeyPathComponentHeader_UnresolvedFieldOffsetPayload
+    }
+    static var unresolvedIndirectOffsetPayload: UInt32 {
+      return _SwiftKeyPathComponentHeader_UnresolvedIndirectOffsetPayload
     }
     static var computedMutatingFlag: UInt32 {
       return _SwiftKeyPathComponentHeader_ComputedMutatingFlag
@@ -1613,10 +1616,13 @@ internal func _getKeyPathClassAndInstanceSize(
     let header = buffer.pop(RawKeyPathComponent.Header.self)
 
     func popOffset() {
-      if header.payload == RawKeyPathComponent.Header.unresolvedOffsetPayload
+      if header.payload == RawKeyPathComponent.Header.unresolvedFieldOffsetPayload
         || header.payload == RawKeyPathComponent.Header.outOfLineOffsetPayload {
         _ = buffer.pop(UInt32.self)
       }
+      if header.payload == RawKeyPathComponent.Header.unresolvedIndirectOffsetPayload {
+        _ = buffer.pop(Int.self)
+      }
     }
 
     switch header.kind {
@@ -1712,6 +1718,12 @@ internal func _instantiateKeyPathBuffer(
     count: origDestData.count - MemoryLayout<KeyPathBuffer.Header>.size)
 
   func pushDest<T>(_ value: T) {
+    // TODO: If key path patterns were better optimized to try to be constant-
+    // memory objects, then it might become profitable to try to avoid writes
+    // here in the case when the dest memory contains the value we want to write
+    // here so that we don't dirty memory. (In practice the current
+    // implementation will always dirty the page when the key path is
+    // instantiated.)
     _sanityCheck(_isPOD(T.self))
     var value2 = value
     let size = MemoryLayout<T>.size
@@ -1729,12 +1741,15 @@ internal func _instantiateKeyPathBuffer(
 
   // Instantiate components that need it.
   var base: Any.Type = rootType
+  // Some pattern forms are pessimistically larger than what we need in the
+  // instantiated key path. Keep track of this.
+  var shrinkage = 0
   while true {
     let componentAddr = destData.baseAddress.unsafelyUnwrapped
     let header = patternBuffer.pop(RawKeyPathComponent.Header.self)
 
     func tryToResolveOffset() {
-      if header.payload == RawKeyPathComponent.Header.unresolvedOffsetPayload {
+      if header.payload == RawKeyPathComponent.Header.unresolvedFieldOffsetPayload {
         // Look up offset in type metadata. The value in the pattern is the
         // offset within the metadata object.
         let metadataPtr = unsafeBitCast(base, to: UnsafeRawPointer.self)
@@ -1749,6 +1764,21 @@ internal func _instantiateKeyPathBuffer(
         return
       }
 
+      if header.payload == RawKeyPathComponent.Header.unresolvedIndirectOffsetPayload {
+        // Look up offset in the indirectly-referenced variable we have a
+        // pointer.
+        let offsetVar = patternBuffer.pop(UnsafeRawPointer.self)
+        let offsetValue = UInt32(offsetVar.load(as: UInt.self))
+        // Rewrite the header for a resolved offset.
+        var newHeader = header
+        newHeader.payload = RawKeyPathComponent.Header.outOfLineOffsetPayload
+        pushDest(newHeader)
+        pushDest(offsetValue)
+        shrinkage += MemoryLayout<UnsafeRawPointer>.size
+                   - MemoryLayout<UInt32>.size
+        return
+      }
+
       // Otherwise, just transfer the pre-resolved component.
       pushDest(header)
       if header.payload == RawKeyPathComponent.Header.outOfLineOffsetPayload {
@@ -1846,10 +1876,11 @@ internal func _instantiateKeyPathBuffer(
   }
 
   // We should have traversed both buffers.
-  _sanityCheck(patternBuffer.data.isEmpty && destData.isEmpty)
+  _sanityCheck(patternBuffer.data.isEmpty && destData.count == shrinkage)
 
   // Write out the header.
-  let destHeader = KeyPathBuffer.Header(size: origPatternBuffer.data.count,
+  let destHeader = KeyPathBuffer.Header(
+    size: origPatternBuffer.data.count - shrinkage,
     trivial: true, // TODO: nontrivial indexes
     hasReferencePrefix: endOfReferencePrefixComponent != nil)