<commit_before>use std;
use std::path::Path;

use sdl2::event::Event;
use sdl2::gfx::primitives::DrawRenderer;
use sdl2::keyboard::Keycode;
use sdl2::pixels::{Color, PixelFormatEnum};
use sdl2::rect::Rect;
use sdl2::render::{BlendMode, Renderer, Texture, TextureQuery};
use sdl2::surface::Surface;
use sdl2::ttf::{Font, Sdl2TtfContext, STYLE_BOLD};

use position::Position;
use text::Text;

const BORDER_COLOR: Color = Color::RGBA(255, 255, 255, 64);

const PADDING: i32 = 10;

const FONT_FILE: &'static str = "assets/FantasqueSansMono-Bold.ttf";
const FONT_COLOR: Color = Color::RGBA(45, 200, 45, 255);
const FONT_SIZE: u16 = 18;

pub struct Console<'a> {
    pub visible: bool,

    leader: Text,
    input_buffer: String,
    last_command: String,
    command_history: Vec<String>,
    history_position: usize,
    cursor_position: usize,
    buffer: Vec<String>,
    backbuffer_y: i32,
    texture: Texture,
    ttf_context: &'a Sdl2TtfContext,
    size: (u32, u32),
    font: Font<'a>,
    line_ending: bool, // Tracks where the next print call should append to
    ctrl: bool, // Tracks the Ctrl key being pressed
}

impl<'a> Console<'a> {
    /// Creates a new empty Console
    pub fn new(ttf_context: &'a Sdl2TtfContext, mut renderer: &mut Renderer) -> Console<'a> {
        let (width, height) = renderer.window().unwrap().size();
        let mut texture =
            renderer.create_texture_streaming(PixelFormatEnum::RGBA8888, width / 2, height)
                .unwrap();
        texture.with_lock(None, |buffer: &mut [u8], pitch: usize| {
                for y in 0..height {
                    for x in 0..width / 2 {
                        let x = x as usize;
                        let y = y as usize;
                        let offset = y * pitch + x * 4;
                        buffer[offset] = 182;
                        buffer[offset + 1] = 0;
                        buffer[offset + 2] = 0;
                        buffer[offset + 3] = 0;
                    }
                }
            })
            .unwrap();

        let mut font = ttf_context.load_font(Path::new(FONT_FILE), FONT_SIZE).unwrap();
        font.set_style(STYLE_BOLD);

        Console {
            visible: false,
            leader: Text::new(ttf_context,
                              &mut renderer,
                              "hakka>",
                              Position::XY(PADDING, height as i32 - FONT_SIZE as i32 - PADDING),
                              FONT_SIZE,
                              FONT_COLOR,
                              FONT_FILE),
            input_buffer: "".into(),
            last_command: "".into(),
            command_history: Vec::new(),
            history_position: 0,
            cursor_position: 0,
            buffer: Vec::new(),
            backbuffer_y: 0,
            texture: texture,
            ttf_context: ttf_context,
            size: (width / 2, height),
            font: font,
            line_ending: true,
            ctrl: false,
        }
    }

    pub fn process(&mut self, event: &Event) {
        match *event {
            Event::TextInput { ref text, .. } => {
                if self.visible {
                    if text == "`" || text == "\\" {
                        return;
                    }
                    self.add_text(text);
                }
            }
            Event::MouseWheel { y, .. } => {
                if self.visible {
                    if self.buffer.len() * FONT_SIZE as usize >
                       (self.size.1 - (FONT_SIZE as u32 * 2)) as usize {
                        self.backbuffer_y += y * 6;
                        if self.backbuffer_y < 0 {
                            self.backbuffer_y = 0;
                        }
                    }
                }
            }
            Event::KeyDown { keycode, .. } => {
                if self.visible {
                    match keycode {
                        Some(Keycode::LCtrl) |
                        Some(Keycode::RCtrl) => self.ctrl = true,
                        Some(Keycode::C) => {
                            if self.ctrl {
                                self.input_buffer.push_str("^C");
                                self.commit(false);
                            }
                        }
                        Some(Keycode::Left) => {
                            self.cursor_left();
                        }
                        Some(Keycode::Right) => {
                            self.cursor_right();
                        }
                        Some(Keycode::Backspace) => {
                            self.backspace();
                        }
                        Some(Keycode::Delete) => {
                            if self.cursor_position < self.input_buffer.len() {
                                self.cursor_position += 1;
                                self.backspace();
                            }
                        }
                        _ => (),
                    }
                }
            }
            Event::KeyUp { keycode, .. } => {
                if self.visible {
                    match keycode {
                        Some(Keycode::LCtrl) |
                        Some(Keycode::RCtrl) => self.ctrl = false,
                        Some(Keycode::Up) => {
                            self.history_navigate_back();
                        }
                        Some(Keycode::Down) => {
                            self.history_navigate_forward();
                        }
                        Some(Keycode::Return) => {
                            self.commit(true);
                        }
                        Some(Keycode::End) => {
                            self.cursor_position = self.input_buffer.len();
                        }
                        Some(Keycode::Home) => {
                            self.cursor_position = 0;
                        },
                        // Toggle console
                        Some(Keycode::Backquote) |
                        Some(Keycode::Backslash) => {
                            self.toggle();
                        }
                        _ => (),
                    }
                }
            }
            _ => (),
        }
    }

    pub fn process_command(&mut self) {
        let command = self.input_buffer.clone();
        if !command.is_empty() {
            self.command_history.push(command.clone());
            self.last_command = command.clone();

            if command == "exit" {
                std::process::exit(0);
            }
        }
    }

    pub fn clear(&mut self) {
        self.buffer.clear();
    }

    fn history_navigate_back(&mut self) {
        if self.history_position > 0 {
            self.input_buffer = self.command_history[self.history_position - 1].clone();
            self.cursor_position = self.input_buffer.len();

            if self.history_position > 0 {
                self.history_position -= 1;
            }
        }
    }

    fn history_navigate_forward(&mut self) {
        if self.history_position < self.command_history.len() - 1 {
            self.input_buffer = self.command_history[self.history_position + 1].clone();
            self.cursor_position = self.input_buffer.len();
            if self.history_position < self.command_history.len() {
                self.history_position += 1;
            }
        }
    }

    pub fn try_process_command(&mut self) -> Option<String> {
        if !self.last_command.is_empty() {
            let cmd = self.last_command.clone();
            self.last_command.clear();
            Some(cmd)
        } else {
            None
        }
    }

    pub fn print<S>(&mut self, text: S)
        where S: Into<String>
    {
        if !self.line_ending {
            let last = self.buffer.last_mut().unwrap();
            last.push_str(&text.into());
        } else {
            self.buffer.push(text.into());
        }
        self.line_ending = false;
    }

    pub fn println<S>(&mut self, text: S)
        where S: Into<String>
    {
        self.buffer.push(text.into());
        self.line_ending = true;
    }

    pub fn print_lines<S>(&mut self, text: S)
        where S: Into<String>
    {
        for line in text.into().lines() {
            self.println(line);
        }
    }

    pub fn wrap_line(&mut self) {
        self.buffer.push("".into());
        self.line_ending = false;
    }

    /// Toggles the visibility of the Console
    pub fn toggle(&mut self) {
        self.visible = !self.visible;
    }

    pub fn add_text(&mut self, input: &str) {
        self.input_buffer.insert(self.cursor_position, input.chars().next().unwrap());
        self.cursor_position += 1;
    }

    pub fn commit(&mut self, execute: bool) {
        self.buffer.push(format!("hakka> {}", self.input_buffer.clone()));
        if execute {
            self.process_command();
        }
        self.input_buffer.clear();
        self.cursor_position = 0;
        self.history_position = self.command_history.len();
    }

    pub fn cursor_left(&mut self) {
        if self.cursor_position > 0 {
            self.cursor_position -= 1;
        }
    }

    pub fn cursor_right(&mut self) {
        if self.cursor_position < self.input_buffer.len() {
            self.cursor_position += 1;
        }
    }

    pub fn backspace(&mut self) {
        if self.visible && self.cursor_position > 0 {
            self.input_buffer.remove(self.cursor_position - 1);
            self.cursor_position -= 1;
        }
    }

    /// Renders the Console
    pub fn render(&mut self, mut renderer: &mut Renderer) {
        if self.visible {

            renderer.set_blend_mode(BlendMode::Blend);
            self.texture.set_blend_mode(BlendMode::Blend);
            renderer.copy(&self.texture,
                      None,
                      Some(Rect::new(0, 0, self.size.0, self.size.1)))
                .unwrap();
            self.generate_backbuffer_texture(&mut renderer);
            self.render_leader(&mut renderer);

            // Insert the cursor via a dodgy vertical line
            let cursor_x =
                60 + PADDING as i16 +
                self.font.size_of(&self.input_buffer[..self.cursor_position]).unwrap().0 as i16;
            // Draw a dodgy cursor
            renderer.thick_line(cursor_x,
                            self.size.1 as i16 - FONT_SIZE as i16 - PADDING as i16,
                            cursor_x,
                            self.size.1 as i16 - PADDING as i16,
                            1,
                            FONT_COLOR)
                .unwrap();

            if !self.input_buffer.is_empty() {
                let text = Text::new(self.ttf_context,
                                     &mut renderer,
                                     &self.input_buffer[..],
                                     Position::XY(60 + PADDING,
                                                  self.size.1 as i32 - FONT_SIZE as i32 - PADDING),
                                     FONT_SIZE,
                                     FONT_COLOR,
                                     FONT_FILE);
                text.render(&mut renderer);
            }

            self.render_border(&mut renderer);
        }
    }

    fn render_border(&self, mut renderer: &mut Renderer) {
        // Render the border
        renderer.set_draw_color(Color::RGBA(255, 255, 255, 255));
        // North
        renderer.thick_line(0, 0, self.size.0 as i16, 0, 1, BORDER_COLOR).unwrap();

        // East
        renderer.thick_line(self.size.0 as i16,
                        0,
                        self.size.0 as i16,
                        self.size.1 as i16,
                        1,
                        BORDER_COLOR)
            .unwrap();

        // South
        renderer.thick_line(0,
                        self.size.1 as i16 - 1,
                        self.size.0 as i16,
                        self.size.1 as i16 - 1,
                        1,
                        BORDER_COLOR)
            .unwrap();
    }

    fn render_leader(&self, mut renderer: &mut Renderer) {
        // Render a black background behind it so the buffer scrolling looks
        // nicer.
        let rect_y = self.size.1 as i32 - FONT_SIZE as i32 - PADDING;
        renderer.set_draw_color(Color::RGBA(0, 0, 0, 255));
        renderer.fill_rect(Rect::new(0, rect_y, self.size.0, rect_y as u32)).unwrap();
        self.leader.render(&mut renderer);
    }

    fn generate_backbuffer_texture(&mut self, mut renderer: &mut Renderer) {
        let mut main_surface = Surface::new(self.size.0,
                                            (self.size.1 - (FONT_SIZE as u32)),
                                            PixelFormatEnum::RGBA8888)
            .unwrap();
        let mut counter = 2;
        for line in self.buffer.iter().rev().take(100) {
            let y_pos = self.size.1 as i32 - (FONT_SIZE as i32 * counter) + self.backbuffer_y;
            counter += 1;

            if line.trim().is_empty() {
                continue;
            }

            let surface = self.font
                .render(line)
                .blended(FONT_COLOR)
                .unwrap();
            surface.blit(None,
                      &mut main_surface,
                      Some(Rect::new(PADDING, y_pos - PADDING, self.size.1, FONT_SIZE as u32)))
                .unwrap();
        }
        let texture = renderer.create_texture_from_surface(&main_surface)
            .unwrap();

        let TextureQuery { width, height, .. } = texture.query();

        renderer.copy(&texture, None, Some(Rect::new(0, 0, width, height))).unwrap();
    }
}
<commit_msg>Check that shift isn't being held down when toggling the console<commit_after>use std;
use std::path::Path;

use sdl2::event::Event;
use sdl2::gfx::primitives::DrawRenderer;
use sdl2::keyboard::Keycode;
use sdl2::pixels::{Color, PixelFormatEnum};
use sdl2::rect::Rect;
use sdl2::render::{BlendMode, Renderer, Texture, TextureQuery};
use sdl2::surface::Surface;
use sdl2::ttf::{Font, Sdl2TtfContext, STYLE_BOLD};

use position::Position;
use text::Text;

const BORDER_COLOR: Color = Color::RGBA(255, 255, 255, 64);

const PADDING: i32 = 10;

const FONT_FILE: &'static str = "assets/FantasqueSansMono-Bold.ttf";
const FONT_COLOR: Color = Color::RGBA(45, 200, 45, 255);
const FONT_SIZE: u16 = 18;

pub struct Console<'a> {
    pub visible: bool,

