<commit_before>//! This module defines [`BVH`] and [`BVHNode`] and functions for building and traversing it.
//!
//! [`BVH`]: struct.BVH.html
//! [`BVHNode`]: struct.BVHNode.html
//!

use EPSILON;
use aabb::{AABB, Bounded};
use bounding_hierarchy::{BoundingHierarchy, BHShape};
use ray::Ray;
use std::f32;
use std::iter::repeat;

/// The [`BVHNode`] enum that describes a node in a [`BVH`].
/// It's either a leaf node and references a shape (by holding its index)
/// or a regular node that has two child nodes.
/// The non-leaf node stores the [`AABB`]s of its children.
///
/// [`AABB`]: ../aabb/struct.AABB.html
/// [`BVH`]: struct.BVH.html
/// [`BVH`]: struct.BVHNode.html
///
#[derive(Debug, Copy, Clone)]
pub enum BVHNode {
    /// Leaf node.
    Leaf {
        /// The node's parent.
        parent: usize,

        /// The node's depth.
        depth: u32,

        /// The shape contained in this leaf.
        shape: usize,
    },
    /// Inner node.
    Node {
        /// The node's parent.
        parent: usize,

        /// The node's depth.
        depth: u32,

        /// The convex hull of the shapes' `AABB`s in child_l.
        child_l_aabb: AABB,

        /// Index of the left subtree's root node.
        child_l: usize,

        /// The convex hull of the shapes' `AABB`s in child_r.
        child_r_aabb: AABB,

        /// Index of the right subtree's root node.
        child_r: usize,
    },
}

impl BVHNode {
    /// The build function sometimes needs to add nodes while their data is not available yet.
    /// A dummy cerated by this function serves the purpose of being changed later on.
    fn create_dummy() -> BVHNode {
        BVHNode::Leaf {
            parent: 0,
            depth: 0,
            shape: 0,
        }
    }

    /// Builds a [`BVHNode`] recursively using SAH partitioning.
    /// Returns the index of the node in the nodes vector.
    ///
    /// [`BVHNode`]: enum.BVHNode.html
    ///
    pub fn build<T: BHShape>(shapes: &mut [T],
                             indices: Vec<usize>,
                             nodes: &mut Vec<BVHNode>,
                             parent: usize,
                             depth: u32)
                             -> usize {
        // Helper function to accumulate the AABB joint and the centroids AABB
        fn grow_convex_hull(convex_hull: (AABB, AABB), shape_aabb: &AABB) -> (AABB, AABB) {
            let center = &shape_aabb.center();
            let convex_hull_aabbs = &convex_hull.0;
            let convex_hull_centroids = &convex_hull.1;
            (convex_hull_aabbs.join(shape_aabb), convex_hull_centroids.grow(center))
        }

        let mut convex_hull = Default::default();
        for index in &indices {
            convex_hull = grow_convex_hull(convex_hull, &shapes[*index].aabb());
        }
        let (aabb_bounds, centroid_bounds) = convex_hull;

        // If there is only one element left, don't split anymore
        if indices.len() == 1 {
            let shape_index = indices[0];
            let node_index = nodes.len();
            nodes.push(BVHNode::Leaf {
                parent: parent,
                depth: depth,
                shape: shape_index,
            });
            // Let the shape know the index of the node that represents it
            shapes[shape_index].set_bh_node_index(node_index);
            return node_index;
        }

        // Find the axis along which the shapes are spread the most
        let split_axis = centroid_bounds.largest_axis();
        let split_axis_size = centroid_bounds.max[split_axis] - centroid_bounds.min[split_axis];

        /// Defines a Bucket utility object.
        #[derive(Copy, Clone)]
        struct Bucket {
            size: usize,
            aabb: AABB,
        }

        impl Bucket {
            /// Returns an empty bucket.
            fn empty() -> Bucket {
                Bucket {
                    size: 0,
                    aabb: AABB::empty(),
                }
            }

            /// Extends this `Bucket` by the given `AABB`.
            fn add_aabb(&mut self, aabb: &AABB) {
                self.size += 1;
                self.aabb = self.aabb.join(aabb);
            }
        }

        /// Joins two `Bucket`s.
        fn join_bucket(a: Bucket, b: &Bucket) -> Bucket {
            Bucket {
                size: a.size + b.size,
                aabb: a.aabb.join(&b.aabb),
            }
        }

        // Create six buckets, and six index assignment vectors
        const NUM_BUCKETS: usize = 6;
        let mut buckets = [Bucket::empty(); NUM_BUCKETS];
        let mut bucket_assignments: [Vec<usize>; NUM_BUCKETS] = Default::default();

        if split_axis_size < EPSILON {
            // Spread the remaining shapes evenly across buckets,
            // instead of using a heuristic
            for (counter, idx) in (&indices).iter().enumerate() {
                let shape = &shapes[*idx];
                let shape_aabb = shape.aabb();

                let bucket_num = counter % NUM_BUCKETS;

                buckets[bucket_num].add_aabb(&shape_aabb);
                bucket_assignments[bucket_num].push(*idx);
            }
        } else {
            // Assign each shape to a bucket
            for idx in &indices {
                let shape = &shapes[*idx];
                let shape_aabb = shape.aabb();
                let shape_center = shape_aabb.center();

                // Get the relative position of the shape centroid [0.0..1.0]
                let bucket_num_relative =
                    (shape_center[split_axis] - centroid_bounds.min[split_axis]) / split_axis_size;

                // Convert that to the actual `Bucket` number
                let bucket_num = (bucket_num_relative * (NUM_BUCKETS as f32 - 0.01)) as usize;

                // Extend the selected `Bucket` and add the index to the actual bucket
                buckets[bucket_num].add_aabb(&shape_aabb);
                bucket_assignments[bucket_num].push(*idx);
            }
        }

        // Compute the costs for each configuration and
        // select the configuration with the minimal costs
        let mut min_bucket = 0;
        let mut min_cost = f32::INFINITY;
        let mut child_l_aabb = AABB::empty();
        let mut child_r_aabb = AABB::empty();
        for i in 0..(NUM_BUCKETS - 1) {
            let (l_buckets, r_buckets) = buckets.split_at(i + 1);
            let child_l = l_buckets.iter().fold(Bucket::empty(), join_bucket);
            let child_r = r_buckets.iter().fold(Bucket::empty(), join_bucket);

            let cost = (child_l.size as f32 * child_l.aabb.surface_area() +
                        child_r.size as f32 * child_r.aabb.surface_area()) /
                       aabb_bounds.surface_area();

            if cost < min_cost {
                min_bucket = i;
                min_cost = cost;
                child_l_aabb = child_l.aabb;
                child_r_aabb = child_r.aabb;
            }
        }

        fn join_vectors(vectors: &mut [Vec<usize>]) -> Vec<usize> {
            let mut result = Vec::new();
            for mut vector in vectors.iter_mut() {
                result.append(&mut vector);
            }
            result
        }

        // Join together all index buckets, and proceed recursively
        let (l_assignments, r_assignments) = bucket_assignments.split_at_mut(min_bucket + 1);
        let child_l_indices = join_vectors(l_assignments);
        let child_r_indices = join_vectors(r_assignments);

        let node_index = nodes.len();
        nodes.push(BVHNode::create_dummy());

        let child_l = BVHNode::build(shapes, child_l_indices, nodes, node_index, depth + 1);
        let child_r = BVHNode::build(shapes, child_r_indices, nodes, node_index, depth + 1);

        // Construct the actual data structure
        nodes[node_index] = BVHNode::Node {
            parent: parent,
            depth: depth,
            child_l_aabb: child_l_aabb,
            child_l: child_l,
            child_r_aabb: child_r_aabb,
            child_r: child_r,
        };

        node_index
    }

    /// Traverses the [`BVH`] recursively and returns all shapes whose [`AABB`] is
    /// intersected by the given [`Ray`].
    ///
    /// [`AABB`]: ../aabb/struct.AABB.html
    /// [`BVH`]: struct.BVH.html
    /// [`Ray`]: ../ray/struct.Ray.html
    ///
    pub fn traverse_recursive(nodes: &Vec<BVHNode>,
                              node_index: usize,
                              ray: &Ray,
                              indices: &mut Vec<usize>) {
        let node = &nodes[node_index];
        match *node {
            BVHNode::Node { ref child_l_aabb, child_l, ref child_r_aabb, child_r, .. } => {
                if ray.intersects_aabb(child_l_aabb) {
                    BVHNode::traverse_recursive(nodes, child_l, ray, indices);
                }
                if ray.intersects_aabb(child_r_aabb) {
                    BVHNode::traverse_recursive(nodes, child_r, ray, indices);
                }
            }
            BVHNode::Leaf { shape, .. } => {
                indices.push(shape);
            }
        }
    }
}

/// The [`BVH`] data structure. Contains the list of [`BVHNode`]s.
///
/// [`BVH`]: struct.BVH.html
///
pub struct BVH {
    /// The list of nodes of the [`BVH`].
    ///
    /// [`BVH`]: struct.BVH.html
    ///
    pub nodes: Vec<BVHNode>,
}

impl BVH {
    /// Creates a new [`BVH`] from the `shapes` slice.
    ///
    /// [`BVH`]: struct.BVH.html
    ///
    pub fn build<Shape: BHShape>(shapes: &mut [Shape]) -> BVH {
        let indices = (0..shapes.len()).collect::<Vec<usize>>();
        let mut nodes = Vec::new();
        BVHNode::build(shapes, indices, &mut nodes, 0, 0);
        BVH { nodes: nodes }
    }

    /// Traverses the [`BVH`].
    /// Returns a subset of `shapes`, in which the [`AABB`]s of the elements were hit by `ray`.
    ///
    /// [`BVH`]: struct.BVH.html
    /// [`AABB`]: ../aabb/struct.AABB.html
    ///
    pub fn traverse<'a, Shape: Bounded>(&'a self, ray: &Ray, shapes: &'a [Shape]) -> Vec<&Shape> {
        let mut indices = Vec::new();
        BVHNode::traverse_recursive(&self.nodes, 0, ray, &mut indices);
        let mut hit_shapes = Vec::new();
        for index in &indices {
            let shape = &shapes[*index];
            hit_shapes.push(shape);
        }
        hit_shapes
    }

    /// Prints the [`BVH`] in a tree-like visualization.
    ///
    /// [`BVH`]: struct.BVH.html
    ///
    pub fn pretty_print(&self) {
        let nodes = &self.nodes;
        fn print_node(nodes: &Vec<BVHNode>, node_index: usize) {
            let node = &nodes[node_index];
            match *node {
                BVHNode::Node { child_l, child_r, depth, child_l_aabb, child_r_aabb, .. } => {
                    let padding: String = repeat(" ").take(depth as usize).collect();
                    println!("{}child_l {:?}", padding, child_l_aabb);
                    print_node(nodes, child_l);
                    println!("{}child_r {:?}", padding, child_r_aabb);
                    print_node(nodes, child_r);
                }
                BVHNode::Leaf { shape, depth, .. } => {
                    let padding: String = repeat(" ").take(depth as usize).collect();
                    println!("{}shape\t{:?}", padding, shape);
                }
            }
        }
        print_node(nodes, 0);
    }
}

