<commit_before>// Copyright 2014 The Gfx-rs Developers.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

extern crate cgmath;
#[macro_use]
extern crate gfx;
extern crate gfx_window_glfw;
extern crate glfw;
extern crate time;
extern crate gfx_gl as gl;

use time::precise_time_s;
use cgmath::FixedArray;
use cgmath::{Matrix, Point3, Vector3, Matrix3, Matrix4};
use cgmath::{Transform, AffineMatrix3, Vector4, Array1};
use gfx::attrib::Floater;
use gfx::traits::{Device, Stream, FactoryExt};
use glfw::Context;
use gl::Gl;
use gl::types::*;
use std::mem;
use std::ptr;
use std::str;
use std::env;
use std::str::FromStr;
use std::sync::mpsc::Receiver;
use std::iter::repeat;
use std::ffi::CString;

gfx_vertex!( Vertex {
    a_Pos@ pos: [Floater<i8>; 3],
});

gfx_parameters!( Params {
    u_Transform@ transform: [[f32; 4]; 4],
});

static VERTEX_SRC: &'static [u8] = b"
    #version 150 core
    in vec3 a_Pos;
    uniform mat4 u_Transform;

    void main() {
        gl_Position = u_Transform * vec4(a_Pos, 1.0);
    }
";

static FRAGMENT_SRC: &'static [u8] = b"
    #version 150 core
    out vec4 o_Color;

    void main() {
        o_Color = vec4(1., 0., 0., 1.);
    }
";

//----------------------------------------

fn gfx_main(mut glfw: glfw::Glfw,
            window: glfw::Window,
            events: Receiver<(f64, glfw::WindowEvent)>,
            dimension: i16) {
    let (mut stream, mut device, mut factory) = gfx_window_glfw::init(window);
    let state = gfx::DrawState::new().depth(gfx::state::Comparison::LessEqual, true);

    let vertex_data = [
        Vertex { pos: Floater::cast3([-1,  1, -1]) },
        Vertex { pos: Floater::cast3([ 1,  1, -1]) },
        Vertex { pos: Floater::cast3([ 1,  1,  1]) },
    ];

    let mesh = factory.create_mesh(&vertex_data);
    let slice = mesh.get_slice();

    let program = factory.link_program(VERTEX_SRC, FRAGMENT_SRC).unwrap();
    let view: AffineMatrix3<f32> = Transform::look_at(
        &Point3::new(0f32, -5.0, 0.0),
        &Point3::new(0f32, 0.0, 0.0),
        &Vector3::unit_z(),
    );
    let aspect = stream.get_aspect_ratio();
    let proj = cgmath::perspective(cgmath::deg(45.0f32), aspect, 1.0, 10.0);

    let batch = gfx::batch::Core::new(mesh, program).unwrap();;

    while !stream.out.window.should_close() {
        glfw.poll_events();
        for (_, event) in glfw::flush_messages(&events) {
            match event {
                glfw::WindowEvent::Key(glfw::Key::Escape, _, glfw::Action::Press, _) =>
                    stream.out.window.set_should_close(true),
                _ => {},
            }
        }

        let start = precise_time_s() * 1000.;
        stream.clear(gfx::ClearData {
            color: [0.3, 0.3, 0.3, 1.0],
            depth: 1.0,
            stencil: 0,
        });

        for x in (-dimension) ..dimension {
            for y in (-dimension) ..dimension {
                let mut model = Matrix4::from(Matrix3::identity().mul_s(0.01f32));
                model.w = Vector4::new(x as f32 * 0.05,
                                       0f32,
                                       y as f32 * 0.05,
                                       1f32);

                let data = Params {
                    transform: proj.mul_m(&view.mat)
                                   .mul_m(&model).into_fixed(),
                    _r: std::marker::PhantomData,
                };
                stream.draw(&batch.with(&slice, &data, &state))
                      .unwrap();
            }
        }

        let pre_submit = precise_time_s() * 1000.;
        stream.flush(&mut device);
        let post_submit = precise_time_s() * 1000.;
        stream.out.window.swap_buffers();
        device.cleanup();
        let swap = precise_time_s() * 1000.;

        println!("total time:\t\t{0:4.2}ms", swap - start);
        println!("\tcreate list:\t{0:4.2}ms", pre_submit - start);
        println!("\tsubmit:\t\t{0:4.2}ms", post_submit - pre_submit);
        println!("\tgpu wait:\t{0:4.2}ms", swap - post_submit)
    }
}

static VS_SRC: &'static str = "
    #version 150 core
    in vec3 a_Pos;
    uniform mat4 u_Transform;

    void main() {
        gl_Position = u_Transform * vec4(a_Pos, 1.0);
    }
";

static FS_SRC: &'static str = "
    #version 150 core
    out vec4 o_Color;

    void main() {
        o_Color = vec4(1., 0., 0., 1.);
    }
";

fn compile_shader(gl: &Gl, src: &str, ty: GLenum) -> GLuint { unsafe {
    let shader = gl.CreateShader(ty);
    // Attempt to compile the shader
    let src = CString::new(src).unwrap();
    gl.ShaderSource(shader, 1, &(src.as_bytes_with_nul().as_ptr() as *const i8), ptr::null());
    gl.CompileShader(shader);

    // Get the compile status
    let mut status = gl::FALSE as GLint;
    gl.GetShaderiv(shader, gl::COMPILE_STATUS, &mut status);

    // Fail on error
    if status != (gl::TRUE as GLint) {
        let mut len: GLint = 0;
        gl.GetShaderiv(shader, gl::INFO_LOG_LENGTH, &mut len);
        let mut buf: Vec<u8> = repeat(0u8).take((len as isize).saturating_sub(1) as usize).collect();     // subtract 1 to skip the trailing null character
        gl.GetShaderInfoLog(shader, len, ptr::null_mut(), buf.as_mut_ptr() as *mut GLchar);
        panic!("{}", str::from_utf8(&buf).ok().expect("ShaderInfoLog not valid utf8"));
    }
    shader
}}

fn link_program(gl: &Gl, vs: GLuint, fs: GLuint) -> GLuint { unsafe {
    let program = gl.CreateProgram();
    gl.AttachShader(program, vs);
    gl.AttachShader(program, fs);
    gl.LinkProgram(program);
    // Get the link status
    let mut status = gl::FALSE as GLint;
    gl.GetProgramiv(program, gl::LINK_STATUS, &mut status);

    // Fail on error
    if status != (gl::TRUE as GLint) {
        let mut len: GLint = 0;
        gl.GetProgramiv(program, gl::INFO_LOG_LENGTH, &mut len);
        let mut buf: Vec<u8> = repeat(0u8).take((len as isize).saturating_sub(1) as usize).collect();     // subtract 1 to skip the trailing null character
        gl.GetProgramInfoLog(program, len, ptr::null_mut(), buf.as_mut_ptr() as *mut GLchar);
        panic!("{}", str::from_utf8(&buf).ok().expect("ProgramInfoLog not valid utf8"));
    }
    program
}}

fn gl_main(mut glfw: glfw::Glfw,
           mut window: glfw::Window,
           _: Receiver<(f64, glfw::WindowEvent),>,
           dimension: i16) {
    let gl = Gl::load_with(|s| window.get_proc_address(s) as *const _);

    // Create GLSL shaders
    let vs = compile_shader(&gl, VS_SRC, gl::VERTEX_SHADER);
    let fs = compile_shader(&gl, FS_SRC, gl::FRAGMENT_SHADER);
    let program = link_program(&gl, vs, fs);

    let mut vao = 0;
    let mut vbo = 0;

    let trans_uniform = unsafe {
        // Create Vertex Array Object
        gl.GenVertexArrays(1, &mut vao);
        gl.BindVertexArray(vao);

        // Create a Vertex Buffer Object and copy the vertex data to it
        gl.GenBuffers(1, &mut vbo);
        gl.BindBuffer(gl::ARRAY_BUFFER, vbo);

        let vertex_data = vec![
            Vertex { pos: Floater::cast3([-1,  1, -1]) },
            Vertex { pos: Floater::cast3([ 1,  1, -1]) },
            Vertex { pos: Floater::cast3([ 1,  1,  1]) },
        ];

        gl.BufferData(gl::ARRAY_BUFFER,
                      (vertex_data.len() * mem::size_of::<Vertex>()) as GLsizeiptr,
                      mem::transmute(&vertex_data[0]),
                      gl::STATIC_DRAW);

        // Use shader program
        gl.UseProgram(program);
        let o_color = CString::new("o_Color").unwrap();
        gl.BindFragDataLocation(program, 0, o_color.as_bytes_with_nul().as_ptr() as *const i8);

        // Specify the layout of the vertex data
        let a_pos = CString::new("a_Pos").unwrap();
        gl.BindFragDataLocation(program, 0, a_pos.as_bytes_with_nul().as_ptr() as *const i8);

        let pos_attr = gl.GetAttribLocation(program, a_pos.as_ptr());
        gl.EnableVertexAttribArray(pos_attr as GLuint);
        gl.VertexAttribPointer(pos_attr as GLuint, 3, gl::BYTE,
                                gl::FALSE as GLboolean, 0, ptr::null());

        let u_transform = CString::new("u_Transform").unwrap();
        gl.GetUniformLocation(program, u_transform.as_bytes_with_nul().as_ptr() as *const i8)
    };

    let (w, h) = window.get_framebuffer_size();
    let view: AffineMatrix3<f32> = Transform::look_at(
        &Point3::new(0f32, -5.0, 0.0),
        &Point3::new(0f32, 0.0, 0.0),
        &Vector3::unit_z(),
    );
    let aspect = w as f32 / h as f32;
    let proj = cgmath::perspective(cgmath::deg(45.0f32), aspect, 1.0, 10.0);

    while !window.should_close() {
        // Poll events
        glfw.poll_events();

        let start = precise_time_s() * 1000.;

        // Clear the screen to black
        unsafe {
            gl.ClearColor(0.3, 0.3, 0.3, 1.0);
            gl.Clear(gl::COLOR_BUFFER_BIT);
        }

        for x in (-dimension) ..dimension {
            for y in (-dimension) ..dimension {
                let mut model = Matrix4::from(Matrix3::identity().mul_s(0.01f32));
                model.w = Vector4::new(x as f32 * 0.05,
                                       0f32,
                                       y as f32 * 0.05,
                                       1f32);

                let mat = proj.mul_m(&view.mat).mul_m(&model);

                unsafe {
                    gl.UniformMatrix4fv(trans_uniform,
                                        1,
                                        gl::FALSE,
                                        mat.x.ptr());
                    gl.DrawArrays(gl::TRIANGLES, 0, 3);
                }

            }
        }

        let submit = precise_time_s() * 1000.;

        // Swap buffers
        window.swap_buffers();
        let swap = precise_time_s() * 1000.;

        println!("total time:\t\t{0:4.2}ms", swap - start);
        println!("\tsubmit:\t\t{0:4.2}ms", submit - start);
        println!("\tgpu wait:\t{0:4.2}ms", swap - submit)

    }

    // Cleanup
    unsafe {
        gl.DeleteProgram(program);
        gl.DeleteShader(fs);
        gl.DeleteShader(vs);
        gl.DeleteBuffers(1, &vbo);
        gl.DeleteVertexArrays(1, &vao);
    }
}

pub fn main() {
    let ref mut args = env::args();
    let args_count = env::args().count();
    if args_count == 1 {
        println!("gfx-perf [gl|gfx] <size>");
        return;
    }

    let mode = args.nth(1).unwrap();
    let count: i32 = if args_count == 3 {
        FromStr::from_str(&args.next().unwrap()).ok()
    } else {
        None
    }.unwrap_or(10000);

    let count = ((count as f64).sqrt() / 2.) as i16;

    let mut glfw = glfw::init(glfw::FAIL_ON_ERRORS)
        .ok().expect("Failed to initialize glfw-rs");

    glfw.window_hint(glfw::WindowHint::ContextVersion(3, 2));
    glfw.window_hint(glfw::WindowHint::OpenGlForwardCompat(true));
    glfw.window_hint(glfw::WindowHint::OpenGlProfile(glfw::OpenGlProfileHint::Core));

    let (mut window, events) = glfw
        .create_window(640, 480, "Cube example", glfw::WindowMode::Windowed)
        .expect("Failed to create GLFW window.");

    window.make_current();
    glfw.set_error_callback(glfw::FAIL_ON_ERRORS);
    window.set_key_polling(true);

    println!("count is {}", count*count*4);
    match mode.as_ref() {
        "gfx" => gfx_main(glfw, window, events, count),
        "gl" => gl_main(glfw, window, events, count),
        x => {
            println!("{} is not a known mode", x)
        }
    }
}
<commit_msg>PSO - ported the performance example and cleaned it up a bit<commit_after>// Copyright 2014 The Gfx-rs Developers.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

