<commit_before>// Copyright (c) 2020 DDN. All rights reserved.
// Use of this source code is governed by a MIT-style
// license that can be found in the LICENSE file.

use futures::TryStreamExt;
use iml_journal::{execute_handlers, get_message_class, ImlJournalError};
use iml_manager_env::get_pool_limit;
use iml_postgres::{
    get_db_pool,
    sqlx::{self, Done, PgPool},
};
use iml_service_queue::service_queue::consume_data;
use iml_tracing::tracing;
use iml_wire_types::JournalMessage;
use lazy_static::lazy_static;
use std::convert::TryInto;

lazy_static! {
    static ref DBLOG_HW: i64 = iml_manager_env::get_dblog_hw() as i64;
    static ref DBLOG_LW: i64 = iml_manager_env::get_dblog_lw() as i64;
}

// Default pool limit if not overridden by POOL_LIMIT
const DEFAULT_POOL_LIMIT: u32 = 2;

async fn purge_excess(pool: &PgPool, num_rows: i64) -> Result<i64, ImlJournalError> {
    if num_rows <= *DBLOG_HW {
        return Ok(num_rows);
    }

    let mut num_remove: i64 = num_rows - *DBLOG_LW;

    while num_remove > 0 {
        let x = sqlx::query!(
            r#"
                DELETE FROM chroma_core_logmessage
                WHERE id in (
                    SELECT id FROM chroma_core_logmessage ORDER BY id LIMIT $1
                )
            "#,
            std::cmp::min(10_000, num_remove)
        )
        .execute(pool)
        .await?
        .rows_affected();

        num_remove -= x as i64;

        tracing::info!("Purged {} rows, {} remain for purging", x, num_remove);
    }

    Ok(num_rows)
}

#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
    iml_tracing::init();

    tracing::info!("Starting");

    let pool = get_db_pool(get_pool_limit().unwrap_or(DEFAULT_POOL_LIMIT)).await?;

    let rabbit_pool = iml_rabbit::connect_to_rabbit(1);

    let conn = iml_rabbit::get_conn(rabbit_pool).await?;

    let ch = iml_rabbit::create_channel(&conn).await?;

    let mut s = consume_data::<Vec<JournalMessage>>(&ch, "rust_agent_journal_rx");

    let mut num_rows = sqlx::query!("SELECT COUNT(*) FROM chroma_core_logmessage")
        .fetch_one(&pool)
        .await?
        .count
        .expect("count returned None");

    while let Some((host, xs)) = s.try_next().await? {
        num_rows = purge_excess(&pool, num_rows).await?;

        struct Row {
            id: i32,
            content_type_id: Option<i32>,
        }

        let row = sqlx::query_as!(Row,
            "select id, content_type_id from chroma_core_managedhost where fqdn = $1 and not_deleted = 't'",
            host.to_string()
        )
        .fetch_optional(&pool)
        .await?;

        let row = match row {
            Some(row) => row,
            None => {
                tracing::warn!("Host '{}' is unknown", host);

                continue;
            }
        };

        let content_type_id = match row.content_type_id.as_ref() {
            Some(x) => *x,
            None => {
                tracing::warn!("Host '{}' content_type_id  is unknown", host);
                continue;
            }
        };

        for x in xs.iter() {
            execute_handlers(&x.message, row.id, content_type_id, &pool).await?;
        }

        num_rows += xs.len() as i64;

        let x = xs.into_iter().fold(
            (vec![], vec![], vec![], vec![], vec![], vec![]),
            |mut acc, x| {
                acc.0.push(x.datetime);
                acc.1.push(x.severity as i16);
                acc.2.push(x.facility);
                acc.3.push(x.source);
                acc.5.push(get_message_class(&x.message) as i16);
                acc.4.push(x.message);

                acc
            },
        );

        let times =
            x.0.into_iter()
                .map(|t| {
                    let t = t.as_secs().try_into()?;

                    Ok(time::OffsetDateTime::from_unix_timestamp(t))
                })
                .collect::<Result<Vec<time::OffsetDateTime>, ImlJournalError>>()?;

        sqlx::query!(
            r#"
             INSERT INTO chroma_core_logmessage
             (datetime, fqdn, severity, facility, tag, message, message_class)
             SELECT datetime, $2, severity, facility, source, message, message_class
             FROM UNNEST($1::timestamptz[], $3::smallint[], $4::smallint[], $5::text[], $6::text[], $7::smallint[])
             AS t(datetime, severity, facility, source, message, message_class)
         "#,
            &times,
            host.to_string(),
            &x.1,
            &x.2,
            &x.3,
            &x.4,
            &x.5
        )
            .execute(&pool)
            .await?;
    }

    Ok(())
}
<commit_msg>Purge loop (#2214)<commit_after>// Copyright (c) 2020 DDN. All rights reserved.
// Use of this source code is governed by a MIT-style
// license that can be found in the LICENSE file.

use futures::TryStreamExt;
use iml_journal::{execute_handlers, get_message_class, ImlJournalError};
use iml_manager_env::get_pool_limit;
use iml_postgres::{
    get_db_pool,
    sqlx::{self, Done, PgPool},
};
use iml_service_queue::service_queue::consume_data;
use iml_tracing::tracing;
use iml_wire_types::JournalMessage;
use lazy_static::lazy_static;
use std::convert::TryInto;

lazy_static! {
    static ref DBLOG_HW: i64 = iml_manager_env::get_dblog_hw() as i64;
    static ref DBLOG_LW: i64 = iml_manager_env::get_dblog_lw() as i64;
}

// Default pool limit if not overridden by POOL_LIMIT
const DEFAULT_POOL_LIMIT: u32 = 2;

async fn purge_excess(pool: &PgPool, mut num_rows: i64) -> Result<i64, ImlJournalError> {
    if num_rows <= *DBLOG_HW {
        return Ok(num_rows);
    }

    while *DBLOG_LW < num_rows {
        let x = sqlx::query!(
            r#"
                DELETE FROM chroma_core_logmessage
                WHERE id in (
                    SELECT id FROM chroma_core_logmessage ORDER BY id LIMIT $1
                )
            "#,
            std::cmp::min(10_000, num_rows - *DBLOG_LW)
        )
        .execute(pool)
        .await?
        .rows_affected();

        num_rows -= x as i64;
        tracing::info!("Purged {} rows, current known row count is {}", x, num_rows);
    }

    Ok(num_rows)
}

#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
    iml_tracing::init();

    tracing::info!("Starting");

    let pool = get_db_pool(get_pool_limit().unwrap_or(DEFAULT_POOL_LIMIT)).await?;

    let rabbit_pool = iml_rabbit::connect_to_rabbit(1);

    let conn = iml_rabbit::get_conn(rabbit_pool).await?;

    let ch = iml_rabbit::create_channel(&conn).await?;

    let mut s = consume_data::<Vec<JournalMessage>>(&ch, "rust_agent_journal_rx");

    let mut num_rows = sqlx::query!("SELECT COUNT(*) FROM chroma_core_logmessage")
        .fetch_one(&pool)
        .await?
        .count
        .expect("count returned None");

    while let Some((host, xs)) = s.try_next().await? {
        num_rows = purge_excess(&pool, num_rows).await?;

        struct Row {
            id: i32,
            content_type_id: Option<i32>,
        }

        let row = sqlx::query_as!(Row,
            "select id, content_type_id from chroma_core_managedhost where fqdn = $1 and not_deleted = 't'",
            host.to_string()
        )
        .fetch_optional(&pool)
        .await?;

        let row = match row {
            Some(row) => row,
            None => {
                tracing::warn!("Host '{}' is unknown", host);

                continue;
            }
        };

        let content_type_id = match row.content_type_id.as_ref() {
            Some(x) => *x,
            None => {
                tracing::warn!("Host '{}' content_type_id  is unknown", host);
                continue;
            }
        };

        for x in xs.iter() {
            execute_handlers(&x.message, row.id, content_type_id, &pool).await?;
        }

        num_rows += xs.len() as i64;

        let x = xs.into_iter().fold(
            (vec![], vec![], vec![], vec![], vec![], vec![]),
            |mut acc, x| {
                acc.0.push(x.datetime);
                acc.1.push(x.severity as i16);
                acc.2.push(x.facility);
                acc.3.push(x.source);
                acc.5.push(get_message_class(&x.message) as i16);
                acc.4.push(x.message);

                acc
            },
        );

        let times =
            x.0.into_iter()
                .map(|t| {
                    let t = t.as_secs().try_into()?;

                    Ok(time::OffsetDateTime::from_unix_timestamp(t))
                })
                .collect::<Result<Vec<time::OffsetDateTime>, ImlJournalError>>()?;

        sqlx::query!(
            r#"
             INSERT INTO chroma_core_logmessage
             (datetime, fqdn, severity, facility, tag, message, message_class)
             SELECT datetime, $2, severity, facility, source, message, message_class
             FROM UNNEST($1::timestamptz[], $3::smallint[], $4::smallint[], $5::text[], $6::text[], $7::smallint[])
             AS t(datetime, severity, facility, source, message, message_class)
         "#,
            &times,
            host.to_string(),
            &x.1,
            &x.2,
            &x.3,
            &x.4,
            &x.5
        )
        .execute(&pool)
        .await?;
    }

    Ok(())
}
<|endoftext|>


<commit_before>//
// Copyright (C) 2017 Kubos Corporation
//
// 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.
//

extern crate getopts;

use getopts::Options;
use std::process::{exit, Command, Stdio};
use std::env;

/// Take a kubos target and convert it
/// to a Rust/Clang target triplet
fn target_converter(kubos_target: &str) -> String {
    match kubos_target {
        "x86-linux-native" => String::from("x86_64-unknown-linux-gnu"),
        "kubos-linux-beaglebone-gcc" => String::from("arm-unknown-linux-gnueabihf"),
        "kubos-linux-pumpkin-mbm2-gcc" => String::from("arm-unknown-linux-gnueabihf"),
        "kubos-linux-isis-gcc" => String::from("arm-unknown-linux-gnueabi"),
        _ => {
            panic!(
                "Target '{}' not supported for cargo/yotta builds\
                 \nCurrently supported targets are:\
                 \nx86-linux-native\
                 \nkubos-linux-beaglebone-gcc\
                 \nkubos-linux-pumpkin-mbm2-gcc\
                 \nkubos-linux-isis-gcc",
                kubos_target
            )
        }
    }
}

/// Perform `cargo 'command'` using the proper Rust/Clang target triplet
fn cargo_command(target: String, command: String, mut extra_params: Vec<String>) {
    let mut params = vec![command, String::from("--target"), target];
    params.append(&mut extra_params);

    let status = Command::new("cargo")
        .args(&params)
        .stdin(Stdio::inherit())
        .stdout(Stdio::inherit())
        .stderr(Stdio::inherit())
        .status()
        .unwrap();

    // Attempt to exit in a way which
    // honors the subprocess exit code
    if status.success() {
        exit(1)
    }
    exit(status.code().unwrap());
}

/// Displays usage message
fn print_usage(opts: Options) {
    let brief = "cargo-kubos is a helper utility for running \
        Cargo commands with a Kubos target attached.\nIt is \
        used when building/running/testing crates which either\
        contain a yotta module or depend on one. \
        \n\nUsage:\
        \n\tcargo kubos -c [cargo command] [options] -- [cargo options]
        \n\tcargo kubos -c build -t x86-linux-native -- -vv";
    print!("{}", opts.usage(&brief));
}

fn main() {
    let args: Vec<String> = env::args().collect();
    let mut opts = Options::new();

    opts.reqopt("c", "command", "cargo command to run", "COMMAND");
    opts.optopt("t", "target", "sets (Kubos) target", "NAME");
    opts.optflag("h", "help", "Displays help");

    let matches = match opts.parse(&args[1..]) {
        Ok(m) => m,
        Err(f) => {
            println!("Error - {}\n", f);
            print_usage(opts);
            return;
        }
    };

    // Collect extra parametrs
    let extra_params = if !matches.free.is_empty() {
        let mut params = matches.free.clone();
        // remove extra kubos pamatere
        params.retain(|x| x != "kubos");
        params
    } else {
        Vec::new()
    };

    if matches.opt_present("h") {
        print_usage(opts);
    } else {
        let k_target = match matches.opt_str("t") {
            Some(t) => t,
            None => String::from("x86-linux-native"),
        };
        let command = matches.opt_str("c").unwrap();
        let c_target = target_converter(&k_target);
        env::set_var("CARGO_KUBOS_TARGET", &k_target);
        cargo_command(c_target, command, extra_params);
    }
}
<commit_msg>Cleaned up duplicate string.<commit_after>//
// Copyright (C) 2017 Kubos Corporation
//
// 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.
//

