<commit_before>use std::f32::consts::PI;
use std::fs::File;
use std::os::raw::c_void;
use std::path::Path;
use std::process;
use std::time::Instant;

use cgmath::{ Point3 };
use gl;
use glutin;
use glutin::{
    Api,
    MouseScrollDelta,
    MouseButton,
    GlContext,
    GlRequest,
    GlProfile,
    VirtualKeyCode,
    WindowEvent,
};
use glutin::ElementState::*;

use gltf_importer;
use gltf_importer::config::ValidationStrategy;
use image::{DynamicImage, ImageFormat};

use controls::{OrbitControls, NavState};
use controls::CameraMovement::*;
use framebuffer::Framebuffer;
use render::*;
use render::math::*;
use utils::{print_elapsed, FrameTimer, gl_check_error, print_context_info};

// TODO!: complete and pass through draw calls? or get rid of multiple shaders?
// How about state ordering anyway?
// struct DrawState {
//     current_shader: ShaderFlags,
//     back_face_culling_enabled: bool
// }

pub struct CameraOptions {
    pub index: Option<u32>,
    pub position: Option<Vector3>,
    pub target: Option<Vector3>,
    pub fovy: f32,
}

pub struct GltfViewer {
    width: u32,
    height: u32,

    orbit_controls: OrbitControls,
    first_mouse: bool,
    last_x: f32,
    last_y: f32,
    events_loop: Option<glutin::EventsLoop>,
    gl_window: Option<glutin::GlWindow>,

    // TODO!: get rid of scene?
    root: Root,
    scene: Scene,

    delta_time: f64, // seconds
    last_frame: Instant,

    render_timer: FrameTimer,
}

/// Note about `headless` and `visible`: True headless rendering doesn't work on
/// all operating systems, but an invisible window usually works
impl GltfViewer {
    pub fn new(source: &str, width: u32, height: u32, headless: bool, visible: bool, camera_options: CameraOptions) -> GltfViewer {
        let gl_request = GlRequest::Specific(Api::OpenGl, (3, 3));
        let gl_profile = GlProfile::Core;
        let (events_loop, gl_window, width, height) =
            if headless {
                let headless_context = glutin::HeadlessRendererBuilder::new(width, height)
                    // .with_gl(gl_request)
                    // .with_gl_profile(gl_profile)
                    .build()
                    .unwrap();
                unsafe { headless_context.make_current().unwrap() }
                gl::load_with(|symbol| headless_context.get_proc_address(symbol) as *const _);
                let framebuffer = Framebuffer::new(width, height);
                framebuffer.bind();
                unsafe { gl::Viewport(0, 0, width as i32, height as i32); }

                (None, None, width, height) // TODO: real height (retina?)
            }
            else {
                // glutin: initialize and configure
                let events_loop = glutin::EventsLoop::new();

                // TODO?: hints for 4.1, core profile, forward compat
                let window = glutin::WindowBuilder::new()
                        .with_title("gltf-viewer")
                        .with_dimensions(width, height)
                        .with_visibility(visible);

                let context = glutin::ContextBuilder::new()
                    .with_gl(gl_request)
                    .with_gl_profile(gl_profile)
                    .with_vsync(true);
                let gl_window = glutin::GlWindow::new(window, context, &events_loop).unwrap();

                // Real dimensions might be much higher on High-DPI displays
                let (real_width, real_height) = gl_window.get_inner_size().unwrap();

                unsafe { gl_window.make_current().unwrap(); }

                // gl: load all OpenGL function pointers
                gl::load_with(|symbol| gl_window.get_proc_address(symbol) as *const _);

                (Some(events_loop), Some(gl_window), real_width, real_height)
            };

        let mut orbit_controls = OrbitControls::new(
            Point3::new(0.0, 0.0, 2.0), width as f32, height as f32
        );
        orbit_controls.camera = Camera::default();
        orbit_controls.camera.fovy = camera_options.fovy;
        orbit_controls.camera.update_aspect_ratio(width as f32 / height as f32); // updates projection matrix

        let first_mouse = true;
        let last_x: f32 = width as f32 / 2.0;
        let last_y: f32 = height as f32 / 2.0;

        unsafe {
            print_context_info();

            gl::ClearColor(0.0, 1.0, 0.0, 1.0); // green for debugging
            gl::Clear(gl::COLOR_BUFFER_BIT | gl::DEPTH_BUFFER_BIT);

            if headless || !visible {
                // transparent background for screenshots
                gl::ClearColor(0.0, 0.0, 0.0, 0.0);
            }
            else {
                gl::ClearColor(0.1, 0.2, 0.3, 1.0);
            }

            gl::Enable(gl::DEPTH_TEST);

            // TODO: keyboard switch?
            // draw in wireframe
            // gl::PolygonMode(gl::FRONT_AND_BACK, gl::LINE);
        };

        let (root, scene) = Self::load(source);
        let mut viewer = GltfViewer {
            width,
            height,

            orbit_controls,
            first_mouse, last_x, last_y,

            events_loop,
            gl_window,

            root,
            scene,

            delta_time: 0.0, // seconds
            last_frame: Instant::now(),

            render_timer: FrameTimer::new("rendering", 300),
        };
        unsafe { gl_check_error!(); };

        if let Some(cam_index) = camera_options.index {
            if cam_index >= viewer.root.camera_nodes.len() as u32 {
                error!("No camera with index {} found in glTF file (max: {})",
                    cam_index, viewer.root.camera_nodes.len() - 1);
                process::exit(2)
            }
            let cam_node = &viewer.root.get_camera_node(cam_index as usize);
            viewer.orbit_controls.set_camera(
                cam_node.camera.as_ref().unwrap(),
                &cam_node.final_transform);

            if camera_options.position.is_some() || camera_options.target.is_some() {
                warn!("Ignoring --cam-pos / --cam-target since --cam-index is given.")
            }
        } else {
            viewer.set_camera_from_bounds();

            if let Some(p) = camera_options.position {
                viewer.orbit_controls.position = Point3::from_vec(p)
            }
            if let Some(target) = camera_options.target {
                viewer.orbit_controls.target = Point3::from_vec(target)
            }
        }

        viewer
    }

    pub fn load(source: &str) -> (Root, Scene) {
        let mut start_time = Instant::now();
        // TODO!: http source
        // let gltf =
        if source.starts_with("http") {
            panic!("not implemented: HTTP support temporarily removed.")
            // let http_source = HttpSource::new(source);
            // let import = gltf::Import::custom(http_source, Default::default());
            // let gltf = import_gltf(import);
            // println!(); // to end the "progress dots"
            // gltf
        }
        //     else {
        let config = gltf_importer::Config { validation_strategy: ValidationStrategy::Complete };
        let (gltf, buffers) = match gltf_importer::import_with_config(source, config) {
            Ok((gltf, buffers)) => (gltf, buffers),
            Err(err) => {
                error!("glTF import failed: {:?}", err);
                if let gltf_importer::Error::Io(_) = err {
                    error!("Hint: Are the .bin file(s) referenced by the .gltf file available?")
                }
                process::exit(1)
            },
        };

        print_elapsed("Imported glTF in ", &start_time);
        start_time = Instant::now();

        // load first scene
        if gltf.scenes().len() > 1 {
            warn!("Found more than 1 scene, can only load first at the moment.")
        }
        let base_path = Path::new(source);
        let mut root = Root::from_gltf(&gltf, &buffers, base_path);
        let scene = Scene::from_gltf(&gltf.scenes().nth(0).unwrap(), &mut root);
        print_elapsed(&format!("Loaded scene with {} nodes, {} meshes in ",
                gltf.nodes().count(), root.meshes.len()), &start_time);

        (root, scene)
    }

    /// determine "nice" camera perspective from bounding box. Inspired by donmccurdy/three-gltf-viewer
    fn set_camera_from_bounds(&mut self) {
        // TODO!!: fix bounds/camera computation (many models NOT centered)
        let bounds = &self.scene.bounds;
        let size = bounds.size().magnitude();
        let center = bounds.center();

        // TODO!: move cam instead?
        let _obj_pos_modifier = -center;

        let _max_distance = size * 10.0;
        // TODO: x,y addition optional, z optionally minus instead
        let cam_pos = Point3::new(
            center.x + size / 2.0,
            center.y + size / 5.0,
            center.z + size / 2.0,
        );
        let _near = size / 100.0;
        let _far = size * 100.0;

        self.orbit_controls.position = cam_pos;
        self.orbit_controls.target = Point3::from_vec(center);

        // TODO!: set near, far, max_distance, obj_pos_modifier...
    }

    pub fn start_render_loop(&mut self) {
        loop {
            // per-frame time logic
            // NOTE: Deliberately ignoring the seconds of `elapsed()`
            self.delta_time = f64::from(self.last_frame.elapsed().subsec_nanos()) / 1_000_000_000.0;
            self.last_frame = Instant::now();

            // events
            let keep_running = process_events(
                &mut self.events_loop.as_mut().unwrap(), self.gl_window.as_mut().unwrap(),
                &mut self.orbit_controls,
                &mut self.width, &mut self.height);
            if !keep_running {
                unsafe { gl_check_error!(); } // final error check so errors don't go unnoticed
                break
            }

            self.orbit_controls.frame_update(self.delta_time); // keyboard navigation

            self.draw();

            self.gl_window.as_ref().unwrap().swap_buffers().unwrap();
        }
    }

    // Returns whether to keep running
    pub fn draw(&mut self) {
        // render
        unsafe {
            self.render_timer.start();

            gl::Clear(gl::COLOR_BUFFER_BIT | gl::DEPTH_BUFFER_BIT);

            let cam_params = self.orbit_controls.camera_params();
            self.scene.draw(&mut self.root, &cam_params);

            self.render_timer.end();
        }
    }

    pub fn screenshot(&mut self, filename: &str, width: u32, height: u32) {
        self.draw();

        let mut img = DynamicImage::new_rgba8(width, height);
        unsafe {
            let pixels = img.as_mut_rgba8().unwrap();
            gl::PixelStorei(gl::PACK_ALIGNMENT, 1);
            gl::ReadPixels(0, 0, width as i32, height as i32, gl::RGBA,
                gl::UNSIGNED_BYTE, pixels.as_mut_ptr() as *mut c_void);
            gl_check_error!();
        }

        let img = img.flipv();

        let mut file = File::create(filename).unwrap();
        if let Err(err) = img.save(&mut file, ImageFormat::PNG) {
            error!("{}", err);
        }
        else {
            println!("Saved {}x{} screenshot to {}", width, height, filename);
        }
    }
    pub fn multiscreenshot(&mut self, filename: &str, width: u32, height: u32, count: u32) {
        let min_angle : f32 = 0.0 ;
        let max_angle : f32 =  2.0 * PI ;
        let increment_angle : f32 = ((max_angle - min_angle)/(count as f32)) as f32;
        for i in 1..(count+1) {
            self.orbit_controls.rotate_object(increment_angle);
            let dot = filename.rfind('.').unwrap_or_else(|| filename.len());
            let mut actual_name = filename.to_string();
            actual_name.insert_str(dot, &format!("_{}", i));
            self.screenshot(&actual_name[..], width,height);
        }
    }
}

