<commit_before>// Copyright(C) 2016 Chris Liebert

extern crate glium;
extern crate nalgebra;
extern crate quick3d;

use glium::glutin;
use glium::glutin::{ElementState, Event, MouseButton, VirtualKeyCode};
use glium::DisplayBuild;

use nalgebra::Matrix4;
use std::io::Error;

use quick3d::camera::Camera;
use quick3d::common::Mesh;
use quick3d::renderer;
use quick3d::shader::Shader;
use quick3d::scene::Scene;

fn main() {
    let screen_width = 600;
    let screen_height = 400;

    let bin_filename = String::from("test.bin.gz");
    let scene: Scene = match Scene::from_compressed_binary_file(bin_filename.clone()) {
        Ok(s) => s,
        Err(e) => panic!("Unable to load {}: {:?}", bin_filename, e),
    };

    let display = glutin::WindowBuilder::new()
        //.resizable()
        //.with_vsync()
        .with_depth_buffer(24)
        .with_title("Rust Window")
        .with_gl_debug_flag(true)
        .with_visibility(false) // Window is shown when scene finishes loading.
        .with_dimensions(screen_width, screen_height)
        .build_glium()
        .unwrap();

    let mut camera: Camera = Camera::new(screen_width as f32, screen_height as f32);
    camera = camera.move_backward(6.0);

    let renderer: renderer::Renderer = match renderer::Renderer::new(&display, scene) {
        Ok(r) => r,
        Err(e) => panic!("Unable to create renderer: {:?}", e),
    };

    let shader_program: glium::program::Program = match Shader::default(&display) {
        Ok(p) => p,
        Err(e) => {
            panic!("Unable to load default shader {:?}", e);
        },
    };

    // Show the window once the data is loaded
    let window = match display.get_window() {
        Some(x) => {
            x.show();
            x
        }
        None => {
            panic!("Error retrieving window");
        }
    };

    // The torus will be movable in the scene if it is found
    let torus: Result<&Mesh, Error> = renderer.get_mesh("Torus");

    let mut torus_x = 0.0f32;
    let mut torus_y = 0.0f32;
    let mut torus_vertical_speed = 0.0f32;
    let mut torus_horizontal_speed = 0.0f32;
    let mut left_button_pressed = false;

    let mut mouse_last_x: i32 = 0;
    let mut mouse_last_y: i32 = 0;

    let mut _mouse_dx: i32 = 0;
    let mut _mouse_dy: i32 = 0;

    let mut _camera_forward_speed = 0.0f32;
    let mut _camera_left_speed = 0.0f32;

    let mut _w_state = false;
    let mut _camera_forward_speed = 0.0;

    let screen_center_x: i32 = (screen_width / 2) as i32;
    let screen_center_y: i32 = (screen_height / 2) as i32;

    let mouse_grab_margin: i32 = screen_center_y / 2;

    'running: loop {
        match renderer.render(&display, &shader_program, &camera) {
            Ok(_) => (),
            Err(e) => panic!("{:?}", e),
        };

        for event in display.poll_events() {
            match event {
                Event::Closed => break 'running,
                Event::KeyboardInput(ElementState::Pressed, _, Some(VirtualKeyCode::Escape)) => {
                    break 'running;
                }
                Event::KeyboardInput(ElementState::Pressed, _, Some(VirtualKeyCode::W)) => {
                    _camera_forward_speed = 1.0;
                }
                Event::KeyboardInput(ElementState::Released, _, Some(VirtualKeyCode::W)) => {
                    _camera_forward_speed = 0.0;
                }
                Event::KeyboardInput(ElementState::Pressed, _, Some(VirtualKeyCode::A)) => {
                    camera = camera.move_left(1.0);
                }
                Event::KeyboardInput(ElementState::Pressed, _, Some(VirtualKeyCode::S)) => {
                    camera = camera.move_backward(1.0);
                }
                Event::KeyboardInput(ElementState::Pressed, _, Some(VirtualKeyCode::D)) => {
                    camera = camera.move_right(1.0);
                }
                Event::KeyboardInput(ElementState::Pressed, _, Some(VirtualKeyCode::I)) => {
                    torus_vertical_speed = 0.001f32;
                }
                Event::KeyboardInput(ElementState::Released, _, Some(VirtualKeyCode::I)) => {
                    torus_vertical_speed = 0.0f32;
                }
                Event::KeyboardInput(ElementState::Pressed, _, Some(VirtualKeyCode::J)) => {
                    torus_horizontal_speed = -0.001f32;
                }
                Event::KeyboardInput(ElementState::Released, _, Some(VirtualKeyCode::J)) => {
                    torus_horizontal_speed = 0.0f32;
                }
                Event::KeyboardInput(ElementState::Pressed, _, Some(VirtualKeyCode::K)) => {
                    torus_vertical_speed = -0.001f32;
                }
                Event::KeyboardInput(ElementState::Released, _, Some(VirtualKeyCode::K)) => {
                    torus_vertical_speed = 0.0f32;
                }
                Event::KeyboardInput(ElementState::Pressed, _, Some(VirtualKeyCode::L)) => {
                    torus_horizontal_speed = 0.001f32;
                }
                Event::KeyboardInput(ElementState::Released, _, Some(VirtualKeyCode::L)) => {
                    torus_horizontal_speed = 0.0f32;
                }
                Event::MouseInput(ElementState::Pressed, MouseButton::Left) => {
                    left_button_pressed = true;
                }
                Event::MouseInput(ElementState::Released, MouseButton::Left) => {
                    left_button_pressed = false;
                }
                Event::MouseMoved(x, y) => {
                    _mouse_dx = mouse_last_x - x;
                    _mouse_dy = mouse_last_y - y;
                    mouse_last_x = x;
                    mouse_last_y = y;
                    if left_button_pressed {
                        // Rotate the camera if the left mouse button is pressed
                        camera = camera.aim(_mouse_dx as f64, _mouse_dy as f64);

                        if x + mouse_grab_margin >= screen_width as i32 || x <= mouse_grab_margin {
                            let _ = window.set_cursor_position(screen_center_x, y);
                            mouse_last_x = screen_center_x;
                            _mouse_dx = 0;
                        } else if y + mouse_grab_margin >= screen_height as i32 ||
                                  y <= mouse_grab_margin {
                            let _ = window.set_cursor_position(x, screen_center_y);
                            mouse_last_y = screen_center_y;
                            _mouse_dy = 0;
                        }
                    }
                }

                _ => (),
            }
        }

        camera = camera.move_forward(_camera_forward_speed * 0.01);

        // Move the torus (if found) based on changes from keyboard input
        match torus {
            Ok(torus) => {
                // Get existing matrix
                let mut matrix: Matrix4<f32> = *torus.matrix.borrow();
                torus_x += torus_horizontal_speed;
                torus_y += torus_vertical_speed;
                matrix.m14 = torus_x;
                matrix.m24 = torus_y;
                // Mutate the matrix
                *torus.matrix.borrow_mut() = matrix;
            }
            Err(_) => {}
        }
    }
}
<commit_msg>added conditional loading support for Rust example so .bin.gz, .bin, and .db files can be specified as an argument.<commit_after>// Copyright(C) 2016 Chris Liebert

extern crate glium;
extern crate nalgebra;
extern crate quick3d;

use glium::glutin;
use glium::glutin::{ElementState, Event, MouseButton, VirtualKeyCode};
use glium::DisplayBuild;

use nalgebra::Matrix4;
use std::io::Error;

use quick3d::camera::Camera;
use quick3d::common::Mesh;
use quick3d::renderer;
use quick3d::shader::Shader;
use quick3d::scene::Scene;

fn load_scene_no_sqlite(filename: String) -> Scene {
    let scene: Scene = match filename.ends_with(".bin.gz") {
        true => {
            match Scene::from_compressed_binary_file(filename.clone()) {
                Ok(s) => s,
                Err(e) => panic!("Unable to load {}: {:?}", filename, e),
            }
        },
        false => {
            match filename.ends_with(".bin") {
                true => {
                    match Scene::from_binary_file(filename.clone()) {
                        Ok(s) => s,
                        Err(e) => panic!("Unable to load {}: {:?}", filename, e),
                    }
                },
                false => {
                    panic!("Error: {} must end in .bin or .bin.gz. If the sqlite feature is enabled, .db is also supported.", filename);
                },
            }
        },
    };

    scene
}

#[cfg(feature = "sqlite")]
fn load_scene(filename: String) -> Scene {

    // If there is an filename argument ending in ".db", and the sqlite feature is enabled, use .db file
    let scene: Scene = match filename.ends_with(".db") {
        true => {
            use quick3d::dbloader::DBLoader;

            let db: DBLoader = match DBLoader::new(&filename){
                Ok(d) => d,
                Err(e) => panic!("Error: {:?}", e),
            };
            match db.load_scene() {
                Ok(s) => s,
                Err(e) => panic!("Unable to load {:?}: {:?}", filename, e),
            }
        },
        false => {
            load_scene_no_sqlite(filename)
        },
    };

    scene
}

#[cfg(not(feature = "sqlite"))]
fn load_scene(filename: String) -> Scene {
    load_scene_no_sqlite(filename)
}

