<commit_before>// Copyright (c) 2017 Chef Software Inc. and/or applicable contributors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use std::net::IpAddr;
use std::os::unix::process::CommandExt;
use std::path::PathBuf;
use std::process::{Command, ExitStatus, Stdio};
use std::sync::Mutex;
use std::sync::atomic::{AtomicUsize, ATOMIC_USIZE_INIT, Ordering};

use hab_core::channel::{BLDR_CHANNEL_ENVVAR, STABLE_CHANNEL};
use hab_core::env;
use hab_core::fs;
use hab_core::url::BLDR_URL_ENVVAR;
use hab_core::AUTH_TOKEN_ENVVAR;

use error::{Error, Result};
use runner::log_pipe::LogPipe;
use runner::{NONINTERACTIVE_ENVVAR, RUNNER_DEBUG_ENVVAR};
use runner::workspace::Workspace;

pub static STUDIO_UID: AtomicUsize = ATOMIC_USIZE_INIT;
pub static STUDIO_GID: AtomicUsize = ATOMIC_USIZE_INIT;
pub const DEBUG_ENVVARS: &'static [&'static str] = &["RUST_LOG", "DEBUG"];
pub const STUDIO_USER: &'static str = "krangschnak";
pub const STUDIO_GROUP: &'static str = "krangschnak";

lazy_static! {
    /// Absolute path to the Studio program
    static ref STUDIO_PROGRAM: PathBuf = fs::resolve_cmd_in_pkg(
        "hab-studio",
        include_str!(concat!(env!("OUT_DIR"), "/STUDIO_PKG_IDENT")),
    );

    pub static ref STUDIO_HOME: Mutex<PathBuf> = {
        Mutex::new(PathBuf::new())
    };
}

pub struct Studio<'a> {
    workspace: &'a Workspace,
    bldr_url: &'a str,
    auth_token: &'a str,
    airlock_enabled: bool,
    networking: Option<(&'a str, &'a IpAddr)>,
}

impl<'a> Studio<'a> {
    /// Creates a new Studio runner for a given `Workspace` and Builder URL.
    pub fn new(
        workspace: &'a Workspace,
        bldr_url: &'a str,
        auth_token: &'a str,
        airlock_enabled: bool,
        networking: Option<(&'a str, &'a IpAddr)>,
    ) -> Self {
        Studio {
            workspace,
            bldr_url,
            auth_token,
            airlock_enabled,
            networking,
        }
    }

    /// Spawns a Studio build command, pipes output streams to the given `LogPipe` and returns the
    /// process' `ExitStatus`.
    ///
    /// # Errors
    ///
    /// * If the child process can't be spawned
    /// * If the calling thread can't wait on the child process
    /// * If the `LogPipe` fails to pipe output
    pub fn build(&self, log_pipe: &mut LogPipe) -> Result<ExitStatus> {
        if self.networking.is_some() {
            self.create_network_namespace()?;
        } else {
            info!("Airlock networking is not configured, skipping network creation");
        }

        let channel = if self.workspace.job.has_channel() {
            self.workspace.job.get_channel()
        } else {
            STABLE_CHANNEL
        };

        let mut cmd = self.studio_command()?;
        cmd.current_dir(self.workspace.src());
        if let Some(val) = env::var_os(RUNNER_DEBUG_ENVVAR) {
            cmd.env("DEBUG", val);
        }
        cmd.env(NONINTERACTIVE_ENVVAR, "true"); // Disables progress bars
        cmd.env("TERM", "xterm-256color"); // Emits ANSI color codes
        // propagate debugging environment variables into Airlock and Studio
        for var in DEBUG_ENVVARS {
            if let Ok(val) = env::var(var) {
                cmd.env(var, val);
            }
        }
        cmd.stdout(Stdio::piped());
        cmd.stderr(Stdio::piped());
        cmd.arg("-k"); // Origin key
        cmd.arg(self.workspace.job.origin());
        cmd.arg("build");
        cmd.arg(build_path(self.workspace.job.get_project().get_plan_path()));
        debug!("building studio build command, cmd={:?}", &cmd);
        debug!(
            "setting studio build command env, {}={}",
            BLDR_CHANNEL_ENVVAR,
            &channel
        );
        cmd.env(BLDR_CHANNEL_ENVVAR, channel);
        debug!(
            "setting studio build command env, {}={}",
            BLDR_URL_ENVVAR,
            self.bldr_url
        );
        cmd.env(BLDR_URL_ENVVAR, self.bldr_url);
        cmd.env(AUTH_TOKEN_ENVVAR, self.auth_token);

        debug!("spawning studio build command");
        let mut child = cmd.spawn().map_err(|e| {
            Error::StudioBuild(self.workspace.studio().to_path_buf(), e)
        })?;
        log_pipe.pipe(&mut child)?;
        let exit_status = child.wait().map_err(|e| {
            Error::StudioBuild(self.workspace.studio().to_path_buf(), e)
        })?;
        debug!("completed studio build command, status={:?}", exit_status);

        if self.networking.is_some() {
            self.destroy_network_namespace()?;
        } else {
            info!("Airlock networking is not configured, skipping network destruction");
        }

        Ok(exit_status)
    }

    fn studio_command(&self) -> Result<Command> {
        if self.airlock_enabled {
            let mut cmd = Command::new("airlock");
            cmd.uid(studio_uid());
            cmd.gid(studio_gid());
            cmd.env_clear();
            cmd.env("HOME", &*STUDIO_HOME.lock().unwrap()); // Sets `$HOME` for build user
            cmd.env("USER", STUDIO_USER); // Sets `$USER` for build user
            cmd.env("PATH", env::var("PATH").unwrap_or(String::from(""))); // Sets `$PATH`
            cmd.arg("run");
            cmd.arg("--fs-root");
            cmd.arg(self.workspace.studio());
            cmd.arg("--no-rm");
            if self.networking.is_some() {
                cmd.arg("--use-userns");
                cmd.arg(self.workspace.ns_dir().join("userns"));
                cmd.arg("--use-netns");
                cmd.arg(self.workspace.ns_dir().join("netns"));
            }
            cmd.arg(&*STUDIO_PROGRAM);

            Ok(cmd)
        } else {
            let mut cmd = Command::new(&*STUDIO_PROGRAM);
            cmd.env_clear();
            cmd.env("NO_ARTIFACT_PATH", "true"); // Disables artifact cache mounting
            cmd.env("HAB_CACHE_KEY_PATH", key_path()); // Sets key cache to build user's home

            info!("Airlock is not enabled, running uncontained Studio");
            Ok(cmd)
        }
    }

    fn create_network_namespace(&self) -> Result<()> {
        let mut cmd = Command::new("airlock");
        cmd.arg("netns");
        cmd.arg("create");
        cmd.arg("--interface");
        cmd.arg(self.networking.unwrap().0);
        cmd.arg("--gateway");
        cmd.arg(self.networking.unwrap().1.to_string());
        cmd.arg("--ns-dir");
        cmd.arg(self.workspace.ns_dir());
        cmd.arg("--user");
        cmd.arg(STUDIO_USER);
        debug!("building airlock networking setup command, cmd={:?}", &cmd);

        debug!("spawning airlock networking setup command");
        let mut child = cmd.spawn().map_err(|e| {
            Error::AirlockNetworking(self.workspace.ns_dir().to_path_buf(), e)
        })?;
        let exit_status = child.wait().map_err(|e| {
            Error::AirlockNetworking(self.workspace.ns_dir().to_path_buf(), e)
        })?;
        debug!(
            "completed airlock networking setup command, status={:?}",
            exit_status
        );

        if exit_status.success() {
            Ok(())
        } else {
            Err(Error::AirlockFailure(exit_status))
        }
    }

    fn destroy_network_namespace(&self) -> Result<()> {
        let mut cmd = Command::new("airlock");
        cmd.arg("netns");
        cmd.arg("destroy");
        cmd.arg("--ns-dir");
        cmd.arg(self.workspace.ns_dir());
        debug!(
            "building airlock networking destroy command, cmd={:?}",
            &cmd
        );

        debug!("spawning airlock networking destroy command");
        let mut child = cmd.spawn().map_err(|e| {
            Error::AirlockNetworking(self.workspace.ns_dir().to_path_buf(), e)
        })?;
        let exit_status = child.wait().map_err(|e| {
            Error::AirlockNetworking(self.workspace.ns_dir().to_path_buf(), e)
        })?;
        debug!(
            "completed airlock networking destroy command, status={:?}",
            exit_status
        );

        if exit_status.success() {
            Ok(())
        } else {
            Err(Error::AirlockFailure(exit_status))
        }
    }
}

pub fn studio_gid() -> u32 {
    STUDIO_GID.load(Ordering::Relaxed) as u32
}

pub fn studio_uid() -> u32 {
    STUDIO_UID.load(Ordering::Relaxed) as u32
}

pub fn set_studio_gid(gid: u32) {
    STUDIO_GID.store(gid as usize, Ordering::Relaxed);
}

pub fn set_studio_uid(uid: u32) {
    STUDIO_UID.store(uid as usize, Ordering::Relaxed);
}

pub fn key_path() -> PathBuf {
    (&*STUDIO_HOME).lock().unwrap().join(format!(
        ".{}",
        fs::CACHE_KEY_PATH
    ))
}