extern crate cgmath;
#[macro_use]
extern crate gfx;
extern crate gfx_window_glfw;
extern crate glfw;
extern crate time;
extern crate gfx_gl as gl;

use time::precise_time_s;
use cgmath::FixedArray;
use cgmath::{Matrix, Point3, Vector3, Matrix3, Matrix4};
use cgmath::{Transform, AffineMatrix3, Vector4, Array1};
use gfx::format::{I8Scaled, DepthStencil, Rgba8};
use glfw::Context;
use gl::Gl;
use gl::types::*;
use std::mem;
use std::ptr;
use std::str;
use std::env;
use std::str::FromStr;
use std::sync::mpsc::Receiver;
use std::iter::repeat;
use std::ffi::CString;

gfx_vertex_struct!( Vertex {
    pos: [I8Scaled; 3] = "a_Pos",
});

gfx_pipeline_init!( PipeData PipeMeta PipeInit {
    vbuf: gfx::VertexBuffer<Vertex> = gfx::PER_VERTEX,
    transform: gfx::Global<[[f32; 4]; 4]> = "u_Transform",
    out_color: gfx::RenderTarget<Rgba8> = ("o_Color", gfx::state::MASK_ALL),
    out_depth: gfx::DepthTarget<DepthStencil> = gfx::state::Depth {
        fun: gfx::state::Comparison::LessEqual,
        write: true,
    },
});

static VERTEX_SRC: &'static [u8] = b"
    #version 150 core
    in vec3 a_Pos;
    uniform mat4 u_Transform;

    void main() {
        gl_Position = u_Transform * vec4(a_Pos, 1.0);
    }
";

static FRAGMENT_SRC: &'static [u8] = b"
    #version 150 core
    out vec4 o_Color;

    void main() {
        o_Color = vec4(1., 0., 0., 1.);
    }
";

//----------------------------------------

fn gfx_main(mut glfw: glfw::Glfw,
            mut window: glfw::Window,
            events: Receiver<(f64, glfw::WindowEvent)>,
            dimension: i16) {
    use gfx::traits::{Device, FactoryExt};

    let (mut device, mut factory, main_color, main_depth) =
        gfx_window_glfw::init_new(&mut window);
    let mut encoder = factory.create_encoder();

    let pso = factory.create_pipeline_simple(
        VERTEX_SRC, FRAGMENT_SRC,
        gfx::state::CullFace::Back,
        PipeInit::new()
        ).unwrap();

    let vertex_data = [
        Vertex { pos: I8Scaled::cast3([-1,  1, -1]) },
        Vertex { pos: I8Scaled::cast3([ 1,  1, -1]) },
        Vertex { pos: I8Scaled::cast3([ 1,  1,  1]) },
    ];
    let (vbuf, slice) = factory.create_vertex_buffer(&vertex_data);

    let view: AffineMatrix3<f32> = Transform::look_at(
        &Point3::new(0f32, -5.0, 0.0),
        &Point3::new(0f32, 0.0, 0.0),
        &Vector3::unit_z(),
    );
    let aspect = {
        let (w, h) = window.get_framebuffer_size();
        w as f32 / h as f32
    };
    let proj = cgmath::perspective(cgmath::deg(45.0f32), aspect, 1.0, 10.0);

    let mut data = PipeData {
        vbuf: vbuf,
        transform: cgmath::Matrix4::identity().into_fixed(),
        out_color: main_color,
        out_depth: main_depth,
    };

    while !window.should_close() {
        glfw.poll_events();
        for (_, event) in glfw::flush_messages(&events) {
            match event {
                glfw::WindowEvent::Key(glfw::Key::Escape, _, glfw::Action::Press, _) =>
                    window.set_should_close(true),
                _ => {},
            }
        }

        let start = precise_time_s() * 1000.;
        encoder.reset();
        encoder.clear_target(&data.out_color, [0.3, 0.3, 0.3, 1.0]);
        encoder.clear_depth(&data.out_depth, 1.0);

        for x in (-dimension) ..dimension {
            for y in (-dimension) ..dimension {
                let mut model = Matrix4::from(Matrix3::identity().mul_s(0.01f32));
                model.w = Vector4::new(x as f32 * 0.05,
                                       0f32,
                                       y as f32 * 0.05,
                                       1f32);
                data.transform = proj.mul_m(&view.mat)
                                     .mul_m(&model).into_fixed();
                encoder.draw_pipeline(&slice, &pso, &data);
            }
        }

        let pre_submit = precise_time_s() * 1000.;
        device.submit(encoder.as_buffer());
        let post_submit = precise_time_s() * 1000.;
        window.swap_buffers();
        device.cleanup();
        let swap = precise_time_s() * 1000.;

        println!("total time:\t\t{0:4.2}ms", swap - start);
        println!("\tcreate list:\t{0:4.2}ms", pre_submit - start);
        println!("\tsubmit:\t\t{0:4.2}ms", post_submit - pre_submit);
        println!("\tgpu wait:\t{0:4.2}ms", swap - post_submit)
    }
}

fn compile_shader(gl: &Gl, src: &[u8], ty: GLenum) -> GLuint { unsafe {
    let shader = gl.CreateShader(ty);
    // Attempt to compile the shader
    gl.ShaderSource(shader, 1,
        &(src.as_ptr() as *const i8),
        &(src.len() as GLint));
    gl.CompileShader(shader);

    // Get the compile status
    let mut status = gl::FALSE as GLint;
    gl.GetShaderiv(shader, gl::COMPILE_STATUS, &mut status);

    // Fail on error
    if status != (gl::TRUE as GLint) {
        let mut len: GLint = 0;
        gl.GetShaderiv(shader, gl::INFO_LOG_LENGTH, &mut len);
        let mut buf: Vec<u8> = repeat(0u8).take((len as isize).saturating_sub(1) as usize).collect();     // subtract 1 to skip the trailing null character
        gl.GetShaderInfoLog(shader, len, ptr::null_mut(), buf.as_mut_ptr() as *mut GLchar);
        panic!("{}", str::from_utf8(&buf).ok().expect("ShaderInfoLog not valid utf8"));
    }
    shader
}}

fn link_program(gl: &Gl, vs: GLuint, fs: GLuint) -> GLuint { unsafe {
    let program = gl.CreateProgram();
    gl.AttachShader(program, vs);
    gl.AttachShader(program, fs);
    gl.LinkProgram(program);
    // Get the link status
    let mut status = gl::FALSE as GLint;
    gl.GetProgramiv(program, gl::LINK_STATUS, &mut status);

    // Fail on error
    if status != (gl::TRUE as GLint) {
        let mut len: GLint = 0;
        gl.GetProgramiv(program, gl::INFO_LOG_LENGTH, &mut len);
        let mut buf: Vec<u8> = repeat(0u8).take((len as isize).saturating_sub(1) as usize).collect();     // subtract 1 to skip the trailing null character
        gl.GetProgramInfoLog(program, len, ptr::null_mut(), buf.as_mut_ptr() as *mut GLchar);
        panic!("{}", str::from_utf8(&buf).ok().expect("ProgramInfoLog not valid utf8"));
    }
    program
}}

fn gl_main(mut glfw: glfw::Glfw,
           mut window: glfw::Window,
           _: Receiver<(f64, glfw::WindowEvent),>,
           dimension: i16) {
    let gl = Gl::load_with(|s| window.get_proc_address(s) as *const _);

    // Create GLSL shaders
    let vs = compile_shader(&gl, VERTEX_SRC, gl::VERTEX_SHADER);
    let fs = compile_shader(&gl, FRAGMENT_SRC, gl::FRAGMENT_SHADER);
    let program = link_program(&gl, vs, fs);

    let mut vao = 0;
    let mut vbo = 0;

    let trans_uniform = unsafe {
        // Create Vertex Array Object
        gl.GenVertexArrays(1, &mut vao);
        gl.BindVertexArray(vao);

        // Create a Vertex Buffer Object and copy the vertex data to it
        gl.GenBuffers(1, &mut vbo);
        gl.BindBuffer(gl::ARRAY_BUFFER, vbo);

        let vertex_data = vec![
            Vertex { pos: I8Scaled::cast3([-1,  1, -1]) },
            Vertex { pos: I8Scaled::cast3([ 1,  1, -1]) },
            Vertex { pos: I8Scaled::cast3([ 1,  1,  1]) },
        ];

        gl.BufferData(gl::ARRAY_BUFFER,
                      (vertex_data.len() * mem::size_of::<Vertex>()) as GLsizeiptr,
                      mem::transmute(&vertex_data[0]),
                      gl::STATIC_DRAW);

        // Use shader program
        gl.UseProgram(program);
        let o_color = CString::new("o_Color").unwrap();
        gl.BindFragDataLocation(program, 0, o_color.as_bytes_with_nul().as_ptr() as *const i8);

        // Specify the layout of the vertex data
        let a_pos = CString::new("a_Pos").unwrap();
        gl.BindFragDataLocation(program, 0, a_pos.as_bytes_with_nul().as_ptr() as *const i8);

        let pos_attr = gl.GetAttribLocation(program, a_pos.as_ptr());
        gl.EnableVertexAttribArray(pos_attr as GLuint);
        gl.VertexAttribPointer(pos_attr as GLuint, 3, gl::BYTE,
                                gl::FALSE as GLboolean, 0, ptr::null());

        let u_transform = CString::new("u_Transform").unwrap();
        gl.GetUniformLocation(program, u_transform.as_bytes_with_nul().as_ptr() as *const i8)
    };

    let (w, h) = window.get_framebuffer_size();
    let view: AffineMatrix3<f32> = Transform::look_at(
        &Point3::new(0f32, -5.0, 0.0),
        &Point3::new(0f32, 0.0, 0.0),
        &Vector3::unit_z(),
    );
    let aspect = w as f32 / h as f32;
    let proj = cgmath::perspective(cgmath::deg(45.0f32), aspect, 1.0, 10.0);

    while !window.should_close() {
        // Poll events
        glfw.poll_events();

        let start = precise_time_s() * 1000.;

        // Clear the screen to black
        unsafe {
            gl.ClearColor(0.3, 0.3, 0.3, 1.0);
            gl.Clear(gl::COLOR_BUFFER_BIT);
        }

        for x in (-dimension) ..dimension {
            for y in (-dimension) ..dimension {
                let mut model = Matrix4::from(Matrix3::identity().mul_s(0.01f32));
                model.w = Vector4::new(x as f32 * 0.05,
                                       0f32,
                                       y as f32 * 0.05,
                                       1f32);

                let mat = proj.mul_m(&view.mat).mul_m(&model);

                unsafe {
                    gl.UniformMatrix4fv(trans_uniform,
                                        1,
                                        gl::FALSE,
                                        mat.x.ptr());
                    gl.DrawArrays(gl::TRIANGLES, 0, 3);
                }

            }
        }

        let submit = precise_time_s() * 1000.;

        // Swap buffers
        window.swap_buffers();
        let swap = precise_time_s() * 1000.;

        println!("total time:\t\t{0:4.2}ms", swap - start);
        println!("\tsubmit:\t\t{0:4.2}ms", submit - start);
        println!("\tgpu wait:\t{0:4.2}ms", swap - submit)

    }

    // Cleanup
    unsafe {
        gl.DeleteProgram(program);
        gl.DeleteShader(fs);
        gl.DeleteShader(vs);
        gl.DeleteBuffers(1, &vbo);
        gl.DeleteVertexArrays(1, &vao);
    }
}