impl BoundingHierarchy for BVH {
    fn build<Shape: BHShape>(shapes: &mut [Shape]) -> BVH {
        BVH::build(shapes)
    }

    fn traverse<'a, Shape: Bounded>(&'a self, ray: &Ray, shapes: &'a [Shape]) -> Vec<&Shape> {
        self.traverse(ray, shapes)
    }

    fn pretty_print(&self) {
        self.pretty_print();
    }
}

#[cfg(test)]
pub mod tests {
    use bvh::BVH;
    use testbase::{build_some_bh, traverse_some_bh, build_1200_triangles_bh, build_12k_triangles_bh,
                   build_120k_triangles_bh, intersect_1200_triangles_bh, intersect_12k_triangles_bh,
                   intersect_120k_triangles_bh};

    #[test]
    /// Tests whether the building procedure succeeds in not failing.
    fn test_build_bvh() {
        build_some_bh::<BVH>();
    }

    #[test]
    /// Runs some primitive tests for intersections of a ray with a fixed scene given as a BVH.
    fn test_traverse_bvh() {
        traverse_some_bh::<BVH>();
    }

    #[bench]
    /// Benchmark the construction of a `BVH` with 1,200 triangles.
    fn bench_build_1200_triangles_bvh(mut b: &mut ::test::Bencher) {
        build_1200_triangles_bh::<BVH>(&mut b);
    }

    #[bench]
    /// Benchmark the construction of a `BVH` with 12,000 triangles.
    fn bench_build_12k_triangles_bvh(mut b: &mut ::test::Bencher) {
        build_12k_triangles_bh::<BVH>(&mut b);
    }

    #[bench]
    /// Benchmark the construction of a `BVH` with 120,000 triangles.
    fn bench_build_120k_triangles_bvh(mut b: &mut ::test::Bencher) {
        build_120k_triangles_bh::<BVH>(&mut b);
    }

    #[bench]
    /// Benchmark intersecting 1,200 triangles using the recursive `BVH`.
    fn bench_intersect_1200_triangles_bvh(mut b: &mut ::test::Bencher) {
        intersect_1200_triangles_bh::<BVH>(&mut b);
    }

    #[bench]
    /// Benchmark intersecting 12,000 triangles using the recursive `BVH`.
    fn bench_intersect_12k_triangles_bvh(mut b: &mut ::test::Bencher) {
        intersect_12k_triangles_bh::<BVH>(&mut b);
    }

    #[bench]
    /// Benchmark intersecting 120,000 triangles using the recursive `BVH`.
    fn bench_intersect_120k_triangles_bvh(mut b: &mut ::test::Bencher) {
        intersect_120k_triangles_bh::<BVH>(&mut b);
    }
}
<commit_msg>Create nodes vector with capacity<commit_after>//! This module defines [`BVH`] and [`BVHNode`] and functions for building and traversing it.
//!
//! [`BVH`]: struct.BVH.html
//! [`BVHNode`]: struct.BVHNode.html
//!

use EPSILON;
use aabb::{AABB, Bounded};
use bounding_hierarchy::{BoundingHierarchy, BHShape};
use ray::Ray;
use std::f32;
use std::iter::repeat;

/// The [`BVHNode`] enum that describes a node in a [`BVH`].
/// It's either a leaf node and references a shape (by holding its index)
/// or a regular node that has two child nodes.
/// The non-leaf node stores the [`AABB`]s of its children.
///
/// [`AABB`]: ../aabb/struct.AABB.html
/// [`BVH`]: struct.BVH.html
/// [`BVH`]: struct.BVHNode.html
///
#[derive(Debug, Copy, Clone)]
pub enum BVHNode {
    /// Leaf node.
    Leaf {
        /// The node's parent.
        parent: usize,

        /// The node's depth.
        depth: u32,

        /// The shape contained in this leaf.
        shape: usize,
    },
    /// Inner node.
    Node {
        /// The node's parent.
        parent: usize,

        /// The node's depth.
        depth: u32,

        /// The convex hull of the shapes' `AABB`s in child_l.
        child_l_aabb: AABB,

        /// Index of the left subtree's root node.
        child_l: usize,

        /// The convex hull of the shapes' `AABB`s in child_r.
        child_r_aabb: AABB,

        /// Index of the right subtree's root node.
        child_r: usize,
    },
}

impl BVHNode {
    /// The build function sometimes needs to add nodes while their data is not available yet.
    /// A dummy cerated by this function serves the purpose of being changed later on.
    fn create_dummy() -> BVHNode {
        BVHNode::Leaf {
            parent: 0,
            depth: 0,
            shape: 0,
        }
    }

    /// Builds a [`BVHNode`] recursively using SAH partitioning.
    /// Returns the index of the node in the nodes vector.
    ///
    /// [`BVHNode`]: enum.BVHNode.html
    ///
    pub fn build<T: BHShape>(shapes: &mut [T],
                             indices: Vec<usize>,
                             nodes: &mut Vec<BVHNode>,
                             parent: usize,
                             depth: u32)
                             -> usize {
        // Helper function to accumulate the AABB joint and the centroids AABB
        fn grow_convex_hull(convex_hull: (AABB, AABB), shape_aabb: &AABB) -> (AABB, AABB) {
            let center = &shape_aabb.center();
            let convex_hull_aabbs = &convex_hull.0;
            let convex_hull_centroids = &convex_hull.1;
            (convex_hull_aabbs.join(shape_aabb), convex_hull_centroids.grow(center))
        }

        let mut convex_hull = Default::default();
        for index in &indices {
            convex_hull = grow_convex_hull(convex_hull, &shapes[*index].aabb());
        }
        let (aabb_bounds, centroid_bounds) = convex_hull;

        // If there is only one element left, don't split anymore
        if indices.len() == 1 {
            let shape_index = indices[0];
            let node_index = nodes.len();
            nodes.push(BVHNode::Leaf {
                parent: parent,
                depth: depth,
                shape: shape_index,
            });
            // Let the shape know the index of the node that represents it
            shapes[shape_index].set_bh_node_index(node_index);
            return node_index;
        }

        // Find the axis along which the shapes are spread the most
        let split_axis = centroid_bounds.largest_axis();
        let split_axis_size = centroid_bounds.max[split_axis] - centroid_bounds.min[split_axis];

        /// Defines a Bucket utility object.
        #[derive(Copy, Clone)]
        struct Bucket {
            size: usize,
            aabb: AABB,
        }

        impl Bucket {
            /// Returns an empty bucket.
            fn empty() -> Bucket {
                Bucket {
                    size: 0,
                    aabb: AABB::empty(),
                }
            }

            /// Extends this `Bucket` by the given `AABB`.
            fn add_aabb(&mut self, aabb: &AABB) {
                self.size += 1;
                self.aabb = self.aabb.join(aabb);
            }
        }

        /// Joins two `Bucket`s.
        fn join_bucket(a: Bucket, b: &Bucket) -> Bucket {
            Bucket {
                size: a.size + b.size,
                aabb: a.aabb.join(&b.aabb),
            }
        }

        // Create six buckets, and six index assignment vectors
        const NUM_BUCKETS: usize = 6;
        let mut buckets = [Bucket::empty(); NUM_BUCKETS];
        let mut bucket_assignments: [Vec<usize>; NUM_BUCKETS] = Default::default();

        if split_axis_size < EPSILON {
            // Spread the remaining shapes evenly across buckets,
            // instead of using a heuristic
            for (counter, idx) in (&indices).iter().enumerate() {
                let shape = &shapes[*idx];
                let shape_aabb = shape.aabb();

                let bucket_num = counter % NUM_BUCKETS;

                buckets[bucket_num].add_aabb(&shape_aabb);
                bucket_assignments[bucket_num].push(*idx);
            }
        } else {
            // Assign each shape to a bucket
            for idx in &indices {
                let shape = &shapes[*idx];
                let shape_aabb = shape.aabb();
                let shape_center = shape_aabb.center();

                // Get the relative position of the shape centroid [0.0..1.0]
                let bucket_num_relative =
                    (shape_center[split_axis] - centroid_bounds.min[split_axis]) / split_axis_size;

                // Convert that to the actual `Bucket` number
                let bucket_num = (bucket_num_relative * (NUM_BUCKETS as f32 - 0.01)) as usize;

                // Extend the selected `Bucket` and add the index to the actual bucket
                buckets[bucket_num].add_aabb(&shape_aabb);
                bucket_assignments[bucket_num].push(*idx);
            }
        }

        // Compute the costs for each configuration and
        // select the configuration with the minimal costs
        let mut min_bucket = 0;
        let mut min_cost = f32::INFINITY;
        let mut child_l_aabb = AABB::empty();
        let mut child_r_aabb = AABB::empty();
        for i in 0..(NUM_BUCKETS - 1) {
            let (l_buckets, r_buckets) = buckets.split_at(i + 1);
            let child_l = l_buckets.iter().fold(Bucket::empty(), join_bucket);
            let child_r = r_buckets.iter().fold(Bucket::empty(), join_bucket);

            let cost = (child_l.size as f32 * child_l.aabb.surface_area() +
                        child_r.size as f32 * child_r.aabb.surface_area()) /
                       aabb_bounds.surface_area();

            if cost < min_cost {
                min_bucket = i;
                min_cost = cost;
                child_l_aabb = child_l.aabb;
                child_r_aabb = child_r.aabb;
            }
        }

        fn join_vectors(vectors: &mut [Vec<usize>]) -> Vec<usize> {
            let mut result = Vec::new();
            for mut vector in vectors.iter_mut() {
                result.append(&mut vector);
            }
            result
        }

        // Join together all index buckets, and proceed recursively
        let (l_assignments, r_assignments) = bucket_assignments.split_at_mut(min_bucket + 1);
        let child_l_indices = join_vectors(l_assignments);
        let child_r_indices = join_vectors(r_assignments);

        let node_index = nodes.len();
        nodes.push(BVHNode::create_dummy());

        let child_l = BVHNode::build(shapes, child_l_indices, nodes, node_index, depth + 1);
        let child_r = BVHNode::build(shapes, child_r_indices, nodes, node_index, depth + 1);

        // Construct the actual data structure
        nodes[node_index] = BVHNode::Node {
            parent: parent,
            depth: depth,
            child_l_aabb: child_l_aabb,
            child_l: child_l,
            child_r_aabb: child_r_aabb,
            child_r: child_r,
        };

        node_index
    }

