<commit_before>/// This module contains the following sub-modules:
///     k
///     ks
///     ksl

/// This module contains functions that work under the assumption that lattice
/// momentum is conserved.
pub mod k {
    use fnv::FnvHashMap;
    use num_complex::Complex;

    use blochfunc::{BlochFunc, BlochFuncSet};
    use common::*;
    use ops;

    fn bloch_states<'a>(nx: Dim, ny: Dim, kx: K, ky: K) -> BlochFuncSet {
        let n = nx * ny;
        let mut sieve = vec![true; 2_usize.pow(n.raw_int())];
        let mut bfuncs: Vec<BlochFunc> = Vec::new();
        let phase = |i, j| {
            let r = 1.;
            let ang1 = 2. * PI * (i * kx.raw_int()) as f64 / nx.raw_int() as f64;
            let ang2 = 2. * PI * (j * ky.raw_int()) as f64 / ny.raw_int() as f64;
            Complex::from_polar(&r, &(ang1 + ang2))
        };

        for dec in 0..2_usize.pow(n.raw_int()) {
            if sieve[dec] {
                // if the corresponding entry of dec in "sieve" is not false,
                // we find all translations of dec and put them in a BlochFunc
                // then mark all corresponding entries in "sieve" as false.

                // "decs" is a hashtable that holds vectors whose entries
                // correspond to Bloch function constituent configurations which
                // are mapped to single decimals that represent the leading states.
                let mut decs: FnvHashMap<BinaryBasis, Complex<f64>> =
                    FnvHashMap::default();
                // "new_dec" represents the configuration we are currently iterating
                // over.
                let mut new_dec = BinaryBasis(dec as u64);
                for j in 0..ny.raw_int() {
                    for i in 0..nx.raw_int() {
                        sieve[new_dec.raw_int() as usize] = false;
                        let new_p = match decs.get(&new_dec) {
                            Some(&p) => p + phase(i, j),
                            None => phase(i, j)
                        };
                        decs.insert(new_dec, new_p);
                        new_dec = translate_x(new_dec, nx, ny);
                    }
                    new_dec = translate_y(new_dec, nx, ny);
                }

                let lead = BinaryBasis(dec as u64);
                let norm = decs.values()
                               .into_iter()
                               .map(|&x| x.norm_sqr())
                               .sum::<f64>()
                               .sqrt();

                if norm > 1e-8 {
                    let mut bfunc = BlochFunc { lead, decs, norm };
                    bfuncs.push(bfunc);
                }
            }
        }

        let mut table = BlochFuncSet::create(nx, ny, bfuncs);
        table.sort();
        table
    }

    pub fn h_ss_z(nx: Dim, ny: Dim, kx: K, ky: K, l: I)
                  -> CoordMatrix<CComplex<f64>> {
        let bfuncs = bloch_states(nx, ny, kx, ky);
        let sites = interacting_sites(nx, ny, l);
        ops::ss_z(&sites, &bfuncs)
    }

    pub fn h_ss_xy(nx: Dim, ny: Dim, kx: K, ky: K, l: I)
                   -> CoordMatrix<CComplex<f64>> {
        let bfuncs = bloch_states(nx, ny, kx, ky);
        let sites = interacting_sites(nx, ny, l);
        ops::ss_xy(&sites, &bfuncs)
    }

    pub fn h_ss_ppmm(nx: Dim, ny: Dim, kx: K, ky: K, l: I)
                     -> CoordMatrix<CComplex<f64>> {
        let bfuncs = bloch_states(nx, ny, kx, ky);
        let sites = interacting_sites(nx, ny, l);
        ops::ss_ppmm(&sites, &bfuncs)
    }

    pub fn h_ss_pmz(nx: Dim, ny: Dim, kx: K, ky: K, l: I)
                    -> CoordMatrix<CComplex<f64>> {
        let bfuncs = bloch_states(nx, ny, kx, ky);
        let sites = interacting_sites(nx, ny, l);
        ops::ss_pmz(&sites, &bfuncs)
    }

    pub fn h_sss_chi(nx: Dim, ny: Dim, kx: K, ky: K) -> CoordMatrix<CComplex<f64>> {
        let bfuncs = bloch_states(nx, ny, kx, ky);
        let sites = triangular_vert_sites(nx, ny);
        ops::sss_chi(&sites, &bfuncs)
    }

    pub fn ss_z(nx: Dim, ny: Dim, kx: K, ky: K, l: I) -> CoordMatrix<CComplex<f64>> {
        let bfuncs = bloch_states(nx, ny, kx, ky);
        let sites = all_sites(nx, ny, l);
        ops::ss_z(&sites, &bfuncs)
    }

    pub fn ss_xy(nx: Dim, ny: Dim, kx: K, ky: K, l: I)
                 -> CoordMatrix<CComplex<f64>> {
        let bfuncs = bloch_states(nx, ny, kx, ky);
        let sites = all_sites(nx, ny, l);
        ops::ss_xy(&sites, &bfuncs)
    }

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

        #[test]
        fn bloch_states_test() {
            let nx = Dim(4);
            let ny = Dim(4);
            let kx = K(1);
            let ky = K(3);
            let bfuncs = bloch_states(nx, ny, kx, ky);
            assert_eq!(bfuncs.nonzero, 4080);
        }
    }
}

/// This module contains functions that work under the assumption that lattice
/// momentum and total Sz are conserved.
pub mod ks {
    use fnv::FnvHashMap;
    use num_complex::Complex;

    use blochfunc::{BlochFunc, BlochFuncSet};
    use common::*;
    use ops;

    fn bloch_states<'a>(nx: Dim, ny: Dim, kx: K, ky: K, nup: u32) -> BlochFuncSet {
        let n = nx * ny;

        let sz_basis_states = sz_basis(n, nup);
        let mut szdec_to_ind: FnvHashMap<BinaryBasis, usize> = FnvHashMap::default();
        let mut ind_to_szdec: FnvHashMap<usize, BinaryBasis> = FnvHashMap::default();
        for (i, &bs) in sz_basis_states.iter().enumerate() {
            ind_to_szdec.insert(i, bs);
            szdec_to_ind.insert(bs, i);
        }

        let mut sieve = vec![true; sz_basis_states.len()];
        let mut bfuncs: Vec<BlochFunc> = Vec::new();
        let phase = |i, j| {
            let r = 1.;
            let ang1 = 2. * PI * (i * kx.raw_int()) as f64 / nx.raw_int() as f64;
            let ang2 = 2. * PI * (j * ky.raw_int()) as f64 / ny.raw_int() as f64;
            Complex::from_polar(&r, &(ang1 + ang2))
        };

        for ind in 0..sieve.len() {
            if sieve[ind] {
                // if the corresponding entry of dec in "sieve" is not false,
                // we find all translations of dec and put them in a BlochFunc
                // then mark all corresponding entries in "sieve" as false.

                // "decs" is a hashtable that holds vectors whose entries
                // correspond to Bloch function constituent configurations which
                // are mapped to single decimals that represent the leading states.
                let mut decs: FnvHashMap<BinaryBasis, Complex<f64>> =
                    FnvHashMap::default();
                // "new_dec" represents the configuration we are currently iterating
                // over.
                let dec = *ind_to_szdec.get(&ind).unwrap();
                let mut new_dec = dec;
                let mut new_ind = ind;
                for j in 0..ny.raw_int() {
                    for i in 0..nx.raw_int() {
                        sieve[new_ind as usize] = false;
                        let new_p = match decs.get(&new_dec) {
                            Some(&p) => p + phase(i, j),
                            None => phase(i, j)
                        };
                        decs.insert(new_dec, new_p);
                        new_dec = translate_x(new_dec, nx, ny);
                        new_ind = *szdec_to_ind.get(&new_dec).unwrap() as usize;
                    }
                    new_dec = translate_y(new_dec, nx, ny);
                }

                let lead = dec;
                let norm = decs.values()
                               .into_iter()
                               .map(|&x| x.norm_sqr())
                               .sum::<f64>()
                               .sqrt();

                if norm > 1e-8 {
                    let mut bfunc = BlochFunc { lead, decs, norm };
                    bfuncs.push(bfunc);
                }
            }
        }

        let mut table = BlochFuncSet::create(nx, ny, bfuncs);
        table.sort();
        table
    }

    pub fn h_ss_z(nx: Dim, ny: Dim, kx: K, ky: K, nup: u32, l: I)
                  -> CoordMatrix<CComplex<f64>> {
        let bfuncs = bloch_states(nx, ny, kx, ky, nup);
        let sites = interacting_sites(nx, ny, l);
        ops::ss_z(&sites, &bfuncs)
    }

    pub fn h_ss_xy(nx: Dim, ny: Dim, kx: K, ky: K, nup: u32, l: I)
                   -> CoordMatrix<CComplex<f64>> {
        let bfuncs = bloch_states(nx, ny, kx, ky, nup);
        let sites = interacting_sites(nx, ny, l);
        ops::ss_xy(&sites, &bfuncs)
    }

    pub fn h_sss_chi(nx: Dim, ny: Dim, kx: K, ky: K, nup: u32)
                     -> CoordMatrix<CComplex<f64>> {
        let bfuncs = bloch_states(nx, ny, kx, ky, nup);
        let sites = triangular_vert_sites(nx, ny);
        ops::sss_chi(&sites, &bfuncs)
    }

    pub fn ss_z(nx: Dim, ny: Dim, kx: K, ky: K, nup: u32, l: I)
                -> CoordMatrix<CComplex<f64>> {
        let bfuncs = bloch_states(nx, ny, kx, ky, nup);
        let sites = all_sites(nx, ny, l);
        ops::ss_z(&sites, &bfuncs)
    }

    pub fn ss_xy(nx: Dim, ny: Dim, kx: K, ky: K, nup: u32, l: I)
                 -> CoordMatrix<CComplex<f64>> {
        let bfuncs = bloch_states(nx, ny, kx, ky, nup);
        let sites = all_sites(nx, ny, l);
        ops::ss_xy(&sites, &bfuncs)
    }
}

/// This module contains functions that work under the assumption that lattice
/// momentum, total Sz and total angular momentum are conserved. This is a
/// special case where kx and ky are 0 which makes k and S compatible quantum
/// numbers.
pub mod ksl {
    use fnv::FnvHashMap;
    use num_complex::Complex;

    use blochfunc::{BlochFunc, BlochFuncSet};
    use common::*;
    use ops;

    fn bloch_states<'a>(nx: Dim, ny: Dim, kx: K, ky: K, nup: u32) -> BlochFuncSet {
        let n = nx * ny;

        let sz_basis_states = sz_basis(n, nup);
        let mut szdec_to_ind: FnvHashMap<BinaryBasis, usize> = FnvHashMap::default();
        let mut ind_to_szdec: FnvHashMap<usize, BinaryBasis> = FnvHashMap::default();
        for (i, &bs) in sz_basis_states.iter().enumerate() {
            ind_to_szdec.insert(i, bs);
            szdec_to_ind.insert(bs, i);
        }

        let mut sieve = vec![true; sz_basis_states.len()];
        let mut bfuncs: Vec<BlochFunc> = Vec::new();
        let phase = |i, j| {
            let r = 1.;
            let ang1 = 2. * PI * (i * kx.raw_int()) as f64 / nx.raw_int() as f64;
            let ang2 = 2. * PI * (j * ky.raw_int()) as f64 / ny.raw_int() as f64;
            Complex::from_polar(&r, &(ang1 + ang2))
        };

