<commit_before>#![cfg(not(windows))]

extern crate libc;

use screen::Key;
use screen::Key::*;
use std::io::{Read, Write};
use std::fs::{File, OpenOptions};
use std::os::unix::io::AsRawFd;
use std::path::*;
use std::process::Command;
use std::process::Stdio;
use std::str;
use std::iter::repeat;
use std::cmp::min;
use ansi;
use ::NEWLINE;

use std::{mem, ptr};
use std::thread;
use std::sync::mpsc::Sender;
use std::sync::mpsc::Receiver;
use std::sync::{mpsc, Arc, Mutex};

use self::libc::{sigaction, SIGWINCH, SA_SIGINFO, sighandler_t, c_int, c_ushort, c_ulong};

pub struct Screen {
    tty: Terminal,
    original_stty_state: Vec<u8>,
    start_line: u16,
    desired_rows: u16,
}

impl Screen {
    pub fn open_screen(desired_rows: u16) -> Screen {
        let mut tty = Terminal::open_terminal();
        let current_stty_state = tty.stty(&["-g"]);
        tty.initialize();
        let (_, rows) = tty.winsize().unwrap();
        let visible_choices = min(desired_rows, rows - 1);
        let start_line = rows - visible_choices - 1;
        for _ in 0..visible_choices {
            tty.write(NEWLINE.as_bytes());
        }
        tty.write(&ansi::save_cursor());

        Screen {
            tty: tty,
            original_stty_state: current_stty_state,
            start_line: start_line,
            desired_rows: desired_rows,
        }
    }

    pub fn visible_choices(&self) -> u16 {
        let (_, rows) = self.tty.winsize().unwrap();
        min(self.desired_rows, rows - 1)
    }

    pub fn width(&self) -> u16 {
        let (cols, _) = self.tty.winsize().unwrap();
        cols
    }

    fn restore_tty(&mut self) {
        self.tty.stty(&[&String::from_utf8(self.original_stty_state.clone()).unwrap()]);
    }

    fn reset_cursor(&mut self) {
        self.tty.write(&ansi::restore_cursor());

        // Writing this carriage return works around a rendering bug in
        // Neovim's terminal emulation. Without it, the cursor flies
        // around all over the place while typing, and the input prompt
        // is not rendered correctly. I have no idea why this workaround
        // is effective, or why others are not, or what the root cause
        // of the issue is (likely something involving save/restore
        // support).
        self.tty.write(b"\r");

        let num_lines = self.visible_choices();
        self.tty.write(&ansi::cursor_up(num_lines));
    }

    pub fn move_cursor_to_prompt_line(&mut self, col: u16) {
        self.reset_cursor();
        self.tty.write(&ansi::cursor_right(col));
    }

    pub fn blank_screen(&mut self) {
        self.reset_cursor();
        let blank_line = repeat(' ').take(self.width() as usize).collect::<String>();
        for _ in 0..self.visible_choices() + 1 {
            self.tty.write(blank_line.as_bytes());
        }
        self.reset_cursor();
    }

    pub fn blank_entire_screen(&mut self){
        self.tty.write(&ansi::blank_screen());
    }

    pub fn show_cursor(&mut self) {
        self.tty.write(&ansi::show_cursor());
    }

    pub fn hide_cursor(&mut self) {
        self.tty.write(&ansi::hide_cursor());
    }

    pub fn write(&mut self, s: &str) {
        self.tty.write(s.as_bytes());
    }

    pub fn write_red_inverted(&mut self, s: &str) {
        self.tty.write(&ansi::red());
        self.tty.write(&ansi::inverse());
        self.tty.write(s.as_bytes());
        self.tty.write(&ansi::reset());
    }

    pub fn write_red(&mut self, s: &str) {
        self.tty.write(&ansi::red());
        self.tty.write(s.as_bytes());
        self.tty.write(&ansi::reset());
    }

    pub fn write_inverted(&mut self, s: &str) {
        self.tty.write(&ansi::inverse());
        self.tty.write(s.as_bytes());
        self.tty.write(&ansi::reset());
    }

    // Return all buffered keystrokes, or the next key if buffer is empty.
    pub fn get_buffered_keys(&mut self) -> Vec<Key> {
        let mut ret = Vec::new();
        while let Ok(bytes) = self.tty.input.try_recv() {
            ret.extend(bytes);
        }
        while ret.is_empty() {
            let bytes = self.tty.input.recv().unwrap();
            if bytes == vec![0] { // SIGWINCH
                self.blank_entire_screen();
                return vec![Nothing];
            } else {
                ret.extend(bytes);
            }
        }
        Terminal::translate_bytes(ret.clone())
    }

    fn more_bytes_needed(bytes: &[u8]) -> bool {
        if let Err(_) = str::from_utf8(bytes) { true } else { false }
    }
}

impl Drop for Screen {
    fn drop(&mut self) {
        self.restore_tty();
    }
}

struct Terminal {
    input: Receiver<Vec<u8>>,
    input_fd: i32,
    output: File,
}

static mut GLOBAL_TX: *const Arc<Mutex<Sender<Vec<u8>>>> = 0 as *const Arc<Mutex<Sender<Vec<u8>>>>;

fn set_global_tx(tx: Sender<Vec<u8>>) {
    unsafe {
        let singleton = Arc::new(Mutex::new(tx));
        GLOBAL_TX = mem::transmute(Box::new(singleton));
    }
}

fn get_global_tx() -> Sender<Vec<u8>> {
    let singleton = unsafe { (*GLOBAL_TX).clone() };
    let tx = singleton.lock().unwrap();
    tx.clone()
}

unsafe extern fn sigwinch_handler(signum: c_int) {
    if signum == SIGWINCH {
        get_global_tx().send([0].to_vec()).unwrap();
    }
}

fn register_sigwinch_handler() {
    unsafe {
        let mut action: sigaction = mem::zeroed();
        action.sa_flags = SA_SIGINFO;
        action.sa_sigaction = sigwinch_handler as sighandler_t;
        sigaction(SIGWINCH, &action, ptr::null_mut());
    }
}

impl Terminal {
    fn open_terminal() -> Terminal {
        let term_path = Path::new("/dev/tty");
        let mut input_file = File::open(&term_path).unwrap();
        let output_file = OpenOptions::new().write(true).open(&term_path).unwrap();
        let input_fd = input_file.as_raw_fd();
        let (tx, rx) = mpsc::channel();
        set_global_tx(tx);

        register_sigwinch_handler();

        thread::spawn(move || {
            loop {
                let mut buf = [0; 255];
                let tx = get_global_tx();
                if let Ok(length) = input_file.read(&mut buf) {
                    tx.send(buf[0..length].to_vec()).unwrap();
                } else {
                    tx.send([0].to_vec()).unwrap();
                }
            }
        });
        Terminal {
            input: rx,
            input_fd: input_fd,
            output: output_file,
        }
    }

    fn initialize(&mut self) {
        self.stty(&["raw", "-echo", "cbreak", "opost", "onlcr"]);
    }

    fn stty(&mut self, args: &[&str]) -> Vec<u8> {
        extern { fn dup2(src: c_int, dst: c_int) -> c_int;  }
        unsafe {
            // This is a hack until a replacement for InheritFd from old_io is available.
            dup2(self.input_fd, 0);
        }
        let mut process = match Command::new("stty")
            .args(args)
            .stdin(Stdio::inherit())
            .stdout(Stdio::piped())
            .spawn()
        {
            Ok(p) => p,
            Err(e) => panic!("Spawn failed: {}", e),
        };

        let exit = process.wait();

        let mut buf = Vec::new();
        if exit.unwrap().success() {
            process.stdout.as_mut().unwrap().read_to_end(&mut buf).unwrap();
            let mut str = String::from_utf8(buf).unwrap();

            // The output from `stty -g` may include a newline, which we have to strip off. Otherwise,
            // when we go to restore the tty, stty (on some platforms) will fail with an "invalid
            // argument" error.
            ::trim(&mut str);

            str.into_bytes()
        } else {
            process.stderr.as_mut().unwrap().read_to_end(&mut buf).unwrap();
            panic!(String::from_utf8(buf).unwrap());
        }
    }

    fn translate_bytes(bytes: Vec<u8>) -> Vec<Key> {
        const BEGIN_PASTE: &'static [u8] = b"\x1B[200~";
        const END_PASTE: &'static [u8] = b"\x1B[201~";

        if bytes == b"\x1B[A" || bytes == b"\x1BOA" { return vec![Up] };
        if bytes == b"\x1B[B" || bytes == b"\x1BOB" { return vec![Down] };
        if bytes == b"\x1B[Z" { return vec![ShiftTab] };

        let bs = if bytes.starts_with(BEGIN_PASTE) && bytes.ends_with(END_PASTE) {
            let start = BEGIN_PASTE.len();
            let end = bytes.len() - END_PASTE.len();
            bytes[start..end].to_vec()
        } else {
            bytes
        };

        String::from_utf8(bs)
            .unwrap()
            .chars()
            .map(Terminal::translate_char)
            .collect()
    }

    fn translate_char(c: char) -> Key {
        let numeric_char = c as u32;
        if c == '\r' {
            Enter
        } else if numeric_char == 9 {
            Tab
        } else if numeric_char == 127 {
            Backspace
        } else if numeric_char & 96 == 0 && numeric_char <= 128u32 {
            let c = c as u8;
            Control((c + 96u8) as char)
        } else {
            Char(c)
        }
    }

    fn write(&mut self, s: &[u8]) {
        self.output.write_all(s).unwrap();
    }

    fn writeln(&mut self, s: &str) {
        self.output.write(s.as_bytes()).unwrap();
        self.output.write(b"\n").unwrap();
    }

    fn winsize(&self) -> Option<(u16, u16)> {
        extern { fn ioctl(fd: c_int, request: c_ulong, ...) -> c_int; }
        #[cfg(any(target_os = "macos", target_os = "freebsd", target_os = "netbsd"))]
        const TIOCGWINSZ: c_ulong = 0x40087468;

        #[cfg(any(target_os = "linux", target_os = "android"))]
        const TIOCGWINSZ: c_ulong = 0x00005413;

        #[repr(C)]
        struct TermSize {
            rows: c_ushort,
            cols: c_ushort,
            x: c_ushort,
            y: c_ushort,
        }

        let size = TermSize { rows: 0, cols: 0, x: 0, y: 0 };
        if unsafe { ioctl(self.output.as_raw_fd(), TIOCGWINSZ, &size) } == 0 {
            Some((size.cols as u16, size.rows as u16))
        } else {
            None
        }
    }
}

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

    #[test]
    fn winsize_test() {
        // Travis-CI builds run without a tty, making this test impossible.
        if option_env!("TRAVIS").is_some() {
            // TODO: It should be made obvious from the output that this test was skipped
            return;
        }
        let term = Terminal::open_terminal();
        let (cols, rows) = term.winsize().expect("Failed to get window size!");
        // We don't know the window size a priori, but we can at least
        // assert that it is within some kind of sensible range.
        assert!(cols > 40);
        assert!(rows > 40);
        assert!(cols < 1000);
        assert!(rows < 1000);
    }
}
<commit_msg>[Unix] Buffer terminal output<commit_after>#![cfg(not(windows))]

extern crate libc;

use screen::Key;
use screen::Key::*;
use std::io::{Read, Write};
use std::fs::{File, OpenOptions};
use std::os::unix::io::AsRawFd;
use std::path::*;
use std::process::Command;
use std::process::Stdio;
use std::str;
use std::iter::repeat;
use std::cmp::min;
use ansi;
use ::NEWLINE;

use std::{mem, ptr};
use std::thread;
use std::sync::mpsc::Sender;
use std::sync::mpsc::Receiver;
use std::sync::{mpsc, Arc, Mutex};

use self::libc::{sigaction, SIGWINCH, SA_SIGINFO, sighandler_t, c_int, c_ushort, c_ulong};

pub struct Screen {
    tty: Terminal,
    original_stty_state: Vec<u8>,
    start_line: u16,
    desired_rows: u16,
}