    /// Traverses the [`BVH`] recursively and returns all shapes whose [`AABB`] is
    /// intersected by the given [`Ray`].
    ///
    /// [`AABB`]: ../aabb/struct.AABB.html
    /// [`BVH`]: struct.BVH.html
    /// [`Ray`]: ../ray/struct.Ray.html
    ///
    pub fn traverse_recursive(nodes: &Vec<BVHNode>,
                              node_index: usize,
                              ray: &Ray,
                              indices: &mut Vec<usize>) {
        let node = &nodes[node_index];
        match *node {
            BVHNode::Node { ref child_l_aabb, child_l, ref child_r_aabb, child_r, .. } => {
                if ray.intersects_aabb(child_l_aabb) {
                    BVHNode::traverse_recursive(nodes, child_l, ray, indices);
                }
                if ray.intersects_aabb(child_r_aabb) {
                    BVHNode::traverse_recursive(nodes, child_r, ray, indices);
                }
            }
            BVHNode::Leaf { shape, .. } => {
                indices.push(shape);
            }
        }
    }
}

/// The [`BVH`] data structure. Contains the list of [`BVHNode`]s.
///
/// [`BVH`]: struct.BVH.html
///
pub struct BVH {
    /// The list of nodes of the [`BVH`].
    ///
    /// [`BVH`]: struct.BVH.html
    ///
    pub nodes: Vec<BVHNode>,
}

impl BVH {
    /// Creates a new [`BVH`] from the `shapes` slice.
    ///
    /// [`BVH`]: struct.BVH.html
    ///
    pub fn build<Shape: BHShape>(shapes: &mut [Shape]) -> BVH {
        let indices = (0..shapes.len()).collect::<Vec<usize>>();
        let expected_node_count = shapes.len() * 2;
        let mut nodes = Vec::with_capacity(expected_node_count);
        BVHNode::build(shapes, indices, &mut nodes, 0, 0);
        BVH { nodes: nodes }
    }

    /// Traverses the [`BVH`].
    /// Returns a subset of `shapes`, in which the [`AABB`]s of the elements were hit by `ray`.
    ///
    /// [`BVH`]: struct.BVH.html
    /// [`AABB`]: ../aabb/struct.AABB.html
    ///
    pub fn traverse<'a, Shape: Bounded>(&'a self, ray: &Ray, shapes: &'a [Shape]) -> Vec<&Shape> {
        let mut indices = Vec::new();
        BVHNode::traverse_recursive(&self.nodes, 0, ray, &mut indices);
        let mut hit_shapes = Vec::new();
        for index in &indices {
            let shape = &shapes[*index];
            hit_shapes.push(shape);
        }
        hit_shapes
    }

    /// Prints the [`BVH`] in a tree-like visualization.
    ///
    /// [`BVH`]: struct.BVH.html
    ///
    pub fn pretty_print(&self) {
        let nodes = &self.nodes;
        fn print_node(nodes: &Vec<BVHNode>, node_index: usize) {
            let node = &nodes[node_index];
            match *node {
                BVHNode::Node { child_l, child_r, depth, child_l_aabb, child_r_aabb, .. } => {
                    let padding: String = repeat(" ").take(depth as usize).collect();
                    println!("{}child_l {:?}", padding, child_l_aabb);
                    print_node(nodes, child_l);
                    println!("{}child_r {:?}", padding, child_r_aabb);
                    print_node(nodes, child_r);
                }
                BVHNode::Leaf { shape, depth, .. } => {
                    let padding: String = repeat(" ").take(depth as usize).collect();
                    println!("{}shape\t{:?}", padding, shape);
                }
            }
        }
        print_node(nodes, 0);
    }
}

impl BoundingHierarchy for BVH {
    fn build<Shape: BHShape>(shapes: &mut [Shape]) -> BVH {
        BVH::build(shapes)
    }

    fn traverse<'a, Shape: Bounded>(&'a self, ray: &Ray, shapes: &'a [Shape]) -> Vec<&Shape> {
        self.traverse(ray, shapes)
    }

    fn pretty_print(&self) {
        self.pretty_print();
    }
}

#[cfg(test)]
pub mod tests {
    use bvh::BVH;
    use testbase::{build_some_bh, traverse_some_bh, build_1200_triangles_bh, build_12k_triangles_bh,
                   build_120k_triangles_bh, intersect_1200_triangles_bh, intersect_12k_triangles_bh,
                   intersect_120k_triangles_bh};

    #[test]
    /// Tests whether the building procedure succeeds in not failing.
    fn test_build_bvh() {
        build_some_bh::<BVH>();
    }

    #[test]
    /// Runs some primitive tests for intersections of a ray with a fixed scene given as a BVH.
    fn test_traverse_bvh() {
        traverse_some_bh::<BVH>();
    }

    #[bench]
    /// Benchmark the construction of a `BVH` with 1,200 triangles.
    fn bench_build_1200_triangles_bvh(mut b: &mut ::test::Bencher) {
        build_1200_triangles_bh::<BVH>(&mut b);
    }

    #[bench]
    /// Benchmark the construction of a `BVH` with 12,000 triangles.
    fn bench_build_12k_triangles_bvh(mut b: &mut ::test::Bencher) {
        build_12k_triangles_bh::<BVH>(&mut b);
    }

    #[bench]
    /// Benchmark the construction of a `BVH` with 120,000 triangles.
    fn bench_build_120k_triangles_bvh(mut b: &mut ::test::Bencher) {
        build_120k_triangles_bh::<BVH>(&mut b);
    }

    #[bench]
    /// Benchmark intersecting 1,200 triangles using the recursive `BVH`.
    fn bench_intersect_1200_triangles_bvh(mut b: &mut ::test::Bencher) {
        intersect_1200_triangles_bh::<BVH>(&mut b);
    }

    #[bench]
    /// Benchmark intersecting 12,000 triangles using the recursive `BVH`.
    fn bench_intersect_12k_triangles_bvh(mut b: &mut ::test::Bencher) {
        intersect_12k_triangles_bh::<BVH>(&mut b);
    }

    #[bench]
    /// Benchmark intersecting 120,000 triangles using the recursive `BVH`.
    fn bench_intersect_120k_triangles_bvh(mut b: &mut ::test::Bencher) {
        intersect_120k_triangles_bh::<BVH>(&mut b);
    }
}
<|endoftext|>


<commit_before>use serde_json::value::Value as Json;

use crate::block::{BlockContext, BlockParams};
use crate::context::Context;
use crate::error::RenderError;
use crate::helpers::{HelperDef, HelperResult};
use crate::json::value::{to_json, JsonTruthy};
use crate::output::Output;
use crate::registry::Registry;
use crate::render::{Helper, RenderContext, Renderable};
use crate::util::copy_on_push_vec;

#[derive(Clone, Copy)]
pub struct EachHelper;

impl HelperDef for EachHelper {
    fn call<'reg: 'rc, 'rc>(
        &self,
        h: &Helper<'reg, 'rc>,
        r: &'reg Registry<'reg>,
        ctx: &'rc Context,
        rc: &mut RenderContext<'reg, 'rc>,
        out: &mut dyn Output,
    ) -> HelperResult {
        let value = h
            .param(0)
            .ok_or_else(|| RenderError::new("Param not found for helper \"each\""))?;

        let template = h.template();

        match template {
            Some(t) => {
                let mut block_context = BlockContext::new();

                if let Some(path) = value.context_path() {
                    *block_context.base_path_mut() = path.to_vec();
                    // TODO: disable base value support for now.
                    // block_context.set_base_value(value.value());
                }

                rc.push_block(block_context);

                let rendered = match (value.value().is_truthy(false), value.value()) {
                    (true, &Json::Array(ref list)) => {
                        let len = list.len();

                        let array_path = value.context_path();

                        for (i, _) in list.iter().enumerate().take(len) {
                            if let Some(ref mut block) = rc.block_mut() {
                                let is_first = i == 0usize;
                                let is_last = i == len - 1;

                                block.set_local_var("@first".to_string(), to_json(is_first));
                                block.set_local_var("@last".to_string(), to_json(is_last));
                                block.set_local_var("@index".to_string(), to_json(i));
                                if let Some(ref p) = array_path {
                                    if is_first {
                                        *block.base_path_mut() = copy_on_push_vec(p, i.to_string());
                                    } else if let Some(ptr) = block.base_path_mut().last_mut() {
                                        *ptr = i.to_string();
                                    }
                                }

                                // TODO: base value

                                if let Some(bp_val) = h.block_param() {
                                    let mut params = BlockParams::new();
                                    params.add_path(bp_val, Vec::with_capacity(0))?;

                                    block.set_block_params(params);
                                } else if let Some((bp_val, bp_index)) = h.block_param_pair() {
                                    let mut params = BlockParams::new();
                                    params.add_path(bp_val, Vec::with_capacity(0))?;
                                    params.add_value(bp_index, to_json(i))?;

                                    block.set_block_params(params);
                                }
                            }

                            t.render(r, ctx, rc, out)?;
                        }

                        Ok(())
                    }
                    (true, &Json::Object(ref obj)) => {
                        let mut is_first = true;
                        let obj_path = value.context_path();

                        for (k, _) in obj.iter() {
                            if let Some(ref mut block) = rc.block_mut() {
                                block.set_local_var("@first".to_string(), to_json(is_first));
                                block.set_local_var("@key".to_string(), to_json(k));

                                if let Some(ref p) = obj_path {
                                    if is_first {
                                        *block.base_path_mut() = copy_on_push_vec(p, k.clone());
                                    } else if let Some(ptr) = block.base_path_mut().last_mut() {
                                        *ptr = k.clone();
                                    }
                                }

                                // TODO

                                if let Some(bp_val) = h.block_param() {
                                    let mut params = BlockParams::new();
                                    params.add_path(bp_val, Vec::with_capacity(0))?;

                                    block.set_block_params(params);
                                } else if let Some((bp_val, bp_key)) = h.block_param_pair() {
                                    let mut params = BlockParams::new();
                                    params.add_path(bp_val, Vec::with_capacity(0))?;
                                    params.add_value(bp_key, to_json(&k))?;

                                    block.set_block_params(params);
                                }
                            }

                            t.render(r, ctx, rc, out)?;

                            if is_first {
                                is_first = false;
                            }
                        }
                        Ok(())
                    }
                    (false, _) => {
                        if let Some(else_template) = h.inverse() {
                            else_template.render(r, ctx, rc, out)?;
                        }
                        Ok(())
                    }
                    _ => Err(RenderError::new(format!(
                        "Param type is not iterable: {:?}",
                        value.value()
                    ))),
                };

                rc.pop_block();
                rendered
            }
            None => Ok(()),
        }
    }
}

pub static EACH_HELPER: EachHelper = EachHelper;

#[cfg(test)]
mod test {
    use crate::json::value::to_json;
    use crate::registry::Registry;
    use serde_json::value::Value as Json;
    use std::collections::BTreeMap;
    use std::str::FromStr;

    #[test]
    fn test_each() {
        let mut handlebars = Registry::new();
        assert!(handlebars
            .register_template_string(
                "t0",
                "{{#each this}}{{@first}}|{{@last}}|{{@index}}:{{this}}|{{/each}}"
            )
            .is_ok());
        assert!(handlebars
            .register_template_string("t1", "{{#each this}}{{@first}}|{{@key}}:{{this}}|{{/each}}")
            .is_ok());

        let r0 = handlebars.render("t0", &vec![1u16, 2u16, 3u16]);
        assert_eq!(
            r0.ok().unwrap(),
            "true|false|0:1|false|false|1:2|false|true|2:3|".to_string()
        );

        let mut m: BTreeMap<String, u16> = BTreeMap::new();
        m.insert("ftp".to_string(), 21);
        m.insert("http".to_string(), 80);
        let r1 = handlebars.render("t1", &m);
        assert_eq!(r1.ok().unwrap(), "true|ftp:21|false|http:80|".to_string());
    }

