<commit_before>// Copyright (c) Aptos
// SPDX-License-Identifier: Apache-2.0

use crate::error_map::generate_error_map;
use crate::natives::code::{
    ModuleMetadata, MoveOption, PackageDep, PackageMetadata, UpgradePolicy,
};
use crate::{zip_metadata, zip_metadata_str, RuntimeModuleMetadata, APTOS_METADATA_KEY};
use aptos_module_verifier::module_init::verify_module_init_function;
use aptos_types::account_address::AccountAddress;
use aptos_types::transaction::EntryABI;
use clap::Parser;
use move_binary_format::CompiledModule;
use move_command_line_common::files::MOVE_COMPILED_EXTENSION;
use move_compiler::compiled_unit::{CompiledUnit, NamedCompiledModule};
use move_core_types::errmap::ErrorMapping;
use move_core_types::metadata::Metadata;
use move_package::compilation::compiled_package::CompiledPackage;
use move_package::compilation::package_layout::CompiledPackageLayout;
use move_package::source_package::manifest_parser::{
    parse_move_manifest_string, parse_source_manifest,
};
use move_package::BuildConfig;
use serde::{Deserialize, Serialize};
use std::collections::{BTreeMap, BTreeSet};
use std::path::{Path, PathBuf};

pub const METADATA_FILE_NAME: &str = "package-metadata.bcs";
pub const UPGRADE_POLICY_CUSTOM_FIELD: &str = "upgrade_policy";

/// Represents a set of options for building artifacts from Move.
#[derive(Debug, Clone, Parser, Serialize, Deserialize)]
pub struct BuildOptions {
    #[clap(long)]
    pub with_srcs: bool,
    #[clap(long)]
    pub with_abis: bool,
    #[clap(long)]
    pub with_source_maps: bool,
    #[clap(long, default_value = "true")]
    pub with_error_map: bool,
    /// Installation directory for compiled artifacts. Defaults to <package>/build.
    #[clap(long, parse(from_os_str))]
    pub install_dir: Option<PathBuf>,
    #[clap(skip)] // TODO: have a parser for this; there is one in the CLI buts its  downstream
    pub named_addresses: BTreeMap<String, AccountAddress>,
}

// Because named_addresses as no parser, we can't use clap's default impl. This must be aligned
// with defaults above.
impl Default for BuildOptions {
    fn default() -> Self {
        Self {
            with_srcs: false,
            with_abis: false,
            with_source_maps: false,
            with_error_map: true,
            install_dir: None,
            named_addresses: Default::default(),
        }
    }
}

/// Represents a built package.  It allows to extract `PackageMetadata`. Can also be used to
/// just build Move code.
pub struct BuiltPackage {
    options: BuildOptions,
    package_path: PathBuf,
    package: CompiledPackage,
}

impl BuiltPackage {
    /// Builds the package and on success delivers a `BuiltPackage`.
    ///
    /// This function currently reports all Move compilation errors and warnings to stdout,
    /// and is not `Ok` if there was an error among those.
    pub fn build(package_path: PathBuf, options: BuildOptions) -> anyhow::Result<Self> {
        let build_config = BuildConfig {
            dev_mode: false,
            additional_named_addresses: options.named_addresses.clone(),
            architecture: None,
            generate_abis: options.with_abis,
            generate_docs: false,
            install_dir: options.install_dir.clone(),
            test_mode: false,
            force_recompilation: false,
            fetch_deps_only: false,
            fetch_latest_git_deps: false,
        };
        let mut package = build_config.compile_package_no_exit(&package_path, &mut Vec::new())?;
        for module in package.root_modules_map().iter_modules().iter() {
            verify_module_init_function(module)?;
        }
        let error_map = if options.with_error_map {
            generate_error_map(&package_path, &options)
        } else {
            None
        };
        if let Some(map) = &error_map {
            inject_module_metadata(
                package_path
                    .join(CompiledPackageLayout::Root.path())
                    .join(package.compiled_package_info.package_name.as_str()),
                &mut package,
                map,
            )?
        }
        Ok(Self {
            options,
            package_path,
            package,
        })
    }

    /// Returns the name of this package.
    pub fn name(&self) -> &str {
        self.package.compiled_package_info.package_name.as_str()
    }

    pub fn package_path(&self) -> &Path {
        self.package_path.as_path()
    }

    pub fn package_artifacts_path(&self) -> PathBuf {
        self.package_path
            .join(CompiledPackageLayout::Root.path())
            .join(self.name())
    }

    /// Extracts the bytecode for the modules of the built package.
    pub fn extract_code(&self) -> Vec<Vec<u8>> {
        self.package
            .root_modules()
            .map(|unit_with_source| unit_with_source.unit.serialize(None))
            .collect()
    }

    /// Returns the abis for this package, if available.
    pub fn extract_abis(&self) -> Option<Vec<EntryABI>> {
        self.package.compiled_abis.as_ref().map(|abis| {
            abis.iter()
                .map(|(_, bytes)| bcs::from_bytes::<EntryABI>(bytes.as_slice()).unwrap())
                .collect()
        })
    }

    /// Returns an iterator for all compiled proper (non-script) modules.
    pub fn modules(&self) -> impl Iterator<Item = &CompiledModule> {
        self.package
            .root_modules()
            .filter_map(|unit| match &unit.unit {
                CompiledUnit::Module(NamedCompiledModule { module, .. }) => Some(module),
                CompiledUnit::Script(_) => None,
            })
    }

    /// Returns the number of scripts in the package.
    pub fn script_count(&self) -> usize {
        self.package.scripts().count()
    }

    /// Returns the serialized bytecode of the scripts in the package.
    pub fn extract_script_code(&self) -> Vec<Vec<u8>> {
        self.package
            .scripts()
            .map(|unit_with_source| unit_with_source.unit.serialize(None))
            .collect()
    }

    /// Extracts metadata, as needed for releasing a package, from the built package.
    pub fn extract_metadata(&self) -> anyhow::Result<PackageMetadata> {
        let source_digest = self
            .package
            .compiled_package_info
            .source_digest
            .map(|s| s.to_string())
            .unwrap_or_default();
        let manifest_file = self.package_path.join("Move.toml");
        let manifest = std::fs::read_to_string(&manifest_file)?;
        let custom_props = extract_custom_fields(&manifest)?;
        let manifest = zip_metadata_str(&manifest)?;
        let upgrade_policy = if let Some(val) = custom_props.get(UPGRADE_POLICY_CUSTOM_FIELD) {
            str::parse::<UpgradePolicy>(val.as_ref())?
        } else {
            UpgradePolicy::compat()
        };
        let mut modules = vec![];
        for u in self.package.root_modules() {
            let name = u.unit.name().to_string();
            let source = if self.options.with_srcs {
                zip_metadata_str(&std::fs::read_to_string(&u.source_path)?)?
            } else {
                vec![]
            };
            let source_map = if self.options.with_source_maps {
                zip_metadata(&u.unit.serialize_source_map())?
            } else {
                vec![]
            };
            modules.push(ModuleMetadata {
                name,
                source,
                source_map,
                extension: MoveOption::default(),
            })
        }
        let deps = self
            .package
            .deps_compiled_units
            .iter()
            .map(|(name, unit)| {
                let package_name = name.as_str().to_string();
                let account = match &unit.unit {
                    CompiledUnit::Module(m) => AccountAddress::new(m.address.into_bytes()),
                    _ => panic!("script not a dependency"),
                };
                PackageDep {
                    account,
                    package_name,
                }
            })
            .collect::<BTreeSet<_>>()
            .into_iter()
            .collect();
        Ok(PackageMetadata {
            name: self.name().to_string(),
            upgrade_policy,
            upgrade_number: 0,
            source_digest,
            manifest,
            modules,
            deps,
            extension: MoveOption::none(),
        })
    }

    pub fn extract_metadata_and_save(&self) -> anyhow::Result<()> {
        let data = self.extract_metadata()?;
        let path = self.package_artifacts_path();
        std::fs::create_dir_all(&path)?;
        std::fs::write(path.join(METADATA_FILE_NAME), &bcs::to_bytes(&data)?)?;
        Ok(())
    }
}

fn extract_custom_fields(toml: &str) -> anyhow::Result<BTreeMap<String, String>> {
    let manifest = parse_source_manifest(parse_move_manifest_string(toml.to_owned())?)?;
    Ok(manifest
        .package
        .custom_properties
        .iter()
        .map(|(s, v)| (s.to_string(), v.to_string()))
        .collect())
}

fn inject_module_metadata(
    package_path: PathBuf,
    pack: &mut CompiledPackage,
    error_map: &ErrorMapping,
) -> anyhow::Result<()> {
    for unit_with_source in pack.root_compiled_units.iter_mut() {
        match &mut unit_with_source.unit {
            CompiledUnit::Module(named_module) => {
                if let Some(module_map) = error_map
                    .module_error_maps
                    .get(&named_module.module.self_id())
                {
                    if !module_map.is_empty() {
                        let serialized_metadata = bcs::to_bytes(&RuntimeModuleMetadata {
                            error_map: module_map.clone(),
                        })
                        .expect("BCS for RuntimeModuleMetadata");
                        named_module.module.metadata.push(Metadata {
                            key: APTOS_METADATA_KEY.clone(),
                            value: serialized_metadata,
                        });

                        // Also need to update the .mv file on disk.
                        let path = package_path
                            .join(CompiledPackageLayout::CompiledModules.path())
                            .join(named_module.name.as_str())
                            .with_extension(MOVE_COMPILED_EXTENSION);
                        if path.is_file() {
                            let bytes = unit_with_source.unit.serialize(None);
                            std::fs::write(path, &bytes)?;
                        }
                    }
                }
            }
            CompiledUnit::Script(_) => {}
        }
    }
    Ok(())
}
<commit_msg>[cli] Add compiler output and dependency warning<commit_after>// Copyright (c) Aptos
// SPDX-License-Identifier: Apache-2.0

use crate::error_map::generate_error_map;
use crate::natives::code::{
    ModuleMetadata, MoveOption, PackageDep, PackageMetadata, UpgradePolicy,
};
use crate::{zip_metadata, zip_metadata_str, RuntimeModuleMetadata, APTOS_METADATA_KEY};
use aptos_module_verifier::module_init::verify_module_init_function;
use aptos_types::account_address::AccountAddress;
use aptos_types::transaction::EntryABI;
use clap::Parser;
use move_binary_format::CompiledModule;
use move_command_line_common::files::MOVE_COMPILED_EXTENSION;
use move_compiler::compiled_unit::{CompiledUnit, NamedCompiledModule};
use move_core_types::errmap::ErrorMapping;
use move_core_types::metadata::Metadata;
use move_package::compilation::compiled_package::CompiledPackage;
use move_package::compilation::package_layout::CompiledPackageLayout;
use move_package::source_package::manifest_parser::{
    parse_move_manifest_string, parse_source_manifest,
};
use move_package::BuildConfig;
use serde::{Deserialize, Serialize};
use std::collections::{BTreeMap, BTreeSet};
use std::path::{Path, PathBuf};

pub const METADATA_FILE_NAME: &str = "package-metadata.bcs";
pub const UPGRADE_POLICY_CUSTOM_FIELD: &str = "upgrade_policy";

/// Represents a set of options for building artifacts from Move.
#[derive(Debug, Clone, Parser, Serialize, Deserialize)]
pub struct BuildOptions {
    #[clap(long)]
    pub with_srcs: bool,
    #[clap(long)]
    pub with_abis: bool,
    #[clap(long)]
    pub with_source_maps: bool,
    #[clap(long, default_value = "true")]
    pub with_error_map: bool,
    /// Installation directory for compiled artifacts. Defaults to <package>/build.
    #[clap(long, parse(from_os_str))]
    pub install_dir: Option<PathBuf>,
    #[clap(skip)] // TODO: have a parser for this; there is one in the CLI buts its  downstream
    pub named_addresses: BTreeMap<String, AccountAddress>,
}

// Because named_addresses as no parser, we can't use clap's default impl. This must be aligned
// with defaults above.
impl Default for BuildOptions {
    fn default() -> Self {
        Self {
            with_srcs: false,
            with_abis: false,
            with_source_maps: false,
            with_error_map: true,
            install_dir: None,
            named_addresses: Default::default(),
        }
    }
}

