steel_core/entity/ai/goal/
try_find_water.rs1use 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}