<commit_before>use time;
use std::thread;
use glutin;
use glium::{self, texture, framebuffer, Surface, DisplayBuild};
use image;
use ::{V2, Rect, AverageDuration};
use ::rgb::{to_srgb};
use super::event::{Event, MouseButton};
use super::{CanvasMagnify};
use super::event_translator::{EventTranslator};

pub struct WindowBuilder {
    title: String,
    size: V2<u32>,
    frame_interval: Option<f64>,
    fullscreen: bool,
    layout_independent_keys: bool,
    magnify: CanvasMagnify,
}

impl WindowBuilder {
    pub fn new() -> WindowBuilder {
        WindowBuilder {
            title: "".to_string(),
            size: V2(640, 360),
            frame_interval: None,
            fullscreen: false,
            layout_independent_keys: true,
            magnify: CanvasMagnify::PixelPerfect,
        }
    }

    /// Set the window title.
    pub fn set_title(mut self, title: &str) -> WindowBuilder {
        self.title = title.to_string();
        self
    }

    /// Set the frame rate.
    pub fn set_frame_interval(mut self, interval_s: f64) -> WindowBuilder {
        assert!(interval_s > 0.00001);
        self.frame_interval = Some(interval_s);
        self
    }

    /// Set the size of the logical canvas.
    pub fn set_size(mut self, width: u32, height: u32) -> WindowBuilder {
        self.size = V2(width, height);
        self
    }

    /// Get the key values from the user's keyboard layout instead of the
    /// hardware keyboard map. Hardware keymap lookup may not work correctly
    /// on all platforms.
    pub fn use_layout_dependent_keys(mut self) -> WindowBuilder {
        self.layout_independent_keys = false;
        self
    }

    /// Set the canvas to start in fullscreen mode.
    /// FIXME: Broken on Linux, https://github.com/tomaka/glutin/issues/148
    pub fn set_fullscreen(mut self, state: bool) -> WindowBuilder {
        self.fullscreen = state;
        self
    }

    pub fn set_magnify(mut self, magnify: CanvasMagnify) -> WindowBuilder {
        self.magnify = magnify;
        self
    }

    /// Build the window object.
    pub fn build(self) -> Window {
        Window::new(self)
    }
}

/// Toplevel application object.
pub struct Window {
    pub display: glium::Display,

    resolution: LogicalResolution,

    /// Shader for blitting the canvas texture to screen.
    blit_shader: glium::Program,

    /// Render target texture.
    buffer: texture::Texture2d,
    depth: framebuffer::DepthRenderBuffer,

    translator: EventTranslator,
    frame_duration: AverageDuration,

    frame_interval: Option<f64>,
    previous_frame_t: f64,
}

impl Window {
    fn new(builder: WindowBuilder) -> Window {
        use glutin::{GlRequest, Api};
        let size = builder.size;
        let title = &builder.title[..];

        let mut glutin = glutin::WindowBuilder::new()
            .with_gl(GlRequest::Specific(Api::OpenGl, (3, 2)))
            .with_title(title.to_string());

        if builder.fullscreen {
            // FIXME: Glutin's X11 fullscreen is broken, this is only enabled
            // for Windows.
            if cfg!(windows) {
                glutin = glutin.with_fullscreen(glutin::get_primary_monitor());
            }
        } else {
            // Zoom up the window to the biggest even pixel multiple that fits
            // the user's monitor.
            let window_border_guesstimate = 32;
            let (w, h) = glutin::get_primary_monitor().get_dimensions();
            let window_size = V2(w, h) - V2(window_border_guesstimate, window_border_guesstimate);

            let (mut x, mut y) = (size.0, size.1);
            while x * 2 <= window_size.0 && y * 2 <= window_size.1 {
                x *= 2;
                y *= 2;
            }

            glutin = glutin.with_dimensions(x, y);
        }

        let display = glutin.build_glium().unwrap();

        let (w, h) = display.get_framebuffer_dimensions();

        let resolution = LogicalResolution::new(builder.magnify, size, V2(w, h));

        let blit_shader = glium::Program::from_source(&display,
            include_str!("blit.vert"),
            include_str!("blit.frag"),
            None).unwrap();

        let buffer = texture::Texture2d::empty(
            &display, size.0, size.1).unwrap();
        let depth = framebuffer::DepthRenderBuffer::new(
            &display, texture::DepthFormat::F32, size.0, size.1);

        Window {
            display: display,
            resolution: resolution,
            blit_shader: blit_shader,
            buffer: buffer,
            depth: depth,
            translator: EventTranslator::new(builder.layout_independent_keys),
            frame_duration: AverageDuration::new(0.1, 0.9),
            frame_interval: builder.frame_interval,
            previous_frame_t: time::precise_time_s(),
        }
    }

    /// Display the screen buffer and do end-of-frame bookkeeping.
    pub fn end_frame(&mut self) {
        self.frame_duration.tick();
        self.show();

        // Stick to a target frame rate if one is set.
        if let Some(target_t) = self.frame_interval {
            let sleepytime = target_t - (time::precise_time_s() - self.previous_frame_t);
            if sleepytime > 0.0 { thread::sleep_ms((sleepytime * 1e3) as u32); }
        }
        self.previous_frame_t = time::precise_time_s();

        // Pull in new input events. This will sleep if there was a suspend
        // event.
        self.translator.pump(&mut self.display, &self.resolution);
    }

    pub fn draw<F>(&self, draw_f: F)
        where F: FnOnce(&mut framebuffer::SimpleFrameBuffer) {
        let mut target = framebuffer::SimpleFrameBuffer::with_depth_buffer(
            &self.display, &self.buffer, &self.depth);
        draw_f(&mut target);
    }

    /// Fill the frame with the given pixel buffer. The pixel data dimensions
    /// must match the logical size of the window.
    pub fn set_frame<'a, T: glium::texture::Texture2dDataSource<'a>>(&mut self, pixels: T) {
        let new_texture = texture::Texture2d::new(&self.display, pixels).unwrap();
        assert!(new_texture.get_width() == self.resolution.canvas.0 &&
                new_texture.get_height() == Some(self.resolution.canvas.1),
            "Pixel data dimensions do not match logical window size");
        self.buffer = new_texture;
    }

    /// Return the logical window resolution.
    #[inline(always)]
    pub fn size(&self) -> V2<u32> { self.resolution.canvas }

    pub fn mouse_pos(&self) -> V2<f32> { self.translator.mouse_pos }

    pub fn mouse_pressed(&self, button: MouseButton) -> bool {
        self.translator.mouse_pressed[button as usize].is_some()
    }

    /// Return the current exponential moving average for frame rendering
    /// duration in seconds.
    pub fn frame_duration(&self) -> f64 { self.frame_duration.value }

    //pub fn events<'a>(&'a mut self) -> EventIterator<'a> { EventIterator { window: self } }
    pub fn events(&mut self) -> Vec<Event> {
        // XXX: Not returning an iterator because of locking crap. May want to
        // do mutable stuff on window in response to events.
        EventIterator { window: self }.collect::<Vec<Event>>()
    }

    /// Map screen position (eg. of a mouse cursor) to canvas position.
    pub fn screen_to_canvas(&self, screen_pos: V2<i32>) -> V2<i32> {
        self.resolution.screen_to_canvas(screen_pos)
    }

    pub fn get_screenshot(&self) -> image::ImageBuffer<image::Rgb<u8>, Vec<u8>> {
        let mut ret = self.buffer.read::<image::DynamicImage>().to_rgb();

        // Convert to sRGB
        // XXX: Probably horribly slow, can we make OpenGL do this?
        for p in ret.pixels_mut() {
            p.data[0] = (to_srgb(p.data[0] as f32 / 255.0) * 255.0).round() as u8;
            p.data[1] = (to_srgb(p.data[1] as f32 / 255.0) * 255.0).round() as u8;
            p.data[2] = (to_srgb(p.data[2] as f32 / 255.0) * 255.0).round() as u8;
        }

        ret
    }

    /// Show the graphics buffer on screen.
    fn show(&mut self) {
        let mut target = self.display.draw();

        // Clear the window.
        target.clear_color(0.0, 0.0, 0.0, 0.0);
        target.clear_depth(1.0);
        let (w, h) = self.display.get_framebuffer_dimensions();
        // Clip viewport dimensions to even to prevent rounding errors in
        // pixel perfect scaling.
        self.resolution.window = V2(w & !1, h & !1);

        // Build the geometry for the on-screen rectangle.
        let Rect(V2(sx, sy), V2(sw, sh)) = self.resolution.screen_rect();

        // XXX: This could use glium::Surface::blit_whole_color_to instead of
        // the handmade blitting, but that was buggy on Windows around
        // 2015-03.

        let vertices = {
            #[derive(Copy, Clone)]
            struct BlitVertex { pos: [f32; 2], tex_coord: [f32; 2] }
            implement_vertex!(BlitVertex, pos, tex_coord);

            glium::VertexBuffer::new(&self.display,
            &[
                BlitVertex { pos: [sx,    sy   ], tex_coord: [0.0, 0.0] },
                BlitVertex { pos: [sx+sw, sy   ], tex_coord: [1.0, 0.0] },
                BlitVertex { pos: [sx+sw, sy+sh], tex_coord: [1.0, 1.0] },
                BlitVertex { pos: [sx,    sy+sh], tex_coord: [0.0, 1.0] },
            ]).unwrap()
        };

        let indices = glium::IndexBuffer::new(&self.display,
            glium::index::PrimitiveType::TrianglesList, &[0u16, 1, 2, 0, 2, 3]).unwrap();

        // Set up the rest of the draw parameters.
        let mut params: glium::DrawParameters = Default::default();
        // Set an explicit viewport to apply the custom resolution that fixes
        // pixel perfect rounding errors.
        params.viewport = Some(glium::Rect{
            left: 0, bottom: 0,
            width: self.resolution.window.0,
            height: self.resolution.window.1 });

        let mag_filter = match self.resolution.magnify {
            CanvasMagnify::Smooth => glium::uniforms::MagnifySamplerFilter::Linear,
            _ => glium::uniforms::MagnifySamplerFilter::Nearest
        };

        let uniforms = glium::uniforms::UniformsStorage::new("tex",
            glium::uniforms::Sampler(&self.buffer, glium::uniforms::SamplerBehavior {
                magnify_filter: mag_filter,
                minify_filter: glium::uniforms::MinifySamplerFilter::Linear,
                .. Default::default() }));

        // Draw the graphics buffer to the window.
        target.draw(&vertices, &indices, &self.blit_shader, &uniforms, &params).unwrap();
        target.finish().unwrap();
    }
}

pub struct LogicalResolution {
    pub magnify: CanvasMagnify,
    /// Logical size, pixel on virtual canvas.
    pub canvas: V2<u32>,
    /// Physical size, window on screen.
    pub window: V2<u32>,
}

impl LogicalResolution {
    pub fn new(magnify: CanvasMagnify, canvas_size: V2<u32>, resolution: V2<u32>) -> LogicalResolution {
        LogicalResolution {
            magnify: magnify,
            canvas: canvas_size,
            window: resolution,
        }
    }

    pub fn screen_rect(&self) -> Rect<f32> {
        match self.magnify {
            CanvasMagnify::PixelPerfect => self.pixel_perfect(),
            _ => self.preserve_aspect(),
        }
    }

