<commit_before>use std::fmt;
use std::io::prelude::*;

use atty;
use termcolor::Color::{Cyan, Green, Red, Yellow};
use termcolor::{self, Color, ColorSpec, StandardStream, WriteColor};

use crate::util::errors::CargoResult;

/// The requested verbosity of output.
#[derive(Debug, Clone, Copy, PartialEq)]
pub enum Verbosity {
    Verbose,
    Normal,
    Quiet,
}

/// An abstraction around a `Write`able object that remembers preferences for output verbosity and
/// color.
pub struct Shell {
    /// the `Write`able object, either with or without color support (represented by different enum
    /// variants)
    err: ShellOut,
    /// How verbose messages should be
    verbosity: Verbosity,
    /// Flag that indicates the current line needs to be cleared before
    /// printing. Used when a progress bar is currently displayed.
    needs_clear: bool,
}

impl fmt::Debug for Shell {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self.err {
            ShellOut::Write(_) => f
                .debug_struct("Shell")
                .field("verbosity", &self.verbosity)
                .finish(),
            ShellOut::Stream { color_choice, .. } => f
                .debug_struct("Shell")
                .field("verbosity", &self.verbosity)
                .field("color_choice", &color_choice)
                .finish(),
        }
    }
}

/// A `Write`able object, either with or without color support
enum ShellOut {
    /// A plain write object without color support
    Write(Box<dyn Write>),
    /// Color-enabled stdio, with information on whether color should be used
    Stream {
        stream: StandardStream,
        tty: bool,
        color_choice: ColorChoice,
    },
}

/// Whether messages should use color output
#[derive(Debug, PartialEq, Clone, Copy)]
pub enum ColorChoice {
    /// Force color output
    Always,
    /// Force disable color output
    Never,
    /// Intelligently guess whether to use color output
    CargoAuto,
}

impl Shell {
    /// Creates a new shell (color choice and verbosity), defaulting to 'auto' color and verbose
    /// output.
    pub fn new() -> Shell {
        Shell {
            err: ShellOut::Stream {
                stream: StandardStream::stderr(ColorChoice::CargoAuto.to_termcolor_color_choice()),
                color_choice: ColorChoice::CargoAuto,
                tty: atty::is(atty::Stream::Stderr),
            },
            verbosity: Verbosity::Verbose,
            needs_clear: false,
        }
    }

    /// Creates a shell from a plain writable object, with no color, and max verbosity.
    pub fn from_write(out: Box<dyn Write>) -> Shell {
        Shell {
            err: ShellOut::Write(out),
            verbosity: Verbosity::Verbose,
            needs_clear: false,
        }
    }

    /// Prints a message, where the status will have `color` color, and can be justified. The
    /// messages follows without color.
    fn print(
        &mut self,
        status: &dyn fmt::Display,
        message: Option<&dyn fmt::Display>,
        color: Color,
        justified: bool,
    ) -> CargoResult<()> {
        match self.verbosity {
            Verbosity::Quiet => Ok(()),
            _ => {
                if self.needs_clear {
                    self.err_erase_line();
                }
                self.err.print(status, message, color, justified)
            }
        }
    }

    /// Sets whether the next print should clear the current line.
    pub fn set_needs_clear(&mut self, needs_clear: bool) {
        self.needs_clear = needs_clear;
    }

    /// Returns `true` if the `needs_clear` flag is unset.
    pub fn is_cleared(&self) -> bool {
        !self.needs_clear
    }

    /// Returns the width of the terminal in spaces, if any.
    pub fn err_width(&self) -> Option<usize> {
        match self.err {
            ShellOut::Stream { tty: true, .. } => imp::stderr_width(),
            _ => None,
        }
    }

    /// Returns `true` if stderr is a tty.
    pub fn is_err_tty(&self) -> bool {
        match self.err {
            ShellOut::Stream { tty, .. } => tty,
            _ => false,
        }
    }

    /// Gets a reference to the underlying writer.
    pub fn err(&mut self) -> &mut dyn Write {
        if self.needs_clear {
            self.err_erase_line();
        }
        self.err.as_write()
    }

    /// Erase from cursor to end of line.
    pub fn err_erase_line(&mut self) {
        if let ShellOut::Stream { tty: true, .. } = self.err {
            imp::err_erase_line(self);
            self.needs_clear = false;
        }
    }

    /// Shortcut to right-align and color green a status message.
    pub fn status<T, U>(&mut self, status: T, message: U) -> CargoResult<()>
    where
        T: fmt::Display,
        U: fmt::Display,
    {
        self.print(&status, Some(&message), Green, true)
    }

    pub fn status_header<T>(&mut self, status: T) -> CargoResult<()>
    where
        T: fmt::Display,
    {
        self.print(&status, None, Cyan, true)
    }

    /// Shortcut to right-align a status message.
    pub fn status_with_color<T, U>(
        &mut self,
        status: T,
        message: U,
        color: Color,
    ) -> CargoResult<()>
    where
        T: fmt::Display,
        U: fmt::Display,
    {
        self.print(&status, Some(&message), color, true)
    }

    /// Runs the callback only if we are in verbose mode.
    pub fn verbose<F>(&mut self, mut callback: F) -> CargoResult<()>
    where
        F: FnMut(&mut Shell) -> CargoResult<()>,
    {
        match self.verbosity {
            Verbosity::Verbose => callback(self),
            _ => Ok(()),
        }
    }

    /// Runs the callback if we are not in verbose mode.
    pub fn concise<F>(&mut self, mut callback: F) -> CargoResult<()>
    where
        F: FnMut(&mut Shell) -> CargoResult<()>,
    {
        match self.verbosity {
            Verbosity::Verbose => Ok(()),
            _ => callback(self),
        }
    }

    /// Prints a red 'error' message.
    pub fn error<T: fmt::Display>(&mut self, message: T) -> CargoResult<()> {
        if self.needs_clear {
            self.err_erase_line();
        }
        self.err.print(&"error:", Some(&message), Red, false)
    }

    /// Prints an amber 'warning' message.
    pub fn warn<T: fmt::Display>(&mut self, message: T) -> CargoResult<()> {
        match self.verbosity {
            Verbosity::Quiet => Ok(()),
            _ => self.print(&"warning:", Some(&message), Yellow, false),
        }
    }

    /// Updates the verbosity of the shell.
    pub fn set_verbosity(&mut self, verbosity: Verbosity) {
        self.verbosity = verbosity;
    }

    /// Gets the verbosity of the shell.
    pub fn verbosity(&self) -> Verbosity {
        self.verbosity
    }

    /// Updates the color choice (always, never, or auto) from a string..
    pub fn set_color_choice(&mut self, color: Option<&str>) -> CargoResult<()> {
        if let ShellOut::Stream {
            ref mut stream,
            ref mut color_choice,
            ..
        } = self.err
        {
            let cfg = match color {
                Some("always") => ColorChoice::Always,
                Some("never") => ColorChoice::Never,

                Some("auto") | None => ColorChoice::CargoAuto,

                Some(arg) => failure::bail!(
                    "argument for --color must be auto, always, or \
                     never, but found `{}`",
                    arg
                ),
            };
            *color_choice = cfg;
            *stream = StandardStream::stderr(cfg.to_termcolor_color_choice());
        }
        Ok(())
    }

    /// Gets the current color choice.
    ///
    /// If we are not using a color stream, this will always return `Never`, even if the color
    /// choice has been set to something else.
    pub fn color_choice(&self) -> ColorChoice {
        match self.err {
            ShellOut::Stream { color_choice, .. } => color_choice,
            ShellOut::Write(_) => ColorChoice::Never,
        }
    }

    /// Whether the shell supports color.
    pub fn supports_color(&self) -> bool {
        match &self.err {
            ShellOut::Write(_) => false,
            ShellOut::Stream { stream, .. } => stream.supports_color(),
        }
    }

    /// Prints a message and translates ANSI escape code into console colors.
    pub fn print_ansi(&mut self, message: &[u8]) -> CargoResult<()> {
        if self.needs_clear {
            self.err_erase_line();
        }
        #[cfg(windows)]
        {
            if let ShellOut::Stream { stream, .. } = &mut self.err {
                ::fwdansi::write_ansi(stream, message)?;
                return Ok(());
            }
        }
        self.err().write_all(message)?;
        Ok(())
    }
}

impl Default for Shell {
    fn default() -> Self {
        Self::new()
    }
}

impl ShellOut {
    /// Prints out a message with a status. The status comes first, and is bold plus the given
    /// color. The status can be justified, in which case the max width that will right align is
    /// 12 chars.
    fn print(
        &mut self,
        status: &dyn fmt::Display,
        message: Option<&dyn fmt::Display>,
        color: Color,
        justified: bool,
    ) -> CargoResult<()> {
        match *self {
            ShellOut::Stream { ref mut stream, .. } => {
                stream.reset()?;
                stream.set_color(ColorSpec::new().set_bold(true).set_fg(Some(color)))?;
                if justified {
                    write!(stream, "{:>12}", status)?;
                } else {
                    write!(stream, "{}", status)?;
                }
                stream.reset()?;
                match message {
                    Some(message) => writeln!(stream, " {}", message)?,
                    None => write!(stream, " ")?,
                }
            }
            ShellOut::Write(ref mut w) => {
                if justified {
                    write!(w, "{:>12}", status)?;
                } else {
                    write!(w, "{}", status)?;
                }
                match message {
                    Some(message) => writeln!(w, " {}", message)?,
                    None => write!(w, " ")?,
                }
            }
        }
        Ok(())
    }

    /// Gets this object as a `io::Write`.
    fn as_write(&mut self) -> &mut dyn Write {
        match *self {
            ShellOut::Stream { ref mut stream, .. } => stream,
            ShellOut::Write(ref mut w) => w,
        }
    }
}

impl ColorChoice {
    /// Converts our color choice to termcolor's version.
    fn to_termcolor_color_choice(self) -> termcolor::ColorChoice {
        match self {
            ColorChoice::Always => termcolor::ColorChoice::Always,
            ColorChoice::Never => termcolor::ColorChoice::Never,
            ColorChoice::CargoAuto => {
                if atty::is(atty::Stream::Stderr) {
                    termcolor::ColorChoice::Auto
                } else {
                    termcolor::ColorChoice::Never
                }
            }
        }
    }
}

#[cfg(any(target_os = "linux", target_os = "macos", target_os = "freebsd"))]
mod imp {
    use std::mem;

    use libc;

    use super::Shell;

    pub fn stderr_width() -> Option<usize> {
        unsafe {
            let mut winsize: libc::winsize = mem::zeroed();
            // The .into() here is needed for FreeBSD which defines TIOCGWINSZ
            // as c_uint but ioctl wants c_ulong.
            if libc::ioctl(libc::STDERR_FILENO, libc::TIOCGWINSZ.into(), &mut winsize) < 0 {
                return None;
            }
            if winsize.ws_col > 0 {
                Some(winsize.ws_col as usize)
            } else {
                None
            }
        }
    }

    pub fn err_erase_line(shell: &mut Shell) {
        // This is the "EL - Erase in Line" sequence. It clears from the cursor
        // to the end of line.
        // https://en.wikipedia.org/wiki/ANSI_escape_code#CSI_sequences
        let _ = shell.err.as_write().write_all(b"\x1B[K");
    }
}

#[cfg(all(
    unix,
    not(any(target_os = "linux", target_os = "macos", target_os = "freebsd"))
))]
mod imp {
    pub(super) use super::default_err_erase_line as err_erase_line;

    pub fn stderr_width() -> Option<usize> {
        None
    }
}

#[cfg(windows)]
mod imp {
    use std::{cmp, mem, ptr};
    use winapi::um::fileapi::*;
    use winapi::um::handleapi::*;
    use winapi::um::processenv::*;
    use winapi::um::winbase::*;
    use winapi::um::wincon::*;
    use winapi::um::winnt::*;

    pub(super) use super::default_err_erase_line as err_erase_line;