/// Represents a built package.  It allows to extract `PackageMetadata`. Can also be used to
/// just build Move code.
pub struct BuiltPackage {
    options: BuildOptions,
    package_path: PathBuf,
    package: CompiledPackage,
}

impl BuiltPackage {
    /// Builds the package and on success delivers a `BuiltPackage`.
    ///
    /// This function currently reports all Move compilation errors and warnings to stdout,
    /// and is not `Ok` if there was an error among those.
    pub fn build(package_path: PathBuf, options: BuildOptions) -> anyhow::Result<Self> {
        let build_config = BuildConfig {
            dev_mode: false,
            additional_named_addresses: options.named_addresses.clone(),
            architecture: None,
            generate_abis: options.with_abis,
            generate_docs: false,
            install_dir: options.install_dir.clone(),
            test_mode: false,
            force_recompilation: false,
            fetch_deps_only: false,
            fetch_latest_git_deps: false,
        };
        eprintln!("Compiling, may take a little while to download git dependencies...");
        let mut package =
            build_config.compile_package_no_exit(&package_path, &mut std::io::stderr())?;
        for module in package.root_modules_map().iter_modules().iter() {
            verify_module_init_function(module)?;
        }
        let error_map = if options.with_error_map {
            generate_error_map(&package_path, &options)
        } else {
            None
        };
        if let Some(map) = &error_map {
            inject_module_metadata(
                package_path
                    .join(CompiledPackageLayout::Root.path())
                    .join(package.compiled_package_info.package_name.as_str()),
                &mut package,
                map,
            )?
        }
        Ok(Self {
            options,
            package_path,
            package,
        })
    }

    /// Returns the name of this package.
    pub fn name(&self) -> &str {
        self.package.compiled_package_info.package_name.as_str()
    }

    pub fn package_path(&self) -> &Path {
        self.package_path.as_path()
    }

    pub fn package_artifacts_path(&self) -> PathBuf {
        self.package_path
            .join(CompiledPackageLayout::Root.path())
            .join(self.name())
    }

    /// Extracts the bytecode for the modules of the built package.
    pub fn extract_code(&self) -> Vec<Vec<u8>> {
        self.package
            .root_modules()
            .map(|unit_with_source| unit_with_source.unit.serialize(None))
            .collect()
    }

    /// Returns the abis for this package, if available.
    pub fn extract_abis(&self) -> Option<Vec<EntryABI>> {
        self.package.compiled_abis.as_ref().map(|abis| {
            abis.iter()
                .map(|(_, bytes)| bcs::from_bytes::<EntryABI>(bytes.as_slice()).unwrap())
                .collect()
        })
    }

    /// Returns an iterator for all compiled proper (non-script) modules.
    pub fn modules(&self) -> impl Iterator<Item = &CompiledModule> {
        self.package
            .root_modules()
            .filter_map(|unit| match &unit.unit {
                CompiledUnit::Module(NamedCompiledModule { module, .. }) => Some(module),
                CompiledUnit::Script(_) => None,
            })
    }

    /// Returns the number of scripts in the package.
    pub fn script_count(&self) -> usize {
        self.package.scripts().count()
    }

    /// Returns the serialized bytecode of the scripts in the package.
    pub fn extract_script_code(&self) -> Vec<Vec<u8>> {
        self.package
            .scripts()
            .map(|unit_with_source| unit_with_source.unit.serialize(None))
            .collect()
    }

    /// Extracts metadata, as needed for releasing a package, from the built package.
    pub fn extract_metadata(&self) -> anyhow::Result<PackageMetadata> {
        let source_digest = self
            .package
            .compiled_package_info
            .source_digest
            .map(|s| s.to_string())
            .unwrap_or_default();
        let manifest_file = self.package_path.join("Move.toml");
        let manifest = std::fs::read_to_string(&manifest_file)?;
        let custom_props = extract_custom_fields(&manifest)?;
        let manifest = zip_metadata_str(&manifest)?;
        let upgrade_policy = if let Some(val) = custom_props.get(UPGRADE_POLICY_CUSTOM_FIELD) {
            str::parse::<UpgradePolicy>(val.as_ref())?
        } else {
            UpgradePolicy::compat()
        };
        let mut modules = vec![];
        for u in self.package.root_modules() {
            let name = u.unit.name().to_string();
            let source = if self.options.with_srcs {
                zip_metadata_str(&std::fs::read_to_string(&u.source_path)?)?
            } else {
                vec![]
            };
            let source_map = if self.options.with_source_maps {
                zip_metadata(&u.unit.serialize_source_map())?
            } else {
                vec![]
            };
            modules.push(ModuleMetadata {
                name,
                source,
                source_map,
                extension: MoveOption::default(),
            })
        }
        let deps = self
            .package
            .deps_compiled_units
            .iter()
            .map(|(name, unit)| {
                let package_name = name.as_str().to_string();
                let account = match &unit.unit {
                    CompiledUnit::Module(m) => AccountAddress::new(m.address.into_bytes()),
                    _ => panic!("script not a dependency"),
                };
                PackageDep {
                    account,
                    package_name,
                }
            })
            .collect::<BTreeSet<_>>()
            .into_iter()
            .collect();
        Ok(PackageMetadata {
            name: self.name().to_string(),
            upgrade_policy,
            upgrade_number: 0,
            source_digest,
            manifest,
            modules,
            deps,
            extension: MoveOption::none(),
        })
    }

    pub fn extract_metadata_and_save(&self) -> anyhow::Result<()> {
        let data = self.extract_metadata()?;
        let path = self.package_artifacts_path();
        std::fs::create_dir_all(&path)?;
        std::fs::write(path.join(METADATA_FILE_NAME), &bcs::to_bytes(&data)?)?;
        Ok(())
    }
}

fn extract_custom_fields(toml: &str) -> anyhow::Result<BTreeMap<String, String>> {
    let manifest = parse_source_manifest(parse_move_manifest_string(toml.to_owned())?)?;
    Ok(manifest
        .package
        .custom_properties
        .iter()
        .map(|(s, v)| (s.to_string(), v.to_string()))
        .collect())
}

fn inject_module_metadata(
    package_path: PathBuf,
    pack: &mut CompiledPackage,
    error_map: &ErrorMapping,
) -> anyhow::Result<()> {
    for unit_with_source in pack.root_compiled_units.iter_mut() {
        match &mut unit_with_source.unit {
            CompiledUnit::Module(named_module) => {
                if let Some(module_map) = error_map
                    .module_error_maps
                    .get(&named_module.module.self_id())
                {
                    if !module_map.is_empty() {
                        let serialized_metadata = bcs::to_bytes(&RuntimeModuleMetadata {
                            error_map: module_map.clone(),
                        })
                        .expect("BCS for RuntimeModuleMetadata");
                        named_module.module.metadata.push(Metadata {
                            key: APTOS_METADATA_KEY.clone(),
                            value: serialized_metadata,
                        });

                        // Also need to update the .mv file on disk.
                        let path = package_path
                            .join(CompiledPackageLayout::CompiledModules.path())
                            .join(named_module.name.as_str())
                            .with_extension(MOVE_COMPILED_EXTENSION);
                        if path.is_file() {
                            let bytes = unit_with_source.unit.serialize(None);
                            std::fs::write(path, &bytes)?;
                        }
                    }
                }
            }
            CompiledUnit::Script(_) => {}
        }
    }
    Ok(())
}
<|endoftext|>


<commit_before>extern crate ktensor as k;
use std::sync::{Arc};

#[test]
fn linear_regression() {
    // Variables
    let input_x = Arc::new(k::Variable::new("input_x".to_string(), k::Vec2(0, 2)));
    let target_y = Arc::new(k::Variable::new("target_y".to_string(), k::Vec2(0, 2)));

    // Initialize
    let mut variable_context = k::Context::<f64>::with_capacity(2);
    k::variable::init_f64(vec![input_x.clone(), target_y.clone()], &mut variable_context);

    // Graph
    let layers: usize = 2;
    let mut states = Vec::<Arc<k::State>>::with_capacity(2 * layers);
    let mut graph_head: Arc<k::Graph<f64>> = input_x.clone();

    {
        let w = Arc::new(k::State::new(format!("weight_w_{}", 1), k::Vec2(2, 2)));
        let b = Arc::new(k::State::new(format!("weight_b_{}", 1), k::Vec2(1, 2)));

        let dot = Arc::new(k::op::dot::<f64>(format!("layer_{}_dot", 1), graph_head, w.clone()));
        let add = Arc::new(k::op::add_f64(format!("layer_{}_add", 1), dot, b.clone()));

        states.push(w);
        states.push(b);
        graph_head = add;
    }

    // initialize states
    let mut state_context = k::Context::<f64>::with_capacity(2 * layers);
    k::state::init_f64(states, &mut state_context);

    let softmax = Arc::new(k::op::softmax_f64("softmax".to_string(), graph_head));
    let xentropy = Arc::new(k::cost::softmax_cross_entropy_f64("xentropy".to_string(), softmax, target_y.clone()));

    // Training
    let mut history = k::Context::<f64>::with_capacity(5 * layers + 2 + 2);
    let print_rate = 4096;

    let training_set = (
        k::Tensor::from_vec(k::Vec2(4, 2), vec![
            0.0, 0.0,
            0.0, 1.0,
            1.0, 0.0,
            1.0, 1.0,
        ]),
        k::Tensor::from_vec(k::Vec2(4, 2), vec![
            0.0, 1.0,
            1.0, 0.0,
            1.0, 0.0,
            0.0, 1.0,
        ]),
    );

    for i in 0..4096 {
        // init batch
        variable_context.set(input_x.get_id(), training_set.0.clone());
        variable_context.set(target_y.get_id(), training_set.1.clone());

        k::train(xentropy.clone(), &mut state_context, &variable_context, &mut history, -0.001);
        if i % print_rate == 0 {

        }
    }

    variable_context.set(input_x.get_id(), training_set.0.clone());
    variable_context.set(target_y.get_id(), training_set.1.clone());
    let final_test = k::execute(xentropy.clone(), &state_context, &variable_context);
    for i in 0..4 {
        println!("{} {}", final_test.get(k::Vec2(i, 0)), final_test.get(k::Vec2(i, 1)));
    }
}
<commit_msg>check forward prop<commit_after>extern crate ktensor as k;
use std::sync::{Arc};

