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}