    pub fn stderr_width() -> Option<usize> {
        unsafe {
            let stdout = GetStdHandle(STD_ERROR_HANDLE);
            let mut csbi: CONSOLE_SCREEN_BUFFER_INFO = mem::zeroed();
            if GetConsoleScreenBufferInfo(stdout, &mut csbi) != 0 {
                return Some((csbi.srWindow.Right - csbi.srWindow.Left) as usize);
            }

            // On mintty/msys/cygwin based terminals, the above fails with
            // INVALID_HANDLE_VALUE. Use an alternate method which works
            // in that case as well.
            let h = CreateFileA(
                "CONOUT$\0".as_ptr() as *const CHAR,
                GENERIC_READ | GENERIC_WRITE,
                FILE_SHARE_READ | FILE_SHARE_WRITE,
                ptr::null_mut(),
                OPEN_EXISTING,
                0,
                ptr::null_mut(),
            );
            if h == INVALID_HANDLE_VALUE {
                return None;
            }

            let mut csbi: CONSOLE_SCREEN_BUFFER_INFO = mem::zeroed();
            let rc = GetConsoleScreenBufferInfo(h, &mut csbi);
            CloseHandle(h);
            if rc != 0 {
                let width = (csbi.srWindow.Right - csbi.srWindow.Left) as usize;
                // Unfortunately cygwin/mintty does not set the size of the
                // backing console to match the actual window size. This
                // always reports a size of 80 or 120 (not sure what
                // determines that). Use a conservative max of 60 which should
                // work in most circumstances. ConEmu does some magic to
                // resize the console correctly, but there's no reasonable way
                // to detect which kind of terminal we are running in, or if
                // GetConsoleScreenBufferInfo returns accurate information.
                return Some(cmp::min(60, width));
            }
            None
        }
    }
}

#[cfg(any(
    all(
        unix,
        not(any(target_os = "linux", target_os = "macos", target_os = "freebsd"))
    ),
    windows,
))]
fn default_err_erase_line(shell: &mut Shell) {
    if let Some(max_width) = imp::stderr_width() {
        let blank = " ".repeat(max_width);
        drop(write!(shell.err.as_write(), "{}\r", blank));
    }
}
<commit_msg>Auto merge of #7294 - SnejUgal:fix-inconsistent-coloring, r=alexcrichton<commit_after>use std::fmt;
use std::io::prelude::*;

use atty;
use termcolor::Color::{Cyan, Green, Red, Yellow};
use termcolor::{self, Color, ColorSpec, StandardStream, WriteColor};

use crate::util::errors::CargoResult;

/// The requested verbosity of output.
#[derive(Debug, Clone, Copy, PartialEq)]
pub enum Verbosity {
    Verbose,
    Normal,
    Quiet,
}

/// An abstraction around a `Write`able object that remembers preferences for output verbosity and
/// color.
pub struct Shell {
    /// the `Write`able object, either with or without color support (represented by different enum
    /// variants)
    err: ShellOut,
    /// How verbose messages should be
    verbosity: Verbosity,
    /// Flag that indicates the current line needs to be cleared before
    /// printing. Used when a progress bar is currently displayed.
    needs_clear: bool,
}

impl fmt::Debug for Shell {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self.err {
            ShellOut::Write(_) => f
                .debug_struct("Shell")
                .field("verbosity", &self.verbosity)
                .finish(),
            ShellOut::Stream { color_choice, .. } => f
                .debug_struct("Shell")
                .field("verbosity", &self.verbosity)
                .field("color_choice", &color_choice)
                .finish(),
        }
    }
}

/// A `Write`able object, either with or without color support
enum ShellOut {
    /// A plain write object without color support
    Write(Box<dyn Write>),
    /// Color-enabled stdio, with information on whether color should be used
    Stream {
        stream: StandardStream,
        tty: bool,
        color_choice: ColorChoice,
    },
}

/// Whether messages should use color output
#[derive(Debug, PartialEq, Clone, Copy)]
pub enum ColorChoice {
    /// Force color output
    Always,
    /// Force disable color output
    Never,
    /// Intelligently guess whether to use color output
    CargoAuto,
}

impl Shell {
    /// Creates a new shell (color choice and verbosity), defaulting to 'auto' color and verbose
    /// output.
    pub fn new() -> Shell {
        Shell {
            err: ShellOut::Stream {
                stream: StandardStream::stderr(ColorChoice::CargoAuto.to_termcolor_color_choice()),
                color_choice: ColorChoice::CargoAuto,
                tty: atty::is(atty::Stream::Stderr),
            },
            verbosity: Verbosity::Verbose,
            needs_clear: false,
        }
    }

    /// Creates a shell from a plain writable object, with no color, and max verbosity.
    pub fn from_write(out: Box<dyn Write>) -> Shell {
        Shell {
            err: ShellOut::Write(out),
            verbosity: Verbosity::Verbose,
            needs_clear: false,
        }
    }

    /// Prints a message, where the status will have `color` color, and can be justified. The
    /// messages follows without color.
    fn print(
        &mut self,
        status: &dyn fmt::Display,
        message: Option<&dyn fmt::Display>,
        color: Color,
        justified: bool,
    ) -> CargoResult<()> {
        match self.verbosity {
            Verbosity::Quiet => Ok(()),
            _ => {
                if self.needs_clear {
                    self.err_erase_line();
                }
                self.err.print(status, message, color, justified)
            }
        }
    }

    /// Sets whether the next print should clear the current line.
    pub fn set_needs_clear(&mut self, needs_clear: bool) {
        self.needs_clear = needs_clear;
    }

    /// Returns `true` if the `needs_clear` flag is unset.
    pub fn is_cleared(&self) -> bool {
        !self.needs_clear
    }

    /// Returns the width of the terminal in spaces, if any.
    pub fn err_width(&self) -> Option<usize> {
        match self.err {
            ShellOut::Stream { tty: true, .. } => imp::stderr_width(),
            _ => None,
        }
    }

    /// Returns `true` if stderr is a tty.
    pub fn is_err_tty(&self) -> bool {
        match self.err {
            ShellOut::Stream { tty, .. } => tty,
            _ => false,
        }
    }

    /// Gets a reference to the underlying writer.
    pub fn err(&mut self) -> &mut dyn Write {
        if self.needs_clear {
            self.err_erase_line();
        }
        self.err.as_write()
    }

    /// Erase from cursor to end of line.
    pub fn err_erase_line(&mut self) {
        if let ShellOut::Stream { tty: true, .. } = self.err {
            imp::err_erase_line(self);
            self.needs_clear = false;
        }
    }

    /// Shortcut to right-align and color green a status message.
    pub fn status<T, U>(&mut self, status: T, message: U) -> CargoResult<()>
    where
        T: fmt::Display,
        U: fmt::Display,
    {
        self.print(&status, Some(&message), Green, true)
    }

    pub fn status_header<T>(&mut self, status: T) -> CargoResult<()>
    where
        T: fmt::Display,
    {
        self.print(&status, None, Cyan, true)
    }

    /// Shortcut to right-align a status message.
    pub fn status_with_color<T, U>(
        &mut self,
        status: T,
        message: U,
        color: Color,
    ) -> CargoResult<()>
    where
        T: fmt::Display,
        U: fmt::Display,
    {
        self.print(&status, Some(&message), color, true)
    }

    /// Runs the callback only if we are in verbose mode.
    pub fn verbose<F>(&mut self, mut callback: F) -> CargoResult<()>
    where
        F: FnMut(&mut Shell) -> CargoResult<()>,
    {
        match self.verbosity {
            Verbosity::Verbose => callback(self),
            _ => Ok(()),
        }
    }

    /// Runs the callback if we are not in verbose mode.
    pub fn concise<F>(&mut self, mut callback: F) -> CargoResult<()>
    where
        F: FnMut(&mut Shell) -> CargoResult<()>,
    {
        match self.verbosity {
            Verbosity::Verbose => Ok(()),
            _ => callback(self),
        }
    }

    /// Prints a red 'error' message.
    pub fn error<T: fmt::Display>(&mut self, message: T) -> CargoResult<()> {
        if self.needs_clear {
            self.err_erase_line();
        }
        self.err.print(&"error", Some(&message), Red, false)
    }

    /// Prints an amber 'warning' message.
    pub fn warn<T: fmt::Display>(&mut self, message: T) -> CargoResult<()> {
        match self.verbosity {
            Verbosity::Quiet => Ok(()),
            _ => self.print(&"warning", Some(&message), Yellow, false),
        }
    }

    /// Updates the verbosity of the shell.
    pub fn set_verbosity(&mut self, verbosity: Verbosity) {
        self.verbosity = verbosity;
    }

    /// Gets the verbosity of the shell.
    pub fn verbosity(&self) -> Verbosity {
        self.verbosity
    }

    /// Updates the color choice (always, never, or auto) from a string..
    pub fn set_color_choice(&mut self, color: Option<&str>) -> CargoResult<()> {
        if let ShellOut::Stream {
            ref mut stream,
            ref mut color_choice,
            ..
        } = self.err
        {
            let cfg = match color {
                Some("always") => ColorChoice::Always,
                Some("never") => ColorChoice::Never,

                Some("auto") | None => ColorChoice::CargoAuto,

                Some(arg) => failure::bail!(
                    "argument for --color must be auto, always, or \
                     never, but found `{}`",
                    arg
                ),
            };
            *color_choice = cfg;
            *stream = StandardStream::stderr(cfg.to_termcolor_color_choice());
        }
        Ok(())
    }

    /// Gets the current color choice.
    ///
    /// If we are not using a color stream, this will always return `Never`, even if the color
    /// choice has been set to something else.
    pub fn color_choice(&self) -> ColorChoice {
        match self.err {
            ShellOut::Stream { color_choice, .. } => color_choice,
            ShellOut::Write(_) => ColorChoice::Never,
        }
    }

    /// Whether the shell supports color.
    pub fn supports_color(&self) -> bool {
        match &self.err {
            ShellOut::Write(_) => false,
            ShellOut::Stream { stream, .. } => stream.supports_color(),
        }
    }

    /// Prints a message and translates ANSI escape code into console colors.
    pub fn print_ansi(&mut self, message: &[u8]) -> CargoResult<()> {
        if self.needs_clear {
            self.err_erase_line();
        }
        #[cfg(windows)]
        {
            if let ShellOut::Stream { stream, .. } = &mut self.err {
                ::fwdansi::write_ansi(stream, message)?;
                return Ok(());
            }
        }
        self.err().write_all(message)?;
        Ok(())
    }
}

impl Default for Shell {
    fn default() -> Self {
        Self::new()
    }
}

impl ShellOut {
    /// Prints out a message with a status. The status comes first, and is bold plus the given
    /// color. The status can be justified, in which case the max width that will right align is
    /// 12 chars.
    fn print(
        &mut self,
        status: &dyn fmt::Display,
        message: Option<&dyn fmt::Display>,
        color: Color,
        justified: bool,
    ) -> CargoResult<()> {
        match *self {
            ShellOut::Stream { ref mut stream, .. } => {
                stream.reset()?;
                stream.set_color(ColorSpec::new().set_bold(true).set_fg(Some(color)))?;
                if justified {
                    write!(stream, "{:>12}", status)?;
                } else {
                    write!(stream, "{}", status)?;
                    stream.set_color(ColorSpec::new().set_bold(true))?;
                    write!(stream, ":")?;
                }
                stream.reset()?;
                match message {
                    Some(message) => writeln!(stream, " {}", message)?,
                    None => write!(stream, " ")?,
                }
            }
            ShellOut::Write(ref mut w) => {
                if justified {
                    write!(w, "{:>12}", status)?;
                } else {
                    write!(w, "{}:", status)?;
                }
                match message {
                    Some(message) => writeln!(w, " {}", message)?,
                    None => write!(w, " ")?,
                }
            }
        }
        Ok(())
    }

    /// Gets this object as a `io::Write`.
    fn as_write(&mut self) -> &mut dyn Write {
        match *self {
            ShellOut::Stream { ref mut stream, .. } => stream,
            ShellOut::Write(ref mut w) => w,
        }
    }
}

