starting refactoring graph states and cache

Author: mbrea-c

crates/bevy_animation_graph/src/core/animation_graph/core.rs

@@ -2,7 +2,6 @@ use super::pin;
 use crate::{
     core::{
         animation_node::AnimationNode,
-        context::{CacheReadFilter, CacheWriteFilter},
         duration_data::DurationData,
         edge_data::AnimationEvent,
         errors::{GraphError, GraphValidationError},
@@ -669,15 +668,10 @@ impl AnimationGraph {
         };
 
         let source_value = match source_pin {
-            SourcePin::NodeData(node_id, _) => {
-                if ctx
-                    .caches()
-                    .get(
-                        |c| c.is_updated(node_id).then_some(()),
-                        CacheReadFilter::for_temp(ctx.temp_cache),
-                    )
-                    .is_none()
-                {
+            SourcePin::NodeData(node_id, node_pin) => {
+                let key = ctx.state_key;
+
+                if !ctx.node_caches().is_updated(node_id.to_owned(), key) {
                     let node = &self.nodes[node_id];
                     let should_debug = node.should_debug;
 
@@ -687,22 +681,11 @@ impl AnimationGraph {
                         node.update(ctx.with_node(node_id, self).with_debugging(should_debug))?;
                     }
 
-                    let is_temp = ctx.temp_cache;
-
-                    ctx.caches_mut().set(
-                        |c| c.set_updated(node_id.clone()),
-                        CacheWriteFilter::for_temp(is_temp),
-                    );
+                    ctx.node_caches_mut().mark_updated(node_id.to_owned(), key);
                 }
 
-                let Some(value) = ctx.caches().get(
-                    |c| c.get_data(source_pin).cloned(),
-                    CacheReadFilter::for_temp(ctx.temp_cache),
-                ) else {
-                    return Err(GraphError::OutputMissing(source_pin.clone()));
-                };
-
-                value
+                ctx.node_caches()
+                    .get_output_data(node_id.clone(), key, node_pin.clone())?
             }
             SourcePin::InputData(pin_id) => ctx
                 .parent_data_back(pin_id)
@@ -740,23 +723,15 @@ impl AnimationGraph {
                 panic!("Incompatible pins connected: {source_pin:?} --> {target_pin:?}")
             }
             SourcePin::NodeTime(node_id) => {
-                if let Some(dur) = ctx.caches().get(
-                    |c| c.get_duration(source_pin),
-                    CacheReadFilter::for_temp(ctx.temp_cache),
-                ) {
+                let key = ctx.state_key;
+
+                if let Ok(dur) = ctx.node_caches().get_duration(node_id.clone(), key) {
                     dur
                 } else {
                     let node = &self.nodes[node_id];
                     let should_debug = node.should_debug;
                     node.duration(ctx.with_node(node_id, self).with_debugging(should_debug))?;
-                    let Some(dur) = ctx.caches().get(
-                        |c| c.get_duration(source_pin),
-                        CacheReadFilter::for_temp(ctx.temp_cache),
-                    ) else {
-                        // TODO: Make a graph error for duration missing
-                        return Err(GraphError::OutputMissing(source_pin.clone()));
-                    };
-                    dur
+                    ctx.node_caches().get_duration(node_id.clone(), key)?
                 }
             }
             SourcePin::InputTime(pin_id) => {
@@ -780,37 +755,22 @@ impl AnimationGraph {
             return Err(GraphError::MissingEdgeToSource(source_pin));
         };
 
+        let key = ctx.state_key;
+
         match target_pin {
             TargetPin::NodeData(_, _) => {
                 panic!("Incompatible pins connected: {source_pin:?} --> {target_pin:?}")
             }
             TargetPin::OutputData(_) => {
                 panic!("Incompatible pins connected: {source_pin:?} --> {target_pin:?}")
             }
-            TargetPin::NodeTime(_, _) => {
-                let Some(time_update) = ctx.caches().get(
-                    |c| c.get_time_update_back(target_pin).cloned(),
-                    CacheReadFilter::for_temp(ctx.temp_cache),
-                ) else {
-                    return Err(GraphError::TimeUpdateMissing(target_pin.clone()));
-                };
-
-                Ok(time_update)
-            }
-            TargetPin::OutputTime => {
-                let Some(time_update) = ctx
-                    .caches()
-                    .get(
-                        |c| c.get_time_update_back(target_pin).cloned(),
-                        CacheReadFilter::for_temp(ctx.temp_cache),
-                    )
-                    .or_else(|| ctx.parent_time_update_fwd().ok())
-                else {
-                    return Err(GraphError::TimeUpdateMissing(target_pin.clone()));
-                };
-
-                Ok(time_update)
-            }
+            TargetPin::NodeTime(target_node, target_pin) => ctx
+                .node_caches()
+                .get_input_time_update(target_node.clone(), key, target_pin.clone()),
+            TargetPin::OutputTime => match ctx.context().query_output_time.clone() {
+                Some(update) => Ok(update),
+                None => ctx.parent_time_update_fwd(),
+            },
         }
     }
 
@@ -855,10 +815,7 @@ impl AnimationGraph {
             entity_map,
             deferred_gizmos,
         );
-        ctx.caches_mut().set(
-            |c| c.set_time_update_back(TargetPin::OutputTime, time_update),
-            CacheWriteFilter::Primary,
-        );
+        ctx.context_mut().query_output_time = Some(time_update);
         let mut outputs = HashMap::new();
         for k in self.output_parameters.keys() {
             let out = self.get_data(TargetPin::OutputData(k.clone()), ctx.clone())?;

crates/bevy_animation_graph/src/core/animation_node.rs

@@ -10,7 +10,7 @@ use crate::{
 use bevy::{
     platform::collections::HashMap,
     prelude::{Deref, DerefMut},
-    reflect::{FromType, prelude::*},
+    reflect::{prelude::*, FromType},
 };
 use std::{any::TypeId, fmt::Debug};
 
@@ -267,15 +267,15 @@ const _: () = {
         */
     {
         #[inline]
-        fn into_any(self: ::std::boxed::Box<Self>) -> ::std::boxed::Box<dyn ::core::any::Any> {
+        fn into_any(self: ::std::boxed::Box<Self>) -> ::std::boxed::Box<dyn::core::any::Any> {
             self
         }
         #[inline]
-        fn as_any(&self) -> &dyn ::core::any::Any {
+        fn as_any(&self) -> &dyn::core::any::Any {
             self
         }
         #[inline]
-        fn as_any_mut(&mut self) -> &mut dyn ::core::any::Any {
+        fn as_any_mut(&mut self) -> &mut dyn::core::any::Any {
             self
         }
         #[inline]

crates/bevy_animation_graph/src/core/context/graph_context.rs

@@ -5,7 +5,7 @@ use crate::{
         prelude::AnimationGraph,
         state_machine::low_level::FSMState,
     },
-    prelude::DataValue,
+    prelude::{DataValue, node_caches::NodeCaches, node_states::NodeStates},
 };
 use bevy::{
     asset::AssetId,
@@ -204,20 +204,25 @@ impl GraphStateStack {
 
 #[derive(Debug, Reflect)]
 pub struct GraphContext {
-    pub caches: GraphStateStack,
+    pub node_states: NodeStates,
+    pub node_caches: NodeCaches,
+    pub query_output_time: Option<TimeUpdate>,
     graph_id: AssetId<AnimationGraph>,
 }
 
 impl GraphContext {
     pub fn new(graph_id: AssetId<AnimationGraph>) -> Self {
         Self {
-            caches: GraphStateStack::default(),
             graph_id,
+            node_states: NodeStates::default(),
+            node_caches: NodeCaches::default(),
+            query_output_time: None,
         }
     }
 
     pub fn next_frame(&mut self) {
-        self.caches.next_frame();
+        self.node_states.next_frame();
+        self.node_caches.next_frame();
     }
 
     pub fn get_graph_id(&self) -> AssetId<AnimationGraph> {

crates/bevy_animation_graph/src/core/context/mod.rs

@@ -1,6 +1,9 @@
 mod deferred_gizmos;
 mod graph_context;
 mod graph_context_arena;
+pub mod node_caches;
+pub mod node_state_box;
+pub mod node_states;
 mod pass_context;
 mod pose_fallback;
 mod spec_context;

crates/bevy_animation_graph/src/core/context/node_caches.rs

@@ -0,0 +1,135 @@
+use bevy::{
+    platform::collections::{HashMap, HashSet},
+    reflect::Reflect,
+};
+
+use crate::{
+    core::{
+        animation_graph::{NodeId, PinId, SourcePin, TargetPin, TimeUpdate},
+        duration_data::DurationData,
+        errors::GraphError,
+    },
+    prelude::{DataValue, node_states::StateKey},
+};
+
+#[derive(Reflect, Default, Debug)]
+pub struct NodeCache {
+    pub output_data: HashMap<(StateKey, PinId), DataValue>,
+    /// Time update coming from the "output time" pin. Perhaps should be called "input time
+    /// update".
+    pub output_time_update: HashMap<StateKey, TimeUpdate>,
+    /// Time updates sent back to nodes via "input time" pins
+    pub input_time_updates: HashMap<(StateKey, PinId), TimeUpdate>,
+    pub duration: HashMap<StateKey, DurationData>,
+    pub updated: HashSet<StateKey>,
+}
+
+#[derive(Reflect, Default, Debug)]
+pub struct NodeCaches {
+    caches: HashMap<NodeId, NodeCache>,
+}
+
+impl NodeCaches {
+    pub fn next_frame(&mut self) {
+        self.caches.clear();
+    }
+
+    pub fn get_duration(&self, node_id: NodeId, key: StateKey) -> Result<DurationData, GraphError> {
+        let error = || GraphError::DurationMissing(SourcePin::NodeTime(node_id.clone()));
+
+        self.caches
+            .get(&node_id)
+            .ok_or_else(&error)
+            .and_then(|c| c.duration.get(&key).ok_or_else(&error))
+            .cloned()
+    }
+
+    pub fn set_duration(&mut self, node_id: NodeId, key: StateKey, duration: DurationData) {
+        self.cache_mut(node_id).duration.insert(key, duration);
+    }
+
+    pub fn get_output_data(
+        &self,
+        node_id: NodeId,
+        key: StateKey,
+        pin: PinId,
+    ) -> Result<DataValue, GraphError> {
+        let error = || GraphError::OutputMissing(SourcePin::NodeData(node_id.clone(), pin.clone()));
+
+        self.caches
+            .get(&node_id)
+            .ok_or_else(&error)
+            .and_then(|c| c.output_data.get(&(key, pin.clone())).ok_or_else(&error))
+            .cloned()
+    }
+
+    pub fn set_output_data(&mut self, node_id: NodeId, key: StateKey, pin: PinId, data: DataValue) {
+        self.cache_mut(node_id).output_data.insert((key, pin), data);
+    }
+
+    pub fn get_output_time_update(
+        &self,
+        node_id: NodeId,
+        key: StateKey,
+    ) -> Result<TimeUpdate, GraphError> {
+        let error = || GraphError::TimeUpdateMissingFwd(SourcePin::NodeTime(node_id.clone()));
+
+        self.caches
+            .get(&node_id)
+            .ok_or_else(&error)
+            .and_then(|c| c.output_time_update.get(&key).ok_or_else(&error))
+            .cloned()
+    }
+
+    pub fn set_output_time_update(&mut self, node_id: NodeId, key: StateKey, update: TimeUpdate) {
+        self.cache_mut(node_id)
+            .output_time_update
+            .insert(key, update);
+    }
+
+    pub fn get_input_time_update(
+        &self,
+        node_id: NodeId,
+        key: StateKey,
+        pin: PinId,
+    ) -> Result<TimeUpdate, GraphError> {
+        let error =
+            || GraphError::TimeUpdateMissingBack(TargetPin::NodeTime(node_id.clone(), pin.clone()));
+
+        self.caches
+            .get(&node_id)
+            .ok_or_else(&error)
+            .and_then(|c| {
+                c.input_time_updates
+                    .get(&(key, pin.clone()))
+                    .ok_or_else(&error)
+            })
+            .cloned()
+    }
+
+    pub fn set_input_time_update(
+        &mut self,
+        node_id: NodeId,
+        key: StateKey,
+        pin: PinId,
+        update: TimeUpdate,
+    ) {
+        self.cache_mut(node_id)
+            .input_time_updates
+            .insert((key, pin), update);
+    }
+
+    pub fn is_updated(&self, node_id: NodeId, key: StateKey) -> bool {
+        self.caches
+            .get(&node_id)
+            .is_some_and(|c| c.updated.contains(&key))
+    }
+
+    pub fn mark_updated(&mut self, node_id: NodeId, key: StateKey) {
+        self.cache_mut(node_id).updated.insert(key);
+    }
+
+    fn cache_mut(&mut self, node_id: NodeId) -> &mut NodeCache {
+        self.caches.entry(node_id).or_default()
+    }
+}