NSArray+Functional additions

Classes/Utility/Categories/NSArray+Functional.h

@@ -12,8 +12,15 @@
 
 /// Actually more like flatmap, but it seems like the objc way to allow returning nil to omit objects.
 /// So, return nil from the block to omit objects, and return an object to include it in the new array.
-- (NSArray *)flex_mapped:(id(^)(T obj, NSUInteger idx))mapFunc;
-- (NSArray<T> *)flex_filtered:(BOOL(^)(T obj, NSUInteger idx))filterFunc;
-- (void)flex_forEach:(id(^)(T obj, NSUInteger idx))block;
+/// Unlike flatmap, however, this will not flatten arrays of arrays into a single array.
+- (__kindof NSArray *)flex_mapped:(id(^)(T obj, NSUInteger idx))mapFunc;
+/// Like flex_mapped, but expects arrays to be returned, and flattens them into one array.
+- (__kindof NSArray *)flex_flatmapped:(NSArray *(^)(id, NSUInteger idx))block;
+- (instancetype)flex_filtered:(BOOL(^)(T obj, NSUInteger idx))filterFunc;
+- (void)flex_forEach:(void(^)(T obj, NSUInteger idx))block;
+
++ (instancetype)flex_forEachUpTo:(NSUInteger)bound map:(T(^)(NSUInteger))block;
+
+- (instancetype)sortedUsingSelector:(SEL)selector;
 
 @end

Classes/Utility/Categories/NSArray+Functional.m

@@ -8,9 +8,11 @@
 
 #import "NSArray+Functional.h"
 
+#define FLEXArrayClassIsMutable(me) ([[self class] isSubclassOfClass:[NSMutableArray class]])
+
 @implementation NSArray (Functional)
 
-- (instancetype)flex_mapped:(id (^)(id, NSUInteger))mapFunc {
+- (__kindof NSArray *)flex_mapped:(id (^)(id, NSUInteger))mapFunc {
     NSMutableArray *map = [NSMutableArray new];
     [self enumerateObjectsUsingBlock:^(id obj, NSUInteger idx, BOOL *stop) {
         id ret = mapFunc(obj, idx);
@@ -19,7 +21,27 @@
         }
     }];
 
-    return map.copy;
+    if (self.count < 2048 && !FLEXArrayClassIsMutable(self)) {
+        return map.copy;
+    }
+
+    return map;
+}
+
+- (__kindof NSArray *)flex_flatmapped:(NSArray *(^)(id, NSUInteger))block {
+    NSMutableArray *array = [NSMutableArray array];
+    [self enumerateObjectsUsingBlock:^(id obj, NSUInteger idx, BOOL *stop) {
+        NSArray *toAdd = block(obj, idx);
+        if (toAdd) {
+            [array addObjectsFromArray:toAdd];
+        }
+    }];
+
+    if (array.count < 2048 && !FLEXArrayClassIsMutable(self)) {
+        return array.copy;
+    }
+
+    return array;
 }
 
 - (NSArray *)flex_filtered:(BOOL (^)(id, NSUInteger))filterFunc {
@@ -28,10 +50,37 @@
     }];
 }
 
-- (void)flex_forEach:(id (^)(id, NSUInteger))block {
+- (void)flex_forEach:(void(^)(id, NSUInteger))block {
     [self enumerateObjectsUsingBlock:^(id obj, NSUInteger idx, BOOL *stop) {
         block(obj, idx);
     }];
 }
 
++ (__kindof NSArray *)flex_forEachUpTo:(NSUInteger)bound map:(id(^)(NSUInteger))block {
+    NSMutableArray *array = [NSMutableArray new];
+    for (NSUInteger i = 0; i < bound; i++) {
+        id obj = block(i);
+        if (obj) {
+            [array addObject:obj];
+        }
+    }
+
+    // For performance reasons, don't copy large arrays
+    if (bound < 2048 && !FLEXArrayClassIsMutable(self)) {
+        return array.copy;
+    }
+
+    return array;
+}
+
+- (instancetype)sortedUsingSelector:(SEL)selector {
+    if (FLEXArrayClassIsMutable(self)) {
+        NSMutableArray *me = (id)self;
+        [me sortUsingSelector:selector];
+        return me;
+    } else {
+        return [self sortedArrayUsingSelector:selector];
+    }
+}
+
 @end