impl ColorChoice {
    /// Converts our color choice to termcolor's version.
    fn to_termcolor_color_choice(self) -> termcolor::ColorChoice {
        match self {
            ColorChoice::Always => termcolor::ColorChoice::Always,
            ColorChoice::Never => termcolor::ColorChoice::Never,
            ColorChoice::CargoAuto => {
                if atty::is(atty::Stream::Stderr) {
                    termcolor::ColorChoice::Auto
                } else {
                    termcolor::ColorChoice::Never
                }
            }
        }
    }
}

#[cfg(any(target_os = "linux", target_os = "macos", target_os = "freebsd"))]
mod imp {
    use std::mem;

    use libc;

    use super::Shell;

    pub fn stderr_width() -> Option<usize> {
        unsafe {
            let mut winsize: libc::winsize = mem::zeroed();
            // The .into() here is needed for FreeBSD which defines TIOCGWINSZ
            // as c_uint but ioctl wants c_ulong.
            if libc::ioctl(libc::STDERR_FILENO, libc::TIOCGWINSZ.into(), &mut winsize) < 0 {
                return None;
            }
            if winsize.ws_col > 0 {
                Some(winsize.ws_col as usize)
            } else {
                None
            }
        }
    }

    pub fn err_erase_line(shell: &mut Shell) {
        // This is the "EL - Erase in Line" sequence. It clears from the cursor
        // to the end of line.
        // https://en.wikipedia.org/wiki/ANSI_escape_code#CSI_sequences
        let _ = shell.err.as_write().write_all(b"\x1B[K");
    }
}

#[cfg(all(
    unix,
    not(any(target_os = "linux", target_os = "macos", target_os = "freebsd"))
))]
mod imp {
    pub(super) use super::default_err_erase_line as err_erase_line;

    pub fn stderr_width() -> Option<usize> {
        None
    }
}

#[cfg(windows)]
mod imp {
    use std::{cmp, mem, ptr};
    use winapi::um::fileapi::*;
    use winapi::um::handleapi::*;
    use winapi::um::processenv::*;
    use winapi::um::winbase::*;
    use winapi::um::wincon::*;
    use winapi::um::winnt::*;

    pub(super) use super::default_err_erase_line as err_erase_line;

    pub fn stderr_width() -> Option<usize> {
        unsafe {
            let stdout = GetStdHandle(STD_ERROR_HANDLE);
            let mut csbi: CONSOLE_SCREEN_BUFFER_INFO = mem::zeroed();
            if GetConsoleScreenBufferInfo(stdout, &mut csbi) != 0 {
                return Some((csbi.srWindow.Right - csbi.srWindow.Left) as usize);
            }

            // On mintty/msys/cygwin based terminals, the above fails with
            // INVALID_HANDLE_VALUE. Use an alternate method which works
            // in that case as well.
            let h = CreateFileA(
                "CONOUT$\0".as_ptr() as *const CHAR,
                GENERIC_READ | GENERIC_WRITE,
                FILE_SHARE_READ | FILE_SHARE_WRITE,
                ptr::null_mut(),
                OPEN_EXISTING,
                0,
                ptr::null_mut(),
            );
            if h == INVALID_HANDLE_VALUE {
                return None;
            }

            let mut csbi: CONSOLE_SCREEN_BUFFER_INFO = mem::zeroed();
            let rc = GetConsoleScreenBufferInfo(h, &mut csbi);
            CloseHandle(h);
            if rc != 0 {
                let width = (csbi.srWindow.Right - csbi.srWindow.Left) as usize;
                // Unfortunately cygwin/mintty does not set the size of the
                // backing console to match the actual window size. This
                // always reports a size of 80 or 120 (not sure what
                // determines that). Use a conservative max of 60 which should
                // work in most circumstances. ConEmu does some magic to
                // resize the console correctly, but there's no reasonable way
                // to detect which kind of terminal we are running in, or if
                // GetConsoleScreenBufferInfo returns accurate information.
                return Some(cmp::min(60, width));
            }
            None
        }
    }
}

#[cfg(any(
    all(
        unix,
        not(any(target_os = "linux", target_os = "macos", target_os = "freebsd"))
    ),
    windows,
))]
fn default_err_erase_line(shell: &mut Shell) {
    if let Some(max_width) = imp::stderr_width() {
        let blank = " ".repeat(max_width);
        drop(write!(shell.err.as_write(), "{}\r", blank));
    }
}
<|endoftext|>


<commit_before>use clap::{Arg, Command};
use reqwest::blocking::Client;
use serde::Deserialize;
use serde_json::json;

use std::collections::HashMap;
use std::env;
use std::fs;

mod settings;
use settings::Settings;

#[derive(Deserialize, Debug)]
struct PathShortened {
    path_shortened: String,
}

fn main() -> Result<(), Box<dyn std::error::Error>> {
    let settings = Settings::new().unwrap();

    let cmd = clap::Command::new("lhc")
        .version("v0.0")
        .propagate_version(true)
        .subcommand_required(true)
        .arg_required_else_help(true)
        .subcommand(
            Command::new("path-shorten")
                .about("Shorten a path")
                .arg(
                    Arg::new("path_to_shorten")
                        .help("the path to shorten")
                        .index(1)
                        .required(true),
                )
                .arg(
                    Arg::new("path_aliases_file")
                        .long("path-aliases")
                        .value_name("FILE")
                        .help("File containing path aliases")
                        .takes_value(true)
                        .required(true),
                ),
        )
        .subcommand(Command::new("ping").about("Check lh server connectivity"))
        .subcommand(Command::new("shutdown").about("Shut down lh server instance"));
    let matches = cmd.get_matches();
    let client = Client::new();
    match matches.subcommand_name() {
        Some("path-shorten") => {
            if let Some(m) = matches.subcommand_matches("path-shorten") {
                // Safe to use unwrap() because of the required() option
                let path_to_shorten = m.value_of("path_to_shorten").unwrap();
                let path_aliases_file = m.value_of("path_aliases_file").unwrap();
                let path_aliases_contents =
                    fs::read_to_string(path_aliases_file).unwrap_or_default();

                let mut subs = HashMap::new();
                let home_dir = env::var("HOME").expect("$HOME not set");

                // This is used to expand the definitions in the path aliases file.
                subs.insert("$HOME", &home_dir);

                // This is used to shrink the /home/... directory to "~" if
                // there is no path aliases match.
                let tilde = "~".to_string();
                subs.insert(&home_dir, &tilde);

                let request_url = format!(
                    "http://{}:{}/path-shorten",
                    settings.server.domain, settings.server.port
                );

                let json_body = json!({
                    "name": path_to_shorten,
                    "aliases_raw": path_aliases_contents,
                    "substitutions": subs,
                });

                let response = client.post(request_url).json(&json_body).send()?;

                let path_shortened: PathShortened = response.json()?;

                print!("{}", path_shortened.path_shortened);
            }
        }
        Some("ping") => {
            let request_url = format!(
                "http://{}:{}/ping",
                settings.server.domain, settings.server.port
            );

            let response = client.get(request_url).send()?;

            println!("{}", response.text()?);
        }
        Some("shutdown") => {
            let request_url = format!(
                "http://{}:{}/shutdown",
                settings.server.domain, settings.server.port
            );

            let response = client.get(request_url).send()?;

            println!("{}", response.text()?);
        }
        _ => unreachable!("clap should ensure we don't get here"),
    }
    Ok(())
}
<commit_msg>lhc: simplify ping request<commit_after>use clap::{Arg, Command};
use reqwest::blocking::Client;
use serde::Deserialize;
use serde_json::json;

use std::collections::HashMap;
use std::env;
use std::fs;

mod settings;
use settings::Settings;

#[derive(Deserialize, Debug)]
struct PathShortened {
    path_shortened: String,
}

fn main() -> Result<(), Box<dyn std::error::Error>> {
    let settings = Settings::new().unwrap();

    let cmd = clap::Command::new("lhc")
        .version("v0.0")
        .propagate_version(true)
        .subcommand_required(true)
        .arg_required_else_help(true)
        .subcommand(
            Command::new("path-shorten")
                .about("Shorten a path")
                .arg(
                    Arg::new("path_to_shorten")
                        .help("the path to shorten")
                        .index(1)
                        .required(true),
                )
                .arg(
                    Arg::new("path_aliases_file")
                        .long("path-aliases")
                        .value_name("FILE")
                        .help("File containing path aliases")
                        .takes_value(true)
                        .required(true),
                ),
        )
        .subcommand(Command::new("ping").about("Check lh server connectivity"))
        .subcommand(Command::new("shutdown").about("Shut down lh server instance"));
    let matches = cmd.get_matches();
    let client = Client::new();
    match matches.subcommand_name() {
        Some("path-shorten") => {
            if let Some(m) = matches.subcommand_matches("path-shorten") {
                // Safe to use unwrap() because of the required() option
                let path_to_shorten = m.value_of("path_to_shorten").unwrap();
                let path_aliases_file = m.value_of("path_aliases_file").unwrap();
                let path_aliases_contents =
                    fs::read_to_string(path_aliases_file).unwrap_or_default();

                let mut subs = HashMap::new();
                let home_dir = env::var("HOME").expect("$HOME not set");

                // This is used to expand the definitions in the path aliases file.
                subs.insert("$HOME", &home_dir);

                // This is used to shrink the /home/... directory to "~" if
                // there is no path aliases match.
                let tilde = "~".to_string();
                subs.insert(&home_dir, &tilde);

                let request_url = format!(
                    "http://{}:{}/path-shorten",
                    settings.server.domain, settings.server.port
                );

                let json_body = json!({
                    "name": path_to_shorten,
                    "aliases_raw": path_aliases_contents,
                    "substitutions": subs,
                });

                let response = client.post(request_url).json(&json_body).send()?;

                let path_shortened: PathShortened = response.json()?;

                print!("{}", path_shortened.path_shortened);
            }
        }
        Some("ping") => {
            let request_url = format!(
                "http://{}:{}/ping",
                settings.server.domain, settings.server.port
            );

            client.get(request_url).send()?;

            println!("OK");
        }
        Some("shutdown") => {
            let request_url = format!(
                "http://{}:{}/shutdown",
                settings.server.domain, settings.server.port
            );

            let response = client.get(request_url).send()?;

            println!("{}", response.text()?);
        }
        _ => unreachable!("clap should ensure we don't get here"),
    }
    Ok(())
}
<|endoftext|>


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

use crate::counters::{FETCHED_TRANSACTION, UNABLE_TO_FETCH_TRANSACTION};
use aptos_rest_client::{Client as RestClient, State, Transaction};
use futures::channel::mpsc;
use futures::{SinkExt, StreamExt};
use serde_json::Value;
use std::time::Duration;
use tokio::task::JoinHandle;
use url::Url;

// TODO: make this configurable
const RETRY_TIME_MILLIS: u64 = 1000;
const TRANSACTION_FETCH_BATCH_SIZE: u16 = 500;
const TRANSACTION_CHANNEL_SIZE: usize = 50_000;
const MAX_THREADS: usize = 10;
const MAX_RETRIES: usize = 5;

#[derive(Debug)]
pub struct Fetcher {
    client: RestClient,
    chain_id: u8,
    current_version: u64,
    highest_known_version: u64,
    transactions_sender: mpsc::Sender<Vec<Transaction>>,
}

impl Fetcher {
    pub fn new(
        client: RestClient,
        current_version: u64,
        transactions_sender: mpsc::Sender<Vec<Transaction>>,
    ) -> Self {
        Self {
            client,
            chain_id: 0,
            current_version,
            highest_known_version: current_version,
            transactions_sender,
        }
    }

    pub async fn fetch_ledger_info(&mut self) -> State {
        self.client
            .get_ledger_information()
            .await
            .expect("ledger info must be present")
            .into_inner()
    }

    pub async fn set_highest_known_version(&mut self) {
        let info = self.client.get_ledger_information().await;
        let res = info.unwrap();
        let state = res.state();
        self.highest_known_version = state.version;
        self.chain_id = state.chain_id;
    }

