[ios, macos] Introduce custom drawing overlays for MGLMapSnapshotterupstream/nb-snapshot-overlay

5 files changed, 189 insertions, 17 deletions

diff --git a/platform/darwin/src/MGLMapSnapshotter.h b/platform/darwin/src/MGLMapSnapshotter.h
index 1ee9bd99bb..0f20cf1bb2 100644
--- a/platform/darwin/src/MGLMapSnapshotter.h
+++ b/platform/darwin/src/MGLMapSnapshotter.h
@@ -6,6 +6,30 @@
 NS_ASSUME_NONNULL_BEGIN
 
 /**
+ An overlay that is placed within a `MGLMapSnapshot`.
+ To access this object, use `-[MGLMapSnapshotter startWithOverlayHandler:completionHandler:]`.
+ */
+
+MGL_EXPORT
+@interface MGLMapSnapshotOverlay : NSObject
+
+/**
+ The current `CGContext` that snapshot is drawing within. You may use this context
+ to perform additional custom drawing.
+ */
+@property (nonatomic, readonly) CGContextRef context;
+
+@end
+
+/**
+A block provided during the snapshot drawing process, enabling the ability to
+draw custom overlays rendered with Core Graphics.
+
+ @param snapshotOverlay The `MGLMapSnapshotOverlay` provided during snapshot drawing.
+ */
+typedef void (^MGLMapSnapshotOverlayHandler)(MGLMapSnapshotOverlay * snapshotOverlay);
+
+/**
  The options to use when creating images with the `MGLMapSnapshotter`.
  */
 MGL_EXPORT