impl Screen {
    pub fn open_screen(desired_rows: u16) -> Screen {
        let mut tty = Terminal::open_terminal();
        let current_stty_state = tty.stty(&["-g"]);
        tty.initialize();
        let (_, rows) = tty.winsize().unwrap();
        let visible_choices = min(desired_rows, rows - 1);
        let start_line = rows - visible_choices - 1;
        for _ in 0..visible_choices {
            tty.write(NEWLINE.as_bytes());
        }
        tty.write(&ansi::save_cursor());

        Screen {
            tty: tty,
            original_stty_state: current_stty_state,
            start_line: start_line,
            desired_rows: desired_rows,
        }
    }

    pub fn visible_choices(&self) -> u16 {
        let (_, rows) = self.tty.winsize().unwrap();
        min(self.desired_rows, rows - 1)
    }

    pub fn width(&self) -> u16 {
        let (cols, _) = self.tty.winsize().unwrap();
        cols
    }

    fn restore_tty(&mut self) {
        self.tty.stty(&[&String::from_utf8(self.original_stty_state.clone()).unwrap()]);
    }

    fn reset_cursor(&mut self) {
        self.tty.write(&ansi::restore_cursor());

        // Writing this carriage return works around a rendering bug in
        // Neovim's terminal emulation. Without it, the cursor flies
        // around all over the place while typing, and the input prompt
        // is not rendered correctly. I have no idea why this workaround
        // is effective, or why others are not, or what the root cause
        // of the issue is (likely something involving save/restore
        // support).
        self.tty.write(b"\r");

        let num_lines = self.visible_choices();
        self.tty.write(&ansi::cursor_up(num_lines));
    }

    pub fn move_cursor_to_prompt_line(&mut self, col: u16) {
        self.reset_cursor();
        self.tty.write(&ansi::cursor_right(col));
    }

    pub fn blank_screen(&mut self) {
        self.reset_cursor();
        let blank_line = repeat(' ').take(self.width() as usize).collect::<String>();
        for _ in 0..self.visible_choices() + 1 {
            self.tty.write(blank_line.as_bytes());
        }
        self.reset_cursor();
    }

    pub fn blank_entire_screen(&mut self){
        self.tty.write(&ansi::blank_screen());
    }

    pub fn show_cursor(&mut self) {
        self.tty.write(&ansi::show_cursor());
        self.tty.flush();
    }

    pub fn hide_cursor(&mut self) {
        self.tty.write(&ansi::hide_cursor());
    }

    pub fn write(&mut self, s: &str) {
        self.tty.write(s.as_bytes());
    }

    pub fn write_red_inverted(&mut self, s: &str) {
        self.tty.write(&ansi::red());
        self.tty.write(&ansi::inverse());
        self.tty.write(s.as_bytes());
        self.tty.write(&ansi::reset());
    }

    pub fn write_red(&mut self, s: &str) {
        self.tty.write(&ansi::red());
        self.tty.write(s.as_bytes());
        self.tty.write(&ansi::reset());
    }

    pub fn write_inverted(&mut self, s: &str) {
        self.tty.write(&ansi::inverse());
        self.tty.write(s.as_bytes());
        self.tty.write(&ansi::reset());
    }

    // Return all buffered keystrokes, or the next key if buffer is empty.
    pub fn get_buffered_keys(&mut self) -> Vec<Key> {
        let mut ret = Vec::new();
        while let Ok(bytes) = self.tty.input.try_recv() {
            ret.extend(bytes);
        }
        while ret.is_empty() {
            let bytes = self.tty.input.recv().unwrap();
            if bytes == vec![0] { // SIGWINCH
                self.blank_entire_screen();
                return vec![Nothing];
            } else {
                ret.extend(bytes);
            }
        }
        Terminal::translate_bytes(ret.clone())
    }

    fn more_bytes_needed(bytes: &[u8]) -> bool {
        if let Err(_) = str::from_utf8(bytes) { true } else { false }
    }
}

impl Drop for Screen {
    fn drop(&mut self) {
        self.restore_tty();
    }
}

struct Terminal {
    input: Receiver<Vec<u8>>,
    input_fd: i32,
    output: File,
    output_buf: Vec<u8>,
}

static mut GLOBAL_TX: *const Arc<Mutex<Sender<Vec<u8>>>> = 0 as *const Arc<Mutex<Sender<Vec<u8>>>>;

fn set_global_tx(tx: Sender<Vec<u8>>) {
    unsafe {
        let singleton = Arc::new(Mutex::new(tx));
        GLOBAL_TX = mem::transmute(Box::new(singleton));
    }
}

fn get_global_tx() -> Sender<Vec<u8>> {
    let singleton = unsafe { (*GLOBAL_TX).clone() };
    let tx = singleton.lock().unwrap();
    tx.clone()
}

unsafe extern fn sigwinch_handler(signum: c_int) {
    if signum == SIGWINCH {
        get_global_tx().send([0].to_vec()).unwrap();
    }
}

fn register_sigwinch_handler() {
    unsafe {
        let mut action: sigaction = mem::zeroed();
        action.sa_flags = SA_SIGINFO;
        action.sa_sigaction = sigwinch_handler as sighandler_t;
        sigaction(SIGWINCH, &action, ptr::null_mut());
    }
}

impl Terminal {
    fn open_terminal() -> Terminal {
        let term_path = Path::new("/dev/tty");
        let mut input_file = File::open(&term_path).unwrap();
        let output_file = OpenOptions::new().write(true).open(&term_path).unwrap();
        let input_fd = input_file.as_raw_fd();
        let (tx, rx) = mpsc::channel();
        set_global_tx(tx);

        register_sigwinch_handler();

        thread::spawn(move || {
            loop {
                let mut buf = [0; 255];
                let tx = get_global_tx();
                if let Ok(length) = input_file.read(&mut buf) {
                    tx.send(buf[0..length].to_vec()).unwrap();
                } else {
                    tx.send([0].to_vec()).unwrap();
                }
            }
        });
        Terminal {
            input: rx,
            input_fd: input_fd,
            output: output_file,
            output_buf: Vec::new(),
        }
    }

    fn initialize(&mut self) {
        self.stty(&["raw", "-echo", "cbreak", "opost", "onlcr"]);
    }

    fn stty(&mut self, args: &[&str]) -> Vec<u8> {
        extern { fn dup2(src: c_int, dst: c_int) -> c_int;  }
        unsafe {
            // This is a hack until a replacement for InheritFd from old_io is available.
            dup2(self.input_fd, 0);
        }
        let mut process = match Command::new("stty")
            .args(args)
            .stdin(Stdio::inherit())
            .stdout(Stdio::piped())
            .spawn()
        {
            Ok(p) => p,
            Err(e) => panic!("Spawn failed: {}", e),
        };

        let exit = process.wait();

        let mut buf = Vec::new();
        if exit.unwrap().success() {
            process.stdout.as_mut().unwrap().read_to_end(&mut buf).unwrap();
            let mut str = String::from_utf8(buf).unwrap();

            // The output from `stty -g` may include a newline, which we have to strip off. Otherwise,
            // when we go to restore the tty, stty (on some platforms) will fail with an "invalid
            // argument" error.
            ::trim(&mut str);

            str.into_bytes()
        } else {
            process.stderr.as_mut().unwrap().read_to_end(&mut buf).unwrap();
            panic!(String::from_utf8(buf).unwrap());
        }
    }

    fn translate_bytes(bytes: Vec<u8>) -> Vec<Key> {
        const BEGIN_PASTE: &'static [u8] = b"\x1B[200~";
        const END_PASTE: &'static [u8] = b"\x1B[201~";

        if bytes == b"\x1B[A" || bytes == b"\x1BOA" { return vec![Up] };
        if bytes == b"\x1B[B" || bytes == b"\x1BOB" { return vec![Down] };
        if bytes == b"\x1B[Z" { return vec![ShiftTab] };

        let bs = if bytes.starts_with(BEGIN_PASTE) && bytes.ends_with(END_PASTE) {
            let start = BEGIN_PASTE.len();
            let end = bytes.len() - END_PASTE.len();
            bytes[start..end].to_vec()
        } else {
            bytes
        };

        String::from_utf8(bs)
            .unwrap()
            .chars()
            .map(Terminal::translate_char)
            .collect()
    }

    fn translate_char(c: char) -> Key {
        let numeric_char = c as u32;
        if c == '\r' {
            Enter
        } else if numeric_char == 9 {
            Tab
        } else if numeric_char == 127 {
            Backspace
        } else if numeric_char & 96 == 0 && numeric_char <= 128u32 {
            let c = c as u8;
            Control((c + 96u8) as char)
        } else {
            Char(c)
        }
    }

    fn write(&mut self, s: &[u8]) {
        self.output_buf.extend_from_slice(s);
    }

    fn writeln(&mut self, s: &str) {
        self.output_buf.extend_from_slice(s.as_bytes());
        self.output_buf.push(b'\n');
    }

    fn flush(&mut self) {
        self.output.write(&self.output_buf).unwrap();
        self.output_buf.clear();
    }

    fn winsize(&self) -> Option<(u16, u16)> {
        extern { fn ioctl(fd: c_int, request: c_ulong, ...) -> c_int; }
        #[cfg(any(target_os = "macos", target_os = "freebsd", target_os = "netbsd"))]
        const TIOCGWINSZ: c_ulong = 0x40087468;

        #[cfg(any(target_os = "linux", target_os = "android"))]
        const TIOCGWINSZ: c_ulong = 0x00005413;

        #[repr(C)]
        struct TermSize {
            rows: c_ushort,
            cols: c_ushort,
            x: c_ushort,
            y: c_ushort,
        }

        let size = TermSize { rows: 0, cols: 0, x: 0, y: 0 };
        if unsafe { ioctl(self.output.as_raw_fd(), TIOCGWINSZ, &size) } == 0 {
            Some((size.cols as u16, size.rows as u16))
        } else {
            None
        }
    }
}

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

    #[test]
    fn winsize_test() {
        // Travis-CI builds run without a tty, making this test impossible.
        if option_env!("TRAVIS").is_some() {
            // TODO: It should be made obvious from the output that this test was skipped
            return;
        }
        let term = Terminal::open_terminal();
        let (cols, rows) = term.winsize().expect("Failed to get window size!");
        // We don't know the window size a priori, but we can at least
        // assert that it is within some kind of sensible range.
        assert!(cols > 40);
        assert!(rows > 40);
        assert!(cols < 1000);
        assert!(rows < 1000);
    }
}
<|endoftext|>


<commit_before>use std::string::String;
use std::vec::Vec;

use byteorder::{BigEndian, WriteBytesExt};

use types::{BERT_LABEL, BertTag};

pub struct Serializer;

pub struct Deserializer {
//    data: Vec<u8>
}

pub trait Serialize<T> {
    fn to_bert(&self, data: T) -> Vec<u8>;
}

trait Deserialize<T> {
    fn to_bert(&self, data: T) -> Vec<u8>;
}

impl Serializer {

    pub fn new() -> Serializer {
        Serializer{}
    }

    pub fn term_to_binary<T>(&self, data: T) -> Vec<u8> where Self: Serialize<T> {
        let mut binary = vec![131u8];
        let serialized_data = self.to_bert(data);
        binary.extend(serialized_data.iter().clone());
        binary
    }

    pub fn generate_term(&self, tag: BertTag, data: Vec<u8>) -> Vec<u8> {
        let mut binary = vec![tag as u8];
        binary.extend(data.iter().clone());
        binary
    }

    pub fn convert_string_to_binary(&self, data: &str) -> Vec<u8> {
        let binary_string = data.as_bytes();
        let binary_length = binary_string.len() as u8;
        let mut binary = vec![0u8, binary_length];
        binary.extend(binary_string.iter().clone());
        binary
    }

    pub fn merge_atoms(&self, atom_1: Vec<u8>, atom_2: Vec<u8>) -> Vec<u8> {
        let mut binary: Vec<u8> = atom_1.clone();
        binary.extend(atom_2.iter().clone());
        binary
    }

    pub fn get_bert_atom(&self) -> Vec<u8> {
        let binary_string = self.convert_string_to_binary(BERT_LABEL);
        self.generate_term(BertTag::Atom, binary_string)
    }
}

