<commit_before>//! FIXME: write short doc here

use std::{
    ffi::OsStr,
    ops,
    path::{Path, PathBuf},
    process::Command,
};

use anyhow::{Context, Result};
use arena::{Arena, Idx};
use base_db::Edition;
use cargo_metadata::{BuildScript, CargoOpt, Message, MetadataCommand, PackageId};
use paths::{AbsPath, AbsPathBuf};
use rustc_hash::FxHashMap;

use crate::cfg_flag::CfgFlag;

/// `CargoWorkspace` represents the logical structure of, well, a Cargo
/// workspace. It pretty closely mirrors `cargo metadata` output.
///
/// Note that internally, rust analyzer uses a different structure:
/// `CrateGraph`. `CrateGraph` is lower-level: it knows only about the crates,
/// while this knows about `Packages` & `Targets`: purely cargo-related
/// concepts.
///
/// We use absolute paths here, `cargo metadata` guarantees to always produce
/// abs paths.
#[derive(Debug, Clone, Eq, PartialEq)]
pub struct CargoWorkspace {
    packages: Arena<PackageData>,
    targets: Arena<TargetData>,
    workspace_root: AbsPathBuf,
}

impl ops::Index<Package> for CargoWorkspace {
    type Output = PackageData;
    fn index(&self, index: Package) -> &PackageData {
        &self.packages[index]
    }
}

impl ops::Index<Target> for CargoWorkspace {
    type Output = TargetData;
    fn index(&self, index: Target) -> &TargetData {
        &self.targets[index]
    }
}

#[derive(Default, Clone, Debug, PartialEq, Eq)]
pub struct CargoConfig {
    /// Do not activate the `default` feature.
    pub no_default_features: bool,

    /// Activate all available features
    pub all_features: bool,

    /// List of features to activate.
    /// This will be ignored if `cargo_all_features` is true.
    pub features: Vec<String>,

    /// Runs cargo check on launch to figure out the correct values of OUT_DIR
    pub load_out_dirs_from_check: bool,

    /// rustc target
    pub target: Option<String>,

    /// Don't load sysroot crates (`std`, `core` & friends). Might be useful
    /// when debugging isolated issues.
    pub no_sysroot: bool,

    /// rustc private crate source
    pub rustc_source: Option<AbsPathBuf>,
}

pub type Package = Idx<PackageData>;

pub type Target = Idx<TargetData>;

#[derive(Debug, Clone, Eq, PartialEq)]
pub struct PackageData {
    pub version: String,
    pub name: String,
    pub manifest: AbsPathBuf,
    pub targets: Vec<Target>,
    pub is_member: bool,
    pub dependencies: Vec<PackageDependency>,
    pub edition: Edition,
    pub features: Vec<String>,
    pub cfgs: Vec<CfgFlag>,
    pub envs: Vec<(String, String)>,
    pub out_dir: Option<AbsPathBuf>,
    pub proc_macro_dylib_path: Option<AbsPathBuf>,
}

#[derive(Debug, Clone, Eq, PartialEq)]
pub struct PackageDependency {
    pub pkg: Package,
    pub name: String,
}

#[derive(Debug, Clone, Eq, PartialEq)]
pub struct TargetData {
    pub package: Package,
    pub name: String,
    pub root: AbsPathBuf,
    pub kind: TargetKind,
    pub is_proc_macro: bool,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum TargetKind {
    Bin,
    /// Any kind of Cargo lib crate-type (dylib, rlib, proc-macro, ...).
    Lib,
    Example,
    Test,
    Bench,
    Other,
}

impl TargetKind {
    fn new(kinds: &[String]) -> TargetKind {
        for kind in kinds {
            return match kind.as_str() {
                "bin" => TargetKind::Bin,
                "test" => TargetKind::Test,
                "bench" => TargetKind::Bench,
                "example" => TargetKind::Example,
                "proc-macro" => TargetKind::Lib,
                _ if kind.contains("lib") => TargetKind::Lib,
                _ => continue,
            };
        }
        TargetKind::Other
    }
}

impl PackageData {
    pub fn root(&self) -> &AbsPath {
        self.manifest.parent().unwrap()
    }
}

impl CargoWorkspace {
    pub fn from_cargo_metadata(
        cargo_toml: &AbsPath,
        config: &CargoConfig,
    ) -> Result<CargoWorkspace> {
        let mut meta = MetadataCommand::new();
        meta.cargo_path(toolchain::cargo());
        meta.manifest_path(cargo_toml.to_path_buf());
        if config.all_features {
            meta.features(CargoOpt::AllFeatures);
        } else {
            if config.no_default_features {
                // FIXME: `NoDefaultFeatures` is mutual exclusive with `SomeFeatures`
                // https://github.com/oli-obk/cargo_metadata/issues/79
                meta.features(CargoOpt::NoDefaultFeatures);
            }
            if !config.features.is_empty() {
                meta.features(CargoOpt::SomeFeatures(config.features.clone()));
            }
        }
        if let Some(parent) = cargo_toml.parent() {
            meta.current_dir(parent.to_path_buf());
        }
        if let Some(target) = config.target.as_ref() {
            meta.other_options(vec![String::from("--filter-platform"), target.clone()]);
        }
        let mut meta = meta.exec().with_context(|| {
            format!("Failed to run `cargo metadata --manifest-path {}`", cargo_toml.display())
        })?;

        let mut out_dir_by_id = FxHashMap::default();
        let mut cfgs = FxHashMap::default();
        let mut envs = FxHashMap::default();
        let mut proc_macro_dylib_paths = FxHashMap::default();
        if config.load_out_dirs_from_check {
            let resources = load_extern_resources(cargo_toml, config)?;
            out_dir_by_id = resources.out_dirs;
            cfgs = resources.cfgs;
            envs = resources.env;
            proc_macro_dylib_paths = resources.proc_dylib_paths;
        }

        let mut pkg_by_id = FxHashMap::default();
        let mut packages = Arena::default();
        let mut targets = Arena::default();

        let ws_members = &meta.workspace_members;

        meta.packages.sort_by(|a, b| a.id.cmp(&b.id));
        for meta_pkg in meta.packages {
            let cargo_metadata::Package { id, edition, name, manifest_path, version, .. } =
                meta_pkg;
            let is_member = ws_members.contains(&id);
            let edition = edition
                .parse::<Edition>()
                .with_context(|| format!("Failed to parse edition {}", edition))?;
            let pkg = packages.alloc(PackageData {
                name,
                version: version.to_string(),
                manifest: AbsPathBuf::assert(manifest_path),
                targets: Vec::new(),
                is_member,
                edition,
                dependencies: Vec::new(),
                features: Vec::new(),
                cfgs: cfgs.get(&id).cloned().unwrap_or_default(),
                envs: envs.get(&id).cloned().unwrap_or_default(),
                out_dir: out_dir_by_id.get(&id).cloned(),
                proc_macro_dylib_path: proc_macro_dylib_paths.get(&id).cloned(),
            });
            let pkg_data = &mut packages[pkg];
            pkg_by_id.insert(id, pkg);
            for meta_tgt in meta_pkg.targets {
                let is_proc_macro = meta_tgt.kind.as_slice() == ["proc-macro"];
                let tgt = targets.alloc(TargetData {
                    package: pkg,
                    name: meta_tgt.name,
                    root: AbsPathBuf::assert(meta_tgt.src_path.clone()),
                    kind: TargetKind::new(meta_tgt.kind.as_slice()),
                    is_proc_macro,
                });
                pkg_data.targets.push(tgt);
            }
        }
        let resolve = meta.resolve.expect("metadata executed with deps");
        for mut node in resolve.nodes {
            let source = match pkg_by_id.get(&node.id) {
                Some(&src) => src,
                // FIXME: replace this and a similar branch below with `.unwrap`, once
                // https://github.com/rust-lang/cargo/issues/7841
                // is fixed and hits stable (around 1.43-is probably?).
                None => {
                    log::error!("Node id do not match in cargo metadata, ignoring {}", node.id);
                    continue;
                }
            };
            node.deps.sort_by(|a, b| a.pkg.cmp(&b.pkg));
            for dep_node in node.deps {
                let pkg = match pkg_by_id.get(&dep_node.pkg) {
                    Some(&pkg) => pkg,
                    None => {
                        log::error!(
                            "Dep node id do not match in cargo metadata, ignoring {}",
                            dep_node.pkg
                        );
                        continue;
                    }
                };
                let dep = PackageDependency { name: dep_node.name, pkg };
                packages[source].dependencies.push(dep);
            }
            packages[source].features.extend(node.features);
        }

        let workspace_root = AbsPathBuf::assert(meta.workspace_root);
        Ok(CargoWorkspace { packages, targets, workspace_root: workspace_root })
    }

    pub fn packages<'a>(&'a self) -> impl Iterator<Item = Package> + ExactSizeIterator + 'a {
        self.packages.iter().map(|(id, _pkg)| id)
    }

    pub fn target_by_root(&self, root: &AbsPath) -> Option<Target> {
        self.packages()
            .filter_map(|pkg| self[pkg].targets.iter().find(|&&it| &self[it].root == root))
            .next()
            .copied()
    }

    pub fn workspace_root(&self) -> &AbsPath {
        &self.workspace_root
    }

    pub fn package_flag(&self, package: &PackageData) -> String {
        if self.is_unique(&*package.name) {
            package.name.clone()
        } else {
            format!("{}:{}", package.name, package.version)
        }
    }

    fn is_unique(&self, name: &str) -> bool {
        self.packages.iter().filter(|(_, v)| v.name == name).count() == 1
    }
}

#[derive(Debug, Clone, Default)]
pub(crate) struct ExternResources {
    out_dirs: FxHashMap<PackageId, AbsPathBuf>,
    proc_dylib_paths: FxHashMap<PackageId, AbsPathBuf>,
    cfgs: FxHashMap<PackageId, Vec<CfgFlag>>,
    env: FxHashMap<PackageId, Vec<(String, String)>>,
}