/// Returns a path argument suitable to pass to a Studio build command.
fn build_path(plan_path: &str) -> String {
    let mut parts: Vec<_> = plan_path.split("/").collect();
    if parts.last().map_or("", |p| *p) == "plan.sh" {
        parts.pop();
    }
    if parts.last().map_or("", |p| *p) == "habitat" {
        parts.pop();
    }

    if parts.is_empty() {
        String::from(".")
    } else {
        parts.join("/")
    }
}

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

    #[test]
    fn build_path_with_plan_sh() {
        assert_eq!(".", build_path("plan.sh"));
    }

    #[test]
    fn build_path_with_habitat_plan_sh() {
        assert_eq!(".", build_path("habitat/plan.sh"));
    }

    #[test]
    fn build_path_with_subdir_plan_sh() {
        assert_eq!("haaay", build_path("haaay/plan.sh"));
    }

    #[test]
    fn build_path_with_subdir_habitat_plan_sh() {
        assert_eq!(
            "components/yep",
            build_path("components/yep/habitat/plan.sh")
        );
    }
}
<commit_msg>[builder-worker] Move `PATH` envvar into common `build` method.<commit_after>// Copyright (c) 2017 Chef Software Inc. and/or applicable contributors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use std::net::IpAddr;
use std::os::unix::process::CommandExt;
use std::path::PathBuf;
use std::process::{Command, ExitStatus, Stdio};
use std::sync::Mutex;
use std::sync::atomic::{AtomicUsize, ATOMIC_USIZE_INIT, Ordering};

use hab_core::channel::{BLDR_CHANNEL_ENVVAR, STABLE_CHANNEL};
use hab_core::env;
use hab_core::fs;
use hab_core::url::BLDR_URL_ENVVAR;
use hab_core::AUTH_TOKEN_ENVVAR;

use error::{Error, Result};
use runner::log_pipe::LogPipe;
use runner::{NONINTERACTIVE_ENVVAR, RUNNER_DEBUG_ENVVAR};
use runner::workspace::Workspace;

pub static STUDIO_UID: AtomicUsize = ATOMIC_USIZE_INIT;
pub static STUDIO_GID: AtomicUsize = ATOMIC_USIZE_INIT;
pub const DEBUG_ENVVARS: &'static [&'static str] = &["RUST_LOG", "DEBUG"];
pub const STUDIO_USER: &'static str = "krangschnak";
pub const STUDIO_GROUP: &'static str = "krangschnak";

lazy_static! {
    /// Absolute path to the Studio program
    static ref STUDIO_PROGRAM: PathBuf = fs::resolve_cmd_in_pkg(
        "hab-studio",
        include_str!(concat!(env!("OUT_DIR"), "/STUDIO_PKG_IDENT")),
    );

    pub static ref STUDIO_HOME: Mutex<PathBuf> = {
        Mutex::new(PathBuf::new())
    };
}

pub struct Studio<'a> {
    workspace: &'a Workspace,
    bldr_url: &'a str,
    auth_token: &'a str,
    airlock_enabled: bool,
    networking: Option<(&'a str, &'a IpAddr)>,
}

impl<'a> Studio<'a> {
    /// Creates a new Studio runner for a given `Workspace` and Builder URL.
    pub fn new(
        workspace: &'a Workspace,
        bldr_url: &'a str,
        auth_token: &'a str,
        airlock_enabled: bool,
        networking: Option<(&'a str, &'a IpAddr)>,
    ) -> Self {
        Studio {
            workspace,
            bldr_url,
            auth_token,
            airlock_enabled,
            networking,
        }
    }

    /// Spawns a Studio build command, pipes output streams to the given `LogPipe` and returns the
    /// process' `ExitStatus`.
    ///
    /// # Errors
    ///
    /// * If the child process can't be spawned
    /// * If the calling thread can't wait on the child process
    /// * If the `LogPipe` fails to pipe output
    pub fn build(&self, log_pipe: &mut LogPipe) -> Result<ExitStatus> {
        if self.networking.is_some() {
            self.create_network_namespace()?;
        } else {
            info!("Airlock networking is not configured, skipping network creation");
        }

        let channel = if self.workspace.job.has_channel() {
            self.workspace.job.get_channel()
        } else {
            STABLE_CHANNEL
        };

        let mut cmd = self.studio_command()?;
        cmd.current_dir(self.workspace.src());
        if let Some(val) = env::var_os(RUNNER_DEBUG_ENVVAR) {
            cmd.env("DEBUG", val);
        }
        cmd.env("PATH", env::var("PATH").unwrap_or(String::from(""))); // Sets `$PATH`
        cmd.env(NONINTERACTIVE_ENVVAR, "true"); // Disables progress bars
        cmd.env("TERM", "xterm-256color"); // Emits ANSI color codes
        // propagate debugging environment variables into Airlock and Studio
        for var in DEBUG_ENVVARS {
            if let Ok(val) = env::var(var) {
                cmd.env(var, val);
            }
        }
        cmd.stdout(Stdio::piped());
        cmd.stderr(Stdio::piped());
        cmd.arg("-k"); // Origin key
        cmd.arg(self.workspace.job.origin());
        cmd.arg("build");
        cmd.arg(build_path(self.workspace.job.get_project().get_plan_path()));
        debug!("building studio build command, cmd={:?}", &cmd);
        debug!(
            "setting studio build command env, {}={}",
            BLDR_CHANNEL_ENVVAR,
            &channel
        );
        cmd.env(BLDR_CHANNEL_ENVVAR, channel);
        debug!(
            "setting studio build command env, {}={}",
            BLDR_URL_ENVVAR,
            self.bldr_url
        );
        cmd.env(BLDR_URL_ENVVAR, self.bldr_url);
        cmd.env(AUTH_TOKEN_ENVVAR, self.auth_token);

        debug!("spawning studio build command");
        let mut child = cmd.spawn().map_err(|e| {
            Error::StudioBuild(self.workspace.studio().to_path_buf(), e)
        })?;
        log_pipe.pipe(&mut child)?;
        let exit_status = child.wait().map_err(|e| {
            Error::StudioBuild(self.workspace.studio().to_path_buf(), e)
        })?;
        debug!("completed studio build command, status={:?}", exit_status);

        if self.networking.is_some() {
            self.destroy_network_namespace()?;
        } else {
            info!("Airlock networking is not configured, skipping network destruction");
        }

        Ok(exit_status)
    }

    fn studio_command(&self) -> Result<Command> {
        if self.airlock_enabled {
            let mut cmd = Command::new("airlock");
            cmd.uid(studio_uid());
            cmd.gid(studio_gid());
            cmd.env_clear();
            cmd.env("HOME", &*STUDIO_HOME.lock().unwrap()); // Sets `$HOME` for build user
            cmd.env("USER", STUDIO_USER); // Sets `$USER` for build user
            cmd.arg("run");
            cmd.arg("--fs-root");
            cmd.arg(self.workspace.studio());
            cmd.arg("--no-rm");
            if self.networking.is_some() {
                cmd.arg("--use-userns");
                cmd.arg(self.workspace.ns_dir().join("userns"));
                cmd.arg("--use-netns");
                cmd.arg(self.workspace.ns_dir().join("netns"));
            }
            cmd.arg(&*STUDIO_PROGRAM);

            Ok(cmd)
        } else {
            let mut cmd = Command::new(&*STUDIO_PROGRAM);
            cmd.env_clear();
            cmd.env("NO_ARTIFACT_PATH", "true"); // Disables artifact cache mounting
            cmd.env("HAB_CACHE_KEY_PATH", key_path()); // Sets key cache to build user's home

            info!("Airlock is not enabled, running uncontained Studio");
            Ok(cmd)
        }
    }

    fn create_network_namespace(&self) -> Result<()> {
        let mut cmd = Command::new("airlock");
        cmd.arg("netns");
        cmd.arg("create");
        cmd.arg("--interface");
        cmd.arg(self.networking.unwrap().0);
        cmd.arg("--gateway");
        cmd.arg(self.networking.unwrap().1.to_string());
        cmd.arg("--ns-dir");
        cmd.arg(self.workspace.ns_dir());
        cmd.arg("--user");
        cmd.arg(STUDIO_USER);
        debug!("building airlock networking setup command, cmd={:?}", &cmd);

        debug!("spawning airlock networking setup command");
        let mut child = cmd.spawn().map_err(|e| {
            Error::AirlockNetworking(self.workspace.ns_dir().to_path_buf(), e)
        })?;
        let exit_status = child.wait().map_err(|e| {
            Error::AirlockNetworking(self.workspace.ns_dir().to_path_buf(), e)
        })?;
        debug!(
            "completed airlock networking setup command, status={:?}",
            exit_status
        );

        if exit_status.success() {
            Ok(())
        } else {
            Err(Error::AirlockFailure(exit_status))
        }
    }

    fn destroy_network_namespace(&self) -> Result<()> {
        let mut cmd = Command::new("airlock");
        cmd.arg("netns");
        cmd.arg("destroy");
        cmd.arg("--ns-dir");
        cmd.arg(self.workspace.ns_dir());
        debug!(
            "building airlock networking destroy command, cmd={:?}",
            &cmd
        );

        debug!("spawning airlock networking destroy command");
        let mut child = cmd.spawn().map_err(|e| {
            Error::AirlockNetworking(self.workspace.ns_dir().to_path_buf(), e)
        })?;
        let exit_status = child.wait().map_err(|e| {
            Error::AirlockNetworking(self.workspace.ns_dir().to_path_buf(), e)
        })?;
        debug!(
            "completed airlock networking destroy command, status={:?}",
            exit_status
        );

        if exit_status.success() {
            Ok(())
        } else {
            Err(Error::AirlockFailure(exit_status))
        }
    }
}

