fmap for Array

NSArray+fmap.m

@@ -0,0 +1,35 @@
+#import "YOLO.ph"
+
+@implementation NSArray (YOLO)
+
+- (NSArray *(^)(id))fmap {
+    return ^(id frock) {
+        NSMethodSignature *sig = YOLOMS(frock);
+        id (^block)(id, NSUInteger) = ^{
+            switch (sig.numberOfArguments){
+                case 2: return ^(id a, NSUInteger b){ return ((id(^)(id))frock)(a); };
+                case 3:
+                    return [sig getArgumentTypeAtIndex:2][0] == '@'
+                        ? ^(id a, NSUInteger b){ return ((id(^)(id, id))frock)(a, @(b)); }
+                        : ^(id a, NSUInteger b){ return ((id(^)(id, NSUInteger))frock)(a, b); };
+                case 4:
+                    return [sig getArgumentTypeAtIndex:2][0] == '@'
+                        ? ^(id a, NSUInteger b){ return ((id(^)(id, id, id))frock)(a, @(b), self); }
+                        : ^(id a, NSUInteger b){ return ((id(^)(id, NSUInteger, id))frock)(a, b, self); };
+                default:
+                    @throw @"Invalid argument count to fmap";
+            }
+        }();
+
+        id mapped[self.count];
+        NSUInteger ii = 0, jj = 0;
+        for (id mappable in self) {
+            id o = block(mappable, ii++);
+            if (o)
+                mapped[jj++] = o;
+        }
+        return [[self.class alloc] initWithObjects:mapped count:jj];
+    };
+}
+
+@end

YOLO.h

@@ -175,6 +175,34 @@
 - (NSArray *(^)(NSArray *(^)(id o)))flatMap;
 #endif
 
+#ifdef YOLO_FMAP
+/**
+ Invokes the given block once for each element of self. Creates a new
+ array containing the values returned by the block.
+
+    id rv = @[@1, @2, @3, @4].fmap(^(NSNumber *n){
+        return @(n.intValue * n.intValue);
+    });
+    // rv => @[@1, @4, @9, @16]
+
+ If the given block returns nil, that element is skipped in the returned
+ array.
+
+ The given block can have up to three parameters, the first is an element
+ in the array, the second that element’s index, and the third the array
+ itself.
+
+ The second parameter can be a primitive (eg. `int`), or an `NSNumber *`:
+
+     @"YOLO".split(@"").fmap(^(NSString *letter, int index){
+         //…
+     });
+
+ @see -map
+*/
+- (NSArray *(^)(id block))fmap;
+#endif
+
 #ifdef YOLO_GROUPBY
 /**
  Groups the collection by result of the given block.

tests

@@ -570,6 +570,32 @@ int main() {
     XCTAssertEqualObjects(aa, bb);
 }
 
+- (void)testArrayFmap {
+    NSArray *aaa = @[@"a", @"b", @"c", @"d"];
+    NSArray *aa = aaa.fmap(^(id o, uint ii, id array){
+        return @[o, @(ii), array];
+    });
+    id bb = @[@[@"a", @0, aaa], @[@"b", @1, aaa], @[@"c", @2, aaa], @[@"d", @3, aaa]];
+    XCTAssertEqualObjects(aa, bb);
+    XCTAssertFalse([aa isKindOfClass:[NSMutableArray class]]);
+
+    aa = @[@"a", @"b", @"c", @"d"];
+    bb = aa.fmap(^(id o){
+        return o;
+    });
+    XCTAssertEqualObjects(aa, bb);
+    
+    bb = aa.fmap(^(id o, id i){
+        return [NSString stringWithFormat:@"%@%@", o, i];
+    });
+    aa = @[@"a0", @"b1", @"c2", @"d3"];
+    XCTAssertEqualObjects(aa, bb);
+
+    bb = ((NSArray *)aa.mutableCopy).fmap(^(id o) { return o; });
+    XCTAssertEqualObjects(aa, bb);
+    XCTAssertTrue([bb isKindOfClass:[NSMutableArray class]]);
+}
+
 - (void)testDictionaryFmap {
     XCTAssertEqualObjects(@{}, (@{}.fmap(^(id obj) { return obj; })));