        for ind in 0..sieve.len() {
            if sieve[ind] {
                // if the corresponding entry of dec in "sieve" is not false,
                // we find all translations of dec and put them in a BlochFunc
                // then mark all corresponding entries in "sieve" as false.

                // "decs" is a hashtable that holds vectors whose entries
                // correspond to Bloch function constituent configurations which
                // are mapped to single decimals that represent the leading states.
                let mut decs: FnvHashMap<BinaryBasis, Complex<f64>> =
                    FnvHashMap::default();
                // "new_dec" represents the configuration we are currently iterating
                // over.
                let dec = *ind_to_szdec.get(&ind).unwrap();
                let mut new_dec = dec;
                let mut new_ind = ind;
                for j in 0..ny.raw_int() {
                    for i in 0..nx.raw_int() {
                        sieve[new_ind as usize] = false;
                        let new_p = match decs.get(&new_dec) {
                            Some(&p) => p + phase(i, j),
                            None => phase(i, j)
                        };
                        decs.insert(new_dec, new_p);
                        new_dec = translate_x(new_dec, nx, ny);
                        new_ind = *szdec_to_ind.get(&new_dec).unwrap() as usize;
                    }
                    new_dec = translate_y(new_dec, nx, ny);
                }

                let lead = dec;
                let norm = decs.values()
                               .into_iter()
                               .map(|&x| x.norm_sqr())
                               .sum::<f64>()
                               .sqrt();

                if norm > 1e-8 {
                    let mut bfunc = BlochFunc { lead, decs, norm };
                    bfuncs.push(bfunc);
                }
            }
        }
        let mut table = BlochFuncSet::create(nx, ny, bfuncs);
        table.sort();
        table
    }
}
<commit_msg>Removed dead code<commit_after>/// This module contains the following sub-modules:
///     k
///     ks
///     ksl

/// This module contains functions that work under the assumption that lattice
/// momentum is conserved.
pub mod k {
    use fnv::FnvHashMap;
    use num_complex::Complex;

    use blochfunc::{BlochFunc, BlochFuncSet};
    use common::*;
    use ops;

    fn bloch_states<'a>(nx: Dim, ny: Dim, kx: K, ky: K) -> BlochFuncSet {
        let n = nx * ny;
        let mut sieve = vec![true; 2_usize.pow(n.raw_int())];
        let mut bfuncs: Vec<BlochFunc> = Vec::new();
        let phase = |i, j| {
            let r = 1.;
            let ang1 = 2. * PI * (i * kx.raw_int()) as f64 / nx.raw_int() as f64;
            let ang2 = 2. * PI * (j * ky.raw_int()) as f64 / ny.raw_int() as f64;
            Complex::from_polar(&r, &(ang1 + ang2))
        };

        for dec in 0..2_usize.pow(n.raw_int()) {
            if sieve[dec] {
                // if the corresponding entry of dec in "sieve" is not false,
                // we find all translations of dec and put them in a BlochFunc
                // then mark all corresponding entries in "sieve" as false.

                // "decs" is a hashtable that holds vectors whose entries
                // correspond to Bloch function constituent configurations which
                // are mapped to single decimals that represent the leading states.
                let mut decs: FnvHashMap<BinaryBasis, Complex<f64>> =
                    FnvHashMap::default();
                // "new_dec" represents the configuration we are currently iterating
                // over.
                let mut new_dec = BinaryBasis(dec as u64);
                for j in 0..ny.raw_int() {
                    for i in 0..nx.raw_int() {
                        sieve[new_dec.raw_int() as usize] = false;
                        let new_p = match decs.get(&new_dec) {
                            Some(&p) => p + phase(i, j),
                            None => phase(i, j)
                        };
                        decs.insert(new_dec, new_p);
                        new_dec = translate_x(new_dec, nx, ny);
                    }
                    new_dec = translate_y(new_dec, nx, ny);
                }

                let lead = BinaryBasis(dec as u64);
                let norm = decs.values()
                               .into_iter()
                               .map(|&x| x.norm_sqr())
                               .sum::<f64>()
                               .sqrt();

                if norm > 1e-8 {
                    let mut bfunc = BlochFunc { lead, decs, norm };
                    bfuncs.push(bfunc);
                }
            }
        }

        let mut table = BlochFuncSet::create(nx, ny, bfuncs);
        table.sort();
        table
    }

    pub fn h_ss_z(nx: Dim, ny: Dim, kx: K, ky: K, l: I)
                  -> CoordMatrix<CComplex<f64>> {
        let bfuncs = bloch_states(nx, ny, kx, ky);
        let sites = interacting_sites(nx, ny, l);
        ops::ss_z(&sites, &bfuncs)
    }

    pub fn h_ss_xy(nx: Dim, ny: Dim, kx: K, ky: K, l: I)
                   -> CoordMatrix<CComplex<f64>> {
        let bfuncs = bloch_states(nx, ny, kx, ky);
        let sites = interacting_sites(nx, ny, l);
        ops::ss_xy(&sites, &bfuncs)
    }

    pub fn h_ss_ppmm(nx: Dim, ny: Dim, kx: K, ky: K, l: I)
                     -> CoordMatrix<CComplex<f64>> {
        let bfuncs = bloch_states(nx, ny, kx, ky);
        let sites = interacting_sites(nx, ny, l);
        ops::ss_ppmm(&sites, &bfuncs)
    }

    pub fn h_ss_pmz(nx: Dim, ny: Dim, kx: K, ky: K, l: I)
                    -> CoordMatrix<CComplex<f64>> {
        let bfuncs = bloch_states(nx, ny, kx, ky);
        let sites = interacting_sites(nx, ny, l);
        ops::ss_pmz(&sites, &bfuncs)
    }

    pub fn h_sss_chi(nx: Dim, ny: Dim, kx: K, ky: K) -> CoordMatrix<CComplex<f64>> {
        let bfuncs = bloch_states(nx, ny, kx, ky);
        let sites = triangular_vert_sites(nx, ny);
        ops::sss_chi(&sites, &bfuncs)
    }

    pub fn ss_z(nx: Dim, ny: Dim, kx: K, ky: K, l: I) -> CoordMatrix<CComplex<f64>> {
        let bfuncs = bloch_states(nx, ny, kx, ky);
        let sites = all_sites(nx, ny, l);
        ops::ss_z(&sites, &bfuncs)
    }

    pub fn ss_xy(nx: Dim, ny: Dim, kx: K, ky: K, l: I)
                 -> CoordMatrix<CComplex<f64>> {
        let bfuncs = bloch_states(nx, ny, kx, ky);
        let sites = all_sites(nx, ny, l);
        ops::ss_xy(&sites, &bfuncs)
    }

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

        #[test]
        fn bloch_states_test() {
            let nx = Dim(4);
            let ny = Dim(4);
            let kx = K(1);
            let ky = K(3);
            let bfuncs = bloch_states(nx, ny, kx, ky);
            assert_eq!(bfuncs.nonzero, 4080);
        }
    }
}

/// This module contains functions that work under the assumption that lattice
/// momentum and total Sz are conserved.
pub mod ks {
    use fnv::FnvHashMap;
    use num_complex::Complex;

    use blochfunc::{BlochFunc, BlochFuncSet};
    use common::*;
    use ops;

    fn bloch_states<'a>(nx: Dim, ny: Dim, kx: K, ky: K, nup: u32) -> BlochFuncSet {
        let n = nx * ny;

        let sz_basis_states = sz_basis(n, nup);
        let mut szdec_to_ind: FnvHashMap<BinaryBasis, usize> = FnvHashMap::default();
        let mut ind_to_szdec: FnvHashMap<usize, BinaryBasis> = FnvHashMap::default();
        for (i, &bs) in sz_basis_states.iter().enumerate() {
            ind_to_szdec.insert(i, bs);
            szdec_to_ind.insert(bs, i);
        }

        let mut sieve = vec![true; sz_basis_states.len()];
        let mut bfuncs: Vec<BlochFunc> = Vec::new();
        let phase = |i, j| {
            let r = 1.;
            let ang1 = 2. * PI * (i * kx.raw_int()) as f64 / nx.raw_int() as f64;
            let ang2 = 2. * PI * (j * ky.raw_int()) as f64 / ny.raw_int() as f64;
            Complex::from_polar(&r, &(ang1 + ang2))
        };

        for ind in 0..sieve.len() {
            if sieve[ind] {
                // if the corresponding entry of dec in "sieve" is not false,
                // we find all translations of dec and put them in a BlochFunc
                // then mark all corresponding entries in "sieve" as false.

                // "decs" is a hashtable that holds vectors whose entries
                // correspond to Bloch function constituent configurations which
                // are mapped to single decimals that represent the leading states.
                let mut decs: FnvHashMap<BinaryBasis, Complex<f64>> =
                    FnvHashMap::default();
                // "new_dec" represents the configuration we are currently iterating
                // over.
                let dec = *ind_to_szdec.get(&ind).unwrap();
                let mut new_dec = dec;
                let mut new_ind = ind;
                for j in 0..ny.raw_int() {
                    for i in 0..nx.raw_int() {
                        sieve[new_ind as usize] = false;
                        let new_p = match decs.get(&new_dec) {
                            Some(&p) => p + phase(i, j),
                            None => phase(i, j)
                        };
                        decs.insert(new_dec, new_p);
                        new_dec = translate_x(new_dec, nx, ny);
                        new_ind = *szdec_to_ind.get(&new_dec).unwrap() as usize;
                    }
                    new_dec = translate_y(new_dec, nx, ny);
                }

                let lead = dec;
                let norm = decs.values()
                               .into_iter()
                               .map(|&x| x.norm_sqr())
                               .sum::<f64>()
                               .sqrt();

                if norm > 1e-8 {
                    let mut bfunc = BlochFunc { lead, decs, norm };
                    bfuncs.push(bfunc);
                }
            }
        }

        let mut table = BlochFuncSet::create(nx, ny, bfuncs);
        table.sort();
        table
    }

    pub fn h_ss_z(nx: Dim, ny: Dim, kx: K, ky: K, nup: u32, l: I)
                  -> CoordMatrix<CComplex<f64>> {
        let bfuncs = bloch_states(nx, ny, kx, ky, nup);
        let sites = interacting_sites(nx, ny, l);
        ops::ss_z(&sites, &bfuncs)
    }

    pub fn h_ss_xy(nx: Dim, ny: Dim, kx: K, ky: K, nup: u32, l: I)
                   -> CoordMatrix<CComplex<f64>> {
        let bfuncs = bloch_states(nx, ny, kx, ky, nup);
        let sites = interacting_sites(nx, ny, l);
        ops::ss_xy(&sites, &bfuncs)
    }

    pub fn h_sss_chi(nx: Dim, ny: Dim, kx: K, ky: K, nup: u32)
                     -> CoordMatrix<CComplex<f64>> {
        let bfuncs = bloch_states(nx, ny, kx, ky, nup);
        let sites = triangular_vert_sites(nx, ny);
        ops::sss_chi(&sites, &bfuncs)
    }

    pub fn ss_z(nx: Dim, ny: Dim, kx: K, ky: K, nup: u32, l: I)
                -> CoordMatrix<CComplex<f64>> {
        let bfuncs = bloch_states(nx, ny, kx, ky, nup);
        let sites = all_sites(nx, ny, l);
        ops::ss_z(&sites, &bfuncs)
    }

    pub fn ss_xy(nx: Dim, ny: Dim, kx: K, ky: K, nup: u32, l: I)
                 -> CoordMatrix<CComplex<f64>> {
        let bfuncs = bloch_states(nx, ny, kx, ky, nup);
        let sites = all_sites(nx, ny, l);
        ops::ss_xy(&sites, &bfuncs)
    }
}
<|endoftext|>