pub fn studio_gid() -> u32 {
    STUDIO_GID.load(Ordering::Relaxed) as u32
}

pub fn studio_uid() -> u32 {
    STUDIO_UID.load(Ordering::Relaxed) as u32
}

pub fn set_studio_gid(gid: u32) {
    STUDIO_GID.store(gid as usize, Ordering::Relaxed);
}

pub fn set_studio_uid(uid: u32) {
    STUDIO_UID.store(uid as usize, Ordering::Relaxed);
}

pub fn key_path() -> PathBuf {
    (&*STUDIO_HOME).lock().unwrap().join(format!(
        ".{}",
        fs::CACHE_KEY_PATH
    ))
}

/// Returns a path argument suitable to pass to a Studio build command.
fn build_path(plan_path: &str) -> String {
    let mut parts: Vec<_> = plan_path.split("/").collect();
    if parts.last().map_or("", |p| *p) == "plan.sh" {
        parts.pop();
    }
    if parts.last().map_or("", |p| *p) == "habitat" {
        parts.pop();
    }

    if parts.is_empty() {
        String::from(".")
    } else {
        parts.join("/")
    }
}

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

    #[test]
    fn build_path_with_plan_sh() {
        assert_eq!(".", build_path("plan.sh"));
    }

    #[test]
    fn build_path_with_habitat_plan_sh() {
        assert_eq!(".", build_path("habitat/plan.sh"));
    }

    #[test]
    fn build_path_with_subdir_plan_sh() {
        assert_eq!("haaay", build_path("haaay/plan.sh"));
    }

    #[test]
    fn build_path_with_subdir_habitat_plan_sh() {
        assert_eq!(
            "components/yep",
            build_path("components/yep/habitat/plan.sh")
        );
    }
}
<|endoftext|>


<commit_before>use std::io::Cursor;
use std::ops::Deref;
use std::time::{SystemTime, UNIX_EPOCH};

use crate::{
    crypto::{self, double_sha256, Digest, KeyPair},
    serializer::*,
    tx::*,
};

#[derive(Clone, Debug, PartialEq)]
pub struct Block {
    pub previous_hash: Digest,
    pub height: u64,
    pub timestamp: u64,
    pub tx_merkle_root: Digest,
    pub transactions: Vec<TxVariant>,
}

impl Block {
    pub fn sign(self, key_pair: &KeyPair) -> SignedBlock {
        let buf = self.calc_hash();
        SignedBlock {
            base: self,
            sig_pair: key_pair.sign(buf.as_ref()),
        }
    }

    pub fn serialize_with_tx(&self, vec: &mut Vec<u8>) {
        self.serialize(vec);

        vec.push_u32(self.transactions.len() as u32);
        for tx in &self.transactions {
            tx.serialize_with_sigs(vec)
        }
    }

    pub fn deserialize_with_tx(cur: &mut Cursor<&[u8]>) -> Option<Self> {
        let previous_hash = Digest::from_slice(&cur.take_bytes().ok()?)?;
        let height = cur.take_u64().ok()?;
        let timestamp = cur.take_u64().ok()?;
        let tx_merkle_root = Digest::from_slice(&cur.take_bytes().ok()?)?;

        let len = cur.take_u32().ok()?;
        let mut transactions = Vec::<TxVariant>::with_capacity(len as usize);
        for _ in 0..len {
            transactions.push(TxVariant::deserialize_with_sigs(cur)?);
        }

        Some(Self {
            previous_hash,
            height,
            timestamp,
            tx_merkle_root,
            transactions,
        })
    }

    fn serialize(&self, vec: &mut Vec<u8>) {
        vec.push_bytes(self.previous_hash.as_ref());
        vec.push_u64(self.height);
        vec.push_u64(self.timestamp);
        vec.push_bytes(self.tx_merkle_root.as_ref());
    }

    pub fn verify_tx_merkle_root(&self) -> bool {
        let mut buf = Vec::with_capacity(4096 * self.transactions.len());
        for tx in &self.transactions {
            tx.serialize_with_sigs(&mut buf)
        }
        let digest = double_sha256(&buf);
        self.tx_merkle_root == digest
    }

    pub fn calc_hash(&self) -> Digest {
        let mut buf = Vec::with_capacity(1024);
        self.serialize(&mut buf);
        double_sha256(&buf)
    }

    pub fn verify_previous_hash(&self, prev_block: &Block) -> bool {
        self.previous_hash == prev_block.calc_hash()
    }
}

#[derive(Clone, Debug, PartialEq)]
pub struct SignedBlock {
    pub base: Block,
    pub sig_pair: crypto::SigPair,
}

impl SignedBlock {
    pub fn new_child(&self, txs: Vec<TxVariant>) -> Block {
        let previous_hash = {
            let mut buf = Vec::with_capacity(1024);
            self.base.serialize(&mut buf);
            double_sha256(&buf)
        };
        let tx_merkle_root = {
            let mut buf = Vec::with_capacity(4096 * txs.len());
            for tx in &txs {
                tx.serialize_with_sigs(&mut buf)
            }
            double_sha256(&buf)
        };
        let timestamp = SystemTime::now()
            .duration_since(UNIX_EPOCH)
            .unwrap()
            .as_millis() as u64;
        Block {
            previous_hash,
            height: self.height + 1,
            timestamp,
            tx_merkle_root,
            transactions: txs,
        }
    }

    pub fn serialize_with_tx(&self, vec: &mut Vec<u8>) {
        self.base.serialize_with_tx(vec);
        vec.push_sig_pair(&self.sig_pair);
    }

    pub fn deserialize_with_tx(cur: &mut Cursor<&[u8]>) -> Option<Self> {
        let block = Block::deserialize_with_tx(cur)?;
        let sig_pair = cur.take_sig_pair().ok()?;
        Some(Self {
            base: block,
            sig_pair,
        })
    }
}

impl Deref for SignedBlock {
    type Target = Block;

    fn deref(&self) -> &Self::Target {
        &self.base
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::{asset::Balance, crypto::KeyPair};

    #[test]
    fn test_serialize_block() {
        let keys = KeyPair::gen_keypair();
        let transactions = {
            let mut vec = Vec::new();
            vec.push(TxVariant::RewardTx(RewardTx {
                base: Tx {
                    tx_type: TxType::REWARD,
                    fee: "0 GOLD".parse().unwrap(),
                    timestamp: 1234567890,
                    signature_pairs: Vec::new(),
                },
                to: keys.0.clone().into(),
                rewards: Balance::default(),
            }));
            vec
        };
        let tx_merkle_root = {
            let mut buf = Vec::new();
            for tx in &transactions {
                tx.serialize_with_sigs(&mut buf)
            }
            double_sha256(&buf)
        };
        let block = (Block {
            previous_hash: Digest::from_slice(&[0u8; 32]).unwrap(),
            height: 123,
            timestamp: 1532992800,
            tx_merkle_root,
            transactions,
        })
        .sign(&keys);

        let mut buf = Vec::new();
        block.serialize_with_tx(&mut buf);

        let mut cur = Cursor::<&[u8]>::new(&buf);
        let dec = SignedBlock::deserialize_with_tx(&mut cur).unwrap();

        assert_eq!(block.previous_hash, dec.previous_hash);
        assert_eq!(block.height, dec.height);
        assert_eq!(block.timestamp, dec.timestamp);
        assert_eq!(block.tx_merkle_root, dec.tx_merkle_root);
        assert_eq!(block.transactions.len(), dec.transactions.len());
        assert_eq!(block.sig_pair, dec.sig_pair);

        assert_eq!(block, dec);
    }
}
<commit_msg>Rename Block::serialize to Block::serialize_header<commit_after>use std::io::Cursor;
use std::ops::Deref;
use std::time::{SystemTime, UNIX_EPOCH};

use crate::{
    crypto::{self, double_sha256, Digest, KeyPair},
    serializer::*,
    tx::*,
};

#[derive(Clone, Debug, PartialEq)]
pub struct Block {
    pub previous_hash: Digest,
    pub height: u64,
    pub timestamp: u64,
    pub tx_merkle_root: Digest,
    pub transactions: Vec<TxVariant>,
}

impl Block {
    pub fn sign(self, key_pair: &KeyPair) -> SignedBlock {
        let buf = self.calc_hash();
        SignedBlock {
            base: self,
            sig_pair: key_pair.sign(buf.as_ref()),
        }
    }

    pub fn serialize_with_tx(&self, buf: &mut Vec<u8>) {
        self.serialize_header(buf);

        buf.push_u32(self.transactions.len() as u32);
        for tx in &self.transactions {
            tx.serialize_with_sigs(buf)
        }
    }

