<commit_before>//! Generates random trees and tree-like plants.

use prop::plant::{Branch, ControlPoint};
use math::{ApproxEq, Basis3, Deg, Euler, InnerSpace, Point3f, Rad, Rotation, Rotation3, Vector1,
           Vector3f};
use rand::{Rand, Rng};
use rand::distributions::range::SampleRange;
use rand::distributions::{self, IndependentSample};
use std::ops::Range;

/// Parameters for the tree generator.
#[derive(Debug)]
struct Preset {
    name: &'static str,
    /// Diameter of the first branch we create (the trunk).
    trunk_diameter: Range<f32>,
    /// Trunk height. Note that branches going upward can increase plant height
    /// beyond this.
    trunk_height: Range<f32>,
    /// Trunk diameter at `trunk_height`. Should be smaller than
    /// `trunk_diameter`.
    trunk_diameter_top: Range<f32>,
    /// Trunk height at which we start creating branches.
    min_branch_height: Range<f32>,
    /// Range of subbranch diameters as a factor of the parent branch.
    branch_diameter_factor: Range<f32>,
    /// Range of subbranch angles in degrees.
    branch_angle_deg: Range<f32>,
    /// Factor by which to reduce segment diameter between consecutive points,
    /// sampled per branch.
    branch_diam_reduction: Range<f32>,
    /// Range of angles to use for rotation of new segments.
    ///
    /// The higher the angle, the more "twisted" branches appear.
    branch_segment_angle: Range<f32>,
    /// Range of segment counts for branches.
    ///
    /// Together with `branch_segment_length`, this defines the overall branch
    /// length.
    branch_segment_count: Range<u32>,
    /// Range of segment lengths to use for branches.
    ///
    /// Each segment will have a random length in this range.
    branch_segment_length: Range<f32>, // FIXME unused
}

static PRESETS: &'static [Preset] = &[Preset {
                                          name: "'Regular' Tree",
                                          trunk_diameter: 0.3..0.5,
                                          trunk_height: 3.0..6.0,
                                          trunk_diameter_top: 0.2..0.4,
                                          min_branch_height: 0.4..0.6,
                                          branch_diameter_factor: 0.3..0.5,
                                          branch_angle_deg: 70.0..110.0,
                                          branch_diam_reduction: 0.75..0.85,
                                          branch_segment_angle: 5.0..15.0,
                                          branch_segment_count: 3..10,
                                          branch_segment_length: 0.30..0.40,
                                      }];

pub struct TreeGen {
    preset: &'static Preset,
    /// Buffer for branches, filled as they're created.
    branches: Vec<Branch>,
}

impl TreeGen {
    /// Generates a new branch at the given start point.
    ///
    /// # Parameters
    ///
    /// * `rng`: The RNG to use
    /// * `start`: The point at which the branch should start
    /// * `dir`: Direction in which the branch should grow
    /// * `depth`: Recursion depth, used to limit the maximum branch depth
    /// * `parent_diam`: Diameter of the node of the parent branch where this
    ///   branch starts
    fn create_branch<R: Rng>(&mut self,
                             rng: &mut R,
                             start: Point3f,
                             dir: Vector3f,
                             depth: u32,
                             parent_diam: f32) {
        if depth > 5 {
            // Limit recursion
            return;
        }

        // Current normalized growing direction, variated slightly as segments are
        // generated
        let mut dir = dir.normalize();

        // Determine starting diameter of the new branch
        let mut diam = range_sample(&self.preset.branch_diameter_factor, rng) * parent_diam;
        // Determine how much segment diameter is reduced
        let diam_factor = range_sample(&self.preset.branch_diam_reduction, rng) * parent_diam;
        // How many segments should this branch get?
        let segment_count = range_sample(&self.preset.branch_segment_count, rng);

        let mut points = vec![ControlPoint {
                                  point: start,
                                  diameter: diam,
                              }];

        {
            let mut last = start;
            // Helper for adding a new point to this branch, which possibly grows a new
            // branch.
            let mut add_point = |dist, diam| {
                // First, get a random angle by which to variate this segment.
                let mut x_angle = range_sample(&self.preset.branch_segment_angle, rng);
                let mut y_angle = range_sample(&self.preset.branch_segment_angle, rng);

                // Invert sign with 50%, to mirror the specified range to the other side
                x_angle = if rng.gen() { -x_angle } else { x_angle };
                y_angle = if rng.gen() { -y_angle } else { y_angle };

                let rotation = Basis3::from(Euler {
                    x: Deg::new(x_angle),
                    y: Deg::new(y_angle),
                    z: Deg::new(0.0),
                });
                dir = rotation.rotate_vector(dir);

                // We need to add a point with a distance of `dist` from `last`.
                let point = last + dir * dist;
                last = point;

                // FIXME Make branch spawn chance configurable
                if rng.gen_weighted_bool(depth * 3) {
                    // Build a vector for the branch direction (Z is up)
                    let dir = self.gen_branch_direction(rng, dir);
                    self.create_branch(rng, point, dir, depth + 1, diam);
                }

                points.push(ControlPoint {
                    point: point,
                    diameter: diam,
                });
            };

            // In a loop, get the length of the next segment from the current diameter.
            for _ in 0..segment_count {
                let length = segment_dist(diam);
                diam *= diam_factor;

                add_point(length, diam);

                if diam < 0.001 {
                    // Bail out at 1mm
                    break;
                }
            }
        }

        assert!(points.len() >= 2,
                "should've generated at least 2 points :(");
        self.branches.push(Branch {
            points: points,
            // FIXME Fixed color for now, we should use a configurable random color (or at least
            // make it brown).
            color: Vector3f::new(0.0, 1.0, 0.0),
        });
    }

    /// Given the growing direction of the parent branch, calculates a growing
    /// direction to use for a new child branch.
    fn gen_branch_direction<R: Rng>(&self, rng: &mut R, parent_dir: Vector3f) -> Vector3f {
        // `branch_angle_deg` specifies the angle range in degrees
        let angle = range_sample(&self.preset.branch_angle_deg, rng);

        // AWESOME hack: Generate *any* vector that is not parallel to `parent_dir`,
        // calculate the cross product between it and `parent_dir`.
        // This gets us a vector perpendicular to `parent_dir`, which we can rotate
        // `parent_dir` around to tilt it by `angle` degrees.
        let mut rand_vec = Vector3f::rand(rng).normalize();
        while rand_vec.angle(parent_dir).approx_eq(&Rad::new(0.0)) {
            // They're parallel, so generate a new vector. This is probably unnecessary,
            // but someone
            // who *actually* knows their math instead of faking it should check.
            rand_vec = Vector3f::rand(rng);
        }

        // Create vector around we'll tilt `parent_dir`
        let rot_vec = parent_dir.cross(rand_vec);

        // Tilt the growing direction of the parent branch by about 90° to get the
        // direction of the
        let tilt_rotation = Basis3::from_axis_angle(rot_vec, Deg::new(angle).into());

        let spin_angle = range_sample(&(0.0..360.0), rng);

        // Then rotate this vector randomly (0-360°) around the parent branch
        let around_we_go = Basis3::from_axis_angle(parent_dir, Deg::new(spin_angle).into());

        around_we_go.rotate_vector(tilt_rotation.rotate_vector(parent_dir))
    }

    fn create_trunk<R: Rng>(&mut self, rng: &mut R) {
        let trunk_diameter = range_sample(&self.preset.trunk_diameter, rng);
        let trunk_height = range_sample(&self.preset.trunk_height, rng);
        let trunk_diameter_top = range_sample(&self.preset.trunk_diameter_top, rng);
        let min_branch_height = range_sample(&self.preset.min_branch_height, rng) * trunk_height;

        debug!("trunk diam {} to {}, height {}, branch start at {}",
               trunk_diameter,
               trunk_diameter_top,
               trunk_height,
               min_branch_height);

        let mut points = Vec::new();

        {
            let mut add_point = |height, diam| {
                let point = Point3f::new(0.0, 0.0, height);
                if height >= min_branch_height {
                    // FIXME Make branch spawn chance configurable
                    let branches = &[0, 0, 1, 1, 1, 2];
                    for _ in 0..*rng.choose(branches).unwrap() {
                        // Build a vector for the branch direction (Z is up)
                        let dir = self.gen_branch_direction(rng, Vector3f::new(0.0, 0.0, 1.0));
                        self.create_branch(rng, point, dir, 1, diam);
                    }
                }

                points.push(ControlPoint {
                    point: point,
                    diameter: diam,
                });
            };

            let diam_start = Vector1::new(trunk_diameter);
            let diam_end = Vector1::new(trunk_diameter_top);

            // Split trunk in segments
            // FIXME Vary the segment direction like we do for normal branches
            // FIXME Make segment count depend on the trunk height
            const SEGMENT_COUNT: u32 = 10;
            for i in 0..SEGMENT_COUNT + 1 {
                let height = i as f32 * trunk_height / SEGMENT_COUNT as f32;
                let height_frac = height / trunk_height;
                let diam = diam_start.lerp(diam_end, height_frac);

                add_point(height, diam.x);
            }
        }

        assert!(points.len() >= 2,
                "should've generated at least 2 points :(");
        self.branches.push(Branch {
            points: points,
            // FIXME Fixed color for now, we should use a configurable random color (or at least
            // make it brown).
            color: Vector3f::new(0.0, 1.0, 0.0),
        });

        debug!("generated tree with {} branches", self.branches.len());
    }

    /// Generates a random tree according to the stored parameters.
    ///
    /// The tree is returned as a list of branches for now.
    pub fn generate<R: Rng>(mut self, rng: &mut R) -> Vec<Branch> {
        // Recursively create the tree and put all branches in a buffer.
        self.create_trunk(rng);
        self.branches
    }
}

impl Rand for TreeGen {
    fn rand<R: Rng>(rng: &mut R) -> Self {
        // Select a random preset that we'll use
        let preset = rng.choose(PRESETS).unwrap().clone();

        TreeGen {
            preset: preset,
            branches: Vec::new(),
        }
    }
}

/// Samples a random element from a range.
fn range_sample<T: SampleRange + PartialOrd + Copy, R: Rng>(range: &Range<T>, rng: &mut R) -> T {
    // Build a `rand` crate Range. We use `std`s Range for the cool `a..b` syntax ;)
    distributions::Range::new(range.start, range.end).ind_sample(rng)
}

/// Approximation of real-world distance of branch segments, depending on the
/// starting branch diameter.
fn segment_dist(diameter: f32) -> f32 {
    diameter * 11.25
}
<commit_msg>glob-import math<commit_after>//! Generates random trees and tree-like plants.

use prop::plant::{Branch, ControlPoint};
use math::*;
use rand::{Rand, Rng};
use rand::distributions::range::SampleRange;
use rand::distributions::{self, IndependentSample};
use std::ops::Range;
use std::cmp;