<commit_before>use std::{env,
          ffi::OsString,
          fs as stdfs,
          path::PathBuf};

use crate::{common::ui::UI,
            hcore::{crypto::CACHE_KEY_PATH_ENV_VAR,
                    env as henv,
                    fs}};

use crate::{config,
            error::Result,
            BLDR_URL_ENVVAR,
            CTL_SECRET_ENVVAR,
            ORIGIN_ENVVAR};

use habitat_core::AUTH_TOKEN_ENVVAR;

pub const ARTIFACT_PATH_ENVVAR: &str = "ARTIFACT_PATH";

const STUDIO_CMD: &str = "hab-studio";
const STUDIO_CMD_ENVVAR: &str = "HAB_STUDIO_BINARY";
const STUDIO_PACKAGE_IDENT: &str = "core/hab-studio";

fn set_env_var_from_config(env_var: &str, config_val: Option<String>, sensitive: bool) {
    if henv::var(env_var).is_err() {
        if let Some(val) = config_val {
            if sensitive {
              debug!("Setting {}=REDACTED (sensitive) via config file", env_var)
            }
            else {
              debug!("Setting {}={} via config file", env_var, val)
            }
            env::set_var(env_var, val);
        }
    }
}

pub fn start(ui: &mut UI, args: &[OsString]) -> Result<()> {
    let config = config::load()?;

    set_env_var_from_config(AUTH_TOKEN_ENVVAR, config.auth_token, true);
    set_env_var_from_config(BLDR_URL_ENVVAR, config.bldr_url, false);
    set_env_var_from_config(CTL_SECRET_ENVVAR, config.ctl_secret, true);
    set_env_var_from_config(ORIGIN_ENVVAR, config.origin, false);

    if henv::var(CACHE_KEY_PATH_ENV_VAR).is_err() {
        let path = fs::cache_key_path(None::<&str>);
        debug!("Setting {}={}", CACHE_KEY_PATH_ENV_VAR, path.display());
        env::set_var(CACHE_KEY_PATH_ENV_VAR, &path);
    };

    let artifact_path = match henv::var(ARTIFACT_PATH_ENVVAR) {
        Ok(p) => PathBuf::from(p),
        Err(_) => {
            let path = fs::cache_artifact_path(None::<&str>);
            debug!("Setting {}={}", ARTIFACT_PATH_ENVVAR, path.display());
            env::set_var(ARTIFACT_PATH_ENVVAR, &path);
            path
        }
    };
    if !artifact_path.is_dir() {
        debug!("Creating artifact_path at: {}", artifact_path.display());
        stdfs::create_dir_all(&artifact_path)?;
    }

    inner::start(ui, args)
}

#[cfg(target_os = "linux")]
mod inner {
    use crate::{command::studio::docker,
                common::ui::{UIWriter,
                             UI},
                error::{Error,
                        Result},
                exec,
                hcore::{crypto::init,
                        env as henv,
                        fs::{am_i_root,
                             find_command},
                        os::process,
                        package::{PackageIdent,
                                  PackageInstall},
                        users::linux as group},
                VERSION};
    use std::{env,
              ffi::OsString,
              path::PathBuf,
              str::FromStr};

    const SUDO_CMD: &str = "sudo";

    pub fn start(ui: &mut UI, args: &[OsString]) -> Result<()> {
        rerun_with_sudo_if_needed(ui, &args)?;
        if is_docker_studio(&args) {
            docker::start_docker_studio(ui, args)
        } else {
            let command = match henv::var(super::STUDIO_CMD_ENVVAR) {
                Ok(command) => PathBuf::from(command),
                Err(_) => {
                    init();
                    let version: Vec<&str> = VERSION.split('/').collect();
                    let ident = PackageIdent::from_str(&format!("{}/{}",
                                                                super::STUDIO_PACKAGE_IDENT,
                                                                version[0]))?;
                    // This is a duplicate of the code in `hab pkg exec` and
                    // should be refactored as part of or after:
                    // https://github.com/habitat-sh/habitat/issues/6633
                    // https://github.com/habitat-sh/habitat/issues/6634
                    let pkg_install = PackageInstall::load(&ident, None)?;
                    let cmd_env = pkg_install.environment_for_command()?;
                    for (key, value) in cmd_env.into_iter() {
                        debug!("Setting: {}='{}'", key, value);
                        env::set_var(key, value);
                    }

                    let mut display_args = super::STUDIO_CMD.to_string();
                    for arg in args {
                        display_args.push(' ');
                        display_args.push_str(arg.to_string_lossy().as_ref());
                    }
                    debug!("Running: {}", display_args);

                    exec::command_from_min_pkg(ui, super::STUDIO_CMD, &ident)?
                }
            };

            if let Some(cmd) = find_command(command.to_string_lossy().as_ref()) {
                process::become_command(cmd, args)?;
                Ok(())
            } else {
                Err(Error::ExecCommandNotFound(command))
            }
        }
    }

    fn is_docker_studio(args: &[OsString]) -> bool {
        if cfg!(not(target_os = "linux")) {
            return false;
        }

        for arg in args.iter() {
            let str_arg = arg.to_string_lossy();
            if str_arg == "-D" {
                return true;
            }
        }

        false
    }

    fn has_docker_group() -> bool {
        let current_user = group::get_current_username().unwrap();
        let docker_members = group::get_members_by_groupname("docker");
        docker_members.map_or(false, |d| d.contains(&current_user))
    }

    fn rerun_with_sudo_if_needed(ui: &mut UI, args: &[OsString]) -> Result<()> {
        // If I have root permissions or if I am executing a docker studio
        // and have the appropriate group - early return, we are done.
        if am_i_root() || (is_docker_studio(args) && has_docker_group()) {
            return Ok(());
        }

        // Otherwise we will try to re-run this program using `sudo`
        match find_command(SUDO_CMD) {
            Some(sudo_prog) => {
                let mut args: Vec<OsString> = vec!["-p".into(),
                                                   "[sudo hab-studio] password for %u: ".into(),
                                                   "-E".into(),];
                args.append(&mut env::args_os().collect());
                process::become_command(sudo_prog, &args)?;
                Ok(())
            }
            None => {
                ui.warn(format!("Could not find the `{}' command, is it in your PATH?",
                                SUDO_CMD))?;
                ui.warn("Running Habitat Studio requires root or administrator privileges. \
                         Please retry this command as a super user or use a privilege-granting \
                         facility such as sudo.")?;
                ui.br()?;
                Err(Error::RootRequired)
            }
        }
    }
}

#[cfg(not(target_os = "linux"))]
mod inner {
    use crate::{command::studio::docker,
                common::ui::UI,
                error::{Error,
                        Result},
                exec,
                hcore::{crypto::init,
                        env as henv,
                        fs::find_command,
                        os::process,
                        package::PackageIdent},
                VERSION};
    use std::{ffi::OsString,
              path::PathBuf,
              str::FromStr};

    pub fn start(_ui: &mut UI, args: &[OsString]) -> Result<()> {
        if is_windows_studio(&args) {
            start_windows_studio(_ui, args)
        } else {
            docker::start_docker_studio(_ui, args)
        }
    }

    pub fn start_windows_studio(ui: &mut UI, args: &[OsString]) -> Result<()> {
        let command = match henv::var(super::STUDIO_CMD_ENVVAR) {
            Ok(command) => PathBuf::from(command),
            Err(_) => {
                init();
                let version: Vec<&str> = VERSION.split('/').collect();
                let ident = PackageIdent::from_str(&format!("{}/{}",
                                                            super::STUDIO_PACKAGE_IDENT,
                                                            version[0]))?;
                exec::command_from_min_pkg(ui, super::STUDIO_CMD, &ident)?
            }
        };

        if let Some(cmd) = find_command(command.to_string_lossy().as_ref()) {
            process::become_command(cmd, args)?;
        } else {
            return Err(Error::ExecCommandNotFound(command));
        }
        Ok(())
    }

    fn is_windows_studio(args: &[OsString]) -> bool {
        if cfg!(not(target_os = "windows")) {
            return false;
        }

        for arg in args.iter() {
            let str_arg = arg.to_string_lossy();
            if str_arg == "-D" {
                return false;
            }
        }

        // -w/--windows is deprecated and should be removed in a post 0.64.0 release
        for arg in args.iter() {
            let str_arg = arg.to_string_lossy().to_lowercase();
            if str_arg == "--windows" || str_arg == "-w" {
                return true;
            }
        }

        true
    }
}
<commit_msg>rustfmt cleanup<commit_after>use std::{env,
          ffi::OsString,
          fs as stdfs,
          path::PathBuf};

use crate::{common::ui::UI,
            hcore::{crypto::CACHE_KEY_PATH_ENV_VAR,
                    env as henv,
                    fs}};

use crate::{config,
            error::Result,
            BLDR_URL_ENVVAR,
            CTL_SECRET_ENVVAR,
            ORIGIN_ENVVAR};

use habitat_core::AUTH_TOKEN_ENVVAR;

pub const ARTIFACT_PATH_ENVVAR: &str = "ARTIFACT_PATH";

const STUDIO_CMD: &str = "hab-studio";
const STUDIO_CMD_ENVVAR: &str = "HAB_STUDIO_BINARY";
const STUDIO_PACKAGE_IDENT: &str = "core/hab-studio";

fn set_env_var_from_config(env_var: &str, config_val: Option<String>, sensitive: bool) {
    if henv::var(env_var).is_err() {
        if let Some(val) = config_val {
            if sensitive {
                debug!("Setting {}=REDACTED (sensitive) via config file", env_var)
            } else {
                debug!("Setting {}={} via config file", env_var, val)
            }
            env::set_var(env_var, val);
        }
    }
}

pub fn start(ui: &mut UI, args: &[OsString]) -> Result<()> {
    let config = config::load()?;

    set_env_var_from_config(AUTH_TOKEN_ENVVAR, config.auth_token, true);
    set_env_var_from_config(BLDR_URL_ENVVAR, config.bldr_url, false);
    set_env_var_from_config(CTL_SECRET_ENVVAR, config.ctl_secret, true);
    set_env_var_from_config(ORIGIN_ENVVAR, config.origin, false);

    if henv::var(CACHE_KEY_PATH_ENV_VAR).is_err() {
        let path = fs::cache_key_path(None::<&str>);
        debug!("Setting {}={}", CACHE_KEY_PATH_ENV_VAR, path.display());
        env::set_var(CACHE_KEY_PATH_ENV_VAR, &path);
    };

    let artifact_path = match henv::var(ARTIFACT_PATH_ENVVAR) {
        Ok(p) => PathBuf::from(p),
        Err(_) => {
            let path = fs::cache_artifact_path(None::<&str>);
            debug!("Setting {}={}", ARTIFACT_PATH_ENVVAR, path.display());
            env::set_var(ARTIFACT_PATH_ENVVAR, &path);
            path
        }
    };
    if !artifact_path.is_dir() {
        debug!("Creating artifact_path at: {}", artifact_path.display());
        stdfs::create_dir_all(&artifact_path)?;
    }

    inner::start(ui, args)
}

