<commit_before>use {Distribution, Generator};

/// A categorical distribution.
#[derive(Clone)]
pub struct Categorical {
    /// The size of the probabilty vector.
    pub k: usize,
    /// The probability vector.
    pub p: Vec<f64>,
}

impl Categorical {
    /// Create a categorical distribution with success probability `p`.
    ///
    /// # Panics
    ///
    /// Panics if `p[i] < 0` or `p[i] > 1` for any `i = 0..k`, or if
    /// `sum(p) != 1`.
    #[inline]
    pub fn new(p: &[f64]) -> Categorical {
        debug_assert!({
            fn is_prob_vec(p: &[f64]) -> bool {
                // Check if p is a probability vector.
                let mut sum = 0.;
                for &p_i in p.iter() {
                    if p_i < 0.|| p_i > 1. { return false; }
                    sum += p_i;
                }
                (sum - 1.).abs() <= 1e-12
            }
            is_prob_vec(&p)
        }, "Categorical::new() is called with p not a probabilty vector");

        Categorical { k: p.len(),  p: p.to_vec()}
    }
}

impl Distribution for Categorical {
    type Value = usize;

    #[inline]
    fn mean(&self) -> f64 {
        // sum_{i=0}^k i p_i
        self.p.iter().enumerate().fold(0., |acc, b| acc + b.0 as f64 * b.1)
    }

    #[inline]
    fn var(&self) -> f64 {
        let mu = self.mean();
        self.p.iter().enumerate().fold(
            0., |acc, b| acc + (b.0 as f64 - mu).powi(2) * b.1
        )
    }

    #[inline]
    fn skewness(&self) -> f64 {
        let mu = self.mean();
        let sigma2 = self.var();
        let skew = self.p.iter().enumerate().fold(
            0., |acc, b| acc + (b.0 as f64 - mu).powi(3) * b.1
        );
        skew / (sigma2 * sigma2.sqrt())
    }

    fn kurtosis(&self) -> f64 {
        let mu = self.mean();
        let sigma2 = self.var();
        let kurt = self.p.iter().enumerate().fold(
            0., |acc, b| acc + (b.0 as f64 - mu).powi(4) * b.1
        );
        kurt / sigma2.powi(2) - 3.
    }

    fn median(&self) -> f64 {
        if self.p[0] > 0.5 { return 0.0; }
        else if self.p[0] == 0.5 { return 0.5; }
        let mut sum = 0.;
        for i in 0..self.k {
            sum += self.p[i];
            if sum == 0.5 { return (2 * i - 1) as f64 / 2.; }
            else if sum > 0.5 {
                return i as f64;
            }
        }
        unreachable!()
    }

    fn modes(&self) -> Vec<Self::Value> {
        let mut m = Vec::new();
        let mut max = 0.;
        for (i, &p_i) in self.p.iter().enumerate() {
            if p_i == max { m.push(i); }
            if p_i > max {
                max = p_i;
                m = vec![i];
            }
        }
        m
    }

    #[inline]
    fn entropy(&self) -> f64 {
        - self.p.iter().fold(0., |acc, p_i| acc + p_i * p_i.ln())
    }

    #[inline]
    fn cdf(&self, x: Self::Value) -> f64 {
        if x >= self.k - 1 { 1. }
        else {
            self.p.iter().take(x + 1).fold(0., |a, b| a + b)
        }
    }

    fn inv_cdf(&self, p: f64) -> Self::Value {
        debug_assert!(0.0 <= p && p <= 1.0, "inv_cdf is called with p outside of [0, 1]");
        if p == 0. {
            // return the first non-zero index
            return self.p.iter().enumerate().find(|&v| *v.1 > 0.).unwrap().0;
        }
        let mut sum = 0.;
        for i in 0..self.k {
            sum += self.p[i];
            if sum >= p || sum == 1. { return i; }
        }
        self.k - 1
    }

    #[inline]
    fn pdf(&self, x: Self::Value) -> f64 {
        self.p[x]
    }

    #[inline]
    fn sample<G: Generator>(&self, generator: &mut G) -> Self::Value {
        self.inv_cdf(generator.gen::<f64>())
    }
}

#[cfg(test)]
mod tests {
    use assert;

    use {Distribution, Sampler};
    use distributions::Categorical;

    macro_rules! new(
        (equal $k:expr) => { Categorical::new(&[1./$k as f64; $k]) };
        ($p:expr) => { Categorical::new(&$p); }
    );

    #[test]
    #[should_panic]
    fn invalid_prob_vec_1() {
        new!([0.6, 0.6, 0.6]);
    }

    #[test]
    #[should_panic]
    fn invalid_prob_vec_2() {
        new!([0.6, 0.6, 0.6, -0.2]);
    }

    #[test]
    #[should_panic]
    fn invalid_prob_vec_3() {
        new!([1.2, -0.2]);
    }

    #[test]
    fn mean() {
        assert_eq!(new!(equal 3).mean(), 1.);
        assert_eq!(new!([0.3, 0.3, 0.4]).mean(), 1.1);
        assert_eq!(new!([1./6., 1./3., 1./3., 1./6.]).mean(), 1.5);
    }

    #[test]
    fn var() {
        assert_eq!(new!(equal 3).var(), 2./3.);
        assert_eq!(new!([1./6., 1./3., 1./3., 1./6.]).var(), 11./12.);
    }

    #[test]
    fn sd() {
        assert_eq!(new!(equal 2).sd(), 0.5);
        assert_eq!(new!([1./6., 1./3., 1./3., 1./6.]).sd(), 0.9574271077563381);
    }

    #[test]
    fn skewness() {
        assert_eq!(new!(equal 6).skewness(), 0.0);
        assert_eq!(new!([1./6., 1./3., 1./3., 1./6.]).skewness(), 0.);
        assert_eq!(new!([0.1, 0.2, 0.3, 0.4]).skewness(), -0.6);
    }

    #[test]
    fn kurtosis() {
        assert_eq!(new!(equal 2).kurtosis(), -2.0);
        assert_eq!(new!([0.1, 0.2, 0.3, 0.4]).kurtosis(), -0.7999999999999998);
    }

    #[test]
    fn median() {
        assert_eq!(new!([0.6, 0.2, 0.2]).median(), 0.0);
        assert_eq!(new!(equal 2).median(), 0.5);
        assert_eq!(new!([0.1, 0.2, 0.3, 0.4]).median(), 2.0);
        assert_eq!(new!([1./6., 1./3., 1./3., 1./6.]).median(), 0.5);
    }

    #[test]
    fn modes() {
        assert_eq!(new!([0.6, 0.2, 0.2]).modes(), vec![0]);
        assert_eq!(new!(equal 2).modes(), vec![0, 1]);
        assert_eq!(new!(equal 3).modes(), vec![0, 1, 2]);
        assert_eq!(new!([0.4, 0.2, 0.4]).modes(), vec![0, 2]);
        assert_eq!(new!([1./6., 1./3., 1./3., 1./6.]).modes(), vec![1, 2]);
    }

    #[test]
    fn entropy() {
        use std::f64::consts::LN_2;
        assert_eq!(new!(equal 2).entropy(), LN_2);
        assert_eq!(new!([0.1, 0.2, 0.3, 0.4]).entropy(), 1.2798542258336676);
    }

    #[test]
    fn pdf() {
        let p = [0.0, 0.75, 0.25, 0.0];
        let d1 = new!(p);
        assert::equal(&(0..4).map(|x| d1.pdf(x)).collect::<Vec<_>>(), &p.to_vec());

        let d2 = new!(equal 3);
        assert::equal(&(0..3).map(|x| d2.pdf(x)).collect::<Vec<_>>(), &vec![1./3.; 3])
    }

    #[test]
    fn cdf() {
        let d1 = new!([0.0, 0.75, 0.25, 0.0]);
        assert::equal(&(0..4).map(|x| d1.cdf(x)).collect::<Vec<_>>(),
                      &vec![0.0, 0.75, 1.0, 1.0]);

        let d2 = new!(equal 3);
        assert::equal(&(0..3).map(|x| d2.cdf(x)).collect::<Vec<_>>(),
                      &vec![1./3., 2./3., 1.]);
    }

    #[test]
    fn inv_cdf() {
        let d1 = new!([0.0, 0.75, 0.25, 0.0]);
        let p1 = vec![0.0, 0.75, 0.7500001, 1.0];
        assert::equal(&p1.iter().map(|&p| d1.inv_cdf(p)).collect::<Vec<_>>(),
                      &vec![1, 1, 2, 2]);

        let d2 = new!(equal 3);
        let p2 = vec![0.0, 0.5, 0.75, 1.0];
        assert::equal(&p2.iter().map(|&p| d2.inv_cdf(p)).collect::<Vec<_>>(),
                      &vec![0, 1, 2, 2]);

    }

    #[test]
    fn sample() {
        let mut generator = ::generator();

        // Discrete Uniform(1, 2)
        let sum = Sampler(&new!([0.0, 0.5, 0.5]), &mut generator)
            .take(100)
            .fold(0, |a, b| a + b);

        assert!(100 <= sum && sum <= 200);
    }
}
<commit_msg>categorical: make use of destructuring<commit_after>use {Distribution, Generator};

