Allow checking if an AttachmentHolder has any attachments, allow remo…

…ving attachments, and allow skipping serialization of certain attachments (#609)
neoforged · Feb 5, 2024 · 5e55ae2 · 5e55ae2
1 parent 77c8899
commit 5e55ae2
Show file tree

Hide file tree

Showing 10 changed files with 150 additions and 8 deletions.
diff --git a/patches/net/minecraft/world/level/block/entity/BlockEntity.java.patch b/patches/net/minecraft/world/level/block/entity/BlockEntity.java.patch
@@ -47,7 +47,7 @@
      }
 
      public boolean triggerEvent(int p_58889_, int p_58890_) {
-@@ -185,6 +_,20 @@
+@@ -185,6 +_,27 @@
 
      public BlockEntityType<?> getType() {
          return this.type;
@@ -65,6 +65,13 @@
 +    public final <T> T setData(net.neoforged.neoforge.attachment.AttachmentType<T> type, T data) {
 +        setChanged();
 +        return super.setData(type, data);
++    }
++
++    @Override
++    @Nullable
++    public final <T> T removeData(net.neoforged.neoforge.attachment.AttachmentType<T> type) {
++        setChanged();
++        return super.removeData(type);
      }
 
      @Deprecated
diff --git a/patches/net/minecraft/world/level/chunk/LevelChunk.java.patch b/patches/net/minecraft/world/level/chunk/LevelChunk.java.patch
@@ -111,7 +111,7 @@
          this.blockEntities.values().forEach(p_187988_ -> {
              Level level = this.level;
              if (level instanceof ServerLevel serverlevel) {
-@@ -646,6 +_,33 @@
+@@ -646,6 +_,45 @@
          return new LevelChunk.BoundTickingBlockEntity<>(p_156376_, p_156377_);
      }
 
@@ -120,6 +120,11 @@
 +    private final net.neoforged.neoforge.common.world.LevelChunkAuxiliaryLightManager auxLightManager = new net.neoforged.neoforge.common.world.LevelChunkAuxiliaryLightManager(this);
 +
 +    @Override
++    public boolean hasAttachments() {
++        return attachmentHolder.hasAttachments();
++    }
++
++    @Override
 +    public boolean hasData(net.neoforged.neoforge.attachment.AttachmentType<?> type) {
 +        return attachmentHolder.hasData(type);
 +    }
@@ -137,6 +142,13 @@
 +    }
 +
 +    @Override
++    @Nullable
++    public <T> T removeData(net.neoforged.neoforge.attachment.AttachmentType<T> type) {
++        setUnsaved(true);
++        return attachmentHolder.removeData(type);
++    }
++
++    @Override
 +    public net.neoforged.neoforge.common.world.LevelChunkAuxiliaryLightManager getAuxLightManager(ChunkPos pos) {
 +        return auxLightManager;
 +    }

diff --git a/src/main/java/net/neoforged/neoforge/attachment/AttachmentHolder.java b/src/main/java/net/neoforged/neoforge/attachment/AttachmentHolder.java
@@ -58,6 +58,11 @@ IAttachmentHolder getExposedHolder() {
         return this;
     }
 
+    @Override
+    public final boolean hasAttachments() {
+        return attachments != null && !attachments.isEmpty();
+    }
+
     @Override
     public final boolean hasData(AttachmentType<?> type) {
         validateAttachmentType(type);
@@ -83,6 +88,16 @@ public final <T> T getData(AttachmentType<T> type) {
         return (T) getAttachmentMap().put(type, data);
     }
 
+    @Override
+    @MustBeInvokedByOverriders
+    public <T> @Nullable T removeData(AttachmentType<T> type) {
+        validateAttachmentType(type);
+        if (attachments == null) {
+            return null;
+        }
+        return (T) attachments.remove(type);
+    }
+
     /**
      * Writes the serializable attachments to a tag.
      * Returns {@code null} if there are no serializable attachments.
@@ -96,9 +111,12 @@ public final CompoundTag serializeAttachments() {
         for (var entry : attachments.entrySet()) {
             var type = entry.getKey();
             if (type.serializer != null) {
-                if (tag == null)
-                    tag = new CompoundTag();
-                tag.put(NeoForgeRegistries.ATTACHMENT_TYPES.getKey(type).toString(), ((IAttachmentSerializer<?, Object>) type.serializer).write(entry.getValue()));
+                Tag serialized = ((IAttachmentSerializer<?, Object>) type.serializer).write(entry.getValue());
+                if (serialized != null) {
+                    if (tag == null)
+                        tag = new CompoundTag();
+                    tag.put(NeoForgeRegistries.ATTACHMENT_TYPES.getKey(type).toString(), serialized);
+                }
             }
         }
         return tag;
@@ -119,6 +137,7 @@ protected final void deserializeAttachments(CompoundTag tag) {
             var type = NeoForgeRegistries.ATTACHMENT_TYPES.get(keyLocation);
             if (type == null || type.serializer == null) {
                 LOGGER.error("Encountered unknown or non-serializable data attachment {}. Skipping.", key);
+                continue;
             }
 
             try {

diff --git a/src/main/java/net/neoforged/neoforge/attachment/AttachmentInternals.java b/src/main/java/net/neoforged/neoforge/attachment/AttachmentInternals.java
@@ -87,7 +87,10 @@ private static <H extends AttachmentHolder> void copyAttachments(H from, H to, P
             @SuppressWarnings("unchecked")
             var serializer = (IAttachmentSerializer<Tag, Object>) type.serializer;
             if (serializer != null && filter.test(type)) {
-                to.getAttachmentMap().put(type, serializer.read(to.getExposedHolder(), serializer.write(entry.getValue())));
+                Tag serialized = serializer.write(entry.getValue());
+                if (serialized != null) {
+                    to.getAttachmentMap().put(type, serializer.read(to.getExposedHolder(), serialized));
+                }
             }
         }
     }

diff --git a/src/main/java/net/neoforged/neoforge/attachment/AttachmentType.java b/src/main/java/net/neoforged/neoforge/attachment/AttachmentType.java
@@ -5,9 +5,11 @@
 
 package net.neoforged.neoforge.attachment;
 
+import com.google.common.base.Predicates;
 import com.mojang.serialization.Codec;
 import java.util.Objects;
 import java.util.function.Function;
+import java.util.function.Predicate;
 import java.util.function.Supplier;
 import net.minecraft.nbt.NbtOps;
 import net.minecraft.nbt.Tag;
@@ -126,6 +128,7 @@ public T read(IAttachmentHolder holder, S tag) {
                 return ret;
             }
 
+            @Nullable
             @Override
             public S write(T attachment) {
                 return attachment.serializeNBT();
@@ -170,6 +173,21 @@ public Builder<T> serialize(IAttachmentSerializer<?, T> serializer) {
          * @param codec The codec to use.
          */
         public Builder<T> serialize(Codec<T> codec) {
+            return serialize(codec, Predicates.alwaysTrue());
+        }
+
+        /**
+         * Requests that this attachment be persisted to disk (on the logical server side), using a {@link Codec}.
+         *
+         * <p>Using a {@link Codec} to serialize attachments is discouraged for item stack attachments,
+         * for performance reasons. Prefer one of the other options.
+         *
+         * <p>Codec-based attachments cannot capture a reference to their holder.
+         *
+         * @param codec           The codec to use.
+         * @param shouldSerialize A check that determines whether serialization of the attachment should occur.
+         */
+        public Builder<T> serialize(Codec<T> codec, Predicate<? super T> shouldSerialize) {
             Objects.requireNonNull(codec);
             // TODO: better error handling
             return serialize(new IAttachmentSerializer<>() {
@@ -178,9 +196,10 @@ public T read(IAttachmentHolder holder, Tag tag) {
                     return codec.parse(NbtOps.INSTANCE, tag).result().get();
                 }
 
+                @Nullable
                 @Override
                 public Tag write(T attachment) {
-                    return codec.encodeStart(NbtOps.INSTANCE, attachment).result().get();
+                    return shouldSerialize.test(attachment) ? codec.encodeStart(NbtOps.INSTANCE, attachment).result().get() : null;
                 }
             });
         }

diff --git a/src/main/java/net/neoforged/neoforge/attachment/IAttachmentHolder.java b/src/main/java/net/neoforged/neoforge/attachment/IAttachmentHolder.java
@@ -12,6 +12,11 @@
  * An object that can hold data attachments.
  */
 public interface IAttachmentHolder {
+    /**
+     * Returns {@code true} if there is any data attachments, {@code false} otherwise.
+     */
+    boolean hasAttachments();
+
     /**
      * Returns {@code true} if there is a data attachment of the give type, {@code false} otherwise.
      */
@@ -56,4 +61,20 @@ default <T> T getData(Supplier<AttachmentType<T>> type) {
     default <T> @Nullable T setData(Supplier<AttachmentType<T>> type, T data) {
         return setData(type.get(), data);
     }
+
+    /**
+     * Removes the data attachment of the given type.
+     *
+     * @return the previous value for that attachment type, if any, or {@code null} if there was none
+     */
+    <T> @Nullable T removeData(AttachmentType<T> type);
+
+    /**
+     * Removes the data attachment of the given type.
+     *
+     * @return the previous value for that attachment type, if any, or {@code null} if there was none
+     */
+    default <T> @Nullable T removeData(Supplier<AttachmentType<T>> type) {
+        return removeData(type.get());
+    }
 }
diff --git a/src/main/java/net/neoforged/neoforge/attachment/IAttachmentSerializer.java b/src/main/java/net/neoforged/neoforge/attachment/IAttachmentSerializer.java
@@ -6,6 +6,7 @@
 package net.neoforged.neoforge.attachment;
 
 import net.minecraft.nbt.Tag;
+import org.jetbrains.annotations.Nullable;
 
 /**
  * Serializer for data attachments.
@@ -39,7 +40,8 @@ default T read(IAttachmentHolder holder, S tag) {
     }
 
     /**
-     * Writes the attachment to NBT.
+     * Writes the attachment to NBT, or returns null if it is should not be serialized.
      */
+    @Nullable
     S write(T attachment);
 }
diff --git a/src/main/java/net/neoforged/neoforge/common/util/INBTSerializable.java b/src/main/java/net/neoforged/neoforge/common/util/INBTSerializable.java
@@ -6,12 +6,14 @@
 package net.neoforged.neoforge.common.util;
 
 import net.minecraft.nbt.Tag;
+import org.jetbrains.annotations.UnknownNullability;
 
 /**
  * An interface designed to unify various things in the Minecraft
  * code base that can be serialized to and from a NBT tag.
  */
 public interface INBTSerializable<T extends Tag> {
+    @UnknownNullability
     T serializeNBT();
 
     void deserializeNBT(T nbt);

diff --git a/src/main/java/net/neoforged/neoforge/common/world/LevelChunkAuxiliaryLightManager.java b/src/main/java/net/neoforged/neoforge/common/world/LevelChunkAuxiliaryLightManager.java
@@ -22,6 +22,7 @@
 import net.neoforged.neoforge.common.util.INBTSerializable;
 import net.neoforged.neoforge.network.payload.AuxiliaryLightDataPayload;
 import org.jetbrains.annotations.ApiStatus;
+import org.jetbrains.annotations.Nullable;
 
 @ApiStatus.Internal
 public final class LevelChunkAuxiliaryLightManager implements AuxiliaryLightManager, INBTSerializable<ListTag> {
@@ -56,6 +57,7 @@ public int getLightAt(BlockPos pos) {
         return lights.getOrDefault(pos, (byte) 0);
     }
 
+    @Nullable
     @Override
     public ListTag serializeNBT() {
         if (lights.isEmpty()) {

diff --git a/tests/src/main/java/net/neoforged/neoforge/debug/attachment/AttachmentTests.java b/tests/src/main/java/net/neoforged/neoforge/debug/attachment/AttachmentTests.java
@@ -264,4 +264,59 @@ static void itemAttachmentCodec(final DynamicTest test, final RegistrationHelper
             helper.succeed();
         });
     }
+
+    @GameTest
+    @EmptyTemplate
+    @TestHolder(description = "Ensures that attachments can opt-out of serializing default values")
+    static void itemAttachmentSkipSerialization(final DynamicTest test, final RegistrationHelper reg) {
+        var attachmentType = reg.registrar(NeoForgeRegistries.Keys.ATTACHMENT_TYPES)
+                .register("test_int", () -> AttachmentType.builder(() -> 0).serialize(Codec.INT, i -> i != 0).build());
+
+        test.onGameTest(helper -> {
+            ItemStack stack = Items.APPLE.getDefaultInstance();
+            stack.setData(attachmentType, 1);
+            helper.assertTrue(stack.serializeAttachments() != null, "Stack should have serialized attachments");
+            stack.setData(attachmentType, 0);
+            helper.assertTrue(stack.serializeAttachments() == null, "None of the stack's attachments should be serialized");
+            helper.assertTrue(stack.hasData(attachmentType), "Stack should have attached data");
+
+            helper.succeed();
+        });
+    }
+
+    @GameTest
+    @EmptyTemplate
+    @TestHolder(description = "Ensures that removing attachments works")
+    static void itemAttachmentRemoval(final DynamicTest test, final RegistrationHelper reg) {
+        var attachmentType = reg.registrar(NeoForgeRegistries.Keys.ATTACHMENT_TYPES)
+                .register("test_int", () -> AttachmentType.builder(() -> 0).serialize(Codec.INT).build());
+
+        test.onGameTest(helper -> {
+            ItemStack stack = Items.APPLE.getDefaultInstance();
+            stack.setData(attachmentType, 1);
+            helper.assertTrue(stack.hasData(attachmentType), "Stack should have attached data");
+            stack.removeData(attachmentType);
+            helper.assertFalse(stack.hasData(attachmentType), "Stack should not have attached data");
+
+            helper.succeed();
+        });
+    }
+
+    @GameTest
+    @EmptyTemplate
+    @TestHolder(description = "Ensures that the presence of non-serializable attachments can be checked")
+    static void itemAttachmentPresence(final DynamicTest test, final RegistrationHelper reg) {
+        var attachmentType = reg.registrar(NeoForgeRegistries.Keys.ATTACHMENT_TYPES)
+                .register("test_int", () -> AttachmentType.builder(() -> 0).build());
+
+        test.onGameTest(helper -> {
+            ItemStack stack = Items.APPLE.getDefaultInstance();
+            stack.setData(attachmentType, 1);
+            helper.assertTrue(stack.hasData(attachmentType), "Stack should have attached data");
+            //Also check we can detect the presence if we don't know what types are attached
+            helper.assertTrue(stack.hasAttachments(), "Stack should have attached data");
+
+            helper.succeed();
+        });
+    }
 }