Re-implement backing state maps (#2883)

* Re-implement backing state maps

This uses a single array combined with efficiently computed indexes to
avoid needing to do entire map equality comparisons or hold large
hash-based tables, improving memory and CPU usage.

* Give a good error message on property mismatch

Author: octylFractal

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

@@ -19,23 +19,12 @@
 
 package com.sk89q.worldedit.world.block;
 
-import com.google.common.collect.ImmutableList;
-import com.google.common.collect.ImmutableMap;
-import com.google.common.collect.Lists;
-import com.google.common.collect.Table;
-import com.google.common.collect.Tables;
-import com.sk89q.worldedit.WorldEdit;
-import com.sk89q.worldedit.extension.platform.Capability;
-import com.sk89q.worldedit.extension.platform.Watchdog;
 import com.sk89q.worldedit.internal.block.BlockStateIdAccess;
 import com.sk89q.worldedit.registry.state.Property;
 import com.sk89q.worldedit.util.concurrency.LazyReference;
-import it.unimi.dsi.fastutil.objects.Reference2ObjectArrayMap;
 import org.enginehub.linbus.tree.LinCompoundTag;
 
-import java.util.Collections;
 import java.util.HashSet;
-import java.util.List;
 import java.util.Map;
 import java.util.Objects;
 import java.util.Set;
@@ -62,141 +51,46 @@ public void setInternalId(BlockState blockState, int internalId) {
 
     private final BlockType blockType;
     private final Map<Property<?>, Object> values;
+    private final int stateListIndex;
 
     private final BaseBlock emptyBaseBlock;
     private final LazyReference<String> lazyStringRepresentation;
 
-    // Neighbouring state table.
-    private Table<Property<?>, Object, BlockState> states;
-
     /**
      * The internal ID of the block state.
      */
     private volatile int internalId = BlockStateIdAccess.invalidId();
 
-    BlockState(BlockType blockType) {
-        this(blockType, Collections.emptyMap());
-    }
-
-    BlockState(BlockType blockType, Map<Property<?>, Object> values) {
+    BlockState(BlockType blockType, Map<Property<?>, Object> values, int stateListIndex) {
         this.blockType = blockType;
         this.values = values;
+        this.stateListIndex = stateListIndex;
         this.emptyBaseBlock = new BaseBlock(this);
         this.lazyStringRepresentation = LazyReference.from(BlockStateHolder.super::getAsString);
     }
 
-    /**
-     * Generates a map of all possible states for a block type.
-     *
-     * @param blockType The block type
-     * @return The map of states
-     */
-    static Map<Map<Property<?>, Object>, BlockState> generateStateMap(BlockType blockType) {
-        List<? extends Property<?>> properties = blockType.getProperties();
-        ImmutableMap.Builder<Map<Property<?>, Object>, BlockState> stateMapBuilder = null;
-
-        if (!properties.isEmpty()) {
-            // Create a list of lists of values, with a copy of the underlying lists
-            List<List<Object>> separatedValues = Lists.newArrayListWithCapacity(properties.size());
-            for (Property<?> prop : properties) {
-                separatedValues.add(ImmutableList.copyOf(prop.values()));
-            }
-
-            List<List<Object>> valueLists = Lists.cartesianProduct(separatedValues);
-            stateMapBuilder = ImmutableMap.builderWithExpectedSize(valueLists.size());
-            for (List<Object> valueList : valueLists) {
-                int valueCount = valueList.size();
-                Map<Property<?>, Object> valueMap = new Reference2ObjectArrayMap<>(valueCount);
-                for (int i = 0; i < valueCount; i++) {
-                    Property<?> property = properties.get(i);
-                    Object value = valueList.get(i);
-                    valueMap.put(property, value);
-                }
-                valueMap = Collections.unmodifiableMap(valueMap);
-                stateMapBuilder.put(valueMap, new BlockState(blockType, valueMap));
-            }
-        }
-
-        ImmutableMap<Map<Property<?>, Object>, BlockState> stateMap;
-
-        if (stateMapBuilder == null) {
-            // No properties.
-            stateMap = ImmutableMap.of(ImmutableMap.of(), new BlockState(blockType));
-        } else {
-            stateMap = stateMapBuilder.build();
-        }
-
-        Watchdog watchdog = WorldEdit.getInstance().getPlatformManager().queryCapability(Capability.GAME_HOOKS)
-            .getWatchdog();
-        long startTime = System.currentTimeMillis();
-
-        for (BlockState state : stateMap.values()) {
-            state.populate(stateMap);
-
-            // Sometimes loading can take a while. This is the perfect spot to let MC know we're working.
-            if (watchdog != null) {
-                watchdog.tick();
-            }
-        }
-        long timeTaken = System.currentTimeMillis() - startTime;
-        if (timeTaken > 5000) {
-            WorldEdit.logger.warn("Took more than 5 seconds to generate complete state map for " + blockType.id() + ". This block is likely improperly using properties. State count: " + stateMap.size() + ". " + timeTaken + "ms elapsed.");
-        }
-
-        return stateMap;
-    }
-
-    /**
-     * Creates the underlying state table for object lookups.
-     *
-     * @param stateMap The state map to generate the table from
-     */
-    private void populate(Map<Map<Property<?>, Object>, BlockState> stateMap) {
-        Table<Property<?>, Object, BlockState> table = Tables.newCustomTable(
-                new Reference2ObjectArrayMap<>(this.values.size()),
-                Reference2ObjectArrayMap::new
-        );
-
-        for (final Map.Entry<Property<?>, Object> entry : this.values.entrySet()) {
-            final Property<Object> property = (Property<Object>) entry.getKey();
-
-            for (Object value : property.values()) {
-                if (value != entry.getValue()) {
-                    BlockState modifiedState = stateMap.get(this.withValue(property, value));
-                    if (modifiedState != null) {
-                        table.put(property, value, modifiedState);
-                    } else {
-                        WorldEdit.logger.warn(stateMap);
-                        WorldEdit.logger.warn("Found a null state at " + this.withValue(property, value));
-                    }
-                }
-            }
-        }
-
-        this.states = Tables.unmodifiableTable(table);
-    }
-
-    private <V> Map<Property<?>, Object> withValue(final Property<V> property, final V value) {
-        final Map<Property<?>, Object> values = new Reference2ObjectArrayMap<>(this.values.size());
-        for (Map.Entry<Property<?>, Object> entry : this.values.entrySet()) {
-            if (entry.getKey().equals(property)) {
-                values.put(entry.getKey(), value);
-            } else {
-                values.put(entry.getKey(), entry.getValue());
-            }
-        }
-        return Collections.unmodifiableMap(values);
-    }
-
     @Override
     public BlockType getBlockType() {
         return this.blockType;
     }
 
     @Override
     public <V> BlockState with(final Property<V> property, final V value) {
-        BlockState result = states.get(property, value);
-        return result == null ? this : result;
+        if (this.stateListIndex == -1) {
+            return this;
+        }
+        Object currentValue = this.values.get(property);
+        if (Objects.equals(currentValue, value)) {
+            return this;
+        }
+
+        int newIndex = blockType.getInternalStateList().updateIndexOrInvalid(
+            this.stateListIndex, property, currentValue, value
+        );
+        if (newIndex == -1) {
+            return this;
+        }
+        return blockType.getInternalStateList().get(newIndex);
     }
 
     @Override

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

@@ -20,7 +20,6 @@
 package com.sk89q.worldedit.world.block;
 
 import com.google.common.collect.ImmutableList;
-import com.google.common.collect.Iterables;
 import com.sk89q.worldedit.WorldEdit;
 import com.sk89q.worldedit.extension.platform.Capability;
 import com.sk89q.worldedit.registry.Keyed;
@@ -32,9 +31,10 @@
 import com.sk89q.worldedit.world.item.ItemTypes;
 import com.sk89q.worldedit.world.registry.BlockMaterial;
 import com.sk89q.worldedit.world.registry.LegacyMapper;
-import it.unimi.dsi.fastutil.objects.Reference2ObjectArrayMap;
+import it.unimi.dsi.fastutil.objects.Object2ObjectArrayMap;
+import it.unimi.dsi.fastutil.objects.Object2ObjectMaps;
 
-import java.util.Collections;
+import java.util.Arrays;
 import java.util.List;
 import java.util.Map;
 import java.util.function.Function;
@@ -46,25 +46,34 @@ public class BlockType implements Keyed {
 
     public static final NamespacedRegistry<BlockType> REGISTRY = new NamespacedRegistry<>("block type", "block_type", "minecraft", true);
 
+    private static Map<String, ? extends Property<?>> computeProperties(BlockType self) {
+        Map<String, ? extends Property<?>> propertiesMap = WorldEdit.getInstance().getPlatformManager()
+            .queryCapability(Capability.GAME_HOOKS).getRegistries().getBlockRegistry().getProperties(self);
+        String[] propertyNames = propertiesMap.keySet().toArray(new String[0]);
+        Arrays.sort(propertyNames);
+        Object[] properties = new Object[propertyNames.length];
+        for (int i = 0; i < propertyNames.length; i++) {
+            properties[i] = propertiesMap.get(propertyNames[i]);
+        }
+        return Object2ObjectMaps.unmodifiable(new Object2ObjectArrayMap<>(propertyNames, properties));
+    }
+
     private final String id;
-    private final Function<BlockState, BlockState> values;
-    private final LazyReference<BlockState> defaultState
-        = LazyReference.from(this::computeDefaultState);
+    private final LazyReference<BlockState> defaultState;
     @SuppressWarnings("this-escape")
     private final LazyReference<FuzzyBlockState> emptyFuzzy
         = LazyReference.from(() -> new FuzzyBlockState(this));
     @SuppressWarnings("this-escape")
     private final LazyReference<Map<String, ? extends Property<?>>> properties
-        = LazyReference.from(this::computeProperties);
+        = LazyReference.from(() -> computeProperties(this));
+    @SuppressWarnings("this-escape")
+    private final LazyReference<BlockTypeStateList> internalStateList =
+        LazyReference.from(() -> BlockTypeStateList.createFor(this));
     @SuppressWarnings("this-escape")
     private final LazyReference<BlockMaterial> blockMaterial
         = LazyReference.from(() -> WorldEdit.getInstance().getPlatformManager()
         .queryCapability(Capability.GAME_HOOKS).getRegistries().getBlockRegistry().getMaterial(this));
     @SuppressWarnings("this-escape")
-    private final LazyReference<Map<Map<Property<?>, Object>, BlockState>> blockStatesMap
-        = LazyReference.from(() -> BlockState.generateStateMap(this));
-
-    @SuppressWarnings("this-escape")
     @Deprecated
     private final LazyReference<String> name = LazyReference.from(() -> WorldEdit.getInstance().getPlatformManager()
         .queryCapability(Capability.GAME_HOOKS).getRegistries().getBlockRegistry().getName(this));
@@ -77,36 +86,25 @@ public BlockType(String id) {
         this(id, null);
     }
 
-    public BlockType(String id, Function<BlockState, BlockState> values) {
+    public BlockType(String id, Function<BlockState, BlockState> applyDefaultValues) {
         // If it has no namespace, assume minecraft.
         if (!id.contains(":")) {
             id = "minecraft:" + id;
         }
         this.id = id;
-        this.values = values;
+        this.defaultState = LazyReference.from(() -> computeDefaultState(applyDefaultValues));
     }
 
-    private Map<String, ? extends Property<?>> computeProperties() {
-        var propertiesMap = WorldEdit.getInstance().getPlatformManager()
-                .queryCapability(Capability.GAME_HOOKS).getRegistries().getBlockRegistry().getProperties(this);
-        List<String> sortedPropertyNames = propertiesMap.keySet().stream().sorted().toList();
-        Map<String, Property<?>> sortedPropertiesMap = new Reference2ObjectArrayMap<>(propertiesMap.size());
-        for (String propertyName : sortedPropertyNames) {
-            sortedPropertiesMap.put(propertyName, propertiesMap.get(propertyName));
-        }
-        return Collections.unmodifiableMap(sortedPropertiesMap);
+    BlockTypeStateList getInternalStateList() {
+        return internalStateList.getValue();
     }
 
-    private BlockState computeDefaultState() {
-        BlockState defaultState = Iterables.getFirst(getBlockStatesMap().values(), null);
-        if (values != null) {
-            defaultState = values.apply(defaultState);
+    private BlockState computeDefaultState(Function<BlockState, BlockState> applyDefaultValues) {
+        BlockState state = getInternalStateList().getFirst();
+        if (applyDefaultValues != null) {
+            state = applyDefaultValues.apply(state);
         }
-        return defaultState;
-    }
-
-    private Map<Map<Property<?>, Object>, BlockState> getBlockStatesMap() {
-        return blockStatesMap.getValue();
+        return state;
     }
 
     /**
@@ -190,7 +188,7 @@ public FuzzyBlockState getFuzzyMatcher() {
      * @return All possible states
      */
     public List<BlockState> getAllStates() {
-        return ImmutableList.copyOf(getBlockStatesMap().values());
+        return getInternalStateList();
     }
 
     /**
@@ -199,9 +197,9 @@ public List<BlockState> getAllStates() {
      * @return The state, if it exists
      */
     public BlockState getState(Map<Property<?>, Object> key) {
-        BlockState state = getBlockStatesMap().get(key);
-        checkArgument(state != null, "%s has no state for %s", this, key);
-        return state;
+        BlockTypeStateList map = getInternalStateList();
+        int index = map.calculateIndex(key);
+        return map.get(index);
     }
 
     /**

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

@@ -0,0 +1,62 @@
+/*
+ * WorldEdit, a Minecraft world manipulation toolkit
+ * Copyright (C) sk89q <http://www.sk89q.com>
+ * Copyright (C) WorldEdit team and contributors
+ *
+ * This program is free software: you can redistribute it and/or modify
+ * it under the terms of the GNU General Public License as published by
+ * the Free Software Foundation, either version 3 of the License, or
+ * (at your option) any later version.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ * GNU General Public License for more details.
+ *
+ * You should have received a copy of the GNU General Public License
+ * along with this program.  If not, see <https://www.gnu.org/licenses/>.
+ */
+
+package com.sk89q.worldedit.world.block;
+
+import com.google.errorprone.annotations.Immutable;
+import com.sk89q.worldedit.registry.state.Property;
+import it.unimi.dsi.fastutil.objects.AbstractObjectList;
+
+import java.util.Collections;
+import java.util.Map;
+
+/**
+ * A specialized list for looking up block states from a block type.
+ */
+// Using AbstractObjectList over AbstractList since we don't need the modCount functionality
+@Immutable
+abstract class BlockTypeStateList extends AbstractObjectList<BlockState> {
+    static BlockTypeStateList createFor(BlockType blockType) {
+        if (blockType.getProperties().isEmpty()) {
+            // Special case, we have only one state: the default state
+            return new SingletonBlockTypeStateList(new BlockState(blockType, Collections.emptyMap(), 0));
+        }
+        return new DefaultBlockTypeStateList(blockType);
+    }
+
+    /**
+     * Calculates the index in the states array for the given property values.
+     * This can later be used to perform fast lookups by replacing only a specific property.
+     *
+     * @param state the map of property values
+     * @return the index in the states
+     */
+    public abstract int calculateIndex(Map<Property<?>, ?> state);
+
+    /**
+     * Updates the current index by changing a single property's value.
+     *
+     * @param currentIndex the current index
+     * @param property the property to change
+     * @param oldValue the old value
+     * @param newValue the new value
+     * @return the updated index, or {@code -1} if the property or value is invalid
+     */
+    public abstract int updateIndexOrInvalid(int currentIndex, Property<?> property, Object oldValue, Object newValue);
+}