<commit_before>use std::io::BufReader;
use std::fs::File;

use serde_json;

use cgmath::*;
use shape::*;
use vec3d::*;
use light::*;

use color;
use color::Color;

use std;

#[derive(Serialize, Deserialize)]
pub struct Scene {
    shapes: Vec<Shape>,
    light: Light,
    camera_pos: Vec3d,
    camera_dir: Vec3d,
    camera_up: Vec3d,
    max_reflection: i32,
}

impl Scene {
    pub fn new(light: Light, camera_pos: Vec3d, camera_dir: Vec3d, camera_up: Vec3d) -> Scene {
        Scene {
            shapes: Vec::new(),
            light: light,
            camera_pos: camera_pos,
            camera_dir: camera_dir,
            camera_up: camera_up,
            max_reflection: 4,
        }
    }

    pub fn from_file(filename: &str) -> Result<Scene, String> {
        let file_in = File::open(filename)
            .map_err(|e| format!("Can't open input file: {}", e))?;
        let reader = BufReader::new(&file_in);

        let s: Scene = serde_json::from_reader(reader)
            .map_err(|e| format!("JSON Parse error: {}", e))?;
        Ok(s)
    }

    pub fn add_shape(&mut self, s: Shape) {
        self.shapes.push(s);
    }

    pub fn step(&mut self) {
        // self.light.translate(&[1.0, 0.0, 0.0]);
    }

    pub fn mv_camera_fwd(&mut self) {
        self.camera_pos = self.camera_pos + self.camera_dir;
    }

    pub fn mv_camera_back(&mut self) {
        self.camera_pos = self.camera_pos - self.camera_dir;
    }

    pub fn rot_camera(&mut self, x_rot: f64, y_rot: f64) {
        self.camera_dir = vec3_roty(self.camera_dir, x_rot);

        // project camera_dir to X plane
        let mut x_proj = self.camera_dir;
        x_proj.x = 0.0;
        x_proj = x_proj.normalize();

        let angle = if x_proj.x < 0.0 {
            -Vec3d::new(0.0, 0.0, 1.0).dot(x_proj).acos()
        } else {
            Vec3d::new(0.0, 0.0, 1.0).dot(x_proj).acos()
        };

        self.camera_dir = vec3_roty(self.camera_dir, -angle);
        self.camera_dir = vec3_rotx(self.camera_dir, y_rot);
        self.camera_dir = vec3_roty(self.camera_dir, angle);

        self.camera_dir = self.camera_dir.normalize();
    }

    /// Returns an iterator that scans a full frame at given resolution
    pub fn draw_iter(&self, h: usize, v: usize) -> DrawIter {
        DrawIter::new(self, h, v)
    }

    /// draw_iter converted to (u8, u8, u8)
    pub fn draw_iter_u8(&self,
                        h: usize,
                        v: usize)
                        -> std::iter::Map<DrawIter, fn(Color) -> (u8, u8, u8)> {
        self.draw_iter(h, v).map(Color::col_to_u8)
    }

    /// Returns an iterator for a line at given resolution
    pub fn line_iter(&self, h: usize, v: usize, l: usize) -> LineIter {
        LineIter::new(self, h, v, l)
    }

    /// line_iter converted to (u8, u8, u8)
    pub fn line_iter_u8(&self,
                        h: usize,
                        v: usize,
                        l: usize)
                        -> std::iter::Map<LineIter, fn(Color) -> (u8, u8, u8)> {
        self.line_iter(h, v, l).map(Color::col_to_u8)
    }

    // Recursively trace lines
    fn trace(&self, point: &Vec3d, dir: &Vec3d, depth: i32) -> Color {
        if depth >= self.max_reflection {
            return color::BLACK;
        }

        if let Some(intersect) = self.closest_q(point, dir) {
            let feeler_d = self.light.get_pos() - intersect.point;
            let dist_light = feeler_d.magnitude();
            let feeler_d_unit = feeler_d.normalize();
            let direct_light = self.is_direct_light(&intersect.point, &feeler_d_unit, dist_light);

            let local = phong(point, &intersect, &self.light, direct_light);

            let tmp = (intersect.point - point).normalize();
            let reflection_dir = tmp - (intersect.normal * 2.0 * tmp.dot(intersect.normal));

            let reflected = self.trace(&intersect.point, &reflection_dir, depth + 1);

            Color::add(&local,
                       &(Color::scale(&reflected, intersect.material.reflectivity)))
        } else {
            color::BLACK
        }
    }

    // Is there anything on the path to the light
    fn is_direct_light(&self, point: &Vec3d, dir: &Vec3d, dl: f64) -> bool {
        !self.shapes
             .iter()
             .map(|&ref x| x.intersect_dist(point, dir))
             .any(|dist_opt| dist_opt.map_or(false, |i| i < dl))
    }

    // Checks against all objects and returns closest intersection
    fn closest_q(&self, point: &Vec3d, dir: &Vec3d) -> Option<Intersection> {
        let find_min_opt = |min, (idx, op_val)| {
            match (min, op_val) {
                (Some((_, min_val)), Some(val)) if val < min_val => Some((idx, val)),
                (None, Some(val)) => Some((idx, val)),
                _ => min,
            }
        };

        let closest = self.shapes
            .iter()
            .map(|x| x.intersect_dist(point, dir))
            .enumerate()
            .fold(None, find_min_opt);

        // get intersection point info
        if let Some((idx, _)) = closest {
            self.shapes[idx].intersect(point, dir)
        } else {
            None
        }
    }
}

/// Iterator that iterates over a full frame
pub struct DrawIter<'a> {
    h_res: usize,
    v_res: usize,
    x: usize,
    y: usize,
    left: Vec3d,
    right_step: Vec3d,
    down_step: Vec3d,
    point: Vec3d,
    scene: &'a Scene,
}

impl<'a> DrawIter<'a> {
    fn new(scene: &Scene, h: usize, v: usize) -> DrawIter {
        let left = scene.camera_dir.cross(scene.camera_up).normalize();
        let up = left.cross(scene.camera_dir).normalize();
        let point = left + up + scene.camera_dir;
        let right_step = left * (-2.0 / h as f64);
        let down_step = up * (-2.0 / v as f64);
        DrawIter {
            h_res: h,
            v_res: v,
            x: 0,
            y: 0,
            left: left,
            point: point,
            right_step: right_step,
            down_step: down_step,
            scene: scene,
        }
    }
}

impl<'a> Iterator for DrawIter<'a> {
    type Item = Color;

    fn next(&mut self) -> Option<Self::Item> {
        if self.y >= self.v_res {
            return None;
        }

        self.point = self.point + self.right_step;

        let c = self.scene
            .trace(&self.scene.camera_pos, &self.point.normalize(), 0);

        self.x += 1;
        if self.x >= self.h_res {
            self.x = 0;
            self.y += 1;
            self.point = self.point + self.left;
            self.point = self.point + self.left;
            self.point = self.point + self.down_step;
        }

        Some(c)
    }
}

/// Iterator that iterates over a single line
pub struct LineIter<'a> {
    h_res: usize,
    x: usize,
    right_step: Vec3d,
    point: Vec3d,
    scene: &'a Scene,
}

impl<'a> LineIter<'a> {
    fn new(scene: &Scene, h: usize, v: usize, l: usize) -> LineIter {
        let left = scene.camera_dir.cross(scene.camera_up).normalize();
        let up = left.cross(scene.camera_dir).normalize();
        let mut point = left + up + scene.camera_dir;
        let right_step = left * (-2.0 / h as f64);
        let down_step = up * (-2.0 / v as f64);

        point = point + down_step * (l as f64);
        LineIter {
            h_res: h,
            x: 0,
            point: point,
            right_step: right_step,
            scene: scene,
        }
    }
}

impl<'a> Iterator for LineIter<'a> {
    type Item = Color;

    fn next(&mut self) -> Option<Self::Item> {
        if self.x >= self.h_res {
            return None;
        }
        self.x += 1;

        self.point = self.point + self.right_step;

        let c = self.scene
            .trace(&self.scene.camera_pos, &self.point.normalize(), 0);

        Some(c)
    }
}

// Calculate color with Phong model
fn phong(view_point: &Vec3d,
         intersection: &Intersection,
         light: &Light,
         direct_light: bool)
         -> Color {

    let point_material = intersection.material;
    let light_pos = light.get_pos();
    let light_color = light.get_color();
    let mut d = (light_pos - intersection.point)
        .normalize()
        .dot(intersection.normal);

    if d < 0.0 {
        d = 0.0;
    }

    let r_diffuse = if direct_light {
        light_color.0 * point_material.diffuse_color.0 * d
    } else {
        0.0
    };
    let g_diffuse = if direct_light {
        light_color.1 * point_material.diffuse_color.1 * d
    } else {
        0.0
    };
    let b_diffuse = if direct_light {
        light_color.2 * point_material.diffuse_color.2 * d
    } else {
        0.0
    };

    let r_ambient = point_material.ambient_color.0;
    let g_ambient = point_material.ambient_color.1;
    let b_ambient = point_material.ambient_color.2;

    let v = (view_point - intersection.point).normalize();
    let lt = (light_pos - intersection.point).normalize();
    let r = intersection.normal * (2.0 * intersection.normal.dot(lt)) - lt;

    let mut s = r.dot(v).powf(point_material.shininess);
    if s < 0.0 {
        s = 0.0;
    }

    let r_specular = if direct_light {
        light_color.0 * point_material.specular_color.0 * s
    } else {
        0.0
    };
    let g_specular = if direct_light {
        light_color.1 * point_material.specular_color.1 * s
    } else {
        0.0
    };
    let b_specular = if direct_light {
        light_color.2 * point_material.specular_color.2 * s
    } else {
        0.0
    };

    let mut r_sum = r_ambient + r_diffuse + r_specular;
    let mut g_sum = g_ambient + g_diffuse + g_specular;
    let mut b_sum = b_ambient + b_diffuse + b_specular;

    if r_sum > 1.0 {
        r_sum = 1.0;
    }

    if g_sum > 1.0 {
        g_sum = 1.0;
    }

    if b_sum > 1.0 {
        b_sum = 1.0;
    }

    Color(r_sum, g_sum, b_sum)
}
<commit_msg>Fix mouse look up/down<commit_after>use std::io::BufReader;
use std::fs::File;

use serde_json;

use cgmath::*;
use shape::*;
use vec3d::*;
use light::*;

use color;
use color::Color;

use std;

#[derive(Serialize, Deserialize)]
pub struct Scene {
    shapes: Vec<Shape>,
    light: Light,
    camera_pos: Vec3d,
    camera_dir: Vec3d,
    camera_up: Vec3d,
    max_reflection: i32,
}

impl Scene {
    pub fn new(light: Light, camera_pos: Vec3d, camera_dir: Vec3d, camera_up: Vec3d) -> Scene {
        Scene {
            shapes: Vec::new(),
            light: light,
            camera_pos: camera_pos,
            camera_dir: camera_dir,
            camera_up: camera_up,
            max_reflection: 4,
        }
    }

    pub fn from_file(filename: &str) -> Result<Scene, String> {
        let file_in = File::open(filename)
            .map_err(|e| format!("Can't open input file: {}", e))?;
        let reader = BufReader::new(&file_in);

        let s: Scene = serde_json::from_reader(reader)
            .map_err(|e| format!("JSON Parse error: {}", e))?;
        Ok(s)
    }

    pub fn add_shape(&mut self, s: Shape) {
        self.shapes.push(s);
    }

    pub fn step(&mut self) {
        // self.light.translate(&[1.0, 0.0, 0.0]);
    }

    pub fn mv_camera_fwd(&mut self) {
        self.camera_pos = self.camera_pos + self.camera_dir;
    }

    pub fn mv_camera_back(&mut self) {
        self.camera_pos = self.camera_pos - self.camera_dir;
    }

    pub fn rot_camera(&mut self, x_rot: f64, y_rot: f64) {
        self.camera_dir = vec3_roty(self.camera_dir, x_rot);

        // project camera_dir to X plane
        let mut x_proj = self.camera_dir;
        x_proj.y = 0.0;
        x_proj = x_proj.normalize();

        let angle = if x_proj.x < 0.0 {
            -Vec3d::new(0.0, 0.0, 1.0).dot(x_proj).acos()
        } else {
            Vec3d::new(0.0, 0.0, 1.0).dot(x_proj).acos()
        };

        self.camera_dir = vec3_roty(self.camera_dir, -angle);
        self.camera_dir = vec3_rotx(self.camera_dir, y_rot);
        self.camera_dir = vec3_roty(self.camera_dir, angle);

        self.camera_dir = self.camera_dir.normalize();
    }

