<commit_before>//  * This file is part of the uutils coreutils package.
//  *
//  * (c) Alan Andrade <alan.andradec@gmail.com>
//  *
//  * For the full copyright and license information, please view the LICENSE
//  * file that was distributed with this source code.

// spell-checker:ignore (ToDO) MAKEWORD addrs hashset

// Clippy bug: https://github.com/rust-lang/rust-clippy/issues/7422
#![allow(clippy::nonstandard_macro_braces)]

#[macro_use]
extern crate uucore;

use clap::{crate_version, App, Arg, ArgMatches};
use std::collections::hash_set::HashSet;
use std::net::ToSocketAddrs;
use std::str;
use uucore::error::{UResult, USimpleError};

#[cfg(windows)]
use winapi::shared::minwindef::MAKEWORD;
#[cfg(windows)]
use winapi::um::winsock2::{WSACleanup, WSAStartup};

static ABOUT: &str = "Display or set the system's host name.";

static OPT_DOMAIN: &str = "domain";
static OPT_IP_ADDRESS: &str = "ip-address";
static OPT_FQDN: &str = "fqdn";
static OPT_SHORT: &str = "short";
static OPT_HOST: &str = "host";

#[uucore_procs::gen_uumain]
pub fn uumain(args: impl uucore::Args) -> UResult<()> {
    #![allow(clippy::let_and_return)]
    #[cfg(windows)]
    unsafe {
        #[allow(deprecated)]
        let mut data = std::mem::uninitialized();
        if WSAStartup(MAKEWORD(2, 2), &mut data as *mut _) != 0 {
            return Err(USimpleError::new(1, format!("Failed to start Winsock 2.2")));
        }
    }
    let result = execute(args);
    #[cfg(windows)]
    unsafe {
        WSACleanup();
    }
    result
}

fn get_usage() -> String {
    format!("{0} [OPTION]... [HOSTNAME]", executable!())
}
fn execute(args: impl uucore::Args) -> UResult<()> {
    let usage = get_usage();
    let matches = uu_app().usage(&usage[..]).get_matches_from(args);

    match matches.value_of(OPT_HOST) {
        None => display_hostname(&matches),
        Some(host) => {
            if let Err(err) = hostname::set(host) {
                return Err(USimpleError::new(1, format!("{}", err)));
            } else {
                Ok(())
            }
        }
    }
}

pub fn uu_app() -> App<'static, 'static> {
    App::new(executable!())
        .version(crate_version!())
        .about(ABOUT)
        .arg(
            Arg::with_name(OPT_DOMAIN)
                .short("d")
                .long("domain")
                .help("Display the name of the DNS domain if possible"),
        )
        .arg(
            Arg::with_name(OPT_IP_ADDRESS)
                .short("i")
                .long("ip-address")
                .help("Display the network address(es) of the host"),
        )
        // TODO: support --long
        .arg(
            Arg::with_name(OPT_FQDN)
                .short("f")
                .long("fqdn")
                .help("Display the FQDN (Fully Qualified Domain Name) (default)"),
        )
        .arg(Arg::with_name(OPT_SHORT).short("s").long("short").help(
            "Display the short hostname (the portion before the first dot) if \
             possible",
        ))
        .arg(Arg::with_name(OPT_HOST))
}

fn display_hostname(matches: &ArgMatches) -> UResult<()> {
    let hostname = hostname::get().unwrap().into_string().unwrap();

    if matches.is_present(OPT_IP_ADDRESS) {
        // XXX: to_socket_addrs needs hostname:port so append a dummy port and remove it later.
        // This was originally supposed to use std::net::lookup_host, but that seems to be
        // deprecated.  Perhaps we should use the dns-lookup crate?
        let hostname = hostname + ":1";
        match hostname.to_socket_addrs() {
            Ok(addresses) => {
                let mut hashset = HashSet::new();
                let mut output = String::new();
                for addr in addresses {
                    // XXX: not sure why this is necessary...
                    if !hashset.contains(&addr) {
                        let mut ip = format!("{}", addr);
                        if ip.ends_with(":1") {
                            let len = ip.len();
                            ip.truncate(len - 2);
                        }
                        output.push_str(&ip);
                        output.push(' ');
                        hashset.insert(addr);
                    }
                }
                let len = output.len();
                if len > 0 {
                    println!("{}", &output[0..len - 1]);
                }

                Ok(())
            }
            Err(f) => {
                return Err(USimpleError::new(1, format!("{}", f)));
            }
        }
    } else {
        if matches.is_present(OPT_SHORT) || matches.is_present(OPT_DOMAIN) {
            let mut it = hostname.char_indices().filter(|&ci| ci.1 == '.');
            if let Some(ci) = it.next() {
                if matches.is_present(OPT_SHORT) {
                    println!("{}", &hostname[0..ci.0]);
                } else {
                    println!("{}", &hostname[ci.0 + 1..]);
                }
                return Ok(());
            }
        }

        println!("{}", hostname);

        Ok(())
    }
}
<commit_msg>Fix clippy warning<commit_after>//  * This file is part of the uutils coreutils package.
//  *
//  * (c) Alan Andrade <alan.andradec@gmail.com>
//  *
//  * For the full copyright and license information, please view the LICENSE
//  * file that was distributed with this source code.

// spell-checker:ignore (ToDO) MAKEWORD addrs hashset

// Clippy bug: https://github.com/rust-lang/rust-clippy/issues/7422
#![allow(clippy::nonstandard_macro_braces)]

#[macro_use]
extern crate uucore;

use clap::{crate_version, App, Arg, ArgMatches};
use std::collections::hash_set::HashSet;
use std::net::ToSocketAddrs;
use std::str;
#[cfg(windows)]
use uucore::error::UUsageError;
use uucore::error::{UResult, USimpleError};

#[cfg(windows)]
use winapi::shared::minwindef::MAKEWORD;
#[cfg(windows)]
use winapi::um::winsock2::{WSACleanup, WSAStartup};

static ABOUT: &str = "Display or set the system's host name.";

static OPT_DOMAIN: &str = "domain";
static OPT_IP_ADDRESS: &str = "ip-address";
static OPT_FQDN: &str = "fqdn";
static OPT_SHORT: &str = "short";
static OPT_HOST: &str = "host";

#[uucore_procs::gen_uumain]
pub fn uumain(args: impl uucore::Args) -> UResult<()> {
    #![allow(clippy::let_and_return)]
    #[cfg(windows)]
    unsafe {
        #[allow(deprecated)]
        let mut data = std::mem::uninitialized();
        if WSAStartup(MAKEWORD(2, 2), &mut data as *mut _) != 0 {
            return Err(UUsageError::new(
                1,
                "Failed to start Winsock 2.2".to_string(),
            ));
        }
    }
    let result = execute(args);
    #[cfg(windows)]
    unsafe {
        WSACleanup();
    }
    result
}

fn get_usage() -> String {
    format!("{0} [OPTION]... [HOSTNAME]", executable!())
}
fn execute(args: impl uucore::Args) -> UResult<()> {
    let usage = get_usage();
    let matches = uu_app().usage(&usage[..]).get_matches_from(args);

    match matches.value_of(OPT_HOST) {
        None => display_hostname(&matches),
        Some(host) => {
            if let Err(err) = hostname::set(host) {
                return Err(USimpleError::new(1, format!("{}", err)));
            } else {
                Ok(())
            }
        }
    }
}

pub fn uu_app() -> App<'static, 'static> {
    App::new(executable!())
        .version(crate_version!())
        .about(ABOUT)
        .arg(
            Arg::with_name(OPT_DOMAIN)
                .short("d")
                .long("domain")
                .help("Display the name of the DNS domain if possible"),
        )
        .arg(
            Arg::with_name(OPT_IP_ADDRESS)
                .short("i")
                .long("ip-address")
                .help("Display the network address(es) of the host"),
        )
        // TODO: support --long
        .arg(
            Arg::with_name(OPT_FQDN)
                .short("f")
                .long("fqdn")
                .help("Display the FQDN (Fully Qualified Domain Name) (default)"),
        )
        .arg(Arg::with_name(OPT_SHORT).short("s").long("short").help(
            "Display the short hostname (the portion before the first dot) if \
             possible",
        ))
        .arg(Arg::with_name(OPT_HOST))
}

fn display_hostname(matches: &ArgMatches) -> UResult<()> {
    let hostname = hostname::get().unwrap().into_string().unwrap();

    if matches.is_present(OPT_IP_ADDRESS) {
        // XXX: to_socket_addrs needs hostname:port so append a dummy port and remove it later.
        // This was originally supposed to use std::net::lookup_host, but that seems to be
        // deprecated.  Perhaps we should use the dns-lookup crate?
        let hostname = hostname + ":1";
        match hostname.to_socket_addrs() {
            Ok(addresses) => {
                let mut hashset = HashSet::new();
                let mut output = String::new();
                for addr in addresses {
                    // XXX: not sure why this is necessary...
                    if !hashset.contains(&addr) {
                        let mut ip = format!("{}", addr);
                        if ip.ends_with(":1") {
                            let len = ip.len();
                            ip.truncate(len - 2);
                        }
                        output.push_str(&ip);
                        output.push(' ');
                        hashset.insert(addr);
                    }
                }
                let len = output.len();
                if len > 0 {
                    println!("{}", &output[0..len - 1]);
                }

                Ok(())
            }
            Err(f) => {
                return Err(USimpleError::new(1, format!("{}", f)));
            }
        }
    } else {
        if matches.is_present(OPT_SHORT) || matches.is_present(OPT_DOMAIN) {
            let mut it = hostname.char_indices().filter(|&ci| ci.1 == '.');
            if let Some(ci) = it.next() {
                if matches.is_present(OPT_SHORT) {
                    println!("{}", &hostname[0..ci.0]);
                } else {
                    println!("{}", &hostname[ci.0 + 1..]);
                }
                return Ok(());
            }
        }

        println!("{}", hostname);

        Ok(())
    }
}
<|endoftext|>


<commit_before>use std::fmt;
use std::error::Error;
use std::any::Any;
use time::Duration;

use time_source::*;
use scheduler::*;

pub enum AbortableWaitError<Token> {
    Empty,
    Overrun(Vec<Token>),
    Aborted
}

impl<Token> PartialEq for AbortableWaitError<Token> where Token: PartialEq<Token> {
    fn eq(&self, other: &Self) -> bool {
        match (self, other) {
            (&AbortableWaitError::Empty, &AbortableWaitError::Empty) => true,
            (&AbortableWaitError::Overrun(ref tokens), &AbortableWaitError::Overrun(ref other_tokens)) => tokens == other_tokens,
            (&AbortableWaitError::Aborted, &AbortableWaitError::Aborted) => true,
            _ => false
        }
    }
}

impl<Token> fmt::Debug for AbortableWaitError<Token> where Token: fmt::Debug {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match self {
            &AbortableWaitError::Empty => write!(f, "AbortableWaitError::Empty"),
            &AbortableWaitError::Overrun(ref tokens) => write!(f, "AbortableWaitError::Overrun({:?})", tokens),
            &AbortableWaitError::Aborted => write!(f, "AbortableWaitError::Aborted")
        }
    }
}

impl<Token> fmt::Display for AbortableWaitError<Token> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match self {
            &AbortableWaitError::Empty => write!(f, "scheduler is empty"),
            &AbortableWaitError::Overrun(ref tokens) => write!(f, "scheduler overrun {} tokens", tokens.len()),
            &AbortableWaitError::Aborted => write!(f, "wait operation was aborted from another thread")
        }
    }
}

impl<Token> Error for AbortableWaitError<Token> where Token: fmt::Debug + Any {
    fn description(&self) -> &str {
        "problem while waiting for next schedule"
    }
}

pub enum AbortableWaitTimeoutError<Token> {
    Empty,
    Timeout,
    Overrun(Vec<Token>),
    Aborted
}