    pub async fn run(&mut self) {
        loop {
            if self.current_version == self.highest_known_version {
                tokio::time::sleep(Duration::from_millis(200)).await;
                self.set_highest_known_version().await;
            }

            let num_missing = self.highest_known_version - self.current_version;
            let num_batches = std::cmp::min(
                (num_missing as f64 / TRANSACTION_FETCH_BATCH_SIZE as f64).ceil() as u64,
                MAX_THREADS as u64,
            ) as usize;
            let mut futures = vec![];
            for i in 0..num_batches {
                futures.push(fetch_nexts(
                    self.client.clone(),
                    self.current_version + (i as u64 * TRANSACTION_FETCH_BATCH_SIZE as u64),
                ));
            }
            let mut res: Vec<Vec<Transaction>> = futures::future::join_all(futures).await;
            res.sort_by(|a, b| {
                a.first()
                    .unwrap()
                    .version()
                    .unwrap()
                    .cmp(&b.first().unwrap().version().unwrap())
            });

            for batch in res {
                self.current_version = batch.last().unwrap().version().unwrap();
                self.transactions_sender.send(batch).await.unwrap();
            }
        }
    }
}

/// Fetches the next version based on its internal version counter
/// Under the hood, it fetches TRANSACTION_FETCH_BATCH_SIZE versions in bulk (when needed), and uses that buffer to feed out
/// In the event it can't fetch, it will keep retrying every RETRY_TIME_MILLIS ms
async fn fetch_nexts(client: RestClient, starting_version: u64) -> Vec<Transaction> {
    let mut retries = 0;
    while retries < MAX_RETRIES {
        retries += 1;

        let res = client
            .get_transactions(Some(starting_version), Some(TRANSACTION_FETCH_BATCH_SIZE))
            .await;

        match res {
            Ok(response) => {
                FETCHED_TRANSACTION.inc();
                return remove_null_bytes_from_txns(response.into_inner());
            }
            Err(err) => {
                let err_str = err.to_string();
                // If it's a 404, then we're all caught up; no need to increment the `UNABLE_TO_FETCH_TRANSACTION` counter
                if err_str.contains("404") {
                    aptos_logger::debug!(
                            "Could not fetch {} transactions starting at {}: all caught up. Will check again in {}ms.",
                            TRANSACTION_FETCH_BATCH_SIZE,
                            starting_version,
                            RETRY_TIME_MILLIS,
                        );
                }
                UNABLE_TO_FETCH_TRANSACTION.inc();
                aptos_logger::error!(
                    "Could not fetch {} transactions starting at {}, will retry in {}ms. Err: {:?}",
                    TRANSACTION_FETCH_BATCH_SIZE,
                    starting_version,
                    RETRY_TIME_MILLIS,
                    err
                );
            }
        }
        tokio::time::sleep(Duration::from_millis(RETRY_TIME_MILLIS)).await;
    }

    panic!(
        "Could not fetch {} transactions starting at {}!",
        TRANSACTION_FETCH_BATCH_SIZE, starting_version
    );
}

#[derive(Debug)]
pub struct TransactionFetcher {
    starting_version: u64,
    client: RestClient,
    fetcher_handle: Option<JoinHandle<()>>,
    transactions_sender: Option<mpsc::Sender<Vec<Transaction>>>,
    transaction_receiver: mpsc::Receiver<Vec<Transaction>>,
}

impl TransactionFetcher {
    pub fn new(node_url: Url, starting_version: Option<u64>) -> Self {
        let (transactions_sender, transaction_receiver) =
            mpsc::channel::<Vec<Transaction>>(TRANSACTION_CHANNEL_SIZE);

        let client = RestClient::new(node_url);

        Self {
            starting_version: starting_version.unwrap_or(0),
            client,
            fetcher_handle: None,
            transactions_sender: Some(transactions_sender),
            transaction_receiver,
        }
    }
}

#[async_trait::async_trait]
impl TransactionFetcherTrait for TransactionFetcher {
    /// Fetches the next batch based on its internal version counter
    /// Under the hood, it fetches TRANSACTION_FETCH_BATCH_SIZE versions in bulk (when needed), and uses that buffer to feed out
    /// In the event it can't fetch, it will keep retrying every RETRY_TIME_MILLIS ms
    async fn fetch_next_batch(&mut self) -> Vec<Transaction> {
        self.transaction_receiver.next().await.unwrap()
    }

    /// fetches one version; this used for error checking/repair/etc
    /// In the event it can't, it will keep retrying every RETRY_TIME_MILLIS ms
    async fn fetch_version(&self, version: u64) -> Transaction {
        loop {
            let res = self.client.get_transaction_by_version(version).await;
            match res {
                Ok(response) => {
                    FETCHED_TRANSACTION.inc();
                    return response.into_inner();
                }
                Err(err) => {
                    UNABLE_TO_FETCH_TRANSACTION.inc();
                    aptos_logger::error!(
                        "Could not fetch version {}, will retry in {}ms. Err: {:?}",
                        version,
                        RETRY_TIME_MILLIS,
                        err
                    );
                    tokio::time::sleep(Duration::from_millis(RETRY_TIME_MILLIS)).await;
                }
            };
        }
    }

    async fn fetch_ledger_info(&mut self) -> State {
        self.client
            .get_ledger_information()
            .await
            .expect("ledger info must be present")
            .into_inner()
    }

    async fn set_version(&mut self, version: u64) {
        if self.fetcher_handle.is_some() {
            panic!("TransactionFetcher already started!");
        }
        self.starting_version = version;
    }

    async fn start(&mut self) {
        if self.fetcher_handle.is_some() {
            panic!("TransactionFetcher already started!");
        }
        let client = self.client.clone();
        let transactions_sender = self.transactions_sender.take().unwrap();
        let starting_version = self.starting_version;
        let fetcher_handle = tokio::spawn(async move {
            let mut fetcher = Fetcher::new(client, starting_version, transactions_sender);
            fetcher.run().await;
        });
        self.fetcher_handle = Some(fetcher_handle);
    }
}

pub fn string_null_byte_replacement(value: &mut str) -> String {
    value.replace('\u{0000}', "").replace("\\u0000", "")
}

pub fn recurse_remove_null_bytes_from_json(sub_json: &mut Value) {
    match sub_json {
        Value::Array(array) => {
            for item in array {
                recurse_remove_null_bytes_from_json(item);
            }
        }
        Value::Object(object) => {
            for (_key, value) in object {
                recurse_remove_null_bytes_from_json(value);
            }
        }
        Value::String(str) => {
            if !str.is_empty() {
                let replacement = string_null_byte_replacement(str);
                *str = replacement;
            }
        }
        _ => {}
    }
}

pub fn remove_null_bytes_from_txns(txns: Vec<Transaction>) -> Vec<Transaction> {
    txns.iter()
        .map(|txn| {
            let mut txn_json = serde_json::to_value(txn).unwrap();
            recurse_remove_null_bytes_from_json(&mut txn_json);
            serde_json::from_value::<Transaction>(txn_json).unwrap()
        })
        .collect::<Vec<Transaction>>()
}

/// For mocking TransactionFetcher in tests
#[async_trait::async_trait]
pub trait TransactionFetcherTrait: Send + Sync {
    async fn fetch_next_batch(&mut self) -> Vec<Transaction>;

    async fn fetch_version(&self, version: u64) -> Transaction;

    async fn fetch_ledger_info(&mut self) -> State;

    async fn set_version(&mut self, version: u64);