    /// Returns an iterator that scans a full frame at given resolution
    pub fn draw_iter(&self, h: usize, v: usize) -> DrawIter {
        DrawIter::new(self, h, v)
    }

    /// draw_iter converted to (u8, u8, u8)
    pub fn draw_iter_u8(&self,
                        h: usize,
                        v: usize)
                        -> std::iter::Map<DrawIter, fn(Color) -> (u8, u8, u8)> {
        self.draw_iter(h, v).map(Color::col_to_u8)
    }

    /// Returns an iterator for a line at given resolution
    pub fn line_iter(&self, h: usize, v: usize, l: usize) -> LineIter {
        LineIter::new(self, h, v, l)
    }

    /// line_iter converted to (u8, u8, u8)
    pub fn line_iter_u8(&self,
                        h: usize,
                        v: usize,
                        l: usize)
                        -> std::iter::Map<LineIter, fn(Color) -> (u8, u8, u8)> {
        self.line_iter(h, v, l).map(Color::col_to_u8)
    }

    // Recursively trace lines
    fn trace(&self, point: &Vec3d, dir: &Vec3d, depth: i32) -> Color {
        if depth >= self.max_reflection {
            return color::BLACK;
        }

        if let Some(intersect) = self.closest_q(point, dir) {
            let feeler_d = self.light.get_pos() - intersect.point;
            let dist_light = feeler_d.magnitude();
            let feeler_d_unit = feeler_d.normalize();
            let direct_light = self.is_direct_light(&intersect.point, &feeler_d_unit, dist_light);

            let local = phong(point, &intersect, &self.light, direct_light);

            let tmp = (intersect.point - point).normalize();
            let reflection_dir = tmp - (intersect.normal * 2.0 * tmp.dot(intersect.normal));

            let reflected = self.trace(&intersect.point, &reflection_dir, depth + 1);

            Color::add(&local,
                       &(Color::scale(&reflected, intersect.material.reflectivity)))
        } else {
            color::BLACK
        }
    }

    // Is there anything on the path to the light
    fn is_direct_light(&self, point: &Vec3d, dir: &Vec3d, dl: f64) -> bool {
        !self.shapes
             .iter()
             .map(|&ref x| x.intersect_dist(point, dir))
             .any(|dist_opt| dist_opt.map_or(false, |i| i < dl))
    }

    // Checks against all objects and returns closest intersection
    fn closest_q(&self, point: &Vec3d, dir: &Vec3d) -> Option<Intersection> {
        let find_min_opt = |min, (idx, op_val)| {
            match (min, op_val) {
                (Some((_, min_val)), Some(val)) if val < min_val => Some((idx, val)),
                (None, Some(val)) => Some((idx, val)),
                _ => min,
            }
        };

        let closest = self.shapes
            .iter()
            .map(|x| x.intersect_dist(point, dir))
            .enumerate()
            .fold(None, find_min_opt);

        // get intersection point info
        if let Some((idx, _)) = closest {
            self.shapes[idx].intersect(point, dir)
        } else {
            None
        }
    }
}

/// Iterator that iterates over a full frame
pub struct DrawIter<'a> {
    h_res: usize,
    v_res: usize,
    x: usize,
    y: usize,
    left: Vec3d,
    right_step: Vec3d,
    down_step: Vec3d,
    point: Vec3d,
    scene: &'a Scene,
}

impl<'a> DrawIter<'a> {
    fn new(scene: &Scene, h: usize, v: usize) -> DrawIter {
        let left = scene.camera_dir.cross(scene.camera_up).normalize();
        let up = left.cross(scene.camera_dir).normalize();
        let point = left + up + scene.camera_dir;
        let right_step = left * (-2.0 / h as f64);
        let down_step = up * (-2.0 / v as f64);
        DrawIter {
            h_res: h,
            v_res: v,
            x: 0,
            y: 0,
            left: left,
            point: point,
            right_step: right_step,
            down_step: down_step,
            scene: scene,
        }
    }
}

impl<'a> Iterator for DrawIter<'a> {
    type Item = Color;

    fn next(&mut self) -> Option<Self::Item> {
        if self.y >= self.v_res {
            return None;
        }

        self.point = self.point + self.right_step;

        let c = self.scene
            .trace(&self.scene.camera_pos, &self.point.normalize(), 0);

        self.x += 1;
        if self.x >= self.h_res {
            self.x = 0;
            self.y += 1;
            self.point = self.point + self.left;
            self.point = self.point + self.left;
            self.point = self.point + self.down_step;
        }

        Some(c)
    }
}

/// Iterator that iterates over a single line
pub struct LineIter<'a> {
    h_res: usize,
    x: usize,
    right_step: Vec3d,
    point: Vec3d,
    scene: &'a Scene,
}

impl<'a> LineIter<'a> {
    fn new(scene: &Scene, h: usize, v: usize, l: usize) -> LineIter {
        let left = scene.camera_dir.cross(scene.camera_up).normalize();
        let up = left.cross(scene.camera_dir).normalize();
        let mut point = left + up + scene.camera_dir;
        let right_step = left * (-2.0 / h as f64);
        let down_step = up * (-2.0 / v as f64);

        point = point + down_step * (l as f64);
        LineIter {
            h_res: h,
            x: 0,
            point: point,
            right_step: right_step,
            scene: scene,
        }
    }
}

impl<'a> Iterator for LineIter<'a> {
    type Item = Color;

    fn next(&mut self) -> Option<Self::Item> {
        if self.x >= self.h_res {
            return None;
        }
        self.x += 1;

        self.point = self.point + self.right_step;

        let c = self.scene
            .trace(&self.scene.camera_pos, &self.point.normalize(), 0);

        Some(c)
    }
}

// Calculate color with Phong model
fn phong(view_point: &Vec3d,
         intersection: &Intersection,
         light: &Light,
         direct_light: bool)
         -> Color {

    let point_material = intersection.material;
    let light_pos = light.get_pos();
    let light_color = light.get_color();
    let mut d = (light_pos - intersection.point)
        .normalize()
        .dot(intersection.normal);

    if d < 0.0 {
        d = 0.0;
    }

    let r_diffuse = if direct_light {
        light_color.0 * point_material.diffuse_color.0 * d
    } else {
        0.0
    };
    let g_diffuse = if direct_light {
        light_color.1 * point_material.diffuse_color.1 * d
    } else {
        0.0
    };
    let b_diffuse = if direct_light {
        light_color.2 * point_material.diffuse_color.2 * d
    } else {
        0.0
    };

    let r_ambient = point_material.ambient_color.0;
    let g_ambient = point_material.ambient_color.1;
    let b_ambient = point_material.ambient_color.2;

    let v = (view_point - intersection.point).normalize();
    let lt = (light_pos - intersection.point).normalize();
    let r = intersection.normal * (2.0 * intersection.normal.dot(lt)) - lt;

    let mut s = r.dot(v).powf(point_material.shininess);
    if s < 0.0 {
        s = 0.0;
    }

    let r_specular = if direct_light {
        light_color.0 * point_material.specular_color.0 * s
    } else {
        0.0
    };
    let g_specular = if direct_light {
        light_color.1 * point_material.specular_color.1 * s
    } else {
        0.0
    };
    let b_specular = if direct_light {
        light_color.2 * point_material.specular_color.2 * s
    } else {
        0.0
    };

    let mut r_sum = r_ambient + r_diffuse + r_specular;
    let mut g_sum = g_ambient + g_diffuse + g_specular;
    let mut b_sum = b_ambient + b_diffuse + b_specular;

    if r_sum > 1.0 {
        r_sum = 1.0;
    }

    if g_sum > 1.0 {
        g_sum = 1.0;
    }

    if b_sum > 1.0 {
        b_sum = 1.0;
    }

    Color(r_sum, g_sum, b_sum)
}
<|endoftext|>


<commit_before>pub type BoxedScene<EventT, RendererT, EngineDataT> = Box<Scene<EventT, RendererT, EngineDataT> + 'static>;
pub type SceneChangeCallback<EventT, RendererT, EngineDataT> = Box<Fn(&mut RendererT, &mut EngineDataT) -> BoxedScene<EventT, RendererT, EngineDataT>>;

pub trait Scene<EventT, RendererT, EngineDataT> {
    fn render(&self, renderer: &mut RendererT, engine_data: &mut EngineDataT, tick: u64);
    fn handle_event(&mut self, event: &EventT, renderer: &mut RendererT, engine_data: &mut EngineDataT, tick: u64) -> Option<SceneChangeEvent<EventT, RendererT, EngineDataT>>;
}

pub enum SceneChangeEvent<EventT, RendererT, EngineDataT> {
    PushScene(SceneChangeCallback<EventT, RendererT, EngineDataT>),
    SwapScene(SceneChangeCallback<EventT, RendererT, EngineDataT>),
    PopScene,
}

pub struct SceneStack<EventT, RendererT, EngineDataT> {
    scenes: Vec<BoxedScene<EventT, RendererT, EngineDataT>>
}

impl<EventT, RendererT, EngineDataT> SceneStack<EventT, RendererT, EngineDataT> {

    pub fn new() -> SceneStack<EngineDataT, EventT, RendererT> {
        SceneStack {
            scenes: vec![]
        }
    }

    pub fn is_empty(&self) -> bool {
        self.scenes.len() == 0
    }

    pub fn push(&mut self, scene: BoxedScene<EventT, RendererT, EngineDataT>) {
        self.scenes.push(scene)
    }

    pub fn swap(&mut self, scene: BoxedScene<EventT, RendererT, EngineDataT>) -> Option<BoxedScene<EventT, RendererT, EngineDataT>> {
        let old_scene = self.scenes.pop();
        self.scenes.push(scene);
        old_scene
    }

    pub fn pop(&mut self) -> Option<BoxedScene<EventT, RendererT, EngineDataT>> {
        self.scenes.pop()
    }

    pub fn render(&mut self, renderer: &mut RendererT, engine_data: &mut EngineDataT, tick: u64) {
        let maybe_last_scene = self.scenes.pop();
        match maybe_last_scene {
            Some(scene) => {
                scene.render(renderer, engine_data, tick);
                self.scenes.push(scene);
            },
            None => ()
        }
    }

    pub fn handle_event(&mut self, event: &EventT, renderer: &mut RendererT, engine_data: &mut EngineDataT, tick: u64) {
        let maybe_last_scene = self.scenes.pop();
        let event = match maybe_last_scene {
            Some(mut scene) => {
                let event = scene.handle_event(event, renderer, engine_data, tick);
                self.scenes.push(scene);
                event
            },
            None => None
        };

        match event {
            Some(SceneChangeEvent::PushScene(callback)) => {
                self.push(callback(renderer, engine_data));
            },
            Some(SceneChangeEvent::SwapScene(callback)) => {
                self.swap(callback(renderer, engine_data));
            },
            Some(SceneChangeEvent::PopScene) => {
                self.pop();
            },
            None => ()
        };
    }
}
<commit_msg>Stricter access to stuff passed in to Render<commit_after>pub type BoxedScene<EventT, RendererT, EngineDataT> = Box<Scene<EventT, RendererT, EngineDataT> + 'static>;
pub type SceneChangeCallback<EventT, RendererT, EngineDataT> = Box<Fn(&mut RendererT, &mut EngineDataT) -> BoxedScene<EventT, RendererT, EngineDataT>>;

pub trait Scene<EventT, RendererT, EngineDataT> {
    fn render(&self, renderer: &mut RendererT, engine_data: &EngineDataT, tick: u64);
    fn handle_event(&mut self, event: &EventT, renderer: &mut RendererT, engine_data: &mut EngineDataT, tick: u64) -> Option<SceneChangeEvent<EventT, RendererT, EngineDataT>>;
}

pub enum SceneChangeEvent<EventT, RendererT, EngineDataT> {
    PushScene(SceneChangeCallback<EventT, RendererT, EngineDataT>),
    SwapScene(SceneChangeCallback<EventT, RendererT, EngineDataT>),
    PopScene,
}

