Skip to main content

steel_core/entity/ai/goal/
selector.rs

1use std::fmt;
2use std::ops::BitOr;
3
4use crate::entity::{PathfinderMob, SharedEntity};
5
6#[derive(Debug, Clone, Copy, PartialEq, Eq)]
7pub enum GoalControl {
8    Move,
9    Look,
10    Jump,
11    Target,
12}
13
14impl GoalControl {
15    const ALL: [Self; 4] = [Self::Move, Self::Look, Self::Jump, Self::Target];
16}
17
18#[derive(Clone, Copy, Default, PartialEq, Eq)]
19pub struct GoalControls(u8);
20
21impl GoalControls {
22    pub const EMPTY: Self = Self(0);
23    pub const MOVE: Self = Self(1 << 0);
24    pub const LOOK: Self = Self(1 << 1);
25    pub const JUMP: Self = Self(1 << 2);
26    pub const TARGET: Self = Self(1 << 3);
27
28    #[must_use]
29    pub const fn from_control(control: GoalControl) -> Self {
30        match control {
31            GoalControl::Move => Self::MOVE,
32            GoalControl::Look => Self::LOOK,
33            GoalControl::Jump => Self::JUMP,
34            GoalControl::Target => Self::TARGET,
35        }
36    }
37
38    #[must_use]
39    pub const fn contains(self, control: GoalControl) -> bool {
40        self.0 & Self::from_control(control).0 != 0
41    }
42
43    #[must_use]
44    pub const fn intersects(self, other: Self) -> bool {
45        self.0 & other.0 != 0
46    }
47
48    pub const fn insert(&mut self, control: GoalControl) {
49        self.0 |= Self::from_control(control).0;
50    }
51
52    pub const fn remove(&mut self, control: GoalControl) {
53        self.0 &= !Self::from_control(control).0;
54    }
55
56    pub fn iter(self) -> impl Iterator<Item = GoalControl> {
57        GoalControl::ALL
58            .into_iter()
59            .filter(move |control| self.contains(*control))
60    }
61}
62
63impl fmt::Debug for GoalControls {
64    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
65        f.debug_set().entries(self.iter()).finish()
66    }
67}
68
69impl BitOr for GoalControls {
70    type Output = Self;
71
72    fn bitor(self, rhs: Self) -> Self::Output {
73        Self(self.0 | rhs.0)
74    }
75}
76
77pub trait Goal: Send {
78    fn controls(&self) -> GoalControls;
79
80    fn can_use(&mut self, mob: &dyn PathfinderMob) -> bool;
81
82    fn can_continue_to_use(&mut self, mob: &dyn PathfinderMob) -> bool {
83        self.can_use(mob)
84    }
85
86    fn is_interruptable(&self) -> bool {
87        true
88    }
89
90    fn is_panic_goal(&self) -> bool {
91        false
92    }
93
94    fn start(&mut self, _mob: &dyn PathfinderMob) {}
95
96    fn stop(&mut self, _mob: &dyn PathfinderMob) {}
97
98    fn requires_update_every_tick(&self) -> bool {
99        false
100    }
101
102    fn tick(&mut self, _mob: &dyn PathfinderMob, _entity: &SharedEntity) {}
103}
104
105struct WrappedGoal {
106    priority: i32,
107    goal: Box<dyn Goal>,
108    running: bool,
109}
110
111impl WrappedGoal {
112    fn new(priority: i32, goal: Box<dyn Goal>) -> Self {
113        Self {
114            priority,
115            goal,
116            running: false,
117        }
118    }
119
120    const fn is_running(&self) -> bool {
121        self.running
122    }
123
124    fn goal(&self) -> &dyn Goal {
125        self.goal.as_ref()
126    }
127
128    fn controls(&self) -> GoalControls {
129        self.goal.controls()
130    }
131
132    fn can_be_replaced_by(&self, candidate_priority: i32) -> bool {
133        self.goal.is_interruptable() && candidate_priority < self.priority
134    }
135
136    fn can_use(&mut self, mob: &dyn PathfinderMob) -> bool {
137        self.goal.can_use(mob)
138    }
139
140    fn can_continue_to_use(&mut self, mob: &dyn PathfinderMob) -> bool {
141        self.goal.can_continue_to_use(mob)
142    }
143
144    fn start(&mut self, mob: &dyn PathfinderMob) {
145        if self.running {
146            return;
147        }
148        self.running = true;
149        self.goal.start(mob);
150    }
151
152    fn stop(&mut self, mob: &dyn PathfinderMob) {
153        if !self.running {
154            return;
155        }
156        self.running = false;
157        self.goal.stop(mob);
158    }
159
160    fn tick(&mut self, mob: &dyn PathfinderMob, entity: &SharedEntity) {
161        self.goal.tick(mob, entity);
162    }
163
164    fn requires_update_every_tick(&self) -> bool {
165        self.goal.requires_update_every_tick()
166    }
167
168    fn is_panic_goal(&self) -> bool {
169        self.goal.is_panic_goal()
170    }
171}
172
173pub struct GoalSelector {
174    available_goals: Vec<WrappedGoal>,
175    disabled_controls: GoalControls,
176}
177
178impl GoalSelector {
179    #[must_use]
180    pub const fn new() -> Self {
181        Self {
182            available_goals: Vec::new(),
183            disabled_controls: GoalControls::EMPTY,
184        }
185    }
186
187    pub fn add_goal<G>(&mut self, priority: i32, goal: G)
188    where
189        G: Goal + 'static,
190    {
191        self.available_goals
192            .push(WrappedGoal::new(priority, Box::new(goal)));
193    }
194
195    /// Stops and removes matching goals. `mob` supplies the stop-hook context.
196    pub fn remove_all_goals(
197        &mut self,
198        mob: &dyn PathfinderMob,
199        mut predicate: impl FnMut(&dyn Goal) -> bool,
200    ) {
201        for wrapped in &mut self.available_goals {
202            if predicate(wrapped.goal()) && wrapped.is_running() {
203                wrapped.stop(mob);
204            }
205        }
206        self.available_goals
207            .retain(|wrapped| !predicate(wrapped.goal()));
208    }
209
210    pub fn tick(&mut self, mob: &dyn PathfinderMob, entity: &SharedEntity) {
211        for index in 0..self.available_goals.len() {
212            let should_stop = {
213                let disabled_controls = self.disabled_controls;
214                let goal = &mut self.available_goals[index];
215                goal.is_running()
216                    && (goal.controls().intersects(disabled_controls)
217                        || !goal.can_continue_to_use(mob))
218            };
219            if should_stop {
220                self.available_goals[index].stop(mob);
221            }
222        }
223
224        for index in 0..self.available_goals.len() {
225            if !self.can_start_goal(index) {
226                continue;
227            }
228            if !self.available_goals[index].can_use(mob) {
229                continue;
230            }
231
232            let controls = self.available_goals[index].controls();
233            for control in controls.iter() {
234                if let Some(current_index) = self.running_goal_index_for(control) {
235                    self.available_goals[current_index].stop(mob);
236                }
237            }
238            self.available_goals[index].start(mob);
239        }
240
241        self.tick_running_goals(mob, entity, true);
242    }
243
244    pub fn tick_running_goals(
245        &mut self,
246        mob: &dyn PathfinderMob,
247        entity: &SharedEntity,
248        force_tick_all_running_goals: bool,
249    ) {
250        for goal in &mut self.available_goals {
251            if goal.is_running()
252                && (force_tick_all_running_goals || goal.requires_update_every_tick())
253            {
254                goal.tick(mob, entity);
255            }
256        }
257    }
258
259    pub const fn disable_control(&mut self, control: GoalControl) {
260        self.disabled_controls.insert(control);
261    }
262
263    pub const fn enable_control(&mut self, control: GoalControl) {
264        self.disabled_controls.remove(control);
265    }
266
267    pub const fn set_control(&mut self, control: GoalControl, enabled: bool) {
268        if enabled {
269            self.enable_control(control);
270        } else {
271            self.disable_control(control);
272        }
273    }
274
275    #[must_use]
276    pub fn running_goal_count(&self) -> usize {
277        self.available_goals
278            .iter()
279            .filter(|goal| goal.is_running())
280            .count()
281    }
282
283    #[must_use]
284    pub const fn available_goal_count(&self) -> usize {
285        self.available_goals.len()
286    }
287
288    #[must_use]
289    pub(crate) fn has_running_panic_goal(&self) -> bool {
290        self.available_goals
291            .iter()
292            .any(|goal| goal.is_running() && goal.is_panic_goal())
293    }
294
295    #[cfg(test)]
296    #[must_use]
297    pub(crate) fn available_goal_priorities(&self) -> Vec<i32> {
298        self.available_goals
299            .iter()
300            .map(|goal| goal.priority)
301            .collect()
302    }
303
304    #[cfg(test)]
305    #[must_use]
306    pub(crate) const fn is_control_disabled(&self, control: GoalControl) -> bool {
307        self.disabled_controls.contains(control)
308    }
309
310    fn can_start_goal(&self, index: usize) -> bool {
311        let goal = &self.available_goals[index];
312        !goal.is_running()
313            && !goal.controls().intersects(self.disabled_controls)
314            && self.goal_can_be_replaced_for_all_controls(index)
315    }
316
317    fn goal_can_be_replaced_for_all_controls(&self, candidate_index: usize) -> bool {
318        let candidate = &self.available_goals[candidate_index];
319        for control in candidate.controls().iter() {
320            if let Some(current_index) = self.running_goal_index_for(control)
321                && !self.available_goals[current_index].can_be_replaced_by(candidate.priority)
322            {
323                return false;
324            }
325        }
326        true
327    }
328
329    fn running_goal_index_for(&self, control: GoalControl) -> Option<usize> {
330        self.available_goals
331            .iter()
332            .position(|goal| goal.is_running() && goal.controls().contains(control))
333    }
334
335    #[cfg(test)]
336    fn is_priority_running(&self, priority: i32) -> bool {
337        self.available_goals
338            .iter()
339            .any(|goal| goal.priority == priority && goal.is_running())
340    }
341}
342
343impl Default for GoalSelector {
344    fn default() -> Self {
345        Self::new()
346    }
347}
348
349impl fmt::Debug for GoalSelector {
350    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
351        f.debug_struct("GoalSelector")
352            .field("available_goals", &self.available_goals.len())
353            .field("running_goals", &self.running_goal_count())
354            .field("disabled_controls", &self.disabled_controls)
355            .finish()
356    }
357}
358
359#[cfg(test)]
360mod tests {
361    use std::sync::Arc;
362
363    use std::sync::Weak;
364    use std::sync::atomic::{AtomicUsize, Ordering};
365
366    use glam::DVec3;
367    use steel_registry::entity_type::EntityTypeRef;
368    use steel_registry::{init_vanilla_registry, vanilla_entities};
369    use steel_utils::locks::SyncMutex;
370
371    use super::*;
372    use crate::entity::{
373        Entity, EntityBase, LivingEntity, LivingEntityBase, Mob, MobBase, PathfinderMob,
374    };
375
376    struct TestPathfinderMob {
377        base: EntityBase,
378        living_base: LivingEntityBase,
379        mob_base: MobBase,
380        mob_flags: SyncMutex<i8>,
381        health: SyncMutex<f32>,
382    }
383
384    impl TestPathfinderMob {
385        fn new() -> Self {
386            init_vanilla_registry();
387            Self {
388                base: EntityBase::new(
389                    1,
390                    DVec3::ZERO,
391                    vanilla_entities::PIG.dimensions,
392                    Weak::new(),
393                ),
394                living_base: LivingEntityBase::new(&vanilla_entities::PIG),
395                mob_base: MobBase::new(),
396                mob_flags: SyncMutex::new(0),
397                health: SyncMutex::new(10.0),
398            }
399        }
400    }
401
402    crate::entity::impl_test_downcast_type!(TestPathfinderMob);
403
404    impl Entity for TestPathfinderMob {
405        fn base(&self) -> &EntityBase {
406            &self.base
407        }
408
409        fn entity_type(&self) -> EntityTypeRef {
410            &vanilla_entities::PIG
411        }
412    }
413
414    impl LivingEntity for TestPathfinderMob {
415        fn living_base(&self) -> &LivingEntityBase {
416            &self.living_base
417        }
418
419        fn get_health(&self) -> f32 {
420            *self.health.lock()
421        }
422
423        fn set_health(&self, health: f32) {
424            *self.health.lock() = health;
425        }
426    }
427
428    impl Mob for TestPathfinderMob {
429        fn mob_base(&self) -> &MobBase {
430            &self.mob_base
431        }
432
433        fn mob_flags(&self) -> i8 {
434            *self.mob_flags.lock()
435        }
436
437        fn set_mob_flags(&self, flags: i8) {
438            *self.mob_flags.lock() = flags;
439        }
440    }
441
442    impl PathfinderMob for TestPathfinderMob {}
443
444    struct StaticGoal {
445        controls: GoalControls,
446        can_use: bool,
447        can_continue: bool,
448        interruptable: bool,
449        requires_update_every_tick: bool,
450        tick_count: Option<&'static AtomicUsize>,
451        can_use_once: bool,
452        panic_goal: bool,
453    }
454
455    impl StaticGoal {
456        const fn new(controls: GoalControls) -> Self {
457            Self {
458                controls,
459                can_use: true,
460                can_continue: true,
461                interruptable: true,
462                requires_update_every_tick: false,
463                tick_count: None,
464                can_use_once: false,
465                panic_goal: false,
466            }
467        }
468
469        const fn non_interruptable(mut self) -> Self {
470            self.interruptable = false;
471            self
472        }
473
474        const fn with_can_continue(mut self, can_continue: bool) -> Self {
475            self.can_continue = can_continue;
476            self
477        }
478
479        const fn with_can_use_once(mut self) -> Self {
480            self.can_use_once = true;
481            self
482        }
483
484        const fn with_update_every_tick(mut self) -> Self {
485            self.requires_update_every_tick = true;
486            self
487        }
488
489        const fn with_tick_counter(mut self, tick_count: &'static AtomicUsize) -> Self {
490            self.tick_count = Some(tick_count);
491            self
492        }
493
494        const fn with_panic_goal(mut self) -> Self {
495            self.panic_goal = true;
496            self
497        }
498    }
499
500    impl Goal for StaticGoal {
501        fn controls(&self) -> GoalControls {
502            self.controls
503        }
504
505        fn can_use(&mut self, _mob: &dyn PathfinderMob) -> bool {
506            if self.can_use_once {
507                if !self.can_use {
508                    return false;
509                }
510                self.can_use = false;
511                return true;
512            }
513            self.can_use
514        }
515
516        fn can_continue_to_use(&mut self, _mob: &dyn PathfinderMob) -> bool {
517            self.can_continue
518        }
519
520        fn is_interruptable(&self) -> bool {
521            self.interruptable
522        }
523
524        fn is_panic_goal(&self) -> bool {
525            self.panic_goal
526        }
527
528        fn requires_update_every_tick(&self) -> bool {
529            self.requires_update_every_tick
530        }
531
532        fn tick(&mut self, _mob: &dyn PathfinderMob, _entity: &SharedEntity) {
533            if let Some(tick_count) = self.tick_count {
534                tick_count.fetch_add(1, Ordering::Relaxed);
535            }
536        }
537    }
538
539    static RUNNING_TICK_COUNT: AtomicUsize = AtomicUsize::new(0);
540
541    #[test]
542    fn lower_priority_goal_replaces_running_goal_for_same_control() {
543        let mob = Arc::new(TestPathfinderMob::new());
544        let mob_entity: SharedEntity = Arc::<TestPathfinderMob>::clone(&mob);
545        let mut selector = GoalSelector::new();
546        selector.add_goal(5, StaticGoal::new(GoalControls::MOVE));
547        selector.tick(mob.as_ref(), &mob_entity);
548
549        selector.add_goal(3, StaticGoal::new(GoalControls::MOVE));
550        selector.tick(mob.as_ref(), &mob_entity);
551
552        assert_eq!(selector.running_goal_count(), 1);
553        assert!(selector.is_priority_running(3));
554    }
555
556    #[test]
557    fn non_interruptable_goal_blocks_replacement() {
558        let mob = Arc::new(TestPathfinderMob::new());
559        let mob_entity: SharedEntity = Arc::<TestPathfinderMob>::clone(&mob);
560        let mut selector = GoalSelector::new();
561        selector.add_goal(5, StaticGoal::new(GoalControls::MOVE).non_interruptable());
562        selector.tick(mob.as_ref(), &mob_entity);
563
564        selector.add_goal(3, StaticGoal::new(GoalControls::MOVE));
565        selector.tick(mob.as_ref(), &mob_entity);
566
567        assert_eq!(selector.running_goal_count(), 1);
568        assert!(selector.is_priority_running(5));
569    }
570
571    #[test]
572    fn disabled_control_stops_running_goal() {
573        let mob = Arc::new(TestPathfinderMob::new());
574        let mob_entity: SharedEntity = Arc::<TestPathfinderMob>::clone(&mob);
575        let mut selector = GoalSelector::new();
576        selector.add_goal(5, StaticGoal::new(GoalControls::MOVE));
577        selector.tick(mob.as_ref(), &mob_entity);
578
579        selector.disable_control(GoalControl::Move);
580        selector.tick(mob.as_ref(), &mob_entity);
581
582        assert_eq!(selector.running_goal_count(), 0);
583    }
584
585    #[test]
586    fn tick_running_goals_respects_requires_update_every_tick() {
587        RUNNING_TICK_COUNT.store(0, Ordering::Relaxed);
588        let mob = Arc::new(TestPathfinderMob::new());
589        let mob_entity: SharedEntity = Arc::<TestPathfinderMob>::clone(&mob);
590        let mut selector = GoalSelector::new();
591        selector.add_goal(
592            5,
593            StaticGoal::new(GoalControls::MOVE)
594                .with_update_every_tick()
595                .with_tick_counter(&RUNNING_TICK_COUNT),
596        );
597        selector.tick(mob.as_ref(), &mob_entity);
598
599        selector.tick_running_goals(mob.as_ref(), &mob_entity, false);
600
601        assert_eq!(RUNNING_TICK_COUNT.load(Ordering::Relaxed), 2);
602    }
603
604    #[test]
605    fn cleanup_stops_goal_that_can_no_longer_continue() {
606        let mob = Arc::new(TestPathfinderMob::new());
607        let mob_entity: SharedEntity = Arc::<TestPathfinderMob>::clone(&mob);
608        let mut selector = GoalSelector::new();
609        selector.add_goal(
610            5,
611            StaticGoal::new(GoalControls::MOVE)
612                .with_can_continue(false)
613                .with_can_use_once(),
614        );
615
616        selector.tick(mob.as_ref(), &mob_entity);
617        selector.tick(mob.as_ref(), &mob_entity);
618
619        assert_eq!(selector.running_goal_count(), 0);
620    }
621
622    #[test]
623    fn running_panic_goal_is_visible_to_pathfinder_mob() {
624        let mob = Arc::new(TestPathfinderMob::new());
625        let mob_entity: SharedEntity = Arc::<TestPathfinderMob>::clone(&mob);
626        mob.mob_base()
627            .goal_selector()
628            .lock()
629            .add_goal(1, StaticGoal::new(GoalControls::MOVE).with_panic_goal());
630
631        assert!(!mob.is_panicking());
632
633        mob.mob_base()
634            .goal_selector()
635            .lock()
636            .tick(mob.as_ref(), &mob_entity);
637
638        assert!(
639            mob.mob_base()
640                .goal_selector()
641                .lock()
642                .has_running_panic_goal()
643        );
644        assert!(mob.is_panicking());
645    }
646
647    #[test]
648    fn remove_all_goals_stops_matching_running_goals_and_keeps_others() {
649        let mob = Arc::new(TestPathfinderMob::new());
650        let mob_entity: SharedEntity = Arc::<TestPathfinderMob>::clone(&mob);
651        let mut selector = GoalSelector::new();
652        selector.add_goal(3, StaticGoal::new(GoalControls::MOVE));
653        selector.add_goal(7, StaticGoal::new(GoalControls::LOOK));
654        selector.tick(mob.as_ref(), &mob_entity);
655        assert_eq!(selector.running_goal_count(), 2);
656
657        selector.remove_all_goals(mob.as_ref(), |goal| goal.controls() == GoalControls::MOVE);
658
659        assert_eq!(selector.available_goal_count(), 1);
660        assert!(
661            selector.is_priority_running(7),
662            "the non-matching goal must keep running through removal"
663        );
664        assert!(
665            !selector.is_priority_running(3),
666            "the removed goal must be gone entirely"
667        );
668
669        selector.tick(mob.as_ref(), &mob_entity);
670        assert_eq!(selector.running_goal_count(), 1);
671        assert!(selector.is_priority_running(7));
672    }
673
674    #[test]
675    fn pathfinder_mob_remove_free_will_clears_all_goals() {
676        let mob = Arc::new(TestPathfinderMob::new());
677        let mob_entity: SharedEntity = Arc::<TestPathfinderMob>::clone(&mob);
678        {
679            let mut selector = mob.mob_base().goal_selector().lock();
680            selector.add_goal(1, StaticGoal::new(GoalControls::MOVE));
681            selector.add_goal(2, StaticGoal::new(GoalControls::LOOK));
682            selector.tick(mob.as_ref(), &mob_entity);
683            assert_eq!(selector.running_goal_count(), 2);
684        }
685
686        mob.remove_free_will();
687
688        let selector = mob.mob_base().goal_selector().lock();
689        assert_eq!(selector.available_goal_count(), 0);
690        assert_eq!(selector.running_goal_count(), 0);
691    }
692}