    async fn start(&mut self);
}
<commit_msg>[ECO] Set indexer max batches to 20 (#3446)<commit_after>// Copyright (c) Aptos
// SPDX-License-Identifier: Apache-2.0

use crate::counters::{FETCHED_TRANSACTION, UNABLE_TO_FETCH_TRANSACTION};
use aptos_rest_client::{Client as RestClient, State, Transaction};
use futures::channel::mpsc;
use futures::{SinkExt, StreamExt};
use serde_json::Value;
use std::time::Duration;
use tokio::task::JoinHandle;
use url::Url;

// TODO: make this configurable
const RETRY_TIME_MILLIS: u64 = 1000;
const TRANSACTION_FETCH_BATCH_SIZE: u16 = 500;
const TRANSACTION_CHANNEL_SIZE: usize = 35;
const MAX_THREADS: usize = 10;
const MAX_RETRIES: usize = 5;

#[derive(Debug)]
pub struct Fetcher {
    client: RestClient,
    chain_id: u8,
    current_version: u64,
    highest_known_version: u64,
    transactions_sender: mpsc::Sender<Vec<Transaction>>,
}

impl Fetcher {
    pub fn new(
        client: RestClient,
        current_version: u64,
        transactions_sender: mpsc::Sender<Vec<Transaction>>,
    ) -> Self {
        Self {
            client,
            chain_id: 0,
            current_version,
            highest_known_version: current_version,
            transactions_sender,
        }
    }

    pub async fn fetch_ledger_info(&mut self) -> State {
        self.client
            .get_ledger_information()
            .await
            .expect("ledger info must be present")
            .into_inner()
    }

    pub async fn set_highest_known_version(&mut self) {
        let info = self.client.get_ledger_information().await;
        let res = info.unwrap();
        let state = res.state();
        self.highest_known_version = state.version;
        self.chain_id = state.chain_id;
    }

    pub async fn run(&mut self) {
        loop {
            if self.current_version == self.highest_known_version {
                tokio::time::sleep(Duration::from_millis(200)).await;
                self.set_highest_known_version().await;
            }

            let num_missing = self.highest_known_version - self.current_version;
            let num_batches = std::cmp::min(
                (num_missing as f64 / TRANSACTION_FETCH_BATCH_SIZE as f64).ceil() as u64,
                MAX_THREADS as u64,
            ) as usize;
            let mut futures = vec![];
            for i in 0..num_batches {
                futures.push(fetch_nexts(
                    self.client.clone(),
                    self.current_version + (i as u64 * TRANSACTION_FETCH_BATCH_SIZE as u64),
                ));
            }
            let mut res: Vec<Vec<Transaction>> = futures::future::join_all(futures).await;
            res.sort_by(|a, b| {
                a.first()
                    .unwrap()
                    .version()
                    .unwrap()
                    .cmp(&b.first().unwrap().version().unwrap())
            });

            for batch in res {
                self.current_version = batch.last().unwrap().version().unwrap();
                self.transactions_sender.send(batch).await.unwrap();
            }
        }
    }
}

/// Fetches the next version based on its internal version counter
/// Under the hood, it fetches TRANSACTION_FETCH_BATCH_SIZE versions in bulk (when needed), and uses that buffer to feed out
/// In the event it can't fetch, it will keep retrying every RETRY_TIME_MILLIS ms
async fn fetch_nexts(client: RestClient, starting_version: u64) -> Vec<Transaction> {
    let mut retries = 0;
    while retries < MAX_RETRIES {
        retries += 1;

        let res = client
            .get_transactions(Some(starting_version), Some(TRANSACTION_FETCH_BATCH_SIZE))
            .await;

        match res {
            Ok(response) => {
                FETCHED_TRANSACTION.inc();
                return remove_null_bytes_from_txns(response.into_inner());
            }
            Err(err) => {
                let err_str = err.to_string();
                // If it's a 404, then we're all caught up; no need to increment the `UNABLE_TO_FETCH_TRANSACTION` counter
                if err_str.contains("404") {
                    aptos_logger::debug!(
                            "Could not fetch {} transactions starting at {}: all caught up. Will check again in {}ms.",
                            TRANSACTION_FETCH_BATCH_SIZE,
                            starting_version,
                            RETRY_TIME_MILLIS,
                        );
                }
                UNABLE_TO_FETCH_TRANSACTION.inc();
                aptos_logger::error!(
                    "Could not fetch {} transactions starting at {}, will retry in {}ms. Err: {:?}",
                    TRANSACTION_FETCH_BATCH_SIZE,
                    starting_version,
                    RETRY_TIME_MILLIS,
                    err
                );
            }
        }
        tokio::time::sleep(Duration::from_millis(RETRY_TIME_MILLIS)).await;
    }

    panic!(
        "Could not fetch {} transactions starting at {}!",
        TRANSACTION_FETCH_BATCH_SIZE, starting_version
    );
}

#[derive(Debug)]
pub struct TransactionFetcher {
    starting_version: u64,
    client: RestClient,
    fetcher_handle: Option<JoinHandle<()>>,
    transactions_sender: Option<mpsc::Sender<Vec<Transaction>>>,
    transaction_receiver: mpsc::Receiver<Vec<Transaction>>,
}

impl TransactionFetcher {
    pub fn new(node_url: Url, starting_version: Option<u64>) -> Self {
        let (transactions_sender, transaction_receiver) =
            mpsc::channel::<Vec<Transaction>>(TRANSACTION_CHANNEL_SIZE);

        let client = RestClient::new(node_url);

        Self {
            starting_version: starting_version.unwrap_or(0),
            client,
            fetcher_handle: None,
            transactions_sender: Some(transactions_sender),
            transaction_receiver,
        }
    }
}

#[async_trait::async_trait]
impl TransactionFetcherTrait for TransactionFetcher {
    /// Fetches the next batch based on its internal version counter
    /// Under the hood, it fetches TRANSACTION_FETCH_BATCH_SIZE versions in bulk (when needed), and uses that buffer to feed out
    /// In the event it can't fetch, it will keep retrying every RETRY_TIME_MILLIS ms
    async fn fetch_next_batch(&mut self) -> Vec<Transaction> {
        self.transaction_receiver.next().await.unwrap()
    }

    /// fetches one version; this used for error checking/repair/etc
    /// In the event it can't, it will keep retrying every RETRY_TIME_MILLIS ms
    async fn fetch_version(&self, version: u64) -> Transaction {
        loop {
            let res = self.client.get_transaction_by_version(version).await;
            match res {
                Ok(response) => {
                    FETCHED_TRANSACTION.inc();
                    return response.into_inner();
                }
                Err(err) => {
                    UNABLE_TO_FETCH_TRANSACTION.inc();
                    aptos_logger::error!(
                        "Could not fetch version {}, will retry in {}ms. Err: {:?}",
                        version,
                        RETRY_TIME_MILLIS,
                        err
                    );
                    tokio::time::sleep(Duration::from_millis(RETRY_TIME_MILLIS)).await;
                }
            };
        }
    }

    async fn fetch_ledger_info(&mut self) -> State {
        self.client
            .get_ledger_information()
            .await
            .expect("ledger info must be present")
            .into_inner()
    }

    async fn set_version(&mut self, version: u64) {
        if self.fetcher_handle.is_some() {
            panic!("TransactionFetcher already started!");
        }
        self.starting_version = version;
    }

    async fn start(&mut self) {
        if self.fetcher_handle.is_some() {
            panic!("TransactionFetcher already started!");
        }
        let client = self.client.clone();
        let transactions_sender = self.transactions_sender.take().unwrap();
        let starting_version = self.starting_version;
        let fetcher_handle = tokio::spawn(async move {
            let mut fetcher = Fetcher::new(client, starting_version, transactions_sender);
            fetcher.run().await;
        });
        self.fetcher_handle = Some(fetcher_handle);
    }
}

pub fn string_null_byte_replacement(value: &mut str) -> String {
    value.replace('\u{0000}', "").replace("\\u0000", "")
}

pub fn recurse_remove_null_bytes_from_json(sub_json: &mut Value) {
    match sub_json {
        Value::Array(array) => {
            for item in array {
                recurse_remove_null_bytes_from_json(item);
            }
        }
        Value::Object(object) => {
            for (_key, value) in object {
                recurse_remove_null_bytes_from_json(value);
            }
        }
        Value::String(str) => {
            if !str.is_empty() {
                let replacement = string_null_byte_replacement(str);
                *str = replacement;
            }
        }
        _ => {}
    }
}

pub fn remove_null_bytes_from_txns(txns: Vec<Transaction>) -> Vec<Transaction> {
    txns.iter()
        .map(|txn| {
            let mut txn_json = serde_json::to_value(txn).unwrap();
            recurse_remove_null_bytes_from_json(&mut txn_json);
            serde_json::from_value::<Transaction>(txn_json).unwrap()
        })
        .collect::<Vec<Transaction>>()
}

/// For mocking TransactionFetcher in tests
#[async_trait::async_trait]
pub trait TransactionFetcherTrait: Send + Sync {
    async fn fetch_next_batch(&mut self) -> Vec<Transaction>;

    async fn fetch_version(&self, version: u64) -> Transaction;

    async fn fetch_ledger_info(&mut self) -> State;

    async fn set_version(&mut self, version: u64);

    async fn start(&mut self);
}
<|endoftext|>


<commit_before>use proc_macro2::TokenStream;
use quote::{quote, ToTokens};
use std::env;
use std::ffi::OsStr;

pub fn fold_quote<F, I, T>(input: impl Iterator<Item = I>, f: F) -> TokenStream
where
    F: Fn(I) -> T,
    T: ToTokens,
{
    input.fold(quote! {}, |acc, x| {
        let y = f(x);
        quote! { #acc #y }
    })
}

pub fn is_unit(v: &syn::Variant) -> bool {
    match v.fields {
        syn::Fields::Unit => true,
        _ => false,
    }
}

#[cfg(feature = "debug-with-rustfmt")]
/// Pretty-print the output of proc macro using rustfmt.
pub fn debug_with_rustfmt(input: &TokenStream) {
    use std::io::Write;
    use std::process::{Command, Stdio};

    let rustfmt_var = env::var_os("RUSTFMT");
    let rustfmt = match &rustfmt_var {
        Some(rustfmt) => rustfmt,
        None => OsStr::new("rustfmt"),
    };
    let mut child = Command::new(rustfmt)
        .stdin(Stdio::piped())
        .stdout(Stdio::piped())
        .spawn()
        .expect("Failed to spawn rustfmt in stdio mode");
    {
        let stdin = child.stdin.as_mut().expect("Failed to get stdin");
        stdin
            .write_all(format!("{}", input).as_bytes())
            .expect("Failed to write to stdin");
    }
    let rustfmt_output = child.wait_with_output().expect("rustfmt has failed");

    eprintln!(
        "{}",
        String::from_utf8(rustfmt_output.stdout).expect("rustfmt returned non-UTF8 string")
    );
}
<commit_msg>refactor: move env/osstr imports to avoid warnings<commit_after>use proc_macro2::TokenStream;
use quote::{quote, ToTokens};

pub fn fold_quote<F, I, T>(input: impl Iterator<Item = I>, f: F) -> TokenStream
where
    F: Fn(I) -> T,
    T: ToTokens,
{
    input.fold(quote! {}, |acc, x| {
        let y = f(x);
        quote! { #acc #y }
    })
}

pub fn is_unit(v: &syn::Variant) -> bool {
    match v.fields {
        syn::Fields::Unit => true,
        _ => false,
    }
}

#[cfg(feature = "debug-with-rustfmt")]
/// Pretty-print the output of proc macro using rustfmt.
pub fn debug_with_rustfmt(input: &TokenStream) {
    use std::io::Write;
    use std::process::{Command, Stdio};
    use std::env;
    use std::ffi::OsStr;

    let rustfmt_var = env::var_os("RUSTFMT");
    let rustfmt = match &rustfmt_var {
        Some(rustfmt) => rustfmt,
        None => OsStr::new("rustfmt"),
    };
    let mut child = Command::new(rustfmt)
        .stdin(Stdio::piped())
        .stdout(Stdio::piped())
        .spawn()
        .expect("Failed to spawn rustfmt in stdio mode");
    {
        let stdin = child.stdin.as_mut().expect("Failed to get stdin");
        stdin
            .write_all(format!("{}", input).as_bytes())
            .expect("Failed to write to stdin");
    }
    let rustfmt_output = child.wait_with_output().expect("rustfmt has failed");

    eprintln!(
        "{}",
        String::from_utf8(rustfmt_output.stdout).expect("rustfmt returned non-UTF8 string")
    );
}
<|endoftext|>


<commit_before>extern crate nalgebra as na;
extern crate piston;
extern crate graphics;
extern crate glutin_window;
extern crate opengl_graphics;
extern crate alga;

use opengl_graphics::{GlGraphics, OpenGL};
use glutin_window::GlutinWindow;
use piston::window::WindowSettings;
use piston::event_loop::{Events, EventSettings};
use piston::input::*;

use alga::linear::Transformation;

struct Edge<'a> {
    p1: &'a na::Point3<f64>,
    p2: &'a na::Point3<f64>
}

struct Model<'a> {
    edges: Vec<Edge<'a>>
}

pub struct App<'a> {
    gl: GlGraphics,
    model: Model<'a>
}

impl<'a> App<'a> {
    fn render(&mut self, args: &RenderArgs) {
        let ref model = self.model;
        self.gl.draw(args.viewport(), |c, gl| {
            // Clear the screen.
            graphics::clear([0.7, 0.6, 0.75, 1.0], gl);

            let proj = na::Orthographic3::new(1.0, 2.0, -3.0, -2.5, 10.0, 900.0);
            let rotation = na::Rotation3::from_euler_angles(1.0, 1.0, 1.0);
            for edge in &model.edges {
                let p1 = proj.project_point(&rotation.transform_point(edge.p1));
                let p2 = proj.project_point(&rotation.transform_point(edge.p2));
                graphics::line([0.0, 0.0, 0.0, 1.0],
                               1.0,
                               [p1[0],p1[1],p2[0],p2[1]],
                               c.transform, gl);
            }
        });
    }
}

fn main() {
    let nodes = vec![
        na::Point3::new(10.0, 50.0, 50.0),
        na::Point3::new(10.0, 50.0, 10.0),
        na::Point3::new(10.0, 10.0, 50.0),
        na::Point3::new(10.0, 10.0, 10.0),
        na::Point3::new(50.0, 50.0, 50.0),
        na::Point3::new(50.0, 50.0, 10.0),
        na::Point3::new(50.0, 10.0, 50.0),
        na::Point3::new(50.0, 10.0, 10.0)
    ];
    let edges = vec![
        Edge { p1: &nodes[0], p2: &nodes[1] },
        Edge { p1: &nodes[2], p2: &nodes[3] },
        Edge { p1: &nodes[4], p2: &nodes[5] },
        Edge { p1: &nodes[6], p2: &nodes[7] },
        Edge { p1: &nodes[0], p2: &nodes[2] },
        Edge { p1: &nodes[4], p2: &nodes[6] },
        Edge { p1: &nodes[1], p2: &nodes[3] },
        Edge { p1: &nodes[5], p2: &nodes[7] },
        Edge { p1: &nodes[0], p2: &nodes[4] },
        Edge { p1: &nodes[2], p2: &nodes[6] },
        Edge { p1: &nodes[1], p2: &nodes[5] },
        Edge { p1: &nodes[3], p2: &nodes[7] }
    ];
    let cube = Model { edges: edges };

    // Rendering stuff

    let opengl = OpenGL::V3_2;

    // Create an Glutin window.
    let mut window: GlutinWindow = WindowSettings::new(
            "Viewer",
            [800, 600]
        )
        .opengl(opengl)
        .exit_on_esc(true)
        .build()
        .unwrap();

    let mut app = App {
        gl: GlGraphics::new(opengl),
        model: cube
    };

    let mut events = Events::new(EventSettings::new());

    while let Some(e) = events.next(&mut window) {
        if let Some(r) = e.render_args() {
            app.render(&r);
        }
    }
}
<commit_msg>Rewrite renderer without nalgebra<commit_after>extern crate piston;
extern crate graphics;
extern crate glutin_window;
extern crate opengl_graphics;

use opengl_graphics::{GlGraphics, OpenGL};
use glutin_window::GlutinWindow;
use piston::window::WindowSettings;
use piston::event_loop::{Events, EventSettings};
use piston::input::*;

struct Point {
    x: f64,
    y: f64,
    z: f64
}