pub(crate) fn load_extern_resources(
    cargo_toml: &Path,
    cargo_features: &CargoConfig,
) -> Result<ExternResources> {
    let mut cmd = Command::new(toolchain::cargo());
    cmd.args(&["check", "--message-format=json", "--manifest-path"]).arg(cargo_toml);

    if let Some(target) = &cargo_features.target {
        cmd.args(&["--target", target]);
    }

    if cargo_features.all_features {
        cmd.arg("--all-features");
    } else {
        if cargo_features.no_default_features {
            // FIXME: `NoDefaultFeatures` is mutual exclusive with `SomeFeatures`
            // https://github.com/oli-obk/cargo_metadata/issues/79
            cmd.arg("--no-default-features");
        }
        if !cargo_features.features.is_empty() {
            cmd.arg("--features");
            cmd.arg(cargo_features.features.join(" "));
        }
    }

    let output = cmd.output()?;

    let mut res = ExternResources::default();

    for message in cargo_metadata::Message::parse_stream(output.stdout.as_slice()) {
        if let Ok(message) = message {
            match message {
                Message::BuildScriptExecuted(BuildScript {
                    package_id,
                    out_dir,
                    cfgs,
                    env,
                    ..
                }) => {
                    let cfgs = {
                        let mut acc = Vec::new();
                        for cfg in cfgs {
                            match cfg.parse::<CfgFlag>() {
                                Ok(it) => acc.push(it),
                                Err(err) => {
                                    anyhow::bail!("invalid cfg from cargo-metadata: {}", err)
                                }
                            };
                        }
                        acc
                    };
                    // cargo_metadata crate returns default (empty) path for
                    // older cargos, which is not absolute, so work around that.
                    if out_dir != PathBuf::default() {
                        let out_dir = AbsPathBuf::assert(out_dir);
                        res.out_dirs.insert(package_id.clone(), out_dir);
                        res.cfgs.insert(package_id.clone(), cfgs);
                    }

                    res.env.insert(package_id, env);
                }
                Message::CompilerArtifact(message) => {
                    if message.target.kind.contains(&"proc-macro".to_string()) {
                        let package_id = message.package_id;
                        // Skip rmeta file
                        if let Some(filename) = message.filenames.iter().find(|name| is_dylib(name))
                        {
                            let filename = AbsPathBuf::assert(filename.clone());
                            res.proc_dylib_paths.insert(package_id, filename);
                        }
                    }
                }
                Message::CompilerMessage(_) => (),
                Message::Unknown => (),
                Message::BuildFinished(_) => {}
                Message::TextLine(_) => {}
            }
        }
    }
    Ok(res)
}

// FIXME: File a better way to know if it is a dylib
fn is_dylib(path: &Path) -> bool {
    match path.extension().and_then(OsStr::to_str).map(|it| it.to_string().to_lowercase()) {
        None => false,
        Some(ext) => matches!(ext.as_str(), "dll" | "dylib" | "so"),
    }
}
<commit_msg>Replicate Cargo environment variables<commit_after>//! FIXME: write short doc here

use std::{
    ffi::OsStr,
    ops,
    path::{Path, PathBuf},
    process::Command,
};

use anyhow::{Context, Result};
use arena::{Arena, Idx};
use base_db::Edition;
use cargo_metadata::{BuildScript, CargoOpt, Message, MetadataCommand, PackageId};
use paths::{AbsPath, AbsPathBuf};
use rustc_hash::FxHashMap;

use crate::cfg_flag::CfgFlag;

/// `CargoWorkspace` represents the logical structure of, well, a Cargo
/// workspace. It pretty closely mirrors `cargo metadata` output.
///
/// Note that internally, rust analyzer uses a different structure:
/// `CrateGraph`. `CrateGraph` is lower-level: it knows only about the crates,
/// while this knows about `Packages` & `Targets`: purely cargo-related
/// concepts.
///
/// We use absolute paths here, `cargo metadata` guarantees to always produce
/// abs paths.
#[derive(Debug, Clone, Eq, PartialEq)]
pub struct CargoWorkspace {
    packages: Arena<PackageData>,
    targets: Arena<TargetData>,
    workspace_root: AbsPathBuf,
}

impl ops::Index<Package> for CargoWorkspace {
    type Output = PackageData;
    fn index(&self, index: Package) -> &PackageData {
        &self.packages[index]
    }
}

impl ops::Index<Target> for CargoWorkspace {
    type Output = TargetData;
    fn index(&self, index: Target) -> &TargetData {
        &self.targets[index]
    }
}

#[derive(Default, Clone, Debug, PartialEq, Eq)]
pub struct CargoConfig {
    /// Do not activate the `default` feature.
    pub no_default_features: bool,

    /// Activate all available features
    pub all_features: bool,

    /// List of features to activate.
    /// This will be ignored if `cargo_all_features` is true.
    pub features: Vec<String>,

    /// Runs cargo check on launch to figure out the correct values of OUT_DIR
    pub load_out_dirs_from_check: bool,

    /// rustc target
    pub target: Option<String>,

    /// Don't load sysroot crates (`std`, `core` & friends). Might be useful
    /// when debugging isolated issues.
    pub no_sysroot: bool,

    /// rustc private crate source
    pub rustc_source: Option<AbsPathBuf>,
}

pub type Package = Idx<PackageData>;

pub type Target = Idx<TargetData>;

#[derive(Debug, Clone, Eq, PartialEq)]
pub struct PackageData {
    pub version: String,
    pub name: String,
    pub manifest: AbsPathBuf,
    pub targets: Vec<Target>,
    pub is_member: bool,
    pub dependencies: Vec<PackageDependency>,
    pub edition: Edition,
    pub features: Vec<String>,
    pub cfgs: Vec<CfgFlag>,
    pub envs: Vec<(String, String)>,
    pub out_dir: Option<AbsPathBuf>,
    pub proc_macro_dylib_path: Option<AbsPathBuf>,
}

#[derive(Debug, Clone, Eq, PartialEq)]
pub struct PackageDependency {
    pub pkg: Package,
    pub name: String,
}

#[derive(Debug, Clone, Eq, PartialEq)]
pub struct TargetData {
    pub package: Package,
    pub name: String,
    pub root: AbsPathBuf,
    pub kind: TargetKind,
    pub is_proc_macro: bool,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum TargetKind {
    Bin,
    /// Any kind of Cargo lib crate-type (dylib, rlib, proc-macro, ...).
    Lib,
    Example,
    Test,
    Bench,
    Other,
}

impl TargetKind {
    fn new(kinds: &[String]) -> TargetKind {
        for kind in kinds {
            return match kind.as_str() {
                "bin" => TargetKind::Bin,
                "test" => TargetKind::Test,
                "bench" => TargetKind::Bench,
                "example" => TargetKind::Example,
                "proc-macro" => TargetKind::Lib,
                _ if kind.contains("lib") => TargetKind::Lib,
                _ => continue,
            };
        }
        TargetKind::Other
    }
}

impl PackageData {
    pub fn root(&self) -> &AbsPath {
        self.manifest.parent().unwrap()
    }
}

impl CargoWorkspace {
    pub fn from_cargo_metadata(
        cargo_toml: &AbsPath,
        config: &CargoConfig,
    ) -> Result<CargoWorkspace> {
        let mut meta = MetadataCommand::new();
        meta.cargo_path(toolchain::cargo());
        meta.manifest_path(cargo_toml.to_path_buf());
        if config.all_features {
            meta.features(CargoOpt::AllFeatures);
        } else {
            if config.no_default_features {
                // FIXME: `NoDefaultFeatures` is mutual exclusive with `SomeFeatures`
                // https://github.com/oli-obk/cargo_metadata/issues/79
                meta.features(CargoOpt::NoDefaultFeatures);
            }
            if !config.features.is_empty() {
                meta.features(CargoOpt::SomeFeatures(config.features.clone()));
            }
        }
        if let Some(parent) = cargo_toml.parent() {
            meta.current_dir(parent.to_path_buf());
        }
        if let Some(target) = config.target.as_ref() {
            meta.other_options(vec![String::from("--filter-platform"), target.clone()]);
        }
        let mut meta = meta.exec().with_context(|| {
            format!("Failed to run `cargo metadata --manifest-path {}`", cargo_toml.display())
        })?;

        let mut out_dir_by_id = FxHashMap::default();
        let mut cfgs = FxHashMap::default();
        let mut envs = FxHashMap::default();
        let mut proc_macro_dylib_paths = FxHashMap::default();
        if config.load_out_dirs_from_check {
            let resources = load_extern_resources(cargo_toml, config)?;
            out_dir_by_id = resources.out_dirs;
            cfgs = resources.cfgs;
            envs = resources.env;
            proc_macro_dylib_paths = resources.proc_dylib_paths;
        }

        let mut pkg_by_id = FxHashMap::default();
        let mut packages = Arena::default();
        let mut targets = Arena::default();

        let ws_members = &meta.workspace_members;

        meta.packages.sort_by(|a, b| a.id.cmp(&b.id));
        for meta_pkg in meta.packages {
            let id = meta_pkg.id.clone();
            inject_cargo_env(&meta_pkg, envs.entry(id).or_default());

            let cargo_metadata::Package { id, edition, name, manifest_path, version, .. } =
                meta_pkg;
            let is_member = ws_members.contains(&id);
            let edition = edition
                .parse::<Edition>()
                .with_context(|| format!("Failed to parse edition {}", edition))?;
            let pkg = packages.alloc(PackageData {
                name,
                version: version.to_string(),
                manifest: AbsPathBuf::assert(manifest_path),
                targets: Vec::new(),
                is_member,
                edition,
                dependencies: Vec::new(),
                features: Vec::new(),
                cfgs: cfgs.get(&id).cloned().unwrap_or_default(),
                envs: envs.get(&id).cloned().unwrap_or_default(),
                out_dir: out_dir_by_id.get(&id).cloned(),
                proc_macro_dylib_path: proc_macro_dylib_paths.get(&id).cloned(),
            });
            let pkg_data = &mut packages[pkg];
            pkg_by_id.insert(id, pkg);
            for meta_tgt in meta_pkg.targets {
                let is_proc_macro = meta_tgt.kind.as_slice() == ["proc-macro"];
                let tgt = targets.alloc(TargetData {
                    package: pkg,
                    name: meta_tgt.name,
                    root: AbsPathBuf::assert(meta_tgt.src_path.clone()),
                    kind: TargetKind::new(meta_tgt.kind.as_slice()),
                    is_proc_macro,
                });
                pkg_data.targets.push(tgt);
            }
        }
        let resolve = meta.resolve.expect("metadata executed with deps");
        for mut node in resolve.nodes {
            let source = match pkg_by_id.get(&node.id) {
                Some(&src) => src,
                // FIXME: replace this and a similar branch below with `.unwrap`, once
                // https://github.com/rust-lang/cargo/issues/7841
                // is fixed and hits stable (around 1.43-is probably?).
                None => {
                    log::error!("Node id do not match in cargo metadata, ignoring {}", node.id);
                    continue;
                }
            };
            node.deps.sort_by(|a, b| a.pkg.cmp(&b.pkg));
            for dep_node in node.deps {
                let pkg = match pkg_by_id.get(&dep_node.pkg) {
                    Some(&pkg) => pkg,
                    None => {
                        log::error!(
                            "Dep node id do not match in cargo metadata, ignoring {}",
                            dep_node.pkg
                        );
                        continue;
                    }
                };
                let dep = PackageDependency { name: dep_node.name, pkg };
                packages[source].dependencies.push(dep);
            }
            packages[source].features.extend(node.features);
        }

        let workspace_root = AbsPathBuf::assert(meta.workspace_root);
        Ok(CargoWorkspace { packages, targets, workspace_root: workspace_root })
    }