    pub fn deserialize_with_tx(cur: &mut Cursor<&[u8]>) -> Option<Self> {
        let previous_hash = Digest::from_slice(&cur.take_bytes().ok()?)?;
        let height = cur.take_u64().ok()?;
        let timestamp = cur.take_u64().ok()?;
        let tx_merkle_root = Digest::from_slice(&cur.take_bytes().ok()?)?;

        let len = cur.take_u32().ok()?;
        let mut transactions = Vec::<TxVariant>::with_capacity(len as usize);
        for _ in 0..len {
            transactions.push(TxVariant::deserialize_with_sigs(cur)?);
        }

        Some(Self {
            previous_hash,
            height,
            timestamp,
            tx_merkle_root,
            transactions,
        })
    }

    fn serialize_header(&self, vec: &mut Vec<u8>) {
        vec.push_bytes(self.previous_hash.as_ref());
        vec.push_u64(self.height);
        vec.push_u64(self.timestamp);
        vec.push_bytes(self.tx_merkle_root.as_ref());
    }

    pub fn verify_tx_merkle_root(&self) -> bool {
        let mut buf = Vec::with_capacity(4096 * self.transactions.len());
        for tx in &self.transactions {
            tx.serialize_with_sigs(&mut buf)
        }
        let digest = double_sha256(&buf);
        self.tx_merkle_root == digest
    }

    pub fn calc_hash(&self) -> Digest {
        let mut buf = Vec::with_capacity(1024);
        self.serialize_header(&mut buf);
        double_sha256(&buf)
    }

    pub fn verify_previous_hash(&self, prev_block: &Block) -> bool {
        self.previous_hash == prev_block.calc_hash()
    }
}

#[derive(Clone, Debug, PartialEq)]
pub struct SignedBlock {
    pub base: Block,
    pub sig_pair: crypto::SigPair,
}

impl SignedBlock {
    pub fn new_child(&self, txs: Vec<TxVariant>) -> Block {
        let previous_hash = {
            let mut buf = Vec::with_capacity(1024);
            self.base.serialize_header(&mut buf);
            double_sha256(&buf)
        };
        let tx_merkle_root = {
            let mut buf = Vec::with_capacity(4096 * txs.len());
            for tx in &txs {
                tx.serialize_with_sigs(&mut buf)
            }
            double_sha256(&buf)
        };
        let timestamp = SystemTime::now()
            .duration_since(UNIX_EPOCH)
            .unwrap()
            .as_millis() as u64;
        Block {
            previous_hash,
            height: self.height + 1,
            timestamp,
            tx_merkle_root,
            transactions: txs,
        }
    }

    pub fn serialize_with_tx(&self, vec: &mut Vec<u8>) {
        self.base.serialize_with_tx(vec);
        vec.push_sig_pair(&self.sig_pair);
    }

    pub fn deserialize_with_tx(cur: &mut Cursor<&[u8]>) -> Option<Self> {
        let block = Block::deserialize_with_tx(cur)?;
        let sig_pair = cur.take_sig_pair().ok()?;
        Some(Self {
            base: block,
            sig_pair,
        })
    }
}

impl Deref for SignedBlock {
    type Target = Block;

    fn deref(&self) -> &Self::Target {
        &self.base
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::{asset::Balance, crypto::KeyPair};

    #[test]
    fn test_serialize_block() {
        let keys = KeyPair::gen_keypair();
        let transactions = {
            let mut vec = Vec::new();
            vec.push(TxVariant::RewardTx(RewardTx {
                base: Tx {
                    tx_type: TxType::REWARD,
                    fee: "0 GOLD".parse().unwrap(),
                    timestamp: 1234567890,
                    signature_pairs: Vec::new(),
                },
                to: keys.0.clone().into(),
                rewards: Balance::default(),
            }));
            vec
        };
        let tx_merkle_root = {
            let mut buf = Vec::new();
            for tx in &transactions {
                tx.serialize_with_sigs(&mut buf)
            }
            double_sha256(&buf)
        };
        let block = (Block {
            previous_hash: Digest::from_slice(&[0u8; 32]).unwrap(),
            height: 123,
            timestamp: 1532992800,
            tx_merkle_root,
            transactions,
        })
        .sign(&keys);

        let mut buf = Vec::new();
        block.serialize_with_tx(&mut buf);

        let mut cur = Cursor::<&[u8]>::new(&buf);
        let dec = SignedBlock::deserialize_with_tx(&mut cur).unwrap();

        assert_eq!(block.previous_hash, dec.previous_hash);
        assert_eq!(block.height, dec.height);
        assert_eq!(block.timestamp, dec.timestamp);
        assert_eq!(block.tx_merkle_root, dec.tx_merkle_root);
        assert_eq!(block.transactions.len(), dec.transactions.len());
        assert_eq!(block.sig_pair, dec.sig_pair);

        assert_eq!(block, dec);
    }
}
<|endoftext|>


<commit_before>//! Fully type-check project and print various stats, like the number of type
//! errors.

use std::{
    path::PathBuf,
    time::{SystemTime, UNIX_EPOCH},
};

use hir::{
    db::{AstDatabase, DefDatabase, HirDatabase},
    AssocItem, Crate, HasSource, HirDisplay, ModuleDef,
};
use hir_def::FunctionId;
use hir_ty::{Ty, TypeWalk};
use ide_db::base_db::{
    salsa::{self, ParallelDatabase},
    SourceDatabaseExt,
};
use itertools::Itertools;
use oorandom::Rand32;
use rayon::prelude::*;
use rustc_hash::FxHashSet;
use stdx::format_to;
use syntax::AstNode;

use crate::cli::{
    load_cargo::load_cargo, print_memory_usage, progress_report::ProgressReport, report_metric,
    Result, Verbosity,
};
use profile::StopWatch;

/// Need to wrap Snapshot to provide `Clone` impl for `map_with`
struct Snap<DB>(DB);
impl<DB: ParallelDatabase> Clone for Snap<salsa::Snapshot<DB>> {
    fn clone(&self) -> Snap<salsa::Snapshot<DB>> {
        Snap(self.0.snapshot())
    }
}

pub struct AnalysisStatsCmd {
    pub randomize: bool,
    pub parallel: bool,
    pub memory_usage: bool,
    pub only: Option<String>,
    pub with_deps: bool,
    pub path: PathBuf,
    pub load_output_dirs: bool,
    pub with_proc_macro: bool,
}

impl AnalysisStatsCmd {
    pub fn run(self, verbosity: Verbosity) -> Result<()> {
        let mut rng = {
            let seed = SystemTime::now().duration_since(UNIX_EPOCH).unwrap().as_millis() as u64;
            Rand32::new(seed)
        };

        let mut db_load_sw = self.stop_watch();
        let (host, vfs) = load_cargo(&self.path, self.load_output_dirs, self.with_proc_macro)?;
        let db = host.raw_database();
        eprintln!("Database loaded {}", db_load_sw.elapsed());

        let mut analysis_sw = self.stop_watch();
        let mut num_crates = 0;
        let mut visited_modules = FxHashSet::default();
        let mut visit_queue = Vec::new();

        let mut krates = Crate::all(db);
        if self.randomize {
            shuffle(&mut rng, &mut krates);
        }
        for krate in krates {
            let module = krate.root_module(db);
            let file_id = module.definition_source(db).file_id;
            let file_id = file_id.original_file(db);
            let source_root = db.file_source_root(file_id);
            let source_root = db.source_root(source_root);
            if !source_root.is_library || self.with_deps {
                num_crates += 1;
                visit_queue.push(module);
            }
        }

        if self.randomize {
            shuffle(&mut rng, &mut visit_queue);
        }

        eprintln!("Crates in this dir: {}", num_crates);
        let mut num_decls = 0;
        let mut funcs = Vec::new();
        while let Some(module) = visit_queue.pop() {
            if visited_modules.insert(module) {
                visit_queue.extend(module.children(db));

                for decl in module.declarations(db) {
                    num_decls += 1;
                    if let ModuleDef::Function(f) = decl {
                        funcs.push(f);
                    }
                }

                for impl_def in module.impl_defs(db) {
                    for item in impl_def.items(db) {
                        num_decls += 1;
                        if let AssocItem::Function(f) = item {
                            funcs.push(f);
                        }
                    }
                }
            }
        }
        eprintln!("Total modules found: {}", visited_modules.len());
        eprintln!("Total declarations: {}", num_decls);
        eprintln!("Total functions: {}", funcs.len());
        eprintln!("Item Collection: {}", analysis_sw.elapsed());

        if self.randomize {
            shuffle(&mut rng, &mut funcs);
        }

        let mut bar = match verbosity {
            Verbosity::Quiet | Verbosity::Spammy => ProgressReport::hidden(),
            _ if self.parallel => ProgressReport::hidden(),
            _ => ProgressReport::new(funcs.len() as u64),
        };

        if self.parallel {
            let mut inference_sw = self.stop_watch();
            let snap = Snap(db.snapshot());
            funcs
                .par_iter()
                .map_with(snap, |snap, &f| {
                    let f_id = FunctionId::from(f);
                    snap.0.body(f_id.into());
                    snap.0.infer(f_id.into());
                })
                .count();
            eprintln!("Parallel Inference: {}", inference_sw.elapsed());
        }

        let mut inference_sw = self.stop_watch();
        bar.tick();
        let mut num_exprs = 0;
        let mut num_exprs_unknown = 0;
        let mut num_exprs_partially_unknown = 0;
        let mut num_type_mismatches = 0;
        for f in funcs {
            let name = f.name(db);
            let full_name = f
                .module(db)
                .path_to_root(db)
                .into_iter()
                .rev()
                .filter_map(|it| it.name(db))
                .chain(Some(f.name(db)))
                .join("::");
            if let Some(only_name) = self.only.as_deref() {
                if name.to_string() != only_name && full_name != only_name {
                    continue;
                }
            }
            let mut msg = format!("processing: {}", full_name);
            if verbosity.is_verbose() {
                if let Some(src) = f.source(db) {
                    let original_file = src.file_id.original_file(db);
                    let path = vfs.file_path(original_file);
                    let syntax_range = src.value.syntax().text_range();
                    format_to!(msg, " ({} {:?})", path, syntax_range);
                }
            }
            if verbosity.is_spammy() {
                bar.println(msg.to_string());
            }
            bar.set_message(&msg);
            let f_id = FunctionId::from(f);
            let body = db.body(f_id.into());
            let inference_result = db.infer(f_id.into());
            let (previous_exprs, previous_unknown, previous_partially_unknown) =
                (num_exprs, num_exprs_unknown, num_exprs_partially_unknown);
            for (expr_id, _) in body.exprs.iter() {
                let ty = &inference_result[expr_id];
                num_exprs += 1;
                if let Ty::Unknown = ty {
                    num_exprs_unknown += 1;
                } else {
                    let mut is_partially_unknown = false;
                    ty.walk(&mut |ty| {
                        if let Ty::Unknown = ty {
                            is_partially_unknown = true;
                        }
                    });
                    if is_partially_unknown {
                        num_exprs_partially_unknown += 1;
                    }
                }
                if self.only.is_some() && verbosity.is_spammy() {
                    // in super-verbose mode for just one function, we print every single expression
                    let (_, sm) = db.body_with_source_map(f_id.into());
                    let src = sm.expr_syntax(expr_id);
                    if let Ok(src) = src {
                        let node = {
                            let root = db.parse_or_expand(src.file_id).unwrap();
                            src.value.to_node(&root)
                        };
                        let original_file = src.file_id.original_file(db);
                        let line_index = host.analysis().file_line_index(original_file).unwrap();
                        let text_range = node.syntax().text_range();
                        let (start, end) = (
                            line_index.line_col(text_range.start()),
                            line_index.line_col(text_range.end()),
                        );
                        bar.println(format!(
                            "{}:{}-{}:{}: {}",
                            start.line + 1,
                            start.col_utf16,
                            end.line + 1,
                            end.col_utf16,
                            ty.display(db)
                        ));
                    } else {
                        bar.println(format!("unknown location: {}", ty.display(db)));
                    }
                }
                if let Some(mismatch) = inference_result.type_mismatch_for_expr(expr_id) {
                    num_type_mismatches += 1;
                    if verbosity.is_verbose() {
                        let (_, sm) = db.body_with_source_map(f_id.into());
                        let src = sm.expr_syntax(expr_id);
                        if let Ok(src) = src {
                            // FIXME: it might be nice to have a function (on Analysis?) that goes from Source<T> -> (LineCol, LineCol) directly
                            // But also, we should just turn the type mismatches into diagnostics and provide these
                            let root = db.parse_or_expand(src.file_id).unwrap();
                            let node = src.map(|e| e.to_node(&root).syntax().clone());
                            let original_range = node.as_ref().original_file_range(db);
                            let path = vfs.file_path(original_range.file_id);
                            let line_index =
                                host.analysis().file_line_index(original_range.file_id).unwrap();
                            let text_range = original_range.range;
                            let (start, end) = (
                                line_index.line_col(text_range.start()),
                                line_index.line_col(text_range.end()),
                            );
                            bar.println(format!(
                                "{} {}:{}-{}:{}: Expected {}, got {}",
                                path,
                                start.line + 1,
                                start.col_utf16,
                                end.line + 1,
                                end.col_utf16,
                                mismatch.expected.display(db),
                                mismatch.actual.display(db)
                            ));
                        } else {
                            bar.println(format!(
                                "{}: Expected {}, got {}",
                                name,
                                mismatch.expected.display(db),
                                mismatch.actual.display(db)
                            ));
                        }
                    }
                }
            }
            if verbosity.is_spammy() {
                bar.println(format!(
                    "In {}: {} exprs, {} unknown, {} partial",
                    full_name,
                    num_exprs - previous_exprs,
                    num_exprs_unknown - previous_unknown,
                    num_exprs_partially_unknown - previous_partially_unknown
                ));
            }
            bar.inc(1);
        }
        bar.finish_and_clear();
        eprintln!("Total expressions: {}", num_exprs);
        eprintln!(
            "Expressions of unknown type: {} ({}%)",
            num_exprs_unknown,
            if num_exprs > 0 { num_exprs_unknown * 100 / num_exprs } else { 100 }
        );
        report_metric("unknown type", num_exprs_unknown, "#");

        eprintln!(
            "Expressions of partially unknown type: {} ({}%)",
            num_exprs_partially_unknown,
            if num_exprs > 0 { num_exprs_partially_unknown * 100 / num_exprs } else { 100 }
        );

        eprintln!("Type mismatches: {}", num_type_mismatches);
        report_metric("type mismatches", num_type_mismatches, "#");

        eprintln!("Inference: {}", inference_sw.elapsed());

        let total_span = analysis_sw.elapsed();
        eprintln!("Total: {}", total_span);
        report_metric("total time", total_span.time.as_millis() as u64, "ms");
        if let Some(instructions) = total_span.instructions {
            report_metric("total instructions", instructions, "#instr");
        }
        if let Some(memory) = total_span.memory {
            report_metric("total memory", memory.allocated.megabytes() as u64, "MB");
        }

        if self.memory_usage {
            print_memory_usage(host, vfs);
        }

        Ok(())
    }