extern crate getopts;

use getopts::Options;
use std::process::{exit, Command, Stdio};
use std::env;

const X86_TARGET_STR: &'static str = "x86-linux-native";

/// Take a kubos target and convert it
/// to a Rust/Clang target triplet
fn target_converter(kubos_target: &str) -> String {
    match kubos_target {
        X86_TARGET_STR => String::from("x86_64-unknown-linux-gnu"),
        "kubos-linux-beaglebone-gcc" => String::from("arm-unknown-linux-gnueabihf"),
        "kubos-linux-pumpkin-mbm2-gcc" => String::from("arm-unknown-linux-gnueabihf"),
        "kubos-linux-isis-gcc" => String::from("arm-unknown-linux-gnueabi"),
        _ => {
            panic!(
                "Target '{}' not supported for cargo/yotta builds\
                 \nCurrently supported targets are:\
                 \nx86-linux-native\
                 \nkubos-linux-beaglebone-gcc\
                 \nkubos-linux-pumpkin-mbm2-gcc\
                 \nkubos-linux-isis-gcc",
                kubos_target
            )
        }
    }
}

/// Perform `cargo 'command'` using the proper Rust/Clang target triplet
fn cargo_command(target: String, command: String, mut extra_params: Vec<String>) {
    let mut params = vec![command, String::from("--target"), target];
    params.append(&mut extra_params);

    let status = Command::new("cargo")
        .args(&params)
        .stdin(Stdio::inherit())
        .stdout(Stdio::inherit())
        .stderr(Stdio::inherit())
        .status()
        .unwrap();

    // Attempt to exit in a way which
    // honors the subprocess exit code
    if status.success() {
        exit(1)
    }
    exit(status.code().unwrap());
}

/// Displays usage message
fn print_usage(opts: Options) {
    let brief = "cargo-kubos is a helper utility for running \
        Cargo commands with a Kubos target attached.\nIt is \
        used when building/running/testing crates which either\
        contain a yotta module or depend on one. \
        \n\nUsage:\
        \n\tcargo kubos -c [cargo command] [options] -- [cargo options]
        \n\tcargo kubos -c build -t x86-linux-native -- -vv";
    print!("{}", opts.usage(&brief));
}

fn main() {
    let args: Vec<String> = env::args().collect();
    let mut opts = Options::new();

    opts.reqopt("c", "command", "cargo command to run", "COMMAND");
    opts.optopt("t", "target", "sets (Kubos) target", "NAME");
    opts.optflag("h", "help", "Displays help");

    let matches = match opts.parse(&args[1..]) {
        Ok(m) => m,
        Err(f) => {
            println!("Error - {}\n", f);
            print_usage(opts);
            return;
        }
    };

    // Collect extra parametrs
    let extra_params = if !matches.free.is_empty() {
        let mut params = matches.free.clone();
        // remove extra kubos pamatere
        params.retain(|x| x != "kubos");
        params
    } else {
        Vec::new()
    };

    if matches.opt_present("h") {
        print_usage(opts);
    } else {
        let k_target = match matches.opt_str("t") {
            Some(t) => t,
            None => String::from(X86_TARGET_STR),
        };
        let command = matches.opt_str("c").unwrap();
        let c_target = target_converter(&k_target);
        env::set_var("CARGO_KUBOS_TARGET", &k_target);
        cargo_command(c_target, command, extra_params);
    }
}
<|endoftext|>


<commit_before>//! Gaussian Processes
//!
//! Provides implementation of gaussian process regression.
//!
//! # Usage
//!
//! ```
//! use rusty_machine::learning::gp;
//! use rusty_machine::learning::SupModel;
//! use rusty_machine::linalg::matrix::Matrix;
//! use rusty_machine::linalg::vector::Vector;
//!
//! let mut gaussp = gp::GaussianProcess::default();
//! gaussp.noise = 10f64;
//!
//! let train_data = Matrix::new(10,1,vec![0.,1.,2.,3.,4.,5.,6.,7.,8.,9.]);
//! let target = Vector::new(vec![0.,1.,2.,3.,4.,4.,3.,2.,1.,0.]);
//!
//! gaussp.train(&train_data, &target);
//!
//! let test_data = Matrix::new(5,1,vec![2.3,4.4,5.1,6.2,7.1]);
//!
//! let output = gaussp.predict(&test_data);
//! ```
//! Alternatively one could use gaussp.get_posterior() which would return both
//! the predictive mean and covariance. However, this is likely to change in
//! a future release.

use learning::toolkit::kernel::{Kernel, SquaredExp};
use learning::SupModel;
use linalg::matrix::Matrix;
use linalg::vector::Vector;

/// Trait for GP mean functions.
pub trait MeanFunc {
    /// Compute the mean function applied elementwise to a matrix.
    fn func(&self, x: Matrix<f64>) -> Vector<f64>;
}

/// Constant mean function
pub struct ConstMean {
    a: f64,
}

impl Default for ConstMean {
    /// Constructs the zero function.
    ///
    /// # Examples
    /// ```
    /// use rusty_machine::learning::gp::ConstMean;
    ///
    /// let zero_m = ConstMean::default();
    /// ```
    fn default() -> ConstMean {
        ConstMean { a: 0f64 }
    }
}

impl MeanFunc for ConstMean {
    fn func(&self, x: Matrix<f64>) -> Vector<f64> {
        Vector::zeros(x.rows()) + self.a
    }
}

/// Gaussian Process struct
///
/// Gaussian process with generic kernel and deterministic mean function.
/// Can be used for gaussian process regression with noise.
/// Currently does not support classification.
pub struct GaussianProcess<T: Kernel, U: MeanFunc> {
    ker: T,
    mean: U,
    /// The observation noise of the GP.
    pub noise: f64,
    alpha: Option<Vector<f64>>,
    train_mat: Option<Matrix<f64>>,
    train_data: Option<Matrix<f64>>,
}

impl Default for GaussianProcess<SquaredExp, ConstMean> {
    /// Construct a default Gaussian Process
    ///
    /// # Examples
    ///
    /// ```
    /// use rusty_machine::learning::gp;
    ///
    /// let gp = gp::GaussianProcess::default();
    /// ```
    fn default() -> GaussianProcess<SquaredExp, ConstMean> {
        GaussianProcess {
            ker: SquaredExp::default(),
            mean: ConstMean::default(),
            noise: 0f64,
            train_mat: None,
            train_data: None,
            alpha: None,
        }
    }
}

impl<T: Kernel, U: MeanFunc> GaussianProcess<T, U> {
    /// Construct a new Gaussian Process.
    ///
    /// # Examples
    ///
    /// ```
    /// use rusty_machine::learning::gp;
    /// use rusty_machine::learning::toolkit::kernel;
    ///
    /// let ker = kernel::SquaredExp::default();
    /// let mean = gp::ConstMean::default();
    /// let gaussp = gp::GaussianProcess::new(ker, mean, 1e-3f64);
    /// ```
    pub fn new(ker: T, mean: U, noise: f64) -> GaussianProcess<T, U> {
        GaussianProcess {
            ker: ker,
            mean: mean,
            noise: noise,
            train_mat: None,
            train_data: None,
            alpha: None,
        }
    }

    /// Construct a kernel matrix
    fn ker_mat(&self, m1: &Matrix<f64>, m2: &Matrix<f64>) -> Matrix<f64> {
        assert_eq!(m1.cols(), m2.cols());
        let cols = m1.cols();

        let dim1 = m1.rows();
        let dim2 = m2.rows();

        let mut ker_data = Vec::with_capacity(dim1 * dim2);

        for i in 0..dim1 {
            for j in 0..dim2 {
                ker_data.push(self.ker.kernel(&m1.data()[i * cols..(i + 1) * cols],
                                              &m2.data()[j * cols..(j + 1) * cols]));
            }
        }

        Matrix::new(dim1, dim2, ker_data)
    }
}

impl<T: Kernel, U: MeanFunc> SupModel<Matrix<f64>, Vector<f64>> for GaussianProcess<T, U> {
    /// Predict output from data.
    fn predict(&self, data: &Matrix<f64>) -> Vector<f64> {

        // Messy referencing for succint syntax
        if let (&Some(ref alpha), &Some(ref t_data)) = (&self.alpha, &self.train_data) {
            let mean = self.mean.func(data.clone());

            let post_mean = self.ker_mat(data, t_data) * alpha;

            return mean + post_mean;

        }

        panic!("The model has not been trained.");
    }

    /// Train the model using data and outputs.
    fn train(&mut self, data: &Matrix<f64>, value: &Vector<f64>) {
        let noise_mat = Matrix::identity(data.rows()) * self.noise;

        let ker_mat = self.ker_mat(data, data);

        let train_mat = (ker_mat + noise_mat).cholesky();

        let x = solve_l_triangular(&train_mat,
                                                            &(value -
                                                              self.mean.func(data.clone())));
        let alpha = solve_u_triangular(&train_mat.transpose(), &x);

        self.train_mat = Some(train_mat);
        self.train_data = Some(data.clone());
        self.alpha = Some(alpha);
    }
}

impl<T: Kernel, U: MeanFunc> GaussianProcess<T, U> {
    /// Compute the posterior distribution [UNSTABLE]
    ///
    /// Requires the model to be trained first.
    /// _Note that this is a messy compromise as GPs do not
    /// fit the SupModel trait as is currently implemented._
    pub fn get_posterior(&self, data: &Matrix<f64>) {
        if let (&Some(ref t_mat), &Some(ref alpha), &Some(ref t_data)) = (&self.train_mat, &self.alpha, &self.train_data) {
            let mean = self.mean.func(data.clone());

            let post_mean = mean + self.ker_mat(data, t_data) * alpha;

            for i in 0..data.rows() {
                let test_point = Vector::new(data.select_rows(&[i]).into_vec());
                let v = solve_l_triangular(t_mat, &test_point);
            }
        }

        panic!("The model has not been trained.");
    }
}

/// Solves an upper triangular linear system.
fn solve_u_triangular(mat: &Matrix<f64>, y: &Vector<f64>) -> Vector<f64> {
    assert!(mat.cols() == y.size(),
            "Matrix and Vector dimensions do not agree.");

    let mut x = vec![0.; y.size()];

    x[y.size() - 1] = y[y.size() - 1] / mat[[y.size() - 1, y.size() - 1]];

    for i in (0..y.size() - 1).rev() {
        let mut holding_u_sum = 0.;
        for j in (i+1..y.size()).rev() {
            holding_u_sum += mat.data()[i * mat.cols() + j] * x[j];
        }

        x[i] = (y[i] - holding_u_sum) /
            mat.data()[i * (mat.cols() + 1)];
    }

    Vector::new(x)
}

/// Solves a lower triangular linear system.
fn solve_l_triangular(mat: &Matrix<f64>, y: &Vector<f64>) -> Vector<f64> {
    assert!(mat.cols() == y.size(),
            "Matrix and Vector dimensions do not agree.");

    let mut x = vec![0.; y.size()];

    x[0] = y[0] / mat[[0, 0]];

    for i in 1..y.size() {
        let mut holding_l_sum = 0.;
        for j in 0..i {
            holding_l_sum += mat.data()[i * mat.cols() + j] * x[j];
        }

        x[i] = (y[i] - holding_l_sum) / mat.data()[i * (mat.cols() + 1)];
    }

    Vector::new(x)
}
<commit_msg>Correct implementation of GP posterior.<commit_after>//! Gaussian Processes
//!
//! Provides implementation of gaussian process regression.
//!
//! # Usage
//!
//! ```
//! use rusty_machine::learning::gp;
//! use rusty_machine::learning::SupModel;
//! use rusty_machine::linalg::matrix::Matrix;
//! use rusty_machine::linalg::vector::Vector;
//!
//! let mut gaussp = gp::GaussianProcess::default();
//! gaussp.noise = 10f64;
//!
//! let train_data = Matrix::new(10,1,vec![0.,1.,2.,3.,4.,5.,6.,7.,8.,9.]);
//! let target = Vector::new(vec![0.,1.,2.,3.,4.,4.,3.,2.,1.,0.]);
//!
//! gaussp.train(&train_data, &target);
//!
//! let test_data = Matrix::new(5,1,vec![2.3,4.4,5.1,6.2,7.1]);
//!
//! let output = gaussp.predict(&test_data);
//! ```
//! Alternatively one could use gaussp.get_posterior() which would return both
//! the predictive mean and covariance. However, this is likely to change in
//! a future release.