#[allow(too_many_arguments)]
fn process_events(
    events_loop: &mut glutin::EventsLoop,
    gl_window: &glutin::GlWindow,
    mut orbit_controls: &mut OrbitControls,
    width: &mut u32,
    height: &mut u32) -> bool
{
    let mut keep_running = true;
    #[allow(single_match)]
    events_loop.poll_events(|event| {
        match event {
            glutin::Event::WindowEvent{ event, .. } => match event {
                WindowEvent::Closed => keep_running = false,
                WindowEvent::Resized(w, h) => {
                    gl_window.resize(w, h);
                    *width = w;
                    *height = h;
                    let w = w as f32;
                    let h = h as f32;
                    orbit_controls.camera.update_aspect_ratio(w / h);
                    orbit_controls.screen_width = w;
                    orbit_controls.screen_height = h;

                    trace!("Resized to {}x{}", w, h);
                },
                WindowEvent::DroppedFile(_path_buf) => (), // TODO: drag file in
                WindowEvent::MouseInput { button, state: Pressed, ..} => {
                    match button {
                        MouseButton::Left => {
                            orbit_controls.state = NavState::Rotating;
                        },
                        MouseButton::Right => {
                            orbit_controls.state = NavState::Panning;
                        },
                        _ => ()
                    }
                },
                WindowEvent::MouseInput { button, state: Released, ..} => {
                    match (button, orbit_controls.state.clone()) {
                        (MouseButton::Left, NavState::Rotating) | (MouseButton::Right, NavState::Panning) => {
                            orbit_controls.state = NavState::None;
                            orbit_controls.handle_mouse_up();
                        },
                        _ => ()
                    }
                }
                WindowEvent::CursorMoved { position: (xpos, ypos), .. } => {
                    let (xpos, ypos) = (xpos as f32, ypos as f32);
                    orbit_controls.handle_mouse_move(xpos, ypos);
                },
                WindowEvent::MouseWheel { delta: MouseScrollDelta::PixelDelta(_xoffset, yoffset), .. } => {
                    // TODO: need to handle LineDelta case too?
                    orbit_controls.process_mouse_scroll(yoffset);
                }
                WindowEvent::KeyboardInput { input, .. } => {
                    keep_running = process_input(input, &mut orbit_controls);
                }
                _ => ()
            },
            _ => ()
        }
    });

    keep_running
}

fn process_input(input: glutin::KeyboardInput, controls: &mut OrbitControls) -> bool {
    let pressed = match input.state {
        Pressed => true,
        Released => false
    };
    if let Some(code) = input.virtual_keycode {
        match code {
            VirtualKeyCode::Escape if pressed => return false,
            VirtualKeyCode::W => controls.process_keyboard(FORWARD, pressed),
            VirtualKeyCode::S => controls.process_keyboard(BACKWARD, pressed),
            VirtualKeyCode::A => controls.process_keyboard(LEFT, pressed),
            VirtualKeyCode::D => controls.process_keyboard(RIGHT, pressed),
            _ => ()
        }
    }
    true
}
<commit_msg>add cursor keys (same as WASD)<commit_after>use std::f32::consts::PI;
use std::fs::File;
use std::os::raw::c_void;
use std::path::Path;
use std::process;
use std::time::Instant;

use cgmath::{ Point3 };
use gl;
use glutin;
use glutin::{
    Api,
    MouseScrollDelta,
    MouseButton,
    GlContext,
    GlRequest,
    GlProfile,
    VirtualKeyCode,
    WindowEvent,
};
use glutin::ElementState::*;

use gltf_importer;
use gltf_importer::config::ValidationStrategy;
use image::{DynamicImage, ImageFormat};

use controls::{OrbitControls, NavState};
use controls::CameraMovement::*;
use framebuffer::Framebuffer;
use render::*;
use render::math::*;
use utils::{print_elapsed, FrameTimer, gl_check_error, print_context_info};

// TODO!: complete and pass through draw calls? or get rid of multiple shaders?
// How about state ordering anyway?
// struct DrawState {
//     current_shader: ShaderFlags,
//     back_face_culling_enabled: bool
// }

pub struct CameraOptions {
    pub index: Option<u32>,
    pub position: Option<Vector3>,
    pub target: Option<Vector3>,
    pub fovy: f32,
}

pub struct GltfViewer {
    width: u32,
    height: u32,

    orbit_controls: OrbitControls,
    first_mouse: bool,
    last_x: f32,
    last_y: f32,
    events_loop: Option<glutin::EventsLoop>,
    gl_window: Option<glutin::GlWindow>,

    // TODO!: get rid of scene?
    root: Root,
    scene: Scene,

    delta_time: f64, // seconds
    last_frame: Instant,

    render_timer: FrameTimer,
}

/// Note about `headless` and `visible`: True headless rendering doesn't work on
/// all operating systems, but an invisible window usually works
impl GltfViewer {
    pub fn new(source: &str, width: u32, height: u32, headless: bool, visible: bool, camera_options: CameraOptions) -> GltfViewer {
        let gl_request = GlRequest::Specific(Api::OpenGl, (3, 3));
        let gl_profile = GlProfile::Core;
        let (events_loop, gl_window, width, height) =
            if headless {
                let headless_context = glutin::HeadlessRendererBuilder::new(width, height)
                    // .with_gl(gl_request)
                    // .with_gl_profile(gl_profile)
                    .build()
                    .unwrap();
                unsafe { headless_context.make_current().unwrap() }
                gl::load_with(|symbol| headless_context.get_proc_address(symbol) as *const _);
                let framebuffer = Framebuffer::new(width, height);
                framebuffer.bind();
                unsafe { gl::Viewport(0, 0, width as i32, height as i32); }

                (None, None, width, height) // TODO: real height (retina?)
            }
            else {
                // glutin: initialize and configure
                let events_loop = glutin::EventsLoop::new();

                // TODO?: hints for 4.1, core profile, forward compat
                let window = glutin::WindowBuilder::new()
                        .with_title("gltf-viewer")
                        .with_dimensions(width, height)
                        .with_visibility(visible);

                let context = glutin::ContextBuilder::new()
                    .with_gl(gl_request)
                    .with_gl_profile(gl_profile)
                    .with_vsync(true);
                let gl_window = glutin::GlWindow::new(window, context, &events_loop).unwrap();

                // Real dimensions might be much higher on High-DPI displays
                let (real_width, real_height) = gl_window.get_inner_size().unwrap();

                unsafe { gl_window.make_current().unwrap(); }

                // gl: load all OpenGL function pointers
                gl::load_with(|symbol| gl_window.get_proc_address(symbol) as *const _);

                (Some(events_loop), Some(gl_window), real_width, real_height)
            };

        let mut orbit_controls = OrbitControls::new(
            Point3::new(0.0, 0.0, 2.0), width as f32, height as f32
        );
        orbit_controls.camera = Camera::default();
        orbit_controls.camera.fovy = camera_options.fovy;
        orbit_controls.camera.update_aspect_ratio(width as f32 / height as f32); // updates projection matrix

        let first_mouse = true;
        let last_x: f32 = width as f32 / 2.0;
        let last_y: f32 = height as f32 / 2.0;

        unsafe {
            print_context_info();

            gl::ClearColor(0.0, 1.0, 0.0, 1.0); // green for debugging
            gl::Clear(gl::COLOR_BUFFER_BIT | gl::DEPTH_BUFFER_BIT);

            if headless || !visible {
                // transparent background for screenshots
                gl::ClearColor(0.0, 0.0, 0.0, 0.0);
            }
            else {
                gl::ClearColor(0.1, 0.2, 0.3, 1.0);
            }

            gl::Enable(gl::DEPTH_TEST);

            // TODO: keyboard switch?
            // draw in wireframe
            // gl::PolygonMode(gl::FRONT_AND_BACK, gl::LINE);
        };

        let (root, scene) = Self::load(source);
        let mut viewer = GltfViewer {
            width,
            height,

            orbit_controls,
            first_mouse, last_x, last_y,

            events_loop,
            gl_window,

            root,
            scene,

            delta_time: 0.0, // seconds
            last_frame: Instant::now(),

            render_timer: FrameTimer::new("rendering", 300),
        };
        unsafe { gl_check_error!(); };

        if let Some(cam_index) = camera_options.index {
            if cam_index >= viewer.root.camera_nodes.len() as u32 {
                error!("No camera with index {} found in glTF file (max: {})",
                    cam_index, viewer.root.camera_nodes.len() - 1);
                process::exit(2)
            }
            let cam_node = &viewer.root.get_camera_node(cam_index as usize);
            viewer.orbit_controls.set_camera(
                cam_node.camera.as_ref().unwrap(),
                &cam_node.final_transform);

            if camera_options.position.is_some() || camera_options.target.is_some() {
                warn!("Ignoring --cam-pos / --cam-target since --cam-index is given.")
            }
        } else {
            viewer.set_camera_from_bounds();

            if let Some(p) = camera_options.position {
                viewer.orbit_controls.position = Point3::from_vec(p)
            }
            if let Some(target) = camera_options.target {
                viewer.orbit_controls.target = Point3::from_vec(target)
            }
        }

        viewer
    }