impl<Token> PartialEq for AbortableWaitTimeoutError<Token> where Token: PartialEq<Token> {
    fn eq(&self, other: &Self) -> bool {
        match (self, other) {
            (&AbortableWaitTimeoutError::Empty, &AbortableWaitTimeoutError::Empty) => true,
            (&AbortableWaitTimeoutError::Timeout, &AbortableWaitTimeoutError::Timeout) => true,
            (&AbortableWaitTimeoutError::Overrun(ref tokens), &AbortableWaitTimeoutError::Overrun(ref other_tokens)) => tokens == other_tokens,
            (&AbortableWaitTimeoutError::Aborted, &AbortableWaitTimeoutError::Aborted) => true,
            _ => false
        }
    }
}

impl<Token> fmt::Debug for AbortableWaitTimeoutError<Token> where Token: fmt::Debug {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match self {
            &AbortableWaitTimeoutError::Empty => write!(f, "AbortableWaitTimeoutError::Empty"),
            &AbortableWaitTimeoutError::Timeout => write!(f, "AbortableWaitTimeoutError::Timeout"),
            &AbortableWaitTimeoutError::Overrun(ref tokens) => write!(f, "AbortableWaitTimeoutError::Overrun({:?})", tokens),
            &AbortableWaitTimeoutError::Aborted => write!(f, "AbortableWaitTimeoutError::Aborted")
        }
    }
}

impl<Token> fmt::Display for AbortableWaitTimeoutError<Token> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match self {
            &AbortableWaitTimeoutError::Empty => write!(f, "scheduler is empty"),
            &AbortableWaitTimeoutError::Timeout => write!(f, "timedout while waiting for tokens"),
            &AbortableWaitTimeoutError::Overrun(ref tokens) => write!(f, "scheduler overrun {} tokens", tokens.len()),
            &AbortableWaitTimeoutError::Aborted => write!(f, "wait operation was aborted from another thread")
        }
    }
}

impl<Token> Error for AbortableWaitTimeoutError<Token> where Token: fmt::Debug + Any {
    fn description(&self) -> &str {
        "problem while waiting for next schedule with timeout"
    }
}

impl<Token, TS> Scheduler<Token, TS> where TS: TimeSource + Wait + AbortableWait, Token: Clone {
    pub fn abort_handle(&self) -> <TS as AbortableWait>::AbortHandle {
        self.time_source.abort_handle()
    }

    pub fn abortable_wait(&mut self) -> Result<Vec<Token>, AbortableWaitError<Token>> {
        match self.next() {
            Some(schedule) => match schedule {
                Schedule::NextIn(duration) => {
                    if let Err(WaitAbortedError) = self.time_source.abortable_wait(duration) {
                        return Err(AbortableWaitError::Aborted);
                    };
                    self.abortable_wait()
                },
                Schedule::Overrun(overrun_tokens) => {
                    Err(AbortableWaitError::Overrun(overrun_tokens))
                },
                Schedule::Current(tokens) => {
                    Ok(tokens)
                }
            },
            None => Err(AbortableWaitError::Empty)
        }
    }

    pub fn abortable_wait_timeout(&mut self, timeout: Duration) -> Result<Vec<Token>, AbortableWaitTimeoutError<Token>> {
        match self.next() {
            Some(schedule) => match schedule {
                Schedule::NextIn(duration) => {
                    if duration > timeout {
                        if let Err(WaitAbortedError) = self.time_source.abortable_wait(timeout) {
                            return Err(AbortableWaitTimeoutError::Aborted);
                        };
                        return Err(AbortableWaitTimeoutError::Timeout);
                    }
                    self.time_source.wait(duration);
                    self.abortable_wait_timeout(Duration::zero())
                },
                Schedule::Overrun(overrun_tokens) => {
                    Err(AbortableWaitTimeoutError::Overrun(overrun_tokens))
                },
                Schedule::Current(tokens) => {
                    Ok(tokens)
                }
            },
            None => Err(AbortableWaitTimeoutError::Empty)
        }
    }
}

#[cfg(test)]
mod test {
    use super::*;
    use time_source::*;
    use steady_time_source::*;
    use test_helpers::*;
    use scheduler::*;
    use time::Duration;
    use std::thread::spawn;

    #[test]
    fn abortable_wait_with_aborted_steady_time_source() {
        let mut scheduler = Scheduler::with_time_source(Duration::milliseconds(100), SteadyTimeSource::new());

        scheduler.after(Duration::milliseconds(100), 0);
        scheduler.after(Duration::seconds(20), 1);
        scheduler.after(Duration::seconds(40), 2);

        let abort_handle = scheduler.abort_handle();

        assert_eq!(scheduler.abortable_wait(), Ok(vec![0]));

        spawn(move || {
            abort_handle.abort();
        });

        assert_eq!(scheduler.abortable_wait(), Err(AbortableWaitError::Aborted));
    }

    #[test]
    fn abortable_wait_with_empty() {
        let mut scheduler = Scheduler::with_time_source(Duration::seconds(1), MockTimeSourceWait::new());
        assert_eq!(scheduler.abortable_wait(), Err(AbortableWaitError::Empty));

        scheduler.after(Duration::seconds(0), 0);
        assert_eq!(scheduler.abortable_wait(), Ok(vec![0]));
        assert_eq!(scheduler.abortable_wait(), Err(AbortableWaitError::Empty));
    }

    #[test]
    fn abortable_wait_with_with_overrun() {
        let mut scheduler = Scheduler::with_time_source(Duration::seconds(1), MockTimeSourceWait::new());

        scheduler.after(Duration::seconds(0), 0);
        scheduler.after(Duration::seconds(1), 1);
        scheduler.after(Duration::seconds(2), 2);

        scheduler.fast_forward(Duration::seconds(2));
        assert_eq!(scheduler.abortable_wait(), Err(AbortableWaitError::Overrun(vec![0, 1])));
        assert_eq!(scheduler.abortable_wait(), Ok(vec![2]));
    }

    #[test]
    fn abortable_wait_timeout_with_aborted_steady_time_source() {
        let mut scheduler = Scheduler::with_time_source(Duration::milliseconds(100), SteadyTimeSource::new());

        scheduler.after(Duration::milliseconds(100), 0);
        scheduler.after(Duration::seconds(20), 1);
        scheduler.after(Duration::seconds(40), 2);

        let abort_handle = scheduler.abort_handle();

        assert_eq!(scheduler.abortable_wait_timeout(Duration::seconds(2)), Ok(vec![0]));

        spawn(move || {
            abort_handle.abort();
        });

        assert_eq!(scheduler.abortable_wait_timeout(Duration::seconds(2)), Err(AbortableWaitTimeoutError::Aborted));
    }

    #[test]
    fn abortable_wait_timeout_with_empty() {
        let mut scheduler = Scheduler::with_time_source(Duration::seconds(1), MockTimeSourceWait::new());
        assert_eq!(scheduler.abortable_wait_timeout(Duration::seconds(4)), Err(AbortableWaitTimeoutError::Empty));

        scheduler.after(Duration::seconds(0), 0);
        assert_eq!(scheduler.abortable_wait_timeout(Duration::seconds(4)), Ok(vec![0]));
        assert_eq!(scheduler.abortable_wait_timeout(Duration::seconds(4)), Err(AbortableWaitTimeoutError::Empty));
    }

    #[test]
    fn abortable_wait_timeout_with_timeout() {
        let mut scheduler = Scheduler::with_time_source(Duration::seconds(1), MockTimeSourceWait::new());

        scheduler.after(Duration::seconds(0), 0);
        scheduler.after(Duration::seconds(1), 1);
        scheduler.after(Duration::seconds(2), 2);

        assert_eq!(scheduler.abortable_wait_timeout(Duration::seconds(2)), Ok(vec![0]));
        assert_eq!(scheduler.abortable_wait_timeout(Duration::milliseconds(500)), Err(AbortableWaitTimeoutError::Timeout));
    }

    #[test]
    fn abortable_wait_timeout_with_overrun() {
        let mut scheduler = Scheduler::with_time_source(Duration::seconds(1), MockTimeSourceWait::new());

        scheduler.after(Duration::seconds(0), 0);
        scheduler.after(Duration::seconds(1), 1);
        scheduler.after(Duration::seconds(2), 2);

        scheduler.fast_forward(Duration::seconds(2));
        assert_eq!(scheduler.abortable_wait_timeout(Duration::seconds(4)), Err(AbortableWaitTimeoutError::Overrun(vec![0, 1])));
        assert_eq!(scheduler.abortable_wait_timeout(Duration::seconds(4)), Ok(vec![2]));
    }
}
<commit_msg>cleanup<commit_after>use std::fmt;
use std::error::Error;
use std::any::Any;
use time::Duration;

use time_source::*;
use scheduler::*;

pub enum AbortableWaitError<Token> {
    Empty,
    Overrun(Vec<Token>),
    Aborted
}

impl<Token> PartialEq for AbortableWaitError<Token> where Token: PartialEq<Token> {
    fn eq(&self, other: &Self) -> bool {
        match (self, other) {
            (&AbortableWaitError::Empty, &AbortableWaitError::Empty) => true,
            (&AbortableWaitError::Overrun(ref tokens), &AbortableWaitError::Overrun(ref other_tokens)) => tokens == other_tokens,
            (&AbortableWaitError::Aborted, &AbortableWaitError::Aborted) => true,
            _ => false
        }
    }
}

impl<Token> fmt::Debug for AbortableWaitError<Token> where Token: fmt::Debug {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match self {
            &AbortableWaitError::Empty => write!(f, "AbortableWaitError::Empty"),
            &AbortableWaitError::Overrun(ref tokens) => write!(f, "AbortableWaitError::Overrun({:?})", tokens),
            &AbortableWaitError::Aborted => write!(f, "AbortableWaitError::Aborted")
        }
    }
}

impl<Token> fmt::Display for AbortableWaitError<Token> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match self {
            &AbortableWaitError::Empty => write!(f, "scheduler is empty"),
            &AbortableWaitError::Overrun(ref tokens) => write!(f, "scheduler overrun {} tokens", tokens.len()),
            &AbortableWaitError::Aborted => write!(f, "wait operation was aborted from another thread")
        }
    }
}

impl<Token> Error for AbortableWaitError<Token> where Token: fmt::Debug + Any {
    fn description(&self) -> &str {
        "problem while waiting for next schedule"
    }
}

pub enum AbortableWaitTimeoutError<Token> {
    Empty,
    Timeout,
    Overrun(Vec<Token>),
    Aborted
}

impl<Token> PartialEq for AbortableWaitTimeoutError<Token> where Token: PartialEq<Token> {
    fn eq(&self, other: &Self) -> bool {
        match (self, other) {
            (&AbortableWaitTimeoutError::Empty, &AbortableWaitTimeoutError::Empty) => true,
            (&AbortableWaitTimeoutError::Timeout, &AbortableWaitTimeoutError::Timeout) => true,
            (&AbortableWaitTimeoutError::Overrun(ref tokens), &AbortableWaitTimeoutError::Overrun(ref other_tokens)) => tokens == other_tokens,
            (&AbortableWaitTimeoutError::Aborted, &AbortableWaitTimeoutError::Aborted) => true,
            _ => false
        }
    }
}

impl<Token> fmt::Debug for AbortableWaitTimeoutError<Token> where Token: fmt::Debug {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match self {
            &AbortableWaitTimeoutError::Empty => write!(f, "AbortableWaitTimeoutError::Empty"),
            &AbortableWaitTimeoutError::Timeout => write!(f, "AbortableWaitTimeoutError::Timeout"),
            &AbortableWaitTimeoutError::Overrun(ref tokens) => write!(f, "AbortableWaitTimeoutError::Overrun({:?})", tokens),
            &AbortableWaitTimeoutError::Aborted => write!(f, "AbortableWaitTimeoutError::Aborted")
        }
    }
}

impl<Token> fmt::Display for AbortableWaitTimeoutError<Token> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match self {
            &AbortableWaitTimeoutError::Empty => write!(f, "scheduler is empty"),
            &AbortableWaitTimeoutError::Timeout => write!(f, "timedout while waiting for tokens"),
            &AbortableWaitTimeoutError::Overrun(ref tokens) => write!(f, "scheduler overrun {} tokens", tokens.len()),
            &AbortableWaitTimeoutError::Aborted => write!(f, "wait operation was aborted from another thread")
        }
    }
}