    leader: Text,
    input_buffer: String,
    last_command: String,
    command_history: Vec<String>,
    history_position: usize,
    cursor_position: usize,
    buffer: Vec<String>,
    backbuffer_y: i32,
    texture: Texture,
    ttf_context: &'a Sdl2TtfContext,
    size: (u32, u32),
    font: Font<'a>,
    line_ending: bool, // Tracks where the next print call should append to
    ctrl: bool, // Tracks the Ctrl key being pressed
    shift: bool, // Tracks the Shift key being pressed
}

impl<'a> Console<'a> {
    /// Creates a new empty Console
    pub fn new(ttf_context: &'a Sdl2TtfContext, mut renderer: &mut Renderer) -> Console<'a> {
        let (width, height) = renderer.window().unwrap().size();
        let mut texture =
            renderer.create_texture_streaming(PixelFormatEnum::RGBA8888, width / 2, height)
                .unwrap();
        texture.with_lock(None, |buffer: &mut [u8], pitch: usize| {
                for y in 0..height {
                    for x in 0..width / 2 {
                        let x = x as usize;
                        let y = y as usize;
                        let offset = y * pitch + x * 4;
                        buffer[offset] = 182;
                        buffer[offset + 1] = 0;
                        buffer[offset + 2] = 0;
                        buffer[offset + 3] = 0;
                    }
                }
            })
            .unwrap();

        let mut font = ttf_context.load_font(Path::new(FONT_FILE), FONT_SIZE).unwrap();
        font.set_style(STYLE_BOLD);

        Console {
            visible: false,
            leader: Text::new(ttf_context,
                              &mut renderer,
                              "hakka>",
                              Position::XY(PADDING, height as i32 - FONT_SIZE as i32 - PADDING),
                              FONT_SIZE,
                              FONT_COLOR,
                              FONT_FILE),
            input_buffer: "".into(),
            last_command: "".into(),
            command_history: Vec::new(),
            history_position: 0,
            cursor_position: 0,
            buffer: Vec::new(),
            backbuffer_y: 0,
            texture: texture,
            ttf_context: ttf_context,
            size: (width / 2, height),
            font: font,
            line_ending: true,
            ctrl: false,
            shift: false,
        }
    }

    pub fn process(&mut self, event: &Event) {
        match *event {
            Event::TextInput { ref text, .. } => {
                if self.visible {
                    if text == "`" || text == "\\" {
                        return;
                    }
                    self.add_text(text);
                }
            }
            Event::MouseWheel { y, .. } => {
                if self.visible {
                    if self.buffer.len() * FONT_SIZE as usize >
                       (self.size.1 - (FONT_SIZE as u32 * 2)) as usize {
                        self.backbuffer_y += y * 6;
                        if self.backbuffer_y < 0 {
                            self.backbuffer_y = 0;
                        }
                    }
                }
            }
            Event::KeyDown { keycode, .. } => {
                if self.visible {
                    match keycode {
                        Some(Keycode::LCtrl) |
                        Some(Keycode::RCtrl) => self.ctrl = true,
                        Some(Keycode::LShift) |
                        Some(Keycode::RShift) => self.shift = true,
                        Some(Keycode::C) => {
                            if self.ctrl {
                                self.input_buffer.push_str("^C");
                                self.commit(false);
                            }
                        }
                        Some(Keycode::Left) => {
                            self.cursor_left();
                        }
                        Some(Keycode::Right) => {
                            self.cursor_right();
                        }
                        Some(Keycode::Backspace) => {
                            self.backspace();
                        }
                        Some(Keycode::Delete) => {
                            if self.cursor_position < self.input_buffer.len() {
                                self.cursor_position += 1;
                                self.backspace();
                            }
                        }
                        _ => (),
                    }
                }
            }
            Event::KeyUp { keycode, .. } => {
                if self.visible {
                    match keycode {
                        Some(Keycode::LCtrl) |
                        Some(Keycode::RCtrl) => self.ctrl = false,
                        Some(Keycode::LShift) |
                        Some(Keycode::RShift) => self.shift = false,
                        Some(Keycode::Up) => {
                            self.history_navigate_back();
                        }
                        Some(Keycode::Down) => {
                            self.history_navigate_forward();
                        }
                        Some(Keycode::Return) => {
                            self.commit(true);
                        }
                        Some(Keycode::End) => {
                            self.cursor_position = self.input_buffer.len();
                        }
                        Some(Keycode::Home) => {
                            self.cursor_position = 0;
                        }
                        // Toggle console
                        Some(Keycode::Backquote) |
                        Some(Keycode::Backslash) => {
                            if !self.shift {
                                self.toggle();
                            }
                        }
                        _ => (),
                    }
                }
            }
            _ => (),
        }
    }

    pub fn process_command(&mut self) {
        let command = self.input_buffer.clone();
        if !command.is_empty() {
            self.command_history.push(command.clone());
            self.last_command = command.clone();

            if command == "exit" {
                std::process::exit(0);
            }
        }
    }

    pub fn clear(&mut self) {
        self.buffer.clear();
    }

    fn history_navigate_back(&mut self) {
        if self.history_position > 0 {
            self.input_buffer = self.command_history[self.history_position - 1].clone();
            self.cursor_position = self.input_buffer.len();

            if self.history_position > 0 {
                self.history_position -= 1;
            }
        }
    }

    fn history_navigate_forward(&mut self) {
        if self.history_position < self.command_history.len() - 1 {
            self.input_buffer = self.command_history[self.history_position + 1].clone();
            self.cursor_position = self.input_buffer.len();
            if self.history_position < self.command_history.len() {
                self.history_position += 1;
            }
        }
    }

    pub fn try_process_command(&mut self) -> Option<String> {
        if !self.last_command.is_empty() {
            let cmd = self.last_command.clone();
            self.last_command.clear();
            Some(cmd)
        } else {
            None
        }
    }

    pub fn print<S>(&mut self, text: S)
        where S: Into<String>
    {
        if !self.line_ending {
            let last = self.buffer.last_mut().unwrap();
            last.push_str(&text.into());
        } else {
            self.buffer.push(text.into());
        }
        self.line_ending = false;
    }

    pub fn println<S>(&mut self, text: S)
        where S: Into<String>
    {
        self.buffer.push(text.into());
        self.line_ending = true;
    }

    pub fn print_lines<S>(&mut self, text: S)
        where S: Into<String>
    {
        for line in text.into().lines() {
            self.println(line);
        }
    }

    pub fn wrap_line(&mut self) {
        self.buffer.push("".into());
        self.line_ending = false;
    }

    /// Toggles the visibility of the Console
    pub fn toggle(&mut self) {
        self.visible = !self.visible;
    }

    pub fn add_text(&mut self, input: &str) {
        self.input_buffer.insert(self.cursor_position, input.chars().next().unwrap());
        self.cursor_position += 1;
    }

    pub fn commit(&mut self, execute: bool) {
        self.buffer.push(format!("hakka> {}", self.input_buffer.clone()));
        if execute {
            self.process_command();
        }
        self.input_buffer.clear();
        self.cursor_position = 0;
        self.history_position = self.command_history.len();
    }

    pub fn cursor_left(&mut self) {
        if self.cursor_position > 0 {
            self.cursor_position -= 1;
        }
    }

    pub fn cursor_right(&mut self) {
        if self.cursor_position < self.input_buffer.len() {
            self.cursor_position += 1;
        }
    }

    pub fn backspace(&mut self) {
        if self.visible && self.cursor_position > 0 {
            self.input_buffer.remove(self.cursor_position - 1);
            self.cursor_position -= 1;
        }
    }

    /// Renders the Console
    pub fn render(&mut self, mut renderer: &mut Renderer) {
        if self.visible {

            renderer.set_blend_mode(BlendMode::Blend);
            self.texture.set_blend_mode(BlendMode::Blend);
            renderer.copy(&self.texture,
                      None,
                      Some(Rect::new(0, 0, self.size.0, self.size.1)))
                .unwrap();
            self.generate_backbuffer_texture(&mut renderer);
            self.render_leader(&mut renderer);

            // Insert the cursor via a dodgy vertical line
            let cursor_x =
                60 + PADDING as i16 +
                self.font.size_of(&self.input_buffer[..self.cursor_position]).unwrap().0 as i16;
            // Draw a dodgy cursor
            renderer.thick_line(cursor_x,
                            self.size.1 as i16 - FONT_SIZE as i16 - PADDING as i16,
                            cursor_x,
                            self.size.1 as i16 - PADDING as i16,
                            1,
                            FONT_COLOR)
                .unwrap();

            if !self.input_buffer.is_empty() {
                let text = Text::new(self.ttf_context,
                                     &mut renderer,
                                     &self.input_buffer[..],
                                     Position::XY(60 + PADDING,
                                                  self.size.1 as i32 - FONT_SIZE as i32 - PADDING),
                                     FONT_SIZE,
                                     FONT_COLOR,
                                     FONT_FILE);
                text.render(&mut renderer);
            }

            self.render_border(&mut renderer);
        }
    }

    fn render_border(&self, mut renderer: &mut Renderer) {
        // Render the border
        renderer.set_draw_color(Color::RGBA(255, 255, 255, 255));
        // North
        renderer.thick_line(0, 0, self.size.0 as i16, 0, 1, BORDER_COLOR).unwrap();

        // East
        renderer.thick_line(self.size.0 as i16,
                        0,
                        self.size.0 as i16,
                        self.size.1 as i16,
                        1,
                        BORDER_COLOR)
            .unwrap();

        // South
        renderer.thick_line(0,
                        self.size.1 as i16 - 1,
                        self.size.0 as i16,
                        self.size.1 as i16 - 1,
                        1,
                        BORDER_COLOR)
            .unwrap();
    }

    fn render_leader(&self, mut renderer: &mut Renderer) {
        // Render a black background behind it so the buffer scrolling looks
        // nicer.
        let rect_y = self.size.1 as i32 - FONT_SIZE as i32 - PADDING;
        renderer.set_draw_color(Color::RGBA(0, 0, 0, 255));
        renderer.fill_rect(Rect::new(0, rect_y, self.size.0, rect_y as u32)).unwrap();
        self.leader.render(&mut renderer);
    }

    fn generate_backbuffer_texture(&mut self, mut renderer: &mut Renderer) {
        let mut main_surface = Surface::new(self.size.0,
                                            (self.size.1 - (FONT_SIZE as u32)),
                                            PixelFormatEnum::RGBA8888)
            .unwrap();
        let mut counter = 2;
        for line in self.buffer.iter().rev().take(100) {
            let y_pos = self.size.1 as i32 - (FONT_SIZE as i32 * counter) + self.backbuffer_y;
            counter += 1;

            if line.trim().is_empty() {
                continue;
            }

            let surface = self.font
                .render(line)
                .blended(FONT_COLOR)
                .unwrap();
            surface.blit(None,
                      &mut main_surface,
                      Some(Rect::new(PADDING, y_pos - PADDING, self.size.1, FONT_SIZE as u32)))
                .unwrap();
        }
        let texture = renderer.create_texture_from_surface(&main_surface)
            .unwrap();

        let TextureQuery { width, height, .. } = texture.query();

        renderer.copy(&texture, None, Some(Rect::new(0, 0, width, height))).unwrap();
    }
}
<|endoftext|>


<commit_before>use std::{
    cell::{Cell, RefCell},
    collections::BTreeMap,
};

use dces::prelude::Entity;

use crate::{prelude::*, render::RenderContext2D, tree::Tree, utils::prelude::*};

use super::{component, component_or_default, component_try_mut, Layout};

/// Stacks visual the children widgets vertical or horizontal.
#[derive(Default)]
pub struct StackLayout {
    desired_size: RefCell<DirtySize>,
    old_alignment: Cell<(Alignment, Alignment)>,
}

impl StackLayout {
    pub fn new() -> Self {
        StackLayout::default()
    }