impl Serialize<u8> for Serializer {
    fn to_bert(&self, data: u8) -> Vec<u8> {
        self.generate_term(BertTag::SmallInteger, vec![data])
    }
}

impl Serialize<i32> for Serializer {
    fn to_bert(&self, data: i32) -> Vec<u8> {
        let mut binary = vec![];
        binary.write_i32::<BigEndian>(data).unwrap();
        self.generate_term(BertTag::Integer, binary)
    }
}

impl Serialize<f64> for Serializer {
    fn to_bert(&self, data: f64) -> Vec<u8> {
        let string_float: String = data.to_string();
        let mut binary = self.to_bert(string_float);
        //binary.write_f64::<BigEndian>(data).unwrap();
        self.generate_term(BertTag::Float, binary)
    }
}

impl Serialize<String> for Serializer {
    fn to_bert(&self, data: String) -> Vec<u8> {
        let binary_string = self.convert_string_to_binary(&data);
        self.generate_term(BertTag::String, binary_string)
    }
}

impl Serialize<bool> for Serializer {
    fn to_bert(&self, data: bool) -> Vec<u8> {
        let boolean_string = data.to_string();
        let binary_boolean = self.convert_string_to_binary(&boolean_string);

        let bert_atom = self.get_bert_atom();
        let boolean_atom = self.generate_term(BertTag::Atom, binary_boolean);

        self.merge_atoms(bert_atom, boolean_atom)
    }
}

#[cfg(test)]
mod test_serializer {
    use super::{Serializer};

    #[test]
    fn test_serialize_u8() {
        let serializer = Serializer::new();

        assert_eq!(
            serializer.term_to_binary(1u8),
            vec![131u8, 97, 1]
        );

        assert_eq!(
            serializer.term_to_binary(255u8),
            vec![131u8, 97, 255]
        );
    }

    #[test]
    fn test_serialize_i32() {
        let serializer = Serializer::new();

        assert_eq!(
            serializer.term_to_binary(-2147483648),
            vec![131u8, 98, 128, 0, 0, 0]
        );

        assert_eq!(
            serializer.term_to_binary(-1i32),
            vec![131u8, 98, 255, 255, 255, 255]
        );

        assert_eq!(
            serializer.term_to_binary(512i32),
            vec![131u8, 98, 0, 0, 2, 0]
        );

        assert_eq!(
            serializer.term_to_binary(2147483647),
            vec![131u8, 98, 127, 255, 255, 255]
        );
    }

    #[test]
    fn test_serialize_bool() {
        let serializer = Serializer::new();

        assert_eq!(
            serializer.term_to_binary(true),
            vec![
                131u8,
                100, 0, 4,  98, 101, 114, 116,  // "bert" as atom
                100, 0, 4, 116, 114, 117, 101   // "true" as atom
            ]
        );

        assert_eq!(
            serializer.term_to_binary(false),
            vec![
                131u8,
                100, 0, 4, 98, 101, 114, 116,      // "bert" as atom
                100, 0, 5, 102, 97, 108, 115, 101  // "false" as atom
            ]
        );
    }
}
<commit_msg>Implemented support for passed strings by reference<commit_after>use std::string::String;
use std::vec::Vec;

use byteorder::{BigEndian, WriteBytesExt};

use types::{BERT_LABEL, BertTag};

// TODO: Wrap complex BERT types into tuples + fix tests
pub struct Serializer;

pub struct Deserializer {
//    data: Vec<u8>
}

pub trait Serialize<T> {
    fn to_bert(&self, data: T) -> Vec<u8>;
}

trait Deserialize<T> {
    fn to_bert(&self, data: T) -> Vec<u8>;
}

impl Serializer {

    pub fn new() -> Serializer {
        Serializer{}
    }

    pub fn term_to_binary<T>(&self, data: T) -> Vec<u8> where Self: Serialize<T> {
        let mut binary = vec![131u8];
        let serialized_data = self.to_bert(data);
        binary.extend(serialized_data.iter().clone());
        binary
    }

    pub fn generate_term(&self, tag: BertTag, data: Vec<u8>) -> Vec<u8> {
        let mut binary = vec![tag as u8];
        binary.extend(data.iter().clone());
        binary
    }

    pub fn convert_string_to_binary(&self, data: &str) -> Vec<u8> {
        let binary_string = data.as_bytes();
        let binary_length = binary_string.len() as u8;
        let mut binary = vec![0u8, binary_length];
        binary.extend(binary_string.iter().clone());
        binary
    }

    pub fn merge_atoms(&self, atom_1: Vec<u8>, atom_2: Vec<u8>) -> Vec<u8> {
        let mut binary: Vec<u8> = atom_1.clone();
        binary.extend(atom_2.iter().clone());
        binary
    }

    pub fn get_bert_atom(&self) -> Vec<u8> {
        let binary_string = self.convert_string_to_binary(BERT_LABEL);
        self.generate_term(BertTag::Atom, binary_string)
    }
}

impl Serialize<u8> for Serializer {
    fn to_bert(&self, data: u8) -> Vec<u8> {
        self.generate_term(BertTag::SmallInteger, vec![data])
    }
}

impl Serialize<i32> for Serializer {
    fn to_bert(&self, data: i32) -> Vec<u8> {
        let mut binary = vec![];
        binary.write_i32::<BigEndian>(data).unwrap();
        self.generate_term(BertTag::Integer, binary)
    }
}

impl Serialize<f64> for Serializer {
    fn to_bert(&self, data: f64) -> Vec<u8> {
        let string_float: String = data.to_string();
        let binary = self.to_bert(string_float);
        self.generate_term(BertTag::Float, binary)
    }
}

impl Serialize<String> for Serializer {
    fn to_bert(&self, data: String) -> Vec<u8> {
        let binary_string = self.convert_string_to_binary(&data);
        self.generate_term(BertTag::String, binary_string)
    }
}

impl<'a> Serialize<&'a str> for Serializer {
    fn to_bert(&self, data: &'a str) -> Vec<u8> {
        let binary_string = self.convert_string_to_binary(data);
        self.generate_term(BertTag::String, binary_string)
    }
}

impl Serialize<bool> for Serializer {
    fn to_bert(&self, data: bool) -> Vec<u8> {
        let boolean_string = data.to_string();
        let binary_boolean = self.convert_string_to_binary(&boolean_string);

        let bert_atom = self.get_bert_atom();
        let boolean_atom = self.generate_term(BertTag::Atom, binary_boolean);

        self.merge_atoms(bert_atom, boolean_atom)
    }
}

#[cfg(test)]
mod test_serializer {
    use super::{Serializer};

    #[test]
    fn test_serialize_u8() {
        let serializer = Serializer::new();

        assert_eq!(
            serializer.term_to_binary(1u8),
            vec![131u8, 97, 1]
        );

        assert_eq!(
            serializer.term_to_binary(255u8),
            vec![131u8, 97, 255]
        );
    }

    #[test]
    fn test_serialize_i32() {
        let serializer = Serializer::new();

        assert_eq!(
            serializer.term_to_binary(-2147483648),
            vec![131u8, 98, 128, 0, 0, 0]
        );

        assert_eq!(
            serializer.term_to_binary(-1i32),
            vec![131u8, 98, 255, 255, 255, 255]
        );

        assert_eq!(
            serializer.term_to_binary(512i32),
            vec![131u8, 98, 0, 0, 2, 0]
        );

        assert_eq!(
            serializer.term_to_binary(2147483647),
            vec![131u8, 98, 127, 255, 255, 255]
        );
    }

    #[test]
    fn test_serialize_bool() {
        let serializer = Serializer::new();

        assert_eq!(
            serializer.term_to_binary(true),
            vec![
                131u8,
                100, 0, 4,  98, 101, 114, 116,  // "bert" as atom
                100, 0, 4, 116, 114, 117, 101   // "true" as atom
            ]
        );

        assert_eq!(
            serializer.term_to_binary(false),
            vec![
                131u8,
                100, 0, 4, 98, 101, 114, 116,      // "bert" as atom
                100, 0, 5, 102, 97, 108, 115, 101  // "false" as atom
            ]
        );
    }

    #[test]
    fn test_serialize_string() {
        let serializer = Serializer::new();

        // string by value
        assert_eq!(
            serializer.term_to_binary(String::from("test")),
            vec![131u8, 107, 0, 4, 116, 101, 115, 116]
        );

        // string by reference
        assert_eq!(
            serializer.term_to_binary("test"),
            vec![131u8, 107, 0, 4, 116, 101, 115, 116]
        );
    }
}
<|endoftext|>


<commit_before>//! Functions and data-structures for working with descriptor tables.
use crate::segmentation::SegmentSelector;
use core::fmt;
use core::mem::size_of;

/// A struct describing a pointer to a descriptor table (GDT / IDT).
/// This is in a format suitable for giving to 'lgdt' or 'lidt'.
#[repr(C, packed)]
pub struct DescriptorTablePointer<Entry> {
    /// Size of the DT.
    pub limit: u16,
    /// Pointer to the memory region containing the DT.
    pub base: *const Entry,
}

impl<T> Default for DescriptorTablePointer<T> {
    fn default() -> DescriptorTablePointer<T> {
        DescriptorTablePointer {
            limit: 0,
            base: core::ptr::null(),
        }
    }
}

impl<T> DescriptorTablePointer<T> {
    pub fn new(tbl: &T) -> Self {
        // GDT, LDT, and IDT all expect the limit to be set to "one less".
        // See Intel 3a, Section 3.5.1 "Segment Descriptor Tables" and
        // Section 6.10 "Interrupt Descriptor Table (IDT)".
        let len = size_of::<T>() - 1;
        assert!(len < 0x10000);
        DescriptorTablePointer {
            base: tbl as *const T,
            limit: len as u16,
        }
    }

    pub fn new_from_slice(slice: &[T]) -> Self {
        // GDT, LDT, and IDT all expect the limit to be set to "one less".
        // See Intel 3a, Section 3.5.1 "Segment Descriptor Tables" and
        // Section 6.10 "Interrupt Descriptor Table (IDT)".
        let len = slice.len() * size_of::<T>() - 1;
        assert!(len < 0x10000);
        DescriptorTablePointer {
            base: slice.as_ptr(),
            limit: len as u16,
        }
    }
}

impl<T> fmt::Debug for DescriptorTablePointer<T> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "DescriptorTablePointer ({} {:?})", { self.limit }, { self.base } )
    }
}

/// Load the GDTR register with the specified base and limit.
///
/// # Safety
/// Needs CPL 0.
pub unsafe fn lgdt<T>(gdt: &DescriptorTablePointer<T>) {
    llvm_asm!("lgdt ($0)" :: "r" (gdt) : "memory");
}

/// Retrieve base and limit from the GDTR register.
///
/// # Safety
/// Needs CPL 0.
pub unsafe fn sgdt<T>(idt: &mut DescriptorTablePointer<T>) {
    llvm_asm!("sgdt ($0)" : "+r" (idt as *mut DescriptorTablePointer<T>) :: "memory");
}

/// Loads the segment selector into the selector field of the local
/// descriptor table register (LDTR).
///
/// After the segment selector is loaded in the LDTR,
/// the processor uses the segment selector to locate
/// the segment descriptor for the LDT in the global
/// descriptor table (GDT).
///
/// # Safety
/// Needs CPL 0.
pub unsafe fn load_ldtr(selector: SegmentSelector) {
    llvm_asm!("lldt $0" :: "r" (selector.bits()) : "memory");
}

/// Returns the segment selector from the local descriptor table register (LDTR).
///
/// The returned segment selector points to the segment descriptor
/// (located in the GDT) for the current LDT.
///
/// # Safety
/// Needs CPL 0.
pub unsafe fn ldtr() -> SegmentSelector {
    let selector: u16;
    llvm_asm!("sldt $0" : "+r"(selector));
    SegmentSelector::from_raw(selector)
}

/// Load the IDTR register with the specified base and limit.
///
/// # Safety
/// Needs CPL 0.
pub unsafe fn lidt<T>(idt: &DescriptorTablePointer<T>) {
    llvm_asm!("lidt ($0)" :: "r" (idt) : "memory");
}