pub fn main() {
    let ref mut args = env::args();
    let args_count = env::args().count();
    if args_count == 1 {
        println!("gfx-perf [gl|gfx] <size>");
        return;
    }

    let mode = args.nth(1).unwrap();
    let count: i32 = if args_count == 3 {
        FromStr::from_str(&args.next().unwrap()).ok()
    } else {
        None
    }.unwrap_or(10000);

    let count = ((count as f64).sqrt() / 2.) as i16;

    let mut glfw = glfw::init(glfw::FAIL_ON_ERRORS)
        .ok().expect("Failed to initialize glfw-rs");

    glfw.window_hint(glfw::WindowHint::ContextVersion(3, 2));
    glfw.window_hint(glfw::WindowHint::OpenGlForwardCompat(true));
    glfw.window_hint(glfw::WindowHint::OpenGlProfile(glfw::OpenGlProfileHint::Core));

    let (mut window, events) = glfw
        .create_window(640, 480, "Cube example", glfw::WindowMode::Windowed)
        .expect("Failed to create GLFW window.");

    window.make_current();
    glfw.set_error_callback(glfw::FAIL_ON_ERRORS);
    window.set_key_polling(true);

    println!("count is {}", count*count*4);
    match mode.as_ref() {
        "gfx" => gfx_main(glfw, window, events, count),
        "gl" => gl_main(glfw, window, events, count),
        x => {
            println!("{} is not a known mode", x)
        }
    }
}
<|endoftext|>


<commit_before>/// This function is where the kernel sets up IRQ handlers
/// It is increcibly unsafe, and should be minimal in nature
/// It must create the IDT with the correct entries, those entries are
/// defined in other files inside of the `arch` module

use core::{mem, slice};
use core::sync::atomic::{AtomicBool, AtomicUsize, Ordering};

use crate::allocator;
#[cfg(feature = "acpi")]
use crate::acpi;
#[cfg(feature = "graphical_debug")]
use arch::x86_64::graphical_debug;
use crate::arch::x86_64::pti;
use crate::device;
use crate::gdt;
use crate::idt;
use crate::interrupt;
use crate::log;
use crate::memory;
use crate::paging;

/// Test of zero values in BSS.
static BSS_TEST_ZERO: usize = 0;
/// Test of non-zero values in data.
static DATA_TEST_NONZERO: usize = 0xFFFF_FFFF_FFFF_FFFF;
/// Test of zero values in thread BSS
#[thread_local]
static mut TBSS_TEST_ZERO: usize = 0;
/// Test of non-zero values in thread data.
#[thread_local]
static mut TDATA_TEST_NONZERO: usize = 0xFFFF_FFFF_FFFF_FFFF;

pub static KERNEL_BASE: AtomicUsize = AtomicUsize::new(0);
pub static KERNEL_SIZE: AtomicUsize = AtomicUsize::new(0);
pub static CPU_COUNT: AtomicUsize = AtomicUsize::new(0);
pub static AP_READY: AtomicBool = AtomicBool::new(false);
static BSP_READY: AtomicBool = AtomicBool::new(false);

#[repr(packed)]
pub struct KernelArgs {
    kernel_base: u64,
    kernel_size: u64,
    stack_base: u64,
    stack_size: u64,
    env_base: u64,
    env_size: u64,

    /// The base 64-bit pointer to an array of saved RSDPs. It's up to the kernel (and possibly
    /// userspace), to decide which RSDP to use. The buffer will be a linked list containing a
    /// 32-bit relative (to this field) next, and the actual struct afterwards.
    ///
    /// This field can be NULL, and if so, the system has not booted with UEFI or in some other way
    /// retrieved the RSDPs. The kernel or a userspace driver will thus try searching the BIOS
    /// memory instead. On UEFI systems, searching is not guaranteed to actually work though.
    acpi_rsdps_base: u64,
    /// The size of the RSDPs region.
    acpi_rsdps_size: u64,
}

/// The entry to Rust, all things must be initialized
#[no_mangle]
pub unsafe extern fn kstart(args_ptr: *const KernelArgs) -> ! {
    let env = {
        let args = &*args_ptr;

        let kernel_base = args.kernel_base as usize;
        let kernel_size = args.kernel_size as usize;
        let stack_base = args.stack_base as usize;
        let stack_size = args.stack_size as usize;
        let env_base = args.env_base as usize;
        let env_size = args.env_size as usize;

        // BSS should already be zero
        {
            assert_eq!(BSS_TEST_ZERO, 0);
            assert_eq!(DATA_TEST_NONZERO, 0xFFFF_FFFF_FFFF_FFFF);
        }

        KERNEL_BASE.store(kernel_base, Ordering::SeqCst);
        KERNEL_SIZE.store(kernel_size, Ordering::SeqCst);

        println!("Kernel: {:X}:{:X}", kernel_base, kernel_base + kernel_size);
        println!("Stack: {:X}:{:X}", stack_base, stack_base + stack_size);
        println!("Env: {:X}:{:X}", env_base, env_base + env_size);
        println!("RSDPs: {:X}:{:X}", args.acpi_rsdps_base, args.acpi_rsdps_base + args.acpi_rsdps_size);

        let ext_mem_ranges = if args.acpi_rsdps_base != 0 && args.acpi_rsdps_size > 0 {
            Some([(args.acpi_rsdps_base as usize, args.acpi_rsdps_size as usize)])
        } else {
            None
        };

        // Set up GDT before paging
        gdt::init();

        // Set up IDT before paging
        idt::init();

        // Initialize memory management
        memory::init(0, kernel_base + ((kernel_size + 4095)/4096) * 4096);

        // Initialize paging
        let (mut active_table, tcb_offset) = paging::init(0, kernel_base, kernel_base + kernel_size, stack_base, stack_base + stack_size, ext_mem_ranges.as_ref().map(|arr| &arr[..]).unwrap_or(&[]));

        // Set up GDT after paging with TLS
        gdt::init_paging(tcb_offset, stack_base + stack_size);

        // Set up IDT
        idt::init_paging();

        // Set up syscall instruction
        interrupt::syscall::init();

        // Test tdata and tbss
        {
            assert_eq!(TBSS_TEST_ZERO, 0);
            TBSS_TEST_ZERO += 1;
            assert_eq!(TBSS_TEST_ZERO, 1);
            assert_eq!(TDATA_TEST_NONZERO, 0xFFFF_FFFF_FFFF_FFFF);
            TDATA_TEST_NONZERO -= 1;
            assert_eq!(TDATA_TEST_NONZERO, 0xFFFF_FFFF_FFFF_FFFE);
        }

        // Reset AP variables
        CPU_COUNT.store(1, Ordering::SeqCst);
        AP_READY.store(false, Ordering::SeqCst);
        BSP_READY.store(false, Ordering::SeqCst);

        // Setup kernel heap
        allocator::init(&mut active_table);

        // Activate memory logging
        log::init();

        // Use graphical debug
        #[cfg(feature="graphical_debug")]
        graphical_debug::init(&mut active_table);

        // Initialize devices
        device::init(&mut active_table);

        // Read ACPI tables, starts APs
        #[cfg(feature = "acpi")]
        {
            acpi::init(&mut active_table, if args.acpi_rsdps_base != 0 && args.acpi_rsdps_size > 0 { Some((args.acpi_rsdps_base, args.acpi_rsdps_size)) } else { None });
            device::init_after_acpi(&mut active_table);
        }

        // Initialize all of the non-core devices not otherwise needed to complete initialization
        device::init_noncore();

        // Initialize memory functions after core has loaded
        memory::init_noncore();

        // Stop graphical debug
        #[cfg(feature="graphical_debug")]
        graphical_debug::fini(&mut active_table);

        BSP_READY.store(true, Ordering::SeqCst);

        slice::from_raw_parts(env_base as *const u8, env_size)
    };

    crate::kmain(CPU_COUNT.load(Ordering::SeqCst), env);
}

#[repr(packed)]
pub struct KernelArgsAp {
    cpu_id: u64,
    page_table: u64,
    stack_start: u64,
    stack_end: u64,
}

/// Entry to rust for an AP
pub unsafe extern fn kstart_ap(args_ptr: *const KernelArgsAp) -> ! {
    let cpu_id = {
        let args = &*args_ptr;
        let cpu_id = args.cpu_id as usize;
        let bsp_table = args.page_table as usize;
        let stack_start = args.stack_start as usize;
        let stack_end = args.stack_end as usize;

        assert_eq!(BSS_TEST_ZERO, 0);
        assert_eq!(DATA_TEST_NONZERO, 0xFFFF_FFFF_FFFF_FFFF);

        // Set up GDT before paging
        gdt::init();

        // Set up IDT before paging
        idt::init();

        // Initialize paging
        let tcb_offset = paging::init_ap(cpu_id, bsp_table, stack_start, stack_end);

        // Set up GDT with TLS
        gdt::init_paging(tcb_offset, stack_end);

        // Set up IDT for AP
        idt::init_paging();

        // Set up syscall instruction
        interrupt::syscall::init();

        // Test tdata and tbss
        {
            assert_eq!(TBSS_TEST_ZERO, 0);
            TBSS_TEST_ZERO += 1;
            assert_eq!(TBSS_TEST_ZERO, 1);
            assert_eq!(TDATA_TEST_NONZERO, 0xFFFF_FFFF_FFFF_FFFF);
            TDATA_TEST_NONZERO -= 1;
            assert_eq!(TDATA_TEST_NONZERO, 0xFFFF_FFFF_FFFF_FFFE);
        }

        // Initialize devices (for AP)
        device::init_ap();

        AP_READY.store(true, Ordering::SeqCst);

        cpu_id
    };

    while ! BSP_READY.load(Ordering::SeqCst) {
        interrupt::pause();
    }

    crate::kmain_ap(cpu_id);
}

#[naked]
pub unsafe fn usermode(ip: usize, sp: usize, arg: usize) -> ! {
    asm!("push r10
          push r11
          push r12
          push r13
          push r14
          push r15"
          : // No output
          :   "{r10}"(gdt::GDT_USER_DATA << 3 | 3), // Data segment
              "{r11}"(sp), // Stack pointer
              "{r12}"(1 << 9), // Flags - Set interrupt enable flag
              "{r13}"(gdt::GDT_USER_CODE << 3 | 3), // Code segment
              "{r14}"(ip), // IP
              "{r15}"(arg) // Argument
          : // No clobbers
          : "intel", "volatile");

    // Unmap kernel
    pti::unmap();

    // Go to usermode
    asm!("mov ds, r14d
         mov es, r14d
         mov fs, r15d
         mov gs, r14d
         xor rax, rax
         xor rbx, rbx
         xor rcx, rcx
         xor rdx, rdx
         xor rsi, rsi
         xor rdi, rdi
         xor rbp, rbp
         xor r8, r8
         xor r9, r9
         xor r10, r10
         xor r11, r11
         xor r12, r12
         xor r13, r13
         xor r14, r14
         xor r15, r15
         fninit
         pop rdi
         iretq"
         : // No output because it never returns
         :   "{r14}"(gdt::GDT_USER_DATA << 3 | 3), // Data segment
             "{r15}"(gdt::GDT_USER_TLS << 3 | 3) // TLS segment
         : // No clobbers because it never returns
         : "intel", "volatile");
    unreachable!();
}
<commit_msg>Fix a page fault.<commit_after>/// This function is where the kernel sets up IRQ handlers
/// It is increcibly unsafe, and should be minimal in nature
/// It must create the IDT with the correct entries, those entries are
/// defined in other files inside of the `arch` module

use core::{mem, slice};
use core::sync::atomic::{AtomicBool, AtomicUsize, Ordering};

use crate::allocator;
#[cfg(feature = "acpi")]
use crate::acpi;
#[cfg(feature = "graphical_debug")]
use arch::x86_64::graphical_debug;
use crate::arch::x86_64::pti;
use crate::device;
use crate::gdt;
use crate::idt;
use crate::interrupt;
use crate::log;
use crate::memory;
use crate::paging;

/// Test of zero values in BSS.
static BSS_TEST_ZERO: usize = 0;
/// Test of non-zero values in data.
static DATA_TEST_NONZERO: usize = 0xFFFF_FFFF_FFFF_FFFF;
/// Test of zero values in thread BSS
#[thread_local]
static mut TBSS_TEST_ZERO: usize = 0;
/// Test of non-zero values in thread data.
#[thread_local]
static mut TDATA_TEST_NONZERO: usize = 0xFFFF_FFFF_FFFF_FFFF;