/// Parameters for the tree generator.
#[derive(Debug)]
struct Preset {
    name: &'static str,
    /// Diameter of the first branch we create (the trunk).
    trunk_diameter: Range<f32>,
    /// Trunk height. Note that branches going upward can increase plant height
    /// beyond this.
    trunk_height: Range<f32>,
    /// Trunk diameter at `trunk_height`. Should be smaller than
    /// `trunk_diameter`.
    trunk_diameter_top: Range<f32>,
    /// Trunk height at which we start creating branches.
    min_branch_height: Range<f32>,
    /// Range of subbranch diameters as a factor of the parent branch.
    branch_diameter_factor: Range<f32>,
    /// Range of subbranch angles in degrees.
    branch_angle_deg: Range<f32>,
    /// Factor by which to reduce segment diameter between consecutive points,
    /// sampled per branch.
    branch_diam_reduction: Range<f32>,
    /// Range of angles to use for rotation of new segments.
    ///
    /// The higher the angle, the more "twisted" branches appear.
    branch_segment_angle: Range<f32>,
    /// Range of segment counts for branches.
    ///
    /// Together with `branch_segment_length`, this defines the overall branch
    /// length.
    branch_segment_count: Range<u32>,
    /// Range of segment lengths to use for branches.
    ///
    /// Each segment will have a random length in this range.
    branch_segment_length: Range<f32>, // FIXME unused
}

static PRESETS: &'static [Preset] = &[Preset {
                                          name: "'Regular' Tree",
                                          trunk_diameter: 0.3..0.5,
                                          trunk_height: 3.0..6.0,
                                          trunk_diameter_top: 0.2..0.4,
                                          min_branch_height: 0.4..0.6,
                                          branch_diameter_factor: 0.3..0.5,
                                          branch_angle_deg: 70.0..110.0,
                                          branch_diam_reduction: 0.75..0.85,
                                          branch_segment_angle: 5.0..15.0,
                                          branch_segment_count: 3..10,
                                          branch_segment_length: 0.30..0.40,
                                      }];

pub struct TreeGen {
    preset: &'static Preset,
    /// Buffer for branches, filled as they're created.
    branches: Vec<Branch>,
}

impl TreeGen {
    /// Generates a new branch at the given start point.
    ///
    /// # Parameters
    ///
    /// * `rng`: The RNG to use
    /// * `start`: The point at which the branch should start
    /// * `dir`: Direction in which the branch should grow
    /// * `depth`: Recursion depth, used to limit the maximum branch depth
    /// * `parent_diam`: Diameter of the node of the parent branch where this
    ///   branch starts
    fn create_branch<R: Rng>(&mut self,
                             rng: &mut R,
                             start: Point3f,
                             dir: Vector3f,
                             depth: u32,
                             parent_diam: f32) {
        if depth > 5 {
            // Limit recursion
            return;
        }

        // Current normalized growing direction, variated slightly as segments are
        // generated
        let mut dir = dir.normalize();

        // Determine starting diameter of the new branch
        let mut diam = range_sample(&self.preset.branch_diameter_factor, rng) * parent_diam;
        // Determine how much segment diameter is reduced
        let diam_factor = range_sample(&self.preset.branch_diam_reduction, rng) * parent_diam;
        // How many segments should this branch get?
        let segment_count = range_sample(&self.preset.branch_segment_count, rng);

        let mut points = vec![ControlPoint {
                                  point: start,
                                  diameter: diam,
                              }];

        {
            let mut last = start;
            // Helper for adding a new point to this branch, which possibly grows a new
            // branch.
            let mut add_point = |dist, diam| {
                // First, get a random angle by which to variate this segment.
                let mut x_angle = range_sample(&self.preset.branch_segment_angle, rng);
                let mut y_angle = range_sample(&self.preset.branch_segment_angle, rng);

                // Invert sign with 50%, to mirror the specified range to the other side
                x_angle = if rng.gen() { -x_angle } else { x_angle };
                y_angle = if rng.gen() { -y_angle } else { y_angle };

                let rotation = Basis3::from(Euler {
                    x: Deg::new(x_angle),
                    y: Deg::new(y_angle),
                    z: Deg::new(0.0),
                });
                dir = rotation.rotate_vector(dir);

                // We need to add a point with a distance of `dist` from `last`.
                let point = last + dir * dist;
                last = point;

                // FIXME Make branch spawn chance configurable
                if rng.gen_weighted_bool(depth * 3) {
                    // Build a vector for the branch direction (Z is up)
                    let dir = self.gen_branch_direction(rng, dir);
                    self.create_branch(rng, point, dir, depth + 1, diam);
                }

                points.push(ControlPoint {
                    point: point,
                    diameter: diam,
                });
            };

            // In a loop, get the length of the next segment from the current diameter.
            for _ in 0..segment_count {
                let length = segment_dist(diam);
                diam *= diam_factor;

                add_point(length, diam);

                if diam < 0.001 {
                    // Bail out at 1mm
                    break;
                }
            }
        }

        assert!(points.len() >= 2,
                "should've generated at least 2 points :(");
        self.branches.push(Branch {
            points: points,
            // FIXME Fixed color for now, we should use a configurable random color (or at least
            // make it brown).
            color: Vector3f::new(0.0, 1.0, 0.0),
        });
    }

    /// Given the growing direction of the parent branch, calculates a growing
    /// direction to use for a new child branch.
    fn gen_branch_direction<R: Rng>(&self, rng: &mut R, parent_dir: Vector3f) -> Vector3f {
        // `branch_angle_deg` specifies the angle range in degrees
        let angle = range_sample(&self.preset.branch_angle_deg, rng);

        // AWESOME hack: Generate *any* vector that is not parallel to `parent_dir`,
        // calculate the cross product between it and `parent_dir`.
        // This gets us a vector perpendicular to `parent_dir`, which we can rotate
        // `parent_dir` around to tilt it by `angle` degrees.
        let mut rand_vec = Vector3f::rand(rng).normalize();
        while rand_vec.angle(parent_dir).approx_eq(&Rad::new(0.0)) {
            // They're parallel, so generate a new vector. This is probably unnecessary,
            // but someone
            // who *actually* knows their math instead of faking it should check.
            rand_vec = Vector3f::rand(rng);
        }

        // Create vector around we'll tilt `parent_dir`
        let rot_vec = parent_dir.cross(rand_vec);

        // Tilt the growing direction of the parent branch by about 90° to get the
        // direction of the
        let tilt_rotation = Basis3::from_axis_angle(rot_vec, Deg::new(angle).into());

        let spin_angle = range_sample(&(0.0..360.0), rng);

        // Then rotate this vector randomly (0-360°) around the parent branch
        let around_we_go = Basis3::from_axis_angle(parent_dir, Deg::new(spin_angle).into());

        around_we_go.rotate_vector(tilt_rotation.rotate_vector(parent_dir))
    }

    fn create_trunk<R: Rng>(&mut self, rng: &mut R) {
        let trunk_diameter = range_sample(&self.preset.trunk_diameter, rng);
        let trunk_height = range_sample(&self.preset.trunk_height, rng);
        let trunk_diameter_top = range_sample(&self.preset.trunk_diameter_top, rng);
        let min_branch_height = range_sample(&self.preset.min_branch_height, rng) * trunk_height;

        debug!("trunk diam {} to {}, height {}, branch start at {}",
               trunk_diameter,
               trunk_diameter_top,
               trunk_height,
               min_branch_height);

        let mut points = Vec::new();

        {
            let mut add_point = |height, diam| {
                let point = Point3f::new(0.0, 0.0, height);
                if height >= min_branch_height {
                    // FIXME Make branch spawn chance configurable
                    let branches = &[0, 0, 1, 1, 1, 2];
                    for _ in 0..*rng.choose(branches).unwrap() {
                        // Build a vector for the branch direction (Z is up)
                        let dir = self.gen_branch_direction(rng, Vector3f::new(0.0, 0.0, 1.0));
                        self.create_branch(rng, point, dir, 1, diam);
                    }
                }

                points.push(ControlPoint {
                    point: point,
                    diameter: diam,
                });
            };

            let diam_start = Vector1::new(trunk_diameter);
            let diam_end = Vector1::new(trunk_diameter_top);

            // Split trunk in segments
            // FIXME Vary the segment direction like we do for normal branches
            // FIXME Make segment count depend on the trunk height
            const SEGMENT_COUNT: u32 = 10;
            for i in 0..SEGMENT_COUNT + 1 {
                let height = i as f32 * trunk_height / SEGMENT_COUNT as f32;
                let height_frac = height / trunk_height;
                let diam = diam_start.lerp(diam_end, height_frac);

                add_point(height, diam.x);
            }
        }

        assert!(points.len() >= 2,
                "should've generated at least 2 points :(");
        self.branches.push(Branch {
            points: points,
            // FIXME Fixed color for now, we should use a configurable random color (or at least
            // make it brown).
            color: Vector3f::new(0.0, 1.0, 0.0),
        });

        debug!("generated tree with {} branches", self.branches.len());
    }

    /// Generates a random tree according to the stored parameters.
    ///
    /// The tree is returned as a list of branches for now.
    pub fn generate<R: Rng>(mut self, rng: &mut R) -> Vec<Branch> {
        // Recursively create the tree and put all branches in a buffer.
        self.create_trunk(rng);
        self.branches
    }
}

impl Rand for TreeGen {
    fn rand<R: Rng>(rng: &mut R) -> Self {
        // Select a random preset that we'll use
        let preset = rng.choose(PRESETS).unwrap().clone();

        TreeGen {
            preset: preset,
            branches: Vec::new(),
        }
    }
}

/// Samples a random element from a range.
fn range_sample<T: SampleRange + cmp::PartialOrd + Copy, R: Rng>(range: &Range<T>,
                                                                 rng: &mut R)
                                                                 -> T {
    // Build a `rand` crate Range. We use `std`s Range for the cool `a..b` syntax ;)
    distributions::Range::new(range.start, range.end).ind_sample(rng)
}

/// Approximation of real-world distance of branch segments, depending on the
/// starting branch diameter.
fn segment_dist(diameter: f32) -> f32 {
    diameter * 11.25
}
<|endoftext|>


<commit_before>#[macro_use]
extern crate toml_edit;

use toml_edit::{Document, Key};

macro_rules! test {
    ($before:expr, $root:ident, $ops:expr, $after:expr) => (
        let before = $before;
        let after = $after;

        let doc = Document::parse(before);
        assert!(doc.is_ok());
        let mut doc = doc.unwrap();

        {
            let $root = doc.root_mut();
            $ops
        }

        assert_eq!(doc.to_string(), after);
    );
}

// insertion

#[test]
fn test_insert_leaf_table() {
    test!(
        r#"
[servers]

[servers.alpha]
ip = "10.0.0.1"
dc = "eqdc10"

[other.table]
"#,
        root,
        {
            let mut servers = root.entry("servers");
            assert!(servers.is_table());
            let servers = servers.as_table_mut().unwrap();
            let mut beta = servers.insert_table(parse_key!("beta"));
            assert!(beta.is_table());
            let beta = beta.as_table_mut().unwrap();
            beta.append_value(parse_key!("ip"), "10.0.0.2");
            beta.append_value(parse_key!("dc"), "eqdc10");
        },
        r#"
[servers]

[servers.alpha]
ip = "10.0.0.1"
dc = "eqdc10"

[servers.beta]
ip = "10.0.0.2"
dc = "eqdc10"

[other.table]
"#
    );
}