/// Retrieve base and limit from the IDTR register.
///
/// # Safety
/// Needs CPL 0.
pub unsafe fn sidt<T>(idt: &mut DescriptorTablePointer<T>) {
    llvm_asm!("sidt ($0)" : "+r" (idt as *mut DescriptorTablePointer<T>) :: "memory");
}
<commit_msg>Add tests for sgdt, sidt<commit_after>//! Functions and data-structures for working with descriptor tables.
use crate::segmentation::SegmentSelector;
use core::fmt;
use core::mem::size_of;

/// A struct describing a pointer to a descriptor table (GDT / IDT).
/// This is in a format suitable for giving to 'lgdt' or 'lidt'.
#[repr(C, packed)]
pub struct DescriptorTablePointer<Entry> {
    /// Size of the DT.
    pub limit: u16,
    /// Pointer to the memory region containing the DT.
    pub base: *const Entry,
}

impl<T> Default for DescriptorTablePointer<T> {
    fn default() -> DescriptorTablePointer<T> {
        DescriptorTablePointer {
            limit: 0,
            base: core::ptr::null(),
        }
    }
}

impl<T> DescriptorTablePointer<T> {
    pub fn new(tbl: &T) -> Self {
        // GDT, LDT, and IDT all expect the limit to be set to "one less".
        // See Intel 3a, Section 3.5.1 "Segment Descriptor Tables" and
        // Section 6.10 "Interrupt Descriptor Table (IDT)".
        let len = size_of::<T>() - 1;
        assert!(len < 0x10000);
        DescriptorTablePointer {
            base: tbl as *const T,
            limit: len as u16,
        }
    }

    pub fn new_from_slice(slice: &[T]) -> Self {
        // GDT, LDT, and IDT all expect the limit to be set to "one less".
        // See Intel 3a, Section 3.5.1 "Segment Descriptor Tables" and
        // Section 6.10 "Interrupt Descriptor Table (IDT)".
        let len = slice.len() * size_of::<T>() - 1;
        assert!(len < 0x10000);
        DescriptorTablePointer {
            base: slice.as_ptr(),
            limit: len as u16,
        }
    }
}

impl<T> fmt::Debug for DescriptorTablePointer<T> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "DescriptorTablePointer ({} {:?})", { self.limit }, { self.base } )
    }
}

/// Load the GDTR register with the specified base and limit.
///
/// # Safety
/// Needs CPL 0.
pub unsafe fn lgdt<T>(gdt: &DescriptorTablePointer<T>) {
    llvm_asm!("lgdt ($0)" :: "r" (gdt) : "memory");
}

/// Retrieve base and limit from the GDTR register.
///
/// # Safety
/// Needs CPL 0.
pub unsafe fn sgdt<T>(idt: &mut DescriptorTablePointer<T>) {
    llvm_asm!("sgdt ($0)" : "+r" (idt as *mut DescriptorTablePointer<T>) :: "memory");
}

/// Loads the segment selector into the selector field of the local
/// descriptor table register (LDTR).
///
/// After the segment selector is loaded in the LDTR,
/// the processor uses the segment selector to locate
/// the segment descriptor for the LDT in the global
/// descriptor table (GDT).
///
/// # Safety
/// Needs CPL 0.
pub unsafe fn load_ldtr(selector: SegmentSelector) {
    llvm_asm!("lldt $0" :: "r" (selector.bits()) : "memory");
}

/// Returns the segment selector from the local descriptor table register (LDTR).
///
/// The returned segment selector points to the segment descriptor
/// (located in the GDT) for the current LDT.
///
/// # Safety
/// Needs CPL 0.
pub unsafe fn ldtr() -> SegmentSelector {
    let selector: u16;
    llvm_asm!("sldt $0" : "+r"(selector));
    SegmentSelector::from_raw(selector)
}

/// Load the IDTR register with the specified base and limit.
///
/// # Safety
/// Needs CPL 0.
pub unsafe fn lidt<T>(idt: &DescriptorTablePointer<T>) {
    llvm_asm!("lidt ($0)" :: "r" (idt) : "memory");
}

/// Retrieve base and limit from the IDTR register.
///
/// # Safety
/// Needs CPL 0.
pub unsafe fn sidt<T>(idt: &mut DescriptorTablePointer<T>) {
    llvm_asm!("sidt ($0)" : "+r" (idt as *mut DescriptorTablePointer<T>) :: "memory");
}

#[cfg(all(test, feature = "utest"))]
mod test {
    use super::*;

    #[test]
    fn check_sgdt() {
        let gdtr: super::DescriptorTablePointer<u64> = Default::default();
        gdtr.limit = 0xdead;
        gdtr.base = 0xbadc0de as *mut u64;
        unsafe {
            sgdt(&mut gdtr);
        }
        assert_ne!(gdtr.limit, 0);
        assert_ne!(gdtr.base, core::ptr::null_mut());
        assert_ne!(gdtr.limit, 0xdead);
        assert_ne!(gdtr.base, 0xbadc0de);
    }

    #[test]
    fn check_sidt() {
        let gdtr: super::DescriptorTablePointer<u64> = Default::default();
        idtr.limit = 0xdead;
        idtr.base = 0xbadc0de as *mut u64;
        unsafe {
            sidt(&mut idtr);
        }
        assert_ne!(idtr.limit, 0);
        assert_ne!(idtr.base, core::ptr::null_mut());
        assert_ne!(idtr.limit, 0xdead);
        assert_ne!(idtr.base, 0xbadc0de);
    }
}
<|endoftext|>


<commit_before>use crate::loom::sync::atomic::AtomicUsize;

use std::fmt;
use std::sync::atomic::Ordering::{AcqRel, Acquire, Release};
use std::usize;

pub(super) struct State {
    val: AtomicUsize,
}

/// Current state value
#[derive(Copy, Clone)]
pub(super) struct Snapshot(usize);

/// The task is currently being run.
const RUNNING: usize = 0b00_0001;

/// The task has been notified by a waker.
const NOTIFIED: usize = 0b00_0010;

/// The task is complete.
///
/// Once this bit is set, it is never unset
const COMPLETE: usize = 0b00_0100;

/// The primary task handle has been dropped.
const RELEASED: usize = 0b00_1000;

/// The join handle is still around
const JOIN_INTEREST: usize = 0b01_0000;

/// A join handle waker has been set
const JOIN_WAKER: usize = 0b10_0000;

/// The task has been forcibly canceled.
const CANCELLED: usize = 0b100_0000;

/// All bits
const LIFECYCLE_MASK: usize =
    RUNNING | NOTIFIED | COMPLETE | RELEASED | JOIN_INTEREST | JOIN_WAKER | CANCELLED;

/// Bits used by the waker ref count portion of the state.
///
/// Ref counts only cover **wakers**. Other handles are tracked with other state
/// bits.
const WAKER_COUNT_MASK: usize = usize::MAX - LIFECYCLE_MASK;

/// Number of positions to shift the ref count
const WAKER_COUNT_SHIFT: usize = WAKER_COUNT_MASK.count_zeros() as usize;

/// One ref count
const WAKER_ONE: usize = 1 << WAKER_COUNT_SHIFT;

/// Initial state
const INITIAL_STATE: usize = NOTIFIED;

/// All transitions are performed via RMW operations. This establishes an
/// unambiguous modification order.
impl State {
    /// Starts with a ref count of 2
    pub(super) fn new_joinable() -> State {
        State {
            val: AtomicUsize::new(INITIAL_STATE | JOIN_INTEREST),
        }
    }

    /// Load the current state, establishes `Acquire` ordering.
    pub(super) fn load(&self) -> Snapshot {
        Snapshot(self.val.load(Acquire))
    }

    /// Transition a task to the `Running` state.
    ///
    /// Returns a snapshot of the state **after** the transition.
    pub(super) fn transition_to_running(&self) -> Snapshot {
        const DELTA: usize = RUNNING | NOTIFIED;

        let prev = Snapshot(self.val.fetch_xor(DELTA, Acquire));
        debug_assert!(prev.is_notified());

        if prev.is_running() {
            // We were signalled to cancel
            //
            // Apply the state
            let prev = self.val.fetch_or(CANCELLED, AcqRel);
            return Snapshot(prev | CANCELLED);
        }

        debug_assert!(!prev.is_running());

        let next = Snapshot(prev.0 ^ DELTA);

        debug_assert!(next.is_running());
        debug_assert!(!next.is_notified());

        next
    }

    /// Transition the task from `Running` -> `Idle`.
    ///
    /// Returns a snapshot of the state **after** the transition.
    pub(super) fn transition_to_idle(&self) -> Snapshot {
        const DELTA: usize = RUNNING;

        let prev = Snapshot(self.val.fetch_xor(DELTA, AcqRel));

        if !prev.is_running() {
            // We were signaled to cancel.
            //
            // Apply the state
            let prev = self.val.fetch_or(CANCELLED, AcqRel);
            return Snapshot(prev | CANCELLED);
        }

        let next = Snapshot(prev.0 ^ DELTA);

        debug_assert!(!next.is_running());

        next
    }

    /// Transition the task from `Running` -> `Complete`.
    ///
    /// Returns a snapshot of the state **after** the transition.
    pub(super) fn transition_to_complete(&self) -> Snapshot {
        const DELTA: usize = RUNNING | COMPLETE;

        let prev = Snapshot(self.val.fetch_xor(DELTA, AcqRel));

        debug_assert!(!prev.is_complete());

        let next = Snapshot(prev.0 ^ DELTA);

        debug_assert!(next.is_complete());

        next
    }

    /// Transition the task from `Running` -> `Released`.
    ///
    /// Returns a snapshot of the state **after** the transition.
    pub(super) fn transition_to_released(&self) -> Snapshot {
        const DELTA: usize = RUNNING | COMPLETE | RELEASED;

        let prev = Snapshot(self.val.fetch_xor(DELTA, AcqRel));

        debug_assert!(prev.is_running());
        debug_assert!(!prev.is_complete());
        debug_assert!(!prev.is_released());

        let next = Snapshot(prev.0 ^ DELTA);

        debug_assert!(!next.is_running());
        debug_assert!(next.is_complete());
        debug_assert!(next.is_released());

        next
    }

    /// Transition the task to the canceled state.
    ///
    /// Returns the snapshot of the state **after** the transition **if** the
    /// transition was made successfully
    ///
    /// # States
    ///
    /// - Notifed: task may be in a queue, caller must not release.
    /// - Running: cannot drop. The poll handle will handle releasing.
    /// - Other prior states do not require cancellation.
    ///
    /// If the task has been notified, then it may still be in a queue. The
    /// caller must not release the task.
    pub(super) fn transition_to_canceled_from_queue(&self) -> Snapshot {
        let prev = Snapshot(self.val.fetch_or(CANCELLED, AcqRel));

        debug_assert!(!prev.is_complete());
        debug_assert!(!prev.is_running() || prev.is_notified());

        Snapshot(prev.0 | CANCELLED)
    }

    pub(super) fn transition_to_canceled_from_list(&self) -> Option<Snapshot> {
        let mut prev = self.load();

        loop {
            if !prev.is_active() {
                return None;
            }

            let mut next = prev;

            // Use the running flag to signal cancellation
            if prev.is_running() {
                next.0 -= RUNNING;
                next.0 |= NOTIFIED;
            } else if prev.is_notified() {
                next.0 += RUNNING;
                next.0 |= NOTIFIED;
            } else {
                next.0 |= CANCELLED;
            }

            let res = self.val.compare_exchange(prev.0, next.0, AcqRel, Acquire);

            match res {
                Ok(_) if next.is_canceled() => return Some(next),
                Ok(_) => return None,
                Err(actual) => {
                    prev = Snapshot(actual);
                }
            }
        }
    }