    pub fn packages<'a>(&'a self) -> impl Iterator<Item = Package> + ExactSizeIterator + 'a {
        self.packages.iter().map(|(id, _pkg)| id)
    }

    pub fn target_by_root(&self, root: &AbsPath) -> Option<Target> {
        self.packages()
            .filter_map(|pkg| self[pkg].targets.iter().find(|&&it| &self[it].root == root))
            .next()
            .copied()
    }

    pub fn workspace_root(&self) -> &AbsPath {
        &self.workspace_root
    }

    pub fn package_flag(&self, package: &PackageData) -> String {
        if self.is_unique(&*package.name) {
            package.name.clone()
        } else {
            format!("{}:{}", package.name, package.version)
        }
    }

    fn is_unique(&self, name: &str) -> bool {
        self.packages.iter().filter(|(_, v)| v.name == name).count() == 1
    }
}

#[derive(Debug, Clone, Default)]
pub(crate) struct ExternResources {
    out_dirs: FxHashMap<PackageId, AbsPathBuf>,
    proc_dylib_paths: FxHashMap<PackageId, AbsPathBuf>,
    cfgs: FxHashMap<PackageId, Vec<CfgFlag>>,
    env: FxHashMap<PackageId, Vec<(String, String)>>,
}

pub(crate) fn load_extern_resources(
    cargo_toml: &Path,
    cargo_features: &CargoConfig,
) -> Result<ExternResources> {
    let mut cmd = Command::new(toolchain::cargo());
    cmd.args(&["check", "--message-format=json", "--manifest-path"]).arg(cargo_toml);

    if let Some(target) = &cargo_features.target {
        cmd.args(&["--target", target]);
    }

    if cargo_features.all_features {
        cmd.arg("--all-features");
    } else {
        if cargo_features.no_default_features {
            // FIXME: `NoDefaultFeatures` is mutual exclusive with `SomeFeatures`
            // https://github.com/oli-obk/cargo_metadata/issues/79
            cmd.arg("--no-default-features");
        }
        if !cargo_features.features.is_empty() {
            cmd.arg("--features");
            cmd.arg(cargo_features.features.join(" "));
        }
    }

    let output = cmd.output()?;

    let mut res = ExternResources::default();

    for message in cargo_metadata::Message::parse_stream(output.stdout.as_slice()) {
        if let Ok(message) = message {
            match message {
                Message::BuildScriptExecuted(BuildScript {
                    package_id,
                    out_dir,
                    cfgs,
                    env,
                    ..
                }) => {
                    let cfgs = {
                        let mut acc = Vec::new();
                        for cfg in cfgs {
                            match cfg.parse::<CfgFlag>() {
                                Ok(it) => acc.push(it),
                                Err(err) => {
                                    anyhow::bail!("invalid cfg from cargo-metadata: {}", err)
                                }
                            };
                        }
                        acc
                    };
                    // cargo_metadata crate returns default (empty) path for
                    // older cargos, which is not absolute, so work around that.
                    if out_dir != PathBuf::default() {
                        let out_dir = AbsPathBuf::assert(out_dir);
                        res.out_dirs.insert(package_id.clone(), out_dir);
                        res.cfgs.insert(package_id.clone(), cfgs);
                    }

                    res.env.insert(package_id, env);
                }
                Message::CompilerArtifact(message) => {
                    if message.target.kind.contains(&"proc-macro".to_string()) {
                        let package_id = message.package_id;
                        // Skip rmeta file
                        if let Some(filename) = message.filenames.iter().find(|name| is_dylib(name))
                        {
                            let filename = AbsPathBuf::assert(filename.clone());
                            res.proc_dylib_paths.insert(package_id, filename);
                        }
                    }
                }
                Message::CompilerMessage(_) => (),
                Message::Unknown => (),
                Message::BuildFinished(_) => {}
                Message::TextLine(_) => {}
            }
        }
    }
    Ok(res)
}

// FIXME: File a better way to know if it is a dylib
fn is_dylib(path: &Path) -> bool {
    match path.extension().and_then(OsStr::to_str).map(|it| it.to_string().to_lowercase()) {
        None => false,
        Some(ext) => matches!(ext.as_str(), "dll" | "dylib" | "so"),
    }
}

/// Recreates the compile-time environment variables that Cargo sets.
///
/// Should be synced with <https://doc.rust-lang.org/cargo/reference/environment-variables.html#environment-variables-cargo-sets-for-crates>
fn inject_cargo_env(package: &cargo_metadata::Package, env: &mut Vec<(String, String)>) {
    // FIXME: Missing variables:
    // CARGO, CARGO_PKG_HOMEPAGE, CARGO_CRATE_NAME, CARGO_BIN_NAME, CARGO_BIN_EXE_<name>

    let mut manifest_dir = package.manifest_path.clone();
    manifest_dir.pop();
    if let Some(cargo_manifest_dir) = manifest_dir.to_str() {
        env.push(("CARGO_MANIFEST_DIR".into(), cargo_manifest_dir.into()));
    }

    env.push(("CARGO_PKG_VERSION".into(), package.version.to_string()));
    env.push(("CARGO_PKG_VERSION_MAJOR".into(), package.version.major.to_string()));
    env.push(("CARGO_PKG_VERSION_MINOR".into(), package.version.minor.to_string()));
    env.push(("CARGO_PKG_VERSION_PATCH".into(), package.version.patch.to_string()));

    let pre = package.version.pre.iter().map(|id| id.to_string()).collect::<Vec<_>>();
    let pre = pre.join(".");
    env.push(("CARGO_PKG_VERSION_PRE".into(), pre));

    let authors = package.authors.join(";");
    env.push(("CARGO_PKG_AUTHORS".into(), authors));

    env.push(("CARGO_PKG_NAME".into(), package.name.clone()));
    env.push(("CARGO_PKG_DESCRIPTION".into(), package.description.clone().unwrap_or_default()));
    //env.push(("CARGO_PKG_HOMEPAGE".into(), package.homepage.clone().unwrap_or_default()));
    env.push(("CARGO_PKG_REPOSITORY".into(), package.repository.clone().unwrap_or_default()));
    env.push(("CARGO_PKG_LICENSE".into(), package.license.clone().unwrap_or_default()));

    let license_file =
        package.license_file.as_ref().map(|buf| buf.display().to_string()).unwrap_or_default();
    env.push(("CARGO_PKG_LICENSE_FILE".into(), license_file));
}
<|endoftext|>


<commit_before>//! FIXME: write short doc here

use hir::db::HirDatabase;
use ra_syntax::{ast::AstNode, ast::Literal, TextRange, TextUnit};
use rustc_lexer;

use crate::{Assist, AssistCtx, AssistId};

pub(crate) fn make_raw_string(mut ctx: AssistCtx<impl HirDatabase>) -> Option<Assist> {
    let literal = ctx.node_at_offset::<Literal>()?;
    if literal.token().kind() != ra_syntax::SyntaxKind::STRING {
        return None;
    }
    let token = literal.token();
    let text = token.text().as_str();
    let usual_string_range = find_usual_string_range(text)?;
    let start_of_inside = usual_string_range.start().to_usize() + 1;
    let end_of_inside = usual_string_range.end().to_usize();
    let inside_str = &text[start_of_inside..end_of_inside];
    let mut unescaped = String::with_capacity(inside_str.len());
    let mut error = Ok(());
    rustc_lexer::unescape::unescape_str(
        inside_str,
        &mut |_, unescaped_char| match unescaped_char {
            Ok(c) => unescaped.push(c),
            Err(_) => error = Err(()),
        },
    );
    if error.is_err() {
        return None;
    }
    ctx.add_action(AssistId("make_raw_string"), "make raw string", |edit| {
        edit.target(literal.syntax().text_range());
        let max_hash_streak = count_hashes(&unescaped);
        let mut hashes = String::with_capacity(max_hash_streak + 1);
        for _ in 0..hashes.capacity() {
            hashes.push('#');
        }
        edit.replace(
            literal.syntax().text_range(),
            format!("r{}\"{}\"{}", hashes, unescaped, hashes),
        );
    });
    ctx.build()
}