#[test]
fn linear_regression() {
    // Variables
    let input_x = Arc::new(k::Variable::new("input_x".to_string(), k::Vec2(0, 2)));
    let target_y = Arc::new(k::Variable::new("target_y".to_string(), k::Vec2(0, 2)));

    // Initialize
    let mut variable_context = k::Context::<f64>::with_capacity(2);
    k::variable::init_f64(vec![input_x.clone(), target_y.clone()], &mut variable_context);

    // Graph
    let layers: usize = 2;
    let mut states = Vec::<Arc<k::State>>::with_capacity(2 * layers);
    let mut graph_head: Arc<k::Graph<f64>> = input_x.clone();

    let w = Arc::new(k::State::new(format!("weight_w_{}", 1), k::Vec2(2, 2)));
    let b = Arc::new(k::State::new(format!("weight_b_{}", 1), k::Vec2(1, 2)));

    let dot = Arc::new(k::op::dot::<f64>(format!("layer_{}_dot", 1), graph_head.clone(), w.clone()));
    let add = Arc::new(k::op::add_f64(format!("layer_{}_add", 1), dot.clone(), b.clone()));

    let relu = Arc::new(k::op::relu_f64(format!("layer_{}_relu", 1), add.clone()));

    graph_head = relu.clone();

    let softmax = Arc::new(k::op::softmax_f64(format!("layer_{}_softmax", 1), graph_head.clone()));
    let xentropy = Arc::new(k::cost::softmax_cross_entropy_f64(format!("layer_{}_xentropy", 1), softmax.clone(), target_y.clone()));

    states.push(w.clone());
    states.push(b.clone());

    let w2 = Arc::new(k::State::new(format!("weight_w_{}", 2), k::Vec2(2, 2)));
    let b2 = Arc::new(k::State::new(format!("weight_b_{}", 2), k::Vec2(1, 2)));

    let dot2 = Arc::new(k::op::dot::<f64>(format!("layer_{}_dot", 2), graph_head.clone(), w2.clone()));
    let add2 = Arc::new(k::op::add_f64(format!("layer_{}_add", 2), dot2.clone(), b2.clone()));

    let relu2 = Arc::new(k::op::relu_f64(format!("layer_{}_relu", 2), add2.clone()));

    states.push(w2.clone());
    states.push(b2.clone());

    graph_head = relu2.clone();

    // initialize states
    let mut state_context = k::Context::<f64>::with_capacity(2 * layers);
    k::state::init_f64(states, &mut state_context);

    let training_set = (
        k::Tensor::from_vec(k::Vec2(4, 2), vec![
        0.0, 0.0,
        0.0, 1.0,
        1.0, 0.0,
        1.0, 1.0,
        ]),
        k::Tensor::from_vec(k::Vec2(4, 2), vec![
        0.0, 1.0,
        1.0, 0.0,
        1.0, 0.0,
        0.0, 1.0,
        ]),
    );

    variable_context.set(input_x.get_id(), training_set.0.clone());
    variable_context.set(target_y.get_id(), training_set.1.clone());
    println!("\ninput x");
    let final_test = k::execute(input_x.clone(), &state_context, &variable_context);
    for i in 0..4 {
        println!("{} {}", final_test.get(k::Vec2(i, 0)), final_test.get(k::Vec2(i, 1)));
    }
    println!("\nweight w 1");
    let final_test = k::execute(w.clone(), &state_context, &variable_context);
    for i in 0..2 {
        println!("{} {}", final_test.get(k::Vec2(i, 0)), final_test.get(k::Vec2(i, 1)));
    }
    println!("\ndot");
    let final_test = k::execute(dot.clone(), &state_context, &variable_context);
    for i in 0..4 {
        println!("{} {}", final_test.get(k::Vec2(i, 0)), final_test.get(k::Vec2(i, 1)));
    }
    println!("\nweight b 1");
    let final_test = k::execute(b.clone(), &state_context, &variable_context);
    for i in 0..1 {
        println!("{} {}", final_test.get(k::Vec2(i, 0)), final_test.get(k::Vec2(i, 1)));
    }
    println!("\nadd");
    let final_test = k::execute(add.clone(), &state_context, &variable_context);
    for i in 0..4 {
        println!("{} {}", final_test.get(k::Vec2(i, 0)), final_test.get(k::Vec2(i, 1)));
    }
    println!("\nrelu");
    let final_test = k::execute(relu.clone(), &state_context, &variable_context);
    for i in 0..4 {
        println!("{} {}", final_test.get(k::Vec2(i, 0)), final_test.get(k::Vec2(i, 1)));
    }
    println!("\ntarget y");
    let final_test = k::execute(target_y.clone(), &state_context, &variable_context);
    for i in 0..4 {
        println!("{} {}", final_test.get(k::Vec2(i, 0)), final_test.get(k::Vec2(i, 1)));
    }
    println!("\nxentropy");
    let final_test = k::execute(xentropy.clone(), &state_context, &variable_context);
    for i in 0..4 {
        println!("{} {}", final_test.get(k::Vec2(i, 0)), final_test.get(k::Vec2(i, 1)));
    }
}
<|endoftext|>


<commit_before>extern crate petgraph;

use petgraph::{
    OGraph,
};

use petgraph::ograph::min_spanning_tree;

#[test]
fn mst() {
    let mut gr = OGraph::<_, f32>::new();
    let a = gr.add_node("A");
    let b = gr.add_node("B");
    let c = gr.add_node("C");
    let d = gr.add_node("D");
    let e = gr.add_node("E");
    let f = gr.add_node("F");
    let g = gr.add_node("G");
    gr.add_edge(a, b, 7.);
    gr.add_edge(a, d, 5.);
    gr.add_edge(d, b, 9.);
    gr.add_edge(b, c, 8.);
    gr.add_edge(b, e, 7.);
    gr.add_edge(c, e, 5.);
    gr.add_edge(d, e, 15.);
    gr.add_edge(d, f, 6.);
    gr.add_edge(f, e, 8.);
    gr.add_edge(f, g, 11.);
    gr.add_edge(e, g, 9.);

    // add a disjoint part
    let h = gr.add_node("H");
    let i = gr.add_node("I");
    let j = gr.add_node("J");
    gr.add_edge(h, i, 1.);
    gr.add_edge(h, j, 3.);
    gr.add_edge(i, j, 1.);

    let mst = min_spanning_tree(&gr);
    println!("MST is:\n{}", mst);
    assert!(mst.node_count() == gr.node_count());
    // |E| = |N| - 2  because ther are two disconnected components.
    assert!(mst.edge_count() == gr.node_count() - 2);

    let one_of = |&: a: bool, b: bool| {
        (a && !b) || (!a && b)
    };
    let have_one = |&: x, y| {
        one_of(mst.find_edge(x, y).is_some(), mst.find_edge(y, x).is_some())
    };

    // check the exact edges are there
    assert!(have_one(a, b));
    assert!(have_one(a, d));
    assert!(have_one(b, e));
    assert!(have_one(e, c));
    assert!(have_one(e, g));
    assert!(have_one(d, f));

    assert!(have_one(h, i));
    assert!(have_one(i, j));

    assert!(mst.find_edge(d, b).is_none());
    assert!(mst.find_edge(b, c).is_none());

}
<commit_msg>Add test for selfloop in OGraph<commit_after>extern crate petgraph;

use petgraph::{
    OGraph,
};

use petgraph::ograph::min_spanning_tree;

#[test]
fn mst() {
    let mut gr = OGraph::<_, f32>::new();
    let a = gr.add_node("A");
    let b = gr.add_node("B");
    let c = gr.add_node("C");
    let d = gr.add_node("D");
    let e = gr.add_node("E");
    let f = gr.add_node("F");
    let g = gr.add_node("G");
    gr.add_edge(a, b, 7.);
    gr.add_edge(a, d, 5.);
    gr.add_edge(d, b, 9.);
    gr.add_edge(b, c, 8.);
    gr.add_edge(b, e, 7.);
    gr.add_edge(c, e, 5.);
    gr.add_edge(d, e, 15.);
    gr.add_edge(d, f, 6.);
    gr.add_edge(f, e, 8.);
    gr.add_edge(f, g, 11.);
    gr.add_edge(e, g, 9.);

    // add a disjoint part
    let h = gr.add_node("H");
    let i = gr.add_node("I");
    let j = gr.add_node("J");
    gr.add_edge(h, i, 1.);
    gr.add_edge(h, j, 3.);
    gr.add_edge(i, j, 1.);

    let mst = min_spanning_tree(&gr);
    println!("MST is:\n{}", mst);
    assert!(mst.node_count() == gr.node_count());
    // |E| = |N| - 2  because ther are two disconnected components.
    assert!(mst.edge_count() == gr.node_count() - 2);

    let one_of = |&: a: bool, b: bool| {
        (a && !b) || (!a && b)
    };
    let have_one = |&: x, y| {
        one_of(mst.find_edge(x, y).is_some(), mst.find_edge(y, x).is_some())
    };

    // check the exact edges are there
    assert!(have_one(a, b));
    assert!(have_one(a, d));
    assert!(have_one(b, e));
    assert!(have_one(e, c));
    assert!(have_one(e, g));
    assert!(have_one(d, f));

    assert!(have_one(h, i));
    assert!(have_one(i, j));

    assert!(mst.find_edge(d, b).is_none());
    assert!(mst.find_edge(b, c).is_none());

}

#[test]
fn selfloop() {
    let mut gr = OGraph::<_, f32>::new();
    let a = gr.add_node("A");
    let b = gr.add_node("B");
    let c = gr.add_node("C");
    gr.add_edge(a, b, 7.);
    gr.add_edge(c, a, 6.);
    let sed = gr.add_edge(a, a, 2.);

    assert!(gr.find_edge(a, b).is_some());
    assert!(gr.find_edge(b, a).is_none());
    assert!(gr.find_any_edge(b, a).is_some());
    assert!(gr.find_edge(a, a).is_some());
    println!("{}", gr);

    gr.remove_edge(sed);
    assert!(gr.find_edge(a, a).is_none());
    println!("{}", gr);
}
<|endoftext|>


<commit_before>// System tests for compiling C code.
//
// Copyright 2016 Mozilla Foundation
//
// 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.

#![allow(dead_code, unused_imports)]

extern crate assert_cmd;
extern crate cc;
extern crate env_logger;
#[macro_use]
extern crate log;
extern crate predicates;
extern crate sccache;
extern crate serde_json;
extern crate tempdir;
extern crate which;

use assert_cmd::prelude::*;
use log::LogLevel::Trace;
use predicates::prelude::*;
use sccache::server::ServerInfo;
use std::collections::HashMap;
use std::env;
use std::ffi::{OsStr,OsString};
use std::fmt;
use std::fs::{self, File};
use std::io::{self, Read, Write};
use std::path::{Path,PathBuf};
use std::process::{
    Command,
    Output,
    Stdio,
};
use std::str;
use tempdir::TempDir;
use which::which_in;

struct Compiler {
    pub name: &'static str,
    pub exe: OsString,
    pub env_vars: Vec<(OsString, OsString)>,
}

// Test GCC + clang on non-OS X platforms.
#[cfg(all(unix, not(target_os="macos")))]
const COMPILERS: &'static [&'static str] = &["gcc", "clang"];

// OS X ships a `gcc` that's just a clang wrapper, so only test clang there.
#[cfg(target_os="macos")]
const COMPILERS: &'static [&'static str] = &["clang"];

//TODO: could test gcc when targeting mingw.

fn stop() {
    trace!("sccache --stop-server");
    drop(Command::main_binary().unwrap()
         .arg("--stop-server")
         .stdout(Stdio::null())
         .stderr(Stdio::null())
         .status());
}

fn zero_stats() {
    trace!("sccache --zero-stats");
    drop(Command::main_binary().unwrap()
         .arg("--zero-stats")
         .stdout(Stdio::null())
         .stderr(Stdio::null())
         .status());
}

macro_rules! vec_from {
    ( $t:ty, $( $x:expr ),* ) => {
        vec!($( Into::<$t>::into(&$x), )*)
    };
}

fn compile_cmdline<T: AsRef<OsStr>>(compiler: &str, exe: T, input: &str, output: &str) -> Vec<OsString> {
    match compiler {
        "gcc" | "clang" => vec_from!(OsString, exe.as_ref(), "-c", input, "-o", output),
        "cl.exe" => vec_from!(OsString, exe, "-c", input, format!("-Fo{}", output)),
        _ => panic!("Unsupported compiler: {}", compiler),
    }
}

fn get_stats<F: 'static + Fn(ServerInfo)>(f: F) {
    Command::main_binary().unwrap()
        .args(&["--show-stats", "--stats-format=json"])
        .assert()
        .success()
        .stdout(predicate::function(move |output: &[u8]| {
            let s = str::from_utf8(output).expect("Output not UTF-8");
            f(serde_json::from_str(s).expect("Failed to parse JSON stats"));
            true
        }));
}

fn write_source(path: &Path, filename: &str, contents: &str) {
    let p = path.join(filename);
    let mut f = File::create(&p).unwrap();
    f.write_all(contents.as_bytes()).unwrap();
}

