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