pub static KERNEL_BASE: AtomicUsize = AtomicUsize::new(0);
pub static KERNEL_SIZE: AtomicUsize = AtomicUsize::new(0);
pub static CPU_COUNT: AtomicUsize = AtomicUsize::new(0);
pub static AP_READY: AtomicBool = AtomicBool::new(false);
static BSP_READY: AtomicBool = AtomicBool::new(false);

#[repr(packed)]
pub struct KernelArgs {
    kernel_base: u64,
    kernel_size: u64,
    stack_base: u64,
    stack_size: u64,
    env_base: u64,
    env_size: u64,

    /// The base 64-bit pointer to an array of saved RSDPs. It's up to the kernel (and possibly
    /// userspace), to decide which RSDP to use. The buffer will be a linked list containing a
    /// 32-bit relative (to this field) next, and the actual struct afterwards.
    ///
    /// This field can be NULL, and if so, the system has not booted with UEFI or in some other way
    /// retrieved the RSDPs. The kernel or a userspace driver will thus try searching the BIOS
    /// memory instead. On UEFI systems, searching is not guaranteed to actually work though.
    acpi_rsdps_base: u64,
    /// The size of the RSDPs region.
    acpi_rsdps_size: u64,
}

/// The entry to Rust, all things must be initialized
#[no_mangle]
pub unsafe extern fn kstart(args_ptr: *const KernelArgs) -> ! {
    let env = {
        let args = &*args_ptr;

        let kernel_base = args.kernel_base as usize;
        let kernel_size = args.kernel_size as usize;
        let stack_base = args.stack_base as usize;
        let stack_size = args.stack_size as usize;
        let env_base = args.env_base as usize;
        let env_size = args.env_size as usize;
        let acpi_rsdps_base = args.acpi_rsdps_base;
        let acpi_rsdps_size = args.acpi_rsdps_size;

        // BSS should already be zero
        {
            assert_eq!(BSS_TEST_ZERO, 0);
            assert_eq!(DATA_TEST_NONZERO, 0xFFFF_FFFF_FFFF_FFFF);
        }

        KERNEL_BASE.store(kernel_base, Ordering::SeqCst);
        KERNEL_SIZE.store(kernel_size, Ordering::SeqCst);

        println!("Kernel: {:X}:{:X}", kernel_base, kernel_base + kernel_size);
        println!("Stack: {:X}:{:X}", stack_base, stack_base + stack_size);
        println!("Env: {:X}:{:X}", env_base, env_base + env_size);
        println!("RSDPs: {:X}:{:X}", acpi_rsdps_base, acpi_rsdps_base + acpi_rsdps_size);

        let ext_mem_ranges = if args.acpi_rsdps_base != 0 && args.acpi_rsdps_size > 0 {
            Some([(acpi_rsdps_base as usize, acpi_rsdps_size as usize)])
        } else {
            None
        };

        // Set up GDT before paging
        gdt::init();

        // Set up IDT before paging
        idt::init();

        // Initialize memory management
        memory::init(0, kernel_base + ((kernel_size + 4095)/4096) * 4096);

        // Initialize paging
        let (mut active_table, tcb_offset) = paging::init(0, kernel_base, kernel_base + kernel_size, stack_base, stack_base + stack_size, ext_mem_ranges.as_ref().map(|arr| &arr[..]).unwrap_or(&[]));

        // Set up GDT after paging with TLS
        gdt::init_paging(tcb_offset, stack_base + stack_size);

        // Set up IDT
        idt::init_paging();

        // Set up syscall instruction
        interrupt::syscall::init();

        // Test tdata and tbss
        {
            assert_eq!(TBSS_TEST_ZERO, 0);
            TBSS_TEST_ZERO += 1;
            assert_eq!(TBSS_TEST_ZERO, 1);
            assert_eq!(TDATA_TEST_NONZERO, 0xFFFF_FFFF_FFFF_FFFF);
            TDATA_TEST_NONZERO -= 1;
            assert_eq!(TDATA_TEST_NONZERO, 0xFFFF_FFFF_FFFF_FFFE);
        }

        // Reset AP variables
        CPU_COUNT.store(1, Ordering::SeqCst);
        AP_READY.store(false, Ordering::SeqCst);
        BSP_READY.store(false, Ordering::SeqCst);

        // Setup kernel heap
        allocator::init(&mut active_table);

        // Activate memory logging
        log::init();

        // Use graphical debug
        #[cfg(feature="graphical_debug")]
        graphical_debug::init(&mut active_table);

        // Initialize devices
        device::init(&mut active_table);

        // Read ACPI tables, starts APs
        #[cfg(feature = "acpi")]
        {
            acpi::init(&mut active_table, if acpi_rsdps_base != 0 && acpi_rsdps_size > 0 { Some((acpi_rsdps_base, acpi_rsdps_size)) } else { None });
            device::init_after_acpi(&mut active_table);
        }

        // Initialize all of the non-core devices not otherwise needed to complete initialization
        device::init_noncore();

        // Initialize memory functions after core has loaded
        memory::init_noncore();

        // Stop graphical debug
        #[cfg(feature="graphical_debug")]
        graphical_debug::fini(&mut active_table);

        BSP_READY.store(true, Ordering::SeqCst);

        slice::from_raw_parts(env_base as *const u8, env_size)
    };

    crate::kmain(CPU_COUNT.load(Ordering::SeqCst), env);
}

#[repr(packed)]
pub struct KernelArgsAp {
    cpu_id: u64,
    page_table: u64,
    stack_start: u64,
    stack_end: u64,
}

/// Entry to rust for an AP
pub unsafe extern fn kstart_ap(args_ptr: *const KernelArgsAp) -> ! {
    let cpu_id = {
        let args = &*args_ptr;
        let cpu_id = args.cpu_id as usize;
        let bsp_table = args.page_table as usize;
        let stack_start = args.stack_start as usize;
        let stack_end = args.stack_end as usize;

        assert_eq!(BSS_TEST_ZERO, 0);
        assert_eq!(DATA_TEST_NONZERO, 0xFFFF_FFFF_FFFF_FFFF);

        // Set up GDT before paging
        gdt::init();

        // Set up IDT before paging
        idt::init();

        // Initialize paging
        let tcb_offset = paging::init_ap(cpu_id, bsp_table, stack_start, stack_end);

        // Set up GDT with TLS
        gdt::init_paging(tcb_offset, stack_end);

        // Set up IDT for AP
        idt::init_paging();

        // Set up syscall instruction
        interrupt::syscall::init();

        // Test tdata and tbss
        {
            assert_eq!(TBSS_TEST_ZERO, 0);
            TBSS_TEST_ZERO += 1;
            assert_eq!(TBSS_TEST_ZERO, 1);
            assert_eq!(TDATA_TEST_NONZERO, 0xFFFF_FFFF_FFFF_FFFF);
            TDATA_TEST_NONZERO -= 1;
            assert_eq!(TDATA_TEST_NONZERO, 0xFFFF_FFFF_FFFF_FFFE);
        }

        // Initialize devices (for AP)
        device::init_ap();

        AP_READY.store(true, Ordering::SeqCst);

        cpu_id
    };

    while ! BSP_READY.load(Ordering::SeqCst) {
        interrupt::pause();
    }

    crate::kmain_ap(cpu_id);
}

#[naked]
pub unsafe fn usermode(ip: usize, sp: usize, arg: usize) -> ! {
    asm!("push r10
          push r11
          push r12
          push r13
          push r14
          push r15"
          : // No output
          :   "{r10}"(gdt::GDT_USER_DATA << 3 | 3), // Data segment
              "{r11}"(sp), // Stack pointer
              "{r12}"(1 << 9), // Flags - Set interrupt enable flag
              "{r13}"(gdt::GDT_USER_CODE << 3 | 3), // Code segment
              "{r14}"(ip), // IP
              "{r15}"(arg) // Argument
          : // No clobbers
          : "intel", "volatile");

    // Unmap kernel
    pti::unmap();

    // Go to usermode
    asm!("mov ds, r14d
         mov es, r14d
         mov fs, r15d
         mov gs, r14d
         xor rax, rax
         xor rbx, rbx
         xor rcx, rcx
         xor rdx, rdx
         xor rsi, rsi
         xor rdi, rdi
         xor rbp, rbp
         xor r8, r8
         xor r9, r9
         xor r10, r10
         xor r11, r11
         xor r12, r12
         xor r13, r13
         xor r14, r14
         xor r15, r15
         fninit
         pop rdi
         iretq"
         : // No output because it never returns
         :   "{r14}"(gdt::GDT_USER_DATA << 3 | 3), // Data segment
             "{r15}"(gdt::GDT_USER_TLS << 3 | 3) // TLS segment
         : // No clobbers because it never returns
         : "intel", "volatile");
    unreachable!();
}
<|endoftext|>


<commit_before>// Copyright 2017 The Exonum Team
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//   http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use futures::{Future, Stream, Sink, IntoFuture};
use futures::future::Either;
use futures::sync::mpsc;
use futures::unsync;
use tokio_core::net::{TcpListener, TcpStream, TcpStreamNew};
use tokio_core::reactor::{Core, Timeout, Handle};
use tokio_io::AsyncRead;
use tokio_retry::{Retry, Action};
use tokio_retry::strategy::{FibonacciBackoff, jitter};

use std::net::SocketAddr;
use std::time::{Duration};
use std::collections::HashMap;

use messages::{Any, Connect, RawMessage};
use node::{ExternalMessage, NodeTimeout};
use helpers::Milliseconds;

use super::EventHandler;
use super::error::{other_error, result_ok, forget_result, into_other, log_error};
use super::codec::MessagesCodec;
use super::EventsAggregator;

#[derive(Debug)]
pub enum NetworkEvent {
    MessageReceived(SocketAddr, RawMessage),
    PeerConnected(SocketAddr, Connect),
    PeerDisconnected(SocketAddr),
}

#[derive(Debug, Clone)]
pub enum NetworkRequest {
    SendMessage(SocketAddr, RawMessage),
    DisconnectWithPeer(SocketAddr),
    Shutdown,
}

#[derive(Serialize, Deserialize, Debug, Clone, Copy)]
pub struct NetworkConfiguration {
    // TODO: think more about config parameters
    pub max_incoming_connections: usize,
    pub max_outgoing_connections: usize,
    pub tcp_nodelay: bool,
    pub tcp_keep_alive: Option<u64>,
    pub tcp_connect_retry_timeout: Milliseconds,
    pub tcp_connect_max_retries: u64,
}

impl Default for NetworkConfiguration {
    fn default() -> NetworkConfiguration {
        NetworkConfiguration {
            max_incoming_connections: 128,
            max_outgoing_connections: 128,
            tcp_keep_alive: None,
            tcp_nodelay: false,
            tcp_connect_retry_timeout: 5000,
            tcp_connect_max_retries: 10,
        }
    }
}

#[derive(Debug)]
pub struct HandlerPart<H: EventHandler> {
    pub core: Core,
    pub handler: H,
    pub timeout_rx: mpsc::Receiver<NodeTimeout>,
    pub network_rx: mpsc::Receiver<NetworkEvent>,
    pub api_rx: mpsc::Receiver<ExternalMessage>,
}

#[derive(Debug)]
pub struct NetworkPart {
    pub listen_address: SocketAddr,
    pub network_config: NetworkConfiguration,
    pub network_requests: (mpsc::Sender<NetworkRequest>, mpsc::Receiver<NetworkRequest>),
    pub network_tx: mpsc::Sender<NetworkEvent>,
}

impl<H: EventHandler> HandlerPart<H> {
    pub fn run(self) -> Result<(), ()> {
        let mut core = self.core;
        let mut handler = self.handler;

        let events_handle = EventsAggregator::new(self.timeout_rx, self.network_rx, self.api_rx)
            .for_each(move |event| {
                handler.handle_event(event);
                Ok(())
            });
        core.run(events_handle)
    }
}

macro_rules! try_future_boxed {
    ($e:expr) =>(
        match $e {
            Ok(v) => v,
            Err(e) => {
                return Err(into_other(e)).into_future().boxed();
            }
        }
    )
}