    /// Final transition to `Released`. Called when primary task handle is
    /// dropped. This is roughly a "ref decrement" operation.
    ///
    /// Returns a snapshot of the state **after** the transition.
    pub(super) fn release_task(&self) -> Snapshot {
        use crate::loom::sync::atomic;

        const DELTA: usize = RELEASED;

        let prev = Snapshot(self.val.fetch_or(DELTA, Release));

        debug_assert!(!prev.is_released());
        debug_assert!(prev.is_terminal(), "state = {:?}", prev);

        let next = Snapshot(prev.0 | DELTA);

        debug_assert!(next.is_released());

        if next.is_final_ref() || (next.has_join_waker() && !next.is_join_interested()) {
            // The final reference to the task was dropped, the caller must free the
            // memory. Establish an acquire ordering.
            atomic::fence(Acquire);
        }

        next
    }

    /// Transition the state to `Scheduled`.
    ///
    /// Returns `true` if the task needs to be submitted to the pool for
    /// execution
    pub(super) fn transition_to_notified(&self) -> bool {
        const MASK: usize = RUNNING | NOTIFIED | COMPLETE | CANCELLED;

        let prev = self.val.fetch_or(NOTIFIED, Release);
        prev & MASK == 0
    }

    /// Optimistically try to swap the state assuming the join handle is
    /// __immediately__ dropped on spawn
    pub(super) fn drop_join_handle_fast(&self) -> bool {
        use std::sync::atomic::Ordering::Relaxed;

        // Relaxed is acceptable as if this function is called and succeeds,
        // then nothing has been done w/ the join handle.
        //
        // The moment the join handle is used (polled), the `JOIN_WAKER` flag is
        // set, at which point the CAS will fail.
        //
        // Given this, there is no risk if this operation is reordered.
        self.val
            .compare_exchange_weak(
                INITIAL_STATE | JOIN_INTEREST,
                INITIAL_STATE,
                Relaxed,
                Relaxed,
            )
            .is_ok()
    }

    /// The join handle has completed by reading the output
    ///
    /// Returns a snapshot of the state **after** the transition.
    pub(super) fn complete_join_handle(&self) -> Snapshot {
        use crate::loom::sync::atomic;

        const DELTA: usize = JOIN_INTEREST;

        let prev = Snapshot(self.val.fetch_sub(DELTA, Release));

        debug_assert!(prev.is_join_interested());

        let next = Snapshot(prev.0 - DELTA);

        if !next.is_final_ref() {
            return next;
        }

        atomic::fence(Acquire);

        next
    }

    /// The join handle is being dropped, this fails if the task has been
    /// completed and the output must be dropped first then
    /// `complete_join_handle` should be called.
    ///
    /// Returns a snapshot of the state **after** the transition.
    pub(super) fn drop_join_handle_slow(&self) -> Result<Snapshot, Snapshot> {
        const MASK: usize = COMPLETE | CANCELLED;

        let mut prev = self.val.load(Acquire);

        loop {
            // Once the complete bit is set, it is never unset.
            if prev & MASK != 0 {
                return Err(Snapshot(prev));
            }

            debug_assert!(prev & JOIN_INTEREST == JOIN_INTEREST);

            let next = (prev - JOIN_INTEREST) & !JOIN_WAKER;

            let res = self.val.compare_exchange(prev, next, AcqRel, Acquire);

            match res {
                Ok(_) => {
                    return Ok(Snapshot(next));
                }
                Err(actual) => {
                    prev = actual;
                }
            }
        }
    }

    /// Store the join waker.
    pub(super) fn store_join_waker(&self) -> Snapshot {
        use crate::loom::sync::atomic;

        const DELTA: usize = JOIN_WAKER;

        let prev = Snapshot(self.val.fetch_xor(DELTA, Release));

        debug_assert!(!prev.has_join_waker());

        let next = Snapshot(prev.0 ^ DELTA);

        debug_assert!(next.has_join_waker());

        if next.is_complete() {
            atomic::fence(Acquire);
        }

        next
    }

    pub(super) fn unset_waker(&self) -> Snapshot {
        const MASK: usize = COMPLETE | CANCELLED;

        let mut prev = self.val.load(Acquire);

        loop {
            // Once the `COMPLETE` bit is set, it is never unset
            if prev & MASK != 0 {
                return Snapshot(prev);
            }

            debug_assert!(Snapshot(prev).has_join_waker());

            let next = prev - JOIN_WAKER;

            let res = self.val.compare_exchange(prev, next, AcqRel, Acquire);

            match res {
                Ok(_) => return Snapshot(next),
                Err(actual) => {
                    prev = actual;
                }
            }
        }
    }

    pub(super) fn ref_inc(&self) {
        use std::process;
        use std::sync::atomic::Ordering::Relaxed;

        // Using a relaxed ordering is alright here, as knowledge of the
        // original reference prevents other threads from erroneously deleting
        // the object.
        //
        // As explained in the [Boost documentation][1], Increasing the
        // reference counter can always be done with memory_order_relaxed: New
        // references to an object can only be formed from an existing
        // reference, and passing an existing reference from one thread to
        // another must already provide any required synchronization.
        //
        // [1]: (www.boost.org/doc/libs/1_55_0/doc/html/atomic/usage_examples.html)
        let prev = self.val.fetch_add(WAKER_ONE, Relaxed);

        // If the reference count overflowed, abort.
        if prev > isize::max_value() as usize {
            process::abort();
        }
    }

    /// Returns `true` if the task should be released.
    pub(super) fn ref_dec(&self) -> bool {
        use crate::loom::sync::atomic;

        let prev = self.val.fetch_sub(WAKER_ONE, Release);
        let next = Snapshot(prev - WAKER_ONE);

        if next.is_final_ref() {
            atomic::fence(Acquire);
        }

        next.is_final_ref()
    }
}

impl Snapshot {
    pub(super) fn is_running(self) -> bool {
        self.0 & RUNNING == RUNNING
    }

    pub(super) fn is_notified(self) -> bool {
        self.0 & NOTIFIED == NOTIFIED
    }

    pub(super) fn is_released(self) -> bool {
        self.0 & RELEASED == RELEASED
    }

    pub(super) fn is_complete(self) -> bool {
        self.0 & COMPLETE == COMPLETE
    }

    pub(super) fn is_canceled(self) -> bool {
        self.0 & CANCELLED == CANCELLED
    }

    /// Used during normal runtime.
    pub(super) fn is_active(self) -> bool {
        self.0 & (COMPLETE | CANCELLED) == 0
    }

    /// Used before dropping the task
    pub(super) fn is_terminal(self) -> bool {
        // When both the notified & running flags are set, the task was canceled
        // after being notified, before it was run.
        //
        // There is a race where:
        // - The task state transitions to notified
        // - The global queue is shutdown
        // - The waker attempts to push into the global queue and fails.
        // - The waker holds the last reference to the task, thus drops it.
        //
        // In this scenario, the cancelled bit will never get set.
        !self.is_active() || (self.is_notified() && self.is_running())
    }

    pub(super) fn is_join_interested(self) -> bool {
        self.0 & JOIN_INTEREST == JOIN_INTEREST
    }

    pub(super) fn has_join_waker(self) -> bool {
        self.0 & JOIN_WAKER == JOIN_WAKER
    }

    pub(super) fn is_final_ref(self) -> bool {
        const MASK: usize = WAKER_COUNT_MASK | RELEASED | JOIN_INTEREST;

        (self.0 & MASK) == RELEASED
    }
}

impl fmt::Debug for State {
    fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
        use std::sync::atomic::Ordering::SeqCst;

        let snapshot = Snapshot(self.val.load(SeqCst));

        fmt.debug_struct("State")
            .field("snapshot", &snapshot)
            .finish()
    }
}

