From f527e4b0690146af4427f5d829cc2c16a9c5865d Mon Sep 17 00:00:00 2001
From: Mark Schall <mark.schall@delinea.com>
Date: Wed, 2 Oct 2024 09:16:19 -0400
Subject: [PATCH] adding a new Iterator

---
 Sources/_CryptoExtras/AES/Nonces.swift     | 42 +++++++++++++++-------
 Sources/_CryptoExtras/AES/Nonces.swift.gyb | 30 +++++++++++++---
 2 files changed, 56 insertions(+), 16 deletions(-)

diff --git a/Sources/_CryptoExtras/AES/Nonces.swift b/Sources/_CryptoExtras/AES/Nonces.swift
index eb6a35b4..68b6ab96 100644
--- a/Sources/_CryptoExtras/AES/Nonces.swift
+++ b/Sources/_CryptoExtras/AES/Nonces.swift
@@ -18,6 +18,30 @@ import Foundation
 // any edits of this file WILL be overwritten and thus discarded
 // see section `gyb` in `README` for details.
 
+fileprivate struct ByteIterator<T>: IteratorProtocol {
+    var currentOffset = 0
+    var pointer: UnsafeRawBufferPointer? = nil
+    let length: Int
+    
+    init(_ bytes: T) {
+        self.length = Mirror(reflecting: bytes).children.count
+        withUnsafeBytes(of: bytes) { pointer in
+            self.pointer = pointer
+        }
+    }
+    
+    @inlinable 
+    public mutating func next() -> UInt8? {
+        guard let pointer,
+              currentOffset < length else { return nil }
+        
+        let next = pointer.load(fromByteOffset: currentOffset, as: UInt8.self)
+        currentOffset += 1
+        return next
+    }
+}
+
+
 // MARK: - AES._CBC + IV
 extension AES._CBC {
     /// A value used once during a cryptographic operation and then discarded.
@@ -109,10 +133,8 @@ extension AES._CBC {
         }
         
         /// Returns an iterator over the elements of the nonce.
-        public func makeIterator() -> Array<UInt8>.Iterator {
-            self.withUnsafeBytes({ (buffPtr) in
-                return Array(buffPtr).makeIterator()
-            })
+        public func makeIterator() -> some IteratorProtocol<UInt8> {
+            ByteIterator(bytes)
         }
     }
 }
@@ -202,10 +224,8 @@ extension AES._CFB {
         }
         
         /// Returns an iterator over the elements of the nonce.
-        public func makeIterator() -> Array<UInt8>.Iterator {
-            self.withUnsafeBytes({ (buffPtr) in
-                return Array(buffPtr).makeIterator()
-            })
+        public func makeIterator() -> some IteratorProtocol<UInt8> {
+            ByteIterator(bytes)
         }
     }
 }
@@ -295,10 +315,8 @@ extension AES._CTR {
         }
         
         /// Returns an iterator over the elements of the nonce.
-        public func makeIterator() -> Array<UInt8>.Iterator {
-            self.withUnsafeBytes({ (buffPtr) in
-                return Array(buffPtr).makeIterator()
-            })
+        public func makeIterator() -> some IteratorProtocol<UInt8> {
+            ByteIterator(bytes)
         }
     }
 }
diff --git a/Sources/_CryptoExtras/AES/Nonces.swift.gyb b/Sources/_CryptoExtras/AES/Nonces.swift.gyb
index a3c936b6..ff8695d6 100644
--- a/Sources/_CryptoExtras/AES/Nonces.swift.gyb
+++ b/Sources/_CryptoExtras/AES/Nonces.swift.gyb
@@ -17,6 +17,30 @@ import Foundation
 // MARK: - Generated file, do NOT edit
 // any edits of this file WILL be overwritten and thus discarded
 // see section `gyb` in `README` for details.
+
+fileprivate struct ByteIterator<T>: IteratorProtocol {
+    var currentOffset = 0
+    var pointer: UnsafeRawBufferPointer? = nil
+    let length: Int
+    
+    init(_ bytes: T) {
+        self.length = Mirror(reflecting: bytes).children.count
+        withUnsafeBytes(of: bytes) { pointer in
+            self.pointer = pointer
+        }
+    }
+    
+    @inlinable 
+    public mutating func next() -> UInt8? {
+        guard let pointer,
+              currentOffset < length else { return nil }
+        
+        let next = pointer.load(fromByteOffset: currentOffset, as: UInt8.self)
+        currentOffset += 1
+        return next
+    }
+}
+
 %{
 ciphers = [
     {"name": "AES._CBC", "nonceName": "IV", "tupleDefinition": """(
@@ -123,10 +147,8 @@ extension ${name} {
         }
         
         /// Returns an iterator over the elements of the nonce.
-        public func makeIterator() -> Array<UInt8>.Iterator {
-            self.withUnsafeBytes({ (buffPtr) in
-                return Array(buffPtr).makeIterator()
-            })
+        public func makeIterator() -> some IteratorProtocol<UInt8> {
+            ByteIterator(bytes)
         }
     }
 }