    pub fn load(source: &str) -> (Root, Scene) {
        let mut start_time = Instant::now();
        // TODO!: http source
        // let gltf =
        if source.starts_with("http") {
            panic!("not implemented: HTTP support temporarily removed.")
            // let http_source = HttpSource::new(source);
            // let import = gltf::Import::custom(http_source, Default::default());
            // let gltf = import_gltf(import);
            // println!(); // to end the "progress dots"
            // gltf
        }
        //     else {
        let config = gltf_importer::Config { validation_strategy: ValidationStrategy::Complete };
        let (gltf, buffers) = match gltf_importer::import_with_config(source, config) {
            Ok((gltf, buffers)) => (gltf, buffers),
            Err(err) => {
                error!("glTF import failed: {:?}", err);
                if let gltf_importer::Error::Io(_) = err {
                    error!("Hint: Are the .bin file(s) referenced by the .gltf file available?")
                }
                process::exit(1)
            },
        };

        print_elapsed("Imported glTF in ", &start_time);
        start_time = Instant::now();

        // load first scene
        if gltf.scenes().len() > 1 {
            warn!("Found more than 1 scene, can only load first at the moment.")
        }
        let base_path = Path::new(source);
        let mut root = Root::from_gltf(&gltf, &buffers, base_path);
        let scene = Scene::from_gltf(&gltf.scenes().nth(0).unwrap(), &mut root);
        print_elapsed(&format!("Loaded scene with {} nodes, {} meshes in ",
                gltf.nodes().count(), root.meshes.len()), &start_time);

        (root, scene)
    }

    /// determine "nice" camera perspective from bounding box. Inspired by donmccurdy/three-gltf-viewer
    fn set_camera_from_bounds(&mut self) {
        // TODO!!: fix bounds/camera computation (many models NOT centered)
        let bounds = &self.scene.bounds;
        let size = bounds.size().magnitude();
        let center = bounds.center();

        // TODO!: move cam instead?
        let _obj_pos_modifier = -center;

        let _max_distance = size * 10.0;
        // TODO: x,y addition optional, z optionally minus instead
        let cam_pos = Point3::new(
            center.x + size / 2.0,
            center.y + size / 5.0,
            center.z + size / 2.0,
        );
        let _near = size / 100.0;
        let _far = size * 100.0;

        self.orbit_controls.position = cam_pos;
        self.orbit_controls.target = Point3::from_vec(center);

        // TODO!: set near, far, max_distance, obj_pos_modifier...
    }

    pub fn start_render_loop(&mut self) {
        loop {
            // per-frame time logic
            // NOTE: Deliberately ignoring the seconds of `elapsed()`
            self.delta_time = f64::from(self.last_frame.elapsed().subsec_nanos()) / 1_000_000_000.0;
            self.last_frame = Instant::now();

            // events
            let keep_running = process_events(
                &mut self.events_loop.as_mut().unwrap(), self.gl_window.as_mut().unwrap(),
                &mut self.orbit_controls,
                &mut self.width, &mut self.height);
            if !keep_running {
                unsafe { gl_check_error!(); } // final error check so errors don't go unnoticed
                break
            }

            self.orbit_controls.frame_update(self.delta_time); // keyboard navigation

            self.draw();

            self.gl_window.as_ref().unwrap().swap_buffers().unwrap();
        }
    }

    // Returns whether to keep running
    pub fn draw(&mut self) {
        // render
        unsafe {
            self.render_timer.start();

            gl::Clear(gl::COLOR_BUFFER_BIT | gl::DEPTH_BUFFER_BIT);

            let cam_params = self.orbit_controls.camera_params();
            self.scene.draw(&mut self.root, &cam_params);

            self.render_timer.end();
        }
    }

    pub fn screenshot(&mut self, filename: &str, width: u32, height: u32) {
        self.draw();

        let mut img = DynamicImage::new_rgba8(width, height);
        unsafe {
            let pixels = img.as_mut_rgba8().unwrap();
            gl::PixelStorei(gl::PACK_ALIGNMENT, 1);
            gl::ReadPixels(0, 0, width as i32, height as i32, gl::RGBA,
                gl::UNSIGNED_BYTE, pixels.as_mut_ptr() as *mut c_void);
            gl_check_error!();
        }

        let img = img.flipv();

        let mut file = File::create(filename).unwrap();
        if let Err(err) = img.save(&mut file, ImageFormat::PNG) {
            error!("{}", err);
        }
        else {
            println!("Saved {}x{} screenshot to {}", width, height, filename);
        }
    }
    pub fn multiscreenshot(&mut self, filename: &str, width: u32, height: u32, count: u32) {
        let min_angle : f32 = 0.0 ;
        let max_angle : f32 =  2.0 * PI ;
        let increment_angle : f32 = ((max_angle - min_angle)/(count as f32)) as f32;
        for i in 1..(count+1) {
            self.orbit_controls.rotate_object(increment_angle);
            let dot = filename.rfind('.').unwrap_or_else(|| filename.len());
            let mut actual_name = filename.to_string();
            actual_name.insert_str(dot, &format!("_{}", i));
            self.screenshot(&actual_name[..], width,height);
        }
    }
}

#[allow(too_many_arguments)]
fn process_events(
    events_loop: &mut glutin::EventsLoop,
    gl_window: &glutin::GlWindow,
    mut orbit_controls: &mut OrbitControls,
    width: &mut u32,
    height: &mut u32) -> bool
{
    let mut keep_running = true;
    #[allow(single_match)]
    events_loop.poll_events(|event| {
        match event {
            glutin::Event::WindowEvent{ event, .. } => match event {
                WindowEvent::Closed => keep_running = false,
                WindowEvent::Resized(w, h) => {
                    gl_window.resize(w, h);
                    *width = w;
                    *height = h;
                    let w = w as f32;
                    let h = h as f32;
                    orbit_controls.camera.update_aspect_ratio(w / h);
                    orbit_controls.screen_width = w;
                    orbit_controls.screen_height = h;

                    trace!("Resized to {}x{}", w, h);
                },
                WindowEvent::DroppedFile(_path_buf) => (), // TODO: drag file in
                WindowEvent::MouseInput { button, state: Pressed, ..} => {
                    match button {
                        MouseButton::Left => {
                            orbit_controls.state = NavState::Rotating;
                        },
                        MouseButton::Right => {
                            orbit_controls.state = NavState::Panning;
                        },
                        _ => ()
                    }
                },
                WindowEvent::MouseInput { button, state: Released, ..} => {
                    match (button, orbit_controls.state.clone()) {
                        (MouseButton::Left, NavState::Rotating) | (MouseButton::Right, NavState::Panning) => {
                            orbit_controls.state = NavState::None;
                            orbit_controls.handle_mouse_up();
                        },
                        _ => ()
                    }
                }
                WindowEvent::CursorMoved { position: (xpos, ypos), .. } => {
                    let (xpos, ypos) = (xpos as f32, ypos as f32);
                    orbit_controls.handle_mouse_move(xpos, ypos);
                },
                WindowEvent::MouseWheel { delta: MouseScrollDelta::PixelDelta(_xoffset, yoffset), .. } => {
                    // TODO: need to handle LineDelta case too?
                    orbit_controls.process_mouse_scroll(yoffset);
                }
                WindowEvent::KeyboardInput { input, .. } => {
                    keep_running = process_input(input, &mut orbit_controls);
                }
                _ => ()
            },
            _ => ()
        }
    });

    keep_running
}

fn process_input(input: glutin::KeyboardInput, controls: &mut OrbitControls) -> bool {
    let pressed = match input.state {
        Pressed => true,
        Released => false
    };
    if let Some(code) = input.virtual_keycode {
        match code {
            VirtualKeyCode::Escape if pressed => return false,
            VirtualKeyCode::W | VirtualKeyCode::Up    => controls.process_keyboard(FORWARD, pressed),
            VirtualKeyCode::S | VirtualKeyCode::Down  => controls.process_keyboard(BACKWARD, pressed),
            VirtualKeyCode::A | VirtualKeyCode::Left  => controls.process_keyboard(LEFT, pressed),
            VirtualKeyCode::D | VirtualKeyCode::Right => controls.process_keyboard(RIGHT, pressed),
            _ => ()
        }
    }
    true
}
<|endoftext|>


<commit_before>/// Parser method generator.
///
/// Invoke inside an impl block. Declare a list of signatures of methods callable on `&mut self`.
/// The macro constructs a `parse(&mut self, input: &str) -> Result<(), String>` method that parses
/// input strings like "method_name 4 foo" into calls of `self.method_name(FromStr::from_str("4")?,
/// FromStr::from_str("foo")?)`.
///
///     # #[macro_use] extern crate calx_alg;
///     # fn main() {
///
///     struct Adder {
///         pub num: u32,
///     }
///
///     impl Adder {
///         fn add(&mut self, amount: u32) {
///             self.num += amount;
///         }
///
///         command_parser!{
///             fn add(&mut self, amount: u32);
///         }
///     }
///
///     let mut a = Adder { num: 1 };
///     a.add(2);
///     assert_eq!(a.num, 3);
///
///     assert!(a.parse("foobie bletch").is_err());
///     assert!(a.parse("add").is_err());       // Too few arguments
///     assert!(a.parse("add derp").is_err());  // Arguments won't parse to unsigned integer.
///     assert!(a.parse("add -1").is_err());
///     assert!(a.parse("add 1.0").is_err());
///
///     // NB: Due to limits of the macro implementation, trailing arguments
///     // are ignored instead of causing an error.
///     // assert!(a.parse("add 1 2").is_err());   // Too many arguments
///
///     assert_eq!(a.num, 3);
///
///     assert!(a.parse("add 3").is_ok());
///     assert_eq!(a.num, 6);
///
///     # }
#[macro_export]
macro_rules! command_parser {
    {
        $(fn $method:ident(&mut self$(, $argname:ident: $argtype:ty)*);)+
    } => {
        fn parse(&mut self, input: &str) -> Result<(), Box<::std::error::Error>> {
            #[derive(Debug)]
            struct ParseError(String);

            impl ::std::fmt::Display for ParseError {
                fn fmt(&self, fmt: &mut ::std::fmt::Formatter) -> Result<(), ::std::fmt::Error> {
                    self.0.fmt(fmt)
                }
            }

            impl ::std::error::Error for ParseError {
                fn description(&self) -> &str { &self.0 }
            }

            fn parse_err(s: String) -> Result<(), Box<::std::error::Error>> {
                Err(Box::new(ParseError(s)))
            }

            use ::std::str::FromStr;
            let input = input.trim_right();
            let mut elts = input.split(" ");
            let cmd = elts.next();
            if let Some(cmd) = cmd {
                match cmd {
                    $(
                    stringify!($method) => {
                        self.$method(
                            $({
                                if let Some(arg) = elts.next() {
                                    let param: $argtype = FromStr::from_str(arg)?;
                                    param
                                } else {
                                    return parse_err(format!("Too few arguments"));
                                }
                            }),*
                        );
                        Ok(())
                    }
                    )+
                    x => parse_err(format!("Unknown command {}", x)),
                }
            } else {
                return parse_err(format!("No input"));
            }
        }
    }
}
<commit_msg>Protect parser from redefined Result<commit_after>/// Parser method generator.
///
/// Invoke inside an impl block. Declare a list of signatures of methods callable on `&mut self`.
/// The macro constructs a `parse(&mut self, input: &str) -> Result<(), String>` method that parses
/// input strings like "method_name 4 foo" into calls of `self.method_name(FromStr::from_str("4")?,
/// FromStr::from_str("foo")?)`.
///
///     # #[macro_use] extern crate calx_alg;
///     # fn main() {
///
///     struct Adder {
///         pub num: u32,
///     }
///
///     impl Adder {
///         fn add(&mut self, amount: u32) {
///             self.num += amount;
///         }
///
///         command_parser!{
///             fn add(&mut self, amount: u32);
///         }
///     }
///
///     let mut a = Adder { num: 1 };
///     a.add(2);
///     assert_eq!(a.num, 3);
///
///     assert!(a.parse("foobie bletch").is_err());
///     assert!(a.parse("add").is_err());       // Too few arguments
///     assert!(a.parse("add derp").is_err());  // Arguments won't parse to unsigned integer.
///     assert!(a.parse("add -1").is_err());
///     assert!(a.parse("add 1.0").is_err());
///
///     // NB: Due to limits of the macro implementation, trailing arguments
///     // are ignored instead of causing an error.
///     // assert!(a.parse("add 1 2").is_err());   // Too many arguments
///
///     assert_eq!(a.num, 3);
///
///     assert!(a.parse("add 3").is_ok());
///     assert_eq!(a.num, 6);
///
///     # }
#[macro_export]
macro_rules! command_parser {
    {
        $(fn $method:ident(&mut self$(, $argname:ident: $argtype:ty)*);)+
    } => {
        fn parse(&mut self, input: &str) -> ::std::result::Result<(), Box<::std::error::Error>> {
            #[derive(Debug)]
            struct ParseError(String);

            impl ::std::fmt::Display for ParseError {
                fn fmt(&self, fmt: &mut ::std::fmt::Formatter) -> ::std::result::Result<(), ::std::fmt::Error> {
                    self.0.fmt(fmt)
                }
            }

            impl ::std::error::Error for ParseError {
                fn description(&self) -> &str { &self.0 }
            }

            fn parse_err(s: String) -> ::std::result::Result<(), Box<::std::error::Error>> {
                Err(Box::new(ParseError(s)))
            }

            use ::std::str::FromStr;
            let input = input.trim_right();
            let mut elts = input.split(" ");
            let cmd = elts.next();
            if let Some(cmd) = cmd {
                match cmd {
                    $(
                    stringify!($method) => {
                        self.$method(
                            $({
                                if let Some(arg) = elts.next() {
                                    let param: $argtype = FromStr::from_str(arg)?;
                                    param
                                } else {
                                    return parse_err(format!("Too few arguments"));
                                }
                            }),*
                        );
                        Ok(())
                    }
                    )+
                    x => parse_err(format!("Unknown command {}", x)),
                }
            } else {
                return parse_err(format!("No input"));
            }
        }
    }
}
<|endoftext|>


