Fix installing modules after the first or removing modules one at a time not working. If the value set in the module tweaker for a config option is no longer valid after changing the number of installed modules, the value will be reset to the default

Author: pupnewfster

src/main/java/mekanism/common/content/gear/Module.java

@@ -28,15 +28,10 @@
 import mekanism.common.util.MekanismUtils;
 import mekanism.common.util.StorageUtils;
 import net.minecraft.MethodsReturnNonnullByDefault;
-import net.minecraft.core.HolderLookup;
-import net.minecraft.nbt.CompoundTag;
-import net.minecraft.nbt.NbtOps;
-import net.minecraft.nbt.Tag;
 import net.minecraft.network.RegistryFriendlyByteBuf;
 import net.minecraft.network.chat.Component;
 import net.minecraft.network.codec.ByteBufCodecs;
 import net.minecraft.network.codec.StreamCodec;
-import net.minecraft.resources.RegistryOps;
 import net.minecraft.resources.ResourceLocation;
 import net.minecraft.util.ExtraCodecs;
 import net.minecraft.world.entity.LivingEntity;
@@ -217,12 +212,31 @@ public boolean isEnabled() {
         return enabled;
     }
 
-    Module<MODULE> withReplacedInstallCount(HolderLookup.Provider provider, int installed) {
-        RegistryOps<Tag> registryOps = provider.createSerializationContext(NbtOps.INSTANCE);
-        //TODO: See if we can come up with a cleaner way to do this
-        CompoundTag tag = (CompoundTag) Module.CODEC.encodeStart(registryOps, this).getOrThrow();
-        tag.putInt(SerializationConstants.AMOUNT, installed);
-        return (Module<MODULE>) Module.CODEC.decode(registryOps, tag).getOrThrow().getFirst();
+    Module<MODULE> withReplacedInstallCount(int installed) {
+        List<ModuleConfig<?>> moduleConfigs = data.defaultConfigs(installed);
+        List<ModuleConfig<?>> copiedConfigs = new ArrayList<>(moduleConfigs.size());
+        for (ModuleConfig<?> moduleConfig : moduleConfigs) {
+            ResourceLocation name = moduleConfig.name();
+            ModuleConfig<?> existingConfig = configItemsByName.get(name);
+            if (existingConfig == null) {
+                copiedConfigs.add(moduleConfig);
+            } else if (moduleConfig.getClass() == existingConfig.getClass()) {
+                copiedConfigs.add(configWithValue(moduleConfig, existingConfig));
+            } else {
+                throw new IllegalStateException("Expected module config " + name + " to have the same class regardless of installed count.");
+            }
+        }
+        return new Module<>(data, installed, List.copyOf(copiedConfigs));
+    }
+
+    @SuppressWarnings("unchecked")
+    private <CONFIG> ModuleConfig<CONFIG> configWithValue(ModuleConfig<CONFIG> defaultConfig, ModuleConfig<?> existing) {
+        try {
+            return defaultConfig.with(((ModuleConfig<CONFIG>) existing).get());
+        } catch (IllegalArgumentException e) {
+            //If the existing value isn't valid for the new config, fallback to the default value
+            return defaultConfig;
+        }
     }
 
     Module<MODULE> withReplacedConfig(ModuleConfig<?> config) {

src/main/java/mekanism/common/content/gear/ModuleContainer.java

@@ -292,7 +292,7 @@ public <MODULE extends ICustomModule<MODULE>> int addModule(HolderLookup.Provide
                 //Nothing to actually install because we are already at the max stack size
                 return 0;
             }
-            module = module.withReplacedInstallCount(provider, module.getInstalledCount() + toInstall);
+            module = module.withReplacedInstallCount(module.getInstalledCount() + toInstall);
         }
         //Add the module to the list of tracked and known modules if necessary or replace the existing value
         SequencedMap<ModuleData<?>, Module<?>> copiedModules = new LinkedHashMap<>(typedModules);
@@ -332,7 +332,7 @@ public <MODULE extends ICustomModule<MODULE>> void removeModule(HolderLookup.Pro
                     }
                 }
             } else {//update the module with the new installed count
-                module = module.withReplacedInstallCount(provider, installed);
+                module = module.withReplacedInstallCount(installed);
                 copiedModules.put(type, module);
                 //Update the level of any corresponding enchantment
                 adjustedEnchantments = updateEnchantment(provider, module, null);