/// A categorical distribution.
#[derive(Clone)]
pub struct Categorical {
    /// The size of the probabilty vector.
    pub k: usize,
    /// The probability vector.
    pub p: Vec<f64>,
}

impl Categorical {
    /// Create a categorical distribution with success probability `p`.
    ///
    /// # Panics
    ///
    /// Panics if `p[i] < 0` or `p[i] > 1` for any `i = 0..k`, or if
    /// `sum(p) != 1`.
    #[inline]
    pub fn new(p: &[f64]) -> Categorical {
        debug_assert!({
            fn is_prob_vec(p: &[f64]) -> bool {
                // Check if p is a probability vector.
                let mut sum = 0.;
                for &p_i in p.iter() {
                    if p_i < 0.|| p_i > 1. { return false; }
                    sum += p_i;
                }
                (sum - 1.).abs() <= 1e-12
            }
            is_prob_vec(&p)
        }, "Categorical::new() is called with p not a probabilty vector");

        Categorical { k: p.len(),  p: p.to_vec()}
    }
}

impl Distribution for Categorical {
    type Value = usize;

    #[inline]
    fn mean(&self) -> f64 {
        self.p.iter().enumerate().fold(0., |sum, (i, p)| sum + i as f64 * p)
    }

    #[inline]
    fn var(&self) -> f64 {
        let mu = self.mean();
        self.p.iter().enumerate().fold(0., |sum, (i, p)| sum + (i as f64 - mu).powi(2) * p)
    }

    #[inline]
    fn skewness(&self) -> f64 {
        let mu = self.mean();
        let sigma2 = self.var();
        let skew = self.p.iter().enumerate()
                         .fold(0., |sum, (i, p)| sum + (i as f64 - mu).powi(3) * p);
        skew / (sigma2 * sigma2.sqrt())
    }

    fn kurtosis(&self) -> f64 {
        let mu = self.mean();
        let sigma2 = self.var();
        let kurt = self.p.iter().enumerate()
                         .fold(0., |sum, (i, p)| sum + (i as f64 - mu).powi(4) * p);
        kurt / sigma2.powi(2) - 3.
    }

    fn median(&self) -> f64 {
        if self.p[0] > 0.5 { return 0.0; }
        else if self.p[0] == 0.5 { return 0.5; }
        let mut sum = 0.;
        for i in 0..self.k {
            sum += self.p[i];
            if sum == 0.5 { return (2 * i - 1) as f64 / 2.; }
            else if sum > 0.5 {
                return i as f64;
            }
        }
        unreachable!()
    }

    fn modes(&self) -> Vec<Self::Value> {
        let mut m = Vec::new();
        let mut max = 0.;
        for (i, &p_i) in self.p.iter().enumerate() {
            if p_i == max { m.push(i); }
            if p_i > max {
                max = p_i;
                m = vec![i];
            }
        }
        m
    }

    #[inline]
    fn entropy(&self) -> f64 {
        - self.p.iter().fold(0., |acc, p_i| acc + p_i * p_i.ln())
    }

    #[inline]
    fn cdf(&self, x: Self::Value) -> f64 {
        if x >= self.k - 1 { 1. }
        else {
            self.p.iter().take(x + 1).fold(0., |a, b| a + b)
        }
    }

    fn inv_cdf(&self, p: f64) -> Self::Value {
        debug_assert!(0.0 <= p && p <= 1.0, "inv_cdf is called with p outside of [0, 1]");
        if p == 0. {
            // return the first non-zero index
            return self.p.iter().enumerate().find(|&(_, &p)| p > 0.).unwrap().0;
        }
        let mut sum = 0.;
        for i in 0..self.k {
            sum += self.p[i];
            if sum >= p || sum == 1. { return i; }
        }
        self.k - 1
    }

    #[inline]
    fn pdf(&self, x: Self::Value) -> f64 {
        self.p[x]
    }

    #[inline]
    fn sample<G: Generator>(&self, generator: &mut G) -> Self::Value {
        self.inv_cdf(generator.gen::<f64>())
    }
}

#[cfg(test)]
mod tests {
    use assert;

    use {Distribution, Sampler};
    use distributions::Categorical;

    macro_rules! new(
        (equal $k:expr) => { Categorical::new(&[1./$k as f64; $k]) };
        ($p:expr) => { Categorical::new(&$p); }
    );

    #[test]
    #[should_panic]
    fn invalid_prob_vec_1() {
        new!([0.6, 0.6, 0.6]);
    }

    #[test]
    #[should_panic]
    fn invalid_prob_vec_2() {
        new!([0.6, 0.6, 0.6, -0.2]);
    }

    #[test]
    #[should_panic]
    fn invalid_prob_vec_3() {
        new!([1.2, -0.2]);
    }

    #[test]
    fn mean() {
        assert_eq!(new!(equal 3).mean(), 1.);
        assert_eq!(new!([0.3, 0.3, 0.4]).mean(), 1.1);
        assert_eq!(new!([1./6., 1./3., 1./3., 1./6.]).mean(), 1.5);
    }

    #[test]
    fn var() {
        assert_eq!(new!(equal 3).var(), 2./3.);
        assert_eq!(new!([1./6., 1./3., 1./3., 1./6.]).var(), 11./12.);
    }

    #[test]
    fn sd() {
        assert_eq!(new!(equal 2).sd(), 0.5);
        assert_eq!(new!([1./6., 1./3., 1./3., 1./6.]).sd(), 0.9574271077563381);
    }

    #[test]
    fn skewness() {
        assert_eq!(new!(equal 6).skewness(), 0.0);
        assert_eq!(new!([1./6., 1./3., 1./3., 1./6.]).skewness(), 0.);
        assert_eq!(new!([0.1, 0.2, 0.3, 0.4]).skewness(), -0.6);
    }

    #[test]
    fn kurtosis() {
        assert_eq!(new!(equal 2).kurtosis(), -2.0);
        assert_eq!(new!([0.1, 0.2, 0.3, 0.4]).kurtosis(), -0.7999999999999998);
    }

    #[test]
    fn median() {
        assert_eq!(new!([0.6, 0.2, 0.2]).median(), 0.0);
        assert_eq!(new!(equal 2).median(), 0.5);
        assert_eq!(new!([0.1, 0.2, 0.3, 0.4]).median(), 2.0);
        assert_eq!(new!([1./6., 1./3., 1./3., 1./6.]).median(), 0.5);
    }

    #[test]
    fn modes() {
        assert_eq!(new!([0.6, 0.2, 0.2]).modes(), vec![0]);
        assert_eq!(new!(equal 2).modes(), vec![0, 1]);
        assert_eq!(new!(equal 3).modes(), vec![0, 1, 2]);
        assert_eq!(new!([0.4, 0.2, 0.4]).modes(), vec![0, 2]);
        assert_eq!(new!([1./6., 1./3., 1./3., 1./6.]).modes(), vec![1, 2]);
    }

    #[test]
    fn entropy() {
        use std::f64::consts::LN_2;
        assert_eq!(new!(equal 2).entropy(), LN_2);
        assert_eq!(new!([0.1, 0.2, 0.3, 0.4]).entropy(), 1.2798542258336676);
    }

    #[test]
    fn pdf() {
        let p = [0.0, 0.75, 0.25, 0.0];
        let d1 = new!(p);
        assert::equal(&(0..4).map(|x| d1.pdf(x)).collect::<Vec<_>>(), &p.to_vec());

        let d2 = new!(equal 3);
        assert::equal(&(0..3).map(|x| d2.pdf(x)).collect::<Vec<_>>(), &vec![1./3.; 3])
    }

    #[test]
    fn cdf() {
        let d1 = new!([0.0, 0.75, 0.25, 0.0]);
        assert::equal(&(0..4).map(|x| d1.cdf(x)).collect::<Vec<_>>(),
                      &vec![0.0, 0.75, 1.0, 1.0]);

        let d2 = new!(equal 3);
        assert::equal(&(0..3).map(|x| d2.cdf(x)).collect::<Vec<_>>(),
                      &vec![1./3., 2./3., 1.]);
    }

    #[test]
    fn inv_cdf() {
        let d1 = new!([0.0, 0.75, 0.25, 0.0]);
        let p1 = vec![0.0, 0.75, 0.7500001, 1.0];
        assert::equal(&p1.iter().map(|&p| d1.inv_cdf(p)).collect::<Vec<_>>(),
                      &vec![1, 1, 2, 2]);

        let d2 = new!(equal 3);
        let p2 = vec![0.0, 0.5, 0.75, 1.0];
        assert::equal(&p2.iter().map(|&p| d2.inv_cdf(p)).collect::<Vec<_>>(),
                      &vec![0, 1, 2, 2]);

    }

    #[test]
    fn sample() {
        let mut generator = ::generator();

        // Discrete Uniform(1, 2)
        let sum = Sampler(&new!([0.0, 0.5, 0.5]), &mut generator)
            .take(100)
            .fold(0, |a, b| a + b);

        assert!(100 <= sum && sum <= 200);
    }
}
<|endoftext|>


<commit_before>extern crate cfg_if;

use cfg_if::cfg_if;

