tawasuyu/taffy

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