[stdlib] AutoreleasingUnsafeMutablePointer: eliminate questionable pointer use

stdlib/private/StdlibUnittest/StdlibUnittest.swift

@@ -154,6 +154,14 @@ public func expectFailure(
   _anyExpectFailed.orAndFetch(startAnyExpectFailed)
 }
 
+/// An opaque function that ignores its argument and returns nothing.
+public func noop<T>(_ value: T) {}
+
+/// An opaque function that simply returns its argument.
+public func identity<T>(_ value: T) -> T {
+  return value
+}
+
 public func identity(_ element: OpaqueValue<Int>) -> OpaqueValue<Int> {
   return element
 }

stdlib/public/core/BridgeObjectiveC.swift

@@ -336,7 +336,10 @@ public func _getBridgedNonVerbatimObjectiveCType<T>(_: T.Type) -> Any.Type?
 
 // -- Pointer argument bridging
 
-/// A mutable pointer-to-ObjC-pointer argument.
+/// A mutable pointer addressing an Objective-C reference that doesn't own its
+/// target.
+///
+/// `Pointee` must be a class type or `Optional<C>` where `C` is a class.
 ///
 /// This type has implicit conversions to allow passing any of the following
 /// to a C or ObjC API:
@@ -368,42 +371,50 @@ public struct AutoreleasingUnsafeMutablePointer<Pointee /* TODO : class */>
     self._rawValue = _rawValue
   }
 
