<commit_before>//! Utilities for generalized iterators over trees of arbitrary types.

use std;

#[repr(u8)]

/// Descriptions of actions that may be taken on a given node when traversing
/// a pattern tree.
pub mod action {
    enum Flag {
        /// Terminate iteration after performing other specified actions.
        TERMINATE = 1,
    }

    pub use self::Flag::*;

    struct Action<U,V> {
        yield_value: Option<U>,
        recurse_on: Option<V>,
        flags: u8
    }
}

pub use self::action::Action;

/// Interface for patterns whose elements may be visited in
/// a runtime-specified manner.
///
/// @tparam Tag Marker type used to discriminate the "type" of walk being
///     implemented.  A particular type may be walked for different purposes,
///     such as to visit *all* elements, or only certain ones.
///
/// @tparam Yielded Type of item the walkable should implement an iterator for.
trait Walkable<Tag,Yielded>: std::marker::PhantomFn<Tag> {
    /// Iterator type that for the implementer's underlying container.
    type BaseIterator;
    type Item = <Self::BaseIterator as Iterator>::Item;
    type Action = Action<Yielded,Self::Item>;

    /// Fetch the initial iterator to use for a walk.
    fn iter(&self) -> Self::BaseIterator;

    /// Determine the action to take for a given node in the implementor's
    /// element tree.
    ///
    /// @param element Pattern element to process.
    ///
    /// @return Tuple containing the optionally-yielded element, and a set of
    ///     flags indicating how to proceed.  Action::RECURSE is only valid
    ///     when `Item` implements Walkable<Yielded>.
    fn action<'a>(element: &'a Self::Item) -> Self::Action
        where Self::Item: Walkable<Tag,Yielded>;
}

// FIXME: should the order of `T` and `Yielded` be reversed?
/// Iterator used to perform a particular walk on a particular type.
///
/// @tparam Tag Walk-type tag.
///
/// @tparam Yielded `Item` type for the iterator implementation of
///     the structure.
///
/// @tparam T Top-level Walkable type for which this structure is an iterator.
struct Walker<Tag,Yielded,T> where T: Walkable<Tag,Yielded> {
    base_iterator: Option<<T as Walkable<Tag,Yielded>>::BaseIterator>,
    sub_iterator: Option<Walker<Tag,Yielded,Yielded>>,
}

impl<Tag,Yielded,T> Walker<Tag,Yielded,T>
where T: Walkable<Tag,Yielded> {
    /// Create a new walker iterator from the given base iterator.
    pub fn new(base: <T as Walkable<Tag,Yielded>>::BaseIterator) -> Walker<Tag,Yielded,T> {
        Walker{base_iterator: base, sub_iterator: None}
    }
}

impl<'a,Tag,Yielded,T> Iterator for Walker<Tag,Yielded,T>
where T: Walkable<Tag,Yielded>, Yielded: Walkable<Tag,Yielded> {
    type Item = &'a Yielded;

    fn next(&mut self) -> Option<<Self as Iterator>::Item> {
        // "can't use type parameter from outer function" Why not?  And it's not a function.
        //type Walk = Walkable<Tag,Yielded>;
        // "huh?  why's there a less-than sign in here?"  Because I put it there.
        //use self::Walkable<Tag,Yielded> as Walk;

        if let Some(ref iter) = self.sub_iterator {
            if let Some(yielded) = (*iter).next() {
                return yielded; }
            else {
                // Sub-iterator is empty; clear it and fall through to the
                // base iterator.
                self.sub_iterator = None;
            } }
        if let Some(ref iter) = self.base_iterator {
            if let Some(ref element) = (*iter).next() {
                let action = <T as Walkable<Tag,Yielded>>::action(element);

                // FIXME: need to provide a way to specify order in which
                // to recurse vs. yield.  The solution should allow for
                // selection of different traversal methods -- need to come
                // up with an API for that.
                if let Some(recurse_item) = action.recurse_on {
                    self.sub_iterator = Some(Walker::new(<Yielded as Walkable<Tag,Yielded>>::iter(recurse_item)));

                    // Recurse immediately if we have no value to yield.
                    if ! action.yield_value.is_some() {
                        self.next()
                    }
                }

                // For now we'll just do depth-first, since it's simplest.
                if let Some(yielded) = action.yield_value {
                    return yielded;
                }

                // Clear the base iterator if we were asked to terminate
                // the walk.  Note that any recursive walks on
                // `self.sub_iterator` will continue until that iterator
                // is exhausted.
                if action.flags & action::TERMINATE {
                    self.base_iterator = None; }
            } else {
                // Base iterator is exhausted; clear it.
                self.base_iterator = None;
            }
        }
        return None;
    }

    fn size_hint(&self) -> (usize, Option<usize>) {
        (1, None)
    }
}
<commit_msg>pattern/walk: implemented `size_hint()` for tree-walk iterator `Walker`<commit_after>//! Utilities for generalized iterators over trees of arbitrary types.

use std;

#[repr(u8)]

/// Descriptions of actions that may be taken on a given node when traversing
/// a pattern tree.
pub mod action {
    enum Flag {
        /// Terminate iteration after performing other specified actions.
        TERMINATE = 1,
    }

    pub use self::Flag::*;

    struct Action<U,V> {
        yield_value: Option<U>,
        recurse_on: Option<V>,
        flags: u8
    }
}

pub use self::action::Action;

/// Interface for patterns whose elements may be visited in
/// a runtime-specified manner.
///
/// @tparam Tag Marker type used to discriminate the "type" of walk being
///     implemented.  A particular type may be walked for different purposes,
///     such as to visit *all* elements, or only certain ones.
///
/// @tparam Yielded Type of item the walkable should implement an iterator for.
trait Walkable<Tag,Yielded>: std::marker::PhantomFn<Tag> {
    /// Iterator type that for the implementer's underlying container.
    type BaseIterator;
    type Item = <Self::BaseIterator as Iterator>::Item;
    type Action = Action<Yielded,Self::Item>;

    /// Fetch the initial iterator to use for a walk.
    fn iter(&self) -> Self::BaseIterator;

    /// Determine the action to take for a given node in the implementor's
    /// element tree.
    ///
    /// @param element Pattern element to process.
    ///
    /// @return Tuple containing the optionally-yielded element, and a set of
    ///     flags indicating how to proceed.  Action::RECURSE is only valid
    ///     when `Item` implements Walkable<Yielded>.
    fn action<'a>(element: &'a Self::Item) -> Self::Action
        where Self::Item: Walkable<Tag,Yielded>;
}

// FIXME: should the order of `T` and `Yielded` be reversed?
/// Iterator used to perform a particular walk on a particular type.
///
/// @tparam Tag Walk-type tag.
///
/// @tparam Yielded `Item` type for the iterator implementation of
///     the structure.
///
/// @tparam T Top-level Walkable type for which this structure is an iterator.
struct Walker<Tag,Yielded,T> where T: Walkable<Tag,Yielded> {
    base_iterator: Option<<T as Walkable<Tag,Yielded>>::BaseIterator>,
    sub_iterator: Option<Walker<Tag,Yielded,Yielded>>,
}

impl<Tag,Yielded,T> Walker<Tag,Yielded,T>
where T: Walkable<Tag,Yielded> {
    /// Create a new walker iterator from the given base iterator.
    pub fn new(base: <T as Walkable<Tag,Yielded>>::BaseIterator) -> Walker<Tag,Yielded,T> {
        Walker{base_iterator: base, sub_iterator: None}
    }
}

impl<'a,Tag,Yielded,T> Iterator for Walker<Tag,Yielded,T>
where T: Walkable<Tag,Yielded>, Yielded: Walkable<Tag,Yielded> {
    type Item = &'a Yielded;

    fn next(&mut self) -> Option<<Self as Iterator>::Item> {
        // "can't use type parameter from outer function" Why not?  And it's not a function.
        //type Walk = Walkable<Tag,Yielded>;
        // "huh?  why's there a less-than sign in here?"  Because I put it there.
        //use self::Walkable<Tag,Yielded> as Walk;

        if let Some(ref iter) = self.sub_iterator {
            if let Some(yielded) = (*iter).next() {
                return yielded; }
            else {
                // Sub-iterator is empty; clear it and fall through to the
                // base iterator.
                self.sub_iterator = None;
            } }
        if let Some(ref iter) = self.base_iterator {
            if let Some(ref element) = (*iter).next() {
                let action = <T as Walkable<Tag,Yielded>>::action(element);

                // FIXME: need to provide a way to specify order in which
                // to recurse vs. yield.  The solution should allow for
                // selection of different traversal methods -- need to come
                // up with an API for that.
                if let Some(recurse_item) = action.recurse_on {
                    self.sub_iterator = Some(Walker::new(<Yielded as Walkable<Tag,Yielded>>::iter(recurse_item)));

                    // Recurse immediately if we have no value to yield.
                    if ! action.yield_value.is_some() {
                        self.next()
                    }
                }

                // For now we'll just do depth-first, since it's simplest.
                if let Some(yielded) = action.yield_value {
                    return yielded;
                }

                // Clear the base iterator if we were asked to terminate
                // the walk.  Note that any recursive walks on
                // `self.sub_iterator` will continue until that iterator
                // is exhausted.
                if action.flags & action::TERMINATE {
                    self.base_iterator = None; }
            } else {
                // Base iterator is exhausted; clear it.
                self.base_iterator = None;
            }
        }
        return None;
    }

    /// We implement `size_hint` as the sum of the hints for the current base-
    /// and sub-iterators, if they're not `None`.  We return an upper bound
    /// only if both iterators are `None`.
    fn size_hint(&self) -> (usize, Option<usize>) {
        let mut min = 0;
        if let Some(iter) = self.sub_iterator { min += iter.size_hint().0 }
        if let Some(iter) = self.base_iterator { min += iter.size_hint().0 }

        (min, if self.sub_iterator.is_none() && self.base_iterator.is_none() { Some(0) } else { None })
    }
}
<|endoftext|>


