Merge pull request swiftlang#19421 from weissi/jw-managed-buffer-realloc

implement ManagerBuffer.reallocated to allow realloc'ing the storage

Author: weissi

stdlib/public/core/Builtin.swift

@@ -658,6 +658,9 @@ internal func _isUnique<T>(_ object: inout T) -> Bool {
   return Bool(Builtin.isUnique(&object))
 }
 
+@_silgen_name("_swift_reallocObject")
+internal func _reallocObject(_ object: UnsafeMutableRawPointer, _ newSizeInBytes: Int) -> UnsafeMutableRawPointer?
+
 /// Returns `true` if `object` is uniquely referenced.
 /// This provides sanity checks on top of the Builtin.
 @_transparent

stdlib/public/core/ManagedBuffer.swift

@@ -136,6 +136,30 @@ extension ManagedBuffer {
   }
 }
 
+@inline(never)
+public func tryReallocateUniquelyReferenced<Header, Element, Buffer: ManagedBuffer<Header, Element>>(
+  buffer: inout Buffer,
+  newMinimumCapacity: Int
+) -> Bool {
+  precondition(_isBitwiseTakable(Header.self))
+  precondition(_isBitwiseTakable(Element.self))
+  precondition(isKnownUniquelyReferenced(&buffer))
+
+  let newSizeInBytes = MemoryLayout<Header>.stride
+    + newMinimumCapacity * MemoryLayout<Element>.stride
+
+  return withUnsafeMutablePointer(to: &buffer) {
+    $0.withMemoryRebound(to: UnsafeMutableRawPointer.self, capacity: 1) {
+      if let reallocdObject = _reallocObject($0.pointee, newSizeInBytes) {
+        $0.pointee = reallocdObject
+        return true
+      } else {
+        return false
+      }
+    }
+  }
+}
+
 /// Contains a buffer object, and provides access to an instance of
 /// `Header` and contiguous storage for an arbitrary number of
 /// `Element` instances stored in that buffer.

stdlib/public/runtime/SwiftObject.h

@@ -88,4 +88,11 @@ NSString *getDescription(OpaqueValue *value, const Metadata *type);
 
 #endif
 
+namespace swift {
+
+SWIFT_CC(swift) SWIFT_RUNTIME_STDLIB_SPI
+HeapObject *_swift_reallocObject(HeapObject *obj, size_t size);
+
+}
+
 #endif

stdlib/public/runtime/SwiftObject.mm

@@ -82,6 +82,24 @@ static uintptr_t computeISAMask() {
 #endif
 }
 
+bool isObjCPinned(HeapObject *obj) {
+  #if SWIFT_OBJC_INTEROP
+    /* future: implement checking the relevant objc runtime bits */
+    return true;
+  #else
+    return false;
+  #endif
+}
+
+// returns non-null if realloc was successful
+SWIFT_CC(swift) SWIFT_RUNTIME_STDLIB_SPI
+HeapObject *swift::_swift_reallocObject(HeapObject *obj, size_t size) {
+ if (isObjCPinned(obj) || obj->refCounts.hasSideTable()) {
+   return nullptr;
+ }
+ return (HeapObject *)realloc(obj, size);
+}
+
 #if SWIFT_OBJC_INTEROP
 
 /// \brief Replacement for ObjC object_isClass(), which is unavailable on

test/stdlib/ManagedBuffer.swift

@@ -40,7 +40,7 @@ extension ManagedBufferPointer
 
 struct CountAndCapacity {
   var count: LifetimeTracked
-  let capacity: Int
+  var capacity: Int
 }
 
 // An example of ManagedBuffer, very similar to what Array will use.
@@ -96,6 +96,23 @@ final class TestManagedBuffer<T> : ManagedBuffer<CountAndCapacity, T> {
     }
     self.count = count + 2
   }
+
+  class func tryGrow(_ buffer: inout TestManagedBuffer<T>, newCapacity: Int) -> Bool {
+    guard isKnownUniquelyReferenced(&buffer) else {
+      return false
+    }
+    guard newCapacity > buffer.capacity else {
+      return false
+    }
+
+    if tryReallocateUniquelyReferenced(buffer: &buffer,
+                                       newMinimumCapacity: newCapacity) {
+      buffer.header.capacity = newCapacity
+      return true
+    } else {
+      return false
+    }
+  }
 }
 
 class MyBuffer<T> {
@@ -241,4 +258,35 @@ tests.test("isKnownUniquelyReferenced") {
   _fixLifetime(s2)
 }
 
+tests.test("canGrowUsingRealloc") {
+  #if os(macOS) || os(iOS) || os(tvOS) || os(watchOS)
+  return // realloc currently unsupported on Darwin
+  #endif
+  func testGrow(_ buffer: inout TestManagedBuffer<LifetimeTracked>,
+                newCapacity: Int,
+                shouldSucceed: Bool = true) {
+    let s = TestManagedBuffer.tryGrow(&buffer, newCapacity: newCapacity)
+    expectEqual(s, shouldSucceed)
+    if shouldSucceed {
+      expectLE(newCapacity, buffer.myCapacity)
+      expectGE((newCapacity + 1) * 2, buffer.myCapacity)
+    }
+    repeat {
+      buffer.append(LifetimeTracked(buffer.count))
+    } while buffer.count < buffer.myCapacity - 2
+  }
+
+  var b = TestManagedBuffer<LifetimeTracked>.create(0)
+  // allow some over-allocation
+  expectLE(0, b.myCapacity)
+  expectGE(3, b.myCapacity)
+
+  testGrow(&b, newCapacity: 5)
+  testGrow(&b, newCapacity: 8)
+  testGrow(&b, newCapacity: 1000)
+  testGrow(&b, newCapacity: 16000)
+  var b2 = b
+  testGrow(&b, newCapacity: 2000, shouldSucceed: false)
+}
+
 runAllTests()