#[cfg(target_os = "linux")]
mod inner {
    use crate::{command::studio::docker,
                common::ui::{UIWriter,
                             UI},
                error::{Error,
                        Result},
                exec,
                hcore::{crypto::init,
                        env as henv,
                        fs::{am_i_root,
                             find_command},
                        os::process,
                        package::{PackageIdent,
                                  PackageInstall},
                        users::linux as group},
                VERSION};
    use std::{env,
              ffi::OsString,
              path::PathBuf,
              str::FromStr};

    const SUDO_CMD: &str = "sudo";

    pub fn start(ui: &mut UI, args: &[OsString]) -> Result<()> {
        rerun_with_sudo_if_needed(ui, &args)?;
        if is_docker_studio(&args) {
            docker::start_docker_studio(ui, args)
        } else {
            let command = match henv::var(super::STUDIO_CMD_ENVVAR) {
                Ok(command) => PathBuf::from(command),
                Err(_) => {
                    init();
                    let version: Vec<&str> = VERSION.split('/').collect();
                    let ident = PackageIdent::from_str(&format!("{}/{}",
                                                                super::STUDIO_PACKAGE_IDENT,
                                                                version[0]))?;
                    // This is a duplicate of the code in `hab pkg exec` and
                    // should be refactored as part of or after:
                    // https://github.com/habitat-sh/habitat/issues/6633
                    // https://github.com/habitat-sh/habitat/issues/6634
                    let pkg_install = PackageInstall::load(&ident, None)?;
                    let cmd_env = pkg_install.environment_for_command()?;
                    for (key, value) in cmd_env.into_iter() {
                        debug!("Setting: {}='{}'", key, value);
                        env::set_var(key, value);
                    }

                    let mut display_args = super::STUDIO_CMD.to_string();
                    for arg in args {
                        display_args.push(' ');
                        display_args.push_str(arg.to_string_lossy().as_ref());
                    }
                    debug!("Running: {}", display_args);

                    exec::command_from_min_pkg(ui, super::STUDIO_CMD, &ident)?
                }
            };

            if let Some(cmd) = find_command(command.to_string_lossy().as_ref()) {
                process::become_command(cmd, args)?;
                Ok(())
            } else {
                Err(Error::ExecCommandNotFound(command))
            }
        }
    }

    fn is_docker_studio(args: &[OsString]) -> bool {
        if cfg!(not(target_os = "linux")) {
            return false;
        }

        for arg in args.iter() {
            let str_arg = arg.to_string_lossy();
            if str_arg == "-D" {
                return true;
            }
        }

        false
    }

    fn has_docker_group() -> bool {
        let current_user = group::get_current_username().unwrap();
        let docker_members = group::get_members_by_groupname("docker");
        docker_members.map_or(false, |d| d.contains(&current_user))
    }

    fn rerun_with_sudo_if_needed(ui: &mut UI, args: &[OsString]) -> Result<()> {
        // If I have root permissions or if I am executing a docker studio
        // and have the appropriate group - early return, we are done.
        if am_i_root() || (is_docker_studio(args) && has_docker_group()) {
            return Ok(());
        }

        // Otherwise we will try to re-run this program using `sudo`
        match find_command(SUDO_CMD) {
            Some(sudo_prog) => {
                let mut args: Vec<OsString> = vec!["-p".into(),
                                                   "[sudo hab-studio] password for %u: ".into(),
                                                   "-E".into(),];
                args.append(&mut env::args_os().collect());
                process::become_command(sudo_prog, &args)?;
                Ok(())
            }
            None => {
                ui.warn(format!("Could not find the `{}' command, is it in your PATH?",
                                SUDO_CMD))?;
                ui.warn("Running Habitat Studio requires root or administrator privileges. \
                         Please retry this command as a super user or use a privilege-granting \
                         facility such as sudo.")?;
                ui.br()?;
                Err(Error::RootRequired)
            }
        }
    }
}

#[cfg(not(target_os = "linux"))]
mod inner {
    use crate::{command::studio::docker,
                common::ui::UI,
                error::{Error,
                        Result},
                exec,
                hcore::{crypto::init,
                        env as henv,
                        fs::find_command,
                        os::process,
                        package::PackageIdent},
                VERSION};
    use std::{ffi::OsString,
              path::PathBuf,
              str::FromStr};

    pub fn start(_ui: &mut UI, args: &[OsString]) -> Result<()> {
        if is_windows_studio(&args) {
            start_windows_studio(_ui, args)
        } else {
            docker::start_docker_studio(_ui, args)
        }
    }

    pub fn start_windows_studio(ui: &mut UI, args: &[OsString]) -> Result<()> {
        let command = match henv::var(super::STUDIO_CMD_ENVVAR) {
            Ok(command) => PathBuf::from(command),
            Err(_) => {
                init();
                let version: Vec<&str> = VERSION.split('/').collect();
                let ident = PackageIdent::from_str(&format!("{}/{}",
                                                            super::STUDIO_PACKAGE_IDENT,
                                                            version[0]))?;
                exec::command_from_min_pkg(ui, super::STUDIO_CMD, &ident)?
            }
        };

        if let Some(cmd) = find_command(command.to_string_lossy().as_ref()) {
            process::become_command(cmd, args)?;
        } else {
            return Err(Error::ExecCommandNotFound(command));
        }
        Ok(())
    }

    fn is_windows_studio(args: &[OsString]) -> bool {
        if cfg!(not(target_os = "windows")) {
            return false;
        }

        for arg in args.iter() {
            let str_arg = arg.to_string_lossy();
            if str_arg == "-D" {
                return false;
            }
        }

        // -w/--windows is deprecated and should be removed in a post 0.64.0 release
        for arg in args.iter() {
            let str_arg = arg.to_string_lossy().to_lowercase();
            if str_arg == "--windows" || str_arg == "-w" {
                return true;
            }
        }

        true
    }
}
<|endoftext|>


<commit_before>use std::{env,
          ffi::OsString,
          fs as stdfs,
          path::PathBuf};

use crate::{common::ui::UI,
            hcore::{crypto::CACHE_KEY_PATH_ENV_VAR,
                    env as henv,
                    fs}};

use crate::{config,
            error::Result,
            BLDR_URL_ENVVAR,
            CTL_SECRET_ENVVAR,
            ORIGIN_ENVVAR};

use habitat_core::AUTH_TOKEN_ENVVAR;

pub const ARTIFACT_PATH_ENVVAR: &str = "ARTIFACT_PATH";

const STUDIO_CMD: &str = "hab-studio";
const STUDIO_CMD_ENVVAR: &str = "HAB_STUDIO_BINARY";
const STUDIO_PACKAGE_IDENT: &str = "core/hab-studio";

#[derive(Clone, Copy)]
enum Sensitivity {
    PrintValue,
    NoPrintValue,
}

fn set_env_var_from_config(env_var: &str, config_val: Option<String>, sensitive: Sensitivity) {
    if henv::var(env_var).is_err() {
        if let Some(val) = config_val {
            match sensitive {
                Sensitivity::NoPrintValue => {
                    debug!("Setting {}=REDACTED (sensitive) via config file", env_var)
                }
                Sensitivity::PrintValue => debug!("Setting {}={} via config file", env_var, val),
            }
            env::set_var(env_var, val);
        }
    }
}

pub fn start(ui: &mut UI, args: &[OsString]) -> Result<()> {
    let config = config::load()?;

    set_env_var_from_config(AUTH_TOKEN_ENVVAR,
                            config.auth_token,
                            Sensitivity::NoPrintValue);
    set_env_var_from_config(BLDR_URL_ENVVAR, config.bldr_url, Sensitivity::PrintValue);
    set_env_var_from_config(CTL_SECRET_ENVVAR,
                            config.ctl_secret,
                            Sensitivity::NoPrintValue);
    set_env_var_from_config(ORIGIN_ENVVAR, config.origin, Sensitivity::PrintValue);

    if henv::var(CACHE_KEY_PATH_ENV_VAR).is_err() {
        let path = fs::cache_key_path(None::<&str>);
        debug!("Setting {}={}", CACHE_KEY_PATH_ENV_VAR, path.display());
        env::set_var(CACHE_KEY_PATH_ENV_VAR, &path);
    };

    let artifact_path = match henv::var(ARTIFACT_PATH_ENVVAR) {
        Ok(p) => PathBuf::from(p),
        Err(_) => {
            let path = fs::cache_artifact_path(None::<&str>);
            debug!("Setting {}={}", ARTIFACT_PATH_ENVVAR, path.display());
            env::set_var(ARTIFACT_PATH_ENVVAR, &path);
            path
        }
    };
    if !artifact_path.is_dir() {
        debug!("Creating artifact_path at: {}", artifact_path.display());
        stdfs::create_dir_all(&artifact_path)?;
    }

    inner::start(ui, args)
}

#[cfg(target_os = "linux")]
mod inner {
    use crate::{command::studio::docker,
                common::ui::{UIWriter,
                             UI},
                error::{Error,
                        Result},
                exec,
                hcore::{crypto::init,
                        env as henv,
                        fs::{am_i_root,
                             find_command},
                        os::process,
                        package::{PackageIdent,
                                  PackageInstall},
                        users::linux as group},
                VERSION};
    use std::{env,
              ffi::OsString,
              path::PathBuf,
              str::FromStr};

    const SUDO_CMD: &str = "sudo";

    pub fn start(ui: &mut UI, args: &[OsString]) -> Result<()> {
        rerun_with_sudo_if_needed(ui, &args)?;
        if is_docker_studio(&args) {
            docker::start_docker_studio(ui, args)
        } else {
            let command = match henv::var(super::STUDIO_CMD_ENVVAR) {
                Ok(command) => PathBuf::from(command),
                Err(_) => {
                    init();
                    let version: Vec<&str> = VERSION.split('/').collect();
                    let ident = PackageIdent::from_str(&format!("{}/{}",
                                                                super::STUDIO_PACKAGE_IDENT,
                                                                version[0]))?;
                    let command = exec::command_from_min_pkg(ui, super::STUDIO_CMD, &ident)?;
                    // This is a duplicate of the code in `hab pkg exec` and
                    // should be refactored as part of or after:
                    // https://github.com/habitat-sh/habitat/issues/6633
                    // https://github.com/habitat-sh/habitat/issues/6634
                    let pkg_install = PackageInstall::load(&ident, None)?;
                    let cmd_env = pkg_install.environment_for_command()?;
                    for (key, value) in cmd_env.into_iter() {
                        debug!("Setting: {}='{}'", key, value);
                        env::set_var(key, value);
                    }

                    let mut display_args = super::STUDIO_CMD.to_string();
                    for arg in args {
                        display_args.push(' ');
                        display_args.push_str(arg.to_string_lossy().as_ref());
                    }
                    debug!("Running: {}", display_args);

                    command
                }
            };

            if let Some(cmd) = find_command(command.to_string_lossy().as_ref()) {
                process::become_command(cmd, args)?;
                Ok(())
            } else {
                Err(Error::ExecCommandNotFound(command))
            }
        }
    }

    fn is_docker_studio(args: &[OsString]) -> bool {
        if cfg!(not(target_os = "linux")) {
            return false;
        }

        for arg in args.iter() {
            let str_arg = arg.to_string_lossy();
            if str_arg == "-D" {
                return true;
            }
        }

        false
    }