#[test]
fn test_insert_nonleaf_table() {
    test!(
        r#"
[other.table]
"#,
        root,
        {
            let mut servers = root.append_table(parse_key!("servers"));
            assert!(servers.is_table());
            let servers = servers.as_table_mut().unwrap();
            let mut alpha = servers.insert_table(parse_key!("alpha"));
            assert!(alpha.is_table());
            let alpha = alpha.as_table_mut().unwrap();
            alpha.append_value(parse_key!("ip"), "10.0.0.1");
            alpha.append_value(parse_key!("dc"), "eqdc10");
        },
        r#"
[other.table]

[servers]

[servers.alpha]
ip = "10.0.0.1"
dc = "eqdc10"
"#
    );
}

#[test]
fn test_insert_array() {
    test!(
        r#"
[package]
title = "withoutarray"
"#,
        root,
        {
            let mut array = root.insert_array(parse_key!("bin"));
            assert!(array.is_array());
            let array = array.as_array_mut().unwrap();
            {
                let mut first = array.append();
                first.append_value(parse_key!("hello"), "world");
            }
            array.append();
        },
        r#"
[package]
title = "withoutarray"

[[bin]]
hello = "world"

[[bin]]
"#
    );
}

#[test]
fn test_insert_values() {
    test!(
        r#"
[tbl]
"#,
        root,
        {
            let mut table = root.entry("tbl");
            let table = table.as_table_mut().unwrap();
            table.append_value(parse_key!("key1"), "value1");
            table.append_value(parse_key!("\"key2\""), 42);
            table.append_value(parse_key!("'key3'"), 3.1415926);
        },

        r#"
[tbl]
key1 = "value1"
"key2" = 42
'key3' = 3.1415926
"#
    );
}

// removal

#[test]
fn test_remove_leaf_table() {
    test!(
        r#"
    [servers]

    # Indentation (tabs and/or spaces) is allowed but not required
[servers.alpha]
    ip = "10.0.0.1"
    dc = "eqdc10"

    [servers.beta]
    ip = "10.0.0.2"
    dc = "eqdc10"
"#,
        root,
        {
            let mut servers = root.entry("servers");
            assert!(servers.is_table());
            let servers = servers.as_table_mut().unwrap();
            assert!(servers.remove_table("alpha"));
        },
        r#"
    [servers]

    [servers.beta]
    ip = "10.0.0.2"
    dc = "eqdc10"
"#
    );
}

#[test]
fn test_remove_nonleaf_table() {
    test!(
        r#"
title = "not relevant"

# comment 1
[a.b.c] # comment 1.1
key1 = 1 # comment 1.2
# comment 2
[b] # comment 2.1
key2 = 2 # comment 2.2

# comment 3
[a] # comment 3.1
key3 = 3 # comment 3.2
[[a.'array']]
b = 1

[[a.b.c.trololololololo]]
c = 2
key3 = 42

   # comment on some other table
   [some.other.table]

# comment 4
[a.b] # comment 4.1
key4 = 4 # comment 4.2
key41 = 41 # comment 4.3

"#,
        root,
        {
            assert!(root.remove_table("a"));
        },
        r#"
title = "not relevant"
# comment 2
[b] # comment 2.1
key2 = 2 # comment 2.2

   # comment on some other table
   [some.other.table]

"#
    );
}

#[test]
fn test_remove_array_entry() {
    test!(
        r#"
    [package]
    name = "hello"
    version = "1.0.0"

    [[bin]]
    name = "world"
    path = "src/bin/world/main.rs"

    [dependencies]
    nom = "4.0" # future is here

    [[bin]]
    name = "delete me please"
    path = "src/bin/dmp/main.rs"
"#,
        root,
        {
            let mut dmp = root.entry("bin");
            assert!(dmp.is_array());
            let dmp = dmp.as_array_mut().unwrap();
            assert_eq!(dmp.len(), 2);
            dmp.remove(1);
        },
        r#"
    [package]
    name = "hello"
    version = "1.0.0"

    [[bin]]
    name = "world"
    path = "src/bin/world/main.rs"

    [dependencies]
    nom = "4.0" # future is here
"#
    );
}

#[test]
fn test_remove_array() {
    test!(
        r#"
    [package]
    name = "hello"
    version = "1.0.0"

    [[bin]]
    name = "world"
    path = "src/bin/world/main.rs"

    [dependencies]
    nom = "4.0" # future is here

    [[bin]]
    name = "delete me please"
    path = "src/bin/dmp/main.rs"
"#,
        root,
        {
            assert!(root.remove_array("bin"));
        },
        r#"
    [package]
    name = "hello"
    version = "1.0.0"

    [dependencies]
    nom = "4.0" # future is here
"#
    );
}

#[test]
fn test_remove_value() {
    test!(
        r#"
    name = "hello"
    # delete this
    version = "1.0.0" # please
    documentation = "https://docs.rs/hello"
"#,
        root,
        {
            let value = root.remove_value("version");
            assert!(value.is_some());
            let value = value.unwrap();
            assert!(value.is_str());
            let value = value.as_str().unwrap();
            assert_eq!(value, "1.0.0");
        },
        r#"
    name = "hello"
    documentation = "https://docs.rs/hello"
"#
    );
}
<commit_msg>fix(clippy): argh, clippy is too smart...<commit_after>#[macro_use]
extern crate toml_edit;

use toml_edit::{Document, Key};

macro_rules! test {
    ($before:expr, $root:ident, $ops:expr, $after:expr) => (
        let before = $before;
        let after = $after;

        let doc = Document::parse(before);
        assert!(doc.is_ok());
        let mut doc = doc.unwrap();

        {
            let $root = doc.root_mut();
            $ops
        }

        assert_eq!(doc.to_string(), after);
    );
}

// insertion

#[test]
fn test_insert_leaf_table() {
    test!(
        r#"
[servers]

[servers.alpha]
ip = "10.0.0.1"
dc = "eqdc10"

[other.table]
"#,
        root,
        {
            let mut servers = root.entry("servers");
            assert!(servers.is_table());
            let servers = servers.as_table_mut().unwrap();
            let mut beta = servers.insert_table(parse_key!("beta"));
            assert!(beta.is_table());
            let beta = beta.as_table_mut().unwrap();
            beta.append_value(parse_key!("ip"), "10.0.0.2");
            beta.append_value(parse_key!("dc"), "eqdc10");
        },
        r#"
[servers]

[servers.alpha]
ip = "10.0.0.1"
dc = "eqdc10"

[servers.beta]
ip = "10.0.0.2"
dc = "eqdc10"

[other.table]
"#
    );
}

#[test]
fn test_insert_nonleaf_table() {
    test!(
        r#"
[other.table]
"#,
        root,
        {
            let mut servers = root.append_table(parse_key!("servers"));
            assert!(servers.is_table());
            let servers = servers.as_table_mut().unwrap();
            let mut alpha = servers.insert_table(parse_key!("alpha"));
            assert!(alpha.is_table());
            let alpha = alpha.as_table_mut().unwrap();
            alpha.append_value(parse_key!("ip"), "10.0.0.1");
            alpha.append_value(parse_key!("dc"), "eqdc10");
        },
        r#"
[other.table]

[servers]

[servers.alpha]
ip = "10.0.0.1"
dc = "eqdc10"
"#
    );
}

#[test]
fn test_insert_array() {
    test!(
        r#"
[package]
title = "withoutarray"
"#,
        root,
        {
            let mut array = root.insert_array(parse_key!("bin"));
            assert!(array.is_array());
            let array = array.as_array_mut().unwrap();
            {
                let mut first = array.append();
                first.append_value(parse_key!("hello"), "world");
            }
            array.append();
        },
        r#"
[package]
title = "withoutarray"

[[bin]]
hello = "world"

[[bin]]
"#
    );
}

#[test]
fn test_insert_values() {
    test!(
        r#"
[tbl]
"#,
        root,
        {
            let mut table = root.entry("tbl");
            let table = table.as_table_mut().unwrap();
            table.append_value(parse_key!("key1"), "value1");
            table.append_value(parse_key!("\"key2\""), 42);
            table.append_value(parse_key!("'key3'"), 8.1415926);
        },

        r#"
[tbl]
key1 = "value1"
"key2" = 42
'key3' = 8.1415926
"#
    );
}

// removal

#[test]
fn test_remove_leaf_table() {
    test!(
        r#"
    [servers]

    # Indentation (tabs and/or spaces) is allowed but not required
[servers.alpha]
    ip = "10.0.0.1"
    dc = "eqdc10"

    [servers.beta]
    ip = "10.0.0.2"
    dc = "eqdc10"
"#,
        root,
        {
            let mut servers = root.entry("servers");
            assert!(servers.is_table());
            let servers = servers.as_table_mut().unwrap();
            assert!(servers.remove_table("alpha"));
        },
        r#"
    [servers]

    [servers.beta]
    ip = "10.0.0.2"
    dc = "eqdc10"
"#
    );
}

#[test]
fn test_remove_nonleaf_table() {
    test!(
        r#"
title = "not relevant"

# comment 1
[a.b.c] # comment 1.1
key1 = 1 # comment 1.2
# comment 2
[b] # comment 2.1
key2 = 2 # comment 2.2

# comment 3
[a] # comment 3.1
key3 = 3 # comment 3.2
[[a.'array']]
b = 1

[[a.b.c.trololololololo]]
c = 2
key3 = 42

   # comment on some other table
   [some.other.table]

# comment 4
[a.b] # comment 4.1
key4 = 4 # comment 4.2
key41 = 41 # comment 4.3

"#,
        root,
        {
            assert!(root.remove_table("a"));
        },
        r#"
title = "not relevant"
# comment 2
[b] # comment 2.1
key2 = 2 # comment 2.2

   # comment on some other table
   [some.other.table]

"#
    );
}

#[test]
fn test_remove_array_entry() {
    test!(
        r#"
    [package]
    name = "hello"
    version = "1.0.0"

    [[bin]]
    name = "world"
    path = "src/bin/world/main.rs"

    [dependencies]
    nom = "4.0" # future is here

    [[bin]]
    name = "delete me please"
    path = "src/bin/dmp/main.rs"
"#,
        root,
        {
            let mut dmp = root.entry("bin");
            assert!(dmp.is_array());
            let dmp = dmp.as_array_mut().unwrap();
            assert_eq!(dmp.len(), 2);
            dmp.remove(1);
        },
        r#"
    [package]
    name = "hello"
    version = "1.0.0"

    [[bin]]
    name = "world"
    path = "src/bin/world/main.rs"

    [dependencies]
    nom = "4.0" # future is here
"#
    );
}