fn count_hashes(s: &str) -> usize {
    let mut max_hash_streak = 0usize;
    for idx in s.match_indices("\"#").map(|(i, _)| i) {
        let (_, sub) = s.split_at(idx + 1);
        let nb_hash = sub.chars().take_while(|c| *c == '#').count();
        if nb_hash > max_hash_streak {
            max_hash_streak = nb_hash;
        }
    }
    max_hash_streak
}

fn find_usual_string_range(s: &str) -> Option<TextRange> {
    Some(TextRange::from_to(
        TextUnit::from(s.find('"')? as u32),
        TextUnit::from(s.rfind('"')? as u32),
    ))
}

pub(crate) fn make_usual_string(mut ctx: AssistCtx<impl HirDatabase>) -> Option<Assist> {
    let literal = ctx.node_at_offset::<Literal>()?;
    if literal.token().kind() != ra_syntax::SyntaxKind::RAW_STRING {
        return None;
    }
    let token = literal.token();
    let text = token.text().as_str();
    let usual_string_range = find_usual_string_range(text)?;
    ctx.add_action(AssistId("make_usual_string"), "make usual string", |edit| {
        edit.target(literal.syntax().text_range());
        // parse inside string to escape `"`
        let start_of_inside = usual_string_range.start().to_usize() + 1;
        let end_of_inside = usual_string_range.end().to_usize();
        let inside_str = &text[start_of_inside..end_of_inside];
        let escaped = inside_str.escape_default().to_string();
        edit.replace(literal.syntax().text_range(), format!("\"{}\"", escaped));
    });
    ctx.build()
}

pub(crate) fn add_hash(mut ctx: AssistCtx<impl HirDatabase>) -> Option<Assist> {
    let literal = ctx.node_at_offset::<Literal>()?;
    if literal.token().kind() != ra_syntax::SyntaxKind::RAW_STRING {
        return None;
    }
    ctx.add_action(AssistId("add_hash"), "add hash to raw string", |edit| {
        edit.target(literal.syntax().text_range());
        edit.insert(literal.syntax().text_range().start() + TextUnit::of_char('r'), "#");
        edit.insert(literal.syntax().text_range().end(), "#");
    });
    ctx.build()
}

pub(crate) fn remove_hash(mut ctx: AssistCtx<impl HirDatabase>) -> Option<Assist> {
    let literal = ctx.node_at_offset::<Literal>()?;
    if literal.token().kind() != ra_syntax::SyntaxKind::RAW_STRING {
        return None;
    }
    let token = literal.token();
    let text = token.text().as_str();
    if text.starts_with("r\"") {
        // no hash to remove
        return None;
    }
    ctx.add_action(AssistId("remove_hash"), "remove hash from raw string", |edit| {
        edit.target(literal.syntax().text_range());
        let result = &text[2..text.len() - 1];
        let result = if result.starts_with("\"") {
            // no more hash, escape
            let internal_str = &result[1..result.len() - 1];
            format!("\"{}\"", internal_str.escape_default().to_string())
        } else {
            result.to_owned()
        };
        edit.replace(literal.syntax().text_range(), format!("r{}", result));
    });
    ctx.build()
}

#[cfg(test)]
mod test {
    use super::*;
    use crate::helpers::{check_assist, check_assist_not_applicable, check_assist_target};

    #[test]
    fn make_raw_string_target() {
        check_assist_target(
            make_raw_string,
            r#"
            fn f() {
                let s = <|>"random\nstring";
            }
            "#,
            r#""random\nstring""#,
        );
    }

    #[test]
    fn make_raw_string_works() {
        check_assist(
            make_raw_string,
            r#"
            fn f() {
                let s = <|>"random\nstring";
            }
            "#,
            r##"
            fn f() {
                let s = <|>r#"random
string"#;
            }
            "##,
        )
    }

    #[test]
    fn make_raw_string_hashes_inside_works() {
        check_assist(
            make_raw_string,
            r###"
            fn f() {
                let s = <|>"#random##\nstring";
            }
            "###,
            r####"
            fn f() {
                let s = <|>r#"#random##
string"#;
            }
            "####,
        )
    }

    #[test]
    fn make_raw_string_closing_hashes_inside_works() {
        check_assist(
            make_raw_string,
            r###"
            fn f() {
                let s = <|>"#random\"##\nstring";
            }
            "###,
            r####"
            fn f() {
                let s = <|>r###"#random"##
string"###;
            }
            "####,
        )
    }

    #[test]
    fn make_raw_string_nothing_to_unescape_works() {
        check_assist(
            make_raw_string,
            r#"
            fn f() {
                let s = <|>"random string";
            }
            "#,
            r##"
            fn f() {
                let s = <|>r#"random string"#;
            }
            "##,
        )
    }

    #[test]
    fn add_hash_target() {
        check_assist_target(
            add_hash,
            r#"
            fn f() {
                let s = <|>r"random string";
            }
            "#,
            r#"r"random string""#,
        );
    }

    #[test]
    fn add_hash_works() {
        check_assist(
            add_hash,
            r#"
            fn f() {
                let s = <|>r"random string";
            }
            "#,
            r##"
            fn f() {
                let s = <|>r#"random string"#;
            }
            "##,
        )
    }

    #[test]
    fn add_more_hash_works() {
        check_assist(
            add_hash,
            r##"
            fn f() {
                let s = <|>r#"random"string"#;
            }
            "##,
            r###"
            fn f() {
                let s = <|>r##"random"string"##;
            }
            "###,
        )
    }

    #[test]
    fn add_hash_not_works() {
        check_assist_not_applicable(
            add_hash,
            r#"
            fn f() {
                let s = <|>"random string";
            }
            "#,
        );
    }

    #[test]
    fn remove_hash_target() {
        check_assist_target(
            remove_hash,
            r##"
            fn f() {
                let s = <|>r#"random string"#;
            }
            "##,
            r##"r#"random string"#"##,
        );
    }

    #[test]
    fn remove_hash_works() {
        check_assist(
            remove_hash,
            r##"
            fn f() {
                let s = <|>r#"random string"#;
            }
            "##,
            r#"
            fn f() {
                let s = <|>r"random string";
            }
            "#,
        )
    }

    #[test]
    fn remove_hash_with_quote_works() {
        check_assist(
            remove_hash,
            r##"
            fn f() {
                let s = <|>r#"random"str"ing"#;
            }
            "##,
            r#"
            fn f() {
                let s = <|>r"random\"str\"ing";
            }
            "#,
        )
    }

    #[test]
    fn remove_more_hash_works() {
        check_assist(
            remove_hash,
            r###"
            fn f() {
                let s = <|>r##"random string"##;
            }
            "###,
            r##"
            fn f() {
                let s = <|>r#"random string"#;
            }
            "##,
        )
    }

    #[test]
    fn remove_hash_not_works() {
        check_assist_not_applicable(
            remove_hash,
            r#"
            fn f() {
                let s = <|>"random string";
            }
            "#,
        );
    }

    #[test]
    fn remove_hash_no_hash_not_works() {
        check_assist_not_applicable(
            remove_hash,
            r#"
            fn f() {
                let s = <|>r"random string";
            }
            "#,
        );
    }

    #[test]
    fn make_usual_string_target() {
        check_assist_target(
            make_usual_string,
            r##"
            fn f() {
                let s = <|>r#"random string"#;
            }
            "##,
            r##"r#"random string"#"##,
        );
    }

    #[test]
    fn make_usual_string_works() {
        check_assist(
            make_usual_string,
            r##"
            fn f() {
                let s = <|>r#"random string"#;
            }
            "##,
            r#"
            fn f() {
                let s = <|>"random string";
            }
            "#,
        )
    }

    #[test]
    fn make_usual_string_with_quote_works() {
        check_assist(
            make_usual_string,
            r##"
            fn f() {
                let s = <|>r#"random"str"ing"#;
            }
            "##,
            r#"
            fn f() {
                let s = <|>"random\"str\"ing";
            }
            "#,
        )
    }

    #[test]
    fn make_usual_string_more_hash_works() {
        check_assist(
            make_usual_string,
            r###"
            fn f() {
                let s = <|>r##"random string"##;
            }
            "###,
            r##"
            fn f() {
                let s = <|>"random string";
            }
            "##,
        )
    }

    #[test]
    fn make_usual_string_not_works() {
        check_assist_not_applicable(
            make_usual_string,
            r#"
            fn f() {
                let s = <|>"random string";
            }
            "#,
        );
    }

    #[test]
    fn count_hashes_test() {
        assert_eq!(0, count_hashes("abc"));
        assert_eq!(0, count_hashes("###"));
        assert_eq!(1, count_hashes("\"#abc"));
        assert_eq!(0, count_hashes("#abc"));
        assert_eq!(2, count_hashes("#ab\"##c"));
        assert_eq!(4, count_hashes("#ab\"##\"####c"));
    }
}
<commit_msg>raw string assists work in macros<commit_after>//! FIXME: write short doc here

use hir::db::HirDatabase;
use ra_syntax::{
    SyntaxKind::{RAW_STRING, STRING},
    SyntaxToken, TextRange, TextUnit,
};
use rustc_lexer;

use crate::{Assist, AssistCtx, AssistId};