<commit_before>use bvh::*;
use bounding_hierarchy::BHShape;
use aabb::AABB;
use std::collections::HashSet;

// TODO Consider: Instead of getting the scene's shapes passed, let leaf nodes store an AABB
// that is updated from the outside, perhaps by passing not only the indices of the changed
// shapes, but also their new AABBs into optimize().

impl BVH {
    /// Optimizes the `BVH` by batch-reorganizing updated nodes.
    /// Based on https://github.com/jeske/SimpleScene/blob/master/SimpleScene/Util/ssBVH/ssBVH.cs
    ///
    /// Needs all the scene's shapes, plus the indices of the shapes that were updated.
    ///
    pub fn optimize<Shape: BHShape>(&mut self,
                                    refit_shape_indices: &HashSet<usize>,
                                    shapes: &[Shape]) {
        let mut refit_node_indices: HashSet<usize> = refit_shape_indices.iter()
            .map(|x| shapes[*x].bh_node_index())
            .collect();

        // As long as we have refit nodes left, take the list of refit nodes
        // with the highest depth (sweep nodes) and try to rotate them all
        while refit_node_indices.len() > 0 {
            let mut max_depth = 0;
            let mut sweep_node_indices: Vec<usize> = Vec::new();

            // Find max_depth and sweep_node_indices in one iteration
            for refit_node_index in refit_node_indices.iter() {
                let depth = match self.nodes[*refit_node_index] {
                    BVHNode::Node { depth, .. } => depth,
                    BVHNode::Leaf { depth, .. } => depth,
                    BVHNode::Dummy => panic!("Dummy node found during BVH optimization!"),
                };

                if depth > max_depth {
                    max_depth = depth;
                    sweep_node_indices.clear();
                    sweep_node_indices.push(*refit_node_index);
                } else if depth == max_depth {
                    sweep_node_indices.push(*refit_node_index);
                }
            }

            // Try to find a useful tree rotation with all nodes previously found
            for sweep_node_index in sweep_node_indices.iter() {
                // This node does not need to be checked again
                refit_node_indices.remove(&sweep_node_index);

                let new_refit_node_index = self.try_rotate(*sweep_node_index, shapes);

                // Instead of finding a useful tree rotation, we found another node
                // that we should check, so we add its index to the refit_node_indices.
                if let Some(index) = new_refit_node_index {
                    refit_node_indices.insert(index);
                }
            }
        }
    }

    /// Checks if there is a way to rotate a child and a grandchild (or two grandchildren) of
    /// the given node (specified by `node_index`) that would improve the `BVH`.
    /// If there is, the best rotation found is performed.
    /// Relies on the children and transitive children of the given node having correct AABBs.
    ///
    /// Returns Some(index_of_node) if a new node was found that should be used for optimization.
    ///
    fn try_rotate<Shape: BHShape>(&mut self, node_index: usize, shapes: &[Shape]) -> Option<usize> {
        let mut nodes = &mut self.nodes;

        let node_clone = nodes[node_index].clone();

        // Contains the surface area that would result from applying the currently favored rotation.
        // The rotation with the smallest SA will be applied in the end.
        // The value is calculated by child_l_aabb.surface_area() + child_r_aabb.surface_area()
        let mut best_surface_area = 0f32;

        // TODO Re-implement without mutability
        let mut parent_index: usize = 0;

        #[derive(Debug, Copy, Clone)]
        struct NodeData {
            index: usize,
            aabb: AABB,
        }

        macro_rules! get_children_node_data {
            ($node_index:expr) => ({
                let node = &nodes[$node_index];
                match *node {
                    BVHNode::Node { child_l, child_r, child_l_aabb, child_r_aabb, .. } => {
                        Some((NodeData{ index: child_l, aabb: child_l_aabb},
                            NodeData{ index: child_r, aabb: child_r_aabb}))
                    }
                    BVHNode::Leaf { .. } => {
                        None
                    }
                    BVHNode::Dummy => panic!("Dummy node found during BVH optimization!"),
                }
            });
        }

        // If this node is not a grandparent, update the AABB,
        // queue the parent for refitting, and bail out.
        // If it is a grandparent, calculate the current best_surface_area.
        let (child_l, child_r) = match node_clone {
            BVHNode::Node { parent, child_l, child_r, child_l_aabb, child_r_aabb, .. } => {
                if let BVHNode::Leaf { shape, .. } = nodes[child_l] {
                    let shape_l_index = shape;
                    if let BVHNode::Leaf { shape, .. } = nodes[child_r] {
                        let shape_r_index = shape;

                        // Update the AABBs saved for the children
                        // since at least one of them changed
                        let mut node = &mut nodes[node_index];
                        match node {
                            &mut BVHNode::Node { ref mut child_l_aabb,
                                                 ref mut child_r_aabb,
                                                 .. } => {
                                *child_l_aabb = shapes[shape_l_index].aabb();
                                *child_r_aabb = shapes[shape_r_index].aabb();
                            }
                            _ => unreachable!(),
                        }

                        return Some(parent);
                    }
                }

                // TODO There are nodes with only one leaf node attached.
                // The following code assumes this to be impossible.

                let left_children = match get_children_node_data!(child_l) {
                    Some(x) => x,
                    _ => unreachable!(),
                };
                let right_children = match get_children_node_data!(child_r) {
                    Some(x) => x,
                    _ => unreachable!(),
                };

                // Update the AABBs saved for the children
                // since at least one of them changed
                let aabb_l = left_children.0.aabb.join(&left_children.1.aabb);
                let aabb_r = right_children.0.aabb.join(&right_children.1.aabb);

                parent_index = parent;
                best_surface_area = child_l_aabb.surface_area() + child_r_aabb.surface_area();
                (NodeData {
                     index: child_l,
                     aabb: aabb_l,
                 },
                 NodeData {
                     index: child_r,
                     aabb: aabb_r,
                 })
            }
            BVHNode::Leaf { parent, .. } => {
                return Some(parent);
            }
            BVHNode::Dummy => panic!("Dummy node found during BVH optimization!"),
        };

        // Stores the Rotation that would result in the surface area best_surface_area,
        // thus being the favored rotation that will be executed after considering all rotations.
        let mut best_rotation: Option<(usize, usize)> = None;

        macro_rules! consider_rotation {
            ($a:expr, $b:expr, $sa:expr) => {
                if $sa < best_surface_area {
                    best_surface_area = $sa;
                    best_rotation = Some(($a.index, $b.index));
                }
            };
        }

        // Get indices and AABBs of all grandchildren
        let left_children_nodes = {
            let node_data: NodeData = child_l;
            let index = node_data.index;
            get_children_node_data!(index)
        };
        let right_children_nodes = {
            let node_data: NodeData = child_r;
            let index = node_data.index;
            get_children_node_data!(index)
        };

        // Child to grandchild rotations
        if let Some((child_rl, child_rr)) = right_children_nodes {
            consider_rotation!(child_l,
                               child_rl,
                               child_rl.aabb.surface_area() +
                               child_l.aabb.join(&child_rr.aabb).surface_area());
            consider_rotation!(child_l,
                               child_rr,
                               child_rr.aabb.surface_area() +
                               child_l.aabb.join(&child_rl.aabb).surface_area());
        }
        if let Some((child_ll, child_lr)) = left_children_nodes {
            consider_rotation!(child_r,
                               child_ll,
                               child_ll.aabb.surface_area() +
                               child_r.aabb.join(&child_lr.aabb).surface_area());
            consider_rotation!(child_r,
                               child_lr,
                               child_lr.aabb.surface_area() +
                               child_r.aabb.join(&child_ll.aabb).surface_area());

            // Grandchild to grandchild rotations
            if let Some((child_rl, child_rr)) = right_children_nodes {
                consider_rotation!(child_ll,
                                   child_rl,
                                   child_rl.aabb.join(&child_lr.aabb).surface_area() +
                                   child_ll.aabb.join(&child_rr.aabb).surface_area());
                consider_rotation!(child_ll,
                                   child_rr,
                                   child_ll.aabb.join(&child_rl.aabb).surface_area() +
                                   child_lr.aabb.join(&child_rr.aabb).surface_area());
            }
        }

        let new_refit_node_index = if parent_index > 0 {
            Some(parent_index)
        } else {
            None
        };

        if let Some(rotation) = best_rotation {
            BVH::rotate(nodes, rotation.0, rotation.1);

            // TODO Update all changed node AABBs

            // Return parent node's index for upcoming refitting,
            // since this node just changed its AABB
            new_refit_node_index
        } else {
            // Update this node's AABBs (child_l_aabb, child_r_aabb)
            // according to the children nodes' AABBs.
            let mut node = &mut nodes[node_index];
            match node {
                &mut BVHNode::Node { ref mut child_l_aabb, ref mut child_r_aabb, .. } => {
                    *child_l_aabb = child_l.aabb;
                    *child_r_aabb = child_r.aabb;
                }
                _ => unreachable!(),
            }

            // Even with no rotation being useful for this node, a parent node's rotation
            // could be beneficial, so queue the parent sometimes.
            // TODO Return None most of the time, randomly
            // (see https://github.com/jeske/SimpleScene/blob/master/SimpleScene/Util/ssBVH/ssBVH_Node.cs#L307)
            new_refit_node_index
        }
    }