use learning::toolkit::kernel::{Kernel, SquaredExp};
use learning::SupModel;
use linalg::matrix::Matrix;
use linalg::vector::Vector;

/// Trait for GP mean functions.
pub trait MeanFunc {
    /// Compute the mean function applied elementwise to a matrix.
    fn func(&self, x: Matrix<f64>) -> Vector<f64>;
}

/// Constant mean function
pub struct ConstMean {
    a: f64,
}

impl Default for ConstMean {
    /// Constructs the zero function.
    ///
    /// # Examples
    /// ```
    /// use rusty_machine::learning::gp::ConstMean;
    ///
    /// let zero_m = ConstMean::default();
    /// ```
    fn default() -> ConstMean {
        ConstMean { a: 0f64 }
    }
}

impl MeanFunc for ConstMean {
    fn func(&self, x: Matrix<f64>) -> Vector<f64> {
        Vector::zeros(x.rows()) + self.a
    }
}

/// Gaussian Process struct
///
/// Gaussian process with generic kernel and deterministic mean function.
/// Can be used for gaussian process regression with noise.
/// Currently does not support classification.
pub struct GaussianProcess<T: Kernel, U: MeanFunc> {
    ker: T,
    mean: U,
    /// The observation noise of the GP.
    pub noise: f64,
    alpha: Option<Vector<f64>>,
    train_mat: Option<Matrix<f64>>,
    train_data: Option<Matrix<f64>>,
}

impl Default for GaussianProcess<SquaredExp, ConstMean> {
    /// Construct a default Gaussian Process
    ///
    /// # Examples
    ///
    /// ```
    /// use rusty_machine::learning::gp;
    ///
    /// let gp = gp::GaussianProcess::default();
    /// ```
    fn default() -> GaussianProcess<SquaredExp, ConstMean> {
        GaussianProcess {
            ker: SquaredExp::default(),
            mean: ConstMean::default(),
            noise: 0f64,
            train_mat: None,
            train_data: None,
            alpha: None,
        }
    }
}

impl<T: Kernel, U: MeanFunc> GaussianProcess<T, U> {
    /// Construct a new Gaussian Process.
    ///
    /// # Examples
    ///
    /// ```
    /// use rusty_machine::learning::gp;
    /// use rusty_machine::learning::toolkit::kernel;
    ///
    /// let ker = kernel::SquaredExp::default();
    /// let mean = gp::ConstMean::default();
    /// let gaussp = gp::GaussianProcess::new(ker, mean, 1e-3f64);
    /// ```
    pub fn new(ker: T, mean: U, noise: f64) -> GaussianProcess<T, U> {
        GaussianProcess {
            ker: ker,
            mean: mean,
            noise: noise,
            train_mat: None,
            train_data: None,
            alpha: None,
        }
    }

    /// Construct a kernel matrix
    fn ker_mat(&self, m1: &Matrix<f64>, m2: &Matrix<f64>) -> Matrix<f64> {
        assert_eq!(m1.cols(), m2.cols());
        let cols = m1.cols();

        let dim1 = m1.rows();
        let dim2 = m2.rows();

        let mut ker_data = Vec::with_capacity(dim1 * dim2);

        for i in 0..dim1 {
            for j in 0..dim2 {
                ker_data.push(self.ker.kernel(&m1.data()[i * cols..(i + 1) * cols],
                                              &m2.data()[j * cols..(j + 1) * cols]));
            }
        }

        Matrix::new(dim1, dim2, ker_data)
    }
}

impl<T: Kernel, U: MeanFunc> SupModel<Matrix<f64>, Vector<f64>> for GaussianProcess<T, U> {
    /// Predict output from data.
    fn predict(&self, data: &Matrix<f64>) -> Vector<f64> {

        // Messy referencing for succint syntax
        if let (&Some(ref alpha), &Some(ref t_data)) = (&self.alpha, &self.train_data) {
            let mean = self.mean.func(data.clone());

            let post_mean = self.ker_mat(data, t_data) * alpha;

            return mean + post_mean;

        }

        panic!("The model has not been trained.");
    }

    /// Train the model using data and outputs.
    fn train(&mut self, data: &Matrix<f64>, value: &Vector<f64>) {
        let noise_mat = Matrix::identity(data.rows()) * self.noise;

        let ker_mat = self.ker_mat(data, data);

        let train_mat = (ker_mat + noise_mat).cholesky();

        let x = solve_l_triangular(&train_mat, &(value - self.mean.func(data.clone())));
        let alpha = solve_u_triangular(&train_mat.transpose(), &x);

        self.train_mat = Some(train_mat);
        self.train_data = Some(data.clone());
        self.alpha = Some(alpha);
    }
}

impl<T: Kernel, U: MeanFunc> GaussianProcess<T, U> {
    /// Compute the posterior distribution [UNSTABLE]
    ///
    /// Requires the model to be trained first.
    ///
    /// Outputs the posterior mean and covariance matrix.
    pub fn get_posterior(&self, data: &Matrix<f64>) -> (Vector<f64>, Matrix<f64>) {
        if let (&Some(ref t_mat), &Some(ref alpha), &Some(ref t_data)) = (&self.train_mat,
                                                                          &self.alpha,
                                                                          &self.train_data) {
            let mean = self.mean.func(data.clone());

            let post_mean = mean + self.ker_mat(data, t_data) * alpha;

            let test_mat = self.ker_mat(data, t_data);
            let mut v_data = Vec::with_capacity(data.rows() * data.cols());
            for i in 0..test_mat.rows() {
                let test_point = Vector::new(test_mat.select_rows(&[i]).into_vec());
                v_data.append(&mut solve_l_triangular(t_mat, &test_point).into_vec());
            }

            let v_mat = Matrix::new(test_mat.rows(), test_mat.cols(), v_data);

            let post_var = self.ker_mat(data, data) - &v_mat * v_mat.transpose();

            return (post_mean, post_var);
        }

        panic!("The model has not been trained.");
    }
}

/// Solves an upper triangular linear system.
fn solve_u_triangular(mat: &Matrix<f64>, y: &Vector<f64>) -> Vector<f64> {
    assert!(mat.cols() == y.size(),
            "Matrix and Vector dimensions do not agree.");

    let mut x = vec![0.; y.size()];

    x[y.size() - 1] = y[y.size() - 1] / mat[[y.size() - 1, y.size() - 1]];

    for i in (0..y.size() - 1).rev() {
        let mut holding_u_sum = 0.;
        for j in (i + 1..y.size()).rev() {
            holding_u_sum += mat.data()[i * mat.cols() + j] * x[j];
        }

        x[i] = (y[i] - holding_u_sum) / mat.data()[i * (mat.cols() + 1)];
    }

    Vector::new(x)
}

/// Solves a lower triangular linear system.
fn solve_l_triangular(mat: &Matrix<f64>, y: &Vector<f64>) -> Vector<f64> {
    assert!(mat.cols() == y.size(),
            "Matrix and Vector dimensions do not agree.");

    let mut x = vec![0.; y.size()];

    x[0] = y[0] / mat[[0, 0]];

    for i in 1..y.size() {
        let mut holding_l_sum = 0.;
        for j in 0..i {
            holding_l_sum += mat.data()[i * mat.cols() + j] * x[j];
        }

        x[i] = (y[i] - holding_l_sum) / mat.data()[i * (mat.cols() + 1)];
    }

    Vector::new(x)
}
<|endoftext|>


<commit_before><commit_msg>Parse register copy Jake's way<commit_after><|endoftext|>


<commit_before><commit_msg>Texture Info is debug<commit_after><|endoftext|>


<commit_before><commit_msg>Refactor: less indenting<commit_after><|endoftext|>