impl<Token> Error for AbortableWaitTimeoutError<Token> where Token: fmt::Debug + Any {
    fn description(&self) -> &str {
        "problem while waiting for next schedule with timeout"
    }
}

impl<Token, TS> Scheduler<Token, TS> where TS: TimeSource + Wait + AbortableWait, Token: Clone {
    pub fn abort_handle(&self) -> <TS as AbortableWait>::AbortHandle {
        self.time_source.abort_handle()
    }

    pub fn abortable_wait(&mut self) -> Result<Vec<Token>, AbortableWaitError<Token>> {
        match self.next() {
            Some(schedule) => match schedule {
                Schedule::NextIn(duration) => {
                    if let Err(WaitAbortedError) = self.time_source.abortable_wait(duration) {
                        return Err(AbortableWaitError::Aborted);
                    };
                    self.abortable_wait()
                },
                Schedule::Overrun(overrun_tokens) => {
                    Err(AbortableWaitError::Overrun(overrun_tokens))
                },
                Schedule::Current(tokens) => {
                    Ok(tokens)
                }
            },
            None => Err(AbortableWaitError::Empty)
        }
    }

    pub fn abortable_wait_timeout(&mut self, timeout: Duration) -> Result<Vec<Token>, AbortableWaitTimeoutError<Token>> {
        match self.next() {
            Some(schedule) => match schedule {
                Schedule::NextIn(duration) => {
                    if duration > timeout {
                        if let Err(WaitAbortedError) = self.time_source.abortable_wait(timeout) {
                            return Err(AbortableWaitTimeoutError::Aborted);
                        };
                        return Err(AbortableWaitTimeoutError::Timeout);
                    }
                    self.time_source.wait(duration);
                    self.abortable_wait_timeout(Duration::zero())
                },
                Schedule::Overrun(overrun_tokens) => {
                    Err(AbortableWaitTimeoutError::Overrun(overrun_tokens))
                },
                Schedule::Current(tokens) => {
                    Ok(tokens)
                }
            },
            None => Err(AbortableWaitTimeoutError::Empty)
        }
    }
}

#[cfg(test)]
mod test {
    use super::*;
    use time_source::*;
    use steady_time_source::*;
    use test_helpers::*;
    use scheduler::*;
    use time::Duration;
    use std::thread::spawn;

    #[test]
    fn abortable_wait_with_aborted_steady_time_source() {
        let mut scheduler = Scheduler::with_time_source(Duration::milliseconds(100), SteadyTimeSource::new());

        scheduler.after(Duration::milliseconds(100), 0);
        scheduler.after(Duration::seconds(20), 1);
        scheduler.after(Duration::seconds(40), 2);

        let abort_handle = scheduler.abort_handle();

        assert_eq!(scheduler.abortable_wait(), Ok(vec![0]));

        spawn(move || {
            abort_handle.abort();
        });

        assert_eq!(scheduler.abortable_wait(), Err(AbortableWaitError::Aborted));
    }

    #[test]
    fn abortable_wait_with_empty() {
        let mut scheduler = Scheduler::with_time_source(Duration::seconds(1), MockTimeSourceWait::new());
        assert_eq!(scheduler.abortable_wait(), Err(AbortableWaitError::Empty));

        scheduler.after(Duration::seconds(0), 0);
        assert_eq!(scheduler.abortable_wait(), Ok(vec![0]));
        assert_eq!(scheduler.abortable_wait(), Err(AbortableWaitError::Empty));
    }

    #[test]
    fn abortable_wait_with_with_overrun() {
        let mut scheduler = Scheduler::with_time_source(Duration::seconds(1), MockTimeSourceWait::new());

        scheduler.after(Duration::seconds(0), 0);
        scheduler.after(Duration::seconds(1), 1);
        scheduler.after(Duration::seconds(2), 2);

        scheduler.fast_forward(Duration::seconds(2));
        assert_eq!(scheduler.abortable_wait(), Err(AbortableWaitError::Overrun(vec![0, 1])));
        assert_eq!(scheduler.abortable_wait(), Ok(vec![2]));
    }

    #[test]
    fn abortable_wait_timeout_with_aborted_steady_time_source() {
        let mut scheduler = Scheduler::with_time_source(Duration::milliseconds(100), SteadyTimeSource::new());

        scheduler.after(Duration::milliseconds(100), 0);
        scheduler.after(Duration::seconds(20), 1);
        scheduler.after(Duration::seconds(40), 2);

        let abort_handle = scheduler.abort_handle();

        assert_eq!(scheduler.abortable_wait_timeout(Duration::seconds(2)), Ok(vec![0]));

        spawn(move || {
            abort_handle.abort();
        });

        assert_eq!(scheduler.abortable_wait_timeout(Duration::seconds(2)), Err(AbortableWaitTimeoutError::Aborted));
    }

    #[test]
    fn abortable_wait_timeout_with_empty() {
        let mut scheduler = Scheduler::with_time_source(Duration::seconds(1), MockTimeSourceWait::new());
        assert_eq!(scheduler.abortable_wait_timeout(Duration::seconds(4)), Err(AbortableWaitTimeoutError::Empty));

        scheduler.after(Duration::seconds(0), 0);
        assert_eq!(scheduler.abortable_wait_timeout(Duration::seconds(4)), Ok(vec![0]));
        assert_eq!(scheduler.abortable_wait_timeout(Duration::seconds(4)), Err(AbortableWaitTimeoutError::Empty));
    }

    #[test]
    fn abortable_wait_timeout_with_overrun() {
        let mut scheduler = Scheduler::with_time_source(Duration::seconds(1), MockTimeSourceWait::new());

        scheduler.after(Duration::seconds(0), 0);
        scheduler.after(Duration::seconds(1), 1);
        scheduler.after(Duration::seconds(2), 2);

        scheduler.fast_forward(Duration::seconds(2));
        assert_eq!(scheduler.abortable_wait_timeout(Duration::seconds(4)), Err(AbortableWaitTimeoutError::Overrun(vec![0, 1])));
        assert_eq!(scheduler.abortable_wait_timeout(Duration::seconds(4)), Ok(vec![2]));
    }

    #[test]
    fn abortable_wait_timeout_with_timeout() {
        let mut scheduler = Scheduler::with_time_source(Duration::seconds(1), MockTimeSourceWait::new());

        scheduler.after(Duration::seconds(0), 0);
        scheduler.after(Duration::seconds(1), 1);
        scheduler.after(Duration::seconds(2), 2);

        assert_eq!(scheduler.abortable_wait_timeout(Duration::seconds(2)), Ok(vec![0]));
        assert_eq!(scheduler.abortable_wait_timeout(Duration::milliseconds(500)), Err(AbortableWaitTimeoutError::Timeout));
    }
}
<|endoftext|>


<commit_before>use num::{self};
use num::traits::{Zero, WrappingAdd, WrappingSub, WrappingMul, CheckedShl, CheckedShr, Signed};
use std::mem;
use std::cmp::{min};
use std::ops::{Add, Sub, Mul, AddAssign, SubAssign, MulAssign, Shl, Shr, Neg, Index};
use std::fmt::Debug;

pub trait Integer:
  'static + num::PrimInt + num::Integer + num::FromPrimitive +
  AddAssign <Self> + SubAssign <Self> + MulAssign <Self> +
  WrappingAdd + WrappingSub + WrappingMul +
  for<'a> Add <&'a Self, Output = Self> + for<'a> Sub <&'a Self, Output = Self> + for<'a> Mul <&'a Self, Output = Self> +
  CheckedShl + CheckedShr + Shl <u32, Output = Self> + Shr <u32, Output = Self> +
  Debug {}
impl <T: 'static + num::PrimInt + num::Integer + num::FromPrimitive + AddAssign <Self> + SubAssign <Self> + MulAssign <Self> + WrappingAdd + WrappingSub + WrappingMul + for<'a> Add <&'a Self, Output = Self> + for<'a> Sub <&'a Self, Output = Self> + for<'a> Mul <&'a Self, Output = Self> + CheckedShl + CheckedShr + Shl <u32, Output = Self> + Shr <u32, Output = Self> + Debug> Integer for T {}

pub trait HasCoordinates
   {
  type Coordinate: Integer + Signed + Vector + HasCoordinates <Coordinate = Self::Coordinate>;}

pub trait Vector:
  'static + Sized + Copy + Clone + HasCoordinates +
  Add <Self, Output = Self> + Sub <Self, Output = Self> + Mul <<Self as HasCoordinates>::Coordinate, Output = Self> +
  for<'a> Add <&'a Self, Output = Self> + for<'a> Sub <&'a Self, Output = Self> + //for<'a> Mul <&'a <Self as HasCoordinates>::Coordinate, Output = Self> +
  AddAssign <Self> + SubAssign <Self> + MulAssign <<Self as HasCoordinates>::Coordinate> +
  for<'a> AddAssign <&'a Self> + for<'a> SubAssign <&'a Self> + //for<'a> MulAssign <&'a <Self as HasCoordinates>::Coordinate> +
  Zero + Neg <Output = Self>
  {
  const DIMENSIONS: usize;
  fn coordinate (&self, which: usize)->Self::Coordinate;
  fn set_coordinate (&mut self, which: usize, value: Self::Coordinate);
}

pub mod impls {
  use super::*;
  use nalgebra::*;
  macro_rules! impl_vector {
    ($([$coordinates: expr, $Vector: ident],)*) => {
      $(
        impl <T: Integer + Signed + Vector + HasCoordinates <Coordinate = T>> HasCoordinates for $Vector <T> {type Coordinate = T;}
        impl <T: Integer + Signed + Vector + HasCoordinates <Coordinate = T>> Vector for $Vector <T> {
          const DIMENSIONS: usize = $coordinates;
          fn coordinate (&self, which: usize)->Self::Coordinate {self [which]}
          fn set_coordinate (&mut self, which: usize, value: Self::Coordinate) {self [which] = value}
        }
      )*
    }
  }
  macro_rules! impl_integer {
    ($($Integer: ident,)*) => {
      $(
        impl HasCoordinates for $Integer {type Coordinate = $Integer;}
        impl Vector for $Integer {
          const DIMENSIONS: usize = 1;
          fn coordinate (&self, which: usize)->Self::Coordinate {*self}
          fn set_coordinate (&mut self, which: usize, value: Self::Coordinate) {*self = value}
        }
      )*
    }
  }
  impl_vector! ([1, Vector1], [2, Vector2], [3, Vector3], [4, Vector4], [5, Vector5], [6, Vector6],);
  impl_integer! (i8, i16, i32, i64, isize,);
}

/// Right-shift an integer, but round to nearest, with ties rounding to even.
///
/// This minimizes error, and avoids a directional bias.
pub fn shr_nicely_rounded <T: Integer> (input: T, shift: u32)->T {
  if shift == 0 {return input}
  let divisor = match T::one().checked_shl ( shift ) {Some (value) => value, None => return T::zero()};
  let mask = divisor.wrapping_sub (&T::one());
  let shifted = input >> shift;
  let half_shift = shift - 1;
  let half = T::one() << half_shift;
  shifted + if (input & mask)+(shifted & T::one()) > half {T::one()} else {T::zero()}
}

/// Right-shift an integer, but round to even.
///
/// This avoids a directional bias.
pub fn shr_round_to_even <T: Integer> (input: T, shift: u32)->T {
  let divisor = match T::one().checked_shl ( shift ) {Some (value) => value, None => return T::zero()};
  let mask = divisor.wrapping_sub (&T::one());
  let shifted = input >> shift;
  shifted + if (input & mask) != T::zero() {shifted & T::one()} else {T::zero()}
}

/// Right-shift an integer, but round towards positive infinity.
pub fn shr_ceil <T: Integer> (input: T, shift: u32)->T {
  let divisor = match T::one().checked_shl ( shift ) {Some (value) => value, None => return T::zero()};
  let mask = divisor.wrapping_sub (&T::one());
  (input >> shift) + if input & mask != T::zero() {T::one()} else {T::zero()}
}

/// Left-shift an integer, returning Some(input*(2^shift)) if it fits within the type, None otherwise.
pub fn overflow_checked_shl <T: Integer> (input: T, shift: u32)->Option <T> {
  if input == T::zero() {return Some (T::zero())}
  let maximum = match T::max_value().checked_shr (shift) {
    None => return None,
    Some (value) => value,
  };
  if input > maximum {return None}
  let minimum = T::min_value() >> shift;
  if input < minimum {return None}
  Some (input << shift)
}