<commit_before>//! Handle network connections for a varlink service

//#![feature(getpid)]
//use std::process;
// FIXME
use libc::getpid;
use std::env;
use std::fs;
use std::io;
use std::io::{Error, ErrorKind, Read, Write};
use std::net::{Shutdown, TcpListener, TcpStream};
use std::os::unix::io::{FromRawFd, IntoRawFd};
use std::os::unix::net::{UnixListener, UnixStream};
use std::sync::{Arc, Mutex};
use std::sync::mpsc;
use std::thread;
use std::time::Duration;
// FIXME: abstract unix domains sockets still not in std
// FIXME: https://github.com/rust-lang/rust/issues/14194
use unix_socket::UnixListener as AbstractUnixListener;

enum VarlinkListener {
    TCP(TcpListener),
    UNIX(UnixListener),
}

enum VarlinkStream {
    TCP(TcpStream),
    UNIX(UnixStream),
}

impl<'a> VarlinkStream {
    pub fn split(&mut self) -> io::Result<(Box<Read + Send + Sync>, Box<Write + Send + Sync>)> {
        match *self {
            VarlinkStream::TCP(ref mut s) => {
                Ok((Box::new(s.try_clone()?), Box::new(s.try_clone()?)))
            }
            VarlinkStream::UNIX(ref mut s) => {
                Ok((Box::new(s.try_clone()?), Box::new(s.try_clone()?)))
            }
        }
    }
    pub fn shutdown(&mut self) -> io::Result<()> {
        match *self {
            VarlinkStream::TCP(ref mut s) => s.shutdown(Shutdown::Both),
            VarlinkStream::UNIX(ref mut s) => s.shutdown(Shutdown::Both),
        }
    }
}

fn activation_listener() -> io::Result<Option<i32>> {
    let nfds: u32;

    match env::var("LISTEN_FDS") {
        Ok(ref n) => match n.parse::<u32>() {
            Ok(n) if n >= 1 => nfds = n,
            _ => return Ok(None),
        },
        _ => return Ok(None),
    }

    unsafe {
        match env::var("LISTEN_PID") {
            //FIXME: replace Ok(getpid()) with Ok(process::id())
            Ok(ref pid) if pid.parse::<i32>() == Ok(getpid()) => {}
            _ => return Ok(None),
        }
    }

    let fdnames: String;

    match env::var("LISTEN_FDNAMES") {
        Ok(n) => {
            fdnames = n;
        }
        _ => if nfds == 1 {
            return Ok(Some(3));
        } else {
            return Ok(None);
        },
    }

    for (i, v) in fdnames.split(":").enumerate() {
        if v == "varlink" {
            return Ok(Some(3 + i as i32));
        }
    }

    Ok(None)
}

//FIXME: add Drop with shutdown() and unix file removal
impl VarlinkListener {
    pub fn new(address: &str) -> io::Result<Self> {
        if let Some(l) = activation_listener()? {
            if address.starts_with("tcp:") {
                unsafe {
                    return Ok(VarlinkListener::TCP(TcpListener::from_raw_fd(l)));
                }
            } else if address.starts_with("unix:") {
                unsafe {
                    return Ok(VarlinkListener::UNIX(UnixListener::from_raw_fd(l)));
                }
            } else {
                return Err(Error::new(ErrorKind::Other, "unknown varlink address"));
            }
        }

        if address.starts_with("tcp:") {
            Ok(VarlinkListener::TCP(TcpListener::bind(&address[4..])?))
        } else if address.starts_with("unix:") {
            let mut addr = String::from(address[5..].split(";").next().unwrap());
            if addr.starts_with("@") {
                addr = addr.replacen("@", "\0", 1);
                let l = AbstractUnixListener::bind(addr)?;
                unsafe {
                    return Ok(VarlinkListener::UNIX(UnixListener::from_raw_fd(
                        l.into_raw_fd(),
                    )));
                }
            }
            // ignore error on non-existant file
            let _ = fs::remove_file(&*addr);
            Ok(VarlinkListener::UNIX(UnixListener::bind(addr)?))
        } else {
            Err(Error::new(ErrorKind::Other, "unknown varlink address"))
        }
    }

    pub fn accept(&self) -> io::Result<VarlinkStream> {
        match self {
            &VarlinkListener::TCP(ref l) => {
                let (mut s, _addr) = l.accept()?;
                Ok(VarlinkStream::TCP(s))
            }
            &VarlinkListener::UNIX(ref l) => {
                let (mut s, _addr) = l.accept()?;
                Ok(VarlinkStream::UNIX(s))
            }
        }
    }
    pub fn set_nonblocking(&self, b: bool) -> io::Result<()> {
        match self {
            &VarlinkListener::TCP(ref l) => l.set_nonblocking(b),
            &VarlinkListener::UNIX(ref l) => l.set_nonblocking(b),
        }
    }
}

enum Message {
    NewJob(Job),
    Terminate,
}

struct ThreadPool {
    workers: Vec<Worker>,
    sender: mpsc::Sender<Message>,
}

trait FnBox {
    fn call_box(self: Box<Self>);
}

impl<F: FnOnce()> FnBox for F {
    fn call_box(self: Box<F>) {
        (*self)()
    }
}

type Job = Box<FnBox + Send + 'static>;

impl ThreadPool {
    /// Create a new ThreadPool.
    ///
    /// The size is the number of threads in the pool.
    ///
    /// # Panics
    ///
    /// The `new` function will panic if the size is zero.
    pub fn new(size: usize) -> ThreadPool {
        assert!(size > 0);

        let (sender, receiver) = mpsc::channel();

        let receiver = Arc::new(Mutex::new(receiver));

        let mut workers = Vec::with_capacity(size);

        for _ in 0..size {
            workers.push(Worker::new(Arc::clone(&receiver)));
        }

        ThreadPool { workers, sender }
    }

    pub fn execute<F>(&self, f: F)
    where
        F: FnOnce() + Send + 'static,
    {
        let job = Box::new(f);

        self.sender.send(Message::NewJob(job)).unwrap();
    }
}

impl Drop for ThreadPool {
    fn drop(&mut self) {
        for _ in &mut self.workers {
            self.sender.send(Message::Terminate).unwrap();
        }

        for worker in &mut self.workers {
            if let Some(thread) = worker.thread.take() {
                thread.join().unwrap();
            }
        }
    }
}

struct Worker {
    thread: Option<thread::JoinHandle<()>>,
}

impl Worker {
    fn new(receiver: Arc<Mutex<mpsc::Receiver<Message>>>) -> Worker {
        let thread = thread::spawn(move || loop {
            let message = receiver.lock().unwrap().recv().unwrap();

            match message {
                Message::NewJob(job) => {
                    job.call_box();
                }
                Message::Terminate => {
                    break;
                }
            }
        });

        Worker {
            thread: Some(thread),
        }
    }
}

pub enum ServerError {
    IoError(Error),
    AcceptTimeout,
}

impl From<Error> for ServerError {
    fn from(e: Error) -> Self {
        ServerError::IoError(e)
    }
}