    /// Switch two nodes by rewiring the involved indices (not by moving them in the nodes slice).
    fn rotate(nodes: &mut Vec<BVHNode>, node_a_index: usize, node_b_index: usize) {
        // This function is not defined with a self parameter to make it easier to call
        // without running into borrow checker issues.
        macro_rules! should_not_happen {
            () => ( panic!("While rotating BVH nodes, something unexpected happened."); );
            ($s:expr) => ( panic!("While rotating BVH nodes, something unexpected happened: {}", $s); );
        }

        #[allow(dead_code)] // The compiler falsely detects dead code here
        fn get_parent_index(nodes: &Vec<BVHNode>, node_index: usize) -> usize {
            let node = &nodes[node_index];

            match *node {
                BVHNode::Node { parent, .. } |
                BVHNode::Leaf { parent, .. } => parent,
                _ => should_not_happen!(),
            }
        }

        // Get parent indices
        let node_a_parent_index = get_parent_index(nodes, node_a_index);
        let node_b_parent_index = get_parent_index(nodes, node_b_index);

        #[allow(dead_code)] // The compiler falsely detects dead code here
        fn get_is_left_child(nodes: &Vec<BVHNode>,
                             node_index: usize,
                             node_parent_index: usize)
                             -> bool {
            let node_parent = &nodes[node_parent_index];

            match *node_parent {
                BVHNode::Node { child_l, .. } => child_l == node_index,
                _ => should_not_happen!(),
            }
        }

        // Get info about the nodes being a left or right child
        let node_a_is_left_child = get_is_left_child(nodes, node_a_index, node_a_parent_index);
        let node_b_is_left_child = get_is_left_child(nodes, node_b_index, node_b_parent_index);

        #[allow(dead_code)] // The compiler falsely detects dead code here
        fn connect_nodes(nodes: &mut Vec<BVHNode>,
                         child_index: usize,
                         parent_index: usize,
                         left_child: bool) {
            // Set parent's child and get its depth
            let parent_depth = {
                let mut parent = &mut nodes[parent_index];
                match parent {
                    &mut BVHNode::Node { ref mut child_l, ref mut child_r, depth, .. } => {
                        if left_child {
                            *child_l = child_index;
                        } else {
                            *child_r = child_index;
                        }
                        depth
                    }
                    _ => should_not_happen!(),
                }
            };

            // Set child's parent and depth
            {
                let mut child = &mut nodes[child_index];
                match child {
                    &mut BVHNode::Node { ref mut parent, ref mut depth, .. } |
                    &mut BVHNode::Leaf { ref mut parent, ref mut depth, .. } => {
                        *parent = parent_index;
                        *depth = parent_depth + 1;
                    }
                    _ => should_not_happen!(),
                };
            }
        }

        // Perform the switch
        connect_nodes(nodes,
                      node_a_index,
                      node_b_parent_index,
                      node_b_is_left_child);
        connect_nodes(nodes,
                      node_b_index,
                      node_a_parent_index,
                      node_a_is_left_child);
    }
}

#[cfg(test)]
pub mod tests {
    use bvh::BVH;
    use std::collections::HashSet;
    use testbase::build_some_bh;

    #[test]
    /// Tests if the optimize function tries to change a fresh BVH even though it shouldn't
    fn test_optimizing_new_bvh() {
        let (mut shapes, mut bvh) = build_some_bh::<BVH>();

        let refit_shape_indices: HashSet<usize> = (0..shapes.len()).collect();
        bvh.optimize(&refit_shape_indices, &mut shapes);
    }

    // TODO Add tests for:
    // * correct parent
    // * correct depth
    // * correct BVH after optimizing
    // * correct parent and depth after optimizing
}
<commit_msg>Fix Todo<commit_after>use bvh::*;
use bounding_hierarchy::BHShape;
use aabb::AABB;
use std::collections::HashSet;

// TODO Consider: Instead of getting the scene's shapes passed, let leaf nodes store an AABB
// that is updated from the outside, perhaps by passing not only the indices of the changed
// shapes, but also their new AABBs into optimize().

impl BVH {
    /// Optimizes the `BVH` by batch-reorganizing updated nodes.
    /// Based on https://github.com/jeske/SimpleScene/blob/master/SimpleScene/Util/ssBVH/ssBVH.cs
    ///
    /// Needs all the scene's shapes, plus the indices of the shapes that were updated.
    ///
    pub fn optimize<Shape: BHShape>(&mut self,
                                    refit_shape_indices: &HashSet<usize>,
                                    shapes: &[Shape]) {
        let mut refit_node_indices: HashSet<usize> = refit_shape_indices.iter()
            .map(|x| shapes[*x].bh_node_index())
            .collect();

        // As long as we have refit nodes left, take the list of refit nodes
        // with the highest depth (sweep nodes) and try to rotate them all
        while refit_node_indices.len() > 0 {
            let mut max_depth = 0;
            let mut sweep_node_indices: Vec<usize> = Vec::new();

            // Find max_depth and sweep_node_indices in one iteration
            for refit_node_index in refit_node_indices.iter() {
                let depth = match self.nodes[*refit_node_index] {
                    BVHNode::Node { depth, .. } => depth,
                    BVHNode::Leaf { depth, .. } => depth,
                    BVHNode::Dummy => panic!("Dummy node found during BVH optimization!"),
                };

                if depth > max_depth {
                    max_depth = depth;
                    sweep_node_indices.clear();
                    sweep_node_indices.push(*refit_node_index);
                } else if depth == max_depth {
                    sweep_node_indices.push(*refit_node_index);
                }
            }

            // Try to find a useful tree rotation with all nodes previously found
            for sweep_node_index in sweep_node_indices.iter() {
                // This node does not need to be checked again
                refit_node_indices.remove(&sweep_node_index);

                let new_refit_node_index = self.try_rotate(*sweep_node_index, shapes);

                // Instead of finding a useful tree rotation, we found another node
                // that we should check, so we add its index to the refit_node_indices.
                if let Some(index) = new_refit_node_index {
                    refit_node_indices.insert(index);
                }
            }
        }
    }

