Add SPI to reroot a keypath for a given superclass

Author: Azoy

stdlib/public/core/KeyPath.swift

@@ -3937,6 +3937,66 @@ public func _createOffsetBasedKeyPath(
   }
 }
 
+@_spi(ObservableRerootKeyPath)
+@available(SwiftStdlib 5.9, *)
+public func _rerootKeyPath<NewRoot>(
+  _ existingKp: AnyKeyPath,
+  to newRoot: NewRoot.Type
+) -> PartialKeyPath<NewRoot> {
+  let (isTrivial, hasReferencePrefix, componentSize) = existingKp.withBuffer {
+    ($0.trivial, $0.hasReferencePrefix, $0.data.count)
+  }
+
+  let existingKpTy = type(of: existingKp)
+
+  func openedRoot<Root>(_: Root.Type) -> AnyKeyPath.Type {
+    func openedValue<Value>(_: Value.Type) -> AnyKeyPath.Type {
+      if existingKpTy == ReferenceWritableKeyPath<Root, Value>.self {
+        return ReferenceWritableKeyPath<NewRoot, Value>.self
+      } else if existingKpTy == KeyPath<Root, Value>.self {
+        return KeyPath<NewRoot, Value>.self
+      } else {
+        fatalError("Unsupported KeyPath type to be rerooted")
+      }
+    }
+
+    return _openExistential(existingKpTy.valueType, do: openedValue(_:))
+  }
+
+  let newKpTy = _openExistential(existingKpTy.rootType, do: openedRoot(_:))
+
+  return newKpTy._create(
+    // This is the buffer header + padding (if needed) + size of components
+    capacityInBytes: MemoryLayout<Int>.size + componentSize
+  ) {
+    var builder = KeyPathBuffer.Builder($0)
+    let header = KeyPathBuffer.Header(
+      size: componentSize,
+      trivial: isTrivial,
+      hasReferencePrefix: hasReferencePrefix
+    )
+
+    builder.pushHeader(header)
+
+    existingKp.withBuffer {
+      var existingBuffer = $0
+
+      while true {
+        let (rawComponent, componentTy) = existingBuffer.next()
+
+        rawComponent.clone(
+          into: &builder.buffer,
+          endOfReferencePrefix: rawComponent.header.endOfReferencePrefix
+        )
+
+        if componentTy == nil {
+          break
+        }
+      }
+    }
+  } as! PartialKeyPath<NewRoot>
+}
+
 #if SWIFT_ENABLE_REFLECTION
 
 @_silgen_name("swift_keyPath_copySymbolName")

test/stdlib/KeyPath.swift

@@ -5,6 +5,9 @@
 // REQUIRES: executable_test
 // UNSUPPORTED: freestanding
 
+@_spi(ObservableRerootKeyPath)
+import Swift
+
 import StdlibUnittest
 
 var keyPath = TestSuite("key paths")
@@ -1086,5 +1089,45 @@ if #available(SwiftStdlib 5.9, *) {
   }
 }
 
+class RerootedSuper {
+  var x = "hello world"
+}
+
+class RerootedSub0: RerootedSuper {}
+class RerootedSub1: RerootedSub0 {}
+
+if #available(SwiftStdlib 5.9, *) {
+  keyPath.test("_rerootKeyPath") {
+    let x = \RerootedSub1.x
+
+    let superValue = RerootedSuper()
+    let sub0 = RerootedSub0()
+    let sub1 = RerootedSub1()
+
+    let sub0Kp = _rerootKeyPath(x, to: RerootedSub0.self)
+
+    expectTrue(type(of: sub0Kp) == ReferenceWritableKeyPath<RerootedSub0, String>.self)
+
+    let superKp = _rerootKeyPath(x, to: RerootedSuper.self)
+
+    expectTrue(type(of: superKp) == ReferenceWritableKeyPath<RerootedSuper, String>.self)
+
+    let x0 = sub1[keyPath: sub0Kp] as! String
+    expectEqual(x0, "hello world")
+
+    let x1 = sub1[keyPath: superKp] as! String
+    expectEqual(x1, "hello world")
+
+    let x2 = sub0[keyPath: sub0Kp] as! String
+    expectEqual(x2, "hello world")
+
+    let x3 = sub0[keyPath: superKp] as! String
+    expectEqual(x3, "hello world")
+
+    let x4 = superValue[keyPath: superKp] as! String
+    expectEqual(x4, "hello world")
+  }
+}
+
 runAllTests()