/// Compute the arithmetic mean of two integers, rounded towards negative infinity. Never overflows.
pub fn mean_floor <T: Integer> (first: T, second: T)->T {
  (first >> 1u32) + (second >> 1u32) + (first & second & T::one())
}

/// Compute the arithmetic mean of two integers, rounded towards positive infinity. Never overflows.
pub fn mean_ceil <T: Integer> (first: T, second: T)->T {
  (first >> 1u32) + (second >> 1u32) + ((first | second) & T::one())
}

pub mod polynomial;

#[cfg (test)]
mod tests {
  use super::*;
  use num::{One, ToPrimitive, Integer};
  use num::bigint::BigInt;
  use num::rational::{Ratio};
  use std::cmp::Ordering;

  fn perfect_shr_nicely_rounded <T: Integer> (input: T, shift: u32)->BigInt where BigInt: From <T> {
    let perfect_result = Ratio::new (BigInt::from (input), BigInt::one() << shift as usize);
    let rounded_down = perfect_result.floor();
    let fraction = & perfect_result - & rounded_down;
    let rounded_down = rounded_down.to_integer();
    match fraction.cmp (& Ratio::new (BigInt::one(), BigInt::one() << 1)) {
      Ordering::Less => rounded_down, Ordering::Greater => rounded_down + BigInt::one() ,
      Ordering::Equal => & rounded_down + rounded_down.mod_floor (& (BigInt::one() << 1)),
    }
  }

  fn perfect_shr_round_to_even <T: Integer> (input: T, shift: u32)->BigInt where BigInt: From <T> {
    let perfect_result = Ratio::new (BigInt::from (input), BigInt::one() << shift as usize);
    let rounded_down = perfect_result.floor();
    if perfect_result == rounded_down {rounded_down.to_integer()}
    else {
      let rounded_down = rounded_down.to_integer();
      if &rounded_down >> 1 << 1 == rounded_down {
        rounded_down
      }
      else {rounded_down + 1}
    }
  }

  #[test]
  fn test_shr_nicely_rounded() {
    let inputs: Vec<(i64, u32, i64)> = vec![
      (0, 0, 0), (0, 5, 0), (5, 0, 5), (1, 3, 0), (4, 3, 0), (5, 3, 1),
      (999, 1, 500), (998, 1, 499), (997, 1, 498)
    ];
    for (input, shift, result) in inputs {
      println!( "{:?}", (input, shift, result));
      assert_eq!(shr_nicely_rounded (input, shift), result);
      assert_eq!(shr_nicely_rounded (-input, shift), -result);
    }
  }

  #[test]
  fn test_shr_round_to_even() {
    let inputs: Vec<(i64, u32, i64)> = vec![
      (0, 0, 0), (0, 5, 0), (5, 0, 5), (1, 3, 0), (4, 3, 0), (5, 3, 0),
      (999, 1, 500), (998, 1, 499), (997, 1, 498)
    ];
    for (input, shift, result) in inputs {
      println!( "{:?}", (input, shift, result));
      assert_eq!(shr_round_to_even (input, shift), result);
      assert_eq!(shr_round_to_even (-input, shift), -result);
    }
  }

  quickcheck! {
    fn quickcheck_shr_nicely_rounded_signed (input: i32, shift: u8)->bool {
      let result = shr_nicely_rounded (input, shift as u32);
      let perfect_result = perfect_shr_nicely_rounded (input, shift as u32);
      println!( "{:?}", (result, & perfect_result.to_str_radix (10)));
      perfect_result == BigInt::from (result)
    }

    fn quickcheck_shr_nicely_rounded_unsigned (input: u32, shift: u8)->bool {
      let result = shr_nicely_rounded (input, shift as u32);
      let perfect_result = perfect_shr_nicely_rounded (input, shift as u32);
      println!( "{:?}", (result, & perfect_result.to_str_radix (10)));
      perfect_result == BigInt::from (result)
    }

    fn quickcheck_shr_round_to_even_signed (input: i32, shift: u8)->bool {
      let result = shr_round_to_even (input, shift as u32);
      let perfect_result = perfect_shr_round_to_even (input, shift as u32);
      println!( "{:?}", (result, & perfect_result.to_str_radix (10)));
      perfect_result == BigInt::from (result)
    }

    fn quickcheck_shr_round_to_even_unsigned (input: u32, shift: u8)->bool {
      let result = shr_round_to_even (input, shift as u32);
      let perfect_result = perfect_shr_round_to_even (input, shift as u32);
      println!( "{:?}", (result, & perfect_result.to_str_radix (10)));
      perfect_result == BigInt::from (result)
    }

    fn quickcheck_overflow_checked_shl (input: i32, shift: u8)->bool {
      let result = overflow_checked_shl (input, shift as u32);
      let perfect_result = BigInt::from (input) << shift as usize;
      result == perfect_result.to_i32()
    }
  }
}
<commit_msg>Made the shr functions more concise again<commit_after>use num::{self};
use num::traits::{Zero, WrappingAdd, WrappingSub, WrappingMul, CheckedShl, CheckedShr, Signed};
use std::mem;
use std::cmp::{min};
use std::ops::{Add, Sub, Mul, AddAssign, SubAssign, MulAssign, Shl, Shr, Neg, Index};
use std::fmt::Debug;

pub trait Integer:
  'static + num::PrimInt + num::Integer + num::FromPrimitive +
  AddAssign <Self> + SubAssign <Self> + MulAssign <Self> +
  WrappingAdd + WrappingSub + WrappingMul +
  for<'a> Add <&'a Self, Output = Self> + for<'a> Sub <&'a Self, Output = Self> + for<'a> Mul <&'a Self, Output = Self> +
  CheckedShl + CheckedShr + Shl <u32, Output = Self> + Shr <u32, Output = Self> +
  Debug {}
impl <T: 'static + num::PrimInt + num::Integer + num::FromPrimitive + AddAssign <Self> + SubAssign <Self> + MulAssign <Self> + WrappingAdd + WrappingSub + WrappingMul + for<'a> Add <&'a Self, Output = Self> + for<'a> Sub <&'a Self, Output = Self> + for<'a> Mul <&'a Self, Output = Self> + CheckedShl + CheckedShr + Shl <u32, Output = Self> + Shr <u32, Output = Self> + Debug> Integer for T {}

pub trait HasCoordinates
   {
  type Coordinate: Integer + Signed + Vector + HasCoordinates <Coordinate = Self::Coordinate>;}

pub trait Vector:
  'static + Sized + Copy + Clone + HasCoordinates +
  Add <Self, Output = Self> + Sub <Self, Output = Self> + Mul <<Self as HasCoordinates>::Coordinate, Output = Self> +
  for<'a> Add <&'a Self, Output = Self> + for<'a> Sub <&'a Self, Output = Self> + //for<'a> Mul <&'a <Self as HasCoordinates>::Coordinate, Output = Self> +
  AddAssign <Self> + SubAssign <Self> + MulAssign <<Self as HasCoordinates>::Coordinate> +
  for<'a> AddAssign <&'a Self> + for<'a> SubAssign <&'a Self> + //for<'a> MulAssign <&'a <Self as HasCoordinates>::Coordinate> +
  Zero + Neg <Output = Self>
  {
  const DIMENSIONS: usize;
  fn coordinate (&self, which: usize)->Self::Coordinate;
  fn set_coordinate (&mut self, which: usize, value: Self::Coordinate);
}

pub mod impls {
  use super::*;
  use nalgebra::*;
  macro_rules! impl_vector {
    ($([$coordinates: expr, $Vector: ident],)*) => {
      $(
        impl <T: Integer + Signed + Vector + HasCoordinates <Coordinate = T>> HasCoordinates for $Vector <T> {type Coordinate = T;}
        impl <T: Integer + Signed + Vector + HasCoordinates <Coordinate = T>> Vector for $Vector <T> {
          const DIMENSIONS: usize = $coordinates;
          fn coordinate (&self, which: usize)->Self::Coordinate {self [which]}
          fn set_coordinate (&mut self, which: usize, value: Self::Coordinate) {self [which] = value}
        }
      )*
    }
  }
  macro_rules! impl_integer {
    ($($Integer: ident,)*) => {
      $(
        impl HasCoordinates for $Integer {type Coordinate = $Integer;}
        impl Vector for $Integer {
          const DIMENSIONS: usize = 1;
          fn coordinate (&self, which: usize)->Self::Coordinate {*self}
          fn set_coordinate (&mut self, which: usize, value: Self::Coordinate) {*self = value}
        }
      )*
    }
  }
  impl_vector! ([1, Vector1], [2, Vector2], [3, Vector3], [4, Vector4], [5, Vector5], [6, Vector6],);
  impl_integer! (i8, i16, i32, i64, isize,);
}

/// Right-shift an integer, but round to nearest, with ties rounding to even.
///
/// This minimizes error, and avoids a directional bias.
pub fn shr_nicely_rounded <T: Integer> (input: T, shift: u32)->T {
  let divisor = match T::one().checked_shl ( shift ) {Some (value) => value, None => return T::zero()};
  let mask = divisor.wrapping_sub (&T::one());
  let shifted = input >> shift;
  let half = divisor >> 1u32;
  shifted + if (input & mask)+(shifted & if shift != 0 {T::one()} else {T::zero()}) > half {T::one()} else {T::zero()}
}

/// Right-shift an integer, but round to even.
///
/// This avoids a directional bias.
pub fn shr_round_to_even <T: Integer> (input: T, shift: u32)->T {
  let divisor = match T::one().checked_shl ( shift ) {Some (value) => value, None => return T::zero()};
  let mask = divisor.wrapping_sub (&T::one());
  let shifted = input >> shift;
  shifted + if (input & mask) != T::zero() {shifted & T::one()} else {T::zero()}
}

/// Right-shift an integer, but round towards positive infinity.
pub fn shr_ceil <T: Integer> (input: T, shift: u32)->T {
  let divisor = match T::one().checked_shl ( shift ) {Some (value) => value, None => return T::zero()};
  let mask = divisor.wrapping_sub (&T::one());
  (input >> shift) + if input & mask != T::zero() {T::one()} else {T::zero()}
}

/// Left-shift an integer, returning Some(input*(2^shift)) if it fits within the type, None otherwise.
pub fn overflow_checked_shl <T: Integer> (input: T, shift: u32)->Option <T> {
  if input == T::zero() {return Some (T::zero())}
  let maximum = match T::max_value().checked_shr (shift) {
    None => return None,
    Some (value) => value,
  };
  if input > maximum {return None}
  let minimum = T::min_value() >> shift;
  if input < minimum {return None}
  Some (input << shift)
}

/// Compute the arithmetic mean of two integers, rounded towards negative infinity. Never overflows.
pub fn mean_floor <T: Integer> (first: T, second: T)->T {
  (first >> 1u32) + (second >> 1u32) + (first & second & T::one())
}

/// Compute the arithmetic mean of two integers, rounded towards positive infinity. Never overflows.
pub fn mean_ceil <T: Integer> (first: T, second: T)->T {
  (first >> 1u32) + (second >> 1u32) + ((first | second) & T::one())
}

pub mod polynomial;

#[cfg (test)]
mod tests {
  use super::*;
  use num::{One, ToPrimitive, Integer};
  use num::bigint::BigInt;
  use num::rational::{Ratio};
  use std::cmp::Ordering;

  fn perfect_shr_nicely_rounded <T: Integer> (input: T, shift: u32)->BigInt where BigInt: From <T> {
    let perfect_result = Ratio::new (BigInt::from (input), BigInt::one() << shift as usize);
    let rounded_down = perfect_result.floor();
    let fraction = & perfect_result - & rounded_down;
    let rounded_down = rounded_down.to_integer();
    match fraction.cmp (& Ratio::new (BigInt::one(), BigInt::one() << 1)) {
      Ordering::Less => rounded_down, Ordering::Greater => rounded_down + BigInt::one() ,
      Ordering::Equal => & rounded_down + rounded_down.mod_floor (& (BigInt::one() << 1)),
    }
  }

