<commit_before>//! Windows specific definitions
use std::io::{self, Stdout, Write};
use std::mem;
use std::sync::atomic;

use kernel32;
use unicode_width::UnicodeWidthChar;
use winapi;

use config::Config;
use consts::{self, KeyPress};
use error;
use Result;
use line_buffer::LineBuffer;
use super::{Position, RawMode, RawReader, Renderer, Term};

const STDIN_FILENO: winapi::DWORD = winapi::STD_INPUT_HANDLE;
const STDOUT_FILENO: winapi::DWORD = winapi::STD_OUTPUT_HANDLE;

fn get_std_handle(fd: winapi::DWORD) -> Result<winapi::HANDLE> {
    let handle = unsafe { kernel32::GetStdHandle(fd) };
    if handle == winapi::INVALID_HANDLE_VALUE {
        try!(Err(io::Error::last_os_error()));
    } else if handle.is_null() {
        try!(Err(io::Error::new(io::ErrorKind::Other,
                                "no stdio handle available for this process")));
    }
    Ok(handle)
}

#[macro_export]
macro_rules! check {
    ($funcall:expr) => {
        {
        let rc = unsafe { $funcall };
        if rc == 0 {
            try!(Err(io::Error::last_os_error()));
        }
        rc
        }
    };
}

fn get_win_size(handle: winapi::HANDLE) -> (usize, usize) {
    let mut info = unsafe { mem::zeroed() };
    match unsafe { kernel32::GetConsoleScreenBufferInfo(handle, &mut info) } {
        0 => (80, 24),
        _ => (info.dwSize.X as usize, (1 + info.srWindow.Bottom - info.srWindow.Top) as usize),
    }
}

fn get_console_mode(handle: winapi::HANDLE) -> Result<winapi::DWORD> {
    let mut original_mode = 0;
    check!(kernel32::GetConsoleMode(handle, &mut original_mode));
    Ok(original_mode)
}

pub type Mode = ConsoleMode;

#[derive(Clone,Copy,Debug)]
pub struct ConsoleMode {
    original_mode: winapi::DWORD,
    stdin_handle: winapi::HANDLE,
}

impl RawMode for Mode {
    /// Disable RAW mode for the terminal.
    fn disable_raw_mode(&self) -> Result<()> {
        check!(kernel32::SetConsoleMode(self.stdin_handle, self.original_mode));
        Ok(())
    }
}

/// Console input reader
pub struct ConsoleRawReader {
    handle: winapi::HANDLE,
    buf: Option<u16>,
}

impl ConsoleRawReader {
    pub fn new() -> Result<ConsoleRawReader> {
        let handle = try!(get_std_handle(STDIN_FILENO));
        Ok(ConsoleRawReader {
               handle: handle,
               buf: None,
           })
    }
}

impl RawReader for ConsoleRawReader {
    fn next_key(&mut self) -> Result<KeyPress> {
        use std::char::decode_utf16;
        use winapi::{LEFT_ALT_PRESSED, LEFT_CTRL_PRESSED, RIGHT_ALT_PRESSED, RIGHT_CTRL_PRESSED};

        let mut rec: winapi::INPUT_RECORD = unsafe { mem::zeroed() };
        let mut count = 0;
        loop {
            // TODO GetNumberOfConsoleInputEvents
            check!(kernel32::ReadConsoleInputW(self.handle,
                                               &mut rec,
                                               1 as winapi::DWORD,
                                               &mut count));

            if rec.EventType == winapi::WINDOW_BUFFER_SIZE_EVENT {
                SIGWINCH.store(true, atomic::Ordering::SeqCst);
                debug!(target: "rustyline", "SIGWINCH");
                return Err(error::ReadlineError::WindowResize);
            } else if rec.EventType != winapi::KEY_EVENT {
                continue;
            }
            let key_event = unsafe { rec.KeyEvent() };
            // writeln!(io::stderr(), "key_event: {:?}", key_event).unwrap();
            if key_event.bKeyDown == 0 &&
               key_event.wVirtualKeyCode != winapi::VK_MENU as winapi::WORD {
                continue;
            }
            // key_event.wRepeatCount seems to be always set to 1 (maybe because we only
            // read one character at a time)

            // let alt_gr = key_event.dwControlKeyState & (LEFT_CTRL_PRESSED |
            // RIGHT_ALT_PRESSED) != 0;
            let alt = key_event.dwControlKeyState & (LEFT_ALT_PRESSED | RIGHT_ALT_PRESSED) != 0;
            let ctrl = key_event.dwControlKeyState & (LEFT_CTRL_PRESSED | RIGHT_CTRL_PRESSED) != 0;
            let meta = alt;

            let utf16 = key_event.UnicodeChar;
            if utf16 == 0 {
                match key_event.wVirtualKeyCode as i32 {
                    winapi::VK_LEFT => {
                        return Ok(if ctrl {
                                      KeyPress::ControlLeft
                                  } else {
                                      KeyPress::Left
                                  })
                    }
                    winapi::VK_RIGHT => {
                        return Ok(if ctrl {
                                      KeyPress::ControlRight
                                  } else {
                                      KeyPress::Right
                                  })
                    }
                    winapi::VK_UP => {
                        return Ok(if ctrl {
                                      KeyPress::ControlUp
                                  } else {
                                      KeyPress::Up
                                  })
                    }
                    winapi::VK_DOWN => {
                        return Ok(if ctrl {
                                      KeyPress::ControlDown
                                  } else {
                                      KeyPress::Down
                                  })
                    }
                    winapi::VK_DELETE => return Ok(KeyPress::Delete),
                    winapi::VK_HOME => return Ok(KeyPress::Home),
                    winapi::VK_END => return Ok(KeyPress::End),
                    winapi::VK_PRIOR => return Ok(KeyPress::PageUp),
                    winapi::VK_NEXT => return Ok(KeyPress::PageDown),
                    winapi::VK_INSERT => return Ok(KeyPress::Insert),
                    winapi::VK_F1 => return Ok(KeyPress::F(1)),
                    winapi::VK_F2 => return Ok(KeyPress::F(2)),
                    winapi::VK_F3 => return Ok(KeyPress::F(3)),
                    winapi::VK_F4 => return Ok(KeyPress::F(4)),
                    winapi::VK_F5 => return Ok(KeyPress::F(5)),
                    winapi::VK_F6 => return Ok(KeyPress::F(6)),
                    winapi::VK_F7 => return Ok(KeyPress::F(7)),
                    winapi::VK_F8 => return Ok(KeyPress::F(8)),
                    winapi::VK_F9 => return Ok(KeyPress::F(9)),
                    winapi::VK_F10 => return Ok(KeyPress::F(10)),
                    winapi::VK_F11 => return Ok(KeyPress::F(11)),
                    winapi::VK_F12 => return Ok(KeyPress::F(12)),
                    // winapi::VK_BACK is correctly handled because the key_event.UnicodeChar is also
                    // set.
                    _ => continue,
                };
            } else if utf16 == 27 {
                return Ok(KeyPress::Esc);
            } else {
                // TODO How to support surrogate pair ?
                self.buf = Some(utf16);
                let orc = decode_utf16(self).next();
                if orc.is_none() {
                    return Err(error::ReadlineError::Eof);
                }
                let c = try!(orc.unwrap());
                if meta {
                    return Ok(KeyPress::Meta(c));
                } else {
                    return Ok(consts::char_to_key_press(c));
                }
            }
        }
    }
}

impl Iterator for ConsoleRawReader {
    type Item = u16;

    fn next(&mut self) -> Option<u16> {
        let buf = self.buf;
        self.buf = None;
        buf
    }
}

pub struct ConsoleRenderer {
    out: Stdout,
    handle: winapi::HANDLE,
    cols: usize, // Number of columns in terminal
}

impl ConsoleRenderer {
    fn new(handle: winapi::HANDLE) -> ConsoleRenderer {
        let (cols, _) = get_win_size(handle);
        ConsoleRenderer {
            out: io::stdout(),
            handle: handle,
            cols: cols,
        }
    }

    fn get_console_screen_buffer_info(&self) -> Result<winapi::CONSOLE_SCREEN_BUFFER_INFO> {
        let mut info = unsafe { mem::zeroed() };
        check!(kernel32::GetConsoleScreenBufferInfo(self.handle, &mut info));
        Ok(info)
    }

    fn set_console_cursor_position(&mut self, pos: winapi::COORD) -> Result<()> {
        check!(kernel32::SetConsoleCursorPosition(self.handle, pos));
        Ok(())
    }

    fn fill_console_output_character(&mut self,
                                     length: winapi::DWORD,
                                     pos: winapi::COORD)
                                     -> Result<()> {
        let mut _count = 0;
        check!(kernel32::FillConsoleOutputCharacterA(self.handle,
                                                     ' ' as winapi::CHAR,
                                                     length,
                                                     pos,
                                                     &mut _count));
        Ok(())
    }
}

impl Renderer for ConsoleRenderer {
    fn move_cursor(&mut self, old: Position, new: Position) -> Result<()> {
        let mut info = try!(self.get_console_screen_buffer_info());
        if new.row > old.row {
            info.dwCursorPosition.Y += (new.row - old.row) as i16;
        } else {
            info.dwCursorPosition.Y -= (old.row - new.row) as i16;
        }
        if new.col > old.col {
            info.dwCursorPosition.X += (new.col - old.col) as i16;
        } else {
            info.dwCursorPosition.X -= (old.col - new.col) as i16;
        }
        self.set_console_cursor_position(info.dwCursorPosition)
    }

    fn refresh_line(&mut self,
                    prompt: &str,
                    prompt_size: Position,
                    line: &LineBuffer,
                    current_row: usize,
                    old_rows: usize)
                    -> Result<(Position, Position)> {
        // calculate the position of the end of the input line
        let end_pos = self.calculate_position(line, prompt_size);
        // calculate the desired position of the cursor
        let cursor = self.calculate_position(&line[..line.pos()], prompt_size);

        // position at the start of the prompt, clear to end of previous input
        let mut info = try!(self.get_console_screen_buffer_info());
        info.dwCursorPosition.X = 0;
        info.dwCursorPosition.Y -= current_row as i16;
        try!(self.set_console_cursor_position(info.dwCursorPosition));
        let mut _count = 0;
        try!(self.fill_console_output_character((info.dwSize.X * (old_rows as i16 + 1)) as u32,
                                                info.dwCursorPosition));
        let mut ab = String::new();
        // display the prompt
        ab.push_str(prompt); // TODO handle ansi escape code (SetConsoleTextAttribute)
        // display the input line
        ab.push_str(&line);
        try!(self.write_and_flush(ab.as_bytes()));

        // position the cursor
        let mut info = try!(self.get_console_screen_buffer_info());
        info.dwCursorPosition.X = cursor.col as i16;
        info.dwCursorPosition.Y -= (end_pos.row - cursor.row) as i16;
        try!(self.set_console_cursor_position(info.dwCursorPosition));
        Ok((cursor, end_pos))
    }

