<commit_before>use crate::agents::SignJws;
use crate::bridges::oidc::ProviderKey;
use crate::config::Config;
use crate::email_address::EmailAddress;
use crate::utils::{base64url, keys::SignError, unix_duration, SecureRandom};
use ring::{
    digest,
    error::Unspecified,
    signature::{self, UnparsedPublicKey},
};
use serde_json as json;
use serde_json::{json, Error as JsonError, Value};
use std::fmt;
use std::iter::Iterator;
use thiserror::Error;

type RsaPublicKey = signature::RsaPublicKeyComponents<Vec<u8>>;

/// Token signing algorithms we support.
#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub enum SigningAlgorithm {
    EdDsa,
    Rs256,
}

impl SigningAlgorithm {
    /// Get the JWA string representation.
    pub fn as_str(self) -> &'static str {
        use SigningAlgorithm::*;
        match self {
            EdDsa => "EdDSA",
            Rs256 => "RS256",
        }
    }

    /// Format a list of algorithms for display.
    pub fn format_list(list: &[Self]) -> String {
        list.iter()
            .map(|alg| alg.as_str())
            .collect::<Vec<_>>()
            .join(", ")
    }
}

impl fmt::Display for SigningAlgorithm {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.write_str(self.as_str())
    }
}

impl std::str::FromStr for SigningAlgorithm {
    type Err = &'static str;
    fn from_str(s: &str) -> Result<SigningAlgorithm, &'static str> {
        use SigningAlgorithm::*;
        match s {
            "EdDSA" => Ok(EdDsa),
            "RS256" => Ok(Rs256),
            _ => Err("unsupported value"),
        }
    }
}

serde_from_str!(SigningAlgorithm);
serde_display!(SigningAlgorithm);

/// The types of public keys we support.
pub enum SupportedPublicKey {
    Ed25519(UnparsedPublicKey<Vec<u8>>),
    Rsa(RsaPublicKey),
}

impl SupportedPublicKey {
    /// Verify a message signature.
    pub fn verify(&self, message: &[u8], signature: &[u8]) -> Result<(), Unspecified> {
        use SupportedPublicKey::*;
        match self {
            Ed25519(ref inner) => inner.verify(message, signature),
            Rsa(ref inner) => {
                inner.verify(&signature::RSA_PKCS1_2048_8192_SHA256, message, signature)
            }
        }
    }
}

/// Helper function to build a session ID for a login attempt.
///
/// Put the email address, the client ID (RP origin) and some randomness into
/// a SHA256 hash, and encode it with URL-safe bas64 encoding. This is used
/// as the key in Redis, as well as the state for OAuth authentication.
pub async fn session_id(email: &EmailAddress, client_id: &str, rng: &SecureRandom) -> String {
    let rand_bytes = rng.generate_async(16).await;
    let mut ctx = digest::Context::new(&digest::SHA256);
    ctx.update(email.as_str().as_bytes());
    ctx.update(client_id.as_bytes());
    ctx.update(&rand_bytes);
    base64url::encode(&ctx.finish())
}

/// Helper function to create a secure nonce.
pub async fn nonce(rng: &SecureRandom) -> String {
    let rand_bytes = rng.generate_async(16).await;
    base64url::encode(&rand_bytes)
}

/// Helper function to create a random string consisting of
/// characters from the z-base-32 set.
pub async fn random_zbase32(len: usize, rng: &SecureRandom) -> String {
    const CHARSET: &[u8] = b"13456789abcdefghijkmnopqrstuwxyz";
    let rand_bytes = rng.generate_async(len).await;
    String::from_utf8(
        rand_bytes
            .into_iter()
            .map(|v| CHARSET[v as usize % CHARSET.len()])
            .collect(),
    )
    .expect("failed to build one-time pad")
}

#[derive(Debug, Error)]
pub enum VerifyError {
    #[error("the token must consist of three dot-separated parts")]
    IncorrectFormat,
    #[error("token part {index} contained invalid base64: {reason}")]
    InvalidPartBase64 {
        index: usize,
        reason: base64::DecodeError,
    },
    #[error("the token header contained invalid JSON: {0}")]
    InvalidHeaderJson(JsonError),
    #[error("did not find a string 'kid' property in the token header")]
    KidMissing,
    #[error("the token 'kid' could not be found in the JWKs document: {kid}")]
    KidNotMatched { kid: String },
    #[error(
        "the '{}' field of the matching JWK contains invalid base64: {}",
        property,
        reason
    )]
    InvalidJwkBase64 {
        property: &'static str,
        reason: base64::DecodeError,
    },
    #[error("the matching JWK is of an unsupported type")]
    UnsupportedKeyType,
    #[error("the token signature did not validate using the matching JWK")]
    BadSignature,
    #[error("the token payload contained invalid JSON: {0}")]
    InvalidPayloadJson(JsonError),
}

/// Verify a JWS signature, returning the payload as a `Value` if successful.
pub fn verify_jws(
    jws: &str,
    key_set: &[ProviderKey],
    signing_alg: SigningAlgorithm,
) -> Result<json::Value, VerifyError> {
    // Split the token up in parts and decode them.
    let parts: Vec<&str> = jws.split('.').collect();
    if parts.len() != 3 {
        return Err(VerifyError::IncorrectFormat);
    }
    let decoded = parts
        .iter()
        .enumerate()
        .map(|(idx, s)| {
            base64url::decode(s).map_err(|reason| VerifyError::InvalidPartBase64 {
                index: idx + 1,
                reason,
            })
        })
        .collect::<Result<Vec<_>, _>>()?;

    // Parse the header and find the key ID.
    let jwt_header: json::Value =
        json::from_slice(&decoded[0]).map_err(VerifyError::InvalidHeaderJson)?;
    let kid = jwt_header
        .get("kid")
        .and_then(Value::as_str)
        .ok_or(VerifyError::KidMissing)?;

    // Look for they key ID in the JWKs.
    let matched_keys: Vec<&ProviderKey> = key_set
        .iter()
        .filter(|key| key.use_ == "sig" && key.kid == kid)
        .collect();

    // Verify that we found exactly one key matching the key ID.
    if matched_keys.len() != 1 {
        return Err(VerifyError::KidNotMatched {
            kid: kid.to_owned(),
        });
    }
    let key = matched_keys.first().unwrap();

    // Then, use the data to build a public key object for verification.
    let pub_key = match (signing_alg, key.alg.as_str(), key.crv.as_str()) {
        (SigningAlgorithm::EdDsa, "EdDSA", "Ed25519") => {
            let x = base64url::decode(&key.x).map_err(|reason| VerifyError::InvalidJwkBase64 {
                property: "x",
                reason,
            })?;
            let key = UnparsedPublicKey::new(&signature::ED25519, x);
            SupportedPublicKey::Ed25519(key)
        }
        (SigningAlgorithm::Rs256, "RS256", _) => {
            let n = base64url::decode(&key.n).map_err(|reason| VerifyError::InvalidJwkBase64 {
                property: "n",
                reason,
            })?;
            let e = base64url::decode(&key.e).map_err(|reason| VerifyError::InvalidJwkBase64 {
                property: "e",
                reason,
            })?;
            let key = RsaPublicKey { n, e };
            SupportedPublicKey::Rsa(key)
        }
        _ => return Err(VerifyError::UnsupportedKeyType),
    };

    // Verify the signature using the public key.
    let message_len = parts[0].len() + parts[1].len() + 1;
    pub_key
        .verify(jws[..message_len].as_bytes(), &decoded[2])
        .map_err(|_err| VerifyError::BadSignature)?;

    // Return the payload.
    Ok(json::from_slice(&decoded[1]).map_err(VerifyError::InvalidPayloadJson)?)
}

/// Helper method to create a JWT for a given email address and audience.
///
/// Builds the JSON payload, then signs it using the last key provided in
/// the configuration object.
pub async fn create_jwt(
    app: &Config,
    email: &str,
    email_addr: &EmailAddress,
    aud: &str,
    nonce: &str,
    signing_alg: SigningAlgorithm,
) -> Result<String, SignError> {
    let now = unix_duration();
    app.key_manager
        .send(SignJws {
            payload: json!({
                "aud": aud,
                "email": email_addr.as_str(),
                "email_verified": true,
                "email_original": email,
                "exp": (now + app.token_ttl).as_secs(),
                "iat": now.as_secs(),
                "iss": &app.public_url,
                "sub": email_addr.as_str(),
                "nonce": nonce,
            }),
            signing_alg,
        })
        .await
}
<commit_msg>Fix new clippy lint<commit_after>use crate::agents::SignJws;
use crate::bridges::oidc::ProviderKey;
use crate::config::Config;
use crate::email_address::EmailAddress;
use crate::utils::{base64url, keys::SignError, unix_duration, SecureRandom};
use ring::{
    digest,
    error::Unspecified,
    signature::{self, UnparsedPublicKey},
};
use serde_json as json;
use serde_json::{json, Error as JsonError, Value};
use std::fmt;
use std::iter::Iterator;
use thiserror::Error;

type RsaPublicKey = signature::RsaPublicKeyComponents<Vec<u8>>;

/// Token signing algorithms we support.
#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub enum SigningAlgorithm {
    EdDsa,
    Rs256,
}

impl SigningAlgorithm {
    /// Get the JWA string representation.
    pub fn as_str(self) -> &'static str {
        use SigningAlgorithm::*;
        match self {
            EdDsa => "EdDSA",
            Rs256 => "RS256",
        }
    }

    /// Format a list of algorithms for display.
    pub fn format_list(list: &[Self]) -> String {
        list.iter()
            .map(|alg| alg.as_str())
            .collect::<Vec<_>>()
            .join(", ")
    }
}

impl fmt::Display for SigningAlgorithm {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.write_str(self.as_str())
    }
}

impl std::str::FromStr for SigningAlgorithm {
    type Err = &'static str;
    fn from_str(s: &str) -> Result<SigningAlgorithm, &'static str> {
        use SigningAlgorithm::*;
        match s {
            "EdDSA" => Ok(EdDsa),
            "RS256" => Ok(Rs256),
            _ => Err("unsupported value"),
        }
    }
}

serde_from_str!(SigningAlgorithm);
serde_display!(SigningAlgorithm);

/// The types of public keys we support.
pub enum SupportedPublicKey {
    Ed25519(UnparsedPublicKey<Vec<u8>>),
    Rsa(RsaPublicKey),
}

impl SupportedPublicKey {
    /// Verify a message signature.
    pub fn verify(&self, message: &[u8], signature: &[u8]) -> Result<(), Unspecified> {
        use SupportedPublicKey::*;
        match self {
            Ed25519(ref inner) => inner.verify(message, signature),
            Rsa(ref inner) => {
                inner.verify(&signature::RSA_PKCS1_2048_8192_SHA256, message, signature)
            }
        }
    }
}

/// Helper function to build a session ID for a login attempt.
///
/// Put the email address, the client ID (RP origin) and some randomness into
/// a SHA256 hash, and encode it with URL-safe bas64 encoding. This is used
/// as the key in Redis, as well as the state for OAuth authentication.
pub async fn session_id(email: &EmailAddress, client_id: &str, rng: &SecureRandom) -> String {
    let rand_bytes = rng.generate_async(16).await;
    let mut ctx = digest::Context::new(&digest::SHA256);
    ctx.update(email.as_str().as_bytes());
    ctx.update(client_id.as_bytes());
    ctx.update(&rand_bytes);
    base64url::encode(&ctx.finish())
}

/// Helper function to create a secure nonce.
pub async fn nonce(rng: &SecureRandom) -> String {
    let rand_bytes = rng.generate_async(16).await;
    base64url::encode(&rand_bytes)
}

/// Helper function to create a random string consisting of
/// characters from the z-base-32 set.
pub async fn random_zbase32(len: usize, rng: &SecureRandom) -> String {
    const CHARSET: &[u8] = b"13456789abcdefghijkmnopqrstuwxyz";
    let rand_bytes = rng.generate_async(len).await;
    String::from_utf8(
        rand_bytes
            .into_iter()
            .map(|v| CHARSET[v as usize % CHARSET.len()])
            .collect(),
    )
    .expect("failed to build one-time pad")
}

#[derive(Debug, Error)]
pub enum VerifyError {
    #[error("the token must consist of three dot-separated parts")]
    IncorrectFormat,
    #[error("token part {index} contained invalid base64: {reason}")]
    InvalidPartBase64 {
        index: usize,
        reason: base64::DecodeError,
    },
    #[error("the token header contained invalid JSON: {0}")]
    InvalidHeaderJson(JsonError),
    #[error("did not find a string 'kid' property in the token header")]
    KidMissing,
    #[error("the token 'kid' could not be found in the JWKs document: {kid}")]
    KidNotMatched { kid: String },
    #[error(
        "the '{}' field of the matching JWK contains invalid base64: {}",
        property,
        reason
    )]
    InvalidJwkBase64 {
        property: &'static str,
        reason: base64::DecodeError,
    },
    #[error("the matching JWK is of an unsupported type")]
    UnsupportedKeyType,
    #[error("the token signature did not validate using the matching JWK")]
    BadSignature,
    #[error("the token payload contained invalid JSON: {0}")]
    InvalidPayloadJson(JsonError),
}