    #[test]
    fn test_each_with_parent() {
        let json_str = r#"{"a":{"a":99,"c":[{"d":100},{"d":200}]}}"#;

        let data = Json::from_str(json_str).unwrap();
        // println!("data: {}", data);
        let mut handlebars = Registry::new();

        // previously, to access the parent in an each block,
        // a user would need to specify ../../b, as the path
        // that is computed includes the array index: ./a.c.[0]
        assert!(handlebars
            .register_template_string("t0", "{{#each a.c}} d={{d}} b={{../a.a}} {{/each}}")
            .is_ok());

        let r1 = handlebars.render("t0", &data);
        assert_eq!(r1.ok().unwrap(), " d=100 b=99  d=200 b=99 ".to_string());
    }

    #[test]
    fn test_nested_each_with_parent() {
        let json_str = r#"{"a": [{"b": [{"d": 100}], "c": 200}]}"#;

        let data = Json::from_str(json_str).unwrap();
        let mut handlebars = Registry::new();
        assert!(handlebars
            .register_template_string(
                "t0",
                "{{#each a}}{{#each b}}{{d}}:{{../c}}{{/each}}{{/each}}"
            )
            .is_ok());

        let r1 = handlebars.render("t0", &data);
        assert_eq!(r1.ok().unwrap(), "100:200".to_string());
    }

    #[test]
    fn test_nested_each() {
        let json_str = r#"{"a": [{"b": true}], "b": [[1, 2, 3],[4, 5]]}"#;

        let data = Json::from_str(json_str).unwrap();

        let mut handlebars = Registry::new();
        assert!(handlebars
            .register_template_string(
                "t0",
                "{{#each b}}{{#if ../a}}{{#each this}}{{this}}{{/each}}{{/if}}{{/each}}"
            )
            .is_ok());

        let r1 = handlebars.render("t0", &data);
        assert_eq!(r1.ok().unwrap(), "12345".to_string());
    }

    #[test]
    fn test_nested_array() {
        let mut handlebars = Registry::new();
        assert!(handlebars
            .register_template_string("t0", "{{#each this.[0]}}{{this}}{{/each}}")
            .is_ok());

        let r0 = handlebars.render("t0", &(vec![vec![1, 2, 3]]));

        assert_eq!(r0.ok().unwrap(), "123".to_string());
    }

    #[test]
    fn test_empty_key() {
        let mut handlebars = Registry::new();
        assert!(handlebars
            .register_template_string("t0", "{{#each this}}{{@key}}-{{value}}\n{{/each}}")
            .is_ok());

        let r0 = handlebars
            .render(
                "t0",
                &json!({
                    "foo": {
                        "value": "bar"
                    },
                    "": {
                        "value": "baz"
                    }
                }),
            )
            .unwrap();

        let mut r0_sp: Vec<_> = r0.split('\n').collect();
        r0_sp.sort();

        assert_eq!(r0_sp, vec!["", "-baz", "foo-bar"]);
    }

    #[test]
    fn test_each_else() {
        let mut handlebars = Registry::new();
        assert!(handlebars
            .register_template_string("t0", "{{#each a}}1{{else}}empty{{/each}}")
            .is_ok());
        let m1 = btreemap! {
            "a".to_string() => Vec::<String>::new(),
        };
        let r0 = handlebars.render("t0", &m1).unwrap();
        assert_eq!(r0, "empty");

        let m2 = btreemap! {
            "b".to_string() => Vec::<String>::new()
        };
        let r1 = handlebars.render("t0", &m2).unwrap();
        assert_eq!(r1, "empty");
    }

    #[test]
    fn test_block_param() {
        let mut handlebars = Registry::new();
        assert!(handlebars
            .register_template_string("t0", "{{#each a as |i|}}{{i}}{{/each}}")
            .is_ok());
        let m1 = btreemap! {
            "a".to_string() => vec![1,2,3,4,5]
        };
        let r0 = handlebars.render("t0", &m1).unwrap();
        assert_eq!(r0, "12345");
    }

    #[test]
    fn test_each_object_block_param() {
        let mut handlebars = Registry::new();
        let template = "{{#each this as |v k|}}\
                        {{#with k as |inner_k|}}{{inner_k}}{{/with}}:{{v}}|\
                        {{/each}}";
        assert!(handlebars.register_template_string("t0", template).is_ok());

        let m = btreemap! {
            "ftp".to_string() => 21,
            "http".to_string() => 80
        };
        let r0 = handlebars.render("t0", &m);
        assert_eq!(r0.ok().unwrap(), "ftp:21|http:80|".to_string());
    }

    #[test]
    fn test_each_object_block_param2() {
        let mut handlebars = Registry::new();
        let template = "{{#each this as |v k|}}\
                        {{#with v as |inner_v|}}{{k}}:{{inner_v}}{{/with}}|\
                        {{/each}}";

        assert!(handlebars.register_template_string("t0", template).is_ok());

        let m = btreemap! {
            "ftp".to_string() => 21,
            "http".to_string() => 80
        };
        let r0 = handlebars.render("t0", &m);
        assert_eq!(r0.ok().unwrap(), "ftp:21|http:80|".to_string());
    }

    fn test_nested_each_with_path_ups() {
        let mut handlebars = Registry::new();
        assert!(handlebars
            .register_template_string(
                "t0",
                "{{#each a.b}}{{#each c}}{{../../d}}{{/each}}{{/each}}"
            )
            .is_ok());

        let data = btreemap! {
            "a".to_string() => to_json(&btreemap! {
                "b".to_string() => vec![btreemap!{"c".to_string() => vec![1]}]
            }),
            "d".to_string() => to_json(&1)
        };

        let r0 = handlebars.render("t0", &data);
        assert_eq!(r0.ok().unwrap(), "1".to_string());
    }

    #[test]
    fn test_nested_each_with_path_up_this() {
        let mut handlebars = Registry::new();
        let template = "{{#each variant}}{{#each ../typearg}}\
                        {{#if @first}}template<{{/if}}{{this}}{{#if @last}}>{{else}},{{/if}}\
                        {{/each}}{{/each}}";
        assert!(handlebars.register_template_string("t0", template).is_ok());
        let data = btreemap! {
            "typearg".to_string() => vec!["T".to_string()],
            "variant".to_string() => vec!["1".to_string(), "2".to_string()]
        };
        let r0 = handlebars.render("t0", &data);
        assert_eq!(r0.ok().unwrap(), "template<T>template<T>".to_string());
    }

    #[test]
    fn test_key_iteration_with_unicode() {
        let mut handlebars = Registry::new();
        assert!(handlebars
            .register_template_string("t0", "{{#each this}}{{@key}}: {{this}}\n{{/each}}")
            .is_ok());
        let data = json!({
            "normal": 1,
            "你好": 2,
            "#special key": 3,
            "😂": 4,
            "me.dot.key": 5
        });
        let r0 = handlebars.render("t0", &data).ok().unwrap();
        assert!(r0.contains("normal: 1"));
        assert!(r0.contains("你好: 2"));
        assert!(r0.contains("#special key: 3"));
        assert!(r0.contains("😂: 4"));
        assert!(r0.contains("me.dot.key: 5"));
    }

    #[test]
    fn test_base_path_after_each() {
        let mut handlebars = Registry::new();
        assert!(handlebars
            .register_template_string("t0", "{{#each a}}{{this}}{{/each}} {{b}}")
            .is_ok());
        let data = json!({
            "a": [1, 2, 3, 4],
            "b": "good",
        });

        let r0 = handlebars.render("t0", &data).ok().unwrap();

        assert_eq!("1234 good", r0);
    }
}
<commit_msg>(test) add failing case for else block in each helper [#344]<commit_after>use serde_json::value::Value as Json;

use crate::block::{BlockContext, BlockParams};
use crate::context::Context;
use crate::error::RenderError;
use crate::helpers::{HelperDef, HelperResult};
use crate::json::value::{to_json, JsonTruthy};
use crate::output::Output;
use crate::registry::Registry;
use crate::render::{Helper, RenderContext, Renderable};
use crate::util::copy_on_push_vec;

#[derive(Clone, Copy)]
pub struct EachHelper;

impl HelperDef for EachHelper {
    fn call<'reg: 'rc, 'rc>(
        &self,
        h: &Helper<'reg, 'rc>,
        r: &'reg Registry<'reg>,
        ctx: &'rc Context,
        rc: &mut RenderContext<'reg, 'rc>,
        out: &mut dyn Output,
    ) -> HelperResult {
        let value = h
            .param(0)
            .ok_or_else(|| RenderError::new("Param not found for helper \"each\""))?;

        let template = h.template();

        match template {
            Some(t) => {
                let mut block_context = BlockContext::new();

                if let Some(path) = value.context_path() {
                    *block_context.base_path_mut() = path.to_vec();
                    // TODO: disable base value support for now.
                    // block_context.set_base_value(value.value());
                }

                rc.push_block(block_context);

                let rendered = match (value.value().is_truthy(false), value.value()) {
                    (true, &Json::Array(ref list)) => {
                        let len = list.len();

                        let array_path = value.context_path();

                        for (i, _) in list.iter().enumerate().take(len) {
                            if let Some(ref mut block) = rc.block_mut() {
                                let is_first = i == 0usize;
                                let is_last = i == len - 1;

                                block.set_local_var("@first".to_string(), to_json(is_first));
                                block.set_local_var("@last".to_string(), to_json(is_last));
                                block.set_local_var("@index".to_string(), to_json(i));
                                if let Some(ref p) = array_path {
                                    if is_first {
                                        *block.base_path_mut() = copy_on_push_vec(p, i.to_string());
                                    } else if let Some(ptr) = block.base_path_mut().last_mut() {
                                        *ptr = i.to_string();
                                    }
                                }

                                // TODO: base value

                                if let Some(bp_val) = h.block_param() {
                                    let mut params = BlockParams::new();
                                    params.add_path(bp_val, Vec::with_capacity(0))?;

                                    block.set_block_params(params);
                                } else if let Some((bp_val, bp_index)) = h.block_param_pair() {
                                    let mut params = BlockParams::new();
                                    params.add_path(bp_val, Vec::with_capacity(0))?;
                                    params.add_value(bp_index, to_json(i))?;

                                    block.set_block_params(params);
                                }
                            }

                            t.render(r, ctx, rc, out)?;
                        }

                        Ok(())
                    }
                    (true, &Json::Object(ref obj)) => {
                        let mut is_first = true;
                        let obj_path = value.context_path();

                        for (k, _) in obj.iter() {
                            if let Some(ref mut block) = rc.block_mut() {
                                block.set_local_var("@first".to_string(), to_json(is_first));
                                block.set_local_var("@key".to_string(), to_json(k));

                                if let Some(ref p) = obj_path {
                                    if is_first {
                                        *block.base_path_mut() = copy_on_push_vec(p, k.clone());
                                    } else if let Some(ptr) = block.base_path_mut().last_mut() {
                                        *ptr = k.clone();
                                    }
                                }

                                // TODO

                                if let Some(bp_val) = h.block_param() {
                                    let mut params = BlockParams::new();
                                    params.add_path(bp_val, Vec::with_capacity(0))?;

                                    block.set_block_params(params);
                                } else if let Some((bp_val, bp_key)) = h.block_param_pair() {
                                    let mut params = BlockParams::new();
                                    params.add_path(bp_val, Vec::with_capacity(0))?;
                                    params.add_value(bp_key, to_json(&k))?;

                                    block.set_block_params(params);
                                }
                            }

                            t.render(r, ctx, rc, out)?;

                            if is_first {
                                is_first = false;
                            }
                        }
                        Ok(())
                    }
                    (false, _) => {
                        if let Some(else_template) = h.inverse() {
                            else_template.render(r, ctx, rc, out)?;
                        }
                        Ok(())
                    }
                    _ => Err(RenderError::new(format!(
                        "Param type is not iterable: {:?}",
                        value.value()
                    ))),
                };

                rc.pop_block();
                rendered
            }
            None => Ok(()),
        }
    }
}