    /// Map screen position (eg. of a mouse cursor) to canvas position.
    pub fn screen_to_canvas(&self, V2(sx, sy): V2<i32>) -> V2<i32> {
        let Rect(V2(rx, ry), V2(rw, rh)) = self.screen_rect();

        // Transform to device coordinates.
        let sx = sx as f32 * 2.0 / self.window.0 as f32 - 1.0;
        let sy = sy as f32 * 2.0 / self.window.1 as f32 - 1.0;

        V2(((sx - rx) * self.canvas.0 as f32 / rw) as i32,
           ((sy - ry) * self.canvas.1 as f32 / rh) as i32)
    }

    /// A pixel perfect centered and scaled rectangle of resolution dim in a
    /// window of size area, mapped to OpenGL device coordinates.
#[inline(always)]
    fn pixel_perfect(&self) -> Rect<f32> {
        // Scale based on whichever of X or Y axis is the tighter fit.
        let mut scale = (self.window.0 as f32 / self.canvas.0 as f32)
            .min(self.window.1 as f32 / self.canvas.1 as f32);

        if scale > 1.0 {
            // Snap to pixel scale if more than 1 window pixel per canvas pixel.
            scale = scale.floor();
        }

        let dim = V2((scale * self.canvas.0 as f32) * 2.0 / self.window.0 as f32,
        (scale * self.canvas.1 as f32) * 2.0 / self.window.1 as f32);
        let offset = -dim / 2.0;
        Rect(offset, dim)
    }

#[inline(always)]
    fn preserve_aspect(&self) -> Rect<f32> {
        // Scale based on whichever of X or Y axis is the tighter fit.
        let scale = (self.window.0 as f32 / self.canvas.0 as f32)
            .min(self.window.1 as f32 / self.canvas.1 as f32);

        let dim = V2((scale * self.canvas.0 as f32) * 2.0 / self.window.0 as f32,
        (scale * self.canvas.1 as f32) * 2.0 / self.window.1 as f32);
        let offset = -dim / 2.0;
        Rect(offset, dim)
    }
}

pub struct EventIterator<'a> {
    window: &'a mut Window
}

impl<'a> Iterator for EventIterator<'a> {
    type Item = Event;
    fn next(&mut self) -> Option<Event> { self.window.translator.next() }
}
<commit_msg>Put the kludge use statement next to the kludge<commit_after>use time;
use std::thread;
use glutin;
use glium::{self, texture, framebuffer, Surface, DisplayBuild};
use image;
use ::{V2, Rect, AverageDuration};
use super::event::{Event, MouseButton};
use super::{CanvasMagnify};
use super::event_translator::{EventTranslator};

pub struct WindowBuilder {
    title: String,
    size: V2<u32>,
    frame_interval: Option<f64>,
    fullscreen: bool,
    layout_independent_keys: bool,
    magnify: CanvasMagnify,
}

impl WindowBuilder {
    pub fn new() -> WindowBuilder {
        WindowBuilder {
            title: "".to_string(),
            size: V2(640, 360),
            frame_interval: None,
            fullscreen: false,
            layout_independent_keys: true,
            magnify: CanvasMagnify::PixelPerfect,
        }
    }

    /// Set the window title.
    pub fn set_title(mut self, title: &str) -> WindowBuilder {
        self.title = title.to_string();
        self
    }

    /// Set the frame rate.
    pub fn set_frame_interval(mut self, interval_s: f64) -> WindowBuilder {
        assert!(interval_s > 0.00001);
        self.frame_interval = Some(interval_s);
        self
    }

    /// Set the size of the logical canvas.
    pub fn set_size(mut self, width: u32, height: u32) -> WindowBuilder {
        self.size = V2(width, height);
        self
    }

    /// Get the key values from the user's keyboard layout instead of the
    /// hardware keyboard map. Hardware keymap lookup may not work correctly
    /// on all platforms.
    pub fn use_layout_dependent_keys(mut self) -> WindowBuilder {
        self.layout_independent_keys = false;
        self
    }

    /// Set the canvas to start in fullscreen mode.
    /// FIXME: Broken on Linux, https://github.com/tomaka/glutin/issues/148
    pub fn set_fullscreen(mut self, state: bool) -> WindowBuilder {
        self.fullscreen = state;
        self
    }

    pub fn set_magnify(mut self, magnify: CanvasMagnify) -> WindowBuilder {
        self.magnify = magnify;
        self
    }

    /// Build the window object.
    pub fn build(self) -> Window {
        Window::new(self)
    }
}

/// Toplevel application object.
pub struct Window {
    pub display: glium::Display,

    resolution: LogicalResolution,

    /// Shader for blitting the canvas texture to screen.
    blit_shader: glium::Program,

    /// Render target texture.
    buffer: texture::Texture2d,
    depth: framebuffer::DepthRenderBuffer,

    translator: EventTranslator,
    frame_duration: AverageDuration,

    frame_interval: Option<f64>,
    previous_frame_t: f64,
}

impl Window {
    fn new(builder: WindowBuilder) -> Window {
        use glutin::{GlRequest, Api};
        let size = builder.size;
        let title = &builder.title[..];

        let mut glutin = glutin::WindowBuilder::new()
            .with_gl(GlRequest::Specific(Api::OpenGl, (3, 2)))
            .with_title(title.to_string());

        if builder.fullscreen {
            // FIXME: Glutin's X11 fullscreen is broken, this is only enabled
            // for Windows.
            if cfg!(windows) {
                glutin = glutin.with_fullscreen(glutin::get_primary_monitor());
            }
        } else {
            // Zoom up the window to the biggest even pixel multiple that fits
            // the user's monitor.
            let window_border_guesstimate = 32;
            let (w, h) = glutin::get_primary_monitor().get_dimensions();
            let window_size = V2(w, h) - V2(window_border_guesstimate, window_border_guesstimate);

            let (mut x, mut y) = (size.0, size.1);
            while x * 2 <= window_size.0 && y * 2 <= window_size.1 {
                x *= 2;
                y *= 2;
            }

            glutin = glutin.with_dimensions(x, y);
        }

        let display = glutin.build_glium().unwrap();

        let (w, h) = display.get_framebuffer_dimensions();

        let resolution = LogicalResolution::new(builder.magnify, size, V2(w, h));

        let blit_shader = glium::Program::from_source(&display,
            include_str!("blit.vert"),
            include_str!("blit.frag"),
            None).unwrap();

        let buffer = texture::Texture2d::empty(
            &display, size.0, size.1).unwrap();
        let depth = framebuffer::DepthRenderBuffer::new(
            &display, texture::DepthFormat::F32, size.0, size.1);

        Window {
            display: display,
            resolution: resolution,
            blit_shader: blit_shader,
            buffer: buffer,
            depth: depth,
            translator: EventTranslator::new(builder.layout_independent_keys),
            frame_duration: AverageDuration::new(0.1, 0.9),
            frame_interval: builder.frame_interval,
            previous_frame_t: time::precise_time_s(),
        }
    }

    /// Display the screen buffer and do end-of-frame bookkeeping.
    pub fn end_frame(&mut self) {
        self.frame_duration.tick();
        self.show();

        // Stick to a target frame rate if one is set.
        if let Some(target_t) = self.frame_interval {
            let sleepytime = target_t - (time::precise_time_s() - self.previous_frame_t);
            if sleepytime > 0.0 { thread::sleep_ms((sleepytime * 1e3) as u32); }
        }
        self.previous_frame_t = time::precise_time_s();

        // Pull in new input events. This will sleep if there was a suspend
        // event.
        self.translator.pump(&mut self.display, &self.resolution);
    }

    pub fn draw<F>(&self, draw_f: F)
        where F: FnOnce(&mut framebuffer::SimpleFrameBuffer) {
        let mut target = framebuffer::SimpleFrameBuffer::with_depth_buffer(
            &self.display, &self.buffer, &self.depth);
        draw_f(&mut target);
    }

    /// Fill the frame with the given pixel buffer. The pixel data dimensions
    /// must match the logical size of the window.
    pub fn set_frame<'a, T: glium::texture::Texture2dDataSource<'a>>(&mut self, pixels: T) {
        let new_texture = texture::Texture2d::new(&self.display, pixels).unwrap();
        assert!(new_texture.get_width() == self.resolution.canvas.0 &&
                new_texture.get_height() == Some(self.resolution.canvas.1),
            "Pixel data dimensions do not match logical window size");
        self.buffer = new_texture;
    }

    /// Return the logical window resolution.
    #[inline(always)]
    pub fn size(&self) -> V2<u32> { self.resolution.canvas }

    pub fn mouse_pos(&self) -> V2<f32> { self.translator.mouse_pos }

    pub fn mouse_pressed(&self, button: MouseButton) -> bool {
        self.translator.mouse_pressed[button as usize].is_some()
    }

    /// Return the current exponential moving average for frame rendering
    /// duration in seconds.
    pub fn frame_duration(&self) -> f64 { self.frame_duration.value }

    //pub fn events<'a>(&'a mut self) -> EventIterator<'a> { EventIterator { window: self } }
    pub fn events(&mut self) -> Vec<Event> {
        // XXX: Not returning an iterator because of locking crap. May want to
        // do mutable stuff on window in response to events.
        EventIterator { window: self }.collect::<Vec<Event>>()
    }

    /// Map screen position (eg. of a mouse cursor) to canvas position.
    pub fn screen_to_canvas(&self, screen_pos: V2<i32>) -> V2<i32> {
        self.resolution.screen_to_canvas(screen_pos)
    }

    pub fn get_screenshot(&self) -> image::ImageBuffer<image::Rgb<u8>, Vec<u8>> {
        use ::rgb::{to_srgb};
        let mut ret = self.buffer.read::<image::DynamicImage>().to_rgb();

        // Convert to sRGB
        // XXX: Probably horribly slow, can we make OpenGL do this?
        for p in ret.pixels_mut() {
            p.data[0] = (to_srgb(p.data[0] as f32 / 255.0) * 255.0).round() as u8;
            p.data[1] = (to_srgb(p.data[1] as f32 / 255.0) * 255.0).round() as u8;
            p.data[2] = (to_srgb(p.data[2] as f32 / 255.0) * 255.0).round() as u8;
        }

        ret
    }