<commit_before><commit_msg>Remove as_bytes (following discussion in #27)<commit_after><|endoftext|>


<commit_before><commit_msg>refactoring<commit_after><|endoftext|>


<commit_before><commit_msg>Specify lifetimes for DebugLoc iterators because lifetime elision is too conservative.<commit_after><|endoftext|>


<commit_before><commit_msg>Fixed padding writing bug.<commit_after><|endoftext|>


<commit_before>// Copyright:: Copyright (c) 2015-2016 The Habitat Maintainers
//
// The terms of the Evaluation Agreement (Habitat) between Chef Software Inc.
// and the party accessing this file ("Licensee") apply to Licensee's use of
// the Software until such time that the Software is made available under an
// open source license such as the Apache 2.0 License.

use std::env;
use std::ffi::OsString;

use hcore::env as henv;

use config;
use error::Result;

pub fn start(args: Vec<OsString>) -> Result<()> {
    // If the `$HAB_ORIGIN` environment variable is not present, then see if a default is set in
    // the CLI config. If so, set it as the `$HAB_ORIGIN` environment variable for the `hab-studio`
    // or `docker` execv call.
    if henv::var("HAB_ORIGIN").is_err() {
        let config = try!(config::load());
        if let Some(default_origin) = config.origin {
            debug!("Setting default origin {} via CLI config", &default_origin);
            env::set_var("HAB_ORIGIN", default_origin);
        }
    }
    inner::start(args)
}

#[cfg(target_os = "linux")]
mod inner {
    use std::ffi::OsString;
    use std::path::PathBuf;
    use std::str::FromStr;

    use hcore::crypto::{init, default_cache_key_path};
    use hcore::env as henv;
    use hcore::fs::find_command;
    use hcore::package::PackageIdent;

    use error::{Error, Result};
    use exec;

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

    pub fn start(args: Vec<OsString>) -> Result<()> {
        let command = match henv::var(STUDIO_CMD_ENVVAR) {
            Ok(command) => PathBuf::from(command),
            Err(_) => {
                init();
                let ident = try!(PackageIdent::from_str(STUDIO_PACKAGE_IDENT));
                try!(exec::command_from_pkg(STUDIO_CMD, &ident, &default_cache_key_path(None), 0))
            }
        };

        if let Some(cmd) = find_command(command.to_string_lossy().as_ref()) {
            try!(exec::exec_command(cmd, args));
        } else {
            return Err(Error::ExecCommandNotFound(command.to_string_lossy().into_owned()));
        }
        Ok(())
    }
}

#[cfg(not(target_os = "linux"))]
mod inner {
    use std::env;
    use std::ffi::OsString;

    use hcore::crypto::default_cache_key_path;
    use hcore::env as henv;
    use hcore::fs::{CACHE_KEY_PATH, find_command};

    use error::{Error, Result};
    use exec;

    const DOCKER_CMD: &'static str = "docker";
    const DOCKER_CMD_ENVVAR: &'static str = "HAB_DOCKER_BINARY";

    const DOCKER_IMAGE: &'static str = "habitat/studio";
    const DOCKER_IMAGE_ENVVAR: &'static str = "HAB_DOCKER_STUDIO_IMAGE";

    pub fn start(args: Vec<OsString>) -> Result<()> {
        let docker = henv::var(DOCKER_CMD_ENVVAR).unwrap_or(DOCKER_CMD.to_string());
        let image = henv::var(DOCKER_IMAGE_ENVVAR).unwrap_or(DOCKER_IMAGE.to_string());

        let cmd = match find_command(&docker) {
            Some(cmd) => cmd,
            None => return Err(Error::ExecCommandNotFound(docker.to_string())),
        };

        let mut cmd_args: Vec<OsString> = vec!["run".into(),
                                               "--rm".into(),
                                               "--tty".into(),
                                               "--interactive".into(),
                                               "--privileged".into()];
        let env_vars = vec!["HAB_DEPOT_URL", "HAB_ORIGIN", "http_proxy", "https_proxy"];
        for var in env_vars {
            if let Ok(val) = henv::var(var) {
                debug!("Propagating environment variable into container: {}={}",
                       var,
                       val);
                cmd_args.push("--env".into());
                cmd_args.push(format!("{}={}", var, val).into());
            }
        }
        cmd_args.push("--volume".into());
        cmd_args.push("/var/run/docker.sock:/var/run/docker.sock".into());
        cmd_args.push("--volume".into());
        cmd_args.push(format!("{}:/{}",
                              default_cache_key_path(None).to_string_lossy(),
                              CACHE_KEY_PATH)
            .into());
        cmd_args.push("--volume".into());
        cmd_args.push(format!("{}:/src", env::current_dir().unwrap().to_string_lossy()).into());
        cmd_args.push(image.into());
        cmd_args.extend_from_slice(args.as_slice());

        for var in vec!["http_proxy", "https_proxy"] {
            if let Ok(_) = henv::var(var) {
                debug!("Unsetting proxy environment variable '{}' before calling `{}'",
                       var,
                       docker);
                env::remove_var(var);
            }
        }

        exec::exec_command(cmd, cmd_args)
    }
}
<commit_msg>[hab/mac] Update go-live Docker Studio image to the Bintray repo.<commit_after>// Copyright:: Copyright (c) 2015-2016 The Habitat Maintainers
//
// The terms of the Evaluation Agreement (Habitat) between Chef Software Inc.
// and the party accessing this file ("Licensee") apply to Licensee's use of
// the Software until such time that the Software is made available under an
// open source license such as the Apache 2.0 License.

use std::env;
use std::ffi::OsString;

use hcore::env as henv;

use config;
use error::Result;

pub fn start(args: Vec<OsString>) -> Result<()> {
    // If the `$HAB_ORIGIN` environment variable is not present, then see if a default is set in
    // the CLI config. If so, set it as the `$HAB_ORIGIN` environment variable for the `hab-studio`
    // or `docker` execv call.
    if henv::var("HAB_ORIGIN").is_err() {
        let config = try!(config::load());
        if let Some(default_origin) = config.origin {
            debug!("Setting default origin {} via CLI config", &default_origin);
            env::set_var("HAB_ORIGIN", default_origin);
        }
    }
    inner::start(args)
}

#[cfg(target_os = "linux")]
mod inner {
    use std::ffi::OsString;
    use std::path::PathBuf;
    use std::str::FromStr;

    use hcore::crypto::{init, default_cache_key_path};
    use hcore::env as henv;
    use hcore::fs::find_command;
    use hcore::package::PackageIdent;

    use error::{Error, Result};
    use exec;

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

    pub fn start(args: Vec<OsString>) -> Result<()> {
        let command = match henv::var(STUDIO_CMD_ENVVAR) {
            Ok(command) => PathBuf::from(command),
            Err(_) => {
                init();
                let ident = try!(PackageIdent::from_str(STUDIO_PACKAGE_IDENT));
                try!(exec::command_from_pkg(STUDIO_CMD, &ident, &default_cache_key_path(None), 0))
            }
        };

        if let Some(cmd) = find_command(command.to_string_lossy().as_ref()) {
            try!(exec::exec_command(cmd, args));
        } else {
            return Err(Error::ExecCommandNotFound(command.to_string_lossy().into_owned()));
        }
        Ok(())
    }
}

#[cfg(not(target_os = "linux"))]
mod inner {
    use std::env;
    use std::ffi::OsString;

    use hcore::crypto::default_cache_key_path;
    use hcore::env as henv;
    use hcore::fs::{CACHE_KEY_PATH, find_command};

    use error::{Error, Result};
    use exec;

    const DOCKER_CMD: &'static str = "docker";
    const DOCKER_CMD_ENVVAR: &'static str = "HAB_DOCKER_BINARY";

    const DOCKER_IMAGE: &'static str = "habitat-docker-registry.bintray.io/studio";
    const DOCKER_IMAGE_ENVVAR: &'static str = "HAB_DOCKER_STUDIO_IMAGE";

    pub fn start(args: Vec<OsString>) -> Result<()> {
        let docker = henv::var(DOCKER_CMD_ENVVAR).unwrap_or(DOCKER_CMD.to_string());
        let image = henv::var(DOCKER_IMAGE_ENVVAR).unwrap_or(DOCKER_IMAGE.to_string());

        let cmd = match find_command(&docker) {
            Some(cmd) => cmd,
            None => return Err(Error::ExecCommandNotFound(docker.to_string())),
        };

        let mut cmd_args: Vec<OsString> = vec!["run".into(),
                                               "--rm".into(),
                                               "--tty".into(),
                                               "--interactive".into(),
                                               "--privileged".into()];
        let env_vars = vec!["HAB_DEPOT_URL", "HAB_ORIGIN", "http_proxy", "https_proxy"];
        for var in env_vars {
            if let Ok(val) = henv::var(var) {
                debug!("Propagating environment variable into container: {}={}",
                       var,
                       val);
                cmd_args.push("--env".into());
                cmd_args.push(format!("{}={}", var, val).into());
            }
        }
        cmd_args.push("--volume".into());
        cmd_args.push("/var/run/docker.sock:/var/run/docker.sock".into());
        cmd_args.push("--volume".into());
        cmd_args.push(format!("{}:/{}",
                              default_cache_key_path(None).to_string_lossy(),
                              CACHE_KEY_PATH)
            .into());
        cmd_args.push("--volume".into());
        cmd_args.push(format!("{}:/src", env::current_dir().unwrap().to_string_lossy()).into());
        cmd_args.push(image.into());
        cmd_args.extend_from_slice(args.as_slice());

        for var in vec!["http_proxy", "https_proxy"] {
            if let Ok(_) = henv::var(var) {
                debug!("Unsetting proxy environment variable '{}' before calling `{}'",
                       var,
                       docker);
                env::remove_var(var);
            }
        }

        exec::exec_command(cmd, cmd_args)
    }
}
<|endoftext|>


<commit_before>/*
 * Copyright (C) 2017-2018 Ryan Huang
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU Affero General Public License as published
 * by the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU Affero General Public License for more details.
 *
 * You should have received a copy of the GNU Affero General Public License
 * along with this program.  If not, see <https://www.gnu.org/licenses/>.
 */

use model::{Color, FieldCoord, HexCoord};

#[derive(Clone, Copy)]
pub struct Board {
    /*
                ____
               /    \
          ____/  +y  \____
         /    \      /    \
        /      \____/  +x  \
        \      /    \      /
         \____/      \____/
         /    \      /    \
        /  -x  \____/      \
        \      /    \      /
         \____/  -y  \____/
              \      /
               \____/

        The hex board uses an axial coordinate system with (0, 0) at the center. The x-axis slopes
        up to the right, and the y-axis goes up and down. The board is stored as a dense 2D array,
        because a ragged array won't work (since each row would have to be a different type).

        See http://www.redblobgames.com/grids/hexagons/#coordinates-axial for more info.
    */
    board: [[Option<Hex>; 5]; 5],
    turn: Color,
    black_pieces: u32,
    black_hexes: u32,
    white_pieces: u32,
    white_hexes: u32,
}

// Fields are numbered clockwise from the top. Even indicies are black, odd indicies are white.
type Hex = [Field; 6];

// There's no need to store the color of the piece on this field, since only white pieces can be on
// white fields, and vice versa.
#[derive(Clone, Copy, Debug, PartialEq)]
pub enum Field {
    Piece,
    Empty,
}

// Public methods
impl Board {
    #[cfg_attr(rustfmt, rustfmt_skip)]
    /// Create a new board with the "Laurentius" starting position.
    pub fn new() -> Board {
        let mut board = Board {
            board: [[None; 5]; 5],
            turn: Color::White,
            black_pieces: 18,
            black_hexes: 0,
            white_pieces: 18,
            white_hexes: 0,
        };

        // (0, 0) is the only empty hex.
        board.set_hex(&HexCoord::new(0, 0), Some([Field::Empty; 6]));

        // All other hexes have exactly two pieces on them in the starting position.
        let piece_locations = [
            (-2,  2, 0, 4),
            (-2,  1, 0, 3),
            (-2,  0, 3, 5),
            (-1,  2, 1, 4),
            (-1,  1, 0, 4),
            (-1,  0, 3, 5),
            (-1, -1, 2, 5),
            ( 0,  2, 1, 5),
            ( 0,  1, 1, 5),
            ( 0, -1, 2, 4),
            ( 0, -2, 2, 4),
            ( 1,  1, 2, 5),
            ( 1,  0, 0, 2),
            ( 1, -1, 1, 3),
            ( 1, -2, 1, 4),
            ( 2,  0, 0, 2),
            ( 2, -1, 0, 3),
            ( 2, -2, 1, 3),
        ];

        for &(x, y, f1, f2) in &piece_locations {
            let mut hex = [Field::Empty; 6];
            hex[f1] = Field::Piece;
            hex[f2] = Field::Piece;
            board.set_hex(&HexCoord::new(x, y), Some(hex));
        }

        board
    }
    pub fn can_move_piece(&self, from: &FieldCoord, to: &FieldCoord) -> bool {
        from.color() == self.turn && self.is_piece_on_field(from) && !self.is_piece_on_field(to)
            && self.get_field_vertex_neighbors(from).contains(to)
    }
    pub fn move_piece(&mut self, from: &FieldCoord, to: &FieldCoord) {
        assert!(self.can_move_piece(from, to));

        self.set_field(from, Field::Empty);
        self.set_field(to, Field::Piece);

        let (capture_count, mut fields_to_check) = self.check_hexes(&from.to_hex());
        fields_to_check.append(&mut self.get_field_edge_neighbors(to));
        self.check_captures(&fields_to_check);

        match self.turn {
            Color::White => self.white_hexes += capture_count,
            Color::Black => self.black_hexes += capture_count,
        }
        self.turn = self.turn.switch();
    }
    pub fn get_available_moves(&self, field: &FieldCoord) -> Vec<FieldCoord> {
        if self.is_piece_on_field(field) {
            self.get_field_vertex_neighbors(field)
                .into_iter()
                .filter(|c| !self.is_piece_on_field(c))
                .collect()
        } else {
            vec![]
        }
    }
    pub fn can_exchange(&self) -> bool {
        2 <= match self.turn {
            Color::Black => self.black_hexes,
            Color::White => self.white_hexes,
        }
    }
    pub fn can_exchange_piece(&self, coord: &FieldCoord) -> bool {
        self.can_exchange() && coord.color() != self.turn && self.is_piece_on_field(coord)
    }
    pub fn exchange_piece(&mut self, coord: &FieldCoord) {
        assert!(self.can_exchange_piece(coord));

        self.remove_piece(coord);
        match self.turn {
            Color::White => self.white_hexes -= 2,
            Color::Black => self.black_hexes -= 2,
        }

        // Players don't collect hexes removed due to an exchange
        let (_, fields_to_check) = self.check_hexes(&coord.to_hex());
        self.check_captures(&fields_to_check);
        self.turn = self.turn.switch();
    }
    pub fn is_piece_on_field(&self, coord: &FieldCoord) -> bool {
        self.get_field(coord) == &Field::Piece
    }
    /// > extant (adj.): Still in existence; not destroyed, lost, or extinct (The Free Dictionary)
    ///
    /// Return the coordinates of the hexes that have not been removed yet.
    pub fn extant_hexes(&self) -> Vec<HexCoord> {
        let mut coords = vec![];
        for x in -2..3 {
            for y in -2..3 {
                coords.push((x, y));
            }
        }
        coords
            .iter()
            .filter_map(|&coord| self.try_hex(coord))
            .collect()
    }
    pub fn turn(&self) -> Color {
        self.turn
    }
    pub fn black_pieces(&self) -> u32 {
        self.black_pieces
    }
    pub fn black_hexes(&self) -> u32 {
        self.black_hexes
    }
    pub fn white_pieces(&self) -> u32 {
        self.white_pieces
    }
    pub fn white_hexes(&self) -> u32 {
        self.white_hexes
    }
}

// Field and piece methods
impl Board {
    fn get_field(&self, coord: &FieldCoord) -> &Field {
        match *self.get_hex(&coord.to_hex()) {
            Some(ref hex) => &hex[coord.f as usize],
            None => panic!(
                "Cannot get field {} on removed hex at {:?}",
                coord.f,
                coord.to_hex()
            ),
        }
    }
    fn set_field(&mut self, coord: &FieldCoord, field: Field) {
        match *self.get_hex(&coord.to_hex()) {
            Some(mut hex) => {
                hex[coord.f as usize] = field;
                self.set_hex(&coord.to_hex(), Some(hex));
            }
            None => panic!(
                "Cannot set field {} on removed hex at {:?}",
                coord.f,
                coord.to_hex()
            ),
        }
    }
    /// Return fields that share an edge with the given field. These fields are always the opposite
    /// color of the given field. If all of a piece's edge neighbors are occupied, that piece might
    /// be capturable.
    fn get_field_edge_neighbors(&self, coord: &FieldCoord) -> Vec<FieldCoord> {
        let hex = coord.to_hex();
        let mut neighbors = vec![
            // There are always two edge neighbors on the same hex as the given field
            hex.to_field((coord.f + 1) % 6),
            hex.to_field((coord.f + 5) % 6),
        ];

        if let Some(hex) = self.get_hex_neighbor(&hex, coord.f) {
            neighbors.push(hex.to_field((coord.f + 3) % 6));
        }

        neighbors
    }
    /// Return fields that share a vertex with the given field and have the same color as the given
    /// field. Pieces can move to fields that are vertex neighbors of the field they are on.
    fn get_field_vertex_neighbors(&self, coord: &FieldCoord) -> Vec<FieldCoord> {
        let hex = coord.to_hex();
        let mut neighbors = vec![
            // There are always two vertex neighbors on the same hex as the given field
            hex.to_field((coord.f + 2) % 6),
            hex.to_field((coord.f + 4) % 6),
        ];
        if let Some(hex) = self.get_hex_neighbor(&hex, coord.f) {
            neighbors.push(hex.to_field((coord.f + 2) % 6));
            neighbors.push(hex.to_field((coord.f + 4) % 6));
        }
        if let Some(hex) = self.get_hex_neighbor(&hex, (coord.f + 1) % 6) {
            neighbors.push(hex.to_field((coord.f + 4) % 6));
        }
        if let Some(hex) = self.get_hex_neighbor(&hex, (coord.f + 5) % 6) {
            neighbors.push(hex.to_field((coord.f + 2) % 6));
        }
        neighbors
    }
    fn remove_piece(&mut self, coord: &FieldCoord) {
        assert!(
            self.is_piece_on_field(coord),
            "There is no piece at {:?} to remove",
            coord
        );
        self.set_field(coord, Field::Empty);
        match coord.color() {
            Color::Black => self.black_pieces -= 1,
            Color::White => self.white_pieces -= 1,
        }
    }
    fn check_captures(&mut self, fields_to_check: &[FieldCoord]) {
        for field in fields_to_check {
            if field.color() != self.turn && self.is_piece_on_field(field)
                && self.get_field_edge_neighbors(field)
                    .iter()
                    .all(|coord| self.is_piece_on_field(coord))
            {
                self.remove_piece(field);
            }
        }
    }
}

// Hex methods
impl Board {
    fn get_hex(&self, coord: &HexCoord) -> &Option<Hex> {
        &self.board[(coord.x + 2) as usize][(coord.y + 2) as usize]
    }
    fn set_hex(&mut self, coord: &HexCoord, hex: Option<Hex>) {
        self.board[(coord.x + 2) as usize][(coord.y + 2) as usize] = hex;
    }
    fn get_hex_neighbor(&self, coord: &HexCoord, direction: u32) -> Option<HexCoord> {
        assert!(direction < 6);

        let neighbors = [
            (coord.x, coord.y + 1),
            (coord.x + 1, coord.y),
            (coord.x + 1, coord.y - 1),
            (coord.x, coord.y - 1),
            (coord.x - 1, coord.y),
            (coord.x - 1, coord.y + 1),
        ];

        self.try_hex(neighbors[direction as usize])
    }
    /// A hex is removable (and must be removed) if it is empty and is "attached to the board by 3
    /// or less adjacent sides."
    fn is_hex_removable(&self, coord: &HexCoord) -> bool {
        match *self.get_hex(coord) {
            Some(hex) => if hex != [Field::Empty; 6] {
                return false;
            },
            None => return false,
        }
        // After repeatedly failing to find an efficient and elegant way to check if a hex is
        // removable, I have opted to just use a lookup table. Even if it is not elegant, at
        // least it is simple and efficient.
        let neighbor_combinations = [
            0b000001, 0b000010, 0b000100, 0b001000, 0b010000, 0b100000, 0b000011, 0b000110,
            0b001100, 0b011000, 0b110000, 0b100001, 0b000111, 0b001110, 0b011100, 0b111000,
            0b110001, 0b100011,
        ];

        let mut neighbors = 0;
        for f in 0..6 {
            if self.get_hex_neighbor(coord, f).is_some() {
                neighbors += 1 << f;
            }
        }
        assert!(
            neighbors != 0,
            "A hex at {:?} is empty and has no neighbors",
            coord
        );
        neighbor_combinations.contains(&neighbors)
    }
    fn remove_hex(&mut self, coord: &HexCoord) -> bool {
        let removable = self.is_hex_removable(coord);

        if removable {
            self.set_hex(coord, None);
        }
        removable
    }
    fn try_hex(&self, coord: (i32, i32)) -> Option<HexCoord> {
        if let Some(coord) = HexCoord::try_new(coord.0, coord.1) {
            if self.get_hex(&coord).is_some() {
                return Some(coord);
            }
        }
        None
    }
    fn check_hexes(&mut self, coord: &HexCoord) -> (u32, Vec<FieldCoord>) {
        let mut remove_count = 0;
        let mut fields = vec![];

        if self.remove_hex(coord) {
            remove_count += 1;
            for f in 0..6 {
                if let Some(neighbor) = self.get_hex_neighbor(coord, f) {
                    let (new_remove_count, mut new_fields) = self.check_hexes(&neighbor);
                    if new_remove_count == 0 {
                        fields.push(neighbor.to_field((f + 3) % 6));
                    } else {
                        remove_count += new_remove_count;
                        fields.append(&mut new_fields);
                    }
                }
            }
        }
        (remove_count, fields)
    }
}
<commit_msg>Simplify get_hex_neighbor<commit_after>/*
 * Copyright (C) 2017-2018 Ryan Huang
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU Affero General Public License as published
 * by the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU Affero General Public License for more details.
 *
 * You should have received a copy of the GNU Affero General Public License
 * along with this program.  If not, see <https://www.gnu.org/licenses/>.
 */

use model::{Color, FieldCoord, HexCoord};

#[derive(Clone, Copy)]
pub struct Board {
    /*
                ____
               /    \
          ____/  +y  \____
         /    \      /    \
        /      \____/  +x  \
        \      /    \      /
         \____/      \____/
         /    \      /    \
        /  -x  \____/      \
        \      /    \      /
         \____/  -y  \____/
              \      /
               \____/

        The hex board uses an axial coordinate system with (0, 0) at the center. The x-axis slopes
        up to the right, and the y-axis goes up and down. The board is stored as a dense 2D array,
        because a ragged array won't work (since each row would have to be a different type).

        See http://www.redblobgames.com/grids/hexagons/#coordinates-axial for more info.
    */
    board: [[Option<Hex>; 5]; 5],
    turn: Color,
    black_pieces: u32,
    black_hexes: u32,
    white_pieces: u32,
    white_hexes: u32,
}

// Fields are numbered clockwise from the top. Even indicies are black, odd indicies are white.
type Hex = [Field; 6];

// There's no need to store the color of the piece on this field, since only white pieces can be on
// white fields, and vice versa.
#[derive(Clone, Copy, Debug, PartialEq)]
pub enum Field {
    Piece,
    Empty,
}

// Public methods
impl Board {
    #[cfg_attr(rustfmt, rustfmt_skip)]
    /// Create a new board with the "Laurentius" starting position.
    pub fn new() -> Board {
        let mut board = Board {
            board: [[None; 5]; 5],
            turn: Color::White,
            black_pieces: 18,
            black_hexes: 0,
            white_pieces: 18,
            white_hexes: 0,
        };

        // (0, 0) is the only empty hex.
        board.set_hex(&HexCoord::new(0, 0), Some([Field::Empty; 6]));

        // All other hexes have exactly two pieces on them in the starting position.
        let piece_locations = [
            (-2,  2, 0, 4),
            (-2,  1, 0, 3),
            (-2,  0, 3, 5),
            (-1,  2, 1, 4),
            (-1,  1, 0, 4),
            (-1,  0, 3, 5),
            (-1, -1, 2, 5),
            ( 0,  2, 1, 5),
            ( 0,  1, 1, 5),
            ( 0, -1, 2, 4),
            ( 0, -2, 2, 4),
            ( 1,  1, 2, 5),
            ( 1,  0, 0, 2),
            ( 1, -1, 1, 3),
            ( 1, -2, 1, 4),
            ( 2,  0, 0, 2),
            ( 2, -1, 0, 3),
            ( 2, -2, 1, 3),
        ];

        for &(x, y, f1, f2) in &piece_locations {
            let mut hex = [Field::Empty; 6];
            hex[f1] = Field::Piece;
            hex[f2] = Field::Piece;
            board.set_hex(&HexCoord::new(x, y), Some(hex));
        }

        board
    }
    pub fn can_move_piece(&self, from: &FieldCoord, to: &FieldCoord) -> bool {
        from.color() == self.turn && self.is_piece_on_field(from) && !self.is_piece_on_field(to)
            && self.get_field_vertex_neighbors(from).contains(to)
    }
    pub fn move_piece(&mut self, from: &FieldCoord, to: &FieldCoord) {
        assert!(self.can_move_piece(from, to));

        self.set_field(from, Field::Empty);
        self.set_field(to, Field::Piece);

        let (capture_count, mut fields_to_check) = self.check_hexes(&from.to_hex());
        fields_to_check.append(&mut self.get_field_edge_neighbors(to));
        self.check_captures(&fields_to_check);

        match self.turn {
            Color::White => self.white_hexes += capture_count,
            Color::Black => self.black_hexes += capture_count,
        }
        self.turn = self.turn.switch();
    }
    pub fn get_available_moves(&self, field: &FieldCoord) -> Vec<FieldCoord> {
        if self.is_piece_on_field(field) {
            self.get_field_vertex_neighbors(field)
                .into_iter()
                .filter(|c| !self.is_piece_on_field(c))
                .collect()
        } else {
            vec![]
        }
    }
    pub fn can_exchange(&self) -> bool {
        2 <= match self.turn {
            Color::Black => self.black_hexes,
            Color::White => self.white_hexes,
        }
    }
    pub fn can_exchange_piece(&self, coord: &FieldCoord) -> bool {
        self.can_exchange() && coord.color() != self.turn && self.is_piece_on_field(coord)
    }
    pub fn exchange_piece(&mut self, coord: &FieldCoord) {
        assert!(self.can_exchange_piece(coord));

        self.remove_piece(coord);
        match self.turn {
            Color::White => self.white_hexes -= 2,
            Color::Black => self.black_hexes -= 2,
        }

        // Players don't collect hexes removed due to an exchange
        let (_, fields_to_check) = self.check_hexes(&coord.to_hex());
        self.check_captures(&fields_to_check);
        self.turn = self.turn.switch();
    }
    pub fn is_piece_on_field(&self, coord: &FieldCoord) -> bool {
        self.get_field(coord) == &Field::Piece
    }
    /// > extant (adj.): Still in existence; not destroyed, lost, or extinct (The Free Dictionary)
    ///
    /// Return the coordinates of the hexes that have not been removed yet.
    pub fn extant_hexes(&self) -> Vec<HexCoord> {
        let mut coords = vec![];
        for x in -2..3 {
            for y in -2..3 {
                coords.push((x, y));
            }
        }
        coords
            .iter()
            .filter_map(|&coord| self.try_hex(coord))
            .collect()
    }
    pub fn turn(&self) -> Color {
        self.turn
    }
    pub fn black_pieces(&self) -> u32 {
        self.black_pieces
    }
    pub fn black_hexes(&self) -> u32 {
        self.black_hexes
    }
    pub fn white_pieces(&self) -> u32 {
        self.white_pieces
    }
    pub fn white_hexes(&self) -> u32 {
        self.white_hexes
    }
}

// Field and piece methods
impl Board {
    fn get_field(&self, coord: &FieldCoord) -> &Field {
        match *self.get_hex(&coord.to_hex()) {
            Some(ref hex) => &hex[coord.f as usize],
            None => panic!(
                "Cannot get field {} on removed hex at {:?}",
                coord.f,
                coord.to_hex()
            ),
        }
    }
    fn set_field(&mut self, coord: &FieldCoord, field: Field) {
        match *self.get_hex(&coord.to_hex()) {
            Some(mut hex) => {
                hex[coord.f as usize] = field;
                self.set_hex(&coord.to_hex(), Some(hex));
            }
            None => panic!(
                "Cannot set field {} on removed hex at {:?}",
                coord.f,
                coord.to_hex()
            ),
        }
    }
    /// Return fields that share an edge with the given field. These fields are always the opposite
    /// color of the given field. If all of a piece's edge neighbors are occupied, that piece might
    /// be capturable.
    fn get_field_edge_neighbors(&self, coord: &FieldCoord) -> Vec<FieldCoord> {
        let hex = coord.to_hex();
        let mut neighbors = vec![
            // There are always two edge neighbors on the same hex as the given field
            hex.to_field((coord.f + 1) % 6),
            hex.to_field((coord.f + 5) % 6),
        ];

        if let Some(hex) = self.get_hex_neighbor(&hex, coord.f) {
            neighbors.push(hex.to_field((coord.f + 3) % 6));
        }

        neighbors
    }
    /// Return fields that share a vertex with the given field and have the same color as the given
    /// field. Pieces can move to fields that are vertex neighbors of the field they are on.
    fn get_field_vertex_neighbors(&self, coord: &FieldCoord) -> Vec<FieldCoord> {
        let hex = coord.to_hex();
        let mut neighbors = vec![
            // There are always two vertex neighbors on the same hex as the given field
            hex.to_field((coord.f + 2) % 6),
            hex.to_field((coord.f + 4) % 6),
        ];
        if let Some(hex) = self.get_hex_neighbor(&hex, coord.f) {
            neighbors.push(hex.to_field((coord.f + 2) % 6));
            neighbors.push(hex.to_field((coord.f + 4) % 6));
        }
        if let Some(hex) = self.get_hex_neighbor(&hex, (coord.f + 1) % 6) {
            neighbors.push(hex.to_field((coord.f + 4) % 6));
        }
        if let Some(hex) = self.get_hex_neighbor(&hex, (coord.f + 5) % 6) {
            neighbors.push(hex.to_field((coord.f + 2) % 6));
        }
        neighbors
    }
    fn remove_piece(&mut self, coord: &FieldCoord) {
        assert!(
            self.is_piece_on_field(coord),
            "There is no piece at {:?} to remove",
            coord
        );
        self.set_field(coord, Field::Empty);
        match coord.color() {
            Color::Black => self.black_pieces -= 1,
            Color::White => self.white_pieces -= 1,
        }
    }
    fn check_captures(&mut self, fields_to_check: &[FieldCoord]) {
        for field in fields_to_check {
            if field.color() != self.turn && self.is_piece_on_field(field)
                && self.get_field_edge_neighbors(field)
                    .iter()
                    .all(|coord| self.is_piece_on_field(coord))
            {
                self.remove_piece(field);
            }
        }
    }
}

// Hex methods
impl Board {
    fn get_hex(&self, coord: &HexCoord) -> &Option<Hex> {
        &self.board[(coord.x + 2) as usize][(coord.y + 2) as usize]
    }
    fn set_hex(&mut self, coord: &HexCoord, hex: Option<Hex>) {
        self.board[(coord.x + 2) as usize][(coord.y + 2) as usize] = hex;
    }
    fn get_hex_neighbor(&self, coord: &HexCoord, direction: u32) -> Option<HexCoord> {
        self.try_hex(match direction {
            0 => (coord.x, coord.y + 1),
            1 => (coord.x + 1, coord.y),
            2 => (coord.x + 1, coord.y - 1),
            3 => (coord.x, coord.y - 1),
            4 => (coord.x - 1, coord.y),
            5 => (coord.x - 1, coord.y + 1),
            _ => panic!("Direction must be less than 6"),
        })
    }
    /// A hex is removable (and must be removed) if it is empty and is "attached to the board by 3
    /// or less adjacent sides."
    fn is_hex_removable(&self, coord: &HexCoord) -> bool {
        match *self.get_hex(coord) {
            Some(hex) => if hex != [Field::Empty; 6] {
                return false;
            },
            None => return false,
        }
        // After repeatedly failing to find an efficient and elegant way to check if a hex is
        // removable, I have opted to just use a lookup table. Even if it is not elegant, at
        // least it is simple and efficient.
        let neighbor_combinations = [
            0b000001, 0b000010, 0b000100, 0b001000, 0b010000, 0b100000, 0b000011, 0b000110,
            0b001100, 0b011000, 0b110000, 0b100001, 0b000111, 0b001110, 0b011100, 0b111000,
            0b110001, 0b100011,
        ];

        let mut neighbors = 0;
        for f in 0..6 {
            if self.get_hex_neighbor(coord, f).is_some() {
                neighbors += 1 << f;
            }
        }
        assert!(
            neighbors != 0,
            "A hex at {:?} is empty and has no neighbors",
            coord
        );
        neighbor_combinations.contains(&neighbors)
    }
    fn remove_hex(&mut self, coord: &HexCoord) -> bool {
        let removable = self.is_hex_removable(coord);

        if removable {
            self.set_hex(coord, None);
        }
        removable
    }
    fn try_hex(&self, coord: (i32, i32)) -> Option<HexCoord> {
        if let Some(coord) = HexCoord::try_new(coord.0, coord.1) {
            if self.get_hex(&coord).is_some() {
                return Some(coord);
            }
        }
        None
    }
    fn check_hexes(&mut self, coord: &HexCoord) -> (u32, Vec<FieldCoord>) {
        let mut remove_count = 0;
        let mut fields = vec![];

        if self.remove_hex(coord) {
            remove_count += 1;
            for f in 0..6 {
                if let Some(neighbor) = self.get_hex_neighbor(coord, f) {
                    let (new_remove_count, mut new_fields) = self.check_hexes(&neighbor);
                    if new_remove_count == 0 {
                        fields.push(neighbor.to_field((f + 3) % 6));
                    } else {
                        remove_count += new_remove_count;
                        fields.append(&mut new_fields);
                    }
                }
            }
        }
        (remove_count, fields)
    }
}
<|endoftext|>


<commit_before>// Copyright 2020 The ChromiumOS Authors
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

use std::env;
use std::ffi::CString;
use std::fs::File;
use std::fs::OpenOptions;
use std::io;
use std::io::BufRead;
use std::io::BufReader;
use std::io::Write;
use std::os::unix::fs::OpenOptionsExt;
use std::path::Path;
use std::path::PathBuf;
use std::process::Child;
use std::process::Command;
use std::process::Stdio;
use std::str::from_utf8;
use std::sync::mpsc::sync_channel;
use std::sync::Once;
use std::thread;
use std::time::Duration;

use anyhow::anyhow;
use anyhow::Result;
use base::syslog;
use libc::O_DIRECT;
use tempfile::TempDir;

const PREBUILT_URL: &str = "https://storage.googleapis.com/chromeos-localmirror/distfiles";

#[cfg(target_arch = "x86_64")]
const ARCH: &str = "x86_64";
#[cfg(target_arch = "arm")]
const ARCH: &str = "arm";
#[cfg(target_arch = "aarch64")]
const ARCH: &str = "aarch64";

/// Timeout for communicating with the VM. If we do not hear back, panic so we
/// do not block the tests.
const VM_COMMUNICATION_TIMEOUT: Duration = Duration::from_secs(10);

fn prebuilt_version() -> &'static str {
    include_str!("../../guest_under_test/PREBUILT_VERSION").trim()
}