pub struct SceneStack<EventT, RendererT, EngineDataT> {
    scenes: Vec<BoxedScene<EventT, RendererT, EngineDataT>>
}

impl<EventT, RendererT, EngineDataT> SceneStack<EventT, RendererT, EngineDataT> {

    pub fn new() -> SceneStack<EngineDataT, EventT, RendererT> {
        SceneStack {
            scenes: vec![]
        }
    }

    pub fn is_empty(&self) -> bool {
        self.scenes.len() == 0
    }

    pub fn push(&mut self, scene: BoxedScene<EventT, RendererT, EngineDataT>) {
        self.scenes.push(scene)
    }

    pub fn swap(&mut self, scene: BoxedScene<EventT, RendererT, EngineDataT>) -> Option<BoxedScene<EventT, RendererT, EngineDataT>> {
        let old_scene = self.scenes.pop();
        self.scenes.push(scene);
        old_scene
    }

    pub fn pop(&mut self) -> Option<BoxedScene<EventT, RendererT, EngineDataT>> {
        self.scenes.pop()
    }

    pub fn render(&mut self, renderer: &mut RendererT, engine_data: &EngineDataT, tick: u64) {
        let maybe_last_scene = self.scenes.pop();
        match maybe_last_scene {
            Some(scene) => {
                scene.render(renderer, engine_data, tick);
                self.scenes.push(scene);
            },
            None => ()
        }
    }

    pub fn handle_event(&mut self, event: &EventT, renderer: &mut RendererT, engine_data: &mut EngineDataT, tick: u64) {
        let maybe_last_scene = self.scenes.pop();
        let event = match maybe_last_scene {
            Some(mut scene) => {
                let event = scene.handle_event(event, renderer, engine_data, tick);
                self.scenes.push(scene);
                event
            },
            None => None
        };

        match event {
            Some(SceneChangeEvent::PushScene(callback)) => {
                self.push(callback(renderer, engine_data));
            },
            Some(SceneChangeEvent::SwapScene(callback)) => {
                self.swap(callback(renderer, engine_data));
            },
            Some(SceneChangeEvent::PopScene) => {
                self.pop();
            },
            None => ()
        };
    }
}
<|endoftext|>


<commit_before>pub fn clear() -> String {
    "\x1b[2J".to_string()
}

pub fn begin() -> String {
    "\x1bPq".to_string()
}

pub fn end() -> String {
    "\x1b\\".to_string()
}

pub fn from(memory: &[u32], width: usize, height: usize) -> String {
    let mut pixels: Vec<u8> = Vec::new();
    let mut row = 0;

    while row < height {
        for x in 0..width {
            let mut byte: u8 = 0;

            for y in 0..6 {
                let offset = x + (y * width);
                if memory[offset] > 0 {
                    byte = byte | (1 << y);
                }
            }

            pixels.push(byte + 63);
        }

        row += 6;
    }

    String::from_utf8(pixels).unwrap()
}

#[cfg(test)]
mod tests {
    #[test]
    fn it_converts_000000_to_question_mark() {
        assert_eq!(super::from(&[0, 0, 0, 0, 0, 0], 1, 6), "?");
    }

    #[test]
    fn it_converts_000111_to_w() {
        assert_eq!(super::from(&[0, 0, 0, 1, 1, 1], 1, 6), "w");
    }

    #[test]
    fn it_converts_111000_to_capital_f() {
        assert_eq!(super::from(&[1, 1, 1, 0, 0, 0], 1, 6), "F");
    }

    #[test]
    fn it_converts_111111_to_tilde() {
        assert_eq!(super::from(&[1, 1, 1, 1, 1, 1], 1, 6), "~");
    }

    #[test]
    fn it_displays_capital_a() {
        assert_eq!(super::from(&[0, 1, 1, 0, 1, 0, 0, 1, 1, 1, 1, 1, 1, 0, 0, 1, 1, 0, 0, 1, 0,
                                 0, 0, 0],
                               4,
                               6),
                   "]DD]");
    }
}
<commit_msg>Break Sixel graphics so they line up in nice rows.<commit_after>pub fn clear() -> String {
    "\x1b[2J".to_string()
}

pub fn begin() -> String {
    "\x1bPq".to_string()
}

pub fn end() -> String {
    "\x1b\\".to_string()
}

pub fn from(memory: &[u32], width: usize, height: usize) -> String {
    let mut pixels: Vec<u8> = Vec::new();
    let mut row = 0;

    while row < height {
        for x in 0..width {
            let mut byte: u8 = 0;

            for y in 0..6 {
                let offset = x + (y * width);
                if memory[offset] > 0 {
                    byte = byte | (1 << y);
                }
            }

            pixels.push(byte + 63);
        }

        // Push "$-" to move to the next row.
        pixels.push(36);
        pixels.push(45);
        row += 6;
    }

    String::from_utf8(pixels).unwrap()
}

#[cfg(test)]
mod tests {
    #[test]
    fn it_converts_000000_to_question_mark() {
        assert_eq!(super::from(&[0, 0, 0, 0, 0, 0], 1, 6), "?$-");
    }

    #[test]
    fn it_converts_000111_to_w() {
        assert_eq!(super::from(&[0, 0, 0, 1, 1, 1], 1, 6), "w$-");
    }

    #[test]
    fn it_converts_111000_to_capital_f() {
        assert_eq!(super::from(&[1, 1, 1, 0, 0, 0], 1, 6), "F$-");
    }

    #[test]
    fn it_converts_111111_to_tilde() {
        assert_eq!(super::from(&[1, 1, 1, 1, 1, 1], 1, 6), "~$-");
    }

    #[test]
    fn it_displays_capital_a() {
        assert_eq!(super::from(&[0, 1, 1, 0, 1, 0, 0, 1, 1, 1, 1, 1, 1, 0, 0, 1, 1, 0, 0, 1, 0,
                                 0, 0, 0],
                               4,
                               6),
                   "]DD]$-");
    }
}
<|endoftext|>


<commit_before>//! Runtime statistics.

use error::P25Error;

/// Tracks stats for an error correction code.
#[derive(Copy, Clone)]
pub struct CodeStats {
    /// Number of symbols per word.
    ///
    /// TODO: make this a const-generic or similar.
    pub size: usize,
    /// Total number of received words.
    pub words: usize,
    /// Number of corrected symbols.
    pub fixed: usize,
    /// Number of unrecoverable words.
    pub err: usize,
}

impl CodeStats {
    /// Create a new `CodeStats` with empty counters for the code with the given number of
    /// symbols per word.
    fn new(size: usize) -> Self {
        CodeStats {
            size: size,
            words: 0,
            err: 0,
            fixed: 0,
        }
    }

    /// Record that a word was received with the given amount of corrected symbols.
    pub fn record_fixes(&mut self, err: usize) {
        debug_assert!(err <= self.size);

        self.words += 1;
        self.fixed += err;
    }

    /// Record that a word was received with an unrecoverable error.
    pub fn record_err(&mut self) {
        self.words += 1;
        self.err += 1;
    }

    /// Merge in the stats from the given object and clear the other stats.
    fn merge(&mut self, other: &mut CodeStats) {
        debug_assert!(self.size == other.size);

        self.words += other.words;
        self.err += other.err;
        self.fixed += other.fixed;

        other.clear();
    }

    /// Clear all stats.
    fn clear(&mut self) {
        self.words = 0;
        self.err = 0;
        self.fixed = 0;
    }
}

/// Records various runtime statistics.
#[derive(Copy, Clone)]
pub struct Stats {
    /// Stats for the BCH code.
    pub bch: CodeStats,
    /// Stats for the cyclic code.
    pub cyclic: CodeStats,
    /// Stats for the standard Golay code.
    pub golay_std: CodeStats,
    /// Stats for the extended Golay code.
    pub golay_ext: CodeStats,
    /// Stats for the short Golay code.
    pub golay_short: CodeStats,
    /// Stats for the standard Hamming code.
    pub hamming_std: CodeStats,
    /// Stats for the short Hamming code.
    pub hamming_short: CodeStats,
    /// Stats for the short RS code.
    pub rs_short: CodeStats,
    /// Stats for the medium RS code.
    pub rs_med: CodeStats,
    /// Stats for the long RS code.
    pub rs_long: CodeStats,
    /// Stats for the dibit Viterbi code.
    pub viterbi_dibit: CodeStats,
    /// Stats for the tribit Viterbi code.
    pub viterbi_tribit: CodeStats,
}

impl Default for Stats {
    fn default() -> Self {
        Stats {
            bch: CodeStats::new(64),
            cyclic: CodeStats::new(16),
            golay_std: CodeStats::new(23),
            golay_ext: CodeStats::new(24),
            golay_short: CodeStats::new(18),
            hamming_std: CodeStats::new(15),
            hamming_short: CodeStats::new(10),
            rs_short: CodeStats::new(24),
            rs_med: CodeStats::new(24),
            rs_long: CodeStats::new(36),
            viterbi_dibit: CodeStats::new(196),
            viterbi_tribit: CodeStats::new(196),
        }
    }
}

impl Stats {
    /// Merge in the stats from the given object and reset the other stats back to
    /// default.
    pub fn merge<T: HasStats>(&mut self, other: &mut T) {
        let stats = other.stats();

        self.bch.merge(&mut stats.bch);
        self.cyclic.merge(&mut stats.cyclic);
        self.golay_std.merge(&mut stats.golay_std);
        self.golay_ext.merge(&mut stats.golay_ext);
        self.golay_short.merge(&mut stats.golay_short);
        self.hamming_std.merge(&mut stats.hamming_std);
        self.hamming_short.merge(&mut stats.hamming_short);
        self.rs_short.merge(&mut stats.rs_short);
        self.rs_med.merge(&mut stats.rs_med);
        self.rs_long.merge(&mut stats.rs_long);
        self.viterbi_dibit.merge(&mut stats.viterbi_dibit);
        self.viterbi_tribit.merge(&mut stats.viterbi_tribit);
    }

    /// Clear all stats.
    pub fn clear(&mut self) {
        *self = Stats::default();
    }

    /// Record the given error into the current stats.
    pub fn record_err(&mut self, err: P25Error) {
        use error::P25Error::*;

        match err {
            BchUnrecoverable => self.bch.record_err(),
            CyclicUnrecoverable => self.cyclic.record_err(),
            GolayStdUnrecoverable => self.golay_std.record_err(),
            GolayExtUnrecoverable => self.golay_ext.record_err(),
            GolayShortUnrecoverable => self.golay_short.record_err(),
            HammingStdUnrecoverable => self.hamming_std.record_err(),
            HammingShortUnrecoverable => self.hamming_short.record_err(),
            RsShortUnrecoverable => self.rs_short.record_err(),
            RsMediumUnrecoverable => self.rs_med.record_err(),
            RsLongUnrecoverable => self.rs_long.record_err(),
            DibitViterbiUnrecoverable => self.viterbi_dibit.record_err(),
            UnknownNid => {},
        }
    }
}

/// Indicates that a type captures statistics.
pub trait HasStats {
    /// Retrieve captured statistics.
    fn stats(&mut self) -> &mut Stats;
}

#[cfg(test)]
mod test {
    use super::*;

    #[test]
    fn test_code_stats() {
        let mut a = CodeStats::new(23);
        let mut b = CodeStats::new(23);

        a.record_fixes(13);
        a.record_err();
        assert_eq!(a.size, 23);
        assert_eq!(a.words, 2);
        assert_eq!(a.fixed, 13);
        assert_eq!(a.err, 1);

        b.record_fixes(11);
        b.record_fixes(19);
        b.record_err();
        b.record_err();
        assert_eq!(b.size, 23);
        assert_eq!(b.words, 4);
        assert_eq!(b.fixed, 30);
        assert_eq!(b.err, 2);

        a.merge(&mut b);
        assert_eq!(a.size, 23);
        assert_eq!(a.words, 6);
        assert_eq!(a.fixed, 43);
        assert_eq!(a.err, 3);
        assert_eq!(b.size, 23);
        assert_eq!(b.words, 0);
        assert_eq!(b.fixed, 0);
        assert_eq!(b.err, 0);

        let mut c = CodeStats::new(7);
        c.record_fixes(3);
        c.record_fixes(2);
        c.record_err();
        assert_eq!(c.size, 7);
        assert_eq!(c.words, 3);
        assert_eq!(c.fixed, 5);
        assert_eq!(c.err, 1);
    }
}
<commit_msg>s/err/errs/<commit_after>//! Runtime statistics.