    pub fn set_dirty(&self, dirty: bool) {
        self.desired_size.borrow_mut().set_dirty(dirty);
    }
}

impl Layout for StackLayout {
    fn measure(
        &self,
        render_context_2_d: &mut RenderContext2D,
        entity: Entity,
        ecm: &mut EntityComponentManager<Tree, StringComponentStore>,
        layouts: &BTreeMap<Entity, Box<dyn Layout>>,
        theme: &ThemeValue,
    ) -> DirtySize {
        if component::<Visibility>(ecm, entity, "visibility") == Visibility::Collapsed {
            let mut desired = self.desired_size.borrow_mut();
            desired.set_size(0.0, 0.0);
            return desired.clone();
        }

        let halign: Alignment = component(ecm, entity, "horizontal_alignment");
        let valign: Alignment = component(ecm, entity, "vertical_alignment");
        let (old_valign, old_halign) = self.old_alignment.get();

        if halign != old_halign || valign != old_valign {
            self.set_dirty(true);
        }

        let orientation: Orientation = component(ecm, entity, "orientation");
        let mut dirty = false;
        let mut desired_size: (f64, f64) = (0.0, 0.0);

        let nchildren = ecm.entity_store().children[&entity].len();
        let spacing: f64 = component_or_default(ecm, entity, "spacing");

        for index in 0..nchildren {
            let child = ecm.entity_store().children[&entity][index];

            if let Some(child_layout) = layouts.get(&child) {
                let child_desired_size =
                    child_layout.measure(render_context_2_d, child, ecm, layouts, theme);

                let mut child_margin = {
                    if child_desired_size.width() > 0.0 && child_desired_size.height() > 0.0 {
                        component(ecm, child, "margin")
                    } else {
                        Thickness::default()
                    }
                };

                if spacing != 0.0 && nchildren > 1 {
                    apply_spacing(&mut child_margin, spacing, orientation, index, nchildren);
                }

                accumulate_desired_size(
                    &mut desired_size,
                    child_desired_size,
                    child_margin,
                    orientation,
                );

                if child_desired_size.dirty() || self.desired_size.borrow().dirty() {
                    dirty = true;
                }
            }
        }

        self.set_dirty(dirty);

        let mut desired = self.desired_size.borrow_mut();
        desired.set_size(desired_size.0, desired_size.1);
        desired.clone()
    }

    fn arrange(
        &self,
        render_context_2_d: &mut RenderContext2D,
        parent_size: (f64, f64),
        entity: Entity,
        ecm: &mut EntityComponentManager<Tree, StringComponentStore>,
        layouts: &BTreeMap<Entity, Box<dyn Layout>>,
        theme: &ThemeValue,
    ) -> (f64, f64) {
        if component::<Visibility>(ecm, entity, "visibility") == Visibility::Collapsed {
            self.desired_size.borrow_mut().set_size(0.0, 0.0);
            return (0.0, 0.0);
        }

        if !self.desired_size.borrow().dirty() {
            return self.desired_size.borrow().size();
        }

        let halign: Alignment = component(ecm, entity, "horizontal_alignment");
        let valign: Alignment = component(ecm, entity, "vertical_alignment");
        let margin: Thickness = component(ecm, entity, "margin");
        let constraint: Constraint = component(ecm, entity, "constraint");
        let orientation: Orientation = component(ecm, entity, "orientation");

        let mut size_counter = 0.0;

        let mut size = constraint.perform((
            halign.align_measure(
                parent_size.0,
                self.desired_size.borrow().width(),
                margin.left(),
                margin.right(),
            ),
            valign.align_measure(
                parent_size.1,
                self.desired_size.borrow().height(),
                margin.top(),
                margin.bottom(),
            ),
        ));

        let available_size = size;
        let nchildren = ecm.entity_store().children[&entity].len();
        let spacing: f64 = component_or_default(ecm, entity, "spacing");

        for index in 0..nchildren {
            let child = ecm.entity_store().children[&entity][index];

            match orientation {
                Orientation::Horizontal => {
                    if let Some(halign) =
                        component_try_mut::<Alignment>(ecm, child, "horizontal_alignment")
                    {
                        *halign = Alignment::Start;
                    }
                }
                _ => {
                    if let Some(halign) =
                        component_try_mut::<Alignment>(ecm, child, "vertical_alignment")
                    {
                        *halign = Alignment::Start;
                    }
                }
            }

            let mut child_desired_size = (0.0, 0.0);
            if let Some(child_layout) = layouts.get(&child) {
                child_desired_size =
                    child_layout.arrange(render_context_2_d, size, child, ecm, layouts, theme);
            }

            let mut child_margin = {
                if child_desired_size.0 > 0.0 && child_desired_size.1 > 0.0 {
                    component(ecm, child, "margin")
                } else {
                    Thickness::default()
                }
            };

            if spacing != 0.0 && nchildren > 1 {
                apply_spacing(&mut child_margin, spacing, orientation, index, nchildren);
            }

            let child_halign: Alignment = component(ecm, child, "horizontal_alignment");
            let child_valign: Alignment = component(ecm, child, "vertical_alignment");

            if let Some(child_bounds) = component_try_mut::<Rectangle>(ecm, child, "bounds") {
                apply_arrangement(
                    child_bounds,
                    &mut size_counter,
                    child_margin,
                    (child_halign, child_valign),
                    orientation,
                    available_size,
                );
            }
        }

        self.set_dirty(false);

        match orientation {
            Orientation::Horizontal => {
                size.0 = size_counter;
            }
            _ => {
                size.1 = size_counter;
            }
        }

        if let Some(bounds) = component_try_mut::<Rectangle>(ecm, entity, "bounds") {
            bounds.set_width(size.0);
            bounds.set_height(size.1);
        }
        size
    }
}

impl From<StackLayout> for Box<dyn Layout> {
    fn from(layout: StackLayout) -> Self {
        Box::new(layout)
    }
}

fn apply_arrangement(
    bounds: &mut Rectangle,
    size_counter: &mut f64,
    margin: Thickness,
    alignment: (Alignment, Alignment),
    orientation: Orientation,
    available_size: (f64, f64),
) {
    let (xpos, ypos, size);

    match orientation {
        Orientation::Horizontal => {
            xpos = *size_counter
                + alignment.0.align_position(
                    available_size.0,
                    bounds.width(),
                    margin.left(),
                    margin.right(),
                );

            ypos = alignment.1.align_position(
                available_size.1,
                bounds.height(),
                margin.top(),
                margin.bottom(),
            );

            size = bounds.width() + margin.left() + margin.right();
        }
        _ => {
            xpos = alignment.0.align_position(
                available_size.0,
                bounds.width(),
                margin.left(),
                margin.right(),
            );

            ypos = *size_counter
                + alignment.1.align_position(
                    available_size.1,
                    bounds.height(),
                    margin.top(),
                    margin.bottom(),
                );

            size = bounds.height() + margin.top() + margin.bottom();
        }
    };

    bounds.set_x(xpos);
    bounds.set_y(ypos);
    *size_counter += size;
}

/// Applies spacing to widgets in a stack, depending upon their position, and the orientation.
fn apply_spacing(
    margins: &mut Thickness,
    spacing: f64,
    orientation: Orientation,
    index: usize,
    nchildren: usize,
) {
    let start = if index == 0 { 0.0 } else { spacing / 2.0 };
    let end = if index == nchildren - 1 {
        0.0
    } else {
        spacing / 2.0
    };

    match orientation {
        Orientation::Vertical => {
            margins.top += start;
            margins.bottom += end;
        }
        Orientation::Horizontal => {
            margins.left += start;
            margins.right += end;
        }
    }
}

fn accumulate_desired_size(
    desired_size: &mut (f64, f64),
    desired: DirtySize,
    margin: Thickness,
    orientation: Orientation,
) {
    let width = desired.width() + margin.left() + margin.right();
    let height = desired.height() + margin.top() + margin.bottom();

    match orientation {
        Orientation::Horizontal => {
            desired_size.0 += width;
            desired_size.1 = height;
        }
        Orientation::Vertical => {
            desired_size.0 = width;
            desired_size.1 += height;
        }
    }
}

#[cfg(test)]
mod tests {
    use orbtk_utils::{Orientation, Thickness};
    use std::iter;

