<commit_before>use libc;
use std::simd;

use stack::Stack;
use mutex::Mutex;

/// Get the offset in bytes of some particular struct member.
///
/// Idea is simply to see what the address of that member would be if the
/// struct was created at address 0x0 (which gives its offset in bytes).
macro_rules! offset_of {
    ($kind:ty, $member:ident) => {
        unsafe {
            let ptr = 0x0usize as *const $kind;
            let member_addr = (&(*ptr).$member as *const _) as usize;

            member_addr
        }
    };
}

#[inline(always)]
pub unsafe fn get_tls_mem<T>(offset: usize) -> T {
    let dest: T;

    // We can't use constant segment offsets here due to some odd
    // asm! behavior, so just use indirect (it's slower, oh well).
    asm!("mov %fs:($1), $0"
         : "=r"(dest)
         : "r" (offset)
         :: "volatile");

    dest
}

#[inline(always)]
pub unsafe fn set_tls_mem<T>(offset: usize, expr: T) {
    asm!("movl $1, %fs:($0)" :
         : "r"(offset), "r"(expr)
         :: "volatile");
}

/// TCB (thread control block). We unfortunately need to match glibc here, so
/// there's a ton of unused vars.
#[repr(C, packed)]
#[derive(Debug)]
pub struct TcbHead {
    pub tcb: *mut libc::c_void,
    pub dtv: *mut libc::c_void,
    pub thread_self: *mut Thread,
    pub multiple_threads: i32,
    pub gscope_flag: i32,
    pub sysinfo: *mut u32,
    pub stack_guard: *mut u32,
    pub pointer_guard: *mut u32,
    pub vgetcpu_cache: [*mut u64; 2],

    __unused_1: [u32; 2],
    __unused_2: [*const libc::c_void; 5],
    __unused_3: u64,
    __unused_4: [simd::f32x4; 32],
    __padding: [*const libc::c_void; 8],
}

#[repr(C, packed)]
#[derive(Debug)]
pub struct Thread {
    pub header: TcbHead,

    // required by glibc
    _list: [*const libc::c_void; 2],

    /// This thread's id
    pub tid: libc::pid_t,
    /// This thread's parent pid
    pub pid: libc::pid_t,

    /// more glibc requirements
    /// TODO: find exact size
    _padding: [*const libc::c_void; 10],

    pub stack: *const Stack,
    pub magic: usize,
    pub mutex: Mutex,
}

pub fn get_current_thread() -> &'static Thread {
    let thread_offset = offset_of!(TcbHead, thread_self);

    let thd_ptr: *mut Thread = unsafe { get_tls_mem(thread_offset) };
    if thd_ptr.is_null() {
        panic!("TLS thread return NULL");
    }

    let thd = unsafe { &*thd_ptr };

    thd
}
<commit_msg>add thread::yield_now<commit_after>/// FIXME: TLS / Thread naming distinction needs to be resolved

use libc;
use std;

use std::simd;
use std::ptr::null_mut;

use stack::Stack;
use mutex::Mutex;
use ffi;

/// Get the offset in bytes of some particular struct member.
///
/// Idea is simply to see what the address of that member would be if the
/// struct was created at address 0x0 (which gives its offset in bytes).
macro_rules! offset_of {
    ($kind:ty, $member:ident) => {
        unsafe {
            let ptr = 0x0usize as *const $kind;
            let member_addr = (&(*ptr).$member as *const _) as usize;

            member_addr
        }
    };
}

#[inline(always)]
pub unsafe fn get_tls_mem<T>(offset: usize) -> T {
    let dest: T;

    // We can't use constant segment offsets here due to some odd
    // asm! behavior, so just use indirect (it's slower, oh well).
    asm!("mov %fs:($1), $0"
         : "=r"(dest)
         : "r" (offset)
         :: "volatile");

    dest
}

#[inline(always)]
pub unsafe fn set_tls_mem<T>(offset: usize, expr: T) {
    asm!("movl $1, %fs:($0)" :
         : "r"(offset), "r"(expr)
         :: "volatile");
}

/// TCB (thread control block). We unfortunately need to match glibc here, so
/// there's a ton of unused vars.
#[repr(C, packed)]
#[derive(Debug)]
pub struct TcbHead {
    pub tcb: *mut libc::c_void,
    pub dtv: *mut libc::c_void,
    pub thread_self: *mut Thread,
    pub multiple_threads: i32,
    pub gscope_flag: i32,
    pub sysinfo: *mut u32,
    pub stack_guard: *mut u32,
    pub pointer_guard: *mut u32,
    pub vgetcpu_cache: [*mut u64; 2],

    __unused_1: [u32; 2],
    __unused_2: [*const libc::c_void; 5],
    __unused_3: u64,
    __unused_4: [simd::f32x4; 32],
    __padding: [*const libc::c_void; 8],
}

#[repr(C, packed)]
#[derive(Debug)]
pub struct Thread {
    pub header: TcbHead,

    // required by glibc
    _list: [*const libc::c_void; 2],

    /// This thread's id
    pub tid: libc::pid_t,
    /// This thread's parent pid
    pub pid: libc::pid_t,

    /// more glibc requirements
    /// TODO: find exact size
    _padding: [*const libc::c_void; 10],

    pub stack: *const Stack,
    pub magic: usize,
    pub mutex: Mutex,
}

pub fn get_current_thread() -> &'static Thread {
    let thread_offset = offset_of!(TcbHead, thread_self);

    let thd_ptr: *mut Thread = unsafe { get_tls_mem(thread_offset) };
    if thd_ptr.is_null() {
        panic!("TLS thread return NULL");
    }

    let thd = unsafe { &*thd_ptr };

    thd
}

/// Have the current thread voluntarily yield the rest of it's scheduled run
/// time.
pub fn yield_now() {
    extern { fn sched_yield() -> libc::c_int; }

    unsafe { sched_yield(); }
}
<|endoftext|>


<commit_before>use crate::exercise::{CompiledExercise, Exercise, Mode, State};
use console::style;
use indicatif::ProgressBar;

pub fn verify<'a>(start_at: impl IntoIterator<Item = &'a Exercise>) -> Result<(), &'a Exercise> {
    for exercise in start_at {
        let compile_result = match exercise.mode {
            Mode::Test => compile_and_test(&exercise, RunMode::Interactive),
            Mode::Compile => compile_and_run_interactively(&exercise),
            Mode::Clippy => compile_only(&exercise),
        };
        if !compile_result.unwrap_or(false) {
            return Err(exercise);
        }
    }
    Ok(())
}

enum RunMode {
    Interactive,
    NonInteractive,
}

pub fn test(exercise: &Exercise) -> Result<(), ()> {
    compile_and_test(exercise, RunMode::NonInteractive)?;
    Ok(())
}

fn compile_only(exercise: &Exercise) -> Result<bool, ()> {
    let progress_bar = ProgressBar::new_spinner();
    progress_bar.set_message(format!("Compiling {}...", exercise).as_str());
    progress_bar.enable_steady_tick(100);

    let _ = compile(&exercise, &progress_bar)?;
    progress_bar.finish_and_clear();

    success!("Successfully compiled {}!", exercise);
    Ok(prompt_for_completion(&exercise, None))
}

fn compile_and_run_interactively(exercise: &Exercise) -> Result<bool, ()> {
    let progress_bar = ProgressBar::new_spinner();
    progress_bar.set_message(format!("Compiling {}...", exercise).as_str());
    progress_bar.enable_steady_tick(100);

    let compilation = compile(&exercise, &progress_bar)?;

    progress_bar.set_message(format!("Running {}...", exercise).as_str());
    let result = compilation.run();
    progress_bar.finish_and_clear();

    let output = match result {
        Ok(output) => output,
        Err(output) => {
            warn!("Ran {} with errors", exercise);
            println!("{}", output.stdout);
            return Err(());
        }
    };

    success!("Successfully ran {}!", exercise);

    Ok(prompt_for_completion(&exercise, Some(output.stdout)))
}

fn compile_and_test(exercise: &Exercise, run_mode: RunMode) -> Result<bool, ()> {
    let progress_bar = ProgressBar::new_spinner();
    progress_bar.set_message(format!("Testing {}...", exercise).as_str());
    progress_bar.enable_steady_tick(100);

    let compilation = compile(exercise, &progress_bar)?;
    let result = compilation.run();
    progress_bar.finish_and_clear();

    match result {
        Ok(_) => {
            success!("Successfully tested {}", &exercise);
            if let RunMode::Interactive = run_mode {
                Ok(prompt_for_completion(&exercise, None))
            } else {
                Ok(true)
            }
        }
        Err(output) => {
            warn!(
                "Testing of {} failed! Please try again. Here's the output:",
                exercise
            );
            println!("{}", output.stdout);
            Err(())
        }
    }
}