    fn write_and_flush(&mut self, buf: &[u8]) -> Result<()> {
        try!(self.out.write_all(buf));
        try!(self.out.flush());
        Ok(())
    }

    /// Characters with 2 column width are correctly handled (not splitted).
    fn calculate_position(&self, s: &str, orig: Position) -> Position {
        let mut pos = orig;
        for c in s.chars() {
            let cw = if c == '\n' {
                pos.col = 0;
                pos.row += 1;
                None
            } else {
                c.width()
            };
            if let Some(cw) = cw {
                pos.col += cw;
                if pos.col > self.cols {
                    pos.row += 1;
                    pos.col = cw;
                }
            }
        }
        if pos.col == self.cols {
            pos.col = 0;
            pos.row += 1;
        }
        pos
    }

    /// Clear the screen. Used to handle ctrl+l
    fn clear_screen(&mut self) -> Result<()> {
        let info = try!(self.get_console_screen_buffer_info());
        let coord = winapi::COORD { X: 0, Y: 0 };
        check!(kernel32::SetConsoleCursorPosition(self.handle, coord));
        let mut _count = 0;
        let n = info.dwSize.X as winapi::DWORD * info.dwSize.Y as winapi::DWORD;
        check!(kernel32::FillConsoleOutputCharacterA(self.handle,
                                                     ' ' as winapi::CHAR,
                                                     n,
                                                     coord,
                                                     &mut _count));
        Ok(())
    }

    fn sigwinch(&self) -> bool {
        SIGWINCH.compare_and_swap(true, false, atomic::Ordering::SeqCst)
    }

    /// Try to get the number of columns in the current terminal,
    /// or assume 80 if it fails.
    fn update_size(&mut self) {
        let (cols, _) = get_win_size(self.handle);
        self.cols = cols;
    }

    fn get_columns(&self) -> usize {
        self.cols
    }

    /// Try to get the number of rows in the current terminal,
    /// or assume 24 if it fails.
    fn get_rows(&self) -> usize {
        let (_, rows) = get_win_size(self.handle);
        rows
    }
}

static SIGWINCH: atomic::AtomicBool = atomic::ATOMIC_BOOL_INIT;

pub type Terminal = Console;

#[derive(Clone,Debug)]
pub struct Console {
    stdin_isatty: bool,
    stdin_handle: winapi::HANDLE,
    stdout_handle: winapi::HANDLE,
}

impl Console {}

impl Term for Console {
    type Reader = ConsoleRawReader;
    type Writer = ConsoleRenderer;
    type Mode = Mode;

    fn new() -> Console {
        use std::ptr;
        let stdin_handle = get_std_handle(STDIN_FILENO);
        let stdin_isatty = match stdin_handle {
            Ok(handle) => {
                // If this function doesn't fail then fd is a TTY
                get_console_mode(handle).is_ok()
            }
            Err(_) => false,
        };

        let stdout_handle = get_std_handle(STDOUT_FILENO).unwrap_or(ptr::null_mut());
        Console {
            stdin_isatty: stdin_isatty,
            stdin_handle: stdin_handle.unwrap_or(ptr::null_mut()),
            stdout_handle: stdout_handle,
        }
    }

    /// Checking for an unsupported TERM in windows is a no-op
    fn is_unsupported(&self) -> bool {
        false
    }

    fn is_stdin_tty(&self) -> bool {
        self.stdin_isatty
    }

    // pub fn install_sigwinch_handler(&mut self) {
    // See ReadConsoleInputW && WINDOW_BUFFER_SIZE_EVENT
    // }

    /// Enable RAW mode for the terminal.
    fn enable_raw_mode(&self) -> Result<Mode> {
        if !self.stdin_isatty {
            try!(Err(io::Error::new(io::ErrorKind::Other,
                                    "no stdio handle available for this process")));
        }
        let original_mode = try!(get_console_mode(self.stdin_handle));
        // Disable these modes
        let raw = original_mode &
                  !(winapi::wincon::ENABLE_LINE_INPUT | winapi::wincon::ENABLE_ECHO_INPUT |
                    winapi::wincon::ENABLE_PROCESSED_INPUT);
        // Enable these modes
        let raw = raw | winapi::wincon::ENABLE_EXTENDED_FLAGS;
        let raw = raw | winapi::wincon::ENABLE_INSERT_MODE;
        let raw = raw | winapi::wincon::ENABLE_QUICK_EDIT_MODE;
        let raw = raw | winapi::wincon::ENABLE_WINDOW_INPUT;
        check!(kernel32::SetConsoleMode(self.stdin_handle, raw));
        Ok(Mode {
               original_mode: original_mode,
               stdin_handle: self.stdin_handle,
           })
    }

    fn create_reader(&self, _: &Config) -> Result<ConsoleRawReader> {
        ConsoleRawReader::new()
    }

    fn create_writer(&self) -> ConsoleRenderer {
        ConsoleRenderer::new(self.stdout_handle)
    }
}
<commit_msg>Not sure how to compute the number of columns<commit_after>//! Windows specific definitions
use std::io::{self, Stdout, Write};
use std::mem;
use std::sync::atomic;

use kernel32;
use unicode_width::UnicodeWidthChar;
use winapi;

use config::Config;
use consts::{self, KeyPress};
use error;
use Result;
use line_buffer::LineBuffer;
use super::{Position, RawMode, RawReader, Renderer, Term};

const STDIN_FILENO: winapi::DWORD = winapi::STD_INPUT_HANDLE;
const STDOUT_FILENO: winapi::DWORD = winapi::STD_OUTPUT_HANDLE;

fn get_std_handle(fd: winapi::DWORD) -> Result<winapi::HANDLE> {
    let handle = unsafe { kernel32::GetStdHandle(fd) };
    if handle == winapi::INVALID_HANDLE_VALUE {
        try!(Err(io::Error::last_os_error()));
    } else if handle.is_null() {
        try!(Err(io::Error::new(io::ErrorKind::Other,
                                "no stdio handle available for this process")));
    }
    Ok(handle)
}

#[macro_export]
macro_rules! check {
    ($funcall:expr) => {
        {
        let rc = unsafe { $funcall };
        if rc == 0 {
            try!(Err(io::Error::last_os_error()));
        }
        rc
        }
    };
}

fn get_win_size(handle: winapi::HANDLE) -> (usize, usize) {
    let mut info = unsafe { mem::zeroed() };
    match unsafe { kernel32::GetConsoleScreenBufferInfo(handle, &mut info) } {
        0 => (80, 24),
        _ => (info.dwSize.X as usize, (1 + info.srWindow.Bottom - info.srWindow.Top) as usize), // (info.srWindow.Right - info.srWindow.Left + 1)
    }
}

fn get_console_mode(handle: winapi::HANDLE) -> Result<winapi::DWORD> {
    let mut original_mode = 0;
    check!(kernel32::GetConsoleMode(handle, &mut original_mode));
    Ok(original_mode)
}

pub type Mode = ConsoleMode;

#[derive(Clone,Copy,Debug)]
pub struct ConsoleMode {
    original_mode: winapi::DWORD,
    stdin_handle: winapi::HANDLE,
}

impl RawMode for Mode {
    /// Disable RAW mode for the terminal.
    fn disable_raw_mode(&self) -> Result<()> {
        check!(kernel32::SetConsoleMode(self.stdin_handle, self.original_mode));
        Ok(())
    }
}

/// Console input reader
pub struct ConsoleRawReader {
    handle: winapi::HANDLE,
    buf: Option<u16>,
}

impl ConsoleRawReader {
    pub fn new() -> Result<ConsoleRawReader> {
        let handle = try!(get_std_handle(STDIN_FILENO));
        Ok(ConsoleRawReader {
               handle: handle,
               buf: None,
           })
    }
}

impl RawReader for ConsoleRawReader {
    fn next_key(&mut self) -> Result<KeyPress> {
        use std::char::decode_utf16;
        use winapi::{LEFT_ALT_PRESSED, LEFT_CTRL_PRESSED, RIGHT_ALT_PRESSED, RIGHT_CTRL_PRESSED};

        let mut rec: winapi::INPUT_RECORD = unsafe { mem::zeroed() };
        let mut count = 0;
        loop {
            // TODO GetNumberOfConsoleInputEvents
            check!(kernel32::ReadConsoleInputW(self.handle,
                                               &mut rec,
                                               1 as winapi::DWORD,
                                               &mut count));

            if rec.EventType == winapi::WINDOW_BUFFER_SIZE_EVENT {
                SIGWINCH.store(true, atomic::Ordering::SeqCst);
                debug!(target: "rustyline", "SIGWINCH");
                return Err(error::ReadlineError::WindowResize);
            } else if rec.EventType != winapi::KEY_EVENT {
                continue;
            }
            let key_event = unsafe { rec.KeyEvent() };
            // writeln!(io::stderr(), "key_event: {:?}", key_event).unwrap();
            if key_event.bKeyDown == 0 &&
               key_event.wVirtualKeyCode != winapi::VK_MENU as winapi::WORD {
                continue;
            }
            // key_event.wRepeatCount seems to be always set to 1 (maybe because we only
            // read one character at a time)

            // let alt_gr = key_event.dwControlKeyState & (LEFT_CTRL_PRESSED |
            // RIGHT_ALT_PRESSED) != 0;
            let alt = key_event.dwControlKeyState & (LEFT_ALT_PRESSED | RIGHT_ALT_PRESSED) != 0;
            let ctrl = key_event.dwControlKeyState & (LEFT_CTRL_PRESSED | RIGHT_CTRL_PRESSED) != 0;
            let meta = alt;

            let utf16 = key_event.UnicodeChar;
            if utf16 == 0 {
                match key_event.wVirtualKeyCode as i32 {
                    winapi::VK_LEFT => {
                        return Ok(if ctrl {
                                      KeyPress::ControlLeft
                                  } else {
                                      KeyPress::Left
                                  })
                    }
                    winapi::VK_RIGHT => {
                        return Ok(if ctrl {
                                      KeyPress::ControlRight
                                  } else {
                                      KeyPress::Right
                                  })
                    }
                    winapi::VK_UP => {
                        return Ok(if ctrl {
                                      KeyPress::ControlUp
                                  } else {
                                      KeyPress::Up
                                  })
                    }
                    winapi::VK_DOWN => {
                        return Ok(if ctrl {
                                      KeyPress::ControlDown
                                  } else {
                                      KeyPress::Down
                                  })
                    }
                    winapi::VK_DELETE => return Ok(KeyPress::Delete),
                    winapi::VK_HOME => return Ok(KeyPress::Home),
                    winapi::VK_END => return Ok(KeyPress::End),
                    winapi::VK_PRIOR => return Ok(KeyPress::PageUp),
                    winapi::VK_NEXT => return Ok(KeyPress::PageDown),
                    winapi::VK_INSERT => return Ok(KeyPress::Insert),
                    winapi::VK_F1 => return Ok(KeyPress::F(1)),
                    winapi::VK_F2 => return Ok(KeyPress::F(2)),
                    winapi::VK_F3 => return Ok(KeyPress::F(3)),
                    winapi::VK_F4 => return Ok(KeyPress::F(4)),
                    winapi::VK_F5 => return Ok(KeyPress::F(5)),
                    winapi::VK_F6 => return Ok(KeyPress::F(6)),
                    winapi::VK_F7 => return Ok(KeyPress::F(7)),
                    winapi::VK_F8 => return Ok(KeyPress::F(8)),
                    winapi::VK_F9 => return Ok(KeyPress::F(9)),
                    winapi::VK_F10 => return Ok(KeyPress::F(10)),
                    winapi::VK_F11 => return Ok(KeyPress::F(11)),
                    winapi::VK_F12 => return Ok(KeyPress::F(12)),
                    // winapi::VK_BACK is correctly handled because the key_event.UnicodeChar is also
                    // set.
                    _ => continue,
                };
            } else if utf16 == 27 {
                return Ok(KeyPress::Esc);
            } else {
                // TODO How to support surrogate pair ?
                self.buf = Some(utf16);
                let orc = decode_utf16(self).next();
                if orc.is_none() {
                    return Err(error::ReadlineError::Eof);
                }
                let c = try!(orc.unwrap());
                if meta {
                    return Ok(KeyPress::Meta(c));
                } else {
                    return Ok(consts::char_to_key_press(c));
                }
            }
        }
    }
}

