bevy_yoetz/
advisor.rs

1use bevy::ecs::query::QueryData;
2use bevy::ecs::system::EntityCommands;
3use bevy::prelude::*;
4
5#[doc(inline)]
6pub use bevy_yoetz_macros::YoetzSuggestion;
7
8/// An action suggestion for the AI agent to consider.
9///
10/// Avoid implementing this trait manually - prefer using the
11/// [`YoetzSuggestion`](bevy_yoetz_macros::YoetzSuggestion) derive macro.
12///
13/// `enum`s that implement this trait are mainly used as the generic parameter for [`YoetzAdvisor`]
14/// and as the data passed to it. A [`YoetzPlugin`](crate::YoetzPlugin) parametrized on them should
15/// also be added to the Bevy application.
16pub trait YoetzSuggestion: 'static + Sized + Send + Sync {
17    /// The key identifies a suggestion even when its data changes. The
18    /// [`YoetzSuggestion`](bevy_yoetz_macros::YoetzSuggestion) derive macro generates a key that
19    /// is a "subset" of the `enum` - it contains all the variants, but each variant only contains
20    /// the fields marked as `#[yoetz(key)]`.
21    type Key: 'static + Send + Sync + Clone + PartialEq;
22
23    /// A query that allows access to all possible behavior components.
24    ///
25    /// The query generated by the [`YoetzSuggestion`](bevy_yoetz_macros::YoetzSuggestion) derive
26    /// macro is unsightly and there it never a reason to use it manually.
27    type OmniQuery: QueryData;
28
29    /// Generate a [`Key`](Self::Key) that identifies the suggestion.
30    fn key(&self) -> Self::Key;
31
32    /// Remove the behavior components that were created by a suggestion with the specified key.
33    fn remove_components(key: &Self::Key, cmd: &mut EntityCommands);
34
35    /// Add behavior components created from the suggestion.
36    fn add_components(self, cmd: &mut EntityCommands);
37
38    /// Update the existing behavior components from the suggestion's data.
39    ///
40    /// The method generated by the [`YoetzSuggestion`](bevy_yoetz_macros::YoetzSuggestion) derive
41    /// macro will only update the fields marked with `#[yoetz(input)]`. Fields marked with
42    /// `#[yoetz(state)]` will not be updated because the action systems are allowed to store their
43    /// own state there (or just maintain the initial state from when the behavior was chosen), and
44    /// fields marked with `#[yoetz(key)]` will not be updated because when they change the
45    /// [`Key`](Self::Key) changes and the components themselves will be re-inserted rather than
46    /// updated.
47    fn update_into_components(
48        self,
49        components: &mut <Self::OmniQuery as QueryData>::Item<'_>,
50    ) -> Result<(), Self>;
51}
52
53/// Controls an entity's AI by listening to [`YoetzSuggestion`]s and updating the entity's behavior
54/// components.
55#[derive(Component)]
56pub struct YoetzAdvisor<S: YoetzSuggestion> {
57    /// Added to score of any suggestion that matches the currently active behavior. This can be
58    /// used to reduce the "flickering" when multiple suggestions are flocking around the same
59    /// score.
60    pub consistency_bonus: f32,
61    active_key: Option<S::Key>,
62    top_suggestion: Option<(f32, S)>,
63}
64
65impl<S: YoetzSuggestion> YoetzAdvisor<S> {
66    /// Create a new advisor with the specified [`consistency_bonus`](Self::consistency_bonus).
67    pub fn new(consistency_bonus: f32) -> Self {
68        Self {
69            consistency_bonus,
70            active_key: None,
71            top_suggestion: None,
72        }
73    }
74
75    /// The [`Key`](YoetzSuggestion::Key) of the currently active behavior.
76    ///
77    /// This can be used to implement a state machine behavior, where the code that suggests a
78    /// behavior can look at the current state.
79    pub fn active_key(&self) -> &Option<S::Key> {
80        &self.active_key
81    }
82
83    /// Suggest a behavior for the AI to consider.
84    ///
85    /// A suggestion should be sent every frame as long as it is valid - once it stops being sent
86    /// it will immediately be replaced by another suggestion.
87    pub fn suggest(&mut self, score: f32, suggestion: S) {
88        if let Some((current_score, _)) = self.top_suggestion.as_ref() {
89            let bonus = if self
90                .active_key
91                .as_ref()
92                .map(|key| *key == suggestion.key())
93                .unwrap_or(false)
94            {
95                self.consistency_bonus
96            } else {
97                0.0
98            };
99            if score + bonus < *current_score {
100                return;
101            }
102        }
103        self.top_suggestion = Some((score, suggestion));
104    }
105}
106
107pub fn update_advisor<S: YoetzSuggestion>(
108    mut query: Query<(Entity, &mut YoetzAdvisor<S>, S::OmniQuery)>,
109    mut commands: Commands,
110) {
111    for (entity, mut advisor, mut components) in query.iter_mut() {
112        let Some((_, mut suggestion)) = advisor.top_suggestion.take() else {
113            continue;
114        };
115        let key = suggestion.key();
116        let mut cmd;
117        if let Some(old_key) = advisor.active_key.as_ref() {
118            if *old_key == key {
119                let update_result = suggestion.update_into_components(&mut components);
120                if let Err(update_result) = update_result {
121                    warn!(
122                        "Components were wrong - will not update, add them with a command instead"
123                    );
124                    suggestion = update_result;
125                } else {
126                    continue;
127                }
128            }
129            cmd = commands.entity(entity);
130            S::remove_components(old_key, &mut cmd)
131        } else {
132            cmd = commands.entity(entity);
133        }
134        suggestion.add_components(&mut cmd);
135        advisor.active_key = Some(key);
136    }
137}