use error::P25Error;

/// Tracks stats for an error correction code.
#[derive(Copy, Clone)]
pub struct CodeStats {
    /// Number of symbols per word.
    ///
    /// TODO: make this a const-generic or similar.
    pub size: usize,
    /// Total number of received words.
    pub words: usize,
    /// Number of corrected symbols.
    pub fixed: usize,
    /// Number of unrecoverable words.
    pub errs: usize,
}

impl CodeStats {
    /// Create a new `CodeStats` with empty counters for the code with the given number of
    /// symbols per word.
    fn new(size: usize) -> Self {
        CodeStats {
            size: size,
            words: 0,
            errs: 0,
            fixed: 0,
        }
    }

    /// Record that a word was received with the given amount of corrected symbols.
    pub fn record_fixes(&mut self, errs: usize) {
        debug_assert!(errs <= self.size);

        self.words += 1;
        self.fixed += errs;
    }

    /// Record that a word was received with an unrecoverable error.
    pub fn record_err(&mut self) {
        self.words += 1;
        self.errs += 1;
    }

    /// Merge in the stats from the given object and clear the other stats.
    fn merge(&mut self, other: &mut CodeStats) {
        debug_assert!(self.size == other.size);

        self.words += other.words;
        self.errs += other.errs;
        self.fixed += other.fixed;

        other.clear();
    }

    /// Clear all stats.
    fn clear(&mut self) {
        self.words = 0;
        self.errs = 0;
        self.fixed = 0;
    }
}

/// Records various runtime statistics.
#[derive(Copy, Clone)]
pub struct Stats {
    /// Stats for the BCH code.
    pub bch: CodeStats,
    /// Stats for the cyclic code.
    pub cyclic: CodeStats,
    /// Stats for the standard Golay code.
    pub golay_std: CodeStats,
    /// Stats for the extended Golay code.
    pub golay_ext: CodeStats,
    /// Stats for the short Golay code.
    pub golay_short: CodeStats,
    /// Stats for the standard Hamming code.
    pub hamming_std: CodeStats,
    /// Stats for the short Hamming code.
    pub hamming_short: CodeStats,
    /// Stats for the short RS code.
    pub rs_short: CodeStats,
    /// Stats for the medium RS code.
    pub rs_med: CodeStats,
    /// Stats for the long RS code.
    pub rs_long: CodeStats,
    /// Stats for the dibit Viterbi code.
    pub viterbi_dibit: CodeStats,
    /// Stats for the tribit Viterbi code.
    pub viterbi_tribit: CodeStats,
}

impl Default for Stats {
    fn default() -> Self {
        Stats {
            bch: CodeStats::new(64),
            cyclic: CodeStats::new(16),
            golay_std: CodeStats::new(23),
            golay_ext: CodeStats::new(24),
            golay_short: CodeStats::new(18),
            hamming_std: CodeStats::new(15),
            hamming_short: CodeStats::new(10),
            rs_short: CodeStats::new(24),
            rs_med: CodeStats::new(24),
            rs_long: CodeStats::new(36),
            viterbi_dibit: CodeStats::new(196),
            viterbi_tribit: CodeStats::new(196),
        }
    }
}

impl Stats {
    /// Merge in the stats from the given object and reset the other stats back to
    /// default.
    pub fn merge<T: HasStats>(&mut self, other: &mut T) {
        let stats = other.stats();

        self.bch.merge(&mut stats.bch);
        self.cyclic.merge(&mut stats.cyclic);
        self.golay_std.merge(&mut stats.golay_std);
        self.golay_ext.merge(&mut stats.golay_ext);
        self.golay_short.merge(&mut stats.golay_short);
        self.hamming_std.merge(&mut stats.hamming_std);
        self.hamming_short.merge(&mut stats.hamming_short);
        self.rs_short.merge(&mut stats.rs_short);
        self.rs_med.merge(&mut stats.rs_med);
        self.rs_long.merge(&mut stats.rs_long);
        self.viterbi_dibit.merge(&mut stats.viterbi_dibit);
        self.viterbi_tribit.merge(&mut stats.viterbi_tribit);
    }

    /// Clear all stats.
    pub fn clear(&mut self) {
        *self = Stats::default();
    }

    /// Record the given error into the current stats.
    pub fn record_err(&mut self, err: P25Error) {
        use error::P25Error::*;

        match err {
            BchUnrecoverable => self.bch.record_err(),
            CyclicUnrecoverable => self.cyclic.record_err(),
            GolayStdUnrecoverable => self.golay_std.record_err(),
            GolayExtUnrecoverable => self.golay_ext.record_err(),
            GolayShortUnrecoverable => self.golay_short.record_err(),
            HammingStdUnrecoverable => self.hamming_std.record_err(),
            HammingShortUnrecoverable => self.hamming_short.record_err(),
            RsShortUnrecoverable => self.rs_short.record_err(),
            RsMediumUnrecoverable => self.rs_med.record_err(),
            RsLongUnrecoverable => self.rs_long.record_err(),
            DibitViterbiUnrecoverable => self.viterbi_dibit.record_err(),
            UnknownNid => {},
        }
    }
}

/// Indicates that a type captures statistics.
pub trait HasStats {
    /// Retrieve captured statistics.
    fn stats(&mut self) -> &mut Stats;
}

#[cfg(test)]
mod test {
    use super::*;

    #[test]
    fn test_code_stats() {
        let mut a = CodeStats::new(23);
        let mut b = CodeStats::new(23);

        a.record_fixes(13);
        a.record_err();
        assert_eq!(a.size, 23);
        assert_eq!(a.words, 2);
        assert_eq!(a.fixed, 13);
        assert_eq!(a.errs, 1);

        b.record_fixes(11);
        b.record_fixes(19);
        b.record_err();
        b.record_err();
        assert_eq!(b.size, 23);
        assert_eq!(b.words, 4);
        assert_eq!(b.fixed, 30);
        assert_eq!(b.errs, 2);

        a.merge(&mut b);
        assert_eq!(a.size, 23);
        assert_eq!(a.words, 6);
        assert_eq!(a.fixed, 43);
        assert_eq!(a.errs, 3);
        assert_eq!(b.size, 23);
        assert_eq!(b.words, 0);
        assert_eq!(b.fixed, 0);
        assert_eq!(b.errs, 0);

        let mut c = CodeStats::new(7);
        c.record_fixes(3);
        c.record_fixes(2);
        c.record_err();
        assert_eq!(c.size, 7);
        assert_eq!(c.words, 3);
        assert_eq!(c.fixed, 5);
        assert_eq!(c.errs, 1);
    }
}
<|endoftext|>


<commit_before>// Copyright © 2016, Canal TP and/or its affiliates. All rights reserved.
//
// This file is part of Navitia,
//     the software to build cool stuff with public transport.
//
// Hope you'll enjoy and contribute to this project,
//     powered by Canal TP (www.canaltp.fr).
// Help us simplify mobility and open public transport:
//     a non ending quest to the responsive locomotion way of traveling!
//
// LICENCE: This program is free software; you can redistribute it
// and/or modify it under the terms of the GNU Affero General Public
// License as published by the Free Software Foundation, either
// version 3 of the License, or (at your option) any later version.
//
// This program is distributed in the hope that it will be useful, but
// WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
// Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public
// License along with this program. If not, see
// <http://www.gnu.org/licenses/>.
//
// Stay tuned using
// twitter @navitia
// IRC #navitia on freenode
// https://groups.google.com/d/forum/navitia
// www.navitia.io

/// In this module we put the code related to stops, that need to draw on 'places', 'mimir',
/// 'common', and 'config' (ie all the workspaces that make up mimirsbrunn).
use futures::stream::{Stream, TryStreamExt};
use mimir::domain::model::configuration::{ContainerConfig, PhysicalModeWeight};
use snafu::{ResultExt, Snafu};
use std::collections::{hash_map::DefaultHasher, HashMap};
use std::hash::BuildHasherDefault;
use std::ops::Deref;
use std::path::PathBuf;
use std::sync::Arc;
use tracing::{info, warn};

use crate::admin_geofinder::AdminGeoFinder;
use crate::labels;
use mimir::adapters::secondary::elasticsearch::{self, ElasticsearchStorage};
use mimir::domain::ports::primary::{generate_index::GenerateIndex, list_documents::ListDocuments};
use places::admin::Admin;
use places::stop::Stop;

#[derive(Debug, Snafu)]
pub enum Error {
    // #[snafu(display("Settings (Configuration or CLI) Error: {}", source))]
    // Settings { source: settings::Error },
    #[snafu(display("Elasticsearch Connection Pool {}", source))]
    ElasticsearchPool {
        source: elasticsearch::remote::Error,
    },

    #[snafu(display("Elasticsearch Connection Pool {}", source))]
    ElasticsearchConnection {
        source: mimir::domain::ports::secondary::remote::Error,
    },

    // #[snafu(display("Cosmogony Error: {}", details))]
    // Cosmogony { details: String },
    #[snafu(display("Index Generation Error {}", source))]
    IndexGeneration {
        source: mimir::domain::model::error::Error,
    },

    // transit_model uses failure::Error, which does not implement std::Error, so
    // we use a String to get the error message instead.
    #[snafu(display("Transit Model Error {}", details))]
    TransitModel { details: String },

    #[snafu(display("Admin Retrieval Error {}", details))]
    AdminRetrieval { details: String },
}

pub fn initialize_weights<'a, It, S: ::std::hash::BuildHasher>(
    stops: It,
    nb_stop_points: &HashMap<String, u32, S>,
) where
    It: Iterator<Item = &'a mut Stop>,
{
    let max = *nb_stop_points.values().max().unwrap_or(&1) as f64;
    for stop in stops {
        stop.weight = if let Some(weight) = nb_stop_points.get(&stop.id) {
            *weight as f64 / max
        } else {
            0.0
        };
    }
}

pub fn make_weight(stop: &mut Stop, physical_mode_weight: &HashMap<String, f64>) {
    // Admin weight
    let mut admin_weight = stop
        .administrative_regions
        .iter()
        .filter(|adm| adm.is_city())
        .map(|adm| adm.weight)
        .next()
        .unwrap_or(0.0);
    // FIXME: 1024, automagic!
    // It's a factor used to bring the stop weight and the admin weight in the same order of
    // magnitude...
    // We then use a log to compress the distance between low admin weight and high ones.
    admin_weight = admin_weight * 1024.0 + 1.0;
    admin_weight = admin_weight.log10();

    let mut result: Vec<f64> = stop
        .physical_modes
        .iter()
        .map(|mode| {
            let pm_w = physical_mode_weight.get(&*mode.id);
            match pm_w {
                Some(value) => *value,
                _ => {
                    warn!(
                        "Physical mode, id: {} name: {}, not found in mimir config.",
                        mode.id, mode.name
                    );
                    0.0
                }
            }
        })
        .filter(|weight| !weight.is_nan())
        .collect();

    result.push(stop.weight);
    result.push(admin_weight);
    let sum: f64 = Iterator::sum(result.iter());
    stop.weight = sum / (result.len() as f64);
}

fn attach_stop(stop: &mut Stop, admins: Vec<Arc<Admin>>) {
    let admins_iter = admins.iter().map(|a| a.deref());
    let country_codes = places::admin::find_country_codes(admins_iter.clone());

    stop.label = labels::format_stop_label(&stop.name, admins_iter, &country_codes);
    stop.zip_codes = places::admin::get_zip_codes_from_admins(&admins);

    stop.country_codes = country_codes;
    stop.administrative_regions = admins;
}