impl Iterator for ConsoleRawReader {
    type Item = u16;

    fn next(&mut self) -> Option<u16> {
        let buf = self.buf;
        self.buf = None;
        buf
    }
}

pub struct ConsoleRenderer {
    out: Stdout,
    handle: winapi::HANDLE,
    cols: usize, // Number of columns in terminal
}

impl ConsoleRenderer {
    fn new(handle: winapi::HANDLE) -> ConsoleRenderer {
        let (cols, _) = get_win_size(handle);
        ConsoleRenderer {
            out: io::stdout(),
            handle: handle,
            cols: cols,
        }
    }

    fn get_console_screen_buffer_info(&self) -> Result<winapi::CONSOLE_SCREEN_BUFFER_INFO> {
        let mut info = unsafe { mem::zeroed() };
        check!(kernel32::GetConsoleScreenBufferInfo(self.handle, &mut info));
        Ok(info)
    }

    fn set_console_cursor_position(&mut self, pos: winapi::COORD) -> Result<()> {
        check!(kernel32::SetConsoleCursorPosition(self.handle, pos));
        Ok(())
    }

    fn fill_console_output_character(&mut self,
                                     length: winapi::DWORD,
                                     pos: winapi::COORD)
                                     -> Result<()> {
        let mut _count = 0;
        check!(kernel32::FillConsoleOutputCharacterA(self.handle,
                                                     ' ' as winapi::CHAR,
                                                     length,
                                                     pos,
                                                     &mut _count));
        Ok(())
    }
}

impl Renderer for ConsoleRenderer {
    fn move_cursor(&mut self, old: Position, new: Position) -> Result<()> {
        let mut info = try!(self.get_console_screen_buffer_info());
        if new.row > old.row {
            info.dwCursorPosition.Y += (new.row - old.row) as i16;
        } else {
            info.dwCursorPosition.Y -= (old.row - new.row) as i16;
        }
        if new.col > old.col {
            info.dwCursorPosition.X += (new.col - old.col) as i16;
        } else {
            info.dwCursorPosition.X -= (old.col - new.col) as i16;
        }
        self.set_console_cursor_position(info.dwCursorPosition)
    }

    fn refresh_line(&mut self,
                    prompt: &str,
                    prompt_size: Position,
                    line: &LineBuffer,
                    current_row: usize,
                    old_rows: usize)
                    -> Result<(Position, Position)> {
        // calculate the position of the end of the input line
        let end_pos = self.calculate_position(line, prompt_size);
        // calculate the desired position of the cursor
        let cursor = self.calculate_position(&line[..line.pos()], prompt_size);

        // position at the start of the prompt, clear to end of previous input
        let mut info = try!(self.get_console_screen_buffer_info());
        info.dwCursorPosition.X = 0;
        info.dwCursorPosition.Y -= current_row as i16;
        try!(self.set_console_cursor_position(info.dwCursorPosition));
        let mut _count = 0;
        try!(self.fill_console_output_character((info.dwSize.X * (old_rows as i16 + 1)) as u32,
                                                info.dwCursorPosition));
        let mut ab = String::new();
        // display the prompt
        ab.push_str(prompt); // TODO handle ansi escape code (SetConsoleTextAttribute)
        // display the input line
        ab.push_str(&line);
        try!(self.write_and_flush(ab.as_bytes()));

        // position the cursor
        let mut info = try!(self.get_console_screen_buffer_info());
        info.dwCursorPosition.X = cursor.col as i16;
        info.dwCursorPosition.Y -= (end_pos.row - cursor.row) as i16;
        try!(self.set_console_cursor_position(info.dwCursorPosition));
        Ok((cursor, end_pos))
    }

    fn write_and_flush(&mut self, buf: &[u8]) -> Result<()> {
        try!(self.out.write_all(buf));
        try!(self.out.flush());
        Ok(())
    }

    /// Characters with 2 column width are correctly handled (not splitted).
    fn calculate_position(&self, s: &str, orig: Position) -> Position {
        let mut pos = orig;
        for c in s.chars() {
            let cw = if c == '\n' {
                pos.col = 0;
                pos.row += 1;
                None
            } else {
                c.width()
            };
            if let Some(cw) = cw {
                pos.col += cw;
                if pos.col > self.cols {
                    pos.row += 1;
                    pos.col = cw;
                }
            }
        }
        if pos.col == self.cols {
            pos.col = 0;
            pos.row += 1;
        }
        pos
    }

    /// Clear the screen. Used to handle ctrl+l
    fn clear_screen(&mut self) -> Result<()> {
        let info = try!(self.get_console_screen_buffer_info());
        let coord = winapi::COORD { X: 0, Y: 0 };
        check!(kernel32::SetConsoleCursorPosition(self.handle, coord));
        let mut _count = 0;
        let n = info.dwSize.X as winapi::DWORD * info.dwSize.Y as winapi::DWORD;
        check!(kernel32::FillConsoleOutputCharacterA(self.handle,
                                                     ' ' as winapi::CHAR,
                                                     n,
                                                     coord,
                                                     &mut _count));
        Ok(())
    }

    fn sigwinch(&self) -> bool {
        SIGWINCH.compare_and_swap(true, false, atomic::Ordering::SeqCst)
    }

    /// Try to get the number of columns in the current terminal,
    /// or assume 80 if it fails.
    fn update_size(&mut self) {
        let (cols, _) = get_win_size(self.handle);
        self.cols = cols;
    }

    fn get_columns(&self) -> usize {
        self.cols
    }

    /// Try to get the number of rows in the current terminal,
    /// or assume 24 if it fails.
    fn get_rows(&self) -> usize {
        let (_, rows) = get_win_size(self.handle);
        rows
    }
}

static SIGWINCH: atomic::AtomicBool = atomic::ATOMIC_BOOL_INIT;

pub type Terminal = Console;

#[derive(Clone,Debug)]
pub struct Console {
    stdin_isatty: bool,
    stdin_handle: winapi::HANDLE,
    stdout_handle: winapi::HANDLE,
}

impl Console {}

impl Term for Console {
    type Reader = ConsoleRawReader;
    type Writer = ConsoleRenderer;
    type Mode = Mode;

    fn new() -> Console {
        use std::ptr;
        let stdin_handle = get_std_handle(STDIN_FILENO);
        let stdin_isatty = match stdin_handle {
            Ok(handle) => {
                // If this function doesn't fail then fd is a TTY
                get_console_mode(handle).is_ok()
            }
            Err(_) => false,
        };

        let stdout_handle = get_std_handle(STDOUT_FILENO).unwrap_or(ptr::null_mut());
        Console {
            stdin_isatty: stdin_isatty,
            stdin_handle: stdin_handle.unwrap_or(ptr::null_mut()),
            stdout_handle: stdout_handle,
        }
    }

    /// Checking for an unsupported TERM in windows is a no-op
    fn is_unsupported(&self) -> bool {
        false
    }

    fn is_stdin_tty(&self) -> bool {
        self.stdin_isatty
    }

    // pub fn install_sigwinch_handler(&mut self) {
    // See ReadConsoleInputW && WINDOW_BUFFER_SIZE_EVENT
    // }

    /// Enable RAW mode for the terminal.
    fn enable_raw_mode(&self) -> Result<Mode> {
        if !self.stdin_isatty {
            try!(Err(io::Error::new(io::ErrorKind::Other,
                                    "no stdio handle available for this process")));
        }
        let original_mode = try!(get_console_mode(self.stdin_handle));
        // Disable these modes
        let raw = original_mode &
                  !(winapi::wincon::ENABLE_LINE_INPUT | winapi::wincon::ENABLE_ECHO_INPUT |
                    winapi::wincon::ENABLE_PROCESSED_INPUT);
        // Enable these modes
        let raw = raw | winapi::wincon::ENABLE_EXTENDED_FLAGS;
        let raw = raw | winapi::wincon::ENABLE_INSERT_MODE;
        let raw = raw | winapi::wincon::ENABLE_QUICK_EDIT_MODE;
        let raw = raw | winapi::wincon::ENABLE_WINDOW_INPUT;
        check!(kernel32::SetConsoleMode(self.stdin_handle, raw));
        Ok(Mode {
               original_mode: original_mode,
               stdin_handle: self.stdin_handle,
           })
    }

    fn create_reader(&self, _: &Config) -> Result<ConsoleRawReader> {
        ConsoleRawReader::new()
    }

    fn create_writer(&self) -> ConsoleRenderer {
        ConsoleRenderer::new(self.stdout_handle)
    }
}
<|endoftext|>


<commit_before>use analysis::symbols;
use regex::{Captures, Regex};

const LANGUAGE_SEP_BEGIN : &'static str = "<!-- language=\"";
const LANGUAGE_SEP_END : &'static str = "\" -->";
const LANGUAGE_BLOCK_BEGIN : &'static str = "|[";
const LANGUAGE_BLOCK_END : &'static str = "\n]|";

pub fn reformat_doc(input: &str, symbols: &symbols::Info) -> String {
    code_blocks_transformation(input, symbols)
}

