Add subscript methods for range-based access in ArrayExtensions

homerooliveira · Jan 22, 2024 · 6260a23 · 6260a23
1 parent 74acc4f
commit 6260a23
Show file tree

Hide file tree

Showing 2 changed files with 38 additions and 0 deletions.
diff --git a/Sources/SwiftExtensions/ArrayExtensions.swift b/Sources/SwiftExtensions/ArrayExtensions.swift
@@ -16,4 +16,24 @@ extension Array {
             self[index.rawValue] = newValue
         }
     }
+
+    public subscript<T>(_ range: Range<T>) -> ArraySlice<Element> where T: RawRepresentable, T: Comparable, T.RawValue == Index {
+        self[range.lowerBound.rawValue..<range.upperBound.rawValue]
+    }
+
+    public subscript<T>(_ range: ClosedRange<T>) -> ArraySlice<Element> where T: RawRepresentable, T.RawValue == Index, T: Comparable {
+        self[range.lowerBound.rawValue...range.upperBound.rawValue]
+    }
+
+    public subscript<T>(_ range: PartialRangeFrom<T>) -> ArraySlice<Element> where T: RawRepresentable, T.RawValue == Index, T: Comparable {
+        self[range.lowerBound.rawValue...]
+    }
+
+    public subscript<T>(_ range: PartialRangeUpTo<T>) -> ArraySlice<Element> where T: RawRepresentable, T.RawValue == Index, T: Comparable {
+        self[..<range.upperBound.rawValue]
+    }
+
+    public subscript<T>(_ range: PartialRangeThrough<T>) -> ArraySlice<Element> where T: RawRepresentable, T.RawValue == Index, T: Comparable {
+        self[...range.upperBound.rawValue]
+    }
 }
diff --git a/Tests/SwiftExtensionsTests/ArrayExtensionsTests.swift b/Tests/SwiftExtensionsTests/ArrayExtensionsTests.swift
@@ -29,4 +29,22 @@ final class ArrayExtensionsTests: XCTestCase {
 
         XCTAssertEqual(numbers, [12, 13])
     }
+
+    func testSubscriptGetWithRange() throws {
+        enum Index: Int, Comparable {
+            case zero
+            case one
+
+            static func < (lhs: Index, rhs: Index) -> Bool {
+                lhs.rawValue < rhs.rawValue
+            }
+        }
+
+        let numbers = [0, 1, 2, 3, 4, 5]
+
+        XCTAssertEqual(numbers[Index.zero...Index.one], [0, 1])
+        XCTAssertEqual(numbers[Index.zero...], [0, 1, 2, 3, 4, 5])
+        XCTAssertEqual(numbers[...Index.one], [0, 1])
+        XCTAssertEqual(numbers[..<Index.one], [0])
+    }
 }