fn run_sccache_command_test(compiler: Compiler, tempdir: &Path) {
    let Compiler { name, exe, env_vars } = compiler;
    trace!("run_sccache_command_test: {}", name);
    // Compile a source file.
    const INPUT: &'static str = "test.c";
    const INPUT_ERR: &'static str = "test_err.c";
    const OUTPUT: &'static str = "test.o";
    // Copy the source files into the tempdir so we can compile with relative paths, since the commandline winds up in the hash key.
    for f in &[INPUT, INPUT_ERR] {
        let original_source_file = Path::new(file!()).parent().unwrap().join(f);
        let source_file = tempdir.join(f);
        trace!("fs::copy({:?}, {:?})", original_source_file, source_file);
        fs::copy(&original_source_file, &source_file).unwrap();
    }

    let out_file = tempdir.join("test.o");
    trace!("compile");
    Command::main_binary().unwrap()
        .args(&compile_cmdline(name, &exe, INPUT, OUTPUT))
        .current_dir(tempdir)
        .envs(env_vars.clone())
        .assert()
        .success();
    assert_eq!(true, fs::metadata(&out_file).and_then(|m| Ok(m.len() > 0)).unwrap());
    trace!("request stats");
    get_stats(|info| {
        assert_eq!(1, info.stats.compile_requests);
        assert_eq!(1, info.stats.requests_executed);
        assert_eq!(0, info.stats.cache_hits);
        assert_eq!(1, info.stats.cache_misses);
    });
    trace!("compile");
    fs::remove_file(&out_file).unwrap();
    Command::main_binary().unwrap()
        .args(&compile_cmdline(name, &exe, INPUT, OUTPUT))
        .current_dir(tempdir)
.envs(env_vars.clone())
        .assert()
        .success();
    assert_eq!(true, fs::metadata(&out_file).and_then(|m| Ok(m.len() > 0)).unwrap());
    trace!("request stats");
    get_stats(|info| {
        assert_eq!(2, info.stats.compile_requests);
        assert_eq!(2, info.stats.requests_executed);
        assert_eq!(1, info.stats.cache_hits);
        assert_eq!(1, info.stats.cache_misses);
    });
    if name == "cl.exe" {
        // Check that -deps works.
        trace!("compile with -deps");
        let mut args = compile_cmdline(name, &exe, INPUT, OUTPUT);
        args.push("-depstest.d".into());
        Command::main_binary().unwrap()
            .args(&args)
            .current_dir(tempdir)
            .envs(env_vars.clone())
            .assert()
            .success();
        // Check the contents
        let mut f = File::open(tempdir.join("test.d")).expect("Failed to open dep file");
        let mut buf = String::new();
        // read_to_string should be safe because we're supplying all the filenames here,
        // and there are no absolute paths.
        f.read_to_string(&mut buf).expect("Failed to read dep file");
        let lines: Vec<_> = buf.lines().map(|l| l.trim_right()).collect();
        let expected = format!("{output}: {input}\n{input}:\n", output=OUTPUT, input=INPUT);
        let expected_lines: Vec<_> = expected.lines().collect();
        assert_eq!(lines, expected_lines);
    }
    if name == "gcc" {
        trace!("test -MP with -Werror");
        let mut args = compile_cmdline(name, &exe, INPUT_ERR, OUTPUT);
        args.extend(vec_from!(OsString, "-MD", "-MP", "-MF", "foo.pp", "-Werror"));
        // This should fail, but the error should be from the #error!
        Command::main_binary().unwrap()
            .args(&args)
            .current_dir(tempdir)
    .envs(env_vars.clone())
            .assert()
            .failure()
            .stderr(predicates::str::contains(
                "to generate dependencies you must specify either -M or -MM").from_utf8().not());
        trace!("test -fprofile-generate with different source inputs");
        zero_stats();
        const SRC: &str = "source.c";
        write_source(&tempdir, SRC, "/*line 1*/
#ifndef UNDEFINED
/*unused line 1*/
#endif

int main(int argc, char** argv) {
  return 0;
}
");
        let mut args = compile_cmdline(name, &exe, SRC, OUTPUT);
        args.extend(vec_from!(OsString, "-fprofile-generate"));
        trace!("compile source.c (1)");
        Command::main_binary().unwrap()
            .args(&args)
            .current_dir(tempdir)
            .envs(env_vars.clone())
            .assert()
            .success();
        get_stats(|info| {
            assert_eq!(0, info.stats.cache_hits);
            assert_eq!(1, info.stats.cache_misses);
        });
        // Compile the same source again to ensure we can get a cache hit.
        trace!("compile source.c (2)");
        Command::main_binary().unwrap()
            .args(&args)
            .current_dir(tempdir)
            .envs(env_vars.clone())
            .assert()
            .success();
        get_stats(|info| {
            assert_eq!(1, info.stats.cache_hits);
            assert_eq!(1, info.stats.cache_misses);
        });
        // Now write out a slightly different source file that will preprocess to the same thing,
        // modulo line numbers. This should not be a cache hit because line numbers are important
        // with -fprofile-generate.
        write_source(&tempdir, SRC, "/*line 1*/
#ifndef UNDEFINED
/*unused line 1*/
/*unused line 2*/
#endif

int main(int argc, char** argv) {
  return 0;
}
");
        trace!("compile source.c (3)");
        Command::main_binary().unwrap()
            .args(&args)
            .current_dir(tempdir)
            .envs(env_vars.clone())
            .assert()
            .success();
        get_stats(|info| {
            assert_eq!(1, info.stats.cache_hits);
            assert_eq!(2, info.stats.cache_misses);
        });
    }
}

#[cfg(unix)]
fn find_compilers() -> Vec<Compiler> {
    let cwd = env::current_dir().unwrap();
    COMPILERS.iter()
        .filter_map(|c| {
            match which_in(c, env::var_os("PATH"), &cwd) {
                Ok(full_path) => match full_path.canonicalize() {
                    Ok(full_path_canon) => Some(Compiler {
                        name: *c,
                        exe: full_path_canon.into_os_string(),
                        env_vars: vec![],
                    }),
                    Err(_) => None,
                },
                Err(_) => None,
            }
        })
        .collect::<Vec<_>>()
}

#[cfg(target_env="msvc")]
fn find_compilers() -> Vec<Compiler> {
    let tool = cc::Build::new()
        .opt_level(1)
        .host("x86_64-pc-windows-msvc")
        .target("x86_64-pc-windows-msvc")
        .debug(false)
        .get_compiler();
    vec![
        Compiler {
            name: "cl.exe",
            exe: tool.path().as_os_str().to_os_string(),
            env_vars: tool.env().to_vec(),
        },
    ]
}

#[test]
#[cfg(any(unix, target_env="msvc"))]
fn test_sccache_command() {
    match env_logger::init() {
        Ok(_) => {},
        Err(_) => {},
    }
    let tempdir = TempDir::new("sccache_system_test").unwrap();
    let compilers = find_compilers();
    if compilers.is_empty() {
        warn!("No compilers found, skipping test");
    } else {
        // Ensure there's no existing sccache server running.
        stop();
        // Create a subdir for the cache.
        let cache = tempdir.path().join("cache");
        fs::create_dir_all(&cache).unwrap();
        // Start a server.
        trace!("start server");
        // Don't run this with run() because on Windows `wait_with_output`
        // will hang because the internal server process is not detached.
        Command::main_binary().unwrap()
            .arg("--start-server")
            .env("SCCACHE_DIR", &cache)
            .status()
            .unwrap()
            .success();
        for compiler in compilers {
            run_sccache_command_test(compiler, tempdir.path());
            zero_stats();
        }
        stop();
    }
}
<commit_msg>split system tests into one function per actual test<commit_after>// System tests for compiling C code.
//
// Copyright 2016 Mozilla Foundation
//
// 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.

#![allow(dead_code, unused_imports)]

extern crate assert_cmd;
extern crate cc;
extern crate env_logger;
#[macro_use]
extern crate log;
extern crate predicates;
extern crate sccache;
extern crate serde_json;
extern crate tempdir;
extern crate which;

use assert_cmd::prelude::*;
use log::LogLevel::Trace;
use predicates::prelude::*;
use sccache::server::ServerInfo;
use std::collections::HashMap;
use std::env;
use std::ffi::{OsStr,OsString};
use std::fmt;
use std::fs::{self, File};
use std::io::{self, Read, Write};
use std::path::{Path,PathBuf};
use std::process::{
    Command,
    Output,
    Stdio,
};
use std::str;
use tempdir::TempDir;
use which::which_in;

#[derive(Clone)]
struct Compiler {
    pub name: &'static str,
    pub exe: OsString,
    pub env_vars: Vec<(OsString, OsString)>,
}

// Test GCC + clang on non-OS X platforms.
#[cfg(all(unix, not(target_os="macos")))]
const COMPILERS: &'static [&'static str] = &["gcc", "clang"];

// OS X ships a `gcc` that's just a clang wrapper, so only test clang there.
#[cfg(target_os="macos")]
const COMPILERS: &'static [&'static str] = &["clang"];

//TODO: could test gcc when targeting mingw.

fn stop() {
    trace!("sccache --stop-server");
    drop(Command::main_binary().unwrap()
         .arg("--stop-server")
         .stdout(Stdio::null())
         .stderr(Stdio::null())
         .status());
}

fn zero_stats() {
    trace!("sccache --zero-stats");
    drop(Command::main_binary().unwrap()
         .arg("--zero-stats")
         .stdout(Stdio::null())
         .stderr(Stdio::null())
         .status());
}

macro_rules! vec_from {
    ( $t:ty, $( $x:expr ),* ) => {
        vec!($( Into::<$t>::into(&$x), )*)
    };
}

fn compile_cmdline<T: AsRef<OsStr>>(compiler: &str, exe: T, input: &str, output: &str) -> Vec<OsString> {
    match compiler {
        "gcc" | "clang" => vec_from!(OsString, exe.as_ref(), "-c", input, "-o", output),
        "cl.exe" => vec_from!(OsString, exe, "-c", input, format!("-Fo{}", output)),
        _ => panic!("Unsupported compiler: {}", compiler),
    }
}

fn get_stats<F: 'static + Fn(ServerInfo)>(f: F) {
    Command::main_binary().unwrap()
        .args(&["--show-stats", "--stats-format=json"])
        .assert()
        .success()
        .stdout(predicate::function(move |output: &[u8]| {
            let s = str::from_utf8(output).expect("Output not UTF-8");
            f(serde_json::from_str(s).expect("Failed to parse JSON stats"));
            true
        }));
}

fn write_source(path: &Path, filename: &str, contents: &str) {
    let p = path.join(filename);
    let mut f = File::create(&p).unwrap();
    f.write_all(contents.as_bytes()).unwrap();
}

const INPUT: &'static str = "test.c";
const INPUT_ERR: &'static str = "test_err.c";
const OUTPUT: &'static str = "test.o";

fn test_basic_compile(compiler: Compiler, tempdir: &Path) {
    let Compiler { name, exe, env_vars } = compiler;
    trace!("run_sccache_command_test: {}", name);
    // Compile a source file.
    // Copy the source files into the tempdir so we can compile with relative paths, since the commandline winds up in the hash key.
    for f in &[INPUT, INPUT_ERR] {
        let original_source_file = Path::new(file!()).parent().unwrap().join(f);
        let source_file = tempdir.join(f);
        trace!("fs::copy({:?}, {:?})", original_source_file, source_file);
        fs::copy(&original_source_file, &source_file).unwrap();
    }

    let out_file = tempdir.join("test.o");
    trace!("compile");
    Command::main_binary().unwrap()
        .args(&compile_cmdline(name, &exe, INPUT, OUTPUT))
        .current_dir(tempdir)
        .envs(env_vars.clone())
        .assert()
        .success();
    assert_eq!(true, fs::metadata(&out_file).and_then(|m| Ok(m.len() > 0)).unwrap());
    trace!("request stats");
    get_stats(|info| {
        assert_eq!(1, info.stats.compile_requests);
        assert_eq!(1, info.stats.requests_executed);
        assert_eq!(0, info.stats.cache_hits);
        assert_eq!(1, info.stats.cache_misses);
    });
    trace!("compile");
    fs::remove_file(&out_file).unwrap();
    Command::main_binary().unwrap()
        .args(&compile_cmdline(name, &exe, INPUT, OUTPUT))
        .current_dir(tempdir)
.envs(env_vars.clone())
        .assert()
        .success();
    assert_eq!(true, fs::metadata(&out_file).and_then(|m| Ok(m.len() > 0)).unwrap());
    trace!("request stats");
    get_stats(|info| {
        assert_eq!(2, info.stats.compile_requests);
        assert_eq!(2, info.stats.requests_executed);
        assert_eq!(1, info.stats.cache_hits);
        assert_eq!(1, info.stats.cache_misses);
    });
}