cfg_if! {
    if #[cfg(all(windows, feature = "neon-sys"))] {
        use std::env::var;
        use std::path::Path;

        const NODE_ROOT_DIR: &'static str = include_str!(concat!(env!("OUT_DIR"), "\\node_root_dir"));
        const NODE_ARCH:     &'static str = include_str!(concat!(env!("OUT_DIR"), "\\node_arch"));
        const NODE_LIB_FILE: &'static str = include_str!(concat!(env!("OUT_DIR"), "\\node_lib_file"));

        /// Set up the build environment by setting Cargo configuration variables.
        pub fn setup() {
            let debug = var("DEBUG").ok().map_or(false, |s| s == "true");
            let configuration = if debug { "Debug" } else { "Release" };
            let node_lib_file_path = Path::new(NODE_LIB_FILE);
            let mut node_lib_path = Path::new(NODE_ROOT_DIR).to_path_buf();
            node_lib_path.push(NODE_ARCH);
            println!("cargo:rustc-link-search={}\\{}", NODE_ROOT_DIR, configuration);
            println!("cargo:rustc-link-search=native={}", &node_lib_path.display());
            println!("cargo:rustc-link-lib={}", &node_lib_file_path.file_stem().unwrap().to_str().unwrap());

            // Link `win_delay_load_hook.obj` for windows electron
            let node_runtime_env = "npm_config_runtime";
            println!("cargo:rerun-if-env-changed={}", node_runtime_env);
            if var(node_runtime_env).map(|s| s == "electron") == Ok(true) {
                println!("cargo:rustc-cdylib-link-arg=win_delay_load_hook.obj");
                println!("cargo:rustc-cdylib-link-arg=delayimp.lib");
                println!("cargo:rustc-cdylib-link-arg=/DELAYLOAD:node.exe");
            }
        }
    } else if #[cfg(windows)] {
        // ^ automatically not neon-sys
        use std::io::{Error, ErrorKind, Result};
        use std::process::Command;
        use std::path::Path;

        fn node_version() -> Result<String> {
            let output = Command::new("node").arg("-v").output()?;
            let stdout = String::from_utf8(output.stdout).map_err(|error| {
                Error::new(ErrorKind::InvalidData, error)
            })?;
            Ok(stdout.trim().to_string())
        }

        fn download_node_lib(version: &str, arch: &str) -> Result<Vec<u8>> {
            let script = format!(r#"
                var url = "https://nodejs.org/dist/{version}/win-{arch}/node.lib";
                // double braces because we're in a format string
                require("https").get(url, function (res) {{
                    res.pipe(process.stdout);
                }});
            "#, version = version, arch = arch);

            let output = Command::new("node").arg("-e").arg(script).output()?;

            Ok(output.stdout)
        }

        /// Set up the build environment by setting Cargo configuration variables.
        pub fn setup() {
            // If the user specified a node.lib path, we do not need to download
            if let Some(node_lib_path) = std::env::var_os("NEON_NODE_LIB") {
                let node_lib_path = Path::new(&node_lib_path);
                // Clearing the file name returns the root+directory name
                let dir = node_lib_path.with_file_name("");
                let basename = node_lib_path.file_stem().expect("Could not parse lib name from NEON_NODE_LIB. Does the path include the full file name?");

                println!("cargo:rustc-link-search=native={}", dir.display());
                // Littul hack to output the OsStr file stem
                println!("cargo:rustc-link-lib={}", Path::new(basename).display());
                return;
            }

            let version = node_version().expect("Could not determine Node.js version");
            let arch = if std::env::var("CARGO_CFG_TARGET_ARCH").unwrap() == "x86" {
                "x86"
            } else {
                "x64"
            };

            let node_lib = download_node_lib(&version, arch).expect("Could not download `node.lib`");
            std::fs::write(format!(r"{}\node-{}.lib", env!("OUT_DIR"), arch), &node_lib).expect("Could not save `node.lib`");

            println!("cargo:rustc-link-search=native={}", env!("OUT_DIR"));
            println!("cargo:rustc-link-lib=node-{}", arch);
        }
    } else if #[cfg(target_os = "macos")] {
        /// Set up the build environment by setting Cargo configuration variables.
        pub fn setup() {
            println!("cargo:rustc-cdylib-link-arg=-undefined");
            println!("cargo:rustc-cdylib-link-arg=dynamic_lookup");
        }
    } else {
        pub fn setup() { }
    }
}
<commit_msg>use @kjvalencik's non-{{ suggestion<commit_after>extern crate cfg_if;

use cfg_if::cfg_if;

cfg_if! {
    if #[cfg(all(windows, feature = "neon-sys"))] {
        use std::env::var;
        use std::path::Path;

        const NODE_ROOT_DIR: &'static str = include_str!(concat!(env!("OUT_DIR"), "\\node_root_dir"));
        const NODE_ARCH:     &'static str = include_str!(concat!(env!("OUT_DIR"), "\\node_arch"));
        const NODE_LIB_FILE: &'static str = include_str!(concat!(env!("OUT_DIR"), "\\node_lib_file"));

        /// Set up the build environment by setting Cargo configuration variables.
        pub fn setup() {
            let debug = var("DEBUG").ok().map_or(false, |s| s == "true");
            let configuration = if debug { "Debug" } else { "Release" };
            let node_lib_file_path = Path::new(NODE_LIB_FILE);
            let mut node_lib_path = Path::new(NODE_ROOT_DIR).to_path_buf();
            node_lib_path.push(NODE_ARCH);
            println!("cargo:rustc-link-search={}\\{}", NODE_ROOT_DIR, configuration);
            println!("cargo:rustc-link-search=native={}", &node_lib_path.display());
            println!("cargo:rustc-link-lib={}", &node_lib_file_path.file_stem().unwrap().to_str().unwrap());

            // Link `win_delay_load_hook.obj` for windows electron
            let node_runtime_env = "npm_config_runtime";
            println!("cargo:rerun-if-env-changed={}", node_runtime_env);
            if var(node_runtime_env).map(|s| s == "electron") == Ok(true) {
                println!("cargo:rustc-cdylib-link-arg=win_delay_load_hook.obj");
                println!("cargo:rustc-cdylib-link-arg=delayimp.lib");
                println!("cargo:rustc-cdylib-link-arg=/DELAYLOAD:node.exe");
            }
        }
    } else if #[cfg(windows)] {
        // ^ automatically not neon-sys
        use std::io::{Error, ErrorKind, Result};
        use std::process::Command;
        use std::path::Path;

        fn node_version() -> Result<String> {
            let output = Command::new("node").arg("-v").output()?;
            if !output.status.success() {
                let hopefully_stack_trace = String::from_utf8(output.stderr)
                    .unwrap_or_else(|_| "<subprocess output garbage>".to_string());
                panic!("Could not download node.lib: {}", hopefully_stack_trace);
            }
            let stdout = String::from_utf8(output.stdout).map_err(|error| {
                Error::new(ErrorKind::InvalidData, error)
            })?;
            Ok(stdout.trim().to_string())
        }

        fn download_node_lib(version: &str, arch: &str) -> Result<Vec<u8>> {
            let script = r#"
                var url = process.argv[1]
                require("https").get(url, function (res) {
                    res.pipe(process.stdout);
                });
            "#;
            let url = format!("https://nodejs.org/dist/{version}/win-{arch}/node.lib", version = version, arch = arch);

            let output = Command::new("node")
                .arg("-e").arg(script)
                .arg(url)
                .output()?;

            Ok(output.stdout)
        }

        /// Set up the build environment by setting Cargo configuration variables.
        pub fn setup() {
            // If the user specified a node.lib path, we do not need to download
            if let Some(node_lib_path) = std::env::var_os("NEON_NODE_LIB") {
                let node_lib_path = Path::new(&node_lib_path);
                // Clearing the file name returns the root+directory name
                let dir = node_lib_path.with_file_name("");
                let basename = node_lib_path.file_stem().expect("Could not parse lib name from NEON_NODE_LIB. Does the path include the full file name?");

                println!("cargo:rustc-link-search=native={}", dir.display());
                // Littul hack to output the OsStr file stem
                println!("cargo:rustc-link-lib={}", Path::new(basename).display());
                return;
            }

            let version = node_version().expect("Could not determine Node.js version");
            let arch = if std::env::var("CARGO_CFG_TARGET_ARCH").unwrap() == "x86" {
                "x86"
            } else {
                "x64"
            };

            let node_lib = download_node_lib(&version, arch).expect("Could not download `node.lib`");
            std::fs::write(format!(r"{}\node-{}.lib", env!("OUT_DIR"), arch), &node_lib).expect("Could not save `node.lib`");

            println!("cargo:rustc-link-search=native={}", env!("OUT_DIR"));
            println!("cargo:rustc-link-lib=node-{}", arch);
        }
    } else if #[cfg(target_os = "macos")] {
        /// Set up the build environment by setting Cargo configuration variables.
        pub fn setup() {
            println!("cargo:rustc-cdylib-link-arg=-undefined");
            println!("cargo:rustc-cdylib-link-arg=dynamic_lookup");
        }
    } else {
        pub fn setup() { }
    }
}
<|endoftext|>


<commit_before>//! Assorted testing utilities.
//!
//! Most notable things are:
//!
//! * Rich text comparison, which outputs a diff.
//! * Extracting markup (mainly, `<|>` markers) out of fixture strings.
//! * marks (see the eponymous module).

