From 6d73cee9fd6fc4ffc669b8a5525ed3aea204b597 Mon Sep 17 00:00:00 2001
From: Adlai Holler <guy.on.internet@icloud.com>
Date: Tue, 6 Oct 2015 16:35:50 -0700
Subject: [PATCH] Give ASMultiplexImageNodeDataSources an opportunity to
 provide PHAssets quicker

---
 AsyncDisplayKit/ASMultiplexImageNode.h  | 12 ++++++++++++
 AsyncDisplayKit/ASMultiplexImageNode.mm | 22 +++++++++++++---------
 2 files changed, 25 insertions(+), 9 deletions(-)

diff --git a/AsyncDisplayKit/ASMultiplexImageNode.h b/AsyncDisplayKit/ASMultiplexImageNode.h
index 1956760899..302aa20ef5 100644
--- a/AsyncDisplayKit/ASMultiplexImageNode.h
+++ b/AsyncDisplayKit/ASMultiplexImageNode.h
@@ -212,6 +212,18 @@ didFinishDownloadingImageWithIdentifier:(id)imageIdentifier
  */
 - (NSURL *)multiplexImageNode:(ASMultiplexImageNode *)imageNode URLForImageIdentifier:(id)imageIdentifier;
 
+/**
+ * @abstract A PHAsset for the specific asset local identifier
+ * @param imageNode The sender.
+ * @param assetLocalIdentifier The local identifier for a PHAsset that this image node is loading.
+ *
+ * @discussion This optional method can improve image performance if your data source already has the PHAsset available.
+ * If this method is not implemented, or returns nil, the image node will request the asset from the Photos framework.
+ * @note This method may be called from any thread.
+ * @return A PHAsset corresponding to `assetLocalIdentifier`, or nil if none is available.
+ */
+- (PHAsset *)multiplexImageNode:(ASMultiplexImageNode *)imageNode assetForLocalIdentifier:(NSString *)assetLocalIdentifier;
+
 @end
 
 #pragma mark - 
diff --git a/AsyncDisplayKit/ASMultiplexImageNode.mm b/AsyncDisplayKit/ASMultiplexImageNode.mm
index 07945f98a0..081a079524 100644
--- a/AsyncDisplayKit/ASMultiplexImageNode.mm
+++ b/AsyncDisplayKit/ASMultiplexImageNode.mm
@@ -546,19 +546,23 @@ typedef void(^ASMultiplexImageLoadCompletionBlock)(UIImage *image, id imageIdent
   // Each ASMultiplexImageNode can have max 1 inflight Photos image request operation
   [_phImageRequestOperation cancel];
   
-  __weak __typeof(self)weakSelf = self;
+  __weak __typeof(self) weakSelf = self;
   NSOperation *newImageRequestOp = [NSBlockOperation blockOperationWithBlock:^{
-    __strong __typeof(weakSelf)strongSelf = weakSelf;
+    __strong __typeof(weakSelf) strongSelf = weakSelf;
     if (strongSelf == nil) { return; }
     
-    // Get the PHAsset itself.
-    [phRequestLock lock];
-    PHAsset *imageAsset;
-    // -[PHFetchResult dealloc] plays a role in the deadlock mentioned above, so we make sure the PHFetchResult is deallocated inside the critical section
-    @autoreleasepool {
-      imageAsset = [PHAsset fetchAssetsWithLocalIdentifiers:@[request.assetIdentifier] options:nil].firstObject;
+    // Try to get the asset immediately from the data source.
+    PHAsset *imageAsset = [strongSelf.dataSource multiplexImageNode:strongSelf assetForLocalIdentifier:request.assetIdentifier];
+    
+    // Fall back to locking and getting the PHAsset.
+    if (imageAsset == nil) {
+      [phRequestLock lock];
+      // -[PHFetchResult dealloc] plays a role in the deadlock mentioned above, so we make sure the PHFetchResult is deallocated inside the critical section
+      @autoreleasepool {
+        imageAsset = [PHAsset fetchAssetsWithLocalIdentifiers:@[request.assetIdentifier] options:nil].firstObject;
+      }
+      [phRequestLock unlock];
     }
-    [phRequestLock unlock];
     
     if (imageAsset == nil) {
       // Error.