    fn stop_watch(&self) -> StopWatch {
        StopWatch::start().memory(self.memory_usage)
    }
}

fn shuffle<T>(rng: &mut Rand32, slice: &mut [T]) {
    for i in 0..slice.len() {
        randomize_first(rng, &mut slice[i..]);
    }

    fn randomize_first<T>(rng: &mut Rand32, slice: &mut [T]) {
        assert!(!slice.is_empty());
        let idx = rng.rand_range(0..slice.len() as u32) as usize;
        slice.swap(0, idx);
    }
}
<commit_msg>Make default memory stats less verbose<commit_after>//! Fully type-check project and print various stats, like the number of type
//! errors.

use std::{
    path::PathBuf,
    time::{SystemTime, UNIX_EPOCH},
};

use hir::{
    db::{AstDatabase, DefDatabase, HirDatabase},
    AssocItem, Crate, HasSource, HirDisplay, ModuleDef,
};
use hir_def::FunctionId;
use hir_ty::{Ty, TypeWalk};
use ide_db::base_db::{
    salsa::{self, ParallelDatabase},
    SourceDatabaseExt,
};
use itertools::Itertools;
use oorandom::Rand32;
use rayon::prelude::*;
use rustc_hash::FxHashSet;
use stdx::format_to;
use syntax::AstNode;

use crate::cli::{
    load_cargo::load_cargo, print_memory_usage, progress_report::ProgressReport, report_metric,
    Result, Verbosity,
};
use profile::StopWatch;

/// Need to wrap Snapshot to provide `Clone` impl for `map_with`
struct Snap<DB>(DB);
impl<DB: ParallelDatabase> Clone for Snap<salsa::Snapshot<DB>> {
    fn clone(&self) -> Snap<salsa::Snapshot<DB>> {
        Snap(self.0.snapshot())
    }
}

pub struct AnalysisStatsCmd {
    pub randomize: bool,
    pub parallel: bool,
    pub memory_usage: bool,
    pub only: Option<String>,
    pub with_deps: bool,
    pub path: PathBuf,
    pub load_output_dirs: bool,
    pub with_proc_macro: bool,
}

impl AnalysisStatsCmd {
    pub fn run(self, verbosity: Verbosity) -> Result<()> {
        let mut rng = {
            let seed = SystemTime::now().duration_since(UNIX_EPOCH).unwrap().as_millis() as u64;
            Rand32::new(seed)
        };

        let mut db_load_sw = self.stop_watch();
        let (host, vfs) = load_cargo(&self.path, self.load_output_dirs, self.with_proc_macro)?;
        let db = host.raw_database();
        eprintln!("Database loaded {}", db_load_sw.elapsed());

        let mut analysis_sw = self.stop_watch();
        let mut num_crates = 0;
        let mut visited_modules = FxHashSet::default();
        let mut visit_queue = Vec::new();

        let mut krates = Crate::all(db);
        if self.randomize {
            shuffle(&mut rng, &mut krates);
        }
        for krate in krates {
            let module = krate.root_module(db);
            let file_id = module.definition_source(db).file_id;
            let file_id = file_id.original_file(db);
            let source_root = db.file_source_root(file_id);
            let source_root = db.source_root(source_root);
            if !source_root.is_library || self.with_deps {
                num_crates += 1;
                visit_queue.push(module);
            }
        }

        if self.randomize {
            shuffle(&mut rng, &mut visit_queue);
        }

        eprintln!("Crates in this dir: {}", num_crates);
        let mut num_decls = 0;
        let mut funcs = Vec::new();
        while let Some(module) = visit_queue.pop() {
            if visited_modules.insert(module) {
                visit_queue.extend(module.children(db));

                for decl in module.declarations(db) {
                    num_decls += 1;
                    if let ModuleDef::Function(f) = decl {
                        funcs.push(f);
                    }
                }

                for impl_def in module.impl_defs(db) {
                    for item in impl_def.items(db) {
                        num_decls += 1;
                        if let AssocItem::Function(f) = item {
                            funcs.push(f);
                        }
                    }
                }
            }
        }
        eprintln!("Total modules found: {}", visited_modules.len());
        eprintln!("Total declarations: {}", num_decls);
        eprintln!("Total functions: {}", funcs.len());
        eprintln!("Item Collection: {}", analysis_sw.elapsed());

        if self.randomize {
            shuffle(&mut rng, &mut funcs);
        }

        let mut bar = match verbosity {
            Verbosity::Quiet | Verbosity::Spammy => ProgressReport::hidden(),
            _ if self.parallel => ProgressReport::hidden(),
            _ => ProgressReport::new(funcs.len() as u64),
        };

        if self.parallel {
            let mut inference_sw = self.stop_watch();
            let snap = Snap(db.snapshot());
            funcs
                .par_iter()
                .map_with(snap, |snap, &f| {
                    let f_id = FunctionId::from(f);
                    snap.0.body(f_id.into());
                    snap.0.infer(f_id.into());
                })
                .count();
            eprintln!("Parallel Inference: {}", inference_sw.elapsed());
        }

        let mut inference_sw = self.stop_watch();
        bar.tick();
        let mut num_exprs = 0;
        let mut num_exprs_unknown = 0;
        let mut num_exprs_partially_unknown = 0;
        let mut num_type_mismatches = 0;
        for f in funcs {
            let name = f.name(db);
            let full_name = f
                .module(db)
                .path_to_root(db)
                .into_iter()
                .rev()
                .filter_map(|it| it.name(db))
                .chain(Some(f.name(db)))
                .join("::");
            if let Some(only_name) = self.only.as_deref() {
                if name.to_string() != only_name && full_name != only_name {
                    continue;
                }
            }
            let mut msg = format!("processing: {}", full_name);
            if verbosity.is_verbose() {
                if let Some(src) = f.source(db) {
                    let original_file = src.file_id.original_file(db);
                    let path = vfs.file_path(original_file);
                    let syntax_range = src.value.syntax().text_range();
                    format_to!(msg, " ({} {:?})", path, syntax_range);
                }
            }
            if verbosity.is_spammy() {
                bar.println(msg.to_string());
            }
            bar.set_message(&msg);
            let f_id = FunctionId::from(f);
            let body = db.body(f_id.into());
            let inference_result = db.infer(f_id.into());
            let (previous_exprs, previous_unknown, previous_partially_unknown) =
                (num_exprs, num_exprs_unknown, num_exprs_partially_unknown);
            for (expr_id, _) in body.exprs.iter() {
                let ty = &inference_result[expr_id];
                num_exprs += 1;
                if let Ty::Unknown = ty {
                    num_exprs_unknown += 1;
                } else {
                    let mut is_partially_unknown = false;
                    ty.walk(&mut |ty| {
                        if let Ty::Unknown = ty {
                            is_partially_unknown = true;
                        }
                    });
                    if is_partially_unknown {
                        num_exprs_partially_unknown += 1;
                    }
                }
                if self.only.is_some() && verbosity.is_spammy() {
                    // in super-verbose mode for just one function, we print every single expression
                    let (_, sm) = db.body_with_source_map(f_id.into());
                    let src = sm.expr_syntax(expr_id);
                    if let Ok(src) = src {
                        let node = {
                            let root = db.parse_or_expand(src.file_id).unwrap();
                            src.value.to_node(&root)
                        };
                        let original_file = src.file_id.original_file(db);
                        let line_index = host.analysis().file_line_index(original_file).unwrap();
                        let text_range = node.syntax().text_range();
                        let (start, end) = (
                            line_index.line_col(text_range.start()),
                            line_index.line_col(text_range.end()),
                        );
                        bar.println(format!(
                            "{}:{}-{}:{}: {}",
                            start.line + 1,
                            start.col_utf16,
                            end.line + 1,
                            end.col_utf16,
                            ty.display(db)
                        ));
                    } else {
                        bar.println(format!("unknown location: {}", ty.display(db)));
                    }
                }
                if let Some(mismatch) = inference_result.type_mismatch_for_expr(expr_id) {
                    num_type_mismatches += 1;
                    if verbosity.is_verbose() {
                        let (_, sm) = db.body_with_source_map(f_id.into());
                        let src = sm.expr_syntax(expr_id);
                        if let Ok(src) = src {
                            // FIXME: it might be nice to have a function (on Analysis?) that goes from Source<T> -> (LineCol, LineCol) directly
                            // But also, we should just turn the type mismatches into diagnostics and provide these
                            let root = db.parse_or_expand(src.file_id).unwrap();
                            let node = src.map(|e| e.to_node(&root).syntax().clone());
                            let original_range = node.as_ref().original_file_range(db);
                            let path = vfs.file_path(original_range.file_id);
                            let line_index =
                                host.analysis().file_line_index(original_range.file_id).unwrap();
                            let text_range = original_range.range;
                            let (start, end) = (
                                line_index.line_col(text_range.start()),
                                line_index.line_col(text_range.end()),
                            );
                            bar.println(format!(
                                "{} {}:{}-{}:{}: Expected {}, got {}",
                                path,
                                start.line + 1,
                                start.col_utf16,
                                end.line + 1,
                                end.col_utf16,
                                mismatch.expected.display(db),
                                mismatch.actual.display(db)
                            ));
                        } else {
                            bar.println(format!(
                                "{}: Expected {}, got {}",
                                name,
                                mismatch.expected.display(db),
                                mismatch.actual.display(db)
                            ));
                        }
                    }
                }
            }
            if verbosity.is_spammy() {
                bar.println(format!(
                    "In {}: {} exprs, {} unknown, {} partial",
                    full_name,
                    num_exprs - previous_exprs,
                    num_exprs_unknown - previous_unknown,
                    num_exprs_partially_unknown - previous_partially_unknown
                ));
            }
            bar.inc(1);
        }
        bar.finish_and_clear();
        eprintln!("Total expressions: {}", num_exprs);
        eprintln!(
            "Expressions of unknown type: {} ({}%)",
            num_exprs_unknown,
            if num_exprs > 0 { num_exprs_unknown * 100 / num_exprs } else { 100 }
        );
        report_metric("unknown type", num_exprs_unknown, "#");

        eprintln!(
            "Expressions of partially unknown type: {} ({}%)",
            num_exprs_partially_unknown,
            if num_exprs > 0 { num_exprs_partially_unknown * 100 / num_exprs } else { 100 }
        );

        eprintln!("Type mismatches: {}", num_type_mismatches);
        report_metric("type mismatches", num_type_mismatches, "#");

        eprintln!("Inference: {}", inference_sw.elapsed());

        let total_span = analysis_sw.elapsed();
        eprintln!("Total: {}", total_span);
        report_metric("total time", total_span.time.as_millis() as u64, "ms");
        if let Some(instructions) = total_span.instructions {
            report_metric("total instructions", instructions, "#instr");
        }
        if let Some(memory) = total_span.memory {
            report_metric("total memory", memory.allocated.megabytes() as u64, "MB");
        }

        if self.memory_usage && verbosity.is_verbose() {
            print_memory_usage(host, vfs);
        }

        Ok(())
    }