fn compile<'a, 'b>(
    exercise: &'a Exercise,
    progress_bar: &'b ProgressBar,
) -> Result<CompiledExercise<'a>, ()> {
    let compilation_result = exercise.compile();

    match compilation_result {
        Ok(compilation) => Ok(compilation),
        Err(output) => {
            progress_bar.finish_and_clear();
            warn!(
                "Compiling of {} failed! Please try again. Here's the output:",
                exercise
            );
            println!("{}", output.stderr);
            Err(())
        }
    }
}

fn prompt_for_completion(exercise: &Exercise, prompt_output: Option<String>) -> bool {
    let context = match exercise.state() {
        State::Done => return true,
        State::Pending(context) => context,
    };

    let success_msg = match exercise.mode {
        Mode::Compile => "The code is compiling!",
        Mode::Test => "The code is compiling, and the tests pass!",
        Mode::Clippy => "The code is compiling, and 📎 Clippy 📎 is happy!",
    };

    println!("");
    println!("🎉 🎉  {} 🎉 🎉", success_msg);
    println!("");

    if let Some(output) = prompt_output {
        println!("Output:");
        println!("{}", separator());
        println!("{}", output);
        println!("{}", separator());
        println!("");
    }

    println!("You can keep working on this exercise,");
    println!(
        "or jump into the next one by removing the {} comment:",
        style("`I AM NOT DONE`").bold()
    );
    println!("");
    for context_line in context {
        let formatted_line = if context_line.important {
            format!("{}", style(context_line.line).bold())
        } else {
            format!("{}", context_line.line)
        };

        println!(
            "{:>2} {}  {}",
            style(context_line.number).blue().bold(),
            style("|").blue(),
            formatted_line
        );
    }

    false
}

fn separator() -> console::StyledObject<&'static str> {
    style("====================").bold()
}
<commit_msg>Use .to_string rather than format macro<commit_after>use crate::exercise::{CompiledExercise, Exercise, Mode, State};
use console::style;
use indicatif::ProgressBar;

pub fn verify<'a>(start_at: impl IntoIterator<Item = &'a Exercise>) -> Result<(), &'a Exercise> {
    for exercise in start_at {
        let compile_result = match exercise.mode {
            Mode::Test => compile_and_test(&exercise, RunMode::Interactive),
            Mode::Compile => compile_and_run_interactively(&exercise),
            Mode::Clippy => compile_only(&exercise),
        };
        if !compile_result.unwrap_or(false) {
            return Err(exercise);
        }
    }
    Ok(())
}

enum RunMode {
    Interactive,
    NonInteractive,
}

pub fn test(exercise: &Exercise) -> Result<(), ()> {
    compile_and_test(exercise, RunMode::NonInteractive)?;
    Ok(())
}

fn compile_only(exercise: &Exercise) -> Result<bool, ()> {
    let progress_bar = ProgressBar::new_spinner();
    progress_bar.set_message(format!("Compiling {}...", exercise).as_str());
    progress_bar.enable_steady_tick(100);

    let _ = compile(&exercise, &progress_bar)?;
    progress_bar.finish_and_clear();

    success!("Successfully compiled {}!", exercise);
    Ok(prompt_for_completion(&exercise, None))
}

fn compile_and_run_interactively(exercise: &Exercise) -> Result<bool, ()> {
    let progress_bar = ProgressBar::new_spinner();
    progress_bar.set_message(format!("Compiling {}...", exercise).as_str());
    progress_bar.enable_steady_tick(100);

    let compilation = compile(&exercise, &progress_bar)?;

    progress_bar.set_message(format!("Running {}...", exercise).as_str());
    let result = compilation.run();
    progress_bar.finish_and_clear();

    let output = match result {
        Ok(output) => output,
        Err(output) => {
            warn!("Ran {} with errors", exercise);
            println!("{}", output.stdout);
            return Err(());
        }
    };

    success!("Successfully ran {}!", exercise);

    Ok(prompt_for_completion(&exercise, Some(output.stdout)))
}

fn compile_and_test(exercise: &Exercise, run_mode: RunMode) -> Result<bool, ()> {
    let progress_bar = ProgressBar::new_spinner();
    progress_bar.set_message(format!("Testing {}...", exercise).as_str());
    progress_bar.enable_steady_tick(100);

    let compilation = compile(exercise, &progress_bar)?;
    let result = compilation.run();
    progress_bar.finish_and_clear();

    match result {
        Ok(_) => {
            success!("Successfully tested {}", &exercise);
            if let RunMode::Interactive = run_mode {
                Ok(prompt_for_completion(&exercise, None))
            } else {
                Ok(true)
            }
        }
        Err(output) => {
            warn!(
                "Testing of {} failed! Please try again. Here's the output:",
                exercise
            );
            println!("{}", output.stdout);
            Err(())
        }
    }
}

fn compile<'a, 'b>(
    exercise: &'a Exercise,
    progress_bar: &'b ProgressBar,
) -> Result<CompiledExercise<'a>, ()> {
    let compilation_result = exercise.compile();

    match compilation_result {
        Ok(compilation) => Ok(compilation),
        Err(output) => {
            progress_bar.finish_and_clear();
            warn!(
                "Compiling of {} failed! Please try again. Here's the output:",
                exercise
            );
            println!("{}", output.stderr);
            Err(())
        }
    }
}

fn prompt_for_completion(exercise: &Exercise, prompt_output: Option<String>) -> bool {
    let context = match exercise.state() {
        State::Done => return true,
        State::Pending(context) => context,
    };

    let success_msg = match exercise.mode {
        Mode::Compile => "The code is compiling!",
        Mode::Test => "The code is compiling, and the tests pass!",
        Mode::Clippy => "The code is compiling, and 📎 Clippy 📎 is happy!",
    };

    println!();
    println!("🎉 🎉  {} 🎉 🎉", success_msg);
    println!();

    if let Some(output) = prompt_output {
        println!("Output:");
        println!("{}", separator());
        println!("{}", output);
        println!("{}", separator());
        println!();
    }

    println!("You can keep working on this exercise,");
    println!(
        "or jump into the next one by removing the {} comment:",
        style("`I AM NOT DONE`").bold()
    );
    println!();
    for context_line in context {
        let formatted_line = if context_line.important {
            format!("{}", style(context_line.line).bold())
        } else {
            context_line.line.to_string()
        };

        println!(
            "{:>2} {}  {}",
            style(context_line.number).blue().bold(),
            style("|").blue(),
            formatted_line
        );
    }

    false
}

fn separator() -> console::StyledObject<&'static str> {
    style("====================").bold()
}
<|endoftext|>


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

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

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

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

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

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

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

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

    delta_time: f64, // seconds
    last_frame: Instant,

    render_timer: FrameTimer,
}

impl GltfViewer {
    pub fn new(source: &str, width: u32, height: u32, headless: bool) -> GltfViewer {
        let (events_loop, gl_window, width, height) =
            if headless {
                let headless_context = glutin::HeadlessRendererBuilder::new(width, height).build().unwrap();
                unsafe { headless_context.make_current().unwrap() }
                gl::load_with(|symbol| headless_context.get_proc_address(symbol) as *const _);
                let framebuffer = Framebuffer::new(width, height);
                framebuffer.bind();
                unsafe { gl::Viewport(0, 0, width as i32, height as i32); }

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

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

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

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

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

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

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

        // TODO!!: tmp duplicated camera controls
        let mut controls = CameraControls {
            position: Point3::new(0.0, 0.0, 2.0),
            camera: Camera::default(),
            ..CameraControls::default()
        };
        controls.camera.fovy = 60.0;
        controls.camera.update_aspect_ratio(width as f32 / height as f32); // updates projection matrix

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

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

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

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

            gl::Enable(gl::DEPTH_TEST);

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

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

            controls,
            orbit_controls,
            first_mouse, last_x, last_y,

            events_loop,
            gl_window,

            root,
            scene,

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

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

        // TODO!!!: completely broken (using gltf camera)
        if false && !viewer.root.camera_nodes.is_empty() {
            // Take first camera node
            let cam_node = &viewer.root.get_camera_node(0);
            viewer.controls.set_camera(
                cam_node.camera.as_ref().unwrap(),
                &cam_node.final_transform);
        } else {
            viewer.set_camera_from_bounds();
        }

        viewer
    }

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

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

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

        (root, scene)
    }

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

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

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

        // TODO!! TMP duplication (orbit controls)
        self.orbit_controls.position = cam_pos;
        self.orbit_controls.target = Point3::from_vec(center);

        self.controls.position = cam_pos;
        self.controls.center = Some(Point3::from_vec(center));

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

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

            // events
            let keep_running = process_events(
                &mut self.events_loop.as_mut().unwrap(), self.gl_window.as_mut().unwrap(),
                &mut self.controls, &mut self.orbit_controls,
                &mut self.width, &mut self.height);
            if !keep_running { break }

            self.controls.update(self.delta_time); // navigation

