diff --git a/Sources/NonEmpty/NonEmpty+Enumerated.swift b/Sources/NonEmpty/NonEmpty+Enumerated.swift new file mode 100644 index 0000000..86efae7 --- /dev/null +++ b/Sources/NonEmpty/NonEmpty+Enumerated.swift @@ -0,0 +1,49 @@ +import Foundation + +extension NonEmpty { + public func enumerated() -> Enumerated { + Enumerated(self) + } +} + +extension NonEmpty { + public struct Enumerated { + private let rawValue: Collection + + init(_ nonEmpty: NonEmpty) { + rawValue = nonEmpty.rawValue + } + + public var first: Iterator.Element { + var iterator = makeIterator() + return iterator.next()! + } + + public func map(_ transform: (Iterator.Element) throws -> T) rethrows -> NonEmpty<[T]> { + NonEmpty<[T]>(rawValue: try makeIterator().map(transform))! + } + + public func flatMap( + _ transform: (Iterator.Element) throws -> NonEmpty + ) rethrows -> NonEmpty<[SegmentOfResult.Element]> where SegmentOfResult: Sequence { + NonEmpty<[SegmentOfResult.Element]>(rawValue: try makeIterator().flatMap(transform))! + } + + public func shuffled(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.Iterator + + public func makeIterator() -> Iterator { + rawValue.enumerated().makeIterator() + } +} diff --git a/Tests/NonEmptyTests/NonEmptyTests.swift b/Tests/NonEmptyTests/NonEmptyTests.swift index ff8dbd6..c30167d 100644 --- a/Tests/NonEmptyTests/NonEmptyTests.swift +++ b/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 {