  fn perfect_shr_round_to_even <T: Integer> (input: T, shift: u32)->BigInt where BigInt: From <T> {
    let perfect_result = Ratio::new (BigInt::from (input), BigInt::one() << shift as usize);
    let rounded_down = perfect_result.floor();
    if perfect_result == rounded_down {rounded_down.to_integer()}
    else {
      let rounded_down = rounded_down.to_integer();
      if &rounded_down >> 1 << 1 == rounded_down {
        rounded_down
      }
      else {rounded_down + 1}
    }
  }

  #[test]
  fn test_shr_nicely_rounded() {
    let inputs: Vec<(i64, u32, i64)> = vec![
      (0, 0, 0), (0, 5, 0), (5, 0, 5), (1, 3, 0), (4, 3, 0), (5, 3, 1),
      (999, 1, 500), (998, 1, 499), (997, 1, 498)
    ];
    for (input, shift, result) in inputs {
      println!( "{:?}", (input, shift, result));
      assert_eq!(shr_nicely_rounded (input, shift), result);
      assert_eq!(shr_nicely_rounded (-input, shift), -result);
    }
  }

  #[test]
  fn test_shr_round_to_even() {
    let inputs: Vec<(i64, u32, i64)> = vec![
      (0, 0, 0), (0, 5, 0), (5, 0, 5), (1, 3, 0), (4, 3, 0), (5, 3, 0),
      (999, 1, 500), (998, 1, 499), (997, 1, 498)
    ];
    for (input, shift, result) in inputs {
      println!( "{:?}", (input, shift, result));
      assert_eq!(shr_round_to_even (input, shift), result);
      assert_eq!(shr_round_to_even (-input, shift), -result);
    }
  }

  quickcheck! {
    fn quickcheck_shr_nicely_rounded_signed (input: i32, shift: u8)->bool {
      let result = shr_nicely_rounded (input, shift as u32);
      let perfect_result = perfect_shr_nicely_rounded (input, shift as u32);
      println!( "{:?}", (result, & perfect_result.to_str_radix (10)));
      perfect_result == BigInt::from (result)
    }

    fn quickcheck_shr_nicely_rounded_unsigned (input: u32, shift: u8)->bool {
      let result = shr_nicely_rounded (input, shift as u32);
      let perfect_result = perfect_shr_nicely_rounded (input, shift as u32);
      println!( "{:?}", (result, & perfect_result.to_str_radix (10)));
      perfect_result == BigInt::from (result)
    }

    fn quickcheck_shr_round_to_even_signed (input: i32, shift: u8)->bool {
      let result = shr_round_to_even (input, shift as u32);
      let perfect_result = perfect_shr_round_to_even (input, shift as u32);
      println!( "{:?}", (result, & perfect_result.to_str_radix (10)));
      perfect_result == BigInt::from (result)
    }

    fn quickcheck_shr_round_to_even_unsigned (input: u32, shift: u8)->bool {
      let result = shr_round_to_even (input, shift as u32);
      let perfect_result = perfect_shr_round_to_even (input, shift as u32);
      println!( "{:?}", (result, & perfect_result.to_str_radix (10)));
      perfect_result == BigInt::from (result)
    }

    fn quickcheck_overflow_checked_shl (input: i32, shift: u8)->bool {
      let result = overflow_checked_shl (input, shift as u32);
      let perfect_result = BigInt::from (input) << shift as usize;
      result == perfect_result.to_i32()
    }
  }
}
<|endoftext|>


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

use std::mem;
use rustc_data_structures::small_vec::SmallVec;
use syntax::ast::{CRATE_NODE_ID, NodeId};
use util::nodemap::FxHashSet;
use ty::context::TyCtxt;
use ty::{AdtDef, VariantDef, FieldDef, TyS};
use ty::{DefId, Substs};
use ty::{AdtKind, Visibility, NodeIdTree};
use ty::TypeVariants::*;

/// Represents a set of nodes closed under the ancestor relation. That is, if a
/// node is in this set then so are all its descendants.
#[derive(Clone)]
pub struct NodeForrest {
    /// The minimal set of nodes required to represent the whole set.
    /// If A and B are nodes in the NodeForrest, and A is a desecendant
    /// of B, then only B will be in root_nodes.
    /// We use a SmallVec here because (for its use in this module) its rare
    /// that this will contain more than one or two nodes.
    root_nodes: SmallVec<[NodeId; 1]>,
}

impl<'a, 'gcx, 'tcx> NodeForrest {
    /// Create an empty set.
    pub fn empty() -> NodeForrest {
        NodeForrest {
            root_nodes: SmallVec::new(),
        }
    }

    /// Create a set containing every node.
    #[inline]
    pub fn full() -> NodeForrest {
        NodeForrest::from_node(CRATE_NODE_ID)
    }

    /// Create a set containing a node and all its descendants.
    pub fn from_node(node: NodeId) -> NodeForrest {
        let mut root_nodes = SmallVec::new();
        root_nodes.push(node);
        NodeForrest {
            root_nodes: root_nodes,
        }
    }

    /// Test whether the set is empty.
    pub fn is_empty(&self) -> bool {
        self.root_nodes.is_empty()
    }

    /// Test whether the set conains a node.
    pub fn contains(&self,
                    tcx: TyCtxt<'a, 'gcx, 'tcx>,
                    node: NodeId) -> bool
    {
        for root_node in self.root_nodes.iter() {
            if tcx.map.is_descendant_of(node, *root_node) {
                return true;
            }
        }
        false
    }

    /// Calculate the intersection of a collection of sets.
    pub fn intersection<I>(tcx: TyCtxt<'a, 'gcx, 'tcx>,
                           iter: I) -> NodeForrest
            where I: IntoIterator<Item=NodeForrest>
    {
        let mut ret = NodeForrest::full();
        let mut next_ret = SmallVec::new();
        let mut old_ret: SmallVec<[NodeId; 1]> = SmallVec::new();
        for next_set in iter {
            for node in ret.root_nodes.drain(..) {
                if next_set.contains(tcx, node) {
                    next_ret.push(node);
                } else {
                    old_ret.push(node);
                }
            }
            ret.root_nodes.extend(old_ret.drain(..));

            for node in next_set.root_nodes {
                if ret.contains(tcx, node) {
                    next_ret.push(node);
                }
            }

            mem::swap(&mut next_ret, &mut ret.root_nodes);
            next_ret.drain(..);
        }
        ret
    }

    /// Calculate the union of a collection of sets.
    pub fn union<I>(tcx: TyCtxt<'a, 'gcx, 'tcx>,
                    iter: I) -> NodeForrest
            where I: IntoIterator<Item=NodeForrest>
    {
        let mut ret = NodeForrest::empty();
        let mut next_ret = SmallVec::new();
        for next_set in iter {
            for node in ret.root_nodes.drain(..) {
                if !next_set.contains(tcx, node) {
                    next_ret.push(node);
                }
            }

            for node in next_set.root_nodes {
                if !next_ret.contains(&node) {
                    next_ret.push(node);
                }
            }

            mem::swap(&mut next_ret, &mut ret.root_nodes);
            next_ret.drain(..);
        }
        ret
    }
}

impl<'a, 'gcx, 'tcx> AdtDef {
    /// Calculate the set of  nodes from which this adt is visibly uninhabited.
    pub fn uninhabited_from(
                &self,
                visited: &mut FxHashSet<(DefId, &'tcx Substs<'tcx>)>,
                tcx: TyCtxt<'a, 'gcx, 'tcx>,
                substs: &'tcx Substs<'tcx>) -> NodeForrest
    {
        if !visited.insert((self.did, substs)) {
            return NodeForrest::empty();
        }

        let ret = NodeForrest::intersection(tcx, self.variants.iter().map(|v| {
            v.uninhabited_from(visited, tcx, substs, self.adt_kind())
        }));
        visited.remove(&(self.did, substs));
        ret
    }
}

impl<'a, 'gcx, 'tcx> VariantDef {
    /// Calculate the set of  nodes from which this variant is visibly uninhabited.
    pub fn uninhabited_from(
                &self,
                visited: &mut FxHashSet<(DefId, &'tcx Substs<'tcx>)>,
                tcx: TyCtxt<'a, 'gcx, 'tcx>,
                substs: &'tcx Substs<'tcx>,
                adt_kind: AdtKind) -> NodeForrest
    {
        match adt_kind {
            AdtKind::Union => {
                NodeForrest::intersection(tcx, self.fields.iter().map(|f| {
                    f.uninhabited_from(visited, tcx, substs, false)
                }))
            },
            AdtKind::Struct => {
                NodeForrest::union(tcx, self.fields.iter().map(|f| {
                    f.uninhabited_from(visited, tcx, substs, false)
                }))
            },
            AdtKind::Enum => {
                NodeForrest::union(tcx, self.fields.iter().map(|f| {
                    f.uninhabited_from(visited, tcx, substs, true)
                }))
            },
        }
    }
}

impl<'a, 'gcx, 'tcx> FieldDef {
    /// Calculate the set of  nodes from which this field is visibly uninhabited.
    pub fn uninhabited_from(
                &self,
                visited: &mut FxHashSet<(DefId, &'tcx Substs<'tcx>)>,
                tcx: TyCtxt<'a, 'gcx, 'tcx>,
                substs: &'tcx Substs<'tcx>,
                is_enum: bool) -> NodeForrest
    {
        if let Visibility::PrivateExternal = self.vis {
            return NodeForrest::empty();
        }

        let data_inhabitedness = self.ty(tcx, substs).uninhabited_from(visited, tcx);
        match self.vis {
            Visibility::Restricted(from) if !is_enum => {
                let node_set = NodeForrest::from_node(from);
                let iter = Some(node_set).into_iter().chain(Some(data_inhabitedness));
                NodeForrest::intersection(tcx, iter)
            },
            _ => data_inhabitedness,
        }
    }
}