    /// Checks if there is a way to rotate a child and a grandchild (or two grandchildren) of
    /// the given node (specified by `node_index`) that would improve the `BVH`.
    /// If there is, the best rotation found is performed.
    /// Relies on the children and transitive children of the given node having correct AABBs.
    ///
    /// Returns Some(index_of_node) if a new node was found that should be used for optimization.
    ///
    fn try_rotate<Shape: BHShape>(&mut self, node_index: usize, shapes: &[Shape]) -> Option<usize> {
        let mut nodes = &mut self.nodes;

        let node_clone = nodes[node_index].clone();

        // Contains the surface area that would result from applying the currently favored rotation.
        // The rotation with the smallest SA will be applied in the end.
        // The value is calculated by child_l_aabb.surface_area() + child_r_aabb.surface_area()
        let mut best_surface_area = 0f32;

        // TODO Re-implement without mutability
        let mut parent_index: usize = 0;

        #[derive(Debug, Copy, Clone)]
        struct NodeData {
            index: usize,
            aabb: AABB,
        }

        macro_rules! get_children_node_data {
            ($node_index:expr) => ({
                let node = &nodes[$node_index];
                match *node {
                    BVHNode::Node { child_l, child_r, child_l_aabb, child_r_aabb, .. } => {
                        Some((NodeData{ index: child_l, aabb: child_l_aabb},
                            NodeData{ index: child_r, aabb: child_r_aabb}))
                    }
                    BVHNode::Leaf { .. } => {
                        None
                    }
                    BVHNode::Dummy => panic!("Dummy node found during BVH optimization!"),
                }
            });
        }

        macro_rules! get_child_aabb {
            ($node_index:expr) => ({
                let node = &nodes[$node_index];
                match *node {
                    BVHNode::Node { child_l_aabb, child_r_aabb, .. } => {
                        child_l_aabb.join(&child_r_aabb)
                    }
                    BVHNode::Leaf { shape, .. } => {
                        shapes[shape].aabb()
                    }
                    BVHNode::Dummy => panic!("Dummy node found during BVH optimization!"),
                }
            });
        }

        // If this node is not a grandparent, update the AABB,
        // queue the parent for refitting, and bail out.
        // If it is a grandparent, calculate the current best_surface_area.
        let (child_l, child_r) = match node_clone {
            BVHNode::Node { parent, child_l, child_r, child_l_aabb, child_r_aabb, .. } => {
                if let BVHNode::Leaf { shape, .. } = nodes[child_l] {
                    let shape_l_index = shape;
                    if let BVHNode::Leaf { shape, .. } = nodes[child_r] {
                        let shape_r_index = shape;

                        // Update the AABBs saved for the children
                        // since at least one of them changed
                        let mut node = &mut nodes[node_index];
                        match node {
                            &mut BVHNode::Node { ref mut child_l_aabb,
                                                 ref mut child_r_aabb,
                                                 .. } => {
                                *child_l_aabb = shapes[shape_l_index].aabb();
                                *child_r_aabb = shapes[shape_r_index].aabb();
                            }
                            _ => unreachable!(),
                        }

                        return Some(parent);
                    }
                }

                // Update the AABBs saved for the children
                // since at least one of them changed
                let aabb_l = get_child_aabb!(child_l);
                let aabb_r = get_child_aabb!(child_r);

                parent_index = parent;
                best_surface_area = child_l_aabb.surface_area() + child_r_aabb.surface_area();
                (NodeData {
                     index: child_l,
                     aabb: aabb_l,
                 },
                 NodeData {
                     index: child_r,
                     aabb: aabb_r,
                 })
            }
            BVHNode::Leaf { parent, .. } => {
                return Some(parent);
            }
            BVHNode::Dummy => panic!("Dummy node found during BVH optimization!"),
        };

        // Stores the Rotation that would result in the surface area best_surface_area,
        // thus being the favored rotation that will be executed after considering all rotations.
        let mut best_rotation: Option<(usize, usize)> = None;

        macro_rules! consider_rotation {
            ($a:expr, $b:expr, $sa:expr) => {
                if $sa < best_surface_area {
                    best_surface_area = $sa;
                    best_rotation = Some(($a.index, $b.index));
                }
            };
        }

        // Get indices and AABBs of all grandchildren
        let left_children_nodes = {
            let node_data: NodeData = child_l;
            let index = node_data.index;
            get_children_node_data!(index)
        };
        let right_children_nodes = {
            let node_data: NodeData = child_r;
            let index = node_data.index;
            get_children_node_data!(index)
        };

        // Child to grandchild rotations
        if let Some((child_rl, child_rr)) = right_children_nodes {
            consider_rotation!(child_l,
                               child_rl,
                               child_rl.aabb.surface_area() +
                               child_l.aabb.join(&child_rr.aabb).surface_area());
            consider_rotation!(child_l,
                               child_rr,
                               child_rr.aabb.surface_area() +
                               child_l.aabb.join(&child_rl.aabb).surface_area());
        }
        if let Some((child_ll, child_lr)) = left_children_nodes {
            consider_rotation!(child_r,
                               child_ll,
                               child_ll.aabb.surface_area() +
                               child_r.aabb.join(&child_lr.aabb).surface_area());
            consider_rotation!(child_r,
                               child_lr,
                               child_lr.aabb.surface_area() +
                               child_r.aabb.join(&child_ll.aabb).surface_area());

            // Grandchild to grandchild rotations
            if let Some((child_rl, child_rr)) = right_children_nodes {
                consider_rotation!(child_ll,
                                   child_rl,
                                   child_rl.aabb.join(&child_lr.aabb).surface_area() +
                                   child_ll.aabb.join(&child_rr.aabb).surface_area());
                consider_rotation!(child_ll,
                                   child_rr,
                                   child_ll.aabb.join(&child_rl.aabb).surface_area() +
                                   child_lr.aabb.join(&child_rr.aabb).surface_area());
            }
        }

        let new_refit_node_index = if parent_index > 0 {
            Some(parent_index)
        } else {
            None
        };

        if let Some(rotation) = best_rotation {
            BVH::rotate(nodes, rotation.0, rotation.1);

            // TODO Update all changed node AABBs

            // Return parent node's index for upcoming refitting,
            // since this node just changed its AABB
            new_refit_node_index
        } else {
            // Update this node's AABBs (child_l_aabb, child_r_aabb)
            // according to the children nodes' AABBs.
            let mut node = &mut nodes[node_index];
            match node {
                &mut BVHNode::Node { ref mut child_l_aabb, ref mut child_r_aabb, .. } => {
                    *child_l_aabb = child_l.aabb;
                    *child_r_aabb = child_r.aabb;
                }
                _ => unreachable!(),
            }

            // Even with no rotation being useful for this node, a parent node's rotation
            // could be beneficial, so queue the parent sometimes.
            // TODO Return None most of the time, randomly
            // (see https://github.com/jeske/SimpleScene/blob/master/SimpleScene/Util/ssBVH/ssBVH_Node.cs#L307)
            new_refit_node_index
        }
    }

    /// Switch two nodes by rewiring the involved indices (not by moving them in the nodes slice).
    fn rotate(nodes: &mut Vec<BVHNode>, node_a_index: usize, node_b_index: usize) {
        // This function is not defined with a self parameter to make it easier to call
        // without running into borrow checker issues.
        macro_rules! should_not_happen {
            () => ( panic!("While rotating BVH nodes, something unexpected happened."); );
            ($s:expr) => ( panic!("While rotating BVH nodes, something unexpected happened: {}", $s); );
        }

        #[allow(dead_code)] // The compiler falsely detects dead code here
        fn get_parent_index(nodes: &Vec<BVHNode>, node_index: usize) -> usize {
            let node = &nodes[node_index];

            match *node {
                BVHNode::Node { parent, .. } |
                BVHNode::Leaf { parent, .. } => parent,
                _ => should_not_happen!(),
            }
        }

        // Get parent indices
        let node_a_parent_index = get_parent_index(nodes, node_a_index);
        let node_b_parent_index = get_parent_index(nodes, node_b_index);

        #[allow(dead_code)] // The compiler falsely detects dead code here
        fn get_is_left_child(nodes: &Vec<BVHNode>,
                             node_index: usize,
                             node_parent_index: usize)
                             -> bool {
            let node_parent = &nodes[node_parent_index];

            match *node_parent {
                BVHNode::Node { child_l, .. } => child_l == node_index,
                _ => should_not_happen!(),
            }
        }

        // Get info about the nodes being a left or right child
        let node_a_is_left_child = get_is_left_child(nodes, node_a_index, node_a_parent_index);
        let node_b_is_left_child = get_is_left_child(nodes, node_b_index, node_b_parent_index);

        #[allow(dead_code)] // The compiler falsely detects dead code here
        fn connect_nodes(nodes: &mut Vec<BVHNode>,
                         child_index: usize,
                         parent_index: usize,
                         left_child: bool) {
            // Set parent's child and get its depth
            let parent_depth = {
                let mut parent = &mut nodes[parent_index];
                match parent {
                    &mut BVHNode::Node { ref mut child_l, ref mut child_r, depth, .. } => {
                        if left_child {
                            *child_l = child_index;
                        } else {
                            *child_r = child_index;
                        }
                        depth
                    }
                    _ => should_not_happen!(),
                }
            };

            // Set child's parent and depth
            {
                let mut child = &mut nodes[child_index];
                match child {
                    &mut BVHNode::Node { ref mut parent, ref mut depth, .. } |
                    &mut BVHNode::Leaf { ref mut parent, ref mut depth, .. } => {
                        *parent = parent_index;
                        *depth = parent_depth + 1;
                    }
                    _ => should_not_happen!(),
                };
            }
        }

        // Perform the switch
        connect_nodes(nodes,
                      node_a_index,
                      node_b_parent_index,
                      node_b_is_left_child);
        connect_nodes(nodes,
                      node_b_index,
                      node_a_parent_index,
                      node_a_is_left_child);
    }
}