pub static EACH_HELPER: EachHelper = EachHelper;

#[cfg(test)]
mod test {
    use crate::json::value::to_json;
    use crate::registry::Registry;
    use serde_json::value::Value as Json;
    use std::collections::BTreeMap;
    use std::str::FromStr;

    #[test]
    fn test_each() {
        let mut handlebars = Registry::new();
        assert!(handlebars
            .register_template_string(
                "t0",
                "{{#each this}}{{@first}}|{{@last}}|{{@index}}:{{this}}|{{/each}}"
            )
            .is_ok());
        assert!(handlebars
            .register_template_string("t1", "{{#each this}}{{@first}}|{{@key}}:{{this}}|{{/each}}")
            .is_ok());

        let r0 = handlebars.render("t0", &vec![1u16, 2u16, 3u16]);
        assert_eq!(
            r0.ok().unwrap(),
            "true|false|0:1|false|false|1:2|false|true|2:3|".to_string()
        );

        let mut m: BTreeMap<String, u16> = BTreeMap::new();
        m.insert("ftp".to_string(), 21);
        m.insert("http".to_string(), 80);
        let r1 = handlebars.render("t1", &m);
        assert_eq!(r1.ok().unwrap(), "true|ftp:21|false|http:80|".to_string());
    }

    #[test]
    fn test_each_with_parent() {
        let json_str = r#"{"a":{"a":99,"c":[{"d":100},{"d":200}]}}"#;

        let data = Json::from_str(json_str).unwrap();
        // println!("data: {}", data);
        let mut handlebars = Registry::new();

        // previously, to access the parent in an each block,
        // a user would need to specify ../../b, as the path
        // that is computed includes the array index: ./a.c.[0]
        assert!(handlebars
            .register_template_string("t0", "{{#each a.c}} d={{d}} b={{../a.a}} {{/each}}")
            .is_ok());

        let r1 = handlebars.render("t0", &data);
        assert_eq!(r1.ok().unwrap(), " d=100 b=99  d=200 b=99 ".to_string());
    }

    #[test]
    fn test_nested_each_with_parent() {
        let json_str = r#"{"a": [{"b": [{"d": 100}], "c": 200}]}"#;

        let data = Json::from_str(json_str).unwrap();
        let mut handlebars = Registry::new();
        assert!(handlebars
            .register_template_string(
                "t0",
                "{{#each a}}{{#each b}}{{d}}:{{../c}}{{/each}}{{/each}}"
            )
            .is_ok());

        let r1 = handlebars.render("t0", &data);
        assert_eq!(r1.ok().unwrap(), "100:200".to_string());
    }

    #[test]
    fn test_nested_each() {
        let json_str = r#"{"a": [{"b": true}], "b": [[1, 2, 3],[4, 5]]}"#;

        let data = Json::from_str(json_str).unwrap();

        let mut handlebars = Registry::new();
        assert!(handlebars
            .register_template_string(
                "t0",
                "{{#each b}}{{#if ../a}}{{#each this}}{{this}}{{/each}}{{/if}}{{/each}}"
            )
            .is_ok());

        let r1 = handlebars.render("t0", &data);
        assert_eq!(r1.ok().unwrap(), "12345".to_string());
    }

    #[test]
    fn test_nested_array() {
        let mut handlebars = Registry::new();
        assert!(handlebars
            .register_template_string("t0", "{{#each this.[0]}}{{this}}{{/each}}")
            .is_ok());

        let r0 = handlebars.render("t0", &(vec![vec![1, 2, 3]]));

        assert_eq!(r0.ok().unwrap(), "123".to_string());
    }

    #[test]
    fn test_empty_key() {
        let mut handlebars = Registry::new();
        assert!(handlebars
            .register_template_string("t0", "{{#each this}}{{@key}}-{{value}}\n{{/each}}")
            .is_ok());

        let r0 = handlebars
            .render(
                "t0",
                &json!({
                    "foo": {
                        "value": "bar"
                    },
                    "": {
                        "value": "baz"
                    }
                }),
            )
            .unwrap();

        let mut r0_sp: Vec<_> = r0.split('\n').collect();
        r0_sp.sort();

        assert_eq!(r0_sp, vec!["", "-baz", "foo-bar"]);
    }

    #[test]
    fn test_each_else() {
        let mut handlebars = Registry::new();
        assert!(handlebars
            .register_template_string("t0", "{{#each a}}1{{else}}empty{{/each}}")
            .is_ok());
        let m1 = btreemap! {
            "a".to_string() => Vec::<String>::new(),
        };
        let r0 = handlebars.render("t0", &m1).unwrap();
        assert_eq!(r0, "empty");

        let m2 = btreemap! {
            "b".to_string() => Vec::<String>::new()
        };
        let r1 = handlebars.render("t0", &m2).unwrap();
        assert_eq!(r1, "empty");
    }

    #[test]
    fn test_block_param() {
        let mut handlebars = Registry::new();
        assert!(handlebars
            .register_template_string("t0", "{{#each a as |i|}}{{i}}{{/each}}")
            .is_ok());
        let m1 = btreemap! {
            "a".to_string() => vec![1,2,3,4,5]
        };
        let r0 = handlebars.render("t0", &m1).unwrap();
        assert_eq!(r0, "12345");
    }

    #[test]
    fn test_each_object_block_param() {
        let mut handlebars = Registry::new();
        let template = "{{#each this as |v k|}}\
                        {{#with k as |inner_k|}}{{inner_k}}{{/with}}:{{v}}|\
                        {{/each}}";
        assert!(handlebars.register_template_string("t0", template).is_ok());

        let m = btreemap! {
            "ftp".to_string() => 21,
            "http".to_string() => 80
        };
        let r0 = handlebars.render("t0", &m);
        assert_eq!(r0.ok().unwrap(), "ftp:21|http:80|".to_string());
    }

    #[test]
    fn test_each_object_block_param2() {
        let mut handlebars = Registry::new();
        let template = "{{#each this as |v k|}}\
                        {{#with v as |inner_v|}}{{k}}:{{inner_v}}{{/with}}|\
                        {{/each}}";

        assert!(handlebars.register_template_string("t0", template).is_ok());

        let m = btreemap! {
            "ftp".to_string() => 21,
            "http".to_string() => 80
        };
        let r0 = handlebars.render("t0", &m);
        assert_eq!(r0.ok().unwrap(), "ftp:21|http:80|".to_string());
    }

    fn test_nested_each_with_path_ups() {
        let mut handlebars = Registry::new();
        assert!(handlebars
            .register_template_string(
                "t0",
                "{{#each a.b}}{{#each c}}{{../../d}}{{/each}}{{/each}}"
            )
            .is_ok());

        let data = btreemap! {
            "a".to_string() => to_json(&btreemap! {
                "b".to_string() => vec![btreemap!{"c".to_string() => vec![1]}]
            }),
            "d".to_string() => to_json(&1)
        };

        let r0 = handlebars.render("t0", &data);
        assert_eq!(r0.ok().unwrap(), "1".to_string());
    }

    #[test]
    fn test_nested_each_with_path_up_this() {
        let mut handlebars = Registry::new();
        let template = "{{#each variant}}{{#each ../typearg}}\
                        {{#if @first}}template<{{/if}}{{this}}{{#if @last}}>{{else}},{{/if}}\
                        {{/each}}{{/each}}";
        assert!(handlebars.register_template_string("t0", template).is_ok());
        let data = btreemap! {
            "typearg".to_string() => vec!["T".to_string()],
            "variant".to_string() => vec!["1".to_string(), "2".to_string()]
        };
        let r0 = handlebars.render("t0", &data);
        assert_eq!(r0.ok().unwrap(), "template<T>template<T>".to_string());
    }

    #[test]
    fn test_key_iteration_with_unicode() {
        let mut handlebars = Registry::new();
        assert!(handlebars
            .register_template_string("t0", "{{#each this}}{{@key}}: {{this}}\n{{/each}}")
            .is_ok());
        let data = json!({
            "normal": 1,
            "你好": 2,
            "#special key": 3,
            "😂": 4,
            "me.dot.key": 5
        });
        let r0 = handlebars.render("t0", &data).ok().unwrap();
        assert!(r0.contains("normal: 1"));
        assert!(r0.contains("你好: 2"));
        assert!(r0.contains("#special key: 3"));
        assert!(r0.contains("😂: 4"));
        assert!(r0.contains("me.dot.key: 5"));
    }

    #[test]
    fn test_base_path_after_each() {
        let mut handlebars = Registry::new();
        assert!(handlebars
            .register_template_string("t0", "{{#each a}}{{this}}{{/each}} {{b}}")
            .is_ok());
        let data = json!({
            "a": [1, 2, 3, 4],
            "b": "good",
        });

        let r0 = handlebars.render("t0", &data).ok().unwrap();

        assert_eq!("1234 good", r0);
    }

    #[test]
    fn test_else_context() {
        let reg = Registry::new();
        let template = "{{#each list}}A{{else}}{{foo}}{{/each}}";
        let input = json!({"list": [], "foo": "bar"});
        let rendered = reg.render_template(template, &input).unwrap();
        assert_eq!("bar", rendered);
    }
}
<|endoftext|>


<commit_before><commit_msg>test: add logging config<commit_after><|endoftext|>


<commit_before>use std::f32::to_str_digits;
use std::num::atan;

fn get_pi(num :uint) ->~str {
    to_str_digits(16f32*atan(1f32/5f32)-4f32*atan(1f32/239f32), num)
}

#[test]
fn get_pi_test() {
    assert!(get_pi(12) == to_str_digits(3.1415927410125732421875f32, 12), "{} {}", get_pi(12), 3.141592741013);
}

#[test]
fn get_pi_test2() {
    assert!(get_pi(0) == to_str_digits(3.1415927410125732421875f32, 0), "{} {}", get_pi(0), 3);
}