/// Attach the stops to administrative regions
///
/// The admins are loaded from Elasticsearch and stored in a quadtree
/// We attach a stop with all the admins that have a boundary containing
/// the coordinate of the stop
async fn attach_stops_to_admins<'a, It: Iterator<Item = &'a mut Stop>>(
    stops: It,
    client: &ElasticsearchStorage,
) -> Result<(), Error> {
    match client.list_documents().await {
        Ok(stream) => {
            let admins: Vec<Admin> = stream.try_collect().await.context(IndexGenerationSnafu)?;

            if admins.is_empty() {
                return Err(Error::AdminRetrieval {
                    details: String::from("no admin retrieved to enrich stops"),
                });
            }
            info!("{} admins retrieved from ES ", admins.len());
            let admins_geofinder = admins.into_iter().collect::<AdminGeoFinder>();

            let mut nb_unmatched = 0u32;
            let mut nb_matched = 0u32;
            // FIXME Opportunity for concurrent work
            for mut stop in stops {
                let admins = admins_geofinder.get(&stop.coord);

                if admins.is_empty() {
                    nb_unmatched += 1;
                } else {
                    nb_matched += 1;
                }

                attach_stop(&mut stop, admins);
            }

            info!(
                "there are {}/{} stops without any admin",
                nb_unmatched,
                nb_matched + nb_unmatched
            );
            Ok(())
        }
        Err(_) => Err(Error::AdminRetrieval {
            details: String::from("Could not retrieve admins to enrich stops"),
        }),
    }
}

/// Stores the stops found in the 'input' directory, in Elasticsearch, with the given
/// configuration.
///
/// The main part of this function is to actually create a list of stops
/// from the information found in the NTFS directory.
pub async fn index_ntfs(
    input: PathBuf,
    config: &ContainerConfig,
    physical_mode_weight: &Option<Vec<PhysicalModeWeight>>,
    client: &ElasticsearchStorage,
) -> Result<(), Error> {
    let navitia = transit_model::ntfs::read(&input).map_err(|err| Error::TransitModel {
        details: format!(
            "Could not read transit model from {}: {}",
            input.display(),
            err.to_string()
        ),
    })?;
    info!("Build number of stops per stoparea");
    let nb_stop_points: HashMap<String, u32, BuildHasherDefault<DefaultHasher>> = navitia
        .stop_areas
        .iter()
        .map(|(idx, sa)| {
            let id = places::utils::normalize_id("stop_area", &sa.id);
            let nb_stop_points = navitia
                .get_corresponding_from_idx::<_, transit_model::objects::StopPoint>(idx)
                .len();
            (id, nb_stop_points as u32)
        })
        .collect();

    info!("Make mimir stops from navitia stops");
    let mut stops: Vec<Stop> = navitia
        .stop_areas
        .iter()
        .map(|(idx, sa)| places::stop::to_mimir(idx, sa, &navitia))
        .collect();

    info!("Initialize stops weights");
    initialize_weights(stops.iter_mut(), &nb_stop_points);

    info!("Attach stops to admins");
    attach_stops_to_admins(stops.iter_mut(), client).await?;

    // FIXME Should be done concurrently (for_each_concurrent....)
    info!("Build stops weight by physical modes and city population");
    let md_weight_hash_map: HashMap<String, f64> = match physical_mode_weight {
        Some(modes) => modes
            .iter()
            .map(|mode| (mode.id.to_string(), mode.weight as f64))
            .collect::<HashMap<String, f64>>(),
        _ => HashMap::new(),
    };
    for stop in &mut stops {
        stop.coverages.push(config.dataset.clone());
        make_weight(stop, &md_weight_hash_map);
    }
    tracing::info!("Beginning to import stops into elasticsearch.");
    import_stops(client, config, futures::stream::iter(stops)).await
}

// FIXME Should not be ElasticsearchStorage, but rather a trait GenerateIndex
pub async fn import_stops<S>(
    client: &ElasticsearchStorage,
    config: &ContainerConfig,
    stops: S,
) -> Result<(), Error>
where
    S: Stream<Item = Stop> + Send + Sync + Unpin + 'static,
{
    client
        .generate_index(config, stops)
        .await
        .context(IndexGenerationSnafu)?;

    Ok(())
}

#[cfg(test)]
mod tests {
    use crate::stops::make_weight;
    use cosmogony::ZoneType;
    use places::admin::Admin;
    use places::stop::{PhysicalMode, Stop};
    use serial_test::serial;
    use std::collections::HashMap;
    use std::sync::Arc;

    fn approx_equal(a: f64, b: f64, dp: u8) -> bool {
        let p = 10f64.powi(-(dp as i32));
        (a - b).abs() < p
    }

    #[tokio::test]
    #[serial]
    async fn test_make_weight_without_physical_mode_weight_and_without_admins() {
        let physical_mode = PhysicalMode {
            id: "physical_mode:Tramway".to_string(),
            name: "Tramway".to_string(),
        };
        let mut stop = Stop {
            id: "123".to_string(),
            physical_modes: vec![physical_mode],
            weight: 0.65,
            ..Default::default()
        };
        make_weight(&mut stop, &HashMap::new());
        assert!(approx_equal(stop.weight, 0.325, 3));
    }

    #[tokio::test]
    #[serial]
    async fn test_make_weight_without_physical_mode_weight_and_with_admins() {
        let physical_mode = PhysicalMode {
            id: "physical_mode:Tramway".to_string(),
            name: "Tramway".to_string(),
        };
        let admin = Admin {
            id: "adm:01".to_string(),
            weight: 0.12,
            zone_type: Some(ZoneType::City),
            ..Default::default()
        };
        let mut stop = Stop {
            id: "123".to_string(),
            physical_modes: vec![physical_mode],
            administrative_regions: vec![Arc::new(admin)],
            weight: 0.65,
            ..Default::default()
        };
        make_weight(&mut stop, &HashMap::new());
        approx_equal(stop.weight, 1.3715, 4);
    }

    #[tokio::test]
    #[serial]
    async fn test_make_weight_with_physical_mode_weight_and_with_admins() {
        let physical_mode = PhysicalMode {
            id: "physical_mode:Tramway".to_string(),
            name: "Tramway".to_string(),
        };
        let mut physical_mode_weight = HashMap::new();
        physical_mode_weight.insert("physical_mode:Tramway".to_string(), 5.0);

        let admin = Admin {
            id: "adm:01".to_string(),
            weight: 0.12,
            zone_type: Some(ZoneType::City),
            ..Default::default()
        };
        let mut stop = Stop {
            id: "123".to_string(),
            physical_modes: vec![physical_mode],
            administrative_regions: vec![Arc::new(admin)],
            weight: 0.65,
            ..Default::default()
        };
        make_weight(&mut stop, &physical_mode_weight);
        approx_equal(stop.weight, 2.581, 4);
    }

    #[tokio::test]
    #[serial]
    async fn test_make_weight_without_physical_mode() {
        let mut physical_mode_weight = HashMap::new();
        physical_mode_weight.insert("physical_mode:Tramway".to_string(), 5.0);
        let admin = Admin {
            id: "adm:01".to_string(),
            weight: 0.12,
            zone_type: Some(ZoneType::City),
            ..Default::default()
        };
        let mut stop = Stop {
            id: "123".to_string(),
            administrative_regions: vec![Arc::new(admin)],
            weight: 0.65,
            ..Default::default()
        };

        make_weight(&mut stop, &physical_mode_weight);
        approx_equal(stop.weight, 1.3715, 4);
    }
}
<commit_msg>correction<commit_after>// Copyright © 2016, Canal TP and/or its affiliates. All rights reserved.
//
// This file is part of Navitia,
//     the software to build cool stuff with public transport.
//
// Hope you'll enjoy and contribute to this project,
//     powered by Canal TP (www.canaltp.fr).
// Help us simplify mobility and open public transport:
//     a non ending quest to the responsive locomotion way of traveling!
//
// LICENCE: This program is free software; you can redistribute it
// and/or modify it under the terms of the GNU Affero General Public
// License as published by the Free Software Foundation, either
// version 3 of the License, or (at your option) any later version.
//
// This program is distributed in the hope that it will be useful, but
// WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
// Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public
// License along with this program. If not, see
// <http://www.gnu.org/licenses/>.
//
// Stay tuned using
// twitter @navitia
// IRC #navitia on freenode
// https://groups.google.com/d/forum/navitia
// www.navitia.io

/// In this module we put the code related to stops, that need to draw on 'places', 'mimir',
/// 'common', and 'config' (ie all the workspaces that make up mimirsbrunn).
use futures::stream::{Stream, TryStreamExt};
use mimir::domain::model::configuration::{ContainerConfig, PhysicalModeWeight};
use snafu::{ResultExt, Snafu};
use std::collections::{hash_map::DefaultHasher, HashMap};
use std::hash::BuildHasherDefault;
use std::ops::Deref;
use std::path::PathBuf;
use std::sync::Arc;
use tracing::{info, warn};

use crate::admin_geofinder::AdminGeoFinder;
use crate::labels;
use mimir::adapters::secondary::elasticsearch::{self, ElasticsearchStorage};
use mimir::domain::ports::primary::{generate_index::GenerateIndex, list_documents::ListDocuments};
use places::admin::Admin;
use places::stop::Stop;

#[derive(Debug, Snafu)]
pub enum Error {
    // #[snafu(display("Settings (Configuration or CLI) Error: {}", source))]
    // Settings { source: settings::Error },
    #[snafu(display("Elasticsearch Connection Pool {}", source))]
    ElasticsearchPool {
        source: elasticsearch::remote::Error,
    },

    #[snafu(display("Elasticsearch Connection Pool {}", source))]
    ElasticsearchConnection {
        source: mimir::domain::ports::secondary::remote::Error,
    },

    // #[snafu(display("Cosmogony Error: {}", details))]
    // Cosmogony { details: String },
    #[snafu(display("Index Generation Error {}", source))]
    IndexGeneration {
        source: mimir::domain::model::error::Error,
    },

    // transit_model uses failure::Error, which does not implement std::Error, so
    // we use a String to get the error message instead.
    #[snafu(display("Transit Model Error {}", details))]
    TransitModel { details: String },

    #[snafu(display("Admin Retrieval Error {}", details))]
    AdminRetrieval { details: String },
}

pub fn initialize_weights<'a, It, S: ::std::hash::BuildHasher>(
    stops: It,
    nb_stop_points: &HashMap<String, u32, S>,
) where
    It: Iterator<Item = &'a mut Stop>,
{
    let max = *nb_stop_points.values().max().unwrap_or(&1) as f64;
    for stop in stops {
        stop.weight = if let Some(weight) = nb_stop_points.get(&stop.id) {
            *weight as f64 / max
        } else {
            0.0
        };
    }
}

pub fn make_weight(stop: &mut Stop, physical_mode_weight: &HashMap<String, f64>) {
    // Admin weight
    let mut admin_weight = stop
        .administrative_regions
        .iter()
        .filter(|adm| adm.is_city())
        .map(|adm| adm.weight)
        .next()
        .unwrap_or(0.0);
    // FIXME: 1024, automagic!
    // It's a factor used to bring the stop weight and the admin weight in the same order of
    // magnitude...
    // We then use a log to compress the distance between low admin weight and high ones.
    admin_weight = admin_weight * 1024.0 + 1.0;
    admin_weight = admin_weight.log10();

    let mut result: Vec<f64> = stop
        .physical_modes
        .iter()
        .map(|mode| {
            let pm_w = physical_mode_weight.get(&*mode.id);
            match pm_w {
                Some(value) => *value,
                _ => {
                    warn!(
                        "Physical mode, id: {} name: {}, not found in mimir config.",
                        mode.id, mode.name
                    );
                    0.0
                }
            }
        })
        .filter(|&weight| weight != 0.0)
        .collect();

    result.push(stop.weight);
    result.push(admin_weight);
    let sum: f64 = Iterator::sum(result.iter());
    stop.weight = sum / (result.len() as f64);
}

fn attach_stop(stop: &mut Stop, admins: Vec<Arc<Admin>>) {
    let admins_iter = admins.iter().map(|a| a.deref());
    let country_codes = places::admin::find_country_codes(admins_iter.clone());

    stop.label = labels::format_stop_label(&stop.name, admins_iter, &country_codes);
    stop.zip_codes = places::admin::get_zip_codes_from_admins(&admins);

    stop.country_codes = country_codes;
    stop.administrative_regions = admins;
}