fn try_split<'a>(src: &'a str, needle: &str) -> (&'a str, Option<&'a str>) {
    match src.find(needle) {
        Some(pos) => (&src[..pos], Some(&src[pos + needle.len()..])),
        None => (src, None),
    }
}

fn code_blocks_transformation(mut input: &str, symbols: &symbols::Info) -> String {
    let mut out = String::with_capacity(input.len());

    loop {
        input = match try_split(input, LANGUAGE_BLOCK_BEGIN) {
            (before, Some(after)) => {
                out.push_str(&format(before, symbols));
                if let (before, Some(after)) = try_split(get_language(after, &mut out),
                                                         LANGUAGE_BLOCK_END) {
                    out.push_str(before);
                    out.push_str("\n```");
                    after
                } else {
                    after
                }
            }
            (before, None) => {
                out.push_str(&format(before, symbols));
                return out
            }
        };
    }
}

fn get_language<'a>(entry: &'a str, out: &mut String) -> &'a str {
    if let (_, Some(after)) = try_split(entry, LANGUAGE_SEP_BEGIN) {
        if let (before, Some(after)) = try_split(after, LANGUAGE_SEP_END) {
            out.push_str(&format!("\n```{}", before));
            return after
        }
    }
    out.push_str("\n```text");
    entry
}

lazy_static! {
    static ref SYMBOL: Regex = Regex::new(r"(^|[^\\])[@#%]([\w]+\b)([:.]+[\w_-]+\b)?") .unwrap();
    static ref FUNCTION: Regex = Regex::new(r"(\b[a-z0-9_]+)\(\)") .unwrap();
    static ref GDK_GTK: Regex = Regex::new(r"(G[dt]k[A-Z][\w]+\b)").unwrap();
    static ref SPACES: Regex = Regex::new(r"[ ][ ]+").unwrap();
}

fn format(mut input: &str, symbols: &symbols::Info) -> String {
    let mut ret = String::with_capacity(input.len());
    loop {
        let (before, after) = try_split(input, "`");
        ret.push_str(&replace_c_types(before, symbols));
        if let Some(after) = after {
            ret.push_str("`");
            let (before, after) = try_split(after, "`");
            // don't touch anything enclosed in backticks
            ret.push_str(before);
            if let Some(after) = after {
                ret.push_str("`");
                input = after;
            } else {
                return ret;
            }
        } else {
            return ret;
        }
    }
}

fn replace_c_types(entry: &str, symbols: &symbols::Info) -> String {
    let lookup = |s: &str| -> String {
        symbols.by_c_name(s)
            .map(|s| s.full_rust_name())
            .unwrap_or(s.into())
    };
    let out = SYMBOL.replace_all(entry, |caps: &Captures| {
        format!("{}`{}{}`", &caps[1], lookup(&caps[2]), caps.at(3).unwrap_or(""))
    });
    let out = FUNCTION.replace_all(&out, |caps: &Captures| {
        format!("`{}`", lookup(&caps[1]))
    });
    let out = GDK_GTK.replace_all(&out, |caps: &Captures| {
        format!("`{}`", lookup(&caps[1]))
    });
    SPACES.replace_all(&out, " ")
}
<commit_msg>Shield XML tags<commit_after>use analysis::symbols;
use regex::{Captures, Regex};

const LANGUAGE_SEP_BEGIN : &'static str = "<!-- language=\"";
const LANGUAGE_SEP_END : &'static str = "\" -->";
const LANGUAGE_BLOCK_BEGIN : &'static str = "|[";
const LANGUAGE_BLOCK_END : &'static str = "\n]|";

pub fn reformat_doc(input: &str, symbols: &symbols::Info) -> String {
    code_blocks_transformation(input, symbols)
}

fn try_split<'a>(src: &'a str, needle: &str) -> (&'a str, Option<&'a str>) {
    match src.find(needle) {
        Some(pos) => (&src[..pos], Some(&src[pos + needle.len()..])),
        None => (src, None),
    }
}

fn code_blocks_transformation(mut input: &str, symbols: &symbols::Info) -> String {
    let mut out = String::with_capacity(input.len());

    loop {
        input = match try_split(input, LANGUAGE_BLOCK_BEGIN) {
            (before, Some(after)) => {
                out.push_str(&format(before, symbols));
                if let (before, Some(after)) = try_split(get_language(after, &mut out),
                                                         LANGUAGE_BLOCK_END) {
                    out.push_str(before);
                    out.push_str("\n```");
                    after
                } else {
                    after
                }
            }
            (before, None) => {
                out.push_str(&format(before, symbols));
                return out
            }
        };
    }
}

fn get_language<'a>(entry: &'a str, out: &mut String) -> &'a str {
    if let (_, Some(after)) = try_split(entry, LANGUAGE_SEP_BEGIN) {
        if let (before, Some(after)) = try_split(after, LANGUAGE_SEP_END) {
            out.push_str(&format!("\n```{}", before));
            return after
        }
    }
    out.push_str("\n```text");
    entry
}

lazy_static! {
    static ref SYMBOL: Regex = Regex::new(r"(^|[^\\])[@#%]([\w]+\b)([:.]+[\w_-]+\b)?") .unwrap();
    static ref FUNCTION: Regex = Regex::new(r"(\b[a-z0-9_]+)\(\)") .unwrap();
    static ref GDK_GTK: Regex = Regex::new(r"G[dt]k[A-Z][\w]+\b").unwrap();
    static ref TAGS: Regex = Regex::new(r"<[\w/-]+>").unwrap();
    static ref SPACES: Regex = Regex::new(r"[ ][ ]+").unwrap();
}

fn format(mut input: &str, symbols: &symbols::Info) -> String {
    let mut ret = String::with_capacity(input.len());
    loop {
        let (before, after) = try_split(input, "`");
        ret.push_str(&replace_c_types(before, symbols));
        if let Some(after) = after {
            ret.push_str("`");
            let (before, after) = try_split(after, "`");
            // don't touch anything enclosed in backticks
            ret.push_str(before);
            if let Some(after) = after {
                ret.push_str("`");
                input = after;
            } else {
                return ret;
            }
        } else {
            return ret;
        }
    }
}

fn replace_c_types(entry: &str, symbols: &symbols::Info) -> String {
    let lookup = |s: &str| -> String {
        symbols.by_c_name(s)
            .map(|s| s.full_rust_name())
            .unwrap_or(s.into())
    };
    let out = SYMBOL.replace_all(entry, |caps: &Captures| {
        format!("{}`{}{}`", &caps[1], lookup(&caps[2]), caps.at(3).unwrap_or(""))
    });
    let out = FUNCTION.replace_all(&out, |caps: &Captures| {
        format!("`{}`", lookup(&caps[1]))
    });
    let out = GDK_GTK.replace_all(&out, |caps: &Captures| {
        format!("`{}`", lookup(&caps[0]))
    });
    let out = TAGS.replace_all(&out, "`$0`");
    SPACES.replace_all(&out, " ")
}
<|endoftext|>


<commit_before>import option;
import str;
import io::reader_util;

import common::config;

export test_props;
export load_props;
export is_test_ignored;

type test_props = {
    // Lines that should be expected, in order, on standard out
    error_patterns: [str],
    // Extra flags to pass to the compiler
    compile_flags: option<str>,
    // If present, the name of a file that this test should match when
    // pretty-printed
    pp_exact: option<str>,
    // Modules from aux directory that should be compiled
    aux_builds: [str],
    // Environment settings to use during execution
    exec_env: [(str,str)]
};

// Load any test directives embedded in the file
fn load_props(testfile: str) -> test_props {
    let mut error_patterns = [];
    let mut aux_builds = [];
    let mut exec_env = [];
    let mut compile_flags = option::none;
    let mut pp_exact = option::none;
    iter_header(testfile) {|ln|
        alt parse_error_pattern(ln) {
          option::some(ep) { error_patterns += [ep]; }
          option::none { }
        };

        if option::is_none(compile_flags) {
            compile_flags = parse_compile_flags(ln);
        }

        if option::is_none(pp_exact) {
            pp_exact = parse_pp_exact(ln, testfile);
        }

        option::iter(parse_aux_build(ln)) {|ab|
            aux_builds += [ab];
        }

        option::iter(parse_exec_env(ln)) {|ee|
            exec_env += [ee];
        }
    };
    ret {
        error_patterns: error_patterns,
        compile_flags: compile_flags,
        pp_exact: pp_exact,
        aux_builds: aux_builds,
        exec_env: exec_env
    };
}

fn is_test_ignored(config: config, testfile: str) -> bool {
    let mut found = false;
    iter_header(testfile) {|ln|
        // FIXME: Can't return or break from iterator
        // (Fix when Issue #1619 is resolved)
        found = found || parse_name_directive(ln, "xfail-test");
        found = found || parse_name_directive(ln, xfail_target());
        if (config.mode == common::mode_pretty) {
            found = found || parse_name_directive(ln, "xfail-pretty");
        }
    };
    ret found;

    fn xfail_target() -> str {
        "xfail-" + os::sysname()
    }
}

fn iter_header(testfile: str, it: fn(str)) {
    let rdr = result::get(io::file_reader(testfile));
    while !rdr.eof() {
        let ln = rdr.read_line();

        // Assume that any directives will be found before the first
        // module or function. This doesn't seem to be an optimization
        // with a warm page cache. Maybe with a cold one.
        if str::starts_with(ln, "fn")
            || str::starts_with(ln, "mod") {
            break;
        } else { it(ln); }
    }
}

fn parse_error_pattern(line: str) -> option<str> {
    parse_name_value_directive(line, "error-pattern")
}

fn parse_aux_build(line: str) -> option<str> {
    parse_name_value_directive(line, "aux-build")
}

fn parse_compile_flags(line: str) -> option<str> {
    parse_name_value_directive(line, "compile-flags")
}