/// Verify a JWS signature, returning the payload as a `Value` if successful.
pub fn verify_jws(
    jws: &str,
    key_set: &[ProviderKey],
    signing_alg: SigningAlgorithm,
) -> Result<json::Value, VerifyError> {
    // Split the token up in parts and decode them.
    let parts: Vec<&str> = jws.split('.').collect();
    if parts.len() != 3 {
        return Err(VerifyError::IncorrectFormat);
    }
    let decoded = parts
        .iter()
        .enumerate()
        .map(|(idx, s)| {
            base64url::decode(s).map_err(|reason| VerifyError::InvalidPartBase64 {
                index: idx + 1,
                reason,
            })
        })
        .collect::<Result<Vec<_>, _>>()?;

    // Parse the header and find the key ID.
    let jwt_header: json::Value =
        json::from_slice(&decoded[0]).map_err(VerifyError::InvalidHeaderJson)?;
    let kid = jwt_header
        .get("kid")
        .and_then(Value::as_str)
        .ok_or(VerifyError::KidMissing)?;

    // Look for they key ID in the JWKs.
    let matched_keys: Vec<&ProviderKey> = key_set
        .iter()
        .filter(|key| key.use_ == "sig" && key.kid == kid)
        .collect();

    // Verify that we found exactly one key matching the key ID.
    if matched_keys.len() != 1 {
        return Err(VerifyError::KidNotMatched {
            kid: kid.to_owned(),
        });
    }
    let key = matched_keys.first().unwrap();

    // Then, use the data to build a public key object for verification.
    let pub_key = match (signing_alg, key.alg.as_str(), key.crv.as_str()) {
        (SigningAlgorithm::EdDsa, "EdDSA", "Ed25519") => {
            let x = base64url::decode(&key.x).map_err(|reason| VerifyError::InvalidJwkBase64 {
                property: "x",
                reason,
            })?;
            let key = UnparsedPublicKey::new(&signature::ED25519, x);
            SupportedPublicKey::Ed25519(key)
        }
        (SigningAlgorithm::Rs256, "RS256", _) => {
            let n = base64url::decode(&key.n).map_err(|reason| VerifyError::InvalidJwkBase64 {
                property: "n",
                reason,
            })?;
            let e = base64url::decode(&key.e).map_err(|reason| VerifyError::InvalidJwkBase64 {
                property: "e",
                reason,
            })?;
            let key = RsaPublicKey { n, e };
            SupportedPublicKey::Rsa(key)
        }
        _ => return Err(VerifyError::UnsupportedKeyType),
    };

    // Verify the signature using the public key.
    let message_len = parts[0].len() + parts[1].len() + 1;
    pub_key
        .verify(jws[..message_len].as_bytes(), &decoded[2])
        .map_err(|_err| VerifyError::BadSignature)?;

    // Return the payload.
    json::from_slice(&decoded[1]).map_err(VerifyError::InvalidPayloadJson)
}

/// Helper method to create a JWT for a given email address and audience.
///
/// Builds the JSON payload, then signs it using the last key provided in
/// the configuration object.
pub async fn create_jwt(
    app: &Config,
    email: &str,
    email_addr: &EmailAddress,
    aud: &str,
    nonce: &str,
    signing_alg: SigningAlgorithm,
) -> Result<String, SignError> {
    let now = unix_duration();
    app.key_manager
        .send(SignJws {
            payload: json!({
                "aud": aud,
                "email": email_addr.as_str(),
                "email_verified": true,
                "email_original": email,
                "exp": (now + app.token_ttl).as_secs(),
                "iat": now.as_secs(),
                "iss": &app.public_url,
                "sub": email_addr.as_str(),
                "nonce": nonce,
            }),
            signing_alg,
        })
        .await
}
<|endoftext|>


<commit_before>//! These are the unrolled butterflies for specific FFT lengths. They are mostly
//! translations of the butterfly functions from KissFFT, copyright Mark Borgerding.

use num::{Complex, Zero, Num, FromPrimitive, Signed};
use std::ops::Neg;

pub unsafe fn butterfly_2<T>(data: &mut [Complex<T>], stride: usize,
                             twiddles: &[Complex<T>], num_ffts: usize)
                             where T: Num + Copy {
    let mut idx_1 = 0usize;
    let mut idx_2 = num_ffts;
    let mut twiddle_idx = 0usize;
    for _ in 0..num_ffts {
        let twiddle = twiddles.get_unchecked(twiddle_idx);
        let temp = data.get_unchecked(idx_2) * twiddle;
        *data.get_unchecked_mut(idx_2) = data.get_unchecked(idx_1) - temp;
        *data.get_unchecked_mut(idx_1) = data.get_unchecked(idx_1) + temp;
        twiddle_idx += stride;
        idx_1 += 1;
        idx_2 += 1;
    }
}

pub unsafe fn butterfly_3<T>(data: &mut [Complex<T>], stride: usize,
                             twiddles: &[Complex<T>], num_ffts: usize)
                             where T: Num + Copy + FromPrimitive {
    let mut idx = 0usize;
    let mut tw_idx_1 = 0usize;
    let mut tw_idx_2 = 0usize;
    let mut scratch: [Complex<T>; 5] = [Zero::zero(); 5];
    for _ in 0..num_ffts {
        scratch[1] = data.get_unchecked(idx + 1 * num_ffts) * twiddles[tw_idx_1];
        scratch[2] = data.get_unchecked(idx + 2 * num_ffts) * twiddles[tw_idx_2];
        scratch[3] = scratch[1] + scratch[2];
        scratch[0] = scratch[1] - scratch[2];

        data.get_unchecked_mut(idx + num_ffts).re = data.get_unchecked(idx).re
                                                    - scratch[3].re
                                                    / FromPrimitive::from_f32(2.0).unwrap();
        data.get_unchecked_mut(idx + num_ffts).im = data.get_unchecked(idx).im
                                                    - scratch[3].im
                                                    / FromPrimitive::from_f32(2.0).unwrap();

        scratch[0].re = scratch[0].re * twiddles[stride * num_ffts].im;
        scratch[0].im = scratch[0].im * twiddles[stride * num_ffts].im;

        *data.get_unchecked_mut(idx) = data.get_unchecked(idx) + scratch[3];
        data.get_unchecked_mut(idx + 2 * num_ffts).re = data.get_unchecked(idx + num_ffts).re
                                                        + scratch[0].im;
        data.get_unchecked_mut(idx + 2 * num_ffts).im = data.get_unchecked(idx + num_ffts).im
                                                        - scratch[0].re;
        data.get_unchecked_mut(idx + num_ffts).re = data.get_unchecked(idx + num_ffts).re
                                                    - scratch[0].im;
        data.get_unchecked_mut(idx + num_ffts).im = data.get_unchecked(idx + num_ffts).im
                                                    + scratch[0].re;

        tw_idx_1 += 1 * stride;
        tw_idx_2 += 2 * stride;
        idx += 1;
    }
}

pub unsafe fn butterfly_4<T>(data: &mut [Complex<T>], stride: usize,
                             twiddles: &[Complex<T>], num_ffts: usize, inverse: bool)
                             where T: Num + Copy {
    let mut idx = 0usize;
    let mut tw_idx_1 = 0usize;
    let mut tw_idx_2 = 0usize;
    let mut tw_idx_3 = 0usize;
    let mut scratch: [Complex<T>; 6] = [Zero::zero(); 6];
    for _ in 0..num_ffts {
        scratch[0] = data.get_unchecked(idx + 1 * num_ffts) * twiddles[tw_idx_1];
        scratch[1] = data.get_unchecked(idx + 2 * num_ffts) * twiddles[tw_idx_2];
        scratch[2] = data.get_unchecked(idx + 3 * num_ffts) * twiddles[tw_idx_3];
        scratch[5] = data.get_unchecked(idx) - scratch[1];
        *data.get_unchecked_mut(idx) = data.get_unchecked(idx) + scratch[1];
        scratch[3] = scratch[0] + scratch[2];
        scratch[4] = scratch[0] - scratch[2];
        *data.get_unchecked_mut(idx + 2 * num_ffts) = data.get_unchecked(idx) - scratch[3];
        *data.get_unchecked_mut(idx) = data.get_unchecked(idx) + scratch[3];
        if inverse {
            data.get_unchecked_mut(idx + num_ffts).re = scratch[5].re - scratch[4].im;
            data.get_unchecked_mut(idx + num_ffts).im = scratch[5].im + scratch[4].re;
            data.get_unchecked_mut(idx + 3 * num_ffts).re = scratch[5].re + scratch[4].im;
            data.get_unchecked_mut(idx + 3 * num_ffts).im = scratch[5].im - scratch[4].re;
        } else {
            data.get_unchecked_mut(idx + num_ffts).re = scratch[5].re + scratch[4].im;
            data.get_unchecked_mut(idx + num_ffts).im = scratch[5].im - scratch[4].re;
            data.get_unchecked_mut(idx + 3 * num_ffts).re = scratch[5].re - scratch[4].im;
            data.get_unchecked_mut(idx + 3 * num_ffts).im = scratch[5].im + scratch[4].re;
        }

        tw_idx_1 += 1 * stride;
        tw_idx_2 += 2 * stride;
        tw_idx_3 += 3 * stride;
        idx += 1;
    }
}

pub unsafe fn butterfly_5<T>(data: &mut [Complex<T>], stride: usize,
                             twiddles: &[Complex<T>], num_ffts: usize)
                             where T: Signed + Copy {
    let mut idx_0 = 0 * num_ffts;
    let mut idx_1 = 1 * num_ffts;
    let mut idx_2 = 2 * num_ffts;
    let mut idx_3 = 3 * num_ffts;
    let mut idx_4 = 4 * num_ffts;
    let mut scratch: [Complex<T>; 13] = [Zero::zero(); 13];
    let ya = twiddles[stride * num_ffts];
    let yb = twiddles[stride * 2 * num_ffts];
    for i in 0..num_ffts {
        scratch[0] = *data.get_unchecked(idx_0);
        scratch[1] = data.get_unchecked(idx_1) * twiddles[1 * i * stride];
        scratch[2] = data.get_unchecked(idx_2) * twiddles[2 * i * stride];
        scratch[3] = data.get_unchecked(idx_3) * twiddles[3 * i * stride];
        scratch[4] = data.get_unchecked(idx_4) * twiddles[4 * i * stride];

        scratch[7] = scratch[1] + scratch[4];
        scratch[10] = scratch[1] - scratch[4];
        scratch[8] = scratch[2] + scratch[3];
        scratch[9] = scratch[2] - scratch[3];

        data.get_unchecked_mut(idx_0).re = data.get_unchecked_mut(idx_0).re
            + scratch[7].re + scratch[8].re;
        data.get_unchecked_mut(idx_0).im = data.get_unchecked_mut(idx_0).im
            + scratch[7].im + scratch[8].im;

        scratch[5].re = scratch[0].re + scratch[7].re * ya.re + scratch[8].re * yb.re;
        scratch[5].im = scratch[0].im + scratch[7].im * ya.re + scratch[8].im * yb.re;

        scratch[6].re = scratch[10].im * ya.im + scratch[9].im * yb.im;
        scratch[6].im = scratch[10].re.neg() * ya.im - scratch[9].re * yb.im;

        *data.get_unchecked_mut(idx_1) = scratch[5] - scratch[6];
        *data.get_unchecked_mut(idx_4) = scratch[5] + scratch[6];

        scratch[11].re = scratch[0].re + scratch[7].re * yb.re + scratch[8].re * ya.re;
        scratch[11].im = scratch[0].im + scratch[7].im * yb.re + scratch[8].im * ya.re;
        scratch[12].re = scratch[10].im.neg() * yb.im + scratch[9].im * ya.im;
        scratch[12].im = scratch[10].re * yb.im - scratch[9].re * ya.im;

        *data.get_unchecked_mut(idx_2) = scratch[11] + scratch[12];
        *data.get_unchecked_mut(idx_3) = scratch[11] - scratch[12];

        idx_0 += 1;
        idx_1 += 1;
        idx_2 += 1;
        idx_3 += 1;
        idx_4 += 1;
    }
}

#[cfg(test)]
mod benches {
    extern crate test;

    use std::f32;
    use std::iter::repeat;
    use self::test::Bencher;
    use num::{Complex, Zero};
    use super::{butterfly_2, butterfly_3, butterfly_4, butterfly_5};

    fn set_up_bench(butterfly_len: usize, stride: usize, num_ffts: usize) -> (Vec<Complex<f32>>, Vec<Complex<f32>>) {
        let len = butterfly_len * stride * num_ffts;
        let twiddles: Vec<Complex<f32>> = (0..len)
            .map(|i| -1. * (i as f32) * 2.0 * f32::consts::PI / (len as f32))
            .map(|phase| Complex::from_polar(&1., &phase))
            .collect();

        let data: Vec<Complex<f32>> = repeat(Zero::zero()).take(len).collect();
        (data, twiddles)
    }

