freya_core/

tree.rs

use std::{
    any::Any,
    borrow::Cow,
    collections::{
        VecDeque,
        hash_map::Entry,
    },
    fmt::Debug,
    rc::Rc,
};

use bitflags::bitflags;
use freya_engine::prelude::{
    FontCollection,
    FontMgr,
};
use futures_channel::mpsc::UnboundedSender;
use itertools::Itertools;
use rustc_hash::{
    FxHashMap,
    FxHashSet,
};
use torin::{
    prelude::{
        Area,
        LayoutMeasurer,
        LayoutNode,
        PostMeasure,
        Size2D,
    },
    torin::{
        DirtyReason,
        Torin,
    },
};

use crate::{
    accessibility::groups::AccessibilityGroups,
    data::{
        AccessibilityState,
        EffectState,
        LayerState,
        TextStyleState,
    },
    element::{
        ElementExt,
        LayoutContext,
        PostMeasureContext,
    },
    elements::rect::RectElement,
    events::{
        data::{
            EventType,
            SizedEventData,
            StyledEventData,
        },
        emittable::EmmitableEvent,
        name::EventName,
    },
    extended_hashmap::ExtendedHashMap,
    integration::{
        AccessibilityDirtyNodes,
        AccessibilityFocusStrategy,
        AccessibilityGenerator,
        EventsChunk,
    },
    layers::Layers,
    node_id::NodeId,
    runner::{
        MutationAdd,
        MutationModified,
        MutationMove,
        MutationRemove,
        Mutations,
    },
    text_cache::TextCache,
    tree_layout_adapter::TreeAdapterFreya,
};

#[derive(Default)]
pub struct Tree {
    pub parents: FxHashMap<NodeId, NodeId>,
    pub children: FxHashMap<NodeId, Vec<NodeId>>,
    pub heights: FxHashMap<NodeId, u16>,

    pub elements: FxHashMap<NodeId, Rc<dyn ElementExt>>,

    // Event listeners
    pub listeners: FxHashMap<EventName, Vec<NodeId>>,

    // Events queued until the next layout measure
    pub events: Vec<EmmitableEvent>,

    // Derived states
    pub layer_state: FxHashMap<NodeId, LayerState>,
    pub accessibility_state: FxHashMap<NodeId, AccessibilityState>,
    pub effect_state: FxHashMap<NodeId, EffectState>,
    pub text_style_state: FxHashMap<NodeId, TextStyleState>,

    // Other
    pub layout: Torin<NodeId>,
    pub layers: Layers,
    pub text_cache: TextCache,

    // Accessibility
    pub accessibility_groups: AccessibilityGroups,
    pub accessibility_diff: AccessibilityDirtyNodes,
    pub accessibility_generator: AccessibilityGenerator,
}

impl Debug for Tree {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("Tree")
            .field("children", &self.children.capacity())
            .field("parents", &self.parents.capacity())
            .field("elements", &self.elements.capacity())
            .field("heights", &self.heights.capacity())
            .field("listeners", &self.listeners.capacity())
            .field("layer_state", &self.layer_state.capacity())
            .field("layout_size", &self.layout.size())
            .field("layers", &self.layers.capacity())
            .field("effect_state", &self.effect_state.capacity())
            .field("accessibility_state", &self.accessibility_state.capacity())
            .field("text_style_state", &self.text_style_state.capacity())
            .field("text_cache", &self.text_cache)
            .finish()
    }
}

impl Tree {
    pub fn size(&self) -> usize {
        self.elements.len()
    }

    pub fn traverse_depth(&self, mut then: impl FnMut(NodeId)) {
        let mut buffer = vec![NodeId::ROOT];
        while let Some(node_id) = buffer.pop() {
            if let Some(children) = self.children.get(&node_id) {
                buffer.extend(children.iter().rev());
            }
            then(node_id);
        }
    }

    pub fn traverse_depth_cancel(&self, mut then: impl FnMut(NodeId) -> bool) {
        let mut buffer = vec![NodeId::ROOT];
        while let Some(node_id) = buffer.pop() {
            if let Some(children) = self.children.get(&node_id) {
                buffer.extend(children.iter().rev());
            }
            if then(node_id) {
                break;
            }
        }
    }