            self.draw();

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

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

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

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

            self.render_timer.end();
        }
    }

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

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

        let img = img.flipv();

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

#[allow(too_many_arguments)]
fn process_events(
    events_loop: &mut glutin::EventsLoop,
    gl_window: &glutin::GlWindow,
    mut controls: &mut CameraControls,
    orbit_controls: &mut OrbitControls,
    width: &mut u32,
    height: &mut u32) -> bool
{
    let mut keep_running = true;
    #[allow(single_match)]
    events_loop.poll_events(|event| {
        match event {
            glutin::Event::WindowEvent{ event, .. } => match event {
                WindowEvent::Closed => keep_running = false,
                WindowEvent::Resized(w, h) => {
                    gl_window.resize(w, h);
                    *width = w;
                    *height = h;
                    let w = w as f32;
                    let h = h as f32;
                    orbit_controls.camera.update_aspect_ratio(w / h);
                    orbit_controls.screen_width = w;
                    orbit_controls.screen_height = h;
                },
                WindowEvent::DroppedFile(_path_buf) => (), // TODO: drag file in
                WindowEvent::MouseInput { button, state: Pressed, ..} => {
                    match button {
                        MouseButton::Left => {
                            orbit_controls.state = NavState::Rotating;
                        },
                        MouseButton::Right => {
                            orbit_controls.state = NavState::Panning;
                        },
                        _ => ()
                    }
                },
                WindowEvent::MouseInput { button, state: Released, ..} => {
                    match (button, orbit_controls.state.clone()) {
                        (MouseButton::Left, NavState::Rotating) | (MouseButton::Right, NavState::Panning) => {
                            orbit_controls.state = NavState::None;
                            orbit_controls.handle_mouse_up();
                        },
                        _ => ()
                    }
                }
                WindowEvent::CursorMoved { position: (xpos, ypos), .. } => {
                    let (xpos, ypos) = (xpos as f32, ypos as f32);
                    orbit_controls.handle_mouse_move(xpos, ypos);
                },
                WindowEvent::MouseWheel { delta: MouseScrollDelta::PixelDelta(_xoffset, yoffset), .. } => {
                    // TODO: need to handle LineDelta case too?
                    orbit_controls.process_mouse_scroll(yoffset);
                }
                WindowEvent::KeyboardInput { input, .. } => {
                    keep_running = process_input(input, &mut controls);
                }
                _ => ()
            },
            _ => ()
        }
    });

    keep_running
}

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

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

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

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

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

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

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

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

    delta_time: f64, // seconds
    last_frame: Instant,

    render_timer: FrameTimer,
}

impl GltfViewer {
    pub fn new(source: &str, width: u32, height: u32, headless: bool) -> GltfViewer {
        let (events_loop, gl_window, width, height) =
            if headless {
                let headless_context = glutin::HeadlessRendererBuilder::new(width, height).build().unwrap();
                unsafe { headless_context.make_current().unwrap() }
                gl::load_with(|symbol| headless_context.get_proc_address(symbol) as *const _);
                let framebuffer = Framebuffer::new(width, height);
                framebuffer.bind();
                unsafe { gl::Viewport(0, 0, width as i32, height as i32); }

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

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

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

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

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

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

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

        // TODO!!: tmp duplicated camera controls
        let mut controls = CameraControls {
            position: Point3::new(0.0, 0.0, 2.0),
            camera: Camera::default(),
            ..CameraControls::default()
        };
        controls.camera.fovy = 60.0;
        controls.camera.update_aspect_ratio(width as f32 / height as f32); // updates projection matrix

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

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

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

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

            gl::Enable(gl::DEPTH_TEST);

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

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

            controls,
            orbit_controls,
            first_mouse, last_x, last_y,

            events_loop,
            gl_window,

            root,
            scene,

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

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

        // TODO!!!: completely broken (using gltf camera)
        if false && !viewer.root.camera_nodes.is_empty() {
            // Take first camera node
            let cam_node = &viewer.root.get_camera_node(0);
            viewer.controls.set_camera(
                cam_node.camera.as_ref().unwrap(),
                &cam_node.final_transform);
        } else {
            viewer.set_camera_from_bounds();
        }

        viewer
    }

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

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

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

        (root, scene)
    }

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

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

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

        // TODO!! TMP duplication (orbit controls)
        self.orbit_controls.position = cam_pos;
        self.orbit_controls.target = Point3::from_vec(center);

        self.controls.position = cam_pos;
        self.controls.center = Some(Point3::from_vec(center));

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

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

            // events
            let keep_running = process_events(
                &mut self.events_loop.as_mut().unwrap(), self.gl_window.as_mut().unwrap(),
                &mut self.controls, &mut self.orbit_controls,
                &mut self.width, &mut self.height);
            if !keep_running { break }

            self.controls.update(self.delta_time); // navigation

            self.draw();

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

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

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

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

            self.render_timer.end();
        }
    }

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

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

        let img = img.flipv();

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

#[allow(too_many_arguments)]
fn process_events(
    events_loop: &mut glutin::EventsLoop,
    gl_window: &glutin::GlWindow,
    mut controls: &mut CameraControls,
    orbit_controls: &mut OrbitControls,
    width: &mut u32,
    height: &mut u32) -> bool
{
    let mut keep_running = true;
    #[allow(single_match)]
    events_loop.poll_events(|event| {
        match event {
            glutin::Event::WindowEvent{ event, .. } => match event {
                WindowEvent::Closed => keep_running = false,
                WindowEvent::Resized(w, h) => {
                    gl_window.resize(w, h);
                    *width = w;
                    *height = h;
                    let w = w as f32;
                    let h = h as f32;
                    orbit_controls.camera.update_aspect_ratio(w / h);
                    orbit_controls.screen_width = w;
                    orbit_controls.screen_height = h;
                },
                WindowEvent::DroppedFile(_path_buf) => (), // TODO: drag file in
                WindowEvent::MouseInput { button, state: Pressed, ..} => {
                    match button {
                        MouseButton::Left => {
                            orbit_controls.state = NavState::Rotating;
                        },
                        MouseButton::Right => {
                            orbit_controls.state = NavState::Panning;
                        },
                        _ => ()
                    }
                },
                WindowEvent::MouseInput { button, state: Released, ..} => {
                    match (button, orbit_controls.state.clone()) {
                        (MouseButton::Left, NavState::Rotating) | (MouseButton::Right, NavState::Panning) => {
                            orbit_controls.state = NavState::None;
                            orbit_controls.handle_mouse_up();
                        },
                        _ => ()
                    }
                }
                WindowEvent::CursorMoved { position: (xpos, ypos), .. } => {
                    let (xpos, ypos) = (xpos as f32, ypos as f32);
                    orbit_controls.handle_mouse_move(xpos, ypos);
                },
                WindowEvent::MouseWheel { delta: MouseScrollDelta::PixelDelta(_xoffset, yoffset), .. } => {
                    // TODO: need to handle LineDelta case too?
                    orbit_controls.process_mouse_scroll(yoffset);
                }
                WindowEvent::KeyboardInput { input, .. } => {
                    keep_running = process_input(input, &mut controls);
                }
                _ => ()
            },
            _ => ()
        }
    });

    keep_running
}

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


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

use std::ptr;
use std;
use std::cell::RefCell;
use std::rc::Rc;
use std::collections::HashMap;

use user32;
use winapi::{
    BOOL, CREATESTRUCTW, HBRUSH, HCURSOR, HICON, HMENU, HWND, INT, LPARAM, LRESULT, RECT, UINT,
    WNDCLASSEXW, WPARAM, c_int,
};

use wchar::ToCU16Str;
use instance::Instance;
use resource::*;

pub struct WndClass {
    pub classname: String,
    pub style: usize,
    pub icon: Option<Image>,
    pub icon_small: Option<Image>,
    pub cursor: Option<Image>,
    pub background: HBRUSH,
    pub menu: MenuResource,
    pub cls_extra: isize,
    pub wnd_extra: isize,
}

impl WndClass {
    pub fn register(&self, instance: Instance) -> bool {
        self.menu.with_menu_p(|menu_p| {
            let clsname_u = self.classname.to_c_u16();
            let wcex = WNDCLASSEXW {
                cbSize: std::mem::size_of::<WNDCLASSEXW>() as UINT,
                style: self.style as UINT,
                lpfnWndProc: Some(main_wnd_proc),
                cbClsExtra: self.cls_extra as INT,
                cbWndExtra: self.wnd_extra as INT,
                hInstance: instance.instance,
                hIcon: self.icon.to_handle() as HICON,
                hCursor: self.cursor.to_handle() as HCURSOR,
                hbrBackground: self.background,
                lpszMenuName: menu_p,
                lpszClassName: clsname_u.as_ptr(),
                hIconSm: self.icon_small.to_handle() as HICON,
            };

            let res = unsafe { user32::RegisterClassExW(&wcex) };
            res != 0
        })
    }
}