impl<'a, 'gcx, 'tcx> TyS<'tcx> {
    /// Calculate the set of  nodes from which this type is visibly uninhabited.
    pub fn uninhabited_from(
                &self,
                visited: &mut FxHashSet<(DefId, &'tcx Substs<'tcx>)>,
                tcx: TyCtxt<'a, 'gcx, 'tcx>) -> NodeForrest
    {
        match tcx.lift_to_global(&self) {
            Some(global_ty) => {
                {
                    let cache = tcx.inhabitedness_cache.borrow();
                    if let Some(closed_node_set) = cache.get(&global_ty) {
                        return closed_node_set.clone();
                    }
                }
                let node_set = global_ty.uninhabited_from_inner(visited, tcx);
                let mut cache = tcx.inhabitedness_cache.borrow_mut();
                cache.insert(global_ty, node_set.clone());
                node_set
            },
            None => {
                let node_set = self.uninhabited_from_inner(visited, tcx);
                node_set
            },
        }
    }

    fn uninhabited_from_inner(
                &self,
                visited: &mut FxHashSet<(DefId, &'tcx Substs<'tcx>)>,
                tcx: TyCtxt<'a, 'gcx, 'tcx>) -> NodeForrest
    {
        match self.sty {
            TyAdt(def, substs) => {
                def.uninhabited_from(visited, tcx, substs)
            },

            TyNever => NodeForrest::full(),
            TyTuple(ref tys) => {
                NodeForrest::union(tcx, tys.iter().map(|ty| {
                    ty.uninhabited_from(visited, tcx)
                }))
            },
            TyArray(ty, len) => {
                if len == 0 {
                    NodeForrest::empty()
                } else {
                    ty.uninhabited_from(visited, tcx)
                }
            }
            TyRef(_, ref tm) => tm.ty.uninhabited_from(visited, tcx),

            _ => NodeForrest::empty(),
        }
    }
}

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

use std::mem;
use rustc_data_structures::small_vec::SmallVec;
use syntax::ast::{CRATE_NODE_ID, NodeId};
use util::nodemap::FxHashSet;
use ty::context::TyCtxt;
use ty::{AdtDef, VariantDef, FieldDef, TyS};
use ty::{DefId, Substs};
use ty::{AdtKind, Visibility, NodeIdTree};
use ty::TypeVariants::*;

/// Represents a set of nodes closed under the ancestor relation. That is, if a
/// node is in this set then so are all its descendants.
#[derive(Clone)]
pub struct NodeForrest {
    /// The minimal set of nodes required to represent the whole set.
    /// If A and B are nodes in the NodeForrest, and A is a desecendant
    /// of B, then only B will be in root_nodes.
    /// We use a SmallVec here because (for its use in this module) its rare
    /// that this will contain more than one or two nodes.
    root_nodes: SmallVec<[NodeId; 1]>,
}

impl<'a, 'gcx, 'tcx> NodeForrest {
    /// Create an empty set.
    pub fn empty() -> NodeForrest {
        NodeForrest {
            root_nodes: SmallVec::new(),
        }
    }

    /// Create a set containing every node.
    #[inline]
    pub fn full() -> NodeForrest {
        NodeForrest::from_node(CRATE_NODE_ID)
    }

    /// Create a set containing a node and all its descendants.
    pub fn from_node(node: NodeId) -> NodeForrest {
        let mut root_nodes = SmallVec::new();
        root_nodes.push(node);
        NodeForrest {
            root_nodes: root_nodes,
        }
    }

    /// Test whether the set is empty.
    pub fn is_empty(&self) -> bool {
        self.root_nodes.is_empty()
    }

    /// Test whether the set conains a node.
    pub fn contains(&self,
                    tcx: TyCtxt<'a, 'gcx, 'tcx>,
                    node: NodeId) -> bool
    {
        for root_node in self.root_nodes.iter() {
            if tcx.map.is_descendant_of(node, *root_node) {
                return true;
            }
        }
        false
    }

    /// Calculate the intersection of a collection of sets.
    pub fn intersection<I>(tcx: TyCtxt<'a, 'gcx, 'tcx>,
                           iter: I) -> NodeForrest
            where I: IntoIterator<Item=NodeForrest>
    {
        let mut ret = NodeForrest::full();
        let mut next_ret = SmallVec::new();
        let mut old_ret: SmallVec<[NodeId; 1]> = SmallVec::new();
        for next_set in iter {
            for node in ret.root_nodes.drain(..) {
                if next_set.contains(tcx, node) {
                    next_ret.push(node);
                } else {
                    old_ret.push(node);
                }
            }
            ret.root_nodes.extend(old_ret.drain(..));

            for node in next_set.root_nodes {
                if ret.contains(tcx, node) {
                    next_ret.push(node);
                }
            }

            mem::swap(&mut next_ret, &mut ret.root_nodes);
            next_ret.drain(..);
        }
        ret
    }

    /// Calculate the union of a collection of sets.
    pub fn union<I>(tcx: TyCtxt<'a, 'gcx, 'tcx>,
                    iter: I) -> NodeForrest
            where I: IntoIterator<Item=NodeForrest>
    {
        let mut ret = NodeForrest::empty();
        let mut next_ret = SmallVec::new();
        for next_set in iter {
            for node in ret.root_nodes.drain(..) {
                if !next_set.contains(tcx, node) {
                    next_ret.push(node);
                }
            }

            for node in next_set.root_nodes {
                if !next_ret.contains(&node) {
                    next_ret.push(node);
                }
            }

            mem::swap(&mut next_ret, &mut ret.root_nodes);
            next_ret.drain(..);
        }
        ret
    }
}

impl<'a, 'gcx, 'tcx> AdtDef {
    /// Calculate the set of  nodes from which this adt is visibly uninhabited.
    pub fn uninhabited_from(
                &self,
                visited: &mut FxHashSet<(DefId, &'tcx Substs<'tcx>)>,
                tcx: TyCtxt<'a, 'gcx, 'tcx>,
                substs: &'tcx Substs<'tcx>) -> NodeForrest
    {
        if !visited.insert((self.did, substs)) {
            return NodeForrest::empty();
        }

        let ret = NodeForrest::intersection(tcx, self.variants.iter().map(|v| {
            v.uninhabited_from(visited, tcx, substs, self.adt_kind())
        }));
        visited.remove(&(self.did, substs));
        ret
    }
}

impl<'a, 'gcx, 'tcx> VariantDef {
    /// Calculate the set of  nodes from which this variant is visibly uninhabited.
    pub fn uninhabited_from(
                &self,
                visited: &mut FxHashSet<(DefId, &'tcx Substs<'tcx>)>,
                tcx: TyCtxt<'a, 'gcx, 'tcx>,
                substs: &'tcx Substs<'tcx>,
                adt_kind: AdtKind) -> NodeForrest
    {
        match adt_kind {
            AdtKind::Union => {
                NodeForrest::intersection(tcx, self.fields.iter().map(|f| {
                    f.uninhabited_from(visited, tcx, substs, false)
                }))
            },
            AdtKind::Struct => {
                NodeForrest::union(tcx, self.fields.iter().map(|f| {
                    f.uninhabited_from(visited, tcx, substs, false)
                }))
            },
            AdtKind::Enum => {
                NodeForrest::union(tcx, self.fields.iter().map(|f| {
                    f.uninhabited_from(visited, tcx, substs, true)
                }))
            },
        }
    }
}

impl<'a, 'gcx, 'tcx> FieldDef {
    /// Calculate the set of  nodes from which this field is visibly uninhabited.
    pub fn uninhabited_from(
                &self,
                visited: &mut FxHashSet<(DefId, &'tcx Substs<'tcx>)>,
                tcx: TyCtxt<'a, 'gcx, 'tcx>,
                substs: &'tcx Substs<'tcx>,
                is_enum: bool) -> NodeForrest
    {
        let mut data_uninhabitedness = move || self.ty(tcx, substs).uninhabited_from(visited, tcx);
        if is_enum {
            data_uninhabitedness()
        } else {
            match self.vis {
                Visibility::PrivateExternal => NodeForrest::empty(),
                Visibility::Restricted(from) => {
                    let node_set = NodeForrest::from_node(from);
                    let iter = Some(node_set).into_iter().chain(Some(data_uninhabitedness()));
                    NodeForrest::intersection(tcx, iter)
                },
                Visibility::Public => data_uninhabitedness(),
            }
        }
    }
}

impl<'a, 'gcx, 'tcx> TyS<'tcx> {
    /// Calculate the set of  nodes from which this type is visibly uninhabited.
    pub fn uninhabited_from(
                &self,
                visited: &mut FxHashSet<(DefId, &'tcx Substs<'tcx>)>,
                tcx: TyCtxt<'a, 'gcx, 'tcx>) -> NodeForrest
    {
        match tcx.lift_to_global(&self) {
            Some(global_ty) => {
                {
                    let cache = tcx.inhabitedness_cache.borrow();
                    if let Some(closed_node_set) = cache.get(&global_ty) {
                        return closed_node_set.clone();
                    }
                }
                let node_set = global_ty.uninhabited_from_inner(visited, tcx);
                let mut cache = tcx.inhabitedness_cache.borrow_mut();
                cache.insert(global_ty, node_set.clone());
                node_set
            },
            None => {
                let node_set = self.uninhabited_from_inner(visited, tcx);
                node_set
            },
        }
    }

    fn uninhabited_from_inner(
                &self,
                visited: &mut FxHashSet<(DefId, &'tcx Substs<'tcx>)>,
                tcx: TyCtxt<'a, 'gcx, 'tcx>) -> NodeForrest
    {
        match self.sty {
            TyAdt(def, substs) => {
                def.uninhabited_from(visited, tcx, substs)
            },

            TyNever => NodeForrest::full(),
            TyTuple(ref tys) => {
                NodeForrest::union(tcx, tys.iter().map(|ty| {
                    ty.uninhabited_from(visited, tcx)
                }))
            },
            TyArray(ty, len) => {
                if len == 0 {
                    NodeForrest::empty()
                } else {
                    ty.uninhabited_from(visited, tcx)
                }
            }
            TyRef(_, ref tm) => tm.ty.uninhabited_from(visited, tcx),

            _ => NodeForrest::empty(),
        }
    }
}

<|endoftext|>


<commit_before>//! Basic syntax highlighting functionality.
//!
//! This module uses librustc_ast's lexer to provide token-based highlighting for
//! the HTML documentation generated by rustdoc.
//!
//! Use the `render_with_highlighting` to highlight some rust code.

use crate::html::escape::Escape;

use std::fmt::Display;
use std::iter::Peekable;

use rustc_lexer::{LiteralKind, TokenKind};
use rustc_span::edition::Edition;
use rustc_span::symbol::Symbol;
use rustc_span::with_default_session_globals;

use super::format::Buffer;

/// Highlights `src`, returning the HTML output.
crate fn render_with_highlighting(
    src: &str,
    out: &mut Buffer,
    class: Option<&str>,
    playground_button: Option<&str>,
    tooltip: Option<(Option<Edition>, &str)>,
    edition: Edition,
) {
    debug!("highlighting: ================\n{}\n==============", src);
    if let Some((edition_info, class)) = tooltip {
        write!(
            out,
            "

ⓘ

",
            class,
            if let Some(edition_info) = edition_info {
                format!(" data-edition=\"{}\"", edition_info)
            } else {
                String::new()
            },
        );
    }

    write_header(out, class);
    write_code(out, &src, edition);
    write_footer(out, playground_button);
}

fn write_header(out: &mut Buffer, class: Option<&str>) {
    writeln!(out, "

", class.unwrap_or_default());
}

fn write_code(out: &mut Buffer, src: &str, edition: Edition) {
    // This replace allows to fix how the code source with DOS backline characters is displayed.
    let src = src.replace("\r\n", "\n");
    Classifier::new(&src, edition).highlight(&mut |highlight| {
        match highlight {
            Highlight::Token { text, class } => string(out, Escape(text), class),
            Highlight::EnterSpan { class } => enter_span(out, class),
            Highlight::ExitSpan => exit_span(out),
        };
    });
}

fn write_footer(out: &mut Buffer, playground_button: Option<&str>) {
    writeln!(out, "
{}
", playground_button.unwrap_or_default());
}

/// How a span of text is classified. Mostly corresponds to token kinds.
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
enum Class {
    Comment,
    DocComment,
    Attribute,
    KeyWord,
    // Keywords that do pointer/reference stuff.
    RefKeyWord,
    Self_,
    Op,
    Macro,
    MacroNonTerminal,
    String,
    Number,
    Bool,
    Ident,
    Lifetime,
    PreludeTy,
    PreludeVal,
    QuestionMark,
}

impl Class {
    /// Returns the css class expected by rustdoc for each `Class`.
    fn as_html(self) -> &'static str {
        match self {
            Class::Comment => "comment",
            Class::DocComment => "doccomment",
            Class::Attribute => "attribute",
            Class::KeyWord => "kw",
            Class::RefKeyWord => "kw-2",
            Class::Self_ => "self",
            Class::Op => "op",
            Class::Macro => "macro",
            Class::MacroNonTerminal => "macro-nonterminal",
            Class::String => "string",
            Class::Number => "number",
            Class::Bool => "bool-val",
            Class::Ident => "ident",
            Class::Lifetime => "lifetime",
            Class::PreludeTy => "prelude-ty",
            Class::PreludeVal => "prelude-val",
            Class::QuestionMark => "question-mark",
        }
    }
}

enum Highlight<'a> {
    Token { text: &'a str, class: Option<Class> },
    EnterSpan { class: Class },
    ExitSpan,
}

struct TokenIter<'a> {
    src: &'a str,
}

impl Iterator for TokenIter<'a> {
    type Item = (TokenKind, &'a str);
    fn next(&mut self) -> Option<(TokenKind, &'a str)> {
        if self.src.is_empty() {
            return None;
        }
        let token = rustc_lexer::first_token(self.src);
        let (text, rest) = self.src.split_at(token.len);
        self.src = rest;
        Some((token.kind, text))
    }
}

fn get_real_ident_class(text: &str, edition: Edition) -> Class {
    match text {
        "ref" | "mut" => Class::RefKeyWord,
        "self" | "Self" => Class::Self_,
        "false" | "true" => Class::Bool,
        _ if Symbol::intern(text).is_reserved(|| edition) => Class::KeyWord,
        _ => Class::Ident,
    }
}

/// Processes program tokens, classifying strings of text by highlighting
/// category (`Class`).
struct Classifier<'a> {
    tokens: Peekable<TokenIter<'a>>,
    in_attribute: bool,
    in_macro: bool,
    in_macro_nonterminal: bool,
    edition: Edition,
    byte_pos: u32,
    src: &'a str,
}

impl<'a> Classifier<'a> {
    fn new(src: &str, edition: Edition) -> Classifier<'_> {
        let tokens = TokenIter { src }.peekable();
        Classifier {
            tokens,
            in_attribute: false,
            in_macro: false,
            in_macro_nonterminal: false,
            edition,
            byte_pos: 0,
            src,
        }
    }