#[macro_use]
pub mod marks;

use std::{
    fs,
    path::{Path, PathBuf},
};

use serde_json::Value;
use text_unit::{TextRange, TextUnit};

pub use difference::Changeset as __Changeset;

pub const CURSOR_MARKER: &str = "<|>";

/// Asserts that two strings are equal, otherwise displays a rich diff between them.
///
/// The diff shows changes from the "original" left string to the "actual" right string.
///
/// All arguments starting from and including the 3rd one are passed to
/// `eprintln!()` macro in case of text inequality.
#[macro_export]
macro_rules! assert_eq_text {
    ($left:expr, $right:expr) => {
        assert_eq_text!($left, $right,)
    };
    ($left:expr, $right:expr, $($tt:tt)*) => {{
        let left = $left;
        let right = $right;
        if left != right {
            if left.trim() == right.trim() {
                eprintln!("Left:\n{:?}\n\nRight:\n{:?}\n\nWhitespace difference\n", left, right);
            } else {
                let changeset = $crate::__Changeset::new(right, left, "\n");
                eprintln!("Left:\n{}\n\nRight:\n{}\n\nDiff:\n{}\n", left, right, changeset);
            }
            eprintln!($($tt)*);
            panic!("text differs");
        }
    }};
}

/// Infallible version of `try_extract_offset()`.
pub fn extract_offset(text: &str) -> (TextUnit, String) {
    match try_extract_offset(text) {
        None => panic!("text should contain cursor marker"),
        Some(result) => result,
    }
}

/// Returns the offset of the first occurence of `<|>` marker and the copy of `text`
/// without the marker.
fn try_extract_offset(text: &str) -> Option<(TextUnit, String)> {
    let cursor_pos = text.find(CURSOR_MARKER)?;
    let mut new_text = String::with_capacity(text.len() - CURSOR_MARKER.len());
    new_text.push_str(&text[..cursor_pos]);
    new_text.push_str(&text[cursor_pos + CURSOR_MARKER.len()..]);
    let cursor_pos = TextUnit::from(cursor_pos as u32);
    Some((cursor_pos, new_text))
}

/// Infallible version of `try_extract_range()`.
pub fn extract_range(text: &str) -> (TextRange, String) {
    match try_extract_range(text) {
        None => panic!("text should contain cursor marker"),
        Some(result) => result,
    }
}

/// Returns `TextRange` between the first two markers `<|>...<|>` and the copy
/// of `text` without both of these markers.
fn try_extract_range(text: &str) -> Option<(TextRange, String)> {
    let (start, text) = try_extract_offset(text)?;
    let (end, text) = try_extract_offset(&text)?;
    Some((TextRange::from_to(start, end), text))
}

pub enum RangeOrOffset {
    Range(TextRange),
    Offset(TextUnit),
}

impl From<RangeOrOffset> for TextRange {
    fn from(selection: RangeOrOffset) -> Self {
        match selection {
            RangeOrOffset::Range(it) => it,
            RangeOrOffset::Offset(it) => TextRange::from_to(it, it),
        }
    }
}

/// Extracts `TextRange` or `TextUnit` depending on the amount of `<|>` markers
/// found in `text`.
///
/// # Panics
/// Panics if no `<|>` marker is present in the `text`.
pub fn extract_range_or_offset(text: &str) -> (RangeOrOffset, String) {
    if let Some((range, text)) = try_extract_range(text) {
        return (RangeOrOffset::Range(range), text);
    }
    let (offset, text) = extract_offset(text);
    (RangeOrOffset::Offset(offset), text)
}

/// Extracts ranges, marked with `<tag> </tag>` pairs from the `text`
pub fn extract_ranges(mut text: &str, tag: &str) -> (Vec<TextRange>, String) {
    let open = format!("<{}>", tag);
    let close = format!("</{}>", tag);
    let mut ranges = Vec::new();
    let mut res = String::new();
    let mut stack = Vec::new();
    loop {
        match text.find('<') {
            None => {
                res.push_str(text);
                break;
            }
            Some(i) => {
                res.push_str(&text[..i]);
                text = &text[i..];
                if text.starts_with(&open) {
                    text = &text[open.len()..];
                    let from = TextUnit::of_str(&res);
                    stack.push(from);
                } else if text.starts_with(&close) {
                    text = &text[close.len()..];
                    let from = stack.pop().unwrap_or_else(|| panic!("unmatched </{}>", tag));
                    let to = TextUnit::of_str(&res);
                    ranges.push(TextRange::from_to(from, to));
                }
            }
        }
    }
    assert!(stack.is_empty(), "unmatched <{}>", tag);
    ranges.sort_by_key(|r| (r.start(), r.end()));
    (ranges, res)
}

/// Inserts `<|>` marker into the `text` at `offset`.
pub fn add_cursor(text: &str, offset: TextUnit) -> String {
    let offset: usize = offset.to_usize();
    let mut res = String::new();
    res.push_str(&text[..offset]);
    res.push_str("<|>");
    res.push_str(&text[offset..]);
    res
}

#[derive(Debug)]
pub struct FixtureEntry {
    pub meta: String,
    pub text: String,
}

/// Parses text which looks like this:
///
///  ```not_rust
///  //- some meta
///  line 1
///  line 2
///  // - other meta
///  ```
pub fn parse_fixture(fixture: &str) -> Vec<FixtureEntry> {
    let margin = fixture
        .lines()
        .filter(|it| it.trim_start().starts_with("//-"))
        .map(|it| it.len() - it.trim_start().len())
        .next()
        .expect("empty fixture");

    let lines = fixture
        .split('\n') // don't use `.lines` to not drop `\r\n`
        .filter_map(|line| {
            if line.len() >= margin {
                assert!(line[..margin].trim().is_empty());
                Some(&line[margin..])
            } else {
                assert!(line.trim().is_empty());
                None
            }
        });

    let mut res = Vec::new();
    let mut buf = String::new();
    let mut meta: Option<&str> = None;

    macro_rules! flush {
        () => {
            if let Some(meta) = meta {
                res.push(FixtureEntry { meta: meta.to_string(), text: buf.clone() });
                buf.clear();
            }
        };
    };

    for line in lines {
        if line.starts_with("//-") {
            flush!();
            buf.clear();
            meta = Some(line["//-".len()..].trim());
            continue;
        }
        buf.push_str(line);
        buf.push('\n');
    }
    flush!();
    res
}

// Comparison functionality borrowed from cargo:

