<commit_before>#![unstable(issue = "none", feature = "windows_stdio")]

use crate::char::decode_utf16;
use crate::cmp;
use crate::io;
use crate::ptr;
use crate::str;
use crate::sys::c;
use crate::sys::cvt;
use crate::sys::handle::Handle;
use core::str::utf8_char_width;

// Don't cache handles but get them fresh for every read/write. This allows us to track changes to
// the value over time (such as if a process calls `SetStdHandle` while it's running). See #40490.
pub struct Stdin {
    surrogate: u16,
}
pub struct Stdout {
    incomplete_utf8: IncompleteUtf8,
}

pub struct Stderr {
    incomplete_utf8: IncompleteUtf8,
}

struct IncompleteUtf8 {
    bytes: [u8; 4],
    len: u8,
}

// Apparently Windows doesn't handle large reads on stdin or writes to stdout/stderr well (see
// #13304 for details).
//
// From MSDN (2011): "The storage for this buffer is allocated from a shared heap for the
// process that is 64 KB in size. The maximum size of the buffer will depend on heap usage."
//
// We choose the cap at 8 KiB because libuv does the same, and it seems to be acceptable so far.
const MAX_BUFFER_SIZE: usize = 8192;

// The standard buffer size of BufReader for Stdin should be able to hold 3x more bytes than there
// are `u16`'s in MAX_BUFFER_SIZE. This ensures the read data can always be completely decoded from
// UTF-16 to UTF-8.
pub const STDIN_BUF_SIZE: usize = MAX_BUFFER_SIZE / 2 * 3;

pub fn get_handle(handle_id: c::DWORD) -> io::Result<c::HANDLE> {
    let handle = unsafe { c::GetStdHandle(handle_id) };
    if handle == c::INVALID_HANDLE_VALUE {
        Err(io::Error::last_os_error())
    } else if handle.is_null() {
        Err(io::Error::from_raw_os_error(c::ERROR_INVALID_HANDLE as i32))
    } else {
        Ok(handle)
    }
}

fn is_console(handle: c::HANDLE) -> bool {
    // `GetConsoleMode` will return false (0) if this is a pipe (we don't care about the reported
    // mode). This will only detect Windows Console, not other terminals connected to a pipe like
    // MSYS. Which is exactly what we need, as only Windows Console needs a conversion to UTF-16.
    let mut mode = 0;
    unsafe { c::GetConsoleMode(handle, &mut mode) != 0 }
}

fn write(
    handle_id: c::DWORD,
    data: &[u8],
    incomplete_utf8: &mut IncompleteUtf8,
) -> io::Result<usize> {
    if data.is_empty() {
        return Ok(0);
    }

    let handle = get_handle(handle_id)?;
    if !is_console(handle) {
        let handle = Handle::new(handle);
        let ret = handle.write(data);
        handle.into_raw(); // Don't close the handle
        return ret;
    }

    match incomplete_utf8.len {
        0 => {}
        1..=3 => {
            if data[0] >> 6 != 0b10 {
                incomplete_utf8.len = 0;
                // not a continuation byte - reject
                return Err(io::Error::new_const(
                    io::ErrorKind::InvalidData,
                    &"Windows stdio in console mode does not support writing non-UTF-8 byte sequences",
                ));
            }
            incomplete_utf8.bytes[incomplete_utf8.len as usize] = data[0];
            incomplete_utf8.len += 1;
            let char_width = utf8_char_width(incomplete_utf8.bytes[0]);
            if (incomplete_utf8.len as usize) < char_width {
                // more bytes needed
                return Ok(1);
            }
            let s = str::from_utf8(&incomplete_utf8.bytes[0..incomplete_utf8.len as usize]);
            incomplete_utf8.len = 0;
            match s {
                Ok(s) => {
                    assert_eq!(char_width, s.len());
                    let written = write_valid_utf8(handle, s)?;
                    assert_eq!(written, s.len()); // guaranteed by write0() for single codepoint writes
                    return Ok(1);
                }
                Err(_) => {
                    return Err(io::Error::new_const(
                        io::ErrorKind::InvalidData,
                        &"Windows stdio in console mode does not support writing non-UTF-8 byte sequences",
                    ));
                }
            }
        }
        _ => {
            panic!("Unexpected number of incomplete UTF-8 chars.");
        }
    }

    // As the console is meant for presenting text, we assume bytes of `data` are encoded as UTF-8,
    // which needs to be encoded as UTF-16.
    //
    // If the data is not valid UTF-8 we write out as many bytes as are valid.
    // If the first byte is invalid it is either first byte of a multi-byte sequence but the
    // provided byte slice is too short or it is the first byte of an invalide multi-byte sequence.
    let len = cmp::min(data.len(), MAX_BUFFER_SIZE / 2);
    let utf8 = match str::from_utf8(&data[..len]) {
        Ok(s) => s,
        Err(ref e) if e.valid_up_to() == 0 => {
            let first_byte_char_width = utf8_char_width(data[0]);
            if first_byte_char_width > 1 && data.len() < first_byte_char_width {
                incomplete_utf8.bytes[0] = data[0];
                incomplete_utf8.len = 1;
                return Ok(1);
            } else {
                return Err(io::Error::new_const(
                    io::ErrorKind::InvalidData,
                    &"Windows stdio in console mode does not support writing non-UTF-8 byte sequences",
                ));
            }
        }
        Err(e) => str::from_utf8(&data[..e.valid_up_to()]).unwrap(),
    };

    write_valid_utf8(handle, utf8)
}

fn write_valid_utf8(handle: c::HANDLE, utf8: &str) -> io::Result<usize> {
    let mut utf16 = [0u16; MAX_BUFFER_SIZE / 2];
    let mut len_utf16 = 0;
    for (chr, dest) in utf8.encode_utf16().zip(utf16.iter_mut()) {
        *dest = chr;
        len_utf16 += 1;
    }
    let utf16 = &utf16[..len_utf16];

    let mut written = write_u16s(handle, &utf16)?;

    // Figure out how many bytes of as UTF-8 were written away as UTF-16.
    if written == utf16.len() {
        Ok(utf8.len())
    } else {
        // Make sure we didn't end up writing only half of a surrogate pair (even though the chance
        // is tiny). Because it is not possible for user code to re-slice `data` in such a way that
        // a missing surrogate can be produced (and also because of the UTF-8 validation above),
        // write the missing surrogate out now.
        // Buffering it would mean we have to lie about the number of bytes written.
        let first_char_remaining = utf16[written];
        if first_char_remaining >= 0xDCEE && first_char_remaining <= 0xDFFF {
            // low surrogate
            // We just hope this works, and give up otherwise
            let _ = write_u16s(handle, &utf16[written..written + 1]);
            written += 1;
        }
        // Calculate the number of bytes of `utf8` that were actually written.
        let mut count = 0;
        for ch in utf16[..written].iter() {
            count += match ch {
                0x0000..=0x007F => 1,
                0x0080..=0x07FF => 2,
                0xDCEE..=0xDFFF => 1, // Low surrogate. We already counted 3 bytes for the other.
                _ => 3,
            };
        }
        debug_assert!(String::from_utf16(&utf16[..written]).unwrap() == utf8[..count]);
        Ok(count)
    }
}

fn write_u16s(handle: c::HANDLE, data: &[u16]) -> io::Result<usize> {
    let mut written = 0;
    cvt(unsafe {
        c::WriteConsoleW(
            handle,
            data.as_ptr() as c::LPCVOID,
            data.len() as u32,
            &mut written,
            ptr::null_mut(),
        )
    })?;
    Ok(written as usize)
}

impl Stdin {
    pub const fn new() -> Stdin {
        Stdin { surrogate: 0 }
    }
}

impl io::Read for Stdin {
    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
        let handle = get_handle(c::STD_INPUT_HANDLE)?;
        if !is_console(handle) {
            let handle = Handle::new(handle);
            let ret = handle.read(buf);
            handle.into_raw(); // Don't close the handle
            return ret;
        }

        if buf.len() == 0 {
            return Ok(0);
        } else if buf.len() < 4 {
            return Err(io::Error::new_const(
                io::ErrorKind::InvalidInput,
                &"Windows stdin in console mode does not support a buffer too small to \
                 guarantee holding one arbitrary UTF-8 character (4 bytes)",
            ));
        }

        let mut utf16_buf = [0u16; MAX_BUFFER_SIZE / 2];
        // In the worst case, an UTF-8 string can take 3 bytes for every `u16` of an UTF-16. So
        // we can read at most a third of `buf.len()` chars and uphold the guarantee no data gets
        // lost.
        let amount = cmp::min(buf.len() / 3, utf16_buf.len());
        let read = read_u16s_fixup_surrogates(handle, &mut utf16_buf, amount, &mut self.surrogate)?;

        utf16_to_utf8(&utf16_buf[..read], buf)
    }
}

// We assume that if the last `u16` is an unpaired surrogate they got sliced apart by our
// buffer size, and keep it around for the next read hoping to put them together.
// This is a best effort, and may not work if we are not the only reader on Stdin.
fn read_u16s_fixup_surrogates(
    handle: c::HANDLE,
    buf: &mut [u16],
    mut amount: usize,
    surrogate: &mut u16,
) -> io::Result<usize> {
    // Insert possibly remaining unpaired surrogate from last read.
    let mut start = 0;
    if *surrogate != 0 {
        buf[0] = *surrogate;
        *surrogate = 0;
        start = 1;
        if amount == 1 {
            // Special case: `Stdin::read` guarantees we can always read at least one new `u16`
            // and combine it with an unpaired surrogate, because the UTF-8 buffer is at least
            // 4 bytes.
            amount = 2;
        }
    }
    let mut amount = read_u16s(handle, &mut buf[start..amount])? + start;

    if amount > 0 {
        let last_char = buf[amount - 1];
        if last_char >= 0xD800 && last_char <= 0xDBFF {
            // high surrogate
            *surrogate = last_char;
            amount -= 1;
        }
    }
    Ok(amount)
}

fn read_u16s(handle: c::HANDLE, buf: &mut [u16]) -> io::Result<usize> {
    // Configure the `pInputControl` parameter to not only return on `\r\n` but also Ctrl-Z, the
    // traditional DOS method to indicate end of character stream / user input (SUB).
    // See #38274 and https://stackoverflow.com/questions/43836040/win-api-readconsole.
    const CTRL_Z: u16 = 0x1A;
    const CTRL_Z_MASK: c::ULONG = 1 << CTRL_Z;
    let mut input_control = c::CONSOLE_READCONSOLE_CONTROL {
        nLength: crate::mem::size_of::<c::CONSOLE_READCONSOLE_CONTROL>() as c::ULONG,
        nInitialChars: 0,
        dwCtrlWakeupMask: CTRL_Z_MASK,
        dwControlKeyState: 0,
    };

    let mut amount = 0;
    cvt(unsafe {
        c::ReadConsoleW(
            handle,
            buf.as_mut_ptr() as c::LPVOID,
            buf.len() as u32,
            &mut amount,
            &mut input_control as c::PCONSOLE_READCONSOLE_CONTROL,
        )
    })?;

    if amount > 0 && buf[amount as usize - 1] == CTRL_Z {
        amount -= 1;
    }
    Ok(amount as usize)
}

#[allow(unused)]
fn utf16_to_utf8(utf16: &[u16], utf8: &mut [u8]) -> io::Result<usize> {
    let mut written = 0;
    for chr in decode_utf16(utf16.iter().cloned()) {
        match chr {
            Ok(chr) => {
                chr.encode_utf8(&mut utf8[written..]);
                written += chr.len_utf8();
            }
            Err(_) => {
                // We can't really do any better than forget all data and return an error.
                return Err(io::Error::new_const(
                    io::ErrorKind::InvalidData,
                    &"Windows stdin in console mode does not support non-UTF-16 input; \
                     encountered unpaired surrogate",
                ));
            }
        }
    }
    Ok(written)
}

impl IncompleteUtf8 {
    pub const fn new() -> IncompleteUtf8 {
        IncompleteUtf8 { bytes: [0; 4], len: 0 }
    }
}

impl Stdout {
    pub const fn new() -> Stdout {
        Stdout { incomplete_utf8: IncompleteUtf8::new() }
    }
}

impl io::Write for Stdout {
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        write(c::STD_OUTPUT_HANDLE, buf, &mut self.incomplete_utf8)
    }

    fn flush(&mut self) -> io::Result<()> {
        Ok(())
    }
}

