From 8df48bc378ee63425a02143326045247048519b1 Mon Sep 17 00:00:00 2001
From: Nassim Jahnke <jahnke.nassim@gmail.com>
Date: Wed, 20 Apr 2022 20:16:08 +0200
Subject: [PATCH] StoredEntityData#remove

---
 .../viaversion/api/data/entity/StoredEntityData.java     | 9 +++++++++
 .../viaversion/data/entity/StoredEntityImpl.java         | 5 +++++
 2 files changed, 14 insertions(+)

diff --git a/api/src/main/java/com/viaversion/viaversion/api/data/entity/StoredEntityData.java b/api/src/main/java/com/viaversion/viaversion/api/data/entity/StoredEntityData.java
index 8f746fc2f..ec96fa3aa 100644
--- a/api/src/main/java/com/viaversion/viaversion/api/data/entity/StoredEntityData.java
+++ b/api/src/main/java/com/viaversion/viaversion/api/data/entity/StoredEntityData.java
@@ -51,6 +51,15 @@ public interface StoredEntityData {
      */
     @Nullable <T> T get(Class<T> objectClass);
 
+    /**
+     * Removes and returns an object from the storage if present.
+     *
+     * @param objectClass class of the object to remove
+     * @param <T>         object type
+     * @return removed object if present
+     */
+    @Nullable <T> T remove(Class<T> objectClass);
+
     /**
      * Stores an object based on its class.
      *
diff --git a/common/src/main/java/com/viaversion/viaversion/data/entity/StoredEntityImpl.java b/common/src/main/java/com/viaversion/viaversion/data/entity/StoredEntityImpl.java
index 8b74b881c..cc1cafa9e 100644
--- a/common/src/main/java/com/viaversion/viaversion/data/entity/StoredEntityImpl.java
+++ b/common/src/main/java/com/viaversion/viaversion/data/entity/StoredEntityImpl.java
@@ -42,6 +42,11 @@ public final class StoredEntityImpl implements StoredEntityData {
         return (T) storedObjects.get(objectClass);
     }
 
+    @Override
+    public <T> @Nullable T remove(Class<T> objectClass) {
+        return (T) storedObjects.remove(objectClass);
+    }
+
     @Override
     public boolean has(Class<?> objectClass) {
         return storedObjects.containsKey(objectClass);