pub static WS_BORDER: u32 = 0x800000;
pub static WS_CAPTION: u32 = 0xc00000;
pub static WS_CHILD: u32 = 0x40000000;
pub static WS_CHILDWINDOW: u32 = 0x40000000;
pub static WS_CLIPCHILDREN: u32 = 0x2000000;
pub static WS_CLIPSIBLINGS: u32 = 0x4000000;
pub static WS_DISABLED: u32 = 0x8000000;
pub static WS_DLGFRAME: u32 = 0x400000;
pub static WS_GROUP: u32 = 0x20000;
pub static WS_HSCROLL: u32 = 0x100000;
pub static WS_ICONIC: u32 = 0x20000000;
pub static WS_MAXIMIZE: u32 = 0x1000000;
pub static WS_MAXIMIZEBOX: u32 = 0x10000;
pub static WS_MINIMIZE: u32 = 0x20000000;
pub static WS_MINIMIZEBOX: u32 = 0x20000;
pub static WS_OVERLAPPED: u32 = 0;
pub static WS_OVERLAPPEDWINDOW: u32 = 0xcf0000;
pub static WS_POPUP: u32 = 0x80000000;
pub static WS_POPUPWINDOW: u32 = 0x80880000;
pub static WS_SIZEBOX: u32 = 0x40000;
pub static WS_SYSMENU: u32 = 0x80000;
pub static WS_TABSTOP: u32 = 0x10000;
pub static WS_THICKFRAME: u32 = 0x40000;
pub static WS_TILED: u32 = 0;
pub static WS_TILEDWINDOW: u32 = 0xcf0000;
pub static WS_VISIBLE: u32 = 0x10000000;
pub static WS_VSCROLL: u32 = 0x200000;

pub static ES_AUTOHSCROLL: u32 = 128;
pub static ES_AUTOVSCROLL: u32 = 64;
pub static ES_CENTER: u32 = 1;
pub static ES_LEFT: u32 = 0;
pub static ES_LOWERCASE: u32 = 16;
pub static ES_MULTILINE: u32 = 4;
pub static ES_NOHIDESEL: u32 = 256;
pub static ES_NUMBER: u32 = 0x2000;
pub static ES_OEMCONVERT: u32 = 0x400;
pub static ES_PASSWORD: u32 = 32;
pub static ES_READONLY: u32 = 0x800;
pub static ES_RIGHT: u32 = 2;
pub static ES_UPPERCASE: u32 = 8;
pub static ES_WANTRETURN: u32 = 4096;

pub struct WindowParams {
    pub window_name: String,
    pub style: u32,
    pub x: isize,
    pub y: isize,
    pub width: isize,
    pub height: isize,
    pub parent: Window,
    pub menu: HMENU,
    pub ex_style: u32,
}

#[derive(PartialEq, Eq, Hash, Copy)]
pub struct Window {
    pub wnd: HWND,
}

// Sending across threads allows, for example, a worker thread to communicate
// with a UI thread via PostMessage.
unsafe impl Send for Window {}

impl Clone for Window {
    fn clone(&self) -> Window {
        Window {
            wnd: self.wnd,
        }
    }
}

impl Window {
    pub fn null() -> Window {
        Window {
            wnd: ptr::null_mut(),
        }
    }

    pub fn new(
        instance: Instance, wproc: Option<Box<WindowImpl + 'static>>, classname: &str, params: &WindowParams
    ) -> Option<Window> {
        KEY_INIT_WND.with(move |f| *f.borrow_mut() = wproc);

        let wnd = unsafe {
            let clsname_u = classname.to_c_u16();
            let title_u = params.window_name.to_c_u16();
            let wnd = user32::CreateWindowExW(
                params.ex_style, clsname_u.as_ptr(), title_u.as_ptr(), params.style,
                params.x as c_int, params.y as c_int,
                params.width as c_int, params.height as c_int,
                params.parent.wnd, params.menu, instance.instance,
                ptr::null_mut()
            );
            wnd
        };

        if wnd != ptr::null_mut() {
            Some(Window { wnd: wnd })
        } else {
            None
        }
    }

    pub fn show(&self, cmd_show: isize) -> bool {
        unsafe { user32::ShowWindow(self.wnd, cmd_show as c_int) == 0 }
    }

    pub fn show_async(&self, cmd_show: isize) -> bool {
        unsafe { user32::ShowWindowAsync(self.wnd, cmd_show as c_int) == 0 }
    }

    pub fn update(&self) -> bool {
        unsafe { user32::UpdateWindow(self.wnd) == 0 }
    }

    pub fn client_rect(&self) -> Option<RECT> {
        let mut rect = RECT {
            left: 0,
            top: 0,
            right: 0,
            bottom: 0,
        };
        let res = unsafe {
            user32::GetClientRect(self.wnd, &mut rect as *mut RECT)
        } != 0;
        match res {
            true => Some(rect),
            false => None,
        }
    }

    pub fn set_window_pos(
        &self, x: isize, y: isize, width: isize, height: isize, flags: UINT
    ) -> bool {
        // TODO: hwndInsertAfter
        unsafe {
            user32::SetWindowPos(
                self.wnd, ptr::null_mut(), x as c_int, y as c_int,
                width as c_int, height as c_int, flags
            ) != 0
        }
    }

    pub fn set_focus(&self) -> Window {
        unsafe {
            Window {
                wnd: user32::SetFocus(self.wnd)
            }
        }
    }

    pub fn send_message(&self, msg: UINT, wparam: WPARAM, lparam: LPARAM) -> LRESULT {
        unsafe {
            user32::SendMessageW(self.wnd, msg, wparam, lparam)
        }
    }

    pub fn post_message(&self, msg: UINT, wparam: WPARAM, lparam: LPARAM) -> bool {
        1 == unsafe {
            user32::PostMessageW(self.wnd, msg, wparam, lparam)
        }
    }

    pub fn invalidate_rect(&self, rect: RECT, erase: bool) -> bool {
        1 == unsafe {
            user32::InvalidateRect(self.wnd, &rect, erase as BOOL)
        }
    }

    pub fn invalidate(&self, erase: bool) -> bool {
        1 == unsafe {
            user32::InvalidateRect(self.wnd, ptr::null(), erase as BOOL)
        }
    }
}

pub trait WindowImpl {
    fn wnd<'a>(&'a self) -> &'a Window;
    fn wnd_mut<'a>(&'a mut self) -> &'a mut Window;
    fn wnd_proc(&self, msg: UINT, w: WPARAM, l: LPARAM) -> LRESULT;
}