#[test]
fn test_remove_array() {
    test!(
        r#"
    [package]
    name = "hello"
    version = "1.0.0"

    [[bin]]
    name = "world"
    path = "src/bin/world/main.rs"

    [dependencies]
    nom = "4.0" # future is here

    [[bin]]
    name = "delete me please"
    path = "src/bin/dmp/main.rs"
"#,
        root,
        {
            assert!(root.remove_array("bin"));
        },
        r#"
    [package]
    name = "hello"
    version = "1.0.0"

    [dependencies]
    nom = "4.0" # future is here
"#
    );
}

#[test]
fn test_remove_value() {
    test!(
        r#"
    name = "hello"
    # delete this
    version = "1.0.0" # please
    documentation = "https://docs.rs/hello"
"#,
        root,
        {
            let value = root.remove_value("version");
            assert!(value.is_some());
            let value = value.unwrap();
            assert!(value.is_str());
            let value = value.as_str().unwrap();
            assert_eq!(value, "1.0.0");
        },
        r#"
    name = "hello"
    documentation = "https://docs.rs/hello"
"#
    );
}
<|endoftext|>


<commit_before>extern crate breakpad_symbols;
extern crate disasm;
extern crate env_logger;
#[macro_use]
extern crate failure;
#[macro_use]
extern crate log;
extern crate minidump;
extern crate minidump_common;
extern crate reqwest;
#[allow(unused_imports)]
#[macro_use]
extern crate structopt;

use breakpad_symbols::{SimpleFrame, HttpSymbolSupplier, Symbolizer};
use disasm::{Color, CpuArch, SourceLocation, SourceLookup};
use failure::Error;
use minidump_common::traits::Module;
use minidump::{Minidump, MinidumpException, MinidumpMemoryList, MinidumpModule, MinidumpModuleList,
               MinidumpSystemInfo};
use minidump::system_info::Cpu;
use reqwest::Client;
use std::env;
use std::fs::{self, File};
use std::path::{Path, PathBuf};
use structopt::StructOpt;

#[derive(StructOpt)]
#[structopt(name = "get-minidump-instructions", about = "Display instructions from a minidump")]
struct GetMinidumpInstructions {
    #[structopt(long = "color", help = "Enable colored output")]
    color: Option<Color>,
    #[structopt(help = "Input minidump", parse(from_os_str))]
    minidump: PathBuf,
    #[structopt(help = "Symbol paths", parse(from_os_str))]
    symbol_paths: Vec<PathBuf>,
}

#[derive(StructOpt)]
#[structopt(name = "dumpstack", about = "Display possible stack frames from a minidump")]
struct DumpStack {
    #[structopt(help = "Input minidump", parse(from_os_str))]
    minidump: PathBuf,
    #[structopt(help = "Symbol paths", parse(from_os_str))]
    symbol_paths: Vec<PathBuf>,
    #[structopt(short = "a", help = "Show stack addresses that don't resolve to functions")]
    all: bool,
}

struct SymLookup {
    modules: MinidumpModuleList,
    symbolizer: Symbolizer,
    client: Client,
}

struct HgWebFile {
    host: String,
    repo: String,
    rev: String,
    path: String,
}

struct GitHubFile {
    repo: String,
    rev: String,
    path: String,
}

trait VCSFile {
    fn raw_url(&self) -> String;
    fn annotate_url(&self, line: u64) -> String;
    fn as_local_filename(&self) -> String;
}

impl VCSFile for HgWebFile {
    fn raw_url(&self) -> String {
        format!("https://{}/{}/raw-file/{}/{}",
                self.host, self.repo, self.rev, self.path)
    }

    fn annotate_url(&self, line: u64) -> String {
        format!("https://{}/{}/annotate/{}/{}#l{}",
                self.host, self.repo, self.rev, self.path, line)
    }
    fn as_local_filename(&self) -> String {
        format!("{}_{}_{}", self.host, self.rev, self.path).replace('/', "_")
    }
}

impl VCSFile for GitHubFile {
    fn raw_url(&self) -> String {
        format!("http://raw.githubusercontent.com/{}/{}/{}", self.repo, self.rev, self.path)
    }
    fn annotate_url(&self, line: u64) -> String {
        format!("https://github.com/{}/blob/{}/{}#L{}", self.repo, self.rev, self.path, line)
    }
    fn as_local_filename(&self) -> String {
        format!("{}_{}_{}", self.repo, self.rev, self.path).replace('/', "_")
    }
}

fn file_exists(path: &Path) -> bool {
    fs::metadata(path).is_ok()
}

fn fetch_url_to_path(client: &Client, url: &str, path: &Path) -> Result<(), Error> {
    debug!("fetch_url_to_path({}, {:?})", url, path);
    let mut res = client.get(url).send()?.error_for_status()?;
    debug!("fetch_url_to_path: HTTP success");
    let mut tmp_path = path.to_owned().into_os_string();
    tmp_path.push(".tmp");
    {
        let mut f = File::create(&tmp_path)?;
        res.copy_to(&mut f)?;
    }
    fs::rename(&tmp_path, &path)?;
    Ok(())
}

fn parse_vcs_info(filename: &str) -> Result<Box<VCSFile>, Error> {
    let mut bits = filename.split(':');
    Ok(match (bits.next(), bits.next(), bits.next(), bits.next()) {
        (Some("hg"), Some(repo), Some(path), Some(rev)) if repo.starts_with("hg.mozilla.org/") => {
            let mut s = repo.splitn(2, '/');
            let host = s.next().unwrap().to_owned();
            let repo = s.next().unwrap().to_owned();
            let path = path.to_owned();
            let rev = rev.to_owned();
            Box::new(HgWebFile {
                host,
                repo,
                rev,
                path,
            })
        }
        (Some("git"), Some(repo), Some(path), Some(rev)) if repo.starts_with("github.com/") => {
            let repo = repo.splitn(2, '/').nth(1).unwrap().to_owned();
            let path = path.to_owned();
            let rev = rev.to_owned();
            Box::new(GitHubFile {
                repo,
                rev,
                path,
            })
        }
        _ => return Err(format_err!("No VCS info in filename")),
    })
}

fn maybe_fetch_source_file(client: &Client, url: &str, path: &Path) -> Result<(), Error> {
    if file_exists(path) {
        Ok(())
    } else {
        fetch_url_to_path(client, url, path)
    }
}

impl SourceLookup for SymLookup {
    fn lookup(&mut self, address: u64) -> Option<SourceLocation> {
        self.modules.module_at_address(address).and_then(|module| {
            let mut frame = SimpleFrame::with_instruction(address);
            self.symbolizer.fill_symbol(module, &mut frame);
            let SimpleFrame { source_file, source_line, .. } = frame;
            if let (Some(file), Some(line)) = (source_file, source_line) {
                let line = line as u64;
                match parse_vcs_info(&file) {
                    Ok(info) => {
                        let url = info.raw_url();
                        let local = env::temp_dir().join(info.as_local_filename());
                        if maybe_fetch_source_file(&self.client, &url, &local).is_ok() {
                            return Some(SourceLocation {
                                file: local.to_string_lossy().into_owned(),
                                file_display: Some(info.annotate_url(line)),
                                line
                            })
                        }
                    }
                    _ => {}
                }
                // If fetching the file doesn't work, just hand back the original filename.
                Some(SourceLocation { file: file.into(), file_display: None, line })
            } else {
                // We didn't find any source info.
                None
            }
        })
    }
}

fn handle_symbol_paths(symbol_paths: Vec<PathBuf>) -> Result<Symbolizer, Error> {
    let tmp_path = env::temp_dir().join("symbols");
    fs::create_dir_all(&tmp_path)?;
    let (symbol_paths, symbol_urls) = if symbol_paths.is_empty() {
        // Use the Mozilla symbol server if no symbol paths are supplied.
        let symbol_urls = vec!["https://symbols.mozilla.org/".to_owned()];
        (symbol_paths, symbol_urls)
    } else {
        let urls = symbol_paths.iter()
            .filter(|p| p.starts_with("http"))
            .filter_map(|p| p.to_str().map(str::to_owned)).collect();
        let paths = symbol_paths.into_iter().filter(|p| !p.starts_with("http")).collect();
        (paths, urls)
    };
    let supplier = HttpSymbolSupplier::new(symbol_urls, tmp_path, symbol_paths);
    Ok(Symbolizer::new(supplier))
}

pub fn get_minidump_instructions() -> Result<(), Error> {
    env_logger::init();
    let GetMinidumpInstructions { color, minidump, symbol_paths } =
        GetMinidumpInstructions::from_args();
    let dump = Minidump::read_path(&minidump)?;
    let modules = dump.get_stream::<MinidumpModuleList>()?;
    let exception = dump.get_stream::<MinidumpException>()?;
    let context = exception.context.as_ref().ok_or(format_err!("Missing exception context"))?;
    let ip = context.get_instruction_pointer();
    let memory_list = dump.get_stream::<MinidumpMemoryList>()?;
    let memory = memory_list.memory_at_address(ip)
        .ok_or(format_err!("Minidump doesn't contain a memory region that contains the instruction pointer from the exception record"))?;
    let sys_info = dump.get_stream::<MinidumpSystemInfo>()?;
    let arch = match sys_info.cpu {
        Cpu::X86 => CpuArch::X86,
        Cpu::X86_64 => CpuArch::X86_64,
        Cpu::Arm64 => CpuArch::ARM64,
        _ => return Err(format_err!("Unsupported CPU architecture: {}", sys_info.cpu)),
    };
    let symbolizer = handle_symbol_paths(symbol_paths)?;
    let client = Client::new();
    let mut lookup = SymLookup { modules, symbolizer, client };
    println!("Faulting instruction pointer: {:#x}", ip);
    info!("Disassembling {} bytes starting at {:#x} of arch {:?}",
          memory.bytes.len(), memory.base_address, arch);
    disasm::disasm_bytes(&memory.bytes,
                         memory.base_address,
                         arch,
                         color.unwrap_or(Color::Auto),
                         Some(ip),
                         &mut lookup)?;
    Ok(())
}

    fn basename(f: &str) -> &str {
        match f.rfind(|c| c == '/' || c == '\\') {
            None => f,
            Some(index) => &f[(index + 1)..],
        }
    }

fn print_frame(module: &MinidumpModule, frame: SimpleFrame) {
    print!("{}", basename(&module.code_file()));
    if let Some(ref func) = frame.function {
        print!("!{}", func);
        if let (Some(ref file), Some(line)) = (frame.source_file, frame.source_line) {
            print!(" [");
            match parse_vcs_info(file) {
                Ok(info) => print!("{}", info.annotate_url(line as u64)),
                _ => print!("{} : {}", basename(file), line),
            }
            if let Some(line_base) = frame.source_line_base {
                print!(" + {}", frame.instruction - line_base);
            }
            print!("]");
        } else if let Some(func_base) = frame.function_base {
            print!(" + {:#x}", frame.instruction - func_base);
        }
    } else {
        print!(" + {:#x}", frame.instruction - module.base_address());
    }
    println!("");
}