    #[cfg_attr(feature = "hotpath", hotpath::measure)]
    pub fn apply_mutations(&mut self, mutations: Mutations) -> MutationsApplyResult {
        let mut needs_render = !mutations.removed.is_empty();
        let mut needs_accessibility = !mutations.removed.is_empty();
        let mut dirty = Vec::<(NodeId, DiffModifies)>::default();

        #[cfg(debug_assertions)]
        tracing::info!("{mutations:?}");

        if let Entry::Vacant(e) = self.elements.entry(NodeId::ROOT) {
            e.insert(Rc::new(RectElement::default()));
            self.heights.insert(NodeId::ROOT, 0);
            dirty.push((NodeId::ROOT, DiffModifies::all()));
        }

        hotpath::measure_block!("mutations run", {
            for remove in mutations.removed.into_iter().sorted() {
                let node_id = remove.node_id();
                let mut buff = vec![remove];
                let Some(parent_id) = self.parents.get(&node_id).copied() else {
                    continue;
                };
                self.layout.invalidate(parent_id);
                needs_render = true;

                while let Some(remove) = buff.pop() {
                    let node_id = remove.node_id();
                    self.layout.raw_remove(node_id);

                    let parent_id = self.parents.remove(&node_id).unwrap();

                    // Remove element
                    let old_element = self.elements.remove(&node_id).unwrap();

                    if let Some(children) = self.children.get_mut(&parent_id) {
                        match remove {
                            MutationRemove::Element { index, .. } => {
                                children.remove(index as usize);
                            }
                            MutationRemove::Scope { .. } => {
                                children.retain(|id| *id != node_id);
                            }
                        }
                    }

                    // Remove its children too
                    if let Some(children) = self.children.remove(&node_id) {
                        buff.extend(children.into_iter().enumerate().map(|(i, e)| {
                            MutationRemove::Element {
                                id: e,
                                index: i as u32,
                            }
                        }));
                    }

                    // Remove old events
                    if let Some(events) = old_element.events_handlers() {
                        for event in events.keys() {
                            self.listeners
                                .entry(*event)
                                .or_default()
                                .retain(|id| *id != node_id);
                        }
                    }

                    // Remove from the layers
                    let layer_state = self.layer_state.remove(&node_id).unwrap();
                    layer_state.remove(node_id, &mut self.layers);

                    // Remove from the accessibility
                    let accessibility_state = self.accessibility_state.remove(&node_id).unwrap();
                    accessibility_state.remove(
                        node_id,
                        parent_id,
                        &mut self.accessibility_diff,
                        &mut self.accessibility_groups,
                    );

                    // Remove from other states
                    self.heights.remove(&node_id);
                    self.effect_state.remove(&node_id);
                    self.text_style_state.remove(&node_id);
                    self.text_cache.remove(&node_id);
                }
            }

            for MutationAdd {
                node_id,
                parent_id,
                index,
                element,
            } in mutations
                .added
                .into_iter()
                .sorted_by_key(|m| (m.parent_id, m.index))
            {
                let parent_height = *self.heights.entry(parent_id).or_default();

                self.parents.insert(node_id, parent_id);
                self.heights.insert(node_id, parent_height + 1);

                let parent = self.children.entry(parent_id).or_default();

                // TODO: Improve this
                if parent.len() < index as usize + 1 {
                    parent.resize(index as usize + 1, NodeId::PLACEHOLDER);

                    parent[index as usize] = node_id;
                } else if parent.get(index as usize) == Some(&NodeId::PLACEHOLDER) {
                    parent[index as usize] = node_id;
                } else {
                    parent.insert(index as usize, node_id);
                }

                // Add events
                if let Some(events) = element.events_handlers() {
                    for event in events.keys() {
                        self.listeners.entry(*event).or_default().push(node_id);
                    }
                }

                self.elements.insert(node_id, element);
                dirty.push((node_id, DiffModifies::all()));
            }

            for (parent_node_id, movements) in mutations.moved {
                let parent = self.children.get_mut(&parent_node_id).unwrap();
                for MutationMove { index: to, node_id } in
                    movements.into_iter().sorted_by_key(|m| m.index)
                {
                    let from = parent.iter().position(|id| *id == node_id).unwrap();
                    parent.remove(from);
                    parent.insert(to as usize, node_id);
                }
                let mut diff = DiffModifies::empty();
                diff.insert(DiffModifies::REORDER_LAYOUT);
                diff.insert(DiffModifies::ACCESSIBILITY);
                diff.insert(DiffModifies::STYLE);
                dirty.push((parent_node_id, diff));
            }

            for MutationModified {
                node_id,
                element,
                flags,
            } in mutations.modified
            {
                dirty.push((node_id, flags));

                let old_element = self.elements.remove(&node_id).unwrap();

                if flags.contains(DiffModifies::EVENT_HANDLERS) {
                    // Remove old events
                    if let Some(events) = old_element.events_handlers() {
                        for event in events.keys() {
                            self.listeners
                                .entry(*event)
                                .or_default()
                                .retain(|id| *id != node_id);
                        }
                    }

                    // Add new events
                    if let Some(events) = element.events_handlers() {
                        for event in events.keys() {
                            self.listeners.entry(*event).or_default().push(node_id);
                        }
                    }
                }

                self.elements.insert(node_id, element);
            }
        });

        // Run states cascades
        let mut layer_cascades: Vec<NodeId> = Vec::new();
        let mut effects_cascades: Vec<NodeId> = Vec::new();
        let mut text_style_cascades: Vec<NodeId> = Vec::new();
        let mut styled_nodes: FxHashSet<NodeId> = FxHashSet::default();

        assert_eq!(dirty.len(), FxHashSet::from_iter(&dirty).len());

        hotpath::measure_block!("dirty run", {
            for (node_id, flags) in dirty {
                let element = self.elements.get(&node_id).unwrap();
                let height_b = self.heights.get(&node_id).unwrap();

                if flags.contains(DiffModifies::REORDER_LAYOUT) {
                    self.layout
                        .invalidate_with_reason(node_id, DirtyReason::Reorder);
                }

                if flags.contains(DiffModifies::INNER_LAYOUT) {
                    self.layout
                        .invalidate_with_reason(node_id, DirtyReason::InnerLayout);
                }

                if flags.contains(DiffModifies::LAYOUT) {
                    self.layout.invalidate(node_id);
                }

                if !needs_render
                    && (flags.intersects(
                        DiffModifies::STYLE
                            | DiffModifies::LAYER
                            | DiffModifies::EFFECT
                            | DiffModifies::TEXT_STYLE
                            | DiffModifies::LAYOUT
                            | DiffModifies::INNER_LAYOUT
                            | DiffModifies::REORDER_LAYOUT,
                    ))
                {
                    needs_render = true;
                }

                if !needs_accessibility && (flags.intersects(DiffModifies::ACCESSIBILITY)) {
                    needs_accessibility = true;
                }

                if flags.intersects(DiffModifies::STYLE | DiffModifies::TEXT_STYLE)
                    && self
                        .listeners
                        .get(&EventName::Styled)
                        .is_some_and(|listeners| listeners.contains(&node_id))
                {
                    styled_nodes.insert(node_id);
                }

                if flags.contains(DiffModifies::ACCESSIBILITY) {
                    match self.accessibility_state.get_mut(&node_id) {
                        Some(accessibility_state) => accessibility_state.update(
                            node_id,
                            element,
                            &mut self.accessibility_diff,
                            &mut self.accessibility_groups,
                        ),
                        None => {
                            self.accessibility_state.insert(
                                node_id,
                                AccessibilityState::create(
                                    node_id,
                                    element,
                                    &mut self.accessibility_diff,
                                    &self.accessibility_generator,
                                    &mut self.accessibility_groups,
                                ),
                            );
                        }
                    }
                }

                let handle_cascade = |cascades: &mut Vec<NodeId>| {
                    // Skip scanning if we already know this node is the a root
                    if cascades.iter_mut().any(|root| {
                        let height_a = self.heights.get(root).unwrap();

                        match height_a.cmp(height_b) {
                            std::cmp::Ordering::Less => {
                                self.balance_heights(&node_id, root) == Some(*root)
                            }
                            std::cmp::Ordering::Greater => {
                                let balanced_root = self.balance_heights(root, &node_id);
                                match balanced_root {
                                    Some(r) if r == node_id => {
                                        // If this node is ascendant than the
                                        // current root we set it as the new root
                                        *root = node_id;
                                        true
                                    }
                                    _ => false,
                                }
                            }
                            std::cmp::Ordering::Equal => false,
                        }
                    }) {
                        return;
                    }
                    cascades.push(node_id);
                };

                if flags.intersects(DiffModifies::LAYER) {
                    handle_cascade(&mut layer_cascades);
                }
                if flags.intersects(DiffModifies::EFFECT | DiffModifies::LAYER) {
                    let element = self.elements.get(&node_id).unwrap();
                    // Has data or the parent has state
                    let run_cascade = element.effect().is_some()
                        || self
                            .parents
                            .get(&node_id)
                            .map(|parent| self.effect_state.contains_key(parent))
                            .unwrap_or_default();
                    if run_cascade {
                        handle_cascade(&mut effects_cascades);
                    }
                }
                if flags.intersects(DiffModifies::TEXT_STYLE) {
                    handle_cascade(&mut text_style_cascades);
                }
            }
        });

        hotpath::measure_block!("layer cascade", {
            // Run the layer state
            for layer_root in layer_cascades {
                let mut buffer = VecDeque::new();
                buffer.push_front(&layer_root);

                while let Some(node_id) = buffer.pop_front() {
                    let element = self.elements.get(node_id).unwrap();
                    if let Some(parent_node_id) = self.parents.get(node_id) {
                        let entries = self
                            .layer_state
                            .get_disjoint_entries([node_id, parent_node_id], |_id| {
                                LayerState::default()
                            });
                        if let Some([layer_state, parent_layer_state]) = entries {
                            layer_state.update(
                                parent_layer_state,
                                *node_id,
                                element,
                                &mut self.layers,
                            );
                        }
                    } else {
                        assert_eq!(*node_id, NodeId::ROOT);
                        self.layer_state.insert(
                            NodeId::ROOT,
                            LayerState::create_for_root(*node_id, &mut self.layers),
                        );
                    }
                    if let Some(children) = self.children.get(node_id) {
                        buffer.extend(children);
                    }
                }
            }
        });

        hotpath::measure_block!("effect cascade", {
            // Run the effect state
            for effect_root in effects_cascades {
                let mut buffer = VecDeque::new();
                buffer.push_front(&effect_root);

                while let Some(node_id) = buffer.pop_front() {
                    let element = self.elements.get(node_id).unwrap();
                    if let Some(parent_node_id) = self.parents.get(node_id) {
                        let entries = self.effect_state.get_disjoint_two_entries(
                            parent_node_id,
                            node_id,
                            |_id| EffectState::default(),
                            |left, _id| left.clone(),
                        );
                        if let [Some(parent_effect_state), Some(effect_state)] = entries {
                            let effect_data = element.effect();
                            let layer = element.layer();
                            effect_state.update(
                                *parent_node_id,
                                parent_effect_state,
                                *node_id,
                                effect_data,
                                layer,
                            );
                        }
                    } else {
                        assert_eq!(*node_id, NodeId::ROOT);
                    }
                    if let Some(children) = self.children.get(node_id) {
                        buffer.extend(children);
                    }
                }
            }
        });

        hotpath::measure_block!("text style cascade", {
            // Run the text style state
            for text_style_root in text_style_cascades {
                let mut buffer = VecDeque::new();
                buffer.push_front(&text_style_root);

                while let Some(node_id) = buffer.pop_front() {
                    let element = self.elements.get(node_id).unwrap();
                    if let Some(parent_node_id) = self.parents.get(node_id) {
                        let entries = self
                            .text_style_state
                            .get_disjoint_entries([node_id, parent_node_id], |_id| {
                                TextStyleState::default()
                            });
                        if let Some([text_style_state, parent_text_style_state]) = entries {
                            let changed = text_style_state.update(
                                *node_id,
                                parent_text_style_state,
                                element,
                                &mut self.layout,
                            );
                            if changed
                                && self
                                    .listeners
                                    .get(&EventName::Styled)
                                    .is_some_and(|listeners| listeners.contains(node_id))
                            {
                                styled_nodes.insert(*node_id);
                            }
                        }
                    } else {
                        assert_eq!(*node_id, NodeId::ROOT);
                        self.text_style_state
                            .insert(NodeId::ROOT, TextStyleState::default());
                    }
                    if let Some(children) = self.children.get(node_id) {
                        buffer.extend(children);
                    }
                }
            }

            #[cfg(all(debug_assertions, feature = "debug-integrity"))]
            self.verify_tree_integrity();
        });

        for node_id in styled_nodes {
            let element = self.elements.get(&node_id).unwrap();
            let text_style_state = self.text_style_state.get(&node_id).unwrap();
            self.events.push(EmmitableEvent {
                name: EventName::Styled,
                source_event: EventName::Styled,
                node_id,
                data: EventType::Styled(StyledEventData {
                    style: element.style().into_owned(),
                    text_style: text_style_state.clone(),
                }),
                bubbles: false,
            });
        }

        MutationsApplyResult {
            needs_render,
            needs_accessibility,
            auto_focus: self.accessibility_diff.requested_auto_focus.take(),
        }
    }

