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