fn parse_exec_env(line: str) -> option<(str, str)> {
    parse_name_value_directive(line, "exec-env").map {|nv|
        // nv is either FOO or FOO=BAR
        let strs = str::splitn_char(nv, '=', 1u);
        alt strs.len() {
          1u { (strs[0], "") }
          2u { (strs[0], strs[1]) }
          n { fail #fmt["Expected 1 or 2 strings, not %u", n]; }
        }
    }
}

fn parse_pp_exact(line: str, testfile: str) -> option<str> {
    alt parse_name_value_directive(line, "pp-exact") {
      option::some(s) { option::some(s) }
      option::none {
        if parse_name_directive(line, "pp-exact") {
            option::some(path::basename(testfile))
        } else {
            option::none
        }
      }
    }
}

fn parse_name_directive(line: str, directive: str) -> bool {
    str::contains(line, directive)
}

fn parse_name_value_directive(line: str,
                              directive: str) -> option<str> unsafe {
    let keycolon = directive + ":";
    alt str::find_str(line, keycolon) {
        option::some(colon) {
            let value = str::slice(line, colon + str::len(keycolon),
                                   str::len(line));
            #debug("%s: %s", directive,  value);
            option::some(value)
        }
        option::none { option::none }
    }
}
<commit_msg>Fix compiletest::header::iter_header to work with new-style for loops<commit_after>import option;
import str;
import io::reader_util;

import common::config;

export test_props;
export load_props;
export is_test_ignored;

type test_props = {
    // Lines that should be expected, in order, on standard out
    error_patterns: [str],
    // Extra flags to pass to the compiler
    compile_flags: option<str>,
    // If present, the name of a file that this test should match when
    // pretty-printed
    pp_exact: option<str>,
    // Modules from aux directory that should be compiled
    aux_builds: [str],
    // Environment settings to use during execution
    exec_env: [(str,str)]
};

// Load any test directives embedded in the file
fn load_props(testfile: str) -> test_props {
    let mut error_patterns = [];
    let mut aux_builds = [];
    let mut exec_env = [];
    let mut compile_flags = option::none;
    let mut pp_exact = option::none;
    for iter_header(testfile) {|ln|
        alt parse_error_pattern(ln) {
          option::some(ep) { error_patterns += [ep]; }
          option::none { }
        };

        if option::is_none(compile_flags) {
            compile_flags = parse_compile_flags(ln);
        }

        if option::is_none(pp_exact) {
            pp_exact = parse_pp_exact(ln, testfile);
        }

        option::iter(parse_aux_build(ln)) {|ab|
            aux_builds += [ab];
        }

        option::iter(parse_exec_env(ln)) {|ee|
            exec_env += [ee];
        }
    };
    ret {
        error_patterns: error_patterns,
        compile_flags: compile_flags,
        pp_exact: pp_exact,
        aux_builds: aux_builds,
        exec_env: exec_env
    };
}

fn is_test_ignored(config: config, testfile: str) -> bool {
    let mut found = false;
    for iter_header(testfile) {|ln|
        if parse_name_directive(ln, "xfail-test") { ret true; }
        if parse_name_directive(ln, xfail_target()) { ret true; }
        if config.mode == common::mode_pretty &&
           parse_name_directive(ln, "xfail-pretty") { ret true; }
    };
    ret found;

    fn xfail_target() -> str {
        "xfail-" + os::sysname()
    }
}

fn iter_header(testfile: str, it: fn(str) -> bool) -> bool {
    let rdr = result::get(io::file_reader(testfile));
    while !rdr.eof() {
        let ln = rdr.read_line();

        // Assume that any directives will be found before the first
        // module or function. This doesn't seem to be an optimization
        // with a warm page cache. Maybe with a cold one.
        if str::starts_with(ln, "fn")
            || str::starts_with(ln, "mod") {
            ret false;
        } else { if !(it(ln)) { ret false; } }
    }
    ret true;
}

fn parse_error_pattern(line: str) -> option<str> {
    parse_name_value_directive(line, "error-pattern")
}

fn parse_aux_build(line: str) -> option<str> {
    parse_name_value_directive(line, "aux-build")
}

fn parse_compile_flags(line: str) -> option<str> {
    parse_name_value_directive(line, "compile-flags")
}

fn parse_exec_env(line: str) -> option<(str, str)> {
    parse_name_value_directive(line, "exec-env").map {|nv|
        // nv is either FOO or FOO=BAR
        let strs = str::splitn_char(nv, '=', 1u);
        alt strs.len() {
          1u { (strs[0], "") }
          2u { (strs[0], strs[1]) }
          n { fail #fmt["Expected 1 or 2 strings, not %u", n]; }
        }
    }
}

fn parse_pp_exact(line: str, testfile: str) -> option<str> {
    alt parse_name_value_directive(line, "pp-exact") {
      option::some(s) { option::some(s) }
      option::none {
        if parse_name_directive(line, "pp-exact") {
            option::some(path::basename(testfile))
        } else {
            option::none
        }
      }
    }
}

fn parse_name_directive(line: str, directive: str) -> bool {
    str::contains(line, directive)
}

fn parse_name_value_directive(line: str,
                              directive: str) -> option<str> unsafe {
    let keycolon = directive + ":";
    alt str::find_str(line, keycolon) {
        option::some(colon) {
            let value = str::slice(line, colon + str::len(keycolon),
                                   str::len(line));
            #debug("%s: %s", directive,  value);
            option::some(value)
        }
        option::none { option::none }
    }
}
<|endoftext|>


<commit_before>use clippy_utils::diagnostics::span_lint_and_sugg;
use clippy_utils::source::snippet_with_macro_callsite;
use clippy_utils::{get_parent_expr_for_hir, in_macro, spans_on_same_line, sugg};
use if_chain::if_chain;
use rustc_errors::Applicability;
use rustc_hir::{Block, BlockCheckMode, ExprKind};
use rustc_lint::{LateContext, LateLintPass};
use rustc_session::{declare_lint_pass, declare_tool_lint};

declare_clippy_lint! {
    /// **What it does:** Looks for blocks of expressions and fires if the last expression returns
    /// `()` but is not followed by a semicolon.
    ///
    /// **Why is this bad?** The semicolon might be optional but when extending the block with new
    /// code, it doesn't require a change in previous last line.
    ///
    /// **Known problems:** None.
    ///
    /// **Example:**
    ///
    /// ```rust
    /// fn main() {
    ///     println!("Hello world")
    /// }
    /// ```
    /// Use instead:
    /// ```rust
    /// fn main() {
    ///     println!("Hello world");
    /// }
    /// ```
    pub SEMICOLON_IF_NOTHING_RETURNED,
    pedantic,
    "add a semicolon if nothing is returned"
}

declare_lint_pass!(SemicolonIfNothingReturned => [SEMICOLON_IF_NOTHING_RETURNED]);

impl LateLintPass<'_> for SemicolonIfNothingReturned {
    fn check_block(&mut self, cx: &LateContext<'tcx>, block: &'tcx Block<'tcx>) {
        if_chain! {
            if !in_macro(block.span);
            if let Some(expr) = block.expr;
            let t_expr = cx.typeck_results().expr_ty(expr);
            if t_expr.is_unit();
            if let snippet = snippet_with_macro_callsite(cx, expr.span, "}");
            if !snippet.ends_with('}');
            then {
                // check if the block is inside a closure or an unsafe block and don't
                // emit if the block is on the same line
                if_chain! {
                    if let Some(parent) = get_parent_expr_for_hir(cx, block.hir_id);

                    if !matches!(block.rules, BlockCheckMode::DefaultBlock) ||
                    matches!(parent.kind, ExprKind::Closure(..) | ExprKind::Block(..));

                    if block.stmts.len() == 0;

                    if spans_on_same_line(cx, parent.span, expr.span);
                    then {
                        return;
                    }
                }

                // filter out the desugared `for` loop
                if let ExprKind::DropTemps(..) = &expr.kind {
                    return;
                }

                let sugg = sugg::Sugg::hir_with_macro_callsite(cx, expr, "..");
                let suggestion = format!("{0};", sugg);
                span_lint_and_sugg(
                    cx,
                    SEMICOLON_IF_NOTHING_RETURNED,
                    expr.span.source_callsite(),
                    "consider adding a `;` to the last statement for consistent formatting",
                    "add a `;` here",
                    suggestion,
                    Applicability::MaybeIncorrect,
                );
            }
        }
    }
}
<commit_msg>Fixed failing tests<commit_after>use clippy_utils::diagnostics::span_lint_and_sugg;
use clippy_utils::source::snippet_with_macro_callsite;
use clippy_utils::{get_parent_expr_for_hir, in_macro, spans_on_same_line, sugg};
use if_chain::if_chain;
use rustc_errors::Applicability;
use rustc_hir::Expr;
use rustc_hir::{Block, BlockCheckMode, ExprKind};
use rustc_lint::{LateContext, LateLintPass};
use rustc_session::{declare_lint_pass, declare_tool_lint};

declare_clippy_lint! {
    /// **What it does:** Looks for blocks of expressions and fires if the last expression returns
    /// `()` but is not followed by a semicolon.
    ///
    /// **Why is this bad?** The semicolon might be optional but when extending the block with new
    /// code, it doesn't require a change in previous last line.
    ///
    /// **Known problems:** None.
    ///
    /// **Example:**
    ///
    /// ```rust
    /// fn main() {
    ///     println!("Hello world")
    /// }
    /// ```
    /// Use instead:
    /// ```rust
    /// fn main() {
    ///     println!("Hello world");
    /// }
    /// ```
    pub SEMICOLON_IF_NOTHING_RETURNED,
    pedantic,
    "add a semicolon if nothing is returned"
}

declare_lint_pass!(SemicolonIfNothingReturned => [SEMICOLON_IF_NOTHING_RETURNED]);

impl LateLintPass<'_> for SemicolonIfNothingReturned {
    fn check_block(&mut self, cx: &LateContext<'tcx>, block: &'tcx Block<'tcx>) {
        if_chain! {
            if !in_macro(block.span);
            if let Some(expr) = block.expr;
            let t_expr = cx.typeck_results().expr_ty(expr);
            if t_expr.is_unit();
            if let snippet = snippet_with_macro_callsite(cx, expr.span, "}");
            if !snippet.ends_with('}');
            if !check_if_inside_block_on_same_line(cx, block, expr);
            then {
                // filter out the desugared `for` loop
                if let ExprKind::DropTemps(..) = &expr.kind {
                    return;
                }

                let sugg = sugg::Sugg::hir_with_macro_callsite(cx, expr, "..");
                let suggestion = format!("{0};", sugg);
                span_lint_and_sugg(
                    cx,
                    SEMICOLON_IF_NOTHING_RETURNED,
                    expr.span.source_callsite(),
                    "consider adding a `;` to the last statement for consistent formatting",
                    "add a `;` here",
                    suggestion,
                    Applicability::MaybeIncorrect,
                );
            }
        }
    }
}

/// Check if this block is inside a closure or an unsafe block or a normal on the same line.
fn check_if_inside_block_on_same_line<'tcx>(
    cx: &LateContext<'tcx>,
    block: &'tcx Block<'tcx>,
    last_expr: &'tcx Expr<'_>,
) -> bool {
    if_chain! {
        if let Some(parent) = get_parent_expr_for_hir(cx, block.hir_id);

        if !matches!(block.rules, BlockCheckMode::DefaultBlock) ||
        matches!(parent.kind, ExprKind::Closure(..) | ExprKind::Block(..));

        if block.stmts.is_empty();
        then {
            return spans_on_same_line(cx, parent.span, last_expr.span);
        }
    }
    false
}
<|endoftext|>


