From 8725b204a07098154be555cbf9b183513e30f7b3 Mon Sep 17 00:00:00 2001
From: weisj <31143295+weisJ@users.noreply.github.com>
Date: Fri, 28 May 2021 15:00:18 +0200
Subject: [PATCH] Memory: Ensure UIUpdater doesn't prevent components from
 being garbage collected.

---
 .../weisj/darklaf/listener/UIUpdater.java     | 16 ++++---
 core/src/test/java/test/MemoryTest.java       | 43 ++++++++++++++++++-
 2 files changed, 52 insertions(+), 7 deletions(-)

diff --git a/core/src/main/java/com/github/weisj/darklaf/listener/UIUpdater.java b/core/src/main/java/com/github/weisj/darklaf/listener/UIUpdater.java
index 0e2c08c9..b605a788 100644
--- a/core/src/main/java/com/github/weisj/darklaf/listener/UIUpdater.java
+++ b/core/src/main/java/com/github/weisj/darklaf/listener/UIUpdater.java
@@ -22,6 +22,7 @@
 package com.github.weisj.darklaf.listener;
 
 import java.awt.*;
+import java.lang.ref.WeakReference;
 
 import javax.swing.*;
 
@@ -66,11 +67,11 @@ public class UIUpdater implements ThemeChangeListener {
     }
 
     private static void removeComponent(final JComponent component, final UIUpdater updater) {
-        component.putClientProperty(KEY_UPDATER, null);
+        if (component != null) component.putClientProperty(KEY_UPDATER, null);
         LafManager.removeThemeChangeListener(updater);
     }
 
-    private final Component component;
+    private final WeakReference<JComponent> component;
 
     /**
      * Creates a new ui updater for the given component. After a new theme has been installed This
@@ -80,8 +81,8 @@ public class UIUpdater implements ThemeChangeListener {
      *
      * @param component the component to update.
      */
-    public UIUpdater(final Component component) {
-        this.component = component;
+    public UIUpdater(final JComponent component) {
+        this.component = new WeakReference<>(component);
     }
 
     @Override
@@ -89,6 +90,11 @@ public class UIUpdater implements ThemeChangeListener {
 
     @Override
     public void themeInstalled(final ThemeChangeEvent e) {
-        if (component != null) SwingUtilities.updateComponentTreeUI(component);
+        JComponent comp = component.get();
+        if (comp != null) {
+            SwingUtilities.updateComponentTreeUI(comp);
+        } else {
+            removeComponent(null, this);
+        }
     }
 }
diff --git a/core/src/test/java/test/MemoryTest.java b/core/src/test/java/test/MemoryTest.java
index f39d16dd..f516b2e3 100644
--- a/core/src/test/java/test/MemoryTest.java
+++ b/core/src/test/java/test/MemoryTest.java
@@ -32,21 +32,36 @@ import javax.swing.JButton;
 import javax.swing.JComponent;
 import javax.swing.JEditorPane;
 import javax.swing.JFrame;
+import javax.swing.JLabel;
 import javax.swing.JPanel;
 import javax.swing.JTextArea;
 import javax.swing.JTextField;
 import javax.swing.JTextPane;
 import javax.swing.SwingUtilities;
 
+import org.junit.jupiter.api.Assertions;
+import org.junit.jupiter.api.BeforeAll;
+import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.Timeout;
 
 import com.github.weisj.darklaf.LafManager;
-import com.github.weisj.darklaf.theme.DarculaTheme;
+import com.github.weisj.darklaf.components.DynamicUI;
+import com.github.weisj.darklaf.listener.UIUpdater;
 
 @Timeout(value = 5)
 class MemoryTest {
 
+    @BeforeAll
+    static void setupLaf() {
+        TestUtils.runOnSwingThreadNotThrowing(LafManager::install);
+    }
+
+    @BeforeEach
+    void checkLaf() {
+        Assertions.assertTrue(LafManager.isInstalled());
+    }
+
     @Test
     void frameGetsGarbageCollectedSimpleContent() {
         testWithContentPane(JPanel::new);
@@ -75,10 +90,34 @@ class MemoryTest {
         });
     }
 
+    @Test
+    void uiUpdaterReleasesMemory() {
+        AtomicReference<JLabel> labelRef = new AtomicReference<>();
+        TestUtils.runOnSwingThreadNotThrowing(() -> {
+            labelRef.set(new JLabel());
+            UIUpdater.registerComponent(labelRef.get());
+        });
+        WeakReference<JLabel> weakRef = new WeakReference<>(labelRef.get());
+        SwingUtilities.invokeLater(() -> labelRef.set(null));
+        Assertions.assertNotNull(weakRef.get());
+        waitForGarbageCollection(weakRef);
+    }
+
+    @Test
+    void dynamicUiReleasesMemory() {
+        AtomicReference<JLabel> labelRef = new AtomicReference<>();
+        TestUtils.runOnSwingThreadNotThrowing(
+                () -> labelRef.set(DynamicUI.withDynamic(new JLabel(), c -> {
+                })));
+        WeakReference<JLabel> weakRef = new WeakReference<>(labelRef.get());
+        SwingUtilities.invokeLater(() -> labelRef.set(null));
+        Assertions.assertNotNull(weakRef.get());
+        waitForGarbageCollection(weakRef);
+    }
+
     private void testWithContentPane(final Supplier<JComponent> content) {
         AtomicReference<JFrame> frame = new AtomicReference<>();
         TestUtils.runOnSwingThreadNotThrowing(() -> {
-            LafManager.install(new DarculaTheme());
             JFrame f = new JFrame();
             f.setContentPane(content.get());
             f.setDefaultCloseOperation(JFrame.DISPOSE_ON_CLOSE);