Clear out raw pointer.

4 files changed, 29 insertions, 11 deletions

diff --git a/include/mbgl/style/layers/custom_layer.hpp b/include/mbgl/style/layers/custom_layer.hpp
index a37ac25b01..3c6d015392 100644
--- a/include/mbgl/style/layers/custom_layer.hpp
+++ b/include/mbgl/style/layers/custom_layer.hpp
@@ -1,5 +1,6 @@
 #pragma once
 
+#include <functional>
 #include <mbgl/style/layer.hpp>
 
 namespace mbgl {
@@ -105,6 +106,8 @@ public:
     std::unique_ptr<Layer> cloneRef(const std::string& id) const final;
 
     CustomLayer(const CustomLayer&) = delete;
+
+    std::function<void()> onDestruction;
 };
 
 template <>
diff --git a/platform/darwin/src/MGLOpenGLStyleLayer.mm b/platform/darwin/src/MGLOpenGLStyleLayer.mm
index c4defddbe2..2641369fd1 100644
--- a/platform/darwin/src/MGLOpenGLStyleLayer.mm
+++ b/platform/darwin/src/MGLOpenGLStyleLayer.mm
@@ -160,7 +160,18 @@ public:
     auto layer = std::make_unique<mbgl::style::CustomLayer>(identifier.UTF8String,
                                                             std::move(context));
 
-    return self = [super initWithPendingLayer:std::move(layer)];
+    self = [super initWithPendingLayer:std::move(layer)];
+
+    if (self) {
+        __weak MGLOpenGLStyleLayer *weakself = self;
+
+        self.rawLayer->onDestruction =^{
+            weakself.rawLayer = NULL;
+        };
+    }
+
+    return self;
+
 }
 
 - (mbgl::style::CustomLayer *)rawLayer {
diff --git a/platform/ios/Integration Tests/MBGLIntegrationTests.mm b/platform/ios/Integration Tests/MBGLIntegrationTests.mm
index a25cc81782..f87e998c7b 100644
--- a/platform/ios/Integration Tests/MBGLIntegrationTests.mm
+++ b/platform/ios/Integration Tests/MBGLIntegrationTests.mm
@@ -507,7 +507,8 @@
 }
 
 - (void)testOpenGLLayerDoesNotLeakWhenStyleChanged {
-    __weak id weakLayer;
+
+    __weak MGLOpenGLStyleLayer *weakLayer;
 
     @autoreleasepool {
         {
@@ -532,22 +533,24 @@
     [self.mapView setStyleURL:styleURL];
     [self waitForExpectations:@[_styleLoadingExpectation] timeout:10];
 
-    // At this point the C++ CustomLayer will have been destroyed, BUT we're holding on
-    // to the obj-c layer. We don't currently clear the rawLayer, so it could now be
-    // pointing at garbage. And the following line will crash (due to description
-    // using .rawLayer)
-
-//    (void)layer2;
-    XCTFail();
-//    [self.style insertLayer:layer2 atIndex:0];
+    // At this point the C++ CustomLayer will have been destroyed, and the rawLayer pointer has been NULLed
+    XCTAssert(weakLayer == layer2);
+    XCTAssertNotNil(weakLayer);
+    XCTAssert(weakLayer.rawLayer == NULL);
 
+    @try {
+        [self.style insertLayer:layer2 atIndex:0];
+        XCTFail();
+    }
+    @catch (NSException *exception) {
+        // Success, we're expecting an exception
+    }
 
     [self waitForManagedLayersToExpire];
 
     // Asking the style for the layer should return nil
     MGLStyleLayer *layer3 = [self.mapView.style layerWithIdentifier:@"gl-layer"];
     XCTAssertNil(layer3);
-//    XCTAssertNil(weakLayer);
 }
 
 
diff --git a/src/mbgl/style/layers/custom_layer.cpp b/src/mbgl/style/layers/custom_layer.cpp
index ee8695b6c3..4f69ce8b52 100644
--- a/src/mbgl/style/layers/custom_layer.cpp
+++ b/src/mbgl/style/layers/custom_layer.cpp
@@ -10,6 +10,7 @@ CustomLayer::~CustomLayer()
     printf("~CustomLayer destructor %p\n", this);
 
     // Need to set rawLayer to nil in the obj-c layer.
+    onDestruction();
 }
 
 CustomLayer::CustomLayer(const std::string& layerID,