stretch.rs

use bevy_ecs::{Changed, Entity, Query, Res, ResMut, With, Without};
use bevy_math::Vec2;
use bevy_render::renderer::RenderResources;
use bevy_transform::prelude::{Children, Parent, Transform};
use bevy_window::Windows;
use std::collections::HashMap;
use stretch::{number::Number, result::Layout, style::Style, Stretch};

#[derive(Debug, Clone, Default, RenderResources)]
pub struct Node {
    pub size: Vec2,
    #[render_resources(ignore)]
    pub style: Style,
    #[render_resources(ignore)]
    pub root: Option<Entity>,
}

#[derive(Default)]
pub struct StretchHierarchies {
    hierarchies: HashMap<Entity, StretchHierarchy>,
}

pub struct StretchHierarchy {
    entity_to_stretch: HashMap<Entity, stretch::node::Node>,
    stretch_to_entity: HashMap<stretch::node::Node, Entity>,
    root_entity: Entity,
    stretch: Stretch,
}

impl StretchHierarchy {
    fn new(root_entity: Entity) -> Self {
        Self {
            entity_to_stretch: Default::default(),
            stretch_to_entity: Default::default(),
            stretch: Stretch::new(),
            root_entity,
        }
    }
    pub fn upsert_node(&mut self, entity: Entity, node: &Node) {
        let mut added = false;
        let stretch = &mut self.stretch;
        let stretch_to_entity = &mut self.stretch_to_entity;
        let stretch_node = self.entity_to_stretch.entry(entity).or_insert_with(|| {
            added = true;
            let stretch_node = stretch.new_node(node.style, Vec::new()).unwrap();
            stretch_to_entity.insert(stretch_node, entity);
            stretch_node
        });

        if !added {
            self.stretch.set_style(*stretch_node, node.style).unwrap();
        }
    }

    pub fn update_children(&mut self, entity: Entity, children: &Children) {
        let mut stretch_children = Vec::with_capacity(children.len());
        for child in children.iter() {
            let stretch_node = self.entity_to_stretch.get(child).unwrap();
            stretch_children.push(*stretch_node);
        }

        let stretch_node = self.entity_to_stretch.get(&entity).unwrap();

        self.stretch
            .set_children(*stretch_node, stretch_children)
            .unwrap();
    }

    pub fn compute_layout(&mut self, width: f32, height: f32) {
        let root_node = self.entity_to_stretch.get(&self.root_entity).unwrap();
        self.stretch
            .compute_layout(
                *root_node,
                stretch::geometry::Size {
                    width: Number::Defined(width),
                    height: Number::Defined(height),
                },
            )
            .unwrap();
    }

    pub fn get_layout(&self, entity: Entity) -> Result<&Layout, stretch::Error> {
        let stretch_node = self.entity_to_stretch.get(&entity).unwrap();
        self.stretch.layout(*stretch_node)
    }
}

// SAFE: as long as MeasureFunc is Send + Sync. https://github.com/vislyhq/stretch/issues/69
unsafe impl Send for StretchHierarchies {}
unsafe impl Sync for StretchHierarchies {}

pub fn stretch_node_system(
    windows: Res<Windows>,
    mut stretch_hierarchies: ResMut<StretchHierarchies>,
    mut root_node_query: Query<With<Node, Without<Parent, Entity>>>,
    mut node_query: Query<(Entity, Changed<Node>)>,
    mut children_query: Query<(Entity, &Node, Changed<Children>)>,
    mut node_transform_query: Query<(Entity, &mut Node, &mut Transform)>,
) {
    // initialize stretch hierarchies
    for root_entity in &mut root_node_query.iter() {
        stretch_hierarchies
            .hierarchies
            .entry(root_entity)
            .or_insert_with(|| StretchHierarchy::new(root_entity));
    }

    // update changed nodes
    for (entity, node) in &mut node_query.iter() {
        // TODO: remove node from old hierarchy if its root has changed
        let root_entity = node.root.unwrap_or(entity);
        let hierarchy = stretch_hierarchies
            .hierarchies
            .get_mut(&root_entity)
            .unwrap();
        hierarchy.upsert_node(entity, &node);
    }

    // TODO: handle removed nodes

    // update children
    for (entity, node, children) in &mut children_query.iter() {
        let root_entity = node.root.unwrap_or(entity);
        let hierarchy = stretch_hierarchies
            .hierarchies
            .get_mut(&root_entity)
            .unwrap();
        hierarchy.update_children(entity, &children);
    }

    // compute layouts
    if let Some(window) = windows.get_primary() {
        for stretch_hierarchy in stretch_hierarchies.hierarchies.values_mut() {
            stretch_hierarchy.compute_layout(window.width as f32, window.height as f32);
        }
    }

    for (entity, mut node, mut transform) in &mut node_transform_query.iter() {
        let root_entity = node.root.unwrap_or(entity);
        let hierarchy = stretch_hierarchies
            .hierarchies
            .get_mut(&root_entity)
            .unwrap();
        let layout = hierarchy.get_layout(entity).unwrap();
        node.size = Vec2::new(layout.size.width, layout.size.height);
        let mut position = transform.value.w_axis();
        position.set_x(layout.location.x);
        position.set_y(layout.location.y);
        transform.value.set_w_axis(position);
    }
}