cushy/src/widgets/array.rs

use std::ops::Deref;

use alot::{LotId, OrderedLots};
use kludgine::figures::units::UPx;
use kludgine::figures::{Point, Rect, Size};

use crate::children::Children;
use crate::context::Context;
use crate::graphics::Graphics;
use crate::tree::ManagedWidget;
use crate::widget::{IntoValue, Value, Widget};
use crate::ConstraintLimit;

#[derive(Debug)]
pub struct Array {
    pub direction: Value<ArrayDirection>,
    pub children: Value<Children>,
    layout: Layout,
    layout_generation: Option<usize>,
    synced_children: Vec<ManagedWidget>,
}

impl Array {
    pub fn new(
        direction: impl IntoValue<ArrayDirection>,
        children: impl IntoValue<Children>,
    ) -> Self {
        let mut direction = direction.into_value();

        let initial_direction = direction.get();

        Self {
            direction,
            children: children.into_value(),
            layout: Layout::new(initial_direction),
            layout_generation: None,
            synced_children: Vec::new(),
        }
    }

    pub fn columns(children: impl IntoValue<Children>) -> Self {
        Self::new(ArrayDirection::columns(), children)
    }

    pub fn rows(children: impl IntoValue<Children>) -> Self {
        Self::new(ArrayDirection::rows(), children)
    }

    fn synchronize_children(&mut self, context: &mut Context<'_, '_>) {
        let current_generation = self.children.generation();
        if current_generation.map_or_else(
            || self.children.map(Children::len) != self.layout.children.len(),
            |gen| Some(gen) != self.layout_generation,
        ) {
            self.layout_generation = self.children.generation();
            self.children.map(|children| {
                for (index, widget) in children.iter().enumerate() {
                    if self
                        .synced_children
                        .get(index)
                        .map_or(true, |child| child != widget)
                    {
                        // These entries do not match. See if we can find the
                        // new id somewhere else, if so we can swap the entries.
                        if let Some((swap_index, _)) = self
                            .synced_children
                            .iter()
                            .enumerate()
                            .skip(index + 1)
                            .find(|(_, child)| *child == widget)
                        {
                            self.synced_children.swap(index, swap_index);
                            self.layout.swap(index, swap_index);
                        } else {
                            // This is a brand new child.
                            self.synced_children
                                .insert(index, context.push_child(widget.clone()));
                            self.layout.insert(index, ArrayDimension::FitContent);
                        }
                    }
                }

                // Any children remaining at the end of this process are ones
                // that have been removed.
                for removed in self.synced_children.drain(children.len()..) {
                    context.remove_child(removed);
                }
                self.layout.truncate(children.len());
            });
        }
    }
}

impl Widget for Array {
    fn redraw(&mut self, graphics: &mut Graphics<'_, '_, '_>, context: &mut Context) {
        self.synchronize_children(context);
        self.layout.update(
            Size::new(
                ConstraintLimit::Known(graphics.size().width),
                ConstraintLimit::Known(graphics.size().height),
            ),
            |child_index, constraints| {
                context
                    .for_other(&self.synced_children[child_index])
                    .measure(constraints, graphics)
            },
        );

        for (index, layout) in self.layout.iter().enumerate() {
            let child = &self.synced_children[index];
            if layout.size > 0 {
                let mut clipped = graphics.clipped_to(Rect::new(
                    self.layout.orientation.make_point(layout.offset, UPx(0)),
                    self.layout
                        .orientation
                        .make_size(layout.size, self.layout.other),
                ));
                context.for_other(child).redraw(&mut clipped);
            }
        }
    }

    fn measure(
        &mut self,
        available_space: Size<ConstraintLimit>,
        graphics: &mut Graphics<'_, '_, '_>,
        context: &mut Context<'_, '_>,
    ) -> Size<UPx> {
        self.synchronize_children(context);

        self.layout
            .update(available_space, |child_index, constraints| {
                context
                    .for_other(&self.synced_children[child_index])
                    .measure(constraints, graphics)
            })
    }
}

#[derive(Debug, Clone, Copy, Eq, PartialEq)]

pub enum ArrayDirection {
    Row { reverse: bool },
    Column { reverse: bool },
}

impl ArrayDirection {
    #[must_use]
    pub const fn columns() -> Self {
        Self::Column { reverse: false }
    }