#[cfg(test)]
pub mod tests {
    use bvh::BVH;
    use std::collections::HashSet;
    use testbase::build_some_bh;

    #[test]
    /// Tests if the optimize function tries to change a fresh BVH even though it shouldn't
    fn test_optimizing_new_bvh() {
        let (mut shapes, mut bvh) = build_some_bh::<BVH>();

        let refit_shape_indices: HashSet<usize> = (0..shapes.len()).collect();
        bvh.optimize(&refit_shape_indices, &mut shapes);
    }

    // TODO Add tests for:
    // * correct parent
    // * correct depth
    // * correct BVH after optimizing
    // * correct parent and depth after optimizing
}
<|endoftext|>


<commit_before>//! Provides functions to read and write segment registers.

use crate::structures::gdt::SegmentSelector;

/// Reload code segment register.
///
/// Note this is special since we can not directly move
/// to %cs. Instead we push the new segment selector
/// and return value on the stack and use lretq
/// to reload cs and continue at 1:.
///
/// ## Safety
///
/// This function is unsafe because the caller must ensure that `sel`
/// is a valid code segment descriptor.
#[inline]
pub unsafe fn set_cs(sel: SegmentSelector) {
    #[cfg(feature = "inline_asm")]
    #[inline(always)]
    unsafe fn inner(sel: SegmentSelector) {
        // FIXME - Use intel syntax
        asm!("pushq {}; leaq 1f(%rip), %rax; pushq %rax; lretq; 1:", in(reg) u64::from(sel.0), out("rax") _, options(att_syntax));
    }

    #[cfg(not(feature = "inline_asm"))]
    #[inline(always)]
    unsafe fn inner(sel: SegmentSelector) {
        crate::asm::x86_64_asm_set_cs(u64::from(sel.0))
    }

    inner(sel)
}

/// Reload stack segment register.
///
/// ## Safety
///
/// This function is unsafe because the caller must ensure that `sel`
/// is a valid stack segment descriptor.
#[inline]
pub unsafe fn load_ss(sel: SegmentSelector) {
    #[cfg(feature = "inline_asm")]
    asm!("mov ss, {0:x}", in(reg) sel.0, options(nostack));

    #[cfg(not(feature = "inline_asm"))]
    crate::asm::x86_64_asm_load_ss(sel.0);
}

/// Reload data segment register.
///
/// ## Safety
///
/// This function is unsafe because the caller must ensure that `sel`
/// is a valid data segment descriptor.
#[inline]
pub unsafe fn load_ds(sel: SegmentSelector) {
    #[cfg(feature = "inline_asm")]
    asm!("mov ds, {0:x}", in(reg) sel.0, options(nostack));

    #[cfg(not(feature = "inline_asm"))]
    crate::asm::x86_64_asm_load_ds(sel.0);
}

/// Reload es segment register.
///
/// ## Safety
///
/// This function is unsafe because the caller must ensure that `sel`
/// is a valid extra segment descriptor.
#[inline]
pub unsafe fn load_es(sel: SegmentSelector) {
    #[cfg(feature = "inline_asm")]
     asm!("mov es, {0:x}", in(reg) sel.0, options(nostack));

    #[cfg(not(feature = "inline_asm"))]
    crate::asm::x86_64_asm_load_es(sel.0);
}

/// Reload fs segment register.
///
/// ## Safety
///
/// This function is unsafe because the caller must ensure that `sel`
/// is a valid fs segment descriptor.
#[inline]
pub unsafe fn load_fs(sel: SegmentSelector) {
    #[cfg(feature = "inline_asm")]
    asm!("mov fs, {0:x}", in(reg) sel.0, options(nostack));

    #[cfg(not(feature = "inline_asm"))]
    crate::asm::x86_64_asm_load_fs(sel.0);
}

/// Reload gs segment register.
///
/// ## Safety
///
/// This function is unsafe because the caller must ensure that `sel`
/// is a valid gs segment descriptor.
#[inline]
pub unsafe fn load_gs(sel: SegmentSelector) {
    #[cfg(feature = "inline_asm")]
    asm!("mov gs, {0:x}", in(reg) sel.0, options(nostack));

    #[cfg(not(feature = "inline_asm"))]
    crate::asm::x86_64_asm_load_gs(sel.0);
}

/// Swap `KernelGsBase` MSR and `GsBase` MSR.
///
/// ## Safety
///
/// This function is unsafe because the caller must ensure that the
/// swap operation cannot lead to undefined behavior.
#[inline]
pub unsafe fn swap_gs() {
    #[cfg(feature = "inline_asm")]
    asm!("swapgs", options(nostack));

    #[cfg(not(feature = "inline_asm"))]
    crate::asm::x86_64_asm_swapgs();
}

/// Returns the current value of the code segment register.
#[inline]
pub fn cs() -> SegmentSelector {
    #[cfg(feature = "inline_asm")]
    {
        let segment: u16;
        unsafe {
            asm!("mov {0:x}, cs", out(reg) segment, options(nostack, nomem))
        };
        SegmentSelector(segment)
    }

    #[cfg(not(feature = "inline_asm"))]
    {
        let segment: u16 = unsafe { crate::asm::x86_64_asm_get_cs() };
        SegmentSelector(segment)
    }
}

/// Writes the FS segment base address
///
/// ## Safety
///
/// If `CR4.FSGSBASE` is not set, this instruction will throw an `#UD`.
///
/// The caller must ensure that this write operation has no unsafe side
/// effects, as the FS segment base address is often used for thread
/// local storage.
#[inline]
pub unsafe fn wrfsbase(val: u64) {
    #[cfg(feature = "inline_asm")]
    #[inline(always)]
    unsafe fn inner(val: u64) {
        llvm_asm!("wrfsbase $0" :: "r"(val) :: "volatile")
    }

    #[cfg(not(feature = "inline_asm"))]
    #[inline(always)]
    unsafe fn inner(val: u64) {
        crate::asm::x86_64_asm_wrfsbase(val)
    }

    inner(val)
}

/// Reads the FS segment base address
///
/// ## Safety
///
/// If `CR4.FSGSBASE` is not set, this instruction will throw an `#UD`.
#[inline]
pub unsafe fn rdfsbase() -> u64 {
    #[cfg(feature = "inline_asm")]
    #[inline(always)]
    unsafe fn inner() -> u64 {
        let val: u64;
        llvm_asm!("rdfsbase $0" : "=r" (val) ::: "volatile");
        val
    }

    #[cfg(not(feature = "inline_asm"))]
    #[inline(always)]
    unsafe fn inner() -> u64 {
        crate::asm::x86_64_asm_rdfsbase()
    }

    inner()
}

/// Writes the GS segment base address
///
/// ## Safety
///
/// If `CR4.FSGSBASE` is not set, this instruction will throw an `#UD`.
///
/// The caller must ensure that this write operation has no unsafe side
/// effects, as the GS segment base address might be in use.
#[inline]
pub unsafe fn wrgsbase(val: u64) {
    #[cfg(feature = "inline_asm")]
    #[inline(always)]
    unsafe fn inner(val: u64) {
        llvm_asm!("wrgsbase $0" :: "r"(val) :: "volatile")
    }

    #[cfg(not(feature = "inline_asm"))]
    #[inline(always)]
    unsafe fn inner(val: u64) {
        crate::asm::x86_64_asm_wrgsbase(val)
    }

    inner(val)
}

/// Reads the GS segment base address
///
/// ## Safety
///
/// If `CR4.FSGSBASE` is not set, this instruction will throw an `#UD`.
#[inline]
pub unsafe fn rdgsbase() -> u64 {
    #[cfg(feature = "inline_asm")]
    #[inline(always)]
    unsafe fn inner() -> u64 {
        let val: u64;
        llvm_asm!("rdgsbase $0" : "=r" (val) ::: "volatile");
        val
    }

    #[cfg(not(feature = "inline_asm"))]
    #[inline(always)]
    unsafe fn inner() -> u64 {
        crate::asm::x86_64_asm_rdgsbase()
    }

    inner()
}
<commit_msg>Make rustfmt happy<commit_after>//! Provides functions to read and write segment registers.

use crate::structures::gdt::SegmentSelector;

/// Reload code segment register.
///
/// Note this is special since we can not directly move
/// to %cs. Instead we push the new segment selector
/// and return value on the stack and use lretq
/// to reload cs and continue at 1:.
///
/// ## Safety
///
/// This function is unsafe because the caller must ensure that `sel`
/// is a valid code segment descriptor.
#[inline]
pub unsafe fn set_cs(sel: SegmentSelector) {
    #[cfg(feature = "inline_asm")]
    #[inline(always)]
    unsafe fn inner(sel: SegmentSelector) {
        // FIXME - Use intel syntax
        asm!("pushq {}; leaq 1f(%rip), %rax; pushq %rax; lretq; 1:", in(reg) u64::from(sel.0), out("rax") _, options(att_syntax));
    }

    #[cfg(not(feature = "inline_asm"))]
    #[inline(always)]
    unsafe fn inner(sel: SegmentSelector) {
        crate::asm::x86_64_asm_set_cs(u64::from(sel.0))
    }

    inner(sel)
}