/// A thread-local global map from windows to the above WindowImpl trait object
/// RefCell is necessary for mutability.
/// Rc is necessary so multiple main_wnd_procs on the same stack can reference it at once.
/// Box is necessary because WindowImpl is unsized, so can't be Rc'ed directly.
thread_local!(static KEY_WIN_MAP: RefCell<HashMap<Window, Rc<Box<WindowImpl + 'static>>>> = RefCell::new(HashMap::new()));

/// A thread-local global pointing to the initial window
thread_local!(static KEY_INIT_WND: RefCell<Option<Box<WindowImpl + 'static>>> = RefCell::new(None));

fn associate_window_impl(win: Window, wnd_impl: Box<WindowImpl + 'static>) {
    KEY_WIN_MAP.with(move |wmap_cell| {
        let mut wmap = wmap_cell.borrow_mut();
        wmap.insert(win, Rc::new(wnd_impl));
    });
}

fn promote_init_wnd(win: Window) {
    KEY_INIT_WND.with(move |maybe_initial_cell| {
        let mut maybe_initial = maybe_initial_cell.borrow_mut();
        if let Some(mut wnd_impl) = maybe_initial.take() {
            wnd_impl.wnd_mut().wnd = win.wnd;
            associate_window_impl(win, wnd_impl);
        }
    });
}

fn lookup_wnd_impl(wnd: HWND) -> Option<Rc<Box<WindowImpl + 'static>>> {
    KEY_WIN_MAP.with(|wmap_cell| {
        if let Some(wnd_impl) = wmap_cell.borrow().get(&Window{wnd: wnd}) {
            return Some(wnd_impl.clone());
        }
        return None;
    })
}

pub unsafe extern "system" fn main_wnd_proc(wnd: HWND,
                                            msg: UINT,
                                            w: WPARAM,
                                            l: LPARAM) -> LRESULT {
    promote_init_wnd(Window { wnd: wnd });

    if let Some(wnd_impl) = lookup_wnd_impl(wnd) {
        wnd_impl.wnd_proc(msg, w, l)
    } else {
        super::def_window_proc(wnd, msg, w, l)
    }
}

pub trait OnCreate {
    fn on_create(&self, _cs: &CREATESTRUCTW) -> bool {
        true
    }
}

pub trait OnDestroy {
    fn on_destroy(&self) {
        unsafe {
            user32::PostQuitMessage(0 as c_int);
        }
    }
}

pub trait OnPaint {
    fn on_paint(&self) {
    }
}

pub trait OnSize {
    fn on_size(&self, _width: isize, _height: isize) {
    }
}

pub trait OnFocus {
    fn on_focus(&self, _prev: Window) {
    }
}

pub trait OnLeftButtonDown {
    fn on_left_button_down(&self, _x: isize, _y: isize, _flags: u32) {
    }
}

pub trait OnLeftButtonUp {
    fn on_left_button_up(&self, _x: isize, _y: isize, _flags: u32) {
    }
}

pub trait OnKeyDown {
    fn on_key_down(&self, _keycode: u8, _flags: u32) -> bool {
        return false;
    }
}

pub trait OnKeyUp {
    fn on_key_up(&self, _keycode: u8, _flags: u32) -> bool {
        return false;
    }
}

pub trait OnEraseBackground {
    fn on_erase_background(&self) -> bool {
        false
    }
}

pub trait OnMessage {
    fn on_message(&self, _message: UINT, _wparam: WPARAM, _lparam: LPARAM) -> Option<LRESULT> {
        None
    }
}
<commit_msg>window text<commit_after>// Copyright 2015 The Rust-Windows Project Developers. See the
// COPYRIGHT file at the top-level directory of this distribution.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

use std::ptr;
use std;
use std::cell::RefCell;
use std::rc::Rc;
use std::collections::HashMap;

use user32;
use winapi::{
    BOOL, CREATESTRUCTW, HBRUSH, HCURSOR, HICON, HMENU, HWND, INT, LPARAM, LRESULT, RECT, UINT,
    WNDCLASSEXW, WPARAM, c_int,
};

use wchar::{FromCU16Str,ToCU16Str};
use instance::Instance;
use resource::*;

pub struct WndClass {
    pub classname: String,
    pub style: usize,
    pub icon: Option<Image>,
    pub icon_small: Option<Image>,
    pub cursor: Option<Image>,
    pub background: HBRUSH,
    pub menu: MenuResource,
    pub cls_extra: isize,
    pub wnd_extra: isize,
}

impl WndClass {
    pub fn register(&self, instance: Instance) -> bool {
        self.menu.with_menu_p(|menu_p| {
            let clsname_u = self.classname.to_c_u16();
            let wcex = WNDCLASSEXW {
                cbSize: std::mem::size_of::<WNDCLASSEXW>() as UINT,
                style: self.style as UINT,
                lpfnWndProc: Some(main_wnd_proc),
                cbClsExtra: self.cls_extra as INT,
                cbWndExtra: self.wnd_extra as INT,
                hInstance: instance.instance,
                hIcon: self.icon.to_handle() as HICON,
                hCursor: self.cursor.to_handle() as HCURSOR,
                hbrBackground: self.background,
                lpszMenuName: menu_p,
                lpszClassName: clsname_u.as_ptr(),
                hIconSm: self.icon_small.to_handle() as HICON,
            };

            let res = unsafe { user32::RegisterClassExW(&wcex) };
            res != 0
        })
    }
}

pub static WS_BORDER: u32 = 0x800000;
pub static WS_CAPTION: u32 = 0xc00000;
pub static WS_CHILD: u32 = 0x40000000;
pub static WS_CHILDWINDOW: u32 = 0x40000000;
pub static WS_CLIPCHILDREN: u32 = 0x2000000;
pub static WS_CLIPSIBLINGS: u32 = 0x4000000;
pub static WS_DISABLED: u32 = 0x8000000;
pub static WS_DLGFRAME: u32 = 0x400000;
pub static WS_GROUP: u32 = 0x20000;
pub static WS_HSCROLL: u32 = 0x100000;
pub static WS_ICONIC: u32 = 0x20000000;
pub static WS_MAXIMIZE: u32 = 0x1000000;
pub static WS_MAXIMIZEBOX: u32 = 0x10000;
pub static WS_MINIMIZE: u32 = 0x20000000;
pub static WS_MINIMIZEBOX: u32 = 0x20000;
pub static WS_OVERLAPPED: u32 = 0;
pub static WS_OVERLAPPEDWINDOW: u32 = 0xcf0000;
pub static WS_POPUP: u32 = 0x80000000;
pub static WS_POPUPWINDOW: u32 = 0x80880000;
pub static WS_SIZEBOX: u32 = 0x40000;
pub static WS_SYSMENU: u32 = 0x80000;
pub static WS_TABSTOP: u32 = 0x10000;
pub static WS_THICKFRAME: u32 = 0x40000;
pub static WS_TILED: u32 = 0;
pub static WS_TILEDWINDOW: u32 = 0xcf0000;
pub static WS_VISIBLE: u32 = 0x10000000;
pub static WS_VSCROLL: u32 = 0x200000;

pub static ES_AUTOHSCROLL: u32 = 128;
pub static ES_AUTOVSCROLL: u32 = 64;
pub static ES_CENTER: u32 = 1;
pub static ES_LEFT: u32 = 0;
pub static ES_LOWERCASE: u32 = 16;
pub static ES_MULTILINE: u32 = 4;
pub static ES_NOHIDESEL: u32 = 256;
pub static ES_NUMBER: u32 = 0x2000;
pub static ES_OEMCONVERT: u32 = 0x400;
pub static ES_PASSWORD: u32 = 32;
pub static ES_READONLY: u32 = 0x800;
pub static ES_RIGHT: u32 = 2;
pub static ES_UPPERCASE: u32 = 8;
pub static ES_WANTRETURN: u32 = 4096;

pub struct WindowParams {
    pub window_name: String,
    pub style: u32,
    pub x: isize,
    pub y: isize,
    pub width: isize,
    pub height: isize,
    pub parent: Window,
    pub menu: HMENU,
    pub ex_style: u32,
}

#[derive(PartialEq, Eq, Hash, Copy)]
pub struct Window {
    pub wnd: HWND,
}

// Sending across threads allows, for example, a worker thread to communicate
// with a UI thread via PostMessage.
unsafe impl Send for Window {}

impl Clone for Window {
    fn clone(&self) -> Window {
        Window {
            wnd: self.wnd,
        }
    }
}

impl Window {
    pub fn null() -> Window {
        Window {
            wnd: ptr::null_mut(),
        }
    }

    pub fn new(
        instance: Instance, wproc: Option<Box<WindowImpl + 'static>>, classname: &str, params: &WindowParams
    ) -> Option<Window> {
        KEY_INIT_WND.with(move |f| *f.borrow_mut() = wproc);

        let wnd = unsafe {
            let clsname_u = classname.to_c_u16();
            let title_u = params.window_name.to_c_u16();
            let wnd = user32::CreateWindowExW(
                params.ex_style, clsname_u.as_ptr(), title_u.as_ptr(), params.style,
                params.x as c_int, params.y as c_int,
                params.width as c_int, params.height as c_int,
                params.parent.wnd, params.menu, instance.instance,
                ptr::null_mut()
            );
            wnd
        };

        if wnd != ptr::null_mut() {
            Some(Window { wnd: wnd })
        } else {
            None
        }
    }

    pub fn show(&self, cmd_show: isize) -> bool {
        unsafe { user32::ShowWindow(self.wnd, cmd_show as c_int) == 0 }
    }

    pub fn show_async(&self, cmd_show: isize) -> bool {
        unsafe { user32::ShowWindowAsync(self.wnd, cmd_show as c_int) == 0 }
    }

    pub fn update(&self) -> bool {
        unsafe { user32::UpdateWindow(self.wnd) == 0 }
    }

    pub fn client_rect(&self) -> Option<RECT> {
        let mut rect = RECT {
            left: 0,
            top: 0,
            right: 0,
            bottom: 0,
        };
        let res = unsafe {
            user32::GetClientRect(self.wnd, &mut rect as *mut RECT)
        } != 0;
        match res {
            true => Some(rect),
            false => None,
        }
    }

    pub fn set_window_pos(
        &self, x: isize, y: isize, width: isize, height: isize, flags: UINT
    ) -> bool {
        // TODO: hwndInsertAfter
        unsafe {
            user32::SetWindowPos(
                self.wnd, ptr::null_mut(), x as c_int, y as c_int,
                width as c_int, height as c_int, flags
            ) != 0
        }
    }

    pub fn set_focus(&self) -> Window {
        unsafe {
            Window {
                wnd: user32::SetFocus(self.wnd)
            }
        }
    }

    pub fn send_message(&self, msg: UINT, wparam: WPARAM, lparam: LPARAM) -> LRESULT {
        unsafe {
            user32::SendMessageW(self.wnd, msg, wparam, lparam)
        }
    }

    pub fn post_message(&self, msg: UINT, wparam: WPARAM, lparam: LPARAM) -> bool {
        1 == unsafe {
            user32::PostMessageW(self.wnd, msg, wparam, lparam)
        }
    }

    pub fn invalidate_rect(&self, rect: RECT, erase: bool) -> bool {
        1 == unsafe {
            user32::InvalidateRect(self.wnd, &rect, erase as BOOL)
        }
    }

    pub fn invalidate(&self, erase: bool) -> bool {
        1 == unsafe {
            user32::InvalidateRect(self.wnd, ptr::null(), erase as BOOL)
        }
    }

    pub fn get_window_text(&self) -> String {
        unsafe {
            let len = user32::GetWindowTextLengthW(self.wnd);
            let mut buf = vec![ 0u16; (len+1) as usize ];

            let read = user32::GetWindowTextW(self.wnd, buf.as_mut_ptr(), (len+1) );
            if read == len {
                match String::from_c_u16(&buf) {
                    None => String::new(),
                    Some(s) => s
                }
            } else {
                String::new()
            }
        }
    }

    pub fn set_window_text(&self, text: &str ) -> bool {
        let text_u = text.to_c_u16();
        1 == unsafe {
            user32::SetWindowTextW(self.wnd, text_u.as_ptr())
        }
    }
}

pub trait WindowImpl {
    fn wnd<'a>(&'a self) -> &'a Window;
    fn wnd_mut<'a>(&'a mut self) -> &'a mut Window;
    fn wnd_proc(&self, msg: UINT, w: WPARAM, l: LPARAM) -> LRESULT;
}