fn kernel_prebuilt_url() -> String {
    format!(
        "{}/crosvm-testing-bzimage-{}-{}",
        PREBUILT_URL,
        ARCH,
        prebuilt_version()
    )
}

fn rootfs_prebuilt_url() -> String {
    format!(
        "{}/crosvm-testing-rootfs-{}-{}",
        PREBUILT_URL,
        ARCH,
        prebuilt_version()
    )
}

/// The kernel bzImage is stored next to the test executable, unless overridden by
/// CROSVM_CARGO_TEST_KERNEL_BINARY
fn kernel_path() -> PathBuf {
    match env::var("CROSVM_CARGO_TEST_KERNEL_BINARY") {
        Ok(value) => PathBuf::from(value),
        Err(_) => env::current_exe()
            .unwrap()
            .parent()
            .unwrap()
            .join("bzImage"),
    }
}

/// The rootfs image is stored next to the test executable, unless overridden by
/// CROSVM_CARGO_TEST_ROOTFS_IMAGE
fn rootfs_path() -> PathBuf {
    match env::var("CROSVM_CARGO_TEST_ROOTFS_IMAGE") {
        Ok(value) => PathBuf::from(value),
        Err(_) => env::current_exe().unwrap().parent().unwrap().join("rootfs"),
    }
}