    fn stop_watch(&self) -> StopWatch {
        StopWatch::start().memory(self.memory_usage)
    }
}

fn shuffle<T>(rng: &mut Rand32, slice: &mut [T]) {
    for i in 0..slice.len() {
        randomize_first(rng, &mut slice[i..]);
    }

    fn randomize_first<T>(rng: &mut Rand32, slice: &mut [T]) {
        assert!(!slice.is_empty());
        let idx = rng.rand_range(0..slice.len() as u32) as usize;
        slice.swap(0, idx);
    }
}
<|endoftext|>


<commit_before>use crate::rc::Rc;
use crate::sync::Arc;
use core::panic::{RefUnwindSafe, UnwindSafe};

// not covered via the Shared impl above b/c the inner contents use
// Cell/AtomicUsize, but the usage here is unwind safe so we can lift the
// impl up one level to Arc/Rc itself
#[stable(feature = "catch_unwind", since = "1.9.0")]
impl<T: RefUnwindSafe + ?Sized> UnwindSafe for Rc<T> {}
#[stable(feature = "catch_unwind", since = "1.9.0")]
impl<T: RefUnwindSafe + ?Sized> UnwindSafe for Arc<T> {}
<commit_msg>Fix comment referring to formerly-above code<commit_after>use crate::rc::Rc;
use crate::sync::Arc;
use core::panic::{RefUnwindSafe, UnwindSafe};

// not covered via the Shared impl in libcore b/c the inner contents use
// Cell/AtomicUsize, but the usage here is unwind safe so we can lift the
// impl up one level to Arc/Rc itself
#[stable(feature = "catch_unwind", since = "1.9.0")]
impl<T: RefUnwindSafe + ?Sized> UnwindSafe for Rc<T> {}
#[stable(feature = "catch_unwind", since = "1.9.0")]
impl<T: RefUnwindSafe + ?Sized> UnwindSafe for Arc<T> {}
<|endoftext|>


<commit_before>#![feature(box_syntax)]
#![feature(asm)]
extern crate e2d2;
extern crate fnv;
extern crate time;
extern crate simd;
extern crate getopts;
extern crate rand;
use e2d2::io::*;
use e2d2::headers::*;
use e2d2::packet_batch::*;
use getopts::Options;
use std::collections::HashMap;
use std::env;
use std::time::Duration;
use std::thread;
use std::sync::Arc;

const CONVERSION_FACTOR: f64 = 1000000000.;

#[allow(dead_code)]
#[inline]
fn rdtscp_unsafe() -> u64 {
    unsafe {
        let low: u32;
        let high: u32;
        asm!("rdtscp"
             : "={eax}" (low), "={edx}" (high)
             :
             : "{rcx} {rdx} {rax}"
             : "volatile");
        ((high as u64) << 32) | (low as u64)
    }
}

#[allow(dead_code)]
#[inline]
fn rdtscp() -> u64 {
    rdtscp_unsafe()
}

#[inline]
fn lat() {
    unsafe {
        asm!("nop"
             :
             :
             :
             : "volatile");
    }
}

#[inline]
fn delay_loop(delay: u64) {
    let mut d = 0;
    while d < delay {
        lat();
        d += 1;
    }
}

fn delay<T: 'static + Batch>(parent: T, delay: u64) -> CompositionBatch {
    parent.parse::<MacHeader>()
          .transform(box move |hdr, _, _| {
              let src = hdr.src.clone();
              hdr.src = hdr.dst;
              hdr.dst = src;
              delay_loop(delay);
          })
          .compose()
}

fn recv_thread(ports: Vec<PortQueue>, core: i32, delay_arg: u64) {
    init_thread(core, core);
    println!("Receiving started");
    for port in &ports {
        println!("Receiving port {} rxq {} txq {} on core {}", port.port.mac_address(), port.rxq(), port.txq(), core);
    }

    let pipelines: Vec<_> = ports.iter()
                                 .map(|port| {
                                     delay(ReceiveBatch::new(port.clone()), delay_arg)
                                         .send(port.clone())
                                         .compose()
                                 })
                                 .collect();
    println!("Running {} pipelines", pipelines.len());
    let mut combined = merge(pipelines);
    loop {
        combined.process();
    }
}

fn main() {
    let args: Vec<String> = env::args().collect();
    let program = args[0].clone();
    let mut opts = Options::new();
    opts.optflag("h", "help", "print this help menu");
    opts.optflag("", "secondary", "run as a secondary process");
    opts.optopt("n", "name", "name to use for the current process", "name");
    opts.optmulti("w", "whitelist", "Whitelist PCI", "PCI");
    opts.optmulti("v", "vdevs", "Virtual Devices to add", "PCI");
    opts.optmulti("c", "core", "Core to use", "core");
    opts.optopt("m", "master", "Master core", "master");
    opts.optopt("d", "delay", "Delay cycles", "cycles");
    let matches = match opts.parse(&args[1..]) {
        Ok(m) => m,
        Err(f) => panic!(f.to_string()),
    };
    if matches.opt_present("h") {
        print!("{}", opts.usage(&format!("Usage: {} [options]", program)));
    }
    let cores_str = matches.opt_strs("c");
    let master_core = matches.opt_str("m")
                             .unwrap_or_else(|| String::from("0"))
                             .parse()
                             .expect("Could not parse master core spec");
    let delay_arg = matches.opt_str("d")
                           .unwrap_or_else(|| String::from("100"))
                           .parse()
                           .expect("Could not parse delay");
    println!("Using master core {}", master_core);
    let name = matches.opt_str("n").unwrap_or_else(|| String::from("recv"));

    let cores: Vec<i32> = cores_str.iter()
                                   .map(|n: &String| n.parse().ok().expect(&format!("Core cannot be parsed {}", n)))
                                   .collect();

    fn extract_cores_for_port(ports: &[String], cores: &[i32]) -> HashMap<String, Vec<i32>> {
        let mut cores_for_port = HashMap::<String, Vec<i32>>::new();
        for (port, core) in ports.iter().zip(cores.iter()) {
            cores_for_port.entry(port.clone()).or_insert(vec![]).push(*core)
        }
        cores_for_port
    }

    let primary = !matches.opt_present("secondary");

    let cores_for_port = if primary {
        extract_cores_for_port(&matches.opt_strs("w"), &cores)
    } else {
        extract_cores_for_port(&matches.opt_strs("v"), &cores)
    };

    if primary {
        let whitelisted: Vec<_> = cores_for_port.keys().map(|p| p.clone()).collect();
        init_system_wl(&name, master_core, &whitelisted[..]);
    } else {
        init_system_secondary(&name, master_core, &[]);
    }

    let ports_to_activate: Vec<_> = cores_for_port.keys().collect();

    let mut queues_by_core = HashMap::<i32, Vec<_>>::with_capacity(cores.len());
    let mut ports = Vec::<Arc<PmdPort>>::with_capacity(ports_to_activate.len());
    for port in &ports_to_activate {
        let cores = cores_for_port.get(*port).unwrap();
        let pmd_port = if primary {
            let port_id = PmdPort::find_port_id(*port);
            if port_id < 0 {
                println!("Could not find port {}", port);
                panic!("Could not find port")
            } else {
                let queues = cores.len() as i32;
                PmdPort::new_mq_port(port_id, queues, queues, cores, cores).expect("Could not initialize port")
            }
        } else {
            assert!(cores.len() == 1);
            PmdPort::new_vdev(*port, cores[0]).expect("Could not initialize vdev")
        };
        for (idx, core) in cores.iter().enumerate() {
            let queue = idx as i32;
            queues_by_core.entry(*core)
                          .or_insert(vec![])
                          .push(PmdPort::new_queue_pair(&pmd_port, queue, queue).unwrap());
        }
        ports.push(pmd_port);
    }

    const _BATCH: usize = 1 << 10;
    const _CHANNEL_SIZE: usize = 256;
    let _thread: Vec<_> = queues_by_core.iter()
                                        .map(|(core, ports)| {
                                            let c = core.clone();
                                            let p: Vec<_> = ports.iter().map(|p| p.clone()).collect();
                                            std::thread::spawn(move || recv_thread(p, c, delay_arg))
                                        })
                                        .collect();
    let mut pkts_so_far = (0, 0);
    let mut last_printed = 0.;
    const MAX_PRINT_INTERVAL: f64 = 15.;
    let mut start = time::precise_time_ns() as f64 / CONVERSION_FACTOR;
    let sleep_time = Duration::from_millis(5000);
    println!("0 OVERALL RX 0.00 TX 0.00 CYCLE_PER_DELAY 0 0 0");
    loop {
        thread::sleep(sleep_time); // Sleep for a bit
        let now = time::precise_time_ns() as f64 / CONVERSION_FACTOR;
        if now - start > 5.0 {
            let mut rx = 0;
            let mut tx = 0;
            for port in &ports {
                let mac = port.mac_address();
                for q in 0..port.rxqs() {
                    let (rp, tp) = port.stats(q);
                    rx += rp;
                    tx += tp;
                    println!("{:.2} PORT {} Q {} RX {} TX {}",
                                   now, mac, q, tx, tx);
                }
            }
            let pkts = (rx, tx);
            let rx_pkts = pkts.0 - pkts_so_far.0;
            if rx_pkts > 0 || now - last_printed > MAX_PRINT_INTERVAL {
                println!("{:.2} OVERALL RX {:.2} TX {:.2}",
                         now - start,
                         rx_pkts as f64 / (now - start),
                         (pkts.1 - pkts_so_far.1) as f64 / (now - start));
                last_printed = now;
                start = now;
                pkts_so_far = pkts;
            }
        }
    }
}
<commit_msg>Eliminate some printouts<commit_after>#![feature(box_syntax)]
#![feature(asm)]
extern crate e2d2;
extern crate fnv;
extern crate time;
extern crate simd;
extern crate getopts;
extern crate rand;
use e2d2::io::*;
use e2d2::headers::*;
use e2d2::packet_batch::*;
use getopts::Options;
use std::collections::HashMap;
use std::env;
use std::time::Duration;
use std::thread;
use std::sync::Arc;

const CONVERSION_FACTOR: f64 = 1000000000.;

#[allow(dead_code)]
#[inline]
fn rdtscp_unsafe() -> u64 {
    unsafe {
        let low: u32;
        let high: u32;
        asm!("rdtscp"
             : "={eax}" (low), "={edx}" (high)
             :
             : "{rcx} {rdx} {rax}"
             : "volatile");
        ((high as u64) << 32) | (low as u64)
    }
}

#[allow(dead_code)]
#[inline]
fn rdtscp() -> u64 {
    rdtscp_unsafe()
}

#[inline]
fn lat() {
    unsafe {
        asm!("nop"
             :
             :
             :
             : "volatile");
    }
}

#[inline]
fn delay_loop(delay: u64) {
    let mut d = 0;
    while d < delay {
        lat();
        d += 1;
    }
}

fn delay<T: 'static + Batch>(parent: T, delay: u64) -> CompositionBatch {
    parent.parse::<MacHeader>()
          .transform(box move |hdr, _, _| {
              let src = hdr.src.clone();
              hdr.src = hdr.dst;
              hdr.dst = src;
              delay_loop(delay);
          })
          .compose()
}

