Fix loading blockstates too early

Author: Swedz

src/main/java/aztech/modern_industrialization/blocks/structure/StructureMultiblockHatchBlockEntity.java

@@ -117,7 +117,7 @@ public HatchFlags getFlags() {
 
     @Override
     public StructureMember getMemberOverride() {
-        return new StructureMember(preview, members, flags);
+        return new StructureMember(() -> preview, members, flags);
     }
 
     @Override

src/main/java/aztech/modern_industrialization/blocks/structure/StructureMultiblockMemberBlockEntity.java

@@ -81,7 +81,7 @@ public List<StructureMemberTest> getMembers() {
 
     @Override
     public StructureMember getMemberOverride() {
-        return new StructureMember(preview, members, null);
+        return new StructureMember(() -> preview, members, null);
     }
 
     @Override

src/main/java/aztech/modern_industrialization/compat/kubejs/machine/ShapeTemplateHelper.java

@@ -31,6 +31,10 @@
 import net.minecraft.resources.ResourceLocation;
 
 public interface ShapeTemplateHelper {
+    default ShapeTemplate structureShape(ResourceLocation id) {
+        return new ShapeTemplate.Structure(id).build();
+    }
+
     default ShapeTemplate.LayeredBuilder layeredShape(String hatchCasing, String[][] layers) {
         return new ShapeTemplate.LayeredBuilder(MachineCasings.get(hatchCasing), layers);
     }

src/main/java/aztech/modern_industrialization/machines/multiblocks/ShapeTemplate.java

@@ -24,11 +24,13 @@
 package aztech.modern_industrialization.machines.multiblocks;
 
 import aztech.modern_industrialization.machines.models.MachineCasing;
+import aztech.modern_industrialization.machines.multiblocks.structure.MIStructureTemplateManager;
 import java.util.HashMap;
 import java.util.HashSet;
 import java.util.Map;
 import java.util.Set;
 import net.minecraft.core.BlockPos;
+import net.minecraft.resources.ResourceLocation;
 import org.jetbrains.annotations.Nullable;
 
 /**
@@ -200,4 +202,24 @@ public ShapeTemplate build() {
             return innerBuilder.build();
         }
     }
+
+    public static class Structure {
+        private final ResourceLocation id;
+
+        public Structure(ResourceLocation id) {
+            this.id = id;
+        }
+
+        public ShapeTemplate build() {
+            var tag = MIStructureTemplateManager.load(id);
+            if (tag == null) {
+                throw new IllegalStateException("Failed to load structure with id %s".formatted(id));
+            }
+            ShapeTemplate template = MIStructureTemplateManager.deserialize(tag);
+            if (template == null) {
+                throw new IllegalStateException("Failed to parse structure with id %s".formatted(id));
+            }
+            return template;
+        }
+    }
 }

src/main/java/aztech/modern_industrialization/machines/multiblocks/structure/MIStructureTemplateManager.java

@@ -46,7 +46,6 @@
 import net.minecraft.FileUtil;
 import net.minecraft.core.BlockPos;
 import net.minecraft.core.Direction;
-import net.minecraft.core.registries.BuiltInRegistries;
 import net.minecraft.nbt.CompoundTag;
 import net.minecraft.nbt.ListTag;
 import net.minecraft.nbt.NbtAccounter;
@@ -158,7 +157,7 @@ public static FromWorldResult fromWorld(Level level, BlockPos controllerPos, Dir
 
             blockTag.put("pos", NbtUtils.writeBlockPos(toTemplatePos(controllerPos, controllerDirection, pos)));
 
-            StructureMember member = new StructureMember(state);
+            StructureMember member = new StructureMember(() -> state);
 
             BlockEntity blockEntity = level.getBlockEntity(pos);
             if (blockEntity instanceof StructureMemberOverride override) {
@@ -207,8 +206,6 @@ public static FromWorldResult fromWorld(Level level, BlockPos controllerPos, Dir
     public static ShapeTemplate deserialize(CompoundTag tag) {
         Objects.requireNonNull(tag);
 
-        var blockRegistry = BuiltInRegistries.BLOCK.asLookup();
-
         ResourceLocation hatchCasingId = ResourceLocation.tryParse(tag.getString("hatch_casing"));
         if (hatchCasingId == null || !MachineCasings.registeredCasings.containsKey(hatchCasingId)) {
             return null;

src/main/java/aztech/modern_industrialization/machines/multiblocks/structure/StructureMultiblockFormatters.java

@@ -98,7 +98,7 @@ public static List<StructureMemberTest> members(String inputMembers) {
             } else {
                 try {
                     var blockResult = BlockStateParser.parseForBlock(registry, part, true);
-                    members.add(new StructureMemberTestState(blockResult.blockState()));
+                    members.add(new StructureMemberTestState(blockResult::blockState));
                 } catch (CommandSyntaxException ignored) {
                     return null;
                 }

src/main/java/aztech/modern_industrialization/machines/multiblocks/structure/member/StructureMember.java

@@ -26,20 +26,49 @@
 import aztech.modern_industrialization.machines.multiblocks.HatchFlags;
 import aztech.modern_industrialization.machines.multiblocks.SimpleMember;
 import java.util.ArrayList;
+import java.util.Collections;
 import java.util.List;
 import java.util.Objects;
+import java.util.function.Supplier;
 import net.minecraft.core.registries.BuiltInRegistries;
 import net.minecraft.nbt.CompoundTag;
 import net.minecraft.nbt.ListTag;
 import net.minecraft.nbt.NbtUtils;
 import net.minecraft.nbt.Tag;
 import net.minecraft.world.level.block.state.BlockState;
 
-public record StructureMember(BlockState preview, List<StructureMemberTest> tests, HatchFlags hatchFlags) implements SimpleMember {
-    public StructureMember(BlockState state) {
+public final class StructureMember implements SimpleMember {
+    private final Supplier<BlockState> previewSupplier;
+    private final List<StructureMemberTest> tests;
+    private final HatchFlags hatchFlags;
+
+    private BlockState preview;
+
+    public StructureMember(Supplier<BlockState> preview, List<StructureMemberTest> tests, HatchFlags hatchFlags) {
+        this.previewSupplier = preview;
+        this.tests = tests;
+        this.hatchFlags = hatchFlags;
+    }
+
+    public StructureMember(Supplier<BlockState> state) {
         this(state, List.of(new StructureMemberTestState(state)), null);
     }
 
+    public BlockState preview() {
+        if (preview == null) {
+            preview = previewSupplier.get();
+        }
+        return preview;
+    }
+
+    public List<StructureMemberTest> tests() {
+        return Collections.unmodifiableList(tests);
+    }
+
+    public HatchFlags hatchFlags() {
+        return hatchFlags;
+    }
+
     @Override
     public boolean matchesState(BlockState state) {
         for (StructureMemberTest test : tests) {
@@ -52,21 +81,21 @@ public boolean matchesState(BlockState state) {
 
     @Override
     public BlockState getPreviewState() {
-        return preview;
+        return this.preview();
     }
 
     @Override
     public boolean equals(Object o) {
         if (o instanceof StructureMember other) {
-            return preview == other.preview() &&
-                    tests.containsAll(other.tests()) && other.tests().containsAll(tests) &&
-                    Objects.equals(hatchFlags, other.hatchFlags());
+            return this.preview() == other.preview() &&
+                    tests.containsAll(other.tests) && other.tests.containsAll(tests) &&
+                    Objects.equals(hatchFlags, other.hatchFlags);
         }
         return false;
     }
 
     public void save(CompoundTag tag) {
-        tag.put("preview", NbtUtils.writeBlockState(preview));
+        tag.put("preview", NbtUtils.writeBlockState(this.preview()));
 
         ListTag testsTag = new ListTag();
         for (StructureMemberTest test : tests) {
@@ -89,9 +118,7 @@ public static StructureMember from(CompoundTag tag) {
             return null;
         }
 
-        var registry = BuiltInRegistries.BLOCK.asLookup();
-
-        BlockState preview = NbtUtils.readBlockState(registry, tag.getCompound("preview"));
+        Supplier<BlockState> preview = () -> NbtUtils.readBlockState(BuiltInRegistries.BLOCK.asLookup(), tag.getCompound("preview"));
 
         ListTag testsTag = tag.getList("tests", Tag.TAG_COMPOUND);
         List<StructureMemberTest> tests = new ArrayList<>();

src/main/java/aztech/modern_industrialization/machines/multiblocks/structure/member/StructureMemberTestState.java

@@ -23,47 +23,57 @@
  */
 package aztech.modern_industrialization.machines.multiblocks.structure.member;
 
+import java.util.function.Supplier;
 import net.minecraft.core.registries.BuiltInRegistries;
 import net.minecraft.nbt.CompoundTag;
 import net.minecraft.nbt.NbtUtils;
 import net.minecraft.world.level.block.state.BlockState;
 
 public class StructureMemberTestState implements StructureMemberTest {
+    private Supplier<BlockState> blockStateSupplier;
+
     private BlockState blockState;
 
-    public StructureMemberTestState(BlockState blockState) {
-        this.blockState = blockState;
+    public StructureMemberTestState(Supplier<BlockState> blockState) {
+        this.blockStateSupplier = blockState;
     }
 
     public StructureMemberTestState() {
         this(null);
     }
 
+    private BlockState blockState() {
+        if (blockState == null) {
+            blockState = blockStateSupplier.get();
+        }
+        return blockState;
+    }
+
     @Override
     public String id() {
         return "state";
     }
 
     @Override
     public boolean matchesState(BlockState state) {
-        return blockState == state;
+        return this.blockState() == state;
     }
 
     @Override
     public void load(CompoundTag tag) {
-        var registry = BuiltInRegistries.BLOCK.asLookup();
-        blockState = NbtUtils.readBlockState(registry, tag.getCompound("state"));
+        blockStateSupplier = () -> NbtUtils.readBlockState(BuiltInRegistries.BLOCK.asLookup(), tag.getCompound("state"));
+        blockState = null;
     }
 
     @Override
     public void save(CompoundTag tag) {
-        tag.put("state", NbtUtils.writeBlockState(blockState));
+        tag.put("state", NbtUtils.writeBlockState(this.blockState()));
     }
 
     @Override
     public boolean equals(Object o) {
         if (o instanceof StructureMemberTestState other) {
-            return blockState == other.blockState;
+            return this.blockState() == other.blockState();
         }
         return false;
     }