    use super::apply_spacing;

    const NUM_WIDGETS: usize = 5;

    #[test]
    fn spacing_vertical() {
        let expected = iter::once(Thickness {
            left: 0.0,
            right: 0.0,
            top: 0.0,
            bottom: 2.0,
        })
        .chain(
            iter::repeat(Thickness {
                left: 0.0,
                right: 0.0,
                top: 2.0,
                bottom: 2.0,
            })
            .take(3),
        )
        .chain(iter::once(Thickness {
            left: 0.0,
            right: 0.0,
            top: 2.0,
            bottom: 0.0,
        }));

        spacing(Orientation::Vertical, 4.0, expected);
    }

    #[test]
    fn spacing_horizontal() {
        let expected = iter::once(Thickness {
            left: 0.0,
            right: 4.0,
            top: 0.0,
            bottom: 0.0,
        })
        .chain(
            iter::repeat(Thickness {
                left: 4.0,
                right: 4.0,
                top: 0.0,
                bottom: 0.0,
            })
            .take(3),
        )
        .chain(iter::once(Thickness {
            left: 4.0,
            right: 0.0,
            top: 0.0,
            bottom: 0.0,
        }));

        spacing(Orientation::Horizontal, 8.0, expected);
    }

    fn spacing(orientation: Orientation, spacing: f64, expected: impl Iterator<Item = Thickness>) {
        let inputs = (0..NUM_WIDGETS).map(|id| (id, Thickness::default()));

        for ((index, mut input), expected) in inputs.zip(expected) {
            apply_spacing(&mut input, spacing, orientation, index, NUM_WIDGETS);
            assert_eq!(input, expected);
        }
    }
}
<commit_msg>Fix stack layout to use the max height/width<commit_after>use std::{
    cell::{Cell, RefCell},
    collections::BTreeMap,
};

use dces::prelude::Entity;

use crate::{prelude::*, render::RenderContext2D, tree::Tree, utils::prelude::*};

use super::{component, component_or_default, component_try_mut, Layout};

/// Stacks visual the children widgets vertical or horizontal.
#[derive(Default)]
pub struct StackLayout {
    desired_size: RefCell<DirtySize>,
    old_alignment: Cell<(Alignment, Alignment)>,
}

impl StackLayout {
    pub fn new() -> Self {
        StackLayout::default()
    }

    pub fn set_dirty(&self, dirty: bool) {
        self.desired_size.borrow_mut().set_dirty(dirty);
    }
}

impl Layout for StackLayout {
    fn measure(
        &self,
        render_context_2_d: &mut RenderContext2D,
        entity: Entity,
        ecm: &mut EntityComponentManager<Tree, StringComponentStore>,
        layouts: &BTreeMap<Entity, Box<dyn Layout>>,
        theme: &ThemeValue,
    ) -> DirtySize {
        if component::<Visibility>(ecm, entity, "visibility") == Visibility::Collapsed {
            let mut desired = self.desired_size.borrow_mut();
            desired.set_size(0.0, 0.0);
            return desired.clone();
        }

        let halign: Alignment = component(ecm, entity, "horizontal_alignment");
        let valign: Alignment = component(ecm, entity, "vertical_alignment");
        let (old_valign, old_halign) = self.old_alignment.get();

        if halign != old_halign || valign != old_valign {
            self.set_dirty(true);
        }

        let orientation: Orientation = component(ecm, entity, "orientation");
        let mut dirty = false;
        let mut desired_size: (f64, f64) = (0.0, 0.0);

        let nchildren = ecm.entity_store().children[&entity].len();
        let spacing: f64 = component_or_default(ecm, entity, "spacing");

        for index in 0..nchildren {
            let child = ecm.entity_store().children[&entity][index];

            if let Some(child_layout) = layouts.get(&child) {
                let child_desired_size =
                    child_layout.measure(render_context_2_d, child, ecm, layouts, theme);

                let mut child_margin = {
                    if child_desired_size.width() > 0.0 && child_desired_size.height() > 0.0 {
                        component(ecm, child, "margin")
                    } else {
                        Thickness::default()
                    }
                };

                if spacing != 0.0 && nchildren > 1 {
                    apply_spacing(&mut child_margin, spacing, orientation, index, nchildren);
                }

                accumulate_desired_size(
                    &mut desired_size,
                    child_desired_size,
                    child_margin,
                    orientation,
                );

                if child_desired_size.dirty() || self.desired_size.borrow().dirty() {
                    dirty = true;
                }
            }
        }

        self.set_dirty(dirty);

        let mut desired = self.desired_size.borrow_mut();
        desired.set_size(desired_size.0, desired_size.1);
        desired.clone()
    }

    fn arrange(
        &self,
        render_context_2_d: &mut RenderContext2D,
        parent_size: (f64, f64),
        entity: Entity,
        ecm: &mut EntityComponentManager<Tree, StringComponentStore>,
        layouts: &BTreeMap<Entity, Box<dyn Layout>>,
        theme: &ThemeValue,
    ) -> (f64, f64) {
        if component::<Visibility>(ecm, entity, "visibility") == Visibility::Collapsed {
            self.desired_size.borrow_mut().set_size(0.0, 0.0);
            return (0.0, 0.0);
        }

        if !self.desired_size.borrow().dirty() {
            return self.desired_size.borrow().size();
        }

        let halign: Alignment = component(ecm, entity, "horizontal_alignment");
        let valign: Alignment = component(ecm, entity, "vertical_alignment");
        let margin: Thickness = component(ecm, entity, "margin");
        let constraint: Constraint = component(ecm, entity, "constraint");
        let orientation: Orientation = component(ecm, entity, "orientation");

        let mut size_counter = 0.0;

        let mut size = constraint.perform((
            halign.align_measure(
                parent_size.0,
                self.desired_size.borrow().width(),
                margin.left(),
                margin.right(),
            ),
            valign.align_measure(
                parent_size.1,
                self.desired_size.borrow().height(),
                margin.top(),
                margin.bottom(),
            ),
        ));

        let available_size = size;
        let nchildren = ecm.entity_store().children[&entity].len();
        let spacing: f64 = component_or_default(ecm, entity, "spacing");

        for index in 0..nchildren {
            let child = ecm.entity_store().children[&entity][index];

            match orientation {
                Orientation::Horizontal => {
                    if let Some(halign) =
                        component_try_mut::<Alignment>(ecm, child, "horizontal_alignment")
                    {
                        *halign = Alignment::Start;
                    }
                }
                _ => {
                    if let Some(halign) =
                        component_try_mut::<Alignment>(ecm, child, "vertical_alignment")
                    {
                        *halign = Alignment::Start;
                    }
                }
            }

            let mut child_desired_size = (0.0, 0.0);
            if let Some(child_layout) = layouts.get(&child) {
                child_desired_size =
                    child_layout.arrange(render_context_2_d, size, child, ecm, layouts, theme);
            }

            let mut child_margin = {
                if child_desired_size.0 > 0.0 && child_desired_size.1 > 0.0 {
                    component(ecm, child, "margin")
                } else {
                    Thickness::default()
                }
            };

            if spacing != 0.0 && nchildren > 1 {
                apply_spacing(&mut child_margin, spacing, orientation, index, nchildren);
            }

            let child_halign: Alignment = component(ecm, child, "horizontal_alignment");
            let child_valign: Alignment = component(ecm, child, "vertical_alignment");

            if let Some(child_bounds) = component_try_mut::<Rectangle>(ecm, child, "bounds") {
                apply_arrangement(
                    child_bounds,
                    &mut size_counter,
                    child_margin,
                    (child_halign, child_valign),
                    orientation,
                    available_size,
                );
            }
        }

        self.set_dirty(false);

        match orientation {
            Orientation::Horizontal => {
                size.0 = size_counter;
            }
            _ => {
                size.1 = size_counter;
            }
        }

        if let Some(bounds) = component_try_mut::<Rectangle>(ecm, entity, "bounds") {
            bounds.set_width(size.0);
            bounds.set_height(size.1);
        }
        size
    }
}

impl From<StackLayout> for Box<dyn Layout> {
    fn from(layout: StackLayout) -> Self {
        Box::new(layout)
    }
}

fn apply_arrangement(
    bounds: &mut Rectangle,
    size_counter: &mut f64,
    margin: Thickness,
    alignment: (Alignment, Alignment),
    orientation: Orientation,
    available_size: (f64, f64),
) {
    let (xpos, ypos, size);

    match orientation {
        Orientation::Horizontal => {
            xpos = *size_counter
                + alignment.0.align_position(
                    available_size.0,
                    bounds.width(),
                    margin.left(),
                    margin.right(),
                );

            ypos = alignment.1.align_position(
                available_size.1,
                bounds.height(),
                margin.top(),
                margin.bottom(),
            );

            size = bounds.width() + margin.left() + margin.right();
        }
        _ => {
            xpos = alignment.0.align_position(
                available_size.0,
                bounds.width(),
                margin.left(),
                margin.right(),
            );

            ypos = *size_counter
                + alignment.1.align_position(
                    available_size.1,
                    bounds.height(),
                    margin.top(),
                    margin.bottom(),
                );

            size = bounds.height() + margin.top() + margin.bottom();
        }
    };

    bounds.set_x(xpos);
    bounds.set_y(ypos);
    *size_counter += size;
}

/// Applies spacing to widgets in a stack, depending upon their position, and the orientation.
fn apply_spacing(
    margins: &mut Thickness,
    spacing: f64,
    orientation: Orientation,
    index: usize,
    nchildren: usize,
) {
    let start = if index == 0 { 0.0 } else { spacing / 2.0 };
    let end = if index == nchildren - 1 {
        0.0
    } else {
        spacing / 2.0
    };

    match orientation {
        Orientation::Vertical => {
            margins.top += start;
            margins.bottom += end;
        }
        Orientation::Horizontal => {
            margins.left += start;
            margins.right += end;
        }
    }
}

fn accumulate_desired_size(
    desired_size: &mut (f64, f64),
    desired: DirtySize,
    margin: Thickness,
    orientation: Orientation,
) {
    let width = desired.width() + margin.left() + margin.right();
    let height = desired.height() + margin.top() + margin.bottom();

    match orientation {
        Orientation::Horizontal => {
            desired_size.0 += width;
            desired_size.1 = desired_size.1.max(height);
        }
        Orientation::Vertical => {
            desired_size.0 = desired_size.0.max(width);
            desired_size.1 += height;
        }
    }
}

#[cfg(test)]
mod tests {
    use orbtk_utils::{Orientation, Thickness};
    use std::iter;

