Added parameter to be able to provide a custom caching layer

MSCachedAsyncViewDrawing.h

@@ -11,14 +11,28 @@
 typedef void (^MSCachedAsyncViewDrawingDrawBlock)(CGRect frame);
 typedef void (^MSCachedAsyncViewDrawingCompletionBlock)(UIImage *drawnImage);
 
+@protocol MSCachedAsyncViewDrawingCache;
+
 @interface MSCachedAsyncViewDrawing : NSObject
 
 /**
  * You can use the shared instance to have a shared cache.
+ * Uses in-memory caching by default.
  * @note It's perfectly valid to create separate instances of `MSCachedAsyncViewDrawing`, they will just have independent caches.
  */
 + (MSCachedAsyncViewDrawing *)sharedInstance;
 
+/**
+ * Initialize instance using in-memory caching.
+ */
+- (instancetype)init;
+
+/**
+ * Designated initializer.
+ * @param cache: must be thread safe.
+ */
+- (instancetype)initWithCache:(id<MSCachedAsyncViewDrawingCache>)cache;
+
 /**
  * This method will call `drawBlock` _on a background thread_ passing a `CGRect` that you can pass to a `drawRect:` method
  * of a view or a layer.
@@ -49,4 +63,14 @@ typedef void (^MSCachedAsyncViewDrawingCompletionBlock)(UIImage *drawnImage);
                             drawBlock:(MSCachedAsyncViewDrawingDrawBlock)drawBlock;
 
 
+@end
+
+@protocol MSCachedAsyncViewDrawingCache <NSObject>
+
+/**
+ * @return `nil` if the image for `key` is not available.
+ */
+- (UIImage *)imageForKey:(NSString *)key;
+- (void)setImage:(UIImage *)image forKey:(NSString *)key;
+
 @end

MSCachedAsyncViewDrawing.m

@@ -20,9 +20,16 @@
     #define MS_dispatch_queue_t_property_qualifier assign
 #endif
 
+@interface _MSCachedAsyncViewDrawingMemoryImageCache : NSObject <MSCachedAsyncViewDrawingCache>
+{
+    NSCache *_cache;
+}
+
+@end
+
 @interface MSCachedAsyncViewDrawing ()
 
-@property (nonatomic, strong) NSCache *cache;
+@property (nonatomic) id<MSCachedAsyncViewDrawingCache> cache;
 
 @property (nonatomic, MS_dispatch_queue_t_property_qualifier) dispatch_queue_t dispatchQueue;
 
@@ -42,13 +49,19 @@ + (MSCachedAsyncViewDrawing *)sharedInstance
     return sharedInstance;
 }
 
-- (id)init
+- (instancetype)init
+{
+    return [self initWithCache:[_MSCachedAsyncViewDrawingMemoryImageCache new]];
+}
+
+- (instancetype)initWithCache:(id<MSCachedAsyncViewDrawingCache>)cache
 {
+    NSParameterAssert(cache);
+
     if ((self = [super init]))
     {
-        self.cache = [[NSCache alloc] init];
-        self.cache.name = @"com.mindsnacks.view_drawing.cache";
-        self.dispatchQueue = dispatch_queue_create("com.mindsnacks.view_drawing.queue", DISPATCH_QUEUE_CONCURRENT);
+        _cache = cache;
+        _dispatchQueue = dispatch_queue_create("com.mindsnacks.view_drawing.queue", DISPATCH_QUEUE_CONCURRENT);
     }
 
     return self;
@@ -70,7 +83,7 @@ - (void)drawViewWithCacheKey:(NSString *)cacheKey
              completionBlock:(MSCachedAsyncViewDrawingCompletionBlock)completionBlock
                waitUntilDone:(BOOL)waitUntilDone
 {
-    UIImage *cachedImage = [self.cache objectForKey:cacheKey];
+    UIImage *cachedImage = [self.cache imageForKey:cacheKey];
 
     if (cachedImage)
     {
@@ -82,7 +95,7 @@ - (void)drawViewWithCacheKey:(NSString *)cacheKey
     completionBlock = [completionBlock copy];
 
     dispatch_block_t loadImageBlock = ^{
-        BOOL opaque = [self colorIsOpaque:backgroundColor];
+        const BOOL opaque = [self colorIsOpaque:backgroundColor];
 
         UIImage *resultImage = nil;
 
@@ -92,7 +105,7 @@ - (void)drawViewWithCacheKey:(NSString *)cacheKey
 
             CGRect rectToDraw = (CGRect){.origin = CGPointZero, .size = imageSize};
 
-            BOOL shouldDrawBackgroundColor = ![backgroundColor isEqual:[UIColor clearColor]];
+            const BOOL shouldDrawBackgroundColor = ![backgroundColor isEqual:[UIColor clearColor]];
 
             if (shouldDrawBackgroundColor)
             {
@@ -110,7 +123,7 @@ - (void)drawViewWithCacheKey:(NSString *)cacheKey
         }
         UIGraphicsEndImageContext();
 
-        [self.cache setObject:resultImage forKey:cacheKey];
+        [self.cache setImage:resultImage forKey:cacheKey];
 
         if (waitUntilDone)
         {
@@ -186,3 +199,28 @@ - (BOOL)colorIsOpaque:(UIColor *)color
 }
 
 @end
+
+@implementation _MSCachedAsyncViewDrawingMemoryImageCache
+
+- (id)init
+{
+    if ((self = [super init]))
+    {
+        _cache = [NSCache new];
+        _cache.name = @"com.mindsnacks.view_drawing.cache";
+    }
+
+    return self;
+}
+
+- (UIImage *)imageForKey:(NSString *)key
+{
+    return [_cache objectForKey:key];
+}
+
+- (void)setImage:(UIImage *)image forKey:(NSString *)key
+{
+    [_cache setObject:image forKey:key];
+}
+
+@end