impl NetworkPart {
    pub fn run(self) -> Result<(), ()> {
        let mut core = Core::new().unwrap();

        // Cancelation token
        let (cancel_sender, cancel_handler) = unsync::mpsc::channel(1);

        // Outgoing connections handler
        let mut outgoing_connections: HashMap<SocketAddr, mpsc::Sender<RawMessage>> =
            HashMap::new();

        // Requests handler
        let handle = core.handle();
        let network_tx = self.network_tx.clone();
        let requests_tx = self.network_requests.0.clone();
        let network_config = self.network_config;
        let requests_handle = self.network_requests.1.for_each(move |request| {
            let cancel_sender = cancel_sender.clone();
            match request {
                NetworkRequest::SendMessage(peer, msg) => {
                    let conn_tx = if let Some(conn_tx) = outgoing_connections.get(&peer).cloned() {
                        conn_tx
                    } else {
                        // Register outgoing channel.
                        let (conn_tx, conn_rx) = mpsc::channel(10);
                        outgoing_connections.insert(peer, conn_tx.clone());
                        // Enable retry feature for outgoing connection.
                        let requests_tx = requests_tx.clone();
                        let timeout = network_config.tcp_connect_retry_timeout;
                        let max_tries = network_config.tcp_connect_max_retries as usize;
                        let connect_handle = Retry::spawn(
                            handle.clone(),
                            FibonacciBackoff::from_millis(timeout)
                                .map(jitter)
                                .take(max_tries),
                            TcpStreamConnectAction {
                                handle: handle.clone(),
                                peer,
                            },
                        ).map_err(into_other)
                            .and_then(move |sock| {
                                try_future_boxed!(sock.set_nodelay(network_config.tcp_nodelay));
                                try_future_boxed!(sock.set_keepalive(
                                    network_config.tcp_keep_alive.map(Duration::from_millis),
                                ));

                                info!("Established connection with peer={}", peer);

                                let stream = sock.framed(MessagesCodec);
                                let (sink, stream) = stream.split();

                                let writer = conn_rx
                                    .map_err(|_| other_error("Can't send data into socket"))
                                    .forward(sink);
                                let reader = stream.for_each(result_ok).map_err(into_other);

                                reader
                                    .select2(writer)
                                    .map_err(|_| other_error("Socket error"))
                                    .and_then(|res| match res {
                                        Either::A((_, _reader)) => Ok(()).into_future(),
                                        Either::B((_, _writer)) => Ok(()).into_future(),
                                    })
                                    .boxed()
                            })
                            .then(move |res| {
                                info!("Connection with peer={} closed, reason={:?}", peer, res);
                                // outgoing_connections.remove(&peer);

                                let request = NetworkRequest::DisconnectWithPeer(peer);
                                requests_tx
                                    .clone()
                                    .send(request)
                                    .map(forget_result)
                                    .map_err(into_other)
                            })
                            .map_err(log_error);
                        handle.spawn(connect_handle);
                        conn_tx
                    };

                    let duration = Duration::from_secs(5);
                    let send_timeout = Timeout::new(duration, &handle)
                        .unwrap()
                        .and_then(result_ok)
                        .map_err(|_| other_error("Can't timeout"));

                    let send_handle = conn_tx.send(msg).map(forget_result).map_err(log_error);

                    let timeouted_connect = send_handle
                        .select2(send_timeout)
                        .map_err(|_| other_error("Unable to send message"))
                        .and_then(move |either| match either {
                            Either::A((send, _timeout_fut)) => Ok(send),
                            Either::B((_, _connect_fut)) => Err(other_error("Send timeout")),
                        })
                        .map_err(log_error);

                    handle.spawn(timeouted_connect);
                }
                NetworkRequest::DisconnectWithPeer(peer) => {
                    outgoing_connections.remove(&peer);

                    let event = NetworkEvent::PeerDisconnected(peer);
                    let event_handle = network_tx.clone().send(event).map(forget_result).map_err(
                        log_error,
                    );
                    handle.spawn(event_handle);
                }
                // Immediately stop the event loop.
                NetworkRequest::Shutdown => {
                    cancel_sender
                        .clone()
                        .send(())
                        .map(forget_result)
                        .map_err(log_error)
                        .wait()?
                }
            }

            Ok(())
        });

        // Incoming connections handler
        let listener = TcpListener::bind(&self.listen_address, &core.handle()).unwrap();
        let network_tx = self.network_tx.clone();
        let handle = core.handle();
        let server = listener
            .incoming()
            .for_each(move |(sock, addr)| {
                info!("Accepted incoming connection with peer={}", addr);

                let stream = sock.framed(MessagesCodec);
                let (_, stream) = stream.split();
                let network_tx = network_tx.clone();
                let connection_handler = stream
                    .into_future()
                    .map_err(|e| e.0)
                    .and_then(move |(raw, stream)| match raw.map(Any::from_raw) {
                        Some(Ok(Any::Connect(msg))) => Ok((msg, stream)),
                        Some(Ok(other)) => Err(other_error(&format!(
                            "First message is not Connect, got={:?}",
                            other
                        ))),
                        Some(Err(e)) => Err(into_other(e)),
                        None => Err(other_error("Incoming socket closed")),
                    })
                    .and_then(move |(connect, stream)| {
                        info!("Received handshake message={:?}", connect);

                        let event = NetworkEvent::PeerConnected(addr, connect);
                        let stream = network_tx
                            .clone()
                            .send(event)
                            .map_err(into_other)
                            .and_then(move |_| Ok(stream))
                            .flatten_stream();

                        stream.for_each(move |raw| {
                            let event = NetworkEvent::MessageReceived(addr, raw);
                            network_tx.clone().send(event).map(forget_result).map_err(
                                into_other,
                            )
                        })
                    })
                    .map(forget_result)
                    .map_err(log_error);
                handle.spawn(connection_handler);
                Ok(())
            })
            .map_err(log_error);

        core.handle().spawn(server);
        core.handle().spawn(requests_handle);
        core.run(
            cancel_handler
                .into_future()
                .map(|_| info!("Network thread shutdown"))
                .map_err(|_| error!("An error during shutdown occured")),
        )
    }
}

// Tcp streams source for retry action.
struct TcpStreamConnectAction {
    handle: Handle,
    peer: SocketAddr,
}

impl Action for TcpStreamConnectAction {
    type Future = TcpStreamNew;
    type Item = <TcpStreamNew as Future>::Item;
    type Error = <TcpStreamNew as Future>::Error;

    fn run(&mut self) -> Self::Future {
        TcpStream::connect(&self.peer, &self.handle)
    }
}<commit_msg>Use fixed interval strategy<commit_after>// Copyright 2017 The Exonum Team
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//   http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use futures::{Future, Stream, Sink, IntoFuture};
use futures::future::Either;
use futures::sync::mpsc;
use futures::unsync;
use tokio_core::net::{TcpListener, TcpStream, TcpStreamNew};
use tokio_core::reactor::{Core, Timeout, Handle};
use tokio_io::AsyncRead;
use tokio_retry::{Retry, Action};
use tokio_retry::strategy::{FixedInterval, jitter};

use std::net::SocketAddr;
use std::time::Duration;
use std::collections::HashMap;

use messages::{Any, Connect, RawMessage};
use node::{ExternalMessage, NodeTimeout};
use helpers::Milliseconds;

use super::EventHandler;
use super::error::{other_error, result_ok, forget_result, into_other, log_error};
use super::codec::MessagesCodec;
use super::EventsAggregator;

#[derive(Debug)]
pub enum NetworkEvent {
    MessageReceived(SocketAddr, RawMessage),
    PeerConnected(SocketAddr, Connect),
    PeerDisconnected(SocketAddr),
}

#[derive(Debug, Clone)]
pub enum NetworkRequest {
    SendMessage(SocketAddr, RawMessage),
    DisconnectWithPeer(SocketAddr),
    Shutdown,
}

#[derive(Serialize, Deserialize, Debug, Clone, Copy)]
pub struct NetworkConfiguration {
    // TODO: think more about config parameters
    pub max_incoming_connections: usize,
    pub max_outgoing_connections: usize,
    pub tcp_nodelay: bool,
    pub tcp_keep_alive: Option<u64>,
    pub tcp_connect_retry_timeout: Milliseconds,
    pub tcp_connect_max_retries: u64,
}

impl Default for NetworkConfiguration {
    fn default() -> NetworkConfiguration {
        NetworkConfiguration {
            max_incoming_connections: 128,
            max_outgoing_connections: 128,
            tcp_keep_alive: None,
            tcp_nodelay: false,
            tcp_connect_retry_timeout: 15000,
            tcp_connect_max_retries: 10,
        }
    }
}

#[derive(Debug)]
pub struct HandlerPart<H: EventHandler> {
    pub core: Core,
    pub handler: H,
    pub timeout_rx: mpsc::Receiver<NodeTimeout>,
    pub network_rx: mpsc::Receiver<NetworkEvent>,
    pub api_rx: mpsc::Receiver<ExternalMessage>,
}

#[derive(Debug)]
pub struct NetworkPart {
    pub listen_address: SocketAddr,
    pub network_config: NetworkConfiguration,
    pub network_requests: (mpsc::Sender<NetworkRequest>, mpsc::Receiver<NetworkRequest>),
    pub network_tx: mpsc::Sender<NetworkEvent>,
}

impl<H: EventHandler> HandlerPart<H> {
    pub fn run(self) -> Result<(), ()> {
        let mut core = self.core;
        let mut handler = self.handler;

        let events_handle = EventsAggregator::new(self.timeout_rx, self.network_rx, self.api_rx)
            .for_each(move |event| {
                handler.handle_event(event);
                Ok(())
            });
        core.run(events_handle)
    }
}

macro_rules! try_future_boxed {
    ($e:expr) =>(
        match $e {
            Ok(v) => v,
            Err(e) => {
                return Err(into_other(e)).into_future().boxed();
            }
        }
    )
}

