Re-load internal state IDs when WORLD_EDITING cap is re-loaded

worldedit-core/src/main/java/com/sk89q/worldedit/extension/platform/Capability.java

@@ -19,6 +19,11 @@
 
 package com.sk89q.worldedit.extension.platform;
 
+import com.sk89q.worldedit.internal.block.BlockStateIdAccess;
+import com.sk89q.worldedit.world.block.BlockState;
+import com.sk89q.worldedit.world.block.BlockType;
+import com.sk89q.worldedit.world.registry.BlockRegistry;
+
 /**
  * A collection of capabilities that a {@link Platform} may support.
  */
@@ -73,7 +78,22 @@ public enum Capability {
     /**
      * The capability of a platform to perform modifications to a world.
      */
-    WORLD_EDITING;
+    WORLD_EDITING {
+        @Override
+        void initialize(PlatformManager platformManager, Platform platform) {
+            BlockRegistry blockRegistry = platform.getRegistries().getBlockRegistry();
+            for (BlockType type : BlockType.REGISTRY) {
+                for (BlockState state : type.getAllStates()) {
+                    BlockStateIdAccess.register(state, blockRegistry.getInternalBlockStateId(state));
+                }
+            }
+        }
+
+        @Override
+        void unload(PlatformManager platformManager, Platform platform) {
+            BlockStateIdAccess.clear();
+        }
+    };
 
     void initialize(PlatformManager platformManager, Platform platform) {
 

worldedit-core/src/main/java/com/sk89q/worldedit/internal/block/BlockStateIdAccess.java

@@ -19,57 +19,46 @@
 
 package com.sk89q.worldedit.internal.block;
 
+import com.google.common.collect.BiMap;
+import com.google.common.collect.HashBiMap;
 import com.sk89q.worldedit.world.block.BlockState;
 
 import javax.annotation.Nullable;
 import java.util.Arrays;
+import java.util.Map;
 import java.util.OptionalInt;
 
 import static com.google.common.base.Preconditions.checkState;
 
 public final class BlockStateIdAccess {
 
-    private BlockStateIdAccess() {
-    }
-
-    public interface Provider {
-
-        OptionalInt getBlockStateId(BlockState holder);
-    }
-
-    private static Provider blockStateStateId;
-
-    public static void setBlockStateStateId(Provider blockStateStateId) {
-        BlockStateIdAccess.blockStateStateId = blockStateStateId;
-    }
+    private static final BiMap<BlockState, Integer> ASSIGNED_IDS = HashBiMap.create(2 << 13);
 
     public static OptionalInt getBlockStateId(BlockState holder) {
-        return blockStateStateId.getBlockStateId(holder);
+        Integer value = ASSIGNED_IDS.get(holder);
+        return value == null ? OptionalInt.empty() : OptionalInt.of(value);
     }
 
     public static @Nullable BlockState getBlockStateById(int id) {
-        return id < blockStates.length ? blockStates[id] : null;
+        return ASSIGNED_IDS.inverse().get(id);
     }
 
-    private static BlockState[] blockStates = new BlockState[2 << 13];
-
-    public static void register(BlockState blockState) {
-        OptionalInt id = getBlockStateId(blockState);
+    public static void register(BlockState blockState, OptionalInt id) {
         if (id.isPresent()) {
             int i = id.getAsInt();
-            if (i >= blockStates.length) {
-                int curLength = blockStates.length;
-                do {
-                    curLength += curLength >> 1;
-                } while (i >= curLength);
-                blockStates = Arrays.copyOf(blockStates, curLength);
-            }
-            BlockState existing = blockStates[i];
+            BlockState existing = ASSIGNED_IDS.inverse().get(i);
             checkState(existing == null || existing == blockState,
                 "BlockState %s is using the same block ID (%s) as BlockState %s",
                 blockState, i, existing);
-            blockStates[i] = blockState;
+            ASSIGNED_IDS.put(blockState, i);
         }
     }
 
+    public static void clear() {
+        ASSIGNED_IDS.clear();
+    }
+
+    private BlockStateIdAccess() {
+    }
+
 }

worldedit-core/src/main/java/com/sk89q/worldedit/world/block/BlockState.java

@@ -27,7 +27,6 @@ import com.google.common.collect.Table;
 import com.sk89q.jnbt.CompoundTag;
 import com.sk89q.worldedit.WorldEdit;
 import com.sk89q.worldedit.extension.platform.Capability;
-import com.sk89q.worldedit.internal.block.BlockStateIdAccess;
 import com.sk89q.worldedit.registry.state.Property;
 import com.sk89q.worldedit.world.registry.BlockRegistry;
 
@@ -38,7 +37,6 @@ import java.util.LinkedHashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.Objects;
-import java.util.OptionalInt;
 import java.util.Set;
 
 /**
@@ -47,13 +45,8 @@ import java.util.Set;
 @SuppressWarnings("unchecked")
 public class BlockState implements BlockStateHolder<BlockState> {
 
-    static {
-        BlockStateIdAccess.setBlockStateStateId(x -> x.internalId);
-    }
-
     private final BlockType blockType;
     private final Map<Property<?>, Object> values;
-    private OptionalInt internalId = OptionalInt.empty();
 
     private BaseBlock emptyBaseBlock;
 
@@ -66,12 +59,6 @@ public class BlockState implements BlockStateHolder<BlockState> {
         this.emptyBaseBlock = new BaseBlock(this);
     }
 
-    BlockState initializeId(BlockRegistry registry) {
-        this.internalId = registry.getInternalBlockStateId(this);
-        BlockStateIdAccess.register(this);
-        return this;
-    }
-
     static Map<Map<Property<?>, Object>, BlockState> generateStateMap(BlockType blockType) {
         BlockRegistry registry = WorldEdit.getInstance().getPlatformManager().queryCapability(Capability.WORLD_EDITING).getRegistries().getBlockRegistry();
         Map<Map<Property<?>, Object>, BlockState> stateMap = new LinkedHashMap<>();
@@ -94,14 +81,13 @@ public class BlockState implements BlockStateHolder<BlockState> {
                     valueMap.put(property, value);
                     stateMaker.setState(property, value);
                 }
-                stateMaker.initializeId(registry);
                 stateMap.put(valueMap, stateMaker);
             }
         }
 
         if (stateMap.isEmpty()) {
             // No properties.
-            stateMap.put(new LinkedHashMap<>(), new BlockState(blockType).initializeId(registry));
+            stateMap.put(new LinkedHashMap<>(), new BlockState(blockType));
         }
 
         for (BlockState state : stateMap.values()) {
@@ -114,11 +100,11 @@ public class BlockState implements BlockStateHolder<BlockState> {
     private void populate(Map<Map<Property<?>, Object>, BlockState> stateMap) {
         final Table<Property<?>, Object, BlockState> states = HashBasedTable.create();
 
-        for(final Map.Entry<Property<?>, Object> entry : this.values.entrySet()) {
+        for (final Map.Entry<Property<?>, Object> entry : this.values.entrySet()) {
             final Property<Object> property = (Property<Object>) entry.getKey();
 
             property.getValues().forEach(value -> {
-                if(value != entry.getValue()) {
+                if (value != entry.getValue()) {
                     BlockState modifiedState = stateMap.get(this.withValue(property, value));
                     if (modifiedState != null) {
                         states.put(property, value, modifiedState);