    /// Show the graphics buffer on screen.
    fn show(&mut self) {
        let mut target = self.display.draw();

        // Clear the window.
        target.clear_color(0.0, 0.0, 0.0, 0.0);
        target.clear_depth(1.0);
        let (w, h) = self.display.get_framebuffer_dimensions();
        // Clip viewport dimensions to even to prevent rounding errors in
        // pixel perfect scaling.
        self.resolution.window = V2(w & !1, h & !1);

        // Build the geometry for the on-screen rectangle.
        let Rect(V2(sx, sy), V2(sw, sh)) = self.resolution.screen_rect();

        // XXX: This could use glium::Surface::blit_whole_color_to instead of
        // the handmade blitting, but that was buggy on Windows around
        // 2015-03.

        let vertices = {
            #[derive(Copy, Clone)]
            struct BlitVertex { pos: [f32; 2], tex_coord: [f32; 2] }
            implement_vertex!(BlitVertex, pos, tex_coord);

            glium::VertexBuffer::new(&self.display,
            &[
                BlitVertex { pos: [sx,    sy   ], tex_coord: [0.0, 0.0] },
                BlitVertex { pos: [sx+sw, sy   ], tex_coord: [1.0, 0.0] },
                BlitVertex { pos: [sx+sw, sy+sh], tex_coord: [1.0, 1.0] },
                BlitVertex { pos: [sx,    sy+sh], tex_coord: [0.0, 1.0] },
            ]).unwrap()
        };

        let indices = glium::IndexBuffer::new(&self.display,
            glium::index::PrimitiveType::TrianglesList, &[0u16, 1, 2, 0, 2, 3]).unwrap();

        // Set up the rest of the draw parameters.
        let mut params: glium::DrawParameters = Default::default();
        // Set an explicit viewport to apply the custom resolution that fixes
        // pixel perfect rounding errors.
        params.viewport = Some(glium::Rect{
            left: 0, bottom: 0,
            width: self.resolution.window.0,
            height: self.resolution.window.1 });

        let mag_filter = match self.resolution.magnify {
            CanvasMagnify::Smooth => glium::uniforms::MagnifySamplerFilter::Linear,
            _ => glium::uniforms::MagnifySamplerFilter::Nearest
        };

        let uniforms = glium::uniforms::UniformsStorage::new("tex",
            glium::uniforms::Sampler(&self.buffer, glium::uniforms::SamplerBehavior {
                magnify_filter: mag_filter,
                minify_filter: glium::uniforms::MinifySamplerFilter::Linear,
                .. Default::default() }));

        // Draw the graphics buffer to the window.
        target.draw(&vertices, &indices, &self.blit_shader, &uniforms, &params).unwrap();
        target.finish().unwrap();
    }
}

pub struct LogicalResolution {
    pub magnify: CanvasMagnify,
    /// Logical size, pixel on virtual canvas.
    pub canvas: V2<u32>,
    /// Physical size, window on screen.
    pub window: V2<u32>,
}

impl LogicalResolution {
    pub fn new(magnify: CanvasMagnify, canvas_size: V2<u32>, resolution: V2<u32>) -> LogicalResolution {
        LogicalResolution {
            magnify: magnify,
            canvas: canvas_size,
            window: resolution,
        }
    }

    pub fn screen_rect(&self) -> Rect<f32> {
        match self.magnify {
            CanvasMagnify::PixelPerfect => self.pixel_perfect(),
            _ => self.preserve_aspect(),
        }
    }

    /// Map screen position (eg. of a mouse cursor) to canvas position.
    pub fn screen_to_canvas(&self, V2(sx, sy): V2<i32>) -> V2<i32> {
        let Rect(V2(rx, ry), V2(rw, rh)) = self.screen_rect();

        // Transform to device coordinates.
        let sx = sx as f32 * 2.0 / self.window.0 as f32 - 1.0;
        let sy = sy as f32 * 2.0 / self.window.1 as f32 - 1.0;

        V2(((sx - rx) * self.canvas.0 as f32 / rw) as i32,
           ((sy - ry) * self.canvas.1 as f32 / rh) as i32)
    }

    /// A pixel perfect centered and scaled rectangle of resolution dim in a
    /// window of size area, mapped to OpenGL device coordinates.
#[inline(always)]
    fn pixel_perfect(&self) -> Rect<f32> {
        // Scale based on whichever of X or Y axis is the tighter fit.
        let mut scale = (self.window.0 as f32 / self.canvas.0 as f32)
            .min(self.window.1 as f32 / self.canvas.1 as f32);

        if scale > 1.0 {
            // Snap to pixel scale if more than 1 window pixel per canvas pixel.
            scale = scale.floor();
        }

        let dim = V2((scale * self.canvas.0 as f32) * 2.0 / self.window.0 as f32,
        (scale * self.canvas.1 as f32) * 2.0 / self.window.1 as f32);
        let offset = -dim / 2.0;
        Rect(offset, dim)
    }

#[inline(always)]
    fn preserve_aspect(&self) -> Rect<f32> {
        // Scale based on whichever of X or Y axis is the tighter fit.
        let scale = (self.window.0 as f32 / self.canvas.0 as f32)
            .min(self.window.1 as f32 / self.canvas.1 as f32);

        let dim = V2((scale * self.canvas.0 as f32) * 2.0 / self.window.0 as f32,
        (scale * self.canvas.1 as f32) * 2.0 / self.window.1 as f32);
        let offset = -dim / 2.0;
        Rect(offset, dim)
    }
}

pub struct EventIterator<'a> {
    window: &'a mut Window
}

impl<'a> Iterator for EventIterator<'a> {
    type Item = Event;
    fn next(&mut self) -> Option<Event> { self.window.translator.next() }
}
<|endoftext|>


<commit_before>// This program watches a configured directory for incoming files.
// When a file is detected the following procedure is used to handle it:
//
//  * First we digest the file, and add the digest to our manifest.
//    This way we can intelligently handle future events for this file
//
//  * Next we attempt to detect a mime type using various heuristics, incl:
//
//    - Quick detection w/ our magic bytes table
//    - Querying external tools (e.g: ffmpeg, imagemagick, et al.)
//    - etc.
//
//  * If we cannot handle the mime type: the file is moved to an exception
//    directory, and an error is logged somewhere the user will see it.
//
//  * If we can handle the mime type: the file is moved to the content store
//    and an entry is created in the database.
//
//    - This should ideally be done atomically so that other aqua utility
//      processes (e.g: sister agnes) don't mistakenly mark the file as an
//      orphan / missing / etc.
//

#[macro_use] extern crate log;

extern crate aqua;
extern crate clap;
extern crate diesel;
extern crate dotenv;
extern crate env_logger;
extern crate image;
extern crate notify;

use aqua::controllers::prelude::hash_file;
use aqua::models::{Entry, NewEntry};
use aqua::schema;
use clap::{Arg, App};
use diesel::prelude::*;
use diesel::pg::PgConnection;
use dotenv::dotenv;
use notify::{DebouncedEvent, Watcher, RecursiveMode, watcher};
use std::{env, fmt};
use std::convert::From;
use std::error::Error;
use std::fs::{self, File, OpenOptions};
use std::io::{self, Read, Write};
use std::path::{Path, PathBuf};
use std::sync::mpsc::channel;
use std::time::Duration;

#[derive(Debug)]
enum ProcessingError {
    DigestFailed,
    DetectionFailed,
    ThumbnailFailed,
    DbConnErr(diesel::ConnectionError),
    DbQueryErr(diesel::result::Error),
    IoErr(io::Error),
    ImageErr(image::ImageError),
}

impl Error for ProcessingError {
    fn description(&self) -> &str {
        match *self {
            // internal errors
            ProcessingError::DigestFailed      => "Unhandled error while generating SHA256 digest",
            ProcessingError::DetectionFailed   => "The file's type could not be detected",
            ProcessingError::ThumbnailFailed   => "The thumbnail could not be generated",

            // external errors
            ProcessingError::DbConnErr(ref inner)  => inner.description(),
            ProcessingError::DbQueryErr(ref inner) => inner.description(),
            ProcessingError::IoErr(ref inner)      => inner.description(),
            ProcessingError::ImageErr(ref inner)   => inner.description(),
        }
    }

    fn cause(&self) -> Option<&Error> {
        match *self {
            ProcessingError::DbConnErr(ref err)  => Some(err),
            ProcessingError::DbQueryErr(ref err) => Some(err),
            ProcessingError::IoErr(ref err)      => Some(err),
            ProcessingError::ImageErr(ref err)   => Some(err),
            _ => None,
        }
    }
}

impl fmt::Display for ProcessingError {
    fn fmt(&self, f: &mut fmt::Formatter) -> Result<(), fmt::Error> {
        match *self {
            // TODO: better display impl.
            _ => write!(f, "{}", self.description()),
        }
    }
}

impl From<diesel::ConnectionError> for ProcessingError {
    fn from(err: diesel::ConnectionError) -> Self { ProcessingError::DbConnErr(err) }
}

impl From<diesel::result::Error> for ProcessingError {
    fn from(err: diesel::result::Error) -> Self { ProcessingError::DbQueryErr(err) }
}

impl From<image::ImageError> for ProcessingError {
    fn from(err: image::ImageError) -> Self { ProcessingError::ImageErr(err) }
}

impl From<io::Error> for ProcessingError {
    fn from(err: io::Error) -> Self { ProcessingError::IoErr(err) }
}

fn main() {
    dotenv().expect("must provide .env file, see README (TODO: haha jk)");
    env_logger::init().expect("could not initialize console logging");

    // read command line arguments
    let matches = App::new("aqua-watch")
        .version("0.1.0")
        .author("himechi <hime@localhost>")
        .about("Watches a directory for new files and moves them to the `aqua` content store.")
        .arg(Arg::with_name("INPUT")
             .help("Determines the input directory to be watched.")
             .required(true)
             .index(1))
        .arg(Arg::with_name("OUTPUT")
             .help("The root of the aqua content store where files will be moved.")
             .required(true)
             .index(2))
        .get_matches();

    let dropbox_path  = matches.value_of("INPUT").unwrap();
    let content_store = matches.value_of("OUTPUT").unwrap();

    // setup fs watcher
    let (fs_tx, fs_rx) = channel();
    let mut fs_watcher = watcher(fs_tx, Duration::from_millis(1000))
        .expect("could not create file system watcher!");

    fs_watcher.watch(dropbox_path, RecursiveMode::NonRecursive)
        .expect("could not enroll dropbox in fs events queue");

    // process filesystem events ...
    loop {
        match fs_rx.recv() {
            Ok(DebouncedEvent::Create(path)) => {
                if path.is_file() {
                    match handle_new_file(path, content_store) {
                        Ok(_res) => info!("file processed successfully ..."),
                        Err(msg) => warn!("could not process file: {}", msg.description()),
                    };
                }
                else { info!("directory created, ignoring ..."); }
            },
            Ok(event) => info!("unhandled evt: {:?}", event),
            Err(msg) => warn!("fs err: {}", msg),
        }
    }
}

fn handle_new_file(path: PathBuf, content_store: &str) -> Result<(), ProcessingError> {
    let digest = hash_file(path.as_path())
        .ok_or(ProcessingError::DigestFailed)?;

    let mut file = OpenOptions::new()
        .read(true)
        .write(false)
        .create_new(false)
        .open(path.as_path())?;

    // read file into memory
    let mut buf = vec![];
    file.read_to_end(&mut buf)?;

    if let Some(image_metadata) = aqua::util::mime_detect(&buf) {
        info!("got an image ...");
        process_image(content_store, &digest, &buf)?;
        move_file(path.as_path(), content_store, &digest, image_metadata.extension())?;

        let db_entry = create_db_entry(&digest, image_metadata.mime())?;
        info!("inserted: {:?} into database", db_entry);

        Ok(())
    } else if let Some(_nil) = ffmpeg_detect(path.as_path()) {
        unreachable!()
    } else {
        Err(ProcessingError::DetectionFailed)
    }
}

fn ffmpeg_detect(path: &Path) -> Option<()> {
    None
}