impl Stderr {
    pub const fn new() -> Stderr {
        Stderr { incomplete_utf8: IncompleteUtf8::new() }
    }
}

impl io::Write for Stderr {
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        write(c::STD_ERROR_HANDLE, buf, &mut self.incomplete_utf8)
    }

    fn flush(&mut self) -> io::Result<()> {
        Ok(())
    }
}

pub fn is_ebadf(err: &io::Error) -> bool {
    err.raw_os_error() == Some(c::ERROR_INVALID_HANDLE as i32)
}

pub fn panic_output() -> Option<impl io::Write> {
    Some(Stderr::new())
}
<commit_msg>correct comment<commit_after>#![unstable(issue = "none", feature = "windows_stdio")]

use crate::char::decode_utf16;
use crate::cmp;
use crate::io;
use crate::ptr;
use crate::str;
use crate::sys::c;
use crate::sys::cvt;
use crate::sys::handle::Handle;
use core::str::utf8_char_width;

// Don't cache handles but get them fresh for every read/write. This allows us to track changes to
// the value over time (such as if a process calls `SetStdHandle` while it's running). See #40490.
pub struct Stdin {
    surrogate: u16,
}
pub struct Stdout {
    incomplete_utf8: IncompleteUtf8,
}

pub struct Stderr {
    incomplete_utf8: IncompleteUtf8,
}

struct IncompleteUtf8 {
    bytes: [u8; 4],
    len: u8,
}

// Apparently Windows doesn't handle large reads on stdin or writes to stdout/stderr well (see
// #13304 for details).
//
// From MSDN (2011): "The storage for this buffer is allocated from a shared heap for the
// process that is 64 KB in size. The maximum size of the buffer will depend on heap usage."
//
// We choose the cap at 8 KiB because libuv does the same, and it seems to be acceptable so far.
const MAX_BUFFER_SIZE: usize = 8192;

// The standard buffer size of BufReader for Stdin should be able to hold 3x more bytes than there
// are `u16`'s in MAX_BUFFER_SIZE. This ensures the read data can always be completely decoded from
// UTF-16 to UTF-8.
pub const STDIN_BUF_SIZE: usize = MAX_BUFFER_SIZE / 2 * 3;

pub fn get_handle(handle_id: c::DWORD) -> io::Result<c::HANDLE> {
    let handle = unsafe { c::GetStdHandle(handle_id) };
    if handle == c::INVALID_HANDLE_VALUE {
        Err(io::Error::last_os_error())
    } else if handle.is_null() {
        Err(io::Error::from_raw_os_error(c::ERROR_INVALID_HANDLE as i32))
    } else {
        Ok(handle)
    }
}

fn is_console(handle: c::HANDLE) -> bool {
    // `GetConsoleMode` will return false (0) if this is a pipe (we don't care about the reported
    // mode). This will only detect Windows Console, not other terminals connected to a pipe like
    // MSYS. Which is exactly what we need, as only Windows Console needs a conversion to UTF-16.
    let mut mode = 0;
    unsafe { c::GetConsoleMode(handle, &mut mode) != 0 }
}

fn write(
    handle_id: c::DWORD,
    data: &[u8],
    incomplete_utf8: &mut IncompleteUtf8,
) -> io::Result<usize> {
    if data.is_empty() {
        return Ok(0);
    }

    let handle = get_handle(handle_id)?;
    if !is_console(handle) {
        let handle = Handle::new(handle);
        let ret = handle.write(data);
        handle.into_raw(); // Don't close the handle
        return ret;
    }

    match incomplete_utf8.len {
        0 => {}
        1..=3 => {
            if data[0] >> 6 != 0b10 {
                incomplete_utf8.len = 0;
                // not a continuation byte - reject
                return Err(io::Error::new_const(
                    io::ErrorKind::InvalidData,
                    &"Windows stdio in console mode does not support writing non-UTF-8 byte sequences",
                ));
            }
            incomplete_utf8.bytes[incomplete_utf8.len as usize] = data[0];
            incomplete_utf8.len += 1;
            let char_width = utf8_char_width(incomplete_utf8.bytes[0]);
            if (incomplete_utf8.len as usize) < char_width {
                // more bytes needed
                return Ok(1);
            }
            let s = str::from_utf8(&incomplete_utf8.bytes[0..incomplete_utf8.len as usize]);
            incomplete_utf8.len = 0;
            match s {
                Ok(s) => {
                    assert_eq!(char_width, s.len());
                    let written = write_valid_utf8(handle, s)?;
                    assert_eq!(written, s.len()); // guaranteed by write_valid_utf8() for single codepoint writes
                    return Ok(1);
                }
                Err(_) => {
                    return Err(io::Error::new_const(
                        io::ErrorKind::InvalidData,
                        &"Windows stdio in console mode does not support writing non-UTF-8 byte sequences",
                    ));
                }
            }
        }
        _ => {
            panic!("Unexpected number of incomplete UTF-8 chars.");
        }
    }

    // As the console is meant for presenting text, we assume bytes of `data` are encoded as UTF-8,
    // which needs to be encoded as UTF-16.
    //
    // If the data is not valid UTF-8 we write out as many bytes as are valid.
    // If the first byte is invalid it is either first byte of a multi-byte sequence but the
    // provided byte slice is too short or it is the first byte of an invalide multi-byte sequence.
    let len = cmp::min(data.len(), MAX_BUFFER_SIZE / 2);
    let utf8 = match str::from_utf8(&data[..len]) {
        Ok(s) => s,
        Err(ref e) if e.valid_up_to() == 0 => {
            let first_byte_char_width = utf8_char_width(data[0]);
            if first_byte_char_width > 1 && data.len() < first_byte_char_width {
                incomplete_utf8.bytes[0] = data[0];
                incomplete_utf8.len = 1;
                return Ok(1);
            } else {
                return Err(io::Error::new_const(
                    io::ErrorKind::InvalidData,
                    &"Windows stdio in console mode does not support writing non-UTF-8 byte sequences",
                ));
            }
        }
        Err(e) => str::from_utf8(&data[..e.valid_up_to()]).unwrap(),
    };

    write_valid_utf8(handle, utf8)
}

fn write_valid_utf8(handle: c::HANDLE, utf8: &str) -> io::Result<usize> {
    let mut utf16 = [0u16; MAX_BUFFER_SIZE / 2];
    let mut len_utf16 = 0;
    for (chr, dest) in utf8.encode_utf16().zip(utf16.iter_mut()) {
        *dest = chr;
        len_utf16 += 1;
    }
    let utf16 = &utf16[..len_utf16];

    let mut written = write_u16s(handle, &utf16)?;

    // Figure out how many bytes of as UTF-8 were written away as UTF-16.
    if written == utf16.len() {
        Ok(utf8.len())
    } else {
        // Make sure we didn't end up writing only half of a surrogate pair (even though the chance
        // is tiny). Because it is not possible for user code to re-slice `data` in such a way that
        // a missing surrogate can be produced (and also because of the UTF-8 validation above),
        // write the missing surrogate out now.
        // Buffering it would mean we have to lie about the number of bytes written.
        let first_char_remaining = utf16[written];
        if first_char_remaining >= 0xDCEE && first_char_remaining <= 0xDFFF {
            // low surrogate
            // We just hope this works, and give up otherwise
            let _ = write_u16s(handle, &utf16[written..written + 1]);
            written += 1;
        }
        // Calculate the number of bytes of `utf8` that were actually written.
        let mut count = 0;
        for ch in utf16[..written].iter() {
            count += match ch {
                0x0000..=0x007F => 1,
                0x0080..=0x07FF => 2,
                0xDCEE..=0xDFFF => 1, // Low surrogate. We already counted 3 bytes for the other.
                _ => 3,
            };
        }
        debug_assert!(String::from_utf16(&utf16[..written]).unwrap() == utf8[..count]);
        Ok(count)
    }
}

fn write_u16s(handle: c::HANDLE, data: &[u16]) -> io::Result<usize> {
    let mut written = 0;
    cvt(unsafe {
        c::WriteConsoleW(
            handle,
            data.as_ptr() as c::LPCVOID,
            data.len() as u32,
            &mut written,
            ptr::null_mut(),
        )
    })?;
    Ok(written as usize)
}

impl Stdin {
    pub const fn new() -> Stdin {
        Stdin { surrogate: 0 }
    }
}

impl io::Read for Stdin {
    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
        let handle = get_handle(c::STD_INPUT_HANDLE)?;
        if !is_console(handle) {
            let handle = Handle::new(handle);
            let ret = handle.read(buf);
            handle.into_raw(); // Don't close the handle
            return ret;
        }

        if buf.len() == 0 {
            return Ok(0);
        } else if buf.len() < 4 {
            return Err(io::Error::new_const(
                io::ErrorKind::InvalidInput,
                &"Windows stdin in console mode does not support a buffer too small to \
                 guarantee holding one arbitrary UTF-8 character (4 bytes)",
            ));
        }

        let mut utf16_buf = [0u16; MAX_BUFFER_SIZE / 2];
        // In the worst case, an UTF-8 string can take 3 bytes for every `u16` of an UTF-16. So
        // we can read at most a third of `buf.len()` chars and uphold the guarantee no data gets
        // lost.
        let amount = cmp::min(buf.len() / 3, utf16_buf.len());
        let read = read_u16s_fixup_surrogates(handle, &mut utf16_buf, amount, &mut self.surrogate)?;

        utf16_to_utf8(&utf16_buf[..read], buf)
    }
}

// We assume that if the last `u16` is an unpaired surrogate they got sliced apart by our
// buffer size, and keep it around for the next read hoping to put them together.
// This is a best effort, and may not work if we are not the only reader on Stdin.
fn read_u16s_fixup_surrogates(
    handle: c::HANDLE,
    buf: &mut [u16],
    mut amount: usize,
    surrogate: &mut u16,
) -> io::Result<usize> {
    // Insert possibly remaining unpaired surrogate from last read.
    let mut start = 0;
    if *surrogate != 0 {
        buf[0] = *surrogate;
        *surrogate = 0;
        start = 1;
        if amount == 1 {
            // Special case: `Stdin::read` guarantees we can always read at least one new `u16`
            // and combine it with an unpaired surrogate, because the UTF-8 buffer is at least
            // 4 bytes.
            amount = 2;
        }
    }
    let mut amount = read_u16s(handle, &mut buf[start..amount])? + start;

    if amount > 0 {
        let last_char = buf[amount - 1];
        if last_char >= 0xD800 && last_char <= 0xDBFF {
            // high surrogate
            *surrogate = last_char;
            amount -= 1;
        }
    }
    Ok(amount)
}

fn read_u16s(handle: c::HANDLE, buf: &mut [u16]) -> io::Result<usize> {
    // Configure the `pInputControl` parameter to not only return on `\r\n` but also Ctrl-Z, the
    // traditional DOS method to indicate end of character stream / user input (SUB).
    // See #38274 and https://stackoverflow.com/questions/43836040/win-api-readconsole.
    const CTRL_Z: u16 = 0x1A;
    const CTRL_Z_MASK: c::ULONG = 1 << CTRL_Z;
    let mut input_control = c::CONSOLE_READCONSOLE_CONTROL {
        nLength: crate::mem::size_of::<c::CONSOLE_READCONSOLE_CONTROL>() as c::ULONG,
        nInitialChars: 0,
        dwCtrlWakeupMask: CTRL_Z_MASK,
        dwControlKeyState: 0,
    };

    let mut amount = 0;
    cvt(unsafe {
        c::ReadConsoleW(
            handle,
            buf.as_mut_ptr() as c::LPVOID,
            buf.len() as u32,
            &mut amount,
            &mut input_control as c::PCONSOLE_READCONSOLE_CONTROL,
        )
    })?;

    if amount > 0 && buf[amount as usize - 1] == CTRL_Z {
        amount -= 1;
    }
    Ok(amount as usize)
}

#[allow(unused)]
fn utf16_to_utf8(utf16: &[u16], utf8: &mut [u8]) -> io::Result<usize> {
    let mut written = 0;
    for chr in decode_utf16(utf16.iter().cloned()) {
        match chr {
            Ok(chr) => {
                chr.encode_utf8(&mut utf8[written..]);
                written += chr.len_utf8();
            }
            Err(_) => {
                // We can't really do any better than forget all data and return an error.
                return Err(io::Error::new_const(
                    io::ErrorKind::InvalidData,
                    &"Windows stdin in console mode does not support non-UTF-16 input; \
                     encountered unpaired surrogate",
                ));
            }
        }
    }
    Ok(written)
}