    #[must_use]
    pub const fn columns_rev() -> Self {
        Self::Column { reverse: true }
    }

    #[must_use]
    pub const fn rows() -> Self {
        Self::Row { reverse: false }
    }

    #[must_use]
    pub const fn rows_rev() -> Self {
        Self::Row { reverse: true }
    }

    pub fn split_size<U>(&self, s: Size<U>) -> (U, U) {
        match self {
            Self::Row { .. } => (s.height, s.width),
            Self::Column { .. } => (s.width, s.height),
        }
    }

    pub fn make_size<U>(&self, measured: U, other: U) -> Size<U> {
        match self {
            Self::Row { .. } => Size::new(other, measured),
            Self::Column { .. } => Size::new(measured, other),
        }
    }

    pub fn make_point<U>(&self, measured: U, other: U) -> Point<U> {
        match self {
            Self::Row { .. } => Point::new(other, measured),
            Self::Column { .. } => Point::new(measured, other),
        }
    }
}

#[derive(Debug, Clone, Copy, Eq, PartialEq)]
pub enum ArrayDimension {
    FitContent,
    Fractional { weight: u8 },
    Exact(UPx),
}

#[derive(Debug)]
struct Layout {
    children: OrderedLots<ArrayDimension>,
    layouts: Vec<ArrayLayout>,
    pub other: UPx,
    total_weights: u32,
    allocated_space: UPx,
    fractional: Vec<(LotId, u8)>,
    measured: Vec<LotId>,
    pub orientation: ArrayDirection,
}

#[derive(Debug, Clone, Copy, Eq, PartialEq)]
struct ArrayLayout {
    pub offset: UPx,
    pub size: UPx,
}

impl Layout {
    pub const fn new(orientation: ArrayDirection) -> Self {
        Self {
            orientation,
            children: OrderedLots::new(),
            layouts: Vec::new(),
            other: UPx(0),
            total_weights: 0,
            allocated_space: UPx(0),
            fractional: Vec::new(),
            measured: Vec::new(),
        }
    }

    #[cfg(test)] // only used in testing
    pub fn push(&mut self, child: ArrayDimension) {
        self.insert(self.len(), child);
    }

    pub fn remove(&mut self, index: usize) -> ArrayDimension {
        let (id, dimension) = self.children.remove_by_index(index).expect("invalid index");
        self.layouts.remove(index);

        match dimension {
            ArrayDimension::FitContent => {
                self.measured.retain(|&measured| measured != id);
            }
            ArrayDimension::Fractional { weight } => {
                self.fractional.retain(|(measured, _)| *measured != id);
                self.total_weights -= u32::from(weight);
            }
            ArrayDimension::Exact(size) => {
                self.allocated_space -= size;
            }
        }

        dimension
    }

    pub fn truncate(&mut self, new_length: usize) {
        while self.len() > new_length {
            self.remove(self.len() - 1);
        }
    }

    pub fn swap(&mut self, a: usize, b: usize) {
        self.children.swap(a, b);
    }

    pub fn insert(&mut self, index: usize, child: ArrayDimension) {
        let id = self.children.insert(index, child);
        let layout = match child {
            ArrayDimension::FitContent => {
                self.measured.push(id);
                UPx(0)
            }
            ArrayDimension::Fractional { weight } => {
                self.total_weights += u32::from(weight);
                self.fractional.push((id, weight));
                UPx(0)
            }
            ArrayDimension::Exact(size) => {
                self.allocated_space += size;
                size
            }
        };
        self.layouts.insert(
            index,
            ArrayLayout {
                offset: UPx(0),
                size: layout,
            },
        );
    }