/// Attach the stops to administrative regions
///
/// The admins are loaded from Elasticsearch and stored in a quadtree
/// We attach a stop with all the admins that have a boundary containing
/// the coordinate of the stop
async fn attach_stops_to_admins<'a, It: Iterator<Item = &'a mut Stop>>(
    stops: It,
    client: &ElasticsearchStorage,
) -> Result<(), Error> {
    match client.list_documents().await {
        Ok(stream) => {
            let admins: Vec<Admin> = stream.try_collect().await.context(IndexGenerationSnafu)?;

            if admins.is_empty() {
                return Err(Error::AdminRetrieval {
                    details: String::from("no admin retrieved to enrich stops"),
                });
            }
            info!("{} admins retrieved from ES ", admins.len());
            let admins_geofinder = admins.into_iter().collect::<AdminGeoFinder>();

            let mut nb_unmatched = 0u32;
            let mut nb_matched = 0u32;
            // FIXME Opportunity for concurrent work
            for mut stop in stops {
                let admins = admins_geofinder.get(&stop.coord);

                if admins.is_empty() {
                    nb_unmatched += 1;
                } else {
                    nb_matched += 1;
                }

                attach_stop(&mut stop, admins);
            }

            info!(
                "there are {}/{} stops without any admin",
                nb_unmatched,
                nb_matched + nb_unmatched
            );
            Ok(())
        }
        Err(_) => Err(Error::AdminRetrieval {
            details: String::from("Could not retrieve admins to enrich stops"),
        }),
    }
}

/// Stores the stops found in the 'input' directory, in Elasticsearch, with the given
/// configuration.
///
/// The main part of this function is to actually create a list of stops
/// from the information found in the NTFS directory.
pub async fn index_ntfs(
    input: PathBuf,
    config: &ContainerConfig,
    physical_mode_weight: &Option<Vec<PhysicalModeWeight>>,
    client: &ElasticsearchStorage,
) -> Result<(), Error> {
    let navitia = transit_model::ntfs::read(&input).map_err(|err| Error::TransitModel {
        details: format!(
            "Could not read transit model from {}: {}",
            input.display(),
            err.to_string()
        ),
    })?;
    info!("Build number of stops per stoparea");
    let nb_stop_points: HashMap<String, u32, BuildHasherDefault<DefaultHasher>> = navitia
        .stop_areas
        .iter()
        .map(|(idx, sa)| {
            let id = places::utils::normalize_id("stop_area", &sa.id);
            let nb_stop_points = navitia
                .get_corresponding_from_idx::<_, transit_model::objects::StopPoint>(idx)
                .len();
            (id, nb_stop_points as u32)
        })
        .collect();

    info!("Make mimir stops from navitia stops");
    let mut stops: Vec<Stop> = navitia
        .stop_areas
        .iter()
        .map(|(idx, sa)| places::stop::to_mimir(idx, sa, &navitia))
        .collect();

    info!("Initialize stops weights");
    initialize_weights(stops.iter_mut(), &nb_stop_points);

    info!("Attach stops to admins");
    attach_stops_to_admins(stops.iter_mut(), client).await?;

    // FIXME Should be done concurrently (for_each_concurrent....)
    info!("Build stops weight by physical modes and city population");
    let md_weight_hash_map: HashMap<String, f64> = match physical_mode_weight {
        Some(modes) => modes
            .iter()
            .map(|mode| (mode.id.to_string(), mode.weight as f64))
            .collect::<HashMap<String, f64>>(),
        _ => HashMap::new(),
    };
    for stop in &mut stops {
        stop.coverages.push(config.dataset.clone());
        make_weight(stop, &md_weight_hash_map);
    }
    tracing::info!("Beginning to import stops into elasticsearch.");
    import_stops(client, config, futures::stream::iter(stops)).await
}

// FIXME Should not be ElasticsearchStorage, but rather a trait GenerateIndex
pub async fn import_stops<S>(
    client: &ElasticsearchStorage,
    config: &ContainerConfig,
    stops: S,
) -> Result<(), Error>
where
    S: Stream<Item = Stop> + Send + Sync + Unpin + 'static,
{
    client
        .generate_index(config, stops)
        .await
        .context(IndexGenerationSnafu)?;

    Ok(())
}

#[cfg(test)]
mod tests {
    use crate::stops::make_weight;
    use cosmogony::ZoneType;
    use places::admin::Admin;
    use places::stop::{PhysicalMode, Stop};
    use serial_test::serial;
    use std::collections::HashMap;
    use std::sync::Arc;

    fn approx_equal(a: f64, b: f64, dp: u8) -> bool {
        let p = 10f64.powi(-(dp as i32));
        (a - b).abs() < p
    }

    #[tokio::test]
    #[serial]
    async fn test_make_weight_without_physical_mode_weight_and_without_admins() {
        let physical_mode = PhysicalMode {
            id: "physical_mode:Tramway".to_string(),
            name: "Tramway".to_string(),
        };
        let mut stop = Stop {
            id: "123".to_string(),
            physical_modes: vec![physical_mode],
            weight: 0.65,
            ..Default::default()
        };
        make_weight(&mut stop, &HashMap::new());
        assert!(approx_equal(stop.weight, 0.325, 3));
    }

    #[tokio::test]
    #[serial]
    async fn test_make_weight_without_physical_mode_weight_and_with_admins() {
        let physical_mode = PhysicalMode {
            id: "physical_mode:Tramway".to_string(),
            name: "Tramway".to_string(),
        };
        let admin = Admin {
            id: "adm:01".to_string(),
            weight: 0.12,
            zone_type: Some(ZoneType::City),
            ..Default::default()
        };
        let mut stop = Stop {
            id: "123".to_string(),
            physical_modes: vec![physical_mode],
            administrative_regions: vec![Arc::new(admin)],
            weight: 0.65,
            ..Default::default()
        };
        make_weight(&mut stop, &HashMap::new());
        approx_equal(stop.weight, 1.3715, 4);
    }

    #[tokio::test]
    #[serial]
    async fn test_make_weight_with_physical_mode_weight_and_with_admins() {
        let physical_mode = PhysicalMode {
            id: "physical_mode:Tramway".to_string(),
            name: "Tramway".to_string(),
        };
        let mut physical_mode_weight = HashMap::new();
        physical_mode_weight.insert("physical_mode:Tramway".to_string(), 5.0);

        let admin = Admin {
            id: "adm:01".to_string(),
            weight: 0.12,
            zone_type: Some(ZoneType::City),
            ..Default::default()
        };
        let mut stop = Stop {
            id: "123".to_string(),
            physical_modes: vec![physical_mode],
            administrative_regions: vec![Arc::new(admin)],
            weight: 0.65,
            ..Default::default()
        };
        make_weight(&mut stop, &physical_mode_weight);
        approx_equal(stop.weight, 2.581, 4);
    }

    #[tokio::test]
    #[serial]
    async fn test_make_weight_without_physical_mode() {
        let mut physical_mode_weight = HashMap::new();
        physical_mode_weight.insert("physical_mode:Tramway".to_string(), 5.0);
        let admin = Admin {
            id: "adm:01".to_string(),
            weight: 0.12,
            zone_type: Some(ZoneType::City),
            ..Default::default()
        };
        let mut stop = Stop {
            id: "123".to_string(),
            administrative_regions: vec![Arc::new(admin)],
            weight: 0.65,
            ..Default::default()
        };

        make_weight(&mut stop, &physical_mode_weight);
        approx_equal(stop.weight, 1.3715, 4);
    }
}
<|endoftext|>


<commit_before>/*!
The types module provides a way of associating glob patterns on file names to
file types.
*/

use std::collections::HashMap;
use std::error::Error as StdError;
use std::fmt;
use std::path::Path;

use regex;

use gitignore::{Match, Pattern};
use globset::{self, GlobBuilder, GlobSet, GlobSetBuilder};

const TYPE_EXTENSIONS: &'static [(&'static str, &'static [&'static str])] = &[
    ("asm", &["*.asm", "*.s", "*.S"]),
    ("awk", &["*.awk"]),
    ("c", &["*.c", "*.h", "*.H"]),
    ("cbor", &["*.cbor"]),
    ("clojure", &["*.clj", "*.cljc", "*.cljs", "*.cljx"]),
    ("cmake", &["*.cmake", "CMakeLists.txt"]),
    ("coffeescript", &["*.coffee"]),
    ("config", &["*.config"]),
    ("cpp", &[
        "*.C", "*.cc", "*.cpp", "*.cxx",
        "*.h", "*.H", "*.hh", "*.hpp",
    ]),
    ("csharp", &["*.cs"]),
    ("css", &["*.css"]),
    ("cython", &["*.pyx"]),
    ("dart", &["*.dart"]),
    ("d", &["*.d"]),
    ("elisp", &["*.el"]),
    ("erlang", &["*.erl", "*.hrl"]),
    ("fortran", &[
        "*.f", "*.F", "*.f77", "*.F77", "*.pfo",
        "*.f90", "*.F90", "*.f95", "*.F95",
    ]),
    ("fsharp", &["*.fs", "*.fsx", "*.fsi"]),
    ("go", &["*.go"]),
    ("groovy", &["*.groovy", "*.gradle"]),
    ("hbs", &["*.hbs"]),
    ("haskell", &["*.hs", "*.lhs"]),
    ("html", &["*.htm", "*.html"]),
    ("java", &["*.java"]),
    ("jinja", &["*.jinja", "*.jinja2"]),
    ("js", &[
        "*.js", "*.jsx", "*.vue",
    ]),
    ("json", &["*.json"]),
    ("jsonl", &["*.jsonl"]),
    ("lisp", &["*.el", "*.jl", "*.lisp", "*.lsp", "*.sc", "*.scm"]),
    ("lua", &["*.lua"]),
    ("m4", &["*.ac", "*.m4"]),
    ("make", &["gnumakefile", "Gnumakefile", "makefile", "Makefile", "*.mk"]),
    ("markdown", &["*.md"]),
    ("md", &["*.md"]),
    ("matlab", &["*.m"]),
    ("mk", &["mkfile"]),
    ("ml", &["*.ml"]),
    ("nim", &["*.nim"]),
    ("objc", &["*.h", "*.m"]),
    ("objcpp", &["*.h", "*.mm"]),
    ("ocaml", &["*.ml", "*.mli", "*.mll", "*.mly"]),
    ("perl", &["*.perl", "*.pl", "*.PL", "*.plh", "*.plx", "*.pm"]),
    ("php", &["*.php", "*.php3", "*.php4", "*.php5", "*.phtml"]),
    ("py", &["*.py"]),
    ("readme", &["README*", "*README"]),
    ("r", &["*.R", "*.r", "*.Rmd", "*.Rnw"]),
    ("rst", &["*.rst"]),
    ("ruby", &["*.rb"]),
    ("rust", &["*.rs"]),
    ("scala", &["*.scala"]),
    ("sh", &["*.bash", "*.csh", "*.ksh", "*.sh", "*.tcsh"]),
    ("spark", &["*.spark"]),
    ("sql", &["*.sql"]),
    ("sv", &["*.v", "*.vg", "*.sv", "*.svh", "*.h"]),
    ("swift", &["*.swift"]),
    ("tcl", &["*.tcl"]),
    ("tex", &["*.tex", "*.cls", "*.sty"]),
    ("ts", &["*.ts", "*.tsx"]),
    ("txt", &["*.txt"]),
    ("toml", &["*.toml", "Cargo.lock"]),
    ("vala", &["*.vala"]),
    ("vb", &["*.vb"]),
    ("vimscript", &["*.vim"]),
    ("xml", &["*.xml"]),
    ("yacc", &["*.y"]),
    ("yaml", &["*.yaml", "*.yml"]),
];

/// Describes all the possible failure conditions for building a file type
/// matcher.
#[derive(Debug)]
pub enum Error {
    /// We tried to select (or negate) a file type that is not defined.
    UnrecognizedFileType(String),
    /// A user specified file type definition could not be parsed.
    InvalidDefinition,
    /// There was an error building the matcher (probably a bad glob).
    Glob(globset::Error),
    /// There was an error compiling a glob as a regex.
    Regex(regex::Error),
}

impl StdError for Error {
    fn description(&self) -> &str {
        match *self {
            Error::UnrecognizedFileType(_) => "unrecognized file type",
            Error::InvalidDefinition => "invalid definition",
            Error::Glob(ref err) => err.description(),
            Error::Regex(ref err) => err.description(),
        }
    }
}