/// Compare a line with an expected pattern.
/// - Use `[..]` as a wildcard to match 0 or more characters on the same line
///   (similar to `.*` in a regex).
pub fn lines_match(expected: &str, actual: &str) -> bool {
    // Let's not deal with / vs \ (windows...)
    // First replace backslash-escaped backslashes with forward slashes
    // which can occur in, for example, JSON output
    let expected = expected.replace(r"\\", "/").replace(r"\", "/");
    let mut actual: &str = &actual.replace(r"\\", "/").replace(r"\", "/");
    for (i, part) in expected.split("[..]").enumerate() {
        match actual.find(part) {
            Some(j) => {
                if i == 0 && j != 0 {
                    return false;
                }
                actual = &actual[j + part.len()..];
            }
            None => return false,
        }
    }
    actual.is_empty() || expected.ends_with("[..]")
}

#[test]
fn lines_match_works() {
    assert!(lines_match("a b", "a b"));
    assert!(lines_match("a[..]b", "a b"));
    assert!(lines_match("a[..]", "a b"));
    assert!(lines_match("[..]", "a b"));
    assert!(lines_match("[..]b", "a b"));

    assert!(!lines_match("[..]b", "c"));
    assert!(!lines_match("b", "c"));
    assert!(!lines_match("b", "cb"));
}

/// Compares JSON object for approximate equality.
/// You can use `[..]` wildcard in strings (useful for OS dependent things such
/// as paths). You can use a `"{...}"` string literal as a wildcard for
/// arbitrary nested JSON. Arrays are sorted before comparison.
pub fn find_mismatch<'a>(expected: &'a Value, actual: &'a Value) -> Option<(&'a Value, &'a Value)> {
    use serde_json::Value::*;
    match (expected, actual) {
        (&Number(ref l), &Number(ref r)) if l == r => None,
        (&Bool(l), &Bool(r)) if l == r => None,
        (&String(ref l), &String(ref r)) if lines_match(l, r) => None,
        (&Array(ref l), &Array(ref r)) => {
            if l.len() != r.len() {
                return Some((expected, actual));
            }

            let mut l = l.iter().collect::<Vec<_>>();
            let mut r = r.iter().collect::<Vec<_>>();

            l.retain(|l| match r.iter().position(|r| find_mismatch(l, r).is_none()) {
                Some(i) => {
                    r.remove(i);
                    false
                }
                None => true,
            });

            if !l.is_empty() {
                assert!(!r.is_empty());
                Some((&l[0], &r[0]))
            } else {
                assert_eq!(r.len(), 0);
                None
            }
        }
        (&Object(ref l), &Object(ref r)) => {
            let same_keys = l.len() == r.len() && l.keys().all(|k| r.contains_key(k));
            if !same_keys {
                return Some((expected, actual));
            }

            l.values().zip(r.values()).filter_map(|(l, r)| find_mismatch(l, r)).nth(0)
        }
        (&Null, &Null) => None,
        // magic string literal "{...}" acts as wildcard for any sub-JSON
        (&String(ref l), _) if l == "{...}" => None,
        _ => Some((expected, actual)),
    }
}

/// Calls callback `f` with input code and file paths of all `.rs` files from `test_data_dir`
/// subdirectories defined by `paths`.
///
/// If the content of the matching `.txt` file differs from the output of `f()`
/// the test will fail.
///
/// If there is no matching `.txt` file it will be created and filled with the
/// output of `f()`, but the test will fail.
pub fn dir_tests<F>(test_data_dir: &Path, paths: &[&str], f: F)
where
    F: Fn(&str, &Path) -> String,
{
    for (path, input_code) in collect_tests(test_data_dir, paths) {
        let parse_tree = f(&input_code, &path);
        let path = path.with_extension("txt");
        if !path.exists() {
            println!("\nfile: {}", path.display());
            println!("No .txt file with expected result, creating...\n");
            println!("{}\n{}", input_code, parse_tree);
            fs::write(&path, &parse_tree).unwrap();
            panic!("No expected result")
        }
        let expected = read_text(&path);
        let expected = expected.as_str();
        let parse_tree = parse_tree.as_str();
        assert_equal_text(expected, parse_tree, &path);
    }
}

/// Collects all `.rs` files from `test_data_dir` subdirectories defined by `paths`.
pub fn collect_tests(test_data_dir: &Path, paths: &[&str]) -> Vec<(PathBuf, String)> {
    paths
        .iter()
        .flat_map(|path| {
            let path = test_data_dir.to_owned().join(path);
            test_from_dir(&path).into_iter()
        })
        .map(|path| {
            let text = read_text(&path);
            (path, text)
        })
        .collect()
}

/// Collects paths to all `.rs` files from `dir` in a sorted `Vec<PathBuf>`.
fn test_from_dir(dir: &Path) -> Vec<PathBuf> {
    let mut acc = Vec::new();
    for file in fs::read_dir(&dir).unwrap() {
        let file = file.unwrap();
        let path = file.path();
        if path.extension().unwrap_or_default() == "rs" {
            acc.push(path);
        }
    }
    acc.sort();
    acc
}

/// Returns the path to the root directory of `rust-analyzer` project.
pub fn project_dir() -> PathBuf {
    let dir = env!("CARGO_MANIFEST_DIR");
    PathBuf::from(dir).parent().unwrap().parent().unwrap().to_owned()
}

/// Read file and normalize newlines.
///
/// `rustc` seems to always normalize `\r\n` newlines to `\n`:
///
/// ```
/// let s = "
/// ";
/// assert_eq!(s.as_bytes(), &[10]);
/// ```
///
/// so this should always be correct.
pub fn read_text(path: &Path) -> String {
    fs::read_to_string(path)
        .unwrap_or_else(|_| panic!("File at {:?} should be valid", path))
        .replace("\r\n", "\n")
}

/// Returns `false` if slow tests should not run, otherwise returns `true` and
/// also creates a file at `./target/.slow_tests_cookie` which serves as a flag
/// that slow tests did run.
pub fn skip_slow_tests() -> bool {
    let should_skip = std::env::var("CI").is_err() && std::env::var("RUN_SLOW_TESTS").is_err();
    if should_skip {
        eprintln!("ignoring slow test")
    } else {
        let path = project_dir().join("./target/.slow_tests_cookie");
        fs::write(&path, ".").unwrap();
    }
    should_skip
}

/// Asserts that `expected` and `actual` strings are equal. If they differ only
/// in trailing or leading whitespace the test won't fail and
/// the contents of `actual` will be written to the file located at `path`.
fn assert_equal_text(expected: &str, actual: &str, path: &Path) {
    if expected == actual {
        return;
    }
    let dir = project_dir();
    let pretty_path = path.strip_prefix(&dir).unwrap_or_else(|_| path);
    if expected.trim() == actual.trim() {
        println!("whitespace difference, rewriting");
        println!("file: {}\n", pretty_path.display());
        fs::write(path, actual).unwrap();
        return;
    }
    const REWRITE: bool = false;
    if REWRITE {
        println!("rewriting {}", pretty_path.display());
        fs::write(path, actual).unwrap();
        return;
    }
    assert_eq_text!(expected, actual, "file: {}", pretty_path.display());
}
<commit_msg>Simplify fixture parsing<commit_after>//! Assorted testing utilities.
//!
//! Most notable things are:
//!
//! * Rich text comparison, which outputs a diff.
//! * Extracting markup (mainly, `<|>` markers) out of fixture strings.
//! * marks (see the eponymous module).

#[macro_use]
pub mod marks;

use std::{
    fs,
    path::{Path, PathBuf},
};

use serde_json::Value;
use text_unit::{TextRange, TextUnit};

pub use difference::Changeset as __Changeset;

pub const CURSOR_MARKER: &str = "<|>";

/// Asserts that two strings are equal, otherwise displays a rich diff between them.
///
/// The diff shows changes from the "original" left string to the "actual" right string.
///
/// All arguments starting from and including the 3rd one are passed to
/// `eprintln!()` macro in case of text inequality.
#[macro_export]
macro_rules! assert_eq_text {
    ($left:expr, $right:expr) => {
        assert_eq_text!($left, $right,)
    };
    ($left:expr, $right:expr, $($tt:tt)*) => {{
        let left = $left;
        let right = $right;
        if left != right {
            if left.trim() == right.trim() {
                eprintln!("Left:\n{:?}\n\nRight:\n{:?}\n\nWhitespace difference\n", left, right);
            } else {
                let changeset = $crate::__Changeset::new(right, left, "\n");
                eprintln!("Left:\n{}\n\nRight:\n{}\n\nDiff:\n{}\n", left, right, changeset);
            }
            eprintln!($($tt)*);
            panic!("text differs");
        }
    }};
}

/// Infallible version of `try_extract_offset()`.
pub fn extract_offset(text: &str) -> (TextUnit, String) {
    match try_extract_offset(text) {
        None => panic!("text should contain cursor marker"),
        Some(result) => result,
    }
}

/// Returns the offset of the first occurence of `<|>` marker and the copy of `text`
/// without the marker.
fn try_extract_offset(text: &str) -> Option<(TextUnit, String)> {
    let cursor_pos = text.find(CURSOR_MARKER)?;
    let mut new_text = String::with_capacity(text.len() - CURSOR_MARKER.len());
    new_text.push_str(&text[..cursor_pos]);
    new_text.push_str(&text[cursor_pos + CURSOR_MARKER.len()..]);
    let cursor_pos = TextUnit::from(cursor_pos as u32);
    Some((cursor_pos, new_text))
}

/// Infallible version of `try_extract_range()`.
pub fn extract_range(text: &str) -> (TextRange, String) {
    match try_extract_range(text) {
        None => panic!("text should contain cursor marker"),
        Some(result) => result,
    }
}

/// Returns `TextRange` between the first two markers `<|>...<|>` and the copy
/// of `text` without both of these markers.
fn try_extract_range(text: &str) -> Option<(TextRange, String)> {
    let (start, text) = try_extract_offset(text)?;
    let (end, text) = try_extract_offset(&text)?;
    Some((TextRange::from_to(start, end), text))
}

pub enum RangeOrOffset {
    Range(TextRange),
    Offset(TextUnit),
}

impl From<RangeOrOffset> for TextRange {
    fn from(selection: RangeOrOffset) -> Self {
        match selection {
            RangeOrOffset::Range(it) => it,
            RangeOrOffset::Offset(it) => TextRange::from_to(it, it),
        }
    }
}

/// Extracts `TextRange` or `TextUnit` depending on the amount of `<|>` markers
/// found in `text`.
///
/// # Panics
/// Panics if no `<|>` marker is present in the `text`.
pub fn extract_range_or_offset(text: &str) -> (RangeOrOffset, String) {
    if let Some((range, text)) = try_extract_range(text) {
        return (RangeOrOffset::Range(range), text);
    }
    let (offset, text) = extract_offset(text);
    (RangeOrOffset::Offset(offset), text)
}

/// Extracts ranges, marked with `<tag> </tag>` pairs from the `text`
pub fn extract_ranges(mut text: &str, tag: &str) -> (Vec<TextRange>, String) {
    let open = format!("<{}>", tag);
    let close = format!("</{}>", tag);
    let mut ranges = Vec::new();
    let mut res = String::new();
    let mut stack = Vec::new();
    loop {
        match text.find('<') {
            None => {
                res.push_str(text);
                break;
            }
            Some(i) => {
                res.push_str(&text[..i]);
                text = &text[i..];
                if text.starts_with(&open) {
                    text = &text[open.len()..];
                    let from = TextUnit::of_str(&res);
                    stack.push(from);
                } else if text.starts_with(&close) {
                    text = &text[close.len()..];
                    let from = stack.pop().unwrap_or_else(|| panic!("unmatched </{}>", tag));
                    let to = TextUnit::of_str(&res);
                    ranges.push(TextRange::from_to(from, to));
                }
            }
        }
    }
    assert!(stack.is_empty(), "unmatched <{}>", tag);
    ranges.sort_by_key(|r| (r.start(), r.end()));
    (ranges, res)
}

/// Inserts `<|>` marker into the `text` at `offset`.
pub fn add_cursor(text: &str, offset: TextUnit) -> String {
    let offset: usize = offset.to_usize();
    let mut res = String::new();
    res.push_str(&text[..offset]);
    res.push_str("<|>");
    res.push_str(&text[offset..]);
    res
}

#[derive(Debug)]
pub struct FixtureEntry {
    pub meta: String,
    pub text: String,
}

/// Parses text which looks like this:
///
///  ```not_rust
///  //- some meta
///  line 1
///  line 2
///  // - other meta
///  ```
pub fn parse_fixture(fixture: &str) -> Vec<FixtureEntry> {
    let margin = fixture
        .lines()
        .filter(|it| it.trim_start().starts_with("//-"))
        .map(|it| it.len() - it.trim_start().len())
        .next()
        .expect("empty fixture");

    let mut lines = fixture
        .split('\n') // don't use `.lines` to not drop `\r\n`
        .filter_map(|line| {
            if line.len() >= margin {
                assert!(line[..margin].trim().is_empty());
                Some(&line[margin..])
            } else {
                assert!(line.trim().is_empty());
                None
            }
        });

    let mut res = Vec::new();
    let mut meta = None;
    loop {
        let mut next_meta = None;
        let mut text = String::new();
        for line in lines.by_ref() {
            if line.starts_with("//-") {
                next_meta = Some(line["//-".len()..].trim().to_string());
                break;
            }
            text.push_str(line);
            text.push('\n');
        }

        if let Some(meta) = meta {
            res.push(FixtureEntry { meta, text });
        }
        meta = next_meta;
        if meta.is_none() {
            break;
        }
    }

    res
}

// Comparison functionality borrowed from cargo:

/// Compare a line with an expected pattern.
/// - Use `[..]` as a wildcard to match 0 or more characters on the same line
///   (similar to `.*` in a regex).
pub fn lines_match(expected: &str, actual: &str) -> bool {
    // Let's not deal with / vs \ (windows...)
    // First replace backslash-escaped backslashes with forward slashes
    // which can occur in, for example, JSON output
    let expected = expected.replace(r"\\", "/").replace(r"\", "/");
    let mut actual: &str = &actual.replace(r"\\", "/").replace(r"\", "/");
    for (i, part) in expected.split("[..]").enumerate() {
        match actual.find(part) {
            Some(j) => {
                if i == 0 && j != 0 {
                    return false;
                }
                actual = &actual[j + part.len()..];
            }
            None => return false,
        }
    }
    actual.is_empty() || expected.ends_with("[..]")
}

#[test]
fn lines_match_works() {
    assert!(lines_match("a b", "a b"));
    assert!(lines_match("a[..]b", "a b"));
    assert!(lines_match("a[..]", "a b"));
    assert!(lines_match("[..]", "a b"));
    assert!(lines_match("[..]b", "a b"));

    assert!(!lines_match("[..]b", "c"));
    assert!(!lines_match("b", "c"));
    assert!(!lines_match("b", "cb"));
}

/// Compares JSON object for approximate equality.
/// You can use `[..]` wildcard in strings (useful for OS dependent things such
/// as paths). You can use a `"{...}"` string literal as a wildcard for
/// arbitrary nested JSON. Arrays are sorted before comparison.
pub fn find_mismatch<'a>(expected: &'a Value, actual: &'a Value) -> Option<(&'a Value, &'a Value)> {
    use serde_json::Value::*;
    match (expected, actual) {
        (&Number(ref l), &Number(ref r)) if l == r => None,
        (&Bool(l), &Bool(r)) if l == r => None,
        (&String(ref l), &String(ref r)) if lines_match(l, r) => None,
        (&Array(ref l), &Array(ref r)) => {
            if l.len() != r.len() {
                return Some((expected, actual));
            }

            let mut l = l.iter().collect::<Vec<_>>();
            let mut r = r.iter().collect::<Vec<_>>();

            l.retain(|l| match r.iter().position(|r| find_mismatch(l, r).is_none()) {
                Some(i) => {
                    r.remove(i);
                    false
                }
                None => true,
            });

            if !l.is_empty() {
                assert!(!r.is_empty());
                Some((&l[0], &r[0]))
            } else {
                assert_eq!(r.len(), 0);
                None
            }
        }
        (&Object(ref l), &Object(ref r)) => {
            let same_keys = l.len() == r.len() && l.keys().all(|k| r.contains_key(k));
            if !same_keys {
                return Some((expected, actual));
            }

            l.values().zip(r.values()).filter_map(|(l, r)| find_mismatch(l, r)).nth(0)
        }
        (&Null, &Null) => None,
        // magic string literal "{...}" acts as wildcard for any sub-JSON
        (&String(ref l), _) if l == "{...}" => None,
        _ => Some((expected, actual)),
    }
}

/// Calls callback `f` with input code and file paths of all `.rs` files from `test_data_dir`
/// subdirectories defined by `paths`.
///
/// If the content of the matching `.txt` file differs from the output of `f()`
/// the test will fail.
///
/// If there is no matching `.txt` file it will be created and filled with the
/// output of `f()`, but the test will fail.
pub fn dir_tests<F>(test_data_dir: &Path, paths: &[&str], f: F)
where
    F: Fn(&str, &Path) -> String,
{
    for (path, input_code) in collect_tests(test_data_dir, paths) {
        let parse_tree = f(&input_code, &path);
        let path = path.with_extension("txt");
        if !path.exists() {
            println!("\nfile: {}", path.display());
            println!("No .txt file with expected result, creating...\n");
            println!("{}\n{}", input_code, parse_tree);
            fs::write(&path, &parse_tree).unwrap();
            panic!("No expected result")
        }
        let expected = read_text(&path);
        let expected = expected.as_str();
        let parse_tree = parse_tree.as_str();
        assert_equal_text(expected, parse_tree, &path);
    }
}

/// Collects all `.rs` files from `test_data_dir` subdirectories defined by `paths`.
pub fn collect_tests(test_data_dir: &Path, paths: &[&str]) -> Vec<(PathBuf, String)> {
    paths
        .iter()
        .flat_map(|path| {
            let path = test_data_dir.to_owned().join(path);
            test_from_dir(&path).into_iter()
        })
        .map(|path| {
            let text = read_text(&path);
            (path, text)
        })
        .collect()
}

/// Collects paths to all `.rs` files from `dir` in a sorted `Vec<PathBuf>`.
fn test_from_dir(dir: &Path) -> Vec<PathBuf> {
    let mut acc = Vec::new();
    for file in fs::read_dir(&dir).unwrap() {
        let file = file.unwrap();
        let path = file.path();
        if path.extension().unwrap_or_default() == "rs" {
            acc.push(path);
        }
    }
    acc.sort();
    acc
}

/// Returns the path to the root directory of `rust-analyzer` project.
pub fn project_dir() -> PathBuf {
    let dir = env!("CARGO_MANIFEST_DIR");
    PathBuf::from(dir).parent().unwrap().parent().unwrap().to_owned()
}

/// Read file and normalize newlines.
///
/// `rustc` seems to always normalize `\r\n` newlines to `\n`:
///
/// ```
/// let s = "
/// ";
/// assert_eq!(s.as_bytes(), &[10]);
/// ```
///
/// so this should always be correct.
pub fn read_text(path: &Path) -> String {
    fs::read_to_string(path)
        .unwrap_or_else(|_| panic!("File at {:?} should be valid", path))
        .replace("\r\n", "\n")
}

/// Returns `false` if slow tests should not run, otherwise returns `true` and
/// also creates a file at `./target/.slow_tests_cookie` which serves as a flag
/// that slow tests did run.
pub fn skip_slow_tests() -> bool {
    let should_skip = std::env::var("CI").is_err() && std::env::var("RUN_SLOW_TESTS").is_err();
    if should_skip {
        eprintln!("ignoring slow test")
    } else {
        let path = project_dir().join("./target/.slow_tests_cookie");
        fs::write(&path, ".").unwrap();
    }
    should_skip
}

/// Asserts that `expected` and `actual` strings are equal. If they differ only
/// in trailing or leading whitespace the test won't fail and
/// the contents of `actual` will be written to the file located at `path`.
fn assert_equal_text(expected: &str, actual: &str, path: &Path) {
    if expected == actual {
        return;
    }
    let dir = project_dir();
    let pretty_path = path.strip_prefix(&dir).unwrap_or_else(|_| path);
    if expected.trim() == actual.trim() {
        println!("whitespace difference, rewriting");
        println!("file: {}\n", pretty_path.display());
        fs::write(path, actual).unwrap();
        return;
    }
    const REWRITE: bool = false;
    if REWRITE {
        println!("rewriting {}", pretty_path.display());
        fs::write(path, actual).unwrap();
        return;
    }
    assert_eq_text!(expected, actual, "file: {}", pretty_path.display());
}
<|endoftext|>


<commit_before>#![deny(unsafe_op_in_unsafe_fn)]

use crate::alloc::{GlobalAlloc, Layout, System};
use crate::ptr;
use crate::sys::c;
use crate::sys_common::alloc::{realloc_fallback, MIN_ALIGN};

#[cfg(test)]
mod tests;

// Heap memory management on Windows is done by using the system Heap API (heapapi.h)
// See https://docs.microsoft.com/windows/win32/api/heapapi/

// Flag to indicate that the memory returned by `HeapAlloc` should be zeroed.
const HEAP_ZERO_MEMORY: c::DWORD = 0x00000008;

extern "system" {
    // Get a handle to the default heap of the current process, or null if the operation fails.
    //
    // See https://docs.microsoft.com/windows/win32/api/heapapi/nf-heapapi-getprocessheap
    fn GetProcessHeap() -> c::HANDLE;

    // Allocate a block of `dwBytes` bytes of memory from a given heap `hHeap`.
    // The allocated memory may be uninitialized, or zeroed if `dwFlags` is
    // set to `HEAP_ZERO_MEMORY`.
    //
    // Returns a pointer to the newly-allocated memory or null if the operation fails.
    // The returned pointer will be aligned to at least `MIN_ALIGN`.
    //
    // SAFETY:
    //  - `hHeap` must be a non-null handle returned by `GetProcessHeap`.
    //  - `dwFlags` must be set to either zero or `HEAP_ZERO_MEMORY`.
    //
    // Note that `dwBytes` is allowed to be zero, contrary to some other allocators.
    //
    // See https://docs.microsoft.com/windows/win32/api/heapapi/nf-heapapi-heapalloc
    fn HeapAlloc(hHeap: c::HANDLE, dwFlags: c::DWORD, dwBytes: c::SIZE_T) -> c::LPVOID;

    // Reallocate a block of memory behind a given pointer `lpMem` from a given heap `hHeap`,
    // to a block of at least `dwBytes` bytes, either shrinking the block in place,
    // or allocating at a new location, copying memory, and freeing the original location.
    //
    // Returns a pointer to the reallocated memory or null if the operation fails.
    // The returned pointer will be aligned to at least `MIN_ALIGN`.
    // If the operation fails the given block will never have been freed.
    //
    // SAFETY:
    //  - `hHeap` must be a non-null handle returned by `GetProcessHeap`.
    //  - `dwFlags` must be set to zero.
    //  - `lpMem` must be a non-null pointer to an allocated block returned by `HeapAlloc` or
    //     `HeapReAlloc`, that has not already been freed.
    // If the block was successfully reallocated at a new location, pointers pointing to
    // the freed memory, such as `lpMem`, must not be dereferenced ever again.
    //
    // Note that `dwBytes` is allowed to be zero, contrary to some other allocators.
    //
    // See https://docs.microsoft.com/windows/win32/api/heapapi/nf-heapapi-heaprealloc
    fn HeapReAlloc(
        hHeap: c::HANDLE,
        dwFlags: c::DWORD,
        lpMem: c::LPVOID,
        dwBytes: c::SIZE_T,
    ) -> c::LPVOID;

    // Free a block of memory behind a given pointer `lpMem` from a given heap `hHeap`.
    // Returns a nonzero value if the operation is successful, and zero if the operation fails.
    //
    // SAFETY:
    //  - `dwFlags` must be set to zero.
    //  - `lpMem` must be a pointer to an allocated block returned by `HeapAlloc` or `HeapReAlloc`,
    //     that has not already been freed.
    // If the block was successfully freed, pointers pointing to the freed memory, such as `lpMem`,
    // must not be dereferenced ever again.
    //
    // Note that both `hHeap` is allowed to be any value, and `lpMem` is allowed to be null,
    // both of which will not cause the operation to fail.
    //
    // See https://docs.microsoft.com/windows/win32/api/heapapi/nf-heapapi-heapfree
    fn HeapFree(hHeap: c::HANDLE, dwFlags: c::DWORD, lpMem: c::LPVOID) -> c::BOOL;
}

// Header containing a pointer to the start of an allocated block.
// SAFETY: size and alignment must be <= `MIN_ALIGN`.
#[repr(C)]
struct Header(*mut u8);

// Allocates a block of optionally zeroed memory for a given `layout`.
// Returns a pointer satisfying the guarantees of `System` about allocated pointers.
#[inline]
unsafe fn allocate(layout: Layout, zeroed: bool) -> *mut u8 {
    let heap = unsafe { GetProcessHeap() };
    if heap.is_null() {
        // Allocation has failed, could not get the current process heap.
        return ptr::null_mut();
    }

    // Allocated memory will be either zeroed or uninitialized.
    let flags = if zeroed { HEAP_ZERO_MEMORY } else { 0 };

    if layout.align() <= MIN_ALIGN {
        // SAFETY: `heap` is a non-null handle returned by `GetProcessHeap`.
        // The returned pointer points to the start of an allocated block.
        unsafe { HeapAlloc(heap, flags, layout.size()) as *mut u8 }
    } else {
        // Allocate extra padding in order to be able to satisfy the alignment.
        let total = layout.align() + layout.size();

        // SAFETY: `heap` is a non-null handle returned by `GetProcessHeap`.
        let ptr = unsafe { HeapAlloc(heap, flags, total) as *mut u8 };
        if ptr.is_null() {
            // Allocation has failed.
            return ptr::null_mut();
        }

        // Create a correctly aligned pointer offset from the start of the allocated block,
        // and write a header before it.

        let offset = layout.align() - (ptr as usize & (layout.align() - 1));
        // SAFETY: `MIN_ALIGN` <= `offset` <= `layout.align()` and the size of the allocated
        // block is `layout.align() + layout.size()`. `aligned` will thus be a correctly aligned
        // pointer inside the allocated block with at least `layout.size()` bytes after it and at
        // least `MIN_ALIGN` bytes of padding before it.
        let aligned = unsafe { ptr.add(offset) };
        // SAFETY: Because the size and alignment of a header is <= `MIN_ALIGN` and `aligned`
        // is aligned to at least `MIN_ALIGN` and has at least `MIN_ALIGN` bytes of padding before
        // it, it is safe to write a header directly before it.
        unsafe { ptr::write((aligned as *mut Header).offset(-1), Header(ptr)) };

        // SAFETY: The returned pointer does not point to the to the start of an allocated block,
        // but there is a header readable directly before it containing the location of the start
        // of the block.
        aligned
    }
}

// All pointers returned by this allocator have, in addition to the guarantees of `GlobalAlloc`, the
// following properties:
//
// If the pointer was allocated or reallocated with a `layout` specifying an alignment <= `MIN_ALIGN`
// the pointer will be aligned to at least `MIN_ALIGN` and point to the start of the allocated block.
//
// If the pointer was allocated or reallocated with a `layout` specifying an alignment > `MIN_ALIGN`
// the pointer will be aligned to the specified alignment and not point to the start of the allocated block.
// Instead there will be a header readable directly before the returned pointer, containing the actual
// location of the start of the block.
#[stable(feature = "alloc_system_type", since = "1.28.0")]
unsafe impl GlobalAlloc for System {
    #[inline]
    unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
        // SAFETY: pointers returned by `allocate` satisfy the guarantees of `System`
        unsafe { allocate(layout, false) }
    }

    #[inline]
    unsafe fn alloc_zeroed(&self, layout: Layout) -> *mut u8 {
        // SAFETY: pointers returned by `allocate` satisfy the guarantees of `System`
        unsafe { allocate(layout, true) }
    }

    #[inline]
    unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
        let block = {
            if layout.align() <= MIN_ALIGN {
                ptr
            } else {
                // The location of the start of the block is stored in the padding before `ptr`.

                // SAFETY: Because of the contract of `System`, `ptr` is guaranteed to be non-null
                // and have a header readable directly before it.
                unsafe { ptr::read((ptr as *mut Header).offset(-1)).0 }
            }
        };

        // SAFETY: `block` is a pointer to the start of an allocated block.
        unsafe {
            let err = HeapFree(GetProcessHeap(), 0, block as c::LPVOID);
            debug_assert!(err != 0, "Failed to free heap memory: {}", c::GetLastError());
        }
    }

    #[inline]
    unsafe fn realloc(&self, ptr: *mut u8, layout: Layout, new_size: usize) -> *mut u8 {
        if layout.align() <= MIN_ALIGN {
            let heap = unsafe { GetProcessHeap() };
            if heap.is_null() {
                // Reallocation has failed, could not get the current process heap.
                return ptr::null_mut();
            }

            // SAFETY: `heap` is a non-null handle returned by `GetProcessHeap`,
            // `ptr` is a pointer to the start of an allocated block.
            // The returned pointer points to the start of an allocated block.
            unsafe { HeapReAlloc(heap, 0, ptr as c::LPVOID, new_size) as *mut u8 }
        } else {
            // SAFETY: `realloc_fallback` is implemented using `dealloc` and `alloc`, which will
            // correctly handle `ptr` and return a pointer satisfying the guarantees of `System`
            unsafe { realloc_fallback(self, ptr, layout, new_size) }
        }
    }
}
<commit_msg>Apply suggestions from code review<commit_after>#![deny(unsafe_op_in_unsafe_fn)]