    fn has_docker_group() -> bool {
        let current_user = group::get_current_username().unwrap();
        let docker_members = group::get_members_by_groupname("docker");
        docker_members.map_or(false, |d| d.contains(&current_user))
    }

    fn rerun_with_sudo_if_needed(ui: &mut UI, args: &[OsString]) -> Result<()> {
        // If I have root permissions or if I am executing a docker studio
        // and have the appropriate group - early return, we are done.
        if am_i_root() || (is_docker_studio(args) && has_docker_group()) {
            return Ok(());
        }

        // Otherwise we will try to re-run this program using `sudo`
        match find_command(SUDO_CMD) {
            Some(sudo_prog) => {
                let mut args: Vec<OsString> = vec!["-p".into(),
                                                   "[sudo hab-studio] password for %u: ".into(),
                                                   "-E".into(),];
                args.append(&mut env::args_os().collect());
                process::become_command(sudo_prog, &args)?;
                Ok(())
            }
            None => {
                ui.warn(format!("Could not find the `{}' command, is it in your PATH?",
                                SUDO_CMD))?;
                ui.warn("Running Habitat Studio requires root or administrator privileges. \
                         Please retry this command as a super user or use a privilege-granting \
                         facility such as sudo.")?;
                ui.br()?;
                Err(Error::RootRequired)
            }
        }
    }
}

#[cfg(not(target_os = "linux"))]
mod inner {
    use crate::{command::studio::docker,
                common::ui::UI,
                error::{Error,
                        Result},
                exec,
                hcore::{crypto::init,
                        env as henv,
                        fs::find_command,
                        os::process,
                        package::PackageIdent},
                VERSION};
    use std::{ffi::OsString,
              path::PathBuf,
              str::FromStr};

    pub fn start(_ui: &mut UI, args: &[OsString]) -> Result<()> {
        if is_windows_studio(&args) {
            start_windows_studio(_ui, args)
        } else {
            docker::start_docker_studio(_ui, args)
        }
    }

    pub fn start_windows_studio(ui: &mut UI, args: &[OsString]) -> Result<()> {
        let command = match henv::var(super::STUDIO_CMD_ENVVAR) {
            Ok(command) => PathBuf::from(command),
            Err(_) => {
                init();
                let version: Vec<&str> = VERSION.split('/').collect();
                let ident = PackageIdent::from_str(&format!("{}/{}",
                                                            super::STUDIO_PACKAGE_IDENT,
                                                            version[0]))?;
                exec::command_from_min_pkg(ui, super::STUDIO_CMD, &ident)?
            }
        };

        if let Some(cmd) = find_command(command.to_string_lossy().as_ref()) {
            process::become_command(cmd, args)?;
        } else {
            return Err(Error::ExecCommandNotFound(command));
        }
        Ok(())
    }

    fn is_windows_studio(args: &[OsString]) -> bool {
        if cfg!(not(target_os = "windows")) {
            return false;
        }

        for arg in args.iter() {
            let str_arg = arg.to_string_lossy();
            if str_arg == "-D" {
                return false;
            }
        }

        true
    }
}
<commit_msg>[rustfmt] Whitespace<commit_after>use std::{env,
          ffi::OsString,
          fs as stdfs,
          path::PathBuf};

use crate::{common::ui::UI,
            hcore::{crypto::CACHE_KEY_PATH_ENV_VAR,
                    env as henv,
                    fs}};

use crate::{config,
            error::Result,
            BLDR_URL_ENVVAR,
            CTL_SECRET_ENVVAR,
            ORIGIN_ENVVAR};

use habitat_core::AUTH_TOKEN_ENVVAR;

pub const ARTIFACT_PATH_ENVVAR: &str = "ARTIFACT_PATH";

const STUDIO_CMD: &str = "hab-studio";
const STUDIO_CMD_ENVVAR: &str = "HAB_STUDIO_BINARY";
const STUDIO_PACKAGE_IDENT: &str = "core/hab-studio";

#[derive(Clone, Copy)]
enum Sensitivity {
    PrintValue,
    NoPrintValue,
}

fn set_env_var_from_config(env_var: &str, config_val: Option<String>, sensitive: Sensitivity) {
    if henv::var(env_var).is_err() {
        if let Some(val) = config_val {
            match sensitive {
                Sensitivity::NoPrintValue => {
                    debug!("Setting {}=REDACTED (sensitive) via config file", env_var)
                }
                Sensitivity::PrintValue => debug!("Setting {}={} via config file", env_var, val),
            }
            env::set_var(env_var, val);
        }
    }
}

pub fn start(ui: &mut UI, args: &[OsString]) -> Result<()> {
    let config = config::load()?;

    set_env_var_from_config(AUTH_TOKEN_ENVVAR,
                            config.auth_token,
                            Sensitivity::NoPrintValue);
    set_env_var_from_config(BLDR_URL_ENVVAR, config.bldr_url, Sensitivity::PrintValue);
    set_env_var_from_config(CTL_SECRET_ENVVAR,
                            config.ctl_secret,
                            Sensitivity::NoPrintValue);
    set_env_var_from_config(ORIGIN_ENVVAR, config.origin, Sensitivity::PrintValue);

    if henv::var(CACHE_KEY_PATH_ENV_VAR).is_err() {
        let path = fs::cache_key_path(None::<&str>);
        debug!("Setting {}={}", CACHE_KEY_PATH_ENV_VAR, path.display());
        env::set_var(CACHE_KEY_PATH_ENV_VAR, &path);
    };

    let artifact_path = match henv::var(ARTIFACT_PATH_ENVVAR) {
        Ok(p) => PathBuf::from(p),
        Err(_) => {
            let path = fs::cache_artifact_path(None::<&str>);
            debug!("Setting {}={}", ARTIFACT_PATH_ENVVAR, path.display());
            env::set_var(ARTIFACT_PATH_ENVVAR, &path);
            path
        }
    };
    if !artifact_path.is_dir() {
        debug!("Creating artifact_path at: {}", artifact_path.display());
        stdfs::create_dir_all(&artifact_path)?;
    }

    inner::start(ui, args)
}

#[cfg(target_os = "linux")]
mod inner {
    use crate::{command::studio::docker,
                common::ui::{UIWriter,
                             UI},
                error::{Error,
                        Result},
                exec,
                hcore::{crypto::init,
                        env as henv,
                        fs::{am_i_root,
                             find_command},
                        os::process,
                        package::{PackageIdent,
                                  PackageInstall},
                        users::linux as group},
                VERSION};
    use std::{env,
              ffi::OsString,
              path::PathBuf,
              str::FromStr};

    const SUDO_CMD: &str = "sudo";

    pub fn start(ui: &mut UI, args: &[OsString]) -> Result<()> {
        rerun_with_sudo_if_needed(ui, &args)?;
        if is_docker_studio(&args) {
            docker::start_docker_studio(ui, args)
        } else {
            let command = match henv::var(super::STUDIO_CMD_ENVVAR) {
                Ok(command) => PathBuf::from(command),
                Err(_) => {
                    init();
                    let version: Vec<&str> = VERSION.split('/').collect();
                    let ident = PackageIdent::from_str(&format!("{}/{}",
                                                                super::STUDIO_PACKAGE_IDENT,
                                                                version[0]))?;
                    let command = exec::command_from_min_pkg(ui, super::STUDIO_CMD, &ident)?;
                    // This is a duplicate of the code in `hab pkg exec` and
                    // should be refactored as part of or after:
                    // https://github.com/habitat-sh/habitat/issues/6633
                    // https://github.com/habitat-sh/habitat/issues/6634
                    let pkg_install = PackageInstall::load(&ident, None)?;
                    let cmd_env = pkg_install.environment_for_command()?;
                    for (key, value) in cmd_env.into_iter() {
                        debug!("Setting: {}='{}'", key, value);
                        env::set_var(key, value);
                    }

                    let mut display_args = super::STUDIO_CMD.to_string();
                    for arg in args {
                        display_args.push(' ');
                        display_args.push_str(arg.to_string_lossy().as_ref());
                    }
                    debug!("Running: {}", display_args);

                    command
                }
            };

            if let Some(cmd) = find_command(command.to_string_lossy().as_ref()) {
                process::become_command(cmd, args)?;
                Ok(())
            } else {
                Err(Error::ExecCommandNotFound(command))
            }
        }
    }

    fn is_docker_studio(args: &[OsString]) -> bool {
        if cfg!(not(target_os = "linux")) {
            return false;
        }

        for arg in args.iter() {
            let str_arg = arg.to_string_lossy();
            if str_arg == "-D" {
                return true;
            }
        }

        false
    }

    fn has_docker_group() -> bool {
        let current_user = group::get_current_username().unwrap();
        let docker_members = group::get_members_by_groupname("docker");
        docker_members.map_or(false, |d| d.contains(&current_user))
    }

    fn rerun_with_sudo_if_needed(ui: &mut UI, args: &[OsString]) -> Result<()> {
        // If I have root permissions or if I am executing a docker studio
        // and have the appropriate group - early return, we are done.
        if am_i_root() || (is_docker_studio(args) && has_docker_group()) {
            return Ok(());
        }

        // Otherwise we will try to re-run this program using `sudo`
        match find_command(SUDO_CMD) {
            Some(sudo_prog) => {
                let mut args: Vec<OsString> = vec!["-p".into(),
                                                   "[sudo hab-studio] password for %u: ".into(),
                                                   "-E".into(),];
                args.append(&mut env::args_os().collect());
                process::become_command(sudo_prog, &args)?;
                Ok(())
            }
            None => {
                ui.warn(format!("Could not find the `{}' command, is it in your PATH?",
                                SUDO_CMD))?;
                ui.warn("Running Habitat Studio requires root or administrator privileges. \
                         Please retry this command as a super user or use a privilege-granting \
                         facility such as sudo.")?;
                ui.br()?;
                Err(Error::RootRequired)
            }
        }
    }
}

#[cfg(not(target_os = "linux"))]
mod inner {
    use crate::{command::studio::docker,
                common::ui::UI,
                error::{Error,
                        Result},
                exec,
                hcore::{crypto::init,
                        env as henv,
                        fs::find_command,
                        os::process,
                        package::PackageIdent},
                VERSION};
    use std::{ffi::OsString,
              path::PathBuf,
              str::FromStr};

    pub fn start(_ui: &mut UI, args: &[OsString]) -> Result<()> {
        if is_windows_studio(&args) {
            start_windows_studio(_ui, args)
        } else {
            docker::start_docker_studio(_ui, args)
        }
    }

    pub fn start_windows_studio(ui: &mut UI, args: &[OsString]) -> Result<()> {
        let command = match henv::var(super::STUDIO_CMD_ENVVAR) {
            Ok(command) => PathBuf::from(command),
            Err(_) => {
                init();
                let version: Vec<&str> = VERSION.split('/').collect();
                let ident = PackageIdent::from_str(&format!("{}/{}",
                                                            super::STUDIO_PACKAGE_IDENT,
                                                            version[0]))?;
                exec::command_from_min_pkg(ui, super::STUDIO_CMD, &ident)?
            }
        };

        if let Some(cmd) = find_command(command.to_string_lossy().as_ref()) {
            process::become_command(cmd, args)?;
        } else {
            return Err(Error::ExecCommandNotFound(command));
        }
        Ok(())
    }

    fn is_windows_studio(args: &[OsString]) -> bool {
        if cfg!(not(target_os = "windows")) {
            return false;
        }

        for arg in args.iter() {
            let str_arg = arg.to_string_lossy();
            if str_arg == "-D" {
                return false;
            }
        }

        true
    }
}
<|endoftext|>


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