/// The crosvm binary is expected to be alongside to the integration tests
/// binary. Alternatively in the parent directory (cargo will put the
/// test binary in target/debug/deps/ but the crosvm binary in target/debug)
fn find_crosvm_binary() -> PathBuf {
    cfg_if::cfg_if! {
        if #[cfg(features="direct")] {
            let binary_name = "crosvm-direct";
        } else {
            let binary_name = "crosvm";
        }
    }

    let exe_dir = env::current_exe().unwrap().parent().unwrap().to_path_buf();
    let first = exe_dir.join(binary_name);
    if first.exists() {
        return first;
    }
    let second = exe_dir.parent().unwrap().join(binary_name);
    if second.exists() {
        return second;
    }
    panic!(
        "Cannot find {} in ./ or ../ alongside test binary.",
        binary_name
    );
}

/// Safe wrapper for libc::mkfifo
fn mkfifo(path: &Path) -> io::Result<()> {
    let cpath = CString::new(path.to_str().unwrap()).unwrap();
    let result = unsafe { libc::mkfifo(cpath.as_ptr(), 0o777) };
    if result == 0 {
        Ok(())
    } else {
        Err(io::Error::last_os_error())
    }
}

/// Run the provided closure, but panic if it does not complete until the timeout has passed.
/// We should panic here, as we cannot gracefully stop the closure from running.
/// `on_timeout` will be called before panic to allow printing debug information.
fn run_with_timeout<F, G, U>(closure: F, timeout: Duration, on_timeout: G) -> U
where
    F: FnOnce() -> U + Send + 'static,
    G: FnOnce(),
    U: Send + 'static,
{
    let (tx, rx) = sync_channel::<()>(1);
    let handle = thread::spawn(move || {
        let result = closure();
        tx.send(()).unwrap();
        result
    });
    if rx.recv_timeout(timeout).is_err() {
        on_timeout();
        panic!("Operation timed out or closure paniced.");
    }
    handle.join().unwrap()
}

fn download_file(url: &str, destination: &Path) -> Result<()> {
    let status = Command::new("curl")
        .arg("--fail")
        .arg("--location")
        .args(&["--output", destination.to_str().unwrap()])
        .arg(url)
        .status();
    match status {
        Ok(exit_code) => {
            if !exit_code.success() {
                Err(anyhow!("Cannot download {}", url))
            } else {
                Ok(())
            }
        }
        Err(error) => Err(anyhow!(error)),
    }
}