/// A thread-local global map from windows to the above WindowImpl trait object
/// RefCell is necessary for mutability.
/// Rc is necessary so multiple main_wnd_procs on the same stack can reference it at once.
/// Box is necessary because WindowImpl is unsized, so can't be Rc'ed directly.
thread_local!(static KEY_WIN_MAP: RefCell<HashMap<Window, Rc<Box<WindowImpl + 'static>>>> = RefCell::new(HashMap::new()));

/// A thread-local global pointing to the initial window
thread_local!(static KEY_INIT_WND: RefCell<Option<Box<WindowImpl + 'static>>> = RefCell::new(None));

fn associate_window_impl(win: Window, wnd_impl: Box<WindowImpl + 'static>) {
    KEY_WIN_MAP.with(move |wmap_cell| {
        let mut wmap = wmap_cell.borrow_mut();
        wmap.insert(win, Rc::new(wnd_impl));
    });
}

fn promote_init_wnd(win: Window) {
    KEY_INIT_WND.with(move |maybe_initial_cell| {
        let mut maybe_initial = maybe_initial_cell.borrow_mut();
        if let Some(mut wnd_impl) = maybe_initial.take() {
            wnd_impl.wnd_mut().wnd = win.wnd;
            associate_window_impl(win, wnd_impl);
        }
    });
}

fn lookup_wnd_impl(wnd: HWND) -> Option<Rc<Box<WindowImpl + 'static>>> {
    KEY_WIN_MAP.with(|wmap_cell| {
        if let Some(wnd_impl) = wmap_cell.borrow().get(&Window{wnd: wnd}) {
            return Some(wnd_impl.clone());
        }
        return None;
    })
}

pub unsafe extern "system" fn main_wnd_proc(wnd: HWND,
                                            msg: UINT,
                                            w: WPARAM,
                                            l: LPARAM) -> LRESULT {
    promote_init_wnd(Window { wnd: wnd });

    if let Some(wnd_impl) = lookup_wnd_impl(wnd) {
        wnd_impl.wnd_proc(msg, w, l)
    } else {
        super::def_window_proc(wnd, msg, w, l)
    }
}

pub trait OnCreate {
    fn on_create(&self, _cs: &CREATESTRUCTW) -> bool {
        true
    }
}

pub trait OnDestroy {
    fn on_destroy(&self) {
        unsafe {
            user32::PostQuitMessage(0 as c_int);
        }
    }
}

pub trait OnPaint {
    fn on_paint(&self) {
    }
}

pub trait OnSize {
    fn on_size(&self, _width: isize, _height: isize) {
    }
}

pub trait OnFocus {
    fn on_focus(&self, _prev: Window) {
    }
}

pub trait OnLeftButtonDown {
    fn on_left_button_down(&self, _x: isize, _y: isize, _flags: u32) {
    }
}

pub trait OnLeftButtonUp {
    fn on_left_button_up(&self, _x: isize, _y: isize, _flags: u32) {
    }
}

pub trait OnKeyDown {
    fn on_key_down(&self, _keycode: u8, _flags: u32) -> bool {
        return false;
    }
}

pub trait OnKeyUp {
    fn on_key_up(&self, _keycode: u8, _flags: u32) -> bool {
        return false;
    }
}

pub trait OnEraseBackground {
    fn on_erase_background(&self) -> bool {
        false
    }
}

pub trait OnMessage {
    fn on_message(&self, _message: UINT, _wparam: WPARAM, _lparam: LPARAM) -> Option<LRESULT> {
        None
    }
}
<|endoftext|>


<commit_before><commit_msg>app: quote field names in errors<commit_after><|endoftext|>


<commit_before>// Copyright (c) The Diem Core Contributors
// SPDX-License-Identifier: Apache-2.0

use std::collections::BTreeMap;

use anyhow::anyhow;
use itertools::zip_eq;

use consensus_types::{
    common::Author, executed_block::ExecutedBlock, experimental::commit_vote::CommitVote,
};
use diem_crypto::ed25519::Ed25519Signature;
use diem_types::{
    account_address::AccountAddress,
    block_info::BlockInfo,
    ledger_info::{LedgerInfo, LedgerInfoWithSignatures},
    validator_verifier::ValidatorVerifier,
};

use crate::state_replication::StateComputerCommitCallBackType;
use diem_crypto::HashValue;

// we differentiate buffer items at different stages
// for better code readability
pub struct OrderedBufferItem {
    pub pending_votes: BTreeMap<AccountAddress, Ed25519Signature>,
    pub callback: StateComputerCommitCallBackType,
    pub ordered_blocks: Vec<ExecutedBlock>,
    pub ordered_proof: LedgerInfoWithSignatures,
}

pub struct ExecutedBufferItem {
    pub executed_blocks: Vec<ExecutedBlock>,
    pub pending_votes: BTreeMap<AccountAddress, Ed25519Signature>,
    pub callback: StateComputerCommitCallBackType,
    pub commit_info: BlockInfo,
    pub ordered_proof: LedgerInfoWithSignatures,
}

impl ExecutedBufferItem {
    pub fn generate_commit_ledger_info(&self) -> LedgerInfo {
        LedgerInfo::new(
            self.commit_info.clone(),
            self.ordered_proof.ledger_info().consensus_data_hash(),
        )
    }
}

pub struct SignedBufferItem {
    pub executed_blocks: Vec<ExecutedBlock>,
    pub commit_proof: LedgerInfoWithSignatures,
    pub callback: StateComputerCommitCallBackType,
    pub commit_vote: CommitVote,
}

pub struct AggregatedBufferItem {
    pub executed_blocks: Vec<ExecutedBlock>,
    pub aggregated_proof: LedgerInfoWithSignatures,
    pub callback: StateComputerCommitCallBackType,
}

pub enum BufferItem {
    Ordered(Box<OrderedBufferItem>),
    Executed(Box<ExecutedBufferItem>),
    Signed(Box<SignedBufferItem>),
    Aggregated(Box<AggregatedBufferItem>),
}

impl BufferItem {
    pub fn new_ordered(
        ordered_blocks: Vec<ExecutedBlock>,
        ordered_proof: LedgerInfoWithSignatures,
        callback: StateComputerCommitCallBackType,
    ) -> Self {
        Self::Ordered(Box::new(OrderedBufferItem {
            pending_votes: BTreeMap::<AccountAddress, Ed25519Signature>::new(),
            callback,
            ordered_blocks,
            ordered_proof,
        }))
    }

    // pipeline functions
    pub fn advance_to_executed(self, executed_blocks: Vec<ExecutedBlock>) -> Self {
        match self {
            Self::Ordered(ordered_item_box) => {
                let ordered_item = *ordered_item_box;
                for (b1, b2) in zip_eq(ordered_item.ordered_blocks.iter(), executed_blocks.iter()) {
                    assert_eq!(b1.id(), b2.id());
                }
                let commit_info = executed_blocks.last().unwrap().block_info();
                Self::Executed(Box::new(ExecutedBufferItem {
                    executed_blocks,
                    pending_votes: ordered_item.pending_votes,
                    callback: ordered_item.callback,
                    commit_info,
                    ordered_proof: ordered_item.ordered_proof,
                }))
            }
            _ => {
                panic!("Only ordered blocks can advance to executed blocks.")
            }
        }
    }

    pub fn advance_to_signed(
        self,
        author: Author,
        signature: Ed25519Signature,
        verifier: &ValidatorVerifier,
    ) -> (Self, CommitVote) {
        match self {
            Self::Executed(executed_item_box) => {
                let executed_item = *executed_item_box;

                let mut valid_sigs = BTreeMap::<AccountAddress, Ed25519Signature>::new();
                valid_sigs.insert(author, signature.clone());

                let commit_ledger_info = executed_item.generate_commit_ledger_info();

                for (author, sig) in executed_item.pending_votes.iter() {
                    if verifier.verify(*author, &commit_ledger_info, sig).is_ok() {
                        valid_sigs.insert(*author, sig.clone());
                    }
                }

                let commit_ledger_info_with_sigs =
                    LedgerInfoWithSignatures::new(commit_ledger_info.clone(), valid_sigs);

                let commit_vote =
                    CommitVote::new_with_signature(author, commit_ledger_info, signature);

                (
                    Self::Signed(Box::new(SignedBufferItem {
                        executed_blocks: executed_item.executed_blocks,
                        callback: executed_item.callback,
                        commit_proof: commit_ledger_info_with_sigs,
                        commit_vote: commit_vote.clone(),
                    })),
                    commit_vote,
                )
            }
            _ => {
                panic!("Only executed buffer items can advance to signed blocks.")
            }
        }
    }

    /// this function assumes block id matches and the validity of ledger_info and that it has the voting power
    /// it returns an updated item
    pub fn try_advance_to_aggregated_with_ledger_info(
        self,
        commit_ledger_info: LedgerInfoWithSignatures,
    ) -> Self {
        match self {
            Self::Signed(signed_item) => {
                let SignedBufferItem {
                    executed_blocks,
                    callback,
                    commit_proof,
                    ..
                } = *signed_item;
                assert_eq!(commit_proof.commit_info(), commit_ledger_info.commit_info(),);
                Self::Aggregated(Box::new(AggregatedBufferItem {
                    executed_blocks,
                    callback,
                    aggregated_proof: commit_ledger_info,
                }))
            }
            Self::Executed(executed_item) => {
                let ExecutedBufferItem {
                    executed_blocks,
                    callback,
                    commit_info,
                    ..
                } = *executed_item;
                assert_eq!(commit_info, *commit_ledger_info.commit_info());
                Self::Aggregated(Box::new(AggregatedBufferItem {
                    executed_blocks,
                    callback,
                    aggregated_proof: commit_ledger_info,
                }))
            }
            Self::Ordered(ordered_item) => {
                let ordered = *ordered_item;
                assert!(ordered
                    .ordered_proof
                    .commit_info()
                    .match_ordered_only(commit_ledger_info.commit_info()));
                Self::Ordered(Box::new(OrderedBufferItem {
                    pending_votes: commit_ledger_info.signatures().clone(),
                    ..ordered
                }))
            }
            Self::Aggregated(_) => {
                unreachable!("Found aggregated buffer item but any aggregated buffer item should get dequeued right away.");
            }
        }
    }

    pub fn try_advance_to_aggregated(self, validator: &ValidatorVerifier) -> Self {
        match self {
            Self::Signed(signed_item_box) => {
                let signed_item = *signed_item_box;
                if signed_item
                    .commit_proof
                    .check_voting_power(validator)
                    .is_ok()
                {
                    Self::Aggregated(Box::new(AggregatedBufferItem {
                        executed_blocks: signed_item.executed_blocks,
                        aggregated_proof: signed_item.commit_proof,
                        callback: signed_item.callback,
                    }))
                } else {
                    Self::Signed(Box::new(signed_item))
                }
            }
            Self::Executed(executed_item_box) => {
                let executed_item = *executed_item_box;
                // TODO this needs to verify the votes
                if validator
                    .check_voting_power(executed_item.pending_votes.keys())
                    .is_ok()
                {
                    let aggregated_proof = LedgerInfoWithSignatures::new(
                        executed_item.generate_commit_ledger_info(),
                        executed_item.pending_votes,
                    );
                    Self::Aggregated(Box::new(AggregatedBufferItem {
                        executed_blocks: executed_item.executed_blocks,
                        aggregated_proof,
                        callback: executed_item.callback,
                    }))
                } else {
                    Self::Executed(Box::new(executed_item))
                }
            }
            _ => self,
        }
    }

    // generic functions
    pub fn get_blocks(&self) -> &Vec<ExecutedBlock> {
        match self {
            Self::Ordered(ordered) => &ordered.ordered_blocks,
            Self::Executed(executed) => &executed.executed_blocks,
            Self::Signed(signed) => &signed.executed_blocks,
            Self::Aggregated(aggregated) => &aggregated.executed_blocks,
        }
    }

    pub fn block_id(&self) -> HashValue {
        self.get_blocks().last().unwrap().id()
    }

    pub fn get_commit_info(&self) -> &BlockInfo {
        match self {
            Self::Ordered(_) => {
                panic!("Ordered buffer item does not contain commit info");
            }
            Self::Executed(executed) => &executed.commit_info,
            Self::Signed(signed) => signed.commit_proof.ledger_info().commit_info(),
            Self::Aggregated(aggregated) => aggregated.aggregated_proof.ledger_info().commit_info(),
        }
    }

    pub fn add_signature_if_matched(
        &mut self,
        target_commit_info: &BlockInfo,
        author: Author,
        sig: Ed25519Signature,
    ) -> anyhow::Result<()> {
        match self {
            Self::Ordered(ordered) => {
                if ordered
                    .ordered_proof
                    .commit_info()
                    .match_ordered_only(target_commit_info)
                {
                    // we optimistically assume the vote will be valid in the future.
                    // when advancing to signed item, we will check if the sigs are valid.
                    // each author at most stores a single sig for each item,
                    // so an adversary will not be able to flood our memory.
                    ordered.pending_votes.insert(author, sig);
                    return Ok(());
                }
            }
            Self::Executed(executed) => {
                if &executed.commit_info == target_commit_info {
                    executed.pending_votes.insert(author, sig);
                    return Ok(());
                }
            }
            Self::Signed(signed) => {
                if signed.commit_proof.commit_info() == target_commit_info {
                    signed.commit_proof.add_signature(author, sig);
                    return Ok(());
                }
            }
            Self::Aggregated(aggregated) => {
                // we do not need to do anything for aggregated
                // but return true is helpful to stop the outer loop early
                if aggregated.aggregated_proof.commit_info() == target_commit_info {
                    return Ok(());
                }
            }
        }
        Err(anyhow!("Inconsistent commit info."))
    }

    pub fn has_been_executed(&self) -> bool {
        !self.is_ordered()
    }

    pub fn is_ordered(&self) -> bool {
        matches!(self, Self::Ordered(_))
    }

    pub fn is_executed(&self) -> bool {
        matches!(self, Self::Executed(_))
    }

    pub fn is_signed(&self) -> bool {
        matches!(self, Self::Signed(_))
    }

    pub fn is_aggregated(&self) -> bool {
        matches!(self, Self::Aggregated(_))
    }
}
<commit_msg>[decoupled-execution] cleanup buffer item<commit_after>// Copyright (c) The Diem Core Contributors
// SPDX-License-Identifier: Apache-2.0

use std::collections::BTreeMap;

use anyhow::anyhow;
use itertools::zip_eq;

use consensus_types::{
    common::Author, executed_block::ExecutedBlock, experimental::commit_vote::CommitVote,
};
use diem_crypto::ed25519::Ed25519Signature;
use diem_types::{
    account_address::AccountAddress,
    block_info::BlockInfo,
    ledger_info::{LedgerInfo, LedgerInfoWithSignatures},
    validator_verifier::ValidatorVerifier,
};

use crate::state_replication::StateComputerCommitCallBackType;
use diem_crypto::HashValue;

// we differentiate buffer items at different stages
// for better code readability
pub struct OrderedBufferItem {
    pub pending_votes: BTreeMap<AccountAddress, Ed25519Signature>,
    pub callback: StateComputerCommitCallBackType,
    pub ordered_blocks: Vec<ExecutedBlock>,
    pub ordered_proof: LedgerInfoWithSignatures,
}

pub struct ExecutedBufferItem {
    pub executed_blocks: Vec<ExecutedBlock>,
    pub pending_votes: BTreeMap<AccountAddress, Ed25519Signature>,
    pub callback: StateComputerCommitCallBackType,
    pub commit_info: BlockInfo,
    pub ordered_proof: LedgerInfoWithSignatures,
}

impl ExecutedBufferItem {
    pub fn generate_commit_ledger_info(&self) -> LedgerInfo {
        LedgerInfo::new(
            self.commit_info.clone(),
            self.ordered_proof.ledger_info().consensus_data_hash(),
        )
    }
}

pub struct SignedBufferItem {
    pub executed_blocks: Vec<ExecutedBlock>,
    pub commit_proof: LedgerInfoWithSignatures,
    pub callback: StateComputerCommitCallBackType,
    pub commit_vote: CommitVote,
}

pub struct AggregatedBufferItem {
    pub executed_blocks: Vec<ExecutedBlock>,
    pub aggregated_proof: LedgerInfoWithSignatures,
    pub callback: StateComputerCommitCallBackType,
}

pub enum BufferItem {
    Ordered(Box<OrderedBufferItem>),
    Executed(Box<ExecutedBufferItem>),
    Signed(Box<SignedBufferItem>),
    Aggregated(Box<AggregatedBufferItem>),
}

impl BufferItem {
    pub fn new_ordered(
        ordered_blocks: Vec<ExecutedBlock>,
        ordered_proof: LedgerInfoWithSignatures,
        callback: StateComputerCommitCallBackType,
    ) -> Self {
        Self::Ordered(Box::new(OrderedBufferItem {
            pending_votes: BTreeMap::<AccountAddress, Ed25519Signature>::new(),
            callback,
            ordered_blocks,
            ordered_proof,
        }))
    }

    // pipeline functions
    pub fn advance_to_executed(self, executed_blocks: Vec<ExecutedBlock>) -> Self {
        match self {
            Self::Ordered(ordered_item) => {
                let OrderedBufferItem {
                    ordered_blocks,
                    pending_votes,
                    callback,
                    ordered_proof,
                } = *ordered_item;
                for (b1, b2) in zip_eq(ordered_blocks.iter(), executed_blocks.iter()) {
                    assert_eq!(b1.id(), b2.id());
                }
                let commit_info = executed_blocks.last().unwrap().block_info();
                Self::Executed(Box::new(ExecutedBufferItem {
                    executed_blocks,
                    pending_votes,
                    callback,
                    commit_info,
                    ordered_proof,
                }))
            }
            _ => {
                panic!("Only ordered blocks can advance to executed blocks.")
            }
        }
    }

    pub fn advance_to_signed(
        self,
        author: Author,
        signature: Ed25519Signature,
        verifier: &ValidatorVerifier,
    ) -> (Self, CommitVote) {
        match self {
            Self::Executed(executed_item) => {
                let commit_ledger_info = executed_item.generate_commit_ledger_info();
                let ExecutedBufferItem {
                    executed_blocks,
                    callback,
                    pending_votes,
                    ..
                } = *executed_item;

                let mut valid_sigs = BTreeMap::<AccountAddress, Ed25519Signature>::new();
                valid_sigs.insert(author, signature.clone());

                for (author, sig) in pending_votes.iter() {
                    if verifier.verify(*author, &commit_ledger_info, sig).is_ok() {
                        valid_sigs.insert(*author, sig.clone());
                    }
                }

                let commit_ledger_info_with_sigs =
                    LedgerInfoWithSignatures::new(commit_ledger_info.clone(), valid_sigs);

                let commit_vote =
                    CommitVote::new_with_signature(author, commit_ledger_info, signature);

                (
                    Self::Signed(Box::new(SignedBufferItem {
                        executed_blocks,
                        callback,
                        commit_proof: commit_ledger_info_with_sigs,
                        commit_vote: commit_vote.clone(),
                    })),
                    commit_vote,
                )
            }
            _ => {
                panic!("Only executed buffer items can advance to signed blocks.")
            }
        }
    }

    /// this function assumes block id matches and the validity of ledger_info and that it has the voting power
    /// it returns an updated item
    pub fn try_advance_to_aggregated_with_ledger_info(
        self,
        commit_ledger_info: LedgerInfoWithSignatures,
    ) -> Self {
        match self {
            Self::Signed(signed_item) => {
                let SignedBufferItem {
                    executed_blocks,
                    callback,
                    commit_proof,
                    ..
                } = *signed_item;
                assert_eq!(commit_proof.commit_info(), commit_ledger_info.commit_info(),);
                Self::Aggregated(Box::new(AggregatedBufferItem {
                    executed_blocks,
                    callback,
                    aggregated_proof: commit_ledger_info,
                }))
            }
            Self::Executed(executed_item) => {
                let ExecutedBufferItem {
                    executed_blocks,
                    callback,
                    commit_info,
                    ..
                } = *executed_item;
                assert_eq!(commit_info, *commit_ledger_info.commit_info());
                Self::Aggregated(Box::new(AggregatedBufferItem {
                    executed_blocks,
                    callback,
                    aggregated_proof: commit_ledger_info,
                }))
            }
            Self::Ordered(ordered_item) => {
                let ordered = *ordered_item;
                assert!(ordered
                    .ordered_proof
                    .commit_info()
                    .match_ordered_only(commit_ledger_info.commit_info()));
                Self::Ordered(Box::new(OrderedBufferItem {
                    pending_votes: commit_ledger_info.signatures().clone(),
                    ..ordered
                }))
            }
            Self::Aggregated(_) => {
                unreachable!("Found aggregated buffer item but any aggregated buffer item should get dequeued right away.");
            }
        }
    }

    pub fn try_advance_to_aggregated(self, validator: &ValidatorVerifier) -> Self {
        match self {
            Self::Signed(signed_item) => {
                if signed_item
                    .commit_proof
                    .check_voting_power(validator)
                    .is_ok()
                {
                    Self::Aggregated(Box::new(AggregatedBufferItem {
                        executed_blocks: signed_item.executed_blocks,
                        aggregated_proof: signed_item.commit_proof,
                        callback: signed_item.callback,
                    }))
                } else {
                    Self::Signed(signed_item)
                }
            }
            Self::Executed(executed_item) => {
                // TODO this needs to verify the votes match the commit info
                if validator
                    .check_voting_power(executed_item.pending_votes.keys())
                    .is_ok()
                {
                    let aggregated_proof = LedgerInfoWithSignatures::new(
                        executed_item.generate_commit_ledger_info(),
                        executed_item.pending_votes,
                    );
                    Self::Aggregated(Box::new(AggregatedBufferItem {
                        executed_blocks: executed_item.executed_blocks,
                        aggregated_proof,
                        callback: executed_item.callback,
                    }))
                } else {
                    Self::Executed(executed_item)
                }
            }
            _ => self,
        }
    }

    // generic functions
    pub fn get_blocks(&self) -> &Vec<ExecutedBlock> {
        match self {
            Self::Ordered(ordered) => &ordered.ordered_blocks,
            Self::Executed(executed) => &executed.executed_blocks,
            Self::Signed(signed) => &signed.executed_blocks,
            Self::Aggregated(aggregated) => &aggregated.executed_blocks,
        }
    }

    pub fn block_id(&self) -> HashValue {
        self.get_blocks().last().unwrap().id()
    }

    pub fn get_commit_info(&self) -> &BlockInfo {
        match self {
            Self::Ordered(_) => {
                panic!("Ordered buffer item does not contain commit info");
            }
            Self::Executed(executed) => &executed.commit_info,
            Self::Signed(signed) => signed.commit_proof.ledger_info().commit_info(),
            Self::Aggregated(aggregated) => aggregated.aggregated_proof.ledger_info().commit_info(),
        }
    }

    pub fn add_signature_if_matched(
        &mut self,
        target_commit_info: &BlockInfo,
        author: Author,
        sig: Ed25519Signature,
    ) -> anyhow::Result<()> {
        match self {
            Self::Ordered(ordered) => {
                if ordered
                    .ordered_proof
                    .commit_info()
                    .match_ordered_only(target_commit_info)
                {
                    // we optimistically assume the vote will be valid in the future.
                    // when advancing to signed item, we will check if the sigs are valid.
                    // each author at most stores a single sig for each item,
                    // so an adversary will not be able to flood our memory.
                    ordered.pending_votes.insert(author, sig);
                    return Ok(());
                }
            }
            Self::Executed(executed) => {
                if &executed.commit_info == target_commit_info {
                    executed.pending_votes.insert(author, sig);
                    return Ok(());
                }
            }
            Self::Signed(signed) => {
                if signed.commit_proof.commit_info() == target_commit_info {
                    signed.commit_proof.add_signature(author, sig);
                    return Ok(());
                }
            }
            Self::Aggregated(aggregated) => {
                // we do not need to do anything for aggregated
                // but return true is helpful to stop the outer loop early
                if aggregated.aggregated_proof.commit_info() == target_commit_info {
                    return Ok(());
                }
            }
        }
        Err(anyhow!("Inconsistent commit info."))
    }

    pub fn has_been_executed(&self) -> bool {
        !self.is_ordered()
    }

    pub fn is_ordered(&self) -> bool {
        matches!(self, Self::Ordered(_))
    }

    pub fn is_executed(&self) -> bool {
        matches!(self, Self::Executed(_))
    }

    pub fn is_signed(&self) -> bool {
        matches!(self, Self::Signed(_))
    }

    pub fn is_aggregated(&self) -> bool {
        matches!(self, Self::Aggregated(_))
    }
}
<|endoftext|>