-  /// Access the `Pointee` instance referenced by `self`.
+  /// Retrieve or set the `Pointee` instance referenced by `self`.
+  ///
+  /// `AutoreleasingUnsafeMutablePointer` is assumed to reference a value with
+  /// `__autoreleasing` ownership semantics, like `NSFoo **` declarations in
+  /// ARC. Setting the pointee autoreleases the new value before trivially
+  /// storing it in the referenced memory.
   ///
   /// - Precondition: the pointee has been initialized with an instance of type
   ///   `Pointee`.
   @inlinable
   public var pointee: Pointee {
-    /// Retrieve the value the pointer points to.
     @_transparent get {
-      // We can do a strong load normally.
-      return UnsafePointer(self).pointee
+      // The memory addressed by this pointer contains a non-owning reference,
+      // therefore we *must not* point an `UnsafePointer<AnyObject>` to
+      // it---otherwise we would allow the compiler to assume it has a +1
+      // refcount, enabling some optimizations that wouldn't be valid.
+      //
+      // Instead, we need to load the pointee as a +0 unmanaged reference. For
+      // an extra twist, `Pointee` is allowed (but not required) to be an
+      // optional type, so we actually need to load it as an optional, and
+      // explicitly handle the nil case.
+      let unmanaged =
+        UnsafePointer<Optional<Unmanaged<AnyObject>>>(_rawValue).pointee
+      return _unsafeReferenceCast(
+        unmanaged?.takeUnretainedValue(),
+        to: Pointee.self)
     }
-    /// Set the value the pointer points to, copying over the previous value.
-    ///
-    /// AutoreleasingUnsafeMutablePointers are assumed to reference a
-    /// value with __autoreleasing ownership semantics, like 'NSFoo**'
-    /// in ARC. This autoreleases the argument before trivially
-    /// storing it to the referenced memory.
+
     @_transparent nonmutating set {
       // Autorelease the object reference.
-      typealias OptionalAnyObject = AnyObject?
-      let newAnyObject = unsafeBitCast(newValue, to: OptionalAnyObject.self)
-      Builtin.retain(newAnyObject)
-      Builtin.autorelease(newAnyObject)
-      // Trivially assign it as an OpaquePointer; the pointer references an
-      // autoreleasing slot, so retains/releases of the original value are
-      // unneeded.
-      typealias OptionalUnmanaged = Unmanaged<AnyObject>?
-      UnsafeMutablePointer<Pointee>(_rawValue).withMemoryRebound(
-        to: OptionalUnmanaged.self, capacity: 1) {
-        if let newAnyObject = newAnyObject {
-          $0.pointee = Unmanaged.passUnretained(newAnyObject)
-        }
-        else {
-          $0.pointee = nil
-        }
+      let object = _unsafeReferenceCast(newValue, to: Optional<AnyObject>.self)
+      Builtin.retain(object)
+      Builtin.autorelease(object)
+
+      // Convert it to an unmanaged reference and trivially assign it to the
+      // memory addressed by this pointer.
+      let unmanaged: Optional<Unmanaged<AnyObject>>
+      if let object = object {
+        unmanaged = Unmanaged.passUnretained(object)
+      } else {
+        unmanaged = nil
       }
+      UnsafeMutablePointer<Optional<Unmanaged<AnyObject>>>(_rawValue).pointee =
+        unmanaged
     }
   }
 
@@ -415,8 +426,7 @@ public struct AutoreleasingUnsafeMutablePointer<Pointee /* TODO : class */>
   public subscript(i: Int) -> Pointee {
     @_transparent
     get {
-      // We can do a strong load normally.
-      return (UnsafePointer<Pointee>(self) + i).pointee
+      return self.advanced(by: i).pointee
     }
   }
 

stdlib/public/core/Unmanaged.swift

@@ -89,7 +89,7 @@ public struct Unmanaged<Instance: AnyObject> {
   /// and you know that you're not responsible for releasing the result.
   ///
   /// - Returns: The object referenced by this `Unmanaged` instance.
-  @inlinable // unsafe-performance
+  @_transparent // unsafe-performance
   public func takeUnretainedValue() -> Instance {
     return _value
   }
@@ -101,7 +101,7 @@ public struct Unmanaged<Instance: AnyObject> {
   /// and you know that you're responsible for releasing the result.
   ///
   /// - Returns: The object referenced by this `Unmanaged` instance.
-  @inlinable // unsafe-performance
+  @_transparent // unsafe-performance
   public func takeRetainedValue() -> Instance {
     let result = _value
     release()

test/stdlib/AutoreleasingUnsafeMutablePointer.swift

@@ -0,0 +1,186 @@
+// RUN: %target-run-simple-swift
+// REQUIRES: executable_test
+// REQUIRES: objc_interop
+
+import StdlibUnittest
+import ObjectiveC // for autoreleasepool
+
+var suite = TestSuite("AutoreleasingUnsafeMutablePointer")
+
+/// Call `body` passing it an AutoreleasingUnsafeMutablePointer whose pointee
+/// has the specified value, allocated in a temporary buffer.
+func withAUMP<Pointee: AnyObject>(
+  as type: Pointee.Type = Pointee.self,
+  initialValue: Optional<Unmanaged<Pointee>> = nil,
+  _ body: (AutoreleasingUnsafeMutablePointer<Optional<Pointee>>) -> Void
+) {
+  let p =
+    UnsafeMutablePointer<Optional<Unmanaged<Pointee>>>.allocate(capacity: 1)
+  p.initialize(to: initialValue)
+  body(AutoreleasingUnsafeMutablePointer(p))
+  p.deallocate()
+}
+
+/// Call `body` passing it an AutoreleasingUnsafeMutablePointer whose pointee
+/// has the specified value, allocated in a temporary buffer.
+func withAUMP<Pointee: AnyObject>(
+  initialValues: [Unmanaged<Pointee>?],
+  _ body: (AutoreleasingUnsafeMutablePointer<Optional<Pointee>>) -> Void
+) {
+  let p = UnsafeMutablePointer<Optional<Unmanaged<Pointee>>>.allocate(
+    capacity: initialValues.count
+  )
+  for i in 0 ..< initialValues.count {
+    (p + i).initialize(to: initialValues[i])
+  }
+
+  let aump = AutoreleasingUnsafeMutablePointer<Optional<Pointee>>(p)
+  body(aump)
+  p.deallocate()
+}
+
+suite.test("helper calls its closure exactly once") {
+  // `withAUMP` is expected to call its closure exactly once, with an argument
+  // pointing to a nil value.
+  var runCount = 0
+  withAUMP(as: LifetimeTracked.self) { p in
+    runCount += 1
+    expectNil(p.pointee)
+  }
+  expectEqual(runCount, 1)
+}
+
+suite.test("init doesn't autorelease") {
+  let originalInstances = LifetimeTracked.instances
+  let unmanaged = Unmanaged.passRetained(LifetimeTracked(42))
+  expectEqual(LifetimeTracked.instances, originalInstances + 1)
+
+  withAUMP(initialValue: unmanaged) { p in noop(p) }
+
+  // Releasing the original reference should immediately deallocate the
+  // object.
+  expectEqual(LifetimeTracked.instances, originalInstances + 1)
+  unmanaged.release()
+  expectEqual(LifetimeTracked.instances, originalInstances)
+}
+
+suite.test("getter initially returns the initial value") {
+  withAUMP(as: LifetimeTracked.self) { p in
+    expectNil(p.pointee)
+  }
+
+  let unmanaged = Unmanaged.passRetained(LifetimeTracked(42))
+  withAUMP(initialValue: unmanaged) { p in
+    expectTrue(p.pointee === unmanaged.takeUnretainedValue())
+  }
+  unmanaged.release()
+}
+
+suite.test("pointee getter returns the last value set") {
+  autoreleasepool {
+    withAUMP(as: LifetimeTracked.self) { p in
+      expectNil(p.pointee)
+      let object = LifetimeTracked(23)
+      p.pointee = object
+      expectTrue(p.pointee === object)
+      let other = LifetimeTracked(42)
+      p.pointee = other
+      expectTrue(p.pointee === other)
+      p.pointee = nil
+      expectNil(p.pointee)
+    }
+  }
+}
+
+suite.test("pointee getter doesn't autorelease") {
+  let originalInstances = LifetimeTracked.instances
+  autoreleasepool {
+    let unmanaged = Unmanaged.passRetained(LifetimeTracked(42))
+    expectEqual(LifetimeTracked.instances, originalInstances + 1)
+
+    withAUMP(initialValue: unmanaged) { p in
+      expectTrue(p.pointee === unmanaged.takeUnretainedValue())
+    }
+    // Releasing the original reference should immediately deallocate the
+    // object.
+    expectEqual(LifetimeTracked.instances, originalInstances + 1)
+    unmanaged.release()
+    expectEqual(LifetimeTracked.instances, originalInstances)
+  }
+}
+
+suite.test("pointee setter autoreleases the new value") {
+  let originalInstances = LifetimeTracked.instances
+  autoreleasepool {
+    withAUMP(as: LifetimeTracked.self) { p in
+      let object = LifetimeTracked(42)
+      // There is one more instance now.
+      expectEqual(LifetimeTracked.instances, originalInstances + 1)
+      p.pointee = object
+    }
+    // The strong reference to `object` is gone, but the autoreleasepool
+    // is still holding onto it.
+    expectEqual(LifetimeTracked.instances, originalInstances + 1)
+  }
+  // Draining the pool deallocates the instance.
+  expectEqual(LifetimeTracked.instances, originalInstances)
+}
+
+suite.test("AutoreleasingUnsafeMutablePointer doesn't retain its value") {
+  let originalInstances = LifetimeTracked.instances
+  withAUMP(as: LifetimeTracked.self) { p in
+    autoreleasepool {
+      let object = LifetimeTracked(42)
+      p.pointee = object
+      expectEqual(LifetimeTracked.instances, originalInstances + 1)
+    }
+    // Draining the pool deallocates the instance, even though
+    // the autoreleasing pointer still points to it.
+    // This is okay as long as the pointer isn't dereferenced.
+    expectEqual(LifetimeTracked.instances, originalInstances)
+  }
+  expectEqual(LifetimeTracked.instances, originalInstances)
+}
+
+suite.test("subscript[0] is the same as pointee.getter") {
+  withAUMP(as: LifetimeTracked.self) { p in
+    expectNil(p[0])
+  }
+
+  let unmanaged = Unmanaged.passRetained(LifetimeTracked(42))
+  withAUMP(initialValue: unmanaged) { p in
+    expectTrue(p[0] === unmanaged.takeUnretainedValue())
+  }
+  unmanaged.release()
+}
+
+
+suite.test("subscript and advanced(by:) works") {
+  let originalInstances = LifetimeTracked.instances
+  let count = 100
+
+  var refs = 0
+  let values: [Unmanaged<LifetimeTracked>?] = (0 ..< count).map { i in
+    if i % 10 == 0 {
+      refs += 1
+      return Unmanaged.passRetained(LifetimeTracked(i))
+    } else {
+      return nil
+    }
+  }
+
+  withAUMP(initialValues: values) { p in
+    for i in 0 ..< count {
+      expectTrue(p[i] === values[i]?.takeUnretainedValue())
+      expectTrue(p.advanced(by: i).pointee === values[i]?.takeUnretainedValue())
+    }
+  }
+
+  expectEqual(LifetimeTracked.instances, originalInstances + refs)
+  for unmanaged in values {
+    unmanaged?.release()
+  }
+  expectEqual(LifetimeTracked.instances, originalInstances)
+}
+
+runAllTests()