From 5b29680a6859e41a4f5a74d604df47e5a89e4e10 Mon Sep 17 00:00:00 2001
From: Peng Liu <peng.liu@mapbox.com>
Date: Thu, 3 Nov 2022 16:03:05 +0200
Subject: [PATCH] Address review comments.

---
 sdk/src/main/java/com/mapbox/maps/MapController.kt        | 6 +++++-
 .../main/java/com/mapbox/maps/renderer/widget/Widget.kt   | 8 ++++----
 2 files changed, 9 insertions(+), 5 deletions(-)

diff --git a/sdk/src/main/java/com/mapbox/maps/MapController.kt b/sdk/src/main/java/com/mapbox/maps/MapController.kt
index a527a4aaab..bde5854ff7 100644
--- a/sdk/src/main/java/com/mapbox/maps/MapController.kt
+++ b/sdk/src/main/java/com/mapbox/maps/MapController.kt
@@ -219,12 +219,16 @@ internal class MapController : MapPluginProviderDelegate, MapControllable {
 
   override fun addWidget(widget: Widget) {
     val weakThis = WeakReference(this)
-    renderer.renderThread.addWidget(widget.also { it.setRepaintTrigger { weakThis.get()?.renderer?.scheduleRepaint() } })
+    widget.setRepaintTrigger {
+      weakThis.get()?.renderer?.scheduleRepaint()
+    }
+    renderer.renderThread.addWidget(widget)
     renderer.scheduleRepaint()
   }
 
   override fun removeWidget(widget: Widget): Boolean {
     val wasRemoved = renderer.renderThread.removeWidget(widget)
+    widget.setRepaintTrigger(null)
     if (wasRemoved) {
       renderer.scheduleRepaint()
     }
diff --git a/sdk/src/main/java/com/mapbox/maps/renderer/widget/Widget.kt b/sdk/src/main/java/com/mapbox/maps/renderer/widget/Widget.kt
index 6b7b3730b5..5dfbd314da 100644
--- a/sdk/src/main/java/com/mapbox/maps/renderer/widget/Widget.kt
+++ b/sdk/src/main/java/com/mapbox/maps/renderer/widget/Widget.kt
@@ -8,14 +8,14 @@ import com.mapbox.maps.MapboxExperimental
 @MapboxExperimental
 abstract class Widget internal constructor() {
   internal abstract val renderer: WidgetRenderer
-  private var repaintTriggerWeakReference: (() -> Unit)? = null
+  private var repaintTrigger: (() -> Unit)? = null
 
   /**
    * Triggers a repaint of the widget, as well as the map.
    * It will do nothing if the widget has not been added to the map.
    */
   protected fun triggerRepaint() {
-    repaintTriggerWeakReference?.invoke()
+    repaintTrigger?.invoke()
   }
 
   /**
@@ -51,7 +51,7 @@ abstract class Widget internal constructor() {
   abstract fun getRotation(): Float
 
   @JvmSynthetic
-  internal fun setRepaintTrigger(trigger: () -> Unit) {
-    repaintTriggerWeakReference = trigger
+  internal fun setRepaintTrigger(trigger: (() -> Unit)?) {
+    repaintTrigger = trigger
   }
 }
\ No newline at end of file