fn establish_connection() -> Result<PgConnection, ProcessingError> {
    let database_url = env::var("DATABASE_URL")
        .expect("DATABASE_URL not set in `.env` file !!!");

    Ok(PgConnection::establish(&database_url)?)
}

// create entry in database
fn create_db_entry(digest: &str, mime_ty: &str) -> Result<Entry, ProcessingError> {
    let pg_conn = establish_connection()?;
    let aqua_entry = NewEntry { hash: &digest, mime: Some(&mime_ty) };
    let entry: Result<Entry, diesel::result::Error> = diesel::insert(&aqua_entry)
        .into(schema::entries::table)
        .get_result(&pg_conn);

    Ok(entry?)
}

// moves the file from `src_path` to the `content_store` based on its digest
fn move_file(src_path: &Path, content_store: &str, digest: &str, file_ext: &str) -> Result<(), ProcessingError> {
    // carve out a bucket based on first byte of SHA256 digest
    // create the bucket if it does not exist
    let file_bucket    = format!("f{}", &digest[0..2]);
    let file_filename  = format!("{}.{}", &digest, file_ext);

    // create destination path
    let dest = PathBuf::from(content_store)
        .join(file_bucket)
        .join(file_filename);

    // TODO: bad error type ...
    let bucket_dir = dest.parent().ok_or(ProcessingError::ThumbnailFailed)?;
    fs::create_dir_all(bucket_dir)?;

    // move the file
    Ok(fs::rename(src_path, &dest)?)
}

// creates a thumbnail in the content store for the specified digest
// this expects an `ImageMeta` structure describing the input.
fn process_image(content_store: &str, digest: &str, buf: &[u8]) -> Result<(), ProcessingError> {
    // create in memory thumbnail
    let image = image::load_from_memory(&buf)?;

    let thumb = image.resize(200, 200, image::FilterType::Nearest);
    let thumb_bucket   = format!("t{}", &digest[0..2]);
    let thumb_filename = format!("{}.thumbnail", &digest);

    // store them in content store
    let dest = PathBuf::from(content_store)
        .join(thumb_bucket)
        .join(thumb_filename);

    // write thumbnail file to disk
    fs::create_dir_all(dest.parent().unwrap())?;
    let mut dest_file = File::create(dest)?;
    thumb.save(&mut dest_file, image::ImageFormat::JPEG)?;
    Ok(dest_file.flush()?)
}
<commit_msg>add detection using ffprobe/ffmpeg<commit_after>// This program watches a configured directory for incoming files.
// When a file is detected the following procedure is used to handle it:
//
//  * First we digest the file, and add the digest to our manifest.
//    This way we can intelligently handle future events for this file
//
//  * Next we attempt to detect a mime type using various heuristics, incl:
//
//    - Quick detection w/ our magic bytes table
//    - Querying external tools (e.g: ffmpeg, imagemagick, et al.)
//    - etc.
//
//  * If we cannot handle the mime type: the file is moved to an exception
//    directory, and an error is logged somewhere the user will see it.
//
//  * If we can handle the mime type: the file is moved to the content store
//    and an entry is created in the database.
//
//    - This should ideally be done atomically so that other aqua utility
//      processes (e.g: sister agnes) don't mistakenly mark the file as an
//      orphan / missing / etc.
//

#[macro_use] extern crate log;
#[macro_use] extern crate serde_derive;

extern crate aqua;
extern crate clap;
extern crate diesel;
extern crate dotenv;
extern crate env_logger;
extern crate image;
extern crate notify;
extern crate serde;
extern crate serde_json;

use aqua::controllers::prelude::hash_file;
use aqua::models::{Entry, NewEntry};
use aqua::schema;
use clap::{Arg, App};
use diesel::prelude::*;
use diesel::pg::PgConnection;
use dotenv::dotenv;
use notify::{DebouncedEvent, Watcher, RecursiveMode, watcher};
use std::{env, fmt, process};
use std::collections::HashMap;
use std::convert::From;
use std::error::Error;
use std::fs::{self, File, OpenOptions};
use std::io::{self, Read, Write};
use std::path::{Path, PathBuf};
use std::sync::mpsc::channel;
use std::time::Duration;

#[derive(Debug)]
enum ProcessingError {
    DigestFailed,
    DetectionFailed,
    ThumbnailFailed,

    DbConnErr(diesel::ConnectionError),
    DbQueryErr(diesel::result::Error),
    IoErr(io::Error),
    ImageErr(image::ImageError),
    Misc(Box<Error>),
}

impl Error for ProcessingError {
    fn description(&self) -> &str {
        match *self {
            // internal errors
            ProcessingError::DigestFailed      => "Unhandled error while generating SHA256 digest",
            ProcessingError::DetectionFailed   => "The file's type could not be detected",
            ProcessingError::ThumbnailFailed   => "The thumbnail could not be generated",

            // external errors
            ProcessingError::DbConnErr(ref inner)  => inner.description(),
            ProcessingError::DbQueryErr(ref inner) => inner.description(),
            ProcessingError::IoErr(ref inner)      => inner.description(),
            ProcessingError::ImageErr(ref inner)   => inner.description(),
            ProcessingError::Misc(ref inner)       => inner.description(),
        }
    }

    fn cause(&self) -> Option<&Error> {
        match *self {
            ProcessingError::DbConnErr(ref err)  => Some(err),
            ProcessingError::DbQueryErr(ref err) => Some(err),
            ProcessingError::IoErr(ref err)      => Some(err),
            ProcessingError::ImageErr(ref err)   => Some(err),
            ProcessingError::Misc(ref err)       => Some(err.as_ref()),
            _ => None,
        }
    }
}

impl fmt::Display for ProcessingError {
    fn fmt(&self, f: &mut fmt::Formatter) -> Result<(), fmt::Error> {
        match *self {
            // TODO: better display impl.
            _ => write!(f, "{}", self.description()),
        }
    }
}

impl From<diesel::ConnectionError> for ProcessingError {
    fn from(err: diesel::ConnectionError) -> Self { ProcessingError::DbConnErr(err) }
}

impl From<diesel::result::Error> for ProcessingError {
    fn from(err: diesel::result::Error) -> Self { ProcessingError::DbQueryErr(err) }
}

impl From<image::ImageError> for ProcessingError {
    fn from(err: image::ImageError) -> Self { ProcessingError::ImageErr(err) }
}

impl From<io::Error> for ProcessingError {
    fn from(err: io::Error) -> Self { ProcessingError::IoErr(err) }
}

impl From<serde_json::Error> for ProcessingError {
    fn from(err: serde_json::Error) -> Self { ProcessingError::Misc(Box::new(err)) }
}

impl From<std::string::FromUtf8Error> for ProcessingError {
    fn from(err: std::string::FromUtf8Error) -> Self { ProcessingError::Misc(Box::new(err)) }
}

fn main() {
    dotenv().expect("must provide .env file, see README (TODO: haha jk)");
    env_logger::init().expect("could not initialize console logging");

    // read command line arguments
    let matches = App::new("aqua-watch")
        .version("0.1.0")
        .author("himechi <hime@localhost>")
        .about("Watches a directory for new files and moves them to the `aqua` content store.")
        .arg(Arg::with_name("INPUT")
             .help("Determines the input directory to be watched.")
             .required(true)
             .index(1))
        .arg(Arg::with_name("OUTPUT")
             .help("The root of the aqua content store where files will be moved.")
             .required(true)
             .index(2))
        .get_matches();

    let dropbox_path  = matches.value_of("INPUT").unwrap();
    let content_store = matches.value_of("OUTPUT").unwrap();

    // setup fs watcher
    let (fs_tx, fs_rx) = channel();
    let mut fs_watcher = watcher(fs_tx, Duration::from_millis(1000))
        .expect("could not create file system watcher!");

    fs_watcher.watch(dropbox_path, RecursiveMode::NonRecursive)
        .expect("could not enroll dropbox in fs events queue");

    // process filesystem events ...
    loop {
        match fs_rx.recv() {
            Ok(DebouncedEvent::Create(path)) => {
                if path.is_file() {
                    match handle_new_file(path, content_store) {
                        Ok(_res) => info!("file processed successfully ..."),
                        Err(msg) => warn!("could not process file: {:?} (inner: {:?}", msg, msg.cause()),
                    };
                }
                else { info!("directory created, ignoring ..."); }
            },
            Ok(event) => info!("unhandled evt: {:?}", event),
            Err(msg) => warn!("fs err: {}", msg),
        }
    }
}

fn handle_new_file(path: PathBuf, content_store: &str) -> Result<(), ProcessingError> {
    let digest = hash_file(path.as_path())
        .ok_or(ProcessingError::DigestFailed)?;

    let mut file = OpenOptions::new()
        .read(true)
        .write(false)
        .create_new(false)
        .open(path.as_path())?;

    // read file into memory
    let mut buf = vec![];
    file.read_to_end(&mut buf)?;

    if let Some(image_metadata) = aqua::util::mime_detect(&buf) {
        info!("got an image ...");
        process_image(content_store, &digest, &buf)?;
        move_file(path.as_path(), content_store, &digest, image_metadata.extension())?;

        let db_entry = create_db_entry(&digest, image_metadata.mime())?;
        info!("inserted: {:?} into database", db_entry);

        Ok(())
    } else if let Some(ffmpeg_metadata) = ffmpeg_detect(path.as_path())? {
        info!("got an video ...");
        process_video(content_store, &digest, &path)?;
        move_file(path.as_path(), content_store, &digest, ffmpeg_metadata.ext)?;

        let db_entry = create_db_entry(&digest, ffmpeg_metadata.mime)?;
        info!("inserted: {:?} into database", db_entry);

        Ok(())
    } else {
        Err(ProcessingError::DetectionFailed)
    }
}

#[derive(Debug, Serialize, Deserialize)]
struct FFProbeResult {
    format: Option<FFProbeFormat>,
}

#[derive(Debug, Serialize, Deserialize)]
struct FFProbeFormat {
    filename:    String,
    nb_streams:  i32,
    format_name: String,
    start_time:  String, // NOTE: these appear to be fixed decimal
    duration:    String, // NOTE: these appear to be fixed decimal
    size:        String, // NOTE: appears to be an integer
    bit_rate:    String, // NOTE: appears to be an integer
    probe_score: i32,    // NOTE: accuracy of the detection? (???)

    tags: HashMap<String, String>, // NOTE: not sure this is correct type
}

struct FFProbeMeta {
    pub mime: &'static str,
    pub ext:  &'static str,
}

fn ffmpeg_detect(path: &Path) -> Result<Option<FFProbeMeta>, ProcessingError> {
    let ffprobe_cmd = process::Command::new("ffprobe")
        .arg("-v").arg("quiet")            // silence debug output
        .arg("-hide_banner")               // don't print ffmpeg configuration
        .arg("-print_format").arg("json") // serialize to json
        .arg("-show_format")              // use the serializer
        .arg("-i").arg(path.as_os_str())  // set the input to current file
        .output()?;

    let json_str = String::from_utf8(ffprobe_cmd.stdout)?;
    let probe_result: FFProbeResult = serde_json::from_str(&json_str)?;
    info!("got result: {:?}", probe_result);

    // see if ffprobe was able to determine the file type ...
    let probe_result = match probe_result.format {
        Some(format_data) => format_data,
        None => return Ok(None),
    };

    // TODO: handle alternate streams
    if probe_result.nb_streams != 1 { return Err(ProcessingError::DetectionFailed) }

    let meta_data = match probe_result.format_name.as_ref() {
        "matroska,webm" => Some(FFProbeMeta { mime: "video/webm", ext: "webm" }),
        _ => None,
    };

    Ok(meta_data) // return the ffprobe results, if we know how to handle them ...
}