    #[bench]
    fn bench_butterfly_2(b: &mut Bencher) {
        let stride = 4usize;
        let num_ffts = 1000usize;
        let (mut data, twiddles) = set_up_bench(2, stride, num_ffts);

        b.iter(|| unsafe {
            butterfly_2(&mut data, stride, &twiddles, num_ffts)
        });
    }

    #[bench]
    fn bench_butterfly_3(b: &mut Bencher) {
        let stride = 4usize;
        let num_ffts = 1000usize;
        let (mut data, twiddles) = set_up_bench(3, stride, num_ffts);

        b.iter(|| unsafe {
            butterfly_3(&mut data, stride, &twiddles, num_ffts)
        });
    }

    #[bench]
    fn bench_butterfly_4(b: &mut Bencher) {
        let stride = 4usize;
        let num_ffts = 1000usize;
        let (mut data, twiddles) = set_up_bench(4, stride, num_ffts);

        b.iter(|| unsafe {
            butterfly_4(&mut data, stride, &twiddles, num_ffts, false)
        });
    }

    #[bench]
    fn bench_butterfly_5(b: &mut Bencher) {
        let stride = 4usize;
        let num_ffts = 1000usize;
        let (mut data, twiddles) = set_up_bench(5, stride, num_ffts);

        b.iter(|| unsafe {
            butterfly_5(&mut data, stride, &twiddles, num_ffts)
        });
    }
}
<commit_msg>Speed up butterfly_2 by doing complex addition ourselves. Matches the performance of kiss_fft's butterfly.<commit_after>//! These are the unrolled butterflies for specific FFT lengths. They are mostly
//! translations of the butterfly functions from KissFFT, copyright Mark Borgerding.

use num::{Complex, Zero, Num, FromPrimitive, Signed};
use std::ops::Neg;

pub unsafe fn butterfly_2<T>(data: &mut [Complex<T>], stride: usize,
                             twiddles: &[Complex<T>], num_ffts: usize)
                             where T: Num + Copy {
    let mut idx_1 = 0usize;
    let mut idx_2 = num_ffts;
    let mut twiddle_idx = 0usize;
    for _ in 0..num_ffts {
        let twiddle = twiddles.get_unchecked(twiddle_idx);
        let temp = data.get_unchecked(idx_2) * twiddle;
        data.get_unchecked_mut(idx_2).re = data.get_unchecked(idx_1).re - temp.re;
        data.get_unchecked_mut(idx_2).im = data.get_unchecked(idx_1).im - temp.im;
        data.get_unchecked_mut(idx_1).re = data.get_unchecked(idx_1).re + temp.re;
        data.get_unchecked_mut(idx_1).im = data.get_unchecked(idx_1).im + temp.im;
        twiddle_idx += stride;
        idx_1 += 1;
        idx_2 += 1;
    }
}

pub unsafe fn butterfly_3<T>(data: &mut [Complex<T>], stride: usize,
                             twiddles: &[Complex<T>], num_ffts: usize)
                             where T: Num + Copy + FromPrimitive {
    let mut idx = 0usize;
    let mut tw_idx_1 = 0usize;
    let mut tw_idx_2 = 0usize;
    let mut scratch: [Complex<T>; 5] = [Zero::zero(); 5];
    for _ in 0..num_ffts {
        scratch[1] = data.get_unchecked(idx + 1 * num_ffts) * twiddles[tw_idx_1];
        scratch[2] = data.get_unchecked(idx + 2 * num_ffts) * twiddles[tw_idx_2];
        scratch[3] = scratch[1] + scratch[2];
        scratch[0] = scratch[1] - scratch[2];

        data.get_unchecked_mut(idx + num_ffts).re = data.get_unchecked(idx).re
                                                    - scratch[3].re
                                                    / FromPrimitive::from_f32(2.0).unwrap();
        data.get_unchecked_mut(idx + num_ffts).im = data.get_unchecked(idx).im
                                                    - scratch[3].im
                                                    / FromPrimitive::from_f32(2.0).unwrap();

        scratch[0].re = scratch[0].re * twiddles[stride * num_ffts].im;
        scratch[0].im = scratch[0].im * twiddles[stride * num_ffts].im;

        *data.get_unchecked_mut(idx) = data.get_unchecked(idx) + scratch[3];
        data.get_unchecked_mut(idx + 2 * num_ffts).re = data.get_unchecked(idx + num_ffts).re
                                                        + scratch[0].im;
        data.get_unchecked_mut(idx + 2 * num_ffts).im = data.get_unchecked(idx + num_ffts).im
                                                        - scratch[0].re;
        data.get_unchecked_mut(idx + num_ffts).re = data.get_unchecked(idx + num_ffts).re
                                                    - scratch[0].im;
        data.get_unchecked_mut(idx + num_ffts).im = data.get_unchecked(idx + num_ffts).im
                                                    + scratch[0].re;

        tw_idx_1 += 1 * stride;
        tw_idx_2 += 2 * stride;
        idx += 1;
    }
}

pub unsafe fn butterfly_4<T>(data: &mut [Complex<T>], stride: usize,
                             twiddles: &[Complex<T>], num_ffts: usize, inverse: bool)
                             where T: Num + Copy {
    let mut idx = 0usize;
    let mut tw_idx_1 = 0usize;
    let mut tw_idx_2 = 0usize;
    let mut tw_idx_3 = 0usize;
    let mut scratch: [Complex<T>; 6] = [Zero::zero(); 6];
    for _ in 0..num_ffts {
        scratch[0] = data.get_unchecked(idx + 1 * num_ffts) * twiddles[tw_idx_1];
        scratch[1] = data.get_unchecked(idx + 2 * num_ffts) * twiddles[tw_idx_2];
        scratch[2] = data.get_unchecked(idx + 3 * num_ffts) * twiddles[tw_idx_3];
        scratch[5] = data.get_unchecked(idx) - scratch[1];
        *data.get_unchecked_mut(idx) = data.get_unchecked(idx) + scratch[1];
        scratch[3] = scratch[0] + scratch[2];
        scratch[4] = scratch[0] - scratch[2];
        *data.get_unchecked_mut(idx + 2 * num_ffts) = data.get_unchecked(idx) - scratch[3];
        *data.get_unchecked_mut(idx) = data.get_unchecked(idx) + scratch[3];
        if inverse {
            data.get_unchecked_mut(idx + num_ffts).re = scratch[5].re - scratch[4].im;
            data.get_unchecked_mut(idx + num_ffts).im = scratch[5].im + scratch[4].re;
            data.get_unchecked_mut(idx + 3 * num_ffts).re = scratch[5].re + scratch[4].im;
            data.get_unchecked_mut(idx + 3 * num_ffts).im = scratch[5].im - scratch[4].re;
        } else {
            data.get_unchecked_mut(idx + num_ffts).re = scratch[5].re + scratch[4].im;
            data.get_unchecked_mut(idx + num_ffts).im = scratch[5].im - scratch[4].re;
            data.get_unchecked_mut(idx + 3 * num_ffts).re = scratch[5].re - scratch[4].im;
            data.get_unchecked_mut(idx + 3 * num_ffts).im = scratch[5].im + scratch[4].re;
        }

        tw_idx_1 += 1 * stride;
        tw_idx_2 += 2 * stride;
        tw_idx_3 += 3 * stride;
        idx += 1;
    }
}

pub unsafe fn butterfly_5<T>(data: &mut [Complex<T>], stride: usize,
                             twiddles: &[Complex<T>], num_ffts: usize)
                             where T: Signed + Copy {
    let mut idx_0 = 0 * num_ffts;
    let mut idx_1 = 1 * num_ffts;
    let mut idx_2 = 2 * num_ffts;
    let mut idx_3 = 3 * num_ffts;
    let mut idx_4 = 4 * num_ffts;
    let mut scratch: [Complex<T>; 13] = [Zero::zero(); 13];
    let ya = twiddles[stride * num_ffts];
    let yb = twiddles[stride * 2 * num_ffts];
    for i in 0..num_ffts {
        scratch[0] = *data.get_unchecked(idx_0);
        scratch[1] = data.get_unchecked(idx_1) * twiddles[1 * i * stride];
        scratch[2] = data.get_unchecked(idx_2) * twiddles[2 * i * stride];
        scratch[3] = data.get_unchecked(idx_3) * twiddles[3 * i * stride];
        scratch[4] = data.get_unchecked(idx_4) * twiddles[4 * i * stride];

        scratch[7] = scratch[1] + scratch[4];
        scratch[10] = scratch[1] - scratch[4];
        scratch[8] = scratch[2] + scratch[3];
        scratch[9] = scratch[2] - scratch[3];

        data.get_unchecked_mut(idx_0).re = data.get_unchecked_mut(idx_0).re
            + scratch[7].re + scratch[8].re;
        data.get_unchecked_mut(idx_0).im = data.get_unchecked_mut(idx_0).im
            + scratch[7].im + scratch[8].im;

        scratch[5].re = scratch[0].re + scratch[7].re * ya.re + scratch[8].re * yb.re;
        scratch[5].im = scratch[0].im + scratch[7].im * ya.re + scratch[8].im * yb.re;

        scratch[6].re = scratch[10].im * ya.im + scratch[9].im * yb.im;
        scratch[6].im = scratch[10].re.neg() * ya.im - scratch[9].re * yb.im;

        *data.get_unchecked_mut(idx_1) = scratch[5] - scratch[6];
        *data.get_unchecked_mut(idx_4) = scratch[5] + scratch[6];

        scratch[11].re = scratch[0].re + scratch[7].re * yb.re + scratch[8].re * ya.re;
        scratch[11].im = scratch[0].im + scratch[7].im * yb.re + scratch[8].im * ya.re;
        scratch[12].re = scratch[10].im.neg() * yb.im + scratch[9].im * ya.im;
        scratch[12].im = scratch[10].re * yb.im - scratch[9].re * ya.im;

        *data.get_unchecked_mut(idx_2) = scratch[11] + scratch[12];
        *data.get_unchecked_mut(idx_3) = scratch[11] - scratch[12];

        idx_0 += 1;
        idx_1 += 1;
        idx_2 += 1;
        idx_3 += 1;
        idx_4 += 1;
    }
}

#[cfg(test)]
mod benches {
    extern crate test;

    use std::f32;
    use std::iter::repeat;
    use self::test::Bencher;
    use num::{Complex, Zero};
    use super::{butterfly_2, butterfly_3, butterfly_4, butterfly_5};

    fn set_up_bench(butterfly_len: usize, stride: usize, num_ffts: usize) -> (Vec<Complex<f32>>, Vec<Complex<f32>>) {
        let len = butterfly_len * stride * num_ffts;
        let twiddles: Vec<Complex<f32>> = (0..len)
            .map(|i| -1. * (i as f32) * 2.0 * f32::consts::PI / (len as f32))
            .map(|phase| Complex::from_polar(&1., &phase))
            .collect();

        let data: Vec<Complex<f32>> = repeat(Zero::zero()).take(len).collect();
        (data, twiddles)
    }

    #[bench]
    fn bench_butterfly_2(b: &mut Bencher) {
        let stride = 4usize;
        let num_ffts = 1000usize;
        let (mut data, twiddles) = set_up_bench(2, stride, num_ffts);

        b.iter(|| unsafe {
            butterfly_2(&mut data, stride, &twiddles, num_ffts)
        });
    }

    #[bench]
    fn bench_butterfly_3(b: &mut Bencher) {
        let stride = 4usize;
        let num_ffts = 1000usize;
        let (mut data, twiddles) = set_up_bench(3, stride, num_ffts);

        b.iter(|| unsafe {
            butterfly_3(&mut data, stride, &twiddles, num_ffts)
        });
    }

    #[bench]
    fn bench_butterfly_4(b: &mut Bencher) {
        let stride = 4usize;
        let num_ffts = 1000usize;
        let (mut data, twiddles) = set_up_bench(4, stride, num_ffts);

        b.iter(|| unsafe {
            butterfly_4(&mut data, stride, &twiddles, num_ffts, false)
        });
    }

    #[bench]
    fn bench_butterfly_5(b: &mut Bencher) {
        let stride = 4usize;
        let num_ffts = 1000usize;
        let (mut data, twiddles) = set_up_bench(5, stride, num_ffts);

        b.iter(|| unsafe {
            butterfly_5(&mut data, stride, &twiddles, num_ffts)
        });
    }
}
<|endoftext|>


<commit_before>use std::path::Path;
use std::slice::Iter;
use std::ops::{Index, IndexMut};
use time::now_utc;
use xi_rope::Rope;

use super::Buffer;

/// Creates, manages and deletes all the buffers in Qork, maintaining the various invariants that
/// we expect from the buffers. Firstly, if a buffer is backed by a file, a second buffer on that
/// file cannot be created. Secondly, all buffers have unique identity.
///
/// Note that a Buffer is very different from a BufferView.
pub struct BufferCollection {
    next_buffer_id: i64,
    buffers: Vec<Buffer>
}

