diff --git a/NSSet+fmap.m b/NSSet+fmap.m new file mode 100644 index 0000000..30e34e6 --- /dev/null +++ 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 index 1977bf8..52ae5a9 100644 --- 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 index 3305d89..b3fee6d 100755 --- 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; })));