fn test_msvc_deps(compiler: Compiler, tempdir: &Path) {
    let Compiler { name, exe, env_vars } = compiler;
    if name == "cl.exe" {
        // Check that -deps works.
        trace!("compile with -deps");
        let mut args = compile_cmdline(name, &exe, INPUT, OUTPUT);
        args.push("-depstest.d".into());
        Command::main_binary().unwrap()
            .args(&args)
            .current_dir(tempdir)
            .envs(env_vars.clone())
            .assert()
            .success();
        // Check the contents
        let mut f = File::open(tempdir.join("test.d")).expect("Failed to open dep file");
        let mut buf = String::new();
        // read_to_string should be safe because we're supplying all the filenames here,
        // and there are no absolute paths.
        f.read_to_string(&mut buf).expect("Failed to read dep file");
        let lines: Vec<_> = buf.lines().map(|l| l.trim_right()).collect();
        let expected = format!("{output}: {input}\n{input}:\n", output=OUTPUT, input=INPUT);
        let expected_lines: Vec<_> = expected.lines().collect();
        assert_eq!(lines, expected_lines);
    }
}

fn test_gcc_mp_werror(compiler: Compiler, tempdir: &Path) {
    let Compiler { name, exe, env_vars } = compiler;
    if name == "gcc" {
        trace!("test -MP with -Werror");
        let mut args = compile_cmdline(name, &exe, INPUT_ERR, OUTPUT);
        args.extend(vec_from!(OsString, "-MD", "-MP", "-MF", "foo.pp", "-Werror"));
        // This should fail, but the error should be from the #error!
        Command::main_binary().unwrap()
            .args(&args)
            .current_dir(tempdir)
            .envs(env_vars.clone())
            .assert()
            .failure()
            .stderr(predicates::str::contains(
                "to generate dependencies you must specify either -M or -MM").from_utf8().not());
    }
}

fn test_gcc_fprofile_generate_source_changes(compiler: Compiler, tempdir: &Path) {
    let Compiler { name, exe, env_vars } = compiler;
    if name == "gcc" {
        trace!("test -fprofile-generate with different source inputs");
        zero_stats();
        const SRC: &str = "source.c";
        write_source(&tempdir, SRC, "/*line 1*/
#ifndef UNDEFINED
/*unused line 1*/
#endif

int main(int argc, char** argv) {
  return 0;
}
");
        let mut args = compile_cmdline(name, &exe, SRC, OUTPUT);
        args.extend(vec_from!(OsString, "-fprofile-generate"));
        trace!("compile source.c (1)");
        Command::main_binary().unwrap()
            .args(&args)
            .current_dir(tempdir)
            .envs(env_vars.clone())
            .assert()
            .success();
        get_stats(|info| {
            assert_eq!(0, info.stats.cache_hits);
            assert_eq!(1, info.stats.cache_misses);
        });
        // Compile the same source again to ensure we can get a cache hit.
        trace!("compile source.c (2)");
        Command::main_binary().unwrap()
            .args(&args)
            .current_dir(tempdir)
            .envs(env_vars.clone())
            .assert()
            .success();
        get_stats(|info| {
            assert_eq!(1, info.stats.cache_hits);
            assert_eq!(1, info.stats.cache_misses);
        });
        // Now write out a slightly different source file that will preprocess to the same thing,
        // modulo line numbers. This should not be a cache hit because line numbers are important
        // with -fprofile-generate.
        write_source(&tempdir, SRC, "/*line 1*/
#ifndef UNDEFINED
/*unused line 1*/
/*unused line 2*/
#endif

int main(int argc, char** argv) {
  return 0;
}
");
        trace!("compile source.c (3)");
        Command::main_binary().unwrap()
            .args(&args)
            .current_dir(tempdir)
            .envs(env_vars.clone())
            .assert()
            .success();
        get_stats(|info| {
            assert_eq!(1, info.stats.cache_hits);
            assert_eq!(2, info.stats.cache_misses);
        });
    }
}

fn run_sccache_command_tests(compiler: Compiler, tempdir: &Path) {
    test_basic_compile(compiler.clone(), tempdir);
    test_msvc_deps(compiler.clone(), tempdir);
    test_gcc_mp_werror(compiler.clone(), tempdir);
    test_gcc_fprofile_generate_source_changes(compiler.clone(), tempdir);
}

#[cfg(unix)]
fn find_compilers() -> Vec<Compiler> {
    let cwd = env::current_dir().unwrap();
    COMPILERS.iter()
        .filter_map(|c| {
            match which_in(c, env::var_os("PATH"), &cwd) {
                Ok(full_path) => match full_path.canonicalize() {
                    Ok(full_path_canon) => Some(Compiler {
                        name: *c,
                        exe: full_path_canon.into_os_string(),
                        env_vars: vec![],
                    }),
                    Err(_) => None,
                },
                Err(_) => None,
            }
        })
        .collect::<Vec<_>>()
}

#[cfg(target_env="msvc")]
fn find_compilers() -> Vec<Compiler> {
    let tool = cc::Build::new()
        .opt_level(1)
        .host("x86_64-pc-windows-msvc")
        .target("x86_64-pc-windows-msvc")
        .debug(false)
        .get_compiler();
    vec![
        Compiler {
            name: "cl.exe",
            exe: tool.path().as_os_str().to_os_string(),
            env_vars: tool.env().to_vec(),
        },
    ]
}

#[test]
#[cfg(any(unix, target_env="msvc"))]
fn test_sccache_command() {
    match env_logger::init() {
        Ok(_) => {},
        Err(_) => {},
    }
    let tempdir = TempDir::new("sccache_system_test").unwrap();
    let compilers = find_compilers();
    if compilers.is_empty() {
        warn!("No compilers found, skipping test");
    } else {
        // Ensure there's no existing sccache server running.
        stop();
        // Create a subdir for the cache.
        let cache = tempdir.path().join("cache");
        fs::create_dir_all(&cache).unwrap();
        // Start a server.
        trace!("start server");
        // Don't run this with run() because on Windows `wait_with_output`
        // will hang because the internal server process is not detached.
        Command::main_binary().unwrap()
            .arg("--start-server")
            .env("SCCACHE_DIR", &cache)
            .status()
            .unwrap()
            .success();
        for compiler in compilers {
            run_sccache_command_tests(compiler, tempdir.path());
            zero_stats();
        }
        stop();
    }
}
<|endoftext|>


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

//! Sampling from random distributions.
//!
//! This is a generalization of `Rand` to allow parameters to control the
//! exact properties of the generated values, e.g. the mean and standard
//! deviation of a normal distribution. The `Sample` trait is the most
//! general, and allows for generating values that change some state
//! internally. The `IndependentSample` trait is for generating values
//! that do not need to record state.

#![experimental]

use core::prelude::*;
use core::num::{Float, Int};

use {Rng, Rand};

pub use self::range::Range;
pub use self::gamma::{Gamma, ChiSquared, FisherF, StudentT};
pub use self::normal::{Normal, LogNormal};
pub use self::exponential::Exp;

pub mod range;
pub mod gamma;
pub mod normal;
pub mod exponential;

/// Types that can be used to create a random instance of `Support`.
pub trait Sample<Support> {
    /// Generate a random value of `Support`, using `rng` as the
    /// source of randomness.
    fn sample<R: Rng>(&mut self, rng: &mut R) -> Support;
}

/// `Sample`s that do not require keeping track of state.
///
/// Since no state is recorded, each sample is (statistically)
/// independent of all others, assuming the `Rng` used has this
/// property.
// FIXME maybe having this separate is overkill (the only reason is to
// take &self rather than &mut self)? or maybe this should be the
// trait called `Sample` and the other should be `DependentSample`.
pub trait IndependentSample<Support>: Sample<Support> {
    /// Generate a random value.
    fn ind_sample<R: Rng>(&self, &mut R) -> Support;
}

/// A wrapper for generating types that implement `Rand` via the
/// `Sample` & `IndependentSample` traits.
pub struct RandSample<Sup>;

impl<Sup: Rand> Sample<Sup> for RandSample<Sup> {
    fn sample<R: Rng>(&mut self, rng: &mut R) -> Sup { self.ind_sample(rng) }
}

impl<Sup: Rand> IndependentSample<Sup> for RandSample<Sup> {
    fn ind_sample<R: Rng>(&self, rng: &mut R) -> Sup {
        rng.gen()
    }
}

/// A value with a particular weight for use with `WeightedChoice`.
pub struct Weighted<T> {
    /// The numerical weight of this item
    pub weight: uint,
    /// The actual item which is being weighted
    pub item: T,
}

/// A distribution that selects from a finite collection of weighted items.
///
/// Each item has an associated weight that influences how likely it
/// is to be chosen: higher weight is more likely.
///
/// The `Clone` restriction is a limitation of the `Sample` and
/// `IndependentSample` traits. Note that `&T` is (cheaply) `Clone` for
/// all `T`, as is `uint`, so one can store references or indices into
/// another vector.
///
/// # Example
///
/// ```rust
/// use std::rand;
/// use std::rand::distributions::{Weighted, WeightedChoice, IndependentSample};
///
/// let mut items = vec!(Weighted { weight: 2, item: 'a' },
///                      Weighted { weight: 4, item: 'b' },
///                      Weighted { weight: 1, item: 'c' });
/// let wc = WeightedChoice::new(items.as_mut_slice());
/// let mut rng = rand::task_rng();
/// for _ in range(0u, 16) {
///      // on average prints 'a' 4 times, 'b' 8 and 'c' twice.
///      println!("{}", wc.ind_sample(&mut rng));
/// }
/// ```
pub struct WeightedChoice<'a, T:'a> {
    items: &'a mut [Weighted<T>],
    weight_range: Range<uint>
}

impl<'a, T: Clone> WeightedChoice<'a, T> {
    /// Create a new `WeightedChoice`.
    ///
    /// Panics if:
    /// - `v` is empty
    /// - the total weight is 0
    /// - the total weight is larger than a `uint` can contain.
    pub fn new<'a>(items: &'a mut [Weighted<T>]) -> WeightedChoice<'a, T> {
        // strictly speaking, this is subsumed by the total weight == 0 case
        assert!(!items.is_empty(), "WeightedChoice::new called with no items");

        let mut running_total = 0u;

        // we convert the list from individual weights to cumulative
        // weights so we can binary search. This *could* drop elements
        // with weight == 0 as an optimisation.
        for item in items.iter_mut() {
            running_total = match running_total.checked_add(item.weight) {
                Some(n) => n,
                None => panic!("WeightedChoice::new called with a total weight \
                               larger than a uint can contain")
            };

            item.weight = running_total;
        }
        assert!(running_total != 0, "WeightedChoice::new called with a total weight of 0");

        WeightedChoice {
            items: items,
            // we're likely to be generating numbers in this range
            // relatively often, so might as well cache it
            weight_range: Range::new(0, running_total)
        }
    }
}

impl<'a, T: Clone> Sample<T> for WeightedChoice<'a, T> {
    fn sample<R: Rng>(&mut self, rng: &mut R) -> T { self.ind_sample(rng) }
}

impl<'a, T: Clone> IndependentSample<T> for WeightedChoice<'a, T> {
    fn ind_sample<R: Rng>(&self, rng: &mut R) -> T {
        // we want to find the first element that has cumulative
        // weight > sample_weight, which we do by binary since the
        // cumulative weights of self.items are sorted.

        // choose a weight in [0, total_weight)
        let sample_weight = self.weight_range.ind_sample(rng);

        // short circuit when it's the first item
        if sample_weight < self.items[0].weight {
            return self.items[0].item.clone();
        }

        let mut idx = 0;
        let mut modifier = self.items.len();

        // now we know that every possibility has an element to the
        // left, so we can just search for the last element that has
        // cumulative weight <= sample_weight, then the next one will
        // be "it". (Note that this greatest element will never be the
        // last element of the vector, since sample_weight is chosen
        // in [0, total_weight) and the cumulative weight of the last
        // one is exactly the total weight.)
        while modifier > 1 {
            let i = idx + modifier / 2;
            if self.items[i].weight <= sample_weight {
                // we're small, so look to the right, but allow this
                // exact element still.
                idx = i;
                // we need the `/ 2` to round up otherwise we'll drop
                // the trailing elements when `modifier` is odd.
                modifier += 1;
            } else {
                // otherwise we're too big, so go left. (i.e. do
                // nothing)
            }
            modifier /= 2;
        }
        return self.items[idx + 1].item.clone();
    }
}