    /// Walk to the ancestor of `base` with the same height of `target`
    fn balance_heights(&self, base: &NodeId, target: &NodeId) -> Option<NodeId> {
        let target_height = self.heights.get(target)?;
        let mut current = base;
        loop {
            if self.heights.get(current)? == target_height {
                break;
            }

            let parent_current = self.parents.get(current);
            if let Some(parent_current) = parent_current {
                current = parent_current;
            }
        }
        Some(*current)
    }

    pub fn measure_layout(
        &mut self,
        size: Size2D,
        font_collection: &mut FontCollection,
        font_manager: &FontMgr,
        events_sender: &UnboundedSender<EventsChunk>,
        scale_factor: f64,
        fallback_fonts: &[Cow<'static, str>],
    ) {
        let mut tree_adapter = TreeAdapterFreya {
            elements: &self.elements,
            parents: &self.parents,
            children: &self.children,
            heights: &self.heights,
            scale_factor,
        };

        let layout_adapter = LayoutMeasurerAdapter {
            elements: &self.elements,
            text_style_state: &self.text_style_state,
            font_collection,
            font_manager,
            events: &mut self.events,
            scale_factor,
            fallback_fonts,
            text_cache: &mut self.text_cache,
        };

        self.layout.find_best_root(&mut tree_adapter);
        self.layout.measure(
            NodeId::ROOT,
            Area::from_size(size),
            &mut Some(layout_adapter),
            &mut tree_adapter,
        );
        events_sender
            .unbounded_send(EventsChunk::Batch(self.events.drain(..).collect()))
            .unwrap();
    }

    pub fn print_ascii(&self, node_id: NodeId, prefix: String, last: bool) {
        let height = self.heights.get(&node_id).unwrap();
        let layer = self.layer_state.get(&node_id).unwrap();

        // Print current node
        println!(
            "{}{}{:?} [{}] ({})",
            prefix,
            if last { "└── " } else { "├── " },
            node_id,
            height,
            layer.layer
        );

        // Get children
        if let Some(children) = self.children.get(&node_id) {
            let len = children.len();
            for (i, child) in children.iter().enumerate() {
                let is_last = i == len - 1;
                // Extend prefix
                let new_prefix = format!("{}{}", prefix, if last { "    " } else { "│   " });
                self.print_ascii(*child, new_prefix, is_last);
            }
        }
    }

    #[cfg(all(debug_assertions, feature = "debug-integrity"))]
    #[cfg_attr(feature = "hotpath", hotpath::measure)]
    pub fn verify_tree_integrity(&self) {
        let mut visited = FxHashSet::default();
        let size = self.elements.len();
        let mut buffer = vec![NodeId::ROOT];
        while let Some(node_id) = buffer.pop() {
            if visited.contains(&node_id) {
                continue;
            }
            visited.insert(node_id);
            if let Some(parent) = self.parents.get(&node_id) {
                buffer.push(*parent);
            }
            if let Some(children) = self.children.get(&node_id) {
                buffer.extend(children);
            }
        }
        assert_eq!(size, visited.len())
    }
}

bitflags! {
    #[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
    pub struct DiffModifies: u32 {
        const LAYOUT = 1 << 0;
        const STYLE = 1 << 1;
        const ACCESSIBILITY = 1 << 2;
        const EVENT_HANDLERS = 1 << 3;
        const LAYER = 1 << 4;
        const TEXT_STYLE = 1 << 5;
        const EFFECT = 1 << 6;
        const INNER_LAYOUT = 1 << 7;
        const REORDER_LAYOUT = 1 << 8;
    }
}

pub struct MutationsApplyResult {
    pub needs_render: bool,
    pub needs_accessibility: bool,
    pub auto_focus: Option<AccessibilityFocusStrategy>,
}

pub struct LayoutMeasurerAdapter<'a> {
    pub font_collection: &'a mut FontCollection,
    pub font_manager: &'a FontMgr,
    elements: &'a FxHashMap<NodeId, Rc<dyn ElementExt>>,
    text_style_state: &'a FxHashMap<NodeId, TextStyleState>,
    events: &'a mut Vec<EmmitableEvent>,
    scale_factor: f64,
    fallback_fonts: &'a [Cow<'static, str>],
    text_cache: &'a mut TextCache,
}