fn establish_connection() -> Result<PgConnection, ProcessingError> {
    let database_url = env::var("DATABASE_URL")
        .expect("DATABASE_URL not set in `.env` file !!!");

    Ok(PgConnection::establish(&database_url)?)
}

// create entry in database
fn create_db_entry(digest: &str, mime_ty: &str) -> Result<Entry, ProcessingError> {
    let pg_conn = establish_connection()?;
    let aqua_entry = NewEntry { hash: &digest, mime: Some(&mime_ty) };
    let entry: Result<Entry, diesel::result::Error> = diesel::insert(&aqua_entry)
        .into(schema::entries::table)
        .get_result(&pg_conn);

    Ok(entry?)
}

// moves the file from `src_path` to the `content_store` based on its digest
fn move_file(src_path: &Path, content_store: &str, digest: &str, file_ext: &str) -> Result<(), ProcessingError> {
    // carve out a bucket based on first byte of SHA256 digest
    // create the bucket if it does not exist
    let file_bucket    = format!("f{}", &digest[0..2]);
    let file_filename  = format!("{}.{}", &digest, file_ext);

    // create destination path
    let dest = PathBuf::from(content_store)
        .join(file_bucket)
        .join(file_filename);

    // TODO: bad error type ...
    let bucket_dir = dest.parent().ok_or(ProcessingError::ThumbnailFailed)?;
    fs::create_dir_all(bucket_dir)?;

    // move the file
    Ok(fs::rename(src_path, &dest)?)
}

// creates a thumbnail in the content store for the specified digest
// this expects an `ImageMeta` structure describing the input.
fn process_image(content_store: &str, digest: &str, buf: &[u8]) -> Result<(), ProcessingError> {
    // create in memory thumbnail
    let image = image::load_from_memory(&buf)?;

    let thumb = image.resize(200, 200, image::FilterType::Nearest);
    let thumb_bucket   = format!("t{}", &digest[0..2]);
    let thumb_filename = format!("{}.thumbnail", &digest);

    // store them in content store
    let dest = PathBuf::from(content_store)
        .join(thumb_bucket)
        .join(thumb_filename);

    // write thumbnail file to disk
    fs::create_dir_all(dest.parent().unwrap())?;
    let mut dest_file = File::create(dest)?;
    thumb.save(&mut dest_file, image::ImageFormat::JPEG)?;
    Ok(dest_file.flush()?)
}

fn process_video(content_store: &str, digest: &str, src: &Path) -> Result<(), ProcessingError> {
    let thumb_bucket   = format!("t{}", &digest[0..2]);
    let thumb_filename = format!("{}.thumbnail", &digest);

    // store them in content store
    let dest = PathBuf::from(content_store)
        .join(thumb_bucket)
        .join(thumb_filename);

    // write thumbnail file to disk
    fs::create_dir_all(dest.parent().unwrap())?;
    let ffmpeg_cmd = process::Command::new("ffmpeg")
        .arg("-i").arg(src.as_os_str())            // the input file
        .arg("-vf").arg("thumbnail,scale=200:200") // have ffmpeg seek for a "thumbnail"
        .arg("-frames:v").arg("1")                 // take a single frame
        .arg("-f").arg("mjpeg")                    // save it as jpeg
        .arg(dest.as_path().as_os_str())           // into the content store
        .output()?;

    debug!("ffmpeg stderr: {}", String::from_utf8_lossy(&ffmpeg_cmd.stderr));
    debug!("ffmpeg stdout: {}", String::from_utf8_lossy(&ffmpeg_cmd.stdout));

    info!("digest? {}", digest);
    info!("dest exists? {:?} => {}", dest, dest.is_file());

    match dest.is_file() {
        true  => Ok(()),
        false => Err(ProcessingError::ThumbnailFailed),
    }
}
<|endoftext|>


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

extern crate mimir;
extern crate mimirsbrunn;
extern crate navitia_model;
#[macro_use]
extern crate slog;
#[macro_use]
extern crate slog_scope;
#[macro_use]
extern crate structopt;

use std::path::PathBuf;
use structopt::StructOpt;
use mimirsbrunn::stops::*;
use navitia_model::objects as navitia;
use navitia_model::collection::Idx;

#[derive(Debug, StructOpt)]
struct Args {
    /// NTFS directory.
    #[structopt(short = "i", long = "input", parse(from_os_str), default_value = ".")]
    input: PathBuf,
    /// Name of the dataset.
    #[structopt(short = "d", long = "dataset", default_value = "fr")]
    dataset: String,
    /// Elasticsearch parameters.
    #[structopt(short = "c", long = "connection-string",
                default_value = "http://localhost:9200/munin")]
    connection_string: String,
    /// Deprecated option.
    #[structopt(short = "C", long = "city-level")]
    city_level: Option<String>,
}

fn to_mimir(
    idx: Idx<navitia::StopArea>,
    stop_area: &navitia::StopArea,
    navitia: &navitia_model::PtObjects,
) -> mimir::Stop {
    let commercial_modes = navitia
        .get_corresponding_from_idx(idx)
        .into_iter()
        .map(|cm_idx| mimir::CommercialMode {
            id: format!("commercial_mode:{}", navitia.commercial_modes[cm_idx].id),
            name: navitia.commercial_modes[cm_idx].name.clone(),
        })
        .collect();
    let physical_modes = navitia
        .get_corresponding_from_idx(idx)
        .into_iter()
        .map(|pm_idx| mimir::PhysicalMode {
            id: format!("physical_mode:{}", navitia.physical_modes[pm_idx].id),
            name: navitia.physical_modes[pm_idx].name.clone(),
        })
        .collect();
    let mut comments: Vec<mimir::Comment> = vec![];
    for comment in navitia.comments.iter_from(&stop_area.comment_links) {
        comments.push(mimir::Comment {
            id: format!("comments:{}", comment.id),
            name: comment.name.clone(),
        })
    }

    let feed_publishers = navitia
        .get_corresponding_from_idx(idx)
        .into_iter()
        .map(|contrib_idx| mimir::FeedPublisher {
            id: navitia.contributors[contrib_idx].id.clone(),
            name: navitia.contributors[contrib_idx].name.clone(),
            license: navitia.contributors[contrib_idx]
                .license
                .clone()
                .unwrap_or_else(|| "".into()),
            url: navitia.contributors[contrib_idx]
                .website
                .clone()
                .unwrap_or_else(|| "".into()),
        })
        .collect();

    mimir::Stop {
        id: format!("stop_area:{}", stop_area.id),
        label: stop_area.name.clone(),
        name: stop_area.name.clone(),
        coord: mimir::Coord::new(stop_area.coord.lon, stop_area.coord.lat),
        commercial_modes: commercial_modes,
        physical_modes: physical_modes,
        administrative_regions: vec![],
        weight: 0.,
        zip_codes: vec![],
        coverages: vec![],
        comments: comments,
        timezone: stop_area.timezone.clone().unwrap_or(format!("")),
        codes: stop_area
            .codes
            .iter()
            .map(|&(ref t, ref v)| mimir::Code {
                name: t.clone(),
                value: v.clone(),
            })
            .collect(),
        properties: stop_area
            .object_properties
            .iter()
            .map(|&(ref k, ref v)| mimir::Property {
                key: k.clone(),
                value: v.clone(),
            })
            .collect(),
        feed_publishers: feed_publishers,
    }
}

fn main() {
    let _guard = mimir::logger_init();
    info!("Launching ntfs2mimir...");

    let args = Args::from_args();
    if args.city_level.is_some() {
        warn!("city-level option is deprecated, it now has no effect.");
    }
    if let Err(err) = run(args) {
        for cause in err.causes() {
            eprintln!("{}", cause);
        }
        std::process::exit(1);
    }
}
fn run(args: Args) -> Result<(), navitia_model::Error> {
    let navitia = navitia_model::ntfs::read(&args.input)?;
    let nb_stop_points = navitia
        .stop_areas
        .iter()
        .map(|(idx, sa)| {
            let id = format!("stop_area:{}", sa.id);
            let nb_stop_points = navitia
                .get_corresponding_from_idx::<_, navitia::StopPoint>(idx)
                .len();
            (id, nb_stop_points as u32)
        })
        .collect();
    let mut stops: Vec<mimir::Stop> = navitia
        .stop_areas
        .iter()
        .map(|(idx, sa)| to_mimir(idx, sa, &navitia))
        .collect();
    set_weights(stops.iter_mut(), &nb_stop_points);
    import_stops(stops, &args.connection_string, &args.dataset);
    Ok(())
}
<commit_msg>Refactoring function (add comment)<commit_after>// Copyright © 2016, Canal TP and/or its affiliates. All rights reserved.
//
// This file is part of Navitia,
//     the software to build cool stuff with public transport.
//
// Hope you'll enjoy and contribute to this project,
//     powered by Canal TP (www.canaltp.fr).
// Help us simplify mobility and open public transport:
//     a non ending quest to the responsive locomotion way of traveling!
//
// LICENCE: This program is free software; you can redistribute it
// and/or modify it under the terms of the GNU Affero General Public
// License as published by the Free Software Foundation, either
// version 3 of the License, or (at your option) any later version.
//
// This program is distributed in the hope that it will be useful, but
// WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
// Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public
// License along with this program. If not, see
// <http://www.gnu.org/licenses/>.
//
// Stay tuned using
// twitter @navitia
// IRC #navitia on freenode
// https://groups.google.com/d/forum/navitia
// www.navitia.io

extern crate mimir;
extern crate mimirsbrunn;
extern crate navitia_model;
#[macro_use]
extern crate slog;
#[macro_use]
extern crate slog_scope;
#[macro_use]
extern crate structopt;

use std::path::PathBuf;
use structopt::StructOpt;
use mimirsbrunn::stops::*;
use navitia_model::objects as navitia;
use navitia_model::collection::Idx;

#[derive(Debug, StructOpt)]
struct Args {
    /// NTFS directory.
    #[structopt(short = "i", long = "input", parse(from_os_str), default_value = ".")]
    input: PathBuf,
    /// Name of the dataset.
    #[structopt(short = "d", long = "dataset", default_value = "fr")]
    dataset: String,
    /// Elasticsearch parameters.
    #[structopt(short = "c", long = "connection-string",
                default_value = "http://localhost:9200/munin")]
    connection_string: String,
    /// Deprecated option.
    #[structopt(short = "C", long = "city-level")]
    city_level: Option<String>,
}

