More preparation for CAS tag write

Signed-off-by: TheMode <themode@outlook.fr>

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

@@ -2,6 +2,7 @@ package net.minestom.server.tag;
 
 import it.unimi.dsi.fastutil.ints.Int2ObjectMap;
 import it.unimi.dsi.fastutil.ints.Int2ObjectOpenHashMap;
+import net.minestom.server.utils.PropertyUtils;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 import org.jetbrains.annotations.UnknownNullability;
@@ -10,24 +11,34 @@ import org.jglrxavpok.hephaistos.nbt.NBTCompound;
 import org.jglrxavpok.hephaistos.nbt.NBTCompoundLike;
 import org.jglrxavpok.hephaistos.nbt.mutable.MutableNBTCompound;
 
+import java.lang.invoke.VarHandle;
 import java.util.function.UnaryOperator;
 
 final class TagHandlerImpl implements TagHandler {
+    private static final boolean CACHE_ENABLE = PropertyUtils.getBoolean("minestom.tag-handler-cache", true);
+
+    private final TagHandlerImpl parent;
     private volatile SPMCMap entries;
     private Cache cache;
 
-    TagHandlerImpl(SPMCMap entries, Cache cache) {
+    TagHandlerImpl(TagHandlerImpl parent, SPMCMap entries, Cache cache) {
+        this.parent = parent;
         this.entries = entries;
         this.cache = cache;
     }
 
-    TagHandlerImpl() {
+    TagHandlerImpl(TagHandlerImpl parent) {
-        this.entries = new SPMCMap();
+        this.parent = parent;
+        this.entries = new SPMCMap(this);
         this.cache = null;
     }
 
+    TagHandlerImpl() {
+        this(null);
+    }
+
     static TagHandlerImpl fromCompound(NBTCompoundLike compound) {
-        TagHandlerImpl handler = new TagHandlerImpl();
+        TagHandlerImpl handler = new TagHandlerImpl(null);
         for (var entry : compound) {
             final Tag<NBT> tag = Tag.NBT(entry.getKey());
             final NBT nbt = entry.getValue();
@@ -46,11 +57,25 @@ final class TagHandlerImpl implements TagHandler {
     public <T> void setTag(@NotNull Tag<T> tag, @Nullable T value) {
         if (tag.isView()) {
             MutableNBTCompound viewCompound = new MutableNBTCompound();
-            tag.writeUnsafe(viewCompound, value);
+            tag.write(viewCompound, value);
             updateContent(viewCompound);
         } else {
             final Entry<?> entry = valueToEntry(tag, value);
-            updateEntry(tag, entry);
+            final int tagIndex = tag.index;
+            final Tag.PathEntry[] paths = tag.path;
+            final boolean present = entry != null;
+            TagHandlerImpl local = this;
+            synchronized (this) {
+                if (paths != null) {
+                    if ((local = traversePathWrite(this, tagIndex, paths, present)) == null)
+                        return;
+                }
+                SPMCMap entries = local.entries;
+                if (present) entries.put(tagIndex, entry);
+                else entries.remove(tagIndex);
+                entries.invalidate();
+                assert !local.entries.rehashed;
+            }
         }
     }
 
@@ -78,7 +103,8 @@ final class TagHandlerImpl implements TagHandler {
 
     @Override
     public synchronized @NotNull TagHandler copy() {
-        return new TagHandlerImpl(entries.clone(), cache);
+        assert parent == null;
+        return new TagHandlerImpl(null, entries.clone(), cache);
     }
 
     @Override
@@ -86,7 +112,7 @@ final class TagHandlerImpl implements TagHandler {
         final TagHandlerImpl converted = fromCompound(compound);
         synchronized (this) {
             this.cache = converted.cache;
-            this.entries = converted.entries;
+            this.entries = new SPMCMap(this, converted.entries);
         }
     }
 
@@ -95,74 +121,70 @@ final class TagHandlerImpl implements TagHandler {
         return updatedCache().compound;
     }
 
-    private synchronized <T> void updateEntry(@NotNull Tag<T> tag, @Nullable Entry<?> value) {
+    private static TagHandlerImpl traversePathWrite(TagHandlerImpl root, int tagIndex,
-        int tagIndex = tag.index;
+                                                    Tag.PathEntry[] paths, boolean present) {
-        TagHandlerImpl local = this;
+        final int length = paths.length;
-
+        TagHandlerImpl local = root;
-        final Tag.PathEntry[] paths = tag.path;
+        TagHandlerImpl[] pathHandlers = new TagHandlerImpl[length];
-        TagHandlerImpl[] pathHandlers = null;
+        for (int i = 0; i < length; i++) {
-        if (paths != null) {
+            final Tag.PathEntry path = paths[i];
-            final int length = paths.length;
+            final int pathIndex = path.index();
-            pathHandlers = new TagHandlerImpl[length];
+            final Entry<?> entry = local.entries.get(pathIndex);
-            for (int i = 0; i < length; i++) {
+            if (entry instanceof PathEntry pathEntry) {
-                final Tag.PathEntry path = paths[i];
+                // Existing path, continue navigating
-                final int pathIndex = path.index();
+                local = pathEntry.value;
-                final Entry<?> entry = local.entries.get(pathIndex);
+            } else {
-                if (entry instanceof PathEntry pathEntry) {
+                if (!present) return null;
-                    // Existing path, continue navigating
+                // Empty path, create a new handler.
-                    local = pathEntry.value;
+                // Slow path is taken if the entry comes from a Structure tag, requiring conversion from NBT
+                TagHandlerImpl tmp = local;
+                local = new TagHandlerImpl(tmp);
+                if (entry != null && entry.updatedNbt() instanceof NBTCompound compound) {
+                    local.updateContent(compound);
+                }
+                tmp.entries.put(pathIndex, new PathEntry(path.name(), local));
+            }
+            pathHandlers[i] = local;
+        }
+        // Handle removal if the tag was present (recursively)
+        if (!present) {
+            // Remove entry
+            TagHandlerImpl targetHandler = pathHandlers[length - 1];
+            if (targetHandler.entries.remove(tagIndex) == null) return null;
+            // Clear empty parents
+            for (int i = length - 1; i >= 0; i--) {
+                final TagHandlerImpl handler = pathHandlers[i];
+                if (!handler.entries.isEmpty()) break;
+                final int pathIndex = paths[i].index();
+                if (i == 0) {
+                    // Remove the root handler
+                    root.entries.remove(pathIndex);
                 } else {
-                    if (value == null) return;
+                    TagHandlerImpl parent = pathHandlers[i - 1];
-                    // Empty path, create a new handler.
+                    parent.entries.remove(pathIndex);
-                    // Slow path is taken if the entry comes from a Structure tag, requiring conversion from NBT
-                    TagHandlerImpl tmp = local;
-                    local = entry != null && entry.updatedNbt() instanceof NBTCompound compound ?
-                            fromCompound(compound) : new TagHandlerImpl();
-                    tmp.entries.put(pathIndex, new PathEntry(path.name(), local));
-                }
-                pathHandlers[i] = local;
-            }
-            // Handle removal if the tag was present (recursively)
-            if (value == null) {
-                // Remove entry
-                if (pathHandlers[length - 1].entries.remove(tagIndex) == null) return;
-                // Clear empty parents
-                for (int i = length - 1; i >= 0; i--) {
-                    final TagHandlerImpl handler = pathHandlers[i];
-                    if (!handler.entries.isEmpty()) break;
-                    final int pathIndex = paths[i].index();
-                    if (i == 0) {
-                        // Remove the root handler
-                        local = this;
-                        tagIndex = pathIndex;
-                    } else {
-                        TagHandlerImpl parent = pathHandlers[i - 1];
-                        parent.entries.remove(pathIndex);
-                    }
                 }
             }
+            targetHandler.entries.invalidate();
+            return null;
         }
-        // Normal tag
+        return local;
-        if (value != null) local.entries.put(tagIndex, value);
-        else local.entries.remove(tagIndex);
-        this.cache = null;
-        if (pathHandlers != null) {
-            for (var handler : pathHandlers) handler.cache = null;
-        }
     }
 
-    private synchronized Cache updatedCache() {
+    private Cache updatedCache() {
-        Cache cache = this.cache;
+        VarHandle.fullFence();
-        if (cache == null) {
+        Cache cache;
-            final var entries = this.entries;
+        if (!CACHE_ENABLE || (cache = this.cache) == null) {
-            if (!entries.isEmpty()) {
+            synchronized (this) {
-                MutableNBTCompound tmp = new MutableNBTCompound();
+                final SPMCMap entries = this.entries;
-                for (Entry<?> entry : entries.values()) {
+                if (!entries.isEmpty()) {
-                    if (entry != null) tmp.put(entry.tag().getKey(), entry.updatedNbt());
+                    MutableNBTCompound tmp = new MutableNBTCompound();
-                }
+                    for (Entry<?> entry : entries.values()) {
-                cache = new Cache(entries.clone(), tmp.toCompound());
+                        if (entry != null) tmp.put(entry.tag().getKey(), entry.updatedNbt());
-            } else cache = Cache.EMPTY;
+                    }
-            this.cache = cache;
+                    cache = new Cache(entries.clone(), tmp.toCompound());
+                } else cache = Cache.EMPTY;
+                this.cache = cache;
+            }
         }
         return cache;
     }
@@ -171,7 +193,7 @@ final class TagHandlerImpl implements TagHandler {
         final Tag.PathEntry[] paths = tag.path;
         if (paths != null) {
             // Must be a path-able entry
-            if ((entries = traversePath(paths, entries)) == null)
+            if ((entries = traversePathRead(paths, entries)) == null)
                 return tag.createDefault();
         }
         final Entry<?> entry;
@@ -191,8 +213,8 @@ final class TagHandlerImpl implements TagHandler {
                 serializerEntry.read(nbt) : tag.createDefault();
     }
 
-    private static Int2ObjectOpenHashMap<Entry<?>> traversePath(Tag.PathEntry[] paths,
+    private static Int2ObjectOpenHashMap<Entry<?>> traversePathRead(Tag.PathEntry[] paths,
-                                                                Int2ObjectOpenHashMap<Entry<?>> entries) {
+                                                                    Int2ObjectOpenHashMap<Entry<?>> entries) {
         for (var path : paths) {
             final Entry<?> entry;
             if ((entry = entries.get(path.index())) == null)
@@ -271,32 +293,27 @@ final class TagHandlerImpl implements TagHandler {
         }
     }
 
-    final class SPMCMap extends Int2ObjectOpenHashMap<Entry<?>> {
+    static final class SPMCMap extends Int2ObjectOpenHashMap<Entry<?>> {
-        boolean rehashed;
+        final TagHandlerImpl handler;
+        volatile boolean rehashed;
 
-        SPMCMap() {
+        SPMCMap(TagHandlerImpl handler) {
-        }
+            super();
-
+            this.handler = handler;
-        SPMCMap(Int2ObjectMap<TagHandlerImpl.Entry<?>> m) {
-            super(m);
-        }
-
-        @Override
-        public TagHandlerImpl.Entry<?> put(int k, TagHandlerImpl.Entry<?> entry) {
             assertState();
-            return super.put(k, entry);
         }
 
-        @Override
+        SPMCMap(TagHandlerImpl handler, Int2ObjectMap<TagHandlerImpl.Entry<?>> m) {
-        public boolean remove(int k, Object v) {
+            super(m.size(), DEFAULT_LOAD_FACTOR);
+            this.handler = handler;
             assertState();
-            return super.remove(k, v);
+            putAll(m);
         }
 
         @Override
         protected void rehash(int newSize) {
             assertState();
-            TagHandlerImpl.this.entries = new SPMCMap(this);
+            this.handler.entries = new SPMCMap(handler, this);
             this.rehashed = true;
         }
 
@@ -305,8 +322,18 @@ final class TagHandlerImpl implements TagHandler {
             return (SPMCMap) super.clone();
         }
 
+        void invalidate() {
+            if (!CACHE_ENABLE) return;
+            TagHandlerImpl tmp = handler;
+            do {
+                tmp.cache = null;
+            } while ((tmp = tmp.parent) != null);
+            VarHandle.fullFence();
+        }
+
         private void assertState() {
             assert !rehashed;
+            assert handler != null;
         }
     }
 }

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

@@ -63,6 +63,19 @@ public class TagNbtTest {
         assertEquals(5, handler.getTag(path));
     }
 
+    @Test
+    public void compoundPathRead() {
+        var handler = TagHandler.newHandler();
+        var nbtTag = Tag.NBT("compound").path("path");
+
+        var nbt = NBT.Compound(Map.of("key", NBT.Int(5)));
+        handler.setTag(nbtTag, nbt);
+        assertEquals(nbt, handler.getTag(nbtTag));
+
+        var path = Tag.Integer("key").path("path", "compound");
+        assertEquals(5, handler.getTag(path));
+    }
+
     @Test
     public void doubleCompoundRead() {
         var handler = TagHandler.newHandler();