pub fn do_listen(
    service: ::VarlinkService,
    addr: &str,
    workers: usize,
    accept_timeout: u64,
) -> Result<(), ServerError> {
    let service = Arc::new(service);
    let listener = Arc::new(VarlinkListener::new(addr)?);
    listener.set_nonblocking(false)?;
    let pool = ThreadPool::new(workers);

    loop {
        let mut stream: VarlinkStream;
        if accept_timeout != 0 {
            let listener = listener.clone();
            let (sender, receiver) = mpsc::channel();

            let _t = thread::spawn(move || {
                match sender.send(listener.accept()) {
                    Ok(()) => {} // everything good
                    Err(_) => {} // we have been released, don't panic
                }
            });

            stream = match receiver.recv_timeout(Duration::from_secs(accept_timeout)) {
                Ok(s) => s?,
                Err(_) => return Err(ServerError::AcceptTimeout),
            };
        } else {
            stream = listener.accept()?;
        }

        let service = service.clone();

        pool.execute(move || {
            let (mut r, mut w) = stream.split().expect("Could not split stream");

            if let Err(_e) = service.handle(&mut r, &mut w) {
                //println!("Handle Error: {}", e);
                let _ = stream.shutdown();
            }
        });
    }
}
<commit_msg>fixup LISTEN_FDNAMES logic<commit_after>//! Handle network connections for a varlink service

//#![feature(getpid)]
//use std::process;
// FIXME
use libc::getpid;
use std::env;
use std::fs;
use std::io;
use std::io::{Error, ErrorKind, Read, Write};
use std::net::{Shutdown, TcpListener, TcpStream};
use std::os::unix::io::{FromRawFd, IntoRawFd};
use std::os::unix::net::{UnixListener, UnixStream};
use std::sync::{Arc, Mutex};
use std::sync::mpsc;
use std::thread;
use std::time::Duration;
// FIXME: abstract unix domains sockets still not in std
// FIXME: https://github.com/rust-lang/rust/issues/14194
use unix_socket::UnixListener as AbstractUnixListener;

enum VarlinkListener {
    TCP(TcpListener),
    UNIX(UnixListener),
}

enum VarlinkStream {
    TCP(TcpStream),
    UNIX(UnixStream),
}

impl<'a> VarlinkStream {
    pub fn split(&mut self) -> io::Result<(Box<Read + Send + Sync>, Box<Write + Send + Sync>)> {
        match *self {
            VarlinkStream::TCP(ref mut s) => {
                Ok((Box::new(s.try_clone()?), Box::new(s.try_clone()?)))
            }
            VarlinkStream::UNIX(ref mut s) => {
                Ok((Box::new(s.try_clone()?), Box::new(s.try_clone()?)))
            }
        }
    }
    pub fn shutdown(&mut self) -> io::Result<()> {
        match *self {
            VarlinkStream::TCP(ref mut s) => s.shutdown(Shutdown::Both),
            VarlinkStream::UNIX(ref mut s) => s.shutdown(Shutdown::Both),
        }
    }
}

fn activation_listener() -> io::Result<Option<i32>> {
    let nfds: u32;

    match env::var("LISTEN_FDS") {
        Ok(ref n) => match n.parse::<u32>() {
            Ok(n) if n >= 1 => nfds = n,
            _ => return Ok(None),
        },
        _ => return Ok(None),
    }

    unsafe {
        match env::var("LISTEN_PID") {
            //FIXME: replace Ok(getpid()) with Ok(process::id())
            Ok(ref pid) if pid.parse::<i32>() == Ok(getpid()) => {}
            _ => return Ok(None),
        }
    }

    if nfds == 1 {
        return Ok(Some(3));
    }

    let fdnames: String;

    match env::var("LISTEN_FDNAMES") {
        Ok(n) => {
            fdnames = n;
        }
        _ => return Ok(None),
    }

    for (i, v) in fdnames.split(":").enumerate() {
        if v == "varlink" {
            return Ok(Some(3 + i as i32));
        }
    }

    Ok(None)
}

//FIXME: add Drop with shutdown() and unix file removal
impl VarlinkListener {
    pub fn new(address: &str) -> io::Result<Self> {
        if let Some(l) = activation_listener()? {
            if address.starts_with("tcp:") {
                unsafe {
                    return Ok(VarlinkListener::TCP(TcpListener::from_raw_fd(l)));
                }
            } else if address.starts_with("unix:") {
                unsafe {
                    return Ok(VarlinkListener::UNIX(UnixListener::from_raw_fd(l)));
                }
            } else {
                return Err(Error::new(ErrorKind::Other, "unknown varlink address"));
            }
        }

        if address.starts_with("tcp:") {
            Ok(VarlinkListener::TCP(TcpListener::bind(&address[4..])?))
        } else if address.starts_with("unix:") {
            let mut addr = String::from(address[5..].split(";").next().unwrap());
            if addr.starts_with("@") {
                addr = addr.replacen("@", "\0", 1);
                let l = AbstractUnixListener::bind(addr)?;
                unsafe {
                    return Ok(VarlinkListener::UNIX(UnixListener::from_raw_fd(
                        l.into_raw_fd(),
                    )));
                }
            }
            // ignore error on non-existant file
            let _ = fs::remove_file(&*addr);
            Ok(VarlinkListener::UNIX(UnixListener::bind(addr)?))
        } else {
            Err(Error::new(ErrorKind::Other, "unknown varlink address"))
        }
    }

    pub fn accept(&self) -> io::Result<VarlinkStream> {
        match self {
            &VarlinkListener::TCP(ref l) => {
                let (mut s, _addr) = l.accept()?;
                Ok(VarlinkStream::TCP(s))
            }
            &VarlinkListener::UNIX(ref l) => {
                let (mut s, _addr) = l.accept()?;
                Ok(VarlinkStream::UNIX(s))
            }
        }
    }
    pub fn set_nonblocking(&self, b: bool) -> io::Result<()> {
        match self {
            &VarlinkListener::TCP(ref l) => l.set_nonblocking(b),
            &VarlinkListener::UNIX(ref l) => l.set_nonblocking(b),
        }
    }
}

enum Message {
    NewJob(Job),
    Terminate,
}

struct ThreadPool {
    workers: Vec<Worker>,
    sender: mpsc::Sender<Message>,
}

trait FnBox {
    fn call_box(self: Box<Self>);
}

impl<F: FnOnce()> FnBox for F {
    fn call_box(self: Box<F>) {
        (*self)()
    }
}

type Job = Box<FnBox + Send + 'static>;

impl ThreadPool {
    /// Create a new ThreadPool.
    ///
    /// The size is the number of threads in the pool.
    ///
    /// # Panics
    ///
    /// The `new` function will panic if the size is zero.
    pub fn new(size: usize) -> ThreadPool {
        assert!(size > 0);

        let (sender, receiver) = mpsc::channel();

        let receiver = Arc::new(Mutex::new(receiver));

        let mut workers = Vec::with_capacity(size);

        for _ in 0..size {
            workers.push(Worker::new(Arc::clone(&receiver)));
        }

        ThreadPool { workers, sender }
    }

    pub fn execute<F>(&self, f: F)
    where
        F: FnOnce() + Send + 'static,
    {
        let job = Box::new(f);

        self.sender.send(Message::NewJob(job)).unwrap();
    }
}

impl Drop for ThreadPool {
    fn drop(&mut self) {
        for _ in &mut self.workers {
            self.sender.send(Message::Terminate).unwrap();
        }

        for worker in &mut self.workers {
            if let Some(thread) = worker.thread.take() {
                thread.join().unwrap();
            }
        }
    }
}

struct Worker {
    thread: Option<thread::JoinHandle<()>>,
}

impl Worker {
    fn new(receiver: Arc<Mutex<mpsc::Receiver<Message>>>) -> Worker {
        let thread = thread::spawn(move || loop {
            let message = receiver.lock().unwrap().recv().unwrap();

            match message {
                Message::NewJob(job) => {
                    job.call_box();
                }
                Message::Terminate => {
                    break;
                }
            }
        });

        Worker {
            thread: Some(thread),
        }
    }
}

pub enum ServerError {
    IoError(Error),
    AcceptTimeout,
}

impl From<Error> for ServerError {
    fn from(e: Error) -> Self {
        ServerError::IoError(e)
    }
}

pub fn do_listen(
    service: ::VarlinkService,
    addr: &str,
    workers: usize,
    accept_timeout: u64,
) -> Result<(), ServerError> {
    let service = Arc::new(service);
    let listener = Arc::new(VarlinkListener::new(addr)?);
    listener.set_nonblocking(false)?;
    let pool = ThreadPool::new(workers);

    loop {
        let mut stream: VarlinkStream;
        if accept_timeout != 0 {
            let listener = listener.clone();
            let (sender, receiver) = mpsc::channel();

            let _t = thread::spawn(move || {
                match sender.send(listener.accept()) {
                    Ok(()) => {} // everything good
                    Err(_) => {} // we have been released, don't panic
                }
            });

            stream = match receiver.recv_timeout(Duration::from_secs(accept_timeout)) {
                Ok(s) => s?,
                Err(_) => return Err(ServerError::AcceptTimeout),
            };
        } else {
            stream = listener.accept()?;
        }

        let service = service.clone();

        pool.execute(move || {
            let (mut r, mut w) = stream.split().expect("Could not split stream");

            if let Err(_e) = service.handle(&mut r, &mut w) {
                //println!("Handle Error: {}", e);
                let _ = stream.shutdown();
            }
        });
    }
}
<|endoftext|>


<commit_before>/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/. */

//! A collection of invalidations due to changes in which stylesheets affect a
//! document.

#![deny(unsafe_code)]

use Atom;
use dom::{TElement, TNode};
use fnv::FnvHashSet;
use invalidation::element::restyle_hints::RestyleHint;
use media_queries::Device;
use selector_parser::SelectorImpl;
use selectors::attr::CaseSensitivity;
use selectors::parser::{Component, Selector};
use shared_lock::SharedRwLockReadGuard;
use stylesheets::{CssRule, StylesheetInDocument};

/// An invalidation scope represents a kind of subtree that may need to be
/// restyled.
#[derive(Debug, Eq, Hash, PartialEq)]
#[cfg_attr(feature = "servo", derive(HeapSizeOf))]
enum InvalidationScope {
    /// All the descendants of an element with a given id.
    ID(Atom),
    /// All the descendants of an element with a given class name.
    Class(Atom),
}

impl InvalidationScope {
    fn is_id(&self) -> bool {
        matches!(*self, InvalidationScope::ID(..))
    }

    fn matches<E>(&self, element: E) -> bool
        where E: TElement,
    {
        match *self {
            InvalidationScope::Class(ref class) => {
                element.has_class(class, CaseSensitivity::CaseSensitive)
            }
            InvalidationScope::ID(ref id) => {
                match element.get_id() {
                    Some(element_id) => element_id == *id,
                    None => false,
                }
            }
        }
    }
}