fn main() {
    let screen_width = 600;
    let screen_height = 400;

    use std::env;
    let mut filename = String::from("test.bin.gz");

    // Prints each argument on a separate line
    let mut i = 0;
    for argument in env::args() {
        if i == 1 {
            filename = argument;
        }
        i += 1;
    }

    let scene: Scene = load_scene(filename);

    let display = glutin::WindowBuilder::new()
        //.resizable()
        //.with_vsync()
        .with_depth_buffer(24)
        .with_title("Rust Window")
        .with_gl_debug_flag(true)
        .with_visibility(false) // Window is shown when scene finishes loading.
        .with_dimensions(screen_width, screen_height)
        .build_glium()
        .unwrap();

    let mut camera: Camera = Camera::new(screen_width as f32, screen_height as f32);
    camera = camera.move_backward(6.0);

    let renderer: renderer::Renderer = match renderer::Renderer::new(&display, scene) {
        Ok(r) => r,
        Err(e) => panic!("Unable to create renderer: {:?}", e),
    };

    let shader_program: glium::program::Program = match Shader::default(&display) {
        Ok(p) => p,
        Err(e) => {
            panic!("Unable to load default shader {:?}", e);
        },
    };

    // Show the window once the data is loaded
    let window = match display.get_window() {
        Some(x) => {
            x.show();
            x
        }
        None => {
            panic!("Error retrieving window");
        }
    };

    // The torus will be movable in the scene if it is found
    let torus: Result<&Mesh, Error> = renderer.get_mesh("Torus");

    let mut torus_x = 0.0f32;
    let mut torus_y = 0.0f32;
    let mut torus_vertical_speed = 0.0f32;
    let mut torus_horizontal_speed = 0.0f32;
    let mut left_button_pressed = false;

    let mut mouse_last_x: i32 = 0;
    let mut mouse_last_y: i32 = 0;

    let mut _mouse_dx: i32 = 0;
    let mut _mouse_dy: i32 = 0;

    let mut _camera_forward_speed = 0.0f32;
    let mut _camera_left_speed = 0.0f32;

    let mut _w_state = false;
    let mut _camera_forward_speed = 0.0;

    let screen_center_x: i32 = (screen_width / 2) as i32;
    let screen_center_y: i32 = (screen_height / 2) as i32;

    let mouse_grab_margin: i32 = screen_center_y / 2;

    'running: loop {
        match renderer.render(&display, &shader_program, &camera) {
            Ok(_) => (),
            Err(e) => panic!("{:?}", e),
        };

        for event in display.poll_events() {
            match event {
                Event::Closed => break 'running,
                Event::KeyboardInput(ElementState::Pressed, _, Some(VirtualKeyCode::Escape)) => {
                    break 'running;
                }
                Event::KeyboardInput(ElementState::Pressed, _, Some(VirtualKeyCode::W)) => {
                    _camera_forward_speed = 1.0;
                }
                Event::KeyboardInput(ElementState::Released, _, Some(VirtualKeyCode::W)) => {
                    _camera_forward_speed = 0.0;
                }
                Event::KeyboardInput(ElementState::Pressed, _, Some(VirtualKeyCode::A)) => {
                    camera = camera.move_left(1.0);
                }
                Event::KeyboardInput(ElementState::Pressed, _, Some(VirtualKeyCode::S)) => {
                    camera = camera.move_backward(1.0);
                }
                Event::KeyboardInput(ElementState::Pressed, _, Some(VirtualKeyCode::D)) => {
                    camera = camera.move_right(1.0);
                }
                Event::KeyboardInput(ElementState::Pressed, _, Some(VirtualKeyCode::I)) => {
                    torus_vertical_speed = 0.001f32;
                }
                Event::KeyboardInput(ElementState::Released, _, Some(VirtualKeyCode::I)) => {
                    torus_vertical_speed = 0.0f32;
                }
                Event::KeyboardInput(ElementState::Pressed, _, Some(VirtualKeyCode::J)) => {
                    torus_horizontal_speed = -0.001f32;
                }
                Event::KeyboardInput(ElementState::Released, _, Some(VirtualKeyCode::J)) => {
                    torus_horizontal_speed = 0.0f32;
                }
                Event::KeyboardInput(ElementState::Pressed, _, Some(VirtualKeyCode::K)) => {
                    torus_vertical_speed = -0.001f32;
                }
                Event::KeyboardInput(ElementState::Released, _, Some(VirtualKeyCode::K)) => {
                    torus_vertical_speed = 0.0f32;
                }
                Event::KeyboardInput(ElementState::Pressed, _, Some(VirtualKeyCode::L)) => {
                    torus_horizontal_speed = 0.001f32;
                }
                Event::KeyboardInput(ElementState::Released, _, Some(VirtualKeyCode::L)) => {
                    torus_horizontal_speed = 0.0f32;
                }
                Event::MouseInput(ElementState::Pressed, MouseButton::Left) => {
                    left_button_pressed = true;
                }
                Event::MouseInput(ElementState::Released, MouseButton::Left) => {
                    left_button_pressed = false;
                }
                Event::MouseMoved(x, y) => {
                    _mouse_dx = mouse_last_x - x;
                    _mouse_dy = mouse_last_y - y;
                    mouse_last_x = x;
                    mouse_last_y = y;
                    if left_button_pressed {
                        // Rotate the camera if the left mouse button is pressed
                        camera = camera.aim(_mouse_dx as f64, _mouse_dy as f64);

                        if x + mouse_grab_margin >= screen_width as i32 || x <= mouse_grab_margin {
                            let _ = window.set_cursor_position(screen_center_x, y);
                            mouse_last_x = screen_center_x;
                            _mouse_dx = 0;
                        } else if y + mouse_grab_margin >= screen_height as i32 ||
                                  y <= mouse_grab_margin {
                            let _ = window.set_cursor_position(x, screen_center_y);
                            mouse_last_y = screen_center_y;
                            _mouse_dy = 0;
                        }
                    }
                }

                _ => (),
            }
        }

        camera = camera.move_forward(_camera_forward_speed * 0.01);

        // Move the torus (if found) based on changes from keyboard input
        match torus {
            Ok(torus) => {
                // Get existing matrix
                let mut matrix: Matrix4<f32> = *torus.matrix.borrow();
                torus_x += torus_horizontal_speed;
                torus_y += torus_vertical_speed;
                matrix.m14 = torus_x;
                matrix.m24 = torus_y;
                // Mutate the matrix
                *torus.matrix.borrow_mut() = matrix;
            }
            Err(_) => {}
        }
    }
}
<|endoftext|>


<commit_before>extern crate libc;

use array::Array;
use defines::MatchType;
use self::libc::{c_void, uint8_t, c_uint, c_int, c_float, c_longlong};

type MutAfArray = *mut self::libc::c_longlong;
type AfArray    = self::libc::c_longlong;
type DimT       = self::libc::c_longlong;
type MutFeat    = *mut *mut self::libc::c_void;
type Feat       = *const self::libc::c_void;

#[allow(dead_code)]
extern {
    fn af_create_features(feat: MutFeat, num: DimT);
    fn af_retain_features(feat: MutFeat, feat: Feat);
    fn af_get_features_num(num: *mut DimT, feat: Feat);
    fn af_get_features_xpos(out: MutAfArray, feat: Feat);
    fn af_get_features_ypos(out: MutAfArray, feat: Feat);
    fn af_get_features_score(out: MutAfArray, feat: Feat);
    fn af_get_features_orientation(out: MutAfArray, feat: Feat);
    fn af_get_features_size(out: MutAfArray, feat: Feat);
    fn af_release_features(feat: *mut c_void);

    fn af_fast(out: MutFeat, input: AfArray, thr: c_float, arc_len: c_uint, non_max: c_int,
               feature_ratio: c_float, edge: c_uint);

    fn af_orb(out: MutFeat, desc: MutAfArray, arr: AfArray, fast_thr: c_float, max_feat: c_uint,
              scl_fctr: c_float, levels: c_uint, blur_img: c_int);

    fn af_hamming_matcher(idx: MutAfArray, dist: MutAfArray,
                          query: AfArray, train: AfArray,
                          dist_dim: DimT, n_dist: c_uint);

    fn af_match_template(out: MutAfArray, search_img: AfArray, template_img: AfArray,
                         mtype: uint8_t);
}

pub struct Features {
    feat: i64,
}

macro_rules! feat_func_def {
    ($fn_name: ident, $ffi_name: ident) => (
        pub fn $fn_name(&self) -> Array {
            unsafe {
                let mut temp: i64 = 0;
                $ffi_name(&mut temp as MutAfArray, self.feat as Feat);
                Array::from(temp)
            }
        }
    )
}

impl Features {
    #[allow(unused_mut)]
    pub fn new(n: u64) -> Features {
        unsafe {
            let mut temp: i64 = 0;
            af_create_features(&mut temp as *mut c_longlong as MutFeat,
                               n as DimT);
            Features {feat: temp}
        }
    }

    pub fn num_features(&self) -> i64 {
        unsafe {
            let mut temp: i64 = 0;
            af_get_features_num(&mut temp as *mut DimT,
                                self.feat as *const c_longlong as Feat);
            temp
        }
    }

    feat_func_def!(xpos, af_get_features_xpos);
    feat_func_def!(ypos, af_get_features_ypos);
    feat_func_def!(score, af_get_features_score);
    feat_func_def!(orientation, af_get_features_orientation);
    feat_func_def!(size, af_get_features_size);

    pub fn get(&self) -> i64 {
        self.feat
    }
}

impl Clone for Features {
    fn clone(&self) -> Features {
        unsafe {
            let mut temp: i64 = 0;
            af_retain_features(&mut temp as *mut c_longlong as MutFeat,
                               self.feat as *const c_longlong as Feat);
            Features {feat: temp}
        }
    }
}

impl Drop for Features {
    fn drop(&mut self) {
        unsafe {
            af_release_features(self.feat as *mut c_longlong as *mut c_void);
        }
    }
}

#[allow(unused_mut)]
pub fn fast(input: &Array, thr: f32, arc_len: u32,
            non_max: bool, feat_ratio: f32, edge: u32) -> Features {
    unsafe {
        let mut temp: i64 = 0;
        af_fast(&mut temp as *mut c_longlong as MutFeat,
                input.get() as AfArray, thr as c_float, arc_len as c_uint,
                non_max as c_int, feat_ratio as c_float, edge as c_uint);
        Features {feat: temp}
    }
}

#[allow(unused_mut)]
pub fn orb(input: &Array, fast_thr: f32, max_feat: u32,
           scl_fctr: f32, levels: u32, blur_img: bool) -> (Features, Array) {
    unsafe {
        let mut f: i64 = 0;
        let mut d: i64 = 0;
        af_orb(&mut f as *mut c_longlong as MutFeat, &mut d as MutAfArray,
               input.get() as AfArray, fast_thr as c_float,
               max_feat as c_uint, scl_fctr as c_float, levels as c_uint, blur_img as c_int);
        (Features {feat: f}, Array::from(d))
    }
}