    /// Concatenate colons and idents as one when possible.
    fn get_full_ident_path(&mut self) -> Vec<(TokenKind, usize, usize)> {
        let start = self.byte_pos as usize;
        let mut pos = start;
        let mut has_ident = false;
        let edition = self.edition;

        loop {
            let mut nb = 0;
            while let Some((TokenKind::Colon, _)) = self.tokens.peek() {
                self.tokens.next();
                nb += 1;
            }
            // Ident path can start with "::" but if we already have content in the ident path,
            // the "::" is mandatory.
            if has_ident && nb == 0 {
                return vec![(TokenKind::Ident, start, pos)];
            } else if nb != 0 && nb != 2 {
                if has_ident {
                    return vec![(TokenKind::Ident, start, pos), (TokenKind::Colon, pos, pos + nb)];
                } else {
                    return vec![(TokenKind::Colon, pos, pos + nb)];
                }
            }

            if let Some((Class::Ident, text)) = self.tokens.peek().map(|(token, text)| {
                if *token == TokenKind::Ident {
                    let class = get_real_ident_class(text, edition);
                    (class, text)
                } else {
                    // Doesn't matter which Class we put in here...
                    (Class::Comment, text)
                }
            }) {
                // We only "add" the colon if there is an ident behind.
                pos += text.len() + nb;
                has_ident = true;
                self.tokens.next();
            } else if nb > 0 && has_ident {
                return vec![(TokenKind::Ident, start, pos), (TokenKind::Colon, pos, pos + nb)];
            } else if nb > 0 {
                return vec![(TokenKind::Colon, pos, pos + nb)];
            } else if has_ident {
                return vec![(TokenKind::Ident, start, pos)];
            } else {
                return Vec::new();
            }
        }
    }

    /// Wraps the tokens iteration to ensure that the byte_pos is always correct.
    fn next(&mut self) -> Option<(TokenKind, &'a str)> {
        if let Some((kind, text)) = self.tokens.next() {
            self.byte_pos += text.len() as u32;
            Some((kind, text))
        } else {
            None
        }
    }

    /// Exhausts the `Classifier` writing the output into `sink`.
    ///
    /// The general structure for this method is to iterate over each token,
    /// possibly giving it an HTML span with a class specifying what flavor of
    /// token is used.
    fn highlight(mut self, sink: &mut dyn FnMut(Highlight<'a>)) {
        with_default_session_globals(|| {
            loop {
                if self
                    .tokens
                    .peek()
                    .map(|t| matches!(t.0, TokenKind::Colon | TokenKind::Ident))
                    .unwrap_or(false)
                {
                    let tokens = self.get_full_ident_path();
                    for (token, start, end) in tokens {
                        let text = &self.src[start..end];
                        self.advance(token, text, sink);
                        self.byte_pos += text.len() as u32;
                    }
                }
                if let Some((token, text)) = self.next() {
                    self.advance(token, text, sink);
                } else {
                    break;
                }
            }
        })
    }

    /// Single step of highlighting. This will classify `token`, but maybe also
    /// a couple of following ones as well.
    fn advance(&mut self, token: TokenKind, text: &'a str, sink: &mut dyn FnMut(Highlight<'a>)) {
        let lookahead = self.peek();
        let no_highlight = |sink: &mut dyn FnMut(_)| sink(Highlight::Token { text, class: None });
        let class = match token {
            TokenKind::Whitespace => return no_highlight(sink),
            TokenKind::LineComment { doc_style } | TokenKind::BlockComment { doc_style, .. } => {
                if doc_style.is_some() {
                    Class::DocComment
                } else {
                    Class::Comment
                }
            }
            // Consider this as part of a macro invocation if there was a
            // leading identifier.
            TokenKind::Bang if self.in_macro => {
                self.in_macro = false;
                sink(Highlight::Token { text, class: None });
                sink(Highlight::ExitSpan);
                return;
            }

            // Assume that '&' or '*' is the reference or dereference operator
            // or a reference or pointer type. Unless, of course, it looks like
            // a logical and or a multiplication operator: `&&` or `* `.
            TokenKind::Star => match lookahead {
                Some(TokenKind::Whitespace) => Class::Op,
                _ => Class::RefKeyWord,
            },
            TokenKind::And => match lookahead {
                Some(TokenKind::And) => {
                    self.next();
                    sink(Highlight::Token { text: "&&", class: Some(Class::Op) });
                    return;
                }
                Some(TokenKind::Eq) => {
                    self.next();
                    sink(Highlight::Token { text: "&=", class: Some(Class::Op) });
                    return;
                }
                Some(TokenKind::Whitespace) => Class::Op,
                _ => Class::RefKeyWord,
            },

            // Operators.
            TokenKind::Minus
            | TokenKind::Plus
            | TokenKind::Or
            | TokenKind::Slash
            | TokenKind::Caret
            | TokenKind::Percent
            | TokenKind::Bang
            | TokenKind::Eq
            | TokenKind::Lt
            | TokenKind::Gt => Class::Op,

            // Miscellaneous, no highlighting.
            TokenKind::Dot
            | TokenKind::Semi
            | TokenKind::Comma
            | TokenKind::OpenParen
            | TokenKind::CloseParen
            | TokenKind::OpenBrace
            | TokenKind::CloseBrace
            | TokenKind::OpenBracket
            | TokenKind::At
            | TokenKind::Tilde
            | TokenKind::Colon
            | TokenKind::Unknown => return no_highlight(sink),

            TokenKind::Question => Class::QuestionMark,

            TokenKind::Dollar => match lookahead {
                Some(TokenKind::Ident) => {
                    self.in_macro_nonterminal = true;
                    Class::MacroNonTerminal
                }
                _ => return no_highlight(sink),
            },

            // This might be the start of an attribute. We're going to want to
            // continue highlighting it as an attribute until the ending ']' is
            // seen, so skip out early. Down below we terminate the attribute
            // span when we see the ']'.
            TokenKind::Pound => {
                match lookahead {
                    // Case 1: #![inner_attribute]
                    Some(TokenKind::Bang) => {
                        self.next();
                        if let Some(TokenKind::OpenBracket) = self.peek() {
                            self.in_attribute = true;
                            sink(Highlight::EnterSpan { class: Class::Attribute });
                        }
                        sink(Highlight::Token { text: "#", class: None });
                        sink(Highlight::Token { text: "!", class: None });
                        return;
                    }
                    // Case 2: #[outer_attribute]
                    Some(TokenKind::OpenBracket) => {
                        self.in_attribute = true;
                        sink(Highlight::EnterSpan { class: Class::Attribute });
                    }
                    _ => (),
                }
                return no_highlight(sink);
            }
            TokenKind::CloseBracket => {
                if self.in_attribute {
                    self.in_attribute = false;
                    sink(Highlight::Token { text: "]", class: None });
                    sink(Highlight::ExitSpan);
                    return;
                }
                return no_highlight(sink);
            }
            TokenKind::Literal { kind, .. } => match kind {
                // Text literals.
                LiteralKind::Byte { .. }
                | LiteralKind::Char { .. }
                | LiteralKind::Str { .. }
                | LiteralKind::ByteStr { .. }
                | LiteralKind::RawStr { .. }
                | LiteralKind::RawByteStr { .. } => Class::String,
                // Number literals.
                LiteralKind::Float { .. } | LiteralKind::Int { .. } => Class::Number,
            },
            TokenKind::Ident | TokenKind::RawIdent if lookahead == Some(TokenKind::Bang) => {
                self.in_macro = true;
                sink(Highlight::EnterSpan { class: Class::Macro });
                sink(Highlight::Token { text, class: None });
                return;
            }
            TokenKind::Ident => match get_real_ident_class(text, self.edition) {
                Class::Ident => match text {
                    "Option" | "Result" => Class::PreludeTy,
                    "Some" | "None" | "Ok" | "Err" => Class::PreludeVal,
                    _ if self.in_macro_nonterminal => {
                        self.in_macro_nonterminal = false;
                        Class::MacroNonTerminal
                    }
                    _ => Class::Ident,
                },
                c => c,
            },
            TokenKind::RawIdent => Class::Ident,
            TokenKind::Lifetime { .. } => Class::Lifetime,
        };
        // Anything that didn't return above is the simple case where we the
        // class just spans a single token, so we can use the `string` method.
        sink(Highlight::Token { text, class: Some(class) });
    }

    fn peek(&mut self) -> Option<TokenKind> {
        self.tokens.peek().map(|(toke_kind, _text)| *toke_kind)
    }
}

/// Called when we start processing a span of text that should be highlighted.
/// The `Class` argument specifies how it should be highlighted.
fn enter_span(out: &mut Buffer, klass: Class) {
    write!(out, "", klass.as_html());
}

/// Called at the end of a span of highlighted text.
fn exit_span(out: &mut Buffer) {
    out.write_str("");
}

/// Called for a span of text. If the text should be highlighted differently
/// from the surrounding text, then the `Class` argument will be a value other
/// than `None`.
///
/// The following sequences of callbacks are equivalent:
/// ```plain
///     enter_span(Foo), string("text", None), exit_span()
///     string("text", Foo)
/// ```
/// The latter can be thought of as a shorthand for the former, which is more
/// flexible.
fn string<T: Display>(out: &mut Buffer, text: T, klass: Option<Class>) {
    match klass {
        None => write!(out, "{}", text),
        Some(klass) => write!(out, "{}", klass.as_html(), text),
    }
}

#[cfg(test)]
mod tests;
<commit_msg>Remove unneeded call to with_default_session_globals in rustdoc highlight<commit_after>//! Basic syntax highlighting functionality.
//!
//! This module uses librustc_ast's lexer to provide token-based highlighting for
//! the HTML documentation generated by rustdoc.
//!
//! Use the `render_with_highlighting` to highlight some rust code.

use crate::html::escape::Escape;

use std::fmt::Display;
use std::iter::Peekable;

use rustc_lexer::{LiteralKind, TokenKind};
use rustc_span::edition::Edition;
use rustc_span::symbol::Symbol;

use super::format::Buffer;

/// Highlights `src`, returning the HTML output.
crate fn render_with_highlighting(
    src: &str,
    out: &mut Buffer,
    class: Option<&str>,
    playground_button: Option<&str>,
    tooltip: Option<(Option<Edition>, &str)>,
    edition: Edition,
) {
    debug!("highlighting: ================\n{}\n==============", src);
    if let Some((edition_info, class)) = tooltip {
        write!(
            out,
            "

ⓘ

",
            class,
            if let Some(edition_info) = edition_info {
                format!(" data-edition=\"{}\"", edition_info)
            } else {
                String::new()
            },
        );
    }

    write_header(out, class);
    write_code(out, &src, edition);
    write_footer(out, playground_button);
}

