From c656eb2309e351c7335eff08a43d5d1a4310c5d8 Mon Sep 17 00:00:00 2001
From: Sollace <sollacea@gmail.com>
Date: Wed, 8 Nov 2023 15:19:07 +0000
Subject: [PATCH] Make beds waterlog

---
 .../unicopia/mixin/MixinBlockItem.java           | 11 +++++++++++
 .../unicopia/mixin/MixinStateManager.java        |  2 +-
 .../server/world/WaterLoggingManager.java        | 16 +++++++++++++++-
 3 files changed, 27 insertions(+), 2 deletions(-)

diff --git a/src/main/java/com/minelittlepony/unicopia/mixin/MixinBlockItem.java b/src/main/java/com/minelittlepony/unicopia/mixin/MixinBlockItem.java
index 06382199..3cc48210 100644
--- a/src/main/java/com/minelittlepony/unicopia/mixin/MixinBlockItem.java
+++ b/src/main/java/com/minelittlepony/unicopia/mixin/MixinBlockItem.java
@@ -1,12 +1,18 @@
 package com.minelittlepony.unicopia.mixin;
 
 import org.spongepowered.asm.mixin.Mixin;
+import org.spongepowered.asm.mixin.injection.At;
+import org.spongepowered.asm.mixin.injection.Inject;
+import org.spongepowered.asm.mixin.injection.callback.CallbackInfoReturnable;
 
 import com.minelittlepony.unicopia.item.toxin.ToxicHolder;
+import com.minelittlepony.unicopia.server.world.WaterLoggingManager;
 
+import net.minecraft.block.BlockState;
 import net.minecraft.entity.player.PlayerEntity;
 import net.minecraft.item.BlockItem;
 import net.minecraft.item.Item;
+import net.minecraft.item.ItemPlacementContext;
 import net.minecraft.item.ItemStack;
 import net.minecraft.util.Hand;
 import net.minecraft.util.TypedActionResult;
@@ -28,4 +34,9 @@ abstract class MixinBlockItem extends Item implements ToxicHolder {
                 .map(t -> t.use(world, player, hand))
                 .orElseGet(() -> super.use(world, player, hand));
     }
+
+    @Inject(method = "getPlacementState", at = @At("RETURN"), cancellable = true)
+    private void onGetPlacementState(ItemPlacementContext context, CallbackInfoReturnable<BlockState> info) {
+        WaterLoggingManager.getInstance().getPlacementState(context, info);
+    }
 }
diff --git a/src/main/java/com/minelittlepony/unicopia/mixin/MixinStateManager.java b/src/main/java/com/minelittlepony/unicopia/mixin/MixinStateManager.java
index c31f2822..ff62db5f 100644
--- a/src/main/java/com/minelittlepony/unicopia/mixin/MixinStateManager.java
+++ b/src/main/java/com/minelittlepony/unicopia/mixin/MixinStateManager.java
@@ -69,4 +69,4 @@ abstract class MixinBlockState extends State<Block, BlockState> {
     private void onGetStateForNeighborUpdate(Direction direction, BlockState neighborState, WorldAccess world, BlockPos pos, BlockPos neighborPos, CallbackInfoReturnable<BlockState> info) {
         WaterLoggingManager.<Block, BlockState>getInstance().getUpdatedState(world, pos, asBlockState(), info);
     }
-}
\ No newline at end of file
+}
diff --git a/src/main/java/com/minelittlepony/unicopia/server/world/WaterLoggingManager.java b/src/main/java/com/minelittlepony/unicopia/server/world/WaterLoggingManager.java
index 2372a3bb..9d96541e 100644
--- a/src/main/java/com/minelittlepony/unicopia/server/world/WaterLoggingManager.java
+++ b/src/main/java/com/minelittlepony/unicopia/server/world/WaterLoggingManager.java
@@ -8,6 +8,8 @@ import net.minecraft.block.*;
 import net.minecraft.block.enums.DoubleBlockHalf;
 import net.minecraft.fluid.FluidState;
 import net.minecraft.fluid.Fluids;
+import net.minecraft.item.ItemPlacementContext;
+import net.minecraft.registry.tag.FluidTags;
 import net.minecraft.state.State;
 import net.minecraft.state.property.Properties;
 import net.minecraft.state.property.Property;
@@ -35,11 +37,22 @@ public class WaterLoggingManager<O, S extends State<O, S>> {
     }
 
     public void getDefaultState(O owner, CallbackInfoReturnable<S> info) {
+        if (owner instanceof BedBlock) {
+            return;
+        }
         if (appliesTo(owner, info.getReturnValue())) {
             info.setReturnValue(info.getReturnValue().with(Properties.WATERLOGGED, true));
         }
     }
 
+    @SuppressWarnings("unchecked")
+    public void getPlacementState(ItemPlacementContext context, CallbackInfoReturnable<BlockState> info) {
+        BlockState state = info.getReturnValue();
+        if (state != null && appliesTo((O)state.getBlock(), (S)state)) {
+            info.setReturnValue(state.with(Properties.WATERLOGGED, context.getWorld().getFluidState(context.getBlockPos()).isIn(FluidTags.WATER)));
+        }
+    }
+
     public void getFluidState(O owner, S state, CallbackInfoReturnable<FluidState> info) {
         if (appliesTo(owner, state)) {
             info.setReturnValue((state.get(Properties.WATERLOGGED) ? Fluids.WATER : Fluids.EMPTY).getDefaultState());
@@ -61,7 +74,8 @@ public class WaterLoggingManager<O, S extends State<O, S>> {
             && (block instanceof SeagrassBlock
                 || block instanceof TallSeagrassBlock
                 || block instanceof KelpBlock
-                || block instanceof KelpPlantBlock);
+                || block instanceof KelpPlantBlock
+                || block instanceof BedBlock);
     }
 
     public boolean shouldPreventRemoval(WorldAccess world, BlockPos pos, AbstractBlock.AbstractBlockState oldState, AbstractBlock.AbstractBlockState newState) {