Unique node mapping (#737)

src/main/java/net/minestom/server/entity/Entity.java

@@ -6,6 +6,7 @@ import net.kyori.adventure.text.event.HoverEvent;
 import net.kyori.adventure.text.event.HoverEvent.ShowEntity;
 import net.kyori.adventure.text.event.HoverEventSource;
 import net.minestom.server.MinecraftServer;
+import net.minestom.server.ServerProcess;
 import net.minestom.server.Tickable;
 import net.minestom.server.Viewable;
 import net.minestom.server.collision.BoundingBox;
@@ -147,7 +148,7 @@ public class Entity implements Viewable, Tickable, Schedulable, Snapshotable, Ev
     protected final Set<Player> viewers = viewEngine.set;
     private final MutableNBTCompound nbtCompound = new MutableNBTCompound();
     private final Scheduler scheduler = Scheduler.newScheduler();
-    private final EventNode<EntityEvent> eventNode = EventNode.lazyMap(this, EventFilter.ENTITY);
+    private final EventNode<EntityEvent> eventNode;
     private final Set<Permission> permissions = new CopyOnWriteArraySet<>();
 
     protected UUID uuid;
@@ -189,6 +190,14 @@ public class Entity implements Viewable, Tickable, Schedulable, Snapshotable, Ev
 
         this.gravityAcceleration = entityType.registry().acceleration();
         this.gravityDragPerTick = entityType.registry().drag();
+
+        final ServerProcess process = MinecraftServer.process();
+        if (process != null) {
+            this.eventNode = process.eventHandler().map(this, EventFilter.ENTITY);
+        } else {
+            // Local nodes require a server process
+            this.eventNode = null;
+        }
     }
 
     public Entity(@NotNull EntityType entityType) {

src/main/java/net/minestom/server/event/EventNode.java

@@ -173,13 +173,6 @@ public sealed interface EventNode<T extends Event> permits EventNodeImpl {
         return create(name, filter, (e, h) -> consumer.test(h.getTag(tag)));
     }
 
-    @ApiStatus.Internal
-    @ApiStatus.Experimental
-    static <E extends Event> @NotNull EventNode<E> lazyMap(@NotNull Object owner,
-                                                           @NotNull EventFilter<E, ?> filter) {
-        return new EventNodeLazyImpl<>(owner, filter);
-    }
-
     private static <E extends Event, V> EventNode<E> create(@NotNull String name,
                                                             @NotNull EventFilter<E, V> filter,
                                                             @Nullable BiPredicate<E, V> predicate) {
@@ -345,20 +338,20 @@ public sealed interface EventNode<T extends Event> permits EventNodeImpl {
      * Be aware that such structure have huge performance penalty as they will
      * always require a map lookup. Use only at last resort.
      *
-     * @param node  the node to map
+     * @param value  the mapped value
-     * @param value the mapped value
+     * @param filter the filter to use
+     * @return the node (which may have already been registered) directly linked to {@code value}
      */
     @ApiStatus.Experimental
-    void map(@NotNull EventNode<? extends T> node, @NotNull Object value);
+    <E extends T, H> @NotNull EventNode<E> map(@NotNull H value, @NotNull EventFilter<E, H> filter);
 
     /**
-     * Undo {@link #map(EventNode, Object)}
+     * Prevents the node from {@link #map(Object, EventFilter)} to be called.
      *
      * @param value the value to unmap
-     * @return true if the value has been unmapped, false if nothing happened
      */
     @ApiStatus.Experimental
-    boolean unmap(@NotNull Object value);
+    void unmap(@NotNull Object value);
 
     @ApiStatus.Experimental
     void register(@NotNull EventBinding<? extends T> binding);

src/main/java/net/minestom/server/event/EventNodeImpl.java

@@ -27,16 +27,17 @@ non-sealed class EventNodeImpl<T extends Event> implements EventNode<T> {
             return new Handle<>((Class<T>) type);
         }
     };
-    private final Map<Class<? extends T>, ListenerEntry<T>> listenerMap = new ConcurrentHashMap<>();
+    final Map<Class<? extends T>, ListenerEntry<T>> listenerMap = new ConcurrentHashMap<>();
-    private final Set<EventNodeImpl<T>> children = new CopyOnWriteArraySet<>();
+    final Set<EventNodeImpl<T>> children = new CopyOnWriteArraySet<>();
-    private final Map<Object, EventNodeImpl<T>> mappedNodeCache = new WeakHashMap<>();
+    final Map<Object, EventNodeImpl<T>> mappedNodeCache = new WeakHashMap<>();
+    final Map<Object, EventNodeImpl<T>> registeredMappedNode = new WeakHashMap<>();
 
-    private final String name;
+    final String name;
-    private final EventFilter<T, ?> filter;
+    final EventFilter<T, ?> filter;
-    private final BiPredicate<T, Object> predicate;
+    final BiPredicate<T, Object> predicate;
-    private final Class<T> eventType;
+    final Class<T> eventType;
-    private volatile int priority;
+    volatile int priority;
-    private volatile EventNodeImpl<? super T> parent;
+    volatile EventNodeImpl<? super T> parent;
 
     EventNodeImpl(@NotNull String name,
                   @NotNull EventFilter<T, ?> filter,
@@ -150,27 +151,36 @@ non-sealed class EventNodeImpl<T extends Event> implements EventNode<T> {
     }
 
     @Override
-    public void map(@NotNull EventNode<? extends T> node, @NotNull Object value) {
+    public @NotNull <E extends T, H> EventNode<E> map(@NotNull H value, @NotNull EventFilter<E, H> filter) {
+        EventNodeImpl<E> node;
         synchronized (GLOBAL_CHILD_LOCK) {
-            final var nodeImpl = (EventNodeImpl<? extends T>) node;
+            node = new EventNodeLazyImpl<>(this, value, filter);
-            Check.stateCondition(nodeImpl.parent != null, "Node already has a parent");
+            Check.stateCondition(node.parent != null, "Node already has a parent");
-            Check.stateCondition(Objects.equals(parent, nodeImpl), "Cannot map to self");
+            Check.stateCondition(Objects.equals(parent, node), "Cannot map to self");
-            EventNodeImpl<T> previous = this.mappedNodeCache.put(value, (EventNodeImpl<T>) nodeImpl);
+            EventNodeImpl<T> previous = this.mappedNodeCache.putIfAbsent(value, (EventNodeImpl<T>) node);
-            if (previous != null) previous.parent = null;
+            if (previous != null) return (EventNode<E>) previous;
-            nodeImpl.parent = this;
+            node.parent = this;
-            nodeImpl.invalidateEventsFor(this);
+        }
+        return node;
+    }
+
+    void mapRegistration(EventNodeLazyImpl<? extends T> node, Object value) {
+        synchronized (GLOBAL_CHILD_LOCK) {
+            var previous = this.registeredMappedNode.putIfAbsent(value, (EventNodeImpl<T>) node);
+            if (previous == null) {
+                node.invalidateEventsFor(this);
+            }
         }
     }
 
     @Override
-    public boolean unmap(@NotNull Object value) {
+    public void unmap(@NotNull Object value) {
         synchronized (GLOBAL_CHILD_LOCK) {
-            final var mappedNode = this.mappedNodeCache.remove(value);
+            final var mappedNode = this.registeredMappedNode.remove(value);
-            if (mappedNode == null) return false; // Mapped node not found
+            if (mappedNode == null) return;
             final var childImpl = (EventNodeImpl<? extends T>) mappedNode;
             childImpl.parent = null;
             childImpl.invalidateEventsFor(this);
-            return true;
         }
     }
 
@@ -263,7 +273,7 @@ non-sealed class EventNodeImpl<T extends Event> implements EventNode<T> {
         }
     }
 
-    private void invalidateEventsFor(EventNodeImpl<? super T> node) {
+    void invalidateEventsFor(EventNodeImpl<? super T> node) {
         for (Class<? extends T> eventType : listenerMap.keySet()) {
             node.invalidateEvent(eventType);
         }
@@ -442,12 +452,12 @@ non-sealed class EventNodeImpl<T extends Event> implements EventNode<T> {
         }
 
         /**
-         * Create a consumer handling {@link EventNode#map(EventNode, Object)}.
+         * Create a consumer handling {@link EventNode#map(Object, EventFilter)}.
          * The goal is to limit the amount of map lookup.
          */
         private @Nullable Consumer<E> mappedConsumer() {
             var node = (EventNodeImpl<E>) EventNodeImpl.this;
-            final var mappedNodeCache = node.mappedNodeCache;
+            final var mappedNodeCache = node.registeredMappedNode;
             if (mappedNodeCache.isEmpty()) return null;
             Set<EventFilter<E, ?>> filters = new HashSet<>(mappedNodeCache.size());
             Map<Object, Handle<E>> handlers = new WeakHashMap<>(mappedNodeCache.size());

src/main/java/net/minestom/server/event/EventNodeLazyImpl.java

@@ -1,6 +1,5 @@
 package net.minestom.server.event;
 
-import net.minestom.server.MinecraftServer;
 import org.jetbrains.annotations.NotNull;
 
 import java.lang.invoke.MethodHandles;
@@ -19,12 +18,15 @@ final class EventNodeLazyImpl<E extends Event> extends EventNodeImpl<E> {
         }
     }
 
+    private final EventNodeImpl<? super E> holder;
     private final WeakReference<Object> owner;
     @SuppressWarnings("unused")
     private boolean mapped;
 
-    EventNodeLazyImpl(@NotNull Object owner, @NotNull EventFilter<E, ?> filter) {
+    EventNodeLazyImpl(@NotNull EventNodeImpl<? super E> holder,
+                      @NotNull Object owner, @NotNull EventFilter<E, ?> filter) {
         super(owner.toString(), filter, null);
+        this.holder = holder;
         this.owner = new WeakReference<>(owner);
     }
 
@@ -47,9 +49,12 @@ final class EventNodeLazyImpl<E extends Event> extends EventNodeImpl<E> {
     }
 
     @Override
-    public void map(@NotNull EventNode<? extends E> node, @NotNull Object value) {
+    public @NotNull <E1 extends E, H> EventNode<E1> map(@NotNull H value, @NotNull EventFilter<E1, H> filter) {
-        ensureMap();
+        final Object owner = retrieveOwner();
-        super.map(node, value);
+        if (owner != value) {
+            throw new IllegalArgumentException("Cannot map an object to an already mapped node.");
+        }
+        return (EventNode<E1>) this;
     }
 
     @Override
@@ -60,11 +65,15 @@ final class EventNodeLazyImpl<E extends Event> extends EventNodeImpl<E> {
 
     private void ensureMap() {
         if (MAPPED.compareAndSet(this, false, true)) {
-            final Object owner = this.owner.get();
+            this.holder.mapRegistration(this, retrieveOwner());
-            if (owner == null) {
-                throw new IllegalStateException("Node handle is null. Be sure to never cache a local node.");
-            }
-            MinecraftServer.getGlobalEventHandler().map(this, owner);
         }
     }
+
+    private Object retrieveOwner() {
+        final Object owner = this.owner.get();
+        if (owner == null) {
+            throw new IllegalStateException("Node handle is null. Be sure to never cache a local node.");
+        }
+        return owner;
+    }
 }

src/test/java/net/minestom/server/api/EnvImpl.java

@@ -4,7 +4,6 @@ import net.minestom.server.ServerProcess;
 import net.minestom.server.event.Event;
 import net.minestom.server.event.EventFilter;
 import net.minestom.server.event.EventListener;
-import net.minestom.server.event.EventNode;
 import org.jetbrains.annotations.NotNull;
 
 import java.util.List;
@@ -35,8 +34,7 @@ final class EnvImpl implements Env {
     @Override
     public @NotNull <E extends Event, H> Collector<E> trackEvent(@NotNull Class<E> eventType, @NotNull EventFilter<? super E, H> filter, @NotNull H actor) {
         var tracker = new EventCollector<E>(actor);
-        var node = EventNode.type("tracker", filter).addListener(eventType, tracker.events::add);
+        this.process.eventHandler().map(actor, filter).addListener(eventType, tracker.events::add);
-        process.eventHandler().map(node, actor);
         return tracker;
     }
 

src/test/java/net/minestom/server/event/EventNodeMapTest.java

@@ -14,17 +14,52 @@ import static net.minestom.server.api.TestUtils.waitUntilCleared;
 import static org.junit.jupiter.api.Assertions.*;
 
 public class EventNodeMapTest {
+
+    @Test
+    public void uniqueMapping() {
+        var item = ItemStack.of(Material.DIAMOND);
+        var node = EventNode.all("main");
+        var itemNode1 = node.map(item, EventFilter.ITEM);
+        var itemNode2 = node.map(item, EventFilter.ITEM);
+        assertNotNull(itemNode1);
+        assertSame(itemNode1, itemNode2);
+
+        // Node should still keep track of the mapping until GCed
+        // This is to ensure that we do not end up with multiple nodes theoretically mapping the same object
+        node.unmap(item);
+        assertSame(itemNode1, node.map(item, EventFilter.ITEM));
+    }
+
+    @Test
+    public void lazyRegistration() {
+        var item = ItemStack.of(Material.DIAMOND);
+        var node = (EventNodeImpl<Event>) EventNode.all("main");
+        var itemNode = node.map(item, EventFilter.ITEM);
+        assertFalse(node.registeredMappedNode.containsKey(item));
+        itemNode.addListener(EventNodeTest.ItemTestEvent.class, event -> {
+        });
+        assertTrue(node.registeredMappedNode.containsKey(item));
+    }
+
+    @Test
+    public void secondMap() {
+        var item = ItemStack.of(Material.DIAMOND);
+        var node = (EventNodeImpl<Event>) EventNode.all("main");
+        var itemNode = node.map(item, EventFilter.ITEM);
+        assertSame(itemNode, itemNode.map(item, EventFilter.ITEM));
+        assertThrows(Exception.class, () -> itemNode.map(ItemStack.AIR, EventFilter.ITEM));
+    }
+
     @Test
     public void map() {
         var item = ItemStack.of(Material.DIAMOND);
         var node = EventNode.all("main");
 
         AtomicBoolean result = new AtomicBoolean(false);
-        var itemNode = EventNode.type("item_node", EventFilter.ITEM);
+        var itemNode = node.map(item, EventFilter.ITEM);
 
         assertFalse(node.hasListener(EventNodeTest.ItemTestEvent.class));
         itemNode.addListener(EventNodeTest.ItemTestEvent.class, event -> result.set(true));
-        assertDoesNotThrow(() -> node.map(itemNode, item));
         assertTrue(node.hasListener(EventNodeTest.ItemTestEvent.class));
 
         node.call(new EventNodeTest.ItemTestEvent(item));
@@ -35,7 +70,7 @@ public class EventNodeMapTest {
         assertFalse(result.get());
 
         result.set(false);
-        assertTrue(node.unmap(item));
+        node.unmap(item);
         node.call(new EventNodeTest.ItemTestEvent(item));
         assertFalse(result.get());
     }
@@ -71,10 +106,9 @@ public class EventNodeMapTest {
         // Ensure that the mapped object gets GCed
         var item = ItemStack.of(Material.DIAMOND);
         var node = EventNode.all("main");
-        var itemNode = EventNode.type("item_node", EventFilter.ITEM);
+        var itemNode = node.map(item, EventFilter.ITEM);
         itemNode.addListener(EventNodeTest.ItemTestEvent.class, event -> {
         });
-        node.map(itemNode, item);
         node.call(new EventNodeTest.ItemTestEvent(item));
 
         var ref = new WeakReference<>(item);