fn recv_thread(ports: Vec<PortQueue>, core: i32, delay_arg: u64) {
    init_thread(core, core);
    println!("Receiving started");
    for port in &ports {
        println!("Receiving port {} rxq {} txq {} on core {}", port.port.mac_address(), port.rxq(), port.txq(), core);
    }

    let pipelines: Vec<_> = ports.iter()
                                 .map(|port| {
                                     delay(ReceiveBatch::new(port.clone()), delay_arg)
                                         .send(port.clone())
                                         .compose()
                                 })
                                 .collect();
    println!("Running {} pipelines", pipelines.len());
    let mut combined = merge(pipelines);
    loop {
        combined.process();
    }
}

fn main() {
    let args: Vec<String> = env::args().collect();
    let program = args[0].clone();
    let mut opts = Options::new();
    opts.optflag("h", "help", "print this help menu");
    opts.optflag("", "secondary", "run as a secondary process");
    opts.optopt("n", "name", "name to use for the current process", "name");
    opts.optmulti("w", "whitelist", "Whitelist PCI", "PCI");
    opts.optmulti("v", "vdevs", "Virtual Devices to add", "PCI");
    opts.optmulti("c", "core", "Core to use", "core");
    opts.optopt("m", "master", "Master core", "master");
    opts.optopt("d", "delay", "Delay cycles", "cycles");
    let matches = match opts.parse(&args[1..]) {
        Ok(m) => m,
        Err(f) => panic!(f.to_string()),
    };
    if matches.opt_present("h") {
        print!("{}", opts.usage(&format!("Usage: {} [options]", program)));
    }
    let cores_str = matches.opt_strs("c");
    let master_core = matches.opt_str("m")
                             .unwrap_or_else(|| String::from("0"))
                             .parse()
                             .expect("Could not parse master core spec");
    let delay_arg = matches.opt_str("d")
                           .unwrap_or_else(|| String::from("100"))
                           .parse()
                           .expect("Could not parse delay");
    println!("Using master core {}", master_core);
    let name = matches.opt_str("n").unwrap_or_else(|| String::from("recv"));

    let cores: Vec<i32> = cores_str.iter()
                                   .map(|n: &String| n.parse().ok().expect(&format!("Core cannot be parsed {}", n)))
                                   .collect();

    fn extract_cores_for_port(ports: &[String], cores: &[i32]) -> HashMap<String, Vec<i32>> {
        let mut cores_for_port = HashMap::<String, Vec<i32>>::new();
        for (port, core) in ports.iter().zip(cores.iter()) {
            cores_for_port.entry(port.clone()).or_insert(vec![]).push(*core)
        }
        cores_for_port
    }

    let primary = !matches.opt_present("secondary");

    let cores_for_port = if primary {
        extract_cores_for_port(&matches.opt_strs("w"), &cores)
    } else {
        extract_cores_for_port(&matches.opt_strs("v"), &cores)
    };

    if primary {
        let whitelisted: Vec<_> = cores_for_port.keys().map(|p| p.clone()).collect();
        init_system_wl(&name, master_core, &whitelisted[..]);
    } else {
        init_system_secondary(&name, master_core, &[]);
    }

    let ports_to_activate: Vec<_> = cores_for_port.keys().collect();

    let mut queues_by_core = HashMap::<i32, Vec<_>>::with_capacity(cores.len());
    let mut ports = Vec::<Arc<PmdPort>>::with_capacity(ports_to_activate.len());
    for port in &ports_to_activate {
        let cores = cores_for_port.get(*port).unwrap();
        let pmd_port = if primary {
            let port_id = PmdPort::find_port_id(*port);
            if port_id < 0 {
                println!("Could not find port {}", port);
                panic!("Could not find port")
            } else {
                let queues = cores.len() as i32;
                PmdPort::new_mq_port(port_id, queues, queues, cores, cores).expect("Could not initialize port")
            }
        } else {
            assert!(cores.len() == 1);
            PmdPort::new_vdev(*port, cores[0]).expect("Could not initialize vdev")
        };
        for (idx, core) in cores.iter().enumerate() {
            let queue = idx as i32;
            queues_by_core.entry(*core)
                          .or_insert(vec![])
                          .push(PmdPort::new_queue_pair(&pmd_port, queue, queue).unwrap());
        }
        ports.push(pmd_port);
    }

    const _BATCH: usize = 1 << 10;
    const _CHANNEL_SIZE: usize = 256;
    let _thread: Vec<_> = queues_by_core.iter()
                                        .map(|(core, ports)| {
                                            let c = core.clone();
                                            let p: Vec<_> = ports.iter().map(|p| p.clone()).collect();
                                            std::thread::spawn(move || recv_thread(p, c, delay_arg))
                                        })
                                        .collect();
    let mut pkts_so_far = (0, 0);
    let mut last_printed = 0.;
    const MAX_PRINT_INTERVAL: f64 = 15.;
    const PRINT_DELAY: f64 = 10.;
    let sleep_delay = (PRINT_DELAY / 2.) as u64;
    let mut start = time::precise_time_ns() as f64 / CONVERSION_FACTOR;
    let sleep_time = Duration::from_millis(sleep_delay);
    println!("0 OVERALL RX 0.00 TX 0.00 CYCLE_PER_DELAY 0 0 0");
    loop {
        thread::sleep(sleep_time); // Sleep for a bit
        let now = time::precise_time_ns() as f64 / CONVERSION_FACTOR;
        if now - start > PRINT_DELAY {
            let mut rx = 0;
            let mut tx = 0;
            for port in &ports {
                for q in 0..port.rxqs() {
                    let (rp, tp) = port.stats(q);
                    rx += rp;
                    tx += tp;
                }
            }
            let pkts = (rx, tx);
            let rx_pkts = pkts.0 - pkts_so_far.0;
            if rx_pkts > 0 || now - last_printed > MAX_PRINT_INTERVAL {
                println!("{:.2} OVERALL RX {:.2} TX {:.2}",
                         now - start,
                         rx_pkts as f64 / (now - start),
                         (pkts.1 - pkts_so_far.1) as f64 / (now - start));
                last_printed = now;
                start = now;
                pkts_so_far = pkts;
            }
        }
    }
}
<|endoftext|>