use crate::alloc::{GlobalAlloc, Layout, System};
use crate::ptr;
use crate::sys::c;
use crate::sys_common::alloc::{realloc_fallback, MIN_ALIGN};

#[cfg(test)]
mod tests;

// Heap memory management on Windows is done by using the system Heap API (heapapi.h)
// See https://docs.microsoft.com/windows/win32/api/heapapi/

// Flag to indicate that the memory returned by `HeapAlloc` should be zeroed.
const HEAP_ZERO_MEMORY: c::DWORD = 0x00000008;

extern "system" {
    // Get a handle to the default heap of the current process, or null if the operation fails.
    //
    // See https://docs.microsoft.com/windows/win32/api/heapapi/nf-heapapi-getprocessheap
    fn GetProcessHeap() -> c::HANDLE;

    // Allocate a block of `dwBytes` bytes of memory from a given heap `hHeap`.
    // The allocated memory may be uninitialized, or zeroed if `dwFlags` is
    // set to `HEAP_ZERO_MEMORY`.
    //
    // Returns a pointer to the newly-allocated memory or null if the operation fails.
    // The returned pointer will be aligned to at least `MIN_ALIGN`.
    //
    // SAFETY:
    //  - `hHeap` must be a non-null handle returned by `GetProcessHeap`.
    //  - `dwFlags` must be set to either zero or `HEAP_ZERO_MEMORY`.
    //
    // Note that `dwBytes` is allowed to be zero, contrary to some other allocators.
    //
    // See https://docs.microsoft.com/windows/win32/api/heapapi/nf-heapapi-heapalloc
    fn HeapAlloc(hHeap: c::HANDLE, dwFlags: c::DWORD, dwBytes: c::SIZE_T) -> c::LPVOID;

    // Reallocate a block of memory behind a given pointer `lpMem` from a given heap `hHeap`,
    // to a block of at least `dwBytes` bytes, either shrinking the block in place,
    // or allocating at a new location, copying memory, and freeing the original location.
    //
    // Returns a pointer to the reallocated memory or null if the operation fails.
    // The returned pointer will be aligned to at least `MIN_ALIGN`.
    // If the operation fails the given block will never have been freed.
    //
    // SAFETY:
    //  - `hHeap` must be a non-null handle returned by `GetProcessHeap`.
    //  - `dwFlags` must be set to zero.
    //  - `lpMem` must be a non-null pointer to an allocated block returned by `HeapAlloc` or
    //     `HeapReAlloc`, that has not already been freed.
    // If the block was successfully reallocated at a new location, pointers pointing to
    // the freed memory, such as `lpMem`, must not be dereferenced ever again.
    //
    // Note that `dwBytes` is allowed to be zero, contrary to some other allocators.
    //
    // See https://docs.microsoft.com/windows/win32/api/heapapi/nf-heapapi-heaprealloc
    fn HeapReAlloc(
        hHeap: c::HANDLE,
        dwFlags: c::DWORD,
        lpMem: c::LPVOID,
        dwBytes: c::SIZE_T,
    ) -> c::LPVOID;

    // Free a block of memory behind a given pointer `lpMem` from a given heap `hHeap`.
    // Returns a nonzero value if the operation is successful, and zero if the operation fails.
    //
    // SAFETY:
    //  - `dwFlags` must be set to zero.
    //  - `lpMem` must be a pointer to an allocated block returned by `HeapAlloc` or `HeapReAlloc`,
    //     that has not already been freed.
    // If the block was successfully freed, pointers pointing to the freed memory, such as `lpMem`,
    // must not be dereferenced ever again.
    //
    // Note that both `hHeap` is allowed to be any value, and `lpMem` is allowed to be null,
    // both of which will not cause the operation to fail.
    //
    // See https://docs.microsoft.com/windows/win32/api/heapapi/nf-heapapi-heapfree
    fn HeapFree(hHeap: c::HANDLE, dwFlags: c::DWORD, lpMem: c::LPVOID) -> c::BOOL;
}