/// Configuration to start `TestVm`.
#[derive(Default)]
pub struct Config {
    /// Extra arguments for the `run` subcommand.
    extra_args: Vec<String>,

    /// Use `O_DIRECT` for the rootfs.
    o_direct: bool,
}

#[cfg(test)]
impl Config {
    /// Creates a new `run` command with `extra_args`.
    pub fn new() -> Self {
        Default::default()
    }

    /// Uses extra arguments for `crosvm run`.
    #[allow(dead_code)]
    pub fn extra_args(mut self, args: Vec<String>) -> Self {
        self.extra_args = args;
        self
    }

    /// Uses `O_DIRECT` for the rootfs.
    pub fn o_direct(mut self) -> Self {
        self.o_direct = true;
        self
    }
}

/// Test fixture to spin up a VM running a guest that can be communicated with.
///
/// After creation, commands can be sent via exec_in_guest. The VM is stopped
/// when this instance is dropped.
#[cfg(test)]
pub struct TestVm {
    /// Maintain ownership of test_dir until the vm is destroyed.
    #[allow(dead_code)]
    test_dir: TempDir,
    from_guest_reader: BufReader<File>,
    to_guest: File,
    control_socket_path: PathBuf,
    process: Option<Child>, // Use `Option` to allow taking the ownership in `Drop::drop()`.
}

impl TestVm {
    /// Magic line sent by the delegate binary when the guest is ready.
    const MAGIC_LINE: &'static str = "\x05Ready";

    /// Downloads prebuilts if needed.
    fn initialize_once() {
        if let Err(e) = syslog::init() {
            panic!("failed to initiailize syslog: {}", e);
        }

        // It's possible the prebuilts downloaded by crosvm-9999.ebuild differ
        // from the version that crosvm was compiled for.
        if let Ok(value) = env::var("CROSVM_CARGO_TEST_PREBUILT_VERSION") {
            if value != prebuilt_version() {
                panic!(
                    "Environment provided prebuilts are version {}, but crosvm was compiled \
                    for prebuilt version {}. Did you update PREBUILT_VERSION everywhere?",
                    value,
                    prebuilt_version()
                );
            }
        }

        let kernel_path = kernel_path();
        if env::var("CROSVM_CARGO_TEST_KERNEL_BINARY").is_err() {
            if !kernel_path.exists() {
                println!("Downloading kernel prebuilt:");
                download_file(&kernel_prebuilt_url(), &kernel_path).unwrap();
            }
        }
        assert!(kernel_path.exists(), "{:?} does not exist", kernel_path);

        let rootfs_path = rootfs_path();
        if env::var("CROSVM_CARGO_TEST_ROOTFS_IMAGE").is_err() {
            if !rootfs_path.exists() {
                println!("Downloading rootfs prebuilt:");
                download_file(&rootfs_prebuilt_url(), &rootfs_path).unwrap();
            }
        }
        assert!(rootfs_path.exists(), "{:?} does not exist", rootfs_path);

        // Check if the test file system is a known compatible one. Needs to support features like O_DIRECT.
        if let Err(e) = OpenOptions::new()
            .custom_flags(O_DIRECT)
            .write(false)
            .read(true)
            .open(rootfs_path)
        {
            panic!(
                "File open with O_DIRECT expected to work but did not: {}",
                e
            );
        }
    }

    // Adds 2 serial devices:
    // - ttyS0: Console device which prints kernel log / debug output of the
    //          delegate binary.
    // - ttyS1: Serial device attached to the named pipes.
    fn configure_serial_devices(
        command: &mut Command,
        from_guest_pipe: &Path,
        to_guest_pipe: &Path,
    ) {
        command.args(&["--serial", "type=syslog"]);

        // Setup channel for communication with the delegate.
        let serial_params = format!(
            "type=file,path={},input={},num=2",
            from_guest_pipe.display(),
            to_guest_pipe.display()
        );
        command.args(&["--serial", &serial_params]);
    }

    /// Configures the VM rootfs to load from the guest_under_test assets.
    fn configure_rootfs(command: &mut Command, o_direct: bool) {
        let rootfs_and_option = format!(
            "{}{},ro,root",
            rootfs_path().to_str().unwrap(),
            if o_direct { ",o_direct=true" } else { "" }
        );
        command
            .args(&["--block", &rootfs_and_option])
            .args(&["--params", "init=/bin/delegate"]);
    }

    /// Instanciate a new crosvm instance. The first call will trigger the download of prebuilt
    /// files if necessary.
    pub fn new(cfg: Config) -> Result<TestVm> {
        static PREP_ONCE: Once = Once::new();
        PREP_ONCE.call_once(TestVm::initialize_once);

        // Create two named pipes to communicate with the guest.
        let test_dir = TempDir::new()?;
        let from_guest_pipe = test_dir.path().join("from_guest");
        let to_guest_pipe = test_dir.path().join("to_guest");
        mkfifo(&from_guest_pipe)?;
        mkfifo(&to_guest_pipe)?;

        let control_socket_path = test_dir.path().join("control");

        let mut command = Command::new(find_crosvm_binary());
        command.args(&["run"]);
        TestVm::configure_serial_devices(&mut command, &from_guest_pipe, &to_guest_pipe);
        command.args(&["--socket", control_socket_path.to_str().unwrap()]);
        TestVm::configure_rootfs(&mut command, cfg.o_direct);
        command.args(cfg.extra_args);
        // Set kernel as the last argument.
        command.arg(kernel_path());
        // Set `Stdio::piped` so we can forward the outputs to stdout later.
        command.stdout(Stdio::piped());
        command.stderr(Stdio::piped());

        println!("$ {:?}", command);

        let mut process = Some(command.spawn()?);

        // Open pipes. Panic if we cannot connect after a timeout.
        let (to_guest, from_guest) = run_with_timeout(
            move || (File::create(to_guest_pipe), File::open(from_guest_pipe)),
            VM_COMMUNICATION_TIMEOUT,
            || {
                let mut process = process.take().unwrap();
                process.kill().unwrap();
                let output = process.wait_with_output().unwrap();

                // Print both the crosvm's stdout/stderr to stdout so that they'll be shown when
                // the test failed.
                println!(
                    "TestVm stdout:\n{}",
                    std::str::from_utf8(&output.stdout).unwrap()
                );
                println!(
                    "TestVm stderr:\n{}",
                    std::str::from_utf8(&output.stderr).unwrap()
                );
            },
        );

        // Wait for magic line to be received, indicating the delegate is ready.
        let mut from_guest_reader = BufReader::new(from_guest?);
        let mut magic_line = String::new();
        from_guest_reader.read_line(&mut magic_line)?;
        assert_eq!(magic_line.trim(), TestVm::MAGIC_LINE);

        Ok(TestVm {
            test_dir,
            from_guest_reader,
            to_guest: to_guest?,
            control_socket_path,
            process,
        })
    }

    /// Executes the shell command `command` and returns the programs stdout.
    pub fn exec_in_guest(&mut self, command: &str) -> Result<String> {
        // Write command to serial port.
        writeln!(&mut self.to_guest, "{}", command)?;

        // We will receive an echo of what we have written on the pipe.
        let mut echo = String::new();
        self.from_guest_reader.read_line(&mut echo)?;
        assert_eq!(echo.trim(), command);

        // Return all remaining lines until we receive the MAGIC_LINE
        let mut output = String::new();
        loop {
            let mut line = String::new();
            self.from_guest_reader.read_line(&mut line)?;
            if line.trim() == TestVm::MAGIC_LINE {
                break;
            }
            output.push_str(&line);
        }
        let trimmed = output.trim();
        println!("<- {:?}", trimmed);

        Ok(trimmed.to_string())
    }

    fn crosvm_command(&self, command: &str, mut args: Vec<String>) -> Result<()> {
        args.push(self.control_socket_path.to_str().unwrap().to_string());

        println!("$ crosvm {} {:?}", command, &args.join(" "));

        let mut cmd = Command::new(find_crosvm_binary());
        cmd.arg(command).args(args);
        cmd.stdout(Stdio::piped());
        cmd.stderr(Stdio::piped());

        let output = cmd.output()?;
        // Print both the crosvm's stdout/stderr to stdout so that they'll be shown when the test
        // is failed.
        println!(
            "`crosvm {}` stdout:\n{}",
            command,
            from_utf8(&output.stdout).unwrap()
        );
        println!(
            "`crosvm {}` stderr:\n{}",
            command,
            from_utf8(&output.stderr).unwrap()
        );

        if !output.status.success() {
            Err(anyhow!("Command failed with exit code {}", output.status))
        } else {
            Ok(())
        }
    }

    pub fn stop(&self) -> Result<()> {
        self.crosvm_command("stop", vec![])
    }

    pub fn suspend(&self) -> Result<()> {
        self.crosvm_command("suspend", vec![])
    }

    pub fn resume(&self) -> Result<()> {
        self.crosvm_command("resume", vec![])
    }

    pub fn disk(&self, args: Vec<String>) -> Result<()> {
        self.crosvm_command("disk", args)
    }
}

impl Drop for TestVm {
    fn drop(&mut self) {
        self.stop().unwrap();
        let output = self.process.take().unwrap().wait_with_output().unwrap();

        // Print both the crosvm's stdout/stderr to stdout so that they'll be shown when the test
        // failed.
        println!(
            "TestVm stdout:\n{}",
            std::str::from_utf8(&output.stdout).unwrap()
        );
        println!(
            "TestVm stderr:\n{}",
            std::str::from_utf8(&output.stderr).unwrap()
        );

        if !output.status.success() {
            panic!("VM exited illegally: {}", output.status);
        }
    }
}
<commit_msg>e2e_tests: Update prebuild URL to crosvm-infra<commit_after>// Copyright 2020 The ChromiumOS Authors
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

use std::env;
use std::ffi::CString;
use std::fs::File;
use std::fs::OpenOptions;
use std::io;
use std::io::BufRead;
use std::io::BufReader;
use std::io::Write;
use std::os::unix::fs::OpenOptionsExt;
use std::path::Path;
use std::path::PathBuf;
use std::process::Child;
use std::process::Command;
use std::process::Stdio;
use std::str::from_utf8;
use std::sync::mpsc::sync_channel;
use std::sync::Once;
use std::thread;
use std::time::Duration;

use anyhow::anyhow;
use anyhow::Result;
use base::syslog;
use libc::O_DIRECT;
use tempfile::TempDir;

const PREBUILT_URL: &str = "https://storage.googleapis.com/crosvm/integration_tests";

#[cfg(target_arch = "x86_64")]
const ARCH: &str = "x86_64";
#[cfg(target_arch = "arm")]
const ARCH: &str = "arm";
#[cfg(target_arch = "aarch64")]
const ARCH: &str = "aarch64";

/// Timeout for communicating with the VM. If we do not hear back, panic so we
/// do not block the tests.
const VM_COMMUNICATION_TIMEOUT: Duration = Duration::from_secs(10);

fn prebuilt_version() -> &'static str {
    include_str!("../../guest_under_test/PREBUILT_VERSION").trim()
}

fn kernel_prebuilt_url() -> String {
    format!(
        "{}/guest-bzimage-{}-{}",
        PREBUILT_URL,
        ARCH,
        prebuilt_version()
    )
}

fn rootfs_prebuilt_url() -> String {
    format!(
        "{}/guest-rootfs-{}-{}",
        PREBUILT_URL,
        ARCH,
        prebuilt_version()
    )
}

/// The kernel bzImage is stored next to the test executable, unless overridden by
/// CROSVM_CARGO_TEST_KERNEL_BINARY
fn kernel_path() -> PathBuf {
    match env::var("CROSVM_CARGO_TEST_KERNEL_BINARY") {
        Ok(value) => PathBuf::from(value),
        Err(_) => env::current_exe()
            .unwrap()
            .parent()
            .unwrap()
            .join("bzImage"),
    }
}

/// The rootfs image is stored next to the test executable, unless overridden by
/// CROSVM_CARGO_TEST_ROOTFS_IMAGE
fn rootfs_path() -> PathBuf {
    match env::var("CROSVM_CARGO_TEST_ROOTFS_IMAGE") {
        Ok(value) => PathBuf::from(value),
        Err(_) => env::current_exe().unwrap().parent().unwrap().join("rootfs"),
    }
}