<commit_before>use DocId;
use schema::Term;
use postings::{FieldSerializer, InvertedIndexSerializer};
use std::io;
use std::collections::HashMap;
use postings::Recorder;
use Result;
use schema::{Field, Schema};
use std::marker::PhantomData;
use std::ops::DerefMut;
use datastruct::stacker::{Heap, TermHashMap};
use postings::{NothingRecorder, TFAndPositionRecorder, TermFrequencyRecorder};
use schema::FieldEntry;
use schema::FieldType;
use tokenizer::Token;
use tokenizer::TokenStream;
use schema::IndexRecordOption;
use postings::UnorderedTermId;

fn posting_from_field_entry<'a>(
    field_entry: &FieldEntry,
    heap: &'a Heap,
) -> Box<PostingsWriter + 'a> {
    match *field_entry.field_type() {
        FieldType::Str(ref text_options) => text_options
            .get_indexing_options()
            .map(|indexing_options| match indexing_options.index_option() {
                IndexRecordOption::Basic => {
                    SpecializedPostingsWriter::<NothingRecorder>::new_boxed(heap)
                }
                IndexRecordOption::WithFreqs => {
                    SpecializedPostingsWriter::<TermFrequencyRecorder>::new_boxed(heap)
                }
                IndexRecordOption::WithFreqsAndPositions => {
                    SpecializedPostingsWriter::<TFAndPositionRecorder>::new_boxed(heap)
                }
            })
            .unwrap_or_else(|| SpecializedPostingsWriter::<NothingRecorder>::new_boxed(heap)),
        FieldType::U64(_) | FieldType::I64(_) | FieldType::HierarchicalFacet => {
            SpecializedPostingsWriter::<NothingRecorder>::new_boxed(heap)
        }
    }
}

pub struct MultiFieldPostingsWriter<'a> {
    heap: &'a Heap,
    term_index: TermHashMap<'a>,
    per_field_postings_writers: Vec<Box<PostingsWriter + 'a>>,
}

impl<'a> MultiFieldPostingsWriter<'a> {
    /// Create a new `MultiFieldPostingsWriter` given
    /// a schema and a heap.
    pub fn new(schema: &Schema, table_bits: usize, heap: &'a Heap) -> MultiFieldPostingsWriter<'a> {
        let term_index = TermHashMap::new(table_bits, heap);
        let per_field_postings_writers: Vec<_> = schema
            .fields()
            .iter()
            .map(|field_entry| posting_from_field_entry(field_entry, heap))
            .collect();

        MultiFieldPostingsWriter {
            heap,
            term_index,
            per_field_postings_writers,
        }
    }

    pub fn index_text(&mut self, doc: DocId, field: Field, token_stream: &mut TokenStream) -> u32 {
        let postings_writer = self.per_field_postings_writers[field.0 as usize].deref_mut();
        postings_writer.index_text(&mut self.term_index, doc, field, token_stream, self.heap)
    }

    pub fn subscribe(&mut self, doc: DocId, term: &Term) -> UnorderedTermId {
        let postings_writer = self.per_field_postings_writers[term.field().0 as usize].deref_mut();
        postings_writer.subscribe(&mut self.term_index, doc, 0u32, term, self.heap)
    }

    /// Serialize the inverted index.
    /// It pushes all term, one field at a time, towards the
    /// postings serializer.
    #[allow(needless_range_loop)]
    pub fn serialize(
        &self,
        serializer: &mut InvertedIndexSerializer,
    ) -> Result<HashMap<Field, HashMap<UnorderedTermId, usize>>> {
        let mut term_offsets: Vec<(&[u8], u32, UnorderedTermId)> = self.term_index.iter().collect();
        term_offsets.sort_by_key(|&(k, _, _)| k);

        let mut offsets: Vec<(Field, usize)> = vec![];
        let term_offsets_it = term_offsets
            .iter()
            .cloned()
            .map(|(key, _, _)| Term::wrap(key).field())
            .enumerate();

        let mut unordered_term_mappings: HashMap<Field, HashMap<UnorderedTermId, usize>> =
            HashMap::new();

        let mut prev_field = Field(u32::max_value());
        for (offset, field) in term_offsets_it {
            if field != prev_field {
                offsets.push((field, offset));
                prev_field = field;
            }
        }
        offsets.push((Field(0), term_offsets.len()));

        for i in 0..(offsets.len() - 1) {
            let (field, start) = offsets[i];
            let (_, stop) = offsets[i + 1];

            // populating the unordered term ord -> ordered term ord mapping
            // for the field.
            let mut mapping = HashMap::new();
            for (term_ord, term_unord_id) in term_offsets[start..stop]
                .iter()
                .map(|&(_, _, bucket)| bucket)
                .enumerate()
            {
                mapping.insert(term_unord_id, term_ord);
            }
            unordered_term_mappings.insert(field, mapping);

            let postings_writer = &self.per_field_postings_writers[field.0 as usize];
            let mut field_serializer = serializer.new_field(field, postings_writer.total_num_tokens())?;
            postings_writer.serialize(
                &term_offsets[start..stop],
                &mut field_serializer,
                self.heap,
            )?;
            field_serializer.close()?;
        }
        Ok(unordered_term_mappings)
    }

    /// Return true iff the term dictionary is saturated.
    pub fn is_term_saturated(&self) -> bool {
        self.term_index.is_saturated()
    }
}

/// The `PostingsWriter` is in charge of receiving documenting
/// and building a `Segment` in anonymous memory.
///
/// `PostingsWriter` writes in a `Heap`.
pub trait PostingsWriter {
    /// Record that a document contains a term at a given position.
    ///
    /// * doc  - the document id
    /// * pos  - the term position (expressed in tokens)
    /// * term - the term
    /// * heap - heap used to store the postings informations as well as the terms
    /// in the hashmap.
    fn subscribe(
        &mut self,
        term_index: &mut TermHashMap,
        doc: DocId,
        pos: u32,
        term: &Term,
        heap: &Heap,
    ) -> UnorderedTermId;

    /// Serializes the postings on disk.
    /// The actual serialization format is handled by the `PostingsSerializer`.
    fn serialize(
        &self,
        term_addrs: &[(&[u8], u32, UnorderedTermId)],
        serializer: &mut FieldSerializer,
        heap: &Heap,
    ) -> io::Result<()>;

    /// Tokenize a text and subscribe all of its token.
    fn index_text(
        &mut self,
        term_index: &mut TermHashMap,
        doc_id: DocId,
        field: Field,
        token_stream: &mut TokenStream,
        heap: &Heap,
    ) -> u32 {
        let mut term = unsafe { Term::with_capacity(100) };
        term.set_field(field);
        let num_tokens = {
            let mut sink = |token: &Token| {
                term.set_text(token.text.as_str());
                self.subscribe(term_index, doc_id, token.position as u32, &term, heap);
            };
            token_stream.process(&mut sink)
        };
        self.add_num_tokens(num_tokens);
        num_tokens
    }

    fn add_num_tokens(&mut self, num_tokens: u32);

    fn total_num_tokens(&self) -> u64;
}

/// The `SpecializedPostingsWriter` is just here to remove dynamic
/// dispatch to the recorder information.
pub struct SpecializedPostingsWriter<'a, Rec: Recorder + 'static> {
    heap: &'a Heap,
    total_num_tokens: u64,
    _recorder_type: PhantomData<Rec>,
}

impl<'a, Rec: Recorder + 'static> SpecializedPostingsWriter<'a, Rec> {
    /// constructor
    pub fn new(heap: &'a Heap) -> SpecializedPostingsWriter<'a, Rec> {
        SpecializedPostingsWriter {
            heap,
            total_num_tokens: 0u64,
            _recorder_type: PhantomData,
        }
    }

    /// Builds a `SpecializedPostingsWriter` storing its data in a heap.
    pub fn new_boxed(heap: &'a Heap) -> Box<PostingsWriter + 'a> {
        Box::new(SpecializedPostingsWriter::<Rec>::new(heap))
    }
}

impl<'a, Rec: Recorder + 'static> PostingsWriter for SpecializedPostingsWriter<'a, Rec> {
    fn subscribe(
        &mut self,
        term_index: &mut TermHashMap,
        doc: DocId,
        position: u32,
        term: &Term,
        heap: &Heap,
    ) -> UnorderedTermId {
        debug_assert!(term.as_slice().len() >= 4);
        let (term_ord, recorder): (UnorderedTermId, &mut Rec) = term_index.get_or_create(term);
        let current_doc = recorder.current_doc();
        if current_doc != doc {
            if current_doc != u32::max_value() {
                recorder.close_doc(heap);
            }
            recorder.new_doc(doc, heap);
        }
        recorder.record_position(position, heap);
        self.add_num_tokens(1u32);
        term_ord
    }

    fn serialize(
        &self,
        term_addrs: &[(&[u8], u32, UnorderedTermId)],
        serializer: &mut FieldSerializer,
        heap: &Heap,
    ) -> io::Result<()> {
        for &(term_bytes, addr, _) in term_addrs {
            let recorder: &mut Rec = self.heap.get_mut_ref(addr);
            serializer.new_term(&term_bytes[4..])?;
            recorder.serialize(addr, serializer, heap)?;
            serializer.close_term()?;
        }
        Ok(())
    }

    fn add_num_tokens(&mut self, num_tokens: u32) {
        self.total_num_tokens += num_tokens as u64;
    }

    fn total_num_tokens(&self) -> u64 {
        self.total_num_tokens
    }
}
<commit_msg>Bugfixed<commit_after>use DocId;
use schema::Term;
use postings::{FieldSerializer, InvertedIndexSerializer};
use std::io;
use std::collections::HashMap;
use postings::Recorder;
use Result;
use schema::{Field, Schema};
use std::marker::PhantomData;
use std::ops::DerefMut;
use datastruct::stacker::{Heap, TermHashMap};
use postings::{NothingRecorder, TFAndPositionRecorder, TermFrequencyRecorder};
use schema::FieldEntry;
use schema::FieldType;
use tokenizer::Token;
use tokenizer::TokenStream;
use schema::IndexRecordOption;
use postings::UnorderedTermId;

fn posting_from_field_entry<'a>(
    field_entry: &FieldEntry,
    heap: &'a Heap,
) -> Box<PostingsWriter + 'a> {
    match *field_entry.field_type() {
        FieldType::Str(ref text_options) => text_options
            .get_indexing_options()
            .map(|indexing_options| match indexing_options.index_option() {
                IndexRecordOption::Basic => {
                    SpecializedPostingsWriter::<NothingRecorder>::new_boxed(heap)
                }
                IndexRecordOption::WithFreqs => {
                    SpecializedPostingsWriter::<TermFrequencyRecorder>::new_boxed(heap)
                }
                IndexRecordOption::WithFreqsAndPositions => {
                    SpecializedPostingsWriter::<TFAndPositionRecorder>::new_boxed(heap)
                }
            })
            .unwrap_or_else(|| SpecializedPostingsWriter::<NothingRecorder>::new_boxed(heap)),
        FieldType::U64(_) | FieldType::I64(_) | FieldType::HierarchicalFacet => {
            SpecializedPostingsWriter::<NothingRecorder>::new_boxed(heap)
        }
    }
}