//! Base classes to work with IDL callbacks.

use dom::bindings::error::{Error, Fallible};
use dom::bindings::global::global_root_from_object;
use dom::bindings::reflector::Reflectable;
use js::jsapi::GetGlobalForObjectCrossCompartment;
use js::jsapi::JSAutoCompartment;
use js::jsapi::{Heap, MutableHandleObject, RootedObject};
use js::jsapi::{IsCallable, JSContext, JSObject, JS_WrapObject};
use js::jsapi::{JSCompartment, JS_EnterCompartment, JS_LeaveCompartment};
use js::jsapi::{JS_GetProperty, JS_IsExceptionPending, JS_ReportPendingException};
use js::jsval::{JSVal, UndefinedValue};
use js::rust::RootedGuard;
use std::default::Default;
use std::ffi::CString;
use std::ptr;
use std::rc::Rc;

/// The exception handling used for a call.
#[derive(Copy, Clone, PartialEq)]
pub enum ExceptionHandling {
    /// Report any exception and don't throw it to the caller code.
    Report,
    /// Throw any exception to the caller code.
    Rethrow,
}

/// A common base class for representing IDL callback function types.
#[derive(JSTraceable, PartialEq)]
pub struct CallbackFunction {
    object: CallbackObject,
}

impl CallbackFunction {
    /// Create a new `CallbackFunction` for this object.
    pub fn new() -> CallbackFunction {
        CallbackFunction {
            object: CallbackObject {
                callback: Heap::default(),
            },
        }
    }

    /// Initialize the callback function with a value.
    /// Should be called once this object is done moving.
    pub fn init(&mut self, callback: *mut JSObject) {
        self.object.callback.set(callback);
    }
}

/// A common base class for representing IDL callback interface types.
#[derive(JSTraceable, PartialEq)]
pub struct CallbackInterface {
    object: CallbackObject,
}

/// A common base class for representing IDL callback function and
/// callback interface types.
#[derive(JSTraceable)]
struct CallbackObject {
    /// The underlying `JSObject`.
    callback: Heap<*mut JSObject>,
}

impl PartialEq for CallbackObject {
    fn eq(&self, other: &CallbackObject) -> bool {
        self.callback.get() == other.callback.get()
    }
}

/// A trait to be implemented by concrete IDL callback function and
/// callback interface types.
pub trait CallbackContainer {
    /// Create a new CallbackContainer object for the given `JSObject`.
    fn new(callback: *mut JSObject) -> Rc<Self>;
    /// Returns the underlying `JSObject`.
    fn callback(&self) -> *mut JSObject;
}

impl CallbackInterface {
    /// Returns the underlying `JSObject`.
    pub fn callback(&self) -> *mut JSObject {
        self.object.callback.get()
    }
}

impl CallbackFunction {
    /// Returns the underlying `JSObject`.
    pub fn callback(&self) -> *mut JSObject {
        self.object.callback.get()
    }
}

impl CallbackInterface {
    /// Create a new CallbackInterface object for the given `JSObject`.
    pub fn new() -> CallbackInterface {
        CallbackInterface {
            object: CallbackObject {
                callback: Heap::default(),
            },
        }
    }

    /// Initialize the callback function with a value.
    /// Should be called once this object is done moving.
    pub fn init(&mut self, callback: *mut JSObject) {
        self.object.callback.set(callback);
    }

    /// Returns the property with the given `name`, if it is a callable object,
    /// or an error otherwise.
    pub fn get_callable_property(&self, cx: *mut JSContext, name: &str) -> Fallible<JSVal> {
        rooted!(in(cx) let mut callable = UndefinedValue());
        rooted!(in(cx) let obj = self.callback());
        unsafe {
            let c_name = CString::new(name).unwrap();
            if !JS_GetProperty(cx, obj.handle(), c_name.as_ptr(), callable.handle_mut()) {
                return Err(Error::JSFailed);
            }

            if !callable.is_object() || !IsCallable(callable.to_object()) {
                return Err(Error::Type(format!("The value of the {} property is not callable",
                                               name)));
            }
        }
        Ok(callable.get())
    }
}

/// Wraps the reflector for `p` into the compartment of `cx`.
pub fn wrap_call_this_object<T: Reflectable>(cx: *mut JSContext,
                                             p: &T,
                                             rval: MutableHandleObject) {
    rval.set(p.reflector().get_jsobject().get());
    assert!(!rval.get().is_null());

    unsafe {
        if !JS_WrapObject(cx, rval) {
            rval.set(ptr::null_mut());
        }
    }
}

/// A class that performs whatever setup we need to safely make a call while
/// this class is on the stack. After `new` returns, the call is safe to make.
pub struct CallSetup<'a> {
    /// The compartment for reporting exceptions.
    exception_compartment: RootedGuard<'a, *mut JSObject>,
    /// The `JSContext` used for the call.
    cx: *mut JSContext,
    /// The compartment we were in before the call.
    old_compartment: *mut JSCompartment,
    /// The exception handling used for the call.
    handling: ExceptionHandling,
}

impl<'a> CallSetup<'a> {
    /// Performs the setup needed to make a call.
    #[allow(unrooted_must_root)]
    pub fn new<T: CallbackContainer>(exception_compartment: &'a mut RootedObject,
                                     callback: &T,
                                     handling: ExceptionHandling)
                                     -> CallSetup<'a> {
        let global = unsafe { global_root_from_object(callback.callback()) };
        let cx = global.r().get_cx();

        exception_compartment.ptr = unsafe {
            GetGlobalForObjectCrossCompartment(callback.callback())
        };
        CallSetup {
            exception_compartment: RootedGuard::new(cx, exception_compartment),
            cx: cx,
            old_compartment: unsafe { JS_EnterCompartment(cx, callback.callback()) },
            handling: handling,
        }
    }

    /// Returns the `JSContext` used for the call.
    pub fn get_context(&self) -> *mut JSContext {
        self.cx
    }
}

impl<'a> Drop for CallSetup<'a> {
    fn drop(&mut self) {
        unsafe {
            JS_LeaveCompartment(self.cx, self.old_compartment);
        }
        let need_to_deal_with_exception = self.handling == ExceptionHandling::Report &&
                                          unsafe { JS_IsExceptionPending(self.cx) };
        if need_to_deal_with_exception {
            unsafe {
                let _ac = JSAutoCompartment::new(self.cx, *self.exception_compartment);
                JS_ReportPendingException(self.cx);
            }
        }
    }
}
<commit_msg>Reuse report_pending_exception in CallSetup.<commit_after>/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/. */

//! Base classes to work with IDL callbacks.

use dom::bindings::error::{Error, Fallible, report_pending_exception};
use dom::bindings::global::global_root_from_object;
use dom::bindings::reflector::Reflectable;
use js::jsapi::GetGlobalForObjectCrossCompartment;
use js::jsapi::JS_GetProperty;
use js::jsapi::{Heap, MutableHandleObject, RootedObject};
use js::jsapi::{IsCallable, JSContext, JSObject, JS_WrapObject};
use js::jsapi::{JSCompartment, JS_EnterCompartment, JS_LeaveCompartment};
use js::jsval::{JSVal, UndefinedValue};
use js::rust::RootedGuard;
use std::default::Default;
use std::ffi::CString;
use std::ptr;
use std::rc::Rc;

/// The exception handling used for a call.
#[derive(Copy, Clone, PartialEq)]
pub enum ExceptionHandling {
    /// Report any exception and don't throw it to the caller code.
    Report,
    /// Throw any exception to the caller code.
    Rethrow,
}

/// A common base class for representing IDL callback function types.
#[derive(JSTraceable, PartialEq)]
pub struct CallbackFunction {
    object: CallbackObject,
}

impl CallbackFunction {
    /// Create a new `CallbackFunction` for this object.
    pub fn new() -> CallbackFunction {
        CallbackFunction {
            object: CallbackObject {
                callback: Heap::default(),
            },
        }
    }

    /// Initialize the callback function with a value.
    /// Should be called once this object is done moving.
    pub fn init(&mut self, callback: *mut JSObject) {
        self.object.callback.set(callback);
    }
}

/// A common base class for representing IDL callback interface types.
#[derive(JSTraceable, PartialEq)]
pub struct CallbackInterface {
    object: CallbackObject,
}

/// A common base class for representing IDL callback function and
/// callback interface types.
#[derive(JSTraceable)]
struct CallbackObject {
    /// The underlying `JSObject`.
    callback: Heap<*mut JSObject>,
}

impl PartialEq for CallbackObject {
    fn eq(&self, other: &CallbackObject) -> bool {
        self.callback.get() == other.callback.get()
    }
}

/// A trait to be implemented by concrete IDL callback function and
/// callback interface types.
pub trait CallbackContainer {
    /// Create a new CallbackContainer object for the given `JSObject`.
    fn new(callback: *mut JSObject) -> Rc<Self>;
    /// Returns the underlying `JSObject`.
    fn callback(&self) -> *mut JSObject;
}

impl CallbackInterface {
    /// Returns the underlying `JSObject`.
    pub fn callback(&self) -> *mut JSObject {
        self.object.callback.get()
    }
}

impl CallbackFunction {
    /// Returns the underlying `JSObject`.
    pub fn callback(&self) -> *mut JSObject {
        self.object.callback.get()
    }
}

impl CallbackInterface {
    /// Create a new CallbackInterface object for the given `JSObject`.
    pub fn new() -> CallbackInterface {
        CallbackInterface {
            object: CallbackObject {
                callback: Heap::default(),
            },
        }
    }

    /// Initialize the callback function with a value.
    /// Should be called once this object is done moving.
    pub fn init(&mut self, callback: *mut JSObject) {
        self.object.callback.set(callback);
    }

    /// Returns the property with the given `name`, if it is a callable object,
    /// or an error otherwise.
    pub fn get_callable_property(&self, cx: *mut JSContext, name: &str) -> Fallible<JSVal> {
        rooted!(in(cx) let mut callable = UndefinedValue());
        rooted!(in(cx) let obj = self.callback());
        unsafe {
            let c_name = CString::new(name).unwrap();
            if !JS_GetProperty(cx, obj.handle(), c_name.as_ptr(), callable.handle_mut()) {
                return Err(Error::JSFailed);
            }

            if !callable.is_object() || !IsCallable(callable.to_object()) {
                return Err(Error::Type(format!("The value of the {} property is not callable",
                                               name)));
            }
        }
        Ok(callable.get())
    }
}

/// Wraps the reflector for `p` into the compartment of `cx`.
pub fn wrap_call_this_object<T: Reflectable>(cx: *mut JSContext,
                                             p: &T,
                                             rval: MutableHandleObject) {
    rval.set(p.reflector().get_jsobject().get());
    assert!(!rval.get().is_null());

    unsafe {
        if !JS_WrapObject(cx, rval) {
            rval.set(ptr::null_mut());
        }
    }
}

/// A class that performs whatever setup we need to safely make a call while
/// this class is on the stack. After `new` returns, the call is safe to make.
pub struct CallSetup<'a> {
    /// The compartment for reporting exceptions.
    exception_compartment: RootedGuard<'a, *mut JSObject>,
    /// The `JSContext` used for the call.
    cx: *mut JSContext,
    /// The compartment we were in before the call.
    old_compartment: *mut JSCompartment,
    /// The exception handling used for the call.
    handling: ExceptionHandling,
}