/// A set of invalidations due to stylesheet additions.
///
/// TODO(emilio): We might be able to do the same analysis for media query
/// changes too (or even selector changes?).
#[cfg_attr(feature = "servo", derive(HeapSizeOf))]
pub struct StylesheetInvalidationSet {
    /// The style scopes we know we have to restyle so far.
    invalid_scopes: FnvHashSet<InvalidationScope>,
    /// Whether the whole document should be invalid.
    fully_invalid: bool,
}

impl StylesheetInvalidationSet {
    /// Create an empty `StylesheetInvalidationSet`.
    pub fn new() -> Self {
        Self {
            invalid_scopes: FnvHashSet::default(),
            fully_invalid: false,
        }
    }

    /// Mark the DOM tree styles' as fully invalid.
    pub fn invalidate_fully(&mut self) {
        debug!("StylesheetInvalidationSet::invalidate_fully");
        self.invalid_scopes.clear();
        self.fully_invalid = true;
    }

    /// Analyze the given stylesheet, and collect invalidations from their
    /// rules, in order to avoid doing a full restyle when we style the document
    /// next time.
    pub fn collect_invalidations_for<S>(
        &mut self,
        device: &Device,
        stylesheet: &S,
        guard: &SharedRwLockReadGuard
    )
    where
        S: StylesheetInDocument,
    {
        debug!("StylesheetInvalidationSet::collect_invalidations_for");
        if self.fully_invalid {
            debug!(" > Fully invalid already");
            return;
        }

        if !stylesheet.enabled() ||
           !stylesheet.is_effective_for_device(device, guard) {
            debug!(" > Stylesheet was not effective");
            return; // Nothing to do here.
        }

        for rule in stylesheet.effective_rules(device, guard) {
            self.collect_invalidations_for_rule(rule, guard);
            if self.fully_invalid {
                self.invalid_scopes.clear();
                break;
            }
        }

        debug!(" > resulting invalidations: {:?}", self.invalid_scopes);
        debug!(" > fully_invalid: {}", self.fully_invalid);
    }

    /// Clears the invalidation set, invalidating elements as needed if
    /// `document_element` is provided.
    ///
    /// Returns true if any invalidations ocurred.
    pub fn flush<E>(&mut self, document_element: Option<E>) -> bool
        where E: TElement,
    {
        let have_invalidations = match document_element {
            Some(e) => self.process_invalidations(e),
            None => false,
        };
        self.clear();
        have_invalidations
    }

    /// Clears the invalidation set without processing.
    pub fn clear(&mut self) {
        self.invalid_scopes.clear();
        self.fully_invalid = false;
    }

    fn process_invalidations<E>(&self, element: E) -> bool
        where E: TElement,
    {
        {
            let mut data = match element.mutate_data() {
                Some(data) => data,
                None => return false,
            };

            if self.fully_invalid {
                debug!("process_invalidations: fully_invalid({:?})",
                       element);
                data.hint.insert(RestyleHint::restyle_subtree());
                return true;
            }
        }

        if self.invalid_scopes.is_empty() {
            debug!("process_invalidations: empty invalidation set");
            return false;
        }

        self.process_invalidations_in_subtree(element)
    }

    /// Process style invalidations in a given subtree, that is, look for all
    /// the relevant scopes in the subtree, and mark as dirty only the relevant
    /// ones.
    ///
    /// Returns whether it invalidated at least one element's style.
    #[allow(unsafe_code)]
    fn process_invalidations_in_subtree<E>(&self, element: E) -> bool
        where E: TElement,
    {
        let mut data = match element.mutate_data() {
            Some(data) => data,
            None => return false,
        };

        if !data.has_styles() {
            return false;
        }

        if data.hint.contains_subtree() {
            debug!("process_invalidations_in_subtree: {:?} was already invalid",
                   element);
            return false;
        }

        for scope in &self.invalid_scopes {
            if scope.matches(element) {
                debug!("process_invalidations_in_subtree: {:?} matched {:?}",
                       element, scope);
                data.hint.insert(RestyleHint::restyle_subtree());
                return true;
            }
        }

        let mut any_children_invalid = false;

        for child in element.as_node().traversal_children() {
            let child = match child.as_element() {
                Some(e) => e,
                None => continue,
            };

            any_children_invalid |= self.process_invalidations_in_subtree(child);
        }

        if any_children_invalid {
            debug!("Children of {:?} changed, setting dirty descendants",
                   element);
            unsafe { element.set_dirty_descendants() }
        }

        return any_children_invalid
    }

    fn scan_component(
        component: &Component<SelectorImpl>,
        scope: &mut Option<InvalidationScope>)
    {
        match *component {
            Component::Class(ref class) => {
                if scope.as_ref().map_or(true, |s| !s.is_id()) {
                    *scope = Some(InvalidationScope::Class(class.clone()));
                }
            }
            Component::ID(ref id) => {
                if scope.is_none() {
                    *scope = Some(InvalidationScope::ID(id.clone()));
                }
            }
            _ => {
                // Ignore everything else, at least for now.
            }
        }
    }

    /// Collect a style scopes for a given selector.
    ///
    /// We look at the outermost class or id selector to the left of an ancestor
    /// combinator, in order to restyle only a given subtree.
    ///
    /// We prefer id scopes to class scopes, and outermost scopes to innermost
    /// scopes (to reduce the amount of traversal we need to do).
    fn collect_scopes(&mut self, selector: &Selector<SelectorImpl>) {
        debug!("StylesheetInvalidationSet::collect_scopes({:?})", selector);

        let mut scope: Option<InvalidationScope> = None;

        let mut scan = true;
        let mut iter = selector.iter();

        loop {
            for component in &mut iter {
                if scan {
                    Self::scan_component(component, &mut scope);
                }
            }
            match iter.next_sequence() {
                None => break,
                Some(combinator) => {
                    scan = combinator.is_ancestor();
                }
            }
        }

        match scope {
            Some(s) => {
                debug!(" > Found scope: {:?}", s);
                self.invalid_scopes.insert(s);
            }
            None => {
                debug!(" > Scope not found");

                // If we didn't find a scope, any element could match this, so
                // let's just bail out.
                self.fully_invalid = true;
            }
        }
    }

    /// Collects invalidations for a given CSS rule.
    fn collect_invalidations_for_rule(
        &mut self,
        rule: &CssRule,
        guard: &SharedRwLockReadGuard)
    {
        use stylesheets::CssRule::*;
        debug!("StylesheetInvalidationSet::collect_invalidations_for_rule");
        debug_assert!(!self.fully_invalid, "Not worth to be here!");

        match *rule {
            Style(ref lock) => {
                let style_rule = lock.read_with(guard);
                for selector in &style_rule.selectors.0 {
                    self.collect_scopes(selector);
                    if self.fully_invalid {
                        return;
                    }
                }
            }
            Document(..) |
            Namespace(..) |
            Import(..) |
            Media(..) |
            Supports(..) => {
                // Do nothing, relevant nested rules are visited as part of the
                // iteration.
            }
            FontFace(..) |
            CounterStyle(..) |
            Keyframes(..) |
            Page(..) |
            Viewport(..) |
            FontFeatureValues(..) => {
                debug!(" > Found unsupported rule, marking the whole subtree \
                       invalid.");

                // TODO(emilio): Can we do better here?
                //
                // At least in `@page`, we could check the relevant media, I
                // guess.
                self.fully_invalid = true;
            }
        }
    }
}
<commit_msg>style: Support more selectors in the style sheet invalidator.<commit_after>/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/. */

//! A collection of invalidations due to changes in which stylesheets affect a
//! document.

#![deny(unsafe_code)]

use Atom;
use LocalName as SelectorLocalName;
use dom::{TElement, TNode};
use fnv::FnvHashSet;
use invalidation::element::restyle_hints::{RESTYLE_SELF, RestyleHint};
use media_queries::Device;
use selector_parser::SelectorImpl;
use selectors::attr::CaseSensitivity;
use selectors::parser::{Component, LocalName, Selector};
use shared_lock::SharedRwLockReadGuard;
use stylesheets::{CssRule, StylesheetInDocument};

/// A style sheet invalidation represents a kind of element or subtree that may
/// need to be restyled. Whether it represents a whole subtree or just a single
/// element is determined by whether the invalidation is stored in the
/// StylesheetInvalidationSet's invalid_scopes or invalid_elements table.
#[derive(Debug, Eq, Hash, PartialEq)]
#[cfg_attr(feature = "servo", derive(HeapSizeOf))]
enum Invalidation {
    /// An element with a given id.
    ID(Atom),
    /// An element with a given class name.
    Class(Atom),
    /// An element with a given local name.
    LocalName { name: SelectorLocalName, lower_name: SelectorLocalName },
}

impl Invalidation {
    fn is_id(&self) -> bool {
        matches!(*self, Invalidation::ID(..))
    }

    fn is_id_or_class(&self) -> bool {
        matches!(*self, Invalidation::ID(..) | Invalidation::Class(..))
    }

    fn matches<E>(&self, element: E) -> bool
        where E: TElement,
    {
        match *self {
            Invalidation::Class(ref class) => {
                // FIXME This should look at the quirks mode of the document to
                // determine case sensitivity.
                element.has_class(class, CaseSensitivity::CaseSensitive)
            }
            Invalidation::ID(ref id) => {
                match element.get_id() {
                    // FIXME This should look at the quirks mode of the document
                    // to determine case sensitivity.
                    Some(element_id) => element_id == *id,
                    None => false,
                }
            }
            Invalidation::LocalName { ref name, ref lower_name } => {
                // This could look at the quirks mode of the document, instead
                // of testing against both names, but it's probably not worth
                // it.
                let local_name = element.get_local_name();
                *local_name == **name || *local_name == **lower_name
            }
        }
    }
}

/// A set of invalidations due to stylesheet additions.
///
/// TODO(emilio): We might be able to do the same analysis for media query
/// changes too (or even selector changes?).
#[cfg_attr(feature = "servo", derive(HeapSizeOf))]
pub struct StylesheetInvalidationSet {
    /// The subtrees we know we have to restyle so far.
    invalid_scopes: FnvHashSet<Invalidation>,
    /// The elements we know we have to restyle so far.
    invalid_elements: FnvHashSet<Invalidation>,
    /// Whether the whole document should be restyled.
    fully_invalid: bool,
}

impl StylesheetInvalidationSet {
    /// Create an empty `StylesheetInvalidationSet`.
    pub fn new() -> Self {
        Self {
            invalid_scopes: FnvHashSet::default(),
            invalid_elements: FnvHashSet::default(),
            fully_invalid: false,
        }
    }