pub struct MultiFieldPostingsWriter<'a> {
    heap: &'a Heap,
    term_index: TermHashMap<'a>,
    per_field_postings_writers: Vec<Box<PostingsWriter + 'a>>,
}

impl<'a> MultiFieldPostingsWriter<'a> {
    /// Create a new `MultiFieldPostingsWriter` given
    /// a schema and a heap.
    pub fn new(schema: &Schema, table_bits: usize, heap: &'a Heap) -> MultiFieldPostingsWriter<'a> {
        let term_index = TermHashMap::new(table_bits, heap);
        let per_field_postings_writers: Vec<_> = schema
            .fields()
            .iter()
            .map(|field_entry| posting_from_field_entry(field_entry, heap))
            .collect();

        MultiFieldPostingsWriter {
            heap,
            term_index,
            per_field_postings_writers,
        }
    }

    pub fn index_text(&mut self, doc: DocId, field: Field, token_stream: &mut TokenStream) -> u32 {
        let postings_writer = self.per_field_postings_writers[field.0 as usize].deref_mut();
        postings_writer.index_text(&mut self.term_index, doc, field, token_stream, self.heap)
    }

    pub fn subscribe(&mut self, doc: DocId, term: &Term) -> UnorderedTermId {
        let postings_writer = self.per_field_postings_writers[term.field().0 as usize].deref_mut();
        postings_writer.subscribe(&mut self.term_index, doc, 0u32, term, self.heap)
    }

    /// Serialize the inverted index.
    /// It pushes all term, one field at a time, towards the
    /// postings serializer.
    #[allow(needless_range_loop)]
    pub fn serialize(
        &self,
        serializer: &mut InvertedIndexSerializer,
    ) -> Result<HashMap<Field, HashMap<UnorderedTermId, usize>>> {
        let mut term_offsets: Vec<(&[u8], u32, UnorderedTermId)> = self.term_index.iter().collect();
        term_offsets.sort_by_key(|&(k, _, _)| k);

        let mut offsets: Vec<(Field, usize)> = vec![];
        let term_offsets_it = term_offsets
            .iter()
            .cloned()
            .map(|(key, _, _)| Term::wrap(key).field())
            .enumerate();

        let mut unordered_term_mappings: HashMap<Field, HashMap<UnorderedTermId, usize>> =
            HashMap::new();

        let mut prev_field = Field(u32::max_value());
        for (offset, field) in term_offsets_it {
            if field != prev_field {
                offsets.push((field, offset));
                prev_field = field;
            }
        }
        offsets.push((Field(0), term_offsets.len()));

        for i in 0..(offsets.len() - 1) {
            let (field, start) = offsets[i];
            let (_, stop) = offsets[i + 1];

            // populating the unordered term ord -> ordered term ord mapping
            // for the field.
            let mut mapping = HashMap::new();
            for (term_ord, term_unord_id) in term_offsets[start..stop]
                .iter()
                .map(|&(_, _, bucket)| bucket)
                .enumerate()
            {
                mapping.insert(term_unord_id, term_ord);
            }
            unordered_term_mappings.insert(field, mapping);

            let postings_writer = &self.per_field_postings_writers[field.0 as usize];
            let mut field_serializer = serializer.new_field(field, postings_writer.total_num_tokens())?;
            postings_writer.serialize(
                &term_offsets[start..stop],
                &mut field_serializer,
                self.heap,
            )?;
            field_serializer.close()?;
        }
        Ok(unordered_term_mappings)
    }

    /// Return true iff the term dictionary is saturated.
    pub fn is_term_saturated(&self) -> bool {
        self.term_index.is_saturated()
    }
}

/// The `PostingsWriter` is in charge of receiving documenting
/// and building a `Segment` in anonymous memory.
///
/// `PostingsWriter` writes in a `Heap`.
pub trait PostingsWriter {
    /// Record that a document contains a term at a given position.
    ///
    /// * doc  - the document id
    /// * pos  - the term position (expressed in tokens)
    /// * term - the term
    /// * heap - heap used to store the postings informations as well as the terms
    /// in the hashmap.
    fn subscribe(
        &mut self,
        term_index: &mut TermHashMap,
        doc: DocId,
        pos: u32,
        term: &Term,
        heap: &Heap,
    ) -> UnorderedTermId;

    /// Serializes the postings on disk.
    /// The actual serialization format is handled by the `PostingsSerializer`.
    fn serialize(
        &self,
        term_addrs: &[(&[u8], u32, UnorderedTermId)],
        serializer: &mut FieldSerializer,
        heap: &Heap,
    ) -> io::Result<()>;

    /// Tokenize a text and subscribe all of its token.
    fn index_text(
        &mut self,
        term_index: &mut TermHashMap,
        doc_id: DocId,
        field: Field,
        token_stream: &mut TokenStream,
        heap: &Heap,
    ) -> u32 {
        let mut term = unsafe { Term::with_capacity(100) };
        term.set_field(field);
        let num_tokens = {
            let mut sink = |token: &Token| {
                term.set_text(token.text.as_str());
                self.subscribe(term_index, doc_id, token.position as u32, &term, heap);
            };
            token_stream.process(&mut sink)
        };
        num_tokens
    }

    fn total_num_tokens(&self) -> u64;
}

/// The `SpecializedPostingsWriter` is just here to remove dynamic
/// dispatch to the recorder information.
pub struct SpecializedPostingsWriter<'a, Rec: Recorder + 'static> {
    heap: &'a Heap,
    total_num_tokens: u64,
    _recorder_type: PhantomData<Rec>,
}

impl<'a, Rec: Recorder + 'static> SpecializedPostingsWriter<'a, Rec> {
    /// constructor
    pub fn new(heap: &'a Heap) -> SpecializedPostingsWriter<'a, Rec> {
        SpecializedPostingsWriter {
            heap,
            total_num_tokens: 0u64,
            _recorder_type: PhantomData,
        }
    }

    /// Builds a `SpecializedPostingsWriter` storing its data in a heap.
    pub fn new_boxed(heap: &'a Heap) -> Box<PostingsWriter + 'a> {
        Box::new(SpecializedPostingsWriter::<Rec>::new(heap))
    }
}

impl<'a, Rec: Recorder + 'static> PostingsWriter for SpecializedPostingsWriter<'a, Rec> {
    fn subscribe(
        &mut self,
        term_index: &mut TermHashMap,
        doc: DocId,
        position: u32,
        term: &Term,
        heap: &Heap,
    ) -> UnorderedTermId {
        debug_assert!(term.as_slice().len() >= 4);
        let (term_ord, recorder): (UnorderedTermId, &mut Rec) = term_index.get_or_create(term);
        let current_doc = recorder.current_doc();
        if current_doc != doc {
            if current_doc != u32::max_value() {
                recorder.close_doc(heap);
            }
            recorder.new_doc(doc, heap);
        }
        self.total_num_tokens += 1;
        recorder.record_position(position, heap);
        term_ord
    }

    fn serialize(
        &self,
        term_addrs: &[(&[u8], u32, UnorderedTermId)],
        serializer: &mut FieldSerializer,
        heap: &Heap,
    ) -> io::Result<()> {
        for &(term_bytes, addr, _) in term_addrs {
            let recorder: &mut Rec = self.heap.get_mut_ref(addr);
            serializer.new_term(&term_bytes[4..])?;
            recorder.serialize(addr, serializer, heap)?;
            serializer.close_term()?;
        }
        Ok(())
    }

    fn total_num_tokens(&self) -> u64 {
        self.total_num_tokens
    }
}
<|endoftext|>


<commit_before>// Implements the `zakopane checksum` subcommand.

use crate::structs::ZakopaneError;

const MAX_TASKS: usize = 8;

struct ChecksumWithPath {
    checksum: String,
    path: std::path::PathBuf,
}

impl ChecksumWithPath {
    pub fn new(checksum: String, path: std::path::PathBuf) -> Self {
        Self { checksum, path }
    }
}

type ChecksumResult = Result<ChecksumWithPath, ZakopaneError>;
type ChecksumTaskJoinHandle = tokio::task::JoinHandle<Vec<ChecksumWithPath>>;

struct ChecksumTaskDispatcherData {
    path: std::path::PathBuf,

    // Rate-limiter for spawned checksum tasks.
    semaphore: std::sync::Arc<tokio::sync::Semaphore>,

    // Counts total number of spawned checksum tasks.
    spawn_counter: std::sync::Arc<std::sync::atomic::AtomicUsize>,