pub(crate) fn make_raw_string(mut ctx: AssistCtx<impl HirDatabase>) -> Option<Assist> {
    let token = ctx.token_at_offset().right_biased()?;
    if token.kind() != STRING {
        return None;
    }
    let text = token.text().as_str();
    let usual_string_range = find_usual_string_range(text)?;
    let start_of_inside = usual_string_range.start().to_usize() + 1;
    let end_of_inside = usual_string_range.end().to_usize();
    let inside_str = &text[start_of_inside..end_of_inside];
    let mut unescaped = String::with_capacity(inside_str.len());
    let mut error = Ok(());
    rustc_lexer::unescape::unescape_str(
        inside_str,
        &mut |_, unescaped_char| match unescaped_char {
            Ok(c) => unescaped.push(c),
            Err(_) => error = Err(()),
        },
    );
    if error.is_err() {
        return None;
    }
    ctx.add_action(AssistId("make_raw_string"), "make raw string", |edit| {
        edit.target(token.text_range());
        let max_hash_streak = count_hashes(&unescaped);
        let mut hashes = String::with_capacity(max_hash_streak + 1);
        for _ in 0..hashes.capacity() {
            hashes.push('#');
        }
        edit.replace(token.text_range(), format!("r{}\"{}\"{}", hashes, unescaped, hashes));
    });
    ctx.build()
}

pub(crate) fn make_usual_string(mut ctx: AssistCtx<impl HirDatabase>) -> Option<Assist> {
    let token = raw_string_token(&ctx)?;
    let text = token.text().as_str();
    let usual_string_range = find_usual_string_range(text)?;
    ctx.add_action(AssistId("make_usual_string"), "make usual string", |edit| {
        edit.target(token.text_range());
        // parse inside string to escape `"`
        let start_of_inside = usual_string_range.start().to_usize() + 1;
        let end_of_inside = usual_string_range.end().to_usize();
        let inside_str = &text[start_of_inside..end_of_inside];
        let escaped = inside_str.escape_default().to_string();
        edit.replace(token.text_range(), format!("\"{}\"", escaped));
    });
    ctx.build()
}

pub(crate) fn add_hash(mut ctx: AssistCtx<impl HirDatabase>) -> Option<Assist> {
    let token = raw_string_token(&ctx)?;
    ctx.add_action(AssistId("add_hash"), "add hash to raw string", |edit| {
        edit.target(token.text_range());
        edit.insert(token.text_range().start() + TextUnit::of_char('r'), "#");
        edit.insert(token.text_range().end(), "#");
    });
    ctx.build()
}

pub(crate) fn remove_hash(mut ctx: AssistCtx<impl HirDatabase>) -> Option<Assist> {
    let token = raw_string_token(&ctx)?;
    let text = token.text().as_str();
    if text.starts_with("r\"") {
        // no hash to remove
        return None;
    }
    ctx.add_action(AssistId("remove_hash"), "remove hash from raw string", |edit| {
        edit.target(token.text_range());
        let result = &text[2..text.len() - 1];
        let result = if result.starts_with("\"") {
            // no more hash, escape
            let internal_str = &result[1..result.len() - 1];
            format!("\"{}\"", internal_str.escape_default().to_string())
        } else {
            result.to_owned()
        };
        edit.replace(token.text_range(), format!("r{}", result));
    });
    ctx.build()
}

fn raw_string_token(ctx: &AssistCtx<impl HirDatabase>) -> Option<SyntaxToken> {
    ctx.token_at_offset().right_biased().filter(|it| it.kind() == RAW_STRING)
}

fn count_hashes(s: &str) -> usize {
    let mut max_hash_streak = 0usize;
    for idx in s.match_indices("\"#").map(|(i, _)| i) {
        let (_, sub) = s.split_at(idx + 1);
        let nb_hash = sub.chars().take_while(|c| *c == '#').count();
        if nb_hash > max_hash_streak {
            max_hash_streak = nb_hash;
        }
    }
    max_hash_streak
}

fn find_usual_string_range(s: &str) -> Option<TextRange> {
    Some(TextRange::from_to(
        TextUnit::from(s.find('"')? as u32),
        TextUnit::from(s.rfind('"')? as u32),
    ))
}

#[cfg(test)]
mod test {
    use super::*;
    use crate::helpers::{check_assist, check_assist_not_applicable, check_assist_target};

    #[test]
    fn make_raw_string_target() {
        check_assist_target(
            make_raw_string,
            r#"
            fn f() {
                let s = <|>"random\nstring";
            }
            "#,
            r#""random\nstring""#,
        );
    }

    #[test]
    fn make_raw_string_works() {
        check_assist(
            make_raw_string,
            r#"
            fn f() {
                let s = <|>"random\nstring";
            }
            "#,
            r##"
            fn f() {
                let s = <|>r#"random
string"#;
            }
            "##,
        )
    }

    #[test]
    fn make_raw_string_works_inside_macros() {
        check_assist(
            make_raw_string,
            r#"
            fn f() {
                format!(<|>"x = {}", 92)
            }
            "#,
            r##"
            fn f() {
                format!(<|>r#"x = {}"#, 92)
            }
            "##,
        )
    }

    #[test]
    fn make_raw_string_hashes_inside_works() {
        check_assist(
            make_raw_string,
            r###"
            fn f() {
                let s = <|>"#random##\nstring";
            }
            "###,
            r####"
            fn f() {
                let s = <|>r#"#random##
string"#;
            }
            "####,
        )
    }

    #[test]
    fn make_raw_string_closing_hashes_inside_works() {
        check_assist(
            make_raw_string,
            r###"
            fn f() {
                let s = <|>"#random\"##\nstring";
            }
            "###,
            r####"
            fn f() {
                let s = <|>r###"#random"##
string"###;
            }
            "####,
        )
    }

    #[test]
    fn make_raw_string_nothing_to_unescape_works() {
        check_assist(
            make_raw_string,
            r#"
            fn f() {
                let s = <|>"random string";
            }
            "#,
            r##"
            fn f() {
                let s = <|>r#"random string"#;
            }
            "##,
        )
    }

    #[test]
    fn add_hash_target() {
        check_assist_target(
            add_hash,
            r#"
            fn f() {
                let s = <|>r"random string";
            }
            "#,
            r#"r"random string""#,
        );
    }

    #[test]
    fn add_hash_works() {
        check_assist(
            add_hash,
            r#"
            fn f() {
                let s = <|>r"random string";
            }
            "#,
            r##"
            fn f() {
                let s = <|>r#"random string"#;
            }
            "##,
        )
    }

    #[test]
    fn add_more_hash_works() {
        check_assist(
            add_hash,
            r##"
            fn f() {
                let s = <|>r#"random"string"#;
            }
            "##,
            r###"
            fn f() {
                let s = <|>r##"random"string"##;
            }
            "###,
        )
    }

    #[test]
    fn add_hash_not_works() {
        check_assist_not_applicable(
            add_hash,
            r#"
            fn f() {
                let s = <|>"random string";
            }
            "#,
        );
    }

    #[test]
    fn remove_hash_target() {
        check_assist_target(
            remove_hash,
            r##"
            fn f() {
                let s = <|>r#"random string"#;
            }
            "##,
            r##"r#"random string"#"##,
        );
    }

    #[test]
    fn remove_hash_works() {
        check_assist(
            remove_hash,
            r##"
            fn f() {
                let s = <|>r#"random string"#;
            }
            "##,
            r#"
            fn f() {
                let s = <|>r"random string";
            }
            "#,
        )
    }

    #[test]
    fn remove_hash_with_quote_works() {
        check_assist(
            remove_hash,
            r##"
            fn f() {
                let s = <|>r#"random"str"ing"#;
            }
            "##,
            r#"
            fn f() {
                let s = <|>r"random\"str\"ing";
            }
            "#,
        )
    }

    #[test]
    fn remove_more_hash_works() {
        check_assist(
            remove_hash,
            r###"
            fn f() {
                let s = <|>r##"random string"##;
            }
            "###,
            r##"
            fn f() {
                let s = <|>r#"random string"#;
            }
            "##,
        )
    }

    #[test]
    fn remove_hash_not_works() {
        check_assist_not_applicable(
            remove_hash,
            r#"
            fn f() {
                let s = <|>"random string";
            }
            "#,
        );
    }

    #[test]
    fn remove_hash_no_hash_not_works() {
        check_assist_not_applicable(
            remove_hash,
            r#"
            fn f() {
                let s = <|>r"random string";
            }
            "#,
        );
    }

    #[test]
    fn make_usual_string_target() {
        check_assist_target(
            make_usual_string,
            r##"
            fn f() {
                let s = <|>r#"random string"#;
            }
            "##,
            r##"r#"random string"#"##,
        );
    }

    #[test]
    fn make_usual_string_works() {
        check_assist(
            make_usual_string,
            r##"
            fn f() {
                let s = <|>r#"random string"#;
            }
            "##,
            r#"
            fn f() {
                let s = <|>"random string";
            }
            "#,
        )
    }

    #[test]
    fn make_usual_string_with_quote_works() {
        check_assist(
            make_usual_string,
            r##"
            fn f() {
                let s = <|>r#"random"str"ing"#;
            }
            "##,
            r#"
            fn f() {
                let s = <|>"random\"str\"ing";
            }
            "#,
        )
    }

    #[test]
    fn make_usual_string_more_hash_works() {
        check_assist(
            make_usual_string,
            r###"
            fn f() {
                let s = <|>r##"random string"##;
            }
            "###,
            r##"
            fn f() {
                let s = <|>"random string";
            }
            "##,
        )
    }

    #[test]
    fn make_usual_string_not_works() {
        check_assist_not_applicable(
            make_usual_string,
            r#"
            fn f() {
                let s = <|>"random string";
            }
            "#,
        );
    }

    #[test]
    fn count_hashes_test() {
        assert_eq!(0, count_hashes("abc"));
        assert_eq!(0, count_hashes("###"));
        assert_eq!(1, count_hashes("\"#abc"));
        assert_eq!(0, count_hashes("#abc"));
        assert_eq!(2, count_hashes("#ab\"##c"));
        assert_eq!(4, count_hashes("#ab\"##\"####c"));
    }
}
<|endoftext|>


<commit_before>mod event;
mod input;

use std::cell::Cell;

use crate::{
    SmolStr,
    syntax_node::syntax_error::{ParseError, SyntaxError},
    parsing::{
    lexer::Token,
    parser_api::Parser,
    parser_impl::{
        event::{Event, EventProcessor},
        input::{InputPosition, ParserInput},
    },
}};

use crate::SyntaxKind::{self, EOF, TOMBSTONE};

pub(crate) trait Sink {
    type Tree;