#[allow(unused_mut)]
pub fn hamming_matcher(query: &Array, train: &Array,
                       dist_dims: i64, n_dist: u32) -> (Array, Array) {
    unsafe {
        let mut idx: i64 = 0;
        let mut dist:i64 = 0;
        af_hamming_matcher(&mut idx as MutAfArray, &mut dist as MutAfArray,
                           query.get() as AfArray, train.get() as AfArray,
                           dist_dims as DimT, n_dist as c_uint);
        (Array::from(idx), Array::from(dist))
    }
}

#[allow(unused_mut)]
pub fn match_template(search_img: &Array, template_img: &Array, mtype: MatchType) -> Array {
    unsafe {
        let mut temp: i64 = 0;
        af_match_template(&mut temp as MutAfArray,
                          search_img.get() as AfArray, template_img.get() as AfArray,
                          mtype as uint8_t);
        Array::from(temp)
    }
}
<commit_msg>Corrected ffi signatures vison header functions<commit_after>extern crate libc;

use array::Array;
use defines::MatchType;
use self::libc::{c_void, uint8_t, c_uint, c_int, c_float, c_longlong};

type MutAfArray = *mut self::libc::c_longlong;
type AfArray    = self::libc::c_longlong;
type DimT       = self::libc::c_longlong;
type MutFeat    = *mut *mut self::libc::c_void;
type Feat       = *const self::libc::c_void;

#[allow(dead_code)]
extern {
    fn af_create_features(feat: MutFeat, num: DimT) -> c_int;
    fn af_retain_features(feat: MutFeat, feat: Feat) -> c_int;
    fn af_get_features_num(num: *mut DimT, feat: Feat) -> c_int;
    fn af_get_features_xpos(out: MutAfArray, feat: Feat) -> c_int;
    fn af_get_features_ypos(out: MutAfArray, feat: Feat) -> c_int;
    fn af_get_features_score(out: MutAfArray, feat: Feat) -> c_int;
    fn af_get_features_orientation(out: MutAfArray, feat: Feat) -> c_int;
    fn af_get_features_size(out: MutAfArray, feat: Feat) -> c_int;
    fn af_release_features(feat: *mut c_void) -> c_int;

    fn af_fast(out: MutFeat, input: AfArray, thr: c_float, arc_len: c_uint, non_max: c_int,
               feature_ratio: c_float, edge: c_uint) -> c_int;

    fn af_orb(out: MutFeat, desc: MutAfArray, arr: AfArray, fast_thr: c_float, max_feat: c_uint,
              scl_fctr: c_float, levels: c_uint, blur_img: c_int) -> c_int;

    fn af_hamming_matcher(idx: MutAfArray, dist: MutAfArray,
                          query: AfArray, train: AfArray,
                          dist_dim: DimT, n_dist: c_uint) -> c_int;

    fn af_match_template(out: MutAfArray, search_img: AfArray, template_img: AfArray,
                         mtype: uint8_t) -> c_int;
}

pub struct Features {
    feat: i64,
}

macro_rules! feat_func_def {
    ($fn_name: ident, $ffi_name: ident) => (
        pub fn $fn_name(&self) -> Array {
            unsafe {
                let mut temp: i64 = 0;
                $ffi_name(&mut temp as MutAfArray, self.feat as Feat);
                Array::from(temp)
            }
        }
    )
}

impl Features {
    #[allow(unused_mut)]
    pub fn new(n: u64) -> Features {
        unsafe {
            let mut temp: i64 = 0;
            af_create_features(&mut temp as *mut c_longlong as MutFeat,
                               n as DimT);
            Features {feat: temp}
        }
    }

    pub fn num_features(&self) -> i64 {
        unsafe {
            let mut temp: i64 = 0;
            af_get_features_num(&mut temp as *mut DimT,
                                self.feat as *const c_longlong as Feat);
            temp
        }
    }

    feat_func_def!(xpos, af_get_features_xpos);
    feat_func_def!(ypos, af_get_features_ypos);
    feat_func_def!(score, af_get_features_score);
    feat_func_def!(orientation, af_get_features_orientation);
    feat_func_def!(size, af_get_features_size);

    pub fn get(&self) -> i64 {
        self.feat
    }
}

impl Clone for Features {
    fn clone(&self) -> Features {
        unsafe {
            let mut temp: i64 = 0;
            af_retain_features(&mut temp as *mut c_longlong as MutFeat,
                               self.feat as *const c_longlong as Feat);
            Features {feat: temp}
        }
    }
}

impl Drop for Features {
    fn drop(&mut self) {
        unsafe {
            af_release_features(self.feat as *mut c_longlong as *mut c_void);
        }
    }
}

#[allow(unused_mut)]
pub fn fast(input: &Array, thr: f32, arc_len: u32,
            non_max: bool, feat_ratio: f32, edge: u32) -> Features {
    unsafe {
        let mut temp: i64 = 0;
        af_fast(&mut temp as *mut c_longlong as MutFeat,
                input.get() as AfArray, thr as c_float, arc_len as c_uint,
                non_max as c_int, feat_ratio as c_float, edge as c_uint);
        Features {feat: temp}
    }
}

#[allow(unused_mut)]
pub fn orb(input: &Array, fast_thr: f32, max_feat: u32,
           scl_fctr: f32, levels: u32, blur_img: bool) -> (Features, Array) {
    unsafe {
        let mut f: i64 = 0;
        let mut d: i64 = 0;
        af_orb(&mut f as *mut c_longlong as MutFeat, &mut d as MutAfArray,
               input.get() as AfArray, fast_thr as c_float,
               max_feat as c_uint, scl_fctr as c_float, levels as c_uint, blur_img as c_int);
        (Features {feat: f}, Array::from(d))
    }
}

#[allow(unused_mut)]
pub fn hamming_matcher(query: &Array, train: &Array,
                       dist_dims: i64, n_dist: u32) -> (Array, Array) {
    unsafe {
        let mut idx: i64 = 0;
        let mut dist:i64 = 0;
        af_hamming_matcher(&mut idx as MutAfArray, &mut dist as MutAfArray,
                           query.get() as AfArray, train.get() as AfArray,
                           dist_dims as DimT, n_dist as c_uint);
        (Array::from(idx), Array::from(dist))
    }
}

#[allow(unused_mut)]
pub fn match_template(search_img: &Array, template_img: &Array, mtype: MatchType) -> Array {
    unsafe {
        let mut temp: i64 = 0;
        af_match_template(&mut temp as MutAfArray,
                          search_img.get() as AfArray, template_img.get() as AfArray,
                          mtype as uint8_t);
        Array::from(temp)
    }
}
<|endoftext|>


<commit_before>
// This pass generates x86-64 assembly from IR.

#![allow(non_upper_case_globals)]

use crate::regalloc::*;
use crate::util::*;
use crate::*;
use std::collections::HashMap;

macro_rules! emit {
    ($fmt:expr) => {
        print!("\t");
        println!($fmt);
    };
    ($fmt:expr,$($x:tt)*) => {
        print!("\t");
        println!($fmt, $( $x )*);
    };
}

lazy_static! {
    static ref NLABEL: Mutex<usize> = Mutex::new(0);
    static ref ESCAPED: Mutex<HashMap<char,char>> = Mutex::new(HashMap::new());
}

fn bump_nlabel() -> usize {
    let mut nlabel = NLABEL.lock().unwrap();
    let ret = *nlabel;
    *nlabel += 1;
    return ret;
}