impl fmt::Display for Error {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match *self {
            Error::UnrecognizedFileType(ref ty) => {
                write!(f, "unrecognized file type: {}", ty)
            }
            Error::InvalidDefinition => {
                write!(f, "invalid definition (format is type:glob, e.g., \
                           html:*.html)")
            }
            Error::Glob(ref err) => err.fmt(f),
            Error::Regex(ref err) => err.fmt(f),
        }
    }
}

impl From<globset::Error> for Error {
    fn from(err: globset::Error) -> Error {
        Error::Glob(err)
    }
}

impl From<regex::Error> for Error {
    fn from(err: regex::Error) -> Error {
        Error::Regex(err)
    }
}

/// A single file type definition.
#[derive(Clone, Debug)]
pub struct FileTypeDef {
    name: String,
    pats: Vec<String>,
}

impl FileTypeDef {
    /// Return the name of this file type.
    pub fn name(&self) -> &str {
        &self.name
    }

    /// Return the glob patterns used to recognize this file type.
    pub fn patterns(&self) -> &[String] {
        &self.pats
    }
}

/// Types is a file type matcher.
#[derive(Clone, Debug)]
pub struct Types {
    defs: Vec<FileTypeDef>,
    selected: Option<GlobSet>,
    negated: Option<GlobSet>,
    has_selected: bool,
    unmatched_pat: Pattern,
}

impl Types {
    /// Creates a new file type matcher from the given Gitignore matcher. If
    /// not Gitignore matcher is provided, then the file type matcher has no
    /// effect.
    ///
    /// If has_selected is true, then at least one file type was selected.
    /// Therefore, any non-matches should be ignored.
    fn new(
        selected: Option<GlobSet>,
        negated: Option<GlobSet>,
        has_selected: bool,
        defs: Vec<FileTypeDef>,
    ) -> Types {
        Types {
            defs: defs,
            selected: selected,
            negated: negated,
            has_selected: has_selected,
            unmatched_pat: Pattern {
                from: Path::new("<filetype>").to_path_buf(),
                original: "<N/A>".to_string(),
                pat: "<N/A>".to_string(),
                whitelist: false,
                only_dir: false,
            },
        }
    }

    /// Creates a new file type matcher that never matches.
    pub fn empty() -> Types {
        Types::new(None, None, false, vec![])
    }

    /// Returns a match for the given path against this file type matcher.
    ///
    /// The path is considered whitelisted if it matches a selected file type.
    /// The path is considered ignored if it matched a negated file type.
    /// If at least one file type is selected and path doesn't match, then
    /// the path is also considered ignored.
    pub fn matched<P: AsRef<Path>>(&self, path: P, is_dir: bool) -> Match {
        // If we don't have any matcher, then we can't do anything.
        if self.negated.is_none() && self.selected.is_none() {
            return Match::None;
        }
        // File types don't apply to directories.
        if is_dir {
            return Match::None;
        }
        let path = path.as_ref();
        let name = match path.file_name() {
            Some(name) => name.to_string_lossy(),
            None if self.has_selected => {
                return Match::Ignored(&self.unmatched_pat);
            }
            None => {
                return Match::None;
            }
        };
        if self.negated.as_ref().map(|s| s.is_match(&*name)).unwrap_or(false) {
            return Match::Ignored(&self.unmatched_pat);
        }
        if self.selected.as_ref().map(|s|s.is_match(&*name)).unwrap_or(false) {
            return Match::Whitelist(&self.unmatched_pat);
        }
        if self.has_selected {
            Match::Ignored(&self.unmatched_pat)
        } else {
            Match::None
        }
    }

    /// Return the set of current file type definitions.
    pub fn definitions(&self) -> &[FileTypeDef] {
        &self.defs
    }
}

/// TypesBuilder builds a type matcher from a set of file type definitions and
/// a set of file type selections.
pub struct TypesBuilder {
    types: HashMap<String, Vec<String>>,
    selected: Vec<String>,
    negated: Vec<String>,
}

impl TypesBuilder {
    /// Create a new builder for a file type matcher.
    pub fn new() -> TypesBuilder {
        TypesBuilder {
            types: HashMap::new(),
            selected: vec![],
            negated: vec![],
        }
    }

    /// Build the current set of file type definitions *and* selections into
    /// a file type matcher.
    pub fn build(&self) -> Result<Types, Error> {
        let selected_globs =
            if self.selected.is_empty() {
                None
            } else {
                let mut bset = GlobSetBuilder::new();
                for name in &self.selected {
                    let globs = match self.types.get(name) {
                        Some(globs) => globs,
                        None => {
                            let msg = name.to_string();
                            return Err(Error::UnrecognizedFileType(msg));
                        }
                    };
                    for glob in globs {
                        let pat = try!(
                            GlobBuilder::new(glob)
                                .literal_separator(true).build());
                        bset.add(pat);
                    }
                }
                Some(try!(bset.build()))
            };
        let negated_globs =
            if self.negated.is_empty() {
                None
            } else {
                let mut bset = GlobSetBuilder::new();
                for name in &self.negated {
                    let globs = match self.types.get(name) {
                        Some(globs) => globs,
                        None => {
                            let msg = name.to_string();
                            return Err(Error::UnrecognizedFileType(msg));
                        }
                    };
                    for glob in globs {
                        let pat = try!(
                            GlobBuilder::new(glob)
                                .literal_separator(true).build());
                        bset.add(pat);
                    }
                }
                Some(try!(bset.build()))
            };
        Ok(Types::new(
            selected_globs,
            negated_globs,
            !self.selected.is_empty(),
            self.definitions(),
        ))
    }

    /// Return the set of current file type definitions.
    pub fn definitions(&self) -> Vec<FileTypeDef> {
        let mut defs = vec![];
        for (ref name, ref pats) in &self.types {
            let mut pats = pats.to_vec();
            pats.sort();
            defs.push(FileTypeDef {
                name: name.to_string(),
                pats: pats,
            });
        }
        defs.sort_by(|def1, def2| def1.name().cmp(def2.name()));
        defs
    }

    /// Select the file type given by `name`.
    ///
    /// If `name` is `all`, then all file types are selected.
    pub fn select(&mut self, name: &str) -> &mut TypesBuilder {
        if name == "all" {
            for name in self.types.keys() {
                self.selected.push(name.to_string());
            }
        } else {
            self.selected.push(name.to_string());
        }
        self
    }

    /// Ignore the file type given by `name`.
    ///
    /// If `name` is `all`, then all file types are negated.
    pub fn negate(&mut self, name: &str) -> &mut TypesBuilder {
        if name == "all" {
            for name in self.types.keys() {
                self.negated.push(name.to_string());
            }
        } else {
            self.negated.push(name.to_string());
        }
        self
    }

    /// Clear any file type definitions for the type given.
    pub fn clear(&mut self, name: &str) -> &mut TypesBuilder {
        self.types.remove(name);
        self
    }

    /// Add a new file type definition. `name` can be arbitrary and `pat`
    /// should be a glob recognizing file paths belonging to the `name` type.
    pub fn add(&mut self, name: &str, pat: &str) -> &mut TypesBuilder {
        self.types.entry(name.to_string())
            .or_insert(vec![]).push(pat.to_string());
        self
    }

    /// Add a new file type definition specified in string form. The format
    /// is `name:glob`. Names may not include a colon.
    pub fn add_def(&mut self, def: &str) -> Result<(), Error> {
        let name: String = def.chars().take_while(|&c| c != ':').collect();
        let pat: String = def.chars().skip(name.chars().count() + 1).collect();
        if name.is_empty() || pat.is_empty() {
            return Err(Error::InvalidDefinition);
        }
        self.add(&name, &pat);
        Ok(())
    }

    /// Add a set of default file type definitions.
    pub fn add_defaults(&mut self) -> &mut TypesBuilder {
        for &(name, exts) in TYPE_EXTENSIONS {
            for ext in exts {
                self.add(name, ext);
            }
        }
        self
    }
}

#[cfg(test)]
mod tests {
    use super::TypesBuilder;

    macro_rules! matched {
        ($name:ident, $types:expr, $sel:expr, $selnot:expr,
         $path:expr) => {
            matched!($name, $types, $sel, $selnot, $path, true);
        };
        (not, $name:ident, $types:expr, $sel:expr, $selnot:expr,
         $path:expr) => {
            matched!($name, $types, $sel, $selnot, $path, false);
        };
        ($name:ident, $types:expr, $sel:expr, $selnot:expr,
         $path:expr, $matched:expr) => {
            #[test]
            fn $name() {
                let mut btypes = TypesBuilder::new();
                for tydef in $types {
                    btypes.add_def(tydef).unwrap();
                }
                for sel in $sel {
                    btypes.select(sel);
                }
                for selnot in $selnot {
                    btypes.negate(selnot);
                }
                let types = btypes.build().unwrap();
                let mat = types.matched($path, false);
                assert_eq!($matched, !mat.is_ignored());
            }
        };
    }

    fn types() -> Vec<&'static str> {
        vec![
            "html:*.html",
            "html:*.htm",
            "rust:*.rs",
            "js:*.js",
        ]
    }

    matched!(match1, types(), vec!["rust"], vec![], "lib.rs");
    matched!(match2, types(), vec!["html"], vec![], "index.html");
    matched!(match3, types(), vec!["html"], vec![], "index.htm");
    matched!(match4, types(), vec!["html", "rust"], vec![], "main.rs");
    matched!(match5, types(), vec![], vec![], "index.html");
    matched!(match6, types(), vec![], vec!["rust"], "index.html");

    matched!(not, matchnot1, types(), vec!["rust"], vec![], "index.html");
    matched!(not, matchnot2, types(), vec![], vec!["rust"], "main.rs");
}
<commit_msg>Add zsh type<commit_after>/*!
The types module provides a way of associating glob patterns on file names to
file types.
*/

use std::collections::HashMap;
use std::error::Error as StdError;
use std::fmt;
use std::path::Path;

use regex;

use gitignore::{Match, Pattern};
use globset::{self, GlobBuilder, GlobSet, GlobSetBuilder};

const TYPE_EXTENSIONS: &'static [(&'static str, &'static [&'static str])] = &[
    ("asm", &["*.asm", "*.s", "*.S"]),
    ("awk", &["*.awk"]),
    ("c", &["*.c", "*.h", "*.H"]),
    ("cbor", &["*.cbor"]),
    ("clojure", &["*.clj", "*.cljc", "*.cljs", "*.cljx"]),
    ("cmake", &["*.cmake", "CMakeLists.txt"]),
    ("coffeescript", &["*.coffee"]),
    ("config", &["*.config"]),
    ("cpp", &[
        "*.C", "*.cc", "*.cpp", "*.cxx",
        "*.h", "*.H", "*.hh", "*.hpp",
    ]),
    ("csharp", &["*.cs"]),
    ("css", &["*.css"]),
    ("cython", &["*.pyx"]),
    ("dart", &["*.dart"]),
    ("d", &["*.d"]),
    ("elisp", &["*.el"]),
    ("erlang", &["*.erl", "*.hrl"]),
    ("fortran", &[
        "*.f", "*.F", "*.f77", "*.F77", "*.pfo",
        "*.f90", "*.F90", "*.f95", "*.F95",
    ]),
    ("fsharp", &["*.fs", "*.fsx", "*.fsi"]),
    ("go", &["*.go"]),
    ("groovy", &["*.groovy", "*.gradle"]),
    ("hbs", &["*.hbs"]),
    ("haskell", &["*.hs", "*.lhs"]),
    ("html", &["*.htm", "*.html"]),
    ("java", &["*.java"]),
    ("jinja", &["*.jinja", "*.jinja2"]),
    ("js", &[
        "*.js", "*.jsx", "*.vue",
    ]),
    ("json", &["*.json"]),
    ("jsonl", &["*.jsonl"]),
    ("lisp", &["*.el", "*.jl", "*.lisp", "*.lsp", "*.sc", "*.scm"]),
    ("lua", &["*.lua"]),
    ("m4", &["*.ac", "*.m4"]),
    ("make", &["gnumakefile", "Gnumakefile", "makefile", "Makefile", "*.mk"]),
    ("markdown", &["*.md"]),
    ("md", &["*.md"]),
    ("matlab", &["*.m"]),
    ("mk", &["mkfile"]),
    ("ml", &["*.ml"]),
    ("nim", &["*.nim"]),
    ("objc", &["*.h", "*.m"]),
    ("objcpp", &["*.h", "*.mm"]),
    ("ocaml", &["*.ml", "*.mli", "*.mll", "*.mly"]),
    ("perl", &["*.perl", "*.pl", "*.PL", "*.plh", "*.plx", "*.pm"]),
    ("php", &["*.php", "*.php3", "*.php4", "*.php5", "*.phtml"]),
    ("py", &["*.py"]),
    ("readme", &["README*", "*README"]),
    ("r", &["*.R", "*.r", "*.Rmd", "*.Rnw"]),
    ("rst", &["*.rst"]),
    ("ruby", &["*.rb"]),
    ("rust", &["*.rs"]),
    ("scala", &["*.scala"]),
    ("sh", &["*.bash", "*.csh", "*.ksh", "*.sh", "*.tcsh"]),
    ("spark", &["*.spark"]),
    ("sql", &["*.sql"]),
    ("sv", &["*.v", "*.vg", "*.sv", "*.svh", "*.h"]),
    ("swift", &["*.swift"]),
    ("tcl", &["*.tcl"]),
    ("tex", &["*.tex", "*.cls", "*.sty"]),
    ("ts", &["*.ts", "*.tsx"]),
    ("txt", &["*.txt"]),
    ("toml", &["*.toml", "Cargo.lock"]),
    ("vala", &["*.vala"]),
    ("vb", &["*.vb"]),
    ("vimscript", &["*.vim"]),
    ("xml", &["*.xml"]),
    ("yacc", &["*.y"]),
    ("yaml", &["*.yaml", "*.yml"]),
    ("zsh", &["*.zsh", ".zshenv", ".zlogin", ".zprofile", ".zshrc"]),
];

