fix: child graph contexts now use the correct graph reference

Author: mbrea-c

crates/bevy_animation_graph_builtin_nodes/src/graph_node.rs

@@ -41,7 +41,7 @@ impl NodeLike for GraphNode {
         };
 
         let sub_ctx = ctx
-            .create_child_context(self.graph.id(), None)
+            .create_child_context(self.graph.id(), graph, None)
             .with_io(&sub_ctx_io);
 
         if graph.io_spec.has_output_time() {
@@ -70,7 +70,7 @@ impl NodeLike for GraphNode {
         };
 
         let mut sub_ctx = ctx
-            .create_child_context(self.graph.id(), None)
+            .create_child_context(self.graph.id(), graph, None)
             .with_io(&sub_ctx_io);
 
         if graph.io_spec.has_output_time() {

crates/bevy_animation_graph_core/src/context/new_context.rs

@@ -292,6 +292,8 @@ impl<'a> NodeContext<'a> {
     pub fn create_child_context(
         &self,
         subgraph_id: AssetId<AnimationGraph>,
+        // We pass the subgraph to avoid a wasted asset lookup in the hot path
+        subgraph: &'a AnimationGraph,
         fsm_state: Option<LowLevelStateId>,
     ) -> GraphContext<'a> {
         let subctx_id = SubContextId {
@@ -306,7 +308,25 @@ impl<'a> NodeContext<'a> {
                 .context_arena
                 .get_mut()
                 .get_sub_context_or_insert_default(subctx_id, subgraph_id),
+            graph: subgraph,
             ..self.graph_context.clone()
         }
     }
+
+    pub fn create_child_context_id(
+        &self,
+        subgraph_id: AssetId<AnimationGraph>,
+        fsm_state: Option<LowLevelStateId>,
+    ) -> GraphContextId {
+        let subctx_id = SubContextId {
+            ctx_id: self.graph_context.context_id,
+            node_id: self.node_id.to_owned(),
+            state_id: fsm_state,
+        };
+
+        self.graph_context
+            .context_arena
+            .get_mut()
+            .get_sub_context_or_insert_default(subctx_id, subgraph_id)
+    }
 }

crates/bevy_animation_graph_core/src/state_machine/low_level/mod.rs

@@ -264,9 +264,15 @@ impl LowLevelStateMachine {
         }
 
         if let Some((graph, state)) = should_clear {
-            ctx.create_child_context(graph, Some(state))
-                .node_states_mut()
-                .clear();
+            let subctx_id = ctx.create_child_context_id(graph, Some(state));
+            if let Some(graph_state) = ctx
+                .graph_context
+                .context_arena
+                .get_mut()
+                .get_context_mut(subctx_id)
+            {
+                graph_state.node_states.clear();
+            }
         }
 
         Ok(())
@@ -351,7 +357,7 @@ impl LowLevelStateMachine {
         );
 
         let sub_ctx = ctx
-            .create_child_context(state.graph.id(), Some(state.id.clone()))
+            .create_child_context(state.graph.id(), graph, Some(state.id.clone()))
             .with_io(&sub_io_env);
 
         for (id, _) in graph.io_spec.iter_output_data() {
@@ -424,7 +430,7 @@ impl LowLevelStateMachine {
         );
 
         let sub_ctx = ctx
-            .create_child_context(state.graph.id(), Some(state.id.clone()))
+            .create_child_context(state.graph.id(), graph, Some(state.id.clone()))
             .with_io(&sub_io_env);
 
         let source_pin = SourcePin::InputTime(GraphInputPin::Passthrough(pin));
@@ -510,7 +516,7 @@ impl<'a> FsmIoEnv<'a> {
 
         let sub_ctx = self
             .node_context
-            .create_child_context(graph_handle.id(), Some(next_state))
+            .create_child_context(graph_handle.id(), graph, Some(next_state))
             .with_state_key(ctx.state_key)
             .with_io(&sub_graph_io);
 
@@ -562,7 +568,7 @@ impl<'a> FsmIoEnv<'a> {
 
         let sub_ctx = self
             .node_context
-            .create_child_context(graph_handle.id(), Some(state))
+            .create_child_context(graph_handle.id(), graph, Some(state))
             .with_state_key(ctx.state_key)
             .with_io(&sub_graph_io);
 
@@ -604,7 +610,7 @@ impl<'a> FsmIoEnv<'a> {
 
             let sub_ctx = self
                 .node_context
-                .create_child_context(graph_handle.id(), Some(next_state))
+                .create_child_context(graph_handle.id(), graph, Some(next_state))
                 .with_state_key(ctx.state_key)
                 .with_io(&sub_graph_io);