impl fmt::Debug for Snapshot {
    fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
        fmt.debug_struct("Snapshot")
            .field("is_running", &self.is_running())
            .field("is_notified", &self.is_notified())
            .field("is_released", &self.is_released())
            .field("is_complete", &self.is_complete())
            .field("is_canceled", &self.is_canceled())
            .field("is_join_interested", &self.is_join_interested())
            .field("has_join_waker", &self.has_join_waker())
            .field("is_final_ref", &self.is_final_ref())
            .finish()
    }
}
<commit_msg>rt: use release ordering in `drop_join_handle_fast` (#2044)<commit_after>use crate::loom::sync::atomic::AtomicUsize;

use std::fmt;
use std::sync::atomic::Ordering::{AcqRel, Acquire, Release};
use std::usize;

pub(super) struct State {
    val: AtomicUsize,
}

/// Current state value
#[derive(Copy, Clone)]
pub(super) struct Snapshot(usize);

/// The task is currently being run.
const RUNNING: usize = 0b00_0001;

/// The task has been notified by a waker.
const NOTIFIED: usize = 0b00_0010;

/// The task is complete.
///
/// Once this bit is set, it is never unset
const COMPLETE: usize = 0b00_0100;

/// The primary task handle has been dropped.
const RELEASED: usize = 0b00_1000;

/// The join handle is still around
const JOIN_INTEREST: usize = 0b01_0000;

/// A join handle waker has been set
const JOIN_WAKER: usize = 0b10_0000;

/// The task has been forcibly canceled.
const CANCELLED: usize = 0b100_0000;

/// All bits
const LIFECYCLE_MASK: usize =
    RUNNING | NOTIFIED | COMPLETE | RELEASED | JOIN_INTEREST | JOIN_WAKER | CANCELLED;

/// Bits used by the waker ref count portion of the state.
///
/// Ref counts only cover **wakers**. Other handles are tracked with other state
/// bits.
const WAKER_COUNT_MASK: usize = usize::MAX - LIFECYCLE_MASK;

/// Number of positions to shift the ref count
const WAKER_COUNT_SHIFT: usize = WAKER_COUNT_MASK.count_zeros() as usize;

/// One ref count
const WAKER_ONE: usize = 1 << WAKER_COUNT_SHIFT;

/// Initial state
const INITIAL_STATE: usize = NOTIFIED;

/// All transitions are performed via RMW operations. This establishes an
/// unambiguous modification order.
impl State {
    /// Starts with a ref count of 2
    pub(super) fn new_joinable() -> State {
        State {
            val: AtomicUsize::new(INITIAL_STATE | JOIN_INTEREST),
        }
    }

    /// Load the current state, establishes `Acquire` ordering.
    pub(super) fn load(&self) -> Snapshot {
        Snapshot(self.val.load(Acquire))
    }

    /// Transition a task to the `Running` state.
    ///
    /// Returns a snapshot of the state **after** the transition.
    pub(super) fn transition_to_running(&self) -> Snapshot {
        const DELTA: usize = RUNNING | NOTIFIED;

        let prev = Snapshot(self.val.fetch_xor(DELTA, Acquire));
        debug_assert!(prev.is_notified());

        if prev.is_running() {
            // We were signalled to cancel
            //
            // Apply the state
            let prev = self.val.fetch_or(CANCELLED, AcqRel);
            return Snapshot(prev | CANCELLED);
        }

        debug_assert!(!prev.is_running());

        let next = Snapshot(prev.0 ^ DELTA);

        debug_assert!(next.is_running());
        debug_assert!(!next.is_notified());

        next
    }

    /// Transition the task from `Running` -> `Idle`.
    ///
    /// Returns a snapshot of the state **after** the transition.
    pub(super) fn transition_to_idle(&self) -> Snapshot {
        const DELTA: usize = RUNNING;

        let prev = Snapshot(self.val.fetch_xor(DELTA, AcqRel));

        if !prev.is_running() {
            // We were signaled to cancel.
            //
            // Apply the state
            let prev = self.val.fetch_or(CANCELLED, AcqRel);
            return Snapshot(prev | CANCELLED);
        }

        let next = Snapshot(prev.0 ^ DELTA);

        debug_assert!(!next.is_running());

        next
    }

    /// Transition the task from `Running` -> `Complete`.
    ///
    /// Returns a snapshot of the state **after** the transition.
    pub(super) fn transition_to_complete(&self) -> Snapshot {
        const DELTA: usize = RUNNING | COMPLETE;

        let prev = Snapshot(self.val.fetch_xor(DELTA, AcqRel));

        debug_assert!(!prev.is_complete());

        let next = Snapshot(prev.0 ^ DELTA);

        debug_assert!(next.is_complete());

        next
    }

    /// Transition the task from `Running` -> `Released`.
    ///
    /// Returns a snapshot of the state **after** the transition.
    pub(super) fn transition_to_released(&self) -> Snapshot {
        const DELTA: usize = RUNNING | COMPLETE | RELEASED;

        let prev = Snapshot(self.val.fetch_xor(DELTA, AcqRel));

        debug_assert!(prev.is_running());
        debug_assert!(!prev.is_complete());
        debug_assert!(!prev.is_released());

        let next = Snapshot(prev.0 ^ DELTA);

        debug_assert!(!next.is_running());
        debug_assert!(next.is_complete());
        debug_assert!(next.is_released());

        next
    }

    /// Transition the task to the canceled state.
    ///
    /// Returns the snapshot of the state **after** the transition **if** the
    /// transition was made successfully
    ///
    /// # States
    ///
    /// - Notifed: task may be in a queue, caller must not release.
    /// - Running: cannot drop. The poll handle will handle releasing.
    /// - Other prior states do not require cancellation.
    ///
    /// If the task has been notified, then it may still be in a queue. The
    /// caller must not release the task.
    pub(super) fn transition_to_canceled_from_queue(&self) -> Snapshot {
        let prev = Snapshot(self.val.fetch_or(CANCELLED, AcqRel));

        debug_assert!(!prev.is_complete());
        debug_assert!(!prev.is_running() || prev.is_notified());

        Snapshot(prev.0 | CANCELLED)
    }

    pub(super) fn transition_to_canceled_from_list(&self) -> Option<Snapshot> {
        let mut prev = self.load();

        loop {
            if !prev.is_active() {
                return None;
            }

            let mut next = prev;

            // Use the running flag to signal cancellation
            if prev.is_running() {
                next.0 -= RUNNING;
                next.0 |= NOTIFIED;
            } else if prev.is_notified() {
                next.0 += RUNNING;
                next.0 |= NOTIFIED;
            } else {
                next.0 |= CANCELLED;
            }

            let res = self.val.compare_exchange(prev.0, next.0, AcqRel, Acquire);

            match res {
                Ok(_) if next.is_canceled() => return Some(next),
                Ok(_) => return None,
                Err(actual) => {
                    prev = Snapshot(actual);
                }
            }
        }
    }

    /// Final transition to `Released`. Called when primary task handle is
    /// dropped. This is roughly a "ref decrement" operation.
    ///
    /// Returns a snapshot of the state **after** the transition.
    pub(super) fn release_task(&self) -> Snapshot {
        use crate::loom::sync::atomic;

        const DELTA: usize = RELEASED;

        let prev = Snapshot(self.val.fetch_or(DELTA, Release));

        debug_assert!(!prev.is_released());
        debug_assert!(prev.is_terminal(), "state = {:?}", prev);

        let next = Snapshot(prev.0 | DELTA);

        debug_assert!(next.is_released());

        if next.is_final_ref() || (next.has_join_waker() && !next.is_join_interested()) {
            // The final reference to the task was dropped, the caller must free the
            // memory. Establish an acquire ordering.
            atomic::fence(Acquire);
        }

        next
    }

    /// Transition the state to `Scheduled`.
    ///
    /// Returns `true` if the task needs to be submitted to the pool for
    /// execution
    pub(super) fn transition_to_notified(&self) -> bool {
        const MASK: usize = RUNNING | NOTIFIED | COMPLETE | CANCELLED;

        let prev = self.val.fetch_or(NOTIFIED, Release);
        prev & MASK == 0
    }

    /// Optimistically try to swap the state assuming the join handle is
    /// __immediately__ dropped on spawn
    pub(super) fn drop_join_handle_fast(&self) -> bool {
        use std::sync::atomic::Ordering::Relaxed;

        // Relaxed is acceptable as if this function is called and succeeds,
        // then nothing has been done w/ the join handle.
        //
        // The moment the join handle is used (polled), the `JOIN_WAKER` flag is
        // set, at which point the CAS will fail.
        //
        // Given this, there is no risk if this operation is reordered.
        self.val
            .compare_exchange_weak(
                INITIAL_STATE | JOIN_INTEREST,
                INITIAL_STATE,
                Release,
                Relaxed,
            )
            .is_ok()
    }

    /// The join handle has completed by reading the output
    ///
    /// Returns a snapshot of the state **after** the transition.
    pub(super) fn complete_join_handle(&self) -> Snapshot {
        use crate::loom::sync::atomic;

        const DELTA: usize = JOIN_INTEREST;

        let prev = Snapshot(self.val.fetch_sub(DELTA, Release));

        debug_assert!(prev.is_join_interested());

        let next = Snapshot(prev.0 - DELTA);

        if !next.is_final_ref() {
            return next;
        }

        atomic::fence(Acquire);

        next
    }

    /// The join handle is being dropped, this fails if the task has been
    /// completed and the output must be dropped first then
    /// `complete_join_handle` should be called.
    ///
    /// Returns a snapshot of the state **after** the transition.
    pub(super) fn drop_join_handle_slow(&self) -> Result<Snapshot, Snapshot> {
        const MASK: usize = COMPLETE | CANCELLED;

        let mut prev = self.val.load(Acquire);

        loop {
            // Once the complete bit is set, it is never unset.
            if prev & MASK != 0 {
                return Err(Snapshot(prev));
            }

            debug_assert!(prev & JOIN_INTEREST == JOIN_INTEREST);

            let next = (prev - JOIN_INTEREST) & !JOIN_WAKER;

            let res = self.val.compare_exchange(prev, next, AcqRel, Acquire);

            match res {
                Ok(_) => {
                    return Ok(Snapshot(next));
                }
                Err(actual) => {
                    prev = actual;
                }
            }
        }
    }

    /// Store the join waker.
    pub(super) fn store_join_waker(&self) -> Snapshot {
        use crate::loom::sync::atomic;

        const DELTA: usize = JOIN_WAKER;

        let prev = Snapshot(self.val.fetch_xor(DELTA, Release));

        debug_assert!(!prev.has_join_waker());

        let next = Snapshot(prev.0 ^ DELTA);

        debug_assert!(next.has_join_waker());

        if next.is_complete() {
            atomic::fence(Acquire);
        }

        next
    }

    pub(super) fn unset_waker(&self) -> Snapshot {
        const MASK: usize = COMPLETE | CANCELLED;

        let mut prev = self.val.load(Acquire);

        loop {
            // Once the `COMPLETE` bit is set, it is never unset
            if prev & MASK != 0 {
                return Snapshot(prev);
            }

            debug_assert!(Snapshot(prev).has_join_waker());

            let next = prev - JOIN_WAKER;

            let res = self.val.compare_exchange(prev, next, AcqRel, Acquire);

            match res {
                Ok(_) => return Snapshot(next),
                Err(actual) => {
                    prev = actual;
                }
            }
        }
    }

    pub(super) fn ref_inc(&self) {
        use std::process;
        use std::sync::atomic::Ordering::Relaxed;

        // Using a relaxed ordering is alright here, as knowledge of the
        // original reference prevents other threads from erroneously deleting
        // the object.
        //
        // As explained in the [Boost documentation][1], Increasing the
        // reference counter can always be done with memory_order_relaxed: New
        // references to an object can only be formed from an existing
        // reference, and passing an existing reference from one thread to
        // another must already provide any required synchronization.
        //
        // [1]: (www.boost.org/doc/libs/1_55_0/doc/html/atomic/usage_examples.html)
        let prev = self.val.fetch_add(WAKER_ONE, Relaxed);

        // If the reference count overflowed, abort.
        if prev > isize::max_value() as usize {
            process::abort();
        }
    }

    /// Returns `true` if the task should be released.
    pub(super) fn ref_dec(&self) -> bool {
        use crate::loom::sync::atomic;

        let prev = self.val.fetch_sub(WAKER_ONE, Release);
        let next = Snapshot(prev - WAKER_ONE);

        if next.is_final_ref() {
            atomic::fence(Acquire);
        }

        next.is_final_ref()
    }
}

impl Snapshot {
    pub(super) fn is_running(self) -> bool {
        self.0 & RUNNING == RUNNING
    }

    pub(super) fn is_notified(self) -> bool {
        self.0 & NOTIFIED == NOTIFIED
    }

    pub(super) fn is_released(self) -> bool {
        self.0 & RELEASED == RELEASED
    }

    pub(super) fn is_complete(self) -> bool {
        self.0 & COMPLETE == COMPLETE
    }

    pub(super) fn is_canceled(self) -> bool {
        self.0 & CANCELLED == CANCELLED
    }

    /// Used during normal runtime.
    pub(super) fn is_active(self) -> bool {
        self.0 & (COMPLETE | CANCELLED) == 0
    }

    /// Used before dropping the task
    pub(super) fn is_terminal(self) -> bool {
        // When both the notified & running flags are set, the task was canceled
        // after being notified, before it was run.
        //
        // There is a race where:
        // - The task state transitions to notified
        // - The global queue is shutdown
        // - The waker attempts to push into the global queue and fails.
        // - The waker holds the last reference to the task, thus drops it.
        //
        // In this scenario, the cancelled bit will never get set.
        !self.is_active() || (self.is_notified() && self.is_running())
    }

    pub(super) fn is_join_interested(self) -> bool {
        self.0 & JOIN_INTEREST == JOIN_INTEREST
    }

    pub(super) fn has_join_waker(self) -> bool {
        self.0 & JOIN_WAKER == JOIN_WAKER
    }

    pub(super) fn is_final_ref(self) -> bool {
        const MASK: usize = WAKER_COUNT_MASK | RELEASED | JOIN_INTEREST;

        (self.0 & MASK) == RELEASED
    }
}

impl fmt::Debug for State {
    fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
        use std::sync::atomic::Ordering::SeqCst;

        let snapshot = Snapshot(self.val.load(SeqCst));

        fmt.debug_struct("State")
            .field("snapshot", &snapshot)
            .finish()
    }
}

impl fmt::Debug for Snapshot {
    fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
        fmt.debug_struct("Snapshot")
            .field("is_running", &self.is_running())
            .field("is_notified", &self.is_notified())
            .field("is_released", &self.is_released())
            .field("is_complete", &self.is_complete())
            .field("is_canceled", &self.is_canceled())
            .field("is_join_interested", &self.is_join_interested())
            .field("has_join_waker", &self.has_join_waker())
            .field("is_final_ref", &self.is_final_ref())
            .finish()
    }
}
<|endoftext|>


<commit_before>use path_parsing::SEP;
use std::str;

struct ExtnameCoords {
  word: bool,
  pred: bool,
  dot: bool,
  start: usize,
  end: usize,
}

impl ExtnameCoords {
  pub fn dec(&mut self) {
    self.start -= 1;
    if !self.word {
      self.end -= 1;
    }
  }
}

pub fn extname(pth: &str) -> &str {
  let path = pth.as_bytes();
  let mut extname = ExtnameCoords {
    word: false,
    pred: false,
    dot: false,
    start: path.len(),
    end: path.len(),
  };

  for &item in path.iter().rev() {
    if (item == b'.' && !extname.dot) || item == SEP {
      if item == SEP && extname.word {
        return ""
      }

      if !extname.pred {
        extname.dec();
      }

      if extname.word {
        extname.dot = true;
      }
    } else {
      if extname.dot {
        extname.pred = true;
        break;
      } else {
        extname.word = true;
      }

      if !extname.pred {
        extname.dec()
      }
    }
  }

  if !extname.pred {
    return "";
  }

  str::from_utf8(&path[extname.start..extname.end]).unwrap_or("")
}
<commit_msg>Faster extname<commit_after>use std::str;
use path_parsing::{SEP, find_last_non_sep_pos};

pub fn extname(path: &str) -> &str {
  let end = match find_last_non_sep_pos(path.as_bytes()) {
    Some(pos) => pos + 1,
    _ => return "",
  };
  let bytes = &path.as_bytes()[..end];
  for (pos, c) in bytes.iter().enumerate().rev() {
    match *c {
      b'.' => {
        let prev = bytes.get(pos - 1);
        if pos == end - 1 || prev == None || prev == Some(&SEP) {
          return "";
        } else {
          return &path[pos..end]
        };
      }
      SEP => return "",
      _ => {}
    }
  };
  ""
}
<|endoftext|>


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