@@ -201,6 +225,16 @@ MGL_EXPORT
 - (void)startWithQueue:(dispatch_queue_t)queue completionHandler:(MGLMapSnapshotCompletionHandler)completionHandler;
 
 /**
+ Starts the snapshot creation and executes the specified blocks with the result
+ on the specified queue. Use this option if you want to add custom drawing on top of the
+ resulting `MGLMapSnapShot`.
+ @param queue The queue to handle the result on.
+ @param overlayHandler The block to handle manipulation of the `MGLMapSnapshotter`'s `CGContext`.
+ @param completionHandler The block to handle the result in.
+ */
+- (void)startWithOverlayHandler:(MGLMapSnapshotOverlayHandler)overlayHandler completionHandler:(MGLMapSnapshotCompletionHandler)completionHandler;
+
+/**
  Cancels the snapshot creation request, if any.
  
  Once you call this method, you cannot resume the snapshot. In order to obtain
diff --git a/platform/darwin/src/MGLMapSnapshotter.mm b/platform/darwin/src/MGLMapSnapshotter.mm
index 0001ebcb86..3a258d146a 100644
--- a/platform/darwin/src/MGLMapSnapshotter.mm
+++ b/platform/darwin/src/MGLMapSnapshotter.mm
@@ -32,6 +32,26 @@
 const CGPoint MGLLogoImagePosition = CGPointMake(8, 8);
 const CGFloat MGLSnapshotterMinimumPixelSize = 64;
 
+
+@interface MGLMapSnapshotOverlay()
+
+- (instancetype)initWithContext:(CGContextRef)context;
+
+@end
+
+@implementation MGLMapSnapshotOverlay
+
+- (instancetype) initWithContext:(CGContextRef)context {
+    self = [super init];
+    if (self) {
+        _context = context;
+    }
+
+    return self;
+}
+
+@end
+
 @implementation MGLMapSnapshotOptions
 
 - (instancetype _Nonnull)initWithStyleURL:(nullable NSURL *)styleURL camera:(MGLMapCamera *)camera size:(CGSize)size
@@ -183,7 +203,15 @@ const CGFloat MGLSnapshotterMinimumPixelSize = 64;
     [self startWithQueue:dispatch_get_main_queue() completionHandler:completion];
 }
 
-- (void)startWithQueue:(dispatch_queue_t)queue completionHandler:(MGLMapSnapshotCompletionHandler)completion
+- (void)startWithQueue:(dispatch_queue_t)queue completionHandler:(MGLMapSnapshotCompletionHandler)completionHandler {
+    [self startWithQueue:queue overlayHandler:nil completionHandler:completionHandler];
+}
+
+- (void)startWithOverlayHandler:(MGLMapSnapshotOverlayHandler)overlayHandler completionHandler:(MGLMapSnapshotCompletionHandler)completion {
+    [self startWithQueue:dispatch_get_main_queue() overlayHandler:overlayHandler completionHandler:completion];
+}
+
+- (void)startWithQueue:(dispatch_queue_t)queue overlayHandler:(MGLMapSnapshotOverlayHandler)overlayHandler completionHandler:(MGLMapSnapshotCompletionHandler)completion
 {
     if (!mbgl::Scheduler::GetCurrent()) {
         [NSException raise:NSInvalidArgumentException
@@ -210,8 +238,8 @@ const CGFloat MGLSnapshotterMinimumPixelSize = 64;
     // capture weakSelf to avoid retain cycle if callback is never called (ie snapshot cancelled)
 
     _snapshotCallback = std::make_unique<mbgl::Actor<mbgl::MapSnapshotter::Callback>>(
-							*mbgl::Scheduler::GetCurrent(),
-							[=](std::exception_ptr mbglError, mbgl::PremultipliedImage image, mbgl::MapSnapshotter::Attributions attributions, mbgl::MapSnapshotter::PointForFn pointForFn, mbgl::MapSnapshotter::LatLngForFn latLngForFn) {
+                                                        *mbgl::Scheduler::GetCurrent(),
+                                                        [=](std::exception_ptr mbglError, mbgl::PremultipliedImage image, mbgl::MapSnapshotter::Attributions attributions, mbgl::MapSnapshotter::PointForFn pointForFn, mbgl::MapSnapshotter::LatLngForFn latLngForFn) {
 
         __typeof__(self) strongSelf = weakSelf;
         // If self had died, _snapshotCallback would have been destroyed and this block would not be executed
@@ -224,7 +252,7 @@ const CGFloat MGLSnapshotterMinimumPixelSize = 64;
             NSString *description = @(mbgl::util::toString(mbglError).c_str());
             NSDictionary *userInfo = @{NSLocalizedDescriptionKey: description};
             NSError *error = [NSError errorWithDomain:MGLErrorDomain code:MGLErrorCodeSnapshotFailed userInfo:userInfo];
-            
+
             // Dispatch to result queue
             dispatch_async(queue, ^{
                 strongSelf.completion(nil, error);
@@ -238,11 +266,12 @@ const CGFloat MGLSnapshotterMinimumPixelSize = 64;
             mglImage.size = NSMakeSize(mglImage.size.width / strongSelf.options.scale,
                                        mglImage.size.height / strongSelf.options.scale);
 #endif
-            [strongSelf drawAttributedSnapshot:attributions snapshotImage:mglImage pointForFn:pointForFn latLngForFn:latLngForFn];
+
+            [strongSelf drawAttributedSnapshot:attributions snapshotImage:mglImage pointForFn:pointForFn latLngForFn:latLngForFn overlayHandler:overlayHandler];
         }
         strongSelf->_snapshotCallback = NULL;
 
-    });
+      });
 
     // Launches snapshot on background Thread owned by mbglMapSnapshotter
     // _snapshotCallback->self() is an ActorRef: if the callback is destroyed, further messages
@@ -250,7 +279,7 @@ const CGFloat MGLSnapshotterMinimumPixelSize = 64;
     _mbglMapSnapshotter->snapshot(_snapshotCallback->self());
 }
 
-+ (MGLImage*)drawAttributedSnapshotWorker:(mbgl::MapSnapshotter::Attributions)attributions snapshotImage:(MGLImage *)mglImage pointForFn:(mbgl::MapSnapshotter::PointForFn)pointForFn latLngForFn:(mbgl::MapSnapshotter::LatLngForFn)latLngForFn scale:(CGFloat)scale size:(CGSize)size {
++ (MGLImage*)drawAttributedSnapshotWorker:(mbgl::MapSnapshotter::Attributions)attributions snapshotImage:(MGLImage *)mglImage pointForFn:(mbgl::MapSnapshotter::PointForFn)pointForFn latLngForFn:(mbgl::MapSnapshotter::LatLngForFn)latLngForFn scale:(CGFloat)scale size:(CGSize)size overlayHandler:(MGLMapSnapshotOverlayHandler)overlayHandler {
 
     NSArray<MGLAttributionInfo *>* attributionInfo = [MGLMapSnapshotter generateAttributionInfos:attributions];
 
@@ -292,7 +321,23 @@ const CGFloat MGLSnapshotterMinimumPixelSize = 64;
     UIGraphicsBeginImageContextWithOptions(mglImage.size, NO, scale);
     
     [mglImage drawInRect:CGRectMake(0, 0, mglImage.size.width, mglImage.size.height)];
-    
+
+    CGContextRef currentContext = UIGraphicsGetCurrentContext();
+
+    if (currentContext && overlayHandler) {
+        MGLMapSnapshotOverlay *snapshotOverlay = [[MGLMapSnapshotOverlay alloc] initWithContext:currentContext];
+        CGContextSaveGState(snapshotOverlay.context);
+        overlayHandler(snapshotOverlay);
+        CGContextRestoreGState(snapshotOverlay.context);
+        currentContext = UIGraphicsGetCurrentContext();
+    }
+
+    if (!currentContext && overlayHandler) {
+        // If the current context has been corrupted by the user,
+        // return nil so we can generate an error later.
+        return nil;
+    }
+
     [logoImage drawInRect:logoImageRect];
     
     UIImage *currentImage = UIGraphicsGetImageFromCurrentImageContext();
@@ -359,6 +404,21 @@ const CGFloat MGLSnapshotterMinimumPixelSize = 64;
     
     [sourceImageRep drawInRect: targetFrame];
     
+    NSGraphicsContext *currentContext = [NSGraphicsContext currentContext];
+    if (currentContext && overlayHandler) {
+        MGLMapSnapshotOverlay *snapshotOverlay = [[MGLMapSnapshotOverlay alloc] initWithContext:currentContext.CGContext];
+        [currentContext saveGraphicsState];
+        overlayHandler(snapshotOverlay);
+        [currentContext restoreGraphicsState];
+        currentContext = [NSGraphicsContext currentContext];
+    }
+    
+    if (!currentContext && overlayHandler) {
+        // If the current context has been corrupted by the user,
+        // return nil so we can generate an error later.
+        return nil;
+    }
+    
     if (logoImage) {
         [logoImage drawInRect:logoImageRect];
     }
@@ -379,8 +439,8 @@ const CGFloat MGLSnapshotterMinimumPixelSize = 64;
 #endif
 }
 
-- (void)drawAttributedSnapshot:(mbgl::MapSnapshotter::Attributions)attributions snapshotImage:(MGLImage *)mglImage pointForFn:(mbgl::MapSnapshotter::PointForFn)pointForFn latLngForFn:(mbgl::MapSnapshotter::LatLngForFn)latLngForFn {
-    
+- (void)drawAttributedSnapshot:(mbgl::MapSnapshotter::Attributions)attributions snapshotImage:(MGLImage *)mglImage pointForFn:(mbgl::MapSnapshotter::PointForFn)pointForFn latLngForFn:(mbgl::MapSnapshotter::LatLngForFn)latLngForFn overlayHandler:(MGLMapSnapshotOverlayHandler)overlayHandler {
+
     // Process image watermark in a work queue
     dispatch_queue_t workQueue = dispatch_get_global_queue(DISPATCH_QUEUE_PRIORITY_DEFAULT, 0);
     dispatch_queue_t resultQueue = self.resultQueue;
@@ -394,19 +454,30 @@ const CGFloat MGLSnapshotterMinimumPixelSize = 64;
 
     dispatch_async(workQueue, ^{
         // Call a class method to ensure we're not accidentally capturing self
-        MGLImage *compositedImage = [MGLMapSnapshotter drawAttributedSnapshotWorker:attributions snapshotImage:mglImage pointForFn:pointForFn latLngForFn:latLngForFn scale:scale size:size];
+        MGLImage *compositedImage = [MGLMapSnapshotter drawAttributedSnapshotWorker:attributions snapshotImage:mglImage pointForFn:pointForFn latLngForFn:latLngForFn scale:scale size:size overlayHandler:overlayHandler];
 
         // Dispatch result to origin queue
         dispatch_async(resultQueue, ^{
             __typeof__(self) strongself = weakself;
 
             if (strongself.completion) {
-                MGLMapSnapshot* snapshot = [[MGLMapSnapshot alloc] initWithImage:compositedImage
-                                                                           scale:scale
-                                                                      pointForFn:pointForFn
-                                                                     latLngForFn:latLngForFn];
-                strongself.completion(snapshot, nil);
-                strongself.completion = nil;
+
+                if (!compositedImage) {
+                    NSDictionary *userInfo = @{NSLocalizedDescriptionKey: @"Failed to generate composited snapshot."};
+                    NSError *error = [NSError errorWithDomain:MGLErrorDomain
+                                                         code:MGLErrorCodeSnapshotFailed
+                                                     userInfo:userInfo];
+
+                    strongself.completion(nil, error);
+                    strongself.completion = nil;
+                } else {
+                    MGLMapSnapshot* snapshot = [[MGLMapSnapshot alloc] initWithImage:compositedImage
+                                                                               scale:scale
+                                                                          pointForFn:pointForFn
+                                                                         latLngForFn:latLngForFn];
+                    strongself.completion(snapshot, nil);
+                    strongself.completion = nil;
+                }
             }
         });
     });
diff --git a/platform/ios/CHANGELOG.md b/platform/ios/CHANGELOG.md
index 7857c6bf65..8515fd467f 100644
--- a/platform/ios/CHANGELOG.md
+++ b/platform/ios/CHANGELOG.md
@@ -2,6 +2,15 @@
 
 Mapbox welcomes participation and contributions from everyone. Please read [CONTRIBUTING.md](../../CONTRIBUTING.md) to get started.
 
+## master
+
+### Styles and rendering
+* Added an `-[MGLMapSnapshotter startWithOverlayHandler:completionHandler:]` method to provide the snapshot's current `CGContext` in order to perform custom drawing on `MGLMapSnapShot` objects. ([#15530](https://github.com/mapbox/mapbox-gl-native/pull/15530))
+
+### Performance improvements
+
+* Newly loaded labels appear faster on the screen. ([#15308](https://github.com/mapbox/mapbox-gl-native/pull/15308))
+
 ## 5.4.0
 
 ### Styles and rendering
diff --git a/platform/ios/Integration Tests/Snapshotter Tests/MGLMapSnapshotterTest.m b/platform/ios/Integration Tests/Snapshotter Tests/MGLMapSnapshotterTest.m
index 39646755ba..7707896203 100644
--- a/platform/ios/Integration Tests/Snapshotter Tests/MGLMapSnapshotterTest.m
+++ b/platform/ios/Integration Tests/Snapshotter Tests/MGLMapSnapshotterTest.m
@@ -393,5 +393,62 @@ MGLMapSnapshotter* snapshotterWithCoordinates(CLLocationCoordinate2D coordinates
     [self waitForExpectations:@[expectation] timeout:10.0];
 }
 
+- (void)testSnapshotWithOverlayHandlerFailure {
+    if (![self validAccessToken]) {
+        return;
+    }
+
+    CGSize size = self.mapView.bounds.size;
+
+    XCTestExpectation *expectation = [self expectationWithDescription:@"snapshot with overlay fails"];
+    expectation.expectedFulfillmentCount = 2;
+
+    CLLocationCoordinate2D coord = CLLocationCoordinate2DMake(30.0, 30.0);
+
+    MGLMapSnapshotter *snapshotter = snapshotterWithCoordinates(coord, size);
+    XCTAssertNotNil(snapshotter);
+
+    [snapshotter startWithOverlayHandler:^(MGLMapSnapshotOverlay * _Nullable snapshotOverlay) {
+        UIGraphicsEndImageContext();
+        [expectation fulfill];
+    } completionHandler:^(MGLMapSnapshot * _Nullable snapshot, NSError * _Nullable error) {
+        XCTAssertNil(snapshot);
+        XCTAssertNotNil(error);
+        [expectation fulfill];
+    }];
+
+    [self waitForExpectations:@[expectation] timeout:10.0];
+}
+
+- (void)testSnapshotWithOverlayHandlerSuccess {
+    if (![self validAccessToken]) {
+        return;
+    }
+
+    CGSize size = self.mapView.bounds.size;
+    CGRect snapshotRect = CGRectMake(0, 0, size.width, size.height);
+
+    XCTestExpectation *expectation = [self expectationWithDescription:@"snapshot with overlay succeeds"];
+    expectation.expectedFulfillmentCount = 2;
+
+    CLLocationCoordinate2D coord = CLLocationCoordinate2DMake(30.0, 30.0);
+
+    MGLMapSnapshotter *snapshotter = snapshotterWithCoordinates(coord, size);
+    XCTAssertNotNil(snapshotter);
+
+    [snapshotter startWithOverlayHandler:^(MGLMapSnapshotOverlay * _Nullable snapshotOverlay) {
+        CGContextSetFillColorWithColor(snapshotOverlay.context, [UIColor.greenColor CGColor]);
+        CGContextSetAlpha(snapshotOverlay.context, 0.2);
+        CGContextAddRect(snapshotOverlay.context, snapshotRect);
+        CGContextFillRect(snapshotOverlay.context, snapshotRect);
+        [expectation fulfill];
+    } completionHandler:^(MGLMapSnapshot * _Nullable snapshot, NSError * _Nullable error) {
+        XCTAssertNil(error);
+        XCTAssertNotNil(snapshot);
+        [expectation fulfill];
+    }];
+
+    [self waitForExpectations:@[expectation] timeout:10.0];
+}
 
 @end
diff --git a/platform/macos/CHANGELOG.md b/platform/macos/CHANGELOG.md
index cbf08f7dd3..da9a179ad3 100644
--- a/platform/macos/CHANGELOG.md
+++ b/platform/macos/CHANGELOG.md
@@ -17,6 +17,7 @@
 * Fixed a rendering issue that non-SDF icon would be treated as SDF icon if they are in the same layer. ([#15456](https://github.com/mapbox/mapbox-gl-native/pull/15456))
 * Fixed a rendering issue of `collisionBox` when `text-translate` or `icon-translate` is enabled. ([#15467](https://github.com/mapbox/mapbox-gl-native/pull/15467))
 * Fixed an issue of integer overflow when converting `tileCoordinates` to `LatLon`, which caused issues such as `queryRenderedFeatures` and `querySourceFeatures` returning incorrect coordinates at zoom levels 20 and higher. ([#15560](https://github.com/mapbox/mapbox-gl-native/pull/15560))
+* Added an `-[MGLMapSnapshotter startWithOverlayHandler:completionHandler:]` method to provide the snapshot's current `CGContext` in order to perform custom drawing on `MGLMapSnapShot` objects. ([#15530](https://github.com/mapbox/mapbox-gl-native/pull/15530))
 
 ### Styles and rendering