impl<'a> CallSetup<'a> {
    /// Performs the setup needed to make a call.
    #[allow(unrooted_must_root)]
    pub fn new<T: CallbackContainer>(exception_compartment: &'a mut RootedObject,
                                     callback: &T,
                                     handling: ExceptionHandling)
                                     -> CallSetup<'a> {
        let global = unsafe { global_root_from_object(callback.callback()) };
        let cx = global.r().get_cx();

        exception_compartment.ptr = unsafe {
            GetGlobalForObjectCrossCompartment(callback.callback())
        };
        CallSetup {
            exception_compartment: RootedGuard::new(cx, exception_compartment),
            cx: cx,
            old_compartment: unsafe { JS_EnterCompartment(cx, callback.callback()) },
            handling: handling,
        }
    }

    /// Returns the `JSContext` used for the call.
    pub fn get_context(&self) -> *mut JSContext {
        self.cx
    }
}

impl<'a> Drop for CallSetup<'a> {
    fn drop(&mut self) {
        unsafe {
            JS_LeaveCompartment(self.cx, self.old_compartment);
            if self.handling == ExceptionHandling::Report {
                report_pending_exception(self.cx, *self.exception_compartment);
            }
        }
    }
}
<|endoftext|>


<commit_before>use proc_macro::TokenStream;
use proc_macro2::Span;
use syn::{
    Token, Ident, Type, Attribute, ReturnType, Expr, Block, Error,
    braced, parenthesized, parse_macro_input,
};
use syn::spanned::Spanned;
use syn::parse::{Result, Parse, ParseStream};
use syn::punctuated::Punctuated;
use syn;
use quote::quote;

#[allow(non_camel_case_types)]
mod kw {
    syn::custom_keyword!(query);
}

/// Ident or a wildcard `_`.
struct IdentOrWild(Ident);

impl Parse for IdentOrWild {
    fn parse(input: ParseStream<'_>) -> Result<Self> {
        Ok(if input.peek(Token![_]) {
            input.parse::<Token![_]>()?;
            IdentOrWild(Ident::new("_", Span::call_site()))
        } else {
            IdentOrWild(input.parse()?)
        })
    }
}

/// A modifier for a query
enum QueryModifier {
    /// The description of the query
    Desc(Option<Ident>, Punctuated<Expr, Token![,]>),

    /// Cache the query to disk if the `Expr` returns true.
    Cache(Option<Ident>, Expr),

    /// Custom code to load the query from disk.
    LoadCached(Ident, Ident, Block),

    /// A cycle error for this query aborting the compilation with a fatal error.
    FatalCycle,
}

impl Parse for QueryModifier {
    fn parse(input: ParseStream<'_>) -> Result<Self> {
        let modifier: Ident = input.parse()?;
        if modifier == "desc" {
            // Parse a description modifier like:
            // `desc { |tcx| "foo {}", tcx.item_path(key) }`
            let attr_content;
            braced!(attr_content in input);
            let tcx = if attr_content.peek(Token![|]) {
                attr_content.parse::<Token![|]>()?;
                let tcx = attr_content.parse()?;
                attr_content.parse::<Token![|]>()?;
                Some(tcx)
            } else {
                None
            };
            let desc = attr_content.parse_terminated(Expr::parse)?;
            Ok(QueryModifier::Desc(tcx, desc))
        } else if modifier == "cache" {
            // Parse a cache modifier like:
            // `cache { |tcx| key.is_local() }`
            let attr_content;
            braced!(attr_content in input);
            let tcx = if attr_content.peek(Token![|]) {
                attr_content.parse::<Token![|]>()?;
                let tcx = attr_content.parse()?;
                attr_content.parse::<Token![|]>()?;
                Some(tcx)
            } else {
                None
            };
            let expr = attr_content.parse()?;
            Ok(QueryModifier::Cache(tcx, expr))
        } else if modifier == "load_cached" {
            // Parse a load_cached modifier like:
            // `load_cached(tcx, id) { tcx.queries.on_disk_cache.try_load_query_result(tcx, id) }`
            let args;
            parenthesized!(args in input);
            let tcx = args.parse()?;
            args.parse::<Token![,]>()?;
            let id = args.parse()?;
            let block = input.parse()?;
            Ok(QueryModifier::LoadCached(tcx, id, block))
        } else if modifier == "fatal_cycle" {
            Ok(QueryModifier::FatalCycle)
        } else {
            Err(Error::new(modifier.span(), "unknown query modifier"))
        }
    }
}

/// Ensures only doc comment attributes are used
fn check_attributes(attrs: Vec<Attribute>) -> Result<()> {
    for attr in attrs {
        if !attr.path.is_ident("doc") {
            return Err(Error::new(attr.span(), "attributes not supported on queries"));
        }
    }
    Ok(())
}

/// A compiler query. `query ... { ... }`
struct Query {
    attrs: List<QueryModifier>,
    name: Ident,
    key: IdentOrWild,
    arg: Type,
    result: ReturnType,
}

impl Parse for Query {
    fn parse(input: ParseStream<'_>) -> Result<Self> {
        check_attributes(input.call(Attribute::parse_outer)?)?;

        // Parse the query declaration. Like `query type_of(key: DefId) -> Ty<'tcx>`
        input.parse::<kw::query>()?;
        let name: Ident = input.parse()?;
        let arg_content;
        parenthesized!(arg_content in input);
        let key = arg_content.parse()?;
        arg_content.parse::<Token![:]>()?;
        let arg = arg_content.parse()?;
        let result = input.parse()?;

        // Parse the query modifiers
        let content;
        braced!(content in input);
        let attrs = content.parse()?;

        Ok(Query {
            attrs,
            name,
            key,
            arg,
            result,
        })
    }
}

/// A type used to greedily parse another type until the input is empty.
struct List<T>(Vec<T>);

impl<T: Parse> Parse for List<T> {
    fn parse(input: ParseStream<'_>) -> Result<Self> {
        let mut list = Vec::new();
        while !input.is_empty() {
            list.push(input.parse()?);
        }
        Ok(List(list))
    }
}

/// A named group containing queries.
struct Group {
    name: Ident,
    queries: List<Query>,
}

impl Parse for Group {
    fn parse(input: ParseStream<'_>) -> Result<Self> {
        let name: Ident = input.parse()?;
        let content;
        braced!(content in input);
        Ok(Group {
            name,
            queries: content.parse()?,
        })
    }
}