//! Interfaces to the operating system provided random number
//! generators.

pub use self::imp::OsRng;

#[cfg(all(unix, not(target_os = "ios")))]
mod imp {
    extern crate libc;

    use self::OsRngInner::*;

    use io::{IoResult, File};
    use path::Path;
    use rand::Rng;
    use rand::reader::ReaderRng;
    use result::Result::{Ok, Err};
    use slice::SliceExt;
    use mem;
    use os::errno;

    #[cfg(all(target_os = "linux",
              any(target_arch = "x86_64", target_arch = "x86", target_arch = "arm")))]
    fn getrandom(buf: &mut [u8]) -> libc::c_long {
        extern "C" {
            fn syscall(number: libc::c_long, ...) -> libc::c_long;
        }

        #[cfg(target_arch = "x86_64")]
        const NR_GETRANDOM: libc::c_long = 318;
        #[cfg(target_arch = "x86")]
        const NR_GETRANDOM: libc::c_long = 355;
        #[cfg(target_arch = "arm")]
        const NR_GETRANDOM: libc::c_long = 384;

        unsafe {
            syscall(NR_GETRANDOM, buf.as_mut_ptr(), buf.len(), 0u)
        }
    }

    #[cfg(not(all(target_os = "linux",
                  any(target_arch = "x86_64", target_arch = "x86", target_arch = "arm"))))]
    fn getrandom(_buf: &mut [u8]) -> libc::c_long { -1 }

    fn getrandom_fill_bytes(v: &mut [u8]) {
        let mut read = 0;
        let len = v.len();
        while read < len {
            let result = getrandom(v[mut read..]);
            if result == -1 {
                let err = errno() as libc::c_int;
                if err == libc::EINTR {
                    continue;
                } else {
                    panic!("unexpected getrandom error: {}", err);
                }
            } else {
                read += result as uint;
            }
        }
    }

    fn getrandom_next_u32() -> u32 {
        let mut buf: [u8, ..4] = [0u8, ..4];
        getrandom_fill_bytes(&mut buf);
        unsafe { mem::transmute::<[u8, ..4], u32>(buf) }
    }

    fn getrandom_next_u64() -> u64 {
        let mut buf: [u8, ..8] = [0u8, ..8];
        getrandom_fill_bytes(&mut buf);
        unsafe { mem::transmute::<[u8, ..8], u64>(buf) }
    }

    #[cfg(all(target_os = "linux",
              any(target_arch = "x86_64", target_arch = "x86", target_arch = "arm")))]
    fn is_getrandom_available() -> bool {
        use sync::atomic::{AtomicBool, INIT_ATOMIC_BOOL, Relaxed};

        static GETRANDOM_CHECKED: AtomicBool = INIT_ATOMIC_BOOL;
        static GETRANDOM_AVAILABLE: AtomicBool = INIT_ATOMIC_BOOL;

        if !GETRANDOM_CHECKED.load(Relaxed) {
            let mut buf: [u8, ..0] = [];
            let result = getrandom(&mut buf);
            let available = if result == -1 {
                let err = errno() as libc::c_int;
                err != libc::ENOSYS
            } else {
                true
            };
            GETRANDOM_AVAILABLE.store(available, Relaxed);
            GETRANDOM_CHECKED.store(true, Relaxed);
            available
        } else {
            GETRANDOM_AVAILABLE.load(Relaxed)
        }
    }

    #[cfg(not(all(target_os = "linux",
                  any(target_arch = "x86_64", target_arch = "x86", target_arch = "arm"))))]
    fn is_getrandom_available() -> bool { false }

    /// A random number generator that retrieves randomness straight from
    /// the operating system. Platform sources:
    ///
    /// - Unix-like systems (Linux, Android, Mac OSX): read directly from
    ///   `/dev/urandom`, or from `getrandom(2)` system call if available.
    /// - Windows: calls `CryptGenRandom`, using the default cryptographic
    ///   service provider with the `PROV_RSA_FULL` type.
    /// - iOS: calls SecRandomCopyBytes as /dev/(u)random is sandboxed.
    ///
    /// This does not block.
    pub struct OsRng {
        inner: OsRngInner,
    }

    enum OsRngInner {
        OsGetrandomRng,
        OsReaderRng(ReaderRng<File>),
    }

    impl OsRng {
        /// Create a new `OsRng`.
        pub fn new() -> IoResult<OsRng> {
            if is_getrandom_available() {
                return Ok(OsRng { inner: OsGetrandomRng });
            }

            let reader = try!(File::open(&Path::new("/dev/urandom")));
            let reader_rng = ReaderRng::new(reader);

            Ok(OsRng { inner: OsReaderRng(reader_rng) })
        }
    }

    impl Rng for OsRng {
        fn next_u32(&mut self) -> u32 {
            match self.inner {
                OsGetrandomRng => getrandom_next_u32(),
                OsReaderRng(ref mut rng) => rng.next_u32(),
            }
        }
        fn next_u64(&mut self) -> u64 {
            match self.inner {
                OsGetrandomRng => getrandom_next_u64(),
                OsReaderRng(ref mut rng) => rng.next_u64(),
            }
        }
        fn fill_bytes(&mut self, v: &mut [u8]) {
            match self.inner {
                OsGetrandomRng => getrandom_fill_bytes(v),
                OsReaderRng(ref mut rng) => rng.fill_bytes(v)
            }
        }
    }
}

#[cfg(target_os = "ios")]
mod imp {
    extern crate libc;

    use io::{IoResult};
    use kinds::marker;
    use mem;
    use os;
    use rand::Rng;
    use result::Result::{Ok};
    use self::libc::{c_int, size_t};
    use slice::SliceExt;

    /// A random number generator that retrieves randomness straight from
    /// the operating system. Platform sources:
    ///
    /// - Unix-like systems (Linux, Android, Mac OSX): read directly from
    ///   `/dev/urandom`, or from `getrandom(2)` system call if available.
    /// - Windows: calls `CryptGenRandom`, using the default cryptographic
    ///   service provider with the `PROV_RSA_FULL` type.
    /// - iOS: calls SecRandomCopyBytes as /dev/(u)random is sandboxed.
    ///
    /// This does not block.
    #[allow(missing_copy_implementations)]
    pub struct OsRng {
        // dummy field to ensure that this struct cannot be constructed outside of this module
        _dummy: (),
    }

    #[repr(C)]
    struct SecRandom;

    #[allow(non_upper_case_globals)]
    static kSecRandomDefault: *const SecRandom = 0 as *const SecRandom;

    #[link(name = "Security", kind = "framework")]
    extern "C" {
        fn SecRandomCopyBytes(rnd: *const SecRandom,
                              count: size_t, bytes: *mut u8) -> c_int;
    }

    impl OsRng {
        /// Create a new `OsRng`.
        pub fn new() -> IoResult<OsRng> {
            Ok(OsRng { _dummy: () })
        }
    }

    impl Rng for OsRng {
        fn next_u32(&mut self) -> u32 {
            let mut v = [0u8, .. 4];
            self.fill_bytes(&mut v);
            unsafe { mem::transmute(v) }
        }
        fn next_u64(&mut self) -> u64 {
            let mut v = [0u8, .. 8];
            self.fill_bytes(&mut v);
            unsafe { mem::transmute(v) }
        }
        fn fill_bytes(&mut self, v: &mut [u8]) {
            let ret = unsafe {
                SecRandomCopyBytes(kSecRandomDefault, v.len() as size_t, v.as_mut_ptr())
            };
            if ret == -1 {
                panic!("couldn't generate random bytes: {}", os::last_os_error());
            }
        }
    }
}

#[cfg(windows)]
mod imp {
    extern crate libc;

    use io::{IoResult, IoError};
    use mem;
    use ops::Drop;
    use os;
    use rand::Rng;
    use result::Result::{Ok, Err};
    use self::libc::{DWORD, BYTE, LPCSTR, BOOL};
    use self::libc::types::os::arch::extra::{LONG_PTR};
    use slice::SliceExt;

    type HCRYPTPROV = LONG_PTR;

    /// A random number generator that retrieves randomness straight from
    /// the operating system. Platform sources:
    ///
    /// - Unix-like systems (Linux, Android, Mac OSX): read directly from
    ///   `/dev/urandom`, or from `getrandom(2)` system call if available.
    /// - Windows: calls `CryptGenRandom`, using the default cryptographic
    ///   service provider with the `PROV_RSA_FULL` type.
    /// - iOS: calls SecRandomCopyBytes as /dev/(u)random is sandboxed.
    ///
    /// This does not block.
    pub struct OsRng {
        hcryptprov: HCRYPTPROV
    }

    static PROV_RSA_FULL: DWORD = 1;
    static CRYPT_SILENT: DWORD = 64;
    static CRYPT_VERIFYCONTEXT: DWORD = 0xF0000000;

    #[allow(non_snake_case)]
    extern "system" {
        fn CryptAcquireContextA(phProv: *mut HCRYPTPROV,
                                pszContainer: LPCSTR,
                                pszProvider: LPCSTR,
                                dwProvType: DWORD,
                                dwFlags: DWORD) -> BOOL;
        fn CryptGenRandom(hProv: HCRYPTPROV,
                          dwLen: DWORD,
                          pbBuffer: *mut BYTE) -> BOOL;
        fn CryptReleaseContext(hProv: HCRYPTPROV, dwFlags: DWORD) -> BOOL;
    }

    impl OsRng {
        /// Create a new `OsRng`.
        pub fn new() -> IoResult<OsRng> {
            let mut hcp = 0;
            let ret = unsafe {
                CryptAcquireContextA(&mut hcp, 0 as LPCSTR, 0 as LPCSTR,
                                     PROV_RSA_FULL,
                                     CRYPT_VERIFYCONTEXT | CRYPT_SILENT)
            };

            if ret == 0 {
                Err(IoError::last_error())
            } else {
                Ok(OsRng { hcryptprov: hcp })
            }
        }
    }

    impl Rng for OsRng {
        fn next_u32(&mut self) -> u32 {
            let mut v = [0u8, .. 4];
            self.fill_bytes(&mut v);
            unsafe { mem::transmute(v) }
        }
        fn next_u64(&mut self) -> u64 {
            let mut v = [0u8, .. 8];
            self.fill_bytes(&mut v);
            unsafe { mem::transmute(v) }
        }
        fn fill_bytes(&mut self, v: &mut [u8]) {
            let ret = unsafe {
                CryptGenRandom(self.hcryptprov, v.len() as DWORD,
                               v.as_mut_ptr())
            };
            if ret == 0 {
                panic!("couldn't generate random bytes: {}", os::last_os_error());
            }
        }
    }

    impl Drop for OsRng {
        fn drop(&mut self) {
            let ret = unsafe {
                CryptReleaseContext(self.hcryptprov, 0)
            };
            if ret == 0 {
                panic!("couldn't release context: {}", os::last_os_error());
            }
        }
    }
}

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

    use super::OsRng;
    use rand::Rng;
    use task;

    #[test]
    fn test_os_rng() {
        let mut r = OsRng::new().unwrap();

        r.next_u32();
        r.next_u64();

        let mut v = [0u8, .. 1000];
        r.fill_bytes(&mut v);
    }

    #[test]
    fn test_os_rng_tasks() {

        let mut txs = vec!();
        for _ in range(0u, 20) {
            let (tx, rx) = channel();
            txs.push(tx);
            task::spawn(move|| {
                // wait until all the tasks are ready to go.
                rx.recv();

                // deschedule to attempt to interleave things as much
                // as possible (XXX: is this a good test?)
                let mut r = OsRng::new().unwrap();
                task::deschedule();
                let mut v = [0u8, .. 1000];

                for _ in range(0u, 100) {
                    r.next_u32();
                    task::deschedule();
                    r.next_u64();
                    task::deschedule();
                    r.fill_bytes(&mut v);
                    task::deschedule();
                }
            })
        }

        // start all the tasks
        for tx in txs.iter() {
            tx.send(())
        }
    }
}
<commit_msg>Fallout from new thread API<commit_after>// Copyright 2013-2014 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