impl IncompleteUtf8 {
    pub const fn new() -> IncompleteUtf8 {
        IncompleteUtf8 { bytes: [0; 4], len: 0 }
    }
}

impl Stdout {
    pub const fn new() -> Stdout {
        Stdout { incomplete_utf8: IncompleteUtf8::new() }
    }
}

impl io::Write for Stdout {
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        write(c::STD_OUTPUT_HANDLE, buf, &mut self.incomplete_utf8)
    }

    fn flush(&mut self) -> io::Result<()> {
        Ok(())
    }
}

impl Stderr {
    pub const fn new() -> Stderr {
        Stderr { incomplete_utf8: IncompleteUtf8::new() }
    }
}

impl io::Write for Stderr {
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        write(c::STD_ERROR_HANDLE, buf, &mut self.incomplete_utf8)
    }

    fn flush(&mut self) -> io::Result<()> {
        Ok(())
    }
}

pub fn is_ebadf(err: &io::Error) -> bool {
    err.raw_os_error() == Some(c::ERROR_INVALID_HANDLE as i32)
}

pub fn panic_output() -> Option<impl io::Write> {
    Some(Stderr::new())
}
<|endoftext|>


<commit_before>extern crate hyper;
extern crate tungsten;

use std::thread;
use std::sync::{Once, ONCE_INIT};
use std::io::Read;
use hyper::status::StatusCode;
use hyper::client::Client;
use hyper::client::response::Response as HyperResponse;
use tungsten::server::Tungsten;
use tungsten::request::Request;
use tungsten::response::Response;
use tungsten::router::HttpMethods;

static TEST_INIT: Once = ONCE_INIT;

struct TestContext {
    req_client: Client
}

impl TestContext {
    fn new() ->  TestContext {
        let ctx = TestContext {
            req_client: Client::new()
        };
        TEST_INIT.call_once(|| {
            let _ = thread::spawn(move || {
                let mut app = Tungsten::new();
                app.get("/".to_string(), |req: Request, res: Response| {
                    res.send(b"Hello World!");
                });
                app.get("/some/random/path".to_string(), |req: Request, res: Response| {
                    res.send(b"You are at /some/random/path");
                });
                app.listen("0.0.0.0:4040");
            });
            loop {
                if ctx.req_client.get("http://0.0.0.0:4040").send().is_ok() {
                    break;
                }
            }
        });
        ctx
    }

    fn request(&self, url: &str) -> HyperResponse {
        self.req_client.get(url).send().unwrap()
    }

    fn body_from_response(&self, res: &mut HyperResponse) -> String {
        let mut body = String::new();
        let _ = res.read_to_string(&mut body);
        body
    }
}

#[test]
fn test_hello_world() {
    let ctx = TestContext::new();
    let mut res = ctx.request("http://0.0.0.0:4040");
    let body = ctx.body_from_response(&mut res);

    assert_eq!(res.status, StatusCode::Ok);
    assert_eq!(body, "Hello World!");
}

#[test]
fn test_some_path() {
    let ctx = TestContext::new();
    let mut res = ctx.request("http://0.0.0.0:4040/some/random/path");
    let body = ctx.body_from_response(&mut res);

    assert_eq!(res.status, StatusCode::Ok);
    assert_eq!(body, "You are at /some/random/path");
}
<commit_msg>tests for response method body<commit_after>extern crate hyper;
extern crate tungsten;

use std::thread;
use std::sync::{Once, ONCE_INIT};
use std::io::{Read, Write};
use hyper::status::StatusCode;
use hyper::client::Client;
use hyper::client::response::Response as HyperResponse;
use tungsten::server::Tungsten;
use tungsten::request::Request;
use tungsten::response::Response;
use tungsten::router::HttpMethods;

static TEST_INIT: Once = ONCE_INIT;

struct TestContext {
    req_client: Client
}

impl TestContext {
    fn new() ->  TestContext {
        let ctx = TestContext {
            req_client: Client::new()
        };
        TEST_INIT.call_once(|| {
            let _ = thread::spawn(move || {
                let mut app = Tungsten::new();
                app.get("/".to_string(), |req: Request, res: Response| {
                    res.send(b"Hello World!");
                });
                app.get("/some/random/path".to_string(), |req: Request, res: Response| {
                    res.send(b"You are at /some/random/path");
                });
                app.get("/throw/error".to_string(), |req: Request, mut res: Response| {
                    *res.status_mut() = StatusCode::InternalServerError;
                    let mut started_res = res.start().unwrap();
                    started_res.write(b"Boom!");
                    started_res.end();
                });
                app.listen("0.0.0.0:4040");
            });
            loop {
                if ctx.req_client.get("http://0.0.0.0:4040").send().is_ok() {
                    break;
                }
            }
        });
        ctx
    }

    fn request(&self, url: &str) -> HyperResponse {
        self.req_client.get(url).send().unwrap()
    }

    fn body_from_response(&self, res: &mut HyperResponse) -> String {
        let mut body = String::new();
        let _ = res.read_to_string(&mut body);
        body
    }
}

#[test]
fn test_hello_world() {
    let ctx = TestContext::new();
    let mut res = ctx.request("http://0.0.0.0:4040");
    let body = ctx.body_from_response(&mut res);

    assert_eq!(res.status, StatusCode::Ok);
    assert_eq!(body, "Hello World!");
}

#[test]
fn test_some_path() {
    let ctx = TestContext::new();
    let mut res = ctx.request("http://0.0.0.0:4040/some/random/path");
    let body = ctx.body_from_response(&mut res);

    assert_eq!(res.status, StatusCode::Ok);
    assert_eq!(body, "You are at /some/random/path");
}

#[test]
fn test_res_methods() {
    let ctx = TestContext::new();
    let mut res = ctx.request("http://0.0.0.0:4040/throw/error");
    let body = ctx.body_from_response(&mut res);

    assert_eq!(res.status, StatusCode::InternalServerError);
    assert_eq!(body, "Boom!");
}
<|endoftext|>


<commit_before>extern crate xml;

use std::collections::HashMap;
use std::env;
use std::fs::File;
use std::io::Read;
use std::io::BufReader;

use xml::reader::{EventReader, XmlEvent};
use xml::attribute::OwnedAttribute;

#[derive(Debug, Clone, Default)]
struct LineAttributes {
    top: i32,
    left: i32,
    height: i32,
}

fn read_attributes(attr_list: &Vec<OwnedAttribute>) -> LineAttributes {
    let mut attributes : LineAttributes = Default::default();

    for attr in attr_list {
        match attr.name.local_name.as_ref() {
            "top" => attributes.top = attr.value.parse().unwrap(),
            "left" => attributes.left = attr.value.parse().unwrap(),
            "height" => attributes.height = attr.value.parse().unwrap(),
            _ => {}
        }
    }

    attributes
}

#[derive(Debug, Clone, Default)]
struct ScriptProperties {
    direction_position: i32,
    dialog_position: i32,
    speaker_direction_position: i32,
    speaker_position: i32,
    intra_paragraph_line_height: i32,
}

fn read_and_analyze_script(reader: Box<Read>) -> (ScriptProperties, Vec<(LineAttributes, String)>) {
    let mut script_properties: ScriptProperties = Default::default();
    let mut lines: Vec<(LineAttributes, String)> = Vec::new();

    // these maps are used to heuristically determine the ScriptProperties
    let mut position_uses: HashMap<i32, u32> = HashMap::new();
    let mut line_height_uses: HashMap<i32, u32> = HashMap::new();

    // states for the streaming xml parsing
    let mut current_line_attributes: LineAttributes = Default::default();
    let mut current_text_buffer = String::new();
    let mut last_line_height = 0;

    let parser = EventReader::new(reader);
    for event in parser {
        match event {
            Ok(XmlEvent::StartElement { name, attributes, .. }) => {
                match name.local_name.as_ref() {
                    "text" => {
                        current_line_attributes = read_attributes(&attributes);

                        current_text_buffer.clear();

                        // increase the count for the current left margin
                        *position_uses.entry(current_line_attributes.left)
                                      .or_insert(0) += 1;

                        // increase the count for the current top position diff
                        *line_height_uses.entry(current_line_attributes.top - last_line_height)
                                         .or_insert(0) += 1;
                        last_line_height = current_line_attributes.top;
                    }
                    "page" => {
                        last_line_height = 0;
                    }
                    _ => {}
                }
            }
            Ok(XmlEvent::Characters(text)) => {
                if current_text_buffer.len() > 0 {
                    current_text_buffer.push(' ');
                }
                current_text_buffer.push_str(text.trim());
            }
            Ok(XmlEvent::EndElement { name, .. }) => {
                if name.local_name == "text" {
                    lines.push((current_line_attributes.clone(), current_text_buffer.clone()));
                }
            }
            Ok(_) => {},
            Err(e) => panic!(e),
        }
    }

    // the position_uses map should at least have 3 different entries
    if position_uses.len() < 3 {
        panic!("Script uses strange layout");
    }

    // copy the position uses map into a vector and sort by value desc
    let mut position_uses_vec = position_uses.into_iter().collect::<Vec<_>>();
    position_uses_vec.sort_by(|&a, &b| b.1.cmp(&a.1));

    // sort the top 3 entries by position again (right-exclusive range)
    position_uses_vec[0..3].sort_by(|&a, &b| b.0.cmp(&a.0));

    script_properties.speaker_position = position_uses_vec[0].0;
    script_properties.dialog_position = position_uses_vec[1].0;
    script_properties.direction_position = position_uses_vec[2].0;

    // take the next value between speaker and dialog as speaker_direction
    for &value in position_uses_vec.iter().skip(3) {
        let (position, _) = value;
        if position > script_properties.dialog_position &&
           position < script_properties.speaker_position {
            script_properties.speaker_direction_position = position;
            break;
        }
    }

    // use to most used line diff value to determine the specific sections
    let mut last_uses: u32 = 0;
    for (&line_height, &uses) in line_height_uses.iter() {
        if uses > last_uses {
            script_properties.intra_paragraph_line_height = line_height;
            last_uses = uses;
        }
    }

    (script_properties, lines)
}

fn condense_script_demo(properties: ScriptProperties, lines: &Vec<(LineAttributes, String)>) {
    let mut last_left_position = 0;
    let mut last_top_position = 0;

    for line in lines.iter() {
        let &(ref attributes, ref line) = line;

        if line.len() == 0 {
            continue;
        }

        // start new section
        if last_left_position != attributes.left ||
           attributes.top - last_top_position > 18 ||
           attributes.top - last_top_position < 0 {
            println!("");

            match attributes.left {
                k if (k == properties.direction_position)
                    => print!("Regie:     "),
                k if (k == properties.dialog_position)
                    => print!("Dialog:         "),
                k if (k == properties.speaker_direction_position)
                    => print!("Dialoganweisung:  "),
                k if (k == properties.speaker_position)
                    => print!("Sprecher:            "),
                _   => print!("------- "),
            }
        }

        last_left_position = attributes.left;
        last_top_position = attributes.top;

        print!("{} ", line);
    }
}

fn main() {
    let xml_file_name = env::args().skip(1).next().expect("No input file given");

    let file_reader = File::open(xml_file_name).expect("Could not open file");
    let buffered_file_reader = Box::new(BufReader::new(file_reader));

    let (properties, lines) = read_and_analyze_script(buffered_file_reader);

    condense_script_demo(properties, &lines);
}
<commit_msg>Produce better formatted demo output<commit_after>extern crate xml;

use std::collections::HashMap;
use std::env;
use std::fs::File;
use std::io::Read;
use std::io::BufReader;

use xml::reader::{EventReader, XmlEvent};
use xml::attribute::OwnedAttribute;

#[derive(Debug, Clone, Default)]
struct LineAttributes {
    top: i32,
    left: i32,
    height: i32,
}

fn read_attributes(attr_list: &Vec<OwnedAttribute>) -> LineAttributes {
    let mut attributes : LineAttributes = Default::default();

    for attr in attr_list {
        match attr.name.local_name.as_ref() {
            "top" => attributes.top = attr.value.parse().unwrap(),
            "left" => attributes.left = attr.value.parse().unwrap(),
            "height" => attributes.height = attr.value.parse().unwrap(),
            _ => {}
        }
    }

    attributes
}

