Skip to main content

steel_core/entity/ai/goal/
try_find_water.rs

1use glam::DVec3;
2use steel_math::fast_floor;
3use steel_registry::blocks::block_state_ext::BlockStateExt as _;
4use steel_utils::BlockPos;
5
6use super::selector::{Goal, GoalControls};
7use crate::entity::PathfinderMob;
8use crate::fluid::FluidStateExt as _;
9
10const WATER_SEARCH_HORIZONTAL_RADIUS: f64 = 2.0;
11const WATER_SEARCH_VERTICAL_BELOW: f64 = 2.0;
12const WATER_MOVE_SPEED_MODIFIER: f64 = 1.0;
13
14pub struct TryFindWaterGoal;
15
16impl TryFindWaterGoal {
17    #[must_use]
18    pub(crate) const fn new() -> Self {
19        Self
20    }
21}
22
23impl Goal for TryFindWaterGoal {
24    fn controls(&self) -> GoalControls {
25        GoalControls::EMPTY
26    }
27
28    fn can_use(&mut self, mob: &dyn PathfinderMob) -> bool {
29        if !mob.on_ground() {
30            return false;
31        }
32
33        let Some(world) = mob.level() else {
34            return false;
35        };
36        !world
37            .get_block_state(mob.block_position())
38            .get_fluid_state()
39            .is_water()
40    }
41
42    fn start(&mut self, mob: &dyn PathfinderMob) {
43        let Some(world) = mob.level() else {
44            return;
45        };
46
47        let (min, max) = water_search_bounds(mob.position());
48        let Some(water_pos) = first_matching_pos_in_closed_box(min, max, |pos| {
49            world.get_block_state(pos).get_fluid_state().is_water()
50        }) else {
51            return;
52        };
53
54        mob.mob_base()
55            .controls()
56            .lock()
57            .move_control
58            .set_wanted_position(
59                DVec3::new(
60                    f64::from(water_pos.x()),
61                    f64::from(water_pos.y()),
62                    f64::from(water_pos.z()),
63                ),
64                WATER_MOVE_SPEED_MODIFIER,
65            );
66    }
67}
68
69fn water_search_bounds(position: DVec3) -> (BlockPos, BlockPos) {
70    (
71        BlockPos::new(
72            fast_floor(position.x - WATER_SEARCH_HORIZONTAL_RADIUS),
73            fast_floor(position.y - WATER_SEARCH_VERTICAL_BELOW),
74            fast_floor(position.z - WATER_SEARCH_HORIZONTAL_RADIUS),
75        ),
76        BlockPos::new(
77            fast_floor(position.x + WATER_SEARCH_HORIZONTAL_RADIUS),
78            fast_floor(position.y),
79            fast_floor(position.z + WATER_SEARCH_HORIZONTAL_RADIUS),
80        ),
81    )
82}
83
84fn first_matching_pos_in_closed_box(
85    min: BlockPos,
86    max: BlockPos,
87    mut predicate: impl FnMut(BlockPos) -> bool,
88) -> Option<BlockPos> {
89    for z in min.z()..=max.z() {
90        for y in min.y()..=max.y() {
91            for x in min.x()..=max.x() {
92                let pos = BlockPos::new(x, y, z);
93                if predicate(pos) {
94                    return Some(pos);
95                }
96            }
97        }
98    }
99    None
100}
101
102#[cfg(test)]
103mod tests {
104    use std::sync::Weak;
105
106    use glam::DVec3;
107    use steel_registry::{init_vanilla_registry, vanilla_entities};
108
109    use super::*;
110    use crate::entity::Entity;
111    use crate::entity::entities::PigEntity;
112
113    #[test]
114    fn try_find_water_goal_claims_no_controls_like_vanilla() {
115        let goal = TryFindWaterGoal::new();
116
117        assert_eq!(goal.controls(), GoalControls::EMPTY);
118    }
119
120    #[test]
121    fn try_find_water_goal_requires_on_ground() {
122        init_vanilla_registry();
123        let mut goal = TryFindWaterGoal::new();
124        let mob = PigEntity::new(&vanilla_entities::PIG, 1, DVec3::ZERO, Weak::new());
125
126        assert!(!goal.can_use(&mob));
127    }
128
129    #[test]
130    fn try_find_water_goal_requires_world_after_on_ground_check() {
131        init_vanilla_registry();
132        let mut goal = TryFindWaterGoal::new();
133        let mob = PigEntity::new(&vanilla_entities::PIG, 1, DVec3::ZERO, Weak::new());
134        mob.set_on_ground(true);
135
136        assert!(!goal.can_use(&mob));
137    }
138
139    #[test]
140    fn water_search_bounds_match_vanilla_offsets() {
141        let (min, max) = water_search_bounds(DVec3::new(-0.25, 64.9, 0.25));
142
143        assert_eq!(min, BlockPos::new(-3, 62, -2));
144        assert_eq!(max, BlockPos::new(1, 64, 2));
145    }
146
147    #[test]
148    fn closed_box_scan_uses_vanilla_x_then_y_then_z_order() {
149        let min = BlockPos::new(10, 20, 30);
150        let max = BlockPos::new(12, 22, 32);
151        let expected = BlockPos::new(11, 20, 30);
152        let later_y = BlockPos::new(10, 21, 30);
153        let later_z = BlockPos::new(10, 20, 31);
154
155        let found = first_matching_pos_in_closed_box(min, max, |pos| {
156            pos == expected || pos == later_y || pos == later_z
157        });
158
159        assert_eq!(found, Some(expected));
160    }
161}