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