    sender: tokio::sync::mpsc::Sender<ChecksumResult>,
}

impl ChecksumTaskDispatcherData {
    pub fn new(
        path: std::path::PathBuf,
        spawn_counter: std::sync::Arc<std::sync::atomic::AtomicUsize>,
        sender: tokio::sync::mpsc::Sender<ChecksumResult>,
    ) -> Self {
        Self {
            path,
            semaphore: std::sync::Arc::new(tokio::sync::Semaphore::new(MAX_TASKS)),
            spawn_counter,
            sender,
        }
    }
}

// Sends a checksum task result downstream to the collector task.
fn checksum_task_send_result(
    result: ChecksumResult,
    sender: tokio::sync::mpsc::Sender<ChecksumResult>,
) {
    if let Err(e) = sender.blocking_send(result) {
        eprintln!("BUG: Sender::send() failed: {}", e);
    }
}

// Separated from `checksum_task()` to ensure the `add_permits()` call is
// always hit.
fn checksum_task_impl(path: std::path::PathBuf, sender: tokio::sync::mpsc::Sender<ChecksumResult>) {
    let contents = match std::fs::read(&path).map_err(ZakopaneError::Io) {
        Ok(contents) => contents,
        Err(e) => return checksum_task_send_result(Err(e), sender),
    };
    let checksum = crypto_hash::hex_digest(crypto_hash::Algorithm::SHA256, &contents);
    checksum_task_send_result(Ok(ChecksumWithPath::new(checksum, path)), sender);
}

// Represents a spawned checksum task.
fn checksum_task(
    path: std::path::PathBuf,
    sender: tokio::sync::mpsc::Sender<ChecksumResult>,
    semaphore_clone: std::sync::Arc<tokio::sync::Semaphore>,
) {
    checksum_task_impl(path, sender);

    // See comment in `ChecksumTaskManager::spawn_task()`.
    semaphore_clone.add_permits(1);
}

// Represents the spawned collection task.
async fn collector_task(
    mut receiver: tokio::sync::mpsc::Receiver<ChecksumResult>,
    spawn_counter: std::sync::Arc<std::sync::atomic::AtomicUsize>,
) -> Vec<ChecksumWithPath> {
    let mut sums: Vec<ChecksumWithPath> = Vec::new();
    let mut errors: usize = 0;
    loop {
        while let Some(result) = receiver.recv().await {
            match result {
                Ok(digest_with_path) => sums.push(digest_with_path),
                Err(_) => errors += 1,
            }
        }
        if sums.len() + errors == spawn_counter.load(std::sync::atomic::Ordering::SeqCst) {
            return sums;
        }
    }
}

// Spawns the collector task that listens for checksum results
// provided by the checksum tasks.
async fn spawn_collector(
    path: std::path::PathBuf,
) -> (ChecksumTaskDispatcherData, ChecksumTaskJoinHandle) {
    let (sender, receiver) = tokio::sync::mpsc::channel(MAX_TASKS);
    let spawn_counter = std::sync::Arc::new(std::sync::atomic::AtomicUsize::new(0));
    let join_handle = tokio::task::spawn(collector_task(receiver, spawn_counter.clone()));
    (
        ChecksumTaskDispatcherData::new(path, spawn_counter, sender),
        join_handle,
    )
}

async fn spawn_checksum_tasks(context: ChecksumTaskDispatcherData) {
    let walk_iter = walkdir::WalkDir::new(&context.path).into_iter();
    // The `filter_entry()` call is crafted s.t.
    // *    we skip and don't descend into hidden directories
    // *    unless the hidden directory is the target directory, because
    //      the target directory is always the first yielded value from
    //      the `WalkDir`.
    for entry in walk_iter.filter_entry(|e| {
        e.path() == &context.path
            || !e
                .file_name()
                .to_str()
                .map(|path| path.starts_with("."))
                .unwrap_or(false)
    }) {
        if let Ok(direntry) = entry {
            let path = direntry.into_path();
            if let Ok(metadata) = std::fs::symlink_metadata(&path) {
                if !metadata.file_type().is_file() {
                    continue;
                }
            } else {
                continue;
            }

            context
                .spawn_counter
                .fetch_add(1, std::sync::atomic::Ordering::SeqCst);
            let permit = context.semaphore.acquire().await.unwrap();

            // Icky workaround for not being able to pass `SemaphorePermit`
            // directly into a spawned task. Note that new permits are added
            // later in the checksum task.
            //
            // See also: https://github.com/tokio-rs/tokio/issues/1998
            permit.forget();

            let sender = context.sender.clone();
            let semaphore_clone = context.semaphore.clone();
            tokio::task::spawn_blocking(move || checksum_task(path, sender, semaphore_clone));
        }
    }
}

// Pretty-prints the sorted `checksums` in a format much like what the
// `sha256sum` binary outputs.
//
// Note that the standard zakopane snapshot header is not added here.
fn pretty_format_checksums(path: std::path::PathBuf, checksums: Vec<ChecksumWithPath>) -> String {
    let mut buffer: Vec<String> = checksums
        .into_iter()
        .map(|e| {
            format!(
                "{}  ./{}",
                e.checksum,
                e.path.strip_prefix(&path).unwrap().to_str().unwrap()
            )
        })
        .collect();
    buffer.push(String::new());
    buffer.join("\n")
}

async fn checksum_impl(path: std::path::PathBuf) -> String {
    let (dispatcher_data, join_handle) = spawn_collector(path.clone()).await;
    spawn_checksum_tasks(dispatcher_data).await;
    let mut checksums: Vec<ChecksumWithPath> = join_handle.await.unwrap();
    checksums.sort_by(|a, b| a.path.cmp(&b.path));
    pretty_format_checksums(path, checksums)
}

pub fn checksum(path: std::path::PathBuf) -> String {
    let runtime = tokio::runtime::Runtime::new().unwrap();
    runtime.block_on(checksum_impl(path))
}
<commit_msg>checksum: Tamp down reads<commit_after>// Implements the `zakopane checksum` subcommand.

use std::io::Read;
use std::io::Write;

use crate::structs::ZakopaneError;

const MAX_TASKS: usize = 8;
const READ_SIZE: usize = 2 << 20;

struct ChecksumWithPath {
    checksum: String,
    path: std::path::PathBuf,
}

impl ChecksumWithPath {
    pub fn new(checksum: String, path: std::path::PathBuf) -> Self {
        Self { checksum, path }
    }
}

type ChecksumResult = Result<ChecksumWithPath, ZakopaneError>;
type ChecksumTaskJoinHandle = tokio::task::JoinHandle<Vec<ChecksumWithPath>>;

struct ChecksumTaskDispatcherData {
    path: std::path::PathBuf,

    // Rate-limiter for spawned checksum tasks.
    semaphore: std::sync::Arc<tokio::sync::Semaphore>,

    // Counts total number of spawned checksum tasks.
    spawn_counter: std::sync::Arc<std::sync::atomic::AtomicUsize>,

    sender: tokio::sync::mpsc::Sender<ChecksumResult>,
}

impl ChecksumTaskDispatcherData {
    pub fn new(
        path: std::path::PathBuf,
        spawn_counter: std::sync::Arc<std::sync::atomic::AtomicUsize>,
        sender: tokio::sync::mpsc::Sender<ChecksumResult>,
    ) -> Self {
        Self {
            path,
            semaphore: std::sync::Arc::new(tokio::sync::Semaphore::new(MAX_TASKS)),
            spawn_counter,
            sender,
        }
    }
}

// Sends a checksum task result downstream to the collector task.
fn checksum_task_send_result(
    result: ChecksumResult,
    sender: tokio::sync::mpsc::Sender<ChecksumResult>,
) {
    if let Err(e) = sender.blocking_send(result) {
        eprintln!("BUG: Sender::send() failed: {}", e);
    }
}

fn checksum_task_impl(path: std::path::PathBuf) -> ChecksumResult {
    let mut hasher = crypto_hash::Hasher::new(crypto_hash::Algorithm::SHA256);
    let mut buffer: Vec<u8> = vec![0; READ_SIZE];
    let mut file = std::fs::File::open(&path).map_err(ZakopaneError::Io)?;
    loop {
        let read_bytes = file.read(&mut buffer).map_err(ZakopaneError::Io)?;
        if read_bytes == 0 {
            let checksum = hasher
                .finish()
                .into_iter()
                .map(|byte| format!("{:02x}", byte))
                .collect::<Vec<String>>()
                .join("");
            return Ok(ChecksumWithPath::new(checksum, path));
        }
        hasher
            .write_all(&buffer[..read_bytes])
            .map_err(ZakopaneError::Io)?;
    }
}

// Represents a spawned checksum task.
fn checksum_task(
    path: std::path::PathBuf,
    sender: tokio::sync::mpsc::Sender<ChecksumResult>,
    semaphore_clone: std::sync::Arc<tokio::sync::Semaphore>,
) {
    let result = checksum_task_impl(path);
    checksum_task_send_result(result, sender);

    // See comment in `ChecksumTaskManager::spawn_task()`.
    semaphore_clone.add_permits(1);
}

// Represents the spawned collection task.
async fn collector_task(
    mut receiver: tokio::sync::mpsc::Receiver<ChecksumResult>,
    spawn_counter: std::sync::Arc<std::sync::atomic::AtomicUsize>,
) -> Vec<ChecksumWithPath> {
    let mut sums: Vec<ChecksumWithPath> = Vec::new();
    let mut errors: usize = 0;
    loop {
        while let Some(result) = receiver.recv().await {
            match result {
                Ok(digest_with_path) => sums.push(digest_with_path),
                Err(_) => errors += 1,
            }
        }
        if sums.len() + errors == spawn_counter.load(std::sync::atomic::Ordering::SeqCst) {
            return sums;
        }
    }
}

// Spawns the collector task that listens for checksum results
// provided by the checksum tasks.
async fn spawn_collector(
    path: std::path::PathBuf,
) -> (ChecksumTaskDispatcherData, ChecksumTaskJoinHandle) {
    let (sender, receiver) = tokio::sync::mpsc::channel(MAX_TASKS);
    let spawn_counter = std::sync::Arc::new(std::sync::atomic::AtomicUsize::new(0));
    let join_handle = tokio::task::spawn(collector_task(receiver, spawn_counter.clone()));
    (
        ChecksumTaskDispatcherData::new(path, spawn_counter, sender),
        join_handle,
    )
}

async fn spawn_checksum_tasks(context: ChecksumTaskDispatcherData) {
    let walk_iter = walkdir::WalkDir::new(&context.path).into_iter();
    // The `filter_entry()` call is crafted s.t.
    // *    we skip and don't descend into hidden directories
    // *    unless the hidden directory is the target directory, because
    //      the target directory is always the first yielded value from
    //      the `WalkDir`.
    for entry in walk_iter.filter_entry(|e| {
        e.path() == &context.path
            || !e
                .file_name()
                .to_str()
                .map(|path| path.starts_with("."))
                .unwrap_or(false)
    }) {
        if let Ok(direntry) = entry {
            let path = direntry.into_path();
            if let Ok(metadata) = std::fs::symlink_metadata(&path) {
                if !metadata.file_type().is_file() {
                    continue;
                }
            } else {
                continue;
            }

            context
                .spawn_counter
                .fetch_add(1, std::sync::atomic::Ordering::SeqCst);
            let permit = context.semaphore.acquire().await.unwrap();

            // Icky workaround for not being able to pass `SemaphorePermit`
            // directly into a spawned task. Note that new permits are added
            // later in the checksum task.
            //
            // See also: https://github.com/tokio-rs/tokio/issues/1998
            permit.forget();

            let sender = context.sender.clone();
            let semaphore_clone = context.semaphore.clone();
            tokio::task::spawn_blocking(move || checksum_task(path, sender, semaphore_clone));
        }
    }
}

// Pretty-prints the sorted `checksums` in a format much like what the
// `sha256sum` binary outputs.
//
// Note that the standard zakopane snapshot header is not added here.
fn pretty_format_checksums(path: std::path::PathBuf, checksums: Vec<ChecksumWithPath>) -> String {
    let mut buffer: Vec<String> = checksums
        .into_iter()
        .map(|e| {
            format!(
                "{}  ./{}",
                e.checksum,
                e.path.strip_prefix(&path).unwrap().to_str().unwrap()
            )
        })
        .collect();
    buffer.push(String::new());
    buffer.join("\n")
}

async fn checksum_impl(path: std::path::PathBuf) -> String {
    let (dispatcher_data, join_handle) = spawn_collector(path.clone()).await;
    spawn_checksum_tasks(dispatcher_data).await;
    let mut checksums: Vec<ChecksumWithPath> = join_handle.await.unwrap();
    checksums.sort_by(|a, b| a.path.cmp(&b.path));
    pretty_format_checksums(path, checksums)
}

pub fn checksum(path: std::path::PathBuf) -> String {
    let runtime = tokio::runtime::Runtime::new().unwrap();
    runtime.block_on(checksum_impl(path))
}
<|endoftext|>