pub fn dump_minidump_stack() -> Result<(), Error> {
    env_logger::init();
    let DumpStack { all, minidump, symbol_paths } =
        DumpStack::from_args();
    info!("Reading minidump {:?}", minidump);
    let dump = Minidump::read_path(&minidump)?;
    let symbolizer = handle_symbol_paths(symbol_paths)?;
    let modules = dump.get_stream::<MinidumpModuleList>()?;
    let memory_list = dump.get_stream::<MinidumpMemoryList>()?;
    let sys_info = dump.get_stream::<MinidumpSystemInfo>()?;
    let wordsize = match sys_info.cpu {
        Cpu::X86 | Cpu::Ppc | Cpu::Sparc | Cpu::Arm => 4,
        Cpu::X86_64 | Cpu::Ppc64 | Cpu::Arm64 => 8,
        Cpu::Unknown(u) => bail!("Unknown cpu: {:#x}", u),
    };
    // TODO: provide a commandline option for the address.
    // Default to the top of the crashing stack.
    let exception = dump.get_stream::<MinidumpException>()?;
    let context = exception.context.as_ref().ok_or(format_err!("Missing exception context"))?;
    let sp = context.get_stack_pointer();
    info!("Stack pointer: {:#x}", sp);
    let memory = memory_list.memory_at_address(sp)
        .ok_or_else(|| format_err!("Minidump doesn't contain a memory region that contains address {:#x}", sp))?;
    for addr in (memory.base_address..memory.base_address + memory.size).step_by(wordsize) {
        let instruction = match wordsize {
            4 => memory.get_memory_at_address::<u32>(addr).map(|i| i as u64),
            8 => memory.get_memory_at_address::<u64>(addr),
            _ => unreachable!(),
        };
        if let Some(instruction) = instruction {
            if let Some(module) = modules.module_at_address(instruction) {
                let mut frame = SimpleFrame::with_instruction(instruction);
                symbolizer.fill_symbol(module, &mut frame);
                if frame.function.is_some() || all {
                    print!("{:#0width$x} ", addr, width=wordsize*2);
                    print_frame(module, frame);
                }
            }
        }
    }
    Ok(())
}
<commit_msg>change get-minidump-instructions to disassemble from crashing ip<commit_after>extern crate breakpad_symbols;
extern crate disasm;
extern crate env_logger;
#[macro_use]
extern crate failure;
#[macro_use]
extern crate log;
extern crate minidump;
extern crate minidump_common;
extern crate reqwest;
#[allow(unused_imports)]
#[macro_use]
extern crate structopt;

use breakpad_symbols::{SimpleFrame, HttpSymbolSupplier, Symbolizer};
use disasm::{Color, CpuArch, SourceLocation, SourceLookup};
use failure::Error;
use minidump_common::traits::Module;
use minidump::{Minidump, MinidumpException, MinidumpMemoryList, MinidumpModule, MinidumpModuleList,
               MinidumpSystemInfo};
use minidump::system_info::Cpu;
use reqwest::Client;
use std::env;
use std::fs::{self, File};
use std::path::{Path, PathBuf};
use structopt::StructOpt;

#[derive(StructOpt)]
#[structopt(name = "get-minidump-instructions", about = "Display instructions from a minidump")]
struct GetMinidumpInstructions {
    #[structopt(long = "color", help = "Enable colored output")]
    color: Option<Color>,
    #[structopt(help = "Input minidump", parse(from_os_str))]
    minidump: PathBuf,
    #[structopt(help = "Symbol paths", parse(from_os_str))]
    symbol_paths: Vec<PathBuf>,
}

#[derive(StructOpt)]
#[structopt(name = "dumpstack", about = "Display possible stack frames from a minidump")]
struct DumpStack {
    #[structopt(help = "Input minidump", parse(from_os_str))]
    minidump: PathBuf,
    #[structopt(help = "Symbol paths", parse(from_os_str))]
    symbol_paths: Vec<PathBuf>,
    #[structopt(short = "a", help = "Show stack addresses that don't resolve to functions")]
    all: bool,
}

struct SymLookup {
    modules: MinidumpModuleList,
    symbolizer: Symbolizer,
    client: Client,
}

struct HgWebFile {
    host: String,
    repo: String,
    rev: String,
    path: String,
}

struct GitHubFile {
    repo: String,
    rev: String,
    path: String,
}

trait VCSFile {
    fn raw_url(&self) -> String;
    fn annotate_url(&self, line: u64) -> String;
    fn as_local_filename(&self) -> String;
}

impl VCSFile for HgWebFile {
    fn raw_url(&self) -> String {
        format!("https://{}/{}/raw-file/{}/{}",
                self.host, self.repo, self.rev, self.path)
    }

    fn annotate_url(&self, line: u64) -> String {
        format!("https://{}/{}/annotate/{}/{}#l{}",
                self.host, self.repo, self.rev, self.path, line)
    }
    fn as_local_filename(&self) -> String {
        format!("{}_{}_{}", self.host, self.rev, self.path).replace('/', "_")
    }
}

impl VCSFile for GitHubFile {
    fn raw_url(&self) -> String {
        format!("http://raw.githubusercontent.com/{}/{}/{}", self.repo, self.rev, self.path)
    }
    fn annotate_url(&self, line: u64) -> String {
        format!("https://github.com/{}/blob/{}/{}#L{}", self.repo, self.rev, self.path, line)
    }
    fn as_local_filename(&self) -> String {
        format!("{}_{}_{}", self.repo, self.rev, self.path).replace('/', "_")
    }
}

fn file_exists(path: &Path) -> bool {
    fs::metadata(path).is_ok()
}

fn fetch_url_to_path(client: &Client, url: &str, path: &Path) -> Result<(), Error> {
    debug!("fetch_url_to_path({}, {:?})", url, path);
    let mut res = client.get(url).send()?.error_for_status()?;
    debug!("fetch_url_to_path: HTTP success");
    let mut tmp_path = path.to_owned().into_os_string();
    tmp_path.push(".tmp");
    {
        let mut f = File::create(&tmp_path)?;
        res.copy_to(&mut f)?;
    }
    fs::rename(&tmp_path, &path)?;
    Ok(())
}

fn parse_vcs_info(filename: &str) -> Result<Box<VCSFile>, Error> {
    let mut bits = filename.split(':');
    Ok(match (bits.next(), bits.next(), bits.next(), bits.next()) {
        (Some("hg"), Some(repo), Some(path), Some(rev)) if repo.starts_with("hg.mozilla.org/") => {
            let mut s = repo.splitn(2, '/');
            let host = s.next().unwrap().to_owned();
            let repo = s.next().unwrap().to_owned();
            let path = path.to_owned();
            let rev = rev.to_owned();
            Box::new(HgWebFile {
                host,
                repo,
                rev,
                path,
            })
        }
        (Some("git"), Some(repo), Some(path), Some(rev)) if repo.starts_with("github.com/") => {
            let repo = repo.splitn(2, '/').nth(1).unwrap().to_owned();
            let path = path.to_owned();
            let rev = rev.to_owned();
            Box::new(GitHubFile {
                repo,
                rev,
                path,
            })
        }
        _ => return Err(format_err!("No VCS info in filename")),
    })
}

fn maybe_fetch_source_file(client: &Client, url: &str, path: &Path) -> Result<(), Error> {
    if file_exists(path) {
        Ok(())
    } else {
        fetch_url_to_path(client, url, path)
    }
}

impl SourceLookup for SymLookup {
    fn lookup(&mut self, address: u64) -> Option<SourceLocation> {
        self.modules.module_at_address(address).and_then(|module| {
            let mut frame = SimpleFrame::with_instruction(address);
            self.symbolizer.fill_symbol(module, &mut frame);
            let SimpleFrame { source_file, source_line, .. } = frame;
            if let (Some(file), Some(line)) = (source_file, source_line) {
                let line = line as u64;
                match parse_vcs_info(&file) {
                    Ok(info) => {
                        let url = info.raw_url();
                        let local = env::temp_dir().join(info.as_local_filename());
                        if maybe_fetch_source_file(&self.client, &url, &local).is_ok() {
                            return Some(SourceLocation {
                                file: local.to_string_lossy().into_owned(),
                                file_display: Some(info.annotate_url(line)),
                                line
                            })
                        }
                    }
                    _ => {}
                }
                // If fetching the file doesn't work, just hand back the original filename.
                Some(SourceLocation { file: file.into(), file_display: None, line })
            } else {
                // We didn't find any source info.
                None
            }
        })
    }
}

fn handle_symbol_paths(symbol_paths: Vec<PathBuf>) -> Result<Symbolizer, Error> {
    let tmp_path = env::temp_dir().join("symbols");
    fs::create_dir_all(&tmp_path)?;
    let (symbol_paths, symbol_urls) = if symbol_paths.is_empty() {
        // Use the Mozilla symbol server if no symbol paths are supplied.
        let symbol_urls = vec!["https://symbols.mozilla.org/".to_owned()];
        (symbol_paths, symbol_urls)
    } else {
        let urls = symbol_paths.iter()
            .filter(|p| p.starts_with("http"))
            .filter_map(|p| p.to_str().map(str::to_owned)).collect();
        let paths = symbol_paths.into_iter().filter(|p| !p.starts_with("http")).collect();
        (paths, urls)
    };
    let supplier = HttpSymbolSupplier::new(symbol_urls, tmp_path, symbol_paths);
    Ok(Symbolizer::new(supplier))
}

pub fn get_minidump_instructions() -> Result<(), Error> {
    env_logger::init();
    let GetMinidumpInstructions { color, minidump, symbol_paths } =
        GetMinidumpInstructions::from_args();
    let dump = Minidump::read_path(&minidump)?;
    let modules = dump.get_stream::<MinidumpModuleList>()?;
    let exception = dump.get_stream::<MinidumpException>()?;
    let context = exception.context.as_ref().ok_or(format_err!("Missing exception context"))?;
    let ip = context.get_instruction_pointer();
    let memory_list = dump.get_stream::<MinidumpMemoryList>()?;
    let memory = memory_list.memory_at_address(ip)
        .ok_or(format_err!("Minidump doesn't contain a memory region that contains the instruction pointer from the exception record"))?;
    let sys_info = dump.get_stream::<MinidumpSystemInfo>()?;
    let arch = match sys_info.cpu {
        Cpu::X86 => CpuArch::X86,
        Cpu::X86_64 => CpuArch::X86_64,
        Cpu::Arm64 => CpuArch::ARM64,
        _ => return Err(format_err!("Unsupported CPU architecture: {}", sys_info.cpu)),
    };
    let symbolizer = handle_symbol_paths(symbol_paths)?;
    let client = Client::new();
    let mut lookup = SymLookup { modules, symbolizer, client };
    println!("Faulting instruction pointer: {:#x}", ip);
    info!("Disassembling {} bytes starting at {:#x} of arch {:?}",
          memory.bytes.len(), memory.base_address, arch);
    disasm::disasm_bytes(&memory.bytes[ip as usize - memory.base_address as usize..],
                         ip,
                         arch,
                         color.unwrap_or(Color::Auto),
                         Some(ip),
                         &mut lookup)?;
    Ok(())
}

    fn basename(f: &str) -> &str {
        match f.rfind(|c| c == '/' || c == '\\') {
            None => f,
            Some(index) => &f[(index + 1)..],
        }
    }

