taffy.rs

  1use crate::{
  2    AbsoluteLength, Bounds, DefiniteLength, Edges, Length, Pixels, Point, Size, Style,
  3    WindowContext,
  4};
  5use collections::{FxHashMap, FxHashSet};
  6use smallvec::SmallVec;
  7use std::fmt::Debug;
  8use taffy::{
  9    AvailableSpace as TaffyAvailableSpace, NodeId, Point as TaffyPoint, Rect as TaffyRect,
 10    Size as TaffySize, TaffyTree, TraversePartialTree,
 11};
 12
 13pub struct TaffyLayoutEngine {
 14    tree: TaffyTree,
 15    styles: FxHashMap<LayoutId, Style>,
 16    absolute_layout_bounds: FxHashMap<LayoutId, Bounds<Pixels>>,
 17    computed_layouts: FxHashSet<LayoutId>,
 18    nodes_to_measure: FxHashMap<
 19        LayoutId,
 20        Box<
 21            dyn FnMut(
 22                Size<Option<Pixels>>,
 23                Size<AvailableSpace>,
 24                &mut WindowContext,
 25            ) -> Size<Pixels>,
 26        >,
 27    >,
 28}
 29
 30static EXPECT_MESSAGE: &str = "we should avoid taffy layout errors by construction if possible";
 31
 32impl TaffyLayoutEngine {
 33    pub fn new() -> Self {
 34        TaffyLayoutEngine {
 35            tree: TaffyTree::new(),
 36            styles: FxHashMap::default(),
 37            absolute_layout_bounds: FxHashMap::default(),
 38            computed_layouts: FxHashSet::default(),
 39            nodes_to_measure: FxHashMap::default(),
 40        }
 41    }
 42
 43    pub fn clear(&mut self) {
 44        self.tree.clear();
 45        self.absolute_layout_bounds.clear();
 46        self.computed_layouts.clear();
 47        self.nodes_to_measure.clear();
 48        self.styles.clear();
 49    }
 50
 51    pub fn requested_style(&self, layout_id: LayoutId) -> Option<&Style> {
 52        self.styles.get(&layout_id)
 53    }
 54
 55    pub fn request_layout(
 56        &mut self,
 57        style: &Style,
 58        rem_size: Pixels,
 59        children: &[LayoutId],
 60    ) -> LayoutId {
 61        let taffy_style = style.to_taffy(rem_size);
 62        let layout_id = if children.is_empty() {
 63            self.tree
 64                .new_leaf(taffy_style)
 65                .expect(EXPECT_MESSAGE)
 66                .into()
 67        } else {
 68            self.tree
 69                // This is safe because LayoutId is repr(transparent) to taffy::tree::NodeId.
 70                .new_with_children(taffy_style, unsafe { std::mem::transmute(children) })
 71                .expect(EXPECT_MESSAGE)
 72                .into()
 73        };
 74        self.styles.insert(layout_id, style.clone());
 75        layout_id
 76    }
 77
 78    pub fn request_measured_layout(
 79        &mut self,
 80        style: Style,
 81        rem_size: Pixels,
 82        measure: impl FnMut(Size<Option<Pixels>>, Size<AvailableSpace>, &mut WindowContext) -> Size<Pixels>
 83            + 'static,
 84    ) -> LayoutId {
 85        let style = style.clone();
 86        let taffy_style = style.to_taffy(rem_size);
 87
 88        let layout_id = self
 89            .tree
 90            .new_leaf_with_context(taffy_style, ())
 91            .expect(EXPECT_MESSAGE)
 92            .into();
 93        self.nodes_to_measure.insert(layout_id, Box::new(measure));
 94        self.styles.insert(layout_id, style.clone());
 95        layout_id
 96    }
 97
 98    // Used to understand performance
 99    #[allow(dead_code)]
100    fn count_all_children(&self, parent: LayoutId) -> anyhow::Result<u32> {
101        let mut count = 0;
102
103        for child in self.tree.children(parent.0)? {
104            // Count this child.
105            count += 1;
106
107            // Count all of this child's children.
108            count += self.count_all_children(LayoutId(child))?
109        }
110
111        Ok(count)
112    }
113
114    // Used to understand performance
115    #[allow(dead_code)]
116    fn max_depth(&self, depth: u32, parent: LayoutId) -> anyhow::Result<u32> {
117        println!(
118            "{parent:?} at depth {depth} has {} children",
119            self.tree.child_count(parent.0)
120        );
121
122        let mut max_child_depth = 0;
123
124        for child in self.tree.children(parent.0)? {
125            max_child_depth = std::cmp::max(max_child_depth, self.max_depth(0, LayoutId(child))?);
126        }
127
128        Ok(depth + 1 + max_child_depth)
129    }
130
131    // Used to understand performance
132    #[allow(dead_code)]
133    fn get_edges(&self, parent: LayoutId) -> anyhow::Result<Vec<(LayoutId, LayoutId)>> {
134        let mut edges = Vec::new();
135
136        for child in self.tree.children(parent.0)? {
137            edges.push((parent, LayoutId(child)));
138
139            edges.extend(self.get_edges(LayoutId(child))?);
140        }
141
142        Ok(edges)
143    }
144
145    pub fn compute_layout(
146        &mut self,
147        id: LayoutId,
148        available_space: Size<AvailableSpace>,
149        cx: &mut WindowContext,
150    ) {
151        // Leaving this here until we have a better instrumentation approach.
152        // println!("Laying out {} children", self.count_all_children(id)?);
153        // println!("Max layout depth: {}", self.max_depth(0, id)?);
154
155        // Output the edges (branches) of the tree in Mermaid format for visualization.
156        // println!("Edges:");
157        // for (a, b) in self.get_edges(id)? {
158        //     println!("N{} --> N{}", u64::from(a), u64::from(b));
159        // }
160        // println!("");
161        //
162
163        if !self.computed_layouts.insert(id) {
164            let mut stack = SmallVec::<[LayoutId; 64]>::new();
165            stack.push(id);
166            while let Some(id) = stack.pop() {
167                self.absolute_layout_bounds.remove(&id);
168                stack.extend(
169                    self.tree
170                        .children(id.into())
171                        .expect(EXPECT_MESSAGE)
172                        .into_iter()
173                        .map(Into::into),
174                );
175            }
176        }
177
178        // let started_at = std::time::Instant::now();
179        self.tree
180            .compute_layout_with_measure(
181                id.into(),
182                available_space.into(),
183                |known_dimensions, available_space, node_id, _context| {
184                    let Some(measure) = self.nodes_to_measure.get_mut(&node_id.into()) else {
185                        return taffy::geometry::Size::default();
186                    };
187
188                    let known_dimensions = Size {
189                        width: known_dimensions.width.map(Pixels),
190                        height: known_dimensions.height.map(Pixels),
191                    };
192
193                    measure(known_dimensions, available_space.into(), cx).into()
194                },
195            )
196            .expect(EXPECT_MESSAGE);
197
198        // println!("compute_layout took {:?}", started_at.elapsed());
199    }
200
201    pub fn layout_bounds(&mut self, id: LayoutId) -> Bounds<Pixels> {
202        if let Some(layout) = self.absolute_layout_bounds.get(&id).cloned() {
203            return layout;
204        }
205
206        let layout = self.tree.layout(id.into()).expect(EXPECT_MESSAGE);
207        let mut bounds = Bounds {
208            origin: layout.location.into(),
209            size: layout.size.into(),
210        };
211
212        if let Some(parent_id) = self.tree.parent(id.0) {
213            let parent_bounds = self.layout_bounds(parent_id.into());
214            bounds.origin += parent_bounds.origin;
215        }
216        self.absolute_layout_bounds.insert(id, bounds);
217
218        bounds
219    }
220}
221
222#[derive(Copy, Clone, Eq, PartialEq, Debug)]
223#[repr(transparent)]
224pub struct LayoutId(NodeId);
225
226impl std::hash::Hash for LayoutId {
227    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
228        u64::from(self.0).hash(state);
229    }
230}
231
232impl From<NodeId> for LayoutId {
233    fn from(node_id: NodeId) -> Self {
234        Self(node_id)
235    }
236}
237
238impl From<LayoutId> for NodeId {
239    fn from(layout_id: LayoutId) -> NodeId {
240        layout_id.0
241    }
242}
243
244trait ToTaffy<Output> {
245    fn to_taffy(&self, rem_size: Pixels) -> Output;
246}
247
248impl ToTaffy<taffy::style::Style> for Style {
249    fn to_taffy(&self, rem_size: Pixels) -> taffy::style::Style {
250        taffy::style::Style {
251            display: self.display,
252            overflow: self.overflow.into(),
253            scrollbar_width: self.scrollbar_width,
254            position: self.position,
255            inset: self.inset.to_taffy(rem_size),
256            size: self.size.to_taffy(rem_size),
257            min_size: self.min_size.to_taffy(rem_size),
258            max_size: self.max_size.to_taffy(rem_size),
259            aspect_ratio: self.aspect_ratio,
260            margin: self.margin.to_taffy(rem_size),
261            padding: self.padding.to_taffy(rem_size),
262            border: self.border_widths.to_taffy(rem_size),
263            align_items: self.align_items,
264            align_self: self.align_self,
265            align_content: self.align_content,
266            justify_content: self.justify_content,
267            gap: self.gap.to_taffy(rem_size),
268            flex_direction: self.flex_direction,
269            flex_wrap: self.flex_wrap,
270            flex_basis: self.flex_basis.to_taffy(rem_size),
271            flex_grow: self.flex_grow,
272            flex_shrink: self.flex_shrink,
273            ..Default::default() // Ignore grid properties for now
274        }
275    }
276}
277
278impl ToTaffy<taffy::style::LengthPercentageAuto> for Length {
279    fn to_taffy(&self, rem_size: Pixels) -> taffy::prelude::LengthPercentageAuto {
280        match self {
281            Length::Definite(length) => length.to_taffy(rem_size),
282            Length::Auto => taffy::prelude::LengthPercentageAuto::Auto,
283        }
284    }
285}
286
287impl ToTaffy<taffy::style::Dimension> for Length {
288    fn to_taffy(&self, rem_size: Pixels) -> taffy::prelude::Dimension {
289        match self {
290            Length::Definite(length) => length.to_taffy(rem_size),
291            Length::Auto => taffy::prelude::Dimension::Auto,
292        }
293    }
294}
295
296impl ToTaffy<taffy::style::LengthPercentage> for DefiniteLength {
297    fn to_taffy(&self, rem_size: Pixels) -> taffy::style::LengthPercentage {
298        match self {
299            DefiniteLength::Absolute(length) => match length {
300                AbsoluteLength::Pixels(pixels) => {
301                    taffy::style::LengthPercentage::Length(pixels.into())
302                }
303                AbsoluteLength::Rems(rems) => {
304                    taffy::style::LengthPercentage::Length((*rems * rem_size).into())
305                }
306            },
307            DefiniteLength::Fraction(fraction) => {
308                taffy::style::LengthPercentage::Percent(*fraction)
309            }
310        }
311    }
312}
313
314impl ToTaffy<taffy::style::LengthPercentageAuto> for DefiniteLength {
315    fn to_taffy(&self, rem_size: Pixels) -> taffy::style::LengthPercentageAuto {
316        match self {
317            DefiniteLength::Absolute(length) => match length {
318                AbsoluteLength::Pixels(pixels) => {
319                    taffy::style::LengthPercentageAuto::Length(pixels.into())
320                }
321                AbsoluteLength::Rems(rems) => {
322                    taffy::style::LengthPercentageAuto::Length((*rems * rem_size).into())
323                }
324            },
325            DefiniteLength::Fraction(fraction) => {
326                taffy::style::LengthPercentageAuto::Percent(*fraction)
327            }
328        }
329    }
330}
331
332impl ToTaffy<taffy::style::Dimension> for DefiniteLength {
333    fn to_taffy(&self, rem_size: Pixels) -> taffy::style::Dimension {
334        match self {
335            DefiniteLength::Absolute(length) => match length {
336                AbsoluteLength::Pixels(pixels) => taffy::style::Dimension::Length(pixels.into()),
337                AbsoluteLength::Rems(rems) => {
338                    taffy::style::Dimension::Length((*rems * rem_size).into())
339                }
340            },
341            DefiniteLength::Fraction(fraction) => taffy::style::Dimension::Percent(*fraction),
342        }
343    }
344}
345
346impl ToTaffy<taffy::style::LengthPercentage> for AbsoluteLength {
347    fn to_taffy(&self, rem_size: Pixels) -> taffy::style::LengthPercentage {
348        match self {
349            AbsoluteLength::Pixels(pixels) => taffy::style::LengthPercentage::Length(pixels.into()),
350            AbsoluteLength::Rems(rems) => {
351                taffy::style::LengthPercentage::Length((*rems * rem_size).into())
352            }
353        }
354    }
355}
356
357impl<T, T2> From<TaffyPoint<T>> for Point<T2>
358where
359    T: Into<T2>,
360    T2: Clone + Default + Debug,
361{
362    fn from(point: TaffyPoint<T>) -> Point<T2> {
363        Point {
364            x: point.x.into(),
365            y: point.y.into(),
366        }
367    }
368}
369
370impl<T, T2> From<Point<T>> for TaffyPoint<T2>
371where
372    T: Into<T2> + Clone + Default + Debug,
373{
374    fn from(val: Point<T>) -> Self {
375        TaffyPoint {
376            x: val.x.into(),
377            y: val.y.into(),
378        }
379    }
380}
381
382impl<T, U> ToTaffy<TaffySize<U>> for Size<T>
383where
384    T: ToTaffy<U> + Clone + Default + Debug,
385{
386    fn to_taffy(&self, rem_size: Pixels) -> TaffySize<U> {
387        TaffySize {
388            width: self.width.to_taffy(rem_size),
389            height: self.height.to_taffy(rem_size),
390        }
391    }
392}
393
394impl<T, U> ToTaffy<TaffyRect<U>> for Edges<T>
395where
396    T: ToTaffy<U> + Clone + Default + Debug,
397{
398    fn to_taffy(&self, rem_size: Pixels) -> TaffyRect<U> {
399        TaffyRect {
400            top: self.top.to_taffy(rem_size),
401            right: self.right.to_taffy(rem_size),
402            bottom: self.bottom.to_taffy(rem_size),
403            left: self.left.to_taffy(rem_size),
404        }
405    }
406}
407
408impl<T, U> From<TaffySize<T>> for Size<U>
409where
410    T: Into<U>,
411    U: Clone + Default + Debug,
412{
413    fn from(taffy_size: TaffySize<T>) -> Self {
414        Size {
415            width: taffy_size.width.into(),
416            height: taffy_size.height.into(),
417        }
418    }
419}
420
421impl<T, U> From<Size<T>> for TaffySize<U>
422where
423    T: Into<U> + Clone + Default + Debug,
424{
425    fn from(size: Size<T>) -> Self {
426        TaffySize {
427            width: size.width.into(),
428            height: size.height.into(),
429        }
430    }
431}
432
433#[derive(Copy, Clone, Default, Debug, Eq, PartialEq)]
434pub enum AvailableSpace {
435    /// The amount of space available is the specified number of pixels
436    Definite(Pixels),
437    /// The amount of space available is indefinite and the node should be laid out under a min-content constraint
438    #[default]
439    MinContent,
440    /// The amount of space available is indefinite and the node should be laid out under a max-content constraint
441    MaxContent,
442}
443
444impl From<AvailableSpace> for TaffyAvailableSpace {
445    fn from(space: AvailableSpace) -> TaffyAvailableSpace {
446        match space {
447            AvailableSpace::Definite(Pixels(value)) => TaffyAvailableSpace::Definite(value),
448            AvailableSpace::MinContent => TaffyAvailableSpace::MinContent,
449            AvailableSpace::MaxContent => TaffyAvailableSpace::MaxContent,
450        }
451    }
452}
453
454impl From<TaffyAvailableSpace> for AvailableSpace {
455    fn from(space: TaffyAvailableSpace) -> AvailableSpace {
456        match space {
457            TaffyAvailableSpace::Definite(value) => AvailableSpace::Definite(Pixels(value)),
458            TaffyAvailableSpace::MinContent => AvailableSpace::MinContent,
459            TaffyAvailableSpace::MaxContent => AvailableSpace::MaxContent,
460        }
461    }
462}
463
464impl From<Pixels> for AvailableSpace {
465    fn from(pixels: Pixels) -> Self {
466        AvailableSpace::Definite(pixels)
467    }
468}
469
470impl From<Size<Pixels>> for Size<AvailableSpace> {
471    fn from(size: Size<Pixels>) -> Self {
472        Size {
473            width: AvailableSpace::Definite(size.width),
474            height: AvailableSpace::Definite(size.height),
475        }
476    }
477}