Enumerated NonEmpty

Sources/NonEmpty/NonEmpty+Enumerated.swift

@@ -0,0 +1,49 @@
+import Foundation
+
+extension NonEmpty {
+  public func enumerated() -> Enumerated<Collection> {
+    Enumerated(self)
+  }
+}
+
+extension NonEmpty {
+  public struct Enumerated<Collection: Swift.Collection> {
+    private let rawValue: Collection
+
+    init(_ nonEmpty: NonEmpty<Collection>) {
+      rawValue = nonEmpty.rawValue
+    }
+
+    public var first: Iterator.Element {
+      var iterator = makeIterator()
+      return iterator.next()!
+    }
+
+    public func map<T>(_ transform: (Iterator.Element) throws -> T) rethrows -> NonEmpty<[T]> {
+      NonEmpty<[T]>(rawValue: try makeIterator().map(transform))!
+    }
+
+    public func flatMap<SegmentOfResult>(
+      _ transform: (Iterator.Element) throws -> NonEmpty<SegmentOfResult>
+    ) rethrows -> NonEmpty<[SegmentOfResult.Element]> where SegmentOfResult: Sequence {
+      NonEmpty<[SegmentOfResult.Element]>(rawValue: try makeIterator().flatMap(transform))!
+    }
+
+    public func shuffled<T>(using generator: inout T) -> NonEmpty<[Iterator.Element]>
+    where T: RandomNumberGenerator {
+      NonEmpty<[Iterator.Element]>(rawValue: makeIterator().shuffled(using: &generator))!
+    }
+
+    public func shuffled() -> NonEmpty<[Iterator.Element]> {
+      NonEmpty<[Iterator.Element]>(rawValue: makeIterator().shuffled())!
+    }
+  }
+}
+
+extension NonEmpty.Enumerated: Sequence {
+  public typealias Iterator = EnumeratedSequence<Collection>.Iterator
+
+  public func makeIterator() -> Iterator {
+    rawValue.enumerated().makeIterator()
+  }
+}

Tests/NonEmptyTests/NonEmptyTests.swift

@@ -240,6 +240,44 @@ final class NonEmptyTests: XCTestCase {
       XCTAssertEqual(.init("A", "C", "B"), xs)
     }
   #endif
+
+  func testEnumerated() {
+    let actual = NonEmptyArray("1", "2", "3").enumerated()
+    let expected = [(offset: 0, element: "1", (offset: 1, element: "2"), (offset: 2, element: "3"))]
+
+    XCTAssertTrue(
+      zip(actual, expected).allSatisfy {
+        $0.offset == $1.offset && $0.element == $1.element
+      }
+    )
+  }
+
+  func testEnumeratedMap() {
+    let xs = NonEmptyArray(1, 2, 3).enumerated()
+
+    let nonEmptyArray: NonEmpty<[String]> = xs.map { String($0.element) }
+    let array: [String] = xs.map { String($0.element) }
+
+    XCTAssertEqual(nonEmptyArray, NonEmpty("1", "2", "3"))
+    XCTAssertEqual(array, ["1", "2", "3"])
+  }
+
+  func testEnumeratedFlatMap() {
+    let xs = NonEmptyArray(NonEmptyArray("1"), NonEmptyArray("2"), NonEmptyArray("3")).enumerated()
+
+    let nonEmptyArray: NonEmpty<[String]> = xs.flatMap { $0.element }
+    let array: [String] = xs.flatMap { $0.element }
+
+    XCTAssertEqual(nonEmptyArray, NonEmpty("1", "2", "3"))
+    XCTAssertEqual(array, ["1", "2", "3"])
+  }
+
+  func testEnumeratedFirst() {
+    let xs = NonEmptyArray("Blob").enumerated()
+
+    XCTAssertEqual(xs.first.offset, 0)
+    XCTAssertEqual(xs.first.element, "Blob")
+  }
 }
 
 struct TrivialHashable: Equatable, Comparable, Hashable {