mod ziggurat_tables;

/// Sample a random number using the Ziggurat method (specifically the
/// ZIGNOR variant from Doornik 2005). Most of the arguments are
/// directly from the paper:
///
/// * `rng`: source of randomness
/// * `symmetric`: whether this is a symmetric distribution, or one-sided with P(x < 0) = 0.
/// * `X`: the $x_i$ abscissae.
/// * `F`: precomputed values of the PDF at the $x_i$, (i.e. $f(x_i)$)
/// * `F_DIFF`: precomputed values of $f(x_i) - f(x_{i+1})$
/// * `pdf`: the probability density function
/// * `zero_case`: manual sampling from the tail when we chose the
///    bottom box (i.e. i == 0)

// the perf improvement (25-50%) is definitely worth the extra code
// size from force-inlining.
#[inline(always)]
fn ziggurat<R: Rng, P, Z>(
            rng: &mut R,
            symmetric: bool,
            x_tab: ziggurat_tables::ZigTable,
            f_tab: ziggurat_tables::ZigTable,
            mut pdf: P,
            mut zero_case: Z)
            -> f64 where P: FnMut(f64) -> f64, Z: FnMut(&mut R, f64) -> f64 {
    static SCALE: f64 = (1u64 << 53) as f64;
    loop {
        // reimplement the f64 generation as an optimisation suggested
        // by the Doornik paper: we have a lot of precision-space
        // (i.e. there are 11 bits of the 64 of a u64 to use after
        // creating a f64), so we might as well reuse some to save
        // generating a whole extra random number. (Seems to be 15%
        // faster.)
        //
        // This unfortunately misses out on the benefits of direct
        // floating point generation if an RNG like dSMFT is
        // used. (That is, such RNGs create floats directly, highly
        // efficiently and overload next_f32/f64, so by not calling it
        // this may be slower than it would be otherwise.)
        // FIXME: investigate/optimise for the above.
        let bits: u64 = rng.gen();
        let i = (bits & 0xff) as uint;
        let f = (bits >> 11) as f64 / SCALE;

        // u is either U(-1, 1) or U(0, 1) depending on if this is a
        // symmetric distribution or not.
        let u = if symmetric {2.0 * f - 1.0} else {f};
        let x = u * x_tab[i];

        let test_x = if symmetric { x.abs() } else {x};

        // algebraically equivalent to |u| < x_tab[i+1]/x_tab[i] (or u < x_tab[i+1]/x_tab[i])
        if test_x < x_tab[i + 1] {
            return x;
        }
        if i == 0 {
            return zero_case(rng, u);
        }
        // algebraically equivalent to f1 + DRanU()*(f0 - f1) < 1
        if f_tab[i + 1] + (f_tab[i] - f_tab[i + 1]) * rng.gen() < pdf(x) {
            return x;
        }
    }
}

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

    use {Rng, Rand};
    use super::{RandSample, WeightedChoice, Weighted, Sample, IndependentSample};

    #[deriving(PartialEq, Show)]
    struct ConstRand(uint);
    impl Rand for ConstRand {
        fn rand<R: Rng>(_: &mut R) -> ConstRand {
            ConstRand(0)
        }
    }

    // 0, 1, 2, 3, ...
    struct CountingRng { i: u32 }
    impl Rng for CountingRng {
        fn next_u32(&mut self) -> u32 {
            self.i += 1;
            self.i - 1
        }
        fn next_u64(&mut self) -> u64 {
            self.next_u32() as u64
        }
    }

    #[test]
    fn test_rand_sample() {
        let mut rand_sample = RandSample::<ConstRand>;

        assert_eq!(rand_sample.sample(&mut ::test::rng()), ConstRand(0));
        assert_eq!(rand_sample.ind_sample(&mut ::test::rng()), ConstRand(0));
    }
    #[test]
    fn test_weighted_choice() {
        // this makes assumptions about the internal implementation of
        // WeightedChoice, specifically: it doesn't reorder the items,
        // it doesn't do weird things to the RNG (so 0 maps to 0, 1 to
        // 1, internally; modulo a modulo operation).

        macro_rules! t (
            ($items:expr, $expected:expr) => {{
                let mut items = $items;
                let wc = WeightedChoice::new(items.as_mut_slice());
                let expected = $expected;

                let mut rng = CountingRng { i: 0 };

                for &val in expected.iter() {
                    assert_eq!(wc.ind_sample(&mut rng), val)
                }
            }}
        );

        t!(vec!(Weighted { weight: 1, item: 10i}), [10]);

        // skip some
        t!(vec!(Weighted { weight: 0, item: 20i},
                Weighted { weight: 2, item: 21i},
                Weighted { weight: 0, item: 22i},
                Weighted { weight: 1, item: 23i}),
           [21,21, 23]);

        // different weights
        t!(vec!(Weighted { weight: 4, item: 30i},
                Weighted { weight: 3, item: 31i}),
           [30,30,30,30, 31,31,31]);

        // check that we're binary searching
        // correctly with some vectors of odd
        // length.
        t!(vec!(Weighted { weight: 1, item: 40i},
                Weighted { weight: 1, item: 41i},
                Weighted { weight: 1, item: 42i},
                Weighted { weight: 1, item: 43i},
                Weighted { weight: 1, item: 44i}),
           [40, 41, 42, 43, 44]);
        t!(vec!(Weighted { weight: 1, item: 50i},
                Weighted { weight: 1, item: 51i},
                Weighted { weight: 1, item: 52i},
                Weighted { weight: 1, item: 53i},
                Weighted { weight: 1, item: 54i},
                Weighted { weight: 1, item: 55i},
                Weighted { weight: 1, item: 56i}),
           [50, 51, 52, 53, 54, 55, 56]);
    }

    #[test] #[should_fail]
    fn test_weighted_choice_no_items() {
        WeightedChoice::<int>::new(&mut []);
    }
    #[test] #[should_fail]
    fn test_weighted_choice_zero_weight() {
        WeightedChoice::new(&mut [Weighted { weight: 0, item: 0i},
                                  Weighted { weight: 0, item: 1i}]);
    }
    #[test] #[should_fail]
    fn test_weighted_choice_weight_overflows() {
        let x = (-1) as uint / 2; // x + x + 2 is the overflow
        WeightedChoice::new(&mut [Weighted { weight: x, item: 0i },
                                  Weighted { weight: 1, item: 1i },
                                  Weighted { weight: x, item: 2i },
                                  Weighted { weight: 1, item: 3i }]);
    }
}
<commit_msg>Remove all shadowed lifetimes.<commit_after>// Copyright 2013 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

//! Sampling from random distributions.
//!
//! This is a generalization of `Rand` to allow parameters to control the
//! exact properties of the generated values, e.g. the mean and standard
//! deviation of a normal distribution. The `Sample` trait is the most
//! general, and allows for generating values that change some state
//! internally. The `IndependentSample` trait is for generating values
//! that do not need to record state.

#![experimental]

use core::prelude::*;
use core::num::{Float, Int};

use {Rng, Rand};

pub use self::range::Range;
pub use self::gamma::{Gamma, ChiSquared, FisherF, StudentT};
pub use self::normal::{Normal, LogNormal};
pub use self::exponential::Exp;

pub mod range;
pub mod gamma;
pub mod normal;
pub mod exponential;

/// Types that can be used to create a random instance of `Support`.
pub trait Sample<Support> {
    /// Generate a random value of `Support`, using `rng` as the
    /// source of randomness.
    fn sample<R: Rng>(&mut self, rng: &mut R) -> Support;
}

/// `Sample`s that do not require keeping track of state.
///
/// Since no state is recorded, each sample is (statistically)
/// independent of all others, assuming the `Rng` used has this
/// property.
// FIXME maybe having this separate is overkill (the only reason is to
// take &self rather than &mut self)? or maybe this should be the
// trait called `Sample` and the other should be `DependentSample`.
pub trait IndependentSample<Support>: Sample<Support> {
    /// Generate a random value.
    fn ind_sample<R: Rng>(&self, &mut R) -> Support;
}

/// A wrapper for generating types that implement `Rand` via the
/// `Sample` & `IndependentSample` traits.
pub struct RandSample<Sup>;

impl<Sup: Rand> Sample<Sup> for RandSample<Sup> {
    fn sample<R: Rng>(&mut self, rng: &mut R) -> Sup { self.ind_sample(rng) }
}

impl<Sup: Rand> IndependentSample<Sup> for RandSample<Sup> {
    fn ind_sample<R: Rng>(&self, rng: &mut R) -> Sup {
        rng.gen()
    }
}

/// A value with a particular weight for use with `WeightedChoice`.
pub struct Weighted<T> {
    /// The numerical weight of this item
    pub weight: uint,
    /// The actual item which is being weighted
    pub item: T,
}

/// A distribution that selects from a finite collection of weighted items.
///
/// Each item has an associated weight that influences how likely it
/// is to be chosen: higher weight is more likely.
///
/// The `Clone` restriction is a limitation of the `Sample` and
/// `IndependentSample` traits. Note that `&T` is (cheaply) `Clone` for
/// all `T`, as is `uint`, so one can store references or indices into
/// another vector.
///
/// # Example
///
/// ```rust
/// use std::rand;
/// use std::rand::distributions::{Weighted, WeightedChoice, IndependentSample};
///
/// let mut items = vec!(Weighted { weight: 2, item: 'a' },
///                      Weighted { weight: 4, item: 'b' },
///                      Weighted { weight: 1, item: 'c' });
/// let wc = WeightedChoice::new(items.as_mut_slice());
/// let mut rng = rand::task_rng();
/// for _ in range(0u, 16) {
///      // on average prints 'a' 4 times, 'b' 8 and 'c' twice.
///      println!("{}", wc.ind_sample(&mut rng));
/// }
/// ```
pub struct WeightedChoice<'a, T:'a> {
    items: &'a mut [Weighted<T>],
    weight_range: Range<uint>
}