impl NetworkPart {
    pub fn run(self) -> Result<(), ()> {
        let mut core = Core::new().unwrap();

        // Cancelation token
        let (cancel_sender, cancel_handler) = unsync::mpsc::channel(1);

        // Outgoing connections handler
        let mut outgoing_connections: HashMap<SocketAddr, mpsc::Sender<RawMessage>> =
            HashMap::new();

        // Requests handler
        let handle = core.handle();
        let network_tx = self.network_tx.clone();
        let requests_tx = self.network_requests.0.clone();
        let network_config = self.network_config;
        let requests_handle = self.network_requests.1.for_each(move |request| {
            let cancel_sender = cancel_sender.clone();
            match request {
                NetworkRequest::SendMessage(peer, msg) => {
                    let conn_tx = if let Some(conn_tx) = outgoing_connections.get(&peer).cloned() {
                        conn_tx
                    } else {
                        // Register outgoing channel.
                        let (conn_tx, conn_rx) = mpsc::channel(10);
                        outgoing_connections.insert(peer, conn_tx.clone());
                        // Enable retry feature for outgoing connection.
                        let requests_tx = requests_tx.clone();
                        let timeout = network_config.tcp_connect_retry_timeout;
                        let max_tries = network_config.tcp_connect_max_retries as usize;
                        let connect_handle = Retry::spawn(
                            handle.clone(),
                            FixedInterval::from_millis(timeout).map(jitter).take(
                                max_tries,
                            ),
                            TcpStreamConnectAction {
                                handle: handle.clone(),
                                peer,
                            },
                        ).map_err(into_other)
                            .and_then(move |sock| {
                                try_future_boxed!(sock.set_nodelay(network_config.tcp_nodelay));
                                try_future_boxed!(sock.set_keepalive(
                                    network_config.tcp_keep_alive.map(Duration::from_millis),
                                ));

                                info!("Established connection with peer={}", peer);

                                let stream = sock.framed(MessagesCodec);
                                let (sink, stream) = stream.split();

                                let writer = conn_rx
                                    .map_err(|_| other_error("Can't send data into socket"))
                                    .forward(sink);
                                let reader = stream.for_each(result_ok).map_err(into_other);

                                reader
                                    .select2(writer)
                                    .map_err(|_| other_error("Socket error"))
                                    .and_then(|res| match res {
                                        Either::A((_, _reader)) => Ok(()).into_future(),
                                        Either::B((_, _writer)) => Ok(()).into_future(),
                                    })
                                    .boxed()
                            })
                            .then(move |res| {
                                info!("Connection with peer={} closed, reason={:?}", peer, res);
                                // outgoing_connections.remove(&peer);

                                let request = NetworkRequest::DisconnectWithPeer(peer);
                                requests_tx
                                    .clone()
                                    .send(request)
                                    .map(forget_result)
                                    .map_err(into_other)
                            })
                            .map_err(log_error);
                        handle.spawn(connect_handle);
                        conn_tx
                    };

                    let duration = Duration::from_secs(5);
                    let send_timeout = Timeout::new(duration, &handle)
                        .unwrap()
                        .and_then(result_ok)
                        .map_err(|_| other_error("Can't timeout"));

                    let send_handle = conn_tx.send(msg).map(forget_result).map_err(log_error);

                    let timeouted_connect = send_handle
                        .select2(send_timeout)
                        .map_err(|_| other_error("Unable to send message"))
                        .and_then(move |either| match either {
                            Either::A((send, _timeout_fut)) => Ok(send),
                            Either::B((_, _connect_fut)) => Err(other_error("Send timeout")),
                        })
                        .map_err(log_error);

                    handle.spawn(timeouted_connect);
                }
                NetworkRequest::DisconnectWithPeer(peer) => {
                    outgoing_connections.remove(&peer);

                    let event = NetworkEvent::PeerDisconnected(peer);
                    let event_handle = network_tx.clone().send(event).map(forget_result).map_err(
                        log_error,
                    );
                    handle.spawn(event_handle);
                }
                // Immediately stop the event loop.
                NetworkRequest::Shutdown => {
                    cancel_sender
                        .clone()
                        .send(())
                        .map(forget_result)
                        .map_err(log_error)
                        .wait()?
                }
            }

            Ok(())
        });

        // Incoming connections handler
        let listener = TcpListener::bind(&self.listen_address, &core.handle()).unwrap();
        let network_tx = self.network_tx.clone();
        let handle = core.handle();
        let server = listener
            .incoming()
            .for_each(move |(sock, addr)| {
                info!("Accepted incoming connection with peer={}", addr);

                let stream = sock.framed(MessagesCodec);
                let (_, stream) = stream.split();
                let network_tx = network_tx.clone();
                let connection_handler = stream
                    .into_future()
                    .map_err(|e| e.0)
                    .and_then(move |(raw, stream)| match raw.map(Any::from_raw) {
                        Some(Ok(Any::Connect(msg))) => Ok((msg, stream)),
                        Some(Ok(other)) => Err(other_error(&format!(
                            "First message is not Connect, got={:?}",
                            other
                        ))),
                        Some(Err(e)) => Err(into_other(e)),
                        None => Err(other_error("Incoming socket closed")),
                    })
                    .and_then(move |(connect, stream)| {
                        info!("Received handshake message={:?}", connect);

                        let event = NetworkEvent::PeerConnected(addr, connect);
                        let stream = network_tx
                            .clone()
                            .send(event)
                            .map_err(into_other)
                            .and_then(move |_| Ok(stream))
                            .flatten_stream();

                        stream.for_each(move |raw| {
                            let event = NetworkEvent::MessageReceived(addr, raw);
                            network_tx.clone().send(event).map(forget_result).map_err(
                                into_other,
                            )
                        })
                    })
                    .map(forget_result)
                    .map_err(log_error);
                handle.spawn(connection_handler);
                Ok(())
            })
            .map_err(log_error);

        core.handle().spawn(server);
        core.handle().spawn(requests_handle);
        core.run(
            cancel_handler
                .into_future()
                .map(|_| info!("Network thread shutdown"))
                .map_err(|_| error!("An error during shutdown occured")),
        )
    }
}

// Tcp streams source for retry action.
struct TcpStreamConnectAction {
    handle: Handle,
    peer: SocketAddr,
}

impl Action for TcpStreamConnectAction {
    type Future = TcpStreamNew;
    type Item = <TcpStreamNew as Future>::Item;
    type Error = <TcpStreamNew as Future>::Error;

    fn run(&mut self) -> Self::Future {
        TcpStream::connect(&self.peer, &self.handle)
    }
}
<|endoftext|>


<commit_before>use bitstream::byteread::*;

pub trait BitReadEndian {
    fn peek_val(&mut self, n:usize) -> u64;
    fn merge_val(msp:u64, lsp:u64, msb:usize, lsb:usize) -> u64;
    fn skip_rem(&mut self, n:usize) -> ();
    fn fill32(&self) -> u64;
    fn fill64(&self) -> u64;
}

pub trait BitReadInternal : BitReadEndian {
    #[inline]
    fn left(&self) -> usize;
    fn refill32(&mut self) -> ();
    fn refill64(&mut self) -> ();

    fn get_val(&mut self, n:usize) -> u64 {
        let ret = self.peek_val(n);

        self.skip_rem(n);

        return ret;
    }

}

pub trait BitRead<'a>: BitReadInternal+Copy {
    fn new(&'a[u8]) -> Self;
    fn consumed(&self) -> usize;
    fn available(&self) -> usize;
    fn can_refill(&self) -> bool;

    fn skip_bits(&mut self, size : usize) -> ();

    #[inline]
    fn get_bit(&mut self) -> bool {
        if self.left() <= 0 {
            self.refill64();
        }

        self.get_val(1) != 0
    }

    #[inline]
    fn get_bits_64(&mut self, mut n:usize) -> u64 {
        let mut left = 0;
        let mut ret = 0;

        if n == 0 {
            return 0;
        }

        if self.left() < n {
            n   -= self.left();
            left = self.left();
            ret  = self.get_val(left);
            self.refill64();
        }

        Self::merge_val(self.get_val(n), ret, left, n)
    }

    #[inline]
    fn get_bits_32(&mut self, n:usize) -> u32 {
        if n == 0 {
            return 0;
        }

        if self.left() <= n {
            self.refill32();
        }

        return self.get_val(n) as u32;
    }

    #[inline]
    fn peek_bit(&mut self) -> bool {
        let mut tmp = *self;

        tmp.get_bit()
    }

    #[inline]
    fn peek_bits_32(&mut self, n:usize) -> u32 {
        let mut tmp = *self;

        tmp.get_bits_32(n)
    }

    #[inline]
    fn peek_bits_64(&self, n:usize) -> u64 {
        let mut tmp = *self;

        tmp.get_bits_64(n)
    }

    #[inline]
    fn align_bits(&mut self) -> () {
        let left = self.left() & 7;

        self.skip_bits(left);
    }
}

macro_rules! endian_reader {
    {$name: ident} => {
        #[derive(Debug, Clone, Copy)]
        pub struct $name<'a> {
            buffer : &'a[u8], /// read buffer, 8-bytes padded
            index : usize,
            cache : u64,
            left : usize,
        }
        impl <'a> BitReadInternal for $name<'a> {
            #[inline]
            fn left(&self) -> usize {
                self.left
            }
            #[inline]
            fn refill32(&mut self) -> () {
                if !self.can_refill() {
                    return;
                }
                let val = self.fill32();

                self.cache  = Self::merge_val(val, self.cache,
                                              self.left, 32 - self.left);
                self.index += 4;
                self.left  += 32;
            }
            #[inline]
            fn refill64(&mut self) -> () {
                if !self.can_refill() {
                    return;
                }

                self.cache  = self.fill64();
                self.index += 8;
                self.left   = 64;
            }
        }

        impl <'a> BitRead<'a> for $name<'a> {
            fn new(buffer: &'a[u8]) -> $name<'a> {
                let mut reader = $name {
                    buffer: buffer,
                    index: 0,
                    cache: 0,
                    left: 0
                };

                reader.refill64();
                return reader;
            }
            #[inline]
            fn consumed(&self) -> usize {
                self.index * 8 - self.left
            }

            #[inline]
            fn available(&self) -> usize {
                (self.buffer.len() - self.index) * 8 + self.left
            }

            #[inline]
            fn can_refill(&self) -> bool {
                self.index < self.buffer.len() - 8
            }

            #[inline]
            fn skip_bits(&mut self, mut n:usize) -> () {
                if n == 0 {
                    return;
                }
                if self.left <= n {
                    n -= self.left;
                    self.skip_rem(n);
                    if n > 64 {
                        let skip = n / 8;

                        n -= skip * 8;
                        self.index += skip;
                    }
                    self.refill64();
                }

                self.skip_rem(n);
            }

        }
    }
}

endian_reader!{ BitReadLE }

impl <'a> BitReadEndian for BitReadLE<'a> {
    #[inline]
    fn peek_val(&mut self, n:usize) -> u64 {
        let v = self.cache & ((1u64 << n) - 1);

        return v;
    }
    #[inline]
    fn skip_rem(&mut self, n:usize) -> () {
        self.cache = self.cache >> n;
        self.left -= n;
    }
    #[inline]
    fn merge_val(msp:u64, lsp:u64, msb:usize, _:usize) -> u64 {
        msp << msb | lsp
    }
    #[inline(always)]
    fn fill32(&self) -> u64 {
        get_u32l(&self.buffer[self.index..]) as u64
    }
    #[inline(always)]
    fn fill64(&self) -> u64 {
        get_u64l(&self.buffer[self.index..])
    }
}

endian_reader!{ BitReadBE }

impl <'a> BitReadEndian for BitReadBE<'a> {
    #[inline]
    fn peek_val(&mut self, n:usize) -> u64 {
        self.cache >> (64 - n)
    }
    #[inline]
    fn skip_rem(&mut self, n:usize) -> () {
        self.cache = self.cache << n;
        self.left -= n;
    }
    #[inline]
    fn merge_val(msp:u64, lsp:u64, _:usize, lsb:usize) -> u64 {
        msp | lsp << lsb
    }
    #[inline(always)]
    fn fill32(&self) -> u64 {
        get_u32b(&self.buffer[self.index..]) as u64
    }
    #[inline(always)]
    fn fill64(&self) -> u64 {
        get_u64b(&self.buffer[self.index..])
    }
}

#[cfg(test)]
mod test {
    pub const CHECKBOARD0101: [u8; 128] = [0b01010101; 128];
    pub const CHECKBOARD0011: [u8; 128] = [0b00110011; 128];

    mod le {
        use super::super::*;
        use super::*;

        #[test]
        fn get_bit() {
            let b = &CHECKBOARD0101;
            let mut reader = BitReadLE::new(b);

            assert!(reader.get_bit());
            assert!(!reader.get_bit());
        }

        #[test]
        fn get_bits_64() {
            let b = &CHECKBOARD0101;
            let mut reader = BitReadLE::new(b);

            assert!(reader.get_bits_64(1) == 1);
            assert!(reader.get_bits_64(2) == 2);
            assert!(reader.get_bits_64(4) == 10);
            assert!(reader.get_bits_64(1) == 0);
            assert!(reader.get_bits_64(8) == 85);
        }

        #[test]
        fn peek_bits_64() {
            let reader = BitReadLE {
                buffer: &CHECKBOARD0101,
                index: 0,
                cache: 0,
                left: 0
            };

            assert!(reader.peek_bits_64(1) == 1);
            assert!(reader.peek_bits_64(1) == 1);
            assert!(reader.peek_bits_64(2) == 1);
            assert!(reader.peek_bits_64(2) == 1);
        }

        #[test]
        fn get_bits_32() {
            let mut reader = BitReadLE {
                buffer: &CHECKBOARD0101,
                index: 0,
                cache: 0,
                left: 0
            };

            assert!(reader.get_bits_32(1) == 1);
            assert!(reader.get_bits_32(2) == 2);
            assert!(reader.get_bits_32(4) == 10);
            assert!(reader.get_bits_32(1) == 0);
            assert!(reader.get_bits_32(8) == 85);
        }

        #[test]
        fn peek_bits_32() {
            let b = &CHECKBOARD0101;
            let mut reader = BitReadLE::new(b);

            assert!(reader.peek_bits_32(1) == 1);
            assert!(reader.peek_bits_32(1) == 1);
            assert!(reader.peek_bits_32(2) == 1);
            assert!(reader.peek_bits_32(2) == 1);
        }

