fmap for Set

mxcl · Dec 6, 2014 · 7c4c62d · 7c4c62d
1 parent caa9e6e
commit 7c4c62d
Show file tree

Hide file tree

Showing 3 changed files with 44 additions and 0 deletions.
diff --git a/NSSet+fmap.m b/NSSet+fmap.m
@@ -0,0 +1,18 @@
+#import "YOLO.ph"
+
+@implementation NSSet (YOLO)
+
+- (NSSet *(^)(id (^)(id)))fmap {
+    return ^(id (^block)(id)) {
+        id mapped[self.count];
+        NSUInteger ii = 0;
+        for (id mappable in self) {
+            id o = block(mappable);
+            if (o)
+                mapped[ii++] = o;
+        }
+        return [[self.class alloc] initWithObjects:mapped count:ii];
+    };
+}
+
+@end
diff --git a/YOLO.h b/YOLO.h
@@ -770,6 +770,13 @@ clear.
 - (NSSet *(^)(void (^)(id o)))each;
 #endif
 
+#ifdef YOLO_FMAP
+/**
+ @see NSArray's -fmap
+*/
+- (NSSet *(^)(id (^)(id obj)))fmap;
+#endif
+
 @end
 
 

diff --git a/tests b/tests
@@ -596,6 +596,25 @@ int main() {
     XCTAssertTrue([bb isKindOfClass:[NSMutableArray class]]);
 }
 
+- (void)testSetFmap {
+    NSSet *aa = [NSSet setWithObjects:@1, @2, @3, nil];
+
+    XCTAssertEqualObjects((aa.fmap(^(id obj) { return obj; })), aa);
+    XCTAssertEqualObjects((aa.fmap(^(id obj) { return (id)nil; })), [NSSet set]);
+    
+    id bb = aa.fmap(^(id obj) {
+        return @([obj intValue] * [obj intValue]);
+    });
+    XCTAssertEqualObjects(bb, ([NSSet setWithObjects:@1, @4, @9, nil]));
+    XCTAssertFalse([bb isKindOfClass:[NSMutableSet class]]);
+    
+    bb = ((NSSet *)aa.mutableCopy).fmap(^(id obj) {
+        return @([obj intValue] * [obj intValue]);
+    });
+    XCTAssertEqualObjects(bb, ([NSSet setWithObjects:@1, @4, @9, nil]));
+    XCTAssertTrue([bb isKindOfClass:[NSMutableSet class]]);
+}
+
 - (void)testDictionaryFmap {
     XCTAssertEqualObjects(@{}, (@{}.fmap(^(id obj) { return obj; })));