fn main() {
    let pi = get_pi(12);
    println!("{}", pi);
}<commit_msg>Added Comment for how to compile Tests.<commit_after>use std::f32::to_str_digits;
use std::num::atan;

fn get_pi(num :uint) ->~str {
    to_str_digits(16f32*atan(1f32/5f32)-4f32*atan(1f32/239f32), num)
}

/* Build Tests with rustc --test *file*.rs*/
#[test]
fn get_pi_test() {
    assert!(get_pi(12) == to_str_digits(3.1415927410125732421875f32, 12), "{} {}", get_pi(12), 3.141592741013);
}

#[test]
fn get_pi_test2() {
    assert!(get_pi(0) == to_str_digits(3.1415927410125732421875f32, 0), "{} {}", get_pi(0), 3);
}

fn main() {
    let pi = get_pi(12);
    println!("{}", pi);
}
<|endoftext|>


<commit_before>use {
    BackEnd,
    Borrowed,
    Clear,
    Draw,
    Field,
    Image,
    Value,
};
use triangulation::{
    rect_tri_list_xy_f32,
    rect_tri_list_rgba_f32,
    rect_tri_list_uv_f32,
};
use internal::{
    CanColor,
    CanSourceRectangle,
    CanTransform,
    CanViewTransform,
    Color,
    HasColor,
    HasSourceRectangle,
    HasTransform,
    HasViewTransform,
    SourceRectangle,
    Matrix2d,
};

/// An image rectangle context.
pub struct ImageColorContext<'a, 'b, I> {
    /// View transformation.
    pub view: Field<'a, Matrix2d>,
    /// Current transformation.
    pub transform: Field<'a, Matrix2d>,
    /// Current image.
    pub image: Field<'a, &'b I>,
    /// Current source rectangle.
    pub source_rect: Field<'a, SourceRectangle>,
    /// Current color.
    pub color: Field<'a, Color>,
}

impl<'a, 'b, I> Clone for ImageColorContext<'a, 'b, I> {
    #[inline(always)]
    fn clone(&self) -> ImageColorContext<'static, 'b, I> {
        ImageColorContext {
            view: Value(*self.view.get()),
            transform: Value(*self.transform.get()),
            image: Value(*self.image.get()),
            source_rect: Value(*self.source_rect.get()),
            color: Value(*self.color.get()),
        }
    }
}

impl<'a, 'b, I> HasTransform<'a, Matrix2d> for ImageColorContext<'a, 'b, I> {
    #[inline(always)]
    fn get_transform(&'a self) -> &'a Matrix2d {
        self.transform.get()
    }
}

impl<'a, 'b, I> CanTransform<'a, ImageColorContext<'a, 'b, I>, Matrix2d>
for ImageColorContext<'a, 'b, I> {
    #[inline(always)]
    fn transform(&'a self, value: Matrix2d) -> ImageColorContext<'a, 'b, I> {
        ImageColorContext {
            view: Borrowed(self.view.get()),
            transform: Value(value),
            image: Borrowed(self.image.get()),
            source_rect: Borrowed(self.source_rect.get()),
            color: Borrowed(self.color.get()),
        }
    }
}

impl<'a, 'b, I> HasViewTransform<'a, Matrix2d> for ImageColorContext<'a, 'b, I> {
    #[inline(always)]
    fn get_view_transform(&'a self) -> &'a Matrix2d {
        self.view.get()
    }
}

impl<'a, 'b, I> CanViewTransform<'a, ImageColorContext<'a, 'b, I>, Matrix2d>
for ImageColorContext<'a, 'b, I> {
    #[inline(always)]
    fn view_transform(&'a self, value: Matrix2d) -> ImageColorContext<'a, 'b, I> {
        ImageColorContext {
            view: Value(value),
            transform: Borrowed(self.transform.get()),
            image: Borrowed(self.image.get()),
            source_rect: Borrowed(self.source_rect.get()),
            color: Borrowed(self.color.get()),
        }
    }
}

impl<'a, 'b, I> HasColor<'a, Color> for ImageColorContext<'a, 'b, I> {
    #[inline(always)]
    fn get_color(&'a self) -> &'a Color {
        self.color.get()
    }
}

impl<'a, 'b, I> CanColor<'a, ImageColorContext<'a, 'b, I>, Color> for ImageColorContext<'a, 'b, I> {
    #[inline(always)]
    fn color(&'a self, value: Color) -> ImageColorContext<'a, 'b, I> {
        ImageColorContext {
            view: Borrowed(self.view.get()),
            transform: Borrowed(self.transform.get()),
            color: Value(value),
            image: Borrowed(self.image.get()),
            source_rect: Borrowed(self.source_rect.get()),
        }
    }
}

impl<'a, 'b, I> HasSourceRectangle<'a, SourceRectangle>
for ImageColorContext<'a, 'b, I> {
    #[inline(always)]
    fn get_source_rectangle(&'a self) -> &'a SourceRectangle {
        self.source_rect.get()
    }
}

impl<'a, 'b, I> CanSourceRectangle<'a, ImageColorContext<'a, 'b, I>, SourceRectangle>
for ImageColorContext<'a, 'b, I> {
    #[inline(always)]
    fn source_rectangle(&'a self, source_rect: SourceRectangle) -> ImageColorContext<'a, 'b, I> {
        ImageColorContext {
            view: Borrowed(self.view.get()),
            transform: Borrowed(self.transform.get()),
            image: Borrowed(self.image.get()),
            source_rect: Value(source_rect),
            color: Borrowed(self.color.get()),
        }
    }
}

impl<'a, 'b, B: BackEnd<I>, I: Image> Draw<'a, B, I> for ImageColorContext<'a, 'b, I> {
    #[inline(always)]
    fn draw(&'a self, back_end: &mut B) {
        if back_end.supports_single_texture()
        && back_end.supports_tri_list_xy_f32_rgba_f32_uv_f32() {
            let color = self.color.get();
            let &texture = self.image.get();
            let source_rect = self.source_rect.get();
            let rect = [0.0, 0.0, source_rect[1] as f64, source_rect[2] as f64];
            // Complete transparency does not need to be rendered.
            if color[3] == 0.0 { return; }
            // Turn on alpha blending if not completely opaque or if the texture has alpha channel.
            let needs_alpha = color[3] != 1.0 || back_end.has_texture_alpha(texture);
            if needs_alpha { back_end.enable_alpha_blend(); }
            back_end.enable_single_texture(texture);
            back_end.tri_list_xy_f32_rgba_f32_uv_f32(
                rect_tri_list_xy_f32(*self.transform.get(), rect),
                rect_tri_list_rgba_f32(*color),
                rect_tri_list_uv_f32(texture, *source_rect)
            );
            back_end.disable_single_texture();
            if needs_alpha { back_end.disable_alpha_blend(); }
        } else {
            unimplemented!();
        }
    }
}

impl<'a, 'b, B: BackEnd<I>, I: Image> Clear<B, I> for ImageColorContext<'a, 'b, I> {
    #[inline(always)]
    fn clear(&self, back_end: &mut B) {
        if back_end.supports_clear_rgba() {
            let color = self.color.get();
            back_end.clear_rgba(color[0], color[1], color[2], color[3]);
        }
    }
}

<commit_msg>Implemented `AddRectangle` for `ImageRectangleColorContext`<commit_after>use {
    AddRectangle,
    BackEnd,
    Borrowed,
    Clear,
    Draw,
    Field,
    Image,
    ImageRectangleColorContext,
    Value,
};
use triangulation::{
    rect_tri_list_xy_f32,
    rect_tri_list_rgba_f32,
    rect_tri_list_uv_f32,
};
use internal::{
    CanColor,
    CanSourceRectangle,
    CanTransform,
    CanViewTransform,
    Color,
    HasColor,
    HasSourceRectangle,
    HasTransform,
    HasViewTransform,
    SourceRectangle,
    Matrix2d,
    Scalar,
};

/// An image rectangle context.
pub struct ImageColorContext<'a, 'b, I> {
    /// View transformation.
    pub view: Field<'a, Matrix2d>,
    /// Current transformation.
    pub transform: Field<'a, Matrix2d>,
    /// Current image.
    pub image: Field<'a, &'b I>,
    /// Current source rectangle.
    pub source_rect: Field<'a, SourceRectangle>,
    /// Current color.
    pub color: Field<'a, Color>,
}

impl<'a, 'b, I> Clone for ImageColorContext<'a, 'b, I> {
    #[inline(always)]
    fn clone(&self) -> ImageColorContext<'static, 'b, I> {
        ImageColorContext {
            view: Value(*self.view.get()),
            transform: Value(*self.transform.get()),
            image: Value(*self.image.get()),
            source_rect: Value(*self.source_rect.get()),
            color: Value(*self.color.get()),
        }
    }
}

impl<'a, 'b, I> HasTransform<'a, Matrix2d> for ImageColorContext<'a, 'b, I> {
    #[inline(always)]
    fn get_transform(&'a self) -> &'a Matrix2d {
        self.transform.get()
    }
}

impl<'a, 'b, I> CanTransform<'a, ImageColorContext<'a, 'b, I>, Matrix2d>
for ImageColorContext<'a, 'b, I> {
    #[inline(always)]
    fn transform(&'a self, value: Matrix2d) -> ImageColorContext<'a, 'b, I> {
        ImageColorContext {
            view: Borrowed(self.view.get()),
            transform: Value(value),
            image: Borrowed(self.image.get()),
            source_rect: Borrowed(self.source_rect.get()),
            color: Borrowed(self.color.get()),
        }
    }
}

impl<'a, 'b, I> HasViewTransform<'a, Matrix2d> for ImageColorContext<'a, 'b, I> {
    #[inline(always)]
    fn get_view_transform(&'a self) -> &'a Matrix2d {
        self.view.get()
    }
}

impl<'a, 'b, I> CanViewTransform<'a, ImageColorContext<'a, 'b, I>, Matrix2d>
for ImageColorContext<'a, 'b, I> {
    #[inline(always)]
    fn view_transform(&'a self, value: Matrix2d) -> ImageColorContext<'a, 'b, I> {
        ImageColorContext {
            view: Value(value),
            transform: Borrowed(self.transform.get()),
            image: Borrowed(self.image.get()),
            source_rect: Borrowed(self.source_rect.get()),
            color: Borrowed(self.color.get()),
        }
    }
}

impl<'a, 'b, I> HasColor<'a, Color> for ImageColorContext<'a, 'b, I> {
    #[inline(always)]
    fn get_color(&'a self) -> &'a Color {
        self.color.get()
    }
}

impl<'a, 'b, I> CanColor<'a, ImageColorContext<'a, 'b, I>, Color> for ImageColorContext<'a, 'b, I> {
    #[inline(always)]
    fn color(&'a self, value: Color) -> ImageColorContext<'a, 'b, I> {
        ImageColorContext {
            view: Borrowed(self.view.get()),
            transform: Borrowed(self.transform.get()),
            color: Value(value),
            image: Borrowed(self.image.get()),
            source_rect: Borrowed(self.source_rect.get()),
        }
    }
}