    use super::apply_spacing;

    const NUM_WIDGETS: usize = 5;

    #[test]
    fn spacing_vertical() {
        let expected = iter::once(Thickness {
            left: 0.0,
            right: 0.0,
            top: 0.0,
            bottom: 2.0,
        })
        .chain(
            iter::repeat(Thickness {
                left: 0.0,
                right: 0.0,
                top: 2.0,
                bottom: 2.0,
            })
            .take(3),
        )
        .chain(iter::once(Thickness {
            left: 0.0,
            right: 0.0,
            top: 2.0,
            bottom: 0.0,
        }));

        spacing(Orientation::Vertical, 4.0, expected);
    }

    #[test]
    fn spacing_horizontal() {
        let expected = iter::once(Thickness {
            left: 0.0,
            right: 4.0,
            top: 0.0,
            bottom: 0.0,
        })
        .chain(
            iter::repeat(Thickness {
                left: 4.0,
                right: 4.0,
                top: 0.0,
                bottom: 0.0,
            })
            .take(3),
        )
        .chain(iter::once(Thickness {
            left: 4.0,
            right: 0.0,
            top: 0.0,
            bottom: 0.0,
        }));

        spacing(Orientation::Horizontal, 8.0, expected);
    }

    fn spacing(orientation: Orientation, spacing: f64, expected: impl Iterator<Item = Thickness>) {
        let inputs = (0..NUM_WIDGETS).map(|id| (id, Thickness::default()));

        for ((index, mut input), expected) in inputs.zip(expected) {
            apply_spacing(&mut input, spacing, orientation, index, NUM_WIDGETS);
            assert_eq!(input, expected);
        }
    }
}
<|endoftext|>


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

use crate::{AptosDB, LedgerPrunerManager, LedgerStore, PrunerManager, TransactionStore};
use aptos_temppath::TempPath;
use proptest::proptest;
use std::sync::Arc;

use aptos_types::{
    account_address::AccountAddress,
    block_metadata::BlockMetadata,
    transaction::{SignedTransaction, Transaction},
};

use accumulator::HashReader;
use aptos_config::config::LedgerPrunerConfig;
use aptos_types::proof::position::Position;
use aptos_types::{
    transaction::{TransactionInfo, Version},
    write_set::WriteSet,
};
use proptest::{collection::vec, prelude::*};
use schemadb::SchemaBatch;
use storage_interface::DbReader;

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

    #[test]
    fn test_txn_store_pruner(txns in vec(
        prop_oneof![
            any::<BlockMetadata>().prop_map(Transaction::BlockMetadata),
            any::<SignedTransaction>().prop_map(Transaction::UserTransaction),
        ], 1..100,),
        txn_infos in vec(any::<TransactionInfo>(),100,),
        step_size in 1usize..20,
    ) {
        verify_txn_store_pruner(txns, txn_infos, step_size)
    }

     #[test]
    fn test_write_set_pruner(
        write_set in vec(any::<WriteSet>(), 100),
        ) {
            verify_write_set_pruner(write_set);
        }
}

fn verify_write_set_pruner(write_sets: Vec<WriteSet>) {
    let tmp_dir = TempPath::new();
    let aptos_db = AptosDB::new_for_test(&tmp_dir);
    let transaction_store = &aptos_db.transaction_store;
    let num_write_sets = write_sets.len();

    let pruner = LedgerPrunerManager::new(
        Arc::clone(&aptos_db.ledger_db),
        Arc::clone(&aptos_db.state_store),
        LedgerPrunerConfig {
            enable: true,
            prune_window: 0,
            batch_size: 1,
            user_pruning_window_offset: 0,
        },
    );

    // write sets
    let mut batch = SchemaBatch::new();
    for (ver, ws) in write_sets.iter().enumerate() {
        transaction_store
            .put_write_set(ver as Version, ws, &mut batch)
            .unwrap();
    }
    aptos_db.ledger_db.write_schemas(batch).unwrap();
    // start pruning write sets in batches of size 2 and verify transactions have been pruned from DB
    for i in (0..=num_write_sets).step_by(2) {
        pruner
            .wake_and_wait_pruner(i as u64 /* latest_version */)
            .unwrap();
        // ensure that all transaction up to i * 2 has been pruned
        for j in 0..i {
            assert!(transaction_store.get_write_set(j as u64).is_err());
        }
        // ensure all other are valid in DB
        for j in i..num_write_sets {
            let write_set_from_db = transaction_store.get_write_set(j as u64).unwrap();
            assert_eq!(write_set_from_db, *write_sets.get(j).unwrap());
        }
    }
}

fn verify_txn_store_pruner(
    txns: Vec<Transaction>,
    txn_infos: Vec<TransactionInfo>,
    step_size: usize,
) {
    let tmp_dir = TempPath::new();
    let aptos_db = AptosDB::new_for_test(&tmp_dir);
    let transaction_store = &aptos_db.transaction_store;
    let ledger_store = LedgerStore::new(Arc::clone(&aptos_db.ledger_db));
    let num_transaction = txns.len();

    let pruner = LedgerPrunerManager::new(
        Arc::clone(&aptos_db.ledger_db),
        Arc::clone(&aptos_db.state_store),
        LedgerPrunerConfig {
            enable: true,
            prune_window: 0,
            batch_size: 1,
            user_pruning_window_offset: 0,
        },
    );

    let ledger_version = num_transaction as Version - 1;
    put_txn_in_store(
        &aptos_db,
        transaction_store,
        &ledger_store,
        &txn_infos,
        &txns,
    );

    // start pruning transactions batches of size step_size and verify transactions have been pruned
    // from DB
    for i in (0..=num_transaction).step_by(step_size) {
        pruner
            .wake_and_wait_pruner(i as u64 /* latest_version */)
            .unwrap();
        // ensure that all transaction up to i * 2 has been pruned
        assert_eq!(
            *pruner.last_version_sent_to_pruner.as_ref().lock(),
            i as u64
        );
        for j in 0..i {
            verify_txn_not_in_store(transaction_store, &txns, j as u64, ledger_version);
            // Ensure that transaction accumulator is pruned in DB. This can be done by trying to
            // read transaction proof.
            // Note: we only prune versions which are odd numbers because the even versions will be
            // pruned in the iteration of even_version + 1. So if the end version, i - 1, is an even
            // version, it will not be pruned.
            if j != i - 1 || j % 2 == 1 {
                assert!(ledger_store
                    .get_transaction_proof(j as u64, ledger_version)
                    .is_err());
            }
        }
        // ensure all other are valid in DB
        for j in i..num_transaction {
            verify_txn_in_store(
                transaction_store,
                &ledger_store,
                &txns,
                j as u64,
                ledger_version,
            );
            aptos_db.get_accumulator_summary(j as Version).unwrap();
        }
        verify_transaction_accumulator_pruned(&ledger_store, i as u64);
    }
}

fn verify_txn_not_in_store(
    transaction_store: &TransactionStore,
    txns: &[Transaction],
    index: u64,
    ledger_version: u64,
) {
    // Ensure that all transaction from transaction schema store has been pruned
    assert!(transaction_store.get_transaction(index).is_err());
    // Ensure that transaction by account store has been pruned
    if let Transaction::UserTransaction(txn) = txns.get(index as usize).unwrap() {
        assert!(transaction_store
            .get_account_transaction_version(txn.sender(), txn.sequence_number(), ledger_version,)
            .unwrap()
            .is_none());
    }
}

fn verify_txn_in_store(
    transaction_store: &TransactionStore,
    ledger_store: &LedgerStore,
    txns: &[Transaction],
    index: u64,
    ledger_version: u64,
) {
    verify_transaction_in_transaction_store(
        transaction_store,
        txns.get(index as usize).unwrap(),
        index as u64,
    );
    if let Transaction::UserTransaction(txn) = txns.get(index as usize).unwrap() {
        verify_transaction_in_account_txn_by_version_index(
            transaction_store,
            index as u64,
            txn.sender(),
            txn.sequence_number(),
            ledger_version,
        );
    }
    // Ensure that transaction accumulator is in DB. This can be done by trying
    // to read transaction proof
    assert!(ledger_store
        .get_transaction_proof(index, ledger_version)
        .is_ok());
}

// Ensure that transaction accumulator has been pruned as well. The idea to verify is get the
// inorder position of the left child of the accumulator root and ensure that all lower index
// position from the DB should be deleted. We need to make several conversion between inorder and
// postorder transaction because the DB stores the indices in postorder, while the APIs for the
// accumulator deals with inorder.
fn verify_transaction_accumulator_pruned(ledger_store: &LedgerStore, least_readable_version: u64) {
    let least_readable_position = if least_readable_version > 0 {
        Position::root_from_leaf_index(least_readable_version).left_child()
    } else {
        Position::root_from_leaf_index(least_readable_version)
    };
    let least_readable_position_postorder = least_readable_position.to_postorder_index();
    for i in 0..least_readable_position_postorder {
        assert!(ledger_store
            .get(Position::from_postorder_index(i).unwrap())
            .is_err())
    }
}

fn put_txn_in_store(
    aptos_db: &AptosDB,
    transaction_store: &TransactionStore,
    ledger_store: &LedgerStore,
    txn_infos: &[TransactionInfo],
    txns: &[Transaction],
) {
    let mut batch = SchemaBatch::new();
    for i in 0..txns.len() {
        transaction_store
            .put_transaction(i as u64, txns.get(i).unwrap(), &mut batch)
            .unwrap();
    }
    ledger_store
        .put_transaction_infos(0, txn_infos, &mut batch)
        .unwrap();
    aptos_db.ledger_db.write_schemas(batch).unwrap();
}

fn verify_transaction_in_transaction_store(
    transaction_store: &TransactionStore,
    expected_value: &Transaction,
    version: Version,
) {
    let txn = transaction_store.get_transaction(version).unwrap();
    assert_eq!(txn, *expected_value)
}

