[Stdlib] Fix swift_setAtWritableKeyPath to check for ReferenceWritableKeyPaths.

Call through to swift_setAtReferenceWritableKeyPath in that case. This fixes an assertion failure (or worse) when upcasting a ReferenceWritableKeyPath and then using subscript(keyPath:) to write a value with it.

rdar://problem/70609888

Author: mikeash

stdlib/public/core/KeyPath.swift

@@ -1920,6 +1920,12 @@ func _setAtWritableKeyPath<Root, Value>(
   keyPath: WritableKeyPath<Root, Value>,
   value: __owned Value
 ) {
+  if type(of: keyPath).kind == .reference {
+    return _setAtReferenceWritableKeyPath(root: root,
+      keyPath: _unsafeUncheckedDowncast(keyPath,
+        to: ReferenceWritableKeyPath<Root, Value>.self),
+      value: value)
+  }
   // TODO: we should be able to do this more efficiently than projecting.
   let (addr, owner) = keyPath._projectMutableAddress(from: &root)
   addr.pointee = value

test/stdlib/KeyPath.swift

@@ -1028,5 +1028,20 @@ keyPath.test("tail allocated c array") {
   expectEqual(4, offset)
 }
 
+keyPath.test("ReferenceWritableKeyPath statically typed as WritableKeyPath") {
+  let inner = C<Int>(x: 42, y: nil, z: 43)
+  var outer = C<C<Int>>(x: 44, y: nil, z: inner)
+  let keyPath = \C<C<Int>>.z.x
+  let upcastKeyPath = keyPath as WritableKeyPath
+
+  expectEqual(outer[keyPath: keyPath], 42)
+  outer[keyPath: keyPath] = 43
+  expectEqual(outer[keyPath: keyPath], 43)
+
+  expectEqual(outer[keyPath: upcastKeyPath], 43)
+  outer[keyPath: upcastKeyPath] = 44
+  expectEqual(outer[keyPath: upcastKeyPath], 44)
+}
+
 runAllTests()