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