/// Reload stack segment register.
///
/// ## Safety
///
/// This function is unsafe because the caller must ensure that `sel`
/// is a valid stack segment descriptor.
#[inline]
pub unsafe fn load_ss(sel: SegmentSelector) {
    #[cfg(feature = "inline_asm")]
    asm!("mov ss, {0:x}", in(reg) sel.0, options(nostack));

    #[cfg(not(feature = "inline_asm"))]
    crate::asm::x86_64_asm_load_ss(sel.0);
}

/// Reload data segment register.
///
/// ## Safety
///
/// This function is unsafe because the caller must ensure that `sel`
/// is a valid data segment descriptor.
#[inline]
pub unsafe fn load_ds(sel: SegmentSelector) {
    #[cfg(feature = "inline_asm")]
    asm!("mov ds, {0:x}", in(reg) sel.0, options(nostack));

    #[cfg(not(feature = "inline_asm"))]
    crate::asm::x86_64_asm_load_ds(sel.0);
}

/// Reload es segment register.
///
/// ## Safety
///
/// This function is unsafe because the caller must ensure that `sel`
/// is a valid extra segment descriptor.
#[inline]
pub unsafe fn load_es(sel: SegmentSelector) {
    #[cfg(feature = "inline_asm")]
    asm!("mov es, {0:x}", in(reg) sel.0, options(nostack));

    #[cfg(not(feature = "inline_asm"))]
    crate::asm::x86_64_asm_load_es(sel.0);
}

/// Reload fs segment register.
///
/// ## Safety
///
/// This function is unsafe because the caller must ensure that `sel`
/// is a valid fs segment descriptor.
#[inline]
pub unsafe fn load_fs(sel: SegmentSelector) {
    #[cfg(feature = "inline_asm")]
    asm!("mov fs, {0:x}", in(reg) sel.0, options(nostack));

    #[cfg(not(feature = "inline_asm"))]
    crate::asm::x86_64_asm_load_fs(sel.0);
}

/// Reload gs segment register.
///
/// ## Safety
///
/// This function is unsafe because the caller must ensure that `sel`
/// is a valid gs segment descriptor.
#[inline]
pub unsafe fn load_gs(sel: SegmentSelector) {
    #[cfg(feature = "inline_asm")]
    asm!("mov gs, {0:x}", in(reg) sel.0, options(nostack));

    #[cfg(not(feature = "inline_asm"))]
    crate::asm::x86_64_asm_load_gs(sel.0);
}

/// Swap `KernelGsBase` MSR and `GsBase` MSR.
///
/// ## Safety
///
/// This function is unsafe because the caller must ensure that the
/// swap operation cannot lead to undefined behavior.
#[inline]
pub unsafe fn swap_gs() {
    #[cfg(feature = "inline_asm")]
    asm!("swapgs", options(nostack));

    #[cfg(not(feature = "inline_asm"))]
    crate::asm::x86_64_asm_swapgs();
}

/// Returns the current value of the code segment register.
#[inline]
pub fn cs() -> SegmentSelector {
    #[cfg(feature = "inline_asm")]
    {
        let segment: u16;
        unsafe { asm!("mov {0:x}, cs", out(reg) segment, options(nostack, nomem)) };
        SegmentSelector(segment)
    }

    #[cfg(not(feature = "inline_asm"))]
    {
        let segment: u16 = unsafe { crate::asm::x86_64_asm_get_cs() };
        SegmentSelector(segment)
    }
}

/// Writes the FS segment base address
///
/// ## Safety
///
/// If `CR4.FSGSBASE` is not set, this instruction will throw an `#UD`.
///
/// The caller must ensure that this write operation has no unsafe side
/// effects, as the FS segment base address is often used for thread
/// local storage.
#[inline]
pub unsafe fn wrfsbase(val: u64) {
    #[cfg(feature = "inline_asm")]
    #[inline(always)]
    unsafe fn inner(val: u64) {
        llvm_asm!("wrfsbase $0" :: "r"(val) :: "volatile")
    }

    #[cfg(not(feature = "inline_asm"))]
    #[inline(always)]
    unsafe fn inner(val: u64) {
        crate::asm::x86_64_asm_wrfsbase(val)
    }

    inner(val)
}

/// Reads the FS segment base address
///
/// ## Safety
///
/// If `CR4.FSGSBASE` is not set, this instruction will throw an `#UD`.
#[inline]
pub unsafe fn rdfsbase() -> u64 {
    #[cfg(feature = "inline_asm")]
    #[inline(always)]
    unsafe fn inner() -> u64 {
        let val: u64;
        llvm_asm!("rdfsbase $0" : "=r" (val) ::: "volatile");
        val
    }

    #[cfg(not(feature = "inline_asm"))]
    #[inline(always)]
    unsafe fn inner() -> u64 {
        crate::asm::x86_64_asm_rdfsbase()
    }

    inner()
}

/// Writes the GS segment base address
///
/// ## Safety
///
/// If `CR4.FSGSBASE` is not set, this instruction will throw an `#UD`.
///
/// The caller must ensure that this write operation has no unsafe side
/// effects, as the GS segment base address might be in use.
#[inline]
pub unsafe fn wrgsbase(val: u64) {
    #[cfg(feature = "inline_asm")]
    #[inline(always)]
    unsafe fn inner(val: u64) {
        llvm_asm!("wrgsbase $0" :: "r"(val) :: "volatile")
    }

    #[cfg(not(feature = "inline_asm"))]
    #[inline(always)]
    unsafe fn inner(val: u64) {
        crate::asm::x86_64_asm_wrgsbase(val)
    }

    inner(val)
}

/// Reads the GS segment base address
///
/// ## Safety
///
/// If `CR4.FSGSBASE` is not set, this instruction will throw an `#UD`.
#[inline]
pub unsafe fn rdgsbase() -> u64 {
    #[cfg(feature = "inline_asm")]
    #[inline(always)]
    unsafe fn inner() -> u64 {
        let val: u64;
        llvm_asm!("rdgsbase $0" : "=r" (val) ::: "volatile");
        val
    }

    #[cfg(not(feature = "inline_asm"))]
    #[inline(always)]
    unsafe fn inner() -> u64 {
        crate::asm::x86_64_asm_rdgsbase()
    }

    inner()
}
<|endoftext|>


<commit_before>#[macro_use]
extern crate clap;
extern crate capnp;

mod hierarchy_capnp;
mod vm_capnp;
mod test_capnp;

use std::fs::File;
use std::path::Path;
use std::io::BufReader;
use std::process;

use hierarchy_capnp::{directories};
use capnp::{serialize, message};

struct ExecuteTest<'a> {
    pub name: String,
    pub test: test_capnp::input_output::Reader<'a>,
}

fn main() {
    let matches = clap_app!(gaslighter =>
        (version: "0.1.0")
        (author: "Stewart Mackenzie <setori88@gmail.com>")
        (about: "Gaslighter - Ethereum Virtual Machine tester.")
        (@arg CAPNPROTO_TYPECHECKED_TEST_BIN: -t --capnp_test_bin +takes_value "Path to a type checked binary compiled by the capnp tool. The source of this artefact is in the tests directory. Please run `$ capnp eval -b tests/mod.capnp all > tests.bin` in the root directory to generate the binary.")
        (@arg TESTS_TO_RUN: -r --run_test +takes_value +required "The format is [directory]/[file]/[test] e.g. `--run_test arith/add/add1` will run the arith/add/add1 test, `--run_test arith/add/` will run every test in the tests/arith/add.capnp file. Likewise `--run_test arith//` will run every test in every file of the `arith` directory. Lastly `--run_test //` will run every single test available.")
        (@arg KEEP_GOING: -k --keep_going "Don't exit the program even if a test fails.")
    ).get_matches();
    let capnp_test_bin = match matches.value_of("CAPNPROTO_TYPECHECKED_TEST_BIN") {
        Some(c) => c,
        None => "",
    };
    let test_to_run = match matches.value_of("TESTS_TO_RUN") {
        Some(c) => c,
        None => "",
    };
    let keep_going = if matches.is_present("KEEP_GOING") { true } else { false };
    let path = Path::new(capnp_test_bin);
    let display = path.display();
    let file = match File::open(&path) {
        Err(_) => panic!("couldn't open {}", display),
        Ok(file) => file,
    };
    let (dir_to_run, file_to_run, test_to_run) = test_scope(test_to_run.into());
    let mut contents = BufReader::new(file);
    let tests_reader = serialize::read_message(&mut contents, message::ReaderOptions::new()).expect("read message failed.");
    let mut tests_to_execute :std::vec::Vec<ExecuteTest> = Vec::new();
    let top_level_tests = tests_reader.get_root::<directories::Reader>().expect("failed to get top level test root.");
    for dir in top_level_tests.get_dirs().expect("failed to directories.").iter() {
        let mut add_dir = false;
        let mut add_file = false;
        let mut add_test = false;
        let mut dirname = dir.get_name().expect("failed to get directory name.");
        let mut filename = "";
        let mut testname = "";
        if dirname == dir_to_run || dir_to_run == "" { add_dir = true; }
        for file in dir.get_files().expect("failed to files.").iter() {
            filename = file.get_name().expect("failed to get filename.");
            if filename == file_to_run || file_to_run == "" { add_file = true; }
            for test in file.get_tests().expect("failed to get tests.").iter() {
                testname = test.get_name().expect("failed to get test name.");
                if testname == test_to_run || test_to_run == "" {
                    add_test = true;
                    if add_dir && add_file && add_test {
                        let execute_test = ExecuteTest {
                            name: format!("{}/{}/{}", dirname, filename, testname),
                            test: test,
                        };
                        tests_to_execute.push(execute_test);
                    }
                }
            }
        }
    }
    if has_all_tests_passed(tests_to_execute, keep_going) {
        process::exit(0);
    } else {
        process::exit(1);
    }
}