#[derive(Debug, Clone, Default)]
struct ScriptProperties {
    direction_position: i32,
    dialog_position: i32,
    speaker_direction_position: i32,
    speaker_position: i32,
    intra_paragraph_line_height: i32,
}

fn read_and_analyze_script(reader: Box<Read>) -> (ScriptProperties, Vec<(LineAttributes, String)>) {
    let mut script_properties: ScriptProperties = Default::default();
    let mut lines: Vec<(LineAttributes, String)> = Vec::new();

    // these maps are used to heuristically determine the ScriptProperties
    let mut position_uses: HashMap<i32, u32> = HashMap::new();
    let mut line_height_uses: HashMap<i32, u32> = HashMap::new();

    // states for the streaming xml parsing
    let mut current_line_attributes: LineAttributes = Default::default();
    let mut current_text_buffer = String::new();
    let mut last_line_height = 0;

    let parser = EventReader::new(reader);
    for event in parser {
        match event {
            Ok(XmlEvent::StartElement { name, attributes, .. }) => {
                match name.local_name.as_ref() {
                    "text" => {
                        current_line_attributes = read_attributes(&attributes);

                        current_text_buffer.clear();

                        // increase the count for the current left margin
                        *position_uses.entry(current_line_attributes.left)
                                      .or_insert(0) += 1;

                        // increase the count for the current top position diff
                        *line_height_uses.entry(current_line_attributes.top - last_line_height)
                                         .or_insert(0) += 1;
                        last_line_height = current_line_attributes.top;
                    }
                    "page" => {
                        last_line_height = 0;
                    }
                    _ => {}
                }
            }
            Ok(XmlEvent::Characters(text)) => {
                if current_text_buffer.len() > 0 {
                    current_text_buffer.push(' ');
                }
                current_text_buffer.push_str(text.trim());
            }
            Ok(XmlEvent::EndElement { name, .. }) => {
                if name.local_name == "text" {
                    lines.push((current_line_attributes.clone(), current_text_buffer.clone()));
                }
            }
            Ok(_) => {},
            Err(e) => panic!(e),
        }
    }

    // the position_uses map should at least have 3 different entries
    if position_uses.len() < 3 {
        panic!("Script uses strange layout");
    }

    // copy the position uses map into a vector and sort by value desc
    let mut position_uses_vec = position_uses.into_iter().collect::<Vec<_>>();
    position_uses_vec.sort_by(|&a, &b| b.1.cmp(&a.1));

    // sort the top 3 entries by position again (right-exclusive range)
    position_uses_vec[0..3].sort_by(|&a, &b| b.0.cmp(&a.0));

    script_properties.speaker_position = position_uses_vec[0].0;
    script_properties.dialog_position = position_uses_vec[1].0;
    script_properties.direction_position = position_uses_vec[2].0;

    // take the next value between speaker and dialog as speaker_direction
    for &value in position_uses_vec.iter().skip(3) {
        let (position, _) = value;
        if position > script_properties.dialog_position &&
           position < script_properties.speaker_position {
            script_properties.speaker_direction_position = position;
            break;
        }
    }

    // use to most used line diff value to determine the specific sections
    let mut last_uses: u32 = 0;
    for (&line_height, &uses) in line_height_uses.iter() {
        if uses > last_uses {
            script_properties.intra_paragraph_line_height = line_height;
            last_uses = uses;
        }
    }

    (script_properties, lines)
}

fn condense_script_demo(properties: ScriptProperties, lines: &Vec<(LineAttributes, String)>) {
    let mut last_left_position = 0;
    let mut last_top_position = 0;

    for line in lines.iter() {
        let &(ref attributes, ref line) = line;

        if line.len() == 0 {
            continue;
        }

        // start new section
        if last_left_position != attributes.left ||
           attributes.top - last_top_position > 18 ||
           attributes.top - last_top_position < 0 {
            if ! (last_left_position == properties.speaker_direction_position ||
                  last_left_position == properties.speaker_position) {
                // empty line after last section
                println!("");
            }
        }

        match attributes.left {
            k if (k == properties.direction_position)
                => println!("{}", line),
            k if (k == properties.dialog_position)
                => println!("          {}", line),
            k if (k == properties.speaker_direction_position)
                => println!("               {}", line),
            k if (k == properties.speaker_position)
                => println!("                    {}", line),
            _   => println!("------- {}", line),
        }

        last_left_position = attributes.left;
        last_top_position = attributes.top;
    }
}

fn main() {
    let xml_file_name = env::args().skip(1).next().expect("No input file given");

    let file_reader = File::open(xml_file_name).expect("Could not open file");
    let buffered_file_reader = Box::new(BufReader::new(file_reader));

    let (properties, lines) = read_and_analyze_script(buffered_file_reader);

    condense_script_demo(properties, &lines);
}
<|endoftext|>


<commit_before>// Functionality that is shared between the cxx_build::bridge entry point and
// the cxxbridge CLI command.

mod builtin;
mod check;
pub(super) mod error;
mod file;
pub(super) mod fs;
mod ifndef;
pub(super) mod include;
mod namespace_organizer;
pub(super) mod out;
mod write;

pub(super) use self::error::Error;
use self::error::{format_err, Result};
use self::file::File;
use self::include::Include;
use crate::syntax::report::Errors;
use crate::syntax::{self, Types};
use std::path::Path;

/// Options for C++ code generation.
///
/// We expect options to be added over time, so this is a non-exhaustive struct.
/// To instantiate one you need to crate a default value and mutate those fields
/// that you want to modify.
///
/// ```
/// # use cxx_gen::Opt;
/// #
/// let impl_annotations = r#"__attribute__((visibility("default")))"#.to_owned();
///
/// let mut opt = Opt::default();
/// opt.cxx_impl_annotations = Some(impl_annotations);
/// ```
#[non_exhaustive]
pub struct Opt {
    /// Any additional headers to #include. The cxxbridge tool does not parse or
    /// even require the given paths to exist; they simply go into the generated
    /// C++ code as #include lines.
    pub include: Vec<Include>,
    /// Optional annotation for implementations of C++ function wrappers that
    /// may be exposed to Rust. You may for example need to provide
    /// `__declspec(dllexport)` or `__attribute__((visibility("default")))` if
    /// Rust code from one shared object or executable depends on these C++
    /// functions in another.
    pub cxx_impl_annotations: Option<String>,

    pub(super) gen_header: bool,
    pub(super) gen_implementation: bool,
    pub(super) allow_dot_includes: bool,
}

/// Results of code generation.
#[derive(Default)]
pub struct GeneratedCode {
    /// The bytes of a C++ header file.
    pub header: Vec<u8>,
    /// The bytes of a C++ implementation file (e.g. .cc, cpp etc.)
    pub implementation: Vec<u8>,
}

impl Default for Opt {
    fn default() -> Self {
        Opt {
            include: Vec::new(),
            cxx_impl_annotations: None,
            gen_header: true,
            gen_implementation: true,
            allow_dot_includes: true,
        }
    }
}

pub(super) fn generate_from_path(path: &Path, opt: &Opt) -> GeneratedCode {
    let source = match read_to_string(path) {
        Ok(source) => source,
        Err(err) => format_err(path, "", err),
    };
    match generate_from_string(&source, opt) {
        Ok(out) => out,
        Err(err) => format_err(path, &source, err),
    }
}

fn read_to_string(path: &Path) -> Result<String> {
    let bytes = if path == Path::new("-") {
        fs::read_stdin()
    } else {
        fs::read(path)
    }?;
    match String::from_utf8(bytes) {
        Ok(string) => Ok(string),
        Err(err) => Err(Error::Utf8(path.to_owned(), err.utf8_error())),
    }
}

fn generate_from_string(source: &str, opt: &Opt) -> Result<GeneratedCode> {
    let mut source = source;
    if source.starts_with("#!") && !source.starts_with("#![") {
        let shebang_end = source.find('\n').unwrap_or(source.len());
        source = &source[shebang_end..];
    }
    proc_macro2::fallback::force();
    let syntax: File = syn::parse_str(source)?;
    generate(syntax, opt)
}

pub(super) fn generate(syntax: File, opt: &Opt) -> Result<GeneratedCode> {
    if syntax.modules.is_empty() {
        return Err(Error::NoBridgeMod);
    }

    let ref mut apis = Vec::new();
    let ref mut errors = Errors::new();
    for bridge in syntax.modules {
        let ref namespace = bridge.namespace;
        let trusted = bridge.unsafety.is_some();
        apis.extend(syntax::parse_items(errors, bridge.content, trusted, namespace));
    }

    let ref types = Types::collect(errors, apis);
    check::precheck(errors, apis, opt);
    errors.propagate()?;
    check::typecheck(errors, apis, types);
    errors.propagate()?;

    // Some callers may wish to generate both header and implementation from the
    // same token stream to avoid parsing twice. Others only need to generate
    // one or the other.
    Ok(GeneratedCode {
        header: if opt.gen_header {
            write::gen(apis, types, opt, true).content()
        } else {
            Vec::new()
        },
        implementation: if opt.gen_implementation {
            write::gen(apis, types, opt, false).content()
        } else {
            Vec::new()
        },
    })
}
<commit_msg>Format with rustfmt 1.4.22-beta<commit_after>// Functionality that is shared between the cxx_build::bridge entry point and
// the cxxbridge CLI command.

mod builtin;
mod check;
pub(super) mod error;
mod file;
pub(super) mod fs;
mod ifndef;
pub(super) mod include;
mod namespace_organizer;
pub(super) mod out;
mod write;

pub(super) use self::error::Error;
use self::error::{format_err, Result};
use self::file::File;
use self::include::Include;
use crate::syntax::report::Errors;
use crate::syntax::{self, Types};
use std::path::Path;

/// Options for C++ code generation.
///
/// We expect options to be added over time, so this is a non-exhaustive struct.
/// To instantiate one you need to crate a default value and mutate those fields
/// that you want to modify.
///
/// ```
/// # use cxx_gen::Opt;
/// #
/// let impl_annotations = r#"__attribute__((visibility("default")))"#.to_owned();
///
/// let mut opt = Opt::default();
/// opt.cxx_impl_annotations = Some(impl_annotations);
/// ```
#[non_exhaustive]
pub struct Opt {
    /// Any additional headers to #include. The cxxbridge tool does not parse or
    /// even require the given paths to exist; they simply go into the generated
    /// C++ code as #include lines.
    pub include: Vec<Include>,
    /// Optional annotation for implementations of C++ function wrappers that
    /// may be exposed to Rust. You may for example need to provide
    /// `__declspec(dllexport)` or `__attribute__((visibility("default")))` if
    /// Rust code from one shared object or executable depends on these C++
    /// functions in another.
    pub cxx_impl_annotations: Option<String>,

    pub(super) gen_header: bool,
    pub(super) gen_implementation: bool,
    pub(super) allow_dot_includes: bool,
}

/// Results of code generation.
#[derive(Default)]
pub struct GeneratedCode {
    /// The bytes of a C++ header file.
    pub header: Vec<u8>,
    /// The bytes of a C++ implementation file (e.g. .cc, cpp etc.)
    pub implementation: Vec<u8>,
}

impl Default for Opt {
    fn default() -> Self {
        Opt {
            include: Vec::new(),
            cxx_impl_annotations: None,
            gen_header: true,
            gen_implementation: true,
            allow_dot_includes: true,
        }
    }
}

pub(super) fn generate_from_path(path: &Path, opt: &Opt) -> GeneratedCode {
    let source = match read_to_string(path) {
        Ok(source) => source,
        Err(err) => format_err(path, "", err),
    };
    match generate_from_string(&source, opt) {
        Ok(out) => out,
        Err(err) => format_err(path, &source, err),
    }
}

fn read_to_string(path: &Path) -> Result<String> {
    let bytes = if path == Path::new("-") {
        fs::read_stdin()
    } else {
        fs::read(path)
    }?;
    match String::from_utf8(bytes) {
        Ok(string) => Ok(string),
        Err(err) => Err(Error::Utf8(path.to_owned(), err.utf8_error())),
    }
}

fn generate_from_string(source: &str, opt: &Opt) -> Result<GeneratedCode> {
    let mut source = source;
    if source.starts_with("#!") && !source.starts_with("#![") {
        let shebang_end = source.find('\n').unwrap_or(source.len());
        source = &source[shebang_end..];
    }
    proc_macro2::fallback::force();
    let syntax: File = syn::parse_str(source)?;
    generate(syntax, opt)
}