const argreg8: [&'static str; 6] = ["dil", "sil", "dl", "cl", "r8b", "r9b"];
const argreg32: [&'static str; 6] = ["edi", "esi", "edx", "ecx", "r8d", "r9d"];
const argreg64: [&'static str; 6] = ["rdi", "rsi", "rdx", "rcx", "r8", "r9"];

fn init_escaped() {
    let mut escaped = ESCAPED.lock().unwrap();
    if escaped.len() > 0 {
        return;
    }

    escaped.insert(char::from(8), 'b');  // \b
    escaped.insert(char::from(12), 'f'); // \f
    escaped.insert(char::from(10), 'n'); // \n
    escaped.insert(char::from(13), 'r'); // \r
    escaped.insert(char::from(9), 't');  // \t
    escaped.insert('\\', '\\');
    escaped.insert('\'', '\'');
    escaped.insert('\"', '"');
}

fn escaped(c: char) -> Option<char> {
    let mut ret = None;
    match ESCAPED.lock() {
        Ok(escaped) => {
            match escaped.get(&c) {
                Some(esc) => {
                    ret = Some(*esc);
                }
                None => { }
            }
        }
        Err(_) => { }
    }
    return ret;
}

fn escape(s: & String) -> String {
    init_escaped();

    let mut buf = String::new();
    let mut chars = s.chars();
    while let Some(c) = chars.next() {
        if let Some(esc) = escaped(c) {
            buf.push('\\');
            buf.push(esc);
        } else if c.is_ascii_graphic() || c == ' ' {
            buf.push(c);
        } else {
            buf.push_str(&format!("\\{:03o}", u32::from(c)));
        }
    }

    //buf.push(char::from(0));
    return buf;
}

fn emit_cmp(insn: &str, ir: &IR) {
    let lhs = ir.lhs as usize;
    let rhs = ir.rhs as usize;

    emit!("cmp {}, {}", regs[lhs], regs[rhs]);
    emit!("{} {}", insn, regs8[lhs]);
    emit!("movzb {}, {}", regs[lhs], regs8[lhs]);
}

fn gen(fun: &IR) {
    let ret = format!(".Lend{}", bump_nlabel());

    println!(".global {}", fun.name);
    println!("{}:", fun.name);
    emit!("push rbp");
    emit!("mov rbp, rsp");
    emit!("sub rsp, {}", roundup(fun.stacksize, 16));
    emit!("push r12");
    emit!("push r13");
    emit!("push r14");
    emit!("push r15");

    for i in 0..fun.ir.len() {
        let ir = &fun.ir[i as usize];
        let lhs = ir.lhs as usize;
        let rhs = ir.rhs as usize;

        match ir.op {
            IRType::IMM => {
                emit!("mov {}, {}", regs[lhs], rhs);
            }
            IRType::BPREL => {
                emit!("lea {}, [rbp-{}]", regs[lhs], rhs);
            }
            IRType::MOV => {
                emit!("mov {}, {}", regs[lhs], regs[rhs]);
            }
            IRType::RETURN => {
                emit!("mov rax, {}", regs[lhs]);
                emit!("jmp {}", ret);
            }
            IRType::CALL => {
                for i in 0..ir.nargs {
                    emit!("mov {}, {}", argreg64[i], regs[ir.args[i] as usize]);
                }

                emit!("push r10");
                emit!("push r11");
                emit!("mov rax, 0");
                emit!("call {}", ir.name);
                emit!("pop r11");
                emit!("pop r10");

                emit!("mov {}, rax", regs[lhs]);
            }
            IRType::LABEL => {
                println!(".L{}:", lhs);
            }
            IRType::LABEL_ADDR => {
                emit!("lea {}, {}", regs[lhs], ir.name);
            }
            IRType::NEG => {
                emit!("neg {}", regs[lhs]);
            }
            IRType::EQ => {
                emit_cmp("sete", ir)
            }
            IRType::NE => {
                emit_cmp("setne", ir)
            }
            IRType::LT => {
                emit_cmp("setl", ir)
            }
            IRType::LE => {
                emit_cmp("setle", ir)
            }
            IRType::AND => {
                emit!("and {}, {}", regs[lhs], regs[rhs]);
            }
            IRType::OR => {
                emit!("or {}, {}", regs[lhs], regs[rhs]);
            }
            IRType::XOR => {
                emit!("xor {}, {}", regs[lhs], regs[rhs]);
            }
            IRType::SHL => {
                emit!("mov cl, {}", regs8[rhs]);
                emit!("shl {}, cl", regs[lhs]);
            }
            IRType::SHR => {
                emit!("mov cl, {}", regs8[rhs]);
                emit!("shr {}, cl", regs[lhs]);
            }
            IRType::JMP => {
                emit!("jmp .L{}", lhs);
            }
            IRType::IF => {
                emit!("cmp {}, 0", regs[lhs]);
                emit!("jne .L{}", rhs);
            }
            IRType::UNLESS => {
                emit!("cmp {}, 0", regs[lhs]);
                emit!("je .L{}", rhs);
            }
            IRType::LOAD8 => {
                emit!("mov {}, [{}]", regs8[lhs], regs[rhs]);
                emit!("movzb {}, {}", regs[lhs], regs8[lhs]);
            }
            IRType::LOAD32 => {
                emit!("mov {}, [{}]", regs32[lhs], regs[rhs]);
            }
            IRType::LOAD64 => {
                emit!("mov {}, [{}]", regs[lhs], regs[rhs]);
            }
            IRType::STORE8 => {
                emit!("mov [{}], {}", regs[lhs], regs8[rhs]);

                // fall through in 9cc
                emit!("mov [{}], {}", regs[lhs], regs32[rhs]);
            }
            IRType::STORE32 => {
                emit!("mov [{}], {}", regs[lhs], regs32[rhs]);
            }
            IRType::STORE64 => {
                emit!("mov [{}], {}", regs[lhs], regs[rhs]);
            }
            IRType::STORE8_ARG => {
                emit!("mov [rbp-{}], {}", lhs, argreg8[rhs]);
            }
            IRType::STORE32_ARG => {
                emit!("mov [rbp-{}], {}", lhs, argreg32[rhs]);
            }
            IRType::STORE64_ARG => {
                emit!("mov [rbp-{}], {}", lhs, argreg64[rhs]);
            }
            IRType::ADD => {
                emit!("add {}, {}", regs[lhs], regs[rhs]);
            }
            IRType::ADD_IMM => {
                emit!("add {}, {}", regs[lhs], rhs);
            }
            IRType::SUB => {
                emit!("sub {}, {}", regs[lhs], regs[rhs]);
            }
            IRType::SUB_IMM => {
                emit!("sub {}, {}", regs[lhs], rhs);
            }
            IRType::MUL => {
                emit!("mov rax, {}", regs[rhs]);
                emit!("mul {}", regs[lhs]);
                emit!("mov {}, rax", regs[lhs]);
            }
            IRType::MUL_IMM => {
                emit!("mov rax, {}", rhs);
                emit!("mul {}", regs[lhs]);
                emit!("mov {}, rax", regs[lhs]);
            }
            IRType::DIV => {
                emit!("mov rax, {}", regs[lhs]);
                emit!("cqo");
                emit!("div {}", regs[rhs]);
                emit!("mov {}, rax", regs[lhs]);
            }
            IRType::MOD => {
                emit!("mov rax, {}", regs[lhs]);
                emit!("cqo");
                emit!("div {}", regs[rhs]);
                emit!("mov {}, rdx", regs[lhs]);
            }
            IRType::NOP => {}
            ref i => {
                panic!("unknown operator {:?}", i);
            }
        }
    }
    println!("{}:", ret);
    emit!("pop r15");
    emit!("pop r14");
    emit!("pop r13");
    emit!("pop r12");
    emit!("mov rsp, rbp");
    emit!("pop rbp");
    emit!("ret");
}

pub fn gen_x86(globals: Vec<Var>, fns: &Vec<IR>) {
    println!(".intel_syntax noprefix");

    println!(".data");
    for var in globals.iter() {
        if var.is_extern {
            continue;
        }
        println!("{}:", var.name);
        emit!(".ascii \"{}\"", escape(&var.data));
    }

    println!(".text");
    for i in 0..fns.len() {
        gen(&fns[i]);
    }
}
<commit_msg>Fix a fall-through bug.<commit_after>
// This pass generates x86-64 assembly from IR.

#![allow(non_upper_case_globals)]

use crate::regalloc::*;
use crate::util::*;
use crate::*;
use std::collections::HashMap;

macro_rules! emit {
    ($fmt:expr) => {
        print!("\t");
        println!($fmt);
    };
    ($fmt:expr,$($x:tt)*) => {
        print!("\t");
        println!($fmt, $( $x )*);
    };
}

lazy_static! {
    static ref NLABEL: Mutex<usize> = Mutex::new(0);
    static ref ESCAPED: Mutex<HashMap<char,char>> = Mutex::new(HashMap::new());
}

fn bump_nlabel() -> usize {
    let mut nlabel = NLABEL.lock().unwrap();
    let ret = *nlabel;
    *nlabel += 1;
    return ret;
}

const argreg8: [&'static str; 6] = ["dil", "sil", "dl", "cl", "r8b", "r9b"];
const argreg32: [&'static str; 6] = ["edi", "esi", "edx", "ecx", "r8d", "r9d"];
const argreg64: [&'static str; 6] = ["rdi", "rsi", "rdx", "rcx", "r8", "r9"];

fn init_escaped() {
    let mut escaped = ESCAPED.lock().unwrap();
    if escaped.len() > 0 {
        return;
    }

    escaped.insert(char::from(8), 'b');  // \b
    escaped.insert(char::from(12), 'f'); // \f
    escaped.insert(char::from(10), 'n'); // \n
    escaped.insert(char::from(13), 'r'); // \r
    escaped.insert(char::from(9), 't');  // \t
    escaped.insert('\\', '\\');
    escaped.insert('\'', '\'');
    escaped.insert('\"', '"');
}

fn escaped(c: char) -> Option<char> {
    let mut ret = None;
    match ESCAPED.lock() {
        Ok(escaped) => {
            match escaped.get(&c) {
                Some(esc) => {
                    ret = Some(*esc);
                }
                None => { }
            }
        }
        Err(_) => { }
    }
    return ret;
}

fn escape(s: & String) -> String {
    init_escaped();

    let mut buf = String::new();
    let mut chars = s.chars();
    while let Some(c) = chars.next() {
        if let Some(esc) = escaped(c) {
            buf.push('\\');
            buf.push(esc);
        } else if c.is_ascii_graphic() || c == ' ' {
            buf.push(c);
        } else {
            buf.push_str(&format!("\\{:03o}", u32::from(c)));
        }
    }

    //buf.push(char::from(0));
    return buf;
}

fn emit_cmp(insn: &str, ir: &IR) {
    let lhs = ir.lhs as usize;
    let rhs = ir.rhs as usize;

    emit!("cmp {}, {}", regs[lhs], regs[rhs]);
    emit!("{} {}", insn, regs8[lhs]);
    emit!("movzb {}, {}", regs[lhs], regs8[lhs]);
}

fn gen(fun: &IR) {
    let ret = format!(".Lend{}", bump_nlabel());

    println!(".global {}", fun.name);
    println!("{}:", fun.name);
    emit!("push rbp");
    emit!("mov rbp, rsp");
    emit!("sub rsp, {}", roundup(fun.stacksize, 16));
    emit!("push r12");
    emit!("push r13");
    emit!("push r14");
    emit!("push r15");

    for i in 0..fun.ir.len() {
        let ir = &fun.ir[i as usize];
        let lhs = ir.lhs as usize;
        let rhs = ir.rhs as usize;

        match ir.op {
            IRType::IMM => {
                emit!("mov {}, {}", regs[lhs], rhs);
            }
            IRType::BPREL => {
                emit!("lea {}, [rbp-{}]", regs[lhs], rhs);
            }
            IRType::MOV => {
                emit!("mov {}, {}", regs[lhs], regs[rhs]);
            }
            IRType::RETURN => {
                emit!("mov rax, {}", regs[lhs]);
                emit!("jmp {}", ret);
            }
            IRType::CALL => {
                for i in 0..ir.nargs {
                    emit!("mov {}, {}", argreg64[i], regs[ir.args[i] as usize]);
                }

                emit!("push r10");
                emit!("push r11");
                emit!("mov rax, 0");
                emit!("call {}", ir.name);
                emit!("pop r11");
                emit!("pop r10");

                emit!("mov {}, rax", regs[lhs]);
            }
            IRType::LABEL => {
                println!(".L{}:", lhs);
            }
            IRType::LABEL_ADDR => {
                emit!("lea {}, {}", regs[lhs], ir.name);
            }
            IRType::NEG => {
                emit!("neg {}", regs[lhs]);
            }
            IRType::EQ => {
                emit_cmp("sete", ir)
            }
            IRType::NE => {
                emit_cmp("setne", ir)
            }
            IRType::LT => {
                emit_cmp("setl", ir)
            }
            IRType::LE => {
                emit_cmp("setle", ir)
            }
            IRType::AND => {
                emit!("and {}, {}", regs[lhs], regs[rhs]);
            }
            IRType::OR => {
                emit!("or {}, {}", regs[lhs], regs[rhs]);
            }
            IRType::XOR => {
                emit!("xor {}, {}", regs[lhs], regs[rhs]);
            }
            IRType::SHL => {
                emit!("mov cl, {}", regs8[rhs]);
                emit!("shl {}, cl", regs[lhs]);
            }
            IRType::SHR => {
                emit!("mov cl, {}", regs8[rhs]);
                emit!("shr {}, cl", regs[lhs]);
            }
            IRType::JMP => {
                emit!("jmp .L{}", lhs);
            }
            IRType::IF => {
                emit!("cmp {}, 0", regs[lhs]);
                emit!("jne .L{}", rhs);
            }
            IRType::UNLESS => {
                emit!("cmp {}, 0", regs[lhs]);
                emit!("je .L{}", rhs);
            }
            IRType::LOAD8 => {
                emit!("mov {}, [{}]", regs8[lhs], regs[rhs]);
                emit!("movzb {}, {}", regs[lhs], regs8[lhs]);
            }
            IRType::LOAD32 => {
                emit!("mov {}, [{}]", regs32[lhs], regs[rhs]);
            }
            IRType::LOAD64 => {
                emit!("mov {}, [{}]", regs[lhs], regs[rhs]);
            }
            IRType::STORE8 => {
                emit!("mov [{}], {}", regs[lhs], regs8[rhs]);
            }
            IRType::STORE32 => {
                emit!("mov [{}], {}", regs[lhs], regs32[rhs]);
            }
            IRType::STORE64 => {
                emit!("mov [{}], {}", regs[lhs], regs[rhs]);
            }
            IRType::STORE8_ARG => {
                emit!("mov [rbp-{}], {}", lhs, argreg8[rhs]);
            }
            IRType::STORE32_ARG => {
                emit!("mov [rbp-{}], {}", lhs, argreg32[rhs]);
            }
            IRType::STORE64_ARG => {
                emit!("mov [rbp-{}], {}", lhs, argreg64[rhs]);
            }
            IRType::ADD => {
                emit!("add {}, {}", regs[lhs], regs[rhs]);
            }
            IRType::ADD_IMM => {
                emit!("add {}, {}", regs[lhs], rhs);
            }
            IRType::SUB => {
                emit!("sub {}, {}", regs[lhs], regs[rhs]);
            }
            IRType::SUB_IMM => {
                emit!("sub {}, {}", regs[lhs], rhs);
            }
            IRType::MUL => {
                emit!("mov rax, {}", regs[rhs]);
                emit!("mul {}", regs[lhs]);
                emit!("mov {}, rax", regs[lhs]);
            }
            IRType::MUL_IMM => {
                if rhs < 256 && rhs.count_ones() == 1 {
                    emit!("shl {}, {}", regs[lhs], rhs.trailing_zeros());
                } else {
                    emit!("mov rax, {}", rhs);
                    emit!("mul {}", regs[lhs]);
                    emit!("mov {}, rax", regs[lhs]);
                }
            }
            IRType::DIV => {
                emit!("mov rax, {}", regs[lhs]);
                emit!("cqo");
                emit!("div {}", regs[rhs]);
                emit!("mov {}, rax", regs[lhs]);
            }
            IRType::MOD => {
                emit!("mov rax, {}", regs[lhs]);
                emit!("cqo");
                emit!("div {}", regs[rhs]);
                emit!("mov {}, rdx", regs[lhs]);
            }
            IRType::NOP => {}
            ref i => {
                panic!("unknown operator {:?}", i);
            }
        }
    }
    println!("{}:", ret);
    emit!("pop r15");
    emit!("pop r14");
    emit!("pop r13");
    emit!("pop r12");
    emit!("mov rsp, rbp");
    emit!("pop rbp");
    emit!("ret");
}

pub fn gen_x86(globals: Vec<Var>, fns: &Vec<IR>) {
    println!(".intel_syntax noprefix");

    println!(".data");
    for var in globals.iter() {
        if var.is_extern {
            continue;
        }
        println!("{}:", var.name);
        emit!(".ascii \"{}\"", escape(&var.data));
    }

    println!(".text");
    for i in 0..fns.len() {
        gen(&fns[i]);
    }
}
<|endoftext|>


<commit_before>use super::{
    LENGTH_BYTES,
};

#[derive(Debug, PartialEq)]
pub enum DecodeError {
    OutOfBounds,
    TooShort,
    TooLong,
}

pub trait Decodable: Sized {
    fn ssz_decode(bytes: &[u8], index: usize) -> Result<Self, DecodeError>;
}

/// Decode the given bytes for the given type
///
/// The single ssz encoded value will be decoded as the given type at the
/// given index.
pub fn decode_ssz<T>(ssz_bytes: &[u8], index: usize)
    -> Result<T, DecodeError>
    where T: Decodable
{
    if index >= ssz_bytes.len() {
        return Err(DecodeError::OutOfBounds)
    }
    T::ssz_decode(ssz_bytes, index)
}

/// Decode the nth element of some ssz list.
///
/// A single ssz encoded value can be considered a list of
/// one element, so this function will work on it too.
pub fn decode_ssz_list_element<T>(ssz_bytes: &[u8], n: usize)
    -> Result<T, DecodeError>
    where T: Decodable
{
    T::ssz_decode(nth_value(ssz_bytes, n)?)
}

/// Return the nth value in some ssz encoded list.
///
/// The four-byte length prefix is not included in the return.
///
/// A single ssz encoded value can be considered a list of
/// one element, so this function will work on it too.
fn nth_value(ssz_bytes: &[u8], n: usize)
    -> Result<&[u8], DecodeError>
{
    let mut c: usize = 0;
    for i in 0..(n + 1) {
        let length = decode_length(&ssz_bytes[c..], LENGTH_BYTES)?;
        let next = c + LENGTH_BYTES + length;

        if i == n {
            return Ok(&ssz_bytes[c + LENGTH_BYTES..next]);
        } else {
            if next >= ssz_bytes.len() {
                return Err(DecodeError::OutOfBounds);
            } else {
                c = next;
            }
        }
    }
    Err(DecodeError::OutOfBounds)
}

/// Given some number of bytes, interpret the first four
/// bytes as a 32-bit big-endian integer and return the
/// result.
fn decode_length(bytes: &[u8], length_bytes: usize)
    -> Result<usize, DecodeError>
{
    if bytes.len() < length_bytes {
        return Err(DecodeError::TooShort);
    };
    let mut len: usize = 0;
    for i in 0..length_bytes {
        let offset = (length_bytes - i - 1) * 8;
        len = ((bytes[i] as usize) << offset) | len;
    };
    Ok(len)
}

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

    #[test]
    fn test_ssz_decode_length() {
        let decoded = decode_length(
            &vec![0, 0, 1],
            LENGTH_BYTES);
        assert_eq!(decoded.unwrap(), 1);

        let decoded = decode_length(
            &vec![0, 1, 0],
            LENGTH_BYTES);
        assert_eq!(decoded.unwrap(), 256);

        let decoded = decode_length(
            &vec![0, 1, 255],
            LENGTH_BYTES);
        assert_eq!(decoded.unwrap(), 511);

        let decoded = decode_length(
            &vec![255, 255, 255],
            LENGTH_BYTES);
        assert_eq!(decoded.unwrap(), 16777215);
    }

    #[test]
    fn test_encode_decode_length() {
        let params: Vec<usize> = vec![
            0, 1, 2, 3, 7, 8, 16,
            2^8,  2^8  + 1,
            2^16, 2^16 + 1,
            2^24, 2^24 + 1,
            2^32,
        ];
        for i in params {
            let decoded = decode_length(
                &encode_length(i, LENGTH_BYTES),
                LENGTH_BYTES).unwrap();
            assert_eq!(i, decoded);
        }
    }

    #[test]
    fn test_ssz_nth_value() {
        let ssz = vec![0, 0, 1, 0];
        let result = nth_value(&ssz, 0).unwrap();
        assert_eq!(result, vec![0].as_slice());

        let ssz = vec![0, 0, 4, 1, 2, 3, 4];
        let result = nth_value(&ssz, 0).unwrap();
        assert_eq!(result, vec![1, 2, 3, 4].as_slice());

        let ssz = vec![0, 0, 1, 0, 0, 0, 1, 1];
        let result = nth_value(&ssz, 1).unwrap();
        assert_eq!(result, vec![1].as_slice());

        let mut ssz = vec![0, 1, 255];
        ssz.append(&mut vec![42; 511]);
        let result = nth_value(&ssz, 0).unwrap();
        assert_eq!(result, vec![42; 511].as_slice());
    }
}
<commit_msg>Remove `decode_ssz_list_element` unused function.<commit_after>use super::{
    LENGTH_BYTES,
};

#[derive(Debug, PartialEq)]
pub enum DecodeError {
    OutOfBounds,
    TooShort,
    TooLong,
}

pub trait Decodable: Sized {
    fn ssz_decode(bytes: &[u8], index: usize) -> Result<Self, DecodeError>;
}

/// Decode the given bytes for the given type
///
/// The single ssz encoded value will be decoded as the given type at the
/// given index.
pub fn decode_ssz<T>(ssz_bytes: &[u8], index: usize)
    -> Result<T, DecodeError>
    where T: Decodable
{
    if index >= ssz_bytes.len() {
        return Err(DecodeError::OutOfBounds)
    }
    T::ssz_decode(ssz_bytes, index)
}

/// Return the nth value in some ssz encoded list.
///
/// The four-byte length prefix is not included in the return.
///
/// A single ssz encoded value can be considered a list of
/// one element, so this function will work on it too.
fn nth_value(ssz_bytes: &[u8], n: usize)
    -> Result<&[u8], DecodeError>
{
    let mut c: usize = 0;
    for i in 0..(n + 1) {
        let length = decode_length(&ssz_bytes[c..], LENGTH_BYTES)?;
        let next = c + LENGTH_BYTES + length;

        if i == n {
            return Ok(&ssz_bytes[c + LENGTH_BYTES..next]);
        } else {
            if next >= ssz_bytes.len() {
                return Err(DecodeError::OutOfBounds);
            } else {
                c = next;
            }
        }
    }
    Err(DecodeError::OutOfBounds)
}

/// Given some number of bytes, interpret the first four
/// bytes as a 32-bit big-endian integer and return the
/// result.
fn decode_length(bytes: &[u8], length_bytes: usize)
    -> Result<usize, DecodeError>
{
    if bytes.len() < length_bytes {
        return Err(DecodeError::TooShort);
    };
    let mut len: usize = 0;
    for i in 0..length_bytes {
        let offset = (length_bytes - i - 1) * 8;
        len = ((bytes[i] as usize) << offset) | len;
    };
    Ok(len)
}

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

    #[test]
    fn test_ssz_decode_length() {
        let decoded = decode_length(
            &vec![0, 0, 1],
            LENGTH_BYTES);
        assert_eq!(decoded.unwrap(), 1);

        let decoded = decode_length(
            &vec![0, 1, 0],
            LENGTH_BYTES);
        assert_eq!(decoded.unwrap(), 256);

        let decoded = decode_length(
            &vec![0, 1, 255],
            LENGTH_BYTES);
        assert_eq!(decoded.unwrap(), 511);

        let decoded = decode_length(
            &vec![255, 255, 255],
            LENGTH_BYTES);
        assert_eq!(decoded.unwrap(), 16777215);
    }

    #[test]
    fn test_encode_decode_length() {
        let params: Vec<usize> = vec![
            0, 1, 2, 3, 7, 8, 16,
            2^8,  2^8  + 1,
            2^16, 2^16 + 1,
            2^24, 2^24 + 1,
            2^32,
        ];
        for i in params {
            let decoded = decode_length(
                &encode_length(i, LENGTH_BYTES),
                LENGTH_BYTES).unwrap();
            assert_eq!(i, decoded);
        }
    }

    #[test]
    fn test_ssz_nth_value() {
        let ssz = vec![0, 0, 1, 0];
        let result = nth_value(&ssz, 0).unwrap();
        assert_eq!(result, vec![0].as_slice());

        let ssz = vec![0, 0, 4, 1, 2, 3, 4];
        let result = nth_value(&ssz, 0).unwrap();
        assert_eq!(result, vec![1, 2, 3, 4].as_slice());

        let ssz = vec![0, 0, 1, 0, 0, 0, 1, 1];
        let result = nth_value(&ssz, 1).unwrap();
        assert_eq!(result, vec![1].as_slice());

        let mut ssz = vec![0, 1, 255];
        ssz.append(&mut vec![42; 511]);
        let result = nth_value(&ssz, 0).unwrap();
        assert_eq!(result, vec![42; 511].as_slice());
    }
}
<|endoftext|>


<commit_before>use crate::{
    buffer::Buffer,
    channels::Channels,
    connection_status::ConnectionState,
    frames::Frames,
    heartbeat::Heartbeat,
    internal_rpc::InternalRPC,
    protocol::{self, AMQPError, AMQPHardError},
    reactor::{ReactorBuilder, ReactorHandle, Slot},
    socket_state::SocketState,
    tcp::HandshakeResult,
    thread::ThreadHandle,
    Configuration, ConnectionStatus, Error, PromiseResolver, Result, TcpStream,
};
use amq_protocol::frame::{gen_frame, parse_frame, AMQPFrame, GenError};
use log::{debug, error, trace};
use std::{
    collections::VecDeque,
    convert::TryFrom,
    io::{self, Write},
    sync::Arc,
    thread::Builder as ThreadBuilder,
    time::Duration,
};

const FRAMES_STORAGE: usize = 32;

// The amount of Read(0) we tolerate before onsidering the connection as broken.
// This might seem high but better safe than sorry.
const BAD_READS_TOLERANCE: u8 = 16;

#[derive(Debug, PartialEq)]
enum Status {
    Initial,
    Connected,
    Stop,
}

pub struct IoLoop {
    connection_status: ConnectionStatus,
    configuration: Configuration,
    channels: Channels,
    internal_rpc: InternalRPC,
    frames: Frames,
    heartbeat: Heartbeat,
    socket_state: SocketState,
    reactor: Box<dyn ReactorHandle + Send>,
    reactor_thread_handle: ThreadHandle,
    connection_io_loop_handle: ThreadHandle,
    stream: TcpStream,
    slot: Slot,
    status: Status,
    frame_size: usize,
    receive_buffer: Buffer,
    send_buffer: Buffer,
    serialized_frames: VecDeque<(u64, Option<PromiseResolver<()>>)>,
    bad_reads: u8,
}

impl IoLoop {
    #[allow(clippy::too_many_arguments)]
    pub(crate) fn new(
        connection_status: ConnectionStatus,
        configuration: Configuration,
        channels: Channels,
        internal_rpc: InternalRPC,
        frames: Frames,
        socket_state: SocketState,
        connection_io_loop_handle: ThreadHandle,
        stream: HandshakeResult,
        reactor_builder: &dyn ReactorBuilder,
    ) -> Result<Self> {
        let mut stream = TcpStream::try_from(stream)?;
        let heartbeat = Heartbeat::new(channels.clone());
        let mut reactor = reactor_builder.build(heartbeat.clone())?;
        let reactor_handle = reactor.handle();
        let frame_size = std::cmp::max(
            protocol::constants::FRAME_MIN_SIZE as usize,
            configuration.frame_max() as usize,
        );
        let slot = reactor.register(stream.inner_mut(), socket_state.handle())?;

        let reactor_thread_handle = reactor.start()?;

        Ok(Self {
            connection_status,
            configuration,
            channels,
            internal_rpc,
            frames,
            heartbeat,
            socket_state,
            reactor: reactor_handle,
            reactor_thread_handle,
            connection_io_loop_handle,
            stream,
            slot,
            status: Status::Initial,
            frame_size,
            receive_buffer: Buffer::with_capacity(FRAMES_STORAGE * frame_size),
            send_buffer: Buffer::with_capacity(FRAMES_STORAGE * frame_size),
            serialized_frames: VecDeque::default(),
            bad_reads: 0,
        })
    }

    fn finish_setup(&mut self) -> Result<bool> {
        if self.connection_status.connected() {
            let frame_max = self.configuration.frame_max() as usize;
            self.frame_size = std::cmp::max(self.frame_size, frame_max);
            self.receive_buffer.grow(FRAMES_STORAGE * self.frame_size);
            self.send_buffer.grow(FRAMES_STORAGE * self.frame_size);
            let heartbeat = self.configuration.heartbeat();
            if heartbeat != 0 {
                let heartbeat = Duration::from_millis(u64::from(heartbeat) * 500); // * 1000 (ms) / 2 (half the negotiated timeout)
                self.heartbeat.set_timeout(heartbeat);
                self.reactor.start_heartbeat();
            }
            let peer = self.stream.inner().peer_addr()?;
            debug!("Connected to {}", peer);
            self.status = Status::Connected;
        }
        Ok(true)
    }

    fn ensure_setup(&mut self) -> Result<bool> {
        match self.status {
            Status::Initial => self.finish_setup(),
            Status::Connected => Ok(true),
            Status::Stop => Ok(false),
        }
    }

    fn has_data(&self) -> bool {
        self.frames.has_pending()
            || self.send_buffer.available_data() > 0
            || !self.serialized_frames.is_empty()
    }

    fn can_write(&mut self) -> bool {
        self.socket_state.writable() && self.has_data() && !self.connection_status.blocked()
    }

    fn can_read(&mut self) -> bool {
        self.socket_state.readable() && self.receive_buffer.available_space() > 0
    }

    fn can_parse(&self) -> bool {
        self.receive_buffer.available_data() > 0
    }

    fn should_continue(&self) -> bool {
        (self.status != Status::Connected
            || self.connection_status.connected()
            || self.connection_status.closing())
            && self.status != Status::Stop
            && !self.connection_status.errored()
    }

    pub fn start(mut self) -> Result<()> {
        let waker = self.socket_state.handle();
        let io_loop_handle = self.connection_io_loop_handle.clone();
        io_loop_handle.register(
            ThreadBuilder::new()
                .name("lapin-io-loop".to_owned())
                .spawn(move || {
                    while self.should_continue() {
                        if let Err(err) = self.run() {
                            self.critical_error(err)?;
                        }
                    }
                    self.heartbeat.cancel();
                    self.reactor.shutdown()?;
                    self.reactor_thread_handle.wait("reactor")
                })?,
        );
        waker.wake();
        Ok(())
    }

    fn poll_internal_rpc(&self) -> Result<()> {
        self.internal_rpc.poll(&self.channels)
    }

    fn poll_socket_events(&mut self) -> Result<()> {
        self.socket_state.poll_events();
        if self.socket_state.error() {
            self.critical_error(io::Error::from(io::ErrorKind::ConnectionAborted).into())?;
        }
        if self.socket_state.closed() {
            self.critical_error(io::Error::from(io::ErrorKind::ConnectionReset).into())?;
        }
        self.poll_internal_rpc()
    }

    fn run(&mut self) -> Result<()> {
        trace!("io_loop run");
        self.poll_socket_events()?;
        if !self.ensure_setup()? {
            return Ok(());
        }
        trace!(
            "io_loop do_run; can_read={}, can_write={}, has_data={}",
            self.socket_state.readable(),
            self.socket_state.writable(),
            self.has_data()
        );
        if !self.can_read() && !self.can_write() {
            self.socket_state.wait();
        }
        self.poll_socket_events()?;
        if self.stream.is_handshaking() {
            self.stream.handshake()?;
            if self.stream.is_handshaking() {
                // We hit WOULDBLOCK while handshaking, wait for the next socket event
                return Ok(());
            }
        }
        self.attempt_flush()?;
        self.write()?;
        if self.connection_status.closed() {
            self.status = Status::Stop;
        }
        if self.should_continue() {
            self.read()?;
        }
        self.handle_frames()?;
        trace!(
            "io_loop do_run done; can_read={}, can_write={}, has_data={}, status={:?}",
            self.socket_state.readable(),
            self.socket_state.writable(),
            self.has_data(),
            self.status
        );
        self.poll_internal_rpc()
    }

    fn critical_error(&mut self, error: Error) -> Result<()> {
        if let Some(resolver) = self.connection_status.connection_resolver() {
            resolver.swear(Err(error.clone()));
        }
        self.status = Status::Stop;
        self.channels.set_connection_error(error.clone())?;
        for (_, resolver) in std::mem::take(&mut self.serialized_frames) {
            if let Some(resolver) = resolver {
                resolver.swear(Err(error.clone()));
            }
        }
        self.reactor.shutdown()?;
        Err(error)
    }

    fn attempt_flush(&mut self) -> Result<()> {
        let res = self.flush();
        if let Err(e) = self.socket_state.handle_flush_result(res) {
            error!("error flushing: {:?}", e);
            self.critical_error(e)?;
        }
        self.poll_internal_rpc()
    }

    fn handle_read_result(&mut self, result: Result<()>) -> Result<()> {
        if let Err(e) = self
            .socket_state
            .handle_read_result(result, &*self.reactor, self.slot)
        {
            error!("error reading: {:?}", e);
            self.critical_error(e)?;
        }
        self.poll_internal_rpc()
    }

    fn handle_write_result(&mut self, result: Result<()>) -> Result<()> {
        if let Err(e) = self
            .socket_state
            .handle_write_result(result, &*self.reactor, self.slot)
        {
            error!("error writing: {:?}", e);
            self.critical_error(e)?;
        }
        self.poll_internal_rpc()
    }

    fn flush(&mut self) -> Result<()> {
        self.stream.flush()?;
        self.poll_internal_rpc()
    }

    fn write(&mut self) -> Result<()> {
        while self.can_write() {
            let res = self.write_to_stream();
            self.handle_write_result(res)?;
        }
        self.poll_internal_rpc()
    }

    fn read(&mut self) -> Result<()> {
        while self.can_read() {
            let res = self.read_from_stream();
            self.handle_read_result(res)?;
        }
        self.poll_internal_rpc()
    }

    fn write_to_stream(&mut self) -> Result<()> {
        self.flush()?;
        self.serialize()?;

        let sz = self.send_buffer.write_to(&mut self.stream)?;

        if sz > 0 {
            self.heartbeat.update_last_write();

            trace!("wrote {} bytes", sz);
            self.send_buffer.consume(sz);

            let mut written = sz as u64;
            while written > 0 {
                if let Some((to_write, resolver)) = self.serialized_frames.pop_front() {
                    if written < to_write {
                        self.serialized_frames
                            .push_front((to_write - written, resolver));
                        trace!("{} to write to complete this frame", to_write - written);
                        written = 0;
                    } else {
                        if let Some(resolver) = resolver {
                            resolver.swear(Ok(()));
                        }
                        written -= to_write;
                    }
                } else {
                    error!(
                        "We've written {} but didn't expect to write anything",
                        written
                    );
                    break;
                }
            }

            if self.send_buffer.available_data() > 0 {
                // We didn't write all the data yet
                trace!("Still {} to write", self.send_buffer.available_data());
            }

            self.flush()?;
        } else {
            error!("Socket was writable but we wrote 0, marking as wouldblock");
            self.handle_write_result(Err(io::Error::from(io::ErrorKind::WouldBlock).into()))?;
        }
        self.poll_internal_rpc()
    }

    fn read_from_stream(&mut self) -> Result<()> {
        match self.connection_status.state() {
            ConnectionState::Closed => Ok(()),
            ConnectionState::Error => Err(Error::InvalidConnectionState(ConnectionState::Error)),
            _ => {
                let sz = self.receive_buffer.read_from(&mut self.stream)?;

                if sz > 0 {
                    self.bad_reads = 0;
                    trace!("read {} bytes", sz);
                    self.receive_buffer.fill(sz);
                } else {
                    self.bad_reads += 1;
                    let io_error_kind = if self.bad_reads > BAD_READS_TOLERANCE {
                        error!("Socket was readable but we read 0 {} times in a row, mark the connection as broken", BAD_READS_TOLERANCE);
                        io::ErrorKind::ConnectionAborted
                    } else {
                        error!("Socket was readable but we read 0, marking as wouldblock");
                        io::ErrorKind::WouldBlock
                    };
                    self.handle_read_result(Err(io::Error::from(io_error_kind).into()))?;
                }
                self.poll_internal_rpc()
            }
        }
    }

    fn serialize(&mut self) -> Result<()> {
        while let Some((next_msg, resolver)) = self.frames.pop(self.channels.flow()) {
            trace!("will write to buffer: {}", next_msg);
            let checkpoint = self.send_buffer.checkpoint();
            let res = gen_frame(&next_msg)((&mut self.send_buffer).into());
            match res.map(|w| w.into_inner().1) {
                Ok(sz) => self.serialized_frames.push_back((sz, resolver)),
                Err(e) => {
                    self.send_buffer.rollback(checkpoint);
                    match e {
                        GenError::BufferTooSmall(_) => {
                            // Requeue msg
                            self.frames.retry((next_msg, resolver));
                            break;
                        }
                        e => {
                            error!("error generating frame: {:?}", e);
                            self.critical_error(Error::SerialisationError(Arc::new(e)))?;
                        }
                    }
                }
            }
        }
        self.poll_internal_rpc()
    }

    fn handle_frames(&mut self) -> Result<()> {
        while self.can_parse() {
            if let Some(frame) = self.parse()? {
                self.channels.handle_frame(frame)?;
            } else {
                break;
            }
        }
        self.poll_internal_rpc()
    }

    fn parse(&mut self) -> Result<Option<AMQPFrame>> {
        match parse_frame(self.receive_buffer.parsing_context()) {
            Ok((i, f)) => {
                let consumed = self.receive_buffer.offset(i);
                let frame_max = self.configuration.frame_max() as usize;
                if frame_max > 0 && consumed > frame_max {
                    error!("received large ({} bytes) frame", consumed);
                    let error = AMQPError::new(
                        AMQPHardError::FRAMEERROR.into(),
                        format!("frame too large: {} bytes", consumed).into(),
                    );
                    self.internal_rpc.handle().close_connection(
                        error.get_id(),
                        error.get_message().to_string(),
                        0,
                        0,
                    );
                    self.poll_internal_rpc()?;
                    self.critical_error(Error::ProtocolError(error))?;
                }
                self.receive_buffer.consume(consumed);
                Ok(Some(f))
            }
            Err(e) => {
                if !e.is_incomplete() {
                    error!("parse error: {:?}", e);
                    self.critical_error(Error::ParsingError(e))?;
                }
                Ok(None)
            }
        }
    }
}
<commit_msg>simplify bad_reads handling<commit_after>use crate::{
    buffer::Buffer,
    channels::Channels,
    connection_status::ConnectionState,
    frames::Frames,
    heartbeat::Heartbeat,
    internal_rpc::InternalRPC,
    protocol::{self, AMQPError, AMQPHardError},
    reactor::{ReactorBuilder, ReactorHandle, Slot},
    socket_state::SocketState,
    tcp::HandshakeResult,
    thread::ThreadHandle,
    Configuration, ConnectionStatus, Error, PromiseResolver, Result, TcpStream,
};
use amq_protocol::frame::{gen_frame, parse_frame, AMQPFrame, GenError};
use log::{debug, error, trace};
use std::{
    collections::VecDeque,
    convert::TryFrom,
    io::{self, Write},
    sync::Arc,
    thread::Builder as ThreadBuilder,
    time::Duration,
};

const FRAMES_STORAGE: usize = 32;

#[derive(Debug, PartialEq)]
enum Status {
    Initial,
    Connected,
    Stop,
}

pub struct IoLoop {
    connection_status: ConnectionStatus,
    configuration: Configuration,
    channels: Channels,
    internal_rpc: InternalRPC,
    frames: Frames,
    heartbeat: Heartbeat,
    socket_state: SocketState,
    reactor: Box<dyn ReactorHandle + Send>,
    reactor_thread_handle: ThreadHandle,
    connection_io_loop_handle: ThreadHandle,
    stream: TcpStream,
    slot: Slot,
    status: Status,
    frame_size: usize,
    receive_buffer: Buffer,
    send_buffer: Buffer,
    serialized_frames: VecDeque<(u64, Option<PromiseResolver<()>>)>,
}

impl IoLoop {
    #[allow(clippy::too_many_arguments)]
    pub(crate) fn new(
        connection_status: ConnectionStatus,
        configuration: Configuration,
        channels: Channels,
        internal_rpc: InternalRPC,
        frames: Frames,
        socket_state: SocketState,
        connection_io_loop_handle: ThreadHandle,
        stream: HandshakeResult,
        reactor_builder: &dyn ReactorBuilder,
    ) -> Result<Self> {
        let mut stream = TcpStream::try_from(stream)?;
        let heartbeat = Heartbeat::new(channels.clone());
        let mut reactor = reactor_builder.build(heartbeat.clone())?;
        let reactor_handle = reactor.handle();
        let frame_size = std::cmp::max(
            protocol::constants::FRAME_MIN_SIZE as usize,
            configuration.frame_max() as usize,
        );
        let slot = reactor.register(stream.inner_mut(), socket_state.handle())?;

        let reactor_thread_handle = reactor.start()?;

        Ok(Self {
            connection_status,
            configuration,
            channels,
            internal_rpc,
            frames,
            heartbeat,
            socket_state,
            reactor: reactor_handle,
            reactor_thread_handle,
            connection_io_loop_handle,
            stream,
            slot,
            status: Status::Initial,
            frame_size,
            receive_buffer: Buffer::with_capacity(FRAMES_STORAGE * frame_size),
            send_buffer: Buffer::with_capacity(FRAMES_STORAGE * frame_size),
            serialized_frames: VecDeque::default(),
        })
    }

    fn finish_setup(&mut self) -> Result<bool> {
        if self.connection_status.connected() {
            let frame_max = self.configuration.frame_max() as usize;
            self.frame_size = std::cmp::max(self.frame_size, frame_max);
            self.receive_buffer.grow(FRAMES_STORAGE * self.frame_size);
            self.send_buffer.grow(FRAMES_STORAGE * self.frame_size);
            let heartbeat = self.configuration.heartbeat();
            if heartbeat != 0 {
                let heartbeat = Duration::from_millis(u64::from(heartbeat) * 500); // * 1000 (ms) / 2 (half the negotiated timeout)
                self.heartbeat.set_timeout(heartbeat);
                self.reactor.start_heartbeat();
            }
            let peer = self.stream.inner().peer_addr()?;
            debug!("Connected to {}", peer);
            self.status = Status::Connected;
        }
        Ok(true)
    }

    fn ensure_setup(&mut self) -> Result<bool> {
        match self.status {
            Status::Initial => self.finish_setup(),
            Status::Connected => Ok(true),
            Status::Stop => Ok(false),
        }
    }

    fn has_data(&self) -> bool {
        self.frames.has_pending()
            || self.send_buffer.available_data() > 0
            || !self.serialized_frames.is_empty()
    }

    fn can_write(&mut self) -> bool {
        self.socket_state.writable() && self.has_data() && !self.connection_status.blocked()
    }

    fn can_read(&mut self) -> bool {
        self.socket_state.readable() && self.receive_buffer.available_space() > 0
    }

    fn can_parse(&self) -> bool {
        self.receive_buffer.available_data() > 0
    }

    fn should_continue(&self) -> bool {
        (self.status != Status::Connected
            || self.connection_status.connected()
            || self.connection_status.closing())
            && self.status != Status::Stop
            && !self.connection_status.errored()
    }

    pub fn start(mut self) -> Result<()> {
        let waker = self.socket_state.handle();
        let io_loop_handle = self.connection_io_loop_handle.clone();
        io_loop_handle.register(
            ThreadBuilder::new()
                .name("lapin-io-loop".to_owned())
                .spawn(move || {
                    while self.should_continue() {
                        if let Err(err) = self.run() {
                            self.critical_error(err)?;
                        }
                    }
                    self.heartbeat.cancel();
                    self.reactor.shutdown()?;
                    self.reactor_thread_handle.wait("reactor")
                })?,
        );
        waker.wake();
        Ok(())
    }

    fn poll_internal_rpc(&self) -> Result<()> {
        self.internal_rpc.poll(&self.channels)
    }

    fn poll_socket_events(&mut self) -> Result<()> {
        self.socket_state.poll_events();
        if self.socket_state.error() {
            self.critical_error(io::Error::from(io::ErrorKind::ConnectionAborted).into())?;
        }
        if self.socket_state.closed() {
            self.critical_error(io::Error::from(io::ErrorKind::ConnectionReset).into())?;
        }
        self.poll_internal_rpc()
    }

    fn run(&mut self) -> Result<()> {
        trace!("io_loop run");
        self.poll_socket_events()?;
        if !self.ensure_setup()? {
            return Ok(());
        }
        trace!(
            "io_loop do_run; can_read={}, can_write={}, has_data={}",
            self.socket_state.readable(),
            self.socket_state.writable(),
            self.has_data()
        );
        if !self.can_read() && !self.can_write() {
            self.socket_state.wait();
        }
        self.poll_socket_events()?;
        if self.stream.is_handshaking() {
            self.stream.handshake()?;
            if self.stream.is_handshaking() {
                // We hit WOULDBLOCK while handshaking, wait for the next socket event
                return Ok(());
            }
        }
        self.attempt_flush()?;
        self.write()?;
        if self.connection_status.closed() {
            self.status = Status::Stop;
        }
        if self.should_continue() {
            self.read()?;
        }
        self.handle_frames()?;
        trace!(
            "io_loop do_run done; can_read={}, can_write={}, has_data={}, status={:?}",
            self.socket_state.readable(),
            self.socket_state.writable(),
            self.has_data(),
            self.status
        );
        self.poll_internal_rpc()
    }

    fn critical_error(&mut self, error: Error) -> Result<()> {
        if let Some(resolver) = self.connection_status.connection_resolver() {
            resolver.swear(Err(error.clone()));
        }
        self.status = Status::Stop;
        self.channels.set_connection_error(error.clone())?;
        for (_, resolver) in std::mem::take(&mut self.serialized_frames) {
            if let Some(resolver) = resolver {
                resolver.swear(Err(error.clone()));
            }
        }
        self.reactor.shutdown()?;
        Err(error)
    }

    fn attempt_flush(&mut self) -> Result<()> {
        let res = self.flush();
        if let Err(e) = self.socket_state.handle_flush_result(res) {
            error!("error flushing: {:?}", e);
            self.critical_error(e)?;
        }
        self.poll_internal_rpc()
    }

    fn handle_read_result(&mut self, result: Result<()>) -> Result<()> {
        if let Err(e) = self
            .socket_state
            .handle_read_result(result, &*self.reactor, self.slot)
        {
            error!("error reading: {:?}", e);
            self.critical_error(e)?;
        }
        self.poll_internal_rpc()
    }

    fn handle_write_result(&mut self, result: Result<()>) -> Result<()> {
        if let Err(e) = self
            .socket_state
            .handle_write_result(result, &*self.reactor, self.slot)
        {
            error!("error writing: {:?}", e);
            self.critical_error(e)?;
        }
        self.poll_internal_rpc()
    }

    fn flush(&mut self) -> Result<()> {
        self.stream.flush()?;
        self.poll_internal_rpc()
    }

    fn write(&mut self) -> Result<()> {
        while self.can_write() {
            let res = self.write_to_stream();
            self.handle_write_result(res)?;
        }
        self.poll_internal_rpc()
    }

    fn read(&mut self) -> Result<()> {
        while self.can_read() {
            let res = self.read_from_stream();
            self.handle_read_result(res)?;
        }
        self.poll_internal_rpc()
    }

    fn write_to_stream(&mut self) -> Result<()> {
        self.flush()?;
        self.serialize()?;

        let sz = self.send_buffer.write_to(&mut self.stream)?;

        if sz > 0 {
            self.heartbeat.update_last_write();

            trace!("wrote {} bytes", sz);
            self.send_buffer.consume(sz);

            let mut written = sz as u64;
            while written > 0 {
                if let Some((to_write, resolver)) = self.serialized_frames.pop_front() {
                    if written < to_write {
                        self.serialized_frames
                            .push_front((to_write - written, resolver));
                        trace!("{} to write to complete this frame", to_write - written);
                        written = 0;
                    } else {
                        if let Some(resolver) = resolver {
                            resolver.swear(Ok(()));
                        }
                        written -= to_write;
                    }
                } else {
                    error!(
                        "We've written {} but didn't expect to write anything",
                        written
                    );
                    break;
                }
            }

            if self.send_buffer.available_data() > 0 {
                // We didn't write all the data yet
                trace!("Still {} to write", self.send_buffer.available_data());
            }

            self.flush()?;
        } else {
            error!("Socket was writable but we wrote 0, marking as wouldblock");
            self.handle_write_result(Err(io::Error::from(io::ErrorKind::WouldBlock).into()))?;
        }
        self.poll_internal_rpc()
    }

    fn read_from_stream(&mut self) -> Result<()> {
        match self.connection_status.state() {
            ConnectionState::Closed => Ok(()),
            ConnectionState::Error => Err(Error::InvalidConnectionState(ConnectionState::Error)),
            _ => {
                let sz = self.receive_buffer.read_from(&mut self.stream)?;

                if sz > 0 {
                    trace!("read {} bytes", sz);
                    self.receive_buffer.fill(sz);
                } else {
                    error!("Socket was readable but we read 0. This usually means that the connection is half closed this mark it as broken");
                    self.handle_read_result(Err(io::Error::from(
                        io::ErrorKind::ConnectionAborted,
                    )
                    .into()))?;
                }
                self.poll_internal_rpc()
            }
        }
    }

    fn serialize(&mut self) -> Result<()> {
        while let Some((next_msg, resolver)) = self.frames.pop(self.channels.flow()) {
            trace!("will write to buffer: {}", next_msg);
            let checkpoint = self.send_buffer.checkpoint();
            let res = gen_frame(&next_msg)((&mut self.send_buffer).into());
            match res.map(|w| w.into_inner().1) {
                Ok(sz) => self.serialized_frames.push_back((sz, resolver)),
                Err(e) => {
                    self.send_buffer.rollback(checkpoint);
                    match e {
                        GenError::BufferTooSmall(_) => {
                            // Requeue msg
                            self.frames.retry((next_msg, resolver));
                            break;
                        }
                        e => {
                            error!("error generating frame: {:?}", e);
                            self.critical_error(Error::SerialisationError(Arc::new(e)))?;
                        }
                    }
                }
            }
        }
        self.poll_internal_rpc()
    }

    fn handle_frames(&mut self) -> Result<()> {
        while self.can_parse() {
            if let Some(frame) = self.parse()? {
                self.channels.handle_frame(frame)?;
            } else {
                break;
            }
        }
        self.poll_internal_rpc()
    }

    fn parse(&mut self) -> Result<Option<AMQPFrame>> {
        match parse_frame(self.receive_buffer.parsing_context()) {
            Ok((i, f)) => {
                let consumed = self.receive_buffer.offset(i);
                let frame_max = self.configuration.frame_max() as usize;
                if frame_max > 0 && consumed > frame_max {
                    error!("received large ({} bytes) frame", consumed);
                    let error = AMQPError::new(
                        AMQPHardError::FRAMEERROR.into(),
                        format!("frame too large: {} bytes", consumed).into(),
                    );
                    self.internal_rpc.handle().close_connection(
                        error.get_id(),
                        error.get_message().to_string(),
                        0,
                        0,
                    );
                    self.poll_internal_rpc()?;
                    self.critical_error(Error::ProtocolError(error))?;
                }
                self.receive_buffer.consume(consumed);
                Ok(Some(f))
            }
            Err(e) => {
                if !e.is_incomplete() {
                    error!("parse error: {:?}", e);
                    self.critical_error(Error::ParsingError(e))?;
                }
                Ok(None)
            }
        }
    }
}
<|endoftext|>