impl Point {
    fn rotate(&self, x: f64, y: f64, z: f64) -> Point {
        let mut rp = Point { x: self.x, y: self.y, z: self.z };

        let xn = rp.x * z.cos() - rp.y * z.sin();
        let yn = rp.x * z.sin() + rp.y * z.cos();
        rp.x = xn;
        rp.y = yn;

        let xn = rp.x * y.cos() - rp.z * y.sin();
        let zn = rp.x * y.sin() + rp.z * y.cos();
        rp.x = xn;
        rp.z = zn;

        let zn = rp.z * x.cos() - rp.y * x.sin();
        let yn = rp.z * x.sin() + rp.y * x.cos();
        rp.z = zn;
        rp.y = yn;

        rp
    }

    fn translate(&self, x: f64, y: f64, z: f64) -> Point {
        let rp = Point { x: self.x + x, y: self.y + y, z: self.z + z };
        rp
    }
}

struct Edge<'a> {
    p1: &'a Point,
    p2: &'a Point
}

struct Model<'a> {
    edges: Vec<Edge<'a>>
}

pub struct App<'a> {
    gl: GlGraphics,
    model: Model<'a>
}

impl<'a> App<'a> {
    fn render(&mut self, args: &RenderArgs) {
        let ref model = self.model;
        self.gl.draw(args.viewport(), |c, gl| {
            // Clear the screen.
            graphics::clear([0.7, 0.6, 0.75, 1.0], gl);

            for edge in &model.edges {
                let mut p1_t = edge.p1.rotate(1.0, 1.0, 1.0);
                let mut p2_t = edge.p2.rotate(1.0, 1.0, 1.0);

                let p1_x = p1_t.x + p1_t.z / 4.0;
                let p1_y = p1_t.y + p1_t.z / 4.0;
                let p2_x = p2_t.x + p2_t.z / 4.0;
                let p2_y = p2_t.y + p2_t.z / 4.0;
                graphics::line([0.0, 0.0, 0.0, 1.0],
                               1.0,
                               [p1_x + 400.0, p1_y + 300.0, p2_x + 400.0, p2_y + 300.0],
                               c.transform, gl);
            }
        });
    }
}

fn main() {
    let nodes = vec![
        Point {x: -50.0, y: 50.0,  z: 50.0},
        Point {x: -50.0, y: 50.0,  z: -50.0},
        Point {x: -50.0, y: -50.0, z: 50.0},
        Point {x: -50.0, y: -50.0, z: -50.0},
        Point {x: 50.0,  y: 50.0,  z: 50.0},
        Point {x: 50.0,  y: 50.0,  z: -50.0},
        Point {x: 50.0,  y: -50.0, z: 50.0},
        Point {x: 50.0,  y: -50.0, z: -50.0}
    ];
    let edges = vec![
        Edge { p1: &nodes[0], p2: &nodes[1] },
        Edge { p1: &nodes[2], p2: &nodes[3] },
        Edge { p1: &nodes[4], p2: &nodes[5] },
        Edge { p1: &nodes[6], p2: &nodes[7] },
        Edge { p1: &nodes[0], p2: &nodes[2] },
        Edge { p1: &nodes[4], p2: &nodes[6] },
        Edge { p1: &nodes[1], p2: &nodes[3] },
        Edge { p1: &nodes[5], p2: &nodes[7] },
        Edge { p1: &nodes[0], p2: &nodes[4] },
        Edge { p1: &nodes[2], p2: &nodes[6] },
        Edge { p1: &nodes[1], p2: &nodes[5] },
        Edge { p1: &nodes[3], p2: &nodes[7] }
    ];
    let cube = Model { edges: edges };

    // Rendering stuff

    let opengl = OpenGL::V3_2;

    // Create an Glutin window.
    let mut window: GlutinWindow = WindowSettings::new(
            "Viewer",
            [800, 600]
        )
        .opengl(opengl)
        .exit_on_esc(true)
        .build()
        .unwrap();

    let mut app = App {
        gl: GlGraphics::new(opengl),
        model: cube
    };

    let mut events = Events::new(EventSettings::new());

    while let Some(e) = events.next(&mut window) {
        if let Some(r) = e.render_args() {
            app.render(&r);
        }
    }
}
<|endoftext|>


<commit_before>use std::mem;
use std::ptr;

use crate::bytecode::{self, BytecodeFunction, ConstPoolEntry, ConstPoolOpcode, SourceTypeOpcode};
use crate::compiler::codegen::should_emit_bytecode;
use crate::compiler::fct::{Code, JitDescriptor};
use crate::gc::Address;
use crate::handle::{handle, Handle};
use crate::object::{
    self, byte_array_from_buffer, int_array_alloc_heap, Int32Array, Obj, Ref, UInt8Array,
};
use crate::sym::NestedSymTable;
use crate::threads::current_thread;
use crate::ty::{SourceType, SourceTypeArray};
use crate::vm::{AnalysisData, Fct, VM};

pub fn compile(vm: &VM, fct: &Fct, src: &AnalysisData, _type_params: &SourceTypeArray) -> Code {
    let bytecode_fct = bytecode::generate(vm, fct, src);

    if should_emit_bytecode(vm, fct) {
        bytecode::dump(vm, Some(fct), &bytecode_fct);
    }

    let compile_name = vm.interner.intern("compile");
    let compile_fct_id = NestedSymTable::new(vm, vm.boots_namespace_id)
        .get_fct(compile_name)
        .expect("compile()-method missing");
    let compile_fct = vm.ensure_compiled(compile_fct_id);

    let encoded_compilation_info = handle(allocate_compilation_info(vm, &bytecode_fct));

    let tld_address = current_thread().tld_address();

    let dora_stub_address = vm.dora_stub();
    let compile_fct_ptr: extern "C" fn(Address, Address, Ref<Obj>) -> Ref<UInt8Array> =
        unsafe { mem::transmute(dora_stub_address) };

    let machine_code = handle(compile_fct_ptr(
        tld_address,
        compile_fct,
        encoded_compilation_info.direct(),
    ));
    let mut machine_code_array = vec![0; machine_code.len()];

    unsafe {
        ptr::copy_nonoverlapping(
            machine_code.data() as *mut u8,
            machine_code_array.as_mut_ptr(),
            machine_code.len(),
        );
    }

    Code::from_optimized_buffer(vm, &machine_code_array, JitDescriptor::DoraFct(fct.id))
}

pub fn bytecode(vm: &VM, name: &str) -> Ref<Obj> {
    let fct_name = vm.interner.intern(name);
    let bc_fct_id = NestedSymTable::new(vm, vm.boots_namespace_id)
        .get_fct(fct_name)
        .expect("compile()-method missing");

    let fct = vm.fcts.idx(bc_fct_id);
    let fct = fct.read();
    let analysis = fct.analysis();

    let bytecode_fct = bytecode::generate(vm, &*fct, &*analysis);

    if should_emit_bytecode(vm, &*fct) {
        bytecode::dump(vm, Some(&*fct), &bytecode_fct);
    }

    allocate_compilation_info(vm, &bytecode_fct)
}

fn allocate_compilation_info(vm: &VM, bytecode_fct: &BytecodeFunction) -> Ref<Obj> {
    let bytecode_array = handle(byte_array_from_buffer(vm, bytecode_fct.code()));
    let constpool_array = handle(allocate_constpool_array(vm, &bytecode_fct));
    let registers_array = handle(allocate_registers_array(vm, &bytecode_fct));

    allocate_encoded_compilation_info(
        vm,
        bytecode_array,
        constpool_array,
        registers_array,
        bytecode_fct.arguments() as i32,
    )
}

fn allocate_registers_array(vm: &VM, fct: &BytecodeFunction) -> Ref<Int32Array> {
    let mut array = int_array_alloc_heap(vm, fct.registers().len());

    for (idx, ty) in fct.registers().iter().enumerate() {
        array.set_at(idx, ty.kind() as u32 as i32);
    }

    array
}

fn allocate_constpool_array(vm: &VM, fct: &BytecodeFunction) -> Ref<UInt8Array> {
    use byteorder::{LittleEndian, WriteBytesExt};
    let mut buffer = Vec::new();

    for const_entry in fct.const_pool_entries() {
        match const_entry {
            ConstPoolEntry::String(ref value) => {
                buffer.push(ConstPoolOpcode::Float32 as u8);
                buffer
                    .write_u32::<LittleEndian>(value.len() as u32)
                    .unwrap();

                for byte in value.bytes() {
                    buffer.push(byte);
                }
            }
            &ConstPoolEntry::Float32(value) => {
                buffer.push(ConstPoolOpcode::Float32 as u8);
                buffer.write_u32::<LittleEndian>(value.to_bits()).unwrap();
            }
            &ConstPoolEntry::Float64(value) => {
                buffer.push(ConstPoolOpcode::Float64 as u8);
                buffer.write_u64::<LittleEndian>(value.to_bits()).unwrap();
            }
            &ConstPoolEntry::Int32(value) => {
                buffer.push(ConstPoolOpcode::Int32 as u8);
                buffer.write_u32::<LittleEndian>(value as u32).unwrap();
            }
            &ConstPoolEntry::Int64(value) => {
                buffer.push(ConstPoolOpcode::Int64 as u8);
                buffer.write_u64::<LittleEndian>(value as u64).unwrap();
            }
            &ConstPoolEntry::Char(value) => {
                buffer.push(ConstPoolOpcode::Char as u8);
                buffer.write_u32::<LittleEndian>(value as u32).unwrap();
            }
            &ConstPoolEntry::Fct(fct_id, ref source_type_array) => {
                buffer.push(ConstPoolOpcode::Fct as u8);
                buffer.write_u32::<LittleEndian>(fct_id.to_usize() as u32).unwrap();
                encode_source_type_array(vm, source_type_array, &mut buffer);
            },
            &ConstPoolEntry::Generic(_, _, _) => unimplemented!(),
            &ConstPoolEntry::Class(cls_id, ref source_type_array) => {
                buffer.push(ConstPoolOpcode::Class as u8);
                buffer.write_u32::<LittleEndian>(cls_id.to_usize() as u32).unwrap();
                encode_source_type_array(vm, source_type_array, &mut buffer);
            }
            &ConstPoolEntry::Field(cls_id, ref source_type_array, field_id) => {
                buffer.push(ConstPoolOpcode::Class as u8);
                buffer.write_u32::<LittleEndian>(cls_id.to_usize() as u32).unwrap();
                encode_source_type_array(vm, source_type_array, &mut buffer);
                buffer.write_u32::<LittleEndian>(field_id.to_usize() as u32).unwrap();
            }
            &ConstPoolEntry::FieldFixed(_, _) => unimplemented!(),
            &ConstPoolEntry::Enum(_, _) => unimplemented!(),
            &ConstPoolEntry::EnumVariant(_, _, _) => unimplemented!(),
            &ConstPoolEntry::Struct(_, _) => unimplemented!(),
            &ConstPoolEntry::StructField(_, _, _) => unimplemented!(),
            &ConstPoolEntry::Trait(_, _, _) => unimplemented!(),
        }
    }

    byte_array_from_buffer(vm, &buffer)
}

fn encode_source_type_array(vm: &VM, sta: &SourceTypeArray, buffer: &mut Vec<u8>) {
    use byteorder::{LittleEndian, WriteBytesExt};
    buffer.write_u32::<LittleEndian>(sta.len() as u32).unwrap();

    for ty in sta.iter() {
        encode_source_type(vm, ty, buffer);
    }
}

fn encode_source_type(_vm: &VM, ty: SourceType, buffer: &mut Vec<u8>) {
    let opcode = match ty {
        SourceType::Error | SourceType::Any | SourceType::Ptr | SourceType::This => unreachable!(),
        SourceType::Unit => SourceTypeOpcode::Unit,
        SourceType::Bool => SourceTypeOpcode::Bool,
        SourceType::Char => SourceTypeOpcode::Char,
        SourceType::UInt8 => SourceTypeOpcode::UInt8,
        SourceType::Int32 => SourceTypeOpcode::Int32,
        SourceType::Int64 => SourceTypeOpcode::Int64,
        SourceType::Float32 => SourceTypeOpcode::Float32,
        SourceType::Float64 => SourceTypeOpcode::Float64,
        _ => unreachable!(),
    };

    buffer.push(opcode as u8);
}