impl BufferCollection {
    fn next_buffer_id(&mut self) -> i64 {
        // No point being clever. No practical scenario in which these will need to be reused.
        let i = self.next_buffer_id;
        self.next_buffer_id += 1;
        i
    }

    pub fn new() -> BufferCollection {
        BufferCollection {
            next_buffer_id: 0,
            buffers: Vec::new()
        }
    }

    /// Creates and returns a new empty buffer.
    pub fn new_empty_buffer(&mut self) -> &mut Buffer {
        let now = now_utc();

        let b = Buffer {
            id: self.next_buffer_id(),
            filename: None,
            title: String::default(),
            data: Rope::from(""),
            is_changed: false,
            created_time_utc: now,
            last_accessed_time_utc: now,
            last_changed_time_utc: now
        };

        let i = self.buffers.len();
        self.buffers.push(b);
        &mut self.buffers[i]
    }

    /*
    /// Creates a buffer from a filename. If there is already a Buffer for the file it is returned,
    /// else the file is opened and loaded if it exists, else if the file does not exist then a
    /// new buffer is created with that filename, but no loading is done (the Buffer is considered
    /// to be backed by a file that does not exist yet, it will be created when you save it.)
    pub fn open_file(&mut self) -> Buffer {
        Buffer::new()
    }
    */

    pub fn len(&mut self) -> usize {
        self.buffers.len()
    }

    pub fn is_empty(&self) -> bool {
        self.buffers.is_empty()
    }

    pub fn iter(&self) -> Iter<Buffer> {
        self.buffers.iter()
    }

    pub fn add(&mut self, buffer: Buffer) {
        self.buffers.push(buffer)
    }

    pub fn find_by_filename<P : AsRef<Path>>(&mut self, filename: P) -> Option<&mut Buffer> {
        self.buffers.iter_mut().find(|b| b.filename.as_ref().map_or(false, |f| f == filename.as_ref()))
    }
}

impl Index<usize> for BufferCollection {
    type Output = Buffer;

    fn index(&self, index: usize) -> &Buffer {
        &self.buffers[index]
    }
}

impl IndexMut<usize> for BufferCollection {
    fn index_mut<'a>(&'a mut self, index: usize) -> &'a mut Buffer {
        &mut self.buffers[index]
    }
}

#[cfg(test)]
mod buffer_collection_tests {
    use super::*;
    use std::path::PathBuf;

    /*
    #[test]
    fn add_works() {
        let mut bc = BufferCollection::new();
        let mut b = Buffer::new();
        b.filename = Some(PathBuf::from("a"));
        bc.add(b);

        let b2 = &bc[0];
        assert_eq!(b2.filename, Some(PathBuf::from("a")));
    }
    */

    #[test]
    fn new_empty_buffer_adds_and_returns_buffer() {
        let now = now_utc();
        let mut bc = BufferCollection::new();
        let b1 = Buffer::new();
        let b2 = Buffer::new();
        bc.add(b1);
        bc.add(b2);

        // let b1 = bc.new_empty_buffer();
        // let b2 = bc.new_empty_buffer();
        // let i = bc.len();
        // assert_eq!(0, b1.id);
        // assert_eq!(1, i);

        // let b1 = Buffer {
        //     id: 2,
        //     filename: None,
        //     title: String::default(),
        //     data: Rope::from(""),
        //     is_changed: false,
        //     created_time_utc: now,
        //     last_accessed_time_utc: now,
        //     last_changed_time_utc: now
        // };

        // let b2 = Buffer {
        //     id: 2,
        //     filename: None,
        //     title: String::default(),
        //     data: Rope::from(""),
        //     is_changed: false,
        //     created_time_utc: now,
        //     last_accessed_time_utc: now,
        //     last_changed_time_utc: now
        // };

        //let b2 = &bc[0];
        //assert_eq!(b1, b2);
    }
}
<commit_msg>Buffers: number ids from 0.<commit_after>use std::path::Path;
use std::slice::Iter;
use std::ops::{Index, IndexMut};
use time::now_utc;
use xi_rope::Rope;

use super::Buffer;

/// Creates, manages and deletes all the buffers in Qork, maintaining the various invariants that
/// we expect from the buffers. Firstly, if a buffer is backed by a file, a second buffer on that
/// file cannot be created. Secondly, all buffers have unique identity.
///
/// Note that a Buffer is very different from a BufferView.
pub struct BufferCollection {
    next_buffer_id: i64,
    buffers: Vec<Buffer>
}

impl BufferCollection {
    fn next_buffer_id(&mut self) -> i64 {
        // No point being clever. No practical scenario in which these will need to be reused.
        self.next_buffer_id += 1;
        self.next_buffer_id
    }

    pub fn new() -> BufferCollection {
        BufferCollection {
            next_buffer_id: -1,
            buffers: Vec::new()
        }
    }

    /// Creates and returns a new empty buffer.
    pub fn new_empty_buffer(&mut self) -> &mut Buffer {
        let now = now_utc();

        let b = Buffer {
            id: self.next_buffer_id(),
            filename: None,
            title: String::default(),
            data: Rope::from(""),
            is_changed: false,
            created_time_utc: now,
            last_accessed_time_utc: now,
            last_changed_time_utc: now
        };

        let i = self.buffers.len();
        self.buffers.push(b);
        &mut self.buffers[i]
    }

    /*
    /// Creates a buffer from a filename. If there is already a Buffer for the file it is returned,
    /// else the file is opened and loaded if it exists, else if the file does not exist then a
    /// new buffer is created with that filename, but no loading is done (the Buffer is considered
    /// to be backed by a file that does not exist yet, it will be created when you save it.)
    pub fn open_file(&mut self) -> Buffer {
        Buffer::new()
    }
    */

    pub fn len(&mut self) -> usize {
        self.buffers.len()
    }

    pub fn is_empty(&self) -> bool {
        self.buffers.is_empty()
    }

    pub fn iter(&self) -> Iter<Buffer> {
        self.buffers.iter()
    }

    pub fn add(&mut self, buffer: Buffer) {
        self.buffers.push(buffer)
    }

    pub fn find_by_filename<P : AsRef<Path>>(&mut self, filename: P) -> Option<&mut Buffer> {
        self.buffers.iter_mut().find(|b| b.filename.as_ref().map_or(false, |f| f == filename.as_ref()))
    }
}

impl Index<usize> for BufferCollection {
    type Output = Buffer;

    fn index(&self, index: usize) -> &Buffer {
        &self.buffers[index]
    }
}

impl IndexMut<usize> for BufferCollection {
    fn index_mut<'a>(&'a mut self, index: usize) -> &'a mut Buffer {
        &mut self.buffers[index]
    }
}

#[cfg(test)]
mod buffer_collection_tests {
    use super::*;
    use std::path::PathBuf;

    /*
    #[test]
    fn add_works() {
        let mut bc = BufferCollection::new();
        let mut b = Buffer::new();
        b.filename = Some(PathBuf::from("a"));
        bc.add(b);

        let b2 = &bc[0];
        assert_eq!(b2.filename, Some(PathBuf::from("a")));
    }
    */

    #[test]
    fn new_empty_buffer_adds_and_returns_buffer() {
        let now = now_utc();
        let mut bc = BufferCollection::new();
        let b1 = Buffer::new();
        let b2 = Buffer::new();
        bc.add(b1);
        bc.add(b2);

        // let b1 = bc.new_empty_buffer();
        // let b2 = bc.new_empty_buffer();
        // let i = bc.len();
        // assert_eq!(0, b1.id);
        // assert_eq!(1, i);

        // let b1 = Buffer {
        //     id: 2,
        //     filename: None,
        //     title: String::default(),
        //     data: Rope::from(""),
        //     is_changed: false,
        //     created_time_utc: now,
        //     last_accessed_time_utc: now,
        //     last_changed_time_utc: now
        // };

        // let b2 = Buffer {
        //     id: 2,
        //     filename: None,
        //     title: String::default(),
        //     data: Rope::from(""),
        //     is_changed: false,
        //     created_time_utc: now,
        //     last_accessed_time_utc: now,
        //     last_changed_time_utc: now
        // };

        //let b2 = &bc[0];
        //assert_eq!(b1, b2);
    }
}
<|endoftext|>


<commit_before>mod init;
mod config;
mod ignore;
mod sweep;
mod burn;
mod start;
mod end;
mod destroy;

use self::init::Init;
use self::config::Config;
use self::ignore::Ignore;
use self::sweep::Sweep;
use self::burn::Burn;
use self::start::Start;
use self::end::End;
use self::destroy::Destroy;

use error::*;
use lib::io::*;
use lib::setting::*;
use std::process;

trait Command {
    fn validation(&self) -> bool;
    fn validate(&self) -> Result<(), EssentialLack> {
        if !working_dir_exists() {
            return Err(EssentialLack { what: EssentialKind::WorkingDir });
        }
        if !storage_dir_exists() {
            return Err(EssentialLack { what: EssentialKind::StorageDir });
        }
        if !config_file_exists() {
            return Err(EssentialLack { what: EssentialKind::ConfigFile });
        }
        if !ignore_file_exists() {
            return Err(EssentialLack { what: EssentialKind::IgnoreFile });
        }

        Ok(())
    }

    fn usage(&self) -> Usage;
    fn help(&self) -> ! {
        println!("{}", self.usage());
        process::exit(1)
    }

    fn main(&self);

    fn exec(&self, help: bool) {
        if self.validation() {
            if let Err(why) = self.validate() {
                print_with_warning(0, why);
            }
        }

        if help {
            self.help();
        } else {
            self.main();
        }
    }
}

pub fn execute(args: Vec<String>) {
    let help = args.iter().any(|a| *a == "-h" || *a == "--help");

    let mut args = args.into_iter();

    let command: Box<Command> = match args.next() {
        Some(first) => match first.as_ref() {
            "init"    => Box::new(Init   ),
            "config"  => Box::new(Config ),
            "ignore"  => Box::new(Ignore { command: args.next(), paths: args.collect() }),
            "sweep"   => Box::new(Sweep  ),
            "burn"    => Box::new(Burn   ),
            "start"   => Box::new(Start  ),
            "end"     => Box::new(End    ),
            "destroy" => Box::new(Destroy),
            _         => print_usage(Usage { kind: UsageKind::Nothing }),
        },
        None => print_usage(Usage { kind: UsageKind::Nothing }),
    };

    command.exec(help);
}

pub fn print_usage(u: Usage) -> ! {
    println!("{}", u);
    process::exit(1)
}
<commit_msg>[Refactoring] Usage print_usage()<commit_after>mod init;
mod config;
mod ignore;
mod sweep;
mod burn;
mod start;
mod end;
mod destroy;

use self::init::Init;
use self::config::Config;
use self::ignore::Ignore;
use self::sweep::Sweep;
use self::burn::Burn;
use self::start::Start;
use self::end::End;
use self::destroy::Destroy;

use error::*;
use lib::io::*;
use lib::setting::*;
use std::process;

trait Command {
    fn validation(&self) -> bool;
    fn validate(&self) -> Result<(), EssentialLack> {
        if !working_dir_exists() {
            return Err(EssentialLack { what: EssentialKind::WorkingDir });
        }
        if !storage_dir_exists() {
            return Err(EssentialLack { what: EssentialKind::StorageDir });
        }
        if !config_file_exists() {
            return Err(EssentialLack { what: EssentialKind::ConfigFile });
        }
        if !ignore_file_exists() {
            return Err(EssentialLack { what: EssentialKind::IgnoreFile });
        }

        Ok(())
    }

    fn usage(&self) -> Usage;
    fn help(&self) -> ! {
        print_usage(self.usage());
    }

    fn main(&self);

    fn exec(&self, help: bool) {
        if self.validation() {
            if let Err(why) = self.validate() {
                print_with_warning(0, why);
            }
        }

        if help {
            self.help();
        } else {
            self.main();
        }
    }
}

pub fn execute(args: Vec<String>) {
    let help = args.iter().any(|a| *a == "-h" || *a == "--help");

    let mut args = args.into_iter();

    let command: Box<Command> = match args.next() {
        Some(first) => match first.as_ref() {
            "init"    => Box::new(Init   ),
            "config"  => Box::new(Config ),
            "ignore"  => Box::new(Ignore { command: args.next(), paths: args.collect() }),
            "sweep"   => Box::new(Sweep  ),
            "burn"    => Box::new(Burn   ),
            "start"   => Box::new(Start  ),
            "end"     => Box::new(End    ),
            "destroy" => Box::new(Destroy),
            _         => print_usage(Usage { kind: UsageKind::Nothing }),
        },
        None => print_usage(Usage { kind: UsageKind::Nothing }),
    };

    command.exec(help);
}

pub fn print_usage(u: Usage) -> ! {
    println!("{}", u);
    process::exit(1)
}
<|endoftext|>


<commit_before>//! Possible ZIP compression methods.

use std::fmt;