impl<'a, 'b, I> HasSourceRectangle<'a, SourceRectangle>
for ImageColorContext<'a, 'b, I> {
    #[inline(always)]
    fn get_source_rectangle(&'a self) -> &'a SourceRectangle {
        self.source_rect.get()
    }
}

impl<'a, 'b, I> CanSourceRectangle<'a, ImageColorContext<'a, 'b, I>, SourceRectangle>
for ImageColorContext<'a, 'b, I> {
    #[inline(always)]
    fn source_rectangle(&'a self, source_rect: SourceRectangle) -> ImageColorContext<'a, 'b, I> {
        ImageColorContext {
            view: Borrowed(self.view.get()),
            transform: Borrowed(self.transform.get()),
            image: Borrowed(self.image.get()),
            source_rect: Value(source_rect),
            color: Borrowed(self.color.get()),
        }
    }
}

impl<'a, 'b, I> AddRectangle<'a, ImageRectangleColorContext<'a, 'b, I>>
for ImageColorContext<'a, 'b, I> {
    #[inline(always)]
    fn rect(&'a self, x: Scalar, y: Scalar, w: Scalar, h: Scalar) -> ImageRectangleColorContext<'a, 'b, I> {
        ImageRectangleColorContext {
            view: Borrowed(self.view.get()),
            transform: Borrowed(self.transform.get()),
            rect: Value([x, y, w, h]),
            image: Borrowed(self.image.get()),
            source_rect: Borrowed(self.source_rect.get()),
            color: Borrowed(self.color.get()),
        }
    }
}

impl<'a, 'b, B: BackEnd<I>, I: Image> Draw<'a, B, I> for ImageColorContext<'a, 'b, I> {
    #[inline(always)]
    fn draw(&'a self, back_end: &mut B) {
        if back_end.supports_single_texture()
        && back_end.supports_tri_list_xy_f32_rgba_f32_uv_f32() {
            let color = self.color.get();
            let &texture = self.image.get();
            let source_rect = self.source_rect.get();
            let rect = [0.0, 0.0, source_rect[1] as f64, source_rect[2] as f64];
            // Complete transparency does not need to be rendered.
            if color[3] == 0.0 { return; }
            // Turn on alpha blending if not completely opaque or if the texture has alpha channel.
            let needs_alpha = color[3] != 1.0 || back_end.has_texture_alpha(texture);
            if needs_alpha { back_end.enable_alpha_blend(); }
            back_end.enable_single_texture(texture);
            back_end.tri_list_xy_f32_rgba_f32_uv_f32(
                rect_tri_list_xy_f32(*self.transform.get(), rect),
                rect_tri_list_rgba_f32(*color),
                rect_tri_list_uv_f32(texture, *source_rect)
            );
            back_end.disable_single_texture();
            if needs_alpha { back_end.disable_alpha_blend(); }
        } else {
            unimplemented!();
        }
    }
}

impl<'a, 'b, B: BackEnd<I>, I: Image> Clear<B, I> for ImageColorContext<'a, 'b, I> {
    #[inline(always)]
    fn clear(&self, back_end: &mut B) {
        if back_end.supports_clear_rgba() {
            let color = self.color.get();
            back_end.clear_rgba(color[0], color[1], color[2], color[3]);
        }
    }
}

<|endoftext|>


<commit_before>use base_db::FileId;
use hir::{EnumVariant, Module, ModuleDef, Name};
use ide_db::{defs::Definition, search::Reference, RootDatabase};
use rustc_hash::FxHashSet;
use syntax::{
    algo::find_node_at_offset,
    ast::{self, edit::IndentLevel, ArgListOwner, AstNode, NameOwner, VisibilityOwner},
    SourceFile, TextRange, TextSize,
};

use crate::{
    assist_context::AssistBuilder, utils::insert_use_statement, AssistContext, AssistId,
    AssistKind, Assists,
};

// Assist: extract_struct_from_enum_variant
//
// Extracts a struct from enum variant.
//
// ```
// enum A { <|>One(u32, u32) }
// ```
// ->
// ```
// struct One(pub u32, pub u32);
//
// enum A { One(One) }
// ```
pub(crate) fn extract_struct_from_enum_variant(
    acc: &mut Assists,
    ctx: &AssistContext,
) -> Option<()> {
    let variant = ctx.find_node_at_offset::<ast::Variant>()?;
    let field_list = match variant.kind() {
        ast::StructKind::Tuple(field_list) => field_list,
        _ => return None,
    };
    let variant_name = variant.name()?.to_string();
    let variant_hir = ctx.sema.to_def(&variant)?;
    if existing_struct_def(ctx.db(), &variant_name, &variant_hir) {
        return None;
    }
    let enum_ast = variant.parent_enum();
    let visibility = enum_ast.visibility();
    let enum_hir = ctx.sema.to_def(&enum_ast)?;
    let variant_hir_name = variant_hir.name(ctx.db());
    let enum_module_def = ModuleDef::from(enum_hir);
    let current_module = enum_hir.module(ctx.db());
    let target = variant.syntax().text_range();
    acc.add(
        AssistId("extract_struct_from_enum_variant", AssistKind::RefactorRewrite),
        "Extract struct from enum variant",
        target,
        |builder| {
            let definition = Definition::ModuleDef(ModuleDef::EnumVariant(variant_hir));
            let res = definition.usages(&ctx.sema).all();
            let start_offset = variant.parent_enum().syntax().text_range().start();
            let mut visited_modules_set = FxHashSet::default();
            visited_modules_set.insert(current_module);
            for reference in res {
                let source_file = ctx.sema.parse(reference.file_range.file_id);
                update_reference(
                    ctx,
                    builder,
                    reference,
                    &source_file,
                    &enum_module_def,
                    &variant_hir_name,
                    &mut visited_modules_set,
                );
            }
            extract_struct_def(
                builder,
                &enum_ast,
                &variant_name,
                &field_list.to_string(),
                start_offset,
                ctx.frange.file_id,
                &visibility,
            );
            let list_range = field_list.syntax().text_range();
            update_variant(builder, &variant_name, ctx.frange.file_id, list_range);
        },
    )
}

fn existing_struct_def(db: &RootDatabase, variant_name: &str, variant: &EnumVariant) -> bool {
    variant
        .parent_enum(db)
        .module(db)
        .scope(db, None)
        .into_iter()
        .any(|(name, _)| name.to_string() == variant_name.to_string())
}

fn insert_import(
    ctx: &AssistContext,
    builder: &mut AssistBuilder,
    path: &ast::PathExpr,
    module: &Module,
    enum_module_def: &ModuleDef,
    variant_hir_name: &Name,
) -> Option<()> {
    let db = ctx.db();
    let mod_path = module.find_use_path(db, enum_module_def.clone());
    if let Some(mut mod_path) = mod_path {
        mod_path.segments.pop();
        mod_path.segments.push(variant_hir_name.clone());
        insert_use_statement(
            path.syntax(),
            &mod_path.to_string(),
            ctx,
            builder.text_edit_builder(),
        );
    }
    Some(())
}

// FIXME: this should use strongly-typed `make`, rather than string manipulation.
fn extract_struct_def(
    builder: &mut AssistBuilder,
    enum_: &ast::Enum,
    variant_name: &str,
    variant_list: &str,
    start_offset: TextSize,
    file_id: FileId,
    visibility: &Option<ast::Visibility>,
) -> Option<()> {
    let visibility_string = if let Some(visibility) = visibility {
        format!("{} ", visibility.to_string())
    } else {
        "".to_string()
    };
    let indent = IndentLevel::from_node(enum_.syntax());
    let struct_def = format!(
        r#"{}struct {}{};

{}"#,
        visibility_string,
        variant_name,
        list_with_visibility(variant_list),
        indent
    );
    builder.edit_file(file_id);
    builder.insert(start_offset, struct_def);
    Some(())
}

fn update_variant(
    builder: &mut AssistBuilder,
    variant_name: &str,
    file_id: FileId,
    list_range: TextRange,
) -> Option<()> {
    let inside_variant_range = TextRange::new(
        list_range.start().checked_add(TextSize::from(1))?,
        list_range.end().checked_sub(TextSize::from(1))?,
    );
    builder.edit_file(file_id);
    builder.replace(inside_variant_range, variant_name);
    Some(())
}

fn update_reference(
    ctx: &AssistContext,
    builder: &mut AssistBuilder,
    reference: Reference,
    source_file: &SourceFile,
    enum_module_def: &ModuleDef,
    variant_hir_name: &Name,
    visited_modules_set: &mut FxHashSet<Module>,
) -> Option<()> {
    let path_expr: ast::PathExpr = find_node_at_offset::<ast::PathExpr>(
        source_file.syntax(),
        reference.file_range.range.start(),
    )?;
    let call = path_expr.syntax().parent().and_then(ast::CallExpr::cast)?;
    let list = call.arg_list()?;
    let segment = path_expr.path()?.segment()?;
    let module = ctx.sema.scope(&path_expr.syntax()).module()?;
    let list_range = list.syntax().text_range();
    let inside_list_range = TextRange::new(
        list_range.start().checked_add(TextSize::from(1))?,
        list_range.end().checked_sub(TextSize::from(1))?,
    );
    builder.edit_file(reference.file_range.file_id);
    if !visited_modules_set.contains(&module) {
        if insert_import(ctx, builder, &path_expr, &module, enum_module_def, variant_hir_name)
            .is_some()
        {
            visited_modules_set.insert(module);
        }
    }
    builder.replace(inside_list_range, format!("{}{}", segment, list));
    Some(())
}

fn list_with_visibility(list: &str) -> String {
    list.split(',')
        .map(|part| {
            let index = if part.chars().next().unwrap() == '(' { 1usize } else { 0 };
            let mut mod_part = part.trim().to_string();
            mod_part.insert_str(index, "pub ");
            mod_part
        })
        .collect::<Vec<String>>()
        .join(", ")
}

#[cfg(test)]
mod tests {

    use crate::{
        tests::{check_assist, check_assist_not_applicable},
        utils::FamousDefs,
    };

    use super::*;

    #[test]
    fn test_extract_struct_several_fields() {
        check_assist(
            extract_struct_from_enum_variant,
            "enum A { <|>One(u32, u32) }",
            r#"struct One(pub u32, pub u32);

enum A { One(One) }"#,
        );
    }

    #[test]
    fn test_extract_struct_one_field() {
        check_assist(
            extract_struct_from_enum_variant,
            "enum A { <|>One(u32) }",
            r#"struct One(pub u32);

enum A { One(One) }"#,
        );
    }

    #[test]
    fn test_extract_struct_pub_visibility() {
        check_assist(
            extract_struct_from_enum_variant,
            "pub enum A { <|>One(u32, u32) }",
            r#"pub struct One(pub u32, pub u32);

pub enum A { One(One) }"#,
        );
    }