pub(super) fn generate(syntax: File, opt: &Opt) -> Result<GeneratedCode> {
    if syntax.modules.is_empty() {
        return Err(Error::NoBridgeMod);
    }

    let ref mut apis = Vec::new();
    let ref mut errors = Errors::new();
    for bridge in syntax.modules {
        let ref namespace = bridge.namespace;
        let trusted = bridge.unsafety.is_some();
        apis.extend(syntax::parse_items(
            errors,
            bridge.content,
            trusted,
            namespace,
        ));
    }

    let ref types = Types::collect(errors, apis);
    check::precheck(errors, apis, opt);
    errors.propagate()?;
    check::typecheck(errors, apis, types);
    errors.propagate()?;

    // Some callers may wish to generate both header and implementation from the
    // same token stream to avoid parsing twice. Others only need to generate
    // one or the other.
    Ok(GeneratedCode {
        header: if opt.gen_header {
            write::gen(apis, types, opt, true).content()
        } else {
            Vec::new()
        },
        implementation: if opt.gen_implementation {
            write::gen(apis, types, opt, false).content()
        } else {
            Vec::new()
        },
    })
}
<|endoftext|>


<commit_before>use lazy_static::lazy_static;
use petgraph::Direction;
use proc_macro2::{Span, TokenStream};
use quote::quote;
use rustc_hash::FxHashMap;
use std::{io::Write, path::Path};
use syn::{spanned::Spanned, Ident};

use super::{
    java_code, map_type::map_type, map_write_err, rust_code, JavaContext, JavaForeignTypeInfo,
    JniForeignMethodSignature,
};
use crate::{
    error::{panic_on_syn_error, DiagnosticError, Result},
    file_cache::FileWriteCache,
    source_registry::SourceId,
    typemap::{
        ast::{DisplayToTokens, TypeName},
        utils::rust_to_foreign_convert_method_inputs,
        ForeignTypeInfo,
    },
    types::ForeignInterface,
    TypeMap, WRITE_TO_MEM_FAILED_MSG,
};

pub(in crate::java_jni) fn generate_interface(
    ctx: &mut JavaContext,
    interface: &ForeignInterface,
) -> Result<()> {
    let f_methods = find_suitable_ftypes_for_interace_methods(ctx.conv_map, interface)?;
    generate_java_code_for_interface(
        &ctx.cfg.output_dir,
        &ctx.cfg.package_name,
        interface,
        &f_methods,
        ctx.cfg.null_annotation_package.as_ref().map(String::as_str),
    )
    .map_err(|err| DiagnosticError::new(interface.src_id, interface.span(), err))?;
    generate_rust_code_for_interface(ctx, interface, &f_methods)?;

    let my_jobj_ti = ctx.conv_map.find_or_alloc_rust_type_with_suffix(
        &parse_type! { jobject },
        &interface.name.to_string(),
        SourceId::none(),
    );
    ctx.conv_map.add_foreign(
        my_jobj_ti,
        TypeName::from_ident(&interface.name, interface.src_id),
    )?;
    Ok(())
}

fn find_suitable_ftypes_for_interace_methods(
    conv_map: &mut TypeMap,
    interace: &ForeignInterface,
) -> Result<Vec<JniForeignMethodSignature>> {
    let void_sym = "void";
    let dummy_ty = parse_type! { () };
    let dummy_rust_ty = conv_map.find_or_alloc_rust_type_no_src_id(&dummy_ty);
    let mut f_methods = Vec::with_capacity(interace.items.len());
    let jobject_ty = conv_map.find_or_alloc_rust_type_no_src_id(&parse_type! { jobject });

    for method in &interace.items {
        let mut input = Vec::<JavaForeignTypeInfo>::with_capacity(method.fn_decl.inputs.len() - 1);
        for arg in method.fn_decl.inputs.iter().skip(1) {
            let named_arg = arg
                .as_named_arg()
                .map_err(|err| DiagnosticError::from_syn_err(interace.src_id, err))?;
            let arg_rust_ty = conv_map.find_or_alloc_rust_type(&named_arg.ty, interace.src_id);
            let arg_span = (interace.src_id, named_arg.ty.span());
            let mut f_arg_type = map_type(conv_map, &arg_rust_ty, Direction::Outgoing, arg_span)?;
            if f_arg_type.java_converter.is_some() {
                // it is hard to use Java code during callback, so may be
                // there is way to convert it to jobject ?
                conv_map.convert_rust_types(
                    arg_rust_ty.to_idx(),
                    jobject_ty.to_idx(),
                    "x",
                    "y",
                    "ret",
                    arg_span,
                )?;
                f_arg_type.java_converter = None;
                f_arg_type.base.correspoding_rust_type = jobject_ty.clone();
            }

            input.push(f_arg_type);
        }
        let output = match method.fn_decl.output {
            syn::ReturnType::Default => ForeignTypeInfo {
                name: void_sym.into(),
                correspoding_rust_type: dummy_rust_ty.clone(),
            }
            .into(),
            _ => unimplemented!(),
        };
        f_methods.push(JniForeignMethodSignature { output, input });
    }
    Ok(f_methods)
}