    /// Adds new leaf to the current branch.
    fn leaf(&mut self, kind: SyntaxKind, text: SmolStr);

    /// Start new branch and make it current.
    fn start_branch(&mut self, kind: SyntaxKind);

    /// Finish current branch and restore previous
    /// branch as current.
    fn finish_branch(&mut self);

    fn error(&mut self, error: SyntaxError);

    /// Complete tree building. Make sure that
    /// `start_branch` and `finish_branch` calls
    /// are paired!
    fn finish(self) -> Self::Tree;
}

/// Parse a sequence of tokens into the representative node tree
pub(crate) fn parse_with<S: Sink>(
    sink: S,
    text: &str,
    tokens: &[Token],
    parser: fn(&mut Parser),
) -> S::Tree {
    let mut events = {
        let input = input::ParserInput::new(text, tokens);
        let parser_impl = ParserImpl::new(&input);
        let mut parser_api = Parser(parser_impl);
        parser(&mut parser_api);
        parser_api.0.into_events()
    };
    EventProcessor::new(sink, text, tokens, &mut events).process().finish()
}

/// Implementation details of `Parser`, extracted
/// to a separate struct in order not to pollute
/// the public API of the `Parser`.
pub(crate) struct ParserImpl<'t> {
    parser_input: &'t ParserInput<'t>,
    pos: InputPosition,
    events: Vec<Event>,
    steps: Cell<u32>,
}

impl<'t> ParserImpl<'t> {
    fn new(inp: &'t ParserInput<'t>) -> ParserImpl<'t> {
        ParserImpl {
            parser_input: inp,
            pos: InputPosition::new(),
            events: Vec::new(),
            steps: Cell::new(0),
        }
    }

    fn into_events(self) -> Vec<Event> {
        assert_eq!(self.nth(0), EOF);
        self.events
    }

    pub(super) fn current2(&self) -> Option<(SyntaxKind, SyntaxKind)> {
        let c1 = self.parser_input.kind(self.pos);
        let c2 = self.parser_input.kind(self.pos + 1);
        if self.parser_input.token_start_at(self.pos + 1)
            == self.parser_input.token_start_at(self.pos) + self.parser_input.token_len(self.pos)
        {
            Some((c1, c2))
        } else {
            None
        }
    }

    pub(super) fn current3(&self) -> Option<(SyntaxKind, SyntaxKind, SyntaxKind)> {
        let c1 = self.parser_input.kind(self.pos);
        let c2 = self.parser_input.kind(self.pos + 1);
        let c3 = self.parser_input.kind(self.pos + 2);
        if self.parser_input.token_start_at(self.pos + 1)
            == self.parser_input.token_start_at(self.pos) + self.parser_input.token_len(self.pos)
            && self.parser_input.token_start_at(self.pos + 2)
                == self.parser_input.token_start_at(self.pos + 1)
                    + self.parser_input.token_len(self.pos + 1)
        {
            Some((c1, c2, c3))
        } else {
            None
        }
    }

    /// Get the syntax kind of the nth token.
    pub(super) fn nth(&self, n: u32) -> SyntaxKind {
        let steps = self.steps.get();
        assert!(steps <= 10_000_000, "the parser seems stuck");
        self.steps.set(steps + 1);

        self.parser_input.kind(self.pos + n)
    }

    pub(super) fn at_kw(&self, t: &str) -> bool {
        self.parser_input.token_text(self.pos) == t
    }

    /// Start parsing right behind the last event.
    pub(super) fn start(&mut self) -> u32 {
        let pos = self.events.len() as u32;
        self.push_event(Event::tombstone());
        pos
    }

    /// Advances the parser by one token unconditionally.
    pub(super) fn bump(&mut self) {
        let kind = self.nth(0);
        if kind == EOF {
            return;
        }
        self.do_bump(kind, 1);
    }

    pub(super) fn bump_remap(&mut self, kind: SyntaxKind) {
        if self.nth(0) == EOF {
            // TODO: panic!?
            return;
        }
        self.do_bump(kind, 1);
    }

    pub(super) fn bump_compound(&mut self, kind: SyntaxKind, n: u8) {
        self.do_bump(kind, n);
    }

    fn do_bump(&mut self, kind: SyntaxKind, n_raw_tokens: u8) {
        self.pos += u32::from(n_raw_tokens);
        self.push_event(Event::Token { kind, n_raw_tokens });
    }

    /// Append one Error event to the back of events.
    pub(super) fn error(&mut self, msg: String) {
        self.push_event(Event::Error { msg: ParseError(msg) })
    }

    /// Complete an event with appending a `Finish` event.
    pub(super) fn complete(&mut self, pos: u32, kind: SyntaxKind) {
        match self.events[pos as usize] {
            Event::Start { kind: ref mut slot, .. } => {
                *slot = kind;
            }
            _ => unreachable!(),
        }
        self.push_event(Event::Finish);
    }

    /// Ignore the dummy `Start` event.
    pub(super) fn abandon(&mut self, pos: u32) {
        let idx = pos as usize;
        if idx == self.events.len() - 1 {
            match self.events.pop() {
                Some(Event::Start { kind: TOMBSTONE, forward_parent: None }) => (),
                _ => unreachable!(),
            }
        }
    }

    /// Save the relative distance of a completed event to its forward_parent.
    pub(super) fn precede(&mut self, pos: u32) -> u32 {
        let new_pos = self.start();
        match self.events[pos as usize] {
            Event::Start { ref mut forward_parent, .. } => {
                *forward_parent = Some(new_pos - pos);
            }
            _ => unreachable!(),
        }
        new_pos
    }

    fn push_event(&mut self, event: Event) {
        self.events.push(event)
    }
}
<commit_msg>tighten visibility<commit_after>mod event;
mod input;

use std::cell::Cell;

use crate::{
    SmolStr,
    syntax_node::syntax_error::{ParseError, SyntaxError},
    parsing::{
    lexer::Token,
    parser_api::Parser,
    parser_impl::{
        event::{Event, EventProcessor},
        input::{InputPosition, ParserInput},
    },
}};

use crate::SyntaxKind::{self, EOF, TOMBSTONE};

pub(super) trait Sink {
    type Tree;

    /// Adds new leaf to the current branch.
    fn leaf(&mut self, kind: SyntaxKind, text: SmolStr);

    /// Start new branch and make it current.
    fn start_branch(&mut self, kind: SyntaxKind);

    /// Finish current branch and restore previous
    /// branch as current.
    fn finish_branch(&mut self);

    fn error(&mut self, error: SyntaxError);

    /// Complete tree building. Make sure that
    /// `start_branch` and `finish_branch` calls
    /// are paired!
    fn finish(self) -> Self::Tree;
}

/// Parse a sequence of tokens into the representative node tree
pub(super) fn parse_with<S: Sink>(
    sink: S,
    text: &str,
    tokens: &[Token],
    parser: fn(&mut Parser),
) -> S::Tree {
    let mut events = {
        let input = input::ParserInput::new(text, tokens);
        let parser_impl = ParserImpl::new(&input);
        let mut parser_api = Parser(parser_impl);
        parser(&mut parser_api);
        parser_api.0.into_events()
    };
    EventProcessor::new(sink, text, tokens, &mut events).process().finish()
}

/// Implementation details of `Parser`, extracted
/// to a separate struct in order not to pollute
/// the public API of the `Parser`.
pub(super) struct ParserImpl<'t> {
    parser_input: &'t ParserInput<'t>,
    pos: InputPosition,
    events: Vec<Event>,
    steps: Cell<u32>,
}

