From e9404d6de347cb2e2c418a3c5ccd07bc636a7068 Mon Sep 17 00:00:00 2001
From: Jozufozu <jozsefaug@gmail.com>
Date: Fri, 5 Nov 2021 16:45:39 -0700
Subject: [PATCH] Revert "Weak hash map in world attached"

This reverts commit 98a3f759b37d63dff2e654c7883b8f388b555427.
---
 .../jozufozu/flywheel/event/ForgeEvents.java  |  9 +++++++
 .../jozufozu/flywheel/util/WorldAttached.java | 26 +++++++++++++++++--
 2 files changed, 33 insertions(+), 2 deletions(-)

diff --git a/src/main/java/com/jozufozu/flywheel/event/ForgeEvents.java b/src/main/java/com/jozufozu/flywheel/event/ForgeEvents.java
index 5d68edb33..4a3a73063 100644
--- a/src/main/java/com/jozufozu/flywheel/event/ForgeEvents.java
+++ b/src/main/java/com/jozufozu/flywheel/event/ForgeEvents.java
@@ -46,4 +46,13 @@ public class ForgeEvents {
 		}
 	}
 
+	@SubscribeEvent
+	public static void onUnloadWorld(WorldEvent.Unload event) {
+		IWorld world = event.getWorld();
+
+		if (Backend.isFlywheelWorld(world)) {
+			WorldAttached.invalidateWorld(world);
+		}
+	}
+
 }
diff --git a/src/main/java/com/jozufozu/flywheel/util/WorldAttached.java b/src/main/java/com/jozufozu/flywheel/util/WorldAttached.java
index bf9080377..21a864314 100644
--- a/src/main/java/com/jozufozu/flywheel/util/WorldAttached.java
+++ b/src/main/java/com/jozufozu/flywheel/util/WorldAttached.java
@@ -1,7 +1,11 @@
 package com.jozufozu.flywheel.util;
 
+import java.lang.ref.WeakReference;
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.Iterator;
+import java.util.List;
 import java.util.Map;
-import java.util.WeakHashMap;
 import java.util.function.BiConsumer;
 import java.util.function.Consumer;
 import java.util.function.Function;
@@ -12,12 +16,30 @@ import net.minecraft.world.IWorld;
 
 public class WorldAttached<T> {
 
+	// weak references to prevent leaking hashmaps when a WorldAttached is GC'd during runtime
+	static List<WeakReference<Map<IWorld, ?>>> allMaps = new ArrayList<>();
 	private final Map<IWorld, T> attached;
 	private final Function<IWorld, T> factory;
 
 	public WorldAttached(Function<IWorld, T> factory) {
 		this.factory = factory;
-		attached = new WeakHashMap<>();
+		attached = new HashMap<>();
+		allMaps.add(new WeakReference<>(attached));
+	}
+
+	public static void invalidateWorld(IWorld world) {
+		Iterator<WeakReference<Map<IWorld, ?>>> i = allMaps.iterator();
+		while (i.hasNext()) {
+			Map<IWorld, ?> map = i.next()
+					.get();
+			if (map == null) {
+				// If the map has been GC'd, remove the weak reference
+				i.remove();
+			} else {
+				// Prevent leaks
+				map.remove(world);
+			}
+		}
 	}
 
 	@Nonnull