impl<'a, T: Clone> WeightedChoice<'a, T> {
    /// Create a new `WeightedChoice`.
    ///
    /// Panics if:
    /// - `v` is empty
    /// - the total weight is 0
    /// - the total weight is larger than a `uint` can contain.
    pub fn new(items: &'a mut [Weighted<T>]) -> WeightedChoice<'a, T> {
        // strictly speaking, this is subsumed by the total weight == 0 case
        assert!(!items.is_empty(), "WeightedChoice::new called with no items");

        let mut running_total = 0u;

        // we convert the list from individual weights to cumulative
        // weights so we can binary search. This *could* drop elements
        // with weight == 0 as an optimisation.
        for item in items.iter_mut() {
            running_total = match running_total.checked_add(item.weight) {
                Some(n) => n,
                None => panic!("WeightedChoice::new called with a total weight \
                               larger than a uint can contain")
            };

            item.weight = running_total;
        }
        assert!(running_total != 0, "WeightedChoice::new called with a total weight of 0");

        WeightedChoice {
            items: items,
            // we're likely to be generating numbers in this range
            // relatively often, so might as well cache it
            weight_range: Range::new(0, running_total)
        }
    }
}

impl<'a, T: Clone> Sample<T> for WeightedChoice<'a, T> {
    fn sample<R: Rng>(&mut self, rng: &mut R) -> T { self.ind_sample(rng) }
}

impl<'a, T: Clone> IndependentSample<T> for WeightedChoice<'a, T> {
    fn ind_sample<R: Rng>(&self, rng: &mut R) -> T {
        // we want to find the first element that has cumulative
        // weight > sample_weight, which we do by binary since the
        // cumulative weights of self.items are sorted.

        // choose a weight in [0, total_weight)
        let sample_weight = self.weight_range.ind_sample(rng);

        // short circuit when it's the first item
        if sample_weight < self.items[0].weight {
            return self.items[0].item.clone();
        }

        let mut idx = 0;
        let mut modifier = self.items.len();

        // now we know that every possibility has an element to the
        // left, so we can just search for the last element that has
        // cumulative weight <= sample_weight, then the next one will
        // be "it". (Note that this greatest element will never be the
        // last element of the vector, since sample_weight is chosen
        // in [0, total_weight) and the cumulative weight of the last
        // one is exactly the total weight.)
        while modifier > 1 {
            let i = idx + modifier / 2;
            if self.items[i].weight <= sample_weight {
                // we're small, so look to the right, but allow this
                // exact element still.
                idx = i;
                // we need the `/ 2` to round up otherwise we'll drop
                // the trailing elements when `modifier` is odd.
                modifier += 1;
            } else {
                // otherwise we're too big, so go left. (i.e. do
                // nothing)
            }
            modifier /= 2;
        }
        return self.items[idx + 1].item.clone();
    }
}

mod ziggurat_tables;

/// Sample a random number using the Ziggurat method (specifically the
/// ZIGNOR variant from Doornik 2005). Most of the arguments are
/// directly from the paper:
///
/// * `rng`: source of randomness
/// * `symmetric`: whether this is a symmetric distribution, or one-sided with P(x < 0) = 0.
/// * `X`: the $x_i$ abscissae.
/// * `F`: precomputed values of the PDF at the $x_i$, (i.e. $f(x_i)$)
/// * `F_DIFF`: precomputed values of $f(x_i) - f(x_{i+1})$
/// * `pdf`: the probability density function
/// * `zero_case`: manual sampling from the tail when we chose the
///    bottom box (i.e. i == 0)

// the perf improvement (25-50%) is definitely worth the extra code
// size from force-inlining.
#[inline(always)]
fn ziggurat<R: Rng, P, Z>(
            rng: &mut R,
            symmetric: bool,
            x_tab: ziggurat_tables::ZigTable,
            f_tab: ziggurat_tables::ZigTable,
            mut pdf: P,
            mut zero_case: Z)
            -> f64 where P: FnMut(f64) -> f64, Z: FnMut(&mut R, f64) -> f64 {
    static SCALE: f64 = (1u64 << 53) as f64;
    loop {
        // reimplement the f64 generation as an optimisation suggested
        // by the Doornik paper: we have a lot of precision-space
        // (i.e. there are 11 bits of the 64 of a u64 to use after
        // creating a f64), so we might as well reuse some to save
        // generating a whole extra random number. (Seems to be 15%
        // faster.)
        //
        // This unfortunately misses out on the benefits of direct
        // floating point generation if an RNG like dSMFT is
        // used. (That is, such RNGs create floats directly, highly
        // efficiently and overload next_f32/f64, so by not calling it
        // this may be slower than it would be otherwise.)
        // FIXME: investigate/optimise for the above.
        let bits: u64 = rng.gen();
        let i = (bits & 0xff) as uint;
        let f = (bits >> 11) as f64 / SCALE;

        // u is either U(-1, 1) or U(0, 1) depending on if this is a
        // symmetric distribution or not.
        let u = if symmetric {2.0 * f - 1.0} else {f};
        let x = u * x_tab[i];

        let test_x = if symmetric { x.abs() } else {x};

        // algebraically equivalent to |u| < x_tab[i+1]/x_tab[i] (or u < x_tab[i+1]/x_tab[i])
        if test_x < x_tab[i + 1] {
            return x;
        }
        if i == 0 {
            return zero_case(rng, u);
        }
        // algebraically equivalent to f1 + DRanU()*(f0 - f1) < 1
        if f_tab[i + 1] + (f_tab[i] - f_tab[i + 1]) * rng.gen() < pdf(x) {
            return x;
        }
    }
}

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

    use {Rng, Rand};
    use super::{RandSample, WeightedChoice, Weighted, Sample, IndependentSample};

    #[deriving(PartialEq, Show)]
    struct ConstRand(uint);
    impl Rand for ConstRand {
        fn rand<R: Rng>(_: &mut R) -> ConstRand {
            ConstRand(0)
        }
    }

    // 0, 1, 2, 3, ...
    struct CountingRng { i: u32 }
    impl Rng for CountingRng {
        fn next_u32(&mut self) -> u32 {
            self.i += 1;
            self.i - 1
        }
        fn next_u64(&mut self) -> u64 {
            self.next_u32() as u64
        }
    }

    #[test]
    fn test_rand_sample() {
        let mut rand_sample = RandSample::<ConstRand>;

        assert_eq!(rand_sample.sample(&mut ::test::rng()), ConstRand(0));
        assert_eq!(rand_sample.ind_sample(&mut ::test::rng()), ConstRand(0));
    }
    #[test]
    fn test_weighted_choice() {
        // this makes assumptions about the internal implementation of
        // WeightedChoice, specifically: it doesn't reorder the items,
        // it doesn't do weird things to the RNG (so 0 maps to 0, 1 to
        // 1, internally; modulo a modulo operation).

        macro_rules! t (
            ($items:expr, $expected:expr) => {{
                let mut items = $items;
                let wc = WeightedChoice::new(items.as_mut_slice());
                let expected = $expected;

                let mut rng = CountingRng { i: 0 };

                for &val in expected.iter() {
                    assert_eq!(wc.ind_sample(&mut rng), val)
                }
            }}
        );

        t!(vec!(Weighted { weight: 1, item: 10i}), [10]);

        // skip some
        t!(vec!(Weighted { weight: 0, item: 20i},
                Weighted { weight: 2, item: 21i},
                Weighted { weight: 0, item: 22i},
                Weighted { weight: 1, item: 23i}),
           [21,21, 23]);

        // different weights
        t!(vec!(Weighted { weight: 4, item: 30i},
                Weighted { weight: 3, item: 31i}),
           [30,30,30,30, 31,31,31]);

        // check that we're binary searching
        // correctly with some vectors of odd
        // length.
        t!(vec!(Weighted { weight: 1, item: 40i},
                Weighted { weight: 1, item: 41i},
                Weighted { weight: 1, item: 42i},
                Weighted { weight: 1, item: 43i},
                Weighted { weight: 1, item: 44i}),
           [40, 41, 42, 43, 44]);
        t!(vec!(Weighted { weight: 1, item: 50i},
                Weighted { weight: 1, item: 51i},
                Weighted { weight: 1, item: 52i},
                Weighted { weight: 1, item: 53i},
                Weighted { weight: 1, item: 54i},
                Weighted { weight: 1, item: 55i},
                Weighted { weight: 1, item: 56i}),
           [50, 51, 52, 53, 54, 55, 56]);
    }

    #[test] #[should_fail]
    fn test_weighted_choice_no_items() {
        WeightedChoice::<int>::new(&mut []);
    }
    #[test] #[should_fail]
    fn test_weighted_choice_zero_weight() {
        WeightedChoice::new(&mut [Weighted { weight: 0, item: 0i},
                                  Weighted { weight: 0, item: 1i}]);
    }
    #[test] #[should_fail]
    fn test_weighted_choice_weight_overflows() {
        let x = (-1) as uint / 2; // x + x + 2 is the overflow
        WeightedChoice::new(&mut [Weighted { weight: x, item: 0i },
                                  Weighted { weight: 1, item: 1i },
                                  Weighted { weight: x, item: 2i },
                                  Weighted { weight: 1, item: 3i }]);
    }
}
<|endoftext|>


<commit_before>/// low level wrapper for the state machine, encoding and decoding from lapin-async
use lapin_async::connection::*;
use lapin_async::format::frame::*;

use nom::{IResult,Offset};
use cookie_factory::GenError;
use bytes::BytesMut;
use std::iter::repeat;
use std::io::{self,Error,ErrorKind};
use std::time::Duration;
use futures::{Async,Poll,Sink,Stream,StartSend,Future};
use tokio_io::{AsyncRead,AsyncWrite};
use tokio_io::codec::{Decoder,Encoder,Framed};
use tokio_timer::{Interval,Timer};
use client::ConnectionOptions;

/// implements tokio-io's Decoder and Encoder
pub struct AMQPCodec;

impl Decoder for AMQPCodec {
    type Item = Frame;
    type Error = io::Error;

    fn decode(&mut self, buf: &mut BytesMut) -> Result<Option<Frame>, io::Error> {
        let (consumed, f) = match frame(buf) {
          IResult::Incomplete(_) => {
            return Ok(None)
          },
          IResult::Error(e) => {
            return Err(io::Error::new(io::ErrorKind::Other, format!("parse error: {:?}", e)))
          },
          IResult::Done(i, frame) => {
            (buf.offset(i), frame)
          }
        };

        trace!("decoded frame: {:?}", f);

        buf.split_to(consumed);

        Ok(Some(f))
    }
}

impl Encoder for AMQPCodec {
    type Item = Frame;
    type Error = io::Error;

    fn encode(&mut self, frame: Frame, buf: &mut BytesMut) -> Result<(), Self::Error> {
      let length = buf.len();
      if length < 8192 {
        //reserve more capacity and intialize it
        buf.extend(repeat(0).take(8192 - length));
      }
      trace!("will encode and write frame: {:?}", frame);

      loop {
        let gen_res = match &frame {
          &Frame::ProtocolHeader => {
            gen_protocol_header((buf, 0)).map(|tup| tup.1)
          },
          &Frame::Heartbeat(_) => {
            gen_heartbeat_frame((buf, 0)).map(|tup| tup.1)
          },
          &Frame::Method(channel, ref method) => {
            gen_method_frame((buf, 0), channel, method).map(|tup| tup.1)
          },
          &Frame::Header(channel_id, class_id, ref header) => {
            gen_content_header_frame((buf, 0), channel_id, class_id, header.body_size, &header.properties).map(|tup| tup.1)
          },
          &Frame::Body(channel_id, ref data) => {
            gen_content_body_frame((buf, 0), channel_id, data).map(|tup| tup.1)
          }
        };

        match gen_res {
          Ok(sz) => {
            buf.truncate(sz);
            trace!("serialized frame: {} bytes", sz);
            return Ok(());
          },
          Err(e) => {
            error!("error generating frame: {:?}", e);
            match e {
              GenError::BufferTooSmall(sz) => {
                buf.extend(repeat(0).take(sz - length));
                //return Err(Error::new(ErrorKind::InvalidData, "send buffer too small"));
              },
              GenError::InvalidOffset | GenError::CustomError(_) | GenError::NotYetImplemented => {
                return Err(Error::new(ErrorKind::InvalidData, "could not generate"));
              }
            }
          }
        }
      }
    }
}

/// Wrappers over a `Framed` stream using `AMQPCodec` and lapin-async's `Connection`
pub struct AMQPTransport<T> {
  pub upstream: Framed<T,AMQPCodec>,
  pub heartbeat: Interval,
  pub conn: Connection,
}