#[allow(deprecated)]
/// Identifies the storage format used to compress a file within a ZIP archive.
///
/// Each file's compression method is stored alongside it, allowing the
/// contents to be read without context.
///
/// When creating ZIP files, you may choose the method to use with
/// [`crate::write::FileOptions::compression_method`]
#[derive(Copy, Clone, PartialEq, Eq, Debug)]
pub enum CompressionMethod {
    /// Store the file as is
    Stored,
    /// Compress the file using Deflate
    #[cfg(any(
        feature = "deflate",
        feature = "deflate-miniz",
        feature = "deflate-zlib"
    ))]
    Deflated,
    /// Compress the file using BZIP2
    #[cfg(feature = "bzip2")]
    Bzip2,
    /// Encrypted using AES.
    ///
    /// The actual compression method has to be taken from the AES extra data field
    /// or from `ZipFileData`.
    #[cfg(feature = "aes-crypto")]
    Aes,
    /// Compress the file using ZStandard
    #[cfg(feature = "zstd")]
    Zstd,
    /// Unsupported compression method
    #[deprecated(since = "0.5.7", note = "use the constants instead")]
    Unsupported(u16),
}
#[allow(deprecated, missing_docs)]
/// All compression methods defined for the ZIP format
impl CompressionMethod {
    pub const STORE: Self = CompressionMethod::Stored;
    pub const SHRINK: Self = CompressionMethod::Unsupported(1);
    pub const REDUCE_1: Self = CompressionMethod::Unsupported(2);
    pub const REDUCE_2: Self = CompressionMethod::Unsupported(3);
    pub const REDUCE_3: Self = CompressionMethod::Unsupported(4);
    pub const REDUCE_4: Self = CompressionMethod::Unsupported(5);
    pub const IMPLODE: Self = CompressionMethod::Unsupported(6);
    #[cfg(any(
        feature = "deflate",
        feature = "deflate-miniz",
        feature = "deflate-zlib"
    ))]
    pub const DEFLATE: Self = CompressionMethod::Deflated;
    #[cfg(not(any(
        feature = "deflate",
        feature = "deflate-miniz",
        feature = "deflate-zlib"
    )))]
    pub const DEFLATE: Self = CompressionMethod::Unsupported(8);
    pub const DEFLATE64: Self = CompressionMethod::Unsupported(9);
    pub const PKWARE_IMPLODE: Self = CompressionMethod::Unsupported(10);
    #[cfg(feature = "bzip2")]
    pub const BZIP2: Self = CompressionMethod::Bzip2;
    #[cfg(not(feature = "bzip2"))]
    pub const BZIP2: Self = CompressionMethod::Unsupported(12);
    pub const LZMA: Self = CompressionMethod::Unsupported(14);
    pub const IBM_ZOS_CMPSC: Self = CompressionMethod::Unsupported(16);
    pub const IBM_TERSE: Self = CompressionMethod::Unsupported(18);
    pub const ZSTD_DEPRECATED: Self = CompressionMethod::Unsupported(20);
    #[cfg(feature = "zstd")]
    pub const ZSTD: Self = CompressionMethod::Zstd;
    #[cfg(not(feature = "zstd"))]
    pub const ZSTD: Self = CompressionMethod::Unsupported(93);
    pub const MP3: Self = CompressionMethod::Unsupported(94);
    pub const XZ: Self = CompressionMethod::Unsupported(95);
    pub const JPEG: Self = CompressionMethod::Unsupported(96);
    pub const WAVPACK: Self = CompressionMethod::Unsupported(97);
    pub const PPMD: Self = CompressionMethod::Unsupported(98);
    #[cfg(feature = "aes-crypto")]
    pub const AES: Self = CompressionMethod::Aes;
    #[cfg(not(feature = "aes-crypto"))]
    pub const AES: Self = CompressionMethod::Unsupported(99);
}
impl CompressionMethod {
    /// Converts an u16 to its corresponding CompressionMethod
    #[deprecated(
        since = "0.5.7",
        note = "use a constant to construct a compression method"
    )]
    pub fn from_u16(val: u16) -> CompressionMethod {
        #[allow(deprecated)]
        match val {
            0 => CompressionMethod::Stored,
            #[cfg(any(
                feature = "deflate",
                feature = "deflate-miniz",
                feature = "deflate-zlib"
            ))]
            8 => CompressionMethod::Deflated,
            #[cfg(feature = "bzip2")]
            12 => CompressionMethod::Bzip2,
            #[cfg(feature = "zstd")]
            93 => CompressionMethod::Zstd,
            99 => CompressionMethod::AES,

            v => CompressionMethod::Unsupported(v),
        }
    }

    /// Converts a CompressionMethod to a u16
    #[deprecated(
        since = "0.5.7",
        note = "to match on other compression methods, use a constant"
    )]
    pub fn to_u16(self) -> u16 {
        #[allow(deprecated)]
        match self {
            CompressionMethod::Stored => 0,
            #[cfg(any(
                feature = "deflate",
                feature = "deflate-miniz",
                feature = "deflate-zlib"
            ))]
            CompressionMethod::Deflated => 8,
            #[cfg(feature = "bzip2")]
            CompressionMethod::Bzip2 => 12,
            CompressionMethod::AES => 99,
            #[cfg(feature = "zstd")]
            CompressionMethod::Zstd => 93,

            CompressionMethod::Unsupported(v) => v,
        }
    }
}

impl fmt::Display for CompressionMethod {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        // Just duplicate what the Debug format looks like, i.e, the enum key:
        write!(f, "{:?}", self)
    }
}

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

    #[test]
    fn from_eq_to() {
        for v in 0..(u16::MAX as u32 + 1) {
            #[allow(deprecated)]
            let from = CompressionMethod::from_u16(v as u16);
            #[allow(deprecated)]
            let to = from.to_u16() as u32;
            assert_eq!(v, to);
        }
    }

    fn methods() -> Vec<CompressionMethod> {
        vec![
            CompressionMethod::Stored,
            #[cfg(any(
                feature = "deflate",
                feature = "deflate-miniz",
                feature = "deflate-zlib"
            ))]
            CompressionMethod::Deflated,
            #[cfg(feature = "bzip2")]
            CompressionMethod::Bzip2,
            #[cfg(feature = "zstd")]
            CompressionMethod::Zstd,
        ]
    }

    #[test]
    fn to_eq_from() {
        fn check_match(method: CompressionMethod) {
            #[allow(deprecated)]
            let to = method.to_u16();
            #[allow(deprecated)]
            let from = CompressionMethod::from_u16(to);
            #[allow(deprecated)]
            let back = from.to_u16();
            assert_eq!(to, back);
        }

        for method in methods() {
            check_match(method);
        }
    }

    #[test]
    fn to_display_fmt() {
        fn check_match(method: CompressionMethod) {
            let debug_str = format!("{:?}", method);
            let display_str = format!("{}", method);
            assert_eq!(debug_str, display_str);
        }

        for method in methods() {
            check_match(method);
        }
    }
}
<commit_msg>remove `CompressionMethod::Aes` enum variant<commit_after>//! Possible ZIP compression methods.

use std::fmt;

#[allow(deprecated)]
/// Identifies the storage format used to compress a file within a ZIP archive.
///
/// Each file's compression method is stored alongside it, allowing the
/// contents to be read without context.
///
/// When creating ZIP files, you may choose the method to use with
/// [`crate::write::FileOptions::compression_method`]
#[derive(Copy, Clone, PartialEq, Eq, Debug)]
pub enum CompressionMethod {
    /// Store the file as is
    Stored,
    /// Compress the file using Deflate
    #[cfg(any(
        feature = "deflate",
        feature = "deflate-miniz",
        feature = "deflate-zlib"
    ))]
    Deflated,
    /// Compress the file using BZIP2
    #[cfg(feature = "bzip2")]
    Bzip2,
    /// Compress the file using ZStandard
    #[cfg(feature = "zstd")]
    Zstd,
    /// Unsupported compression method
    #[deprecated(since = "0.5.7", note = "use the constants instead")]
    Unsupported(u16),
}
#[allow(deprecated, missing_docs)]
/// All compression methods defined for the ZIP format
impl CompressionMethod {
    pub const STORE: Self = CompressionMethod::Stored;
    pub const SHRINK: Self = CompressionMethod::Unsupported(1);
    pub const REDUCE_1: Self = CompressionMethod::Unsupported(2);
    pub const REDUCE_2: Self = CompressionMethod::Unsupported(3);
    pub const REDUCE_3: Self = CompressionMethod::Unsupported(4);
    pub const REDUCE_4: Self = CompressionMethod::Unsupported(5);
    pub const IMPLODE: Self = CompressionMethod::Unsupported(6);
    #[cfg(any(
        feature = "deflate",
        feature = "deflate-miniz",
        feature = "deflate-zlib"
    ))]
    pub const DEFLATE: Self = CompressionMethod::Deflated;
    #[cfg(not(any(
        feature = "deflate",
        feature = "deflate-miniz",
        feature = "deflate-zlib"
    )))]
    pub const DEFLATE: Self = CompressionMethod::Unsupported(8);
    pub const DEFLATE64: Self = CompressionMethod::Unsupported(9);
    pub const PKWARE_IMPLODE: Self = CompressionMethod::Unsupported(10);
    #[cfg(feature = "bzip2")]
    pub const BZIP2: Self = CompressionMethod::Bzip2;
    #[cfg(not(feature = "bzip2"))]
    pub const BZIP2: Self = CompressionMethod::Unsupported(12);
    pub const LZMA: Self = CompressionMethod::Unsupported(14);
    pub const IBM_ZOS_CMPSC: Self = CompressionMethod::Unsupported(16);
    pub const IBM_TERSE: Self = CompressionMethod::Unsupported(18);
    pub const ZSTD_DEPRECATED: Self = CompressionMethod::Unsupported(20);
    #[cfg(feature = "zstd")]
    pub const ZSTD: Self = CompressionMethod::Zstd;
    #[cfg(not(feature = "zstd"))]
    pub const ZSTD: Self = CompressionMethod::Unsupported(93);
    pub const MP3: Self = CompressionMethod::Unsupported(94);
    pub const XZ: Self = CompressionMethod::Unsupported(95);
    pub const JPEG: Self = CompressionMethod::Unsupported(96);
    pub const WAVPACK: Self = CompressionMethod::Unsupported(97);
    pub const PPMD: Self = CompressionMethod::Unsupported(98);
    /// Encrypted using AES.
    ///
    /// The actual compression method has to be taken from the AES extra data field
    /// or from `ZipFileData`.
    pub const AES: Self = CompressionMethod::Unsupported(99);
}
impl CompressionMethod {
    /// Converts an u16 to its corresponding CompressionMethod
    #[deprecated(
        since = "0.5.7",
        note = "use a constant to construct a compression method"
    )]
    pub fn from_u16(val: u16) -> CompressionMethod {
        #[allow(deprecated)]
        match val {
            0 => CompressionMethod::Stored,
            #[cfg(any(
                feature = "deflate",
                feature = "deflate-miniz",
                feature = "deflate-zlib"
            ))]
            8 => CompressionMethod::Deflated,
            #[cfg(feature = "bzip2")]
            12 => CompressionMethod::Bzip2,
            #[cfg(feature = "zstd")]
            93 => CompressionMethod::Zstd,
            99 => CompressionMethod::AES,

            v => CompressionMethod::Unsupported(v),
        }
    }

    /// Converts a CompressionMethod to a u16
    #[deprecated(
        since = "0.5.7",
        note = "to match on other compression methods, use a constant"
    )]
    pub fn to_u16(self) -> u16 {
        #[allow(deprecated)]
        match self {
            CompressionMethod::Stored => 0,
            #[cfg(any(
                feature = "deflate",
                feature = "deflate-miniz",
                feature = "deflate-zlib"
            ))]
            CompressionMethod::Deflated => 8,
            #[cfg(feature = "bzip2")]
            CompressionMethod::Bzip2 => 12,
            CompressionMethod::AES => 99,
            #[cfg(feature = "zstd")]
            CompressionMethod::Zstd => 93,

            CompressionMethod::Unsupported(v) => v,
        }
    }
}

impl fmt::Display for CompressionMethod {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        // Just duplicate what the Debug format looks like, i.e, the enum key:
        write!(f, "{:?}", self)
    }
}

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

    #[test]
    fn from_eq_to() {
        for v in 0..(u16::MAX as u32 + 1) {
            #[allow(deprecated)]
            let from = CompressionMethod::from_u16(v as u16);
            #[allow(deprecated)]
            let to = from.to_u16() as u32;
            assert_eq!(v, to);
        }
    }

    fn methods() -> Vec<CompressionMethod> {
        vec![
            CompressionMethod::Stored,
            #[cfg(any(
                feature = "deflate",
                feature = "deflate-miniz",
                feature = "deflate-zlib"
            ))]
            CompressionMethod::Deflated,
            #[cfg(feature = "bzip2")]
            CompressionMethod::Bzip2,
            #[cfg(feature = "zstd")]
            CompressionMethod::Zstd,
        ]
    }

    #[test]
    fn to_eq_from() {
        fn check_match(method: CompressionMethod) {
            #[allow(deprecated)]
            let to = method.to_u16();
            #[allow(deprecated)]
            let from = CompressionMethod::from_u16(to);
            #[allow(deprecated)]
            let back = from.to_u16();
            assert_eq!(to, back);
        }

        for method in methods() {
            check_match(method);
        }
    }

    #[test]
    fn to_display_fmt() {
        fn check_match(method: CompressionMethod) {
            let debug_str = format!("{:?}", method);
            let display_str = format!("{}", method);
            assert_eq!(debug_str, display_str);
        }

        for method in methods() {
            check_match(method);
        }
    }
}
<|endoftext|>


