Tag path API (#800)

src/main/java/net/minestom/server/tag/Tag.java

@@ -9,6 +9,8 @@ import org.jetbrains.annotations.Nullable;
 import org.jglrxavpok.hephaistos.nbt.*;
 import org.jglrxavpok.hephaistos.nbt.mutable.MutableNBTCompound;
 
+import java.util.Arrays;
+import java.util.List;
 import java.util.function.Function;
 import java.util.function.Supplier;
 
@@ -23,6 +25,9 @@ import java.util.function.Supplier;
 public class Tag<T> {
     private static final IndexMap<String> INDEX_MAP = new IndexMap<>();
 
+    record PathEntry(String name, int index) {
+    }
+
     private final String key;
     final Function<NBT, T> readFunction;
     final Function<T, NBT> writeFunction;
@@ -30,21 +35,25 @@ public class Tag<T> {
 
     final int index;
 
+    final List<PathEntry> path;
+
     protected Tag(@NotNull String key,
                   @NotNull Function<NBT, T> readFunction,
                   @NotNull Function<T, NBT> writeFunction,
-                  @Nullable Supplier<T> defaultValue) {
+                  @Nullable Supplier<T> defaultValue,
+                  @Nullable List<PathEntry> path) {
+        this.index = INDEX_MAP.get(key);
         this.key = key;
         this.readFunction = readFunction;
         this.writeFunction = writeFunction;
         this.defaultValue = defaultValue;
-        this.index = INDEX_MAP.get(key);
+        this.path = path;
     }
 
     static <T, N extends NBT> Tag<T> tag(@NotNull String key,
                                          @NotNull Function<N, T> readFunction,
                                          @NotNull Function<T, N> writeFunction) {
-        return new Tag<>(key, (Function<NBT, T>) readFunction, (Function<T, NBT>) writeFunction, null);
+        return new Tag<>(key, (Function<NBT, T>) readFunction, (Function<T, NBT>) writeFunction, null, null);
     }
 
     /**
@@ -58,7 +67,7 @@ public class Tag<T> {
 
     @Contract(value = "_ -> new", pure = true)
     public Tag<T> defaultValue(@NotNull Supplier<T> defaultValue) {
-        return new Tag<>(key, readFunction, writeFunction, defaultValue);
+        return new Tag<>(key, readFunction, writeFunction, defaultValue, path);
     }
 
     @Contract(value = "_ -> new", pure = true)
@@ -78,7 +87,15 @@ public class Tag<T> {
                 // Write
                 writeMap.andThen(writeFunction),
                 // Default value
-                () -> readMap.apply(createDefault()));
+                () -> readMap.apply(createDefault()),
+                path);
+    }
+
+    @ApiStatus.Experimental
+    @Contract(value = "_ -> new", pure = true)
+    public Tag<T> path(@NotNull String @Nullable ... path) {
+        final List<PathEntry> entries = path != null ? Arrays.stream(path).map(s -> new PathEntry(s, INDEX_MAP.get(s))).toList() : null;
+        return new Tag<>(key, readFunction, writeFunction, defaultValue, entries);
     }
 
     public @Nullable T read(@NotNull NBTCompoundLike nbt) {

src/main/java/net/minestom/server/tag/TagHandlerImpl.java

@@ -24,16 +24,78 @@ final class TagHandlerImpl implements TagHandler {
     @Override
     public synchronized <T> void setTag(@NotNull Tag<T> tag, @Nullable T value) {
         VarHandle.acquireFence();
-        final int index = tag.index;
+        int tagIndex = tag.index;
+        TagHandlerImpl local = this;
         Entry<?>[] entries = this.entries;
-        final Entry<T> entry = value != null ? new Entry<>(tag, value) : null;
+
-        if (index >= entries.length) {
+        var paths = tag.path;
-            if (value == null)
+        TagHandlerImpl[] pathHandlers = null;
-                return; // no need to create/remove an entry
+        if (paths != null) {
-            this.entries = entries = Arrays.copyOf(entries, index + 1);
+            pathHandlers = new TagHandlerImpl[paths.size()];
+            // Path-able tag
+            int in = 0;
+            for (var path : paths) {
+                final int pathIndex = path.index();
+                if (pathIndex >= entries.length) {
+                    if (value == null) return;
+                    local.entries = entries = Arrays.copyOf(entries, pathIndex + 1);
+                }
+                Entry<?> entry = entries[pathIndex];
+                if (entry == null) {
+                    if (value == null) return;
+                    var updated = new TagHandlerImpl();
+                    entries[pathIndex] = new Entry<>(Tag.tag(path.name(), null, null), updated);
+                    local = updated;
+                } else if (entry.value instanceof TagHandlerImpl handler) {
+                    local = handler;
+                } else {
+                    throw new IllegalStateException("Cannot set a path-able tag on a non-path-able entry");
+                }
+                entries = local.entries;
+                pathHandlers[in++] = local;
+            }
+            // Handle removal if the tag was present (recursively)
+            if (value == null) {
+                pathHandlers[pathHandlers.length - 1].entries[tagIndex] = null;
+                boolean empty = false;
+                removalLoop:
+                for (int i = pathHandlers.length - 1; i >= 0; i--) {
+                    TagHandlerImpl handler = pathHandlers[i];
+                    Entry<?>[] entr = handler.entries;
+                    // Verify if the handler is empty
+                    empty = tagIndex >= entr.length;
+                    if (!empty) {
+                        empty = true;
+                        for (var entry : entr) {
+                            if (entry != null) {
+                                empty = false;
+                                continue removalLoop;
+                            }
+                        }
+                    }
+                    if (i > 0) {
+                        TagHandlerImpl parent = pathHandlers[i - 1];
+                        parent.entries[paths.get(i).index()] = null;
+                    }
+                }
+                if (empty) {
+                    // Remove the root handler
+                    local = this;
+                    entries = this.entries;
+                    tagIndex = paths.get(0).index();
+                }
+            }
         }
-        entries[index] = entry;
+        // Normal tag
+        if (tagIndex >= entries.length) {
+            if (value == null) return;
+            local.entries = entries = Arrays.copyOf(entries, tagIndex + 1);
+        }
+        entries[tagIndex] = value != null ? new Entry<>(tag, value) : null;
         this.cache = null;
+        if (pathHandlers != null) {
+            for (var handler : pathHandlers) handler.cache = null;
+        }
         VarHandle.releaseFence();
     }
 
@@ -76,7 +138,13 @@ final class TagHandlerImpl implements TagHandler {
                 for (Entry<?> entry : entries) {
                     if (entry == null) continue;
                     final Tag<?> tag = entry.tag;
-                    tag.writeUnsafe(tmp, entry.value);
+                    final Object value = entry.value;
+                    if (value instanceof TagHandler handler) {
+                        // Path-able entry
+                        tmp.put(tag.getKey(), handler.asCompound());
+                    } else {
+                        tag.writeUnsafe(tmp, value);
+                    }
                 }
                 cache = !tmp.isEmpty() ? new Cache(entries, tmp.toCompound()) : Cache.EMPTY;
             } else {
@@ -90,7 +158,7 @@ final class TagHandlerImpl implements TagHandler {
 
     private static final class Entry<T> {
         final Tag<T> tag;
-        final T value;
+        final T value; // TagHandler type for path-able tags
         volatile NBT nbt;
 
         Entry(Tag<T> tag, T value) {
@@ -110,19 +178,35 @@ final class TagHandlerImpl implements TagHandler {
 
     private static <T> T read(Entry<?>[] entries, Tag<T> tag) {
         final int index = tag.index;
-        final Entry<?> entry;
+        Entry<?> entry;
+        if (tag.path != null) {
+            // Must be a path-able entry
+            var paths = tag.path;
+            for (var path : paths) {
+                final int pathIndex = path.index();
+                if (pathIndex >= entries.length || (entry = entries[pathIndex]) == null) {
+                    return tag.createDefault();
+                }
+                if (entry.value instanceof TagHandlerImpl handler) {
+                    entries = handler.entries;
+                }
+            }
+        }
         if (index >= entries.length || (entry = entries[index]) == null) {
             return tag.createDefault();
         }
+        final Object value = entry.value;
+        if (value instanceof TagHandlerImpl)
+            throw new IllegalStateException("Cannot read path-able tag " + tag.getKey());
         final Tag entryTag = entry.tag;
         if (entryTag == tag) {
             // Tag is the same, return the value
             //noinspection unchecked
-            return (T) entry.value;
+            return (T) value;
         }
         // Value must be parsed from nbt if the tag is different
         NBT nbt = entry.nbt;
-        if (nbt == null) entry.nbt = nbt = (NBT) entryTag.writeFunction.apply(entry.value);
+        if (nbt == null) entry.nbt = nbt = (NBT) entryTag.writeFunction.apply(value);
         try {
             return tag.readFunction.apply(nbt);
         } catch (ClassCastException e) {

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

@@ -1,9 +1,13 @@
 package net.minestom.server.api;
 
+import org.jglrxavpok.hephaistos.nbt.NBTCompound;
+import org.jglrxavpok.hephaistos.nbt.NBTException;
+import org.jglrxavpok.hephaistos.parser.SNBTParser;
+
+import java.io.StringReader;
 import java.lang.ref.WeakReference;
 
-import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.*;
-import static org.junit.jupiter.api.Assertions.assertTrue;
 
 public class TestUtils {
     public static void waitUntilCleared(WeakReference<?> ref) {
@@ -16,6 +20,15 @@ public class TestUtils {
         }
     }
 
+    public static void assertEqualsSNBT(String snbt, NBTCompound compound) {
+        try {
+            final var converted = (NBTCompound) new SNBTParser(new StringReader(snbt)).parse();
+            assertEquals(converted, compound);
+        } catch (NBTException e) {
+            fail(e);
+        }
+    }
+
     public static void assertEqualsIgnoreSpace(String s1, String s2, boolean matchCase) {
         final String val1 = stripExtraSpaces(s1);
         final String val2 = stripExtraSpaces(s2);

src/test/java/net/minestom/server/tag/TagPathTest.java

@@ -0,0 +1,267 @@
+package net.minestom.server.tag;
+
+import org.jetbrains.annotations.NotNull;
+import org.jetbrains.annotations.Nullable;
+import org.junit.jupiter.api.Test;
+
+import static net.minestom.server.api.TestUtils.assertEqualsSNBT;
+import static org.junit.jupiter.api.Assertions.*;
+
+public class TagPathTest {
+
+    @Test
+    public void basic() {
+        var handler = TagHandler.newHandler();
+        var tag = Tag.Integer("number");
+        var path = tag.path("display");
+        handler.setTag(path, 5);
+        assertEquals(5, handler.getTag(path));
+        assertNull(handler.getTag(tag));
+
+        handler.setTag(path, 6);
+        assertEquals(6, handler.getTag(path));
+        assertNull(handler.getTag(tag));
+
+        handler.removeTag(path);
+        assertNull(handler.getTag(path));
+        assertNull(handler.getTag(tag));
+    }
+
+    @Test
+    public void emptyRemoval() {
+        var handler = TagHandler.newHandler();
+        var tag = Tag.Integer("number").path("display");
+        handler.removeTag(tag);
+        assertNull(handler.getTag(tag));
+        assertEqualsSNBT("{}", handler.asCompound());
+    }
+
+    @Test
+    public void snbt() {
+        var handler = TagHandler.newHandler();
+        var tag = Tag.Integer("number").path("display");
+        handler.setTag(tag, 5);
+        assertEqualsSNBT("""
+                {
+                  "display": {
+                    "number":5
+                  }
+                }
+                """, handler.asCompound());
+
+        handler.removeTag(tag);
+        assertEqualsSNBT("{}", handler.asCompound());
+    }
+
+    @Test
+    public void doubleSnbt() {
+        var handler = TagHandler.newHandler();
+        var tag = Tag.Integer("number").path("display");
+        var tag1 = Tag.String("string").path("display");
+        handler.setTag(tag, 5);
+        handler.setTag(tag1, "test");
+
+        assertEqualsSNBT("""
+                {
+                  "display": {
+                    "string":"test",
+                    "number":5
+                  }
+                }
+                """, handler.asCompound());
+
+        handler.removeTag(tag);
+        assertEqualsSNBT("""
+                {
+                  "display": {
+                    "string":"test"
+                  }
+                }
+                """, handler.asCompound());
+
+        handler.removeTag(tag1);
+        assertEqualsSNBT("{}", handler.asCompound());
+    }
+
+    @Test
+    public void differentPath() {
+        var handler = TagHandler.newHandler();
+        var tag = Tag.Integer("number");
+        var path = tag.path("display");
+        handler.setTag(tag, 5);
+        assertEqualsSNBT("""
+                {
+                  "number":5
+                }
+                """, handler.asCompound());
+
+        handler.setTag(path, 5);
+        assertEqualsSNBT("""
+                {
+                  "number":5,
+                  "display": {
+                    "number":5
+                  }
+                }
+                """, handler.asCompound());
+
+        handler.removeTag(tag);
+        assertEqualsSNBT("""
+                {
+                  "display": {
+                    "number":5
+                  }
+                }
+                """, handler.asCompound());
+    }
+
+    @Test
+    public void overrideSnbt() {
+        var handler = TagHandler.newHandler();
+        var tag = Tag.Integer("key");
+        var tag1 = Tag.Integer("value").path("key");
+        handler.setTag(tag, 5);
+        assertEqualsSNBT("""
+                {
+                  "key":5
+                }
+                """, handler.asCompound());
+
+        assertThrows(IllegalStateException.class, () -> handler.setTag(tag1, 5));
+    }
+
+    @Test
+    public void forgetPath() {
+        var handler = TagHandler.newHandler();
+        var tag = Tag.Integer("key");
+        var path = Tag.Integer("value").path("key");
+        handler.setTag(path, 5);
+        assertThrows(IllegalStateException.class, () -> handler.getTag(tag));
+    }
+
+    @Test
+    public void chaining() {
+        var handler = TagHandler.newHandler();
+        var tag = Tag.Integer("key");
+        var path = Tag.Integer("key").path("first", "second");
+        handler.setTag(path, 5);
+        assertEqualsSNBT("""
+                {
+                  "first": {
+                    "second": {
+                      "key":5
+                    }
+                  }
+                }
+                """, handler.asCompound());
+
+        assertEquals(5, handler.getTag(path));
+        assertNull(handler.getTag(tag));
+
+        handler.removeTag(path);
+        assertEqualsSNBT("{}", handler.asCompound());
+    }
+
+    @Test
+    public void chainingDouble() {
+        var handler = TagHandler.newHandler();
+        var path = Tag.Integer("key").path("first", "second");
+        var path1 = Tag.Integer("key").path("first");
+        handler.setTag(path, 5);
+        assertEqualsSNBT("""
+                {
+                  "first": {
+                    "second": {
+                      "key":5
+                    }
+                  }
+                }
+                """, handler.asCompound());
+        assertEquals(5, handler.getTag(path));
+
+        handler.setTag(path1, 5);
+        assertEqualsSNBT("""
+                {
+                  "first": {
+                    "key":5,
+                    "second": {
+                      "key":5
+                    }
+                  }
+                }
+                """, handler.asCompound());
+        assertEquals(5, handler.getTag(path));
+        assertEquals(5, handler.getTag(path1));
+
+        handler.removeTag(path);
+        assertEqualsSNBT("""
+                {
+                  "first": {
+                    "key":5
+                  }
+                }
+                """, handler.asCompound());
+
+        handler.removeTag(path1);
+        assertEqualsSNBT("{}", handler.asCompound());
+    }
+
+    @Test
+    public void structureObstruction() {
+        record Entry(int value) {
+        }
+
+        var handler = TagHandler.newHandler();
+        var tag = Tag.Integer("value");
+        var struct = Tag.Structure("struct", new TagSerializer<Entry>() {
+            private static final Tag<Integer> VALUE_TAG = Tag.Integer("value");
+
+            @Override
+            public @Nullable Entry read(@NotNull TagReadable reader) {
+                final Integer value = reader.getTag(VALUE_TAG);
+                return value != null ? new Entry(value) : null;
+            }
+
+            @Override
+            public void write(@NotNull TagWritable writer, @Nullable Entry value) {
+                writer.setTag(VALUE_TAG, value != null ? value.value : null);
+            }
+        });
+
+        handler.setTag(struct, new Entry(5));
+        assertEqualsSNBT("""
+                {
+                  "struct": {
+                    "value":5
+                  }
+                }
+                """, handler.asCompound());
+
+        handler.setTag(tag, 5);
+        assertEqualsSNBT("""
+                {
+                  value:5,
+                  "struct": {
+                    "value":5
+                  }
+                }
+                """, handler.asCompound());
+
+        // Cannot enter a structure from a path tag
+        assertThrows(IllegalStateException.class, () -> handler.setTag(tag.path("struct"), 5));
+    }
+
+    @Test
+    public void tagObstruction() {
+        var handler = TagHandler.newHandler();
+        var tag = Tag.Integer("key");
+        var path = Tag.Integer("value").path("key", "second");
+        handler.setTag(tag, 5);
+        assertEqualsSNBT("""
+                {
+                  "key":5
+                }
+                """, handler.asCompound());
+        assertThrows(IllegalStateException.class, () -> handler.setTag(path, 5));
+    }
+}