    /// Mark the DOM tree styles' as fully invalid.
    pub fn invalidate_fully(&mut self) {
        debug!("StylesheetInvalidationSet::invalidate_fully");
        self.invalid_scopes.clear();
        self.invalid_elements.clear();
        self.fully_invalid = true;
    }

    /// Analyze the given stylesheet, and collect invalidations from their
    /// rules, in order to avoid doing a full restyle when we style the document
    /// next time.
    pub fn collect_invalidations_for<S>(
        &mut self,
        device: &Device,
        stylesheet: &S,
        guard: &SharedRwLockReadGuard
    )
    where
        S: StylesheetInDocument,
    {
        debug!("StylesheetInvalidationSet::collect_invalidations_for");
        if self.fully_invalid {
            debug!(" > Fully invalid already");
            return;
        }

        if !stylesheet.enabled() ||
           !stylesheet.is_effective_for_device(device, guard) {
            debug!(" > Stylesheet was not effective");
            return; // Nothing to do here.
        }

        for rule in stylesheet.effective_rules(device, guard) {
            self.collect_invalidations_for_rule(rule, guard);
            if self.fully_invalid {
                self.invalid_scopes.clear();
                self.invalid_elements.clear();
                break;
            }
        }

        debug!(" > resulting subtree invalidations: {:?}", self.invalid_scopes);
        debug!(" > resulting self invalidations: {:?}", self.invalid_elements);
        debug!(" > fully_invalid: {}", self.fully_invalid);
    }

    /// Clears the invalidation set, invalidating elements as needed if
    /// `document_element` is provided.
    ///
    /// Returns true if any invalidations ocurred.
    pub fn flush<E>(&mut self, document_element: Option<E>) -> bool
        where E: TElement,
    {
        let have_invalidations = match document_element {
            Some(e) => self.process_invalidations(e),
            None => false,
        };
        self.clear();
        have_invalidations
    }

    /// Clears the invalidation set without processing.
    pub fn clear(&mut self) {
        self.invalid_scopes.clear();
        self.invalid_elements.clear();
        self.fully_invalid = false;
    }

    fn process_invalidations<E>(&self, element: E) -> bool
        where E: TElement,
    {
        {
            let mut data = match element.mutate_data() {
                Some(data) => data,
                None => return false,
            };

            if self.fully_invalid {
                debug!("process_invalidations: fully_invalid({:?})",
                       element);
                data.hint.insert(RestyleHint::restyle_subtree());
                return true;
            }
        }

        if self.invalid_scopes.is_empty() && self.invalid_elements.is_empty() {
            debug!("process_invalidations: empty invalidation set");
            return false;
        }

        self.process_invalidations_in_subtree(element)
    }

    /// Process style invalidations in a given subtree. This traverses the
    /// subtree looking for elements that match the invalidations in
    /// invalid_scopes and invalid_elements.
    ///
    /// Returns whether it invalidated at least one element's style.
    #[allow(unsafe_code)]
    fn process_invalidations_in_subtree<E>(&self, element: E) -> bool
        where E: TElement,
    {
        let mut data = match element.mutate_data() {
            Some(data) => data,
            None => return false,
        };

        if !data.has_styles() {
            return false;
        }

        if data.hint.contains_subtree() {
            debug!("process_invalidations_in_subtree: {:?} was already invalid",
                   element);
            return false;
        }

        for invalidation in &self.invalid_scopes {
            if invalidation.matches(element) {
                debug!("process_invalidations_in_subtree: {:?} matched subtree {:?}",
                       element, invalidation);
                data.hint.insert(RestyleHint::restyle_subtree());
                return true;
            }
        }

        let mut self_invalid = false;

        if !data.hint.contains(RESTYLE_SELF) {
            for invalidation in &self.invalid_elements {
                if invalidation.matches(element) {
                    debug!("process_invalidations_in_subtree: {:?} matched self {:?}",
                           element, invalidation);
                    data.hint.insert(RESTYLE_SELF);
                    self_invalid = true;
                    break;
                }
            }
        }

        let mut any_children_invalid = false;

        for child in element.as_node().traversal_children() {
            let child = match child.as_element() {
                Some(e) => e,
                None => continue,
            };

            any_children_invalid |= self.process_invalidations_in_subtree(child);
        }

        if any_children_invalid {
            debug!("Children of {:?} changed, setting dirty descendants",
                   element);
            unsafe { element.set_dirty_descendants() }
        }

        return self_invalid || any_children_invalid
    }

    fn scan_component(
        component: &Component<SelectorImpl>,
        invalidation: &mut Option<Invalidation>)
    {
        match *component {
            Component::LocalName(LocalName { ref name, ref lower_name }) => {
                if invalidation.as_ref().map_or(true, |s| !s.is_id_or_class()) {
                    *invalidation = Some(Invalidation::LocalName {
                        name: name.clone(),
                        lower_name: lower_name.clone(),
                    });
                }
            }
            Component::Class(ref class) => {
                if invalidation.as_ref().map_or(true, |s| !s.is_id()) {
                    *invalidation = Some(Invalidation::Class(class.clone()));
                }
            }
            Component::ID(ref id) => {
                if invalidation.is_none() {
                    *invalidation = Some(Invalidation::ID(id.clone()));
                }
            }
            _ => {
                // Ignore everything else, at least for now.
            }
        }
    }

    /// Collect invalidations for a given selector.
    ///
    /// We look at the outermost local name, class, or ID selector to the left
    /// of an ancestor combinator, in order to restyle only a given subtree.
    ///
    /// If the selector has no ancestor combinator, then we do the same for
    /// the only sequence it has, but record it as an element invalidation
    /// instead of a subtree invalidation.
    ///
    /// We prefer IDs to classs, and classes to local names, on the basis
    /// that the former should be more specific than the latter. We also
    /// prefer to generate subtree invalidations for the outermost part
    /// of the selector, to reduce the amount of traversal we need to do
    /// when flushing invalidations.
    fn collect_invalidations(&mut self, selector: &Selector<SelectorImpl>) {
        debug!("StylesheetInvalidationSet::collect_invalidations({:?})", selector);

        let mut element_invalidation: Option<Invalidation> = None;
        let mut subtree_invalidation: Option<Invalidation> = None;

        let mut scan_for_element_invalidation = true;
        let mut scan_for_subtree_invalidation = false;

        let mut iter = selector.iter();

        loop {
            for component in &mut iter {
                if scan_for_element_invalidation {
                    Self::scan_component(component, &mut element_invalidation);
                } else if scan_for_subtree_invalidation {
                    Self::scan_component(component, &mut subtree_invalidation);
                }
            }
            match iter.next_sequence() {
                None => break,
                Some(combinator) => {
                    scan_for_subtree_invalidation = combinator.is_ancestor();
                }
            }
            scan_for_element_invalidation = false;
        }

        if let Some(s) = subtree_invalidation {
            debug!(" > Found subtree invalidation: {:?}", s);
            self.invalid_scopes.insert(s);
        } else if let Some(s) = element_invalidation {
            debug!(" > Found element invalidation: {:?}", s);
            self.invalid_elements.insert(s);
        } else {
            // The selector was of a form that we can't handle. Any element
            // could match it, so let's just bail out.
            debug!(" > Can't handle selector, marking fully invalid");
            self.fully_invalid = true;
        }
    }

    /// Collects invalidations for a given CSS rule.
    fn collect_invalidations_for_rule(
        &mut self,
        rule: &CssRule,
        guard: &SharedRwLockReadGuard)
    {
        use stylesheets::CssRule::*;
        debug!("StylesheetInvalidationSet::collect_invalidations_for_rule");
        debug_assert!(!self.fully_invalid, "Not worth to be here!");

        match *rule {
            Style(ref lock) => {
                let style_rule = lock.read_with(guard);
                for selector in &style_rule.selectors.0 {
                    self.collect_invalidations(selector);
                    if self.fully_invalid {
                        return;
                    }
                }
            }
            Document(..) |
            Namespace(..) |
            Import(..) |
            Media(..) |
            Supports(..) => {
                // Do nothing, relevant nested rules are visited as part of the
                // iteration.
            }
            FontFace(..) |
            CounterStyle(..) |
            Keyframes(..) |
            Page(..) |
            Viewport(..) |
            FontFeatureValues(..) => {
                debug!(" > Found unsupported rule, marking the whole subtree \
                       invalid.");

                // TODO(emilio): Can we do better here?
                //
                // At least in `@page`, we could check the relevant media, I
                // guess.
                self.fully_invalid = true;
            }
        }
    }
}
<|endoftext|>


<commit_before>// Copyright 2012-2015 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

#![feature(plugin_registrar, rustc_private)]

extern crate syntax;
extern crate syntax_ext;
extern crate rustc_plugin;

use syntax::ast::{MetaItem, Expr, BinOpKind};
use syntax::ast;
use syntax::codemap::Span;
use syntax::ext::base::{ExtCtxt, Annotatable};
use syntax::ext::build::AstBuilder;
use syntax_ext::deriving::generic::*;
use syntax_ext::deriving::generic::ty::*;
use syntax::parse::token::InternedString;
use syntax::ptr::P;
use syntax::ext::base::MultiDecorator;
use syntax::parse::token;

use rustc_plugin::Registry;

macro_rules! pathvec {
    ($($x:ident)::+) => (
        vec![ $( stringify!($x) ),+ ]
    )
}

macro_rules! path {
    ($($x:tt)*) => (
        ::syntax_ext::deriving::generic::ty::Path::new( pathvec!( $($x)* ) )
    )
}

macro_rules! path_local {
    ($x:ident) => (
        ::syntax_ext::deriving::generic::ty::Path::new_local(stringify!($x))
    )
}

macro_rules! pathvec_std {
    ($cx:expr, $first:ident :: $($rest:ident)::+) => ({
        let mut v = pathvec!($($rest)::+);
        if let Some(s) = $cx.crate_root {
            v.insert(0, s);
        }
        v
    })
}