<commit_before>//! Learning benchmark domains module.

use geometry::Space;
use geometry::dimensions;
use geometry::dimensions::Dimension;
use std::collections::HashSet;

/// Container class for data associated with a domain observation.
pub enum Observation<S, A> {
    /// Fully observed state of the environment.
    Full {
        /// Current state of the environment.
        state: S,

        /// Set of available actions.
        actions: HashSet<A>,
    },

    /// Partially observed state of the environment.
    Partial {
        /// Current state of the environment.
        state: S,

        /// Set of available actions.
        actions: HashSet<A>,
    },

    /// Terminal state of the environment.
    Terminal(S),
}

impl<S, A> Observation<S, A> {
    /// Helper function returning a reference to the state values for the given
    /// observation.
    pub fn state(&self) -> &S {
        use self::Observation::*;

        match self {
            &Full { ref state, .. } | &Partial { ref state, .. } | &Terminal(ref state) => state,
        }
    }
}

/// Container class for data associated with a domain transition.
pub struct Transition<S, A> {
    /// State transitioned _from_, `s`.
    pub from: Observation<S, A>,

    /// Action taken to initiate the transition.
    pub action: A,

    /// Reward obtained from the transition.
    pub reward: f64,

    /// State transitioned _to_, `s'`.
    pub to: Observation<S, A>,
}

/// An interface for constructing reinforcement learning problem domains.
pub trait Domain {
    /// State space representation type class.
    type StateSpace: Space;

    /// Action space representation type class.
    type ActionSpace: Space;

    /// Emit an observation of the current state of the environment.
    fn emit(
        &self,
    ) -> Observation<<Self::StateSpace as Space>::Repr, <Self::ActionSpace as Space>::Repr>;

    /// Transition the environment forward a single step given an action, `a`.
    fn step(
        &mut self,
        a: <dimensions::Discrete as Dimension>::Value,
    ) -> Transition<<Self::StateSpace as Space>::Repr, <Self::ActionSpace as Space>::Repr>;

    /// Returns true if the current state is terminal.
    fn is_terminal(&self) -> bool;

    /// Compute the reward associated with a transition from one state to
    /// another.
    fn reward(
        &self,
        from: &Observation<<Self::StateSpace as Space>::Repr, <Self::ActionSpace as Space>::Repr>,
        to: &Observation<<Self::StateSpace as Space>::Repr, <Self::ActionSpace as Space>::Repr>,
    ) -> f64;

    /// Returns an instance of the state space type class.
    fn state_space(&self) -> Self::StateSpace;

    /// Returns an instance of the action space type class.
    fn action_space(&self) -> Self::ActionSpace;
}

mod ode;
use self::ode::*;

mod mountain_car;
pub use self::mountain_car::MountainCar;

mod cart_pole;
pub use self::cart_pole::CartPole;

mod acrobat;
pub use self::acrobat::Acrobat;

mod hiv;
pub use self::hiv::HIVTreatment;

mod cliff_walk;
mod grid_world;
pub use self::cliff_walk::CliffWalk;

mod openai;
pub use self::openai::OpenAIGym;
<commit_msg>feat(domains): implement Into methods for Observation and Transition<commit_after>//! Learning benchmark domains module.

use geometry::Space;
use geometry::dimensions;
use geometry::dimensions::Dimension;
use std::collections::HashSet;

macro_rules! impl_into {
    (Observation<S, $type:ty> => Observation<S, ()>) => (
        impl<S> Into<Observation<S, ()>> for Observation<S, $type> {
            fn into(self) -> Observation<S, ()> {
                use self::Observation::*;

                match self {
                    Full { state, .. } => Full { state: state, actions: HashSet::new() },
                    Partial { state, .. } => Partial { state: state, actions: HashSet::new() },
                    Terminal(state) => Terminal(state),
                }
            }
        }

        impl<'a, S: Clone> From<&'a Observation<S, $type>> for Observation<S, ()> {
            fn from(obs: &'a Observation<S, $type>) -> Observation<S, ()> {
                use self::Observation::*;

                match obs {
                    Full { state, .. } => Full { state: state.clone(), actions: HashSet::new() },
                    Partial { state, .. } => Partial { state: state.clone(), actions: HashSet::new() },
                    Terminal(state) => Terminal(state.clone()),
                }
            }
        }
    );
    (Transition<S, $type:ty> => Transition<S, ()>) => (
        impl<S> Into<Transition<S, ()>> for Transition<S, $type> {
            fn into(self) -> Transition<S, ()> {
                Transition {
                    from: self.from.into(),
                    action: (),
                    reward: self.reward,
                    to: self.to.into(),
                }
            }
        }

        impl<'a, S: Clone> From<&'a Transition<S, $type>> for Transition<S, ()> {
            fn from(t: &'a Transition<S, $type>) -> Transition<S, ()> {
                Transition {
                    from: (&t.from).into(),
                    action: (),
                    reward: t.reward,
                    to: (&t.to).into(),
                }
            }
        }
    )
}

/// Container class for data associated with a domain observation.
#[derive(Clone)]
pub enum Observation<S, A> {
    /// Fully observed state of the environment.
    Full {
        /// Current state of the environment.
        state: S,

        /// Set of available actions.
        actions: HashSet<A>,
    },

    /// Partially observed state of the environment.
    Partial {
        /// Current state of the environment.
        state: S,

        /// Set of available actions.
        actions: HashSet<A>,
    },

    /// Terminal state of the environment.
    Terminal(S),
}

impl<S, A> Observation<S, A> {
    /// Helper function returning a reference to the state values for the given
    /// observation.
    pub fn state(&self) -> &S {
        use self::Observation::*;

        match self {
            &Full { ref state, .. } | &Partial { ref state, .. } | &Terminal(ref state) => state,
        }
    }
}

impl_into!(Observation<S, u8> => Observation<S, ()>);
impl_into!(Observation<S, u16> => Observation<S, ()>);
impl_into!(Observation<S, u32> => Observation<S, ()>);
impl_into!(Observation<S, u64> => Observation<S, ()>);
impl_into!(Observation<S, usize> => Observation<S, ()>);
impl_into!(Observation<S, i8> => Observation<S, ()>);
impl_into!(Observation<S, i16> => Observation<S, ()>);
impl_into!(Observation<S, i32> => Observation<S, ()>);
impl_into!(Observation<S, i64> => Observation<S, ()>);
impl_into!(Observation<S, isize> => Observation<S, ()>);
impl_into!(Observation<S, f32> => Observation<S, ()>);
impl_into!(Observation<S, f64> => Observation<S, ()>);

/// Container class for data associated with a domain transition.
#[derive(Clone)]
pub struct Transition<S, A> {
    /// State transitioned _from_, `s`.
    pub from: Observation<S, A>,

    /// Action taken to initiate the transition (control tasks).
    pub action: A,

    /// Reward obtained from the transition.
    pub reward: f64,

    /// State transitioned _to_, `s'`.
    pub to: Observation<S, A>,
}

impl_into!(Transition<S, u8> => Transition<S, ()>);
impl_into!(Transition<S, u16> => Transition<S, ()>);
impl_into!(Transition<S, u32> => Transition<S, ()>);
impl_into!(Transition<S, u64> => Transition<S, ()>);
impl_into!(Transition<S, usize> => Transition<S, ()>);
impl_into!(Transition<S, i8> => Transition<S, ()>);
impl_into!(Transition<S, i16> => Transition<S, ()>);
impl_into!(Transition<S, i32> => Transition<S, ()>);
impl_into!(Transition<S, i64> => Transition<S, ()>);
impl_into!(Transition<S, isize> => Transition<S, ()>);
impl_into!(Transition<S, f32> => Transition<S, ()>);
impl_into!(Transition<S, f64> => Transition<S, ()>);

/// An interface for constructing reinforcement learning problem domains.
pub trait Domain {
    /// State space representation type class.
    type StateSpace: Space;

    /// Action space representation type class.
    type ActionSpace: Space;

    /// Emit an observation of the current state of the environment.
    fn emit(
        &self,
    ) -> Observation<<Self::StateSpace as Space>::Repr, <Self::ActionSpace as Space>::Repr>;

    /// Transition the environment forward a single step given an action, `a`.
    fn step(
        &mut self,
        a: <dimensions::Discrete as Dimension>::Value,
    ) -> Transition<<Self::StateSpace as Space>::Repr, <Self::ActionSpace as Space>::Repr>;

    /// Returns true if the current state is terminal.
    fn is_terminal(&self) -> bool;

    /// Compute the reward associated with a transition from one state to
    /// another.
    fn reward(
        &self,
        from: &Observation<<Self::StateSpace as Space>::Repr, <Self::ActionSpace as Space>::Repr>,
        to: &Observation<<Self::StateSpace as Space>::Repr, <Self::ActionSpace as Space>::Repr>,
    ) -> f64;

    /// Returns an instance of the state space type class.
    fn state_space(&self) -> Self::StateSpace;

    /// Returns an instance of the action space type class.
    fn action_space(&self) -> Self::ActionSpace;
}

mod ode;
use self::ode::*;

mod mountain_car;
pub use self::mountain_car::MountainCar;

mod cart_pole;
pub use self::cart_pole::CartPole;

mod acrobat;
pub use self::acrobat::Acrobat;

mod hiv;
pub use self::hiv::HIVTreatment;

mod cliff_walk;
mod grid_world;
pub use self::cliff_walk::CliffWalk;

mod openai;
pub use self::openai::OpenAIGym;
<|endoftext|>


<commit_before>use std::time::{SystemTime, Duration};
use std::thread::sleep;

//John Renner
//

/// trait for conversion to milliseconds
trait AsMillis {
    /// convert the self (duration) to milliseconds
    fn as_millis(&self) -> u64;
}

impl AsMillis for Duration {
    fn as_millis(&self) -> u64 {
        self.as_secs() * 1000 + u64::from(self.subsec_nanos() / 1_000_000u32)
    }
}

pub struct FrameTimer {
    n_frame_per_second_length: Duration,
    //I allocate the frame_duration varible here so that I can avoid an allocation each frame
    frame_duration: Duration,
    sleep_duration: Duration,
    /// the instant the frame started
    frame_start: SystemTime,
    /// the instant the frame ended
    frame_end: SystemTime,
    /// the time elapsed per frame in milliseconds
    delta_time: u64
}

impl FrameTimer {
    pub fn new() -> Self {
        let sleep_duration = Duration::from_millis(0);
        //TODO: Make this more then just a 60 FPS timer
        let sixty_frames_per_second_length = Duration::from_millis(16);

        let frame_duration = Duration::from_millis(0);
        let frame_start = SystemTime::now();
        let frame_end = SystemTime::now();

        FrameTimer {
            n_frame_per_second_length: sixty_frames_per_second_length,
            frame_duration: frame_duration,
            sleep_duration: sleep_duration,
            frame_start: frame_start,
            frame_end: frame_end,
            delta_time: 0u64
        }
    }

    pub fn frame_start(&mut self) {
        self.frame_start = SystemTime::now();
    }

    pub fn delta_time(&self) -> f32 {
        (self.delta_time as f32 / 1000f32)
    }

    pub fn frame_end(&mut self) {
        self.frame_end = SystemTime::now();
        self.frame_duration = self.frame_start.duration_since(self.frame_start).unwrap();
        self.delta_time = self.frame_duration.as_millis();
        if self.frame_duration.as_secs() > 0 {
            return;
        }

        if self.sleep_duration.subsec_nanos() < self.n_frame_per_second_length.subsec_nanos() {
            self.sleep_duration = self.n_frame_per_second_length - self.frame_duration;
            sleep(self.sleep_duration);
        }
    }
}
<commit_msg>Fixed pretty serious bug with the sleep check<commit_after>use std::time::{SystemTime, Duration};
use std::thread::sleep;

//John Renner
//

/// trait for conversion to milliseconds
trait AsMillis {
    /// convert the self (duration) to milliseconds
    fn as_millis(&self) -> u64;
}

impl AsMillis for Duration {
    fn as_millis(&self) -> u64 {
        self.as_secs() * 1000 + u64::from(self.subsec_nanos() / 1_000_000u32)
    }
}

pub struct FrameTimer {
    n_frame_per_second_length: Duration,
    //I allocate the frame_duration varible here so that I can avoid an allocation each frame
    frame_duration: Duration,
    sleep_duration: Duration,
    /// the instant the frame started
    frame_start: SystemTime,
    /// the instant the frame ended
    frame_end: SystemTime,
    /// the time elapsed per frame in milliseconds
    delta_time: u64
}