fn verify_transaction_in_account_txn_by_version_index(
    transaction_store: &TransactionStore,
    expected_value: Version,
    address: AccountAddress,
    sequence_number: u64,
    ledger_version: Version,
) {
    let transaction = transaction_store
        .get_account_transaction_version(address, sequence_number, ledger_version)
        .unwrap()
        .unwrap();
    assert_eq!(transaction, expected_value)
}
<commit_msg>[Storage] Add back the test for min_readable_version. (#3984)<commit_after>// Copyright (c) Aptos
// SPDX-License-Identifier: Apache-2.0

use crate::{AptosDB, LedgerPrunerManager, LedgerStore, PrunerManager, TransactionStore};
use aptos_temppath::TempPath;
use proptest::proptest;
use std::sync::Arc;

use aptos_types::{
    account_address::AccountAddress,
    block_metadata::BlockMetadata,
    transaction::{SignedTransaction, Transaction},
};

use accumulator::HashReader;
use aptos_config::config::LedgerPrunerConfig;
use aptos_types::{
    proof::position::Position,
    transaction::{TransactionInfo, Version},
    write_set::WriteSet,
};
use proptest::{collection::vec, prelude::*};
use schemadb::SchemaBatch;
use storage_interface::DbReader;

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

    #[test]
    fn test_txn_store_pruner(txns in vec(
        prop_oneof![
            any::<BlockMetadata>().prop_map(Transaction::BlockMetadata),
            any::<SignedTransaction>().prop_map(Transaction::UserTransaction),
        ], 1..100,),
        txn_infos in vec(any::<TransactionInfo>(),100,),
        step_size in 1usize..20,
    ) {
        verify_txn_store_pruner(txns, txn_infos, step_size)
    }

     #[test]
    fn test_write_set_pruner(
        write_set in vec(any::<WriteSet>(), 100),
        ) {
            verify_write_set_pruner(write_set);
        }
}

fn verify_write_set_pruner(write_sets: Vec<WriteSet>) {
    let tmp_dir = TempPath::new();
    let aptos_db = AptosDB::new_for_test(&tmp_dir);
    let transaction_store = &aptos_db.transaction_store;
    let num_write_sets = write_sets.len();

    let pruner = LedgerPrunerManager::new(
        Arc::clone(&aptos_db.ledger_db),
        Arc::clone(&aptos_db.state_store),
        LedgerPrunerConfig {
            enable: true,
            prune_window: 0,
            batch_size: 1,
            user_pruning_window_offset: 0,
        },
    );

    // write sets
    let mut batch = SchemaBatch::new();
    for (ver, ws) in write_sets.iter().enumerate() {
        transaction_store
            .put_write_set(ver as Version, ws, &mut batch)
            .unwrap();
    }
    aptos_db.ledger_db.write_schemas(batch).unwrap();
    // start pruning write sets in batches of size 2 and verify transactions have been pruned from DB
    for i in (0..=num_write_sets).step_by(2) {
        pruner
            .wake_and_wait_pruner(i as u64 /* latest_version */)
            .unwrap();
        // ensure that all transaction up to i * 2 has been pruned
        for j in 0..i {
            assert!(transaction_store.get_write_set(j as u64).is_err());
        }
        // ensure all other are valid in DB
        for j in i..num_write_sets {
            let write_set_from_db = transaction_store.get_write_set(j as u64).unwrap();
            assert_eq!(write_set_from_db, *write_sets.get(j).unwrap());
        }
    }
}

fn verify_txn_store_pruner(
    txns: Vec<Transaction>,
    txn_infos: Vec<TransactionInfo>,
    step_size: usize,
) {
    let tmp_dir = TempPath::new();
    let aptos_db = AptosDB::new_for_test(&tmp_dir);
    let transaction_store = &aptos_db.transaction_store;
    let ledger_store = LedgerStore::new(Arc::clone(&aptos_db.ledger_db));
    let num_transaction = txns.len();

    let ledger_version = num_transaction as Version - 1;
    put_txn_in_store(
        &aptos_db,
        transaction_store,
        &ledger_store,
        &txn_infos,
        &txns,
    );

    // start pruning transactions batches of size step_size and verify transactions have been pruned
    // from DB
    for i in (0..=num_transaction).step_by(step_size) {
        // Initialize a pruner in every iteration to test the min_readable_version initialization
        // logic.
        let pruner = LedgerPrunerManager::new(
            Arc::clone(&aptos_db.ledger_db),
            Arc::clone(&aptos_db.state_store),
            LedgerPrunerConfig {
                enable: true,
                prune_window: 0,
                batch_size: 1,
                user_pruning_window_offset: 0,
            },
        );
        pruner
            .wake_and_wait_pruner(i as u64 /* latest_version */)
            .unwrap();
        // ensure that all transaction up to i * 2 has been pruned
        assert_eq!(
            *pruner.last_version_sent_to_pruner.as_ref().lock(),
            i as u64
        );
        for j in 0..i {
            verify_txn_not_in_store(transaction_store, &txns, j as u64, ledger_version);
            // Ensure that transaction accumulator is pruned in DB. This can be done by trying to
            // read transaction proof.
            // Note: we only prune versions which are odd numbers because the even versions will be
            // pruned in the iteration of even_version + 1. So if the end version, i - 1, is an even
            // version, it will not be pruned.
            if j != i - 1 || j % 2 == 1 {
                assert!(ledger_store
                    .get_transaction_proof(j as u64, ledger_version)
                    .is_err());
            }
        }
        // ensure all other are valid in DB
        for j in i..num_transaction {
            verify_txn_in_store(
                transaction_store,
                &ledger_store,
                &txns,
                j as u64,
                ledger_version,
            );
            aptos_db.get_accumulator_summary(j as Version).unwrap();
        }
        verify_transaction_accumulator_pruned(&ledger_store, i as u64);
    }
}

fn verify_txn_not_in_store(
    transaction_store: &TransactionStore,
    txns: &[Transaction],
    index: u64,
    ledger_version: u64,
) {
    // Ensure that all transaction from transaction schema store has been pruned
    assert!(transaction_store.get_transaction(index).is_err());
    // Ensure that transaction by account store has been pruned
    if let Transaction::UserTransaction(txn) = txns.get(index as usize).unwrap() {
        assert!(transaction_store
            .get_account_transaction_version(txn.sender(), txn.sequence_number(), ledger_version,)
            .unwrap()
            .is_none());
    }
}

fn verify_txn_in_store(
    transaction_store: &TransactionStore,
    ledger_store: &LedgerStore,
    txns: &[Transaction],
    index: u64,
    ledger_version: u64,
) {
    verify_transaction_in_transaction_store(
        transaction_store,
        txns.get(index as usize).unwrap(),
        index as u64,
    );
    if let Transaction::UserTransaction(txn) = txns.get(index as usize).unwrap() {
        verify_transaction_in_account_txn_by_version_index(
            transaction_store,
            index as u64,
            txn.sender(),
            txn.sequence_number(),
            ledger_version,
        );
    }
    // Ensure that transaction accumulator is in DB. This can be done by trying
    // to read transaction proof
    assert!(ledger_store
        .get_transaction_proof(index, ledger_version)
        .is_ok());
}

// Ensure that transaction accumulator has been pruned as well. The idea to verify is get the
// inorder position of the left child of the accumulator root and ensure that all lower index
// position from the DB should be deleted. We need to make several conversion between inorder and
// postorder transaction because the DB stores the indices in postorder, while the APIs for the
// accumulator deals with inorder.
fn verify_transaction_accumulator_pruned(ledger_store: &LedgerStore, least_readable_version: u64) {
    let least_readable_position = if least_readable_version > 0 {
        Position::root_from_leaf_index(least_readable_version).left_child()
    } else {
        Position::root_from_leaf_index(least_readable_version)
    };
    let least_readable_position_postorder = least_readable_position.to_postorder_index();
    for i in 0..least_readable_position_postorder {
        assert!(ledger_store
            .get(Position::from_postorder_index(i).unwrap())
            .is_err())
    }
}

fn put_txn_in_store(
    aptos_db: &AptosDB,
    transaction_store: &TransactionStore,
    ledger_store: &LedgerStore,
    txn_infos: &[TransactionInfo],
    txns: &[Transaction],
) {
    let mut batch = SchemaBatch::new();
    for i in 0..txns.len() {
        transaction_store
            .put_transaction(i as u64, txns.get(i).unwrap(), &mut batch)
            .unwrap();
    }
    ledger_store
        .put_transaction_infos(0, txn_infos, &mut batch)
        .unwrap();
    aptos_db.ledger_db.write_schemas(batch).unwrap();
}

fn verify_transaction_in_transaction_store(
    transaction_store: &TransactionStore,
    expected_value: &Transaction,
    version: Version,
) {
    let txn = transaction_store.get_transaction(version).unwrap();
    assert_eq!(txn, *expected_value)
}

fn verify_transaction_in_account_txn_by_version_index(
    transaction_store: &TransactionStore,
    expected_value: Version,
    address: AccountAddress,
    sequence_number: u64,
    ledger_version: Version,
) {
    let transaction = transaction_store
        .get_account_transaction_version(address, sequence_number, ledger_version)
        .unwrap()
        .unwrap();
    assert_eq!(transaction, expected_value)
}
<|endoftext|>


<commit_before>mod utils;
mod cpp_type;
mod cpp_method;
mod cpp_ffi_data;
mod cpp_operator;
mod cpp_parser;
mod full_run;

use std::path::{Path, PathBuf};
use common::file_utils::{create_dir_all, PathBufWithAdded};

pub enum TempTestDir {
  System(::tempdir::TempDir),
  Custom(PathBuf),
}

impl TempTestDir {
  pub fn new(name: &str) -> TempTestDir {
    if let Ok(value) = ::std::env::var("CPP_TO_RUST_TEMP_TEST_DIR") {
      let path = PathBuf::from(value).with_added(name);
      create_dir_all(&path).unwrap();
      TempTestDir::Custom(path)
    } else {
      TempTestDir::System(::tempdir::TempDir::new(name).unwrap())
    }
  }

  pub fn path(&self) -> &Path {
    match *self {
      TempTestDir::System(ref dir) => dir.path(),
      TempTestDir::Custom(ref path) => path,
    }
  }
}

// Testing plan:
//
// - qt_specific
// - main (refactor)
// - cpp_data
// - cpp_ffi_generator
// - cpp_code_generator
//
// - utils (file operations)
//
// - rust_type
// - rust_info
// - rust_generator
// - rust_code_generator
//
//
<commit_msg>new ci config 3<commit_after>mod cpp_type;
mod cpp_method;
mod cpp_ffi_data;
mod cpp_operator;
mod cpp_parser;
mod full_run;

use std::path::{Path, PathBuf};
use common::file_utils::{create_dir_all, PathBufWithAdded};

pub enum TempTestDir {
  System(::tempdir::TempDir),
  Custom(PathBuf),
}

impl TempTestDir {
  pub fn new(name: &str) -> TempTestDir {
    if let Ok(value) = ::std::env::var("CPP_TO_RUST_TEMP_TEST_DIR") {
      let path = PathBuf::from(value).with_added(name);
      create_dir_all(&path).unwrap();
      TempTestDir::Custom(path)
    } else {
      TempTestDir::System(::tempdir::TempDir::new(name).unwrap())
    }
  }

  pub fn path(&self) -> &Path {
    match *self {
      TempTestDir::System(ref dir) => dir.path(),
      TempTestDir::Custom(ref path) => path,
    }
  }
}