impl LayoutMeasurer<NodeId> for LayoutMeasurerAdapter<'_> {
    fn measure(
        &mut self,
        node_id: NodeId,
        torin_node: &torin::node::Node,
        area_size: &Size2D,
    ) -> Option<(Size2D, Rc<dyn Any>)> {
        self.elements.get(&node_id)?.measure(LayoutContext {
            node_id,
            torin_node,
            area_size,
            font_collection: self.font_collection,
            font_manager: self.font_manager,
            text_style_state: self.text_style_state.get(&node_id).unwrap(),
            scale_factor: self.scale_factor,
            fallback_fonts: self.fallback_fonts,
            text_cache: self.text_cache,
        })
    }

    fn should_hook_measurement(&mut self, node_id: NodeId) -> bool {
        if let Some(element) = self.elements.get(&node_id) {
            element.should_hook_measurement()
        } else {
            false
        }
    }

    fn should_measure_inner_children(&mut self, node_id: NodeId) -> bool {
        if let Some(element) = self.elements.get(&node_id) {
            element.should_measure_inner_children()
        } else {
            false
        }
    }

    fn should_post_measure(&mut self, node_id: NodeId) -> bool {
        self.elements
            .get(&node_id)
            .is_some_and(|element| element.needs_post_measure())
    }

    fn post_measure(
        &mut self,
        node_id: NodeId,
        node_layout: &LayoutNode,
        children: &[NodeId],
        layout: &Torin<NodeId>,
    ) -> PostMeasure<NodeId> {
        self.elements
            .get(&node_id)
            .unwrap()
            .post_measure(PostMeasureContext {
                node_layout,
                children,
                layout,
                font_collection: self.font_collection,
                text_style_state: self.text_style_state.get(&node_id).unwrap(),
                fallback_fonts: self.fallback_fonts,
                scale_factor: self.scale_factor,
            })
    }

    fn notify_layout_references(
        &mut self,
        node_id: NodeId,
        area: Area,
        visible_area: Area,
        inner_sizes: Size2D,
    ) {
        let mut data = SizedEventData::new(area, visible_area, inner_sizes);
        data.div(self.scale_factor as f32);
        self.events.push(EmmitableEvent {
            node_id,
            name: EventName::Sized,
            data: EventType::Sized(data),
            bubbles: false,
            source_event: EventName::Sized,
        });
    }
}