support cancelation in ASGraphicsCreateImage

Source/Details/ASGraphicsContext.h

@@ -30,6 +30,23 @@ NS_ASSUME_NONNULL_BEGIN
  */
 AS_EXTERN UIImage *ASGraphicsCreateImageWithOptions(CGSize size, BOOL opaque, CGFloat scale, UIImage * _Nullable sourceImage, asdisplaynode_iscancelled_block_t NS_NOESCAPE _Nullable isCancelled, void (NS_NOESCAPE ^work)(void)) ASDISPLAYNODE_DEPRECATED_MSG("Use ASGraphicsCreateImageWithTraitCollectionAndOptions instead");
 
+/**
+* A wrapper for the UIKit drawing APIs. If you are in ASExperimentalDrawingGlobal, and you have iOS >= 10, we will create
+* a UIGraphicsRenderer with an appropriate format. Otherwise, we will use UIGraphicsBeginImageContext et al.
+*
+* @param traitCollection Trait collection. The `work` block will be executed with this trait collection, so it will affect dynamic colors, etc.
+* @param size The size of the context.
+* @param opaque Whether the context should be opaque or not.
+* @param scale The scale of the context. 0 uses main screen scale.
+* @param sourceImage If you are planning to render a UIImage into this context, provide it here and we will use its
+*   preferred renderer format if we are using UIGraphicsImageRenderer.
+* @param isCancelled An optional block for canceling the drawing before forming the image.
+* @param work A block, wherein the current UIGraphics context is set based on the arguments.
+*
+* @return The rendered image. You can also render intermediary images using UIGraphicsGetImageFromCurrentImageContext.
+*/
+AS_EXTERN UIImage *ASGraphicsCreateImage(ASPrimitiveTraitCollection traitCollection, CGSize size, BOOL opaque, CGFloat scale, UIImage * _Nullable sourceImage, asdisplaynode_iscancelled_block_t _Nullable NS_NOESCAPE isCancelled, void (NS_NOESCAPE ^work)(void));
+
 /**
 * A wrapper for the UIKit drawing APIs.
 *

Source/Details/ASGraphicsContext.mm

@@ -14,7 +14,7 @@
 
 
 #if AS_AT_LEAST_IOS13
-#define PERFORM_WORK_WITH_TRAIT_COLLECTION(work, traitCollection) \
+#define ASPerformBlockWithTraitCollection(work, traitCollection) \
     if (@available(iOS 13.0, *)) { \
       UITraitCollection *uiTraitCollection = ASPrimitiveTraitCollectionToUITraitCollection(traitCollection); \
       [uiTraitCollection performAsCurrentTraitCollection:^{ \
@@ -24,7 +24,7 @@
       work(); \
     }
 #else
-#define PERFORM_WORK_WITH_TRAIT_COLLECTION(work, traitCollection) work();
+#define ASPerformBlockWithTraitCollection(work, traitCollection) work();
 #endif
 
 
@@ -43,10 +43,10 @@ NS_INLINE void ASConfigureExtendedRange(UIGraphicsImageRendererFormat *format)
                                           asdisplaynode_iscancelled_block_t NS_NOESCAPE isCancelled,
                                           void (^NS_NOESCAPE work)())
 {
-  return ASGraphicsCreateImageWithTraitCollectionAndOptions(ASPrimitiveTraitCollectionMakeDefault(), size, opaque, scale, sourceImage, work);
+  return ASGraphicsCreateImage(ASPrimitiveTraitCollectionMakeDefault(), size, opaque, scale, sourceImage, isCancelled, work);
 }
 
-UIImage *ASGraphicsCreateImageWithTraitCollectionAndOptions(ASPrimitiveTraitCollection traitCollection, CGSize size, BOOL opaque, CGFloat scale, UIImage * sourceImage, void (NS_NOESCAPE ^work)()) {
+UIImage *ASGraphicsCreateImage(ASPrimitiveTraitCollection traitCollection, CGSize size, BOOL opaque, CGFloat scale, UIImage * sourceImage, asdisplaynode_iscancelled_block_t NS_NOESCAPE isCancelled, void (NS_NOESCAPE ^work)()) {
   if (AS_AVAILABLE_IOS_TVOS(10, 10)) {
     if (ASActivateExperimentalFeature(ASExperimentalDrawingGlobal)) {
       // If they used default scale, reuse one of two preferred formats.
@@ -98,17 +98,39 @@ NS_INLINE void ASConfigureExtendedRange(UIGraphicsImageRendererFormat *format)
         ASConfigureExtendedRange(format);
       }
 
-      return [[[UIGraphicsImageRenderer alloc] initWithSize:size format:format] imageWithActions:^(UIGraphicsImageRendererContext *rendererContext) {
-        ASDisplayNodeCAssert(UIGraphicsGetCurrentContext(), @"Should have a context!");
-        PERFORM_WORK_WITH_TRAIT_COLLECTION(work, traitCollection)
-      }];
+      // Avoid using the imageWithActions: method because it does not support cancellation at the
+      // last moment i.e. before actually creating the resulting image.
+      __block UIImage *image;
+      NSError *error;
+      [[[UIGraphicsImageRenderer alloc] initWithSize:size format:format]
+          runDrawingActions:^(UIGraphicsImageRendererContext *rendererContext) {
+            ASDisplayNodeCAssert(UIGraphicsGetCurrentContext(), @"Should have a context!");
+            ASPerformBlockWithTraitCollection(work, traitCollection);
+          }
+          completionActions:^(UIGraphicsImageRendererContext *rendererContext) {
+            if (isCancelled == nil || !isCancelled()) {
+              image = rendererContext.currentImage;
+            }
+          }
+          error:&error];
+      if (error) {
+        NSCAssert(NO, @"Error drawing: %@", error);
+      }
+      return image;
     }
   }
 
   // Bad OS or experiment flag. Use UIGraphics* API.
   UIGraphicsBeginImageContextWithOptions(size, opaque, scale);
-  PERFORM_WORK_WITH_TRAIT_COLLECTION(work, traitCollection)
-  UIImage *image = UIGraphicsGetImageFromCurrentImageContext();
+  ASPerformBlockWithTraitCollection(work, traitCollection)
+  UIImage *image = nil;
+  if (isCancelled == nil || !isCancelled()) {
+    image = UIGraphicsGetImageFromCurrentImageContext();
+  }
   UIGraphicsEndImageContext();
   return image;
 }
+
+UIImage *ASGraphicsCreateImageWithTraitCollectionAndOptions(ASPrimitiveTraitCollection traitCollection, CGSize size, BOOL opaque, CGFloat scale, UIImage * sourceImage, void (NS_NOESCAPE ^work)()) {
+  return ASGraphicsCreateImage(traitCollection, size, opaque, scale, sourceImage, nil, work);
+}

Tests/ASGraphicsContextTests.mm

@@ -28,6 +28,44 @@ - (void)setUp
 
 
 #if AS_AT_LEAST_IOS13
+- (void)testCanceled
+{
+  if (AS_AVAILABLE_IOS_TVOS(13, 13)) {
+    CGSize size = CGSize{.width=100, .height=100};
+
+    XCTestExpectation *expectationCancelled = [self expectationWithDescription:@"canceled"];
+
+    asdisplaynode_iscancelled_block_t isCancelledBlock =^BOOL{
+      [expectationCancelled fulfill];
+      return true;
+    };
+
+    ASPrimitiveTraitCollection traitCollection = ASPrimitiveTraitCollectionMakeDefault();
+    UIImage *canceledImage = ASGraphicsCreateImage(traitCollection, size, false, 0, nil, isCancelledBlock, ^{});
+
+    XCTAssertNil(canceledImage);
+
+    [self waitForExpectations:@[expectationCancelled] timeout:1];
+  }
+}
+
+- (void)testCanceledNil
+{
+  if (AS_AVAILABLE_IOS_TVOS(13, 13)) {
+    CGSize size = CGSize{.width=100, .height=100};
+    ASPrimitiveTraitCollection traitCollection = ASPrimitiveTraitCollectionMakeDefault();
+
+    XCTestExpectation *expectation = [self expectationWithDescription:@"normal"];
+    UIImage *image = ASGraphicsCreateImage(traitCollection, size, false, 0, nil, nil, ^{
+      [expectation fulfill];
+    });
+
+    XCTAssert(image);
+
+    [self waitForExpectations:@[expectation] timeout:1];
+  }
+}
+
 - (void)testTraitCollectionPassedToWork
 {
   if (AS_AVAILABLE_IOS_TVOS(13, 13)) {
@@ -36,7 +74,7 @@ - (void)testTraitCollectionPassedToWork
     XCTestExpectation *expectationDark = [self expectationWithDescription:@"trait collection dark"];
     ASPrimitiveTraitCollection traitCollectionDark = ASPrimitiveTraitCollectionMakeDefault();
     traitCollectionDark.userInterfaceStyle = UIUserInterfaceStyleDark;
-    ASGraphicsCreateImageWithTraitCollectionAndOptions(traitCollectionDark, size, false, 0, nil, ^{
+    ASGraphicsCreateImage(traitCollectionDark, size, false, 0, nil, nil, ^{
       UITraitCollection *currentTraitCollection = [UITraitCollection currentTraitCollection];
       XCTAssertEqual(currentTraitCollection.userInterfaceStyle, UIUserInterfaceStyleDark);
       [expectationDark fulfill];
@@ -45,7 +83,7 @@ - (void)testTraitCollectionPassedToWork
     XCTestExpectation *expectationLight = [self expectationWithDescription:@"trait collection light"];
     ASPrimitiveTraitCollection traitCollectionLight = ASPrimitiveTraitCollectionMakeDefault();
     traitCollectionLight.userInterfaceStyle = UIUserInterfaceStyleLight;
-    ASGraphicsCreateImageWithTraitCollectionAndOptions(traitCollectionLight, size, false, 0, nil, ^{
+    ASGraphicsCreateImage(traitCollectionLight, size, false, 0, nil, nil, ^{
       UITraitCollection *currentTraitCollection = [UITraitCollection currentTraitCollection];
       XCTAssertEqual(currentTraitCollection.userInterfaceStyle, UIUserInterfaceStyleLight);
       [expectationLight fulfill];