fn print_frame(module: &MinidumpModule, frame: SimpleFrame) {
    print!("{}", basename(&module.code_file()));
    if let Some(ref func) = frame.function {
        print!("!{}", func);
        if let (Some(ref file), Some(line)) = (frame.source_file, frame.source_line) {
            print!(" [");
            match parse_vcs_info(file) {
                Ok(info) => print!("{}", info.annotate_url(line as u64)),
                _ => print!("{} : {}", basename(file), line),
            }
            if let Some(line_base) = frame.source_line_base {
                print!(" + {}", frame.instruction - line_base);
            }
            print!("]");
        } else if let Some(func_base) = frame.function_base {
            print!(" + {:#x}", frame.instruction - func_base);
        }
    } else {
        print!(" + {:#x}", frame.instruction - module.base_address());
    }
    println!("");
}

pub fn dump_minidump_stack() -> Result<(), Error> {
    env_logger::init();
    let DumpStack { all, minidump, symbol_paths } =
        DumpStack::from_args();
    info!("Reading minidump {:?}", minidump);
    let dump = Minidump::read_path(&minidump)?;
    let symbolizer = handle_symbol_paths(symbol_paths)?;
    let modules = dump.get_stream::<MinidumpModuleList>()?;
    let memory_list = dump.get_stream::<MinidumpMemoryList>()?;
    let sys_info = dump.get_stream::<MinidumpSystemInfo>()?;
    let wordsize = match sys_info.cpu {
        Cpu::X86 | Cpu::Ppc | Cpu::Sparc | Cpu::Arm => 4,
        Cpu::X86_64 | Cpu::Ppc64 | Cpu::Arm64 => 8,
        Cpu::Unknown(u) => bail!("Unknown cpu: {:#x}", u),
    };
    // TODO: provide a commandline option for the address.
    // Default to the top of the crashing stack.
    let exception = dump.get_stream::<MinidumpException>()?;
    let context = exception.context.as_ref().ok_or(format_err!("Missing exception context"))?;
    let sp = context.get_stack_pointer();
    info!("Stack pointer: {:#x}", sp);
    let memory = memory_list.memory_at_address(sp)
        .ok_or_else(|| format_err!("Minidump doesn't contain a memory region that contains address {:#x}", sp))?;
    for addr in (memory.base_address..memory.base_address + memory.size).step_by(wordsize) {
        let instruction = match wordsize {
            4 => memory.get_memory_at_address::<u32>(addr).map(|i| i as u64),
            8 => memory.get_memory_at_address::<u64>(addr),
            _ => unreachable!(),
        };
        if let Some(instruction) = instruction {
            if let Some(module) = modules.module_at_address(instruction) {
                let mut frame = SimpleFrame::with_instruction(instruction);
                symbolizer.fill_symbol(module, &mut frame);
                if frame.function.is_some() || all {
                    print!("{:#0width$x} ", addr, width=wordsize*2);
                    print_frame(module, frame);
                }
            }
        }
    }
    Ok(())
}
<|endoftext|>


<commit_before>use super::branchify::generate_branchified_method;
use super::get_writer;
use std::io::IoResult;

