tawasuyu/taffy

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