    pub fn update(
        &mut self,
        available: Size<ConstraintLimit>,
        mut measure: impl FnMut(usize, Size<ConstraintLimit>) -> Size<UPx>,
    ) -> Size<UPx> {
        let (space_constraint, other_constraint) = self.orientation.split_size(available);
        let available_space = space_constraint.max();
        let mut remaining = available_space.saturating_sub(self.allocated_space);

        // Measure the children that fit their content
        for &id in &self.measured {
            let index = self.children.index_of_id(id).expect("child not found");
            if remaining > 0 {
                let (measured, _) = self.orientation.split_size(measure(
                    index,
                    self.orientation
                        .make_size(ConstraintLimit::ClippedAfter(remaining), other_constraint),
                ));
                self.layouts[index].size = measured;
                remaining = remaining.saturating_sub(measured);
            } else {
                self.layouts[index].size = UPx(0);
            }
        }

        // Measure the weighted children within the remaining space
        if self.total_weights > 0 {
            let space_per_weight = remaining / self.total_weights;
            remaining %= self.total_weights;
            for (fractional_index, &(id, weight)) in self.fractional.iter().enumerate() {
                let index = self.children.index_of_id(id).expect("child not found");
                let size = space_per_weight * u32::from(weight);
                self.layouts[index].size = size;

                // If we have fractional amounts remaining, divide the pixels
                if remaining > 0 {
                    let from_end = u32::try_from(self.fractional.len() - fractional_index)
                        .expect("too many items");
                    if remaining >= from_end {
                        let amount = (remaining + from_end - 1) / from_end;
                        remaining -= amount;
                        self.layouts[index].size += amount;
                    }
                }
            }
        }

        // Now that we know the constrained sizes, we can measure the children
        // to get the other measurement using the constrainted measurement.
        self.other = UPx(0);
        let mut offset = UPx(0);
        for index in 0..self.children.len() {
            self.layouts[index].offset = offset;
            offset += self.layouts[index].size;
            let (_, measured) = self.orientation.split_size(measure(
                index,
                self.orientation.make_size(
                    ConstraintLimit::Known(self.layouts[index].size),
                    other_constraint,
                ),
            ));
            self.other = self.other.max(measured);
        }

        self.other = match other_constraint {
            ConstraintLimit::Known(max) => self.other.max(max),
            ConstraintLimit::ClippedAfter(clip_limit) => self.other.min(clip_limit),
        };

        self.orientation.make_size(available_space, self.other)
    }
}

impl Deref for Layout {
    type Target = [ArrayLayout];

    fn deref(&self) -> &Self::Target {
        &self.layouts
    }
}

#[cfg(test)]
mod tests {
    use std::cmp::Ordering;

    use kludgine::figures::units::UPx;
    use kludgine::figures::Size;

    use super::{ArrayDimension, ArrayDirection, Layout};
    use crate::ConstraintLimit;

    struct Child {
        size: UPx,
        dimension: ArrayDimension,
        other: UPx,
        divisible_by: Option<UPx>,
    }

    impl Child {
        pub fn new(size: impl Into<UPx>, other: impl Into<UPx>) -> Self {
            Self {
                size: size.into(),
                dimension: ArrayDimension::FitContent,
                other: other.into(),
                divisible_by: None,
            }
        }

        pub fn fixed_size(mut self, size: UPx) -> Self {
            self.dimension = ArrayDimension::Exact(size);
            self
        }

        pub fn weighted(mut self, weight: u8) -> Self {
            self.dimension = ArrayDimension::Fractional { weight };
            self
        }

        pub fn divisible_by(mut self, split_at: impl Into<UPx>) -> Self {
            self.divisible_by = Some(split_at.into());
            self
        }
    }

    fn assert_measured_children_in_orientation(
        orientation: ArrayDirection,
        children: &[Child],
        available: Size<ConstraintLimit>,
        expected: &[UPx],
        expected_size: Size<UPx>,
    ) {
        assert_eq!(children.len(), expected.len());
        let mut flex = Layout::new(orientation);
        for child in children {
            flex.push(child.dimension);
        }

        let computed_size = flex.update(available, |index, constraints| {
            let (measured_constraint, _other_constraint) = orientation.split_size(constraints);
            let child = &children[index];
            let maximum_measured = measured_constraint.max();
            let (measured, other) = match (child.size.cmp(&maximum_measured), child.divisible_by) {
                (Ordering::Greater, Some(divisible_by)) => {
                    let available_divided = maximum_measured / divisible_by;
                    let rows = ((child.size + divisible_by - 1) / divisible_by + available_divided
                        - 1)
                        / available_divided;
                    (available_divided * divisible_by, child.other * rows)
                }
                _ => (child.size, child.other),
            };
            orientation.make_size(measured, other)
        });
        assert_eq!(computed_size, expected_size);
        let mut offset = UPx(0);
        for ((index, &child), &expected) in flex.iter().enumerate().zip(expected) {
            assert_eq!(
                child.size,
                expected,
                "child {index} measured to {}, expected {}",
                child.size,
                expected // TODO Display for UPx
            );
            assert_eq!(child.offset, offset);
            offset += child.size;
        }
    }