    #[test]
    fn test_extract_struct_with_complex_imports() {
        check_assist(
            extract_struct_from_enum_variant,
            r#"mod my_mod {
    fn another_fn() {
        let m = my_other_mod::MyEnum::MyField(1, 1);
    }

    pub mod my_other_mod {
        fn another_fn() {
            let m = MyEnum::MyField(1, 1);
        }

        pub enum MyEnum {
            <|>MyField(u8, u8),
        }
    }
}

fn another_fn() {
    let m = my_mod::my_other_mod::MyEnum::MyField(1, 1);
}"#,
            r#"use my_mod::my_other_mod::MyField;

mod my_mod {
    use my_other_mod::MyField;

    fn another_fn() {
        let m = my_other_mod::MyEnum::MyField(MyField(1, 1));
    }

    pub mod my_other_mod {
        fn another_fn() {
            let m = MyEnum::MyField(MyField(1, 1));
        }

        pub struct MyField(pub u8, pub u8);

        pub enum MyEnum {
            MyField(MyField),
        }
    }
}

fn another_fn() {
    let m = my_mod::my_other_mod::MyEnum::MyField(MyField(1, 1));
}"#,
        );
    }

    fn check_not_applicable(ra_fixture: &str) {
        let fixture =
            format!("//- /main.rs crate:main deps:core\n{}\n{}", ra_fixture, FamousDefs::FIXTURE);
        check_assist_not_applicable(extract_struct_from_enum_variant, &fixture)
    }

    #[test]
    fn test_extract_enum_not_applicable_for_element_with_no_fields() {
        check_not_applicable("enum A { <|>One }");
    }

    #[test]
    fn test_extract_enum_not_applicable_if_struct_exists() {
        check_not_applicable(
            r#"struct One;
        enum A { <|>One(u8) }"#,
        );
    }
}
<commit_msg>Minor<commit_after>use base_db::FileId;
use hir::{EnumVariant, Module, ModuleDef, Name};
use ide_db::{defs::Definition, search::Reference, RootDatabase};
use itertools::Itertools;
use rustc_hash::FxHashSet;
use syntax::{
    algo::find_node_at_offset,
    ast::{self, edit::IndentLevel, ArgListOwner, AstNode, NameOwner, VisibilityOwner},
    SourceFile, TextRange, TextSize,
};

use crate::{
    assist_context::AssistBuilder, utils::insert_use_statement, AssistContext, AssistId,
    AssistKind, Assists,
};

// Assist: extract_struct_from_enum_variant
//
// Extracts a struct from enum variant.
//
// ```
// enum A { <|>One(u32, u32) }
// ```
// ->
// ```
// struct One(pub u32, pub u32);
//
// enum A { One(One) }
// ```
pub(crate) fn extract_struct_from_enum_variant(
    acc: &mut Assists,
    ctx: &AssistContext,
) -> Option<()> {
    let variant = ctx.find_node_at_offset::<ast::Variant>()?;
    let field_list = match variant.kind() {
        ast::StructKind::Tuple(field_list) => field_list,
        _ => return None,
    };
    let variant_name = variant.name()?.to_string();
    let variant_hir = ctx.sema.to_def(&variant)?;
    if existing_struct_def(ctx.db(), &variant_name, &variant_hir) {
        return None;
    }
    let enum_ast = variant.parent_enum();
    let visibility = enum_ast.visibility();
    let enum_hir = ctx.sema.to_def(&enum_ast)?;
    let variant_hir_name = variant_hir.name(ctx.db());
    let enum_module_def = ModuleDef::from(enum_hir);
    let current_module = enum_hir.module(ctx.db());
    let target = variant.syntax().text_range();
    acc.add(
        AssistId("extract_struct_from_enum_variant", AssistKind::RefactorRewrite),
        "Extract struct from enum variant",
        target,
        |builder| {
            let definition = Definition::ModuleDef(ModuleDef::EnumVariant(variant_hir));
            let res = definition.usages(&ctx.sema).all();
            let start_offset = variant.parent_enum().syntax().text_range().start();
            let mut visited_modules_set = FxHashSet::default();
            visited_modules_set.insert(current_module);
            for reference in res {
                let source_file = ctx.sema.parse(reference.file_range.file_id);
                update_reference(
                    ctx,
                    builder,
                    reference,
                    &source_file,
                    &enum_module_def,
                    &variant_hir_name,
                    &mut visited_modules_set,
                );
            }
            extract_struct_def(
                builder,
                &enum_ast,
                &variant_name,
                &field_list.to_string(),
                start_offset,
                ctx.frange.file_id,
                &visibility,
            );
            let list_range = field_list.syntax().text_range();
            update_variant(builder, &variant_name, ctx.frange.file_id, list_range);
        },
    )
}

fn existing_struct_def(db: &RootDatabase, variant_name: &str, variant: &EnumVariant) -> bool {
    variant
        .parent_enum(db)
        .module(db)
        .scope(db, None)
        .into_iter()
        .any(|(name, _)| name.to_string() == variant_name.to_string())
}

fn insert_import(
    ctx: &AssistContext,
    builder: &mut AssistBuilder,
    path: &ast::PathExpr,
    module: &Module,
    enum_module_def: &ModuleDef,
    variant_hir_name: &Name,
) -> Option<()> {
    let db = ctx.db();
    let mod_path = module.find_use_path(db, enum_module_def.clone());
    if let Some(mut mod_path) = mod_path {
        mod_path.segments.pop();
        mod_path.segments.push(variant_hir_name.clone());
        insert_use_statement(
            path.syntax(),
            &mod_path.to_string(),
            ctx,
            builder.text_edit_builder(),
        );
    }
    Some(())
}

// FIXME: this should use strongly-typed `make`, rather than string manipulation.
fn extract_struct_def(
    builder: &mut AssistBuilder,
    enum_: &ast::Enum,
    variant_name: &str,
    variant_list: &str,
    start_offset: TextSize,
    file_id: FileId,
    visibility: &Option<ast::Visibility>,
) -> Option<()> {
    let visibility_string = if let Some(visibility) = visibility {
        format!("{} ", visibility.to_string())
    } else {
        "".to_string()
    };
    let indent = IndentLevel::from_node(enum_.syntax());
    let struct_def = format!(
        r#"{}struct {}{};

{}"#,
        visibility_string,
        variant_name,
        list_with_visibility(variant_list),
        indent
    );
    builder.edit_file(file_id);
    builder.insert(start_offset, struct_def);
    Some(())
}

fn update_variant(
    builder: &mut AssistBuilder,
    variant_name: &str,
    file_id: FileId,
    list_range: TextRange,
) -> Option<()> {
    let inside_variant_range = TextRange::new(
        list_range.start().checked_add(TextSize::from(1))?,
        list_range.end().checked_sub(TextSize::from(1))?,
    );
    builder.edit_file(file_id);
    builder.replace(inside_variant_range, variant_name);
    Some(())
}

fn update_reference(
    ctx: &AssistContext,
    builder: &mut AssistBuilder,
    reference: Reference,
    source_file: &SourceFile,
    enum_module_def: &ModuleDef,
    variant_hir_name: &Name,
    visited_modules_set: &mut FxHashSet<Module>,
) -> Option<()> {
    let path_expr: ast::PathExpr = find_node_at_offset::<ast::PathExpr>(
        source_file.syntax(),
        reference.file_range.range.start(),
    )?;
    let call = path_expr.syntax().parent().and_then(ast::CallExpr::cast)?;
    let list = call.arg_list()?;
    let segment = path_expr.path()?.segment()?;
    let module = ctx.sema.scope(&path_expr.syntax()).module()?;
    let list_range = list.syntax().text_range();
    let inside_list_range = TextRange::new(
        list_range.start().checked_add(TextSize::from(1))?,
        list_range.end().checked_sub(TextSize::from(1))?,
    );
    builder.edit_file(reference.file_range.file_id);
    if !visited_modules_set.contains(&module) {
        if insert_import(ctx, builder, &path_expr, &module, enum_module_def, variant_hir_name)
            .is_some()
        {
            visited_modules_set.insert(module);
        }
    }
    builder.replace(inside_list_range, format!("{}{}", segment, list));
    Some(())
}

fn list_with_visibility(list: &str) -> String {
    list.split(',')
        .map(|part| {
            let index = if part.chars().next().unwrap() == '(' { 1usize } else { 0 };
            let mut mod_part = part.trim().to_string();
            mod_part.insert_str(index, "pub ");
            mod_part
        })
        .join(", ")
}

#[cfg(test)]
mod tests {
    use crate::{
        tests::{check_assist, check_assist_not_applicable},
        utils::FamousDefs,
    };

    use super::*;

    #[test]
    fn test_extract_struct_several_fields() {
        check_assist(
            extract_struct_from_enum_variant,
            "enum A { <|>One(u32, u32) }",
            r#"struct One(pub u32, pub u32);

enum A { One(One) }"#,
        );
    }

    #[test]
    fn test_extract_struct_one_field() {
        check_assist(
            extract_struct_from_enum_variant,
            "enum A { <|>One(u32) }",
            r#"struct One(pub u32);

enum A { One(One) }"#,
        );
    }

    #[test]
    fn test_extract_struct_pub_visibility() {
        check_assist(
            extract_struct_from_enum_variant,
            "pub enum A { <|>One(u32, u32) }",
            r#"pub struct One(pub u32, pub u32);

pub enum A { One(One) }"#,
        );
    }

    #[test]
    fn test_extract_struct_with_complex_imports() {
        check_assist(
            extract_struct_from_enum_variant,
            r#"mod my_mod {
    fn another_fn() {
        let m = my_other_mod::MyEnum::MyField(1, 1);
    }

    pub mod my_other_mod {
        fn another_fn() {
            let m = MyEnum::MyField(1, 1);
        }

        pub enum MyEnum {
            <|>MyField(u8, u8),
        }
    }
}

fn another_fn() {
    let m = my_mod::my_other_mod::MyEnum::MyField(1, 1);
}"#,
            r#"use my_mod::my_other_mod::MyField;

mod my_mod {
    use my_other_mod::MyField;

    fn another_fn() {
        let m = my_other_mod::MyEnum::MyField(MyField(1, 1));
    }

    pub mod my_other_mod {
        fn another_fn() {
            let m = MyEnum::MyField(MyField(1, 1));
        }

        pub struct MyField(pub u8, pub u8);

        pub enum MyEnum {
            MyField(MyField),
        }
    }
}

fn another_fn() {
    let m = my_mod::my_other_mod::MyEnum::MyField(MyField(1, 1));
}"#,
        );
    }

    fn check_not_applicable(ra_fixture: &str) {
        let fixture =
            format!("//- /main.rs crate:main deps:core\n{}\n{}", ra_fixture, FamousDefs::FIXTURE);
        check_assist_not_applicable(extract_struct_from_enum_variant, &fixture)
    }

    #[test]
    fn test_extract_enum_not_applicable_for_element_with_no_fields() {
        check_not_applicable("enum A { <|>One }");
    }

    #[test]
    fn test_extract_enum_not_applicable_if_struct_exists() {
        check_not_applicable(
            r#"struct One;
        enum A { <|>One(u8) }"#,
        );
    }
}
<|endoftext|>