iris/core/src/primitive/layer.rs

use std::ops::{Index, IndexMut};

use crate::{
    render::{MaskIdx, Primitive, PrimitiveHandle, PrimitiveInst, Primitives},
    util::to_mut,
};

pub type LayerId = usize;

struct LayerNode<T> {
    next: Ptr,
    prev: Ptr,
    child: Option<Child>,
    depth: usize,
    data: T,
}

#[derive(Clone, Copy, Debug)]
enum Ptr {
    /// continue on same level
    Next(usize),
    /// go back to parent
    Parent(usize),
    /// end
    None,
}

/// TODO: currently this does not ever free layers
/// is that realistically desired?
pub struct Layers<T> {
    vec: Vec<LayerNode<T>>,
    /// index of last layer at top level (start at first = 0)
    last: usize,
}

#[derive(Clone, Copy)]
struct Child {
    head: usize,
    tail: usize,
}

pub type PrimitiveLayers = Layers<Primitives>;

impl<T: Default> Layers<T> {
    pub fn new() -> Layers<T> {
        Self {
            vec: vec![LayerNode::head()],
            last: 0,
        }
    }

    pub fn clear(&mut self) {
        self.vec.clear();
        self.vec.push(LayerNode::head());
    }

    fn push(&mut self, node: LayerNode<T>) -> LayerId {
        let i = self.vec.len();
        self.vec.push(node);
        i
    }

    pub fn next(&mut self, i: LayerId) -> LayerId {
        if let Ptr::Next(i) = self.vec[i].next {
            return i;
        }
        let i_new = self.push(LayerNode::new(
            T::default(),
            self.vec[i].next,
            Ptr::Next(i),
            self.vec[i].depth,
        ));
        self.vec[i].next = Ptr::Next(i_new);
        self.vec[i_new].prev = Ptr::Next(i);
        match self.vec[i_new].next {
            Ptr::Next(i) => self.vec[i].prev = Ptr::Next(i_new),
            Ptr::Parent(i) => self.vec[i].child.as_mut().unwrap().tail = i_new,
            Ptr::None => self.last = i_new,
        }
        i_new
    }

    pub fn child(&mut self, i: LayerId) -> LayerId {
        if let Some(c) = self.vec[i].child {
            return c.head;
        }
        let i_child = self.push(LayerNode::new(
            T::default(),
            Ptr::Parent(i),
            Ptr::Parent(i),
            self.vec[i].depth + 1,
        ));
        self.vec[i].child = Some(Child {
            head: i_child,
            tail: i_child,
        });
        i_child
    }

    pub fn iter_mut(&mut self) -> LayerIteratorMut<'_, T> {
        LayerIteratorMut::new(&mut self.vec, self.last)
    }

    pub fn iter_orderless_mut(&mut self) -> impl Iterator<Item = (usize, &mut T)> {
        self.vec.iter_mut().map(|n| &mut n.data).enumerate()
    }

    pub fn iter(&self) -> impl Iterator<Item = (LayerId, &T)> {
        self.indices().map(|i| (i, &self.vec[i].data))
    }

    pub fn iter_depth(&self) -> impl Iterator<Item = ((LayerId, usize), &T)> {
        self.indices()
            .map(|i| ((i, self.vec[i].depth), &self.vec[i].data))
    }

    pub fn indices(&self) -> LayerIndexIterator<'_, T> {
        LayerIndexIterator::new(&self.vec, self.last)
    }
}

impl PrimitiveLayers {
    pub fn write<P: Primitive>(
        &mut self,
        layer: LayerId,
        info: PrimitiveInst<P>,
    ) -> PrimitiveHandle {
        self[layer].write(layer, info)
    }

    pub fn free(&mut self, h: &PrimitiveHandle) -> MaskIdx {
        self[h.layer].free(h)
    }
}

impl<T: Default> Default for Layers<T> {
    fn default() -> Self {
        Self::new()
    }
}

impl<T> Index<LayerId> for Layers<T> {
    type Output = T;

    fn index(&self, index: LayerId) -> &Self::Output {
        &self.vec[index].data
    }
}

impl<T> IndexMut<LayerId> for Layers<T> {
    fn index_mut(&mut self, index: LayerId) -> &mut Self::Output {
        &mut self.vec[index].data
    }
}

impl<T: Default> LayerNode<T> {
    pub fn new(data: T, next: Ptr, prev: Ptr, depth: usize) -> Self {
        Self {
            next,
            prev,
            child: None,
            data,
            depth,
        }
    }

    pub fn head() -> Self {
        Self::new(T::default(), Ptr::None, Ptr::None, 0)
    }
}

pub struct LayerIteratorMut<'a, T> {
    inner: LayerIndexIterator<'a, T>,
}

impl<'a, T> Iterator for LayerIteratorMut<'a, T> {
    type Item = (usize, &'a mut T);

    fn next(&mut self) -> Option<Self::Item> {
        let i = self.inner.next()?;
        // SAFETY: requires index iterator to work properly
        let layer = unsafe { to_mut(&self.inner.vec[i].data) };
        Some((i, layer))
    }
}

impl<'a, T> DoubleEndedIterator for LayerIteratorMut<'a, T> {
    fn next_back(&mut self) -> Option<Self::Item> {
        let i = self.inner.next_back()?;
        // SAFETY: requires index iterator to work properly
        let layer = unsafe { to_mut(&self.inner.vec[i].data) };
        Some((i, layer))
    }
}

impl<'a, T> LayerIteratorMut<'a, T> {
    fn new(vec: &'a mut Vec<LayerNode<T>>, last: usize) -> Self {
        Self {
            inner: LayerIndexIterator::new(vec, last),
        }
    }
}

pub struct LayerIndexIterator<'a, T> {
    next: Option<usize>,
    next_back: Option<usize>,
    vec: &'a Vec<LayerNode<T>>,
}

impl<'a, T> Iterator for LayerIndexIterator<'a, T> {
    type Item = usize;

    fn next(&mut self) -> Option<Self::Item> {
        let ret_i = self.next?;
        let node = &self.vec[ret_i];
        self.next = if let Some(c) = node.child {
            Some(c.head)
        } else if let Ptr::Next(i) = node.next {
            Some(i)
        } else if let Ptr::Parent(i) = node.next {
            let mut node = &self.vec[i];
            while let Ptr::Parent(i) = node.next {
                node = &self.vec[i];
            }
            if let Ptr::Next(i) = node.next {
                Some(i)
            } else {
                None
            }
        } else {
            None
        };
        if self.next_back.unwrap() == ret_i {
            self.next = None;
            self.next_back = None;
        }
        Some(ret_i)
    }
}

impl<'a, T> DoubleEndedIterator for LayerIndexIterator<'a, T> {
    fn next_back(&mut self) -> Option<Self::Item> {
        let ret_i = self.next_back?;
        let node = &self.vec[ret_i];
        self.next_back = if let Ptr::Next(mut i) = node.prev {
            while let Some(c) = self.vec[i].child {
                i = c.tail
            }
            Some(i)
        } else if let Ptr::Parent(i) = node.prev {
            Some(i)
        } else {
            None
        };
        if self.next.unwrap() == ret_i {
            self.next = None;
            self.next_back = None;
        }
        Some(ret_i)
    }
}

impl<'a, T> LayerIndexIterator<'a, T> {
    fn new(vec: &'a Vec<LayerNode<T>>, last: usize) -> Self {
        let mut last = last;
        while let Some(c) = vec[last].child {
            last = c.tail;
        }
        Self {
            next: Some(0),
            next_back: Some(last),
            vec,
        }
    }
}