/// Describes all the possible failure conditions for building a file type
/// matcher.
#[derive(Debug)]
pub enum Error {
    /// We tried to select (or negate) a file type that is not defined.
    UnrecognizedFileType(String),
    /// A user specified file type definition could not be parsed.
    InvalidDefinition,
    /// There was an error building the matcher (probably a bad glob).
    Glob(globset::Error),
    /// There was an error compiling a glob as a regex.
    Regex(regex::Error),
}

impl StdError for Error {
    fn description(&self) -> &str {
        match *self {
            Error::UnrecognizedFileType(_) => "unrecognized file type",
            Error::InvalidDefinition => "invalid definition",
            Error::Glob(ref err) => err.description(),
            Error::Regex(ref err) => err.description(),
        }
    }
}

impl fmt::Display for Error {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match *self {
            Error::UnrecognizedFileType(ref ty) => {
                write!(f, "unrecognized file type: {}", ty)
            }
            Error::InvalidDefinition => {
                write!(f, "invalid definition (format is type:glob, e.g., \
                           html:*.html)")
            }
            Error::Glob(ref err) => err.fmt(f),
            Error::Regex(ref err) => err.fmt(f),
        }
    }
}

impl From<globset::Error> for Error {
    fn from(err: globset::Error) -> Error {
        Error::Glob(err)
    }
}

impl From<regex::Error> for Error {
    fn from(err: regex::Error) -> Error {
        Error::Regex(err)
    }
}

/// A single file type definition.
#[derive(Clone, Debug)]
pub struct FileTypeDef {
    name: String,
    pats: Vec<String>,
}

impl FileTypeDef {
    /// Return the name of this file type.
    pub fn name(&self) -> &str {
        &self.name
    }

    /// Return the glob patterns used to recognize this file type.
    pub fn patterns(&self) -> &[String] {
        &self.pats
    }
}

/// Types is a file type matcher.
#[derive(Clone, Debug)]
pub struct Types {
    defs: Vec<FileTypeDef>,
    selected: Option<GlobSet>,
    negated: Option<GlobSet>,
    has_selected: bool,
    unmatched_pat: Pattern,
}

impl Types {
    /// Creates a new file type matcher from the given Gitignore matcher. If
    /// not Gitignore matcher is provided, then the file type matcher has no
    /// effect.
    ///
    /// If has_selected is true, then at least one file type was selected.
    /// Therefore, any non-matches should be ignored.
    fn new(
        selected: Option<GlobSet>,
        negated: Option<GlobSet>,
        has_selected: bool,
        defs: Vec<FileTypeDef>,
    ) -> Types {
        Types {
            defs: defs,
            selected: selected,
            negated: negated,
            has_selected: has_selected,
            unmatched_pat: Pattern {
                from: Path::new("<filetype>").to_path_buf(),
                original: "<N/A>".to_string(),
                pat: "<N/A>".to_string(),
                whitelist: false,
                only_dir: false,
            },
        }
    }

    /// Creates a new file type matcher that never matches.
    pub fn empty() -> Types {
        Types::new(None, None, false, vec![])
    }

    /// Returns a match for the given path against this file type matcher.
    ///
    /// The path is considered whitelisted if it matches a selected file type.
    /// The path is considered ignored if it matched a negated file type.
    /// If at least one file type is selected and path doesn't match, then
    /// the path is also considered ignored.
    pub fn matched<P: AsRef<Path>>(&self, path: P, is_dir: bool) -> Match {
        // If we don't have any matcher, then we can't do anything.
        if self.negated.is_none() && self.selected.is_none() {
            return Match::None;
        }
        // File types don't apply to directories.
        if is_dir {
            return Match::None;
        }
        let path = path.as_ref();
        let name = match path.file_name() {
            Some(name) => name.to_string_lossy(),
            None if self.has_selected => {
                return Match::Ignored(&self.unmatched_pat);
            }
            None => {
                return Match::None;
            }
        };
        if self.negated.as_ref().map(|s| s.is_match(&*name)).unwrap_or(false) {
            return Match::Ignored(&self.unmatched_pat);
        }
        if self.selected.as_ref().map(|s|s.is_match(&*name)).unwrap_or(false) {
            return Match::Whitelist(&self.unmatched_pat);
        }
        if self.has_selected {
            Match::Ignored(&self.unmatched_pat)
        } else {
            Match::None
        }
    }

    /// Return the set of current file type definitions.
    pub fn definitions(&self) -> &[FileTypeDef] {
        &self.defs
    }
}

/// TypesBuilder builds a type matcher from a set of file type definitions and
/// a set of file type selections.
pub struct TypesBuilder {
    types: HashMap<String, Vec<String>>,
    selected: Vec<String>,
    negated: Vec<String>,
}

impl TypesBuilder {
    /// Create a new builder for a file type matcher.
    pub fn new() -> TypesBuilder {
        TypesBuilder {
            types: HashMap::new(),
            selected: vec![],
            negated: vec![],
        }
    }

    /// Build the current set of file type definitions *and* selections into
    /// a file type matcher.
    pub fn build(&self) -> Result<Types, Error> {
        let selected_globs =
            if self.selected.is_empty() {
                None
            } else {
                let mut bset = GlobSetBuilder::new();
                for name in &self.selected {
                    let globs = match self.types.get(name) {
                        Some(globs) => globs,
                        None => {
                            let msg = name.to_string();
                            return Err(Error::UnrecognizedFileType(msg));
                        }
                    };
                    for glob in globs {
                        let pat = try!(
                            GlobBuilder::new(glob)
                                .literal_separator(true).build());
                        bset.add(pat);
                    }
                }
                Some(try!(bset.build()))
            };
        let negated_globs =
            if self.negated.is_empty() {
                None
            } else {
                let mut bset = GlobSetBuilder::new();
                for name in &self.negated {
                    let globs = match self.types.get(name) {
                        Some(globs) => globs,
                        None => {
                            let msg = name.to_string();
                            return Err(Error::UnrecognizedFileType(msg));
                        }
                    };
                    for glob in globs {
                        let pat = try!(
                            GlobBuilder::new(glob)
                                .literal_separator(true).build());
                        bset.add(pat);
                    }
                }
                Some(try!(bset.build()))
            };
        Ok(Types::new(
            selected_globs,
            negated_globs,
            !self.selected.is_empty(),
            self.definitions(),
        ))
    }

    /// Return the set of current file type definitions.
    pub fn definitions(&self) -> Vec<FileTypeDef> {
        let mut defs = vec![];
        for (ref name, ref pats) in &self.types {
            let mut pats = pats.to_vec();
            pats.sort();
            defs.push(FileTypeDef {
                name: name.to_string(),
                pats: pats,
            });
        }
        defs.sort_by(|def1, def2| def1.name().cmp(def2.name()));
        defs
    }

    /// Select the file type given by `name`.
    ///
    /// If `name` is `all`, then all file types are selected.
    pub fn select(&mut self, name: &str) -> &mut TypesBuilder {
        if name == "all" {
            for name in self.types.keys() {
                self.selected.push(name.to_string());
            }
        } else {
            self.selected.push(name.to_string());
        }
        self
    }

    /// Ignore the file type given by `name`.
    ///
    /// If `name` is `all`, then all file types are negated.
    pub fn negate(&mut self, name: &str) -> &mut TypesBuilder {
        if name == "all" {
            for name in self.types.keys() {
                self.negated.push(name.to_string());
            }
        } else {
            self.negated.push(name.to_string());
        }
        self
    }

    /// Clear any file type definitions for the type given.
    pub fn clear(&mut self, name: &str) -> &mut TypesBuilder {
        self.types.remove(name);
        self
    }

    /// Add a new file type definition. `name` can be arbitrary and `pat`
    /// should be a glob recognizing file paths belonging to the `name` type.
    pub fn add(&mut self, name: &str, pat: &str) -> &mut TypesBuilder {
        self.types.entry(name.to_string())
            .or_insert(vec![]).push(pat.to_string());
        self
    }

    /// Add a new file type definition specified in string form. The format
    /// is `name:glob`. Names may not include a colon.
    pub fn add_def(&mut self, def: &str) -> Result<(), Error> {
        let name: String = def.chars().take_while(|&c| c != ':').collect();
        let pat: String = def.chars().skip(name.chars().count() + 1).collect();
        if name.is_empty() || pat.is_empty() {
            return Err(Error::InvalidDefinition);
        }
        self.add(&name, &pat);
        Ok(())
    }

    /// Add a set of default file type definitions.
    pub fn add_defaults(&mut self) -> &mut TypesBuilder {
        for &(name, exts) in TYPE_EXTENSIONS {
            for ext in exts {
                self.add(name, ext);
            }
        }
        self
    }
}

#[cfg(test)]
mod tests {
    use super::TypesBuilder;

    macro_rules! matched {
        ($name:ident, $types:expr, $sel:expr, $selnot:expr,
         $path:expr) => {
            matched!($name, $types, $sel, $selnot, $path, true);
        };
        (not, $name:ident, $types:expr, $sel:expr, $selnot:expr,
         $path:expr) => {
            matched!($name, $types, $sel, $selnot, $path, false);
        };
        ($name:ident, $types:expr, $sel:expr, $selnot:expr,
         $path:expr, $matched:expr) => {
            #[test]
            fn $name() {
                let mut btypes = TypesBuilder::new();
                for tydef in $types {
                    btypes.add_def(tydef).unwrap();
                }
                for sel in $sel {
                    btypes.select(sel);
                }
                for selnot in $selnot {
                    btypes.negate(selnot);
                }
                let types = btypes.build().unwrap();
                let mat = types.matched($path, false);
                assert_eq!($matched, !mat.is_ignored());
            }
        };
    }

    fn types() -> Vec<&'static str> {
        vec![
            "html:*.html",
            "html:*.htm",
            "rust:*.rs",
            "js:*.js",
        ]
    }

    matched!(match1, types(), vec!["rust"], vec![], "lib.rs");
    matched!(match2, types(), vec!["html"], vec![], "index.html");
    matched!(match3, types(), vec!["html"], vec![], "index.htm");
    matched!(match4, types(), vec!["html", "rust"], vec![], "main.rs");
    matched!(match5, types(), vec![], vec![], "index.html");
    matched!(match6, types(), vec![], vec!["rust"], "index.html");

    matched!(not, matchnot1, types(), vec!["rust"], vec![], "index.html");
    matched!(not, matchnot2, types(), vec![], vec!["rust"], "main.rs");
}
<|endoftext|>