From 94d7ff14193d1c94a541b97c7c22ce47bd23c838 Mon Sep 17 00:00:00 2001
From: Sollace <sollacea@gmail.com>
Date: Tue, 12 Sep 2023 12:27:47 +0100
Subject: [PATCH] Ensure zap apple trees have the correct state when generated

---
 .../unicopia/block/ZapAppleLeavesBlock.java   | 20 +++++++++++++++++++
 .../block/ZapAppleLeavesPlaceholderBlock.java | 15 ++++++++++++++
 2 files changed, 35 insertions(+)

diff --git a/src/main/java/com/minelittlepony/unicopia/block/ZapAppleLeavesBlock.java b/src/main/java/com/minelittlepony/unicopia/block/ZapAppleLeavesBlock.java
index 573e30d3..b036532a 100644
--- a/src/main/java/com/minelittlepony/unicopia/block/ZapAppleLeavesBlock.java
+++ b/src/main/java/com/minelittlepony/unicopia/block/ZapAppleLeavesBlock.java
@@ -3,8 +3,11 @@ package com.minelittlepony.unicopia.block;
 import com.minelittlepony.unicopia.server.world.ZapAppleStageStore;
 
 import net.minecraft.block.*;
+import net.minecraft.server.world.ServerWorld;
 import net.minecraft.state.StateManager;
 import net.minecraft.state.property.*;
+import net.minecraft.util.math.BlockPos;
+import net.minecraft.world.World;
 
 public class ZapAppleLeavesBlock extends BaseZapAppleLeavesBlock {
     public static final EnumProperty<ZapAppleStageStore.Stage> STAGE = EnumProperty.of("stage", ZapAppleStageStore.Stage.class);
@@ -13,6 +16,23 @@ public class ZapAppleLeavesBlock extends BaseZapAppleLeavesBlock {
         setDefaultState(getDefaultState().with(STAGE, ZapAppleStageStore.Stage.HIBERNATING));
     }
 
+    @Override
+    public void onBlockAdded(BlockState state, World world, BlockPos pos, BlockState oldState, boolean notify) {
+        if (oldState.isOf(state.getBlock())
+                || oldState.isOf(UBlocks.ZAP_LEAVES)
+                || oldState.isOf(UBlocks.FLOWERING_ZAP_LEAVES)
+                || oldState.isOf(UBlocks.ZAP_LEAVES_PLACEHOLDER)
+                || !(world instanceof ServerWorld sw)) {
+            return;
+        }
+
+        ZapAppleStageStore store = ZapAppleStageStore.get(sw);
+        ZapAppleStageStore.Stage currentStage = store.getStage();
+        if (currentStage != getStage(state)) {
+            world.setBlockState(pos, currentStage.getNewState(state));
+        }
+    }
+
     @Override
     protected ZapAppleStageStore.Stage getStage(BlockState state) {
         return state.get(STAGE);
diff --git a/src/main/java/com/minelittlepony/unicopia/block/ZapAppleLeavesPlaceholderBlock.java b/src/main/java/com/minelittlepony/unicopia/block/ZapAppleLeavesPlaceholderBlock.java
index eb2f1526..5b6a7cdc 100644
--- a/src/main/java/com/minelittlepony/unicopia/block/ZapAppleLeavesPlaceholderBlock.java
+++ b/src/main/java/com/minelittlepony/unicopia/block/ZapAppleLeavesPlaceholderBlock.java
@@ -6,6 +6,7 @@ import net.minecraft.block.*;
 import net.minecraft.server.world.ServerWorld;
 import net.minecraft.util.math.*;
 import net.minecraft.util.math.random.Random;
+import net.minecraft.world.WorldAccess;
 
 public class ZapAppleLeavesPlaceholderBlock extends AirBlock {
 
@@ -18,6 +19,20 @@ public class ZapAppleLeavesPlaceholderBlock extends AirBlock {
         return true;
     }
 
+    @Override
+    public BlockState getStateForNeighborUpdate(BlockState state, Direction direction, BlockState neighborState, WorldAccess world, BlockPos pos, BlockPos neighborPos) {
+
+        if (world instanceof ServerWorld sw) {
+            ZapAppleStageStore store = ZapAppleStageStore.get(sw);
+            ZapAppleStageStore.Stage currentStage = store.getStage();
+            if (currentStage != ZapAppleStageStore.Stage.HIBERNATING) {
+                return currentStage.getNewState(state);
+            }
+        }
+
+        return state;
+    }
+
     @Deprecated
     @Override
     public void scheduledTick(BlockState state, ServerWorld world, BlockPos pos, Random random) {