fn to_mimir(
    idx: Idx<navitia::StopArea>,
    stop_area: &navitia::StopArea,
    navitia: &navitia_model::PtObjects,
) -> mimir::Stop {
    let commercial_modes = navitia
        .get_corresponding_from_idx(idx)
        .into_iter()
        .map(|cm_idx| mimir::CommercialMode {
            id: format!("commercial_mode:{}", navitia.commercial_modes[cm_idx].id),
            name: navitia.commercial_modes[cm_idx].name.clone(),
        })
        .collect();
    let physical_modes = navitia
        .get_corresponding_from_idx(idx)
        .into_iter()
        .map(|pm_idx| mimir::PhysicalMode {
            id: format!("physical_mode:{}", navitia.physical_modes[pm_idx].id),
            name: navitia.physical_modes[pm_idx].name.clone(),
        })
        .collect();
    let comments = navitia
        .comments
        .iter_from(&stop_area.comment_links)
        .map(|comment| mimir::Comment {
            id: format!("comments:{}", comment.id),
            name: comment.name.clone(),
        })
        .collect();
    let feed_publishers = navitia
        .get_corresponding_from_idx(idx)
        .into_iter()
        .map(|contrib_idx| mimir::FeedPublisher {
            id: navitia.contributors[contrib_idx].id.clone(),
            name: navitia.contributors[contrib_idx].name.clone(),
            license: navitia.contributors[contrib_idx]
                .license
                .clone()
                .unwrap_or_else(|| "".into()),
            url: navitia.contributors[contrib_idx]
                .website
                .clone()
                .unwrap_or_else(|| "".into()),
        })
        .collect();

    mimir::Stop {
        id: format!("stop_area:{}", stop_area.id),
        label: stop_area.name.clone(),
        name: stop_area.name.clone(),
        coord: mimir::Coord::new(stop_area.coord.lon, stop_area.coord.lat),
        commercial_modes: commercial_modes,
        physical_modes: physical_modes,
        administrative_regions: vec![],
        weight: 0.,
        zip_codes: vec![],
        coverages: vec![],
        comments: comments,
        timezone: stop_area.timezone.clone().unwrap_or(format!("")),
        codes: stop_area
            .codes
            .iter()
            .map(|&(ref t, ref v)| mimir::Code {
                name: t.clone(),
                value: v.clone(),
            })
            .collect(),
        properties: stop_area
            .object_properties
            .iter()
            .map(|&(ref k, ref v)| mimir::Property {
                key: k.clone(),
                value: v.clone(),
            })
            .collect(),
        feed_publishers: feed_publishers,
    }
}

fn main() {
    let _guard = mimir::logger_init();
    info!("Launching ntfs2mimir...");

    let args = Args::from_args();
    if args.city_level.is_some() {
        warn!("city-level option is deprecated, it now has no effect.");
    }
    if let Err(err) = run(args) {
        for cause in err.causes() {
            eprintln!("{}", cause);
        }
        std::process::exit(1);
    }
}
fn run(args: Args) -> Result<(), navitia_model::Error> {
    let navitia = navitia_model::ntfs::read(&args.input)?;
    let nb_stop_points = navitia
        .stop_areas
        .iter()
        .map(|(idx, sa)| {
            let id = format!("stop_area:{}", sa.id);
            let nb_stop_points = navitia
                .get_corresponding_from_idx::<_, navitia::StopPoint>(idx)
                .len();
            (id, nb_stop_points as u32)
        })
        .collect();
    let mut stops: Vec<mimir::Stop> = navitia
        .stop_areas
        .iter()
        .map(|(idx, sa)| to_mimir(idx, sa, &navitia))
        .collect();
    set_weights(stops.iter_mut(), &nb_stop_points);
    import_stops(stops, &args.connection_string, &args.dataset);
    Ok(())
}
<|endoftext|>


<commit_before>use enum_primitive::FromPrimitive;
use super::Mapper;

struct BankSelection(u8);

impl BankSelection {
    fn set_mode(&mut self, mode: BankMode) {
        self.0 = (self.0 & 0b0111_1111) | ((mode as u8) << 7)
    }

    fn mode(&self) -> BankMode {
        BankMode::from_u8((self.0 & 0b1000_0000) >> 7).unwrap()
    }

    fn set_upper(&mut self, val: u8) {
        self.0 = (self.0 & 0b1001_1111) | ((val << 5) & 0b0110_0000);
    }

    fn upper(&self) -> u8 {
        (self.0 & 0b0110_0000) >> 5
    }

    fn set_lower(&mut self, val: u8) {
        self.0 = (self.0 & 0b1110_0000) | (val & 0b0001_1111);
    }

    fn lower(&self) -> u8 {
        self.0 & 0b0001_1111
    }

    fn rom_bank(&self) -> usize {
        let bank = if self.mode() == BankMode::Rom {
            (self.upper() << 5) | self.lower()
        } else {
            self.lower()
        };

        (match bank {
            0 | 0x20 | 0x40 | 0x60 => bank + 1,
            _ => bank
        }) as usize
    }

    fn ram_bank(&self) -> usize {
        if self.mode() == BankMode::Ram {
            self.upper() as usize
        } else {
            0
        }
    }
}

enum_from_primitive! {
#[derive(PartialEq)]
enum BankMode {
    Rom = 0,
    Ram = 1
}
}

pub struct Mbc1 {
    rom_banks: usize,
    ram_banks: usize,
    ram_enabled: bool,
    ram_data: Vec<u8>,
    bank_selection: BankSelection
}

impl Mbc1 {
    pub fn new(rom_banks: usize, ram_banks: usize) -> Self {
        Mbc1 {
            rom_banks: rom_banks,
            ram_banks: ram_banks,
            ram_enabled: false,
            ram_data: vec![0; (ram_banks * 0x2000)],
            bank_selection: BankSelection(0)
        }
    }

    fn rom_index(&self, addr: u16) -> usize {
        let bank = self.bank_selection.rom_bank() & (self.rom_banks - 1);
        ((addr as usize) - 0x4000) + (bank * 0x4000)
    }

    fn ram_index(&self, addr: u16) -> usize {
        let bank = self.bank_selection.ram_bank() & (self.ram_banks - 1);

        ((addr as usize) - 0xA000) + (bank * 0x2000)
    }
}

impl Mapper for Mbc1 {
    fn read(&self, rom: &Vec<u8>, addr: u16) -> u8 {
        match addr {
            0x0000...0x3FFF => {
                // ROM bank 0
                rom[addr as usize]
            },
            0x4000...0x7FFF => {
                // Switchable ROM bank
                rom[self.rom_index(addr)]
            },
            0xA000...0xBFFF => {
                // Switchable RAM bank
                if self.ram_enabled {
                    self.ram_data[self.ram_index(addr)]
                } else {
                    0xFF
                }
            },
            _ => panic!("Address {:#X} not handled by MBC1")
        }
    }

    fn write(&mut self, addr: u16, val: u8) {
        match addr {
            0x0000...0x1FFF => {
                // Writing to this space toggles RAM
                // Any value with A in the lower bits enables it
                self.ram_enabled = (val & 0x0A) == 0x0A;
            },
            0x2000...0x3FFF => {
                // Writing to this space switches ROM bank lower bits
                self.bank_selection.set_lower(val);
            },
            0x4000...0x5FFF => {
                // Writing to this space switches the ROM bank upper bits/RAM bank
                self.bank_selection.set_upper(val);
            },
            0x6000...0x7FFF => {
                // Writing to this space toggles RAM/ROM bank mode. 1 = RAM mode
                self.bank_selection.set_mode(BankMode::from_u8(val & 1).unwrap_or_else(|| panic!("Unknown bank mode: {:#X}", val)));
            },
            0xA000...0xBFFF => {
                // Write to RAM Bank
                if self.ram_enabled {
                    let index = self.ram_index(addr);
                    self.ram_data[index] = val;
                }
            },
             _ => panic!("Address {:#X} not handled by MBC1")
        }
    }
}<commit_msg>MBC1 crashes Fixes underflow crash when an MBC1 cart has zero ROM or RAM banks.<commit_after>use enum_primitive::FromPrimitive;
use super::Mapper;

struct BankSelection(u8);

impl BankSelection {
    fn set_mode(&mut self, mode: BankMode) {
        self.0 = (self.0 & 0b0111_1111) | ((mode as u8) << 7)
    }

    fn mode(&self) -> BankMode {
        BankMode::from_u8((self.0 & 0b1000_0000) >> 7).unwrap()
    }

    fn set_upper(&mut self, val: u8) {
        self.0 = (self.0 & 0b1001_1111) | ((val << 5) & 0b0110_0000);
    }

    fn upper(&self) -> u8 {
        (self.0 & 0b0110_0000) >> 5
    }

    fn set_lower(&mut self, val: u8) {
        self.0 = (self.0 & 0b1110_0000) | (val & 0b0001_1111);
    }

    fn lower(&self) -> u8 {
        self.0 & 0b0001_1111
    }

    fn rom_bank(&self) -> usize {
        let bank = if self.mode() == BankMode::Rom {
            (self.upper() << 5) | self.lower()
        } else {
            self.lower()
        };

        (match bank {
            0 | 0x20 | 0x40 | 0x60 => bank + 1,
            _ => bank
        }) as usize
    }

    fn ram_bank(&self) -> usize {
        if self.mode() == BankMode::Ram {
            self.upper() as usize
        } else {
            0
        }
    }
}

enum_from_primitive! {
#[derive(PartialEq)]
enum BankMode {
    Rom = 0,
    Ram = 1
}
}

pub struct Mbc1 {
    rom_banks: usize,
    ram_banks: usize,
    ram_enabled: bool,
    ram_data: Vec<u8>,
    bank_selection: BankSelection
}

impl Mbc1 {
    pub fn new(rom_banks: usize, ram_banks: usize) -> Self {
        Mbc1 {
            rom_banks: rom_banks,
            ram_banks: ram_banks,
            ram_enabled: false,
            ram_data: vec![0; (ram_banks * 0x2000)],
            bank_selection: BankSelection(0)
        }
    }

    fn rom_index(&self, addr: u16) -> usize {
        let bank = if self.rom_banks == 0 {
            1
        } else {
            self.bank_selection.rom_bank() & (self.rom_banks - 1)
        };
        ((addr as usize) - 0x4000) + (bank * 0x4000)
    }

    fn ram_index(&self, addr: u16) -> usize {
        let bank = if self.ram_banks == 0 {
            1
        } else {
            self.bank_selection.ram_bank() & (self.ram_banks - 1)
        };
        ((addr as usize) - 0xA000) + (bank * 0x2000)
    }
}

impl Mapper for Mbc1 {
    fn read(&self, rom: &Vec<u8>, addr: u16) -> u8 {
        match addr {
            0x0000...0x3FFF => {
                // ROM bank 0
                rom[addr as usize]
            },
            0x4000...0x7FFF => {
                // Switchable ROM bank
                rom[self.rom_index(addr)]
            },
            0xA000...0xBFFF => {
                // Switchable RAM bank
                if self.ram_enabled {
                    self.ram_data[self.ram_index(addr)]
                } else {
                    0xFF
                }
            },
            _ => panic!("Address {:#X} not handled by MBC1")
        }
    }

