Skip to main content

mittens_engine/engine/ecs/system/
armature_visualization_system.rs

1use crate::engine::ecs::component::{
2    ColorComponent, GLTFComponent, OverlayComponent, RaycastableComponent, RenderableComponent,
3    SignalRouteUpwardComponent, TransformComponent,
4};
5use crate::engine::ecs::system::GLTFSystem;
6use crate::engine::ecs::{ComponentId, IntentValue, SignalEmitter, World};
7use crate::engine::graphics::VisualWorld;
8use crate::engine::user_input::InputState;
9use std::collections::HashMap;
10
11const VIZ_BOX_SCALE: f32 = 0.03;
12
13#[derive(Debug, Default)]
14pub struct ArmatureVisualizationSystem {
15    visualization_roots: HashMap<ComponentId, Vec<ComponentId>>,
16}
17
18impl ArmatureVisualizationSystem {
19    pub fn new() -> Self {
20        Self::default()
21    }
22
23    pub fn registry_entry(&self, component_id: ComponentId) -> Option<&Vec<ComponentId>> {
24        self.visualization_roots.get(&component_id)
25    }
26
27    pub fn tick_with_queue(
28        &mut self,
29        world: &mut World,
30        gltf_system: &GLTFSystem,
31        _visuals: &mut VisualWorld,
32        emit: &mut dyn SignalEmitter,
33        _dt_sec: f32,
34    ) {
35        self.cleanup_dead_entries(world);
36
37        for component_id in gltf_system.tracked_components() {
38            let Some(gltf) = world.get_component_by_id_as::<GLTFComponent>(component_id) else {
39                continue;
40            };
41            if !gltf.spawned {
42                continue;
43            }
44
45            let wants_visible = gltf.armature_visible;
46            let has_visualizations = self
47                .visualization_roots
48                .get(&component_id)
49                .is_some_and(|roots| !roots.is_empty());
50
51            if wants_visible {
52                if has_visualizations {
53                    continue;
54                }
55                let joint_transforms = gltf.armature_joint_transforms.clone();
56                self.ensure_visualizations(world, emit, component_id, &joint_transforms);
57            } else {
58                if !has_visualizations {
59                    continue;
60                }
61                self.remove_visualizations(emit, component_id);
62            }
63        }
64    }
65
66    fn cleanup_dead_entries(&mut self, world: &World) {
67        self.visualization_roots
68            .retain(|component_id, _| world.get_component_record(*component_id).is_some());
69    }
70
71    fn ensure_visualizations(
72        &mut self,
73        world: &mut World,
74        emit: &mut dyn SignalEmitter,
75        component_id: ComponentId,
76        joint_transforms: &[ComponentId],
77    ) {
78        let existing = self.visualization_roots.entry(component_id).or_default();
79        if !existing.is_empty() {
80            existing.retain(|marker_root| world.get_component_record(*marker_root).is_some());
81        }
82        if !existing.is_empty() || joint_transforms.is_empty() {
83            return;
84        }
85
86        for &joint_transform in joint_transforms {
87            if world
88                .get_component_by_id_as::<TransformComponent>(joint_transform)
89                .is_none()
90            {
91                continue;
92            }
93
94            let marker_root = spawn_joint_marker(world, emit, joint_transform);
95            existing.push(marker_root);
96        }
97    }
98
99    fn remove_visualizations(&mut self, emit: &mut dyn SignalEmitter, component_id: ComponentId) {
100        let Some(existing) = self.visualization_roots.get_mut(&component_id) else {
101            return;
102        };
103
104        for &marker_root in existing.iter() {
105            emit.push_intent_now(
106                marker_root,
107                IntentValue::RemoveSubtree {
108                    component_ids: vec![marker_root],
109                },
110            );
111        }
112        existing.clear();
113    }
114}
115
116fn spawn_joint_marker(
117    world: &mut World,
118    emit: &mut dyn SignalEmitter,
119    joint_transform: ComponentId,
120) -> ComponentId {
121    let marker_root = world.add_component_boxed_named(
122        "armature_joint_marker",
123        Box::new(TransformComponent::new().with_scale(VIZ_BOX_SCALE, VIZ_BOX_SCALE, VIZ_BOX_SCALE)),
124    );
125    let route_up = world.add_component_boxed_named(
126        "armature_joint_marker_route",
127        Box::new(SignalRouteUpwardComponent::new(
128            "update_transform",
129            "transform",
130        )),
131    );
132    let overlay = world.add_component_boxed_named(
133        "armature_joint_marker_overlay",
134        Box::new(OverlayComponent::new()),
135    );
136    let renderable = world.add_component_boxed_named(
137        "armature_joint_marker_renderable",
138        Box::new(RenderableComponent::cube()),
139    );
140    let raycastable = world.add_component_boxed_named(
141        "armature_joint_marker_raycastable",
142        Box::new(RaycastableComponent::enabled()),
143    );
144    let color = world.add_component_boxed_named(
145        "armature_joint_marker_color",
146        Box::new(ColorComponent::rgba(1.0, 1.0, 1.0, 1.0)),
147    );
148
149    let _ = world.add_child(joint_transform, marker_root);
150    let _ = world.add_child(marker_root, route_up);
151    let _ = world.add_child(marker_root, overlay);
152    let _ = world.add_child(overlay, renderable);
153    let _ = world.add_child(renderable, raycastable);
154    let _ = world.add_child(renderable, color);
155    world.init_component_tree(marker_root, emit);
156    marker_root
157}
158
159impl crate::engine::ecs::system::System for ArmatureVisualizationSystem {
160    fn tick(
161        &mut self,
162        _world: &mut World,
163        _visuals: &mut VisualWorld,
164        _input: &InputState,
165        _dt_sec: f32,
166    ) {
167    }
168}
169
170#[cfg(test)]
171mod tests {
172    use super::*;
173    use crate::engine::ecs::CommandQueue;
174    use crate::engine::ecs::component::GLTFComponent;
175    use crate::engine::ecs::system::SystemWorld;
176    use crate::engine::graphics::{RenderAssets, VisualWorld};
177
178    #[test]
179    fn armature_visibility_spawns_once_and_removes_idempotently() {
180        let mut world = World::default();
181        let mut visuals = VisualWorld::default();
182        let mut render_assets = RenderAssets::new();
183        let mut systems = SystemWorld::default();
184        let mut queue = CommandQueue::new();
185
186        let editor_root = world.add_component(TransformComponent::new());
187        let gltf_id = world.add_component(GLTFComponent::new("cat.glb"));
188        let joint_a = world.add_component(TransformComponent::new());
189        let joint_b = world.add_component(TransformComponent::new());
190        let _ = world.add_child(editor_root, gltf_id);
191        let _ = world.add_child(editor_root, joint_a);
192        let _ = world.add_child(editor_root, joint_b);
193
194        let gltf = world
195            .get_component_by_id_as_mut::<GLTFComponent>(gltf_id)
196            .expect("gltf");
197        gltf.spawned = true;
198        gltf.armature_visible = true;
199        gltf.armature_joint_transforms = vec![joint_a, joint_b];
200        systems.gltf.register_component(gltf_id);
201
202        systems.armature_visualization.tick_with_queue(
203            &mut world,
204            &systems.gltf,
205            &mut visuals,
206            &mut queue,
207            0.016,
208        );
209        assert_eq!(
210            systems
211                .armature_visualization
212                .registry_entry(gltf_id)
213                .map(Vec::len),
214            Some(2)
215        );
216        assert_eq!(world.children_of(joint_a).len(), 1);
217        assert_eq!(world.children_of(joint_b).len(), 1);
218
219        systems.armature_visualization.tick_with_queue(
220            &mut world,
221            &systems.gltf,
222            &mut visuals,
223            &mut queue,
224            0.016,
225        );
226        assert_eq!(
227            systems
228                .armature_visualization
229                .registry_entry(gltf_id)
230                .map(Vec::len),
231            Some(2)
232        );
233        assert_eq!(world.children_of(joint_a).len(), 1);
234        assert_eq!(world.children_of(joint_b).len(), 1);
235
236        world
237            .get_component_by_id_as_mut::<GLTFComponent>(gltf_id)
238            .expect("gltf")
239            .armature_visible = false;
240        systems.armature_visualization.tick_with_queue(
241            &mut world,
242            &systems.gltf,
243            &mut visuals,
244            &mut queue,
245            0.016,
246        );
247        systems.process_commands(&mut world, &mut visuals, &mut render_assets, &mut queue);
248        assert_eq!(
249            systems
250                .armature_visualization
251                .registry_entry(gltf_id)
252                .map(Vec::len),
253            Some(0)
254        );
255        assert!(world.children_of(joint_a).is_empty());
256        assert!(world.children_of(joint_b).is_empty());
257
258        systems.armature_visualization.tick_with_queue(
259            &mut world,
260            &systems.gltf,
261            &mut visuals,
262            &mut queue,
263            0.016,
264        );
265        systems.process_commands(&mut world, &mut visuals, &mut render_assets, &mut queue);
266        assert_eq!(
267            systems
268                .armature_visualization
269                .registry_entry(gltf_id)
270                .map(Vec::len),
271            Some(0)
272        );
273    }
274}