/// Add the impl of QueryDescription for the query to `impls` if one is requested
fn add_query_description_impl(query: &Query, impls: &mut proc_macro2::TokenStream) {
    let name = &query.name;
    let arg = &query.arg;
    let key = &query.key.0;

    // Find custom code to load the query from disk
    let load_cached = query.attrs.0.iter().find_map(|attr| match attr {
        QueryModifier::LoadCached(tcx, id, block) => Some((tcx, id, block)),
        _ => None,
    });

    // Find out if we should cache the query on disk
    let cache = query.attrs.0.iter().find_map(|attr| match attr {
        QueryModifier::Cache(tcx, expr) => Some((tcx, expr)),
        _ => None,
    }).map(|(tcx, expr)| {
        let try_load_from_disk = if let Some((tcx, id, block)) = load_cached {
            quote! {
                #[inline]
                fn try_load_from_disk(
                    #tcx: TyCtxt<'_, 'tcx, 'tcx>,
                    #id: SerializedDepNodeIndex
                ) -> Option<Self::Value> {
                    #block
                }
            }
        } else {
            quote! {
                #[inline]
                fn try_load_from_disk(
                    tcx: TyCtxt<'_, 'tcx, 'tcx>,
                    id: SerializedDepNodeIndex
                ) -> Option<Self::Value> {
                    tcx.queries.on_disk_cache.try_load_query_result(tcx, id)
                }
            }
        };

        let tcx = tcx.as_ref().map(|t| quote! { #t }).unwrap_or(quote! { _ });
        quote! {
            #[inline]
            fn cache_on_disk(#tcx: TyCtxt<'_, 'tcx, 'tcx>, #key: Self::Key) -> bool {
                #expr
            }

            #try_load_from_disk
        }
    });

    if cache.is_none() && load_cached.is_some() {
        panic!("load_cached modifier on query `{}` without a cache modifier", name);
    }

    let desc = query.attrs.0.iter().find_map(|attr| match attr {
        QueryModifier::Desc(tcx, desc) => Some((tcx, desc)),
        _ => None,
    }).map(|(tcx, desc)| {
        let tcx = tcx.as_ref().map(|t| quote! { #t }).unwrap_or(quote! { _ });
        quote! {
            fn describe(
                #tcx: TyCtxt<'_, '_, '_>,
                #key: #arg,
            ) -> Cow<'static, str> {
                format!(#desc).into()
            }
        }
    });

    if desc.is_some() || cache.is_some() {
        let cache = cache.unwrap_or(quote! {});
        let desc = desc.unwrap_or(quote! {});

        impls.extend(quote! {
            impl<'tcx> QueryDescription<'tcx> for queries::#name<'tcx> {
                #desc
                #cache
            }
        });
    }
}

pub fn rustc_queries(input: TokenStream) -> TokenStream {
    let groups = parse_macro_input!(input as List<Group>);

    let mut query_stream = quote! {};
    let mut query_description_stream = quote! {};
    let mut dep_node_def_stream = quote! {};
    let mut dep_node_force_stream = quote! {};

    for group in groups.0 {
        let mut group_stream = quote! {};
        for query in &group.queries.0 {
            let name = &query.name;
            let arg = &query.arg;
            let result_full = &query.result;
            let result = match query.result {
                ReturnType::Default => quote! { -> () },
                _ => quote! { #result_full },
            };

            // Look for a fatal_cycle modifier to pass on
            let fatal_cycle = query.attrs.0.iter().find_map(|attr| match attr {
                QueryModifier::FatalCycle => Some(()),
                _ => None,
            }).map(|_| quote! { fatal_cycle }).unwrap_or(quote! {});

            // Add the query to the group
            group_stream.extend(quote! {
                [#fatal_cycle] fn #name: #name(#arg) #result,
            });

            add_query_description_impl(query, &mut query_description_stream);

            // Create a dep node for the query
            dep_node_def_stream.extend(quote! {
                [] #name(#arg),
            });

            // Add a match arm to force the query given the dep node
            dep_node_force_stream.extend(quote! {
                DepKind::#name => {
                    if let Some(key) = RecoverKey::recover($tcx, $dep_node) {
                        force_ex!($tcx, #name, key);
                    } else {
                        return false;
                    }
                }
            });
        }
        let name = &group.name;
        query_stream.extend(quote! {
            #name { #group_stream },
        });
    }
    TokenStream::from(quote! {
        macro_rules! rustc_query_append {
            ([$($macro:tt)*][$($other:tt)*]) => {
                $($macro)* {
                    $($other)*

                    #query_stream

                }
            }
        }
        macro_rules! rustc_dep_node_append {
            ([$($macro:tt)*][$($other:tt)*]) => {
                $($macro)*(
                    $($other)*

                    #dep_node_def_stream
                );
            }
        }
        macro_rules! rustc_dep_node_force {
            ([$dep_node:expr, $tcx:expr] $($other:tt)*) => {
                match $dep_node.kind {
                    $($other)*

                    #dep_node_force_stream
                }
            }
        }
        #query_description_stream
    })
}
<commit_msg>Preprocess query modifiers<commit_after>use proc_macro::TokenStream;
use syn::{
    Token, Ident, Type, Attribute, ReturnType, Expr, Block, Error,
    braced, parenthesized, parse_macro_input,
};
use syn::spanned::Spanned;
use syn::parse::{Result, Parse, ParseStream};
use syn::punctuated::Punctuated;
use syn;
use quote::quote;

#[allow(non_camel_case_types)]
mod kw {
    syn::custom_keyword!(query);
}

/// Ident or a wildcard `_`.
struct IdentOrWild(Ident);

impl Parse for IdentOrWild {
    fn parse(input: ParseStream<'_>) -> Result<Self> {
        Ok(if input.peek(Token![_]) {
            let underscore = input.parse::<Token![_]>()?;
            IdentOrWild(Ident::new("_", underscore.span()))
        } else {
            IdentOrWild(input.parse()?)
        })
    }
}

/// A modifier for a query
enum QueryModifier {
    /// The description of the query.
    Desc(Option<Ident>, Punctuated<Expr, Token![,]>),

    /// Cache the query to disk if the `Expr` returns true.
    Cache(Option<Ident>, Expr),

    /// Custom code to load the query from disk.
    LoadCached(Ident, Ident, Block),

    /// A cycle error for this query aborting the compilation with a fatal error.
    FatalCycle,
}

impl Parse for QueryModifier {
    fn parse(input: ParseStream<'_>) -> Result<Self> {
        let modifier: Ident = input.parse()?;
        if modifier == "desc" {
            // Parse a description modifier like:
            // `desc { |tcx| "foo {}", tcx.item_path(key) }`
            let attr_content;
            braced!(attr_content in input);
            let tcx = if attr_content.peek(Token![|]) {
                attr_content.parse::<Token![|]>()?;
                let tcx = attr_content.parse()?;
                attr_content.parse::<Token![|]>()?;
                Some(tcx)
            } else {
                None
            };
            let desc = attr_content.parse_terminated(Expr::parse)?;
            Ok(QueryModifier::Desc(tcx, desc))
        } else if modifier == "cache" {
            // Parse a cache modifier like:
            // `cache { |tcx| key.is_local() }`
            let attr_content;
            braced!(attr_content in input);
            let tcx = if attr_content.peek(Token![|]) {
                attr_content.parse::<Token![|]>()?;
                let tcx = attr_content.parse()?;
                attr_content.parse::<Token![|]>()?;
                Some(tcx)
            } else {
                None
            };
            let expr = attr_content.parse()?;
            Ok(QueryModifier::Cache(tcx, expr))
        } else if modifier == "load_cached" {
            // Parse a load_cached modifier like:
            // `load_cached(tcx, id) { tcx.queries.on_disk_cache.try_load_query_result(tcx, id) }`
            let args;
            parenthesized!(args in input);
            let tcx = args.parse()?;
            args.parse::<Token![,]>()?;
            let id = args.parse()?;
            let block = input.parse()?;
            Ok(QueryModifier::LoadCached(tcx, id, block))
        } else if modifier == "fatal_cycle" {
            Ok(QueryModifier::FatalCycle)
        } else {
            Err(Error::new(modifier.span(), "unknown query modifier"))
        }
    }
}

/// Ensures only doc comment attributes are used
fn check_attributes(attrs: Vec<Attribute>) -> Result<()> {
    for attr in attrs {
        if !attr.path.is_ident("doc") {
            return Err(Error::new(attr.span(), "attributes not supported on queries"));
        }
    }
    Ok(())
}

/// A compiler query. `query ... { ... }`
struct Query {
    modifiers: List<QueryModifier>,
    name: Ident,
    key: IdentOrWild,
    arg: Type,
    result: ReturnType,
}

impl Parse for Query {
    fn parse(input: ParseStream<'_>) -> Result<Self> {
        check_attributes(input.call(Attribute::parse_outer)?)?;

        // Parse the query declaration. Like `query type_of(key: DefId) -> Ty<'tcx>`
        input.parse::<kw::query>()?;
        let name: Ident = input.parse()?;
        let arg_content;
        parenthesized!(arg_content in input);
        let key = arg_content.parse()?;
        arg_content.parse::<Token![:]>()?;
        let arg = arg_content.parse()?;
        let result = input.parse()?;

        // Parse the query modifiers
        let content;
        braced!(content in input);
        let modifiers = content.parse()?;

        Ok(Query {
            modifiers,
            name,
            key,
            arg,
            result,
        })
    }
}

/// A type used to greedily parse another type until the input is empty.
struct List<T>(Vec<T>);

impl<T: Parse> Parse for List<T> {
    fn parse(input: ParseStream<'_>) -> Result<Self> {
        let mut list = Vec::new();
        while !input.is_empty() {
            list.push(input.parse()?);
        }
        Ok(List(list))
    }
}

/// A named group containing queries.
struct Group {
    name: Ident,
    queries: List<Query>,
}

impl Parse for Group {
    fn parse(input: ParseStream<'_>) -> Result<Self> {
        let name: Ident = input.parse()?;
        let content;
        braced!(content in input);
        Ok(Group {
            name,
            queries: content.parse()?,
        })
    }
}

struct QueryModifiers {
    /// The description of the query.
    desc: Option<(Option<Ident>, Punctuated<Expr, Token![,]>)>,

    /// Cache the query to disk if the `Expr` returns true.
    cache: Option<(Option<Ident>, Expr)>,

    /// Custom code to load the query from disk.
    load_cached: Option<(Ident, Ident, Block)>,

    /// A cycle error for this query aborting the compilation with a fatal error.
    fatal_cycle: bool,
}

/// Process query modifiers into a struct, erroring on duplicates
fn process_modifiers(query: &mut Query) -> QueryModifiers {
    let mut load_cached = None;
    let mut cache = None;
    let mut desc = None;
    let mut fatal_cycle = false;
    for modifier in query.modifiers.0.drain(..) {
        match modifier {
            QueryModifier::LoadCached(tcx, id, block) => {
                if load_cached.is_some() {
                    panic!("duplicate modifier `load_cached` for query `{}`", query.name);
                }
                load_cached = Some((tcx, id, block));
            }
            QueryModifier::Cache(tcx, expr) => {
                if cache.is_some() {
                    panic!("duplicate modifier `cache` for query `{}`", query.name);
                }
                cache = Some((tcx, expr));
            }
            QueryModifier::Desc(tcx, list) => {
                if desc.is_some() {
                    panic!("duplicate modifier `desc` for query `{}`", query.name);
                }
                desc = Some((tcx, list));
            }
            QueryModifier::FatalCycle => {
                if fatal_cycle {
                    panic!("duplicate modifier `fatal_cycle` for query `{}`", query.name);
                }
                fatal_cycle = true;
            }
        }
    }
    QueryModifiers {
        load_cached,
        cache,
        desc,
        fatal_cycle,
    }
}

/// Add the impl of QueryDescription for the query to `impls` if one is requested
fn add_query_description_impl(
    query: &Query,
    modifiers: QueryModifiers,
    impls: &mut proc_macro2::TokenStream
) {
    let name = &query.name;
    let arg = &query.arg;
    let key = &query.key.0;

    // Find out if we should cache the query on disk
    let cache = modifiers.cache.as_ref().map(|(tcx, expr)| {
        let try_load_from_disk = if let Some((tcx, id, block)) = modifiers.load_cached.as_ref() {
            // Use custom code to load the query from disk
            quote! {
                #[inline]
                fn try_load_from_disk(
                    #tcx: TyCtxt<'_, 'tcx, 'tcx>,
                    #id: SerializedDepNodeIndex
                ) -> Option<Self::Value> {
                    #block
                }
            }
        } else {
            // Use the default code to load the query from disk
            quote! {
                #[inline]
                fn try_load_from_disk(
                    tcx: TyCtxt<'_, 'tcx, 'tcx>,
                    id: SerializedDepNodeIndex
                ) -> Option<Self::Value> {
                    tcx.queries.on_disk_cache.try_load_query_result(tcx, id)
                }
            }
        };

        let tcx = tcx.as_ref().map(|t| quote! { #t }).unwrap_or(quote! { _ });
        quote! {
            #[inline]
            fn cache_on_disk(#tcx: TyCtxt<'_, 'tcx, 'tcx>, #key: Self::Key) -> bool {
                #expr
            }

            #try_load_from_disk
        }
    });

    if cache.is_none() && modifiers.load_cached.is_some() {
        panic!("load_cached modifier on query `{}` without a cache modifier", name);
    }

    let desc = modifiers.desc.as_ref().map(|(tcx, desc)| {
        let tcx = tcx.as_ref().map(|t| quote! { #t }).unwrap_or(quote! { _ });
        quote! {
            fn describe(
                #tcx: TyCtxt<'_, '_, '_>,
                #key: #arg,
            ) -> Cow<'static, str> {
                format!(#desc).into()
            }
        }
    });

    if desc.is_some() || cache.is_some() {
        let cache = cache.unwrap_or(quote! {});
        let desc = desc.unwrap_or(quote! {});

        impls.extend(quote! {
            impl<'tcx> QueryDescription<'tcx> for queries::#name<'tcx> {
                #desc
                #cache
            }
        });
    }
}

pub fn rustc_queries(input: TokenStream) -> TokenStream {
    let groups = parse_macro_input!(input as List<Group>);

    let mut query_stream = quote! {};
    let mut query_description_stream = quote! {};
    let mut dep_node_def_stream = quote! {};
    let mut dep_node_force_stream = quote! {};

    for group in groups.0 {
        let mut group_stream = quote! {};
        for mut query in group.queries.0 {
            let modifiers = process_modifiers(&mut query);
            let name = &query.name;
            let arg = &query.arg;
            let result_full = &query.result;
            let result = match query.result {
                ReturnType::Default => quote! { -> () },
                _ => quote! { #result_full },
            };

            // Pass on the fatal_cycle modifier
            let fatal_cycle = if modifiers.fatal_cycle {
                quote! { fatal_cycle }
            } else {
                quote! {}
            };

            // Add the query to the group
            group_stream.extend(quote! {
                [#fatal_cycle] fn #name: #name(#arg) #result,
            });

            add_query_description_impl(&query, modifiers, &mut query_description_stream);

            // Create a dep node for the query
            dep_node_def_stream.extend(quote! {
                [] #name(#arg),
            });

            // Add a match arm to force the query given the dep node
            dep_node_force_stream.extend(quote! {
                DepKind::#name => {
                    if let Some(key) = RecoverKey::recover($tcx, $dep_node) {
                        force_ex!($tcx, #name, key);
                    } else {
                        return false;
                    }
                }
            });
        }
        let name = &group.name;
        query_stream.extend(quote! {
            #name { #group_stream },
        });
    }
    TokenStream::from(quote! {
        macro_rules! rustc_query_append {
            ([$($macro:tt)*][$($other:tt)*]) => {
                $($macro)* {
                    $($other)*

                    #query_stream

                }
            }
        }
        macro_rules! rustc_dep_node_append {
            ([$($macro:tt)*][$($other:tt)*]) => {
                $($macro)*(
                    $($other)*

                    #dep_node_def_stream
                );
            }
        }
        macro_rules! rustc_dep_node_force {
            ([$dep_node:expr, $tcx:expr] $($other:tt)*) => {
                match $dep_node.kind {
                    $($other)*

                    #dep_node_force_stream
                }
            }
        }
        #query_description_stream
    })
}
<|endoftext|>