    fn write(&mut self, addr: u16, val: u8) {
        match addr {
            0x0000...0x1FFF => {
                // Writing to this space toggles RAM
                // Any value with A in the lower bits enables it
                self.ram_enabled = (val & 0x0A) == 0x0A;
            },
            0x2000...0x3FFF => {
                // Writing to this space switches ROM bank lower bits
                self.bank_selection.set_lower(val);
            },
            0x4000...0x5FFF => {
                // Writing to this space switches the ROM bank upper bits/RAM bank
                self.bank_selection.set_upper(val);
            },
            0x6000...0x7FFF => {
                // Writing to this space toggles RAM/ROM bank mode. 1 = RAM mode
                self.bank_selection.set_mode(BankMode::from_u8(val & 1).unwrap_or_else(|| panic!("Unknown bank mode: {:#X}", val)));
            },
            0xA000...0xBFFF => {
                // Write to RAM Bank
                if self.ram_enabled {
                    let index = self.ram_index(addr);
                    self.ram_data[index] = val;
                }
            },
             _ => panic!("Address {:#X} not handled by MBC1")
        }
    }
}<|endoftext|>


<commit_before>// Copyright (c) Aptos
// SPDX-License-Identifier: Apache-2.0

use crate::{
    get_first_seq_num_and_limit,
    pruner::{
        ledger_pruner_manager::LedgerPrunerManager, state_pruner_manager::StatePrunerManager,
    },
    test_helper,
    test_helper::{arb_blocks_to_commit, put_as_state_root, put_transaction_info},
    AptosDB, PrunerManager, StaleNodeIndexSchema, ROCKSDB_PROPERTIES,
};
use aptos_config::config::{
    EpochSnapshotPrunerConfig, LedgerPrunerConfig, PrunerConfig, RocksdbConfigs,
    StateMerklePrunerConfig, DEFAULT_MAX_NUM_NODES_PER_LRU_CACHE_SHARD, TARGET_SNAPSHOT_SIZE,
};
use aptos_crypto::{hash::CryptoHash, HashValue};
use aptos_temppath::TempPath;
use aptos_types::ledger_info::LedgerInfoWithSignatures;
use aptos_types::state_store::state_storage_usage::StateStorageUsage;
use aptos_types::transaction::{TransactionToCommit, Version};
use aptos_types::{
    proof::SparseMerkleLeafNode,
    state_store::{state_key::StateKey, state_value::StateValue},
    transaction::{ExecutionStatus, TransactionInfo},
};
use proptest::prelude::*;
use std::collections::HashSet;
use std::{sync::Arc, time::Duration};
use storage_interface::{DbReader, DbWriter, ExecutedTrees, Order};
use test_helper::{test_save_blocks_impl, test_sync_transactions_impl};

proptest! {
    #![proptest_config(ProptestConfig::with_cases(10))]

    #[test]
    fn test_save_blocks(input in arb_blocks_to_commit(), threshold in 10..20usize) {
        test_save_blocks_impl(input, threshold);
    }

    #[test]
    fn test_sync_transactions(input in arb_blocks_to_commit(), threshold in 10..20usize) {
        test_sync_transactions_impl(input, threshold);
    }
}

#[test]
fn test_get_first_seq_num_and_limit() {
    assert!(get_first_seq_num_and_limit(Order::Ascending, 0, 0).is_err());

    // ascending
    assert_eq!(
        get_first_seq_num_and_limit(Order::Ascending, 0, 4).unwrap(),
        (0, 4)
    );
    assert_eq!(
        get_first_seq_num_and_limit(Order::Ascending, 0, 1).unwrap(),
        (0, 1)
    );

    // descending
    assert_eq!(
        get_first_seq_num_and_limit(Order::Descending, 2, 1).unwrap(),
        (2, 1)
    );
    assert_eq!(
        get_first_seq_num_and_limit(Order::Descending, 2, 2).unwrap(),
        (1, 2)
    );
    assert_eq!(
        get_first_seq_num_and_limit(Order::Descending, 2, 3).unwrap(),
        (0, 3)
    );
    assert_eq!(
        get_first_seq_num_and_limit(Order::Descending, 2, 4).unwrap(),
        (0, 3)
    );
}

#[test]
fn test_too_many_requested() {
    let tmp_dir = TempPath::new();
    let db = AptosDB::new_for_test(&tmp_dir);

    assert!(db.get_transactions(0, 1001 /* limit */, 0, true).is_err());
    assert!(db.get_transaction_outputs(0, 1001 /* limit */, 0).is_err());
}

#[test]
fn test_pruner_config() {
    let tmp_dir = TempPath::new();
    let aptos_db = AptosDB::new_for_test(&tmp_dir);
    for enable in [false, true] {
        let state_pruner = StatePrunerManager::<StaleNodeIndexSchema>::new(
            Arc::clone(&aptos_db.state_merkle_db),
            StateMerklePrunerConfig {
                enable,
                prune_window: 20,
                batch_size: 1,
            },
        );
        assert_eq!(state_pruner.is_pruner_enabled(), enable);
        assert_eq!(state_pruner.get_prune_window(), 20);

        let ledger_pruner = LedgerPrunerManager::new(
            Arc::clone(&aptos_db.ledger_db),
            Arc::clone(&aptos_db.state_store),
            LedgerPrunerConfig {
                enable,
                prune_window: 100,
                batch_size: 1,
                user_pruning_window_offset: 0,
            },
        );
        assert_eq!(ledger_pruner.is_pruner_enabled(), enable);
        assert_eq!(ledger_pruner.get_prune_window(), 100);
    }
}

#[test]
fn test_error_if_version_pruned() {
    let tmp_dir = TempPath::new();
    let db = AptosDB::new_for_test(&tmp_dir);
    db.state_store
        .state_db
        .state_pruner
        .testonly_update_min_version(5);
    db.ledger_pruner.testonly_update_min_version(10);
    assert_eq!(
        db.error_if_state_merkle_pruned("State", 4)
            .unwrap_err()
            .to_string(),
        "Version 4 is not epoch ending."
    );
    assert!(db.error_if_state_merkle_pruned("State", 5).is_ok());
    assert_eq!(
        db.error_if_ledger_pruned("Transaction", 9)
            .unwrap_err()
            .to_string(),
        "Transaction at version 9 is pruned, min available version is 10."
    );
    assert!(db.error_if_ledger_pruned("Transaction", 10).is_ok());
}

#[test]
fn test_get_latest_executed_trees() {
    let tmp_dir = TempPath::new();
    let db = AptosDB::new_for_test(&tmp_dir);

    // entirely emtpy db
    let empty = db.get_latest_executed_trees().unwrap();
    assert!(empty.is_same_view(&ExecutedTrees::new_empty()));

    // bootstrapped db (any transaction info is in)
    let key = StateKey::Raw(String::from("test_key").into_bytes());
    let value = StateValue::from(String::from("test_val").into_bytes());
    let hash = SparseMerkleLeafNode::new(key.hash(), value.hash()).hash();
    put_as_state_root(&db, 0, key, value);
    let txn_info = TransactionInfo::new(
        HashValue::random(),
        HashValue::random(),
        HashValue::random(),
        Some(hash),
        0,
        ExecutionStatus::MiscellaneousError(None),
    );
    put_transaction_info(&db, 0, &txn_info);

    let bootstrapped = db.get_latest_executed_trees().unwrap();
    assert!(
        bootstrapped.is_same_view(&ExecutedTrees::new_at_state_checkpoint(
            txn_info.state_checkpoint_hash().unwrap(),
            StateStorageUsage::new_untracked(),
            vec![txn_info.hash()],
            1,
        ))
    );
}

#[test]
fn test_rocksdb_properties_reporter() {
    fn get_metric() -> i64 {
        ROCKSDB_PROPERTIES
            .get_metric_with_label_values(&[
                "transaction_info",
                "aptos_rocksdb_is-file-deletions-enabled",
            ])
            .unwrap()
            .get()
    }

    assert_eq!(get_metric(), 0);
    let tmp_dir = TempPath::new();
    let _db = AptosDB::new_for_test(&tmp_dir);
    std::thread::sleep(Duration::from_secs(1));
    assert_eq!(get_metric(), 1);
}

pub fn test_state_merkle_pruning_impl(
    input: Vec<(Vec<TransactionToCommit>, LedgerInfoWithSignatures)>,
) {
    // set up DB with state prune window 5 and epoch ending state prune window 10
    let tmp_dir = TempPath::new();
    let db = AptosDB::open(
        &tmp_dir,
        false, /* is_read_only */
        PrunerConfig {
            ledger_pruner_config: LedgerPrunerConfig {
                enable: true,
                prune_window: 10,
                batch_size: 1,
                user_pruning_window_offset: 0,
            },
            state_merkle_pruner_config: StateMerklePrunerConfig {
                enable: true,
                prune_window: 5,
                batch_size: 1,
            },
            epoch_snapshot_pruner_config: EpochSnapshotPrunerConfig {
                enable: true,
                prune_window: 10,
                batch_size: 1,
            },
        },
        RocksdbConfigs::default(),
        false, /* enable_indexer */
        TARGET_SNAPSHOT_SIZE,
        DEFAULT_MAX_NUM_NODES_PER_LRU_CACHE_SHARD,
    )
    .unwrap();

    // augment DB in blocks
    let mut in_memory_state = db
        .state_store
        .buffered_state()
        .lock()
        .current_state()
        .clone();
    let _ancester = in_memory_state.current.clone();
    let mut next_ver: Version = 0;
    let mut snapshot_versions = vec![];
    for (txns_to_commit, ledger_info_with_sigs) in input.iter() {
        test_helper::update_in_memory_state(&mut in_memory_state, txns_to_commit.as_slice());
        db.save_transactions(
            txns_to_commit,
            next_ver,                /* first_version */
            next_ver.checked_sub(1), /* base_state_version */
            Some(ledger_info_with_sigs),
            true, /* sync_commit */
            in_memory_state.clone(),
        )
        .unwrap();

        next_ver += txns_to_commit.len() as u64;

        let last_version = next_ver - 1;
        let is_epoch_ending = ledger_info_with_sigs.ledger_info().ends_epoch();
        snapshot_versions.push((last_version, is_epoch_ending));

        let state_merkle_min_readable = last_version.saturating_sub(5);
        let epoch_snapshot_min_readable = last_version.saturating_sub(10);
        let snapshots: Vec<_> = snapshot_versions
            .iter()
            .filter(|(v, _is_epoch_ending)| *v >= state_merkle_min_readable)
            .map(|(v, _)| *v)
            .collect();
        let epoch_snapshots: Vec<_> = snapshot_versions
            .iter()
            .filter(|(v, is_epoch_ending)| *is_epoch_ending && *v >= epoch_snapshot_min_readable)
            .map(|(v, _)| *v)
            .collect();

        // Prune till the oldest snapshot readable.
        let pruner = &db.state_store.state_db.state_pruner;
        let epoch_snapshot_pruner = &db.state_store.state_db.epoch_snapshot_pruner;
        pruner
            .pruner_worker
            .set_target_db_version(*snapshots.first().unwrap());
        epoch_snapshot_pruner
            .pruner_worker
            .set_target_db_version(std::cmp::min(
                *snapshots.first().unwrap(),
                *epoch_snapshots.first().unwrap_or(&Version::MAX),
            ));
        pruner.wait_for_pruner().unwrap();
        epoch_snapshot_pruner.wait_for_pruner().unwrap();

        // Check strictly that all trees in the window accessible and all those nodes not needed
        // must be gone.
        let non_pruned_versions: HashSet<_> = snapshots
            .into_iter()
            .chain(epoch_snapshots.into_iter())
            .collect();

        let expected_nodes: HashSet<_> = non_pruned_versions
            .iter()
            .flat_map(|v| db.state_store.get_all_jmt_nodes_referenced(*v).unwrap())
            .collect();
        let all_nodes: HashSet<_> = db
            .state_store
            .get_all_jmt_nodes()
            .unwrap()
            .into_iter()
            .collect();

        assert_eq!(expected_nodes, all_nodes);
    }
}