impl FrameTimer {
    pub fn new() -> Self {
        let sleep_duration = Duration::from_millis(0);
        //TODO: Make this more then just a 60 FPS timer
        let sixty_frames_per_second_length = Duration::from_millis(16);

        let frame_duration = Duration::from_millis(0);
        let frame_start = SystemTime::now();
        let frame_end = SystemTime::now();

        FrameTimer {
            n_frame_per_second_length: sixty_frames_per_second_length,
            frame_duration: frame_duration,
            sleep_duration: sleep_duration,
            frame_start: frame_start,
            frame_end: frame_end,
            delta_time: 0u64
        }
    }

    pub fn frame_start(&mut self) {
        self.frame_start = SystemTime::now();
    }

    pub fn delta_time(&self) -> f32 {
        (self.delta_time as f32 / 1000f32)
    }

    pub fn frame_end(&mut self) {
        self.frame_end = SystemTime::now();
        self.frame_duration = self.frame_start.duration_since(self.frame_start).unwrap();
        self.delta_time = self.frame_duration.as_millis();
        if self.frame_duration.as_secs() > 0 {
            return;
        }

        if self.frame_duration.subsec_nanos() < self.n_frame_per_second_length.subsec_nanos() {
            self.sleep_duration = self.n_frame_per_second_length - self.frame_duration;
            sleep(self.sleep_duration);
        }
    }
}
<|endoftext|>


<commit_before>//! Xen backend for Uni.rs
//!
//! Note: This backend is enabled by the feature named *xen* and thus might
//! not be available depending on your build's configuration.

use io::Write;

use alloc_uni;

use hal;

use hal::arch::utils::wmb;
use hal::xen::store::XenStore;

use thread::Scheduler;

// In order to use stdout, heap must be initialized, so "raw" console is used
// before that
macro_rules! raw_println {
    ($fmt:expr) => {
        raw_print!(concat!($fmt, "\r\n"))
    };
    ($fmt:expr, $($arg:tt)*) => {
        raw_print!(concat!($fmt, "\r\n"), $($arg)*)
    }
}

macro_rules! raw_print {
    ($($arg:tt)*) => {{
        use $crate::io::Write;

        $crate::hal::console().write_fmt(format_args!($($arg)*)).unwrap();
    }}
}

mod hypercall;

pub mod arch;
pub mod defs;
pub mod boot;
pub mod store;
pub mod memory;
pub mod event;
pub mod sched;
pub mod console;

extern "C" {
    // This symbol must be present in code using libxen
    pub static mut shared_info: self::defs::SharedInfo;
}

pub fn enable_upcalls() -> u8 {
    unsafe {
        let ret = shared_info.vcpu_info[0].evtchn_upcall_mask;

        wmb();
        shared_info.vcpu_info[0].evtchn_upcall_mask = 0;
        wmb();

        ret
    }
}

pub fn disable_upcalls() -> u8 {
    unsafe {
        let ret = shared_info.vcpu_info[0].evtchn_upcall_mask;

        shared_info.vcpu_info[0].evtchn_upcall_mask = 1;
        wmb();

        ret
    }
}

pub fn set_upcalls_state(state: u8) {
    unsafe {
        wmb();
        shared_info.vcpu_info[0].evtchn_upcall_mask = state;
        wmb();
    }
}

extern {
    fn main(_: isize, _: *const *const u8) -> isize;
}

#[no_mangle]
/// Entry point of the application called by boot assembly
pub extern "C" fn uni_rust_entry() -> ! {
    alloc_uni::init();

    boot::init();

    raw_println!("Uni.rs is booting");

    // Memory initialization is unsafe
    boot::init_memory();

    event::init();

    unsafe {
        console::console().init_input();
        XenStore::init_event();
    }

    hal::local_irq_enable();

    // Spawn main thread
    Scheduler::spawn(|| {
        println!("Main thread started");
        println!("Uni.rs is now ready");
        println!("Control is now transfered to the application");
        println!("");

        let app_ret = unsafe {
            main(0, ::core::ptr::null())
        };

        hal::local_irq_disable();

        hal::console().flush().unwrap();

        hal::app::exit(app_ret);

        panic!("Failed to poweroff the machine !");
    });

    Scheduler::schedule();

    unreachable!();
}
<commit_msg>misc: Remove old comment<commit_after>//! Xen backend for Uni.rs
//!
//! Note: This backend is enabled by the feature named *xen* and thus might
//! not be available depending on your build's configuration.

use io::Write;

use alloc_uni;

use hal;

use hal::arch::utils::wmb;
use hal::xen::store::XenStore;

use thread::Scheduler;

// In order to use stdout, heap must be initialized, so "raw" console is used
// before that
macro_rules! raw_println {
    ($fmt:expr) => {
        raw_print!(concat!($fmt, "\r\n"))
    };
    ($fmt:expr, $($arg:tt)*) => {
        raw_print!(concat!($fmt, "\r\n"), $($arg)*)
    }
}

macro_rules! raw_print {
    ($($arg:tt)*) => {{
        use $crate::io::Write;

        $crate::hal::console().write_fmt(format_args!($($arg)*)).unwrap();
    }}
}

mod hypercall;

pub mod arch;
pub mod defs;
pub mod boot;
pub mod store;
pub mod memory;
pub mod event;
pub mod sched;
pub mod console;

extern "C" {
    // This symbol must be present in code using libxen
    pub static mut shared_info: self::defs::SharedInfo;
}

pub fn enable_upcalls() -> u8 {
    unsafe {
        let ret = shared_info.vcpu_info[0].evtchn_upcall_mask;

        wmb();
        shared_info.vcpu_info[0].evtchn_upcall_mask = 0;
        wmb();

        ret
    }
}

pub fn disable_upcalls() -> u8 {
    unsafe {
        let ret = shared_info.vcpu_info[0].evtchn_upcall_mask;

        shared_info.vcpu_info[0].evtchn_upcall_mask = 1;
        wmb();

        ret
    }
}

pub fn set_upcalls_state(state: u8) {
    unsafe {
        wmb();
        shared_info.vcpu_info[0].evtchn_upcall_mask = state;
        wmb();
    }
}

extern {
    fn main(_: isize, _: *const *const u8) -> isize;
}

#[no_mangle]
/// Entry point of the application called by boot assembly
pub extern "C" fn uni_rust_entry() -> ! {
    alloc_uni::init();

    boot::init();

    raw_println!("Uni.rs is booting");

    boot::init_memory();

    event::init();

    unsafe {
        console::console().init_input();
        XenStore::init_event();
    }

    hal::local_irq_enable();

    // Spawn main thread
    Scheduler::spawn(|| {
        println!("Main thread started");
        println!("Uni.rs is now ready");
        println!("Control is now transfered to the application");
        println!("");

        let app_ret = unsafe {
            main(0, ::core::ptr::null())
        };

        hal::local_irq_disable();

        hal::console().flush().unwrap();

        hal::app::exit(app_ret);

        panic!("Failed to poweroff the machine !");
    });

    Scheduler::schedule();

    unreachable!();
}
<|endoftext|>


<commit_before>//! Implementation of Xen's network driver
<commit_msg>xen/net: Add definition of all network related types<commit_after>//! Implementation of Xen's network driver

use sync::Arc;

use vec::Vec;

use sync::spin::{InterruptSpinLock, SpinLock, RwLock};

use net::{Interface, Packet};

use hal::xen::grant::Ref as GrantRef;
use hal::xen::ring::{SharedRing, FrontRing};

use hal::xen::defs::EvtchnPort;

#[repr(i16)]
#[derive(Debug, PartialEq)]
#[allow(dead_code)]
/// XEN_NETIF_RSP_*
enum NetifRsp {
    /// XEN_NETIF_RSP_OKAY
    Okay = 0,
    /// XEN_NETIF_RSP_ERROR
    Error = -1,
    /// XEN_NETIF_RSP_DROPPED
    Dropped = -2,
    /// XEN_NETIF_RSP_NULL
    Null = 1,
}

#[repr(u16)]
#[derive(Debug)]
#[allow(dead_code)]
/// XEN_NETTXF_*
enum NetTxFlags {
    CsumBlank = 1 << 0,
    DataValidated = 1 << 1,
    MoreData = 1 << 2,
    ExtraInfo = 1 << 3,
}

#[repr(u16)]
#[derive(Debug)]
#[allow(dead_code)]
/// XEN_NETRXF_*
enum NetRxFlags {
    DataValidated = 1 << 0,
    CsumBlank = 1 << 1,
    MoreData = 1 << 2,
    ExtraInfo = 1 << 3,
    GsoPrefix = 1 << 4,
}

#[repr(C)]
#[derive(Debug)]
/// struct xen_netif_tx_request
struct NetifTxRequest {
    gref: GrantRef,
    offset: u16,
    flags: NetTxFlags,
    id: u16,
    size: u16,
}

#[repr(C)]
#[derive(Debug)]
/// struct xen_netif_tx_response
struct NetifTxResponse {
    id: u16,
    status: NetifRsp,
}

#[repr(C)]
#[derive(Debug)]
/// struct xen_netif_rx_request
struct NetifRxRequest {
    id: u16,
    gref: GrantRef,
}

#[repr(C)]
#[derive(Debug)]
/// struct xen_netif_rx_response
struct NetifRxResponse {
    id: u16,
    offset: u16,
    flags: NetRxFlags,
    status: i16,
}

type TxSharedRing = SharedRing<NetifTxRequest, NetifTxResponse>;
type RxSharedRing = SharedRing<NetifRxRequest, NetifRxResponse>;

type TxFrontRing = FrontRing<NetifTxRequest, NetifTxResponse>;
type RxFrontRing = FrontRing<NetifRxRequest, NetifRxResponse>;

/// Buffer allocated for packet reception
struct RxBuffer {
    pub id: u16,
    pub page: *mut u8,
    pub grant_ref: GrantRef,
}

/// Buffer allocated for packet transmission
struct TxBuffer {
    pub id: u16,
    pub pkt: Option<Packet>,
    pub grant_ref: GrantRef,
}

/// A Xen vif (virtual interface)
pub struct XenNetDevice {
    evtchn: EvtchnPort,
    backend_id: u16,
    tx_ring: TxFrontRing,
    rx_ring: RxFrontRing,
    tx_buffer: SpinLock<Vec<TxBuffer>>,
    rx_buffer: InterruptSpinLock<Vec<RxBuffer>>,
    intf: *const Arc<RwLock<Interface>>,
}
<|endoftext|>


<commit_before>/*!
helper functions used by this crate
whose utility surpasses the confines of this crate.

generic things that are also useful in other crates.
*/

mod repeatedly;
pub use self::repeatedly::{Repeatedly, repeatedly};

// this needs to be public otherwise the impl is not exported
pub mod normalize;
pub use self::normalize::Normalize;

mod partial_max;
pub use self::partial_max::PartialMaxIteratorExt;

mod random;
pub use self::random::random01;
<commit_msg>add ToImage to helpers<commit_after>/*!
helper functions used by this crate
whose utility surpasses the confines of this crate.

generic things that are also useful in other crates.
*/

mod repeatedly;
pub use self::repeatedly::{Repeatedly, repeatedly};

// this needs to be public otherwise the impl is not exported
pub mod normalize;
pub use self::normalize::Normalize;

mod partial_max;
pub use self::partial_max::PartialMaxIteratorExt;

mod random;
pub use self::random::random01;

pub mod to_image;
pub use self::to_image::ToImage;
<|endoftext|>


<commit_before><commit_msg>Clean up temporary files<commit_after><|endoftext|>


<commit_before>use std::str::FromStr;

#[derive(Debug, PartialEq)]
pub enum Instruction {
    NOP,
    MOV { src: Operand, dst: Operand },
    SWP,
    SAV,
    ADD { addend: Operand },
    SUB { subtrahend: Operand} ,
    NEG,
    J { cond: Condition, dst: Label},
    JRO { dst: Operand },
}

pub static BAD_OPCODE_ERR: &'static str = "Bad opcode for # of arguments";
pub static NUM_ARGS_ERR: &'static str = "Wrong number of arguments";
pub static LIT_DST_ERR: &'static str = "Literal not allowed as dst operand";

impl FromStr for Instruction {
    type Err = &'static str;