fn test_scope(test_to_run: String) -> (String, String, String) {
    let vec: Vec<&str> = test_to_run.split("/").collect();
    (vec[0].into(), vec[1].into(), vec[2].into())
}

fn has_all_tests_passed(tests_to_execute: std::vec::Vec<ExecuteTest>, keep_going: bool) -> bool {
    println!("Executing tests:");
    let mut has_all_tests_passed = true;
    for test in tests_to_execute {
        print!("--> {:?} ", test.name);
        let test = test.test;
        let eo = test.get_expected_output().expect("failed to get expected output");
        let io = test.get_input_output().expect("failed to get actual input");
        let ao_gas = io.get_output().expect("").get_gas();
        let eo_gas = eo.get_gas();
        let ao_code = io.get_output().expect("").get_code().expect("").iter();
        let eo_code = eo.get_code().expect("").iter();
        let mut ao_vec = Vec::new();
        let mut eo_vec = Vec::new();
        let mut has_this_test_failed = false;
        for ao_char in ao_code {
            ao_vec.push(ao_char.expect("character expected")[0]);
        }
        for eo_char in eo_code {
            eo_vec.push(eo_char.expect("character expected")[0]);
        }
        let length = eo_vec.len();
        let matching = ao_vec.iter().zip(eo_vec.iter()).filter(|&(a, b)| a == b).count();
        if matching != length {
            has_this_test_failed = true;
            print!("\n\n code equality fail: only {} actual output opcodes matched the {} opcodes of the expected output.\n", matching, length);
            println!(" actual code output:\t{:?}", ao_vec);
            println!(" expected code output:\t{:?}", eo_vec);
        }
        if eo_gas != ao_gas {
            has_this_test_failed = true;
            print!("\n gas equality fail: actual output gas value of {} doesn't equal the expected output gas value of {}.", ao_gas, eo_gas);
            print!("\n actual gas output:\t{}", ao_gas);
            print!("\n expected gas output:\t{}\n", eo_gas);
        }
        if has_this_test_failed {
            print!("\n");
            has_all_tests_passed = false;
            if !keep_going {
                return has_all_tests_passed;
            }
        } else {
            print!("... ok\n")
        }
    }
    has_all_tests_passed
}
<commit_msg>Problem: gaslighter incorrectly filters tests<commit_after>#[macro_use]
extern crate clap;
extern crate capnp;

mod hierarchy_capnp;
mod vm_capnp;
mod test_capnp;

use std::fs::File;
use std::path::Path;
use std::io::BufReader;
use std::process;

use hierarchy_capnp::{directories};
use capnp::{serialize, message};

struct ExecuteTest<'a> {
    pub name: String,
    pub test: test_capnp::input_output::Reader<'a>,
}

fn main() {
    let matches = clap_app!(gaslighter =>
        (version: "0.1.0")
        (author: "Stewart Mackenzie <setori88@gmail.com>")
        (about: "Gaslighter - Ethereum Virtual Machine tester.")
        (@arg CAPNPROTO_TYPECHECKED_TEST_BIN: -t --capnp_test_bin +takes_value "Path to a type checked binary compiled by the capnp tool. The source of this artefact is in the tests directory. Please run `$ capnp eval -b tests/mod.capnp all > tests.bin` in the root directory to generate the binary.")
        (@arg TESTS_TO_RUN: -r --run_test +takes_value +required "The format is [directory]/[file]/[test] e.g. `--run_test arith/add/add1` will run the arith/add/add1 test, `--run_test arith/add/` will run every test in the tests/arith/add.capnp file. Likewise `--run_test arith//` will run every test in every file of the `arith` directory. Lastly `--run_test //` will run every single test available.")
        (@arg KEEP_GOING: -k --keep_going "Don't exit the program even if a test fails.")
    ).get_matches();
    let capnp_test_bin = match matches.value_of("CAPNPROTO_TYPECHECKED_TEST_BIN") {
        Some(c) => c,
        None => "",
    };
    let test_to_run = match matches.value_of("TESTS_TO_RUN") {
        Some(c) => c,
        None => "",
    };
    let keep_going = if matches.is_present("KEEP_GOING") { true } else { false };
    let path = Path::new(capnp_test_bin);
    let display = path.display();
    let file = match File::open(&path) {
        Err(_) => panic!("couldn't open {}", display),
        Ok(file) => file,
    };
    let (dir_to_run, file_to_run, test_to_run) = test_scope(test_to_run.into());
    let mut contents = BufReader::new(file);
    let tests_reader = serialize::read_message(&mut contents, message::ReaderOptions::new()).expect("read message failed.");
    let mut tests_to_execute :std::vec::Vec<ExecuteTest> = Vec::new();
    let top_level_tests = tests_reader.get_root::<directories::Reader>().expect("failed to get top level test root.");
    for dir in top_level_tests.get_dirs().expect("failed to directories.").iter() {
        let mut add_dir = false;
        let mut dirname = dir.get_name().expect("failed to get directory name.");
        let mut filename = "";
        let mut testname = "";
        if dirname == dir_to_run || dir_to_run == "" { add_dir = true; }
        for file in dir.get_files().expect("failed to files.").iter() {
            let mut add_file = false;
            filename = file.get_name().expect("failed to get filename.");
            if filename == file_to_run || file_to_run == "" { add_file = true; }
            for test in file.get_tests().expect("failed to get tests.").iter() {
                let mut add_test = false;
                testname = test.get_name().expect("failed to get test name.");
                if testname == test_to_run || test_to_run == "" {
                    add_test = true;
                    if add_dir && add_file && add_test {
                        let execute_test = ExecuteTest {
                            name: format!("{}::{}::{}", dirname, filename, testname),
                            test: test,
                        };
                        tests_to_execute.push(execute_test);
                    }
                }
            }
        }
    }
    if has_all_tests_passed(tests_to_execute, keep_going) {
        process::exit(0);
    } else {
        process::exit(1);
    }
}

fn test_scope(test_to_run: String) -> (String, String, String) {
    let vec: Vec<&str> = test_to_run.split("/").collect();
    (vec[0].into(), vec[1].into(), vec[2].into())
}

fn has_all_tests_passed(tests_to_execute: std::vec::Vec<ExecuteTest>, keep_going: bool) -> bool {
    println!("running {} tests", tests_to_execute.len());
    let mut has_all_tests_passed = true;
    for test in tests_to_execute {
        print!("sputnikvm test {} ", test.name);
        let test = test.test;
        let eo = test.get_expected_output().expect("failed to get expected output");
        let io = test.get_input_output().expect("failed to get actual input");
        let ao_gas = io.get_output().expect("").get_gas();
        let eo_gas = eo.get_gas();
        let ao_code = io.get_output().expect("").get_code().expect("").iter();
        let eo_code = eo.get_code().expect("").iter();
        let mut ao_vec = Vec::new();
        let mut eo_vec = Vec::new();
        let mut has_this_test_failed = false;
        for ao_char in ao_code {
            ao_vec.push(ao_char.expect("character expected")[0]);
        }
        for eo_char in eo_code {
            eo_vec.push(eo_char.expect("character expected")[0]);
        }
        let length = eo_vec.len();
        let matching = ao_vec.iter().zip(eo_vec.iter()).filter(|&(a, b)| a == b).count();
        if matching != length {
            has_this_test_failed = true;
            print!("\n\n code equality fail: only {} actual output opcodes matched the {} opcodes of the expected output.\n", matching, length);
            println!(" actual code output:\t{:?}", ao_vec);
            println!(" expected code output:\t{:?}", eo_vec);
        }
        if eo_gas != ao_gas {
            has_this_test_failed = true;
            print!("\n gas equality fail: actual output gas value of {} doesn't equal the expected output gas value of {}.", ao_gas, eo_gas);
            print!("\n actual gas output:\t{}", ao_gas);
            print!("\n expected gas output:\t{}\n", eo_gas);
        }
        if has_this_test_failed {
            print!("\n");
            has_all_tests_passed = false;
            if !keep_going {
                return has_all_tests_passed;
            }
        } else {
            print!("... ok\n")
        }
    }
    has_all_tests_passed
}
<|endoftext|>