impl<'t> ParserImpl<'t> {
    fn new(inp: &'t ParserInput<'t>) -> ParserImpl<'t> {
        ParserImpl {
            parser_input: inp,
            pos: InputPosition::new(),
            events: Vec::new(),
            steps: Cell::new(0),
        }
    }

    fn into_events(self) -> Vec<Event> {
        assert_eq!(self.nth(0), EOF);
        self.events
    }

    pub(super) fn current2(&self) -> Option<(SyntaxKind, SyntaxKind)> {
        let c1 = self.parser_input.kind(self.pos);
        let c2 = self.parser_input.kind(self.pos + 1);
        if self.parser_input.token_start_at(self.pos + 1)
            == self.parser_input.token_start_at(self.pos) + self.parser_input.token_len(self.pos)
        {
            Some((c1, c2))
        } else {
            None
        }
    }

    pub(super) fn current3(&self) -> Option<(SyntaxKind, SyntaxKind, SyntaxKind)> {
        let c1 = self.parser_input.kind(self.pos);
        let c2 = self.parser_input.kind(self.pos + 1);
        let c3 = self.parser_input.kind(self.pos + 2);
        if self.parser_input.token_start_at(self.pos + 1)
            == self.parser_input.token_start_at(self.pos) + self.parser_input.token_len(self.pos)
            && self.parser_input.token_start_at(self.pos + 2)
                == self.parser_input.token_start_at(self.pos + 1)
                    + self.parser_input.token_len(self.pos + 1)
        {
            Some((c1, c2, c3))
        } else {
            None
        }
    }

    /// Get the syntax kind of the nth token.
    pub(super) fn nth(&self, n: u32) -> SyntaxKind {
        let steps = self.steps.get();
        assert!(steps <= 10_000_000, "the parser seems stuck");
        self.steps.set(steps + 1);

        self.parser_input.kind(self.pos + n)
    }

    pub(super) fn at_kw(&self, t: &str) -> bool {
        self.parser_input.token_text(self.pos) == t
    }

    /// Start parsing right behind the last event.
    pub(super) fn start(&mut self) -> u32 {
        let pos = self.events.len() as u32;
        self.push_event(Event::tombstone());
        pos
    }

    /// Advances the parser by one token unconditionally.
    pub(super) fn bump(&mut self) {
        let kind = self.nth(0);
        if kind == EOF {
            return;
        }
        self.do_bump(kind, 1);
    }

    pub(super) fn bump_remap(&mut self, kind: SyntaxKind) {
        if self.nth(0) == EOF {
            // TODO: panic!?
            return;
        }
        self.do_bump(kind, 1);
    }

    pub(super) fn bump_compound(&mut self, kind: SyntaxKind, n: u8) {
        self.do_bump(kind, n);
    }

    fn do_bump(&mut self, kind: SyntaxKind, n_raw_tokens: u8) {
        self.pos += u32::from(n_raw_tokens);
        self.push_event(Event::Token { kind, n_raw_tokens });
    }

    /// Append one Error event to the back of events.
    pub(super) fn error(&mut self, msg: String) {
        self.push_event(Event::Error { msg: ParseError(msg) })
    }

    /// Complete an event with appending a `Finish` event.
    pub(super) fn complete(&mut self, pos: u32, kind: SyntaxKind) {
        match self.events[pos as usize] {
            Event::Start { kind: ref mut slot, .. } => {
                *slot = kind;
            }
            _ => unreachable!(),
        }
        self.push_event(Event::Finish);
    }

    /// Ignore the dummy `Start` event.
    pub(super) fn abandon(&mut self, pos: u32) {
        let idx = pos as usize;
        if idx == self.events.len() - 1 {
            match self.events.pop() {
                Some(Event::Start { kind: TOMBSTONE, forward_parent: None }) => (),
                _ => unreachable!(),
            }
        }
    }

    /// Save the relative distance of a completed event to its forward_parent.
    pub(super) fn precede(&mut self, pos: u32) -> u32 {
        let new_pos = self.start();
        match self.events[pos as usize] {
            Event::Start { ref mut forward_parent, .. } => {
                *forward_parent = Some(new_pos - pos);
            }
            _ => unreachable!(),
        }
        new_pos
    }

    fn push_event(&mut self, event: Event) {
        self.events.push(event)
    }
}
<|endoftext|>


<commit_before><commit_msg>Add FromIterator<Comparator> for VersionReq<commit_after><|endoftext|>


<commit_before>#![feature(plugin)]
#![feature(custom_derive)]
#![plugin(rocket_codegen)]

#[macro_use]
extern crate clap;
extern crate config;
extern crate diesel;
#[macro_use]
extern crate failure;
extern crate r2d2;
extern crate r2d2_diesel;
// extern crate ring;
extern crate log;
extern crate rocket;
extern crate rocket_contrib;
extern crate serde;
extern crate yaml_rust;
extern crate rocket_i18n;

#[cfg(feature = "log-simple")]
extern crate simple_logging;
#[cfg(feature = "log-syslog")]
extern crate syslog;
#[cfg(feature = "use-systemd")]
extern crate systemd;

extern crate _aardwolf as aardwolf;

mod common;
use common::{configure, db_conn_string};

use failure::Error;
// use ring::rand::SystemRandom;
use clap::App;
use diesel::pg::PgConnection;
use log::LevelFilter;
use r2d2_diesel::ConnectionManager;
use rocket::Rocket;
use rocket_contrib::Template;

type Pool = r2d2::Pool<ConnectionManager<PgConnection>>;

fn db_pool(rocket: &Rocket) -> Result<Pool, Error> {
    let database_url = rocket.config().get_str("database_url")?;
    let manager = ConnectionManager::<PgConnection>::new(database_url);
    Ok(r2d2::Pool::builder().build(manager)?)
}

fn app(config: config::Config) -> Result<Rocket, Error> {
    let db_url = db_conn_string(&config)?;

    let c = rocket::Config::build(rocket::config::Environment::Development)
        .address(config.get_str("Web.Listen.address")?)
        .port(config.get::<u16>("Web.Listen.port")?)
        .extra("database_url", db_url.as_str())
        .unwrap();

    let mut routes = routes![
        aardwolf::routes::app::home,
        aardwolf::routes::app::home_redirect,
    ];

    #[cfg(debug_assertions)]
    routes.extend(routes![
        // webroot/favicon
        aardwolf::routes::app::webroot,
        // emoji
        aardwolf::routes::app::emoji,
        // themes
        aardwolf::routes::app::themes,
    ]);

    let auth = routes![
        aardwolf::routes::auth::sign_up_form,
        aardwolf::routes::auth::sign_up_form_with_error,
        aardwolf::routes::auth::sign_in_form,
        aardwolf::routes::auth::sign_in_form_with_error,
        aardwolf::routes::auth::sign_up,
        aardwolf::routes::auth::sign_in,
        aardwolf::routes::auth::confirm,
        aardwolf::routes::auth::sign_out,
    ];

    let personas = routes![
        aardwolf::routes::personas::new,
        aardwolf::routes::personas::create,
        aardwolf::routes::personas::delete,
        aardwolf::routes::personas::switch,
    ];

    let r = rocket::custom(c, true)
        .mount("/auth", auth)
        .mount("/personas", personas)
        .mount(
            "/api/v1",
            routes![aardwolf::routes::applications::register_application],
        )
        .mount("/", routes)
        .attach(Template::fairing());
    // .manage(SystemRandom::new());

    // we need an instance of the app to access the config values in Rocket.toml,
    // so we pass it to the db_pool function, get the pool, and _then_ return the instance
    let pool = db_pool(&r)?;
    Ok(r.manage(pool))

    // Just for giggles, what happens if I put the rocket_i18n fairing here....

    // Register the fairing. The parameter is the domain you want to use (the name of your app most of the time)
    .attach(rocket_i18n::I18n::new("my_app"))
    // Eventually register the Tera filters (only works with the master branch of Rocket)
    .attach(rocket_contrib::Template::custom(|engines| {
        rocket_i18n::tera(&mut engines.tera);
    }))
}

fn cli<'a, 'b>(yaml: &'a yaml_rust::yaml::Yaml) -> App<'a, 'b> {
    App::from_yaml(yaml)
        .name(env!("CARGO_PKG_NAME"))
        .version(env!("CARGO_PKG_VERSION"))
        .author(env!("CARGO_PKG_AUTHORS"))
        .about(env!("CARGO_PKG_DESCRIPTION"))
}

#[cfg(feature = "log-simple")]
fn begin_log(config: &config::Config) {
    match config.get_str("log_file").unwrap().as_ref() {
        "_CONSOLE_" => (),
        l => simple_logging::log_to_file(l, LevelFilter::Info).unwrap(),
    }
}

#[cfg(feature = "log-syslog")]
fn begin_log(config: &config::Config) {
    // TODO: Implement log-syslog:begin_log()
}

#[cfg(feature = "use-systemd")]
fn begin_log(config: &config::Config) {
    // TODO: Implement use-systemd:begin_log()
}

fn main() {
    let yaml = load_yaml!("cli.yml");
    let cli = cli(&yaml);
    let config = configure(cli).unwrap();
    begin_log(&config);

    app(config).unwrap().launch();
}
<commit_msg>Move rocket_i18n fairings<commit_after>#![feature(plugin)]
#![feature(custom_derive)]
#![plugin(rocket_codegen)]

#[macro_use]
extern crate clap;
extern crate config;
extern crate diesel;
#[macro_use]
extern crate failure;
extern crate r2d2;
extern crate r2d2_diesel;
// extern crate ring;
extern crate log;
extern crate rocket;
extern crate rocket_contrib;
extern crate serde;
extern crate yaml_rust;
extern crate rocket_i18n;

#[cfg(feature = "log-simple")]
extern crate simple_logging;
#[cfg(feature = "log-syslog")]
extern crate syslog;
#[cfg(feature = "use-systemd")]
extern crate systemd;

extern crate _aardwolf as aardwolf;

mod common;
use common::{configure, db_conn_string};

use failure::Error;
// use ring::rand::SystemRandom;
use clap::App;
use diesel::pg::PgConnection;
use log::LevelFilter;
use r2d2_diesel::ConnectionManager;
use rocket::Rocket;
use rocket_contrib::Template;

type Pool = r2d2::Pool<ConnectionManager<PgConnection>>;

fn db_pool(rocket: &Rocket) -> Result<Pool, Error> {
    let database_url = rocket.config().get_str("database_url")?;
    let manager = ConnectionManager::<PgConnection>::new(database_url);
    Ok(r2d2::Pool::builder().build(manager)?)
}