impl<T> AMQPTransport<T>
   where T: AsyncRead+AsyncWrite,
         T: 'static               {

  /// starts the connection process
  ///
  /// returns a future of a `AMQPTransport` that is connected
  pub fn connect(upstream: Framed<T,AMQPCodec>, options: &ConnectionOptions) -> Box<Future<Item = AMQPTransport<T>, Error = io::Error>> {
    let mut conn = Connection::new();
    conn.set_credentials(&options.username, &options.password);
    conn.set_vhost(&options.vhost);
    conn.set_heartbeat(options.heartbeat);
    conn.connect();

    let mut t = AMQPTransport {
      upstream:  upstream,
      heartbeat: Timer::default().interval(Duration::from_secs(conn.configuration.heartbeat as u64)),
      conn:      conn,
    };

    t.send_and_handle_frames();
    t.poll_children();

    let mut connector = AMQPTransportConnector {
      transport: Some(t)
    };

    trace!("pre-poll");
    connector.poll();
    trace!("post-poll");

    Box::new(connector)
  }

  fn poll_heartbeat(&mut self) {
    if let Ok(Async::Ready(_)) = self.heartbeat.poll() {
      debug!("Heartbeat");
      self.start_send(Frame::Heartbeat(0));
      self.poll_complete();
    }
  }

  fn poll_upstream(&mut self) -> Poll<Option<()>, io::Error> {
    let value = match self.upstream.poll() {
      Ok(Async::Ready(t)) => t,
      Ok(Async::NotReady) => {
        trace!("upstream poll gave NotReady");
        return Ok(Async::NotReady);
      },
      Err(e) => {
        error!("upstream poll gave error: {:?}", e);
        return Err(From::from(e));
      },
    };

    if let Some(frame) = value {
      trace!("upstream poll gave frame: {:?}", frame);
      self.conn.handle_frame(frame);
      self.send_frames();
      self.upstream.poll_complete();
      Ok(Async::Ready(Some(())))
    } else {
      error!("upstream poll gave Ready(None)");
      Ok(Async::Ready(None))
    }
  }

  // poll heartbeat and upstream and return what upstream gave us
  pub fn poll_children(&mut self) -> Poll<Option<()>, io::Error> {
    self.poll_heartbeat();
    self.poll_upstream()
  }

  pub fn send_and_handle_frames(&mut self) {
    self.send_frames();
    self.handle_frames();
  }

  pub fn send_frames(&mut self) {
    //FIXME: find a way to use a future here
    while let Some(f) = self.conn.next_frame() {
      self.upstream.start_send(f);
      self.upstream.poll_complete();
    }
    //self.upstream.poll_complete();
  }

  pub fn handle_frames(&mut self) {
    while let Ok(Async::Ready(Some(_))) = self.poll() {}
  }
}

/// implements a future of `AMQPTransport`
///
/// this structure is used to perform the AMQP handshake and provide
/// a connected transport afterwards
pub struct AMQPTransportConnector<T> {
  pub transport: Option<AMQPTransport<T>>,
}

impl<T> Future for AMQPTransportConnector<T>
    where T: AsyncRead + AsyncWrite,
          T: 'static {

  type Item  = AMQPTransport<T>;
  type Error = io::Error;

  fn poll(&mut self) -> Poll<Self::Item, Self::Error> {
    debug!("AMQPTransportConnector poll transport is none? {}", self.transport.is_none());
    let mut transport = self.transport.take().unwrap();

    //we might have received a frame before here
    transport.handle_frames();
    transport.send_frames();

    debug!("conn state: {:?}", transport.conn.state);
    if transport.conn.state == ConnectionState::Connected {
      debug!("already connected");
      return Ok(Async::Ready(transport))
    }

    trace!("waiting before poll");
    transport.poll_children().map(|value| {
      match value {
        Async::Ready(Some(_)) => {
          if transport.conn.state == ConnectionState::Connected {
            // Upstream had frames available and we're connected, the transport is ready
            Async::Ready(transport)
          } else {
            // Upstream had frames but we're not yet connected, continue polling
            transport.poll_children();
            self.transport = Some(transport);
            Async::NotReady
          }
        },
        _ => {
          // Upstream had no frames
          self.transport = Some(transport);
          Async::NotReady
        },
      }
    })
  }
}

impl<T> Stream for AMQPTransport<T>
    where T: AsyncRead + AsyncWrite,
          T: 'static {
    type Item = ();
    type Error = io::Error;

    fn poll(&mut self) -> Poll<Option<()>, io::Error> {
      trace!("stream poll");
      self.poll_children()
    }
}

impl<T> Sink for AMQPTransport<T>
    where T: AsyncWrite {
    type SinkItem = Frame;
    type SinkError = io::Error;

    fn start_send(&mut self, item: Frame) -> StartSend<Frame, io::Error> {
        trace!("sink start send");
        self.upstream.start_send(item)
    }

    fn poll_complete(&mut self) -> Poll<(), io::Error> {
        trace!("sink poll_complete");
        self.upstream.poll_complete()
    }
}

<commit_msg>futures: use send_and_handle_frames<commit_after>/// low level wrapper for the state machine, encoding and decoding from lapin-async
use lapin_async::connection::*;
use lapin_async::format::frame::*;

use nom::{IResult,Offset};
use cookie_factory::GenError;
use bytes::BytesMut;
use std::iter::repeat;
use std::io::{self,Error,ErrorKind};
use std::time::Duration;
use futures::{Async,Poll,Sink,Stream,StartSend,Future};
use tokio_io::{AsyncRead,AsyncWrite};
use tokio_io::codec::{Decoder,Encoder,Framed};
use tokio_timer::{Interval,Timer};
use client::ConnectionOptions;

/// implements tokio-io's Decoder and Encoder
pub struct AMQPCodec;

impl Decoder for AMQPCodec {
    type Item = Frame;
    type Error = io::Error;

    fn decode(&mut self, buf: &mut BytesMut) -> Result<Option<Frame>, io::Error> {
        let (consumed, f) = match frame(buf) {
          IResult::Incomplete(_) => {
            return Ok(None)
          },
          IResult::Error(e) => {
            return Err(io::Error::new(io::ErrorKind::Other, format!("parse error: {:?}", e)))
          },
          IResult::Done(i, frame) => {
            (buf.offset(i), frame)
          }
        };

        trace!("decoded frame: {:?}", f);

        buf.split_to(consumed);

        Ok(Some(f))
    }
}

impl Encoder for AMQPCodec {
    type Item = Frame;
    type Error = io::Error;

    fn encode(&mut self, frame: Frame, buf: &mut BytesMut) -> Result<(), Self::Error> {
      let length = buf.len();
      if length < 8192 {
        //reserve more capacity and intialize it
        buf.extend(repeat(0).take(8192 - length));
      }
      trace!("will encode and write frame: {:?}", frame);

      loop {
        let gen_res = match &frame {
          &Frame::ProtocolHeader => {
            gen_protocol_header((buf, 0)).map(|tup| tup.1)
          },
          &Frame::Heartbeat(_) => {
            gen_heartbeat_frame((buf, 0)).map(|tup| tup.1)
          },
          &Frame::Method(channel, ref method) => {
            gen_method_frame((buf, 0), channel, method).map(|tup| tup.1)
          },
          &Frame::Header(channel_id, class_id, ref header) => {
            gen_content_header_frame((buf, 0), channel_id, class_id, header.body_size, &header.properties).map(|tup| tup.1)
          },
          &Frame::Body(channel_id, ref data) => {
            gen_content_body_frame((buf, 0), channel_id, data).map(|tup| tup.1)
          }
        };

        match gen_res {
          Ok(sz) => {
            buf.truncate(sz);
            trace!("serialized frame: {} bytes", sz);
            return Ok(());
          },
          Err(e) => {
            error!("error generating frame: {:?}", e);
            match e {
              GenError::BufferTooSmall(sz) => {
                buf.extend(repeat(0).take(sz - length));
                //return Err(Error::new(ErrorKind::InvalidData, "send buffer too small"));
              },
              GenError::InvalidOffset | GenError::CustomError(_) | GenError::NotYetImplemented => {
                return Err(Error::new(ErrorKind::InvalidData, "could not generate"));
              }
            }
          }
        }
      }
    }
}

/// Wrappers over a `Framed` stream using `AMQPCodec` and lapin-async's `Connection`
pub struct AMQPTransport<T> {
  pub upstream: Framed<T,AMQPCodec>,
  pub heartbeat: Interval,
  pub conn: Connection,
}

impl<T> AMQPTransport<T>
   where T: AsyncRead+AsyncWrite,
         T: 'static               {

  /// starts the connection process
  ///
  /// returns a future of a `AMQPTransport` that is connected
  pub fn connect(upstream: Framed<T,AMQPCodec>, options: &ConnectionOptions) -> Box<Future<Item = AMQPTransport<T>, Error = io::Error>> {
    let mut conn = Connection::new();
    conn.set_credentials(&options.username, &options.password);
    conn.set_vhost(&options.vhost);
    conn.set_heartbeat(options.heartbeat);
    conn.connect();

    let mut t = AMQPTransport {
      upstream:  upstream,
      heartbeat: Timer::default().interval(Duration::from_secs(conn.configuration.heartbeat as u64)),
      conn:      conn,
    };

    t.send_and_handle_frames();
    t.poll_children();

    let mut connector = AMQPTransportConnector {
      transport: Some(t)
    };

    trace!("pre-poll");
    connector.poll();
    trace!("post-poll");

    Box::new(connector)
  }

  fn poll_heartbeat(&mut self) {
    if let Ok(Async::Ready(_)) = self.heartbeat.poll() {
      debug!("Heartbeat");
      self.start_send(Frame::Heartbeat(0));
      self.poll_complete();
    }
  }

  fn poll_upstream(&mut self) -> Poll<Option<()>, io::Error> {
    let value = match self.upstream.poll() {
      Ok(Async::Ready(t)) => t,
      Ok(Async::NotReady) => {
        trace!("upstream poll gave NotReady");
        return Ok(Async::NotReady);
      },
      Err(e) => {
        error!("upstream poll gave error: {:?}", e);
        return Err(From::from(e));
      },
    };

    if let Some(frame) = value {
      trace!("upstream poll gave frame: {:?}", frame);
      self.conn.handle_frame(frame);
      self.send_frames();
      self.upstream.poll_complete();
      Ok(Async::Ready(Some(())))
    } else {
      error!("upstream poll gave Ready(None)");
      Ok(Async::Ready(None))
    }
  }

  // poll heartbeat and upstream and return what upstream gave us
  pub fn poll_children(&mut self) -> Poll<Option<()>, io::Error> {
    self.poll_heartbeat();
    self.poll_upstream()
  }

  pub fn send_and_handle_frames(&mut self) {
    self.send_frames();
    self.handle_frames();
  }

  pub fn send_frames(&mut self) {
    //FIXME: find a way to use a future here
    while let Some(f) = self.conn.next_frame() {
      self.upstream.start_send(f);
      self.upstream.poll_complete();
    }
    //self.upstream.poll_complete();
  }

  pub fn handle_frames(&mut self) {
    while let Ok(Async::Ready(Some(_))) = self.poll() {}
  }
}

/// implements a future of `AMQPTransport`
///
/// this structure is used to perform the AMQP handshake and provide
/// a connected transport afterwards
pub struct AMQPTransportConnector<T> {
  pub transport: Option<AMQPTransport<T>>,
}

impl<T> Future for AMQPTransportConnector<T>
    where T: AsyncRead + AsyncWrite,
          T: 'static {

  type Item  = AMQPTransport<T>;
  type Error = io::Error;

  fn poll(&mut self) -> Poll<Self::Item, Self::Error> {
    debug!("AMQPTransportConnector poll transport is none? {}", self.transport.is_none());
    let mut transport = self.transport.take().unwrap();

    //we might have received a frame before here
    transport.send_and_handle_frames();

    debug!("conn state: {:?}", transport.conn.state);
    if transport.conn.state == ConnectionState::Connected {
      debug!("already connected");
      return Ok(Async::Ready(transport))
    }

    trace!("waiting before poll");
    transport.poll_children().map(|value| {
      match value {
        Async::Ready(Some(_)) => {
          if transport.conn.state == ConnectionState::Connected {
            // Upstream had frames available and we're connected, the transport is ready
            Async::Ready(transport)
          } else {
            // Upstream had frames but we're not yet connected, continue polling
            transport.poll_children();
            self.transport = Some(transport);
            Async::NotReady
          }
        },
        _ => {
          // Upstream had no frames
          self.transport = Some(transport);
          Async::NotReady
        },
      }
    })
  }
}

impl<T> Stream for AMQPTransport<T>
    where T: AsyncRead + AsyncWrite,
          T: 'static {
    type Item = ();
    type Error = io::Error;

    fn poll(&mut self) -> Poll<Option<()>, io::Error> {
      trace!("stream poll");
      self.poll_children()
    }
}

impl<T> Sink for AMQPTransport<T>
    where T: AsyncWrite {
    type SinkItem = Frame;
    type SinkError = io::Error;

    fn start_send(&mut self, item: Frame) -> StartSend<Frame, io::Error> {
        trace!("sink start send");
        self.upstream.start_send(item)
    }

    fn poll_complete(&mut self) -> Poll<(), io::Error> {
        trace!("sink poll_complete");
        self.upstream.poll_complete()
    }
}

<|endoftext|>