fn allocate_encoded_compilation_info(
    vm: &VM,
    bytecode_array: Handle<UInt8Array>,
    constpool_array: Handle<UInt8Array>,
    registers_array: Handle<Int32Array>,
    arguments: i32,
) -> Ref<Obj> {
    let cls_id = vm.cls_def_by_name(vm.boots_namespace_id, "EncodedCompilationInfo");
    let obj = object::alloc(vm, cls_id);

    let fid = vm.field_in_class(cls_id, "code");
    object::write_ref(vm, obj, cls_id, fid, bytecode_array.direct().cast::<Obj>());

    let fid = vm.field_in_class(cls_id, "constpool");
    object::write_ref(vm, obj, cls_id, fid, constpool_array.direct().cast::<Obj>());

    let fid = vm.field_in_class(cls_id, "registers");
    object::write_ref(vm, obj, cls_id, fid, registers_array.direct().cast::<Obj>());

    let fid = vm.field_in_class(cls_id, "arguments");
    object::write_int32(vm, obj, cls_id, fid, arguments);

    obj
}
<commit_msg>reformat<commit_after>use std::mem;
use std::ptr;

use crate::bytecode::{self, BytecodeFunction, ConstPoolEntry, ConstPoolOpcode, SourceTypeOpcode};
use crate::compiler::codegen::should_emit_bytecode;
use crate::compiler::fct::{Code, JitDescriptor};
use crate::gc::Address;
use crate::handle::{handle, Handle};
use crate::object::{
    self, byte_array_from_buffer, int_array_alloc_heap, Int32Array, Obj, Ref, UInt8Array,
};
use crate::sym::NestedSymTable;
use crate::threads::current_thread;
use crate::ty::{SourceType, SourceTypeArray};
use crate::vm::{AnalysisData, Fct, VM};

pub fn compile(vm: &VM, fct: &Fct, src: &AnalysisData, _type_params: &SourceTypeArray) -> Code {
    let bytecode_fct = bytecode::generate(vm, fct, src);

    if should_emit_bytecode(vm, fct) {
        bytecode::dump(vm, Some(fct), &bytecode_fct);
    }

    let compile_name = vm.interner.intern("compile");
    let compile_fct_id = NestedSymTable::new(vm, vm.boots_namespace_id)
        .get_fct(compile_name)
        .expect("compile()-method missing");
    let compile_fct = vm.ensure_compiled(compile_fct_id);

    let encoded_compilation_info = handle(allocate_compilation_info(vm, &bytecode_fct));

    let tld_address = current_thread().tld_address();

    let dora_stub_address = vm.dora_stub();
    let compile_fct_ptr: extern "C" fn(Address, Address, Ref<Obj>) -> Ref<UInt8Array> =
        unsafe { mem::transmute(dora_stub_address) };

    let machine_code = handle(compile_fct_ptr(
        tld_address,
        compile_fct,
        encoded_compilation_info.direct(),
    ));
    let mut machine_code_array = vec![0; machine_code.len()];

    unsafe {
        ptr::copy_nonoverlapping(
            machine_code.data() as *mut u8,
            machine_code_array.as_mut_ptr(),
            machine_code.len(),
        );
    }

    Code::from_optimized_buffer(vm, &machine_code_array, JitDescriptor::DoraFct(fct.id))
}

pub fn bytecode(vm: &VM, name: &str) -> Ref<Obj> {
    let fct_name = vm.interner.intern(name);
    let bc_fct_id = NestedSymTable::new(vm, vm.boots_namespace_id)
        .get_fct(fct_name)
        .expect("compile()-method missing");

    let fct = vm.fcts.idx(bc_fct_id);
    let fct = fct.read();
    let analysis = fct.analysis();

    let bytecode_fct = bytecode::generate(vm, &*fct, &*analysis);

    if should_emit_bytecode(vm, &*fct) {
        bytecode::dump(vm, Some(&*fct), &bytecode_fct);
    }

    allocate_compilation_info(vm, &bytecode_fct)
}

fn allocate_compilation_info(vm: &VM, bytecode_fct: &BytecodeFunction) -> Ref<Obj> {
    let bytecode_array = handle(byte_array_from_buffer(vm, bytecode_fct.code()));
    let constpool_array = handle(allocate_constpool_array(vm, &bytecode_fct));
    let registers_array = handle(allocate_registers_array(vm, &bytecode_fct));

    allocate_encoded_compilation_info(
        vm,
        bytecode_array,
        constpool_array,
        registers_array,
        bytecode_fct.arguments() as i32,
    )
}

fn allocate_registers_array(vm: &VM, fct: &BytecodeFunction) -> Ref<Int32Array> {
    let mut array = int_array_alloc_heap(vm, fct.registers().len());

    for (idx, ty) in fct.registers().iter().enumerate() {
        array.set_at(idx, ty.kind() as u32 as i32);
    }

    array
}

fn allocate_constpool_array(vm: &VM, fct: &BytecodeFunction) -> Ref<UInt8Array> {
    use byteorder::{LittleEndian, WriteBytesExt};
    let mut buffer = Vec::new();

    for const_entry in fct.const_pool_entries() {
        match const_entry {
            ConstPoolEntry::String(ref value) => {
                buffer.push(ConstPoolOpcode::Float32 as u8);
                buffer
                    .write_u32::<LittleEndian>(value.len() as u32)
                    .unwrap();

                for byte in value.bytes() {
                    buffer.push(byte);
                }
            }
            &ConstPoolEntry::Float32(value) => {
                buffer.push(ConstPoolOpcode::Float32 as u8);
                buffer.write_u32::<LittleEndian>(value.to_bits()).unwrap();
            }
            &ConstPoolEntry::Float64(value) => {
                buffer.push(ConstPoolOpcode::Float64 as u8);
                buffer.write_u64::<LittleEndian>(value.to_bits()).unwrap();
            }
            &ConstPoolEntry::Int32(value) => {
                buffer.push(ConstPoolOpcode::Int32 as u8);
                buffer.write_u32::<LittleEndian>(value as u32).unwrap();
            }
            &ConstPoolEntry::Int64(value) => {
                buffer.push(ConstPoolOpcode::Int64 as u8);
                buffer.write_u64::<LittleEndian>(value as u64).unwrap();
            }
            &ConstPoolEntry::Char(value) => {
                buffer.push(ConstPoolOpcode::Char as u8);
                buffer.write_u32::<LittleEndian>(value as u32).unwrap();
            }
            &ConstPoolEntry::Fct(fct_id, ref source_type_array) => {
                buffer.push(ConstPoolOpcode::Fct as u8);
                buffer
                    .write_u32::<LittleEndian>(fct_id.to_usize() as u32)
                    .unwrap();
                encode_source_type_array(vm, source_type_array, &mut buffer);
            }
            &ConstPoolEntry::Generic(_, _, _) => unimplemented!(),
            &ConstPoolEntry::Class(cls_id, ref source_type_array) => {
                buffer.push(ConstPoolOpcode::Class as u8);
                buffer
                    .write_u32::<LittleEndian>(cls_id.to_usize() as u32)
                    .unwrap();
                encode_source_type_array(vm, source_type_array, &mut buffer);
            }
            &ConstPoolEntry::Field(cls_id, ref source_type_array, field_id) => {
                buffer.push(ConstPoolOpcode::Class as u8);
                buffer
                    .write_u32::<LittleEndian>(cls_id.to_usize() as u32)
                    .unwrap();
                encode_source_type_array(vm, source_type_array, &mut buffer);
                buffer
                    .write_u32::<LittleEndian>(field_id.to_usize() as u32)
                    .unwrap();
            }
            &ConstPoolEntry::FieldFixed(_, _) => unimplemented!(),
            &ConstPoolEntry::Enum(_, _) => unimplemented!(),
            &ConstPoolEntry::EnumVariant(_, _, _) => unimplemented!(),
            &ConstPoolEntry::Struct(_, _) => unimplemented!(),
            &ConstPoolEntry::StructField(_, _, _) => unimplemented!(),
            &ConstPoolEntry::Trait(_, _, _) => unimplemented!(),
        }
    }

    byte_array_from_buffer(vm, &buffer)
}

fn encode_source_type_array(vm: &VM, sta: &SourceTypeArray, buffer: &mut Vec<u8>) {
    use byteorder::{LittleEndian, WriteBytesExt};
    buffer.write_u32::<LittleEndian>(sta.len() as u32).unwrap();

    for ty in sta.iter() {
        encode_source_type(vm, ty, buffer);
    }
}

fn encode_source_type(_vm: &VM, ty: SourceType, buffer: &mut Vec<u8>) {
    let opcode = match ty {
        SourceType::Error | SourceType::Any | SourceType::Ptr | SourceType::This => unreachable!(),
        SourceType::Unit => SourceTypeOpcode::Unit,
        SourceType::Bool => SourceTypeOpcode::Bool,
        SourceType::Char => SourceTypeOpcode::Char,
        SourceType::UInt8 => SourceTypeOpcode::UInt8,
        SourceType::Int32 => SourceTypeOpcode::Int32,
        SourceType::Int64 => SourceTypeOpcode::Int64,
        SourceType::Float32 => SourceTypeOpcode::Float32,
        SourceType::Float64 => SourceTypeOpcode::Float64,
        _ => unreachable!(),
    };

    buffer.push(opcode as u8);
}

fn allocate_encoded_compilation_info(
    vm: &VM,
    bytecode_array: Handle<UInt8Array>,
    constpool_array: Handle<UInt8Array>,
    registers_array: Handle<Int32Array>,
    arguments: i32,
) -> Ref<Obj> {
    let cls_id = vm.cls_def_by_name(vm.boots_namespace_id, "EncodedCompilationInfo");
    let obj = object::alloc(vm, cls_id);

    let fid = vm.field_in_class(cls_id, "code");
    object::write_ref(vm, obj, cls_id, fid, bytecode_array.direct().cast::<Obj>());

    let fid = vm.field_in_class(cls_id, "constpool");
    object::write_ref(vm, obj, cls_id, fid, constpool_array.direct().cast::<Obj>());

    let fid = vm.field_in_class(cls_id, "registers");
    object::write_ref(vm, obj, cls_id, fid, registers_array.direct().cast::<Obj>());

    let fid = vm.field_in_class(cls_id, "arguments");
    object::write_int32(vm, obj, cls_id, fid, arguments);

    obj
}
<|endoftext|>


<commit_before>/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/. */

use dom::bindings::codegen::Bindings::StorageBinding;
use dom::bindings::codegen::Bindings::StorageBinding::StorageMethods;
use dom::bindings::error::{Error, ErrorResult};
use dom::bindings::inheritance::Castable;
use dom::bindings::js::Root;
use dom::bindings::refcounted::Trusted;
use dom::bindings::reflector::{DomObject, Reflector, reflect_dom_object};
use dom::bindings::str::DOMString;
use dom::event::{Event, EventBubbles, EventCancelable};
use dom::storageevent::StorageEvent;
use dom::window::Window;
use dom_struct::dom_struct;
use ipc_channel::ipc::{self, IpcSender};
use net_traits::IpcSend;
use net_traits::storage_thread::{StorageThreadMsg, StorageType};
use script_traits::ScriptMsg;
use servo_url::ServoUrl;
use task::Task;
use task_source::TaskSource;

#[dom_struct]
pub struct Storage {
    reflector_: Reflector,
    storage_type: StorageType
}

impl Storage {
    fn new_inherited(storage_type: StorageType) -> Storage {
        Storage {
            reflector_: Reflector::new(),
            storage_type: storage_type
        }
    }

    pub fn new(global: &Window, storage_type: StorageType) -> Root<Storage> {
        reflect_dom_object(box Storage::new_inherited(storage_type), global, StorageBinding::Wrap)
    }

    fn get_url(&self) -> ServoUrl {
        self.global().get_url()
    }

    fn get_storage_thread(&self) -> IpcSender<StorageThreadMsg> {
        self.global().resource_threads().sender()
    }

}