    fn assert_measured_children(
        children: &[Child],
        main_constraint: ConstraintLimit,
        other_constraint: ConstraintLimit,
        expected: &[UPx],
        expected_measured: UPx,
        expected_other: UPx,
    ) {
        assert_measured_children_in_orientation(
            ArrayDirection::rows(),
            children,
            ArrayDirection::rows().make_size(main_constraint, other_constraint),
            expected,
            ArrayDirection::rows().make_size(expected_measured, expected_other),
        );
        assert_measured_children_in_orientation(
            ArrayDirection::columns(),
            children,
            ArrayDirection::columns().make_size(main_constraint, other_constraint),
            expected,
            ArrayDirection::columns().make_size(expected_measured, expected_other),
        );
    }

    #[test]
    fn size_to_fit() {
        assert_measured_children(
            &[Child::new(3, 1), Child::new(3, 1), Child::new(3, 1)],
            ConstraintLimit::ClippedAfter(UPx(10)),
            ConstraintLimit::ClippedAfter(UPx(10)),
            &[UPx(3), UPx(3), UPx(3)],
            UPx(10),
            UPx(1),
        );
    }

    #[test]
    fn wrapping() {
        // This tests some fun rounding edge cases. Because the total weights is
        // 4 and the size is 10, we have inexact math to determine the pixel
        // width of each child.
        //
        // In this particular example, it shows the weights are clamped so that
        // each is credited for 2px. This is why the first child ends up with
        // 4px. However, with 4 total weight, that leaves a remaining 2px to be
        // assigned. The flex algorithm divides the remaining pixels amongst the
        // remaining children.
        assert_measured_children(
            &[
                Child::new(20, 1).divisible_by(3).weighted(2),
                Child::new(3, 1).weighted(1),
                Child::new(3, 1).weighted(1),
            ],
            ConstraintLimit::Known(UPx(10)),
            ConstraintLimit::ClippedAfter(UPx(10)),
            &[UPx(4), UPx(3), UPx(3)],
            UPx(10),
            UPx(7), // 20 / 3 = 6.666, rounded up is 7
        );
        // Same as above, but with an 11px box. This creates a leftover of 3 px
        // (11 % 4), adding 1px to all three children.
        assert_measured_children(
            &[
                Child::new(20, 1).divisible_by(3).weighted(2),
                Child::new(3, 1).weighted(1),
                Child::new(3, 1).weighted(1),
            ],
            ConstraintLimit::Known(UPx(11)),
            ConstraintLimit::ClippedAfter(UPx(11)),
            &[UPx(5), UPx(3), UPx(3)],
            UPx(11),
            UPx(7), // 20 / 3 = 6.666, rounded up is 7
        );
        // 12px box. This creates no leftover.
        assert_measured_children(
            &[
                Child::new(20, 1).divisible_by(3).weighted(2),
                Child::new(3, 1).weighted(1),
                Child::new(3, 1).weighted(1),
            ],
            ConstraintLimit::Known(UPx(12)),
            ConstraintLimit::ClippedAfter(UPx(12)),
            &[UPx(6), UPx(3), UPx(3)],
            UPx(12),
            UPx(4), // 20 / 6 = 3.666, rounded up is 4
        );
        // 13px box. This creates a leftover of 1 px (13 % 4), adding 1px only
        // to the final child
        assert_measured_children(
            &[
                Child::new(20, 1).divisible_by(3).weighted(2),
                Child::new(3, 1).weighted(1),
                Child::new(3, 1).weighted(1),
            ],
            ConstraintLimit::Known(UPx(13)),
            ConstraintLimit::ClippedAfter(UPx(13)),
            &[UPx(6), UPx(3), UPx(4)],
            UPx(13),
            UPx(4), // 20 / 6 = 3.666, rounded up is 4
        );
    }

    #[test]
    fn fixed_size() {
        assert_measured_children(
            &[
                Child::new(3, 1).fixed_size(UPx(7)),
                Child::new(3, 1).weighted(1),
                Child::new(3, 1).weighted(1),
            ],
            ConstraintLimit::Known(UPx(15)),
            ConstraintLimit::ClippedAfter(UPx(15)),
            &[UPx(7), UPx(4), UPx(4)],
            UPx(15),
            UPx(1),
        );
    }
}