proptest! {
    #![proptest_config(ProptestConfig::with_cases(10))]

    #[test]
    fn test_state_merkle_pruning(input in arb_blocks_to_commit()) {
        test_state_merkle_pruning_impl(input);
    }
}
<commit_msg>[storage] remove flaky test<commit_after>// Copyright (c) Aptos
// SPDX-License-Identifier: Apache-2.0

use crate::{
    get_first_seq_num_and_limit,
    pruner::{
        ledger_pruner_manager::LedgerPrunerManager, state_pruner_manager::StatePrunerManager,
    },
    test_helper,
    test_helper::{arb_blocks_to_commit, put_as_state_root, put_transaction_info},
    AptosDB, PrunerManager, StaleNodeIndexSchema,
};
use aptos_config::config::{
    EpochSnapshotPrunerConfig, LedgerPrunerConfig, PrunerConfig, RocksdbConfigs,
    StateMerklePrunerConfig, DEFAULT_MAX_NUM_NODES_PER_LRU_CACHE_SHARD, TARGET_SNAPSHOT_SIZE,
};
use aptos_crypto::{hash::CryptoHash, HashValue};
use aptos_temppath::TempPath;
use aptos_types::ledger_info::LedgerInfoWithSignatures;
use aptos_types::state_store::state_storage_usage::StateStorageUsage;
use aptos_types::transaction::{TransactionToCommit, Version};
use aptos_types::{
    proof::SparseMerkleLeafNode,
    state_store::{state_key::StateKey, state_value::StateValue},
    transaction::{ExecutionStatus, TransactionInfo},
};
use proptest::prelude::*;
use std::collections::HashSet;
use std::sync::Arc;
use storage_interface::{DbReader, DbWriter, ExecutedTrees, Order};
use test_helper::{test_save_blocks_impl, test_sync_transactions_impl};

proptest! {
    #![proptest_config(ProptestConfig::with_cases(10))]

    #[test]
    fn test_save_blocks(input in arb_blocks_to_commit(), threshold in 10..20usize) {
        test_save_blocks_impl(input, threshold);
    }

    #[test]
    fn test_sync_transactions(input in arb_blocks_to_commit(), threshold in 10..20usize) {
        test_sync_transactions_impl(input, threshold);
    }
}

#[test]
fn test_get_first_seq_num_and_limit() {
    assert!(get_first_seq_num_and_limit(Order::Ascending, 0, 0).is_err());

    // ascending
    assert_eq!(
        get_first_seq_num_and_limit(Order::Ascending, 0, 4).unwrap(),
        (0, 4)
    );
    assert_eq!(
        get_first_seq_num_and_limit(Order::Ascending, 0, 1).unwrap(),
        (0, 1)
    );

    // descending
    assert_eq!(
        get_first_seq_num_and_limit(Order::Descending, 2, 1).unwrap(),
        (2, 1)
    );
    assert_eq!(
        get_first_seq_num_and_limit(Order::Descending, 2, 2).unwrap(),
        (1, 2)
    );
    assert_eq!(
        get_first_seq_num_and_limit(Order::Descending, 2, 3).unwrap(),
        (0, 3)
    );
    assert_eq!(
        get_first_seq_num_and_limit(Order::Descending, 2, 4).unwrap(),
        (0, 3)
    );
}

#[test]
fn test_too_many_requested() {
    let tmp_dir = TempPath::new();
    let db = AptosDB::new_for_test(&tmp_dir);

    assert!(db.get_transactions(0, 1001 /* limit */, 0, true).is_err());
    assert!(db.get_transaction_outputs(0, 1001 /* limit */, 0).is_err());
}

#[test]
fn test_pruner_config() {
    let tmp_dir = TempPath::new();
    let aptos_db = AptosDB::new_for_test(&tmp_dir);
    for enable in [false, true] {
        let state_pruner = StatePrunerManager::<StaleNodeIndexSchema>::new(
            Arc::clone(&aptos_db.state_merkle_db),
            StateMerklePrunerConfig {
                enable,
                prune_window: 20,
                batch_size: 1,
            },
        );
        assert_eq!(state_pruner.is_pruner_enabled(), enable);
        assert_eq!(state_pruner.get_prune_window(), 20);

        let ledger_pruner = LedgerPrunerManager::new(
            Arc::clone(&aptos_db.ledger_db),
            Arc::clone(&aptos_db.state_store),
            LedgerPrunerConfig {
                enable,
                prune_window: 100,
                batch_size: 1,
                user_pruning_window_offset: 0,
            },
        );
        assert_eq!(ledger_pruner.is_pruner_enabled(), enable);
        assert_eq!(ledger_pruner.get_prune_window(), 100);
    }
}

#[test]
fn test_error_if_version_pruned() {
    let tmp_dir = TempPath::new();
    let db = AptosDB::new_for_test(&tmp_dir);
    db.state_store
        .state_db
        .state_pruner
        .testonly_update_min_version(5);
    db.ledger_pruner.testonly_update_min_version(10);
    assert_eq!(
        db.error_if_state_merkle_pruned("State", 4)
            .unwrap_err()
            .to_string(),
        "Version 4 is not epoch ending."
    );
    assert!(db.error_if_state_merkle_pruned("State", 5).is_ok());
    assert_eq!(
        db.error_if_ledger_pruned("Transaction", 9)
            .unwrap_err()
            .to_string(),
        "Transaction at version 9 is pruned, min available version is 10."
    );
    assert!(db.error_if_ledger_pruned("Transaction", 10).is_ok());
}

#[test]
fn test_get_latest_executed_trees() {
    let tmp_dir = TempPath::new();
    let db = AptosDB::new_for_test(&tmp_dir);

    // entirely emtpy db
    let empty = db.get_latest_executed_trees().unwrap();
    assert!(empty.is_same_view(&ExecutedTrees::new_empty()));

    // bootstrapped db (any transaction info is in)
    let key = StateKey::Raw(String::from("test_key").into_bytes());
    let value = StateValue::from(String::from("test_val").into_bytes());
    let hash = SparseMerkleLeafNode::new(key.hash(), value.hash()).hash();
    put_as_state_root(&db, 0, key, value);
    let txn_info = TransactionInfo::new(
        HashValue::random(),
        HashValue::random(),
        HashValue::random(),
        Some(hash),
        0,
        ExecutionStatus::MiscellaneousError(None),
    );
    put_transaction_info(&db, 0, &txn_info);

    let bootstrapped = db.get_latest_executed_trees().unwrap();
    assert!(
        bootstrapped.is_same_view(&ExecutedTrees::new_at_state_checkpoint(
            txn_info.state_checkpoint_hash().unwrap(),
            StateStorageUsage::new_untracked(),
            vec![txn_info.hash()],
            1,
        ))
    );
}

pub fn test_state_merkle_pruning_impl(
    input: Vec<(Vec<TransactionToCommit>, LedgerInfoWithSignatures)>,
) {
    // set up DB with state prune window 5 and epoch ending state prune window 10
    let tmp_dir = TempPath::new();
    let db = AptosDB::open(
        &tmp_dir,
        false, /* is_read_only */
        PrunerConfig {
            ledger_pruner_config: LedgerPrunerConfig {
                enable: true,
                prune_window: 10,
                batch_size: 1,
                user_pruning_window_offset: 0,
            },
            state_merkle_pruner_config: StateMerklePrunerConfig {
                enable: true,
                prune_window: 5,
                batch_size: 1,
            },
            epoch_snapshot_pruner_config: EpochSnapshotPrunerConfig {
                enable: true,
                prune_window: 10,
                batch_size: 1,
            },
        },
        RocksdbConfigs::default(),
        false, /* enable_indexer */
        TARGET_SNAPSHOT_SIZE,
        DEFAULT_MAX_NUM_NODES_PER_LRU_CACHE_SHARD,
    )
    .unwrap();

    // augment DB in blocks
    let mut in_memory_state = db
        .state_store
        .buffered_state()
        .lock()
        .current_state()
        .clone();
    let _ancester = in_memory_state.current.clone();
    let mut next_ver: Version = 0;
    let mut snapshot_versions = vec![];
    for (txns_to_commit, ledger_info_with_sigs) in input.iter() {
        test_helper::update_in_memory_state(&mut in_memory_state, txns_to_commit.as_slice());
        db.save_transactions(
            txns_to_commit,
            next_ver,                /* first_version */
            next_ver.checked_sub(1), /* base_state_version */
            Some(ledger_info_with_sigs),
            true, /* sync_commit */
            in_memory_state.clone(),
        )
        .unwrap();

        next_ver += txns_to_commit.len() as u64;

        let last_version = next_ver - 1;
        let is_epoch_ending = ledger_info_with_sigs.ledger_info().ends_epoch();
        snapshot_versions.push((last_version, is_epoch_ending));

        let state_merkle_min_readable = last_version.saturating_sub(5);
        let epoch_snapshot_min_readable = last_version.saturating_sub(10);
        let snapshots: Vec<_> = snapshot_versions
            .iter()
            .filter(|(v, _is_epoch_ending)| *v >= state_merkle_min_readable)
            .map(|(v, _)| *v)
            .collect();
        let epoch_snapshots: Vec<_> = snapshot_versions
            .iter()
            .filter(|(v, is_epoch_ending)| *is_epoch_ending && *v >= epoch_snapshot_min_readable)
            .map(|(v, _)| *v)
            .collect();

        // Prune till the oldest snapshot readable.
        let pruner = &db.state_store.state_db.state_pruner;
        let epoch_snapshot_pruner = &db.state_store.state_db.epoch_snapshot_pruner;
        pruner
            .pruner_worker
            .set_target_db_version(*snapshots.first().unwrap());
        epoch_snapshot_pruner
            .pruner_worker
            .set_target_db_version(std::cmp::min(
                *snapshots.first().unwrap(),
                *epoch_snapshots.first().unwrap_or(&Version::MAX),
            ));
        pruner.wait_for_pruner().unwrap();
        epoch_snapshot_pruner.wait_for_pruner().unwrap();

        // Check strictly that all trees in the window accessible and all those nodes not needed
        // must be gone.
        let non_pruned_versions: HashSet<_> = snapshots
            .into_iter()
            .chain(epoch_snapshots.into_iter())
            .collect();

        let expected_nodes: HashSet<_> = non_pruned_versions
            .iter()
            .flat_map(|v| db.state_store.get_all_jmt_nodes_referenced(*v).unwrap())
            .collect();
        let all_nodes: HashSet<_> = db
            .state_store
            .get_all_jmt_nodes()
            .unwrap()
            .into_iter()
            .collect();

        assert_eq!(expected_nodes, all_nodes);
    }
}

proptest! {
    #![proptest_config(ProptestConfig::with_cases(10))]

    #[test]
    fn test_state_merkle_pruning(input in arb_blocks_to_commit()) {
        test_state_merkle_pruning_impl(input);
    }
}
<|endoftext|>