pub fn generate(output_dir: &Path) -> IoResult<()> {
    let mut writer = get_writer(output_dir, "read_method.rs");
    try!(writer.write(b"\
// This automatically generated file is included in request.rs.
use std::io::{Stream, IoResult};
use method::{Method, Connect, Delete, Get, Head, Options, Patch, Post, Put, Trace, ExtensionMethod};
use server::request::MAX_METHOD_LEN;
use rfc2616::{SP, is_token_item};
use buffer::BufferedStream;

#[inline]
pub fn read_method<S: Stream>(stream: &mut BufferedStream<S>) -> IoResult<Method> {
"));

    try!(generate_branchified_method(
        &mut writer,
        branchify!(case sensitive,
            "CONNECT" => Connect,
            "DELETE"  => Delete,
            "GET"     => Get,
            "HEAD"    => Head,
            "OPTIONS" => Options,
            "PATCH"   => Patch,
            "POST"    => Post,
            "PUT"     => Put,
            "TRACE"   => Trace
        ),
        1,
        "stream.read_byte()",
        "SP",
        "MAX_METHOD_LEN",
        "is_token_item(b)",
        "ExtensionMethod({})"));
    writer.write(b"}\n")
}
<commit_msg>Fix trait object use with DST<commit_after>use super::branchify::generate_branchified_method;
use super::get_writer;
use std::io::IoResult;

pub fn generate(output_dir: &Path) -> IoResult<()> {
    let mut writer = get_writer(output_dir, "read_method.rs");
    try!(writer.write(b"\
// This automatically generated file is included in request.rs.
use std::io::{Stream, IoResult};
use method::{Method, Connect, Delete, Get, Head, Options, Patch, Post, Put, Trace, ExtensionMethod};
use server::request::MAX_METHOD_LEN;
use rfc2616::{SP, is_token_item};
use buffer::BufferedStream;

#[inline]
pub fn read_method<S: Stream>(stream: &mut BufferedStream<S>) -> IoResult<Method> {
"));

    try!(generate_branchified_method(
        &mut *writer,
        branchify!(case sensitive,
            "CONNECT" => Connect,
            "DELETE"  => Delete,
            "GET"     => Get,
            "HEAD"    => Head,
            "OPTIONS" => Options,
            "PATCH"   => Patch,
            "POST"    => Post,
            "PUT"     => Put,
            "TRACE"   => Trace
        ),
        1,
        "stream.read_byte()",
        "SP",
        "MAX_METHOD_LEN",
        "is_token_item(b)",
        "ExtensionMethod({})"));
    writer.write(b"}\n")
}
<|endoftext|>


<commit_before>#![cfg(all(feature = "os-poll", feature = "tcp"))]

use mio::net::{TcpListener, TcpStream};
use mio::{event, Events, Interest, Poll, Registry, Token};
use std::net;
use std::sync::{Arc, Barrier};
use std::thread::{self, sleep};
use std::time::Duration;
use std::{fmt, io};

mod util;
use util::{any_local_address, assert_send, assert_sync, init};

#[test]
fn is_send_and_sync() {
    assert_send::<Events>();
    assert_sync::<Events>();

    assert_sync::<Poll>();
    assert_send::<Poll>();

    assert_sync::<Registry>();
    assert_send::<Registry>();
}

#[test]
fn run_once_with_nothing() {
    init();

    let mut events = Events::with_capacity(16);
    let mut poll = Poll::new().unwrap();
    poll.poll(&mut events, Some(Duration::from_millis(100)))
        .unwrap();
}

#[test]
fn add_then_drop() {
    init();

    let mut events = Events::with_capacity(16);
    let mut l = TcpListener::bind(any_local_address()).unwrap();
    let mut poll = Poll::new().unwrap();
    poll.registry()
        .register(&mut l, Token(1), Interest::READABLE | Interest::WRITABLE)
        .unwrap();
    drop(l);
    poll.poll(&mut events, Some(Duration::from_millis(100)))
        .unwrap();
}

#[test]
fn zero_duration_polls_events() {
    init();

    let mut poll = Poll::new().unwrap();
    let mut events = Events::with_capacity(16);

    let listener = net::TcpListener::bind(any_local_address()).unwrap();
    let addr = listener.local_addr().unwrap();

    let streams: Vec<TcpStream> = (0..3)
        .map(|n| {
            let mut stream = TcpStream::connect(addr).unwrap();
            poll.registry()
                .register(&mut stream, Token(n), Interest::WRITABLE)
                .unwrap();
            stream
        })
        .collect();

    // Ensure the TcpStreams have some time to connection and for the events to
    // show up.
    sleep(Duration::from_millis(10));

    // Even when passing a zero duration timeout we still want do the system
    // call.
    poll.poll(&mut events, Some(Duration::from_nanos(0)))
        .unwrap();
    assert!(!events.is_empty());

    // Both need to live until here.
    drop(streams);
    drop(listener);
}

#[test]
fn poll_closes_fd() {
    init();

    for _ in 0..2000 {
        let mut poll = Poll::new().unwrap();
        let mut events = Events::with_capacity(4);

        poll.poll(&mut events, Some(Duration::from_millis(0)))
            .unwrap();

        drop(poll);
    }
}

#[test]
fn drop_cancels_interest_and_shuts_down() {
    init();

    use mio::net::TcpStream;
    use std::io;
    use std::io::Read;
    use std::net::TcpListener;
    use std::thread;

    let l = TcpListener::bind("127.0.0.1:0").unwrap();
    let addr = l.local_addr().unwrap();

    let t = thread::spawn(move || {
        let mut s = l.incoming().next().unwrap().unwrap();
        s.set_read_timeout(Some(Duration::from_secs(5)))
            .expect("set_read_timeout");
        let r = s.read(&mut [0; 16]);
        match r {
            Ok(_) => (),
            Err(e) => {
                if e.kind() != io::ErrorKind::UnexpectedEof {
                    panic!(e);
                }
            }
        }
    });

    let mut poll = Poll::new().unwrap();
    let mut s = TcpStream::connect(addr).unwrap();

    poll.registry()
        .register(&mut s, Token(1), Interest::READABLE | Interest::WRITABLE)
        .unwrap();
    let mut events = Events::with_capacity(16);
    'outer: loop {
        poll.poll(&mut events, None).unwrap();
        for event in &events {
            if event.token() == Token(1) {
                // connected
                break 'outer;
            }
        }
    }

    let mut b = [0; 1024];
    match s.read(&mut b) {
        Ok(_) => panic!("unexpected ok"),
        Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => (),
        Err(e) => panic!("unexpected error: {:?}", e),
    }

    drop(s);
    t.join().unwrap();
}

#[test]
fn registry_behind_arc() {
    // `Registry` should work behind an `Arc`, being `Sync` and `Send`.
    init();

    let mut poll = Poll::new().unwrap();
    let registry = Arc::new(poll.registry().try_clone().unwrap());
    let mut events = Events::with_capacity(128);

    let mut listener = TcpListener::bind(any_local_address()).unwrap();
    let addr = listener.local_addr().unwrap();
    let barrier = Arc::new(Barrier::new(3));

    let registry2 = Arc::clone(&registry);
    let registry3 = Arc::clone(&registry);
    let barrier2 = Arc::clone(&barrier);
    let barrier3 = Arc::clone(&barrier);

    let handle1 = thread::spawn(move || {
        registry2
            .register(&mut listener, Token(0), Interest::READABLE)
            .unwrap();
        barrier2.wait();
    });
    let handle2 = thread::spawn(move || {
        let mut stream = TcpStream::connect(addr).unwrap();
        registry3
            .register(
                &mut stream,
                Token(1),
                Interest::READABLE | Interest::WRITABLE,
            )
            .unwrap();
        barrier3.wait();
    });

    poll.poll(&mut events, Some(Duration::from_millis(1000)))
        .unwrap();
    assert!(events.iter().count() >= 1);

    // Let the threads return.
    barrier.wait();

    handle1.join().unwrap();
    handle2.join().unwrap();
}

// On kqueue platforms registering twice (not *re*registering) works.
#[test]
#[cfg(any(target_os = "linux", target_os = "windows"))]
pub fn double_register() {
    init();
    let poll = Poll::new().unwrap();

    let mut l = TcpListener::bind("127.0.0.1:0".parse().unwrap()).unwrap();

    poll.registry()
        .register(&mut l, Token(0), Interest::READABLE)
        .unwrap();

    assert!(poll
        .registry()
        .register(&mut l, Token(1), Interest::READABLE)
        .is_err());
}

struct TestEventSource {
    registrations: Vec<(Token, Interest)>,
    reregistrations: Vec<(Token, Interest)>,
    deregister_count: usize,
}

impl TestEventSource {
    fn new() -> TestEventSource {
        TestEventSource {
            registrations: Vec::new(),
            reregistrations: Vec::new(),
            deregister_count: 0,
        }
    }
}

impl event::Source for TestEventSource {
    fn register(
        &mut self,
        _registry: &Registry,
        token: Token,
        interests: Interest,
    ) -> io::Result<()> {
        self.registrations.push((token, interests));
        Ok(())
    }

    fn reregister(
        &mut self,
        _registry: &Registry,
        token: Token,
        interests: Interest,
    ) -> io::Result<()> {
        self.reregistrations.push((token, interests));
        Ok(())
    }

    fn deregister(&mut self, _registry: &Registry) -> io::Result<()> {
        self.deregister_count += 1;
        Ok(())
    }
}

#[test]
fn poll_registration() {
    init();
    let poll = Poll::new().unwrap();
    let registry = poll.registry();

    let mut source = TestEventSource::new();
    let token = Token(0);
    let interests = Interest::READABLE;
    registry.register(&mut source, token, interests).unwrap();
    assert_eq!(source.registrations.len(), 1);
    assert_eq!(source.registrations.get(0), Some(&(token, interests)));
    assert!(source.reregistrations.is_empty());
    assert_eq!(source.deregister_count, 0);

    let re_token = Token(0);
    let re_interests = Interest::READABLE;
    registry
        .reregister(&mut source, re_token, re_interests)
        .unwrap();
    assert_eq!(source.registrations.len(), 1);
    assert_eq!(source.reregistrations.len(), 1);
    assert_eq!(
        source.reregistrations.get(0),
        Some(&(re_token, re_interests))
    );
    assert_eq!(source.deregister_count, 0);

    registry.deregister(&mut source).unwrap();
    assert_eq!(source.registrations.len(), 1);
    assert_eq!(source.reregistrations.len(), 1);
    assert_eq!(source.deregister_count, 1);
}

struct ErroneousTestEventSource;

impl event::Source for ErroneousTestEventSource {
    fn register(
        &mut self,
        _registry: &Registry,
        _token: Token,
        _interests: Interest,
    ) -> io::Result<()> {
        Err(io::Error::new(io::ErrorKind::Other, "register"))
    }

    fn reregister(
        &mut self,
        _registry: &Registry,
        _token: Token,
        _interests: Interest,
    ) -> io::Result<()> {
        Err(io::Error::new(io::ErrorKind::Other, "reregister"))
    }

    fn deregister(&mut self, _registry: &Registry) -> io::Result<()> {
        Err(io::Error::new(io::ErrorKind::Other, "deregister"))
    }
}

#[test]
fn poll_erroneous_registration() {
    init();
    let poll = Poll::new().unwrap();
    let registry = poll.registry();

    let mut source = ErroneousTestEventSource;
    let token = Token(0);
    let interests = Interest::READABLE;
    assert_error(registry.register(&mut source, token, interests), "register");
    assert_error(
        registry.reregister(&mut source, token, interests),
        "reregister",
    );
    assert_error(registry.deregister(&mut source), "deregister");
}

/// Assert that `result` is an error and the formatted error (via
/// `fmt::Display`) equals `expected_msg`.
pub fn assert_error<T, E: fmt::Display>(result: Result<T, E>, expected_msg: &str) {
    match result {
        Ok(_) => panic!("unexpected OK result"),
        Err(err) => assert!(
            err.to_string().contains(expected_msg),
            "wanted: {}, got: {}",
            err,
            expected_msg
        ),
    }
}
<commit_msg>added a test with registry ops on multiple threads<commit_after>#![cfg(all(feature = "os-poll", feature = "tcp"))]

use mio::event::Source;
use mio::net::{TcpListener, TcpStream, UdpSocket};
use mio::{event, Events, Interest, Poll, Registry, Token};
use std::net;
use std::sync::{Arc, Barrier};
use std::thread::{self, sleep};
use std::time::Duration;
use std::{fmt, io};

mod util;

use util::{
    any_local_address, assert_send, assert_sync, expect_events, init, init_with_poll, ExpectEvent,
};

const ID1: Token = Token(1);
const ID2: Token = Token(2);
const ID3: Token = Token(3);

#[test]
fn is_send_and_sync() {
    assert_send::<Events>();
    assert_sync::<Events>();

    assert_sync::<Poll>();
    assert_send::<Poll>();

    assert_sync::<Registry>();
    assert_send::<Registry>();
}

#[test]
fn run_once_with_nothing() {
    init();

    let mut events = Events::with_capacity(16);
    let mut poll = Poll::new().unwrap();
    poll.poll(&mut events, Some(Duration::from_millis(100)))
        .unwrap();
}

#[test]
fn add_then_drop() {
    init();

    let mut events = Events::with_capacity(16);
    let mut l = TcpListener::bind(any_local_address()).unwrap();
    let mut poll = Poll::new().unwrap();
    poll.registry()
        .register(&mut l, Token(1), Interest::READABLE | Interest::WRITABLE)
        .unwrap();
    drop(l);
    poll.poll(&mut events, Some(Duration::from_millis(100)))
        .unwrap();
}

#[test]
fn zero_duration_polls_events() {
    init();

    let mut poll = Poll::new().unwrap();
    let mut events = Events::with_capacity(16);

    let listener = net::TcpListener::bind(any_local_address()).unwrap();
    let addr = listener.local_addr().unwrap();

    let streams: Vec<TcpStream> = (0..3)
        .map(|n| {
            let mut stream = TcpStream::connect(addr).unwrap();
            poll.registry()
                .register(&mut stream, Token(n), Interest::WRITABLE)
                .unwrap();
            stream
        })
        .collect();

    // Ensure the TcpStreams have some time to connection and for the events to
    // show up.
    sleep(Duration::from_millis(10));

    // Even when passing a zero duration timeout we still want do the system
    // call.
    poll.poll(&mut events, Some(Duration::from_nanos(0)))
        .unwrap();
    assert!(!events.is_empty());

    // Both need to live until here.
    drop(streams);
    drop(listener);
}

#[test]
fn poll_closes_fd() {
    init();

    for _ in 0..2000 {
        let mut poll = Poll::new().unwrap();
        let mut events = Events::with_capacity(4);

        poll.poll(&mut events, Some(Duration::from_millis(0)))
            .unwrap();

        drop(poll);
    }
}

#[test]
fn drop_cancels_interest_and_shuts_down() {
    init();

    use mio::net::TcpStream;
    use std::io;
    use std::io::Read;
    use std::net::TcpListener;
    use std::thread;

    let l = TcpListener::bind("127.0.0.1:0").unwrap();
    let addr = l.local_addr().unwrap();

    let t = thread::spawn(move || {
        let mut s = l.incoming().next().unwrap().unwrap();
        s.set_read_timeout(Some(Duration::from_secs(5)))
            .expect("set_read_timeout");
        let r = s.read(&mut [0; 16]);
        match r {
            Ok(_) => (),
            Err(e) => {
                if e.kind() != io::ErrorKind::UnexpectedEof {
                    panic!(e);
                }
            }
        }
    });

    let mut poll = Poll::new().unwrap();
    let mut s = TcpStream::connect(addr).unwrap();

    poll.registry()
        .register(&mut s, Token(1), Interest::READABLE | Interest::WRITABLE)
        .unwrap();
    let mut events = Events::with_capacity(16);
    'outer: loop {
        poll.poll(&mut events, None).unwrap();
        for event in &events {
            if event.token() == Token(1) {
                // connected
                break 'outer;
            }
        }
    }

    let mut b = [0; 1024];
    match s.read(&mut b) {
        Ok(_) => panic!("unexpected ok"),
        Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => (),
        Err(e) => panic!("unexpected error: {:?}", e),
    }

    drop(s);
    t.join().unwrap();
}

#[test]
fn registry_behind_arc() {
    // `Registry` should work behind an `Arc`, being `Sync` and `Send`.
    init();

    let mut poll = Poll::new().unwrap();
    let registry = Arc::new(poll.registry().try_clone().unwrap());
    let mut events = Events::with_capacity(128);

    let mut listener = TcpListener::bind(any_local_address()).unwrap();
    let addr = listener.local_addr().unwrap();
    let barrier = Arc::new(Barrier::new(3));

    let registry2 = Arc::clone(&registry);
    let registry3 = Arc::clone(&registry);
    let barrier2 = Arc::clone(&barrier);
    let barrier3 = Arc::clone(&barrier);

    let handle1 = thread::spawn(move || {
        registry2
            .register(&mut listener, Token(0), Interest::READABLE)
            .unwrap();
        barrier2.wait();
    });
    let handle2 = thread::spawn(move || {
        let mut stream = TcpStream::connect(addr).unwrap();
        registry3
            .register(
                &mut stream,
                Token(1),
                Interest::READABLE | Interest::WRITABLE,
            )
            .unwrap();
        barrier3.wait();
    });

    poll.poll(&mut events, Some(Duration::from_millis(1000)))
        .unwrap();
    assert!(events.iter().count() >= 1);

    // Let the threads return.
    barrier.wait();

    handle1.join().unwrap();
    handle2.join().unwrap();
}

/// Call all registration operations, ending with `source` being registered with `token` and `final_interests`.
pub fn registry_ops_flow(
    registry: &Registry,
    source: &mut dyn Source,
    token: Token,
    init_interests: Interest,
    final_interests: Interest,
) -> io::Result<()> {
    registry.register(source, token, init_interests).unwrap();
    registry.deregister(source).unwrap();

    registry.register(source, token, init_interests).unwrap();
    registry.reregister(source, token, final_interests)
}

#[test]
fn registry_operations_are_thread_safe() {
    let (mut poll, mut events) = init_with_poll();

    let registry = Arc::new(poll.registry().try_clone().unwrap());
    let registry1 = Arc::clone(&registry);
    let registry2 = Arc::clone(&registry);
    let registry3 = Arc::clone(&registry);

    let barrier = Arc::new(Barrier::new(4));
    let barrier1 = Arc::clone(&barrier);
    let barrier2 = Arc::clone(&barrier);
    let barrier3 = Arc::clone(&barrier);

    let mut listener = TcpListener::bind(any_local_address()).unwrap();
    let addr = listener.local_addr().unwrap();

    // Expect that multiple register/deregister/reregister work fine on multiple
    // threads. Main thread will wait before the expect_events for all other 3
    // threads to do their work. Otherwise expect_events timeout might be too short
    // for all threads to complete, and call might fail.

    let handle1 = thread::spawn(move || {
        registry_ops_flow(
            &registry1,
            &mut listener,
            ID1,
            Interest::READABLE,
            Interest::READABLE,
        )
        .unwrap();

        barrier1.wait();
        barrier1.wait();
    });

    let handle2 = thread::spawn(move || {
        let mut udp_socket = UdpSocket::bind(any_local_address()).unwrap();
        registry_ops_flow(
            &registry2,
            &mut udp_socket,
            ID2,
            Interest::WRITABLE,
            Interest::WRITABLE.add(Interest::READABLE),
        )
        .unwrap();

        barrier2.wait();
        barrier2.wait();
    });

    let handle3 = thread::spawn(move || {
        let mut stream = TcpStream::connect(addr).unwrap();
        registry_ops_flow(
            &registry3,
            &mut stream,
            ID3,
            Interest::READABLE,
            Interest::READABLE | Interest::WRITABLE,
        )
        .unwrap();

        barrier3.wait();
        barrier3.wait();
    });

    // wait for threads to finish before expect_events
    barrier.wait();
    expect_events(
        &mut poll,
        &mut events,
        vec![
            ExpectEvent::new(ID1, Interest::READABLE),
            ExpectEvent::new(ID2, Interest::WRITABLE),
            ExpectEvent::new(ID3, Interest::WRITABLE),
        ],
    );

    // Let the threads return.
    barrier.wait();

    handle1.join().unwrap();
    handle2.join().unwrap();
    handle3.join().unwrap();
}

// On kqueue platforms registering twice (not *re*registering) works.
#[test]
#[cfg(any(target_os = "linux", target_os = "windows"))]
pub fn double_register() {
    init();
    let poll = Poll::new().unwrap();

    let mut l = TcpListener::bind("127.0.0.1:0".parse().unwrap()).unwrap();

    poll.registry()
        .register(&mut l, Token(0), Interest::READABLE)
        .unwrap();

    assert!(poll
        .registry()
        .register(&mut l, Token(1), Interest::READABLE)
        .is_err());
}

struct TestEventSource {
    registrations: Vec<(Token, Interest)>,
    reregistrations: Vec<(Token, Interest)>,
    deregister_count: usize,
}

impl TestEventSource {
    fn new() -> TestEventSource {
        TestEventSource {
            registrations: Vec::new(),
            reregistrations: Vec::new(),
            deregister_count: 0,
        }
    }
}

impl event::Source for TestEventSource {
    fn register(
        &mut self,
        _registry: &Registry,
        token: Token,
        interests: Interest,
    ) -> io::Result<()> {
        self.registrations.push((token, interests));
        Ok(())
    }

    fn reregister(
        &mut self,
        _registry: &Registry,
        token: Token,
        interests: Interest,
    ) -> io::Result<()> {
        self.reregistrations.push((token, interests));
        Ok(())
    }

    fn deregister(&mut self, _registry: &Registry) -> io::Result<()> {
        self.deregister_count += 1;
        Ok(())
    }
}

#[test]
fn poll_registration() {
    init();
    let poll = Poll::new().unwrap();
    let registry = poll.registry();

    let mut source = TestEventSource::new();
    let token = Token(0);
    let interests = Interest::READABLE;
    registry.register(&mut source, token, interests).unwrap();
    assert_eq!(source.registrations.len(), 1);
    assert_eq!(source.registrations.get(0), Some(&(token, interests)));
    assert!(source.reregistrations.is_empty());
    assert_eq!(source.deregister_count, 0);

    let re_token = Token(0);
    let re_interests = Interest::READABLE;
    registry
        .reregister(&mut source, re_token, re_interests)
        .unwrap();
    assert_eq!(source.registrations.len(), 1);
    assert_eq!(source.reregistrations.len(), 1);
    assert_eq!(
        source.reregistrations.get(0),
        Some(&(re_token, re_interests))
    );
    assert_eq!(source.deregister_count, 0);

    registry.deregister(&mut source).unwrap();
    assert_eq!(source.registrations.len(), 1);
    assert_eq!(source.reregistrations.len(), 1);
    assert_eq!(source.deregister_count, 1);
}

struct ErroneousTestEventSource;

impl event::Source for ErroneousTestEventSource {
    fn register(
        &mut self,
        _registry: &Registry,
        _token: Token,
        _interests: Interest,
    ) -> io::Result<()> {
        Err(io::Error::new(io::ErrorKind::Other, "register"))
    }

    fn reregister(
        &mut self,
        _registry: &Registry,
        _token: Token,
        _interests: Interest,
    ) -> io::Result<()> {
        Err(io::Error::new(io::ErrorKind::Other, "reregister"))
    }

    fn deregister(&mut self, _registry: &Registry) -> io::Result<()> {
        Err(io::Error::new(io::ErrorKind::Other, "deregister"))
    }
}

#[test]
fn poll_erroneous_registration() {
    init();
    let poll = Poll::new().unwrap();
    let registry = poll.registry();

    let mut source = ErroneousTestEventSource;
    let token = Token(0);
    let interests = Interest::READABLE;
    assert_error(registry.register(&mut source, token, interests), "register");
    assert_error(
        registry.reregister(&mut source, token, interests),
        "reregister",
    );
    assert_error(registry.deregister(&mut source), "deregister");
}

/// Assert that `result` is an error and the formatted error (via
/// `fmt::Display`) equals `expected_msg`.
pub fn assert_error<T, E: fmt::Display>(result: Result<T, E>, expected_msg: &str) {
    match result {
        Ok(_) => panic!("unexpected OK result"),
        Err(err) => assert!(
            err.to_string().contains(expected_msg),
            "wanted: {}, got: {}",
            err,
            expected_msg
        ),
    }
}
<|endoftext|>


<commit_before>
extern crate ndarray;

use ndarray::{RcArray, Dimension};

#[test]
fn broadcast_1()
{
    let a_dim = (2, 4, 2, 2);
    let b_dim = (2, 1, 2, 1);
    let a = RcArray::linspace(0., 1., a_dim.size()).reshape(a_dim);
    let b = RcArray::linspace(0., 1., b_dim.size()).reshape(b_dim);
    assert!(b.broadcast(a.dim()).is_some());

    let c_dim = (2, 1);
    let c = RcArray::linspace(0., 1., c_dim.size()).reshape(c_dim);
    assert!(c.broadcast(1).is_none());
    assert!(c.broadcast(()).is_none());
    assert!(c.broadcast((2, 1)).is_some());
    assert!(c.broadcast((2, 2)).is_some());
    assert!(c.broadcast((32, 2, 1)).is_some());
    assert!(c.broadcast((32, 1, 2)).is_none());

    /* () can be broadcast to anything */
    let z = RcArray::<f32,_>::zeros(());
    assert!(z.broadcast(()).is_some());
    assert!(z.broadcast(1).is_some());
    assert!(z.broadcast(3).is_some());
    assert!(z.broadcast((7,2,9)).is_some());
}

#[test]
fn test_add()
{
    let a_dim = (2, 4, 2, 2);
    let b_dim = (2, 1, 2, 1);
    let mut a = RcArray::linspace(0.0, 1., a_dim.size()).reshape(a_dim);
    let b = RcArray::linspace(0.0, 1., b_dim.size()).reshape(b_dim);
    a += &b;
    let t = RcArray::from_elem((), 1.0f32);
    a += &t;
}

#[test] #[should_panic]
fn test_add_incompat()
{
    let a_dim = (2, 4, 2, 2);
    let mut a = RcArray::linspace(0.0, 1., a_dim.size()).reshape(a_dim);
    let incompat = RcArray::from_elem(3, 1.0f32);
    a += &incompat;
}
<commit_msg>Add test for broadcast<commit_after>
extern crate ndarray;

use ndarray::prelude::*;

#[test]
fn broadcast_1()
{
    let a_dim = (2, 4, 2, 2);
    let b_dim = (2, 1, 2, 1);
    let a = RcArray::linspace(0., 1., a_dim.size()).reshape(a_dim);
    let b = RcArray::linspace(0., 1., b_dim.size()).reshape(b_dim);
    assert!(b.broadcast(a.dim()).is_some());

    let c_dim = (2, 1);
    let c = RcArray::linspace(0., 1., c_dim.size()).reshape(c_dim);
    assert!(c.broadcast(1).is_none());
    assert!(c.broadcast(()).is_none());
    assert!(c.broadcast((2, 1)).is_some());
    assert!(c.broadcast((2, 2)).is_some());
    assert!(c.broadcast((32, 2, 1)).is_some());
    assert!(c.broadcast((32, 1, 2)).is_none());

    /* () can be broadcast to anything */
    let z = RcArray::<f32,_>::zeros(());
    assert!(z.broadcast(()).is_some());
    assert!(z.broadcast(1).is_some());
    assert!(z.broadcast(3).is_some());
    assert!(z.broadcast((7,2,9)).is_some());
}

#[test]
fn test_add()
{
    let a_dim = (2, 4, 2, 2);
    let b_dim = (2, 1, 2, 1);
    let mut a = RcArray::linspace(0.0, 1., a_dim.size()).reshape(a_dim);
    let b = RcArray::linspace(0.0, 1., b_dim.size()).reshape(b_dim);
    a += &b;
    let t = RcArray::from_elem((), 1.0f32);
    a += &t;
}

#[test] #[should_panic]
fn test_add_incompat()
{
    let a_dim = (2, 4, 2, 2);
    let mut a = RcArray::linspace(0.0, 1., a_dim.size()).reshape(a_dim);
    let incompat = RcArray::from_elem(3, 1.0f32);
    a += &incompat;
}

#[test]
fn test_broadcast() {
    let (_, n, k) = (16, 16, 16);
    let x1 = 1.;
    // b0 broadcast 1 -> n, k
    let x = Array::from_vec(vec![x1]);
    let b0 = x.broadcast((n, k)).unwrap();
    // b1 broadcast n -> n, k
    let b1 = Array::from_elem(n, x1);
    let b1 = b1.broadcast((n, k)).unwrap();
    // b2 is n, k
    let b2 = Array::from_elem((n, k), x1);

    println!("b0=\n{:?}", b0);
    println!("b1=\n{:?}", b1);
    println!("b2=\n{:?}", b2);
    assert_eq!(b0, b1);
    assert_eq!(b0, b2);
}
<|endoftext|>