// Header containing a pointer to the start of an allocated block.
// SAFETY: size and alignment must be <= `MIN_ALIGN`.
#[repr(C)]
struct Header(*mut u8);

// Allocates a block of optionally zeroed memory for a given `layout`.
// Returns a pointer satisfying the guarantees of `System` about allocated pointers.
#[inline]
unsafe fn allocate(layout: Layout, zeroed: bool) -> *mut u8 {
    let heap = unsafe { GetProcessHeap() };
    if heap.is_null() {
        // Allocation has failed, could not get the current process heap.
        return ptr::null_mut();
    }

    // Allocated memory will be either zeroed or uninitialized.
    let flags = if zeroed { HEAP_ZERO_MEMORY } else { 0 };

    if layout.align() <= MIN_ALIGN {
        // SAFETY: `heap` is a non-null handle returned by `GetProcessHeap`.
        // The returned pointer points to the start of an allocated block.
        unsafe { HeapAlloc(heap, flags, layout.size()) as *mut u8 }
    } else {
        // Allocate extra padding in order to be able to satisfy the alignment.
        let total = layout.align() + layout.size();

        // SAFETY: `heap` is a non-null handle returned by `GetProcessHeap`.
        let ptr = unsafe { HeapAlloc(heap, flags, total) as *mut u8 };
        if ptr.is_null() {
            // Allocation has failed.
            return ptr::null_mut();
        }

        // Create a correctly aligned pointer offset from the start of the allocated block,
        // and write a header before it.

        let offset = layout.align() - (ptr as usize & (layout.align() - 1));
        // SAFETY: `MIN_ALIGN` <= `offset` <= `layout.align()` and the size of the allocated
        // block is `layout.align() + layout.size()`. `aligned` will thus be a correctly aligned
        // pointer inside the allocated block with at least `layout.size()` bytes after it and at
        // least `MIN_ALIGN` bytes of padding before it.
        let aligned = unsafe { ptr.add(offset) };
        // SAFETY: Because the size and alignment of a header is <= `MIN_ALIGN` and `aligned`
        // is aligned to at least `MIN_ALIGN` and has at least `MIN_ALIGN` bytes of padding before
        // it, it is safe to write a header directly before it.
        unsafe { ptr::write((aligned as *mut Header).offset(-1), Header(ptr)) };

        // SAFETY: The returned pointer does not point to the to the start of an allocated block,
        // but there is a header readable directly before it containing the location of the start
        // of the block.
        aligned
    }
}