<commit_before>/*
    Copyright 2020 Andrew Medworth <github@medworth.org.uk>

    This file is part of Dots-and-Boxes Engine.

    Dots-and-Boxes Engine is free software: you can redistribute it and/or modify
    it under the terms of the GNU Affero General Public License as published by
    the Free Software Foundation, either version 3 of the License, or
    (at your option) any later version.

    Dots-and-Boxes Engine is distributed in the hope that it will be useful,
    but WITHOUT ANY WARRANTY; without even the implied warranty of
    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
    GNU Affero General Public License for more details.

    You should have received a copy of the GNU Affero General Public License
    along with Dots-and-Boxes Engine.  If not, see <http://www.gnu.org/licenses/>.
*/
use criterion::{criterion_group, criterion_main, Criterion};
use dabengine::{game, eval};

pub fn eval_benchmark(c: &mut Criterion) {
    c.bench_function("2x2", |b| b.iter(|| eval::eval(&game::SimplePosition::new_game(2, 2))));
}

criterion_group!(eval_benches, eval_benchmark);
criterion_main!(eval_benches);
<commit_msg>Add more benchmarks<commit_after>/*
    Copyright 2020 Andrew Medworth <github@medworth.org.uk>

    This file is part of Dots-and-Boxes Engine.

    Dots-and-Boxes Engine is free software: you can redistribute it and/or modify
    it under the terms of the GNU Affero General Public License as published by
    the Free Software Foundation, either version 3 of the License, or
    (at your option) any later version.

    Dots-and-Boxes Engine is distributed in the hope that it will be useful,
    but WITHOUT ANY WARRANTY; without even the implied warranty of
    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
    GNU Affero General Public License for more details.

    You should have received a copy of the GNU Affero General Public License
    along with Dots-and-Boxes Engine.  If not, see <http://www.gnu.org/licenses/>.
*/
use criterion::{criterion_group, criterion_main, Criterion};
use dabengine::{game, eval, examples};
use std::time::Duration;

pub fn eval_benchmark(c: &mut Criterion) {
    let mut fast_group = c.benchmark_group("fast-evaluations");
    fast_group.measurement_time(Duration::from_secs(30));
    fast_group.bench_function("Eval simple 2x2", |b| b.iter(|| eval::eval(&game::SimplePosition::new_game(2, 2))));
    fast_group.bench_function("Eval composite OLMT", |b| b.iter(|| eval::eval(&examples::one_long_multi_three(5, 4))));
    fast_group.finish();

    let mut slow_group = c.benchmark_group("slow-evaluations");
    slow_group.sample_size(10);
    slow_group.measurement_time(Duration::from_secs(240));
    slow_group.bench_function("Eval simple 3x3", |b| b.iter(|| eval::eval(&game::SimplePosition::new_game(3, 3))));
    slow_group.finish();
}

criterion_group!(eval_benches, eval_benchmark);
criterion_main!(eval_benches);
<|endoftext|>