fn generate_java_code_for_interface(
    output_dir: &Path,
    package_name: &str,
    interface: &ForeignInterface,
    methods_sign: &[JniForeignMethodSignature],
    use_null_annotation: Option<&str>,
) -> std::result::Result<(), String> {
    let path = output_dir.join(format!("{}.java", interface.name));
    let mut file = FileWriteCache::new(&path);
    let imports = java_code::get_null_annotation_imports(use_null_annotation, methods_sign);
    let interface_comments =
        java_code::doc_comments_to_java_comments(&interface.doc_comments, true);
    writeln!(
        file,
        r#"// Automatically generated by rust_swig
package {package_name};
{imports}
{doc_comments}
public interface {interface_name} {{"#,
        package_name = package_name,
        interface_name = interface.name,
        doc_comments = interface_comments,
        imports = imports,
    )
    .expect(WRITE_TO_MEM_FAILED_MSG);

    for (method, f_method) in interface.items.iter().zip(methods_sign) {
        writeln!(
            file,
            r#"
{doc_comments}
    void {method_name}({single_args_with_types});"#,
            method_name = method.name,
            doc_comments = java_code::doc_comments_to_java_comments(&method.doc_comments, false),
            single_args_with_types = java_code::args_with_java_types(
                f_method,
                method.arg_names_without_self(),
                java_code::ArgsFormatFlags::EXTERNAL,
                use_null_annotation.is_some()
            ),
        )
        .expect(WRITE_TO_MEM_FAILED_MSG);
    }

    file.write_all(
        br#"
}
"#,
    )
    .expect(WRITE_TO_MEM_FAILED_MSG);
    file.update_file_if_necessary().map_err(&map_write_err)?;
    Ok(())
}

fn generate_rust_code_for_interface(
    ctx: &mut JavaContext,
    interface: &ForeignInterface,
    methods_sign: &[JniForeignMethodSignature],
) -> Result<()> {
    use std::fmt::Write;

    let mut new_conv_code = format!(
        r#"
#[swig_from_foreigner_hint = "{interface_name}"]
impl SwigFrom<jobject> for Box<{trait_name}> {{
    fn swig_from(this: jobject, env: *mut JNIEnv) -> Self {{
        let mut cb = JavaCallback::new(this, env);
        cb.methods.reserve({methods_len});
        let class = unsafe {{ (**env).GetObjectClass.unwrap()(env, cb.this) }};
        assert!(!class.is_null(), "GetObjectClass return null class for {interface_name}");
"#,
        interface_name = interface.name,
        trait_name = DisplayToTokens(&interface.self_type),
        methods_len = interface.items.len(),
    );
    for (method, f_method) in interface.items.iter().zip(methods_sign) {
        writeln!(
            &mut new_conv_code,
            r#"
        let method_id: jmethodID = unsafe {{
            (**env).GetMethodID.unwrap()(env, class, swig_c_str!("{method_name}"),
                                         swig_c_str!("{method_sig}"))
        }};
        assert!(!method_id.is_null(), "Can not find {method_name} id");
        cb.methods.push(method_id);"#,
            method_name = method.name,
            method_sig =
                rust_code::jni_method_signature(f_method, &ctx.cfg.package_name, ctx.conv_map),
        )
        .unwrap();
    }
    new_conv_code.push_str(
        r#"
        Box::new(cb)
    }
}
"#,
    );
    ctx.conv_map
        .merge(SourceId::none(), &new_conv_code, ctx.pointer_target_width)?;

    let mut trait_impl_funcs = Vec::<TokenStream>::with_capacity(interface.items.len());

    for (method_idx, (method, f_method)) in interface.items.iter().zip(methods_sign).enumerate() {
        let func_name = &method
            .rust_name
            .segments
            .last()
            .ok_or_else(|| {
                DiagnosticError::new(
                    interface.src_id,
                    method.rust_name.span(),
                    "Empty trait function name",
                )
            })?
            .value()
            .ident;

        let self_arg: TokenStream = method.fn_decl.inputs[0]
            .as_self_arg(interface.src_id)?
            .into();
        let mut args_with_types = Vec::with_capacity(method.fn_decl.inputs.len());
        args_with_types.push(self_arg);
        args_with_types.extend(
            method
                .fn_decl
                .inputs
                .iter()
                .skip(1)
                .enumerate()
                .map(|(i, v)| {
                    let arg_ty = &v.as_named_arg().unwrap().ty;
                    let arg_name = Ident::new(&format!("a{}", i), Span::call_site());
                    quote!(#arg_name: #arg_ty)
                }),
        );
        assert!(!method.fn_decl.inputs.is_empty());
        let n_args = method.fn_decl.inputs.len() - 1;
        let (args, type_size_asserts) = convert_args_for_variadic_function_call(f_method);

        let (mut conv_deps, convert_args_code) = rust_to_foreign_convert_method_inputs(
            ctx.conv_map,
            interface.src_id,
            method,
            f_method,
            (0..n_args).map(|v| format!("a{}", v)),
            "()",
        )?;
        ctx.rust_code.append(&mut conv_deps);
        let convert_args: TokenStream = syn::parse_str(&convert_args_code).unwrap_or_else(|err| {
            panic_on_syn_error(
                "java/jni internal parse failed for convert arguments code",
                convert_args_code,
                err,
            )
        });

        trait_impl_funcs.push(quote! {
            #[allow(unused_mut)]
            fn #func_name(#(#args_with_types),*) {
                #type_size_asserts
                let env = self.get_jni_env();
                if let Some(env) = env.env {
                    #convert_args
                    unsafe {
                        (**env).CallVoidMethod.unwrap()(env, self.this,
                                                        self.methods[#method_idx],
                                                        #(#args),*);
                        if (**env).ExceptionCheck.unwrap()(env) != 0 {
                            error!(concat!(stringify!(#func_name), ": java throw exception"));
                            (**env).ExceptionDescribe.unwrap()(env);
                            (**env).ExceptionClear.unwrap()(env);
                        }
                    };
                }
            }
        });
    }

    let self_type_name = &interface.self_type;
    let tt: TokenStream = quote! {
        impl #self_type_name for JavaCallback {
            #(#trait_impl_funcs)*
        }
    };
    ctx.rust_code.push(tt);
    Ok(())
}

lazy_static! {
    static ref JNI_FOR_VARIADIC_C_FUNC_CALL: FxHashMap<&'static str, &'static str> = {
        let mut m = FxHashMap::default();
        m.insert("jboolean", "::std::os::raw::c_uint");
        m.insert("jbyte", "::std::os::raw::c_int");
        m.insert("jshort", "::std::os::raw::c_int");
        m.insert("jfloat", "f64");
        m
    };
}

// To use `C` function with variable number of arguments,
// we need automatic type conversation, see
// http://en.cppreference.com/w/c/language/conversion#Default_argument_promotions
// for more details.
// return arg with conversation plus asserts
fn convert_args_for_variadic_function_call(
    f_method: &JniForeignMethodSignature,
) -> (Vec<TokenStream>, TokenStream) {
    let mut ret = Vec::with_capacity(f_method.input.len());
    for (i, arg) in f_method.input.iter().enumerate() {
        let arg_name = Ident::new(&format!("a{}", i), Span::call_site());
        if let Some(conv_type_str) = JNI_FOR_VARIADIC_C_FUNC_CALL
            .get(&*arg.as_ref().correspoding_rust_type.normalized_name.as_str())
        {
            let conv_type: TokenStream = syn::parse_str(*conv_type_str).unwrap_or_else(|err| {
                panic_on_syn_error(
                    "java/jni internal error: can not parse type for variable conversation",
                    conv_type_str.to_string(),
                    err,
                )
            });
            ret.push(quote!(#arg_name as #conv_type));
        } else {
            ret.push(quote!(#arg_name));
        }
    }
    let check_sizes = quote! {
        swig_assert_eq_size!(::std::os::raw::c_uint, u32);
        swig_assert_eq_size!(::std::os::raw::c_int, i32);
    };
    (ret, check_sizes)
}
<commit_msg>infra: may be fix on azure/windows 64bit<commit_after>use lazy_static::lazy_static;
use petgraph::Direction;
use proc_macro2::{Span, TokenStream};
use quote::quote;
use rustc_hash::FxHashMap;
use std::{io::Write, path::Path};
use syn::{spanned::Spanned, Ident};

use super::{
    java_code, map_type::map_type, map_write_err, rust_code, JavaContext, JavaForeignTypeInfo,
    JniForeignMethodSignature,
};
use crate::{
    error::{panic_on_syn_error, DiagnosticError, Result},
    file_cache::FileWriteCache,
    source_registry::SourceId,
    typemap::{
        ast::{DisplayToTokens, TypeName},
        utils::rust_to_foreign_convert_method_inputs,
        ForeignTypeInfo,
    },
    types::ForeignInterface,
    TypeMap, WRITE_TO_MEM_FAILED_MSG,
};

pub(in crate::java_jni) fn generate_interface(
    ctx: &mut JavaContext,
    interface: &ForeignInterface,
) -> Result<()> {
    let f_methods = find_suitable_ftypes_for_interace_methods(ctx.conv_map, interface)?;
    generate_java_code_for_interface(
        &ctx.cfg.output_dir,
        &ctx.cfg.package_name,
        interface,
        &f_methods,
        ctx.cfg.null_annotation_package.as_ref().map(String::as_str),
    )
    .map_err(|err| DiagnosticError::new(interface.src_id, interface.span(), err))?;
    generate_rust_code_for_interface(ctx, interface, &f_methods)?;

    let my_jobj_ti = ctx.conv_map.find_or_alloc_rust_type_with_suffix(
        &parse_type! { jobject },
        &interface.name.to_string(),
        SourceId::none(),
    );
    ctx.conv_map.add_foreign(
        my_jobj_ti,
        TypeName::from_ident(&interface.name, interface.src_id),
    )?;
    Ok(())
}

fn find_suitable_ftypes_for_interace_methods(
    conv_map: &mut TypeMap,
    interace: &ForeignInterface,
) -> Result<Vec<JniForeignMethodSignature>> {
    let void_sym = "void";
    let dummy_ty = parse_type! { () };
    let dummy_rust_ty = conv_map.find_or_alloc_rust_type_no_src_id(&dummy_ty);
    let mut f_methods = Vec::with_capacity(interace.items.len());
    let jobject_ty = conv_map.find_or_alloc_rust_type_no_src_id(&parse_type! { jobject });

    for method in &interace.items {
        let mut input = Vec::<JavaForeignTypeInfo>::with_capacity(method.fn_decl.inputs.len() - 1);
        for arg in method.fn_decl.inputs.iter().skip(1) {
            let named_arg = arg
                .as_named_arg()
                .map_err(|err| DiagnosticError::from_syn_err(interace.src_id, err))?;
            let arg_rust_ty = conv_map.find_or_alloc_rust_type(&named_arg.ty, interace.src_id);
            let arg_span = (interace.src_id, named_arg.ty.span());
            let mut f_arg_type = map_type(conv_map, &arg_rust_ty, Direction::Outgoing, arg_span)?;
            if f_arg_type.java_converter.is_some() {
                // it is hard to use Java code during callback, so may be
                // there is way to convert it to jobject ?
                conv_map.convert_rust_types(
                    arg_rust_ty.to_idx(),
                    jobject_ty.to_idx(),
                    "x",
                    "y",
                    "ret",
                    arg_span,
                )?;
                f_arg_type.java_converter = None;
                f_arg_type.base.correspoding_rust_type = jobject_ty.clone();
            }

            input.push(f_arg_type);
        }
        let output = match method.fn_decl.output {
            syn::ReturnType::Default => ForeignTypeInfo {
                name: void_sym.into(),
                correspoding_rust_type: dummy_rust_ty.clone(),
            }
            .into(),
            _ => unimplemented!(),
        };
        f_methods.push(JniForeignMethodSignature { output, input });
    }
    Ok(f_methods)
}

fn generate_java_code_for_interface(
    output_dir: &Path,
    package_name: &str,
    interface: &ForeignInterface,
    methods_sign: &[JniForeignMethodSignature],
    use_null_annotation: Option<&str>,
) -> std::result::Result<(), String> {
    let path = output_dir.join(format!("{}.java", interface.name));
    let mut file = FileWriteCache::new(&path);
    let imports = java_code::get_null_annotation_imports(use_null_annotation, methods_sign);
    let interface_comments =
        java_code::doc_comments_to_java_comments(&interface.doc_comments, true);
    writeln!(
        file,
        r#"// Automatically generated by rust_swig
package {package_name};
{imports}
{doc_comments}
public interface {interface_name} {{"#,
        package_name = package_name,
        interface_name = interface.name,
        doc_comments = interface_comments,
        imports = imports,
    )
    .expect(WRITE_TO_MEM_FAILED_MSG);

    for (method, f_method) in interface.items.iter().zip(methods_sign) {
        writeln!(
            file,
            r#"
{doc_comments}
    void {method_name}({single_args_with_types});"#,
            method_name = method.name,
            doc_comments = java_code::doc_comments_to_java_comments(&method.doc_comments, false),
            single_args_with_types = java_code::args_with_java_types(
                f_method,
                method.arg_names_without_self(),
                java_code::ArgsFormatFlags::EXTERNAL,
                use_null_annotation.is_some()
            ),
        )
        .expect(WRITE_TO_MEM_FAILED_MSG);
    }

    file.write_all(b"\n}\n").expect(WRITE_TO_MEM_FAILED_MSG);
    file.update_file_if_necessary().map_err(&map_write_err)?;
    Ok(())
}

fn generate_rust_code_for_interface(
    ctx: &mut JavaContext,
    interface: &ForeignInterface,
    methods_sign: &[JniForeignMethodSignature],
) -> Result<()> {
    use std::fmt::Write;

    let mut new_conv_code = format!(
        r#"
#[swig_from_foreigner_hint = "{interface_name}"]
impl SwigFrom<jobject> for Box<{trait_name}> {{
    fn swig_from(this: jobject, env: *mut JNIEnv) -> Self {{
        let mut cb = JavaCallback::new(this, env);
        cb.methods.reserve({methods_len});
        let class = unsafe {{ (**env).GetObjectClass.unwrap()(env, cb.this) }};
        assert!(!class.is_null(), "GetObjectClass return null class for {interface_name}");
"#,
        interface_name = interface.name,
        trait_name = DisplayToTokens(&interface.self_type),
        methods_len = interface.items.len(),
    );
    for (method, f_method) in interface.items.iter().zip(methods_sign) {
        writeln!(
            &mut new_conv_code,
            r#"
        let method_id: jmethodID = unsafe {{
            (**env).GetMethodID.unwrap()(env, class, swig_c_str!("{method_name}"),
                                         swig_c_str!("{method_sig}"))
        }};
        assert!(!method_id.is_null(), "Can not find {method_name} id");
        cb.methods.push(method_id);"#,
            method_name = method.name,
            method_sig =
                rust_code::jni_method_signature(f_method, &ctx.cfg.package_name, ctx.conv_map),
        )
        .unwrap();
    }
    new_conv_code.push_str(
        r#"
        Box::new(cb)
    }
}
"#,
    );
    ctx.conv_map
        .merge(SourceId::none(), &new_conv_code, ctx.pointer_target_width)?;

    let mut trait_impl_funcs = Vec::<TokenStream>::with_capacity(interface.items.len());

    for (method_idx, (method, f_method)) in interface.items.iter().zip(methods_sign).enumerate() {
        let func_name = &method
            .rust_name
            .segments
            .last()
            .ok_or_else(|| {
                DiagnosticError::new(
                    interface.src_id,
                    method.rust_name.span(),
                    "Empty trait function name",
                )
            })?
            .value()
            .ident;

        let self_arg: TokenStream = method.fn_decl.inputs[0]
            .as_self_arg(interface.src_id)?
            .into();
        let mut args_with_types = Vec::with_capacity(method.fn_decl.inputs.len());
        args_with_types.push(self_arg);
        args_with_types.extend(
            method
                .fn_decl
                .inputs
                .iter()
                .skip(1)
                .enumerate()
                .map(|(i, v)| {
                    let arg_ty = &v.as_named_arg().unwrap().ty;
                    let arg_name = Ident::new(&format!("a{}", i), Span::call_site());
                    quote!(#arg_name: #arg_ty)
                }),
        );
        assert!(!method.fn_decl.inputs.is_empty());
        let n_args = method.fn_decl.inputs.len() - 1;
        let (args, type_size_asserts) = convert_args_for_variadic_function_call(f_method);

        let (mut conv_deps, convert_args_code) = rust_to_foreign_convert_method_inputs(
            ctx.conv_map,
            interface.src_id,
            method,
            f_method,
            (0..n_args).map(|v| format!("a{}", v)),
            "()",
        )?;
        ctx.rust_code.append(&mut conv_deps);
        let convert_args: TokenStream = syn::parse_str(&convert_args_code).unwrap_or_else(|err| {
            panic_on_syn_error(
                "java/jni internal parse failed for convert arguments code",
                convert_args_code,
                err,
            )
        });

        trait_impl_funcs.push(quote! {
            #[allow(unused_mut)]
            fn #func_name(#(#args_with_types),*) {
                #type_size_asserts
                let env = self.get_jni_env();
                if let Some(env) = env.env {
                    #convert_args
                    unsafe {
                        (**env).CallVoidMethod.unwrap()(env, self.this,
                                                        self.methods[#method_idx],
                                                        #(#args),*);
                        if (**env).ExceptionCheck.unwrap()(env) != 0 {
                            error!(concat!(stringify!(#func_name), ": java throw exception"));
                            (**env).ExceptionDescribe.unwrap()(env);
                            (**env).ExceptionClear.unwrap()(env);
                        }
                    };
                }
            }
        });
    }

    let self_type_name = &interface.self_type;
    let tt: TokenStream = quote! {
        impl #self_type_name for JavaCallback {
            #(#trait_impl_funcs)*
        }
    };
    ctx.rust_code.push(tt);
    Ok(())
}

lazy_static! {
    static ref JNI_FOR_VARIADIC_C_FUNC_CALL: FxHashMap<&'static str, &'static str> = {
        let mut m = FxHashMap::default();
        m.insert("jboolean", "::std::os::raw::c_uint");
        m.insert("jbyte", "::std::os::raw::c_int");
        m.insert("jshort", "::std::os::raw::c_int");
        m.insert("jfloat", "f64");
        m
    };
}

// To use `C` function with variable number of arguments,
// we need automatic type conversation, see
// http://en.cppreference.com/w/c/language/conversion#Default_argument_promotions
// for more details.
// return arg with conversation plus asserts
fn convert_args_for_variadic_function_call(
    f_method: &JniForeignMethodSignature,
) -> (Vec<TokenStream>, TokenStream) {
    let mut ret = Vec::with_capacity(f_method.input.len());
    for (i, arg) in f_method.input.iter().enumerate() {
        let arg_name = Ident::new(&format!("a{}", i), Span::call_site());
        if let Some(conv_type_str) = JNI_FOR_VARIADIC_C_FUNC_CALL
            .get(&*arg.as_ref().correspoding_rust_type.normalized_name.as_str())
        {
            let conv_type: TokenStream = syn::parse_str(*conv_type_str).unwrap_or_else(|err| {
                panic_on_syn_error(
                    "java/jni internal error: can not parse type for variable conversation",
                    conv_type_str.to_string(),
                    err,
                )
            });
            ret.push(quote!(#arg_name as #conv_type));
        } else {
            ret.push(quote!(#arg_name));
        }
    }
    let check_sizes = quote! {
        swig_assert_eq_size!(::std::os::raw::c_uint, u32);
        swig_assert_eq_size!(::std::os::raw::c_int, i32);
    };
    (ret, check_sizes)
}
<|endoftext|>


<commit_before>//
//  SOS: the Stupid Operating System
//  by Eliza Weisman (hi@hawkweisman.me)
//
//  Copyright (c) 2016-2017 Eliza Weisman
//  Released under the terms of the MIT license. See `LICENSE` in the root
//  directory of this repository for more information.
//
//! Initial 32-bit bootloader for x86_64
#![crate_name = "boot"]
#![feature(asm)]
#![feature(lang_items)]
#![feature(naked_functions)]
#![feature(static_recursion)]
#![feature(linkage)]
#![feature(struct_field_attributes)]
#![feature(stmt_expr_attributes)]
#![no_std]

const TABLE_LENGTH: usize = 512;

type Table = [u64; TABLE_LENGTH];

use core::convert;

macro_rules! set_flags {
    (%$register:ident $( |= $body:expr);+ ) => {
        let mut $register: usize;
        asm!( concat!("mov $0, ", stringify!($register))
            : "=r"($register)
            ::: "intel");
        $($register |= $body;)+
        asm!( concat!("mov ", stringify!($register), ", $0")
            :: "r"($register)
            :: "intel");
    }
}

#[repr(C, packed)]
pub struct Gdt { _null: u64
               , code: u64
               }

#[repr(C, packed)]
pub struct GdtPointer { /// the length of the GDT
                        pub limit: u16
                      , /// pointer to the GDT
                        pub base: &'static Gdt
                      }

impl GdtPointer {
    #[cold] #[inline(always)]
    unsafe fn load (&self) {
        asm!("lgdt ($0)" :: "r"(self) : "memory");
    }
}

impl convert::From<&'static Gdt> for GdtPointer {
    #[cold] #[inline(always)]
    fn from(gdt: &'static Gdt) -> Self {
        use core::mem::size_of_val;
        GdtPointer { limit: size_of_val(gdt) as u16 - 1
                   , base: gdt }
    }
}

#[link_name = ".gdt64"]
#[link_section = ".gdt"]
#[no_mangle]
pub static GDT: Gdt
    = Gdt { _null: 0
          , code: (1<<44) | (1<<47) | (1<<41) | (1<<43) | (1<<53)
        //   , data: (1 << 44) | (1 << 47) | (1 << 41)
          };

#[inline(always)]
fn boot_write(s: &[u8]) {
    unsafe {
        use core::ptr::write_volatile;
        let vga_buf = 0xb8000 as *mut u16;
    	for (n, c) in s.iter().enumerate() {
    		write_volatile(vga_buf.offset(n as isize), 0x0200 + *c as u16);
    	}

    }

}

extern "C" {
    static mut pml4_table: Table;
    static mut pdp_table: Table;
    static mut pd_table: Table;
}

#[cold]
#[inline(always)]
#[naked]
unsafe fn create_page_tables() {
    // 3. if everything is okay, create the page tables and start long mode
    const HUGE_PAGE_SIZE: u64 = 2 * 1024 * 1024; // 2 MiB

    //-- map the PML4 and PDP tables -----------------------------------------
    // recursive map last PML4 entry
    pml4_table[511] = (&pml4_table as *const _ as u64) | 0b11;
    // map first PML4 entry to PDP table
    pml4_table[0] = (&pdp_table as *const _ as u64) | 0b11;
    // map first PDPT entry to PD table
    pdp_table[0] = (&pd_table as *const _ as u64) | 0b11;

    boot_write(b"3.1");

    //-- map the PD table ----------------------------------------------------
    for (number, entry) in pd_table.iter_mut().enumerate() {
        // set each PD entry equal to the start address of the page (the page
        // number times the page's size)
        let addr = number as u64 * HUGE_PAGE_SIZE;
        // with the appropriate flags (present + writable + huge)
        // TODO: do we want to do this using bitflags, or is that too
        //       heavyweight for the boot module?
        //          - eliza, 1/23/2017
        *entry = addr | 0b10000011;
    }

    boot_write(b"3.2");
}
#[cold]
#[inline(always)]
#[naked]
unsafe fn set_long_mode() {
    // load PML4 addr to cr3
    asm!("mov cr3, $0" :: "r"(&pml4_table) :: "intel");
    boot_write(b"3.3");

    // // enable PAE flag in cr4
    set_flags!(%cr4 |= 1 << 5 );
    boot_write(b"3.4");

    // set the long mode bit in EFER MSR (model specific register)
    asm!( "mov   ecx, 0xC0000080
           rdmsr
           or    eax, 1 << 8
           wrmsr"
        :::: "intel");
    boot_write(b"3.5");

    // enable paging in cr0
    set_flags!(%cr0 |= 1 << 31;
                    |= 1 << 16 );
}

#[cold]
#[no_mangle]
#[naked]
pub unsafe extern "C" fn _start() {
    boot_write(b"0");
    asm!("cli");

    // 1. Move Multiboot info pointer to edi
    asm!("mov edi, ebx" :::: "intel");
    boot_write(b"1");

    // 2. make sure the system supports SOS
    // TODO: port this from boot.asm
    boot_write(b"2");

    create_page_tables();
    set_long_mode();

    // 4. load the 64-bit GDT
    GdtPointer::from(&GDT).load();
    boot_write(b"4");

    // 6. jump to the 64-bit boot subroutine.
    asm!("ljmpl $$8, $$long_mode_init");

}
<commit_msg>[boot] add test for multiboot support back in<commit_after>//
//  SOS: the Stupid Operating System
//  by Eliza Weisman (hi@hawkweisman.me)
//
//  Copyright (c) 2016-2017 Eliza Weisman
//  Released under the terms of the MIT license. See `LICENSE` in the root
//  directory of this repository for more information.
//
//! Initial 32-bit bootloader for x86_64
#![crate_name = "boot"]
#![feature(asm)]
#![feature(lang_items)]
#![feature(naked_functions)]
#![feature(static_recursion)]
#![feature(linkage)]
#![feature(struct_field_attributes)]
#![feature(stmt_expr_attributes)]
#![no_std]

const TABLE_LENGTH: usize = 512;

type Table = [u64; TABLE_LENGTH];

use core::convert;

macro_rules! set_flags {
    (%$register:ident $( |= $body:expr);+ ) => {
        let mut $register: usize;
        asm!( concat!("mov $0, ", stringify!($register))
            : "=r"($register)
            ::: "intel");
        $($register |= $body;)+
        asm!( concat!("mov ", stringify!($register), ", $0")
            :: "r"($register)
            :: "intel");
    }
}

#[repr(C, packed)]
pub struct Gdt { _null: u64
               , code: u64
               }

#[repr(C, packed)]
pub struct GdtPointer { /// the length of the GDT
                        pub limit: u16
                      , /// pointer to the GDT
                        pub base: &'static Gdt
                      }

impl GdtPointer {
    #[cold] #[inline(always)]
    unsafe fn load (&self) {
        asm!("lgdt ($0)" :: "r"(self) : "memory");
    }
}

impl convert::From<&'static Gdt> for GdtPointer {
    #[cold] #[inline(always)]
    fn from(gdt: &'static Gdt) -> Self {
        use core::mem::size_of_val;
        GdtPointer { limit: size_of_val(gdt) as u16 - 1
                   , base: gdt }
    }
}

#[link_name = ".gdt64"]
#[link_section = ".gdt"]
#[no_mangle]
pub static GDT: Gdt
    = Gdt { _null: 0
          , code: (1<<44) | (1<<47) | (1<<41) | (1<<43) | (1<<53)
        //   , data: (1 << 44) | (1 << 47) | (1 << 41)
          };

#[inline(always)]
fn boot_write(s: &[u8]) {
    unsafe {
        use core::ptr::write_volatile;
        let vga_buf = 0xb8000 as *mut u16;
    	for (n, c) in s.iter().enumerate() {
    		write_volatile(vga_buf.offset(n as isize), 0x0200 + *c as u16);
    	}

    }

}

extern "C" {
    static mut pml4_table: Table;
    static mut pdp_table: Table;
    static mut pd_table: Table;
}

#[cold]
#[inline(always)]
#[naked]
unsafe fn create_page_tables() {
    // 3. if everything is okay, create the page tables and start long mode
    const HUGE_PAGE_SIZE: u64 = 2 * 1024 * 1024; // 2 MiB

    //-- map the PML4 and PDP tables -----------------------------------------
    // recursive map last PML4 entry
    pml4_table[511] = (&pml4_table as *const _ as u64) | 0b11;
    // map first PML4 entry to PDP table
    pml4_table[0] = (&pdp_table as *const _ as u64) | 0b11;
    // map first PDPT entry to PD table
    pdp_table[0] = (&pd_table as *const _ as u64) | 0b11;

    boot_write(b"3.1");

    //-- map the PD table ----------------------------------------------------
    for (number, entry) in pd_table.iter_mut().enumerate() {
        // set each PD entry equal to the start address of the page (the page
        // number times the page's size)
        let addr = number as u64 * HUGE_PAGE_SIZE;
        // with the appropriate flags (present + writable + huge)
        // TODO: do we want to do this using bitflags, or is that too
        //       heavyweight for the boot module?
        //          - eliza, 1/23/2017
        *entry = addr | 0b10000011;
    }

    boot_write(b"3.2");
}
#[cold]
#[inline(always)]
#[naked]
unsafe fn set_long_mode() {
    // load PML4 addr to cr3
    asm!("mov cr3, $0" :: "r"(&pml4_table) :: "intel");
    boot_write(b"3.3");

    // // enable PAE flag in cr4
    set_flags!(%cr4 |= 1 << 5 );
    boot_write(b"3.4");

    // set the long mode bit in EFER MSR (model specific register)
    asm!( "mov   ecx, 0xC0000080
           rdmsr
           or    eax, 1 << 8
           wrmsr"
        :::: "intel");
    boot_write(b"3.5");

    // enable paging in cr0
    set_flags!(%cr0 |= 1 << 31;
                    |= 1 << 16 );
}

/// Test whether or not this system supports Multiboot 2
#[cold]
#[inline(always)]
unsafe fn is_multiboot_supported() -> bool {
    let eax: usize;
    asm!("mov eax, $0" : "=r"(eax) ::: "intel");
    eax == 0x36d76289
}

#[cold]
#[no_mangle]
#[naked]
pub unsafe extern "C" fn _start() {
    boot_write(b"0");
    asm!("cli");

    // 1. Move Multiboot info pointer to edi
    asm!("mov edi, ebx" :::: "intel");
    boot_write(b"1");

    // 2. make sure the system supports SOS
    // TODO: port this from boot.asm
    if !is_multiboot_supported() {
        loop { boot_write(b"ERROR: multiboot not supported!"); }
    }
    boot_write(b"2");

    create_page_tables();
    set_long_mode();

    // 4. load the 64-bit GDT
    GdtPointer::from(&GDT).load();
    boot_write(b"4");

    // 6. jump to the 64-bit boot subroutine.
    asm!("ljmpl $$8, $$long_mode_init");

}
<|endoftext|>


<commit_before>//
//  SOS: the Stupid Operating System
//  by Eliza Weisman (hi@hawkweisman.me)
//
//  Copyright (c) 2016-2017 Eliza Weisman
//  Released under the terms of the MIT license. See `LICENSE` in the root
//  directory of this repository for more information.
//
//! Initial 32-bit bootloader for x86_64
#![crate_name = "boot"]
// #![crate_type = "staticlib"]
#![feature(asm)]
#![feature(lang_items)]
#![feature(naked_functions)]
#![feature(static_recursion)]
#![feature(linkage)]
#![feature(struct_field_attributes)]
#![feature(stmt_expr_attributes)]
#![no_std]
// #![no_main]

const TABLE_LENGTH: usize = 512;

type Table = [u64; TABLE_LENGTH];

use core::ptr;

// #[lang = "panic_fmt"]
// #[cold]
// unsafe extern fn panic_fmt() -> ! {
//     let vga_buf: *mut u16 = 0xb8000 as *mut u16;
//     for (i, c) in b"boot panic occurred!".iter().enumerate() {
//         ptr::write_volatile( vga_buf.offset(i as isize)
//                            , 0x0200 + *c as u16);
//     }
//     loop { asm!("hlt") };
// }

#[repr(C, packed)]
pub struct Gdt { _null: u64
               , code: u64
               }

#[repr(C, packed)]
pub struct GdtPointer { /// the length of the GDT
                        pub limit: u16
                      , /// pointer to the GDT
                        pub base: &'static Gdt
                      }

impl GdtPointer {
    #[cold]
    #[inline(always)]
    #[naked]
    unsafe fn load (&self) {
        asm!("lgdt ($0)" :: "r"(self) : "memory");
    }
}

#[link_name = ".gdt64"]
#[link_section = ".gdt"]
#[no_mangle]
pub static GDT: Gdt
    = Gdt { _null: 0
          , code: (1<<44) | (1<<47) | (1<<41) | (1<<43) | (1<<53)
        //   , data: (1 << 44) | (1 << 47) | (1 << 41)
          };

// #[no_mangle]
// static mut GDT_PTR: GdtPointer
//     = GdtPointer { limit: 15
//                  , base: &GDT
//                  };

#[cfg(feature = "log")]
#[inline(always)]
#[naked]
fn boot_write(s: &[u8]) {
	// use core::ptr;
    let mut offset = 0;
    unsafe {
        let vga_buf: *mut u16 = 0xb8000 as *mut u16;
    	for c in s {
    		ptr::write_volatile( vga_buf.offset(offset)
                               , 0x0200 + *c as u16);
    		offset += 1;
    	}

    }

}

#[cfg(not(feature = "log"))]
fn boot_write(_s: &[u8]) { }

extern "C" {
    static mut pml4_table: Table;
    static mut pdp_table: Table;
    static mut pd_table: Table;    // static mut page_table: Table;
    // #[link_name = "__vga_buffer"]
    // static mut vga_buf: *mut u16;
}

// #[naked]
// #[inline(always)]
// unsafe fn set_cs(cs: u16) {
//     asm!("ljmpl $0, $$fake_label
//           fake_label: \n\t"
//         :
//         : "i"(cs)
//         : "ax"
//         : "volatile"
//         );
// }
#[cold]
#[inline(always)]
#[naked]
unsafe fn create_page_tables() {
    // 3. if everything is okay, create the page tables and start long mode
    const HUGE_PAGE_SIZE: u64 = 2 * 1024 * 1024; // 2 MiB

    //-- map the PML4 and PDP tables -----------------------------------------
    // recursive map last PML4 entry
    pml4_table[511] = (&pml4_table as *const _ as u64) | 0b11;
    // map first PML4 entry to PDP table
    pml4_table[0] = (&pdp_table as *const _ as u64) | 0b11;
    // map first PDPT entry to PD table
    pdp_table[0] = (&pd_table as *const _ as u64) | 0b11;

    boot_write(b"3.1");

    //-- map the PD table ----------------------------------------------------
    for (number, entry) in pd_table.iter_mut().enumerate() {
        // set each PD entry equal to the start address of the page (the page
        // number times the page's size)
        let addr = number as u64 * HUGE_PAGE_SIZE;
        // with the appropriate flags (present + writable + huge)
        // TODO: do we want to do this using bitflags, or is that too
        //       heavyweight for the boot module?
        //          - eliza, 1/23/2017
        *entry = addr | 0b10000011;
    }

    boot_write(b"3.2");
}
#[cold]
#[inline(always)]
#[naked]
unsafe fn set_long_mode() {

    // load PML4 addr to cr3
    asm!( "mov   cr3, $0" :: "r"(&pml4_table) :: "intel");
    boot_write(b"3.3");

    // enable PAE flag in cr4
    let mut cr4: usize;
    asm!("mov $0, cr4" : "=r"(cr4) ::: "intel");
    cr4 |= 1 << 5;
    asm!("mov cr4, $0" :: "r"(cr4) :: "intel");
    boot_write(b"3.4");

    // set the long mode bit in EFER MSR (model specific register)
    asm!( "mov   ecx, 0xC0000080
           rdmsr
           or    eax, 1 << 8
           wrmsr"
        :::: "intel");
    boot_write(b"3.5");

    // enable paging in cr0
    // let mut cr0: usize;
    // asm!( "mov $0, cr0" : "=r"(cr0) ::: "intel");
    // cr0 |= 1 << 31;
    // cr0 |= 1 << 16;
    // asm!( "mov cr0, $0" :: "r"(cr0) :: "intel");
    // boot_write(b"3.6");
    asm!( "mov eax, cr0
           or  eax, 1 << 31
           or  eax, 1 << 16
           mov cr0, eax"
        :::: "intel");
}

#[cold]
#[no_mangle]
#[naked]
pub unsafe extern "C" fn _start() {
    use core::mem;
    boot_write(b"0");
    asm!("cli");

    // 1. Move Multiboot info pointer to edi
    // let multiboot: usize;
    asm!("mov edi, ebx"
         : //"=r"(multiboot)
         ::: "intel"
    );
    boot_write(b"1");

    // 2. make sure the system supports SOS
    // TODO: port this from boot.asm
    boot_write(b"2");

    create_page_tables();
    set_long_mode();

    // 4. load the 64-bit GDT
    let gdt_ptr = GdtPointer { limit: mem::size_of_val(&GDT) as u16 - 1
                             , base: &GDT };

    // asm!("lgdt ($0)" :: "r"(&gdt_ptr) : "memory" );
    gdt_ptr.load();
    boot_write(b"4");

    // // 5. update selectors
    // asm!("mov ax, 0x10" :::: "intel");
    // boot_write(b"5.1");
    // // stack selector
    // asm!("mov ss, ax" :::: "intel");
    // boot_write(b"5.2");
    // // data selector
    // asm!("mov ds, ax" :::: "intel");
    // boot_write(b"5.3");
    // // extra selector
    // asm!("mov es, ax" :::: "intel");
    // boot_write(b"5.4");
    // loop { }
    // 6. jump to the 64-bit boot subroutine.
    // 0. Move the stack pointer to the top of the stack.
    asm!("ljmpl $$8, $$long_mode_init");

}
<commit_msg>[boot] make GDT loading more idiomatic<commit_after>//
//  SOS: the Stupid Operating System
//  by Eliza Weisman (hi@hawkweisman.me)
//
//  Copyright (c) 2016-2017 Eliza Weisman
//  Released under the terms of the MIT license. See `LICENSE` in the root
//  directory of this repository for more information.
//
//! Initial 32-bit bootloader for x86_64
#![crate_name = "boot"]
// #![crate_type = "staticlib"]
#![feature(asm)]
#![feature(lang_items)]
#![feature(naked_functions)]
#![feature(static_recursion)]
#![feature(linkage)]
#![feature(struct_field_attributes)]
#![feature(stmt_expr_attributes)]
#![no_std]
// #![no_main]

const TABLE_LENGTH: usize = 512;

type Table = [u64; TABLE_LENGTH];

use core::ptr;
use core::convert;

#[repr(C, packed)]
pub struct Gdt { _null: u64
               , code: u64
               }

#[repr(C, packed)]
pub struct GdtPointer { /// the length of the GDT
                        pub limit: u16
                      , /// pointer to the GDT
                        pub base: &'static Gdt
                      }

impl GdtPointer {
    #[cold] #[inline(always)]
    unsafe fn load (&self) {
        asm!("lgdt ($0)" :: "r"(self) : "memory");
    }
}

impl convert::From<&'static Gdt> for GdtPointer {
    #[cold] #[inline(always)]
    fn from(gdt: &'static Gdt) -> Self {
        use core::mem::size_of_val;
        GdtPointer { limit: size_of_val(gdt) as u16 - 1
                   , base: gdt }
    }
}

#[link_name = ".gdt64"]
#[link_section = ".gdt"]
#[no_mangle]
pub static GDT: Gdt
    = Gdt { _null: 0
          , code: (1<<44) | (1<<47) | (1<<41) | (1<<43) | (1<<53)
        //   , data: (1 << 44) | (1 << 47) | (1 << 41)
          };

#[cfg(feature = "log")]
#[inline(always)]
#[naked]
fn boot_write(s: &[u8]) {
	// use core::ptr;
    let mut offset = 0;
    unsafe {
        let vga_buf: *mut u16 = 0xb8000 as *mut u16;
    	for c in s {
    		ptr::write_volatile( vga_buf.offset(offset)
                               , 0x0200 + *c as u16);
    		offset += 1;
    	}

    }

}

#[cfg(not(feature = "log"))]
fn boot_write(_s: &[u8]) { }

extern "C" {
    static mut pml4_table: Table;
    static mut pdp_table: Table;
    static mut pd_table: Table;    // static mut page_table: Table;
    // #[link_name = "__vga_buffer"]
    // static mut vga_buf: *mut u16;
}

#[cold]
#[inline(always)]
#[naked]
unsafe fn create_page_tables() {
    // 3. if everything is okay, create the page tables and start long mode
    const HUGE_PAGE_SIZE: u64 = 2 * 1024 * 1024; // 2 MiB

    //-- map the PML4 and PDP tables -----------------------------------------
    // recursive map last PML4 entry
    pml4_table[511] = (&pml4_table as *const _ as u64) | 0b11;
    // map first PML4 entry to PDP table
    pml4_table[0] = (&pdp_table as *const _ as u64) | 0b11;
    // map first PDPT entry to PD table
    pdp_table[0] = (&pd_table as *const _ as u64) | 0b11;

    boot_write(b"3.1");

    //-- map the PD table ----------------------------------------------------
    for (number, entry) in pd_table.iter_mut().enumerate() {
        // set each PD entry equal to the start address of the page (the page
        // number times the page's size)
        let addr = number as u64 * HUGE_PAGE_SIZE;
        // with the appropriate flags (present + writable + huge)
        // TODO: do we want to do this using bitflags, or is that too
        //       heavyweight for the boot module?
        //          - eliza, 1/23/2017
        *entry = addr | 0b10000011;
    }

    boot_write(b"3.2");
}
#[cold]
#[inline(always)]
#[naked]
unsafe fn set_long_mode() {

    // load PML4 addr to cr3
    asm!( "mov   cr3, $0" :: "r"(&pml4_table) :: "intel");
    boot_write(b"3.3");

    // enable PAE flag in cr4
    let mut cr4: usize;
    asm!("mov $0, cr4" : "=r"(cr4) ::: "intel");
    cr4 |= 1 << 5;
    asm!("mov cr4, $0" :: "r"(cr4) :: "intel");
    boot_write(b"3.4");

    // set the long mode bit in EFER MSR (model specific register)
    asm!( "mov   ecx, 0xC0000080
           rdmsr
           or    eax, 1 << 8
           wrmsr"
        :::: "intel");
    boot_write(b"3.5");

    // enable paging in cr0
    // let mut cr0: usize;
    // asm!( "mov $0, cr0" : "=r"(cr0) ::: "intel");
    // cr0 |= 1 << 31;
    // cr0 |= 1 << 16;
    // asm!( "mov cr0, $0" :: "r"(cr0) :: "intel");
    // boot_write(b"3.6");
    asm!( "mov eax, cr0
           or  eax, 1 << 31
           or  eax, 1 << 16
           mov cr0, eax"
        :::: "intel");
}

#[cold]
#[no_mangle]
#[naked]
pub unsafe extern "C" fn _start() {
    boot_write(b"0");
    asm!("cli");

    // 1. Move Multiboot info pointer to edi
    // let multiboot: usize;
    asm!("mov edi, ebx"
         : //"=r"(multiboot)
         ::: "intel"
    );
    boot_write(b"1");

    // 2. make sure the system supports SOS
    // TODO: port this from boot.asm
    boot_write(b"2");

    create_page_tables();
    set_long_mode();

    // 4. load the 64-bit GDT
    GdtPointer::from(&GDT).load();

    boot_write(b"4");

    // // 5. update selectors
    // asm!("mov ax, 0x10" :::: "intel");
    // boot_write(b"5.1");
    // // stack selector
    // asm!("mov ss, ax" :::: "intel");
    // boot_write(b"5.2");
    // // data selector
    // asm!("mov ds, ax" :::: "intel");
    // boot_write(b"5.3");
    // // extra selector
    // asm!("mov es, ax" :::: "intel");
    // boot_write(b"5.4");
    // loop { }
    // 6. jump to the 64-bit boot subroutine.
    // 0. Move the stack pointer to the top of the stack.
    asm!("ljmpl $$8, $$long_mode_init");

}
<|endoftext|>