    fn from_str(insn: &str) -> Result<Instruction, Self::Err> {
        //Nightly: let words: Vec<&str> = insn.split_whitespace().collect();
        let words: Vec<&str> = insn.split(' ').filter(|s| *s != "").collect();

        match words.len() {
            1 => match words[0] {
                "NOP" => Ok(Instruction::NOP),
                "SWP" => Ok(Instruction::SWP),
                "SAV" => Ok(Instruction::SAV),
                "NEG" => Ok(Instruction::NEG),
                _ => Err(BAD_OPCODE_ERR),
            },

            2 => match words[0] {
                "ADD" => Operand::from_str(words[1]).map(|o| Instruction::ADD { addend: o }),
                "SUB" => Operand::from_str(words[1]).map(|o| Instruction::SUB { subtrahend: o }),
                "JMP" => Ok(Instruction::J { cond: Condition::Unconditional, dst: words[1].to_string() }),
                "JEZ" => Ok(Instruction::J { cond: Condition::Ez,            dst: words[1].to_string() }),
                "JNZ" => Ok(Instruction::J { cond: Condition::Nz,            dst: words[1].to_string() }),
                "JGZ" => Ok(Instruction::J { cond: Condition::Gz,            dst: words[1].to_string() }),
                "JLZ" => Ok(Instruction::J { cond: Condition::Lz,            dst: words[1].to_string() }),
                "JRO" => Operand::from_str(words[1]).map(|o| Instruction::JRO { dst: o }),
                _ => Err(BAD_OPCODE_ERR),
            },

            3 => match words[0] {
                "MOV" => Operand::from_str(words[1]).and_then(|s| Operand::from_str(words[2])
                    .and_then(|d| match d {
                        Operand::Lit(_) => Err(LIT_DST_ERR),
                        _ => Ok(Instruction::MOV { src: s, dst: d })
                        })),
                _ => Err(BAD_OPCODE_ERR),
            },

            _ => Err(NUM_ARGS_ERR)
        }
    }
}

#[test]
fn instruction_from_str() {
    fn i(s: &str) -> Instruction {
        Instruction::from_str(s).unwrap()
    }
    assert_eq!(i("NOP"), Instruction::NOP);
    assert_eq!(i("SWP"), Instruction::SWP);
    assert_eq!(i("SAV"), Instruction::SAV);
    assert_eq!(i("NEG"), Instruction::NEG);
    assert_eq!(Instruction::from_str("BOGUS").unwrap_err(), BAD_OPCODE_ERR);

    assert_eq!(i("ADD 10"), Instruction::ADD { addend: Operand::Lit(10) });
    assert_eq!(i("SUB 10"), Instruction::SUB { subtrahend: Operand::Lit(10) });
    assert_eq!(i("JMP LOC"), Instruction::J { cond: Condition::Unconditional,   dst: "LOC".to_string() });
    assert_eq!(i("JEZ LOC"), Instruction::J { cond: Condition::Ez,              dst: "LOC".to_string() });
    assert_eq!(i("JNZ LOC"), Instruction::J { cond: Condition::Nz,              dst: "LOC".to_string() });
    assert_eq!(i("JGZ LOC"), Instruction::J { cond: Condition::Gz,              dst: "LOC".to_string() });
    assert_eq!(i("JLZ LOC"), Instruction::J { cond: Condition::Lz,              dst: "LOC".to_string() });
    assert_eq!(i("JRO 1"), Instruction::JRO { dst: Operand::Lit(1) });
    assert_eq!(Instruction::from_str("JEQ LOC").unwrap_err(), BAD_OPCODE_ERR);

    assert_eq!(i("MOV UP DOWN"),    Instruction::MOV { src: Operand::Port(Port::Up), dst: Operand::Port(Port::Down) });
    assert_eq!(i("MOV  UP  DOWN"),  Instruction::MOV { src: Operand::Port(Port::Up), dst: Operand::Port(Port::Down) });
    assert_eq!(i("MOV UP ACC"),     Instruction::MOV { src: Operand::Port(Port::Up), dst: Operand::ACC });
    assert_eq!(i("MOV ACC ACC"),    Instruction::MOV { src: Operand::ACC, dst: Operand::ACC });
    assert_eq!(Instruction::from_str("MV UP ACC").unwrap_err(), BAD_OPCODE_ERR);
    assert_eq!(Instruction::from_str("MOV UP 10").unwrap_err(), LIT_DST_ERR);

    assert_eq!(Instruction::from_str("1 2 3 4").unwrap_err(), NUM_ARGS_ERR);
}

#[derive(Clone, Debug, PartialEq)]
pub enum Port {
    Up,
    Down,
    Left,
    Right,
}

impl FromStr for Port {
    type Err = &'static str;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        match s {
            "UP" => Ok(Port::Up),
            "DOWN" => Ok(Port::Down),
            "LEFT" => Ok(Port::Left),
            "RIGHT" => Ok(Port::Right),
            _ => Err("bad port"),
        }
    }
}

#[test]
fn port_from_str() {
    assert_eq!(Port::from_str("UP").unwrap(),   Port::Up);
    assert_eq!(Port::from_str("DOWN").unwrap(), Port::Down);
    assert_eq!(Port::from_str("LEFT").unwrap(), Port::Left);
    assert_eq!(Port::from_str("RIGHT").unwrap(), Port::Right);
    assert_eq!(Port::from_str("OTHER").unwrap_err(), "bad port");
}

#[derive(Debug, PartialEq)]
pub enum Operand {
    Lit(i32),
    Port(Port),
    ACC,
}

impl FromStr for Operand {
    type Err = &'static str;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        if s == "ACC" {
            return Ok(Operand::ACC);
        }

        let as_int = i32::from_str(s);
        if as_int.is_ok() {
            return Ok(Operand::Lit(as_int.unwrap()));
        }

        let as_port = Port::from_str(s);
        if as_port.is_ok() {
            return Ok(Operand::Port(as_port.unwrap()));
        }

        Err("Invalid operand")
    }
}

#[test]
fn operand_from_str() {
    assert_eq!(Operand::from_str("ACC").unwrap(), Operand::ACC);
    assert_eq!(Operand::from_str("32").unwrap(), Operand::Lit(32));
    assert_eq!(Operand::from_str("UP").unwrap(), Operand::Port(Port::Up));
    assert_eq!(Operand::from_str("FOO").unwrap_err(), "Invalid operand");
}

#[derive(Debug, PartialEq)]
pub enum Condition {
    Unconditional,
    Ez,
    Nz,
    Gz,
    Lz,
}

pub type Label = String;
<commit_msg>instruction: use if/let for parsing operand<commit_after>use std::str::FromStr;

#[derive(Debug, PartialEq)]
pub enum Instruction {
    NOP,
    MOV { src: Operand, dst: Operand },
    SWP,
    SAV,
    ADD { addend: Operand },
    SUB { subtrahend: Operand} ,
    NEG,
    J { cond: Condition, dst: Label},
    JRO { dst: Operand },
}

pub static BAD_OPCODE_ERR: &'static str = "Bad opcode for # of arguments";
pub static NUM_ARGS_ERR: &'static str = "Wrong number of arguments";
pub static LIT_DST_ERR: &'static str = "Literal not allowed as dst operand";

impl FromStr for Instruction {
    type Err = &'static str;

    fn from_str(insn: &str) -> Result<Instruction, Self::Err> {
        //Nightly: let words: Vec<&str> = insn.split_whitespace().collect();
        let words: Vec<&str> = insn.split(' ').filter(|s| *s != "").collect();

        match words.len() {
            1 => match words[0] {
                "NOP" => Ok(Instruction::NOP),
                "SWP" => Ok(Instruction::SWP),
                "SAV" => Ok(Instruction::SAV),
                "NEG" => Ok(Instruction::NEG),
                _ => Err(BAD_OPCODE_ERR),
            },

            2 => match words[0] {
                "ADD" => Operand::from_str(words[1]).map(|o| Instruction::ADD { addend: o }),
                "SUB" => Operand::from_str(words[1]).map(|o| Instruction::SUB { subtrahend: o }),
                "JMP" => Ok(Instruction::J { cond: Condition::Unconditional, dst: words[1].to_string() }),
                "JEZ" => Ok(Instruction::J { cond: Condition::Ez,            dst: words[1].to_string() }),
                "JNZ" => Ok(Instruction::J { cond: Condition::Nz,            dst: words[1].to_string() }),
                "JGZ" => Ok(Instruction::J { cond: Condition::Gz,            dst: words[1].to_string() }),
                "JLZ" => Ok(Instruction::J { cond: Condition::Lz,            dst: words[1].to_string() }),
                "JRO" => Operand::from_str(words[1]).map(|o| Instruction::JRO { dst: o }),
                _ => Err(BAD_OPCODE_ERR),
            },

            3 => match words[0] {
                "MOV" => Operand::from_str(words[1]).and_then(|s| Operand::from_str(words[2])
                    .and_then(|d| match d {
                        Operand::Lit(_) => Err(LIT_DST_ERR),
                        _ => Ok(Instruction::MOV { src: s, dst: d })
                        })),
                _ => Err(BAD_OPCODE_ERR),
            },

            _ => Err(NUM_ARGS_ERR)
        }
    }
}

#[test]
fn instruction_from_str() {
    fn i(s: &str) -> Instruction {
        Instruction::from_str(s).unwrap()
    }
    assert_eq!(i("NOP"), Instruction::NOP);
    assert_eq!(i("SWP"), Instruction::SWP);
    assert_eq!(i("SAV"), Instruction::SAV);
    assert_eq!(i("NEG"), Instruction::NEG);
    assert_eq!(Instruction::from_str("BOGUS").unwrap_err(), BAD_OPCODE_ERR);

    assert_eq!(i("ADD 10"), Instruction::ADD { addend: Operand::Lit(10) });
    assert_eq!(i("SUB 10"), Instruction::SUB { subtrahend: Operand::Lit(10) });
    assert_eq!(i("JMP LOC"), Instruction::J { cond: Condition::Unconditional,   dst: "LOC".to_string() });
    assert_eq!(i("JEZ LOC"), Instruction::J { cond: Condition::Ez,              dst: "LOC".to_string() });
    assert_eq!(i("JNZ LOC"), Instruction::J { cond: Condition::Nz,              dst: "LOC".to_string() });
    assert_eq!(i("JGZ LOC"), Instruction::J { cond: Condition::Gz,              dst: "LOC".to_string() });
    assert_eq!(i("JLZ LOC"), Instruction::J { cond: Condition::Lz,              dst: "LOC".to_string() });
    assert_eq!(i("JRO 1"), Instruction::JRO { dst: Operand::Lit(1) });
    assert_eq!(Instruction::from_str("JEQ LOC").unwrap_err(), BAD_OPCODE_ERR);

    assert_eq!(i("MOV UP DOWN"),    Instruction::MOV { src: Operand::Port(Port::Up), dst: Operand::Port(Port::Down) });
    assert_eq!(i("MOV  UP  DOWN"),  Instruction::MOV { src: Operand::Port(Port::Up), dst: Operand::Port(Port::Down) });
    assert_eq!(i("MOV UP ACC"),     Instruction::MOV { src: Operand::Port(Port::Up), dst: Operand::ACC });
    assert_eq!(i("MOV ACC ACC"),    Instruction::MOV { src: Operand::ACC, dst: Operand::ACC });
    assert_eq!(Instruction::from_str("MV UP ACC").unwrap_err(), BAD_OPCODE_ERR);
    assert_eq!(Instruction::from_str("MOV UP 10").unwrap_err(), LIT_DST_ERR);

    assert_eq!(Instruction::from_str("1 2 3 4").unwrap_err(), NUM_ARGS_ERR);
}

#[derive(Clone, Debug, PartialEq)]
pub enum Port {
    Up,
    Down,
    Left,
    Right,
}

impl FromStr for Port {
    type Err = &'static str;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        match s {
            "UP" => Ok(Port::Up),
            "DOWN" => Ok(Port::Down),
            "LEFT" => Ok(Port::Left),
            "RIGHT" => Ok(Port::Right),
            _ => Err("bad port"),
        }
    }
}

#[test]
fn port_from_str() {
    assert_eq!(Port::from_str("UP").unwrap(),   Port::Up);
    assert_eq!(Port::from_str("DOWN").unwrap(), Port::Down);
    assert_eq!(Port::from_str("LEFT").unwrap(), Port::Left);
    assert_eq!(Port::from_str("RIGHT").unwrap(), Port::Right);
    assert_eq!(Port::from_str("OTHER").unwrap_err(), "bad port");
}

#[derive(Debug, PartialEq)]
pub enum Operand {
    Lit(i32),
    Port(Port),
    ACC,
}

impl FromStr for Operand {
    type Err = &'static str;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        if s == "ACC" {
            Ok(Operand::ACC)
        } else if let Ok(as_int) = i32::from_str(s) {
            Ok(Operand::Lit(as_int))
        } else if let Ok(as_port) = Port::from_str(s) {
            Ok(Operand::Port(as_port))
        } else {
            Err("Invalid operand")
        }
    }
}

#[test]
fn operand_from_str() {
    assert_eq!(Operand::from_str("ACC").unwrap(), Operand::ACC);
    assert_eq!(Operand::from_str("32").unwrap(), Operand::Lit(32));
    assert_eq!(Operand::from_str("UP").unwrap(), Operand::Port(Port::Up));
    assert_eq!(Operand::from_str("FOO").unwrap_err(), "Invalid operand");
}

#[derive(Debug, PartialEq)]
pub enum Condition {
    Unconditional,
    Ez,
    Nz,
    Gz,
    Lz,
}

pub type Label = String;
<|endoftext|>