        #[test]
        fn skip_bits() {
            let b = &CHECKBOARD0101;
            let mut reader = BitReadLE::new(b);

            reader.skip_bits(0);
            assert!(reader.peek_bits_32(1) == 1);
            reader.skip_bits(2);
            assert!(reader.peek_bits_32(1) == 1);
            reader.skip_bits(2);
            assert!(reader.peek_bits_32(1) == 1);
        }

        #[test]
        fn align_bits() {
            let b = &CHECKBOARD0011;
            let mut reader = BitReadLE::new(b);

            reader.align_bits();
            assert!(reader.get_bits_64(3) == 3);
            reader.align_bits();
            assert!(reader.get_bits_64(4) == 3);
            reader.skip_bits(1);
            reader.align_bits();
            assert!(reader.get_bits_64(4) == 3);
        }
    }
    mod be {
        use super::super::*;
        use super::*;

        #[test]
        fn get_bit() {
            let b = &CHECKBOARD0101;
            let mut reader = BitReadBE::new(b);

            assert!(!reader.get_bit());
            assert!(reader.get_bit());
        }

        #[test]
        fn get_bits_64() {
            let b = &CHECKBOARD0101;
            let mut reader = BitReadBE::new(b);

            assert!(reader.get_bits_64(1) == 0);
            assert!(reader.get_bits_64(2) == 2);
            assert!(reader.get_bits_64(4) == 10);
            assert!(reader.get_bits_64(1) == 1);
            assert!(reader.get_bits_64(8) == 85);
        }

        #[test]
        fn peek_bits_64() {
            let b = &CHECKBOARD0101;
            let reader = BitReadBE::new(b);

            assert!(reader.peek_bits_64(1) == 0);
            assert!(reader.peek_bits_64(1) == 0);
            assert!(reader.peek_bits_64(2) == 1);
            assert!(reader.peek_bits_64(2) == 1);
        }

        #[test]
        fn get_bits_32() {
            let b = &CHECKBOARD0101;
            let mut reader = BitReadBE::new(b);

            assert!(reader.get_bits_32(1) == 0);
            assert!(reader.get_bits_32(2) == 2);
            assert!(reader.get_bits_32(4) == 10);
            assert!(reader.get_bits_32(1) == 1);
            assert!(reader.get_bits_32(8) == 85);
        }

        #[test]
        fn peek_bits_32() {
            let b = &CHECKBOARD0101;
            let mut reader = BitReadBE::new(b);

            assert!(reader.peek_bits_32(1) == 0);
            assert!(reader.peek_bits_32(1) == 0);
            assert!(reader.peek_bits_32(2) == 1);
            assert!(reader.peek_bits_32(2) == 1);
        }

        #[test]
        fn skip_bits() {
            let b = &CHECKBOARD0101;
            let mut reader = BitReadBE::new(b);

            reader.skip_bits(0);
            assert!(reader.peek_bits_32(1) == 0);
            reader.skip_bits(2);
            assert!(reader.peek_bits_32(1) == 0);
            reader.skip_bits(2);
            assert!(reader.peek_bits_32(1) == 0);
        }

        #[test]
        fn align_bits() {
            let b = &CHECKBOARD0011;
            let mut reader = BitReadBE::new(b);

            reader.align_bits();
            assert!(reader.get_bits_64(3) == 1);
            reader.align_bits();
            assert!(reader.get_bits_64(4) == 3);
            reader.skip_bits(1);
            reader.align_bits();
            assert!(reader.get_bits_64(4) == 3);
        }
    }
}
<commit_msg>bitread: Handle overflowing corner cases<commit_after>use bitstream::byteread::*;

pub trait BitReadEndian {
    fn peek_val(&mut self, n:usize) -> u64;
    fn merge_val(msp:u64, lsp:u64, msb:usize, lsb:usize) -> u64;
    fn skip_rem(&mut self, n:usize) -> ();
    fn fill32(&self) -> u64;
    fn fill64(&self) -> u64;
}

pub trait BitReadInternal : BitReadEndian {
    #[inline]
    fn left(&self) -> usize;
    fn refill32(&mut self) -> ();
    fn refill64(&mut self) -> ();

    fn get_val(&mut self, n:usize) -> u64 {
        let ret = self.peek_val(n);

        self.skip_rem(n);

        return ret;
    }

}

pub trait BitRead<'a>: BitReadInternal+Copy {
    fn new(&'a[u8]) -> Self;
    fn consumed(&self) -> usize;
    fn available(&self) -> usize;
    fn can_refill(&self) -> bool;

    fn skip_bits(&mut self, size : usize) -> ();

    #[inline]
    fn get_bit(&mut self) -> bool {
        if self.left() <= 0 {
            self.refill64();
        }

        self.get_val(1) != 0
    }

    #[inline]
    fn get_bits_64(&mut self, mut n:usize) -> u64 {
        let mut left = 0;
        let mut ret = 0;

        if n == 0 {
            return 0;
        }

        if self.left() < n {
            n   -= self.left();
            left = self.left();
            ret  = self.get_val(left);
            self.refill64();
        }

        Self::merge_val(self.get_val(n), ret, left, n)
    }

    #[inline]
    fn get_bits_32(&mut self, n:usize) -> u32 {
        if n == 0 {
            return 0;
        }

        if self.left() <= n {
            self.refill32();
        }

        return self.get_val(n) as u32;
    }

    #[inline]
    fn peek_bit(&mut self) -> bool {
        let mut tmp = *self;

        tmp.get_bit()
    }

    #[inline]
    fn peek_bits_32(&mut self, n:usize) -> u32 {
        let mut tmp = *self;

        tmp.get_bits_32(n)
    }

    #[inline]
    fn peek_bits_64(&self, n:usize) -> u64 {
        let mut tmp = *self;

        tmp.get_bits_64(n)
    }

    #[inline]
    fn align_bits(&mut self) -> () {
        let left = self.left() & 7;

        self.skip_bits(left);
    }
}

macro_rules! endian_reader {
    {$name: ident} => {
        #[derive(Debug, Clone, Copy)]
        pub struct $name<'a> {
            buffer : &'a[u8], /// read buffer, 8-bytes padded
            index : usize,
            cache : u64,
            left : usize,
        }
        impl <'a> BitReadInternal for $name<'a> {
            #[inline]
            fn left(&self) -> usize {
                self.left
            }
            #[inline]
            fn refill32(&mut self) -> () {
                if !self.can_refill() {
                    return;
                }
                let val = self.fill32();

                self.cache  = Self::merge_val(val, self.cache,
                                              self.left, 32 - self.left);
                self.index += 4;
                self.left  += 32;
            }
            #[inline]
            fn refill64(&mut self) -> () {
                if !self.can_refill() {
                    return;
                }

                self.cache  = self.fill64();
                self.index += 8;
                self.left   = 64;
            }
        }

        impl <'a> BitRead<'a> for $name<'a> {
            fn new(buffer: &'a[u8]) -> $name<'a> {
                let mut reader = $name {
                    buffer: buffer,
                    index: 0,
                    cache: 0,
                    left: 0
                };

                reader.refill64();
                return reader;
            }
            #[inline]
            fn consumed(&self) -> usize {
                self.index * 8 - self.left
            }

            #[inline]
            fn available(&self) -> usize {
                (self.buffer.len() - self.index) * 8 + self.left
            }

            #[inline]
            fn can_refill(&self) -> bool {
                self.index + 8 <= self.buffer.len()
            }

            #[inline]
            fn skip_bits(&mut self, mut n:usize) -> () {
                if self.left < n {
                    n -= self.left;
                    if n > 64 {
                        let skip = n / 8;

                        n -= skip * 8;
                        self.index += skip;
                    }
                    self.skip_rem(n);
                    self.refill64();
                }

                self.skip_rem(n);
            }

        }
    }
}

endian_reader!{ BitReadLE }

impl <'a> BitReadEndian for BitReadLE<'a> {
    #[inline]
    fn peek_val(&mut self, n:usize) -> u64 {
        let v = self.cache & ((1u64 << n) - 1);

        return v;
    }
    #[inline]
    fn skip_rem(&mut self, n:usize) -> () {
        self.cache = self.cache >> n;
        self.left = self.left.saturating_sub(n);
    }
    #[inline]
    fn merge_val(msp:u64, lsp:u64, msb:usize, _:usize) -> u64 {
        msp << msb | lsp
    }
    #[inline(always)]
    fn fill32(&self) -> u64 {
        get_u32l(&self.buffer[self.index..]) as u64
    }
    #[inline(always)]
    fn fill64(&self) -> u64 {
        get_u64l(&self.buffer[self.index..])
    }
}

endian_reader!{ BitReadBE }

impl <'a> BitReadEndian for BitReadBE<'a> {
    #[inline]
    fn peek_val(&mut self, n:usize) -> u64 {
        self.cache >> (64 - n)
    }
    #[inline]
    fn skip_rem(&mut self, n:usize) -> () {
        self.cache = self.cache << n;
        self.left = self.left.saturating_sub(n);
    }
    #[inline]
    fn merge_val(msp:u64, lsp:u64, _:usize, lsb:usize) -> u64 {
        msp | lsp << lsb
    }
    #[inline(always)]
    fn fill32(&self) -> u64 {
        get_u32b(&self.buffer[self.index..]) as u64
    }
    #[inline(always)]
    fn fill64(&self) -> u64 {
        get_u64b(&self.buffer[self.index..])
    }
}

#[cfg(test)]
mod test {
    pub const CHECKBOARD0101: [u8; 128] = [0b01010101; 128];
    pub const CHECKBOARD0011: [u8; 128] = [0b00110011; 128];

    mod le {
        use super::super::*;
        use super::*;

        #[test]
        fn get_bit() {
            let b = &CHECKBOARD0101;
            let mut reader = BitReadLE::new(b);

            assert!(reader.get_bit());
            assert!(!reader.get_bit());
        }

        #[test]
        fn get_bits_64() {
            let b = &CHECKBOARD0101;
            let mut reader = BitReadLE::new(b);

            assert!(reader.get_bits_64(1) == 1);
            assert!(reader.get_bits_64(2) == 2);
            assert!(reader.get_bits_64(4) == 10);
            assert!(reader.get_bits_64(1) == 0);
            assert!(reader.get_bits_64(8) == 85);
        }

        #[test]
        fn peek_bits_64() {
            let reader = BitReadLE {
                buffer: &CHECKBOARD0101,
                index: 0,
                cache: 0,
                left: 0
            };

            assert!(reader.peek_bits_64(1) == 1);
            assert!(reader.peek_bits_64(1) == 1);
            assert!(reader.peek_bits_64(2) == 1);
            assert!(reader.peek_bits_64(2) == 1);
        }

        #[test]
        fn get_bits_32() {
            let mut reader = BitReadLE {
                buffer: &CHECKBOARD0101,
                index: 0,
                cache: 0,
                left: 0
            };

            assert!(reader.get_bits_32(1) == 1);
            assert!(reader.get_bits_32(2) == 2);
            assert!(reader.get_bits_32(4) == 10);
            assert!(reader.get_bits_32(1) == 0);
            assert!(reader.get_bits_32(8) == 85);
        }

        #[test]
        fn peek_bits_32() {
            let b = &CHECKBOARD0101;
            let mut reader = BitReadLE::new(b);

            assert!(reader.peek_bits_32(1) == 1);
            assert!(reader.peek_bits_32(1) == 1);
            assert!(reader.peek_bits_32(2) == 1);
            assert!(reader.peek_bits_32(2) == 1);
        }

        #[test]
        fn skip_bits() {
            let b = &CHECKBOARD0101;
            let mut reader = BitReadLE::new(b);

            reader.skip_bits(0);
            assert!(reader.peek_bits_32(1) == 1);
            reader.skip_bits(2);
            assert!(reader.peek_bits_32(1) == 1);
            reader.skip_bits(2);
            assert!(reader.peek_bits_32(1) == 1);
        }

        #[test]
        fn align_bits() {
            let b = &CHECKBOARD0011;
            let mut reader = BitReadLE::new(b);

            reader.align_bits();
            assert!(reader.get_bits_64(3) == 3);
            reader.align_bits();
            assert!(reader.get_bits_64(4) == 3);
            reader.skip_bits(1);
            reader.align_bits();
            assert!(reader.get_bits_64(4) == 3);
        }

