1use core::ops::Drop;
2
3use crate::forest::Forest;
4use crate::geometry::Size;
5use crate::id::{Allocator, Id, NodeId};
6use crate::number::Number;
7use crate::result::Layout;
8use crate::style::Style;
9#[cfg(any(feature = "std", feature = "alloc"))]
10use crate::sys::Box;
11use crate::sys::{new_map_with_capacity, ChildrenVec, Map, Vec};
12use crate::Error;
13
14pub enum MeasureFunc {
15 Raw(fn(Size<Number>) -> Size<f32>),
16 #[cfg(any(feature = "std", feature = "alloc"))]
17 Boxed(Box<dyn Fn(Size<Number>) -> Size<f32>>),
18}
19
20/// Global stretch instance id allocator.
21static INSTANCE_ALLOCATOR: Allocator = Allocator::new();
22
23#[derive(Copy, Clone, Debug, PartialEq, Eq, Hash)]
24#[cfg_attr(not(any(feature = "std", feature = "alloc")), derive(hash32_derive::Hash32))]
25pub struct Node {
26 instance: Id,
27 local: Id,
28}
29
30pub struct Stretch {
31 id: Id,
32 nodes: Allocator,
33 nodes_to_ids: Map<Node, NodeId>,
34 ids_to_nodes: Map<NodeId, Node>,
35 forest: Forest,
36}
37
38impl Default for Stretch {
39 fn default() -> Self {
40 Self::with_capacity(16)
41 }
42}
43
44impl Stretch {
45 pub fn new() -> Self {
46 Default::default()
47 }
48
49 pub fn with_capacity(capacity: usize) -> Self {
50 Self {
51 id: INSTANCE_ALLOCATOR.allocate(),
52 nodes: Allocator::new(),
53 nodes_to_ids: new_map_with_capacity(capacity),
54 ids_to_nodes: new_map_with_capacity(capacity),
55 forest: Forest::with_capacity(capacity),
56 }
57 }
58
59 fn allocate_node(&mut self) -> Node {
60 let local = self.nodes.allocate();
61 Node { instance: self.id, local }
62 }
63
64 fn add_node(&mut self, node: Node, id: NodeId) {
65 let _ = self.nodes_to_ids.insert(node, id);
66 let _ = self.ids_to_nodes.insert(id, node);
67 }
68
69 // Find node in the forest.
70 fn find_node(&self, node: Node) -> Result<NodeId, Error> {
71 match self.nodes_to_ids.get(&node) {
72 Some(id) => Ok(*id),
73 None => Err(Error::InvalidNode(node)),
74 }
75 }
76
77 pub fn new_leaf(&mut self, style: Style, measure: MeasureFunc) -> Result<Node, Error> {
78 let node = self.allocate_node();
79 let id = self.forest.new_leaf(style, measure);
80 self.add_node(node, id);
81 Ok(node)
82 }
83
84 pub fn new_node(&mut self, style: Style, children: &[Node]) -> Result<Node, Error> {
85 let node = self.allocate_node();
86 let children =
87 children.iter().map(|child| self.find_node(*child)).collect::<Result<ChildrenVec<_>, Error>>()?;
88 let id = self.forest.new_node(style, children);
89 self.add_node(node, id);
90 Ok(node)
91 }
92
93 /// Removes all nodes.
94 ///
95 /// All associated nodes will be invalid.
96 pub fn clear(&mut self) {
97 for node in self.nodes_to_ids.keys() {
98 self.nodes.free(&[node.local]);
99 }
100 self.nodes_to_ids.clear();
101 self.ids_to_nodes.clear();
102 self.forest.clear();
103 }
104
105 /// Remove nodes.
106 pub fn remove(&mut self, node: Node) {
107 let id = if let Ok(id) = self.find_node(node) { id } else { return };
108
109 self.nodes_to_ids.remove(&node);
110 self.ids_to_nodes.remove(&id);
111
112 if let Some(new_id) = self.forest.swap_remove(id) {
113 let new = self.ids_to_nodes.remove(&new_id).unwrap();
114 let _ = self.nodes_to_ids.insert(new, id);
115 let _ = self.ids_to_nodes.insert(id, new);
116 }
117 }
118
119 pub fn set_measure(&mut self, node: Node, measure: Option<MeasureFunc>) -> Result<(), Error> {
120 let id = self.find_node(node)?;
121 self.forest.nodes[id].measure = measure;
122 self.forest.mark_dirty(id);
123 Ok(())
124 }
125
126 pub fn add_child(&mut self, node: Node, child: Node) -> Result<(), Error> {
127 let node_id = self.find_node(node)?;
128 let child_id = self.find_node(child)?;
129
130 self.forest.add_child(node_id, child_id);
131 Ok(())
132 }
133
134 pub fn set_children(&mut self, node: Node, children: &[Node]) -> Result<(), Error> {
135 let node_id = self.find_node(node)?;
136 let children_id = children.iter().map(|child| self.find_node(*child)).collect::<Result<ChildrenVec<_>, _>>()?;
137
138 // Remove node as parent from all its current children.
139 for child in &self.forest.children[node_id] {
140 self.forest.parents[*child].retain(|p| *p != node_id);
141 }
142
143 // Build up relation node <-> child
144 for child in &children_id {
145 self.forest.parents[*child].push(node_id);
146 }
147 self.forest.children[node_id] = children_id;
148
149 self.forest.mark_dirty(node_id);
150 Ok(())
151 }
152
153 pub fn remove_child(&mut self, node: Node, child: Node) -> Result<Node, Error> {
154 let node_id = self.find_node(node)?;
155 let child_id = self.find_node(child)?;
156
157 let prev_id = self.forest.remove_child(node_id, child_id);
158 Ok(self.ids_to_nodes[&prev_id])
159 }
160
161 pub fn remove_child_at_index(&mut self, node: Node, index: usize) -> Result<Node, Error> {
162 let node_id = self.find_node(node)?;
163 // TODO: index check
164
165 let prev_id = self.forest.remove_child_at_index(node_id, index);
166 Ok(self.ids_to_nodes[&prev_id])
167 }
168
169 pub fn replace_child_at_index(&mut self, node: Node, index: usize, child: Node) -> Result<Node, Error> {
170 let node_id = self.find_node(node)?;
171 let child_id = self.find_node(child)?;
172 // TODO: index check
173
174 self.forest.parents[child_id].push(node_id);
175 let old_child = core::mem::replace(&mut self.forest.children[node_id][index], child_id);
176 self.forest.parents[old_child].retain(|p| *p != node_id);
177
178 self.forest.mark_dirty(node_id);
179
180 Ok(self.ids_to_nodes[&old_child])
181 }
182
183 pub fn children(&self, node: Node) -> Result<Vec<Node>, Error> {
184 let id = self.find_node(node)?;
185 Ok(self.forest.children[id].iter().map(|child| self.ids_to_nodes[child]).collect())
186 }
187
188 pub fn child_at_index(&self, node: Node, index: usize) -> Result<Node, Error> {
189 let id = self.find_node(node)?;
190 Ok(self.ids_to_nodes[&self.forest.children[id][index]])
191 }
192
193 pub fn child_count(&self, node: Node) -> Result<usize, Error> {
194 let id = self.find_node(node)?;
195 Ok(self.forest.children[id].len())
196 }
197
198 pub fn set_style(&mut self, node: Node, style: Style) -> Result<(), Error> {
199 let id = self.find_node(node)?;
200 self.forest.nodes[id].style = style;
201 self.forest.mark_dirty(id);
202 Ok(())
203 }
204
205 pub fn style(&self, node: Node) -> Result<&Style, Error> {
206 let id = self.find_node(node)?;
207 Ok(&self.forest.nodes[id].style)
208 }
209
210 /// Return this node layout relative to its parent
211 pub fn layout(&self, node: Node) -> Result<&Layout, Error> {
212 let id = self.find_node(node)?;
213 Ok(&self.forest.nodes[id].layout)
214 }
215
216 pub fn mark_dirty(&mut self, node: Node) -> Result<(), Error> {
217 let id = self.find_node(node)?;
218 self.forest.mark_dirty(id);
219 Ok(())
220 }
221
222 pub fn dirty(&self, node: Node) -> Result<bool, Error> {
223 let id = self.find_node(node)?;
224 Ok(self.forest.nodes[id].is_dirty)
225 }
226
227 pub fn compute_layout(&mut self, node: Node, size: Size<Number>) -> Result<(), Error> {
228 let id = self.find_node(node)?;
229 self.forest.compute_layout(id, size);
230 Ok(())
231 }
232}
233
234impl Drop for Stretch {
235 fn drop(&mut self) {
236 INSTANCE_ALLOCATOR.free(&[self.id]);
237 }
238}