fn app(config: config::Config) -> Result<Rocket, Error> {
    let db_url = db_conn_string(&config)?;

    let c = rocket::Config::build(rocket::config::Environment::Development)
        .address(config.get_str("Web.Listen.address")?)
        .port(config.get::<u16>("Web.Listen.port")?)
        .extra("database_url", db_url.as_str())
        .unwrap();

    let mut routes = routes![
        aardwolf::routes::app::home,
        aardwolf::routes::app::home_redirect,
    ];

    #[cfg(debug_assertions)]
    routes.extend(routes![
        // webroot/favicon
        aardwolf::routes::app::webroot,
        // emoji
        aardwolf::routes::app::emoji,
        // themes
        aardwolf::routes::app::themes,
    ]);

    let auth = routes![
        aardwolf::routes::auth::sign_up_form,
        aardwolf::routes::auth::sign_up_form_with_error,
        aardwolf::routes::auth::sign_in_form,
        aardwolf::routes::auth::sign_in_form_with_error,
        aardwolf::routes::auth::sign_up,
        aardwolf::routes::auth::sign_in,
        aardwolf::routes::auth::confirm,
        aardwolf::routes::auth::sign_out,
    ];

    let personas = routes![
        aardwolf::routes::personas::new,
        aardwolf::routes::personas::create,
        aardwolf::routes::personas::delete,
        aardwolf::routes::personas::switch,
    ];

    let r = rocket::custom(c, true)
        .mount("/auth", auth)
        .mount("/personas", personas)
        .mount(
            "/api/v1",
            routes![aardwolf::routes::applications::register_application],
        )
        .mount("/", routes)
        .attach(Template::fairing())//;
    // .manage(SystemRandom::new());

        // Just for giggles, what happens if I put the rocket_i18n fairing here....

        // Register the fairing. The parameter is the domain you want to use (the name of your app most of the time)
        .attach(rocket_i18n::I18n::new("my_app"))
        // Eventually register the Tera filters (only works with the master branch of Rocket)
        .attach(rocket_contrib::Template::custom(|engines| {
                rocket_i18n::tera(&mut engines.tera);
        }));

    // we need an instance of the app to access the config values in Rocket.toml,
    // so we pass it to the db_pool function, get the pool, and _then_ return the instance
    let pool = db_pool(&r)?;
    Ok(r.manage(pool))
}

fn cli<'a, 'b>(yaml: &'a yaml_rust::yaml::Yaml) -> App<'a, 'b> {
    App::from_yaml(yaml)
        .name(env!("CARGO_PKG_NAME"))
        .version(env!("CARGO_PKG_VERSION"))
        .author(env!("CARGO_PKG_AUTHORS"))
        .about(env!("CARGO_PKG_DESCRIPTION"))
}

#[cfg(feature = "log-simple")]
fn begin_log(config: &config::Config) {
    match config.get_str("log_file").unwrap().as_ref() {
        "_CONSOLE_" => (),
        l => simple_logging::log_to_file(l, LevelFilter::Info).unwrap(),
    }
}

#[cfg(feature = "log-syslog")]
fn begin_log(config: &config::Config) {
    // TODO: Implement log-syslog:begin_log()
}

#[cfg(feature = "use-systemd")]
fn begin_log(config: &config::Config) {
    // TODO: Implement use-systemd:begin_log()
}

fn main() {
    let yaml = load_yaml!("cli.yml");
    let cli = cli(&yaml);
    let config = configure(cli).unwrap();
    begin_log(&config);

    app(config).unwrap().launch();
}
<|endoftext|>


<commit_before><commit_msg>refactor: Remove redundant check<commit_after><|endoftext|>


<commit_before><commit_msg>Convert get_password() Mac implementation to use the keychain<commit_after><|endoftext|>


<commit_before><commit_msg>Fixed bug with typechecking Apps<commit_after><|endoftext|>


<commit_before>mod header;

use std::net::IpAddr;
use std::io::{Write, ErrorKind};
use mio::*;
use mio::tcp::TcpStream;
use mio::unix::UnixReady;

use self::header::*;
use network::ClientResult;
use network::Readiness;
use network::protocol::ProtocolResult;
use network::socket::SocketHandler;

pub struct ProxyProtocol<Front:SocketHandler> {
  pub header:     Option<Vec<u8>>,
  pub frontend:   Front,
  pub backend:    Option<TcpStream>,
  frontend_token: Option<Token>,
  backend_token:  Option<Token>,
  pub readiness:  Readiness,
  cursor_header:  usize,
}

impl <Front:SocketHandler>ProxyProtocol<Front> {
  pub fn new(frontend: Front, backend: Option<TcpStream>) -> Self {
    ProxyProtocol {
      header: None,
      frontend,
      backend,
      frontend_token: None,
      backend_token:  None,
      readiness: Readiness {
        front_interest:  UnixReady::hup() | UnixReady::error(),
        back_interest:   UnixReady::from(Ready::writable()) | UnixReady::hup() | UnixReady::error(),
        front_readiness: UnixReady::from(Ready::empty()),
        back_readiness:  UnixReady::from(Ready::empty()),
      },
      cursor_header: 0,
    }
  }

  // The header is send immediately at once upon the connection is establish
  // and prepended before any data.
  pub fn back_writable(&mut self) -> (ProtocolResult, ClientResult) {
    debug!("Writing proxy protocol header");

    if let Some(ref mut socket) = self.backend {
      if let Some(ref mut header) = self.header {
        loop {
          match socket.write(&mut header[self.cursor_header..]) {
            Ok(sz) => {
              self.cursor_header += sz;

              if self.cursor_header == header.len() {
                debug!("Proxy protocol sent, upgrading");
                return (ProtocolResult::Upgrade, ClientResult::Continue)
              }
            },
            Err(e) => {
              debug!("PROXY PROTOCOL {}", e);
              break;
            },
          }
        }
      }
    }
    (ProtocolResult::Continue, ClientResult::Continue)
  }

  pub fn front_socket(&self) -> &TcpStream {
    self.frontend.socket_ref()
  }

  pub fn back_socket(&self) -> Option<&TcpStream> {
    self.backend.as_ref()
  }

  pub fn set_back_socket(&mut self, socket: TcpStream) {
    self.gen_proxy_protocol_header(&socket);
    self.backend = Some(socket);
  }

  pub fn front_token(&self) -> Option<Token> {
    self.frontend_token
  }

  pub fn set_front_token(&mut self, token: Token) {
    self.frontend_token = Some(token);
  }

  pub fn back_token(&self) -> Option<Token> {
    self.backend_token
  }

  pub fn set_back_token(&mut self, token: Token) {
    self.backend_token = Some(token);
  }

  pub fn readiness(&mut self) -> &mut Readiness {
    &mut self.readiness
  }

  fn gen_proxy_protocol_header(&mut self, socket: &TcpStream) {
    let addr_frontend = self.frontend.socket_ref().peer_addr().unwrap();
    let addr_backend = socket.peer_addr().unwrap();

    let protocol_header = ProxyProtocolHeader::V1(HeaderV1::new(addr_frontend, addr_backend));
    self.header = Some(protocol_header.into_bytes());
  }
}
<commit_msg>Send a V2 proxy protocol header instead of V1<commit_after>mod header;

use std::net::IpAddr;
use std::io::{Write, ErrorKind};
use mio::*;
use mio::tcp::TcpStream;
use mio::unix::UnixReady;

use self::header::*;
use network::ClientResult;
use network::Readiness;
use network::protocol::ProtocolResult;
use network::socket::SocketHandler;

pub struct ProxyProtocol<Front:SocketHandler> {
  pub header:     Option<Vec<u8>>,
  pub frontend:   Front,
  pub backend:    Option<TcpStream>,
  frontend_token: Option<Token>,
  backend_token:  Option<Token>,
  pub readiness:  Readiness,
  cursor_header:  usize,
}

impl <Front:SocketHandler>ProxyProtocol<Front> {
  pub fn new(frontend: Front, backend: Option<TcpStream>) -> Self {
    ProxyProtocol {
      header: None,
      frontend,
      backend,
      frontend_token: None,
      backend_token:  None,
      readiness: Readiness {
        front_interest:  UnixReady::hup() | UnixReady::error(),
        back_interest:   UnixReady::from(Ready::writable()) | UnixReady::hup() | UnixReady::error(),
        front_readiness: UnixReady::from(Ready::empty()),
        back_readiness:  UnixReady::from(Ready::empty()),
      },
      cursor_header: 0,
    }
  }

  // The header is send immediately at once upon the connection is establish
  // and prepended before any data.
  pub fn back_writable(&mut self) -> (ProtocolResult, ClientResult) {
    debug!("Writing proxy protocol header");

    if let Some(ref mut socket) = self.backend {
      if let Some(ref mut header) = self.header {
        loop {
          match socket.write(&mut header[self.cursor_header..]) {
            Ok(sz) => {
              self.cursor_header += sz;

              if self.cursor_header == header.len() {
                debug!("Proxy protocol sent, upgrading");
                return (ProtocolResult::Upgrade, ClientResult::Continue)
              }
            },
            Err(e) => {
              debug!("PROXY PROTOCOL {}", e);
              break;
            },
          }
        }
      }
    }
    (ProtocolResult::Continue, ClientResult::Continue)
  }

  pub fn front_socket(&self) -> &TcpStream {
    self.frontend.socket_ref()
  }

  pub fn back_socket(&self) -> Option<&TcpStream> {
    self.backend.as_ref()
  }

  pub fn set_back_socket(&mut self, socket: TcpStream) {
    self.gen_proxy_protocol_header(&socket);
    self.backend = Some(socket);
  }

  pub fn front_token(&self) -> Option<Token> {
    self.frontend_token
  }

  pub fn set_front_token(&mut self, token: Token) {
    self.frontend_token = Some(token);
  }

  pub fn back_token(&self) -> Option<Token> {
    self.backend_token
  }

  pub fn set_back_token(&mut self, token: Token) {
    self.backend_token = Some(token);
  }

  pub fn readiness(&mut self) -> &mut Readiness {
    &mut self.readiness
  }

  fn gen_proxy_protocol_header(&mut self, socket: &TcpStream) {
    let addr_frontend = self.frontend.socket_ref().peer_addr().unwrap();
    let addr_backend = socket.peer_addr().unwrap();

    let protocol_header = ProxyProtocolHeader::V2(HeaderV2::new(addr_frontend, addr_backend));
    self.header = Some(protocol_header.into_bytes());
  }
}
<|endoftext|>