        #[test]
        fn overread() {
            let b = &CHECKBOARD0011;
            let mut reader = BitReadLE::new(b);

            reader.skip_bits(128*8+2);
            reader.get_bits_64(6);
        }
    }
    mod be {
        use super::super::*;
        use super::*;

        #[test]
        fn get_bit() {
            let b = &CHECKBOARD0101;
            let mut reader = BitReadBE::new(b);

            assert!(!reader.get_bit());
            assert!(reader.get_bit());
        }

        #[test]
        fn get_bits_64() {
            let b = &CHECKBOARD0101;
            let mut reader = BitReadBE::new(b);

            assert!(reader.get_bits_64(1) == 0);
            assert!(reader.get_bits_64(2) == 2);
            assert!(reader.get_bits_64(4) == 10);
            assert!(reader.get_bits_64(1) == 1);
            assert!(reader.get_bits_64(8) == 85);
        }

        #[test]
        fn peek_bits_64() {
            let b = &CHECKBOARD0101;
            let reader = BitReadBE::new(b);

            assert!(reader.peek_bits_64(1) == 0);
            assert!(reader.peek_bits_64(1) == 0);
            assert!(reader.peek_bits_64(2) == 1);
            assert!(reader.peek_bits_64(2) == 1);
        }

        #[test]
        fn get_bits_32() {
            let b = &CHECKBOARD0101;
            let mut reader = BitReadBE::new(b);

            assert!(reader.get_bits_32(1) == 0);
            assert!(reader.get_bits_32(2) == 2);
            assert!(reader.get_bits_32(4) == 10);
            assert!(reader.get_bits_32(1) == 1);
            assert!(reader.get_bits_32(8) == 85);
        }

        #[test]
        fn peek_bits_32() {
            let b = &CHECKBOARD0101;
            let mut reader = BitReadBE::new(b);

            assert!(reader.peek_bits_32(1) == 0);
            assert!(reader.peek_bits_32(1) == 0);
            assert!(reader.peek_bits_32(2) == 1);
            assert!(reader.peek_bits_32(2) == 1);
        }

        #[test]
        fn skip_bits() {
            let b = &CHECKBOARD0101;
            let mut reader = BitReadBE::new(b);

            reader.skip_bits(0);
            assert!(reader.peek_bits_32(1) == 0);
            reader.skip_bits(2);
            assert!(reader.peek_bits_32(1) == 0);
            reader.skip_bits(2);
            assert!(reader.peek_bits_32(1) == 0);
        }

        #[test]
        fn align_bits() {
            let b = &CHECKBOARD0011;
            let mut reader = BitReadBE::new(b);

            reader.align_bits();
            assert!(reader.get_bits_64(3) == 1);
            reader.align_bits();
            assert!(reader.get_bits_64(4) == 3);
            reader.skip_bits(1);
            reader.align_bits();
            assert!(reader.get_bits_64(4) == 3);
        }

        #[test]
        fn overread() {
            let b = &CHECKBOARD0011;
            let mut reader = BitReadBE::new(b);

            reader.skip_bits(128*8+2);
            reader.get_bits_64(6);
        }
    }
}
<|endoftext|>


<commit_before>extern crate futures;
#[macro_use]
extern crate log;
extern crate tokio_io;
extern crate tokio_process;

use std::io;
use std::process::{Stdio, ExitStatus, Command};

use futures::future::Future;
use futures::stream::{self, Stream};
use tokio_io::io::{read_until, write_all, read_to_end};
use tokio_process::{CommandExt, Child};

mod support;

fn cat() -> Command {
    let mut cmd = support::cmd("cat");
    cmd.stdin(Stdio::piped())
       .stdout(Stdio::piped());
    cmd
}

fn feed_cat(mut cat: Child, n: usize) -> Box<Future<Item = ExitStatus, Error = io::Error>> {
    let stdin = cat.stdin().take().unwrap();
    let stdout = cat.stdout().take().unwrap();

    debug!("starting to feed");
    // Produce n lines on the child's stdout.
    let numbers = stream::iter_ok(0..n);
    let write = numbers.fold(stdin, |stdin, i| {
        debug!("sending line {} to child", i);
        write_all(stdin, format!("line {}\n", i).into_bytes()).map(|p| p.0)
    }).map(|_| ());

    // Try to read `n + 1` lines, ensuring the last one is empty
    // (i.e. EOF is reached after `n` lines.
    let reader = io::BufReader::new(stdout);
    let expected_numbers = stream::iter_ok(0..=n);
    let read = expected_numbers.fold((reader, 0), move |(reader, i), _| {
        let done = i >= n;
        debug!("starting read from child");
        read_until(reader, b'\n', Vec::new()).and_then(move |(reader, vec)| {
            debug!("read line {} from child ({} bytes, done: {})",
                   i, vec.len(), done);
            match (done, vec.len()) {
                (false, 0) => {
                    Err(io::Error::new(io::ErrorKind::BrokenPipe, "broken pipe"))
                },
                (true, n) if n != 0 => {
                    Err(io::Error::new(io::ErrorKind::Other, "extraneous data"))
                },
                _ => {
                    let s = std::str::from_utf8(&vec).unwrap();
                    let expected = format!("line {}\n", i);
                    if done || s == expected {
                        Ok((reader, i + 1))
                    } else {
                        Err(io::Error::new(io::ErrorKind::Other, "unexpected data"))
                    }
                }
            }
        })
    });

    // Compose reading and writing concurrently.
    Box::new(write.join(read).and_then(|_| cat))
}

/// Check for the following properties when feeding stdin and
/// consuming stdout of a cat-like process:
///
/// - A number of lines that amounts to a number of bytes exceeding a
///   typical OS buffer size can be fed to the child without
///   deadlock. This tests that we also consume the stdout
///   concurrently; otherwise this would deadlock.
///
/// - We read the same lines from the child that we fed it.
///
/// - The child does produce EOF on stdout after the last line.
#[test]
fn feed_a_lot() {
    let child = cat().spawn_async().unwrap();
    let status = support::run_with_timeout(feed_cat(child, 10000)).unwrap();
    assert_eq!(status.code(), Some(0));
}

// FIXME: delete this test once we have a resolution for #51
// This test's setup is flaky, and setting up a consistent test is nearly
// impossible: right now we invoke `cat` and immediately kill it, expecting
// that it didn't write anything, but if there's something wrong with the
// command itself (e.g. redirection issues, it doesn't actually print anything
// out, etc.) this test can falsely pass. Attempting a solution which writes
// some data, *then* kill the child, write more data, and assert that only the
// first write is echoed back seems like a good approach, however, due to the
// ordering of context switches or how the kernel buffers data we can get
// inconsistent results. We can keep this test around for now, but as soon as
// we have a solution for #51, we may have a better avenue for testing this
// functionality.
#[test]
fn drop_kills() {
    let mut child = cat().spawn_async().unwrap();
    let stdin = child.stdin().take().unwrap();
    let stdout = child.stdout().take().unwrap();
    drop(child);

    // Ignore all write errors since we expect a broken pipe here
    let writer = write_all(stdin, b"1234").then(|_| Ok(()));
    let reader = read_to_end(stdout, Vec::new());

    let (_, output) = support::CurrentThreadRuntime::new()
        .expect("failed to get rt")
        .spawn(writer)
        .block_on(support::with_timeout(reader))
        .expect("failed to get output");

    assert_eq!(output.len(), 0);
}

#[test]
fn wait_with_output_captures() {
    let mut child = cat().spawn_async().unwrap();
    let stdin = child.stdin().take().unwrap();
    let out = child.wait_with_output();

    let future = write_all(stdin, b"1234").map(|p| p.1).join(out);
    let ret = support::run_with_timeout(future).unwrap();
    let (written, output) = ret;

    assert!(output.status.success());
    assert_eq!(output.stdout, written);
    assert_eq!(output.stderr.len(), 0);
}

#[test]
fn status_closes_any_pipes() {
    // Cat will open a pipe between the parent and child.
    // If `status_async` doesn't ensure the handles are closed,
    // we would end up blocking forever (and time out).
    let child = cat().status_async().expect("failed to spawn child");

    support::run_with_timeout(child)
        .expect("time out exceeded! did we get stuck waiting on the child?");
}
<commit_msg>process: Delete flaky and (now) unused test<commit_after>extern crate futures;
#[macro_use]
extern crate log;
extern crate tokio_io;
extern crate tokio_process;

use std::io;
use std::process::{Stdio, ExitStatus, Command};

use futures::future::Future;
use futures::stream::{self, Stream};
use tokio_io::io::{read_until, write_all};
use tokio_process::{CommandExt, Child};

mod support;

fn cat() -> Command {
    let mut cmd = support::cmd("cat");
    cmd.stdin(Stdio::piped())
       .stdout(Stdio::piped());
    cmd
}

fn feed_cat(mut cat: Child, n: usize) -> Box<Future<Item = ExitStatus, Error = io::Error>> {
    let stdin = cat.stdin().take().unwrap();
    let stdout = cat.stdout().take().unwrap();

    debug!("starting to feed");
    // Produce n lines on the child's stdout.
    let numbers = stream::iter_ok(0..n);
    let write = numbers.fold(stdin, |stdin, i| {
        debug!("sending line {} to child", i);
        write_all(stdin, format!("line {}\n", i).into_bytes()).map(|p| p.0)
    }).map(|_| ());

    // Try to read `n + 1` lines, ensuring the last one is empty
    // (i.e. EOF is reached after `n` lines.
    let reader = io::BufReader::new(stdout);
    let expected_numbers = stream::iter_ok(0..=n);
    let read = expected_numbers.fold((reader, 0), move |(reader, i), _| {
        let done = i >= n;
        debug!("starting read from child");
        read_until(reader, b'\n', Vec::new()).and_then(move |(reader, vec)| {
            debug!("read line {} from child ({} bytes, done: {})",
                   i, vec.len(), done);
            match (done, vec.len()) {
                (false, 0) => {
                    Err(io::Error::new(io::ErrorKind::BrokenPipe, "broken pipe"))
                },
                (true, n) if n != 0 => {
                    Err(io::Error::new(io::ErrorKind::Other, "extraneous data"))
                },
                _ => {
                    let s = std::str::from_utf8(&vec).unwrap();
                    let expected = format!("line {}\n", i);
                    if done || s == expected {
                        Ok((reader, i + 1))
                    } else {
                        Err(io::Error::new(io::ErrorKind::Other, "unexpected data"))
                    }
                }
            }
        })
    });

    // Compose reading and writing concurrently.
    Box::new(write.join(read).and_then(|_| cat))
}

/// Check for the following properties when feeding stdin and
/// consuming stdout of a cat-like process:
///
/// - A number of lines that amounts to a number of bytes exceeding a
///   typical OS buffer size can be fed to the child without
///   deadlock. This tests that we also consume the stdout
///   concurrently; otherwise this would deadlock.
///
/// - We read the same lines from the child that we fed it.
///
/// - The child does produce EOF on stdout after the last line.
#[test]
fn feed_a_lot() {
    let child = cat().spawn_async().unwrap();
    let status = support::run_with_timeout(feed_cat(child, 10000)).unwrap();
    assert_eq!(status.code(), Some(0));
}

#[test]
fn wait_with_output_captures() {
    let mut child = cat().spawn_async().unwrap();
    let stdin = child.stdin().take().unwrap();
    let out = child.wait_with_output();

    let future = write_all(stdin, b"1234").map(|p| p.1).join(out);
    let ret = support::run_with_timeout(future).unwrap();
    let (written, output) = ret;

    assert!(output.status.success());
    assert_eq!(output.stdout, written);
    assert_eq!(output.stderr.len(), 0);
}

#[test]
fn status_closes_any_pipes() {
    // Cat will open a pipe between the parent and child.
    // If `status_async` doesn't ensure the handles are closed,
    // we would end up blocking forever (and time out).
    let child = cat().status_async().expect("failed to spawn child");

    support::run_with_timeout(child)
        .expect("time out exceeded! did we get stuck waiting on the child?");
}
<|endoftext|>