Skip to main content

khora_infra/ui/taffy/
taffy_layout.rs

1// Copyright 2025 eraflo
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15//! # UI Layout System
16//!
17//! Computes layouts for all entities with a `UiNode` and writes the results to `UiTransform`.
18
19use khora_core::ecs::entity::EntityId;
20use khora_core::math::Vec2;
21use khora_core::ui::layout::UiLayoutView;
22use khora_core::ui::types::{UiFlexDirection, UiNode, UiTransform, UiVal};
23use khora_core::ui::LayoutSystem;
24use std::any::Any;
25use std::collections::HashMap;
26use taffy::prelude::*;
27
28/// Layout system implementation using the Taffy layout engine.
29pub struct TaffyLayoutSystem {
30    taffy: TaffyTree,
31    entity_to_node: HashMap<EntityId, NodeId>,
32}
33
34impl Default for TaffyLayoutSystem {
35    fn default() -> Self {
36        Self::new()
37    }
38}
39
40impl TaffyLayoutSystem {
41    /// Creates a new TaffyLayoutSystem.
42    pub fn new() -> Self {
43        Self {
44            taffy: TaffyTree::new(),
45            entity_to_node: HashMap::new(),
46        }
47    }
48}
49
50impl LayoutSystem for TaffyLayoutSystem {
51    fn compute_layouts(&mut self, view: &mut dyn UiLayoutView) {
52        self.taffy.clear();
53        self.entity_to_node.clear();
54
55        // Real layout viewport, resolved from the view (falls back to the
56        // trait default when the view has no live surface size).
57        let (viewport_w, viewport_h) = view.viewport_size();
58
59        // 1. First pass: Create taffy nodes for all UI entities
60        let entities = view.get_all_ui_entities();
61        for entity in entities.iter() {
62            if let Some(ui_node) = view.get_node(*entity) {
63                let style = self.convert_style(&ui_node);
64                match self.taffy.new_leaf(style) {
65                    Ok(node) => {
66                        self.entity_to_node.insert(*entity, node);
67                    }
68                    Err(e) => {
69                        log::error!("TaffyLayoutSystem: failed to create layout node: {e}");
70                    }
71                }
72            }
73        }
74
75        // 2. Second pass: Build hierarchy
76        let mut roots = Vec::new();
77
78        for &entity in entities.iter() {
79            if view.has_parent(entity) {
80                continue;
81            }
82            roots.push(entity);
83            self.attach_children(entity, view);
84        }
85
86        // 3. Compute layout against the real viewport.
87        for root in roots {
88            if let Some(&root_node) = self.entity_to_node.get(&root) {
89                let available_space = Size {
90                    width: AvailableSpace::Definite(viewport_w as f32),
91                    height: AvailableSpace::Definite(viewport_h as f32),
92                };
93
94                if let Err(e) = self.taffy.compute_layout(root_node, available_space) {
95                    log::error!("TaffyLayoutSystem: compute_layout failed: {e}");
96                    continue;
97                }
98                self.update_transforms(root, view, 0); // 0 Z-Index base
99            }
100        }
101    }
102
103    fn as_any(&self) -> &dyn Any {
104        self
105    }
106
107    fn as_any_mut(&mut self) -> &mut dyn Any {
108        self
109    }
110}
111
112// SAFETY: TaffyTree contains *const () which makes it !Send/!Sync.
113// However, we only access it from the ISA thread in a deterministic way.
114// In Taffy 0.5+ it should be Send+Sync if we use the default generic.
115// If it still complains, we might need these:
116unsafe impl Send for TaffyLayoutSystem {}
117unsafe impl Sync for TaffyLayoutSystem {}
118
119impl TaffyLayoutSystem {
120    fn attach_children(&mut self, entity: EntityId, view: &dyn UiLayoutView) {
121        let children = view.get_children(entity);
122        if !children.is_empty() {
123            let mut taffy_children = Vec::new();
124            for &child_id in &children {
125                if let Some(&child_node) = self.entity_to_node.get(&child_id) {
126                    self.attach_children(child_id, view);
127                    taffy_children.push(child_node);
128                }
129            }
130
131            if !taffy_children.is_empty() {
132                if let Some(&parent_node) = self.entity_to_node.get(&entity) {
133                    if let Err(e) = self.taffy.set_children(parent_node, &taffy_children) {
134                        log::error!("TaffyLayoutSystem: failed to set children: {e}");
135                    }
136                }
137            }
138        }
139    }
140
141    fn update_transforms(&self, entity: EntityId, view: &mut dyn UiLayoutView, z_index: i32) {
142        if let Some(&node_id) = self.entity_to_node.get(&entity) {
143            if let Ok(layout) = self.taffy.layout(node_id) {
144                let transform = UiTransform {
145                    pos: Vec2::new(layout.location.x, layout.location.y),
146                    size: Vec2::new(layout.size.width, layout.size.height),
147                    z_index,
148                };
149
150                // Add parent's absolute position logic?
151                // Taffy layout is relative to parent by default in some configurations.
152                // In our implementation we assume absolute screen coordinates are needed.
153                // However, Taffy gives local coords.
154
155                // Writing transform back
156                view.set_transform(entity, transform);
157            }
158        }
159
160        let children = view.get_children(entity);
161        for child_id in children {
162            self.update_transforms(child_id, view, z_index + 1);
163        }
164    }
165
166    fn convert_style(&self, node: &UiNode) -> Style {
167        Style {
168            size: Size {
169                width: self.convert_val(node.width),
170                height: self.convert_val(node.height),
171            },
172            min_size: Size {
173                width: self.convert_val(node.min_width),
174                height: self.convert_val(node.min_height),
175            },
176            max_size: Size {
177                width: self.convert_val(node.max_width),
178                height: self.convert_val(node.max_height),
179            },
180            padding: Rect {
181                left: self.convert_length_percentage(node.padding.left),
182                right: self.convert_length_percentage(node.padding.right),
183                top: self.convert_length_percentage(node.padding.top),
184                bottom: self.convert_length_percentage(node.padding.bottom),
185            },
186            margin: Rect {
187                left: self.convert_length_percentage_auto(node.margin.left),
188                right: self.convert_length_percentage_auto(node.margin.right),
189                top: self.convert_length_percentage_auto(node.margin.top),
190                bottom: self.convert_length_percentage_auto(node.margin.bottom),
191            },
192            flex_direction: match node.flex_direction {
193                UiFlexDirection::Column => FlexDirection::Column,
194                UiFlexDirection::Row => FlexDirection::Row,
195            },
196            flex_grow: node.flex_grow,
197            flex_shrink: node.flex_shrink,
198            ..Default::default()
199        }
200    }
201
202    fn convert_val(&self, val: UiVal) -> Dimension {
203        match val {
204            UiVal::Px(v) => length(v),
205            UiVal::Percent(v) => percent(v / 100.0),
206            UiVal::Auto => auto(),
207        }
208    }
209
210    fn convert_length_percentage(&self, val: UiVal) -> LengthPercentage {
211        match val {
212            UiVal::Px(v) => length(v),
213            UiVal::Percent(v) => percent(v / 100.0),
214            UiVal::Auto => length(0.0_f32),
215        }
216    }
217
218    fn convert_length_percentage_auto(&self, val: UiVal) -> LengthPercentageAuto {
219        match val {
220            UiVal::Px(v) => length(v),
221            UiVal::Percent(v) => percent(v / 100.0),
222            UiVal::Auto => auto(),
223        }
224    }
225}