// All pointers returned by this allocator have, in addition to the guarantees of `GlobalAlloc`, the
// following properties:
//
// If the pointer was allocated or reallocated with a `layout` specifying an alignment <= `MIN_ALIGN`
// the pointer will be aligned to at least `MIN_ALIGN` and point to the start of the allocated block.
//
// If the pointer was allocated or reallocated with a `layout` specifying an alignment > `MIN_ALIGN`
// the pointer will be aligned to the specified alignment and not point to the start of the allocated block.
// Instead there will be a header readable directly before the returned pointer, containing the actual
// location of the start of the block.
#[stable(feature = "alloc_system_type", since = "1.28.0")]
unsafe impl GlobalAlloc for System {
    #[inline]
    unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
        // SAFETY: pointers returned by `allocate` satisfy the guarantees of `System`
        let zeroed = false;
        unsafe { allocate(layout, zeroed) }
    }

    #[inline]
    unsafe fn alloc_zeroed(&self, layout: Layout) -> *mut u8 {
        // SAFETY: pointers returned by `allocate` satisfy the guarantees of `System`
        let zeroed = true;
        unsafe { allocate(layout, zeroed) }
    }

    #[inline]
    unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
        let block = {
            if layout.align() <= MIN_ALIGN {
                ptr
            } else {
                // The location of the start of the block is stored in the padding before `ptr`.

                // SAFETY: Because of the contract of `System`, `ptr` is guaranteed to be non-null
                // and have a header readable directly before it.
                unsafe { ptr::read((ptr as *mut Header).offset(-1)).0 }
            }
        };

        // SAFETY: `block` is a pointer to the start of an allocated block.
        unsafe {
            let err = HeapFree(GetProcessHeap(), 0, block as c::LPVOID);
            debug_assert!(err != 0, "Failed to free heap memory: {}", c::GetLastError());
        }
    }

    #[inline]
    unsafe fn realloc(&self, ptr: *mut u8, layout: Layout, new_size: usize) -> *mut u8 {
        if layout.align() <= MIN_ALIGN {
            let heap = unsafe { GetProcessHeap() };
            if heap.is_null() {
                // Reallocation has failed, could not get the current process heap.
                return ptr::null_mut();
            }

            // SAFETY: `heap` is a non-null handle returned by `GetProcessHeap`,
            // `ptr` is a pointer to the start of an allocated block.
            // The returned pointer points to the start of an allocated block.
            unsafe { HeapReAlloc(heap, 0, ptr as c::LPVOID, new_size) as *mut u8 }
        } else {
            // SAFETY: `realloc_fallback` is implemented using `dealloc` and `alloc`, which will
            // correctly handle `ptr` and return a pointer satisfying the guarantees of `System`
            unsafe { realloc_fallback(self, ptr, layout, new_size) }
        }
    }
}
<|endoftext|>


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

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

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

pub struct Stderr {
    incomplete_utf8: IncompleteUtf8,
}

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

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

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

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

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

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

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

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

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

    write_valid_utf8(handle, utf8)
}

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

pub struct Stderr {
    incomplete_utf8: IncompleteUtf8,
}

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

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

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

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

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

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

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

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

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

    write_valid_utf8(handle, utf8)
}

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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