// Testing plan:
//
// - qt_specific
// - main (refactor)
// - cpp_data
// - cpp_ffi_generator
// - cpp_code_generator
//
// - utils (file operations)
//
// - rust_type
// - rust_info
// - rust_generator
// - rust_code_generator
//
//
<|endoftext|>


<commit_before>//! Provides a timer which runs in the tail-end of each slot and maybe advances the state of the
//! head block forward a single slot.
//!
//! This provides an optimization with the following benefits:
//!
//! 1. Removes the burden of a single, mandatory `per_slot_processing` call from the leading-edge of
//!    block processing. This helps import blocks faster.
//! 2. Allows the node to learn of the shuffling for the next epoch, before the first block from
//!    that epoch has arrived. This helps reduce gossip block propagation times.
//!
//! The downsides to this optimization are:
//!
//! 1. We are required to store an additional `BeaconState` for the head block. This consumes
//!    memory.
//! 2. There's a possibility that the head block is never built upon, causing wasted CPU cycles.
use crate::validator_monitor::HISTORIC_EPOCHS as VALIDATOR_MONITOR_HISTORIC_EPOCHS;
use crate::{
    beacon_chain::BLOCK_PROCESSING_CACHE_LOCK_TIMEOUT, snapshot_cache::StateAdvance, BeaconChain,
    BeaconChainError, BeaconChainTypes,
};
use slog::{debug, error, warn, Logger};
use slot_clock::SlotClock;
use state_processing::per_slot_processing;
use std::sync::{
    atomic::{AtomicBool, Ordering},
    Arc,
};
use task_executor::TaskExecutor;
use tokio::time::sleep;
use types::{EthSpec, Hash256, Slot};

/// If the head slot is more than `MAX_ADVANCE_DISTANCE` from the current slot, then don't perform
/// the state advancement.
///
/// This avoids doing unnecessary work whilst the node is syncing or has perhaps been put to sleep
/// for some period of time.
const MAX_ADVANCE_DISTANCE: u64 = 4;

#[derive(Debug)]
enum Error {
    BeaconChain(BeaconChainError),
    HeadMissingFromSnapshotCache(Hash256),
    MaxDistanceExceeded { current_slot: Slot, head_slot: Slot },
    StateAlreadyAdvanced { block_root: Hash256 },
    BadStateSlot { state_slot: Slot, block_slot: Slot },
}

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

/// Provides a simple thread-safe lock to be used for task co-ordination. Practically equivalent to
/// `Mutex<()>`.
#[derive(Clone)]
struct Lock(Arc<AtomicBool>);

impl Lock {
    /// Instantiate an unlocked self.
    pub fn new() -> Self {
        Self(Arc::new(AtomicBool::new(false)))
    }

    /// Lock self, returning `true` if the lock was already set.
    pub fn lock(&self) -> bool {
        self.0.fetch_or(true, Ordering::SeqCst)
    }

    /// Unlock self.
    pub fn unlock(&self) {
        self.0.store(false, Ordering::SeqCst);
    }
}

/// Spawns the timer described in the module-level documentation.
pub fn spawn_state_advance_timer<T: BeaconChainTypes>(
    executor: TaskExecutor,
    beacon_chain: Arc<BeaconChain<T>>,
    log: Logger,
) {
    executor.spawn(
        state_advance_timer(executor.clone(), beacon_chain, log),
        "state_advance_timer",
    );
}

/// Provides the timer described in the module-level documentation.
async fn state_advance_timer<T: BeaconChainTypes>(
    executor: TaskExecutor,
    beacon_chain: Arc<BeaconChain<T>>,
    log: Logger,
) {
    let is_running = Lock::new();
    let slot_clock = &beacon_chain.slot_clock;
    let slot_duration = slot_clock.slot_duration();

    loop {
        match beacon_chain.slot_clock.duration_to_next_slot() {
            Some(duration) => sleep(duration + (slot_duration / 4) * 3).await,
            None => {
                error!(log, "Failed to read slot clock");
                // If we can't read the slot clock, just wait another slot.
                sleep(slot_duration).await;
                continue;
            }
        };

        // Only start spawn the state advance task if the lock was previously free.
        if !is_running.lock() {
            let log = log.clone();
            let beacon_chain = beacon_chain.clone();
            let is_running = is_running.clone();

            executor.spawn_blocking(
                move || {
                    match advance_head(&beacon_chain, &log) {
                        Ok(()) => (),
                        Err(Error::BeaconChain(e)) => error!(
                            log,
                            "Failed to advance head state";
                            "error" => ?e
                        ),
                        Err(Error::StateAlreadyAdvanced { block_root }) => debug!(
                            log,
                            "State already advanced on slot";
                            "block_root" => ?block_root
                        ),
                        Err(Error::MaxDistanceExceeded {
                            current_slot,
                            head_slot,
                        }) => debug!(
                            log,
                            "Refused to advance head state";
                            "head_slot" => head_slot,
                            "current_slot" => current_slot,
                        ),
                        other => warn!(
                            log,
                            "Did not advance head state";
                            "reason" => ?other
                        ),
                    };

                    // Permit this blocking task to spawn again, next time the timer fires.
                    is_running.unlock();
                },
                "state_advance_blocking",
            );
        } else {
            warn!(
                log,
                "State advance routine overloaded";
                "msg" => "system resources may be overloaded"
            )
        }
    }
}