impl StorageMethods for Storage {
    // https://html.spec.whatwg.org/multipage/#dom-storage-length
    fn Length(&self) -> u32 {
        let (sender, receiver) = ipc::channel().unwrap();

        self.get_storage_thread().send(StorageThreadMsg::Length(sender, self.get_url(), self.storage_type)).unwrap();
        receiver.recv().unwrap() as u32
    }

    // https://html.spec.whatwg.org/multipage/#dom-storage-key
    fn Key(&self, index: u32) -> Option<DOMString> {
        let (sender, receiver) = ipc::channel().unwrap();

        self.get_storage_thread()
            .send(StorageThreadMsg::Key(sender, self.get_url(), self.storage_type, index))
            .unwrap();
        receiver.recv().unwrap().map(DOMString::from)
    }

    // https://html.spec.whatwg.org/multipage/#dom-storage-getitem
    fn GetItem(&self, name: DOMString) -> Option<DOMString> {
        let (sender, receiver) = ipc::channel().unwrap();
        let name = String::from(name);

        let msg = StorageThreadMsg::GetItem(sender, self.get_url(), self.storage_type, name);
        self.get_storage_thread().send(msg).unwrap();
        receiver.recv().unwrap().map(DOMString::from)
    }

    // https://html.spec.whatwg.org/multipage/#dom-storage-setitem
    fn SetItem(&self, name: DOMString, value: DOMString) -> ErrorResult {
        let (sender, receiver) = ipc::channel().unwrap();
        let name = String::from(name);
        let value = String::from(value);

        let msg = StorageThreadMsg::SetItem(sender, self.get_url(), self.storage_type, name.clone(), value.clone());
        self.get_storage_thread().send(msg).unwrap();
        match receiver.recv().unwrap() {
            Err(_) => Err(Error::QuotaExceeded),
            Ok((changed, old_value)) => {
              if changed {
                  self.broadcast_change_notification(Some(name), old_value, Some(value));
              }
              Ok(())
            }
        }
    }

    // https://html.spec.whatwg.org/multipage/#dom-storage-removeitem
    fn RemoveItem(&self, name: DOMString) {
        let (sender, receiver) = ipc::channel().unwrap();
        let name = String::from(name);

        let msg = StorageThreadMsg::RemoveItem(sender, self.get_url(), self.storage_type, name.clone());
        self.get_storage_thread().send(msg).unwrap();
        if let Some(old_value) = receiver.recv().unwrap() {
            self.broadcast_change_notification(Some(name), Some(old_value), None);
        }
    }

    // https://html.spec.whatwg.org/multipage/#dom-storage-clear
    fn Clear(&self) {
        let (sender, receiver) = ipc::channel().unwrap();

        self.get_storage_thread().send(StorageThreadMsg::Clear(sender, self.get_url(), self.storage_type)).unwrap();
        if receiver.recv().unwrap() {
            self.broadcast_change_notification(None, None, None);
        }
    }

    // https://html.spec.whatwg.org/multipage/#the-storage-interface:supported-property-names
    fn SupportedPropertyNames(&self) -> Vec<DOMString> {
        let (sender, receiver) = ipc::channel().unwrap();

        self.get_storage_thread().send(StorageThreadMsg::Keys(sender, self.get_url(), self.storage_type)).unwrap();
        receiver.recv()
                .unwrap()
                .into_iter()
                .map(DOMString::from)
                .collect()
    }

    // check-tidy: no specs after this line
    fn NamedGetter(&self, name: DOMString) -> Option<DOMString> {
        self.GetItem(name)
    }

    fn NamedSetter(&self, name: DOMString, value: DOMString) -> ErrorResult {
        self.SetItem(name, value)
    }

    fn NamedDeleter(&self, name: DOMString) {
        self.RemoveItem(name);
    }
}

impl Storage {
    /// https://html.spec.whatwg.org/multipage/#send-a-storage-notification
    fn broadcast_change_notification(&self, key: Option<String>, old_value: Option<String>,
                                     new_value: Option<String>) {
        let storage = self.storage_type;
        let url = self.get_url();
        let msg = ScriptMsg::BroadcastStorageEvent(storage, url, key, old_value, new_value);
        self.global().script_to_constellation_chan().send(msg).unwrap();
    }

    /// https://html.spec.whatwg.org/multipage/#send-a-storage-notification
    pub fn queue_storage_event(
        &self,
        url: ServoUrl,
        key: Option<String>,
        old_value: Option<String>,
        new_value: Option<String>,
    ) {
        let global = self.global();
        global.as_window().dom_manipulation_task_source().queue(
            box StorageEventTask {
                element: Trusted::new(self),
                url,
                key,
                old_value,
                new_value,
            },
            global.upcast(),
        ).unwrap();
        struct StorageEventTask {
            element: Trusted<Storage>,
            url: ServoUrl,
            key: Option<String>,
            old_value: Option<String>,
            new_value: Option<String>
        }
        impl Task for StorageEventTask {
            fn run(self: Box<Self>) {
                let this = *self;
                let storage = this.element.root();
                let global = storage.global();
                let storage_event = StorageEvent::new(
                    global.as_window(),
                    atom!("storage"),
                    EventBubbles::DoesNotBubble,
                    EventCancelable::NotCancelable,
                    this.key.map(DOMString::from),
                    this.old_value.map(DOMString::from),
                    this.new_value.map(DOMString::from),
                    DOMString::from(this.url.into_string()),
                    Some(&storage)
                );
                storage_event.upcast::<Event>().fire(global.upcast());
            }
        }
    }
}
<commit_msg>Use task! to send a storage notification<commit_after>/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/. */

use dom::bindings::codegen::Bindings::StorageBinding;
use dom::bindings::codegen::Bindings::StorageBinding::StorageMethods;
use dom::bindings::error::{Error, ErrorResult};
use dom::bindings::inheritance::Castable;
use dom::bindings::js::Root;
use dom::bindings::refcounted::Trusted;
use dom::bindings::reflector::{DomObject, Reflector, reflect_dom_object};
use dom::bindings::str::DOMString;
use dom::event::{Event, EventBubbles, EventCancelable};
use dom::storageevent::StorageEvent;
use dom::window::Window;
use dom_struct::dom_struct;
use ipc_channel::ipc::{self, IpcSender};
use net_traits::IpcSend;
use net_traits::storage_thread::{StorageThreadMsg, StorageType};
use script_traits::ScriptMsg;
use servo_url::ServoUrl;
use task_source::TaskSource;

#[dom_struct]
pub struct Storage {
    reflector_: Reflector,
    storage_type: StorageType
}

impl Storage {
    fn new_inherited(storage_type: StorageType) -> Storage {
        Storage {
            reflector_: Reflector::new(),
            storage_type: storage_type
        }
    }

    pub fn new(global: &Window, storage_type: StorageType) -> Root<Storage> {
        reflect_dom_object(box Storage::new_inherited(storage_type), global, StorageBinding::Wrap)
    }

    fn get_url(&self) -> ServoUrl {
        self.global().get_url()
    }

    fn get_storage_thread(&self) -> IpcSender<StorageThreadMsg> {
        self.global().resource_threads().sender()
    }

}

impl StorageMethods for Storage {
    // https://html.spec.whatwg.org/multipage/#dom-storage-length
    fn Length(&self) -> u32 {
        let (sender, receiver) = ipc::channel().unwrap();

        self.get_storage_thread().send(StorageThreadMsg::Length(sender, self.get_url(), self.storage_type)).unwrap();
        receiver.recv().unwrap() as u32
    }

    // https://html.spec.whatwg.org/multipage/#dom-storage-key
    fn Key(&self, index: u32) -> Option<DOMString> {
        let (sender, receiver) = ipc::channel().unwrap();

        self.get_storage_thread()
            .send(StorageThreadMsg::Key(sender, self.get_url(), self.storage_type, index))
            .unwrap();
        receiver.recv().unwrap().map(DOMString::from)
    }

    // https://html.spec.whatwg.org/multipage/#dom-storage-getitem
    fn GetItem(&self, name: DOMString) -> Option<DOMString> {
        let (sender, receiver) = ipc::channel().unwrap();
        let name = String::from(name);

        let msg = StorageThreadMsg::GetItem(sender, self.get_url(), self.storage_type, name);
        self.get_storage_thread().send(msg).unwrap();
        receiver.recv().unwrap().map(DOMString::from)
    }

    // https://html.spec.whatwg.org/multipage/#dom-storage-setitem
    fn SetItem(&self, name: DOMString, value: DOMString) -> ErrorResult {
        let (sender, receiver) = ipc::channel().unwrap();
        let name = String::from(name);
        let value = String::from(value);

        let msg = StorageThreadMsg::SetItem(sender, self.get_url(), self.storage_type, name.clone(), value.clone());
        self.get_storage_thread().send(msg).unwrap();
        match receiver.recv().unwrap() {
            Err(_) => Err(Error::QuotaExceeded),
            Ok((changed, old_value)) => {
              if changed {
                  self.broadcast_change_notification(Some(name), old_value, Some(value));
              }
              Ok(())
            }
        }
    }

    // https://html.spec.whatwg.org/multipage/#dom-storage-removeitem
    fn RemoveItem(&self, name: DOMString) {
        let (sender, receiver) = ipc::channel().unwrap();
        let name = String::from(name);

        let msg = StorageThreadMsg::RemoveItem(sender, self.get_url(), self.storage_type, name.clone());
        self.get_storage_thread().send(msg).unwrap();
        if let Some(old_value) = receiver.recv().unwrap() {
            self.broadcast_change_notification(Some(name), Some(old_value), None);
        }
    }

    // https://html.spec.whatwg.org/multipage/#dom-storage-clear
    fn Clear(&self) {
        let (sender, receiver) = ipc::channel().unwrap();

        self.get_storage_thread().send(StorageThreadMsg::Clear(sender, self.get_url(), self.storage_type)).unwrap();
        if receiver.recv().unwrap() {
            self.broadcast_change_notification(None, None, None);
        }
    }

    // https://html.spec.whatwg.org/multipage/#the-storage-interface:supported-property-names
    fn SupportedPropertyNames(&self) -> Vec<DOMString> {
        let (sender, receiver) = ipc::channel().unwrap();

        self.get_storage_thread().send(StorageThreadMsg::Keys(sender, self.get_url(), self.storage_type)).unwrap();
        receiver.recv()
                .unwrap()
                .into_iter()
                .map(DOMString::from)
                .collect()
    }

    // check-tidy: no specs after this line
    fn NamedGetter(&self, name: DOMString) -> Option<DOMString> {
        self.GetItem(name)
    }

    fn NamedSetter(&self, name: DOMString, value: DOMString) -> ErrorResult {
        self.SetItem(name, value)
    }

    fn NamedDeleter(&self, name: DOMString) {
        self.RemoveItem(name);
    }
}

impl Storage {
    /// https://html.spec.whatwg.org/multipage/#send-a-storage-notification
    fn broadcast_change_notification(&self, key: Option<String>, old_value: Option<String>,
                                     new_value: Option<String>) {
        let storage = self.storage_type;
        let url = self.get_url();
        let msg = ScriptMsg::BroadcastStorageEvent(storage, url, key, old_value, new_value);
        self.global().script_to_constellation_chan().send(msg).unwrap();
    }

    /// https://html.spec.whatwg.org/multipage/#send-a-storage-notification
    pub fn queue_storage_event(
        &self,
        url: ServoUrl,
        key: Option<String>,
        old_value: Option<String>,
        new_value: Option<String>,
    ) {
        let global = self.global();
        let this = Trusted::new(self);
        global.as_window().dom_manipulation_task_source().queue(
            box task!(send_storage_notification: move || {
                let this = this.root();
                let global = this.global();
                let event = StorageEvent::new(
                    global.as_window(),
                    atom!("storage"),
                    EventBubbles::DoesNotBubble,
                    EventCancelable::NotCancelable,
                    key.map(DOMString::from),
                    old_value.map(DOMString::from),
                    new_value.map(DOMString::from),
                    DOMString::from(url.into_string()),
                    Some(&this),
                );
                event.upcast::<Event>().fire(global.upcast());
            }),
            global.upcast(),
        ).unwrap();
    }
}
<|endoftext|>