//! Interfaces to the operating system provided random number
//! generators.

pub use self::imp::OsRng;

#[cfg(all(unix, not(target_os = "ios")))]
mod imp {
    extern crate libc;

    use self::OsRngInner::*;

    use io::{IoResult, File};
    use path::Path;
    use rand::Rng;
    use rand::reader::ReaderRng;
    use result::Result::{Ok, Err};
    use slice::SliceExt;
    use mem;
    use os::errno;

    #[cfg(all(target_os = "linux",
              any(target_arch = "x86_64", target_arch = "x86", target_arch = "arm")))]
    fn getrandom(buf: &mut [u8]) -> libc::c_long {
        extern "C" {
            fn syscall(number: libc::c_long, ...) -> libc::c_long;
        }

        #[cfg(target_arch = "x86_64")]
        const NR_GETRANDOM: libc::c_long = 318;
        #[cfg(target_arch = "x86")]
        const NR_GETRANDOM: libc::c_long = 355;
        #[cfg(target_arch = "arm")]
        const NR_GETRANDOM: libc::c_long = 384;

        unsafe {
            syscall(NR_GETRANDOM, buf.as_mut_ptr(), buf.len(), 0u)
        }
    }

    #[cfg(not(all(target_os = "linux",
                  any(target_arch = "x86_64", target_arch = "x86", target_arch = "arm"))))]
    fn getrandom(_buf: &mut [u8]) -> libc::c_long { -1 }

    fn getrandom_fill_bytes(v: &mut [u8]) {
        let mut read = 0;
        let len = v.len();
        while read < len {
            let result = getrandom(v[mut read..]);
            if result == -1 {
                let err = errno() as libc::c_int;
                if err == libc::EINTR {
                    continue;
                } else {
                    panic!("unexpected getrandom error: {}", err);
                }
            } else {
                read += result as uint;
            }
        }
    }

    fn getrandom_next_u32() -> u32 {
        let mut buf: [u8, ..4] = [0u8, ..4];
        getrandom_fill_bytes(&mut buf);
        unsafe { mem::transmute::<[u8, ..4], u32>(buf) }
    }

    fn getrandom_next_u64() -> u64 {
        let mut buf: [u8, ..8] = [0u8, ..8];
        getrandom_fill_bytes(&mut buf);
        unsafe { mem::transmute::<[u8, ..8], u64>(buf) }
    }

    #[cfg(all(target_os = "linux",
              any(target_arch = "x86_64", target_arch = "x86", target_arch = "arm")))]
    fn is_getrandom_available() -> bool {
        use sync::atomic::{AtomicBool, INIT_ATOMIC_BOOL, Relaxed};

        static GETRANDOM_CHECKED: AtomicBool = INIT_ATOMIC_BOOL;
        static GETRANDOM_AVAILABLE: AtomicBool = INIT_ATOMIC_BOOL;

        if !GETRANDOM_CHECKED.load(Relaxed) {
            let mut buf: [u8, ..0] = [];
            let result = getrandom(&mut buf);
            let available = if result == -1 {
                let err = errno() as libc::c_int;
                err != libc::ENOSYS
            } else {
                true
            };
            GETRANDOM_AVAILABLE.store(available, Relaxed);
            GETRANDOM_CHECKED.store(true, Relaxed);
            available
        } else {
            GETRANDOM_AVAILABLE.load(Relaxed)
        }
    }

    #[cfg(not(all(target_os = "linux",
                  any(target_arch = "x86_64", target_arch = "x86", target_arch = "arm"))))]
    fn is_getrandom_available() -> bool { false }

    /// A random number generator that retrieves randomness straight from
    /// the operating system. Platform sources:
    ///
    /// - Unix-like systems (Linux, Android, Mac OSX): read directly from
    ///   `/dev/urandom`, or from `getrandom(2)` system call if available.
    /// - Windows: calls `CryptGenRandom`, using the default cryptographic
    ///   service provider with the `PROV_RSA_FULL` type.
    /// - iOS: calls SecRandomCopyBytes as /dev/(u)random is sandboxed.
    ///
    /// This does not block.
    pub struct OsRng {
        inner: OsRngInner,
    }

    enum OsRngInner {
        OsGetrandomRng,
        OsReaderRng(ReaderRng<File>),
    }

    impl OsRng {
        /// Create a new `OsRng`.
        pub fn new() -> IoResult<OsRng> {
            if is_getrandom_available() {
                return Ok(OsRng { inner: OsGetrandomRng });
            }

            let reader = try!(File::open(&Path::new("/dev/urandom")));
            let reader_rng = ReaderRng::new(reader);

            Ok(OsRng { inner: OsReaderRng(reader_rng) })
        }
    }

    impl Rng for OsRng {
        fn next_u32(&mut self) -> u32 {
            match self.inner {
                OsGetrandomRng => getrandom_next_u32(),
                OsReaderRng(ref mut rng) => rng.next_u32(),
            }
        }
        fn next_u64(&mut self) -> u64 {
            match self.inner {
                OsGetrandomRng => getrandom_next_u64(),
                OsReaderRng(ref mut rng) => rng.next_u64(),
            }
        }
        fn fill_bytes(&mut self, v: &mut [u8]) {
            match self.inner {
                OsGetrandomRng => getrandom_fill_bytes(v),
                OsReaderRng(ref mut rng) => rng.fill_bytes(v)
            }
        }
    }
}

#[cfg(target_os = "ios")]
mod imp {
    extern crate libc;

    use io::{IoResult};
    use kinds::marker;
    use mem;
    use os;
    use rand::Rng;
    use result::Result::{Ok};
    use self::libc::{c_int, size_t};
    use slice::SliceExt;

    /// A random number generator that retrieves randomness straight from
    /// the operating system. Platform sources:
    ///
    /// - Unix-like systems (Linux, Android, Mac OSX): read directly from
    ///   `/dev/urandom`, or from `getrandom(2)` system call if available.
    /// - Windows: calls `CryptGenRandom`, using the default cryptographic
    ///   service provider with the `PROV_RSA_FULL` type.
    /// - iOS: calls SecRandomCopyBytes as /dev/(u)random is sandboxed.
    ///
    /// This does not block.
    #[allow(missing_copy_implementations)]
    pub struct OsRng {
        // dummy field to ensure that this struct cannot be constructed outside of this module
        _dummy: (),
    }

    #[repr(C)]
    struct SecRandom;

    #[allow(non_upper_case_globals)]
    static kSecRandomDefault: *const SecRandom = 0 as *const SecRandom;

    #[link(name = "Security", kind = "framework")]
    extern "C" {
        fn SecRandomCopyBytes(rnd: *const SecRandom,
                              count: size_t, bytes: *mut u8) -> c_int;
    }

    impl OsRng {
        /// Create a new `OsRng`.
        pub fn new() -> IoResult<OsRng> {
            Ok(OsRng { _dummy: () })
        }
    }

    impl Rng for OsRng {
        fn next_u32(&mut self) -> u32 {
            let mut v = [0u8, .. 4];
            self.fill_bytes(&mut v);
            unsafe { mem::transmute(v) }
        }
        fn next_u64(&mut self) -> u64 {
            let mut v = [0u8, .. 8];
            self.fill_bytes(&mut v);
            unsafe { mem::transmute(v) }
        }
        fn fill_bytes(&mut self, v: &mut [u8]) {
            let ret = unsafe {
                SecRandomCopyBytes(kSecRandomDefault, v.len() as size_t, v.as_mut_ptr())
            };
            if ret == -1 {
                panic!("couldn't generate random bytes: {}", os::last_os_error());
            }
        }
    }
}

#[cfg(windows)]
mod imp {
    extern crate libc;

    use io::{IoResult, IoError};
    use mem;
    use ops::Drop;
    use os;
    use rand::Rng;
    use result::Result::{Ok, Err};
    use self::libc::{DWORD, BYTE, LPCSTR, BOOL};
    use self::libc::types::os::arch::extra::{LONG_PTR};
    use slice::SliceExt;

    type HCRYPTPROV = LONG_PTR;

    /// A random number generator that retrieves randomness straight from
    /// the operating system. Platform sources:
    ///
    /// - Unix-like systems (Linux, Android, Mac OSX): read directly from
    ///   `/dev/urandom`, or from `getrandom(2)` system call if available.
    /// - Windows: calls `CryptGenRandom`, using the default cryptographic
    ///   service provider with the `PROV_RSA_FULL` type.
    /// - iOS: calls SecRandomCopyBytes as /dev/(u)random is sandboxed.
    ///
    /// This does not block.
    pub struct OsRng {
        hcryptprov: HCRYPTPROV
    }

    static PROV_RSA_FULL: DWORD = 1;
    static CRYPT_SILENT: DWORD = 64;
    static CRYPT_VERIFYCONTEXT: DWORD = 0xF0000000;

    #[allow(non_snake_case)]
    extern "system" {
        fn CryptAcquireContextA(phProv: *mut HCRYPTPROV,
                                pszContainer: LPCSTR,
                                pszProvider: LPCSTR,
                                dwProvType: DWORD,
                                dwFlags: DWORD) -> BOOL;
        fn CryptGenRandom(hProv: HCRYPTPROV,
                          dwLen: DWORD,
                          pbBuffer: *mut BYTE) -> BOOL;
        fn CryptReleaseContext(hProv: HCRYPTPROV, dwFlags: DWORD) -> BOOL;
    }

    impl OsRng {
        /// Create a new `OsRng`.
        pub fn new() -> IoResult<OsRng> {
            let mut hcp = 0;
            let ret = unsafe {
                CryptAcquireContextA(&mut hcp, 0 as LPCSTR, 0 as LPCSTR,
                                     PROV_RSA_FULL,
                                     CRYPT_VERIFYCONTEXT | CRYPT_SILENT)
            };

            if ret == 0 {
                Err(IoError::last_error())
            } else {
                Ok(OsRng { hcryptprov: hcp })
            }
        }
    }

    impl Rng for OsRng {
        fn next_u32(&mut self) -> u32 {
            let mut v = [0u8, .. 4];
            self.fill_bytes(&mut v);
            unsafe { mem::transmute(v) }
        }
        fn next_u64(&mut self) -> u64 {
            let mut v = [0u8, .. 8];
            self.fill_bytes(&mut v);
            unsafe { mem::transmute(v) }
        }
        fn fill_bytes(&mut self, v: &mut [u8]) {
            let ret = unsafe {
                CryptGenRandom(self.hcryptprov, v.len() as DWORD,
                               v.as_mut_ptr())
            };
            if ret == 0 {
                panic!("couldn't generate random bytes: {}", os::last_os_error());
            }
        }
    }

    impl Drop for OsRng {
        fn drop(&mut self) {
            let ret = unsafe {
                CryptReleaseContext(self.hcryptprov, 0)
            };
            if ret == 0 {
                panic!("couldn't release context: {}", os::last_os_error());
            }
        }
    }
}

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

    use super::OsRng;
    use rand::Rng;
    use thread::Thread;

    #[test]
    fn test_os_rng() {
        let mut r = OsRng::new().unwrap();

        r.next_u32();
        r.next_u64();

        let mut v = [0u8, .. 1000];
        r.fill_bytes(&mut v);
    }

    #[test]
    fn test_os_rng_tasks() {

        let mut txs = vec!();
        for _ in range(0u, 20) {
            let (tx, rx) = channel();
            txs.push(tx);

            Thread::spawn(move|| {
                // wait until all the tasks are ready to go.
                rx.recv();

                // deschedule to attempt to interleave things as much
                // as possible (XXX: is this a good test?)
                let mut r = OsRng::new().unwrap();
                Thread::yield_now();
                let mut v = [0u8, .. 1000];

                for _ in range(0u, 100) {
                    r.next_u32();
                    Thread::yield_now();
                    r.next_u64();
                    Thread::yield_now();
                    r.fill_bytes(&mut v);
                    Thread::yield_now();
                }
            });
        }

        // start all the tasks
        for tx in txs.iter() {
            tx.send(())
        }
    }
}
<|endoftext|>