/// Reads the `snapshot_cache` from the `beacon_chain` and attempts to take a clone of the
/// `BeaconState` of the head block. If it obtains this clone, the state will be advanced a single
/// slot then placed back in the `snapshot_cache` to be used for block verification.
///
/// See the module-level documentation for rationale.
fn advance_head<T: BeaconChainTypes>(
    beacon_chain: &BeaconChain<T>,
    log: &Logger,
) -> Result<(), Error> {
    let current_slot = beacon_chain.slot()?;

    // These brackets ensure that the `head_slot` value is dropped before we run fork choice and
    // potentially invalidate it.
    //
    // Fork-choice is not run *before* this function to avoid unnecessary calls whilst syncing.
    {
        let head_slot = beacon_chain.head_info()?.slot;

        // Don't run this when syncing or if lagging too far behind.
        if head_slot + MAX_ADVANCE_DISTANCE < current_slot {
            return Err(Error::MaxDistanceExceeded {
                current_slot,
                head_slot,
            });
        }
    }

    // Run fork choice so we get the latest view of the head.
    //
    // This is useful since it's quite likely that the last time we ran fork choice was shortly
    // after receiving the latest gossip block, but not necessarily after we've received the
    // majority of attestations.
    beacon_chain.fork_choice()?;

    let head_root = beacon_chain.head_info()?.block_root;

    let (head_slot, head_state_root, mut state) = match beacon_chain
        .snapshot_cache
        .try_write_for(BLOCK_PROCESSING_CACHE_LOCK_TIMEOUT)
        .ok_or(BeaconChainError::SnapshotCacheLockTimeout)?
        .get_for_state_advance(head_root)
    {
        StateAdvance::AlreadyAdvanced => {
            return Err(Error::StateAlreadyAdvanced {
                block_root: head_root,
            })
        }
        StateAdvance::BlockNotFound => return Err(Error::HeadMissingFromSnapshotCache(head_root)),
        StateAdvance::State {
            state,
            state_root,
            block_slot,
        } => (block_slot, state_root, *state),
    };

    let initial_slot = state.slot;
    let initial_epoch = state.current_epoch();

    let state_root = if state.slot == head_slot {
        Some(head_state_root)
    } else {
        // Protect against advancing a state more than a single slot.
        //
        // Advancing more than one slot without storing the intermediate state would corrupt the
        // database. Future works might store temporary, intermediate states inside this function.
        return Err(Error::BadStateSlot {
            block_slot: head_slot,
            state_slot: state.slot,
        });
    };

    // Advance the state a single slot.
    if let Some(summary) = per_slot_processing(&mut state, state_root, &beacon_chain.spec)
        .map_err(BeaconChainError::from)?
    {
        // Only notify the validator monitor for recent blocks.
        if state.current_epoch() + VALIDATOR_MONITOR_HISTORIC_EPOCHS as u64
            >= current_slot.epoch(T::EthSpec::slots_per_epoch())
        {
            // Potentially create logs/metrics for locally monitored validators.
            beacon_chain
                .validator_monitor
                .read()
                .process_validator_statuses(state.current_epoch(), &summary.statuses);
        }
    }

    debug!(
        log,
        "Advanced head state one slot";
        "head_root" => ?head_root,
        "state_slot" => state.slot,
        "current_slot" => current_slot,
    );

    // If the `pre_state` is in a later epoch than `state`, pre-emptively add the proposer
    // shuffling for the next epoch into the cache.
    if initial_epoch > state.current_epoch() {
        debug!(
            log,
            "Priming proposer cache";
            "head_root" => ?head_root,
            "state_epoch" => state.current_epoch(),
            "current_epoch" => current_slot.epoch(T::EthSpec::slots_per_epoch()),
        );
        beacon_chain
            .beacon_proposer_cache
            .lock()
            .insert(
                state.current_epoch(),
                head_root,
                state
                    .get_beacon_proposer_indices(&beacon_chain.spec)
                    .map_err(BeaconChainError::from)?,
                state.fork,
            )
            .map_err(BeaconChainError::from)?;
    }

    let final_slot = state.slot;

    // Insert the advanced state back into the snapshot cache.
    beacon_chain
        .snapshot_cache
        .try_write_for(BLOCK_PROCESSING_CACHE_LOCK_TIMEOUT)
        .ok_or(BeaconChainError::SnapshotCacheLockTimeout)?
        .update_pre_state(head_root, state)
        .ok_or(Error::HeadMissingFromSnapshotCache(head_root))?;

    // If we have moved into the next slot whilst processing the state then this function is going
    // to become ineffective and likely become a hindrance as we're stealing the tree hash cache
    // from the snapshot cache (which may force the next block to rebuild a new one).
    //
    // If this warning occurs very frequently on well-resourced machines then we should consider
    // starting it earlier in the slot. Otherwise, it's a good indication that the machine is too
    // slow/overloaded and will be useful information for the user.
    let starting_slot = current_slot;
    let current_slot = beacon_chain.slot()?;
    if starting_slot < current_slot {
        warn!(
            log,
            "State advance too slow";
            "head_root" => %head_root,
            "advanced_slot" => final_slot,
            "current_slot" => current_slot,
            "starting_slot" => starting_slot,
            "msg" => "system resources may be overloaded",
        );
    }

    debug!(
        log,
        "Completed state advance";
        "head_root" => ?head_root,
        "advanced_slot" => final_slot,
        "initial_slot" => initial_slot,
    );

    Ok(())
}

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

    #[test]
    fn lock() {
        let lock = Lock::new();
        assert_eq!(lock.lock(), false);
        assert_eq!(lock.lock(), true);
        assert_eq!(lock.lock(), true);
        lock.unlock();
        assert_eq!(lock.lock(), false);
        assert_eq!(lock.lock(), true);
    }
}
<commit_msg>Fix proposer cache priming upon state advance (#2252)<commit_after>//! Provides a timer which runs in the tail-end of each slot and maybe advances the state of the
//! head block forward a single slot.
//!
//! This provides an optimization with the following benefits:
//!
//! 1. Removes the burden of a single, mandatory `per_slot_processing` call from the leading-edge of
//!    block processing. This helps import blocks faster.
//! 2. Allows the node to learn of the shuffling for the next epoch, before the first block from
//!    that epoch has arrived. This helps reduce gossip block propagation times.
//!
//! The downsides to this optimization are:
//!
//! 1. We are required to store an additional `BeaconState` for the head block. This consumes
//!    memory.
//! 2. There's a possibility that the head block is never built upon, causing wasted CPU cycles.
use crate::validator_monitor::HISTORIC_EPOCHS as VALIDATOR_MONITOR_HISTORIC_EPOCHS;
use crate::{
    beacon_chain::BLOCK_PROCESSING_CACHE_LOCK_TIMEOUT, snapshot_cache::StateAdvance, BeaconChain,
    BeaconChainError, BeaconChainTypes,
};
use slog::{debug, error, warn, Logger};
use slot_clock::SlotClock;
use state_processing::per_slot_processing;
use std::sync::{
    atomic::{AtomicBool, Ordering},
    Arc,
};
use task_executor::TaskExecutor;
use tokio::time::sleep;
use types::{EthSpec, Hash256, Slot};

/// If the head slot is more than `MAX_ADVANCE_DISTANCE` from the current slot, then don't perform
/// the state advancement.
///
/// This avoids doing unnecessary work whilst the node is syncing or has perhaps been put to sleep
/// for some period of time.
const MAX_ADVANCE_DISTANCE: u64 = 4;

#[derive(Debug)]
enum Error {
    BeaconChain(BeaconChainError),
    HeadMissingFromSnapshotCache(Hash256),
    MaxDistanceExceeded { current_slot: Slot, head_slot: Slot },
    StateAlreadyAdvanced { block_root: Hash256 },
    BadStateSlot { state_slot: Slot, block_slot: Slot },
}

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

/// Provides a simple thread-safe lock to be used for task co-ordination. Practically equivalent to
/// `Mutex<()>`.
#[derive(Clone)]
struct Lock(Arc<AtomicBool>);

impl Lock {
    /// Instantiate an unlocked self.
    pub fn new() -> Self {
        Self(Arc::new(AtomicBool::new(false)))
    }

    /// Lock self, returning `true` if the lock was already set.
    pub fn lock(&self) -> bool {
        self.0.fetch_or(true, Ordering::SeqCst)
    }

    /// Unlock self.
    pub fn unlock(&self) {
        self.0.store(false, Ordering::SeqCst);
    }
}

/// Spawns the timer described in the module-level documentation.
pub fn spawn_state_advance_timer<T: BeaconChainTypes>(
    executor: TaskExecutor,
    beacon_chain: Arc<BeaconChain<T>>,
    log: Logger,
) {
    executor.spawn(
        state_advance_timer(executor.clone(), beacon_chain, log),
        "state_advance_timer",
    );
}

/// Provides the timer described in the module-level documentation.
async fn state_advance_timer<T: BeaconChainTypes>(
    executor: TaskExecutor,
    beacon_chain: Arc<BeaconChain<T>>,
    log: Logger,
) {
    let is_running = Lock::new();
    let slot_clock = &beacon_chain.slot_clock;
    let slot_duration = slot_clock.slot_duration();

    loop {
        match beacon_chain.slot_clock.duration_to_next_slot() {
            Some(duration) => sleep(duration + (slot_duration / 4) * 3).await,
            None => {
                error!(log, "Failed to read slot clock");
                // If we can't read the slot clock, just wait another slot.
                sleep(slot_duration).await;
                continue;
            }
        };

        // Only start spawn the state advance task if the lock was previously free.
        if !is_running.lock() {
            let log = log.clone();
            let beacon_chain = beacon_chain.clone();
            let is_running = is_running.clone();

            executor.spawn_blocking(
                move || {
                    match advance_head(&beacon_chain, &log) {
                        Ok(()) => (),
                        Err(Error::BeaconChain(e)) => error!(
                            log,
                            "Failed to advance head state";
                            "error" => ?e
                        ),
                        Err(Error::StateAlreadyAdvanced { block_root }) => debug!(
                            log,
                            "State already advanced on slot";
                            "block_root" => ?block_root
                        ),
                        Err(Error::MaxDistanceExceeded {
                            current_slot,
                            head_slot,
                        }) => debug!(
                            log,
                            "Refused to advance head state";
                            "head_slot" => head_slot,
                            "current_slot" => current_slot,
                        ),
                        other => warn!(
                            log,
                            "Did not advance head state";
                            "reason" => ?other
                        ),
                    };

                    // Permit this blocking task to spawn again, next time the timer fires.
                    is_running.unlock();
                },
                "state_advance_blocking",
            );
        } else {
            warn!(
                log,
                "State advance routine overloaded";
                "msg" => "system resources may be overloaded"
            )
        }
    }
}

/// Reads the `snapshot_cache` from the `beacon_chain` and attempts to take a clone of the
/// `BeaconState` of the head block. If it obtains this clone, the state will be advanced a single
/// slot then placed back in the `snapshot_cache` to be used for block verification.
///
/// See the module-level documentation for rationale.
fn advance_head<T: BeaconChainTypes>(
    beacon_chain: &BeaconChain<T>,
    log: &Logger,
) -> Result<(), Error> {
    let current_slot = beacon_chain.slot()?;

    // These brackets ensure that the `head_slot` value is dropped before we run fork choice and
    // potentially invalidate it.
    //
    // Fork-choice is not run *before* this function to avoid unnecessary calls whilst syncing.
    {
        let head_slot = beacon_chain.head_info()?.slot;

        // Don't run this when syncing or if lagging too far behind.
        if head_slot + MAX_ADVANCE_DISTANCE < current_slot {
            return Err(Error::MaxDistanceExceeded {
                current_slot,
                head_slot,
            });
        }
    }

    // Run fork choice so we get the latest view of the head.
    //
    // This is useful since it's quite likely that the last time we ran fork choice was shortly
    // after receiving the latest gossip block, but not necessarily after we've received the
    // majority of attestations.
    beacon_chain.fork_choice()?;

    let head_root = beacon_chain.head_info()?.block_root;

    let (head_slot, head_state_root, mut state) = match beacon_chain
        .snapshot_cache
        .try_write_for(BLOCK_PROCESSING_CACHE_LOCK_TIMEOUT)
        .ok_or(BeaconChainError::SnapshotCacheLockTimeout)?
        .get_for_state_advance(head_root)
    {
        StateAdvance::AlreadyAdvanced => {
            return Err(Error::StateAlreadyAdvanced {
                block_root: head_root,
            })
        }
        StateAdvance::BlockNotFound => return Err(Error::HeadMissingFromSnapshotCache(head_root)),
        StateAdvance::State {
            state,
            state_root,
            block_slot,
        } => (block_slot, state_root, *state),
    };

    let initial_slot = state.slot;
    let initial_epoch = state.current_epoch();

    let state_root = if state.slot == head_slot {
        Some(head_state_root)
    } else {
        // Protect against advancing a state more than a single slot.
        //
        // Advancing more than one slot without storing the intermediate state would corrupt the
        // database. Future works might store temporary, intermediate states inside this function.
        return Err(Error::BadStateSlot {
            block_slot: head_slot,
            state_slot: state.slot,
        });
    };

    // Advance the state a single slot.
    if let Some(summary) = per_slot_processing(&mut state, state_root, &beacon_chain.spec)
        .map_err(BeaconChainError::from)?
    {
        // Only notify the validator monitor for recent blocks.
        if state.current_epoch() + VALIDATOR_MONITOR_HISTORIC_EPOCHS as u64
            >= current_slot.epoch(T::EthSpec::slots_per_epoch())
        {
            // Potentially create logs/metrics for locally monitored validators.
            beacon_chain
                .validator_monitor
                .read()
                .process_validator_statuses(state.current_epoch(), &summary.statuses);
        }
    }

    debug!(
        log,
        "Advanced head state one slot";
        "head_root" => ?head_root,
        "state_slot" => state.slot,
        "current_slot" => current_slot,
    );

    // If the advanced state is in a later epoch than where it started, pre-emptively add the
    // proposer shuffling for the new epoch into the cache.
    if state.current_epoch() > initial_epoch {
        debug!(
            log,
            "Priming proposer cache";
            "head_root" => ?head_root,
            "state_epoch" => state.current_epoch(),
            "current_epoch" => current_slot.epoch(T::EthSpec::slots_per_epoch()),
        );
        beacon_chain
            .beacon_proposer_cache
            .lock()
            .insert(
                state.current_epoch(),
                head_root,
                state
                    .get_beacon_proposer_indices(&beacon_chain.spec)
                    .map_err(BeaconChainError::from)?,
                state.fork,
            )
            .map_err(BeaconChainError::from)?;
    }

    let final_slot = state.slot;

    // Insert the advanced state back into the snapshot cache.
    beacon_chain
        .snapshot_cache
        .try_write_for(BLOCK_PROCESSING_CACHE_LOCK_TIMEOUT)
        .ok_or(BeaconChainError::SnapshotCacheLockTimeout)?
        .update_pre_state(head_root, state)
        .ok_or(Error::HeadMissingFromSnapshotCache(head_root))?;

    // If we have moved into the next slot whilst processing the state then this function is going
    // to become ineffective and likely become a hindrance as we're stealing the tree hash cache
    // from the snapshot cache (which may force the next block to rebuild a new one).
    //
    // If this warning occurs very frequently on well-resourced machines then we should consider
    // starting it earlier in the slot. Otherwise, it's a good indication that the machine is too
    // slow/overloaded and will be useful information for the user.
    let starting_slot = current_slot;
    let current_slot = beacon_chain.slot()?;
    if starting_slot < current_slot {
        warn!(
            log,
            "State advance too slow";
            "head_root" => %head_root,
            "advanced_slot" => final_slot,
            "current_slot" => current_slot,
            "starting_slot" => starting_slot,
            "msg" => "system resources may be overloaded",
        );
    }

    debug!(
        log,
        "Completed state advance";
        "head_root" => ?head_root,
        "advanced_slot" => final_slot,
        "initial_slot" => initial_slot,
    );

    Ok(())
}

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

    #[test]
    fn lock() {
        let lock = Lock::new();
        assert_eq!(lock.lock(), false);
        assert_eq!(lock.lock(), true);
        assert_eq!(lock.lock(), true);
        lock.unlock();
        assert_eq!(lock.lock(), false);
        assert_eq!(lock.lock(), true);
    }
}
<|endoftext|>