KeyPaths: Correctly instantiate offsets for final stored properties in NSObject subclasses.

We need to use the ivar offset variables in this case, since the Swift field offset vector doesn't pick up the adjusted offsets from the ObjC runtime. Fixes SR-5036 | rdar://problem/32488871.

Author: jckarter

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

@@ -506,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: {

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)
 

test/IRGen/keypaths_objc.sil

@@ -6,6 +6,7 @@ import Foundation
 
 class C: NSObject {
   dynamic var x: NSString { get }
+  @sil_stored final var stored: Int
   override init()
 }
 
@@ -18,6 +19,11 @@ sil @x_get : $@convention(thin) (@in C) -> @out NSString
 // CHECK-SAME: i32 536870914
 // CHECK-SAME: i8** @"\01L_selector(x)"
 
+// CHECK: [[KEYPATH_B:@keypath.*]] = private global
+// --             0x5ffffffd: class stored property with indirect offset
+// CHECK-SAME: i32 1610612733,
+// CHECK-SAME: @_T013keypaths_objc1CC6storedSivWvd
+
 // CHECK-LABEL: define swiftcc void @objc_only_property()
 sil @objc_only_property : $@convention(thin) () -> () {
 entry:
@@ -35,3 +41,10 @@ sil hidden @_T013keypaths_objc1CCACycfcTo : $@convention(objc_method) (@objc_met
 entry(%0 : $@objc_metatype C.Type):
   unreachable
 }
+
+sil @objc_stored_property : $@convention(thin) () -> () {
+entry:
+  // CHECK: call %swift.refcounted* @swift_getKeyPath({{.*}} [[KEYPATH_B]]
+  %b = keypath $KeyPath<C, Int>, (objc "stored"; root $C; stored_property #C.stored : $Int)
+  unreachable
+}

test/stdlib/KeyPathObjC.swift

@@ -22,12 +22,34 @@ class Foo: NSObject {
   @objc subscript(x: Bar) -> Foo { return self }
 
   dynamic var dynamic: Bar { fatalError() }
+
+  let storedLet = LifetimeTracked(0)
+}
+
+// We just need some non-empty ObjC-defined class here to ensure we get the
+// right offset for a 'let' or final stored property after the ObjC runtime
+// slides offsets
+class MyWeirdFormatter: DateFormatter {
+  let storedLet = LifetimeTracked(1)
 }
 
 class Bar: NSObject {
   @objc var foo: Foo { fatalError() }
 }
 
+var testStoredProperties = TestSuite("stored properties in ObjC subclasses")
+
+testStoredProperties.test("final stored properties in ObjC subclasses") {
+  let fooLet = \Foo.storedLet
+  let formatterLet = \MyWeirdFormatter.storedLet
+
+  let foo = Foo()
+  let formatter = MyWeirdFormatter()
+
+  expectTrue(foo[keyPath: fooLet] === foo.storedLet)
+  expectTrue(formatter[keyPath: formatterLet] === formatter.storedLet)
+}
+
 var testKVCStrings = TestSuite("KVC strings")
 
 testKVCStrings.test("KVC strings") {