From 0f367b91b895068ce27b1061c4f807292c498266 Mon Sep 17 00:00:00 2001
From: William Taylor <git@will.au>
Date: Wed, 14 May 2025 22:46:04 +1000
Subject: [PATCH] Fix AsyncBufferSequence leaking base sequence

---
 .../Buffer/AsyncBufferSequence.swift          | 48 ++++++++++++++-----
 .../Buffer/BoundedBufferStorage.swift         |  4 +-
 .../Buffer/UnboundedBufferStorage.swift       |  4 +-
 3 files changed, 41 insertions(+), 15 deletions(-)

diff --git a/Sources/AsyncAlgorithms/Buffer/AsyncBufferSequence.swift b/Sources/AsyncAlgorithms/Buffer/AsyncBufferSequence.swift
index 9a0794cc..1a053da2 100644
--- a/Sources/AsyncAlgorithms/Buffer/AsyncBufferSequence.swift
+++ b/Sources/AsyncAlgorithms/Buffer/AsyncBufferSequence.swift
@@ -112,19 +112,45 @@ public struct AsyncBufferSequence<Base: AsyncSequence & Sendable>: AsyncSequence
   }
 
   public struct Iterator: AsyncIteratorProtocol {
-    var storageType: StorageType
+    final class InternalClass {
+      private var storageType: StorageType
 
-    public mutating func next() async rethrows -> Element? {
-      switch self.storageType {
-      case .transparent(var iterator):
-        let element = try await iterator.next()
-        self.storageType = .transparent(iterator)
-        return element
-      case .bounded(let storage):
-        return try await storage.next()?._rethrowGet()
-      case .unbounded(let storage):
-        return try await storage.next()?._rethrowGet()
+      fileprivate init(storageType: StorageType) {
+        self.storageType = storageType
+      }
+
+      deinit {
+        switch self.storageType {
+        case .transparent: break
+        case .bounded(let storage):
+          storage.iteratorDeinitialized()
+        case .unbounded(let storage):
+          storage.iteratorDeinitialized()
+        }
       }
+
+      public func next() async rethrows -> Element? {
+        switch self.storageType {
+        case .transparent(var iterator):
+          let element = try await iterator.next()
+          self.storageType = .transparent(iterator)
+          return element
+        case .bounded(let storage):
+          return try await storage.next()?._rethrowGet()
+        case .unbounded(let storage):
+          return try await storage.next()?._rethrowGet()
+        }
+      }
+    }
+
+    let internalClass: InternalClass
+
+    fileprivate init(storageType: StorageType) {
+      internalClass = InternalClass(storageType: storageType)
+    }
+
+    public mutating func next() async rethrows -> Element? {
+      try await internalClass.next()
     }
   }
 }
diff --git a/Sources/AsyncAlgorithms/Buffer/BoundedBufferStorage.swift b/Sources/AsyncAlgorithms/Buffer/BoundedBufferStorage.swift
index ce89cd5d..f706f4e8 100644
--- a/Sources/AsyncAlgorithms/Buffer/BoundedBufferStorage.swift
+++ b/Sources/AsyncAlgorithms/Buffer/BoundedBufferStorage.swift
@@ -10,7 +10,7 @@
 //===----------------------------------------------------------------------===//
 
 @available(AsyncAlgorithms 1.0, *)
-final class BoundedBufferStorage<Base: AsyncSequence>: Sendable where Base: Sendable {
+struct BoundedBufferStorage<Base: AsyncSequence>: Sendable where Base: Sendable {
   private let stateMachine: ManagedCriticalState<BoundedBufferStateMachine<Base>>
 
   init(base: Base, limit: Int) {
@@ -152,7 +152,7 @@ final class BoundedBufferStorage<Base: AsyncSequence>: Sendable where Base: Send
     }
   }
 
-  deinit {
+  func iteratorDeinitialized() {
     self.interrupted()
   }
 }
diff --git a/Sources/AsyncAlgorithms/Buffer/UnboundedBufferStorage.swift b/Sources/AsyncAlgorithms/Buffer/UnboundedBufferStorage.swift
index 219b5f50..17a0ce63 100644
--- a/Sources/AsyncAlgorithms/Buffer/UnboundedBufferStorage.swift
+++ b/Sources/AsyncAlgorithms/Buffer/UnboundedBufferStorage.swift
@@ -10,7 +10,7 @@
 //===----------------------------------------------------------------------===//
 
 @available(AsyncAlgorithms 1.0, *)
-final class UnboundedBufferStorage<Base: AsyncSequence>: Sendable where Base: Sendable {
+struct UnboundedBufferStorage<Base: AsyncSequence>: Sendable where Base: Sendable {
   private let stateMachine: ManagedCriticalState<UnboundedBufferStateMachine<Base>>
 
   init(base: Base, policy: UnboundedBufferStateMachine<Base>.Policy) {
@@ -120,7 +120,7 @@ final class UnboundedBufferStorage<Base: AsyncSequence>: Sendable where Base: Se
     }
   }
 
-  deinit {
+  func iteratorDeinitialized() {
     self.interrupted()
   }
 }