/// The crosvm binary is expected to be alongside to the integration tests
/// binary. Alternatively in the parent directory (cargo will put the
/// test binary in target/debug/deps/ but the crosvm binary in target/debug)
fn find_crosvm_binary() -> PathBuf {
    cfg_if::cfg_if! {
        if #[cfg(features="direct")] {
            let binary_name = "crosvm-direct";
        } else {
            let binary_name = "crosvm";
        }
    }

    let exe_dir = env::current_exe().unwrap().parent().unwrap().to_path_buf();
    let first = exe_dir.join(binary_name);
    if first.exists() {
        return first;
    }
    let second = exe_dir.parent().unwrap().join(binary_name);
    if second.exists() {
        return second;
    }
    panic!(
        "Cannot find {} in ./ or ../ alongside test binary.",
        binary_name
    );
}

/// Safe wrapper for libc::mkfifo
fn mkfifo(path: &Path) -> io::Result<()> {
    let cpath = CString::new(path.to_str().unwrap()).unwrap();
    let result = unsafe { libc::mkfifo(cpath.as_ptr(), 0o777) };
    if result == 0 {
        Ok(())
    } else {
        Err(io::Error::last_os_error())
    }
}

/// Run the provided closure, but panic if it does not complete until the timeout has passed.
/// We should panic here, as we cannot gracefully stop the closure from running.
/// `on_timeout` will be called before panic to allow printing debug information.
fn run_with_timeout<F, G, U>(closure: F, timeout: Duration, on_timeout: G) -> U
where
    F: FnOnce() -> U + Send + 'static,
    G: FnOnce(),
    U: Send + 'static,
{
    let (tx, rx) = sync_channel::<()>(1);
    let handle = thread::spawn(move || {
        let result = closure();
        tx.send(()).unwrap();
        result
    });
    if rx.recv_timeout(timeout).is_err() {
        on_timeout();
        panic!("Operation timed out or closure paniced.");
    }
    handle.join().unwrap()
}

fn download_file(url: &str, destination: &Path) -> Result<()> {
    let status = Command::new("curl")
        .arg("--fail")
        .arg("--location")
        .args(&["--output", destination.to_str().unwrap()])
        .arg(url)
        .status();
    match status {
        Ok(exit_code) => {
            if !exit_code.success() {
                Err(anyhow!("Cannot download {}", url))
            } else {
                Ok(())
            }
        }
        Err(error) => Err(anyhow!(error)),
    }
}

/// Configuration to start `TestVm`.
#[derive(Default)]
pub struct Config {
    /// Extra arguments for the `run` subcommand.
    extra_args: Vec<String>,

    /// Use `O_DIRECT` for the rootfs.
    o_direct: bool,
}

#[cfg(test)]
impl Config {
    /// Creates a new `run` command with `extra_args`.
    pub fn new() -> Self {
        Default::default()
    }

    /// Uses extra arguments for `crosvm run`.
    #[allow(dead_code)]
    pub fn extra_args(mut self, args: Vec<String>) -> Self {
        self.extra_args = args;
        self
    }

    /// Uses `O_DIRECT` for the rootfs.
    pub fn o_direct(mut self) -> Self {
        self.o_direct = true;
        self
    }
}

/// Test fixture to spin up a VM running a guest that can be communicated with.
///
/// After creation, commands can be sent via exec_in_guest. The VM is stopped
/// when this instance is dropped.
#[cfg(test)]
pub struct TestVm {
    /// Maintain ownership of test_dir until the vm is destroyed.
    #[allow(dead_code)]
    test_dir: TempDir,
    from_guest_reader: BufReader<File>,
    to_guest: File,
    control_socket_path: PathBuf,
    process: Option<Child>, // Use `Option` to allow taking the ownership in `Drop::drop()`.
}

impl TestVm {
    /// Magic line sent by the delegate binary when the guest is ready.
    const MAGIC_LINE: &'static str = "\x05Ready";

    /// Downloads prebuilts if needed.
    fn initialize_once() {
        if let Err(e) = syslog::init() {
            panic!("failed to initiailize syslog: {}", e);
        }

        // It's possible the prebuilts downloaded by crosvm-9999.ebuild differ
        // from the version that crosvm was compiled for.
        if let Ok(value) = env::var("CROSVM_CARGO_TEST_PREBUILT_VERSION") {
            if value != prebuilt_version() {
                panic!(
                    "Environment provided prebuilts are version {}, but crosvm was compiled \
                    for prebuilt version {}. Did you update PREBUILT_VERSION everywhere?",
                    value,
                    prebuilt_version()
                );
            }
        }

        let kernel_path = kernel_path();
        if env::var("CROSVM_CARGO_TEST_KERNEL_BINARY").is_err() {
            if !kernel_path.exists() {
                println!("Downloading kernel prebuilt:");
                download_file(&kernel_prebuilt_url(), &kernel_path).unwrap();
            }
        }
        assert!(kernel_path.exists(), "{:?} does not exist", kernel_path);

        let rootfs_path = rootfs_path();
        if env::var("CROSVM_CARGO_TEST_ROOTFS_IMAGE").is_err() {
            if !rootfs_path.exists() {
                println!("Downloading rootfs prebuilt:");
                download_file(&rootfs_prebuilt_url(), &rootfs_path).unwrap();
            }
        }
        assert!(rootfs_path.exists(), "{:?} does not exist", rootfs_path);

        // Check if the test file system is a known compatible one. Needs to support features like O_DIRECT.
        if let Err(e) = OpenOptions::new()
            .custom_flags(O_DIRECT)
            .write(false)
            .read(true)
            .open(rootfs_path)
        {
            panic!(
                "File open with O_DIRECT expected to work but did not: {}",
                e
            );
        }
    }

    // Adds 2 serial devices:
    // - ttyS0: Console device which prints kernel log / debug output of the
    //          delegate binary.
    // - ttyS1: Serial device attached to the named pipes.
    fn configure_serial_devices(
        command: &mut Command,
        from_guest_pipe: &Path,
        to_guest_pipe: &Path,
    ) {
        command.args(&["--serial", "type=syslog"]);

        // Setup channel for communication with the delegate.
        let serial_params = format!(
            "type=file,path={},input={},num=2",
            from_guest_pipe.display(),
            to_guest_pipe.display()
        );
        command.args(&["--serial", &serial_params]);
    }

    /// Configures the VM rootfs to load from the guest_under_test assets.
    fn configure_rootfs(command: &mut Command, o_direct: bool) {
        let rootfs_and_option = format!(
            "{}{},ro,root",
            rootfs_path().to_str().unwrap(),
            if o_direct { ",o_direct=true" } else { "" }
        );
        command
            .args(&["--block", &rootfs_and_option])
            .args(&["--params", "init=/bin/delegate"]);
    }

    /// Instanciate a new crosvm instance. The first call will trigger the download of prebuilt
    /// files if necessary.
    pub fn new(cfg: Config) -> Result<TestVm> {
        static PREP_ONCE: Once = Once::new();
        PREP_ONCE.call_once(TestVm::initialize_once);

        // Create two named pipes to communicate with the guest.
        let test_dir = TempDir::new()?;
        let from_guest_pipe = test_dir.path().join("from_guest");
        let to_guest_pipe = test_dir.path().join("to_guest");
        mkfifo(&from_guest_pipe)?;
        mkfifo(&to_guest_pipe)?;

        let control_socket_path = test_dir.path().join("control");

        let mut command = Command::new(find_crosvm_binary());
        command.args(&["run"]);
        TestVm::configure_serial_devices(&mut command, &from_guest_pipe, &to_guest_pipe);
        command.args(&["--socket", control_socket_path.to_str().unwrap()]);
        TestVm::configure_rootfs(&mut command, cfg.o_direct);
        command.args(cfg.extra_args);
        // Set kernel as the last argument.
        command.arg(kernel_path());
        // Set `Stdio::piped` so we can forward the outputs to stdout later.
        command.stdout(Stdio::piped());
        command.stderr(Stdio::piped());

        println!("$ {:?}", command);

        let mut process = Some(command.spawn()?);

        // Open pipes. Panic if we cannot connect after a timeout.
        let (to_guest, from_guest) = run_with_timeout(
            move || (File::create(to_guest_pipe), File::open(from_guest_pipe)),
            VM_COMMUNICATION_TIMEOUT,
            || {
                let mut process = process.take().unwrap();
                process.kill().unwrap();
                let output = process.wait_with_output().unwrap();

                // Print both the crosvm's stdout/stderr to stdout so that they'll be shown when
                // the test failed.
                println!(
                    "TestVm stdout:\n{}",
                    std::str::from_utf8(&output.stdout).unwrap()
                );
                println!(
                    "TestVm stderr:\n{}",
                    std::str::from_utf8(&output.stderr).unwrap()
                );
            },
        );

        // Wait for magic line to be received, indicating the delegate is ready.
        let mut from_guest_reader = BufReader::new(from_guest?);
        let mut magic_line = String::new();
        from_guest_reader.read_line(&mut magic_line)?;
        assert_eq!(magic_line.trim(), TestVm::MAGIC_LINE);

        Ok(TestVm {
            test_dir,
            from_guest_reader,
            to_guest: to_guest?,
            control_socket_path,
            process,
        })
    }

    /// Executes the shell command `command` and returns the programs stdout.
    pub fn exec_in_guest(&mut self, command: &str) -> Result<String> {
        // Write command to serial port.
        writeln!(&mut self.to_guest, "{}", command)?;

        // We will receive an echo of what we have written on the pipe.
        let mut echo = String::new();
        self.from_guest_reader.read_line(&mut echo)?;
        assert_eq!(echo.trim(), command);

        // Return all remaining lines until we receive the MAGIC_LINE
        let mut output = String::new();
        loop {
            let mut line = String::new();
            self.from_guest_reader.read_line(&mut line)?;
            if line.trim() == TestVm::MAGIC_LINE {
                break;
            }
            output.push_str(&line);
        }
        let trimmed = output.trim();
        println!("<- {:?}", trimmed);

        Ok(trimmed.to_string())
    }

    fn crosvm_command(&self, command: &str, mut args: Vec<String>) -> Result<()> {
        args.push(self.control_socket_path.to_str().unwrap().to_string());

        println!("$ crosvm {} {:?}", command, &args.join(" "));

        let mut cmd = Command::new(find_crosvm_binary());
        cmd.arg(command).args(args);
        cmd.stdout(Stdio::piped());
        cmd.stderr(Stdio::piped());

        let output = cmd.output()?;
        // Print both the crosvm's stdout/stderr to stdout so that they'll be shown when the test
        // is failed.
        println!(
            "`crosvm {}` stdout:\n{}",
            command,
            from_utf8(&output.stdout).unwrap()
        );
        println!(
            "`crosvm {}` stderr:\n{}",
            command,
            from_utf8(&output.stderr).unwrap()
        );

        if !output.status.success() {
            Err(anyhow!("Command failed with exit code {}", output.status))
        } else {
            Ok(())
        }
    }

    pub fn stop(&self) -> Result<()> {
        self.crosvm_command("stop", vec![])
    }

    pub fn suspend(&self) -> Result<()> {
        self.crosvm_command("suspend", vec![])
    }

    pub fn resume(&self) -> Result<()> {
        self.crosvm_command("resume", vec![])
    }

    pub fn disk(&self, args: Vec<String>) -> Result<()> {
        self.crosvm_command("disk", args)
    }
}

impl Drop for TestVm {
    fn drop(&mut self) {
        self.stop().unwrap();
        let output = self.process.take().unwrap().wait_with_output().unwrap();

        // Print both the crosvm's stdout/stderr to stdout so that they'll be shown when the test
        // failed.
        println!(
            "TestVm stdout:\n{}",
            std::str::from_utf8(&output.stdout).unwrap()
        );
        println!(
            "TestVm stderr:\n{}",
            std::str::from_utf8(&output.stderr).unwrap()
        );

        if !output.status.success() {
            panic!("VM exited illegally: {}", output.status);
        }
    }
}
<|endoftext|>