<commit_before>extern crate hyper;
extern crate hyper_rustls;
extern crate make_hyper_great_again;
#[macro_use]
extern crate serde_derive;
#[macro_use]
extern crate serde_json;
extern crate scraper;
extern crate chrono;
extern crate dotenv;

use make_hyper_great_again::Client;

use scraper::{Html, Selector};

use chrono::{Utc, TimeZone};

use std::collections::HashMap;
use std::fs::{File, OpenOptions};
use std::io::Read;
use std::env;

const NEWS_URL: &'static str = "http://na.finalfantasyxiv.com/lodestone/news/";

struct NewsScraper {
  client: Client<hyper_rustls::HttpsConnector>,
  database: NewsDatabase
}

#[derive(Debug, Default, Serialize, Deserialize)]
struct NewsDatabase {
  items: HashMap<String, NewsItem>
}

#[derive(Debug, Serialize, Deserialize)]
struct NewsItem {
  title: String,
  url: String,
  kind: NewsKind,
  time: i64,
  tag: Option<String>
}

#[derive(Debug, Serialize, Deserialize)]
enum NewsKind {
  SpecialNotice,
  News,
  Topic
}

impl ToString for NewsKind {
  fn to_string(&self) -> String {
    match *self {
      NewsKind::News => "News",
      NewsKind::Topic => "Topic",
      NewsKind::SpecialNotice => "Special notice"
    }.to_string()
  }
}

fn main() {
  dotenv::dotenv().ok();

  let webhook_url = match env::var("LB_NEWS_WEBHOOK") {
    Ok(u) => u,
    Err(_) => {
      println!("No webhook specified");
      return;
    }
  };

  let database = match NewsDatabase::load() {
    Some(db) => db,
    None => NewsDatabase::default()
  };
  let mut scraper = NewsScraper::new(database);

  let news = match scraper.download_news() {
    Some(n) => n,
    None => {
      println!("could not download news");
      return;
    }
  };

  let downloaded_news = scraper.parse_news(&news);
  // Ignore clippy false positive
  #[cfg_attr(feature = "cargo-clippy", allow(map_entry))]
  for (id, item) in downloaded_news {
    if !scraper.database.items.contains_key(&id) {
      let mut embed = json!({
        "type": "rich",
        "timestamp": Utc.timestamp(item.time, 0).to_rfc3339(),
        "fields": [
          {
            "name": "Title",
            "value": item.title,
            "inline": false
          },
          {
            "name": "Link",
            "value": item.url,
            "inline": false
          },
          {
            "name": "Kind",
            "value": item.kind.to_string(),
            "inline": true
          }
        ]
      });
      if let Some(ref tag) = item.tag {
        embed["fields"].as_array_mut().unwrap().push(json!({
          "name": "Tag",
          "value": tag,
          "inline": true
        }));
      }
      let data = json!({
        "embeds": [embed]
      });
      let res = scraper.client.post(&webhook_url)
        .header(hyper::header::ContentType::json())
        .body(data.to_string())
        .send();
      let mut data = match res {
        Ok(r) => r,
        Err(e) => {
          println!("error sending webhook: {}", e);
          continue;
        }
      };
      let mut content = String::new();
      if let Err(e) = data.read_to_string(&mut content) {
        println!("could not read webhook response: {}", e);
        continue;
      }
      if data.status().is_success() {
        println!("discord says no to {}", item.title);
        println!("{}", content);
      } else {
        println!("webhook sent: {}", item.title);
        scraper.database.items.insert(id, item);
      }
    }
  }
  if scraper.database.save().is_none() {
    println!("could not save database");
  }
}

impl NewsDatabase {
  fn load() -> Option<Self> {
    let news_file = match File::open("./lodestone_news.json") {
      Ok(n) => n,
      Err(_) => return None
    };
    serde_json::from_reader(news_file).ok()
  }

  fn save(&self) -> Option<()> {
    let news_file = match OpenOptions::new().create(true).write(true).open("./lodestone_news.json") {
      Ok(n) => n,
      Err(_) => return None
    };
    serde_json::to_writer(news_file, self).ok()
  }
}

impl NewsScraper {
  fn new(database: NewsDatabase) -> Self {
    NewsScraper {
      client: Client::create_connector(|c| hyper_rustls::HttpsConnector::new(4, &c.handle())).unwrap(),
      database: database
    }
  }

  fn download_news(&self) -> Option<String> {
    let mut response = match self.client.get(NEWS_URL).send() {
      Ok(r) => r,
      Err(_) => return None
    };

    let mut content = String::new();
    if response.read_to_string(&mut content).is_err() {
      return None;
    }
    Some(content)
  }

  fn parse_news(&self, news: &str) -> HashMap<String, NewsItem> {
    let html = Html::parse_document(news);
    let special_notices_selector = Selector::parse("div.news__content.parts__space--add > ul:nth-of-type(1) > li").unwrap();
    let news_selector = Selector::parse("div.news__content.parts__space--add > ul:nth-of-type(2) > li").unwrap();
    let topics_selector = Selector::parse("div.news__content.parts__space--add > ul:nth-of-type(3) > li").unwrap();
    let title_selector = Selector::parse("p.news__list--title").unwrap();
    let time_script_selector = Selector::parse("time.news__list--time > script").unwrap();
    let mut lis: Vec<_> = html.select(&news_selector).map(|x| (NewsKind::News, x)).collect();
    lis.append(&mut html.select(&topics_selector).map(|x| (NewsKind::Topic, x)).collect());
    lis.append(&mut html.select(&special_notices_selector).map(|x| (NewsKind::SpecialNotice, x)).collect());
    let mut items = HashMap::with_capacity(lis.len());
    for (kind, li) in lis {
      let child = match kind {
        NewsKind::News | NewsKind::SpecialNotice => li.first_child().and_then(|v| v.value().as_element()),
        NewsKind::Topic => li.select(&title_selector).next().and_then(|v| v.first_child().and_then(|x| x.value().as_element())),
      };

      let child = match child {
        Some(c) => c,
        None => {
          println!("could not get news item child");
          continue
        }
      };

      let href = match child.attr("href") {
        Some(h) => h,
        None => {
          println!("invalid link in news item");
          continue;
        }
      };

      let (title, tag) = match kind {
        NewsKind::News | NewsKind::SpecialNotice => {
          let title = match li.select(&title_selector).next() {
            Some(t) => t,
            None => {
              println!("missing title in news item");
              continue;
            }
          };

          let tag: Option<String> = title.first_child()
            .and_then(scraper::ElementRef::wrap)
            .map(|c| c.text().collect::<String>())
            .map(|tag| tag[1..tag.len() - 1].to_string());

          let text_iter = title.text();
          let title: String = if tag.is_some() {
            text_iter.skip(1).collect()
          } else {
            text_iter.collect()
          };

          (title, tag)
        },
        NewsKind::Topic => {
          let text = li.select(&title_selector).next()
            .and_then(|v| v.first_child())
            .and_then(scraper::ElementRef::wrap)
            .map(|v| v.text().collect());
          match text {
            Some(t) => (t, None),
            None => {
              println!("invalid topic/special notice: no title");
              continue;
            }
          }
        }
      };

      let time_script = match li.select(&time_script_selector).next() {
        Some(ts) => ts,
        None => {
          println!("news item missing time script");
          continue
        }
      };

      let time_script: String = time_script.text().collect();
      let time_string = match time_script.split("strftime(").nth(1).and_then(|v| v.split(',').next()) {
        Some(time) => time,
        None => {
          println!("invalid script in news item");
          continue
        }
      };
      let time: i64 = match time_string.parse() {
        Ok(t) => t,
        Err(_) => {
          println!("invalid time in time script");
          continue
        }
      };

      let id = match href.split('/').last() {
        Some(i) => i,
        None => {
          println!("invalid href in news item");
          continue;
        }
      };

      let news_item = NewsItem {
        title: title.trim().to_string(),
        url: format!("http://na.finalfantasyxiv.com{}", href),
        kind: kind,
        time: time,
        tag: tag.map(|x| x.trim().to_string())
      };
      items.insert(id.to_string(), news_item);
    }

    items
  }
}
<commit_msg>refactor: remove lodestone_news from this package<commit_after><|endoftext|>