fn write_header(out: &mut Buffer, class: Option<&str>) {
    writeln!(out, "

", class.unwrap_or_default());
}

fn write_code(out: &mut Buffer, src: &str, edition: Edition) {
    // This replace allows to fix how the code source with DOS backline characters is displayed.
    let src = src.replace("\r\n", "\n");
    Classifier::new(&src, edition).highlight(&mut |highlight| {
        match highlight {
            Highlight::Token { text, class } => string(out, Escape(text), class),
            Highlight::EnterSpan { class } => enter_span(out, class),
            Highlight::ExitSpan => exit_span(out),
        };
    });
}

fn write_footer(out: &mut Buffer, playground_button: Option<&str>) {
    writeln!(out, "
{}
", playground_button.unwrap_or_default());
}

/// How a span of text is classified. Mostly corresponds to token kinds.
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
enum Class {
    Comment,
    DocComment,
    Attribute,
    KeyWord,
    // Keywords that do pointer/reference stuff.
    RefKeyWord,
    Self_,
    Op,
    Macro,
    MacroNonTerminal,
    String,
    Number,
    Bool,
    Ident,
    Lifetime,
    PreludeTy,
    PreludeVal,
    QuestionMark,
}

impl Class {
    /// Returns the css class expected by rustdoc for each `Class`.
    fn as_html(self) -> &'static str {
        match self {
            Class::Comment => "comment",
            Class::DocComment => "doccomment",
            Class::Attribute => "attribute",
            Class::KeyWord => "kw",
            Class::RefKeyWord => "kw-2",
            Class::Self_ => "self",
            Class::Op => "op",
            Class::Macro => "macro",
            Class::MacroNonTerminal => "macro-nonterminal",
            Class::String => "string",
            Class::Number => "number",
            Class::Bool => "bool-val",
            Class::Ident => "ident",
            Class::Lifetime => "lifetime",
            Class::PreludeTy => "prelude-ty",
            Class::PreludeVal => "prelude-val",
            Class::QuestionMark => "question-mark",
        }
    }
}

enum Highlight<'a> {
    Token { text: &'a str, class: Option<Class> },
    EnterSpan { class: Class },
    ExitSpan,
}

struct TokenIter<'a> {
    src: &'a str,
}

impl Iterator for TokenIter<'a> {
    type Item = (TokenKind, &'a str);
    fn next(&mut self) -> Option<(TokenKind, &'a str)> {
        if self.src.is_empty() {
            return None;
        }
        let token = rustc_lexer::first_token(self.src);
        let (text, rest) = self.src.split_at(token.len);
        self.src = rest;
        Some((token.kind, text))
    }
}

fn get_real_ident_class(text: &str, edition: Edition) -> Class {
    match text {
        "ref" | "mut" => Class::RefKeyWord,
        "self" | "Self" => Class::Self_,
        "false" | "true" => Class::Bool,
        _ if Symbol::intern(text).is_reserved(|| edition) => Class::KeyWord,
        _ => Class::Ident,
    }
}

/// Processes program tokens, classifying strings of text by highlighting
/// category (`Class`).
struct Classifier<'a> {
    tokens: Peekable<TokenIter<'a>>,
    in_attribute: bool,
    in_macro: bool,
    in_macro_nonterminal: bool,
    edition: Edition,
    byte_pos: u32,
    src: &'a str,
}

impl<'a> Classifier<'a> {
    fn new(src: &str, edition: Edition) -> Classifier<'_> {
        let tokens = TokenIter { src }.peekable();
        Classifier {
            tokens,
            in_attribute: false,
            in_macro: false,
            in_macro_nonterminal: false,
            edition,
            byte_pos: 0,
            src,
        }
    }

    /// Concatenate colons and idents as one when possible.
    fn get_full_ident_path(&mut self) -> Vec<(TokenKind, usize, usize)> {
        let start = self.byte_pos as usize;
        let mut pos = start;
        let mut has_ident = false;
        let edition = self.edition;

        loop {
            let mut nb = 0;
            while let Some((TokenKind::Colon, _)) = self.tokens.peek() {
                self.tokens.next();
                nb += 1;
            }
            // Ident path can start with "::" but if we already have content in the ident path,
            // the "::" is mandatory.
            if has_ident && nb == 0 {
                return vec![(TokenKind::Ident, start, pos)];
            } else if nb != 0 && nb != 2 {
                if has_ident {
                    return vec![(TokenKind::Ident, start, pos), (TokenKind::Colon, pos, pos + nb)];
                } else {
                    return vec![(TokenKind::Colon, pos, pos + nb)];
                }
            }

            if let Some((Class::Ident, text)) = self.tokens.peek().map(|(token, text)| {
                if *token == TokenKind::Ident {
                    let class = get_real_ident_class(text, edition);
                    (class, text)
                } else {
                    // Doesn't matter which Class we put in here...
                    (Class::Comment, text)
                }
            }) {
                // We only "add" the colon if there is an ident behind.
                pos += text.len() + nb;
                has_ident = true;
                self.tokens.next();
            } else if nb > 0 && has_ident {
                return vec![(TokenKind::Ident, start, pos), (TokenKind::Colon, pos, pos + nb)];
            } else if nb > 0 {
                return vec![(TokenKind::Colon, pos, pos + nb)];
            } else if has_ident {
                return vec![(TokenKind::Ident, start, pos)];
            } else {
                return Vec::new();
            }
        }
    }

    /// Wraps the tokens iteration to ensure that the byte_pos is always correct.
    fn next(&mut self) -> Option<(TokenKind, &'a str)> {
        if let Some((kind, text)) = self.tokens.next() {
            self.byte_pos += text.len() as u32;
            Some((kind, text))
        } else {
            None
        }
    }

    /// Exhausts the `Classifier` writing the output into `sink`.
    ///
    /// The general structure for this method is to iterate over each token,
    /// possibly giving it an HTML span with a class specifying what flavor of
    /// token is used.
    fn highlight(mut self, sink: &mut dyn FnMut(Highlight<'a>)) {
        loop {
            if self
                .tokens
                .peek()
                .map(|t| matches!(t.0, TokenKind::Colon | TokenKind::Ident))
                .unwrap_or(false)
            {
                let tokens = self.get_full_ident_path();
                for (token, start, end) in tokens {
                    let text = &self.src[start..end];
                    self.advance(token, text, sink);
                    self.byte_pos += text.len() as u32;
                }
            }
            if let Some((token, text)) = self.next() {
                self.advance(token, text, sink);
            } else {
                break;
            }
        }
    }

    /// Single step of highlighting. This will classify `token`, but maybe also
    /// a couple of following ones as well.
    fn advance(&mut self, token: TokenKind, text: &'a str, sink: &mut dyn FnMut(Highlight<'a>)) {
        let lookahead = self.peek();
        let no_highlight = |sink: &mut dyn FnMut(_)| sink(Highlight::Token { text, class: None });
        let class = match token {
            TokenKind::Whitespace => return no_highlight(sink),
            TokenKind::LineComment { doc_style } | TokenKind::BlockComment { doc_style, .. } => {
                if doc_style.is_some() {
                    Class::DocComment
                } else {
                    Class::Comment
                }
            }
            // Consider this as part of a macro invocation if there was a
            // leading identifier.
            TokenKind::Bang if self.in_macro => {
                self.in_macro = false;
                sink(Highlight::Token { text, class: None });
                sink(Highlight::ExitSpan);
                return;
            }

            // Assume that '&' or '*' is the reference or dereference operator
            // or a reference or pointer type. Unless, of course, it looks like
            // a logical and or a multiplication operator: `&&` or `* `.
            TokenKind::Star => match lookahead {
                Some(TokenKind::Whitespace) => Class::Op,
                _ => Class::RefKeyWord,
            },
            TokenKind::And => match lookahead {
                Some(TokenKind::And) => {
                    self.next();
                    sink(Highlight::Token { text: "&&", class: Some(Class::Op) });
                    return;
                }
                Some(TokenKind::Eq) => {
                    self.next();
                    sink(Highlight::Token { text: "&=", class: Some(Class::Op) });
                    return;
                }
                Some(TokenKind::Whitespace) => Class::Op,
                _ => Class::RefKeyWord,
            },

            // Operators.
            TokenKind::Minus
            | TokenKind::Plus
            | TokenKind::Or
            | TokenKind::Slash
            | TokenKind::Caret
            | TokenKind::Percent
            | TokenKind::Bang
            | TokenKind::Eq
            | TokenKind::Lt
            | TokenKind::Gt => Class::Op,

            // Miscellaneous, no highlighting.
            TokenKind::Dot
            | TokenKind::Semi
            | TokenKind::Comma
            | TokenKind::OpenParen
            | TokenKind::CloseParen
            | TokenKind::OpenBrace
            | TokenKind::CloseBrace
            | TokenKind::OpenBracket
            | TokenKind::At
            | TokenKind::Tilde
            | TokenKind::Colon
            | TokenKind::Unknown => return no_highlight(sink),

            TokenKind::Question => Class::QuestionMark,

            TokenKind::Dollar => match lookahead {
                Some(TokenKind::Ident) => {
                    self.in_macro_nonterminal = true;
                    Class::MacroNonTerminal
                }
                _ => return no_highlight(sink),
            },

            // This might be the start of an attribute. We're going to want to
            // continue highlighting it as an attribute until the ending ']' is
            // seen, so skip out early. Down below we terminate the attribute
            // span when we see the ']'.
            TokenKind::Pound => {
                match lookahead {
                    // Case 1: #![inner_attribute]
                    Some(TokenKind::Bang) => {
                        self.next();
                        if let Some(TokenKind::OpenBracket) = self.peek() {
                            self.in_attribute = true;
                            sink(Highlight::EnterSpan { class: Class::Attribute });
                        }
                        sink(Highlight::Token { text: "#", class: None });
                        sink(Highlight::Token { text: "!", class: None });
                        return;
                    }
                    // Case 2: #[outer_attribute]
                    Some(TokenKind::OpenBracket) => {
                        self.in_attribute = true;
                        sink(Highlight::EnterSpan { class: Class::Attribute });
                    }
                    _ => (),
                }
                return no_highlight(sink);
            }
            TokenKind::CloseBracket => {
                if self.in_attribute {
                    self.in_attribute = false;
                    sink(Highlight::Token { text: "]", class: None });
                    sink(Highlight::ExitSpan);
                    return;
                }
                return no_highlight(sink);
            }
            TokenKind::Literal { kind, .. } => match kind {
                // Text literals.
                LiteralKind::Byte { .. }
                | LiteralKind::Char { .. }
                | LiteralKind::Str { .. }
                | LiteralKind::ByteStr { .. }
                | LiteralKind::RawStr { .. }
                | LiteralKind::RawByteStr { .. } => Class::String,
                // Number literals.
                LiteralKind::Float { .. } | LiteralKind::Int { .. } => Class::Number,
            },
            TokenKind::Ident | TokenKind::RawIdent if lookahead == Some(TokenKind::Bang) => {
                self.in_macro = true;
                sink(Highlight::EnterSpan { class: Class::Macro });
                sink(Highlight::Token { text, class: None });
                return;
            }
            TokenKind::Ident => match get_real_ident_class(text, self.edition) {
                Class::Ident => match text {
                    "Option" | "Result" => Class::PreludeTy,
                    "Some" | "None" | "Ok" | "Err" => Class::PreludeVal,
                    _ if self.in_macro_nonterminal => {
                        self.in_macro_nonterminal = false;
                        Class::MacroNonTerminal
                    }
                    _ => Class::Ident,
                },
                c => c,
            },
            TokenKind::RawIdent => Class::Ident,
            TokenKind::Lifetime { .. } => Class::Lifetime,
        };
        // Anything that didn't return above is the simple case where we the
        // class just spans a single token, so we can use the `string` method.
        sink(Highlight::Token { text, class: Some(class) });
    }

    fn peek(&mut self) -> Option<TokenKind> {
        self.tokens.peek().map(|(toke_kind, _text)| *toke_kind)
    }
}

/// Called when we start processing a span of text that should be highlighted.
/// The `Class` argument specifies how it should be highlighted.
fn enter_span(out: &mut Buffer, klass: Class) {
    write!(out, "", klass.as_html());
}

/// Called at the end of a span of highlighted text.
fn exit_span(out: &mut Buffer) {
    out.write_str("");
}

/// Called for a span of text. If the text should be highlighted differently
/// from the surrounding text, then the `Class` argument will be a value other
/// than `None`.
///
/// The following sequences of callbacks are equivalent:
/// ```plain
///     enter_span(Foo), string("text", None), exit_span()
///     string("text", Foo)
/// ```
/// The latter can be thought of as a shorthand for the former, which is more
/// flexible.
fn string<T: Display>(out: &mut Buffer, text: T, klass: Option<Class>) {
    match klass {
        None => write!(out, "{}", text),
        Some(klass) => write!(out, "{}", klass.as_html(), text),
    }
}

#[cfg(test)]
mod tests;
<|endoftext|>