From e713779569b434b697c7c32cd45030fdb5f312b9 Mon Sep 17 00:00:00 2001
From: fumoboy007 <2100868+fumoboy007@users.noreply.github.com>
Date: Sat, 22 Apr 2023 22:07:40 -0700
Subject: [PATCH] Add `NSLock.withLock(_:)` method.

Fixes #4735.

Xcode 14 introduced the `NSLocking.withLock(_:)` method to the Darwin implementation of Foundation. This commit adds the new method to the open-source implementation.

(Note: The [documentation](https://developer.apple.com/documentation/foundation/nslocking/4059821-withlock) incorrectly specifies the availability of the method as iOS 16+, etc. but it is actually available on all OS versions because it is marked [`@_alwaysEmitIntoClient`](https://github.com/apple/swift/blob/main/docs/ReferenceGuides/UnderscoredAttributes.md#_alwaysemitintoclient). The documentation bug is tracked by [FB11968310](https://feedbackassistant.apple.com/feedback/11968310).)

I copied the implementation from `Xcode.app/Contents/Developer/Platforms/iPhoneOS.platform/Developer/SDKs/iPhoneOS.sdk/System/Library/Frameworks/Foundation.framework/Modules/Foundation.swiftmodule/arm64e-apple-ios.swiftinterface`.
---
 Sources/Foundation/NSLock.swift         | 13 +++++++++++
 Tests/Foundation/Tests/TestNSLock.swift | 30 +++++++++++++++++++++++++
 2 files changed, 43 insertions(+)

diff --git a/Sources/Foundation/NSLock.swift b/Sources/Foundation/NSLock.swift
index 59d334f1ad..f4384b70ac 100644
--- a/Sources/Foundation/NSLock.swift
+++ b/Sources/Foundation/NSLock.swift
@@ -22,6 +22,19 @@ public protocol NSLocking {
     func unlock()
 }
 
+extension NSLocking {
+    @_alwaysEmitIntoClient
+    @_disfavoredOverload
+    public func withLock<R>(_ body: () throws -> R) rethrows -> R {
+        self.lock()
+        defer {
+            self.unlock()
+        }
+
+        return try body()
+    }
+}
+
 #if os(Windows)
 private typealias _MutexPointer = UnsafeMutablePointer<SRWLOCK>
 private typealias _RecursiveMutexPointer = UnsafeMutablePointer<CRITICAL_SECTION>
diff --git a/Tests/Foundation/Tests/TestNSLock.swift b/Tests/Foundation/Tests/TestNSLock.swift
index 3e40a8a38b..58a1b00060 100644
--- a/Tests/Foundation/Tests/TestNSLock.swift
+++ b/Tests/Foundation/Tests/TestNSLock.swift
@@ -14,6 +14,7 @@ class TestNSLock: XCTestCase {
             ("test_lockWait", test_lockWait),
             ("test_threadsAndLocks", test_threadsAndLocks),
             ("test_recursiveLock", test_recursiveLock),
+            ("test_withLock", test_withLock),
             
         ]
     }
@@ -187,4 +188,33 @@ class TestNSLock: XCTestCase {
         
         threadCompletedCondition.unlock()
     }
+
+    func test_withLock() {
+        let lock = NSLock()
+
+        var counter = 0
+        let counterIncrementPerThread = 10_000
+
+        let threadCount = 10
+
+        let threadCompletedExpectation = expectation(description: "Expected threads to complete.")
+        threadCompletedExpectation.expectedFulfillmentCount = threadCount
+
+        for _ in 0..<threadCount {
+            let thread = Thread {
+                for _ in 0..<counterIncrementPerThread {
+                    lock.withLock {
+                        counter += 1
+                    }
+                }
+
+                threadCompletedExpectation.fulfill()
+            }
+            thread.start()
+        }
+
+        wait(for: [threadCompletedExpectation], timeout: 10)
+
+        XCTAssertEqual(counter, counterIncrementPerThread * threadCount)
+    }
 }