select appropriate class

mxcl · Dec 17, 2014 · 8fbd9ff · 8fbd9ff
1 parent 420dc5b
commit 8fbd9ff
Showing 1 changed file with 5 additions and 1 deletion.
diff --git a/NSArray+map.m b/NSArray+map.m
@@ -28,7 +28,11 @@ @implementation NSArray (YOLO)
             if (o)
                 mapped[jj++] = o;
         }
-        return [[self.class alloc] initWithObjects:mapped count:jj];
+
+        Class immutableClass = [NSArray class];
+        Class mutableClass = [NSMutableArray class];
+        Class cls = [self isKindOfClass:mutableClass] ? mutableClass : immutableClass;
+        return [[cls alloc] initWithObjects:mapped count:jj];
     };
 }