tree.rs

use std::cmp;
use cgmath::Vector2;

/// Tree that subdivides a 2D area.
#[derive(Debug, Clone)]
pub struct Tree {
    // Dimensions of the root.
    root_dimensions: Vector2<u32>,

    // Binary tree of the space partition within the texture. The first element represents the area
    // of the whole texture. Further elements are only relevant if the first element is `Split`.
    // See the docs of `Elem` for the rest.
    tree: Vec<Elem>,
}

#[derive(Debug, Clone, PartialEq, Eq)]
enum Elem {
    // Area is not accessible in the tree. It just a reserved space without any actual meaning.
    Inaccessible,
    // The area is free to be used.
    Empty,
    // The area contains one texture.
    One,
    // The space is partitionned in two at the given pixel.
    // The partitionning alternates between horizontal and vertical.
    // If this element is at offset `n` in the tree, the two children are at
    // offsets `n*2+1` and `n*2+2`.
    Split { pixel: u32 },
}

impl Tree {
    /// Builds a new tree.
    pub fn new(root_dimensions: Vector2<u32>) -> Tree {
        assert!(root_dimensions.x != 0);
        assert!(root_dimensions.y != 0);

        Tree {
            root_dimensions: root_dimensions,
            tree: vec![Elem::Empty],
        }
    }

    /// Tries to insert an element of the given dimensions in the tree so that it doesn't overlap
    /// on any other element.
    ///
    /// Returns `None` if the tree has no space for it. Otherwise returns the offset where to put
    /// the element.
    pub fn insert(&mut self, dims: Vector2<u32>) -> Option<Vector2<u32>> {
        assert!(dims.x != 0);
        assert!(dims.y != 0);

        let possible_slots_iter = self.tree
            .iter()
            .enumerate()
            .filter_map(|(n, p)| if *p == Elem::Empty { Some(n) } else { None })
            .collect::<Vec<_>>()
            .into_iter();

        for possible_slot in possible_slots_iter {
            debug_assert_eq!(self.tree[possible_slot], Elem::Empty);

            let (slot_offset, slot_size) = self.infos(possible_slot);

            if slot_size.x < dims.x || slot_size.y < dims.y {
                continue;
            }

            if slot_size.x == dims.x && slot_size.y == dims.y {
                self.tree[possible_slot] = Elem::One;
                return Some(slot_offset);
            }

            let necessary_tree_size = (possible_slot * 2 + 1) * 2 + 3;
            if self.tree.len() < necessary_tree_size {
                self.tree.resize(necessary_tree_size, Elem::Inaccessible);
            }

            debug_assert_eq!(self.tree[possible_slot * 2 + 1], Elem::Inaccessible);
            debug_assert_eq!(self.tree[possible_slot * 2 + 2], Elem::Inaccessible);
            debug_assert_eq!(self.tree[(possible_slot * 2 + 1) * 2 + 1],
                             Elem::Inaccessible);
            debug_assert_eq!(self.tree[(possible_slot * 2 + 1) * 2 + 2],
                             Elem::Inaccessible);

            if ((possible_slot + 2).next_power_of_two().trailing_zeros() % 2) != 0 {
                self.tree[possible_slot] = Elem::Split { pixel: dims.x };
                if slot_size.y == dims.y {
                    self.tree[possible_slot * 2 + 1] = Elem::One;
                    self.tree[possible_slot * 2 + 2] = Elem::Empty;
                } else {
                    self.tree[possible_slot * 2 + 1] = Elem::Split { pixel: dims.y };
                    self.tree[possible_slot * 2 + 2] = Elem::Empty;
                    self.tree[(possible_slot * 2 + 1) * 2 + 1] = Elem::One;
                    self.tree[(possible_slot * 2 + 1) * 2 + 2] = Elem::Empty;
                }
            } else {
                self.tree[possible_slot] = Elem::Split { pixel: dims.y };
                if slot_size.x == dims.x {
                    self.tree[possible_slot * 2 + 1] = Elem::One;
                    self.tree[possible_slot * 2 + 2] = Elem::Empty;
                } else {
                    self.tree[possible_slot * 2 + 1] = Elem::Split { pixel: dims.x };
                    self.tree[possible_slot * 2 + 2] = Elem::Empty;
                    self.tree[(possible_slot * 2 + 1) * 2 + 1] = Elem::One;
                    self.tree[(possible_slot * 2 + 1) * 2 + 2] = Elem::Empty;
                }
            }

            return Some(slot_offset);
        }

        None
    }

    // Returns the offset and size of an element.
    fn infos(&self, id: usize) -> (Vector2<u32>, Vector2<u32>) {
        let mut offset_x = 0;
        let mut offset_y = 0;
        let mut size_x = self.root_dimensions.x;
        let mut size_y = self.root_dimensions.y;

        let mut slot_id = id;

        debug_assert_ne!(self.tree[slot_id], Elem::Inaccessible);

        while slot_id >= 1 {
            let parent_id = (slot_id - 1) / 2;

            let split = match self.tree[parent_id] {
                Elem::Split { pixel } => pixel,
                _ => panic!("binary tree is in a wrong state slot"),
            };

            // True if the value of `split` retreived above is a split on the x axis. False if the
            // y axis.
            let split_is_x = ((slot_id + 2).next_power_of_two().trailing_zeros() % 2) == 0;

            if (slot_id % 2) == 0 {
                if split_is_x {
                    offset_x += split;
                    if split <= size_x {
                        size_x = cmp::min(size_x, size_x - split);
                    }
                } else {
                    offset_y += split;
                    if split <= size_y {
                        size_y = cmp::min(size_y, size_y - split);
                    }
                }
            } else {
                if split_is_x {
                    size_x = cmp::min(size_x, split);
                } else {
                    size_y = cmp::min(size_y, split);
                }
            }

            slot_id = parent_id;
        }

        (Vector2::new(offset_x, offset_y), Vector2::new(size_x, size_y))
    }
}

#[cfg(test)]
mod tests {
    use cgmath::Vector2;
    use batch::tree::Tree;

    #[test]
    fn basic() {
        let mut tree = Tree::new(Vector2::new(1024, 1024));

        assert_eq!(tree.insert(Vector2::new(100, 50)), Some(Vector2::new(0, 0)));
        assert_eq!(tree.insert(Vector2::new(100, 60)),
                   Some(Vector2::new(100, 0)));
        assert_eq!(tree.insert(Vector2::new(30, 40)), Some(Vector2::new(0, 50)));
        assert_eq!(tree.insert(Vector2::new(30, 40)),
                   Some(Vector2::new(100, 60)));
        assert_eq!(tree.insert(Vector2::new(1000, 1000)), None);
    }
}