pub fn expand_deriving_from_primitive(cx: &mut ExtCtxt,
                                      span: Span,
                                      mitem: &MetaItem,
                                      item: &Annotatable,
                                      push: &mut FnMut(Annotatable))
{
    let inline = cx.meta_word(span, InternedString::new("inline"));
    let attrs = vec!(cx.attribute(span, inline));
    let trait_def = TraitDef {
        is_unsafe: false,
        span: span,
        attributes: Vec::new(),
        path: path!(num::FromPrimitive),
        additional_bounds: Vec::new(),
        generics: LifetimeBounds::empty(),
        methods: vec!(
            MethodDef {
                name: "from_i64",
                is_unsafe: false,
                unify_fieldless_variants: false,
                generics: LifetimeBounds::empty(),
                explicit_self: None,
                args: vec!(Literal(path_local!(i64))),
                ret_ty: Literal(Path::new_(pathvec_std!(cx, core::option::Option),
                                           None,
                                           vec!(Box::new(Self_)),
                                           true)),
                // #[inline] liable to cause code-bloat
                attributes: attrs.clone(),
                combine_substructure: combine_substructure(Box::new(|c, s, sub| {
                    cs_from("i64", c, s, sub)
                })),
            },
            MethodDef {
                name: "from_u64",
                is_unsafe: false,
                unify_fieldless_variants: false,
                generics: LifetimeBounds::empty(),
                explicit_self: None,
                args: vec!(Literal(path_local!(u64))),
                ret_ty: Literal(Path::new_(pathvec_std!(cx, core::option::Option),
                                           None,
                                           vec!(Box::new(Self_)),
                                           true)),
                // #[inline] liable to cause code-bloat
                attributes: attrs,
                combine_substructure: combine_substructure(Box::new(|c, s, sub| {
                    cs_from("u64", c, s, sub)
                })),
            }
        ),
        associated_types: Vec::new(),
    };

    trait_def.expand(cx, mitem, &item, push)
}

fn cs_from(name: &str, cx: &mut ExtCtxt, trait_span: Span, substr: &Substructure) -> P<Expr> {
    if substr.nonself_args.len() != 1 {
        cx.span_bug(trait_span, "incorrect number of arguments in `derive(FromPrimitive)`")
    }

    let n = &substr.nonself_args[0];

    match *substr.fields {
        StaticStruct(..) => {
            cx.span_err(trait_span, "`FromPrimitive` cannot be derived for structs");
            return cx.expr_fail(trait_span, InternedString::new(""));
        }
        StaticEnum(enum_def, _) => {
            if enum_def.variants.is_empty() {
                cx.span_err(trait_span,
                            "`FromPrimitive` cannot be derived for enums with no variants");
                return cx.expr_fail(trait_span, InternedString::new(""));
            }

            let mut arms = Vec::new();

            for variant in &enum_def.variants {
                match variant.node.data {
                    ast::VariantData::Unit(..) => {
                        let span = variant.span;

                        // expr for `$n == $variant as $name`
                        let path = cx.path(span, vec![substr.type_ident, variant.node.name]);
                        let variant = cx.expr_path(path);
                        let ty = cx.ty_ident(span, cx.ident_of(name));
                        let cast = cx.expr_cast(span, variant.clone(), ty);
                        let guard = cx.expr_binary(span, BinOpKind::Eq, n.clone(), cast);

                        // expr for `Some($variant)`
                        let body = cx.expr_some(span, variant);

                        // arm for `_ if $guard => $body`
                        let arm = ast::Arm {
                            attrs: vec!(),
                            pats: vec!(cx.pat_wild(span)),
                            guard: Some(guard),
                            body: body,
                        };

                        arms.push(arm);
                    }
                    ast::VariantData::Tuple(..) => {
                        cx.span_err(trait_span,
                                    "`FromPrimitive` cannot be derived for \
                                    enum variants with arguments");
                        return cx.expr_fail(trait_span,
                                            InternedString::new(""));
                    }
                    ast::VariantData::Struct(..) => {
                        cx.span_err(trait_span,
                                    "`FromPrimitive` cannot be derived for enums \
                                    with struct variants");
                        return cx.expr_fail(trait_span,
                                            InternedString::new(""));
                    }
                }
            }

            // arm for `_ => None`
            let arm = ast::Arm {
                attrs: vec!(),
                pats: vec!(cx.pat_wild(trait_span)),
                guard: None,
                body: cx.expr_none(trait_span),
            };
            arms.push(arm);

            cx.expr_match(trait_span, n.clone(), arms)
        }
        _ => cx.span_bug(trait_span, "expected StaticEnum in derive(FromPrimitive)")
    }
}

#[plugin_registrar]
#[doc(hidden)]
pub fn plugin_registrar(reg: &mut Registry) {
    reg.register_syntax_extension(
        token::intern("derive_NumFromPrimitive"),
        MultiDecorator(Box::new(expand_deriving_from_primitive)));
}
<commit_msg>Auto merge of #229 - bluss:fix-macros, r=hauleth<commit_after>// Copyright 2012-2015 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

#![feature(plugin_registrar, rustc_private)]

extern crate syntax;
extern crate syntax_ext;
extern crate rustc_plugin;

use syntax::ast::{MetaItem, Expr, BinOpKind};
use syntax::ast;
use syntax::codemap::Span;
use syntax::ext::base::{ExtCtxt, Annotatable};
use syntax::ext::build::AstBuilder;
use syntax_ext::deriving::generic::*;
use syntax_ext::deriving::generic::ty::*;
use syntax::parse::token::InternedString;
use syntax::ptr::P;
use syntax::ext::base::MultiDecorator;
use syntax::parse::token;

use rustc_plugin::Registry;

macro_rules! pathvec {
    ($($x:ident)::+) => (
        vec![ $( stringify!($x) ),+ ]
    )
}

macro_rules! path {
    ($($x:tt)*) => (
        ::syntax_ext::deriving::generic::ty::Path::new( pathvec!( $($x)* ) )
    )
}

macro_rules! path_local {
    ($x:ident) => (
        ::syntax_ext::deriving::generic::ty::Path::new_local(stringify!($x))
    )
}

macro_rules! pathvec_std {
    ($cx:expr, $first:ident :: $($rest:ident)::+) => ({
        let mut v = pathvec!($($rest)::+);
        if let Some(s) = $cx.crate_root {
            v.insert(0, s);
        }
        v
    })
}

pub fn expand_deriving_from_primitive(cx: &mut ExtCtxt,
                                      span: Span,
                                      mitem: &MetaItem,
                                      item: &Annotatable,
                                      push: &mut FnMut(Annotatable))
{
    let inline = cx.meta_word(span, InternedString::new("inline"));
    let attrs = vec!(cx.attribute(span, inline));
    let trait_def = TraitDef {
        is_unsafe: false,
        span: span,
        attributes: Vec::new(),
        path: path!(num::FromPrimitive),
        additional_bounds: Vec::new(),
        generics: LifetimeBounds::empty(),
        methods: vec!(
            MethodDef {
                name: "from_i64",
                is_unsafe: false,
                unify_fieldless_variants: false,
                generics: LifetimeBounds::empty(),
                explicit_self: None,
                args: vec!(Literal(path_local!(i64))),
                ret_ty: Literal(Path::new_(pathvec_std!(cx, core::option::Option),
                                           None,
                                           vec!(Box::new(Self_)),
                                           true)),
                // #[inline] liable to cause code-bloat
                attributes: attrs.clone(),
                combine_substructure: combine_substructure(Box::new(|c, s, sub| {
                    cs_from("i64", c, s, sub)
                })),
            },
            MethodDef {
                name: "from_u64",
                is_unsafe: false,
                unify_fieldless_variants: false,
                generics: LifetimeBounds::empty(),
                explicit_self: None,
                args: vec!(Literal(path_local!(u64))),
                ret_ty: Literal(Path::new_(pathvec_std!(cx, core::option::Option),
                                           None,
                                           vec!(Box::new(Self_)),
                                           true)),
                // #[inline] liable to cause code-bloat
                attributes: attrs,
                combine_substructure: combine_substructure(Box::new(|c, s, sub| {
                    cs_from("u64", c, s, sub)
                })),
            }
        ),
        associated_types: Vec::new(),
        supports_unions: false,
    };

    trait_def.expand(cx, mitem, &item, push)
}

fn cs_from(name: &str, cx: &mut ExtCtxt, trait_span: Span, substr: &Substructure) -> P<Expr> {
    if substr.nonself_args.len() != 1 {
        cx.span_bug(trait_span, "incorrect number of arguments in `derive(FromPrimitive)`")
    }

    let n = &substr.nonself_args[0];

    match *substr.fields {
        StaticStruct(..) => {
            cx.span_err(trait_span, "`FromPrimitive` cannot be derived for structs");
            return cx.expr_fail(trait_span, InternedString::new(""));
        }
        StaticEnum(enum_def, _) => {
            if enum_def.variants.is_empty() {
                cx.span_err(trait_span,
                            "`FromPrimitive` cannot be derived for enums with no variants");
                return cx.expr_fail(trait_span, InternedString::new(""));
            }

            let mut arms = Vec::new();

            for variant in &enum_def.variants {
                match variant.node.data {
                    ast::VariantData::Unit(..) => {
                        let span = variant.span;

                        // expr for `$n == $variant as $name`
                        let path = cx.path(span, vec![substr.type_ident, variant.node.name]);
                        let variant = cx.expr_path(path);
                        let ty = cx.ty_ident(span, cx.ident_of(name));
                        let cast = cx.expr_cast(span, variant.clone(), ty);
                        let guard = cx.expr_binary(span, BinOpKind::Eq, n.clone(), cast);

                        // expr for `Some($variant)`
                        let body = cx.expr_some(span, variant);

                        // arm for `_ if $guard => $body`
                        let arm = ast::Arm {
                            attrs: vec!(),
                            pats: vec!(cx.pat_wild(span)),
                            guard: Some(guard),
                            body: body,
                        };

                        arms.push(arm);
                    }
                    ast::VariantData::Tuple(..) => {
                        cx.span_err(trait_span,
                                    "`FromPrimitive` cannot be derived for \
                                    enum variants with arguments");
                        return cx.expr_fail(trait_span,
                                            InternedString::new(""));
                    }
                    ast::VariantData::Struct(..) => {
                        cx.span_err(trait_span,
                                    "`FromPrimitive` cannot be derived for enums \
                                    with struct variants");
                        return cx.expr_fail(trait_span,
                                            InternedString::new(""));
                    }
                }
            }

            // arm for `_ => None`
            let arm = ast::Arm {
                attrs: vec!(),
                pats: vec!(cx.pat_wild(trait_span)),
                guard: None,
                body: cx.expr_none(trait_span),
            };
            arms.push(arm);

            cx.expr_match(trait_span, n.clone(), arms)
        }
        _ => cx.span_bug(trait_span, "expected StaticEnum in derive(FromPrimitive)")
    }
}

#[plugin_registrar]
#[doc(hidden)]
pub fn plugin_registrar(reg: &mut Registry) {
    reg.register_syntax_extension(
        token::intern("derive_NumFromPrimitive"),
        MultiDecorator(Box::new(expand_deriving_from_primitive)));
}
<|endoftext|>